import random
import math
import json

INPUTS = [[0,0],[0,1],[1,0],[1,1]]
OUTPUTS = [[0],[1],[1],[0]]
EPOCHS = 1000000
ALPHAS = 20


WEPOCHS = EPOCHS // 100


VARIANCE_W = 0.5
VARIANCE_B = 0

w11 = random.uniform(-VARIANCE_W,VARIANCE_W)
w21 = random.uniform(-VARIANCE_W,VARIANCE_W)
b1 = VARIANCE_B

w12 = random.uniform(-VARIANCE_W,VARIANCE_W)
w22 = random.uniform(-VARIANCE_W,VARIANCE_W)
b2 = VARIANCE_B

w13 = random.uniform(-VARIANCE_W,VARIANCE_W)
w23 = random.uniform(-VARIANCE_W,VARIANCE_W)
b3 = VARIANCE_B

o1 = random.uniform(-VARIANCE_W,VARIANCE_W)
o2 = random.uniform(-VARIANCE_W,VARIANCE_W)
o3 = random.uniform(-VARIANCE_W,VARIANCE_W)
ob = VARIANCE_B

## Tudo a 0.5
# VARIANCE_W = 0.5
# VARIANCE_B = 1
# w11 = VARIANCE_W
# w21 = VARIANCE_W
# b1 = VARIANCE_B

# w12 = VARIANCE_W
# w22 = VARIANCE_W
# b2 = VARIANCE_B

# w13 = VARIANCE_W
# w23 = VARIANCE_W
# b3 = VARIANCE_B

# o1 = VARIANCE_W
# o2 = VARIANCE_W
# o3 = VARIANCE_W
# ob = VARIANCE_B


def sigmoid(x):
    return 1.0 / (1.0 + math.exp(-x))


def sigmoid_prime(x): # x already sigmoided
    return x * (1 - x)

def relu(x):
    return max(0,x)

def relu_prime(x):
    return 1 if x>0 else 0

def tanh(x):
    return math.tanh(x)

def tanh_prime(x):
    return 1 - x**2

def softmax(x):
    return math.exp(x) / (math.exp(x) + 1)

def softmax_prime(x):
    return x * (1 - x)

def predict(i1, i2, activation=sigmoid):    
    s1 = w11 * i1 + w21 * i2 + b1
    # s1 = sigmoid(s1)
    s1 = activation(s1)
    s2 = w12 * i1 + w22 * i2 + b2
    # s2 = sigmoid(s2)
    s2 = activation(s2)
    s3 = w13 * i1 + w23 * i2 + b3
    # s3 = sigmoid(s3)
    s3 = activation(s3)
    
    output = s1 * o1 + s2 * o2 + s3 * o3 + ob
    # output = sigmoid(output)
    output = activation(output)
    
    return output


def learn(i1,i2,target, activation, activation_prime, alpha=0.2):
    global w11,w21,b1,w12,w22,b2,w13,w23,b3
    global o1,o2,o3,ob
    
    s1 = w11 * i1 + w21 * i2 + b1
    # s1 = sigmoid(s1)
    s1 = activation(s1)
    s2 = w12 * i1 + w22 * i2 + b2
    # s2 = sigmoid(s2)
    s2 = activation(s2)
    s3 = w13 * i1 + w23 * i2 + b3
    # s3 = sigmoid(s3)
    s3 = activation(s3)
    
    output = s1 * o1 + s2 * o2 + s3 * o3 + ob
    # output = sigmoid(output)
    output = activation(output)
    
    error = target - output
    # derror = error * sigmoid_prime(output)
    derror = error * activation_prime(output)
    
    # ds1 = derror * o1 * sigmoid_prime(s1)
    ds1 = derror * o1 * activation_prime(s1)
    # ds2 = derror * o2 * sigmoid_prime(s2)
    ds2 = derror * o2 * activation_prime(s2)
    # ds3 = derror * o3 * sigmoid_prime(s3)
    ds3 = derror * o3 * activation_prime(s3)
    
    o1 += alpha * s1 * derror
    o2 += alpha * s2 * derror
    o3 += alpha * s3 * derror
    ob += alpha * derror
    
    w11 += alpha * i1 * ds1
    w21 += alpha * i2 * ds1
    b1 += alpha * ds1
    w12 += alpha * i1 * ds2
    w22 += alpha * i2 * ds2
    b2 += alpha * ds2
    w13 += alpha * i1 * ds3
    w23 += alpha * i2 * ds3
    b3 += alpha * ds3   


def train(epochs=EPOCHS, alpha=ALPHAS):
    modelo = None
    for epoch in range(1,epochs+1):
        indexes = [0,1,2,3]
        random.shuffle(indexes)
        for j in indexes:
            learn(INPUTS[j][0],INPUTS[j][1],OUTPUTS[j][0], activation=sigmoid, activation_prime=sigmoid_prime, alpha=alpha)
        
        if epoch%WEPOCHS  == 0:
            cost = 0
            for j in range(4):
                o = predict(INPUTS[j][0],INPUTS[j][1], activation=sigmoid)
                cost += (OUTPUTS[j][0] - o) ** 2
            cost /= 4
            print("epoch", epoch, "mean squared error:", cost)
    
    modelo = {
        "w11": w11,
        "w21": w21,
        "b1": b1,
        "w12": w12,
        "w22": w22,
        "b2": b2,
        "w13": w13,
        "w23": w23,
        "b3": b3,
        "o1": o1,
        "o2": o2,
        "o3": o3,
        "ob": ob
    }
    return modelo

def save_model(modelo, filename):
    with open(filename, 'w') as json_file:
        json.dump(modelo, json_file)

## Main
def main():
    # Train model
    modelo = train()    
    
    print(modelo)
    # Save model to file
    save_model(modelo, "modelo.json")

    for i in range(4):
        result = 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")
        # print("for input", INPUTS[i], "expected", OUTPUTS[i][0], "predicted", result, "which is", "correct" if round(result)==OUTPUTS[i][0] else "incorrect")


if __name__ == "__main__":
    main()