# import csv
import sys
import numpy as np
import time
from sklearn.externals import joblib
# import argparse
import os
from sklearn.ensemble import ExtraTreesClassifier
import cv2

from sklearn.model_selection import train_test_split
import tensorflow as tf
import random
from datetime import datetime
TIMESTAMP = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())

# csv.field_size_limit(sys.maxsize)


# def load_data(path_name, n_rows, n_cols):
#
#     count = 0
#     X = np.zeros((n_rows, n_cols), dtype=np.float)
#     y = np.zeros(n_rows, dtype=np.uint8)
#     pai_id = np.zeros(n_rows, dtype=np.uint8)
#
#     with open(path_name, 'r') as csvfile:
#         reader = csv.reader(csvfile, delimiter=',')
#         for row in reader:
#             x = np.fromstring(row[0], dtype=np.float, sep=' ')
#             X[count] = x
#             y[count] = int(row[1])
#             pai_id[count] = int(row[2])
#             count += 1
#             if count >= N_ROWS:
#                 break
#
#     return X, y, pai_id


if __name__ == "__main__":
    imgs = []
    labs = []
    size = 64
    faces_number = 2
    # ap = argparse.ArgumentParser()
    # ap.add_argument("-n", "--name", required=True, help="name of trained model")
    # ap.add_argument("-i", "--csv_file_to_train", required=True, help="input csv file to train the classifier")
    # ap.add_argument("-f", "--folder_to_store_classifier", required=True, help="directory to store the trained classifier")
    # ap.add_argument("-r", "--rows", required=True, help="total number of rows in the csv file")
    # ap.add_argument("-c", "--cols", required=True, help="dimension of the feature vector")
    # args = vars(ap.parse_args())

    # N_ROWS = int(args["rows"])  # Total number of rows in csv file
    # N_COLS = int(args["cols"])  # Dimension of the Feature Vector
    # training_file = args["csv_file_to_train"]  # Input csv file with the training data
    # output_folder = args["folder_to_store_classifier"]  # Directory to store the trained classifier
    true_path = "../dataset_processed/ClientFace"
    fake_path = "../dataset_processed/ImposterFace"
    for file_dir_name in os.listdir(true_path):
        next_files = true_path + '/' + file_dir_name #左斜杠linux与windows都兼容
        for filename in os.listdir(next_files):
            if filename.endswith('.bmp') or filename.endswith('.BMP') or filename.endswith('.PNG') or filename.endswith('.png') or filename.endswith('.JPG') or filename.endswith('.jpg'):
            #三种图像格式,bmp/png/jpg
                filename = next_files + '/' + filename

                img = cv2.imread(filename)

                imgs.append(img)
                labs.append(true_path)
    for file_dir_name in os.listdir(fake_path):
        next_files = fake_path + '/' + file_dir_name #左斜杠linux与windows都兼容
        for filename in os.listdir(next_files):
            if filename.endswith('.bmp') or filename.endswith('.BMP') or filename.endswith('.PNG') or filename.endswith('.png') or filename.endswith('.JPG') or filename.endswith('.jpg'):
            #三种图像格式,bmp/png/jpg
                filename = next_files + '/' + filename

                img = cv2.imread(filename)

                imgs.append(img)
                labs.append(fake_path)
    imgs = np.array(imgs)
    labs = np.array([[0, 1] if lab == true_path else [1, 0] for lab in labs])

    # 随机划分测试集与训练集
    train_x, test_x, train_y, test_y = train_test_split(imgs, labs, test_size=0.05, random_state=random.randint(0, 100))
    # 参数：图片数据的总数，图片的高、宽、通道
    train_x = train_x.reshape(train_x.shape[0], size, size, 3)
    test_x = test_x.reshape(test_x.shape[0], size, size, 3)
    # 将数据转换成小于1的数
    train_x = train_x.astype('float32') / 255.0
    test_x = test_x.astype('float32') / 255.0

    print('train size:%s, test size:%s' % (len(train_x), len(test_x)))
    # 图片块，每次取100张图片
    batch_size = 100
    num_batch = len(train_x) // batch_size

    x = tf.placeholder(tf.float32, [None, size, size, 3])
    y_ = tf.placeholder(tf.float32, [None, faces_number])
    # 这里也是根据识别人脸个数
    keep_prob_5 = tf.placeholder(tf.float32)
    keep_prob_75 = tf.placeholder(tf.float32)

    v1 = tf.get_variable("v1", shape=[3], initializer=tf.zeros_initializer)  # 测试参数传递
    inc_v1 = v1.assign(v1 + 1)


    def weightVariable(shape):
        init = tf.random_normal(shape, stddev=0.01)
        return tf.Variable(init)


    def biasVariable(shape):
        init = tf.random_normal(shape)
        return tf.Variable(init)


    def conv2d(x, W):
        return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')


    def maxPool(x):
        return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')


    def dropout(x, keep):
        return tf.nn.dropout(x, keep)


    def cnnLayer():
        # 第一层
        # W1 = weightVariable([3,3,3,32]) # 卷积核大小(3,3)， 输入通道(3)， 输出通道(32)
        W1 = tf.get_variable("W1", shape=[3, 3, 3, 32], initializer=tf.random_normal_initializer(stddev=0.01))
        # b1 = biasVariable([32])
        b1 = tf.get_variable("b1", shape=[32], initializer=tf.random_normal_initializer)
        # 卷积
        conv1 = tf.nn.relu(conv2d(x, W1) + b1)
        # 池化
        pool1 = maxPool(conv1)
        # 减少过拟合，随机让某些权重不更新
        drop1 = dropout(pool1, keep_prob_5)

        # 第二层
        # W2 = weightVariable([3,3,32,64])
        W2 = tf.get_variable("W2", shape=[3, 3, 32, 64], initializer=tf.random_normal_initializer(stddev=0.01))
        # b2 = biasVariable([64])
        b2 = tf.get_variable("b2", shape=[64], initializer=tf.random_normal_initializer)
        conv2 = tf.nn.relu(conv2d(drop1, W2) + b2)
        pool2 = maxPool(conv2)
        drop2 = dropout(pool2, keep_prob_5)

        # 第三层
        # W3 = weightVariable([3,3,64,64])
        W3 = tf.get_variable("W3", shape=[3, 3, 64, 64], initializer=tf.random_normal_initializer(stddev=0.01))
        # b3 = biasVariable([64])
        b3 = tf.get_variable("b3", shape=[64], initializer=tf.random_normal_initializer)
        conv3 = tf.nn.relu(conv2d(drop2, W3) + b3)
        pool3 = maxPool(conv3)
        drop3 = dropout(pool3, keep_prob_5)

        # 全连接层
        # Wf = weightVariable([8*8*64, 512])
        Wf = tf.get_variable("Wf", shape=[8 * 8 * 64, 512], initializer=tf.random_normal_initializer(stddev=0.01))
        # bf = biasVariable([512])
        bf = tf.get_variable("bf", shape=[512], initializer=tf.random_normal_initializer)
        drop3_flat = tf.reshape(drop3, [-1, 8 * 8 * 64])
        dense = tf.nn.relu(tf.matmul(drop3_flat, Wf) + bf)
        dropf = dropout(dense, keep_prob_75)

        # 输出层
        # Wout = weightVariable([512, 2])
        Wout = tf.get_variable("Wout", shape=[512, faces_number], initializer=tf.random_normal_initializer(stddev=0.01))
        # bout = biasVariable([2])
        bout = tf.get_variable("bout", shape=[faces_number], initializer=tf.random_normal_initializer)
        # out = tf.matmul(dropf, Wout) + bout
        out = tf.add(tf.matmul(dropf, Wout), bout)
        return out


    # 输出层个数根据标签决定
    def cnnTrain():
        out = cnnLayer()

        cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=out, labels=y_))
        # learning rate = 0.01
        train_step = tf.train.AdamOptimizer(0.001).minimize(cross_entropy)
        # 比较标签是否相等，再求的所有数的平均值，tf.cast(强制转换类型)
        # accuracy = tf.reduce_mean(tf.cast(tf.equal(out, y_), tf.float32))
        accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(out, 1), tf.argmax(y_, 1)), tf.float32))
        # 将loss与accuracy保存以供tensorboard使用
        tf.summary.scalar('loss', cross_entropy)
        tf.summary.scalar('accuracy', accuracy)
        merged_summary_op = tf.summary.merge_all()
        # 数据保存器的初始化
        saver = tf.train.Saver()

        with tf.Session() as sess:

            sess.run(tf.global_variables_initializer())

            inc_v1.op.run()  # 测试参数传递

            train_writer = tf.summary.FileWriter('./tmp/train' + TIMESTAMP, graph=tf.get_default_graph())
            test_writer = tf.summary.FileWriter("./tmp/test" + TIMESTAMP, graph=tf.get_default_graph())

            for n in range(5000):
                for i in range(num_batch):
                    batch_x = train_x[i * batch_size: (i + 1) * batch_size]
                    batch_y = train_y[i * batch_size: (i + 1) * batch_size]
                    # 开始训练数据，同时训练三个变量，返回三个数据
                    _, loss, train_result = sess.run([train_step, cross_entropy, merged_summary_op],
                                                     feed_dict={x: batch_x, y_: batch_y, keep_prob_5: 0.5,
                                                                keep_prob_75: 0.75})
                    train_writer.add_summary(train_result, n * num_batch + i)
                    # 打印损失
                    print(n * num_batch + i, "loss", loss)

                    # if (n*num_batch+i) % 100 == 0:
                    #     # 获取测试数据的准确率
                    #     acc = accuracy.eval({x:test_x, y_:test_y, keep_prob_5:1.0, keep_prob_75:1.0})
                    #     print(n*num_batch+i,"accuracy", acc)
                    #     # 准确率大于0.98时保存并退出
                    #     if acc > 0.98 and n > 2:
                    #         saver.save(sess, './train_faces.model', global_step=n*num_batch+i)
                    #         sys.exit(0)

                    # n1 = n
                    # num_batch1 = num_batch
                    # i1 = i
                    acc, test_result = sess.run([accuracy, merged_summary_op],
                                                feed_dict={x: test_x, y_: test_y, keep_prob_5: 1.0, keep_prob_75: 1.0})
                    # acc = accuracy.eval({x: test_x, y_: test_y, keep_prob_5: 1.0, keep_prob_75: 1.0})
                    test_writer.add_summary(test_result, n * num_batch + i)
                    # saver.save(sess, './train_faces.model', global_step=n1 * num_batch1 + i1)
                    print(n * num_batch + i, 'accuracy', acc)
                    if acc > 0.99:
                        saver.save(sess, './tmp/model.ckpt')
                        sys.exit(0)

            saver.save(sess, './tmp/model.ckpt')


    cnnTrain()

    # X_train = imgs #四维
    # y_train = labs
    # # LOAD DATA
    # X_train, y_train, PAI_ID = load_data(training_file, N_ROWS, N_COLS)

    # # Random Forest CLASSIFIER
    # print("Training data ETC")
    # clf = ExtraTreesClassifier(n_estimators=20, min_samples_leaf=10, n_jobs=8)
    # start_time = time.time()
    # clf.fit(X_train, y_train) #二维
    # elapsed_time = time.time() - start_time
    # print("ExtraTreesClassifier %.3f s" % elapsed_time)
    # # joblib.dump(clf, output_folder + "/" + args["name"] + "_extraTreesClassifier" + '.pkl')
    # # print("Generated:", output_folder + "/" + args["name"] + "_extraTreesClassifier" + '.pkl')
    # joblib.dump(clf, "/root/zlf_projects/PAD/train_result/extraTreesClassifier.pkl")
    # print("Generated: extraTreesClassifier.pkl")
