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

import logging

import time

import tensorflow as tf

from build_rnn import AFD_RNN

from utils import parser_cfg_file

from data_load import DataLoad

class AFD_RNN_Train(object):

    def __init__(self, train_config):

        self.learing_rate = float(train_config['learning_rate'])

        self.train_iterior = int(train_config['train_iteration'])

        self._train_logger_init()

        net_config = parser_cfg_file('./config/rnn_net.cfg')

        self.rnn_net = AFD_RNN(net_config)

        self.predict = self.rnn_net.build_net_graph()

        self.label = tf.placeholder(tf.float32, [None, self.rnn_net.time_step, self.rnn_net.class_num])

    def _compute_loss(self):

        with tf.name_scope('loss'):

            # [batchszie, time_step, class_num] ==> [time_step][batchsize, class_num]

            predict = tf.unstack(self.predict, axis=0)

            label = tf.unstack(self.label, axis=1)

            loss = [tf.nn.softmax_cross_entropy_with_logits(labels=label[i], logits=predict[i]) for i in range(self.rnn_net.time_step) ]

            loss = tf.reduce_mean(loss)

            train_op = tf.train.AdamOptimizer(self.learing_rate).minimize(loss)

        return loss, train_op

    def train_rnn(self):

        loss, train_op = self._compute_loss()

        with tf.name_scope('accuracy'):

            predict = tf.transpose(self.predict, [1,0,2])

            correct_pred = tf.equal(tf.argmax(self.label, 2), tf.argmax(predict, axis=2))

            accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

        dataset = DataLoad('./dataset/train/', time_step=self.rnn_net.time_step, class_num= self.rnn_net.class_num)

        saver = tf.train.Saver()

        with tf.Session() as sess:

            sess.run(tf.global_variables_initializer())

            for step in range(1, self.train_iterior+1):

                x, y = dataset.get_batch(self.rnn_net.batch_size)

                if step == 1:

                    feed_dict = {self.rnn_net.input_tensor: x, self.label: y}

                else:

                    feed_dict = {self.rnn_net.input_tensor: x, self.label: y, self.rnn_net.cell_state:state}

                _, compute_loss, state = sess.run([train_op, loss, self.rnn_net.cell_state], feed_dict=feed_dict)

                if step%10 == 0:

                    compute_accuracy = sess.run(accuracy, feed_dict=feed_dict)

                    self.train_logger.info('train step = %d,loss = %f,accuracy = %f'%(step, compute_loss, compute_accuracy))

                if step%1000 == 0:

                    save_path = saver.save(sess, './model/model.ckpt')

                    self.train_logger.info("train step = %d ,model save to =%s" % (step, save_path))

    def _train_logger_init(self):

        """

        初始化log日志

        :return:

        """

        self.train_logger = logging.getLogger('train')

        self.train_logger.setLevel(logging.DEBUG)

        # 添加文件输出

        log_file = './train_logs/' + time.strftime('%Y%m%d%H%M', time.localtime(time.time())) + '.logs'

        file_handler = logging.FileHandler(log_file, mode='w')

        file_handler.setLevel(logging.DEBUG)

        file_formatter = logging.Formatter('%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s')

        file_handler.setFormatter(file_formatter)

        self.train_logger.addHandler(file_handler)

        # 添加控制台输出

        consol_handler = logging.StreamHandler()

        consol_handler.setLevel(logging.DEBUG)

        consol_formatter = logging.Formatter('%(message)s')

        consol_handler.setFormatter(consol_formatter)

        self.train_logger.addHandler(consol_handler)

if __name__ == '__main__':

    train_config = parser_cfg_file('./config/train.cfg')

    train = AFD_RNN_Train(train_config)

    train.train_rnn()

    # a = tf.zeros([1,2,3])

    # b = tf.unstack(a, axis=1)

    # c = tf.zeros([2,1,3])

    # sess = tf.Session()

    # d = b[0]

    # print(sess.run(b[0]))

    #

    # pass