# coding:utf8
import tensorflow as tf
from PIL import Image,ImageFilter
import matplotlib.pyplot as plt
import time
from Jump import Jump
import numpy as np

type2 = input("1=训练，2=玩游戏")

model_path = "./model/model.ckpt" #模型文件



    
sess = tf.InteractiveSession()

with tf.name_scope('input'):
    x = tf.placeholder("float", shape=[None, 224,224,1])
    y_ = tf.placeholder("float", shape=[None,4])

#读取二进制数据
def read_and_decode(filename):
    # 创建文件队列,不限读取的数量
    filename_queue = tf.train.string_input_producer([filename])
    # create a reader from file queue
    reader = tf.TFRecordReader()
    # reader从文件队列中读入一个序列化的样本
    _, serialized_example = reader.read(filename_queue)
    # get feature from serialized example
    # 解析符号化的样本
    features = tf.parse_single_example(
        serialized_example,
        features={
            'label': tf.FixedLenFeature([], tf.string),
            'img_raw': tf.FixedLenFeature([], tf.string)
        }
    )
    label = features['label']
    img = features['img_raw']    
    img = tf.decode_raw(img, tf.uint8)
    img = tf.reshape(img, [224, 224, 1])
    img = tf.cast(img, tf.float32) * (1. / 255) - 0.5
    
    label = tf.decode_raw(label, tf.uint8)
    label = tf.reshape(label, [4])
    label = tf.cast(label, tf.float32)
    return img, label


def print_activations(t):
    '''
    展示一个Tensor的name和shape
    :param t:
    :return:
    '''
    print(t.op.name, ' ', t.get_shape().as_list())



'''
定义了AlexNet的前五个卷积层(FC计算速度较快,这里不做考虑)
:param images: 输入图像Tensor
:return:  返回最后一层pool5和parameters
'''
# ====================================模型开始======================================
def weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(initial)
 
def bias_variable(shape):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial)
    

# 第1个卷积层
with tf.name_scope('conv1') as scope:
    kernel = tf.Variable(tf.truncated_normal([11, 11, 1, 64], dtype=tf.float32,
                                             stddev=1e-1), name='weights')
    conv = tf.nn.conv2d(x, kernel, [1, 4, 4, 1], padding='SAME')
    biases = tf.Variable(tf.constant(0.0, shape=[64], dtype=tf.float32),
                         trainable=True, name='biases')
    bias = tf.nn.bias_add(conv, biases)
    conv1 = tf.nn.relu(bias, name=scope)
    print_activations(conv1)
    # 添加LRN层和最大池化层
    lrn1 = tf.nn.lrn(conv1, 4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, name='lrn1')
    pool1 = tf.nn.max_pool(lrn1,
                           ksize=[1, 3, 3, 1],
                           strides=[1, 2, 2, 1],
                           padding='VALID',
                           name='pool1')
    print_activations(pool1)

# 设计第2个卷积层
with tf.name_scope('conv2') as scope:
    kernel = tf.Variable(tf.truncated_normal([5, 5, 64, 192], dtype=tf.float32,
                                             stddev=1e-1), name='weights')
    conv = tf.nn.conv2d(pool1, kernel, [1, 1, 1, 1], padding='SAME')
    biases = tf.Variable(tf.constant(0.0, shape=[192], dtype=tf.float32),
                         trainable=True, name='biases')
    bias = tf.nn.bias_add(conv, biases)
    conv2 = tf.nn.relu(bias, name=scope)
    print_activations(conv2)
    # 对第2个卷积层的输出进行处理，同样也是先做LRN处理再做最大化池处理。
    lrn2 = tf.nn.lrn(conv2, 4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, name='lrn2')
    pool2 = tf.nn.max_pool(lrn2,
                           ksize=[1, 3, 3, 1],
                           strides=[1, 2, 2, 1],
                           padding='VALID',
                           name='pool2')
    print_activations(pool2)

# 设计第3个卷积层
with tf.name_scope('conv3') as scope:
    kernel = tf.Variable(tf.truncated_normal([3, 3, 192, 384],
                                             dtype=tf.float32,
                                             stddev=1e-1), name='weights')
    conv = tf.nn.conv2d(pool2, kernel, [1, 1, 1, 1], padding='SAME')
    biases = tf.Variable(tf.constant(0.0, shape=[384], dtype=tf.float32),
                         trainable=True, name='biases')
    bias = tf.nn.bias_add(conv, biases)
    conv3 = tf.nn.relu(bias, name=scope)
    print_activations(conv3)

# 设计第4个卷积层
with tf.name_scope('conv4') as scope:
    kernel = tf.Variable(tf.truncated_normal([3, 3, 384, 256],
                                             dtype=tf.float32,
                                             stddev=1e-1), name='weights')
    conv = tf.nn.conv2d(conv3, kernel, [1, 1, 1, 1], padding='SAME')
    biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
                         trainable=True, name='biases')
    bias = tf.nn.bias_add(conv, biases)
    conv4 = tf.nn.relu(bias, name=scope)
    print_activations(conv4)

# 设计第5个卷积层
with tf.name_scope('conv5') as scope:
    kernel = tf.Variable(tf.truncated_normal([3, 3, 256, 256],
                                             dtype=tf.float32,
                                             stddev=1e-1), name='weights')
    conv = tf.nn.conv2d(conv4, kernel, [1, 1, 1, 1], padding='SAME')
    biases = tf.Variable(tf.constant(0.0, shape=[256], dtype=tf.float32),
                         trainable=True, name='biases')
    bias = tf.nn.bias_add(conv, biases)
    conv5 = tf.nn.relu(bias, name=scope)
    print_activations(conv5)
    # 最大池化层
    pool5 = tf.nn.max_pool(conv5,
           ksize=[1, 3, 3, 1],
           strides=[1, 2, 2, 1],
           padding='VALID',
           name='pool5')
    print_activations(pool5)
    
with tf.name_scope('full-connect-1'):
    W_fc1 = weight_variable([6 * 6 * 256, 4096])
    b_fc1 = bias_variable([4096])
    h_pool5_flat = tf.reshape(pool5, [-1, 6 * 6 * 256])
    h_fc1 = tf.nn.relu(tf.matmul(h_pool5_flat, W_fc1) + b_fc1)
    keep_prob = tf.placeholder(tf.float32)
    h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)
    
with tf.name_scope('full-connect-2'):
    W_fc2 = weight_variable([4096, 4])
    b_fc2 = bias_variable([4])
    y_conv=tf.matmul(h_fc1_drop, W_fc2) + b_fc2
# ====================================模型结束======================================

#损失函数
with tf.name_scope('loss'):
    cross_entropy = tf.reduce_mean(tf.square(y_- y_conv))

#梯度下降
train_step = tf.train.AdamOptimizer(0.0001).minimize(cross_entropy)
    
pred = y_conv



saver = tf.train.Saver()
sess = tf.InteractiveSession()
if type2 == '1':
	sess.run(tf.global_variables_initializer())#第一次运行不要屏蔽
else:
    saver.restore(sess, model_path)#恢复模型  第一次运行屏蔽


# 训练
def train():
    batch_size = 32
    
    img, label = read_and_decode("./data/train.tfrecords")
    testimg, testlabel = read_and_decode("./data/test.tfrecords")
    img_batch, label_batch = tf.train.shuffle_batch([img, label],
                                                batch_size=batch_size, capacity=2000,
                                                min_after_dequeue=1000)
                                                
    testimg_batch, testlabel_batch = tf.train.shuffle_batch([testimg, testlabel],
                                                batch_size=10, capacity=2000,
                                                min_after_dequeue=1000)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(coord=coord)
    for i in range(200000):
        train_batch_x, train_batch_y = sess.run([img_batch, label_batch])
        step,loss = sess.run([train_step,cross_entropy],feed_dict={x:train_batch_x, y_: train_batch_y,keep_prob: 0.6})
        print(i,loss)
        if i % 1000 == 0:
            save_path = saver.save(sess, model_path)#保存模型
            #打印测试数组
            result = sess.run(pred, feed_dict={x: train_batch_x, y_: train_batch_y,keep_prob: 0.6})
            print("=========正确数据============")
            print(train_batch_y)
            print("=========预测数据==============")
            print(result)
        if i % 100 == 0:
            #打印测试验证集
            test_batch_x, test_batch_y = sess.run([testimg_batch, testlabel_batch])
            result,loss = sess.run([pred,cross_entropy], feed_dict={x: test_batch_x,y_: test_batch_y,keep_prob: 1})
            print("=========正确数据============")
            print(test_batch_y)
            print("=========预测数据============== loss=",loss)
            print(result)
    
    coord.request_stop()
    coord.join(threads)

    
#预测坐标
def predict(img):
    img = img.crop((0, 500, 1080, 1500))
    img = img.convert('L')
    img = img.resize((224, 224))
    img = np.array(img)
    img = tf.constant(img)
    img = tf.reshape(img, [224, 224, 1])
    img = tf.cast(img, tf.float32) * (1. / 255) - 0.5
    img = tf.expand_dims(img, 0)
    img = sess.run(img)
    #print('img.shape',img.shape)
    #print('img',img)
    #plt.imshow(img.reshape(224,224),cmap='gray')
    #plt.show()
    result = sess.run([pred],feed_dict={x:img,keep_prob: 1})  
      
    return result
    


#自动跳
def jump():
    # uimodel = Uimodel()
    jump = Jump()
    while 1:
        time.sleep(2)
        jump.pull_screenshot()
        img = Image.open("./1.png")
        result = predict(img)
        print(result[0][0])
        jump.jump(jump.calc(result[0][0]))

    
if __name__ == '__main__':
    #predict(Image.open('./1.png'))
    if type2 == '1':
	    train()
    else:
        jump()