import numpy as np
np.set_printoptions(threshold=np.inf)
import logging

def load_ckpt(checkpoint_path):
	import tensorflow as tf
	if tf.__version__.startswith('2'):
		import tensorflow.compat.v1 as tf
		tf.disable_v2_behavior()
		reader = tf.compat.v1.train.NewCheckpointReader(checkpoint_path)
	else:
		from tensorflow.python import pywrap_tensorflow
		reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
	var_to_shape_map = reader.get_variable_to_shape_map()

	var_dict = dict()
	for var_name in sorted(var_to_shape_map):
		if "Adam" in var_name:
			continue
		tensor = reader.get_tensor(var_name)
		# print("in ckpt: {}, {}".format(var_name, tensor.shape))
		# print(tensor)
		var_dict[var_name] = tensor

	return var_dict



def load_tf_pb_dict(pb_model):
	import tensorflow as tf
	if tf.__version__.startswith('2'):
		import tensorflow.compat.v1 as tf
		tf.disable_v2_behavior()
		# import tensorflow_addons as tfa
		# from tensorflow_addons.seq2seq.python.ops import beam_search_ops
	else:
		from tensorflow.contrib.seq2seq.python.ops import beam_search_ops
	from tensorflow.python.ops import lookup_ops as lookup
	from tensorflow.python.framework import tensor_util
	from tensorflow.python.platform import gfile
	
	sess = tf.Session()
	with gfile.FastGFile(pb_model, 'rb') as f:
		graph_def = tf.GraphDef()
		graph_def.ParseFromString(f.read())
		sess.graph.as_default()
		tf.import_graph_def(graph_def, name='')
	
	var_dict = dict()
	for node in sess.graph_def.node:
		if node.op == 'Const':
			value = tensor_util.MakeNdarray(node.attr['value'].tensor)
			if len(value.shape) >= 1:
				var_dict[node.name] = value
	return var_dict

def load_tf_dict(pb_model):
	if "model.ckpt-" in pb_model:
		var_dict = load_ckpt(pb_model)
	else:
		var_dict = load_tf_pb_dict(pb_model)
	return var_dict