# -*- coding:utf-8 -*-
"""
 the idea of this script came from LUNA2016 champion paper.
 This model composed of three network,namely Archi-1(size of 10x10x6),Archi-2(size of 30x30x10),Archi-3(size of 40x40x26)

"""
import tensorflow as tf
from data_prepare import get_train_batch,get_all_filename,get_test_batch
import random
import time

class model(object):

    def __init__(self,learning_rate,keep_prob,batch_size,epoch):
        print(" network begin...")
        self.learning_rate = learning_rate
        self.keep_prob = keep_prob
        self.batch_size = batch_size
        self.epoch = epoch

        self.cubic_shape = [[6, 20, 20], [10, 30, 30], [26, 40, 40]]

    def archi_1(self,input,keep_prob):
        with tf.name_scope("Archi-1"):
            # input size is batch_sizex20x20x6
            # 5x5x3 is the kernel size of conv1,1 is the input depth,64 is the number output channel
            w_conv1 = tf.Variable(tf.random_normal([3,5,5,1,32],stddev=0.001),dtype=tf.float32,name='w_conv1')
            b_conv1 = tf.Variable(tf.constant(0.01,shape=[32]),dtype=tf.float32,name='b_conv1')
            out_conv1 = tf.nn.relu(tf.add(tf.nn.conv3d(input,w_conv1,strides=[1,1,1,1,1],padding='SAME'),b_conv1))
            out_conv1 = tf.nn.dropout(out_conv1,keep_prob)
            out_conv1_shape = tf.shape(out_conv1)

            # max pooling ,pooling layer has no effect on the data size
            hidden_conv1 = tf.nn.max_pool3d(out_conv1,strides=[1,2,2,2,1],ksize=[1,2,2,2,1],padding='SAME')
            hidden_conv1_shape = tf.shape(hidden_conv1)

            # after conv1 ,the output size is batch_sizex4x16x16x64([batch_size,in_deep,width,height,output_deep])
            w_conv2 = tf.Variable(tf.random_normal([3,5, 5, 32,64], stddev=0.001), dtype=tf.float32,name='w_conv2')
            b_conv2 = tf.Variable(tf.constant(0.01, shape=[64]), dtype=tf.float32, name='b_conv2')
            out_conv2 = tf.nn.relu(tf.add(tf.nn.conv3d(hidden_conv1, w_conv2, strides=[1, 1, 1,1, 1], padding='SAME'), b_conv2))
            out_conv2 = tf.nn.dropout(out_conv2, keep_prob)
            out_conv2_shape = tf.shape(out_conv2)

            hidden_conv2 = tf.nn.max_pool3d(out_conv2,strides=[1,2,2,2,1],ksize=[1,2,2,2,1],padding='SAME')
            hidden_conv2_shape = tf.shape(hidden_conv2)

            # after conv2 ,the output size is batch_sizex2x12x12x64([batch_size,in_deep,width,height,output_deep])
            w_conv3 = tf.Variable(tf.random_normal([1,5, 5, 64,64], stddev=0.001), dtype=tf.float32, name='w_conv3')
            b_conv3 = tf.Variable(tf.constant(0.01, shape=[64]), dtype=tf.float32, name='b_conv3')
            out_conv3 = tf.nn.relu(tf.add(tf.nn.conv3d(hidden_conv2, w_conv3, strides=[1, 1, 1, 1,1], padding='SAME'),b_conv3))
            out_conv3 = tf.nn.dropout(out_conv3, keep_prob)
            out_conv3_shape = tf.shape(out_conv3)
            tf.summary.scalar('out_conv3_shape', out_conv3_shape[0])

            # after conv2 ,the output size is batch_sizex2x8x8x64([batch_size,in_deep,width,height,output_deep])
            # all feature map flatten to one dimension vector,this vector will be much long
            out_conv3 = tf.reshape(out_conv3,[-1,64*10*10*7])
            w_fc1 = tf.Variable(tf.random_normal([64*10*10*7,2048],stddev=0.001),name='w_fc1')
            b_fc1 = tf.Variable(tf.constant(0.01, shape=[2048]), dtype=tf.float32, name='b_fc1')
            out_fc1 = tf.nn.relu(tf.add(tf.matmul(out_conv3,w_fc1),b_fc1))
            out_fc1 = tf.nn.dropout(out_fc1,keep_prob)
            out_fc1_shape = tf.shape(out_fc1)
            tf.summary.scalar('out_fc1_shape', out_fc1_shape[0])

            w_fc2 = tf.Variable(tf.random_normal([2048, 256], stddev=0.001), name='w_fc2')
            b_fc2 = tf.Variable(tf.constant(0.01, shape=[256]), dtype=tf.float32, name='b_fc2')
            out_fc2 = tf.nn.relu(tf.add(tf.matmul(out_fc1, w_fc2), b_fc2))
            out_fc2 = tf.nn.dropout(out_fc2, keep_prob)

            w_fc3 = tf.Variable(tf.random_normal([256, 2], stddev=0.001), name='w_fc3')
            b_fc3 = tf.Variable(tf.constant(0.01, shape=[2]), dtype=tf.float32, name='b_fc3')
            out_fc3 = tf.nn.relu(tf.add(tf.matmul(out_fc2, w_fc3), b_fc3))
            out_fc3 = tf.nn.dropout(out_fc3, keep_prob)

            w_sm = tf.Variable(tf.random_normal([2, 2], stddev=0.001), name='w_sm')
            b_sm = tf.constant(0.001, shape=[2])
            out_sm = tf.nn.softmax(tf.add(tf.matmul(out_fc3, w_sm), b_sm))

            return out_sm

    def inference(self,npy_path,test_path,model_index,train_flag=True):

        # some statistic index
        highest_acc = 0.0
        highest_iterator = 1

        all_filenames = get_all_filename(npy_path,self.cubic_shape[model_index][1])
        
        # how many time should one epoch should loop to feed all data
        times = int(len(all_filenames) / self.batch_size)
        if (len(all_filenames) % self.batch_size) != 0:
            times = times + 1

        # keep_prob used for dropout
        keep_prob = tf.placeholder(tf.float32)
        # take placeholder as input
        x = tf.placeholder(tf.float32, [None, self.cubic_shape[model_index][0], self.cubic_shape[model_index][1], self.cubic_shape[model_index][2]])
        x_image = tf.reshape(x, [-1, self.cubic_shape[model_index][0], self.cubic_shape[model_index][1], self.cubic_shape[model_index][2], 1])
        net_out = self.archi_1(x_image,keep_prob)

        saver = tf.train.Saver()  # default to save all variable,save mode or restore from path

        if train_flag:
            # softmax layer
            real_label = tf.placeholder(tf.float32, [None, 2])
            cross_entropy = tf.reduce_sum(tf.nn.softmax_cross_entropy_with_logits(logits=net_out, labels=real_label))
            net_loss = tf.reduce_mean(cross_entropy)

            train_step = tf.train.MomentumOptimizer(self.learning_rate, 0.9).minimize(net_loss)

            correct_prediction = tf.equal(tf.argmax(net_out, 1), tf.argmax(real_label, 1))
            accruacy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

            merged = tf.summary.merge_all()

            with tf.Session() as sess:
                sess.run(tf.global_variables_initializer())
                train_writer = tf.summary.FileWriter('./tensorboard/', sess.graph)
                # loop epoches
                for i in range(self.epoch):
                    epoch_start =time.time()
                    #  the data will be shuffled by every epoch
                    random.shuffle(all_filenames)
                    for t in range(times):
                        batch_files = all_filenames[t*self.batch_size:(t+1)*self.batch_size]
                        batch_data, batch_label = get_train_batch(batch_files)
                        feed_dict = {x: batch_data, real_label: batch_label,
                                     keep_prob: self.keep_prob}
                        _,summary = sess.run([train_step, merged],feed_dict =feed_dict)
                        train_writer.add_summary(summary, i)
                        saver.save(sess, './ckpt/archi-1', global_step=i + 1)

                    epoch_end = time.time()
                    test_batch,test_label = get_test_batch(test_path)
                    test_dict = {x: test_batch, real_label: test_label, keep_prob:self.keep_prob}
                    acc_test,loss = sess.run([accruacy,net_loss],feed_dict=test_dict)
                    print('accuracy  is %f' % acc_test)
                    print("loss is ", loss)
                    print(" epoch %d time consumed %f seconds"%(i,(epoch_end-epoch_start)))
                    if (acc_test > highest_acc):
                        highest_acc = acc_test
                        highest_iterator = i

            print("training finshed..highest accuracy is %f,the iterator is %d " % (highest_acc, highest_iterator))