import numpy as np

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

class NeuralNetwork:
    def __init__(self, n_input, n_hidden, n_output,m,iters):
        print("Входной слой: ",n_input," элемента")
        print("Скрытый слой: ",n_hidden," элемента")
        print("Выходной слой: ",n_output," элемента")
        W1 = np.random.randn(n_hidden, n_input)
        b1 = np.zeros((n_hidden, 1))
        W2 = np.random.randn(n_output, n_hidden)
        b2 = np.zeros((n_output, 1))
        self.parameters = {
            "weights1": W1,
            "shifts1" : b1,
            "weights2": W2,
            "shifts2" : b2
        }
        print("Начальные веса между входным и скрытым:")
        print(W1)
        print("Начальные веса между скрытым и выходным:")
        print(W2)
        print("Коэффииенты сдвига скрытых элементов:")
        print(b1)
        print("Коэффициенты сдвига выходного элемента:")
        print(b2)
        self.learning_rate=0.3
        print("норма обучения: ",self.learning_rate)
        self.num_of_iters=iters
        self.m = m

    def forward_prop(self,X):
        W1 = self.parameters["weights1"]
        b1 = self.parameters["shifts1"]
        W2 = self.parameters["weights2"]
        b2 = self.parameters["shifts2"]
        Z1 = np.dot(W1, X) + b1
        A1 = np.tanh(Z1)
        Z2 = np.dot(W2, A1) + b2
        A2 = sigmoid(Z2)
        cache = {
            "A1": A1,
            "A2": A2
        }
        return A2, cache

    def calculate_cost(self,A2, Y):
        cost = -np.sum(np.multiply(Y, np.log(A2)) +  np.multiply(1-Y, np.log(1-A2)))/m
        cost = np.squeeze(cost)
        return cost

    def backward_prop(self,X, Y, cache):
        A1 = cache["A1"]
        A2 = cache["A2"]
        W2 = self.parameters["weights2"]
        dZ2 = A2 - Y
        dW2 = np.dot(dZ2, A1.T)/m
        db2 = np.sum(dZ2, axis=1, keepdims=True)/m
        dZ1 = np.multiply(np.dot(W2.T, dZ2), 1-np.power(A1, 2))
        dW1 = np.dot(dZ1, X.T)/m
        db1 = np.sum(dZ1, axis=1, keepdims=True)/m
        grads = {
            "dW1": dW1,
            "db1": db1,
            "dW2": dW2,
            "db2": db2
        }
        return grads

    def update_parameters(self, grads):
        W1 = self.parameters["weights1"]
        b1 = self.parameters["shifts1"]
        W2 = self.parameters["weights2"]
        b2 = self.parameters["shifts2"]
        dW1 = grads["dW1"]
        db1 = grads["db1"]
        dW2 = grads["dW2"]
        db2 = grads["db2"]
        W1 = W1 - self.learning_rate*dW1
        b1 = b1 - self.learning_rate*db1
        W2 = W2 - self.learning_rate*dW2
        b2 = b2 - self.learning_rate*db2
        new_parameters = {
            "weights1": W1,
            "weights2": W2,
            "shifts1" : b1,
            "shifts2" : b2
        }
        self.parameters=new_parameters


    def train(self,X, Y):
        print("Начинаем обучаться на данных")
        print(X)
        print(Y)
        for i in range(0, self.num_of_iters+1):
            a2, cache = self.forward_prop(X)
            cost = self.calculate_cost(a2, Y)
            grads = self.backward_prop(X, Y, cache)
            parameters = self.update_parameters(grads)
            if(i%200 == 0):
                print('Функция стоимости после итерации# {:d}: {:f}'.format(i, cost))
        print("Параметры после обучения")
        print("Начальные веса между входным и скрытым:")
        print(self.parameters["weights1"])
        print("Начальные веса между скрытым и выходным:")
        print(self.parameters["weights2"])
        print("Коэфициенты сдвига скрытых элементов:")
        print(self.parameters["shifts1"])
        print("Коэфициенты сдвига выходного элемента:")
        print(self.parameters["shifts2"])


    def getResult(self,X):
        a2, cache = self.forward_prop(X)
        yhat = a2
        yhat = np.squeeze(yhat)
        if(yhat >= 0.5):
            y_predict = 1
        else:
            y_predict = 0

        return y_predict

def test(x1,x2,n):
    X_test = np.array([[x1], [x2]])
    y_predict = n.getResult(X_test)
    print('На вход подано ({:d}, {:d}) на выходе получили {:d}'.format(
        X_test[0][0], X_test[1][0], y_predict))


def getTrain(operation):
    if(operation=='AND'):
        X = np.array([[0, 0, 1, 1], [0, 1, 0, 1]])
        Y = np.array([[0, 0, 0, 1]])
        m = X.shape[1]
        return X,Y,m
    if(operation=='OR'):
        X = np.array([[0, 0, 1, 1], [0, 1, 0, 1]])
        Y = np.array([[0, 1, 1, 1]])
        m = X.shape[1]
        return X,Y,m
    if(operation=='XOR'):
        X = np.array([[0, 0, 1, 1], [0, 1, 0, 1]])
        Y = np.array([[0, 1, 1, 0]])
        m = X.shape[1]
        return X,Y,m

OPERATION = 'XOR'
print("Тернируемся для операции ",OPERATION)
np.random.seed(2)
X,Y,m = getTrain(OPERATION)
n = NeuralNetwork(2,2,1,m,1000)
n.train(X,Y)
print("Проверяем результат обучения")
test(1,0,n)
test(0,0,n)
test(0,1,n)
test(1,1,n)