import random
import json
from typing import Any
from utils import sigmoid, sigmoid_prime

EPOCHS = 1000000
ALPHAS = 20

WEPOCHS = EPOCHS // 100
VARIANCE_W = 0.5
VARIANCE_B = 1

class NeuralNetwork:
    def __init__(self):
        self._w11 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._w21 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._b1 = VARIANCE_B

        self._w12 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._w22 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._b2 = VARIANCE_B

        self._w13 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._w23 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._b3 = VARIANCE_B

        self._o1 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._o2 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._o3 = random.uniform(-VARIANCE_W,VARIANCE_W)
        self._ob = VARIANCE_B
    
    def predict(self, i1, i2, activation=sigmoid):    
        s1 = self._w11 * i1 + self._w21 * i2 + self._b1
        s1 = activation(s1)
        s2 = self._w12 * i1 + self._w22 * i2 + self._b2
        s2 = activation(s2)
        s3 = self._w13 * i1 + self._w23 * i2 + self._b3
        s3 = activation(s3)
        
        output = s1 * self._o1 + s2 * self._o2 + s3 * self._o3 + self._ob
        output = activation(output)
        
        return output


    def learn(self, i1, i2, target, activation=sigmoid, activation_prime=sigmoid_prime, alpha=0.2):
        s1 = self._w11 * i1 + self._w21 * i2 + self._b1
        s1 = activation(s1)
        s2 = self._w12 * i1 + self._w22 * i2 + self._b2
        s2 = activation(s2)
        s3 = self._w13 * i1 + self._w23 * i2 + self._b3
        s3 = activation(s3)
        
        output = s1 * self._o1 + s2 * self._o2 + s3 * self._o3 + self._ob
        output = activation(output)
        
        error = target - output
    
        derror = error * activation_prime(output)
        ds1 = derror * self._o1 * activation_prime(s1)
        ds2 = derror * self._o2 * activation_prime(s2)
        ds3 = derror * self._o3 * activation_prime(s3)
        
        self._o1 += alpha * s1 * derror
        self._o2 += alpha * s2 * derror
        self._o3 += alpha * s3 * derror
        self._ob += alpha * derror
        
        self._w11 += alpha * i1 * ds1
        self._w21 += alpha * i2 * ds1
        self._b1 += alpha * ds1
        self._w12 += alpha * i1 * ds2
        self._w22 += alpha * i2 * ds2
        self._b2 += alpha * ds2
        self._w13 += alpha * i1 * ds3
        self._w23 += alpha * i2 * ds3
        self._b3 += alpha * ds3   


    def train(self, inputs=[], outputs=[], epochs=EPOCHS, alpha=ALPHAS):
        if len(inputs) > 0 and len(outputs) > 0:
            for epoch in range(1,epochs+1):
                indexes = [0,1,2,3]
                random.shuffle(indexes)
                for j in indexes:
                    self.learn(inputs[j][0],inputs[j][1],outputs[j][0], activation=sigmoid, activation_prime=sigmoid_prime, alpha=alpha)
            
                # Print cost every 100 epochs for debug
                if epoch%WEPOCHS  == 0:
                    cost = 0
                    for j in range(4):
                        o = self.predict(inputs[j][0],inputs[j][1], activation=sigmoid)
                        cost += (outputs[j][0] - o) ** 2
                    cost /= 4
                    print("epoch", epoch, "mean squared error:", cost)

    def save_model(self, filename):
        model = self.getModelJson()
        with open(filename, 'w') as json_file:
            json.dump(model, json_file)

    def getModelJson(self):
        return {
            "w11": self._w11,
            "w21": self._w21,
            "b1": self._b1,
            "w12": self._w12,
            "w22": self._w22,
            "b2": self._b2,
            "w13": self._w13,
            "w23": self._w23,
            "b3": self._b3,
            "o1": self._o1,
            "o2": self._o2,
            "o3": self._o3,
            "ob": self._ob
        }
    
    def __str__(self):
        return self.getModelJson()
    

## Main
def main():
    INPUTS = [[0,0],[0,1],[1,0],[1,1]]
    OUTPUTS = [[0],[1],[1],[0]]
    # Train model
    nn = NeuralNetwork()
    print("Neural Network created")
    nn.train(inputs=INPUTS, outputs=OUTPUTS, epochs=EPOCHS, alpha=ALPHAS)
    print("Model trained")
    print("Printing Model:")
    print(nn._getModelJson())

    # Save model to files
    nn.save_model("model.json")
    print("Model saved to model.json")

    for i in range(4):
        result = nn.predict(INPUTS[i][0],INPUTS[i][1], activation=sigmoid)
        print("for input", INPUTS[i], "expected", OUTPUTS[i][0], "predicted", f"{result:4.4}", "which is", "correct" if round(result)==OUTPUTS[i][0] else "incorrect")

if __name__ == "__main__":
    main()