#coding: utf8
"""
1. 读取training_index.txt，整理出训练数据，实现next_batch
2. 数据处理
3. model建立
4. 训练和预测
"""

import tensorflow as tf
import numpy as np
import math
import cv2
import sys
import os
from scipy import ndimage


#-------read the training_index.txt--------#


training_index = './training_index.txt'
total_image = 6297*3
global_step = 0
global_idx = np.arange(total_image)
images_g = np.array((total_image, 784))
labels_g = np.array((total_image,36))
max_epoch = 3500

def read_traing_list():
	train_image_dir = []
	train_label_dir = []
	reader = open(training_index)
	while 1:
		line = reader.readline()
		tmp = line.split(" ")
		# print tmp
		if not line:
			break
		train_image_dir.append(tmp[0])
		train_label_dir.append(tmp[1][0:-1])
	print train_image_dir[1:10]
	print train_label_dir[1:10]
	return train_image_dir, train_label_dir

#------run the preprocess-------#

class Dadaset(object):
	def __init__(self, image, label,dtype=tf.float32):
		self._image = image
		self.label = label
	@property
	def image(self):
	    return self._image
	@property
	def label(self):
	    return self._label
	

def read_image():
	train_image_dir, train_label_dir = read_traing_list()
	images = []
	labels = []
	index = 0
	global images_g
	global labels_g
	#print len(train_image_dir)
	for idx in range(len(train_image_dir)):
		image_tmp = cv2.imread(str(train_image_dir[idx]),cv2.IMREAD_GRAYSCALE)
		# print "jjgjg",image_tmp.shape
		label_tmp = int(str(train_label_dir[idx]))
		label_tmp2 = [0]*36
		label_tmp2[label_tmp] = 1
		images.append(image_tmp.flatten()/255.0)
		# print label_tmp
		# print "label:", label_tmp2
		# print "label len:", len(labels)
		labels.append(label_tmp2)
		index += 1
	
	images = np.array(images)
	labels = np.array(labels)
	images_g = images
	labels_g = labels
	return images, labels

# read_image()



def next_batch(batch_size):
	global global_step
	global labels_g
	global images_g
	print global_step
	global global_idx
	# print labels_g.shape
	start = batch_size * global_step
	end = batch_size * (global_step+1)
	global_step += 1
	if(end >= total_image):
		np.random.shuffle(global_idx)
		global_step = 0
		start = 0
		end = batch_size
		images_g[:] = images_g[global_idx]
		labels_g[:] = labels_g[global_idx]
		print images_g.shape

	return images_g[start:end], labels_g[start:end]

def generate_test_image(test_size):

	global images_g
	global labels_g
	global global_idx
	np.random.shuffle(global_idx)
	images_g[:] = images_g[global_idx]
	labels_g[:] = labels_g[global_idx]
	return images_g[0:test_size], labels_g[0:test_size]

#--------run CNN model-----#

def CNNmodel():
	x = tf.placeholder(tf.float32,[None,28*28])
	y_ = tf.placeholder(tf.float32,[None,36])
	def weight_variable(shape):
		init = tf.truncated_normal(shape,stddev = 0.1)
		return tf.Variable(init)
	def bias_variable(shape):
		init = tf.constant(0.1, shape = shape)
		return tf.Variable(init)

	## 声明卷积操作和pool操作
	## 在这里声明的卷积操作是步长为1，padding为0的vanilla版本
	## pool操作是2X2的max pool
	def conv2d(x,W):
		# strides:[batch, in_height, in_width, in_channels]
		return tf.nn.conv2d(x,W,strides = [1,1,1,1],padding = 'SAME')

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

	## model构建过程
	# 第一层是[一个卷积接一个max pooling], 卷积层的patch_size 是 5X5的，输入的通道数目是1(因为是灰度图)，输出是32个feature maps
	# [5,5,1,32]: patch_size是5x5，输入通道数目是1，输出通道的数目是32(此处的32是根据网络定义得来，非计算得到)
	x_image = tf.reshape(x,[-1,28,28,1]) #把mnist变成需要的格式

	W_conv1 = weight_variable([5,5,1,32])
	b_conv1 = bias_variable([32])

	#做相应的操作,conv, relu, maxpool
	h_conv1 = tf.nn.relu(conv2d(x_image,W_conv1)+b_conv1)
	h_pool1 = maxpool2d(h_conv1)

	# 第二层[一个卷积加一个maxpool]
	W_conv2 = weight_variable([5,5,32,64])
	b_conv2 = bias_variable([64])

	h_conv2 = tf.nn.relu(conv2d(h_pool1,W_conv2)+b_conv2)
	h_pool2 = maxpool2d(h_conv2)

	# 全连接层，共有1024个神经元，此时图片进行了两次2x2的maxpool，每次的步长是2，此时图片已经变为了7X7

	W_fc1 = weight_variable([7*7*64,1024])
	b_fc1 = bias_variable([1024])

	h_pool2_flat = tf.reshape(h_pool2,[-1,7*7*64])
	h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat,W_fc1)+b_fc1)

	# 在训练时加入dropout，在测试的时候记得要关闭哦。。。
	# keep_prob 表示的是保留参数的可能性，当等于1.0时表示不进行dropout
	keep_prob = tf.placeholder("float") #要输入的值

	h_fc1_drop = tf.nn.dropout(h_fc1,keep_prob)

	#添加 softmax层

	W_fc2 = weight_variable([1024,36])
	b_fc2 = bias_variable([36])

	y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop,W_fc2)+b_fc2)

	# loss 

	cross_entropy = -tf.reduce_sum(y_*tf.log(y_conv))

	train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

	correct_prediction = tf.equal(tf.argmax(y_conv,1),tf.argmax(y_,1))

	accuracy = tf.reduce_mean(tf.cast(correct_prediction,"float"))

	predict = [tf.reduce_max(y_conv),tf.argmax(y_conv,1)[0]]

	saver = tf.train.Saver()
	checkpoint_dir = "./tmp/model.cpkt"

	sess = tf.InteractiveSession()

	sess.run(tf.initialize_all_variables())

	for idx in range(max_epoch):
		imagess, labelss = next_batch(100)
		if idx%10 == 0:
			acc = sess.run(accuracy,feed_dict={x:imagess,y_: labelss,keep_prob:1.0})
			print "At epoch %d, acc is %lf" % (idx,acc)
		sess.run(train_step,feed_dict={x:imagess,y_: labelss,keep_prob:0.5})
	saver.save(sess,checkpoint_dir)
	test_image,test_label = generate_test_image(900)
	print "Total test acc: ", sess.run(accuracy,feed_dict={x:test_image,y_:test_label,keep_prob:1.0})

if __name__ == '__main__':
	read_image()
	CNNmodel()


