# -*- coding:utf-8 -*-
"""
you do not need to update a existing instance of class ConvNetValue. because when you generate a instance of class
ConvNetValue, this instance must be a instance who has the latest params of matconvnet.
so just generate a new instance from class ConvNetValue when you need it

"""
import numpy as np
import tensorflow as tf
import scipy.io
import os

class ConvNetValue():
    def __init__(self, matconvnet = None, Flag=False):

        if matconvnet == None:

            self.name = None
            # TODO: <> params are stored in these dictionary
            self.W_d = {}
            self.b_d = {}
            self.bn_gamma_d = {}
            self.bn_moments_d = {}
            self.bn_beta_d = {}
            self.bn_moving_mean_d = {}
            self.bn_moving_variance_d = {}

            self.bnorm_adjust = False

            self.bn_gamma_adj = None
            self.bn_moments_adj = None
            self.bn_beta_adj = None
            self.bn_moving_mean_adj = None
            self.bn_moving_variance_adj = None

            self.num_layers = 5
            self.bnorm_yn_n = np.array([1, 1, 1, 1, 0], dtype=bool)
            self.conv_stride_n = np.array([2, 1, 1, 1, 1])
            self.filtergroup_yn_n = np.array([0, 1, 0, 1, 1], dtype=bool)
            self.bnorm_yn_n = np.array([1, 1, 1, 1, 0], dtype=bool)
            self.relu_yn_n = np.array([1, 1, 1, 1, 0], dtype=bool)
            self.trainable_n = np.array([0, 0, 0, 0, 1])
            self.pool_stride_n = np.array([2, 1, 0, 0, 0]) # 0 means no pool
            self.pool_sz = 3
            self.Flag = Flag
            #self.update_value()
        else :
            self.name = matconvnet.name
            # TODO: <> params are stored in these dictionary
            self.W_d = matconvnet.W_d
            self.b_d = matconvnet.b_d
            self.bn_gamma_d = matconvnet.bn_gamma_d
            self.bn_moments_d = matconvnet.bn_moments_d
            self.bn_beta_d = matconvnet.bn_beta_d
            self.bn_moving_mean_d = matconvnet.bn_moving_mean_d
            self.bn_moving_variance_d = matconvnet.bn_moving_variance_d

            self.bnorm_adjust = matconvnet.bnorm_adjust

            self.bn_gamma_adj = matconvnet.bn_gamma_adj
            self.bn_moments_adj = matconvnet.bn_moments_adj
            self.bn_beta_adj = matconvnet.bn_beta_adj
            self.bn_moving_mean_adj = matconvnet.bn_moving_mean_adj
            self.bn_moving_variance_adj = matconvnet.bn_moving_variance_adj

            self.num_layers = matconvnet.num_layers
            self.bnorm_yn_n = matconvnet.bnorm_yn_n
            self.conv_stride_n = matconvnet.conv_stride_n
            self.filtergroup_yn_n = matconvnet.filtergroup_yn_n
            self.bnorm_yn_n = matconvnet.bnorm_yn_n
            self.relu_yn_n = matconvnet.relu_yn_n
            self.trainable_n = matconvnet.trainable_n
            self.pool_stride_n = matconvnet.pool_stride_n
            self.pool_sz = matconvnet.pool_sz
            self.Flag = matconvnet.Flag

    def update_value(self, instance=None):
        # self.Flag = Flag
        if instance==None:
            self.check_file()
            # TODO: <v3> update the members of instance

            # TODO: <v3> Flag = True: import from matconvnet_ckpt; Flag = False: import from matconvnet.mat
            # TODO: <v3> you can check the number of files in the dictionary to judge while there is a ckpt file instead of
            # TODO: <v3> setting Flag
            if self.Flag:
                #self.net_path = "Model/siamese_network.ckpt"
                self.net_path = "Model/siamese_network.ckpt"
                #self.net_meta_path = "../Model/siamese_network.ckpt.meta"
                self.net_meta_path = "Model/siamese_network.ckpt.meta"
                params_names_list, params_values_list = self._import_from_matconvnet_ckpt()
            else:
                #self.net_path = '../pretrained/baseline-conv5_e55.mat'
                self.net_path = "pretrained/baseline-conv5_e55.mat"
                params_names_list, params_values_list = self._import_from_matconvnet()

            # loop through the flag arrays and re-construct network, reading parameters of conv and bnorm layers
            # 通过获取的MatConvNet的参数名称、值列表，逐层重新构建网络

            if not self.Flag:
                for i in range(self.num_layers):
                    conv_W_name = self._find_params('conv' + str(i + 1) + 'f', params_names_list)[0]
                    conv_b_name = self._find_params('conv' + str(i + 1) + 'b', params_names_list)[0]

                    # 获取指定参数名称在参数名称列表中的index，这个index和指定参数的参数值在参数值列表的index是一一对应的
                    conv_W = params_values_list[params_names_list.index(conv_W_name)]
                    self.W_d[i] = conv_W
                    conv_b = params_values_list[params_names_list.index(conv_b_name)]
                    self.b_d[i] = conv_b

                    if self.bnorm_yn_n[i]:
                        bn_beta_name = self._find_params('bn' + str(i + 1) + 'b', params_names_list)[0]
                        bn_gamma_name = self._find_params('bn' + str(i + 1) + 'm', params_names_list)[0]
                        bn_moments_name = self._find_params('bn' + str(i + 1) + 'x', params_names_list)[0]
                        bn_beta = params_values_list[params_names_list.index(bn_beta_name)]
                        self.bn_beta_d[i] = bn_beta
                        bn_gamma = params_values_list[params_names_list.index(bn_gamma_name)]
                        self.bn_gamma_d[i] = bn_gamma
                        bn_moments = params_values_list[params_names_list.index(bn_moments_name)]
                        self.bn_moments_d[i] = bn_moments
                        # 滑动平均
                        self.bn_moving_mean_d[i] = self.bn_moments_d[i][:, 0]
                        # 滑动方差
                        self.bn_moving_variance_d[i] = self.bn_moments_d[i][:, 1] ** 2  # saved as std in matconvnet
                    else:
                        self.bn_beta_d[i] = self.bn_gamma_d[i] = self.bn_moving_mean_d[i] = \
                            self.bn_moving_variance_d[i] = []

                    if self.bnorm_adjust:
                        self.bn_beta_adj = params_values_list[params_names_list.index('fin_adjust_bnb')]
                        self.bn_gamma_adj = params_values_list[params_names_list.index('fin_adjust_bnm')]
                        self.bn_moments_adj = params_values_list[params_names_list.index('fin_adjust_bnx')]
                        self.bn_moving_mean_adj = self.bn_moments_adj[:, 0]
                        self.bn_moving_variance_adj = self.bn_moments_adj[:, 1] ** 2
            else:

                for i in range(self.num_layers):
                    # conv
                    conv_W_name = self._find_params('conv' + str(i + 1) + '/W:0', params_names_list)[0]
                    conv_b_name = self._find_params('conv' + str(i + 1) + '/b:0', params_names_list)[0]

                    # 获取指定参数名称在参数名称列表中的index，这个index和指定参数的参数值在参数值列表的index是一一对应的
                    conv_W = params_values_list[params_names_list.index(conv_W_name)]
                    self.W_d[i] = conv_W
                    conv_b = params_values_list[params_names_list.index(conv_b_name)]
                    self.b_d[i] = np.transpose(conv_b)

                    # TODO: <> actually bn_params do not need to update
                    # TODO: <> you should know that function import_from_ckpt likes a updating mechanism
                    if self.bnorm_yn_n[i]:
                        bn_beta_name = self._find_params('conv' + str(i + 1) + '/batch_normalization/beta:0', params_names_list)[0]
                        bn_gamma_name = self._find_params('conv' + str(i + 1) + '/batch_normalization/gamma:0', params_names_list)[0]
                        #bn_moments_name = self._find_params('conv' + str(i + 1) + '/batch_normalization/gamma:0', params_names_list)[0]
                        bn_beta = params_values_list[params_names_list.index(bn_beta_name)]
                        self.bn_beta_d[i] = bn_beta
                        bn_gamma = params_values_list[params_names_list.index(bn_gamma_name)]
                        self.bn_gamma_d[i] = bn_gamma
                        #bn_moments = params_values_list[params_names_list.index(bn_moments_name)]
                        #self.bn_moments_d[i] = bn_moments
                        # 滑动平均
                        bn_moving_mean_name =  self._find_params('conv' + str(i + 1) + '/batch_normalization/moving_mean:0', params_names_list)[0]
                        self.bn_moving_mean_d[i] = params_values_list[params_names_list.index(bn_moving_mean_name)]
                        # 滑动方差
                        bn_moving_variance_name = self._find_params('conv' + str(i + 1) + '/batch_normalization/moving_variance:0', params_names_list)[0]
                        self.bn_moving_variance_d[i] = params_values_list[params_names_list.index(bn_moving_variance_name)]
                    else:
                        self.bn_beta_d[i] = self.bn_gamma_d[i] = self.bn_moving_mean_d[i] = \
                            self.bn_moving_variance_d[i] = []

                    if self.bnorm_adjust:
                    #if 0:

                        self.bn_beta_adj = params_values_list[params_names_list.index('fin_adjust_bnb')]
                        self.bn_gamma_adj = params_values_list[params_names_list.index('fin_adjust_bnm')]
                        self.bn_moments_adj = params_values_list[params_names_list.index('fin_adjust_bnx')]
                        self.bn_moving_mean_adj = self.bn_moments_adj[:, 0]
                        self.bn_moving_variance_adj = self.bn_moments_adj[:, 1] ** 2

            return self
        else:
            # TODO: <> it seems that it makes no sense to take this strategy
            # TODO: <> when method upadte() has transferred parameter: a instance of class
            instance.net_path = "Model/siamese_network.ckpt"
            # self.net_meta_path = "../Model/siamese_network.ckpt.meta"
            instance.net_meta_path = "Model/siamese_network.ckpt.meta"
            params_names_list, params_values_list = instance._import_from_matconvnet_ckpt()

            for i in range(instance.num_layers):
                # conv
                conv_W_name = instance._find_params('conv' + str(i + 1) + '/W:0', params_names_list)[0]
                conv_b_name = instance._find_params('conv' + str(i + 1) + '/b:0', params_names_list)[0]

                # 获取指定参数名称在参数名称列表中的index，这个index和指定参数的参数值在参数值列表的index是一一对应的
                conv_W = params_values_list[params_names_list.index(conv_W_name)]
                instance.W_d[i] = conv_W
                conv_b = params_values_list[params_names_list.index(conv_b_name)]
                # TODO: <> transpose!!!!!!
                instance.b_d[i] = np.transpose(conv_b)

                # TODO: <> actually bn_params do not need to update
                # TODO: <> you should know that function import_from_ckpt likes a updating mechanism
                if instance.bnorm_yn_n[i]:
                    bn_beta_name = \
                    instance._find_params('conv' + str(i + 1) + '/batch_normalization/beta:0', params_names_list)[0]
                    bn_gamma_name = \
                    instance._find_params('conv' + str(i + 1) + '/batch_normalization/gamma:0', params_names_list)[0]
                    # bn_moments_name = self._find_params('conv' + str(i + 1) + '/batch_normalization/gamma:0', params_names_list)[0]
                    bn_beta = params_values_list[params_names_list.index(bn_beta_name)]
                    instance.bn_beta_d[i] = bn_beta
                    bn_gamma = params_values_list[params_names_list.index(bn_gamma_name)]
                    instance.bn_gamma_d[i] = bn_gamma
                    # bn_moments = params_values_list[params_names_list.index(bn_moments_name)]
                    # self.bn_moments_d[i] = bn_moments
                    # 滑动平均
                    bn_moving_mean_name = \
                    instance._find_params('conv' + str(i + 1) + '/batch_normalization/moving_mean:0', params_names_list)[0]
                    instance.bn_moving_mean_d[i] = params_values_list[params_names_list.index(bn_moving_mean_name)]
                    # 滑动方差
                    bn_moving_variance_name = \
                    instance._find_params('conv' + str(i + 1) + '/batch_normalization/moving_variance:0',
                                      params_names_list)[0]
                    instance.bn_moving_variance_d[i] = params_values_list[params_names_list.index(bn_moving_variance_name)]
                else:
                    instance.bn_beta_d[i] = instance.bn_gamma_d[i] = instance.bn_moving_mean_d[i] = \
                        instance.bn_moving_variance_d[i] = []

                if instance.bnorm_adjust:
                    # if 0:

                    instance.bn_beta_adj = params_values_list[params_names_list.index('fin_adjust_bnb')]
                    instance.bn_gamma_adj = params_values_list[params_names_list.index('fin_adjust_bnm')]
                    instance.bn_moments_adj = params_values_list[params_names_list.index('fin_adjust_bnx')]
                    instance.bn_moving_mean_adj = instance.bn_moments_adj[:, 0]
                    instance.bn_moving_variance_adj = instance.bn_moments_adj[:, 1] ** 2
            return instance

    # 导入MatConvNet模型
    def _import_from_matconvnet(self):
        # the method to restore the model here is just simple get the value from .mat file and regard the value as
        # the initial value to initialize the variables in the model to achieve the destination that restore the model
        # 使用scipy.io.loadmat()函数对mat数据进行读取
        mat = scipy.io.loadmat(self.net_path)
        net_dot_mat = mat.get('net')
        # organize parameters to import
        params = net_dot_mat['params']
        params = params[0][0]
        params_names = params['name'][0]
        # 网络参数名称列表
        params_names_list = [params_names[p][0] for p in range(params_names.size)]
        params_values = params['value'][0]
        # 网络参数值列表
        params_values_list = [params_values[p] for p in range(params_values.size)]
        # 返回网络参数名称列表、网络参数值列表
        return params_names_list, params_values_list


    # find all parameters matching the codename (there should be only one)
    # 用于从获取得到的网络参数名称、值列表中查找指定参数并获取其值
    def _find_params(self, x, params):
        matching = [s for s in params if x in s]
        # assert 机制要保证参数名称的唯一性 否则触发异常
        assert len(matching) == 1, ('Ambiguous param name found')
        return matching

    def _import_from_matconvnet_ckpt(self):
        graph_cnvc = tf.Graph()
        with graph_cnvc.as_default():
            sess1 = tf.Session()
            saver_r = tf.train.import_meta_graph(self.net_meta_path)
            saver_r.restore(sess1, self.net_path)
            params_names_list = []
            params_values_list = []
            for i in tf.global_variables():
                params_names_list.append(i.name)
                j = sess1.run(graph_cnvc.get_tensor_by_name(i.name))
                params_values_list.append(j)

            sess1.close()
        self.params_name_list = params_names_list
        self.params_value_list = params_values_list
        del graph_cnvc
        return params_names_list, params_values_list


    def check_file(self):
        """
        in this method, we will check whether there is a ckpt file existing to set the self.Flag for the method to load
        the network
        """
        #dataset_folder = "../Model"
        dataset_folder = "Model"
        # self.Flag = len(os.listdir(dataset_folder))
        #print os.listdir(dataset_folder)
        if len(os.listdir(dataset_folder))>0:
            self.Flag = True
        elif len(os.listdir(dataset_folder)) <=0:
            self.Flag = False

        return self

#matconvnet_value = ConvNetValue()
