import scipy.io as sio
import tensorflow as tf
import numpy as np
from PIL import Image
import os
matfile='imagenet-vgg-verydeep-19.mat'

"""该数据包含很多信息，我们需要的信息是每层神经网络的kernels和bias。
kernels的获取方式是data['layers'][0][第i层][0][0][0][0]，形状为[width, height, in_channels, out_channels]
bias的获取方式是data['layers'][0][第i层][0][0][0][0]，形状为[1,out_channels]
对于VGG-19的卷积，全部采用了3X3的filters，所以width为3，height为3。
VGG-19的pooling采用了长宽为2X2的max-pooling

注意，这里面的层数i，指的是最细粒度的层数，包括conv、relu、pool、fc各种操作。
    因此，i=0为卷积核，i=1为relu，i=2为卷积核，i=3为relu，i=4为pool，i=5为卷积核，……，i=37为全连接层，以此类推。
注意：# matconvnet: weights are [width, height, in_channels, out_channels]
     # tensorflow: weights are [height, width, in_channels, out_channels]
"""

VGG19_LAYERS = (
    'conv1_1', 'relu1_1', 'conv1_2', 'relu1_2', 'pool1',

    'conv2_1', 'relu2_1', 'conv2_2', 'relu2_2', 'pool2',

    'conv3_1', 'relu3_1', 'conv3_2', 'relu3_2', 'conv3_3', 'relu3_3', 'conv3_4', 'relu3_4', 'pool3',

    'conv4_1', 'relu4_1', 'conv4_2', 'relu4_2', 'conv4_3', 'relu4_3', 'conv4_4', 'relu4_4', 'pool4',

    'conv5_1', 'relu5_1', 'conv5_2', 'relu5_2', 'conv5_3', 'relu5_3', 'conv5_4', 'relu5_4', 'pool5',

    'fcon1',   'relu_fc1','fcon2',   'relu_fc2','fcon3',   'softmax'
)
LABEL_NAME=('daisy','dandelion','roses','sunflowers','tulips')

def preprocess(image,mean_pixel):
    return image-mean_pixel

def unprocess(image,mean_pixel):
    return image+mean_pixel
def write_feature_to_txt(file_name,feature):
    with open(file_name, 'w') as f:

        data_str ='0'
        for i in range(1):
            for j in range(7):
                for k in range(7):
                    for l in range(512):
                        data_str += ' '+str(feature[i, j, k,l])
        f.write(data_str)
        f.close()
def read_feature_from_txt(file_name):
    # 从txt  返回 特征矩阵
    f = open(file_name, 'r')
    lines = f.read()
    feature_value = [float(x) for x in lines.split(' ')]
    feature_map = np.reshape(feature_value[1:], [7, 7, 512])
    f.close()
    return feature_map
def get_or_create_feature_from_txt(sess,network,txt_file_name,image_dir,txt_dir,mean_pixel):
    # 从给定的txt文件 返回特征矩阵，如果不存在则 创建 txt 返回特征矩阵
    txt_file_path=os.path.join(txt_dir,txt_file_name)
    if not os.path.exists(txt_file_path):
        # 如果文件不存在， 从txt文件名获取 对应image名与image路径，得到对应的特征矩阵 并存入txt文件
        image_name = txt_file_name.split('.')[0] + '.jpg'
        image_path = os.path.join(image_dir, image_name)

        imdata = get_and_resize_image(image_name=image_path, resize_H=224, resize_W=224)
        imdata = preprocess(imdata, mean_pixel)
        #imtensor = tf.convert_to_tensor(imdata, dtype=tf.float32)

        #network = VGG19(imtensor, weights=weights, trainable=False)
        networks = sess.run(network, feed_dict={input_feature: imdata})
        #sess.run(tf.global_variables_initializer())
        feature = networks['pool5']

        write_feature_to_txt(txt_file_path,feature)
        print(txt_file_path,'*** save finish,return feature')
    else:

        feature= read_feature_from_txt(txt_file_path)
        print('get feature from ***',txt_file_path)
    return feature


def read_label_aslist(file_name):
    '''读取txt文件，每行为一个分类标签，返回标签列表'''
    f=open(file_name)
    lines=f.readlines()
    return lines

#PIL读取图片
def get_and_resize_image(image_name,resize_H,resize_W):
    imdata=Image.open(image_name)
    imdata=imdata.resize((resize_H, resize_W), Image.ANTIALIAS)
    #imdata.show()
    #print(imdate.size)  # （宽，高）
    #转换为numpy数组
    im_array = np.array(imdata)
    #print(im_array.shape)   #(行，列，通道)
    return im_array

def add_conv_layer(input_tensor,weights_value,bias_value,trainable,name):
    with tf.name_scope(name)as scope:
        kernels=tf.Variable(tf.constant(weights_value,dtype=tf.float32),trainable=trainable,name='kernels')
        bias=tf.Variable(tf.constant(bias_value,dtype=tf.float32),trainable=trainable,name='bias')
        conv=tf.nn.conv2d(input_tensor,kernels,strides=(1,1,1,1),padding='SAME',name='conv')
        add_bias=tf.nn.bias_add(conv,bias,name=scope)
    return add_bias

def add_pool_layer(input_tensor,name):
    with tf.name_scope(name):
        pool=tf.nn.max_pool(input_tensor,ksize=(1,2,2,1),strides=(1,2,2,1),padding='SAME',name='pool')
    return pool

def add_fc_layer(input_tensor,weights_value,bias_value,trainable,name):
    shape=input_tensor.get_shape().as_list()
    dim=1
    for d in shape[1:]:
        dim *= d
    input_tensor=tf.reshape(input_tensor,[-1,dim])
    with tf.name_scope(name)as scope:
        weights=tf.Variable(tf.constant(weights_value,dtype=tf.float32),trainable=trainable,name='weights')
        bias = tf.Variable(tf.constant(bias_value,dtype=tf.float32), trainable=trainable, name='bias')
        fc=tf.nn.bias_add(tf.matmul(input_tensor,weights),bias,name=scope)
    return fc

def add_softmax_layer(input_tensor):
    preds=tf.nn.softmax(input_tensor,name='prediction')
    return preds
def VGG19(input_image,weights,trainable):
    network={}
    net=input_image
    for i,name in enumerate(VGG19_LAYERS):
        layer_type=name[:4]
        if layer_type =='conv':
            kernels,bias=weights[i][0][0][0][0]
            kernels=np.transpose(kernels,(1,0,2,3))
            bias=bias.reshape(-1)
            net=add_conv_layer(net,weights_value=kernels,bias_value=bias,trainable=trainable,name=name)
            #print(name,'***',net.get_shape())
        elif layer_type == 'relu':
            net=tf.nn.relu(net)
            #print(name, '***', net.get_shape())
        elif layer_type == 'pool':
            net=add_pool_layer(net,name=name)
            #print(name, '***', net.get_shape())
        elif layer_type == 'soft':
            net=add_softmax_layer(net)
            #print(name, '***', net.get_shape())
        elif layer_type == 'fcon':
            kernels1, bias1 = weights[i][0][0][0][0]

            kernels1 = kernels1.reshape(-1, kernels1.shape[-1])
            bias1 = bias1.reshape(-1)
            net = add_fc_layer(net,weights_value=kernels1,bias_value=bias1,trainable=trainable,name=name)
            #print(name, '***', net.get_shape())
        network[name]=net
    return network

if __name__ == '__main__':
    matfile = 'imagenet-vgg-verydeep-19.mat'
    label_file = 'vgg19-1000.txt'
    lines = read_label_aslist(label_file)
    data=sio.loadmat(matfile)
    #mean=data['normalization'][0][0][0]
    #mean_pixel=np.mean(mean,axis=(0,1))
    mean_pixel = np.array([123.68, 116.779, 103.939]).reshape((1, 1, 1, 3))
    weights=data['layers'][0]

    input_feature=tf.placeholder(dtype=tf.float32,shape=[None, 224, 224, 3])
    input_label=tf.placeholder(dtype=tf.int64,shape=[None])

    network = VGG19(input_image=input_feature, weights=weights, trainable=False)

    with tf.name_scope('final_training_ops'):
        input_feature_=tf.reshape(input_feature,[-1,7*7*512])
        kernels = tf.Variable(tf.truncated_normal([7*7*512,5], stddev=0.001))
        biases = tf.Variable(tf.zeros([5]))
        logits = tf.matmul(input_feature_, kernels) + biases

    cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=input_label)
    cross_entropy_mean = tf.reduce_mean(cross_entropy)
    train_step = tf.train.GradientDescentOptimizer(0.01).minimize(cross_entropy_mean)

    saver=tf.train.Saver()
    num=0
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        image_dir='test_image'
        for image_name in os.listdir(image_dir):
            input_image=os.path.join(image_dir,image_name)

            imdata = get_and_resize_image(image_name=input_image, resize_H=224, resize_W=224)
            imdata = preprocess(imdata, mean_pixel)
            #imtensor = tf.convert_to_tensor(imdata, dtype=tf.float32)

            #print(imtensor.get_shape().as_list())
            #network = VGG19(input_image=imtensor, weights=weights, trainable=False)
            #networks=sess.run(network,feed_dict={input_feature:imdata})
            txt_file_name=image_name.split('.')[0]+'.txt'
            feature_map=get_or_create_feature_from_txt(sess=sess,network=network,txt_file_name=txt_file_name,image_dir='train_image',txt_dir='features',mean_pixel=mean_pixel)
            #feature=network['pool5']
            #write_feature_to_txt(file_name=txt_file_name,feature=feature)
            #print(input_image,'********************')