import numpy as np


def one_in_k_encoding(vec, k):
    """ One-in-k encoding of vector to k classes

    Args:
       vec: numpy array - data to encode
       k: int - number of classes to encode to (0,...,k-1)
    """
    n = vec.shape[0]
    enc = np.zeros((n, k))
    enc[np.arange(n), vec] = 1
    return enc


def softmax(X):
    """
    You can take this from handin I
    Compute the softmax of each row of an input matrix (2D numpy array).

    the numpy functions amax, log, exp, sum may come in handy as well as the keepdims=True option and the axis option.
    Remember to handle the numerical problems as discussed in the description.
    You should compute lg softmax first and then exponentiate

    More precisely this is what you must do.

    For each row x do:
    compute max of x
    compute the log of the denominator sum for softmax but subtracting out the max i.e (log sum exp x-max) + max
    compute log of the softmax: x - logsum
    exponentiate that

    You can do all of it without for loops using numpys vectorized operations.

    Args:
        X: numpy array shape (n, d) each row is a data point
    Returns:
        res: numpy array shape (n, d)  where each row is the softmax transformation of the corresponding row in X i.e res[i, :] = softmax(X[i, :])
    """
    res = np.zeros(X.shape)
    toAssert = res
    ### YOUR CODE HERE no for loops please

    max = np.amax(X, axis=1, keepdims=True)
    exp = np.exp(X - max)
    np_sum = np.sum(exp, axis=1, keepdims=True)
    logsum = np.log(np_sum) + max
    res = np.exp(X - logsum)

    assert toAssert.shape == res.shape, "Shapes mismatch in softmax. Expected {0} - Got {1}".format(toAssert, res)
    ### END CODE
    return res


def relu(x):
    """ Compute the relu activation function on every element of the input

        Args:
            x: np.array
        Returns:
            res: np.array same shape as x
        Beware of np.max and look at np.maximum
    """
    ### YOUR CODE HERE
    zeros = np.zeros(x.shape)
    res = np.maximum(x, zeros)
    ### END CODE
    return res


def make_dict(W1, b1, W2, b2):
    """ Trivial helper function """
    return {'W1': W1, 'b1': b1, 'W2': W2, 'b2': b2}


def get_init_params(input_dim, hidden_size, output_size):
    """ Initializer function using he et al Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification

    Args:
      input_dim: int
      hidden_size: int
      output_size: int
    Returns:
       dict of randomly initialized parameter matrices.
    """
    W1 = np.random.normal(0, np.sqrt(1. / (input_dim + hidden_size)), size=(input_dim, hidden_size))
    b1 = np.zeros((1, hidden_size))
    W2 = np.random.normal(0, np.sqrt(2. / (hidden_size + output_size)), size=(hidden_size, output_size))
    b2 = np.zeros((1, output_size))
    return {'W1': W1, 'b1': b1, 'W2': W2, 'b2': b2}


def test_shapes_fit(name, vec, test_vec):
    if vec.shape == test_vec.shape:
        # print("SHAPES FIT \"{2}\" : VEC: {0} , TEST_VEC: {1}".format(vec.shape, test_vec.shape,name))
        A = None
    else:
        B = None
        # print("SHAPES NOT FITTING: \"{2}\" : VEC: {0} , TEST_VEC: {1}".format(vec.shape, test_vec.shape,name))

    # assert vec.shape == test_vec.shape, "SHAPES NOT FITTING: VEC: {0} , TEST_VEC: {1}".format(vec, test_vec)


class NetClassifier():

    def __init__(self):
        """ Trivial Init """
        self.params = None
        self.hist = None

    def predict(self, X, params=None):
        """ Compute class prediction for all data points in class X

        Args:
            X: np.array shape n, d
            params: dict of params to use (if none use stored params)
        Returns:
            np.array shape n, 1
        """
        if params is None:
            params = self.params
        pred = None
        ### YOUR CODE HERE
        W1 = params['W1']
        b1 = params['b1']
        W2 = params['W2']
        b2 = params['b2']
        hidden_layer = relu(X @ W1 + b1)
        pred = hidden_layer @ W2 + b2
        ### END CODE
        return pred

    def score(self, X, y, params=None):
        """ Compute accuracy of model on data X with labels y

        Args:
            X: np.array shape n, d
            y: np.array shape n, 1
            params: dict of params to use (if none use stored params)

        Returns:
            np.array shape n, 1
        """
        if params is None:
            params = self.params
        acc = None
        ### YOUR CODE HERE
        acc = np.mean(np.equal(self.predict(X, params).T, y))

        ### END CODE
        return acc

    @staticmethod
    def cost_grad(X, y, params, reg=0.0):
        """ Compute cost and gradient of neural net on data X with labels y using weight decay parameter c
        You should implement a forward pass and store the intermediate results
        and the implement the backwards pass using the intermediate stored results

        Use the derivative for cost as a function for input to softmax as derived above

        Args:
            X: np.array shape n, self.input_size
            y: np.array shape n, 1
            params: dict with keys (W1, W2, b1, b2)
            reg: float - weight decay regularization weight
            params: dict of params to use for the computation

        Returns
            cost: scalar - average cross entropy cost
            dict with keys
            d_w1: np.array shape w1.shape, entry d_w1[i, j] = \partial cost/ \partial w1[i, j]
            d_w2: np.array shape w2.shape, entry d_w2[i, j] = \partial cost/ \partial w2[i, j]
            d_b1: np.array shape b1.shape, entry d_b1[1, j] = \partial cost/ \partial b1[1, j]
            d_b2: np.array shape b2.shape, entry d_b2[1, j] = \partial cost/ \partial b2[1, j]

        """

        W1 = params['W1']
        b1 = params['b1']
        W2 = params['W2']
        b2 = params['b2']
        labels = one_in_k_encoding(y, W2.shape[1])  # shape n x k

        ### INIT ###
        cost = 0
        d_W1 = d_W1_TEST = np.zeros(W1.shape)
        d_W2 = d_W2_TEST = np.zeros(W2.shape)
        d_b1 = d_b1_TEST = np.zeros(b1.shape)
        d_b2 = d_b2_TEST = np.zeros(b2.shape)
        n = X.shape[0]
        d1 = W1.shape[0]
        d2 = W1.shape[1]
        d3 = W2.shape[1]
        ### YOUR CODE HERE - FORWARD PASS - compute regularized cost and store relevant values for backprop

        hidden_layer = X @ W1  # => n x d2
        test_shapes_fit("hidden_layer", hidden_layer, np.zeros((n, d2)))

        biased_hidden_layer = hidden_layer + b1  # => n x d2
        test_shapes_fit("biased_hidden_layer", biased_hidden_layer, np.zeros((n, d2)))

        activated_hidden_layer = relu(biased_hidden_layer)  # => n x d2
        test_shapes_fit("activated_hidden_layer", activated_hidden_layer, np.zeros((n, d2)))

        output_layer = activated_hidden_layer @ W2  # => n x d3
        test_shapes_fit("output_layer", output_layer, np.zeros((n, d3)))

        biased_output_layer = output_layer + b2  # => n x d3
        test_shapes_fit("biased_output_layer", biased_output_layer, np.zeros((n, d3)))

        smax = softmax(biased_output_layer)  # => n x d3
        test_shapes_fit("smax", smax, np.zeros((n, d3)))

        smax_where_y_is_one = smax[labels == 1]  # => n x 1
        # test_shapes_fit("smax_where_y_is_one", smax_where_y_is_one, np.zeros((n, 1)))

        L = - np.log(smax_where_y_is_one)  # => n x 1
        # test_shapes_fit("L", L, np.zeros((n, 1)))

        L_mean = np.mean(L)  # => 1 x 1
        # test_shapes_fit("L_mean", L_mean, np.zeros((1, 1)))

        # weight decay
        W1_squared = W1 ** 2.0  # => d1 x d2
        test_shapes_fit("W1_squared", W1_squared, np.zeros((d1, d2)))

        W2_squared = W2 ** 2.0  # => d2 x d3
        test_shapes_fit("W2_squared", W2_squared, np.zeros((d2, d3)))

        summ = np.sum(W1_squared) + np.sum(W2_squared)  # => 1 x 1
        # test_shapes_fit("summ", summ, np.zeros((1, 1)))

        weight_decay = reg * summ  # => 1 x 1
        # test_shapes_fit("weight_decay", weight_decay, np.zeros((1, 1)))

        cost = L_mean + weight_decay  # => 1 x 1
        # test_shapes_fit("cost", cost, np.zeros((1, 1)))

        ### END CODE

        ### YOUR CODE HERE - BACKWARDS PASS - compute derivatives of all (regularized) weights and bias, store them in d_w1, d_w2' d_w2, d_b1, d_b2

        d_cost = 1.0  # => 1 x 1
        d_L_mean = 1 / n * d_cost  # => 1 x 1
        d_L = (- labels + smax) * d_L_mean  # => n x d3
        test_shapes_fit("d_L", d_L, np.zeros((n, d3)))

        d_b2 = 1.0 * d_L.sum(axis=0, keepdims=True)  # => 1 x d3
        test_shapes_fit("d_b2", d_b2, np.zeros((1, d3)))

        d_biased_output_layer = 1.0 * d_L  # => n x d3
        test_shapes_fit("d_biased_output_layer", d_biased_output_layer, np.zeros((n, d3)))

        d_activated_hidden_layer = d_biased_output_layer @ W2.T  # => n x d2
        test_shapes_fit("d_activated_hidden_layer", d_activated_hidden_layer, np.zeros((n, d2)))

        d_W2 = activated_hidden_layer.T @ d_biased_output_layer  # => d2 x d3
        test_shapes_fit("d_W2", d_W2, np.zeros((d2, d3)))

        # d_activated_hidden_layer = (biased_hidden_layer > 0).astype(int) * d_path_to_relu  # => n * d2
        d_biased_hidden_layer = np.multiply((biased_hidden_layer > 0).astype(int),
                                            d_activated_hidden_layer)  # => n * d2 ???????????
        test_shapes_fit("d_biased_hidden_layer", d_biased_hidden_layer, np.zeros((n, d2)))

        d_b1 = 1.0 * d_biased_hidden_layer.sum(axis=0, keepdims=True)  # => 1 x d2
        test_shapes_fit("d_b1", d_b1, np.zeros((1, d2)))

        d_hidden_layer = 1.0 * d_biased_hidden_layer  # => n x d2
        test_shapes_fit("d_hidden_layer", d_hidden_layer, np.zeros((n, d2)))

        d_W1 = X.T @ d_hidden_layer  # => d1 x d2
        test_shapes_fit("d_W1", d_W1, np.zeros((d1, d2)))

        # Weight decay
        d_weight_decay = d_cost * reg  # 1 x 1

        d_W1_squared = 1.0 * d_weight_decay  # 1 x 1
        d_W2_squared = 1.0 * d_weight_decay  # 1 x 1

        d_W1 += 2.0 * W1 * d_W1_squared  # w1
        test_shapes_fit("d_W1", d_W1, W1)

        d_W2 += 2.0 * W2 * d_W2_squared  # w2
        test_shapes_fit("d_W2", d_W2, W2)

        test_shapes_fit("W1", d_W1, d_W1_TEST)
        test_shapes_fit("W2", d_W2, d_W2_TEST)
        test_shapes_fit("b1", d_b1, d_b1_TEST)
        test_shapes_fit("b2", d_b2, d_b2_TEST)

        ### END CODE
        # the return signature
        return cost, {'d_w1': d_W1, 'd_w2': d_W2, 'd_b1': d_b1, 'd_b2': d_b2}

    def fit(self, X_train, y_train, X_val, y_val, init_params, batch_size=32, lr=0.1, reg=1e-4, epochs=30):
        """ Run Mini-Batch Gradient Descent on data X, Y to minimize the in sample error (1/n)Cross Entropy for Neural Net classification
        Printing the performance every epoch is a good idea to see if the algorithm is working

        Args:
           X_train: numpy array shape (n, d) - the training data each row is a data point
           y_train: numpy array shape (n,) int - training target labels numbers in {0, 1,..., k-1}
           X_val: numpy array shape (n, d) - the validation data each row is a data point
           y_val: numpy array shape (n,) int - validation target labels numbers in {0, 1,..., k-1}
           init_params: dict - has initial setting of parameters
           lr: scalar - initial learning rate
           batch_size: scalar - size of mini-batch
           epochs: scalar - number of iterations through the data to use

        Sets:
           params: dict with keys {W1, W2, b1, b2} parameters for neural net
           history: dict:{keys: train_loss, train_acc, val_loss, val_acc} each an np.array of size epochs of the the given cost after every epoch
        """

        W1 = init_params['W1']
        b1 = init_params['b1']
        W2 = init_params['W2']
        b2 = init_params['b2']
        params = init_params

        d_W1 = d_W2 = d_b1 = d_b2 = None

        train_loss = []
        train_acc = []
        val_loss = []
        val_acc = []
        n = len(X_train)
        b = batch_size
        ### YOUR CODE HERE
        for i in range(1, epochs):
            perm = np.random.permutation(n)
            mini_batches = [perm[k:k + b] for k in range(0, n, b)]

            cost = 0
            for mini_batch in mini_batches:
                xs = np.array([X_train[index] for index in mini_batch])
                ys = np.array([y_train[index] for index in mini_batch])
                _, dictWb = self.cost_grad(xs, ys, params, reg=reg)
                d_W1 = dictWb['d_w1']
                d_b1 = dictWb['d_b1']
                d_W2 = dictWb['d_w2']
                d_b2 = dictWb['d_b2']

                W1 = - lr * (1 / b * d_W1)
                W2 = - lr * (1 / b * d_W2)
                b1 = - lr * (1 / b * d_b1)
                b2 = - lr * (1 / b * d_b2)
                params = make_dict(W1, b1, W2, b2)

            # Calculate the cost over all datapoints, to get a better history
            # Remove for better effeciency
            lr = lr * 0.99
            tester, _ = self.cost_grad(X_train, y_train, params, reg)
            train_loss.append(tester)
            test1 = np.sum(self.predict(X_train, params) == y_train) / len(X_train)
            train_acc.append(test1)
            tester2, _ = self.cost_grad(X_val, y_val, params, reg)
            val_loss.append(tester2)
            test2 = np.sum(self.predict(X_val, params) == y_val) / len(X_val)
            val_acc.append(test2)

        ### END CODE
        # hist dict should look like this with something different than none
        self.history = {
            'train_loss': np.array(train_loss),
            'train_acc': np.array(train_acc),
            'val_loss': np.array(val_loss),
            'val_acc': np.array(val_acc),
        }
        ## self.params should look like this with something better than none, i.e. the best parameters found.
        self.params = params


def numerical_grad_check(f, x, key):
    """ Numerical Gradient Checker """
    eps = 1e-6
    h = 1e-5
    # d = x.shape[0]
    cost, grad = f(x)
    grad = grad[key]
    it = np.nditer(x, flags=['multi_index'])
    while not it.finished:
        dim = it.multi_index
        print(dim)
        tmp = x[dim]
        x[dim] = tmp + h
        cplus, _ = f(x)
        x[dim] = tmp - h
        cminus, _ = f(x)
        x[dim] = tmp
        num_grad = (cplus - cminus) / (2 * h)
        # print('cplus: {0}, cminus: {1}, cplus-cminus: {2}'.format(cplus, cminus, cplus - cminus))
        # print('dim: {0}, grad: {1}, num_grad: {2}, grad-num_grad: {3}'.format(dim, grad[dim], num_grad,
        #                                                                      grad[dim] - num_grad))
        np_abs = np.abs(num_grad - grad[dim])
        abs_eps = np_abs < eps
        print("np_abs: ", np_abs, " np_abs < eps =", abs_eps)
        assert abs_eps, 'numerical gradient error index {0}, numerical gradient {1}, computed gradient {2}'.format(dim,
                                                                                                                   num_grad,
                                                                                                                   grad[
                                                                                                                       dim])
        it.iternext()


def test_grad():
    stars = '*' * 5
    print(stars, 'Testing  Cost and Gradient Together', stars)
    input_dim = 7
    hidden_size = 1
    output_size = 3
    nc = NetClassifier()
    params = get_init_params(input_dim, hidden_size, output_size)

    nc = NetClassifier()
    X = np.random.randn(7, input_dim)
    y = np.array([0, 1, 2, 0, 1, 2, 0])

    f = lambda z: nc.cost_grad(X, y, params, reg=1.0)
    print('\n', stars, 'Test Cost and Gradient of b2', stars)
    numerical_grad_check(f, params['b2'], 'd_b2')
    print(stars, 'Test Success of b2', stars)

    print('\n', stars, 'Test Cost and Gradient of w2', stars)
    numerical_grad_check(f, params['W2'], 'd_w2')
    print('Test Success w2')

    print('\n', stars, 'Test Cost and Gradient of b1', stars)
    numerical_grad_check(f, params['b1'], 'd_b1')
    print(stars, 'Test Success b1', stars)

    print('\n', stars, 'Test Cost and Gradient of w1', stars)
    numerical_grad_check(f, params['W1'], 'd_w1')
    print('Test Success w1')


if __name__ == '__main__':
    input_dim = 3
    hidden_size = 5
    output_size = 4
    batch_size = 7
    nc = NetClassifier()
    params = get_init_params(input_dim, hidden_size, output_size)
    X = np.random.randn(batch_size, input_dim)
    Y = np.array([0, 1, 2, 0, 1, 2, 0])
    nc.cost_grad(X, Y, params, reg=0)
    test_grad()