from layer import *
import numpy as np
import pickle
import os

class CheckPoint:
    def __init__(self,param,accuracy,loss,iteration):
        super().__init__()
        self.param = param
        self.accuracy = accuracy
        self.loss = loss
        self.iteration = iteration

class Model:
    def __init__(self, layerDim:[int]):
        super().__init__()
        gen:np.random.Generator = np.random.default_rng()
        self.layerDim = layerDim
        self.param = []
        self.checkpoints = []
        self.iteration = 0
        #...
        front = 784
        for sd in layerDim:
            back = sd
            weight = Variable(gen.normal(0,1,size=(front,back)))
            bias = Variable(gen.normal(0,1,size=(back)))
            self.param.append((weight,bias))
            front = back
    
    def caculate(self,input_x,y):
        input_var = Variable(input_x)
        Z = input_var
        for i,(w,b) in enumerate(self.param):
            U = Z @ w + b
            if i < len(self.param) - 1:
                Z = relu(U)
            else:
                Z = U
        J = SoftmaxWithNegativeLogLikelihood(Z,y)
        return J
    def train_one_iterate(self,input_x,y,eta):
        #forward pass
        J = self.caculate(input_x,y)
        #backpropagation
        J.backprop(np.ones(()))
        for i,(w,b) in enumerate(self.param):
            w = Variable(w.numpy() - (w.grad) * eta)
            b = Variable(b.numpy() - (b.grad) * eta)
            self.param[i] = (w,b)
        self.iteration += 1
        return J
    
    def get_loss_and_confusion(self,input_x,y):
        J = self.caculate(input_x,y)
        s = J.softmax_numpy()
        s = np.round(s)
        confusion = (np.transpose(y)@s)
        return J.numpy(), confusion

    def set_checkpoint(self,dev_x,dev_y):
        J = self.caculate(dev_x,dev_y)
        loss = np.average(J.numpy())
        print(f"check point #{len(self.checkpoints)}")
        print(self.iteration,'iteration : avg loss : ',loss)

        confusion = get_confusion(J)
        accuracy = get_accuracy_from_confusion(confusion)
        print('accuracy : {:.2f}%'.format(accuracy * 100))
        self.checkpoints.append(CheckPoint(
            self.param,
            accuracy*100,
            loss,
            self.iteration
        ))

def get_confusion(J:SoftmaxWithNegativeLogLikelihood):
    s = J.softmax_numpy()
    s = np.eye(10)[np.argmax(s,axis=len(s.shape)-1)]
    confusion = (np.transpose(J.y)@s)
    return confusion

def get_accuracy_from_confusion(confusion):
    return np.trace(confusion).sum() / np.sum(confusion)

def model_filename(layerDim:[int]):
    return f"model{layerDim}.pickle"

def save_model(model:Model):
    with open(model_filename(model.layerDim),"wb") as model_file:
        pickle.dump(model,model_file)

def load_or_create_model(layerDim:list):
    model_name = model_filename(layerDim)
    if os.path.exists(model_name):
        with open(model_name,"rb") as model_file:
            return pickle.load(model_file)
    else:
        return Model(layerDim)