import numpy as np
import json
import matplotlib.pyplot as plt
import itertools


def sigmoid(x):
    return 1/(1 + np.exp(-x))


def per(size, ones):
    return list(filter(lambda x: sum(x) == ones, itertools.product([0, 1], repeat=size)))


def toNum(l):
    s = "".join(map(str, l))
    return int(s, base=2)


class Ubal:
    """" Vanila UBAL """

    def __init__(self, numInp, numHid, numOut, f_beta_input, gamma_1, alpha=0.25):

        # sizes
        self.numInp = numInp
        self.numHid = numHid
        self.numOut = numOut

        # meta arg
        self.lr = alpha
        self.f_beta_input = f_beta_input
        self.f_beta_hidden = f_beta_input
        self.f_beta_output = f_beta_input

        self.b_beta_input = 1 - f_beta_input
        self.b_beta_hidden = 1 - self.f_beta_hidden
        self.b_beta_output = 1 - self.f_beta_output

        self.gamma_1 = gamma_1
        self.gamma_2 = gamma_1
        self.gamma_3 = gamma_1
        self.gamma_4 = gamma_1

        # init weights
        self.w1 = np.random.uniform(size=(numInp, numHid)) / 10
        self.b1 = np.random.uniform(size=(1, numHid)) / 10
        self.w2 = np.random.uniform(size=(numHid, numOut)) / 10
        self.b2 = np.random.uniform(size=(1, numOut)) / 10

        self.m1 = np.random.uniform(size=(numOut, numHid)) / 10
        self.d1 = np.random.uniform(size=(1, numHid)) / 10
        self.m2 = np.random.uniform(size=(numHid, numInp)) / 10
        self.d2 = np.random.uniform(size=(1, numInp)) / 10

    def train(self, X, y, epochs=50000, track=False):

        # init statistic
        fse, bse, fsr, bsr, fce, bce = [], [], [], [], [], []

        # training
        for epoch in range(epochs):
            f_sum_error, b_sum_error, f_class_error, = 0, 0, 0
            f_msqre, b_msqre, b_class_error = 0, 0, 0
            # forward phase
            qhfp = sigmoid(np.dot(X, self.w1) + self.b1)  # hidden prediction
            pofe = sigmoid(np.dot(qhfp, self.m2) + self.d2)  # output echo

            qofp = sigmoid(np.dot(qhfp, self.w2) + self.b2)  # output prediction
            phfe = sigmoid(np.dot(qofp, self.m1) + self.d1)  # hidden echo

            # backward phase
            phbp = sigmoid(np.dot(y, self.m1) + self.d1)  # hidden prediction
            qobe = sigmoid(np.dot(phbp, self.w2) + self.b2)  # output echo

            pobp = sigmoid(np.dot(phbp, self.m2) + self.d2)  # output prediction
            qhbe = sigmoid(np.dot(pobp, self.w1) + self.b1)  # hidden echo

            # learning rules hidden
            tfh = self.f_beta_hidden * qhfp + (1 - self.b_beta_hidden) * phbp
            efh = self.gamma_4 * qhfp + (1 - self.gamma_4) * qhbe
            tbh = self.b_beta_hidden * phbp + (1 - self.f_beta_hidden) * qhfp
            ebh = self.gamma_4 * phbp + (1 - self.gamma_4) * phfe

            # learning rules output
            tfo = self.f_beta_output * qofp + (1 - self.f_beta_output) * y
            efo = self.gamma_4 * qofp + (1 - self.gamma_4) * qobe
            tbo = self.b_beta_input * pobp + (1 - self.b_beta_input) * X
            ebo = self.gamma_4 * pobp + (1 - self.gamma_4) * pofe

            # update weights and biases
            delta_h_f = self.lr * np.dot(tbo.T, (tfh - efh))
            self.w1 += delta_h_f
            self.b1 += np.sum(tfh - efh, axis=0)

            delta_o_f = self.lr * np.dot(tbh.T, (tfo - efo))
            self.w2 += delta_o_f
            self.b2 += np.sum(tfo - efo, axis=0)

            delta_h_b = self.lr * np.dot(tfo.T, (tbh - ebh))
            self.m1 += delta_h_b
            self.d1 += np.sum(tbh - ebh, axis=0)

            delta_o_b = self.lr * np.dot(tfh.T, (tbo - ebo))
            self.m2 += delta_o_b
            self.d2 += np.sum(tbo - ebo, axis=0)

            # sum_error
            f_error = (np.where(qofp > .5, 1, 0) - y) ** 2
            b_error = (np.where(pobp > .5, 1, 0) - X) ** 2
            f_sum_error += np.sum(f_error)
            b_sum_error += np.sum(b_error)

            f_class_error += len(list(filter(lambda x: np.sum(x) != 0, f_error)))
            b_class_error += len(list(filter(lambda x: np.sum(x) != 0, b_error)))

            f_msqre += np.sum((qofp - y) ** 2) / len(y)
            b_msqre += np.sum((pobp - X) ** 2) / len(X)

            # add errors
            fse.append(f_sum_error)
            bse.append(b_sum_error)
            fsr.append(f_msqre)
            bsr.append(b_msqre)
            fce.append((f_class_error / len(f_error)) * 100)
            bce.append((b_class_error / len(b_error)) * 100)

            if track and epoch % 500 == 0 and epoch != 0:
                print('epoch: {}   '
                      'forward error: {}   '
                      'backward error: {}   '
                      'forward loss: {}   '
                      'backward loss: {}   '
                      'forwars class error {}%  '
                      'backward class error {}%  '.format(epoch, f_sum_error, b_sum_error, f_msqre, b_msqre,
                                                          fce[-1], bce[-1]))
                plt.subplot(311).set_title("mean squared error")
                plt.plot(fsr[-499:], label="Forward")
                plt.plot(bsr[-499:], label="Backward")
                plt.legend()
                plt.subplot(313).set_title("class error")
                plt.plot(fce[-499:], label="Forward")
                plt.plot(bce[-499:], label="Backward")
                plt.legend()
                plt.show()

        return fsr, fce

if __name__ == "__main__":
    """ 
        Teaching raw data
    """
    f = open("../data.data")
    rawData = json.load(f)
    X = np.array([i['proprio'] for i in rawData])
    y = np.array([i["touch"] for i in rawData])

    for speed in [0.1, 0.2, 0.3]:
        for gam in [0, 0.25, 0.5, 0.75, 1]:
            for bet in [0, 0.25, 0.5, 0.75, 1]:
                model = Ubal(len(X[0]), 80, len(y[0]), 0.5, 0.5, 0.2)
                msqre, tmp = model.train(X, y, 1001, track=False)
                print("alfa", speed, "beta", "{0:.2f}".format(bet), "gama", "{0:.2f}".format(gam), "res", min(msqre), "max", max(msqre))


