# Libraries
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D, MaxPooling2D
import time

# Import training data
dataset_train = pd.read_csv('../input/train.csv')
dataset_test = pd.read_csv('../input/test.csv')
x_train = dataset_train.iloc[:, 1:]
y_train = dataset_train.iloc[:,0]
x_test = dataset_test.iloc[:, :]

# Convert training data to numpy array
x_train_np = x_train.values
y_train_np = y_train.values
x_test_np = x_test.values

# Create Stritified K-Fold Splitter for splitting train and test sets
from sklearn.model_selection import StratifiedKFold
skf = StratifiedKFold(n_splits=3, shuffle=True, random_state=int(time.time()))

# Reshape 1D picture data into 2D
x_train_final = np.zeros((x_train_np.shape[0], 28, 28))
for i in range(0,x_train_np.shape[0]):
    x_train_final[i] = np.reshape(x_train_np[i],(28,28))

x_train_final = x_train_final.reshape(x_train_final.shape[0],28,28,1)

x_test_final = np.zeros((x_test_np.shape[0], 28, 28))
for i in range(0,x_test_np.shape[0]):
    x_test_final[i] = np.reshape(x_test_np[i],(28,28))
x_test_final = x_test_final.reshape(x_test_final.shape[0],28,28,1)

# Categorize output
from keras.utils import to_categorical
y_train = to_categorical(y_train_np)

# Split data into training and test sets using skf
x_train_indices = []
x_test_indices = []
for tr, te in skf.split(x_train_final, y_train_np):
    x_train_indices = tr
    x_test_indices = te
x_train_sampled = []
x_test_sampled = []
y_train_sampled = []
y_test_sampled = []
for i in x_train_indices:
    x_train_sampled.append(x_train_final[i])
    y_train_sampled.append(y_train[i])
for i in x_test_indices:
    x_test_sampled.append(x_train_final[i])
    y_test_sampled.append(y_train[i])

# Back to numpy array
x_train_sampled = np.array(x_train_sampled)
y_train_sampled = np.array(y_train_sampled)
x_test_sampled = np.array(x_test_sampled)
y_test_sampled = np.array(y_test_sampled)

# Scale data
from sklearn.preprocessing import MinMaxScaler
scaler = MinMaxScaler()
d2_train = x_train_sampled.reshape((x_train_sampled.shape[0],28*28))
d2_test = x_test_sampled.reshape((x_test_sampled.shape[0],28*28))
scaler.fit(d2_train)
d2_train = scaler.transform(d2_train)
d2_test = scaler.transform(d2_test)
x_train_sampled = d2_train.reshape((x_train_sampled.shape[0],28,28,1))
x_test_sampled = d2_test.reshape((x_test_sampled.shape[0],28,28,1))

# Construct Model
model = Sequential()
model.add(Conv2D(64, kernel_size=3, activation='relu', input_shape=(28,28,1)))
model.add(Conv2D(32, kernel_size=3, activation='relu'))
model.add(Flatten())
model.add(Dense(10, activation='softmax'))

# Compile Model
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(x_train_sampled, y_train_sampled, validation_data=(x_test_sampled, y_test_sampled), epochs=3)

# Input test data and print output to file
f=open("results.csv","w+")
f.write("ImageId,Label\n")
d2_test_final = x_test_final.reshape((x_test_final.shape[0],28*28))
d2_test_final = scaler.transform(d2_test_final)
x_test_final = d2_test_final.reshape((x_test_final.shape[0],28,28,1))
result = model.predict_classes(x_test_final)
for i in range(0,x_test_final.shape[0]):
    f.write("%s,%s\n" % (i+1,result[i]))
f.close()