#-*- coding=utf-8 -*-
from __future__ import division, print_function, absolute_import

from base_model import BaseModel
from helper import *
import tensorflow as tf
import pickle
import numpy as np
import time

class Vgg16(BaseModel):
    default_param = {
        "loss" : "square_loss",
        "metrics" : ["loss"],
        "optimizer" : "sgd",
        "learning_rate" : 1e-2,
        "batch_size" : 100,
        "num_epochs" : 25,
        "keep_prob":0.75
    }

    def __init__(self,classnum):
        self.class_num = classnum
        self.model = None
        self.sess = tf.Session()
        self.scope = {}
        self.summary = []
    def conv2d(self,layer_name,inputs, out_channels, kernel_size, strides=1, padding='SAME'):
        in_channels = inputs.get_shape()[-1]
        with tf.variable_scope(layer_name) as scope:
            self.scope[layer_name] = scope
            w = tf.get_variable(name='weights',
                                trainable=True,
                                shape=[kernel_size, kernel_size, in_channels, out_channels],
                                initializer=tf.contrib.layers.xavier_initializer())
            b = tf.get_variable(name='biases',
                                trainable=True,
                                shape=[out_channels],
                                initializer=tf.constant_initializer(0.0))
            inputs = tf.nn.conv2d(inputs, w, [1, strides, strides, 1], padding=padding, name='conv')
            inputs = tf.nn.bias_add(inputs, b, name='bias_add')
            inputs = tf.nn.relu(inputs, name='relu')
            return inputs

    def max_pool(self, layer_name, inputs, pool_size, strides, padding='SAME'):
        with tf.name_scope(layer_name):
            return tf.nn.max_pool(inputs, [1, pool_size, pool_size, 1], [1, strides, strides, 1], padding=padding,
                                  name=layer_name)

    def avg_pool(self, layer_name, inputs, pool_size, strides, padding='SAME'):
        with tf.name_scope(layer_name):
            return tf.nn.avg_pool(inputs, [1, pool_size, pool_size, 1], [1, strides, strides, 1], padding=padding,
                                  name=layer_name)

    def lrn(self, layer_name, inputs, depth_radius=5, alpha=0.0001, beta=0.75):
        with tf.name_scope(layer_name):
            return tf.nn.local_response_normalization(name='pool1_norm1', input=inputs, depth_radius=depth_radius,
                                                      alpha=alpha, beta=beta)

    def concat(self, layer_name, inputs):
        with tf.name_scope(layer_name):
            one_by_one = inputs[0]
            three_by_three = inputs[1]
            five_by_five = inputs[2]
            pooling = inputs[3]
            return tf.concat([one_by_one, three_by_three, five_by_five, pooling], axis=3)

    def dropout(self, layer_name, inputs, keep_prob):
        # dropout_rate = 1 - keep_prob
        with tf.name_scope(layer_name):
            return tf.nn.dropout(name=layer_name, x=inputs, keep_prob=keep_prob)

    def bn(self, layer_name, inputs, epsilon=1e-3):
        with tf.name_scope(layer_name):
            batch_mean, batch_var = tf.nn.moments(inputs, [0])
            inputs = tf.nn.batch_normalization(inputs, mean=batch_mean, variance=batch_var, offset=None,
                                               scale=None, variance_epsilon=epsilon)
            return inputs

    def fc(self, layer_name, inputs, out_nodes):
        shape = inputs.get_shape()
        if len(shape) == 4:  # x is 4D tensor
            size = shape[1].value * shape[2].value * shape[3].value
        else:  # x has already flattened
            size = shape[-1].value
        with tf.variable_scope(layer_name) as scope:
            self.scope[layer_name] = scope
            w = tf.get_variable('weights',
                                shape=[size, out_nodes],
                                initializer=tf.contrib.layers.xavier_initializer())
            b = tf.get_variable('biases',
                                shape=[out_nodes],
                                initializer=tf.constant_initializer(0.0))
            flat_x = tf.reshape(inputs, [-1, size])
            inputs = tf.nn.bias_add(tf.matmul(flat_x, w), b)
            inputs = tf.nn.relu(inputs)
            return inputs

    def build_model(self):

        # 训练数据
        self.inputs = tf.placeholder(tf.float32, shape=[None, 224, 224, 3])

        # 训练标签数据
        self.labels = tf.placeholder(tf.float32, shape=[None, self.class_num])

        # dropout
        self.keep_prob = tf.placeholder(tf.float32)

        self.conv1_1 = self.conv2d("conv1_1",self.inputs,64,3)
        self.conv1_2 = self.conv2d("conv1_2",self.conv1_1, 64,3)
        self.pool1 = self.max_pool('pool1',self.conv1_2,pool_size=2,strides=2)
        #112*112*64

        self.conv2_1 = self.conv2d("conv2_1",self.pool1, 128,3)
        self.conv2_2 = self.conv2d( "conv2_2",self.conv2_1, 128,3)
        self.pool2 = self.max_pool("pool2",self.conv2_2,pool_size=2,strides=2)
        #56*56*128

        self.conv3_1 = self.conv2d("conv3_1",self.pool2, 256,3)
        self.conv3_2 = self.conv2d("conv3_2",self.conv3_1, 256,3)
        self.conv3_3 = self.conv2d("conv3_3",self.conv3_2, 256, 3)
        self.pool3 = self.max_pool("pool3",self.conv3_3,pool_size=2,strides=2)
        #28*28*256

        self.conv4_1 = self.conv2d("conv4_1",self.pool3, 512, 3)
        self.conv4_2 = self.conv2d("conv4_2",self.conv4_1, 512, 3)
        self.conv4_3 = self.conv2d("conv4_3",self.conv4_2, 512, 3)
        self.pool4 = self.max_pool("pool4",self.conv4_3, pool_size=2,strides=2)
        #14*14*512

        self.conv5_1 = self.conv2d("conv5_1",self.pool4, 512, 3)
        self.conv5_2 = self.conv2d("conv5_2",self.conv5_1, 512, 3)
        self.conv5_3 = self.conv2d("conv5_3",self.conv5_2, 512, 3)
        self.pool5 = self.max_pool( 'pool5',self.conv5_3,pool_size=2,strides=2)
        #7*7*512

        self.fc6 = self.fc("fc6",self.pool5,4096)  # 25088 = 7*7*512

        self.relu6 = tf.nn.dropout(self.fc6, self.keep_prob)

        self.fc7 = self.fc("fc7",self.relu6,4096)

        self.relu7 = tf.nn.dropout(self.fc7, self.keep_prob)

        self.pred = self.fc("fc8",self.relu7, self.class_num)


    def set_parameter(self, param):
        for name in self.default_param:
            if name not in param:
                param[name] = self.default_param[name]

        self.build_model()
        # 定义交叉熵损失函数
        self.keep_prob_value = param["keep_prob"]

        loss_fun = param["loss"]
        self.loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=self.pred, labels=self.labels))


        optimizer = param["optimizer"]
        self.learning_rate = param["learning_rate"]

        self.optimizer = tf.train.RMSPropOptimizer(self.learning_rate).minimize(self.loss)

        self.correct_prediction = tf.equal(tf.argmax(self.pred, 1), tf.argmax(self.labels, 1))
        self.accuracy = tf.reduce_mean(tf.cast(self.correct_prediction, tf.float32))

        self.batch_size = param["batch_size"]
        self.num_epochs = param["num_epochs"]

    def get_batch(self, feed_data):
        X = feed_data["inputs"]
        Y = feed_data["labels"]
        totalbatch = int(len(X)/self.batch_size)+1
        if (totalbatch * self.batch_size == len(X)):
            totalbatch = totalbatch - 1

        for i in range(0,totalbatch):
            startindex = i*self.batch_size
            endindex = (i+1)*self.batch_size
            batch_xs = X[startindex:endindex]
            batch_ys = Y[startindex:endindex]
            yield { "batch_xs" : batch_xs, "batch_ys" : batch_ys }

    def train(self, feed_data):
        self.sess.run(tf.global_variables_initializer())
        trainstep = 0
        for epoch in range(self.num_epochs):
            avg_cost = 0.0
            totalaccuracy = 0.0
            for batch in self.get_batch(feed_data):
                feed_dict = {
                    self.inputs : batch["batch_xs"],
                    self.labels : batch["batch_ys"],
                    self.keep_prob: self.keep_prob_value,
                }
                _, loss, acc = self.sess.run([self.optimizer, self.loss,self.accuracy], feed_dict=feed_dict)
                totalaccuracy += acc*len(batch["batch_xs"])
                avg_cost +=  loss
                trainstep = trainstep + 1

            totalaccuracy /= len(feed_data['inputs'])
            print("train_step"+"\t"+str(trainstep)+"\t"+"epoch:"+"\t"+str(epoch+1)+"\t"+"accuracy:"+"\t"+str(totalaccuracy)+"\t"+"loss:"+"\t"+str(avg_cost))

    def model_load(self,path):
        saver = tf.train.Saver()
        saver.restore(self.sess, path)
        return

    def model_save(self,path):
        saver = tf.train.Saver()
        saver.save(self.sess, path)
        return

    def evaluate(self, feed_data):
        avg_loss = 0.0
        totalaccuracy = 0.0
        totallen = len(feed_data["inputs"])

        for batch in self.get_batch(feed_data):
            feed_dict = {
                self.inputs: batch["batch_xs"],
                self.labels: batch["batch_ys"],
                self.keep_prob:self.keep_prob_value
            }
            loss, acc = self.sess.run([self.loss, self.accuracy], feed_dict=feed_dict)
            totalaccuracy += acc * len(batch["batch_xs"])
            avg_loss += loss
        avg_loss /= totallen
        totalaccuracy /= len(feed_data['inputs'])

        res = {"accuracy":totalaccuracy,"loss":avg_loss}
        return res

    def predict(self, feed_data):
        res = []
        for batch in self.get_batch(feed_data):
            feed_dict = {
                self.inputs: batch["batch_xs"]
            }
            pred = self.sess.run(self.pred, feed_dict=feed_dict)
            res.extend(pred.tolist())

        return res


[((25, 20, 25, 32), 'tensorflow.Session', 'tf.Session', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((102, 22, 102, 75), 'tensorflow.placeholder', 'tf.placeholder', (), '', True, 'import tensorflow as tf\n'), ((105, 22, 105, 78), 'tensorflow.placeholder', 'tf.placeholder', (), '', True, 'import tensorflow as tf\n'), ((108, 25, 108, 51), 'tensorflow.placeholder', 'tf.placeholder', ({(108, 40, 108, 50): 'tf.float32'}, {}), '(tf.float32)', True, 'import tensorflow as tf\n'), ((140, 21, 140, 60), 'tensorflow.nn.dropout', 'tf.nn.dropout', ({(140, 35, 140, 43): 'self.fc6', (140, 45, 140, 59): 'self.keep_prob'}, {}), '(self.fc6, self.keep_prob)', True, 'import tensorflow as tf\n'), ((144, 21, 144, 60), 'tensorflow.nn.dropout', 'tf.nn.dropout', ({(144, 35, 144, 43): 'self.fc7', (144, 45, 144, 59): 'self.keep_prob'}, {}), '(self.fc7, self.keep_prob)', True, 'import tensorflow as tf\n'), ((208, 16, 208, 32), 'tensorflow.train.Saver', 'tf.train.Saver', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((213, 16, 213, 32), 'tensorflow.train.Saver', 'tf.train.Saver', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((30, 13, 30, 42), 'tensorflow.variable_scope', 'tf.variable_scope', ({(30, 31, 30, 41): 'layer_name'}, {}), '(layer_name)', True, 'import tensorflow as tf\n'), ((40, 21, 40, 100), 'tensorflow.nn.conv2d', 'tf.nn.conv2d', (), '', True, 'import tensorflow as tf\n'), ((41, 21, 41, 63), 'tensorflow.nn.bias_add', 'tf.nn.bias_add', (), '', True, 'import tensorflow as tf\n'), ((42, 21, 42, 52), 'tensorflow.nn.relu', 'tf.nn.relu', (), '', True, 'import tensorflow as tf\n'), ((46, 13, 46, 38), 'tensorflow.name_scope', 'tf.name_scope', ({(46, 27, 46, 37): 'layer_name'}, {}), '(layer_name)', True, 'import tensorflow as tf\n'), ((47, 19, 48, 50), 'tensorflow.nn.max_pool', 'tf.nn.max_pool', (), '', True, 'import tensorflow as tf\n'), ((51, 13, 51, 38), 'tensorflow.name_scope', 'tf.name_scope', ({(51, 27, 51, 37): 'layer_name'}, {}), '(layer_name)', True, 'import tensorflow as tf\n'), ((52, 19, 53, 50), 'tensorflow.nn.avg_pool', 'tf.nn.avg_pool', (), '', True, 'import tensorflow as tf\n'), ((56, 13, 56, 38), 'tensorflow.name_scope', 'tf.name_scope', ({(56, 27, 56, 37): 'layer_name'}, {}), '(layer_name)', True, 'import tensorflow as tf\n'), ((57, 19, 58, 77), 'tensorflow.nn.local_response_normalization', 'tf.nn.local_response_normalization', (), '', True, 'import tensorflow as tf\n'), ((61, 13, 61, 38), 'tensorflow.name_scope', 'tf.name_scope', ({(61, 27, 61, 37): 'layer_name'}, {}), '(layer_name)', True, 'import tensorflow as tf\n'), ((66, 19, 66, 89), 'tensorflow.concat', 'tf.concat', (), '', True, 'import tensorflow as tf\n'), ((70, 13, 70, 38), 'tensorflow.name_scope', 'tf.name_scope', ({(70, 27, 70, 37): 'layer_name'}, {}), '(layer_name)', True, 'import tensorflow as tf\n'), ((71, 19, 71, 80), 'tensorflow.nn.dropout', 'tf.nn.dropout', (), '', True, 'import tensorflow as tf\n'), ((74, 13, 74, 38), 'tensorflow.name_scope', 'tf.name_scope', ({(74, 27, 74, 37): 'layer_name'}, {}), '(layer_name)', True, 'import tensorflow as tf\n'), ((75, 36, 75, 62), 'tensorflow.nn.moments', 'tf.nn.moments', ({(75, 50, 75, 56): 'inputs', (75, 58, 75, 61): '[0]'}, {}), '(inputs, [0])', True, 'import tensorflow as tf\n'), ((76, 21, 77, 84), 'tensorflow.nn.batch_normalization', 'tf.nn.batch_normalization', (), '', True, 'import tensorflow as tf\n'), ((86, 13, 86, 42), 'tensorflow.variable_scope', 'tf.variable_scope', ({(86, 31, 86, 41): 'layer_name'}, {}), '(layer_name)', True, 'import tensorflow as tf\n'), ((94, 21, 94, 51), 'tensorflow.reshape', 'tf.reshape', ({(94, 32, 94, 38): 'inputs', (94, 40, 94, 50): '[-1, size]'}, {}), '(inputs, [-1, size])', True, 'import tensorflow as tf\n'), ((96, 21, 96, 39), 'tensorflow.nn.relu', 'tf.nn.relu', ({(96, 32, 96, 38): 'inputs'}, {}), '(inputs)', True, 'import tensorflow as tf\n'), ((159, 35, 159, 112), 'tensorflow.nn.softmax_cross_entropy_with_logits', 'tf.nn.softmax_cross_entropy_with_logits', (), '', True, 'import tensorflow as tf\n'), ((167, 43, 167, 66), 'tensorflow.argmax', 'tf.argmax', ({(167, 53, 167, 62): 'self.pred', (167, 64, 167, 65): '1'}, {}), '(self.pred, 1)', True, 'import tensorflow as tf\n'), ((167, 68, 167, 93), 'tensorflow.argmax', 'tf.argmax', ({(167, 78, 167, 89): 'self.labels', (167, 91, 167, 92): '1'}, {}), '(self.labels, 1)', True, 'import tensorflow as tf\n'), ((168, 39, 168, 83), 'tensorflow.cast', 'tf.cast', ({(168, 47, 168, 70): 'self.correct_prediction', (168, 72, 168, 82): 'tf.float32'}, {}), '(self.correct_prediction, tf.float32)', True, 'import tensorflow as tf\n'), ((188, 22, 188, 55), 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((95, 36, 95, 56), 'tensorflow.matmul', 'tf.matmul', ({(95, 46, 95, 52): 'flat_x', (95, 54, 95, 55): 'w'}, {}), '(flat_x, w)', True, 'import tensorflow as tf\n'), ((165, 25, 165, 70), 'tensorflow.train.RMSPropOptimizer', 'tf.train.RMSPropOptimizer', ({(165, 51, 165, 69): 'self.learning_rate'}, {}), '(self.learning_rate)', True, 'import tensorflow as tf\n'), ((35, 44, 35, 82), 'tensorflow.contrib.layers.xavier_initializer', 'tf.contrib.layers.xavier_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((39, 44, 39, 72), 'tensorflow.constant_initializer', 'tf.constant_initializer', ({(39, 68, 39, 71): '0.0'}, {}), '(0.0)', True, 'import tensorflow as tf\n'), ((90, 44, 90, 82), 'tensorflow.contrib.layers.xavier_initializer', 'tf.contrib.layers.xavier_initializer', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((93, 44, 93, 72), 'tensorflow.constant_initializer', 'tf.constant_initializer', ({(93, 68, 93, 71): '0.0'}, {}), '(0.0)', True, 'import tensorflow as tf\n')]