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.0001

MiniBatchN = 32

model = load_or_create_model([300,10])

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

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

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

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()