from __future__ import print_function
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cross_validation import train_test_split
import tensorflow as tf
from sklearn import preprocessing
import csv

DATA_DIR = '../input'

test_df = pd.read_csv(DATA_DIR + '/test.csv', header=0)
test_image = test_df.iloc[0]
test_image = np.reshape(test_image, [28, 28])

train_df = pd.read_csv(DATA_DIR + '/train.csv', header=0)
train_image = train_df.iloc[0][1:]
train_image = np.reshape(train_image, [28, 28])

train_labels = train_df.loc[:, 'label']
train_images = train_df.loc[:, 'pixel0':]

train_labels = np.reshape(train_labels, [-1, 1])
num_train_labels = train_labels.shape[0]
index_offset = np.arange(num_train_labels) * 10
train_labels_one_hot = np.zeros((num_train_labels, 10))
train_labels_one_hot.flat[index_offset + train_labels.ravel()] = 1

train_labels = np.array(train_labels_one_hot, dtype=np.float32)
train_images = np.reshape(train_images, [-1, 28 * 28])
train_images = np.array(train_images, dtype=np.float32)
train_images = np.multiply(train_images, 1.0/255)

test_images = test_df.loc[:, 'pixel0':]
test_images = np.reshape(test_images, [-1, 28 * 28])
test_images = np.array(test_images, dtype=np.float32)
test_images = np.multiply(test_images, 1.0/255)
# n_samples_test = int(n_samples*0.3)
# indices = np.random.choice(n_samples, n_samples_test)
X_train, X_test, y_train, y_test = train_test_split(train_images, train_labels, test_size=0.3, random_state=42)
# print X_train.shape, y_train.shape
# print X_test.shape, y_test.shape
# input data
n_samples = X_train.shape[0]
n_samples_batch = 128
with tf.name_scope('data'):
    X_data = tf.placeholder(tf.float32, [None, 28*28], name='X_data')
    y_data = tf.placeholder(tf.float32, [None, 10], name='y_data')

# conv1 5x5x1x32 filter size 5x5 input size 28x28x1 output size 28x28x32
with tf.variable_scope('conv1') as scope:
    X_data_images = tf.reshape(X_data, [-1, 28, 28, 1])
    kernel1 = tf.get_variable('kernel1', [5, 5, 1, 32], initializer=tf.truncated_normal_initializer())
    bias1 = tf.get_variable('bias1', [32], initializer=tf.random_normal_initializer())
    conv1 = tf.nn.conv2d(X_data_images, kernel1, strides=[1, 1, 1, 1], padding='SAME')
    conv1 = tf.nn.relu(conv1+bias1, name=scope.name)

# pool1 max pool input size 28x28x32 output size 14x14x32
with tf.variable_scope('pool1') as scope:
    pool1 = tf.nn.max_pool(conv1, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME')

# conv2 5x5x1x32 filter size 5x5 input size 14x14x32 output size 14x14x64
with tf.variable_scope('conv2') as scope:
    kernel2 = tf.get_variable('kernel2', [5, 5, 32, 64], initializer=tf.truncated_normal_initializer())
    bias2 = tf.get_variable('bias2', [64], initializer=tf.random_normal_initializer())
    conv2 = tf.nn.conv2d(pool1, kernel2, strides=[1, 1, 1, 1], padding='SAME')
    conv2 = tf.nn.relu(conv2+bias2, name=scope.name)

# pool2 max pool input size 14x14x64 output size 7x7x64
with tf.variable_scope('pool2') as scope:
    pool2 = tf.nn.max_pool(conv2, [1, 2, 2, 1], [1, 2, 2, 1], padding='SAME')

# fc fully connected input size 7x7x64 output size 1024
with tf.variable_scope('fc') as scope:
    input_features = 7*7*64
    pool2 = tf.reshape(pool2, [-1, input_features])
    W1 = tf.get_variable('weights1', [input_features, 1024], initializer=tf.truncated_normal_initializer())
    b1 = tf.get_variable('bias1', [1024], initializer=tf.constant_initializer(0.0))
    fc = tf.nn.relu(tf.matmul(pool2, W1) + b1, name=scope.name)

# out input size 1024 output size 10
with tf.variable_scope('out') as scope:
    W2 = tf.get_variable('weights2', [1024, 10], initializer=tf.truncated_normal_initializer())
    b2 = tf.get_variable('bias2', [10], initializer=tf.random_normal_initializer())
    logits = tf.matmul(fc, W2) + b2

with tf.name_scope('loss'):
    entropy = tf.nn.softmax_cross_entropy_with_logits(labels=y_data, logits=logits)
    loss = tf.reduce_mean(entropy, name='loss')

LEARNING_RATE = 0.001
optimizer = tf.train.AdamOptimizer(learning_rate=LEARNING_RATE).minimize(loss)

pred_labels = tf.argmax(logits, 1)
correct_pred = tf.equal(tf.argmax(logits, 1), tf.argmax(y_data, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    step = 1
    for _ in range(500):
        step += 1
        indices = np.random.choice(n_samples, n_samples_batch)
        X_train_batch, y_train_batch = X_train[indices], y_train[indices]
        sess.run([optimizer], feed_dict={X_data: X_train_batch, y_data: y_train_batch})
        if step % 100 == 0:
            fit_acc = sess.run([accuracy], feed_dict={X_data: X_train_batch, y_data: y_train_batch})
            print('fit_acc:', fit_acc)
        # if step % 100 == 0:
        #     valid_acc = sess.run([accuracy], feed_dict={X_data: X_test, y_data: y_test})
        #     print('valid_acc:', valid_acc)
    pred_labels_lists = sess.run([pred_labels], feed_dict={X_data: test_images})
    # result_csv = open('result.txt', 'wb')
    # result_csv.writelines(["%s\n" % pred_labels_list for pred_labels_list in pred_labels_lists[0]])
    # result_csv.close()
    # print 'pred_labels_lists:', pred_labels_lists
# print train_labels[0]
# print train_images[0]

# plt.subplot(121)
# plt.imshow(test_image, cmap='gray')
# plt.subplot(122)
# plt.imshow(train_image, cmap='gray')
# plt.show()