from sklearn import datasets
import numpy as np
from layer import *
import os
import pickle
import matplotlib
import matplotlib.pyplot as plt
import random
import itertools
import math
import mnist_load
from p4_model import *
#matplotlib.use("TkAgg")

train_set, dev_set, test_set = mnist_load.load_mnistdata()

train_x,train_y = train_set
dev_x,dev_y = dev_set
test_x,test_y = test_set

gen:np.random.Generator = np.random.default_rng()
eta = 0.00001

MiniBatchN = 32

model = load_or_create_model([300,10])

end_n = math.floor(3500*17 /MiniBatchN)

for epoch in range(1):
    #one epoch
    for iteration in range(0,end_n):
        choiced_index = gen.choice(range(0,len(train_x)),MiniBatchN)
        batch_x = train_x[choiced_index]
        batch_y = train_y[choiced_index]
        #batch_x = train_x[MiniBatchN*iteration:MiniBatchN*(iteration+1)]
        #batch_y = train_y[MiniBatchN*iteration:MiniBatchN*(iteration+1)]
        model.train_one_iterate(batch_x,batch_y,eta)
        if (model.iteration-1) % 200 == 0:
            model.set_checkpoint(dev_x,dev_y)
        if (model.iteration) % 10 == 0:
            print(f"iteration {model.iteration+1}")

J = model.caculate(dev_x,dev_y)
loss = np.average(J.numpy())
print('testset : avg loss : ',loss)

confusion = get_confusion(J)
accuracy = get_accuracy_from_confusion(confusion)
print('accuracy : {:.2f}%'.format(accuracy * 100))
if True:
    save_model(model)

plt.subplot(1,2,1)
plt.title("accuracy")
plt.plot([*map(lambda x: x.iteration,model.checkpoints)],
    [*map(lambda x: x.accuracy,model.checkpoints)]
)
plt.subplot(1,2,2)
plt.title("loss")
plt.plot([*map(lambda x: x.iteration,model.checkpoints)],
    [*map(lambda x: x.loss,model.checkpoints)])
plt.show()

plt.title("confusion matrix")
plt.imshow(confusion,cmap='Blues')
plt.colorbar()
for i,j in itertools.product(range(confusion.shape[0]),range(confusion.shape[1])):
    plt.text(j,i,"{:}".format(confusion[i,j]),horizontalalignment="center",color="white" if i == j else "black")
plt.show()