# python keras_mnist.py --output ./output/keras_mnist.png

# import the necessary packages
import sklearn
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from keras.models import Sequential
from keras.layers.core import Dense
from keras.optimizers import SGD
from sklearn import datasets
import matplotlib.pyplot as plt
import numpy as np
import argparse

def keras_mnist(output):
	# grab the MNISY dataset
	print("[INFO] loading MNIST (full) dataset...")
	datasets = sklearn.datasets.fetch_mldata("MNIST Original",data_home='C:/kuaipan/cv-data/MNIST/')

	# scale the raw pixel intensities to the range [0, 1.0] then
	# construct the training and testing splits
	data = datasets.data.astype("float") / 255.0
	(trainX, testX, trainY, testY) = train_test_split(data, datasets.target, test_size=0.25)
	print(trainY)
	print(testY)
	# covert the labels from integers to vectors
	lb = LabelBinarizer()
	trainY = lb.fit_transform(trainY)
	testY = lb.transform(testY)

	# define the 784-256-10 architecture using Keras
	model = Sequential()
	model.add(Dense(256, input_shape=(784,), activation="sigmoid"))
	model.add(Dense(128, activation="sigmoid"))
	model.add(Dense(10, activation="softmax"))

	# train the model using SGD
	print("[INFO] training network...")
	sgd = SGD(0.01)
	model.compile(loss="categorical_crossentropy", optimizer=sgd, metrics=["accuracy"])
	H = model.fit(trainX, trainY, validation_data=(testX, testY), epochs=100, batch_size=128)

	# evaluate the network
	print("[INFO] evaluating network...")
	predictions = model.predict(testX, batch_size=128)
	print(classification_report(testY.argmax(axis=1), predictions.argmax(axis=1), target_names=[str(x) for x in lb.classes_]))


	# plot the training loss and accuracy
	plt.style.use("ggplot")
	plt.figure()
	plt.plot(np.arange(0, 100), H.history["loss"], label="train_loss")
	plt.plot(np.arange(0, 100), H.history["val_loss"], label="val_loss")
	plt.plot(np.arange(0, 100), H.history["accuracy"], label="train_acc")
	plt.plot(np.arange(0, 100), H.history["val_accuracy"], label="val_acc")
	plt.title("Training Loss and Accuracy")
	plt.xlabel("Epoch #")
	plt.ylabel("Loss/Accuracy")
	plt.legend()
	plt.savefig(output)

	print("Done.")

if __name__ == '__main__':
	'''
	# construct the argument parse and parse the arguments
	ap = argparse.ArgumentParser()
	ap.add_argument("-o", "--output", required=True, help="path to the output loss/accuary plot")
	args = vars(ap.parse_args())
	
	output = args["output"]
		
	'''
	output = 'c:/temp/output'

	keras_mnist(output)