import time
import numpy as np
import matplotlib.pyplot as plt
from NeuralNet.neural_net import *
from SVM.gradient_check import eval_numerical_gradient, eval_numerical_gradient_array
from FullyConnectedNet import get_CIFAR10_data
from solver import Solver

def rel_error(x, y):
    """ returns relative error """
    return np.max(np.abs(x - y) / (np.maximum(1e-8, np.abs(x) + np.abs(y))))

def print_mean_std(x,axis=0):
    print('  means: ', x.mean(axis=axis))
    print('  stds:  ', x.std(axis=axis))
    print()

def plot_training_history(title, label, baseline, bn_solvers, plot_fn, bl_marker='.', bn_marker='.', labels=None):
    """utility function for plotting training history"""
    plt.title(title)
    plt.xlabel(label)
    bn_plots = [plot_fn(bn_solver) for bn_solver in bn_solvers]
    bl_plot = plot_fn(baseline)
    num_bn = len(bn_plots)
    for i in range(num_bn):
        label = 'with_norm'
        if labels is not None:
            label += str(labels[i])
        plt.plot(bn_plots[i], bn_marker, label=label)
    label = 'baseline'
    if labels is not None:
        label += str(labels[0])
    plt.plot(bl_plot, bl_marker, label=label)
    plt.legend(loc='lower center', ncol=num_bn+1)

def run_batchsize_experiments(normalization_mode, small_data):
    np.random.seed(231)
    hidden_dims = [100,100,100,100,100]
    n_epochs = 10
    weight_scale = 2e-2
    batch_sizes = [5,10,50]
    lr = 10*(-3.5)
    solver_bsize = batch_sizes[0]

    print('No normalization: batch size = ', solver_bsize)
    model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale, normalization=None)
    solver = Solver(model, small_data,
                    num_epochs=n_epochs, batch_size=solver_bsize,
                    update_rule='adam',
                    optim_config={
                        'learning_rate': lr,
                    },
                    verbose=False)
    solver.train()

    bn_solvers = []
    for i in range(len(batch_sizes)):
        b_size = batch_sizes[i]
        print('Batch normalization: batch size = ', b_size)
        bn_model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale, normalization='batchnorm')
        bn_solver = Solver(bn_model, small_data,
                        num_epochs=n_epochs, batch_size=b_size,
                        update_rule='adam',
                        optim_config={
                            'learning_rate': lr,
                        },
                        verbose=False)
        bn_solver.train()
        bn_solvers.append(bn_solver)

    return bn_solvers, solver, batch_sizes



if __name__ == '__main__':

    data = get_CIFAR10_data()
    num_train = 1000
    small_data = {
        'X_train': data['X_train'][:num_train],
        'y_train': data['Y_train'][:num_train],
        'X_val': data['X_val'],
        'y_val': data['Y_val'],
    }


    ##---------------------test bachnorm_forward: pass!-----------------------
    # np.random.seed(231)
    # N, D1, D2, D3 = 200, 50, 60, 3
    # X = np.random.randn(N, D1)
    # W1 = np.random.randn(D1, D2)
    # W2 = np.random.randn(D2, D3)
    # a = np.maximum(0, X.dot(W1)).dot(W2)
    # print('Before batch normalization: ')
    # print_mean_std(a, axis=0)
    #
    # gamma = np.ones((D3,))
    # beta = np.zeros((D3,))
    #
    # print('After batch normalization (gamma=1, beta=0)')
    # a_norm, _ = batchnorm_forward(a, gamma, beta, {'mode': 'train'})
    # print_mean_std(a_norm, axis=0)
    #
    # gamma = np.asarray([1.0, 2.0, 3.0])
    # beta = np.asarray([11.0, 12.0, 13.0])
    # print('After batch normalization (gamma=', gamma, ', beta=', beta, ')')
    # a_norm, _ = batchnorm_forward(a, gamma, beta, {'mode': 'train'})
    # print_mean_std(a_norm, axis=0)

    # np.random.seed(231)
    # N, D1, D2, D3 = 200, 50, 60, 3
    # W1 = np.random.randn(D1, D2)
    # W2 = np.random.randn(D2, D3)
    #
    # bn_param = {'mode': 'train'}
    # gamma = np.ones(D3)
    # beta = np.zeros(D3)
    #
    # for t in range(50):
    #     X = np.random.randn(N, D1)
    #     a = np.maximum(0, X.dot(W1)).dot(W2)
    #     batchnorm_forward(a, gamma, beta, bn_param)
    # bn_param['mode'] = 'test'
    # X = np.random.randn(N, D1)
    # a = np.maximum(0, X.dot(W1)).dot(W2)
    # a_norm, _ = batchnorm_forward(a, gamma, beta, bn_param)
    #
    # print('After batch normalization (test-time):')
    # print_mean_std(a_norm, axis=0)

    ##------------------------test batchnorm_backward: pass!---------------------
    # np.random.seed(231)
    # N, D = 4,5
    # x = 5*np.random.randn(N, D) + 12
    # gamma = np.random.randn(D)
    # beta = np.random.randn(D)
    # dout = np.random.randn(N, D)
    #
    # bn_param = {'mode': 'train'}
    # fx = lambda x:batchnorm_forward(x, gamma, beta, bn_param)[0]
    # fg = lambda a:batchnorm_forward(x, a, beta, bn_param)[0]
    # fb = lambda b:batchnorm_forward(x, gamma, b, bn_param)[0]
    #
    # dx_num = eval_numerical_gradient_array(fx, x, dout)
    # da_num = eval_numerical_gradient_array(fg, gamma.copy(), dout)
    # db_num = eval_numerical_gradient_array(fb, beta.copy(), dout)
    #
    # _, cache = batchnorm_forward(x, gamma, beta, bn_param)
    # dx, dgamma, dbeta = batchnorm_backward(dout, cache)
    # # You should expect to see relative errors between 1e-13 and 1e-8
    # print('dx error: ', rel_error(dx_num, dx))
    # print('dgamma error: ', rel_error(da_num, dgamma))
    # print('dbeta error: ', rel_error(db_num, dbeta))

    ##------------------------test batchnorm_backward_alt: pass!------------------------
    # np.random.seed(231)
    # N, D = 100, 500
    # x = 5 * np.random.randn(N, D) + 12
    # gamma = np.random.randn(D)
    # beta = np.random.randn(D)
    # dout = np.random.randn(N, D)
    #
    # bn_param = {'mode': 'train'}
    # out, cache = batchnorm_forward(x, gamma, beta, bn_param)
    #
    # t1 = time.time()
    # dx1, dgamma1, dbeta1 = batchnorm_backward(dout, cache)
    # t2 = time.time()
    # dx2, dgamma2, dbeta2 = batchnorm_backward_alt(dout, cache)
    # t3 = time.time()
    #
    # print('dx difference: ', rel_error(dx1, dx2))
    # print('dgamma difference: ', rel_error(dgamma1, dgamma2))
    # print('dbeta difference: ', rel_error(dbeta1, dbeta2))
    # print('speedup: %.2fx' % ((t2 - t1) / (t3 - t2)))

    # np.random.seed(231)
    # N, D, H1, H2, C = 2, 15, 20, 30, 10
    # X = np.random.randn(N, D)
    # y = np.random.randint(C, size=(N,))
    #
    # for reg in [0, 3.14]:
    #     print('Running check with reg=', reg)
    #     model = FullyConnectedNet([H1, H2], input_dim=D, num_classes=C,
    #                               reg=reg, weight_scale=5e-2, dtype=np.float64,
    #                               normalization='batchnorm')
    #     loss, grads = model.loss(X, y)
    #     print('Initial loss: ', loss)
    #
    #     for name in sorted(grads):
    #         f = lambda _: model.loss(X, y)[0]
    #         grad_num = eval_numerical_gradient(f, model.params[name], verbose=False, h=1e-5)
    #         print('%s relative error: %.2e' % (name, rel_error(grad_num, grads[name])))
    #     if reg == 0:print()


    ##-------------------------------batch normalization for six layers: pass!------------------------------
    # np.random.seed(231)
    # hidden_dims = [100, 100, 100, 100, 100]
    #
    # weight_scale = 2e-2
    # bn_model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale,
    #                              normalization='batchnorm')
    # model = FullyConnectedNet(hidden_dims, weight_scale=weight_scale,
    #                           normalization=None)
    #
    # bn_solver = Solver(bn_model, small_data,
    #                    num_epochs=10, batch_size=50,
    #                    update_rule='adam',
    #                    optim_config={
    #                        'learning_rate': 1e-3,
    #                    },
    #                    verbose=True, print_every=20
    #                    )
    # bn_solver.train()
    #
    # solver = Solver(model, small_data,
    #                    num_epochs=10, batch_size=50,
    #                    update_rule='adam',
    #                    optim_config={
    #                        'learning_rate': 1e-3,
    #                    },
    #                    verbose=True, print_every=20
    #                    )
    # solver.train()
    # plt.subplot(3, 1, 1)
    # plot_training_history('Training loss', 'Iteration', solver, [bn_solver],
    #                       lambda x: x.loss_history, bl_marker='o', bn_marker='o')
    # plt.subplot(3, 1, 2)
    # plot_training_history('Training accuracy', 'Epoch', solver, [bn_solver],
    #                       lambda x: x.train_acc_history, bl_marker='-o', bn_marker='-o')
    # plt.subplot(3, 1, 3)
    # plot_training_history('Validation accuracy', 'Epoch', solver, [bn_solver],
    #                       lambda x: x.val_acc_history, bl_marker='-o', bn_marker='-o')
    #
    # plt.gcf().set_size_inches(15, 15)
    # plt.show()


    # bn_solvers, solver, batch_sizes = run_batchsize_experiments('batchnorm', small_data)
    # plt.subplot(2,1,1)
    # plot_training_history('Training accuracy (Batch normalization)', 'Epoch', solver, bn_solvers,
    #                       lambda x: x.train_acc_history, bl_marker='-^', bn_marker='-o', labels=batch_sizes)
    #
    # plt.subplot(2,1,2)
    # plot_training_history('Validation accuracy (Batch normalization)', 'Epoch', solver, bn_solvers,
    #                       lambda x:x.val_acc_history, bl_marker='-^', bn_marker='-o', labels=batch_sizes)
    #
    # plt.gcf().set_size_inches(15, 10)
    # plt.show()


    ##--------------------------------test layernorm forward: pass!----------------------------
    # np.random.seed(231)
    # N, D1, D2, D3 = 4, 50, 60, 3
    # X = np.random.randn(N, D1)
    # W1 = np.random.randn(D1, D2)
    # W2 = np.random.randn(D2, D3)
    # a = np.maximum(0, X.dot(W1)).dot(W2)
    #
    # print('Before layer normalizatioN: ')
    # print_mean_std(a, axis=1)
    #
    # gamma = np.ones(D3)
    # beta = np.zeros(D3)
    #
    # print('After layer normalization (gamma=1, beta=0)')
    # a_norm, _ = layernorm_forward(a, gamma, beta, {'mode': 'train'})
    # print_mean_std(a_norm, axis=1)
    #
    # gamma = np.asarray([3.0,3.0,3.0])
    # beta = np.asarray([5.0,5.0,5.0])
    # print('After layer normalization (gamma=', gamma, ', beta=', beta, ')')
    # a_norm, _ = layernorm_forward(a, gamma, beta, {'mode': 'train'})
    # print_mean_std(a_norm, axis=1)


    ##-------------------------test layernorm backward: pass!---------------------------
    np.random.seed(231)
    N, D = 4,5
    x = 5*np.random.randn(N, D) + 12
    gamma = np.random.randn(D)
    beta = np.random.randn(D)
    dout = np.random.randn(N, D)

    ln_param = {}
    fx = lambda x: layernorm_forward(x, gamma, beta, ln_param)[0]
    fg = lambda a: layernorm_forward(x, a, beta, ln_param)[0]
    fb = lambda b: layernorm_forward(x, gamma, b, ln_param)[0]

    dx_num = eval_numerical_gradient_array(fx, x, dout)
    da_num = eval_numerical_gradient_array(fg, gamma.copy(), dout)
    db_num = eval_numerical_gradient_array(fb, beta.copy(), dout)

    _, cache = layernorm_forward(x, gamma, beta, ln_param)
    dx, dgamma, dbeta = layernorm_backward(dout, cache)

    print('dx error: ', rel_error(dx_num, dx))
    print('dgamma error: ', rel_error(da_num, dgamma))
    print('dbeta error: ', rel_error(db_num, dbeta))

##6/16/2018






