import numpy as np
from scipy.misc import imread,imresize
import glob
import matplotlib.pyplot as plt
import os
import tensorflow as tf
from tensorflow.python.framework import ops
ops.reset_default_graph()
from tensorflow.python.framework.graph_util import convert_variables_to_constants


def read_image
 
(path):  # 讀取要訓練的圖片
    imgs = []
    labels = []
    # 資料夾下的檔名
    cate = [path + x for x in os.listdir(path)]  # 獲取絕對路徑
    for idx, folder in (enumerate(cate)):
        print (idx, folder)
        for im in glob.glob(folder + '/*.png'):  # glob查詢檔案路徑
            # print ('reading the images:%s'%(im))
 

            img = imread(im)

            img = imresize(img, (w, h))
            imgs.append(img)
            labels.append(idx)
    return np.asarray(imgs), np.asarray(labels, np.int32)

#網路引數設定
def cnn_param_set():
    conv1_weight = tf.Variable(tf.truncated_normal([3, 3, num_channels, conv1_features],
                                                   stddev=0.1
 
, dtype=tf.float32), name='conv1_weight')
    conv1_bias = tf.Variable(tf.zeros([conv1_features], dtype=tf.float32), name='conv1_bias')
    conv2_weight = tf.Variable(tf.truncated_normal([3, 3, conv1_features, conv2_features],
                                                   stddev=0.1, dtype=tf.float32), name='conv2_weight')
    conv2_bias = tf.Variable(tf.zeros([conv2_features], dtype=tf.float32), name='conv2_bias')

    # fully connected variables 全連線層
    resulting_width = image_width // (max_pool_size1 * max_pool_size2)
    print(resulting_width)
    resulting_height = image_height // (max_pool_size1 * max_pool_size2)
    print (image_height)
    print (resulting_height)
    full1_input_size = resulting_width * resulting_height * conv2_features
    print (full1_input_size)
    full1_weight = tf.Variable(tf.truncated_normal([full1_input_size, fully_connected_size1],
                                                   stddev=0.1, dtype=tf.float32), name='full1_weight')
    full1_bias = tf.Variable(tf.truncated_normal([fully_connected_size1], stddev=0.1, dtype=tf.float32)
                             , name='full1_bias')

    full2_weight = tf.Variable(tf.truncated_normal([fully_connected_size1, target_size],
                                                   stddev=0.1, dtype=tf.float32), name='full2_weight')
    full2_bias = tf.Variable(tf.truncated_normal([target_size], stddev=0.1, dtype=tf.float32), name='full2_bias')


def my_conv_net(input_data):
    # First Conv-ReLU-MaxPool Layer
    conv1 = tf.nn.conv2d(input_data, conv1_weight, strides=[1, 1, 1, 1], padding='SAME', name='conv1')

    relu1 = tf.nn.relu(tf.nn.bias_add(conv1, conv1_bias))
    max_pool1 = tf.nn.max_pool(relu1, ksize=[1, max_pool_size1, max_pool_size1, 1],
                               strides=[1, max_pool_size1, max_pool_size1, 1], padding='SAME', name='max_pool1')
    print (max_pool1.shape)
    # Second Conv-ReLU-MaxPool Layer
    conv2 = tf.nn.conv2d(max_pool1, conv2_weight, strides=[1, 1, 1, 1], padding='SAME', name='conv2')
    relu2 = tf.nn.relu(tf.nn.bias_add(conv2, conv2_bias))
    max_pool2 = tf.nn.max_pool(relu2, ksize=[1, max_pool_size2, max_pool_size2, 1],
                               strides=[1, max_pool_size2, max_pool_size2, 1], padding='SAME', name='max_pool2')
    print (max_pool2.shape)
    # third Conv-ReLU-MaxPool Layer

    # Transform Output into a 1xN layer for next fully connected layer
    final_conv_shape = max_pool2.get_shape().as_list()
    print(final_conv_shape)
    final_shape = final_conv_shape[1] * final_conv_shape[2] * final_conv_shape[3]  # [100, 7, 7, 50]
    flat_output = tf.reshape(max_pool2, [-1, final_shape])  # 100*2450

    # First Fully Connected Layer
    fully_connected1 = tf.nn.relu(tf.add(tf.matmul(flat_output, full1_weight), full1_bias),
                                  name='fully_connected1')  # 矩陣相乘
    print('full1: ' + str(fully_connected1))
    # Second Fully Connected Layer
    final_model_output = tf.add(tf.matmul(fully_connected1, full2_weight), full2_bias, name='final_model_output')
    print (final_model_output)
    return (final_model_output)
#準確度計算
def get_accuracy(logits,targets):
    batch_predictions = np.argmax(logits,axis=1)
    num_correct = np.sum(np.equal(batch_predictions,targets))
    return(100.* num_correct/batch_predictions.shape[0])

def train():
    # training
    import datetime
    begin = datetime.datetime.now()

    train_loss = []
    train_acc = []
    test_acc = []
    for i in range(generations):
        rand_index = np.random.choice(len(x_train), size=batch_size)  # 隨機選擇size個數
        print
        rand_x = x_train[rand_index]  # (100, 28, 28)
        rand_y = y_train[rand_index]
        train_dict = {x_input: rand_x, y_target: rand_y}

        sess.run(train_step, feed_dict=train_dict)
        temp_train_loss, temp_train_preds = sess.run([loss, prediction], feed_dict=train_dict)
        loss1_temp = sess.run([loss1], feed_dict=train_dict)
        temp_train_acc = get_accuracy(temp_train_preds, rand_y)
        if (i) % eval_every == 0:
            eval_index = np.random.choice(len(x_val), size=evaluation_size)
            eval_x = x_val[eval_index]
            eval_y = y_val[eval_index]
            test_dict = {eval_input: eval_x, eval_target: eval_y}
            temp_test_loss = sess.run(loss_test, feed_dict=test_dict)
            test_preds = sess.run(test_prediction, feed_dict=test_dict)
            temp_test_acc = get_accuracy(test_preds, eval_y)

            # Record and print results
            train_loss.append(temp_train_loss)
            train_acc.append(temp_train_acc)
            test_acc.append(temp_test_acc)
            acc_and_loss = [(i + 1), temp_train_loss, temp_train_acc, temp_test_acc]
            acc_and_loss = [np.round(x, 2) for x in acc_and_loss]
            print('Generation # {}. Train Loss: {:.2f}. Train Acc (Test Acc): {:.2f} ({:.2f})'.format(*acc_and_loss))
            print('Generation # test Loss: ', temp_test_loss)
    end = datetime.datetime.now()

    print(end - begin)

def loss_accuray_plt():
    # Matlotlib code to plot the loss and accuracies
    eval_indices = range(0, generations, eval_every)
    # Plot loss over time
    plt.plot(eval_indices, train_loss, 'k-')
    plt.title('Softmax Loss per Generation')
    plt.xlabel('Generation')
    plt.ylabel('Softmax Loss')
    plt.show()

    # Plot train and test accuracy
    plt.plot(eval_indices, train_acc, 'k-', label='Train Set Accuracy')
    plt.plot(eval_indices, test_acc, 'r--', label='Test Set Accuracy')
    plt.title('Train and Test Accuracy')
    plt.xlabel('Generation')
    plt.ylabel('Accuracy')
    plt.legend(loc='lower right')
    plt.show()

if __name__ == "__main__":
    path = u'F:/PycharmWorkspace/ipython/tensorflow/station_data/'
    # 將所有的圖片resize成32*64
    w = 32
    h = 64
    c = 4
    #讀取圖片
    data, label = read_image(path)
    #將圖片順序打亂
    num_example = data.shape[0]
    print (num_example)
    print (label)
    print ((set(label)))
    arr = np.arange(num_example)
    np.random.shuffle(arr)
    data = data[arr]
    label = label[arr]

    # 將所有資料分為訓練集和驗證集
    ratio = 0.8
    s = np.int(num_example * ratio)
    x_train = data[:s]
    y_train = label[:s]
    x_val = data[s:]
    y_val = label[s:]

    # 設定模型引數
    generations = 1000
    batch_size = 200
    evaluation_size = 500
    image_width = w
    image_height = h
    print(image_height)
    num_channels = 4  # greyscale = 1 channel 灰度影象通道是1
    target_size = len(set(label))
    print (target_size)
    learning_rate = 0.001
    eval_every = 5
    conv1_features = 32  # conv1卷積核數
    conv2_features = 64  # conv2卷積核數
    max_pool_size1 = 2  # NxN window for 1st max pool layer
    max_pool_size2 = 2  # NxN window for 2nd max pool layer
    max_pool_size3 = 2  # NxN window for 2nd max pool layer
    fully_connected_size1 = 512

    #cnn佔位符
    x_input_shape = (None, w, h, c)
    x_input = tf.placeholder(tf.float32, shape=x_input_shape, name='x_input')
    y_target = tf.placeholder(tf.int32, shape=[None, ], name='y_target')
    print(x_input)
    print(y_target)
    eval_input_shape = (None, w, h, c)
    eval_input = tf.placeholder(tf.float32, shape=eval_input_shape, name='eval_input')
    eval_target = tf.placeholder(tf.int32, shape=[None, ], name='eval_target')

    # 網路引數設定
    cnn_param_set()
    #網路模型：
    model_output = my_conv_net(x_input)
    test_model_output = my_conv_net(eval_input)

    #損失函式
    loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=model_output, labels=y_target),name='loss')
    loss_test = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=test_model_output, labels=eval_target))
    #預測
    prediction = tf.nn.softmax(model_output,name="softmax")
    print (prediction)
    prediction_labels = tf.argmax(prediction, axis=1, name='output')
    print(prediction_labels)
    test_prediction = tf.nn.softmax(test_model_output)

    # Create an optimizer
    my_optimizer = tf.train.AdamOptimizer(learning_rate)
    train_step = my_optimizer.minimize(loss)

    # Initialize Variables
    init = tf.global_variables_initializer()
    sess.run(init)
    #訓練
    train()
    #畫出損失和預測的變化
    loss_accuray_plt()

    #儲存訓練的模型
    graph = convert_variables_to_constants(sess, sess.graph_def, ["softmax"])  # out為儲存網路的最後輸出節點名稱
    tf.train.write_graph(graph, '.', 'graph/graph_train.pb', as_text=False)
    sess.close()