# -*- coding:utf-8 -*-

'''
DenseNet
参考：https://blog.csdn.net/u014380165/article/details/75142664

sigmoid 编码 round取整 [0,1,1,0,1] #有多个1
softmax 编码 最大值只会有一个 argmax [0,0,1,0,0] #只有一个1 可以直接用一个索引值代替 如：这里是2

如果最后编码层 shape [1,m*256]
使用sigmoid 编码 为[1,m*256] # {0,1}^(m*256)

使用softmax 编码 为[1,m*256]-->[1,m,256]-->[m,256]
每一个[1,256]对应一个索引值 --->[m,1] 如[8,9,10,……]其中数字为每一行的索引值

'''

import tensorflow as tf
from tensorflow.contrib import slim
import numpy as np
import cv2
import os
import argparse

os.environ['CUDA_VISIBLE_DEVICES'] = '1'  # 指定GPU from 0,1,2....

parse = argparse.ArgumentParser()
parse.add_argument('--mode', type=int, default=1, help='1 train ,0 test,2 编码验证测试,3 编码测试')
parse.add_argument('-ta', '--trainable', type=bool, default=True, help='trainable or not')
parse.add_argument('--lr', type=float, default=1e-4, help='learning rate')
parse.add_argument('--epochs', type=int, default=10, help='epochs')
arg = parse.parse_args()

# Hyperparameter
train = arg.mode  # 1 train ,0 test,2 输出编码测试
flag = arg.trainable
lr = arg.lr  # 1e-3~1e-6
batch_size = 64  # 逐步改变 128~128*4
img_Pixels_h = 128
img_Pixels_w = 64
img_Pixels_c = 3
num_classes = 10
epochs = arg.epochs
log_dir = './model'
keep_rata = 0.7

m = 16
k = 2 #256


# # 数据增强
# https://blog.csdn.net/medium_hao/article/details/79227056
def distort_color(image, color_ordering=0):
    if color_ordering == 0:
        image = tf.image.random_brightness(image, max_delta=32. / 255.)  # 亮度
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)  # 饱和度
        image = tf.image.random_hue(image, max_delta=0.2)  # 色相
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)  # 对比度
    if color_ordering == 1:
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
        image = tf.image.random_hue(image, max_delta=0.2)
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
    if color_ordering == 2:
        image = tf.image.random_hue(image, max_delta=0.2)
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
    if color_ordering == 3:
        image = tf.image.random_contrast(image, lower=0.5, upper=1.5)
        image = tf.image.random_brightness(image, max_delta=32. / 255.)
        image = tf.image.random_saturation(image, lower=0.5, upper=1.5)
        image = tf.image.random_hue(image, max_delta=0.2)
    return tf.clip_by_value(image, 0.0, 1.0)

def preprocess_for_train(image, height, width, bbox):
    if bbox is None:
        bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4])
    # if image.dytpe != tf.float32:
    #    image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    # bbox_begin, bbox_size, _ = tf.image.sample_distorted_bounding_box(tf.shape(image), bounding_boxes=bbox)
    # distorted_image = tf.slice(image, bbox_begin, bbox_size)
    # distorted_image = tf.image.resize_images(distorted_image, height, width, method=np.random.randint(4))
    distorted_image = tf.image.random_flip_left_right(image)
    distorted_image = distort_color(distorted_image, np.random.randint(4))
    return distorted_image

class Net():
    def __init__(self, images, dropout, trainable=True, reuse=None):
        self.images = images
        self.trainable = trainable
        self.reuse = reuse
        self.dropout = dropout
        self.forward_Net()

    def forward_Net(self):
        # net_work={}
        net = self.images  # [n,128,64,3]

        '''
        x1 = []
        # x2=[]
        for i in range(net.shape[0]):
            # boxes = tf.constant([[[0.05, 0.05, 0.9, 0.7], [0.35, 0.47, 0.5, 0.56]]])
            x2 = tf.squeeze(tf.slice(net, [i, 0, 0, 0], [1, img_Pixels_h, img_Pixels_w, img_Pixels_c]), 0)
            x1.append(preprocess_for_train(x2, img_Pixels_h, img_Pixels_w, None))
            
            # 数据增强
        net = tf.convert_to_tensor(x1, tf.float32)
        '''
        net = slim.conv2d(net, 32, 7, (2, 1), padding='SAME', activation_fn=tf.nn.leaky_relu,
                          normalizer_fn=slim.layers.batch_norm, reuse=self.reuse, trainable=self.trainable,
                          scope='conv1')  # [n,64,64,32]
        net = tf.layers.max_pooling2d(net, 3, 2, 'same', name='pool1')  # [n,32,32,32]

        # Dense Block-1
        branch_1 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse, trainable=self.trainable,
                               scope='db1_1_1')
        branch_1 = slim.conv2d(branch_1, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db1_1_2')  # [n,32,32,32]

        net = tf.concat((net, branch_1), -1)  # [n,32,32,64]

        branch_2 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db1_2_1')  # [n,32,32,32]
        branch_2 = slim.conv2d(branch_2, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db1_2_2')  # [n,32,32,64]

        net = tf.concat((net, branch_2), -1)  # [n,32,32,128]

        branch_3 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db1_3_1')  # [n,32,32,64]
        branch_3 = slim.conv2d(branch_3, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm,
                               reuse=self.reuse,
                               trainable=self.trainable, scope='db1_3_2')  # [n,32,32,128]

        net = tf.concat((net, branch_3), -1)  # [n,32,32,256]

        # transition layer
        net = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                          normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                          trainable=self.trainable, scope='tl1_1')  # [n,32,32,128]
        net = slim.avg_pool2d(net, 2, 2, 'same', scope='pool2')  # [n,16,16,128]

        net = tf.layers.dropout(net, self.dropout)

        # Dense Block-2
        branch_1 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse, trainable=self.trainable,
                               scope='db2_1_1')
        branch_1 = slim.conv2d(branch_1, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db2_1_2')  # [n,16,16,128]

        net = tf.concat((net, branch_1), -1)  # [n,32,32,256]

        branch_2 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db2_2_1')  # [n,16,16,128]
        branch_2 = slim.conv2d(branch_2, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db2_2_2')  # [n,16,16,256]

        net = tf.concat((net, branch_2), -1)  # [n,16,16,512]

        branch_3 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db2_3_1')  # [n,16,16,256]
        branch_3 = slim.conv2d(branch_3, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm,
                               reuse=self.reuse,
                               trainable=self.trainable, scope='db2_3_2')  # [n,16,16,512]

        net = tf.concat((net, branch_3), -1)  # [n,16,16,1024]

        # transition layer
        net = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                          normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                          trainable=self.trainable, scope='tl2_1')  # [n,16,16,512]
        net = slim.avg_pool2d(net, 2, 2, 'same', scope='pool2')  # [n,8,8,512]

        net = tf.layers.dropout(net, self.dropout)

        # Dense Block-3
        branch_1 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse, trainable=self.trainable,
                               scope='db3_1_1')
        branch_1 = slim.conv2d(branch_1, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db3_1_2')  # [n,16,16,128]

        net = tf.concat((net, branch_1), -1)  # [n,32,32,256]

        branch_2 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db3_2_1')  # [n,16,16,128]
        branch_2 = slim.conv2d(branch_2, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db3_2_2')  # [n,16,16,256]

        net = tf.concat((net, branch_2), -1)  # [n,16,16,512]

        branch_3 = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                               trainable=self.trainable, scope='db3_3_1')  # [n,16,16,256]
        branch_3 = slim.conv2d(branch_3, net.shape[-1], 3, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                               normalizer_fn=slim.layers.batch_norm,
                               reuse=self.reuse,
                               trainable=self.trainable, scope='db3_3_2')  # [n,16,16,512]

        net = tf.concat((net, branch_3), -1)  # [n,16,16,1024]

        # transition layer
        net = slim.conv2d(net, net.shape[-1] // 2, 1, 1, padding='same', activation_fn=tf.nn.leaky_relu,
                          normalizer_fn=slim.layers.batch_norm, reuse=self.reuse,
                          trainable=self.trainable, scope='tl3_1')  # [n,16,16,512]
        net = slim.avg_pool2d(net, 2, 2, 'same', scope='pool3')  # [n,4,4,512]

        net = tf.layers.dropout(net, self.dropout)

        # Classifications
        net = slim.avg_pool2d(net, 4, 1, 'valid', scope='pool2')  # [n,1,1,512]

        net = tf.reshape(net, [-1, np.int(np.prod(net.get_shape()[1:]))], name='flatten1')  # [n,512]

        self.fc7 = net
        # net = slim.fully_connected(net, 512, activation_fn=tf.nn.relu,trainable=self.trainable, scope='fc1')

        net = slim.fully_connected(net, m * k, activation_fn=tf.nn.relu, scope='coding_layer') # sigmoid ,softmax

        feat = tf.split(net, num_or_size_splits=m, axis=1)
        out = [None] * m
        for i in range(m):
            out[i] = tf.nn.softmax(feat[i])

        self.net = tf.concat(out, axis=1)  # [n,m*k]

        # self.net = net  # [n,m*k]

        self.pred = slim.fully_connected(net, num_classes, activation_fn=tf.nn.softmax, scope='output')

def sign(x):
    '''
    :param x: 矩阵
    :return: 矩阵
    '''
    # y = np.zeros_like(x,np.uint8)
    shape = x.shape
    x = x.flatten()
    y = np.zeros_like(x, np.uint8)

    for i in range(len(x)):
        if x[i] <= 0:
            y[i] = 0
        elif x[i] > 0 and x[i] <= 0.01:
            y[i] = 1
        elif x[i] > 0.01 and x[i] <= 0.1:
            y[i] = 2
        elif x[i] > 0.1 and x[i] <= 0.2:
            y[i] = 3
        elif x[i] > 0.2 and x[i] <= 0.3:
            y[i] = 4
        elif x[i] > 0.3 and x[i] <= 0.4:
            y[i] = 5
        elif x[i] > 0.4 and x[i] <= 0.5:
            y[i] = 6
        elif x[i] > 0.5 and x[i] <= 0.6:
            y[i] = 7
        elif x[i] > 0.6 and x[i] <= 0.7:
            y[i] = 8
        elif x[i] > 0.7 and x[i] <= 0.8:
            y[i] = 9
        elif x[i] > 0.8 and x[i] <= 0.9:
            y[i] = 10
        elif x[i] > 0.9 and x[i] <= 1:
            y[i] = 11
    return np.reshape(y, shape)

if __name__ == "__main__":
    # data = Data('train.data', batch_size)
    # data.Next_batch()

    x = tf.placeholder(tf.float32, (None, img_Pixels_h, img_Pixels_w, img_Pixels_c))
    y = tf.placeholder(tf.int64, (None,))
    dropout = tf.placeholder(tf.float32)

    net = Net(x, dropout, flag)
    coding = net.net
    pred = net.pred
    fc7 = net.fc7

    # Initializing the variables
    init = tf.global_variables_initializer()

    # Launch the graph
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = 0.7  # 占用GPU70%的显存

    saver = tf.train.Saver(max_to_keep=1)

    # steps = data.num_images // batch_size
    # index2 = np.arange(0, batch_size)

    with tf.Session(config=config) as sess:
        sess.run(init)
        ckpt = tf.train.get_checkpoint_state(log_dir)
        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)

        # 输入一张查询图片，输出查询结果
        if train ==1:
            # 每个类别计算ap，所有类别取平均得到mAp
            def voc_ap(rec, prec, use_07_metric=False):
                """Compute VOC AP given precision and recall. If use_07_metric is true, uses
                the VOC 07 11-point method (default:False).
                """
                if use_07_metric:
                    # 11 point metric
                    ap = 0.
                    for t in np.arange(0., 1.1, 0.1):
                        if np.sum(rec >= t) == 0:
                            p = 0
                        else:
                            p = np.max(prec[rec >= t])
                        ap = ap + p / 11.
                else:
                    # correct AP calculation
                    # first append sentinel values at the end
                    mrec = np.concatenate(([0.], rec, [1.]))
                    mpre = np.concatenate(([0.], prec, [0.]))

                    # compute the precision envelope
                    for i in range(mpre.size - 1, 0, -1):
                        mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

                    # to calculate area under PR curve, look for points
                    # where X axis (recall) changes value
                    i = np.where(mrec[1:] != mrec[:-1])[0]

                    # and sum (\Delta recall) * prec
                    ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
                return ap

            import glob
            image_paths=glob.glob('./train_1/*/*') # ['/xxx/xxx/xxx.jpg','/xxx/xxx/xxx.jpg',...]
            # np.random.shuffle(image_paths)
            # 选择第一张作为查询图片
            # query_image=image_paths[0]
            # print('查询影像：',query_image)

            def load_image(image_paths):
                # 剩下的为查询的源图片
                images = []
                labels = []

                for path in image_paths:
                    labels.append(int(path.split('/')[-2]))
                    images.append(cv2.imread(path))

                # 特征归一化处理
                imgs = np.asarray(images, np.float32)
                imgs = (imgs - np.min(imgs, 0)) * 1. / (np.max(imgs, 0) - np.min(imgs, 0))

                return imgs,np.asarray(labels, np.int64)

            N = len(image_paths)
            feat = coding
            steps = N // batch_size
            # N=steps*batch_size
            # 样本较多，按批进行
            encoding = np.zeros((N, m * k))  # 记录每张图片的编码  [64,2048]
            # labels = np.zeros((N,), np.int64)  # 记录每张图的真实标签 [64,]
            # score = np.zeros((N, num_classes))  # 记录每张图片的分成每一类的概率值  [64,10]
            # cols=c7.shape[-1]
            cols = fc7.get_shape().as_list()[-1]
            cols = int(cols)
            fc7_Eu_ = np.zeros((N, cols))  # [64,2048]

            imgs, labels = load_image(image_paths)


            start=0;end=0

            for step in range(steps):
                end=min((start+batch_size,len(image_paths)))
                batch_x, batch_y = imgs[start:end],labels[start:end]
                if start>=len(image_paths):start=0

                feat_1 = sess.run(feat, {x: batch_x, dropout: 1.})  # 编码
                pred_1 = sess.run(pred, {x: batch_x, dropout: 1.})  # 预测类别
                fc7_1 = sess.run(fc7, {x: batch_x, dropout: 1.})
                encoding[step * batch_size:(step + 1) * batch_size] = feat_1
                fc7_Eu_[step * batch_size:(step + 1) * batch_size] = fc7_1

            score_1 = np.zeros([N - 1, ])

            encoding0 = encoding.copy()

            # print(sign(encoding0))

            e=np.round(encoding).astype(np.uint8)
            with open('coding.txt','w') as fp:
                for i in range(len(encoding0)):
                    fp.write(str(e[i])+'\n')
            with open('name.txt','w') as fp:
                for i in range(len(encoding0)):
                    fp.write(str(image_paths[i])+'\n')
            # exit(-1)


            # 根据编码粗搜索
            encoding = np.round(encoding).astype(np.uint8)  # [N,m*k]

            encoding_1 = np.bitwise_xor(encoding[1:], encoding[0])

            encoding_2 = np.asarray(np.sum(encoding_1, 1) < int(m * k * 2 * 0.18), np.uint8)  # [n-1,] 0.18～0.2

            encoding_8=np.sum(encoding_1, 1) # 编码的汉明距离

            encoding_2_1 = encoding_2.copy()

            # 精搜索
            # 使用编码层求欧式距离
            fc8 = []
            index = []
            for i, x in enumerate(encoding_2):
                if x:
                    fc8.append(encoding0[1:][i])
                    index.append(i)

            # fc8与fc7[0]做欧式距离，fc8为0的一行跳过
            fc8 = np.asarray(fc8)
            if len(fc8) == 0: exit(-1)

            error = fc8 - encoding0[0]

            Eu_dis = np.sum(error ** 2, 1) / (m * k)  # [N-1,]

            Eu_dis_2=np.ones([N-1,])*10 # 也可以按照该字段 从小到大排序
            for i in range(len(Eu_dis)):
                Eu_dis_2[index[i]]=Eu_dis[i]

            Eu_dis = np.asarray((Eu_dis < 0.1), np.uint8)  # [x,] 0.065～0.1
            # print('Eu_dis', Eu_dis)
            # 结合欧式距离，返回到源图的序号上
            encoding_3 = encoding_2.copy()
            # --------------------------------------------------------

            # 记录查询到的图片
            results=[]
            encoding_results=[] # 记录查询到的图片编码
            for i in range(len(Eu_dis)):
                if not Eu_dis[i]:
                    encoding_3[index[i]] = 0
                else:
                    results.append(labels[index[i]])
                    encoding_results.append(encoding[index[i]])

            # 选择第一张作为查询图片
            print('查询影像：\n', image_paths[0])
            print('查询影像的编码：\n', encoding[0])

            # print('查询影像：\n', image_paths[0])
            # print('查询影像的编码：\n', encoding0[0])
            # print('查询影像的label：\n', labels[0])


            print('Query result:\n', results)
            print('Query result code:\n',encoding_results)
            exit(-1)



            labels_1 = np.asarray((labels[1:] - labels[0]) == 0, np.uint8)  # [N-1,]

            encoding_4 = encoding_3 * labels_1  # 编码一样，还有保证对应的真实标签是一样的（否则说明编码不对）

            comp1_cls_test = np.zeros((N - 1, 7))
            comp1_cls_test[:, 0] = score_1
            comp1_cls_test[:, 1] = encoding_4
            comp1_cls_test[:, 2] = encoding_2_1
            comp1_cls_test[:, 3] = labels_1
            comp1_cls_test[:, 4] = encoding_3
            comp1_cls_test[:,5]=Eu_dis_2
            comp1_cls_test[:,6]=encoding_8

            # 按score 字段排序
            import pandas as pd

            comp1_cls_test = pd.DataFrame(comp1_cls_test,
                                          columns=['score', 'encoding', 'True encoding', 'True label', 'Eu_dis','Eu_error','Hamming_error'])

            # comp1_cls_test = comp1_cls_test.sort_values('score', ascending=False)
            comp1_cls_test = comp1_cls_test.sort_values('Eu_error', ascending=True)
            print(comp1_cls_test[:50])
            # exit(-1)
            # 根据这个排序表，计算Precision，Recall ，计算top-10的结果
            Precision = np.sum(comp1_cls_test.values[:10][:, 1]) / 10
            print('Precision', Precision)

            Recall = np.sum(comp1_cls_test.values[:10][:, 1]) / np.sum(comp1_cls_test.values[:, 1])
            print('Recall', Recall)

            # 计算ap
            ap = voc_ap([Recall], [Precision])
            print('ap', ap)
