from __future__ import print_function
import mxnet as mx
import mxnet.ndarray as nd
import time
import logging
from utils import *


def calc_potential(exe, params, label_name, noise_precision, prior_precision):
    exe.copy_params_from(params)
    exe.forward(is_train=False)
    ret = 0.0
    ret += (nd.norm(
        exe.outputs[0] - exe.arg_dict[label_name]).asscalar() ** 2) / 2.0 * noise_precision
    for v in params.values():
        ret += (nd.norm(v).asscalar() ** 2) / 2.0 * prior_precision
    return ret


def calc_grad(exe, exe_grads, params, X, Y, label_name=None, outgrad_f=None):
    exe.copy_params_from(params)
    exe.arg_dict['data'][:] = X
    if outgrad_f is None:
        exe.arg_dict[label_name][:] = Y
        exe.forward(is_train=True)
        exe.backward()
    else:
        exe.forward(is_train=True)
        exe.backward(outgrad_f(exe.outpus, Y))
    for k, v in exe_grads.items():
        v.wait_to_read()


def step_HMC(exe, exe_params, exe_grads, label_key, noise_precision, prior_precision, L=10,
             eps=1E-6):
    init_params = {k: v.copyto(v.context) for k, v in exe_params.items()}
    end_params = {k: v.copyto(v.context) for k, v in exe_params.items()}
    init_momentums = {k: mx.random.normal(0, 1, v.shape) for k, v in init_params.items()}
    end_momentums = {k: v.copyto(v.context) for k, v in init_momentums.items()}
    init_potential = calc_potential(exe, init_params, label_key, noise_precision, prior_precision)

    # 0. Calculate Initial Energy and Kinetic
    init_kinetic = sum([nd.sum(nd.square(momentum)) / 2.0
                        for momentum in init_momentums.values()]).asscalar()
    # 1. Make a half step for momentum at the beginning
    exe.copy_params_from(end_params)
    exe.forward(is_train=True)
    exe.backward()
    for k, v in exe_grads.items():
        v.wait_to_read()
    for k, momentum in end_momentums.items():
        momentum[:] = momentum - (eps / 2) * exe_grads[k]
    # 2. Alternate full steps for position and momentum
    for i in range(L):
        # 2.1 Full step for position
        for k, param in exe_params.items():
            param[:] = param + eps * end_momentums[k]
        # 2.2 Full step for the momentum, except at the end of trajectory we perform a half step
        exe.forward(is_train=True)
        exe.backward()
        for v in exe_grads.values():
            v.wait_to_read()
        if i != L - 1:
            for k, momentum in end_momentums.items():
                momentum[:] = momentum - eps * exe_grads[k]
        else:
            for k, momentum in end_momentums.items():
                # We should reverse the sign of the momentum at the end
                momentum[:] = -(momentum - eps / 2.0 * exe_grads[k])
    copy_param(exe, end_params)
    # 3. Calculate acceptance ratio and accept/reject the move
    end_potential = calc_potential(exe, end_params, label_key, noise_precision, prior_precision)
    end_kinetic = sum([nd.sum(nd.square(momentum)) / 2.0
                       for momentum in end_momentums.values()]).asscalar()
    # print init_potential, init_kinetic, end_potential, end_kinetic
    r = numpy.random.rand(1)
    if r < numpy.exp(-(end_potential + end_kinetic) + (init_potential + init_kinetic)):
        exe.copy_params_from(end_params)
        return end_params, 1
    else:
        exe.copy_params_from(init_params)
        return init_params, 0


def HMC(sym, data_inputs, X, Y, X_test, Y_test, sample_num,
        initializer=None, noise_precision=1 / 9.0, prior_precision=0.1,
        learning_rate=1E-6, L=10, dev=mx.gpu()):
    label_key = list(set(data_inputs.keys()) - set(['data']))[0]
    exe, exe_params, exe_grads, _ = get_executor(sym, dev, data_inputs, initializer)
    exe.arg_dict['data'][:] = X
    exe.arg_dict[label_key][:] = Y
    sample_pool = []
    accept_num = 0
    start = time.time()
    for i in range(sample_num):
        sample_params, is_accept = step_HMC(exe, exe_params, exe_grads, label_key, noise_precision,
                                            prior_precision, L, learning_rate)
        accept_num += is_accept

        if (i + 1) % 10 == 0:
            sample_pool.append(sample_params)
            if (i + 1) % 100000 == 0:
                end = time.time()
                print("Current Iter Num: %d" % (i + 1), "Time Spent: %f" % (end - start), "MSE:",
                      sample_test_regression(exe, X=X_test, Y=Y_test, sample_pool=sample_pool,
                                             minibatch_size=Y.shape[0],
                                             save_path='regression_HMC.txt'))
                start = time.time()
        exe.copy_params_from(sample_params)
    print('accept ratio', accept_num / float(sample_num))
    return sample_pool


def SGD(sym, data_inputs, X, Y, X_test, Y_test, total_iter_num,
        lr=None,
        lr_scheduler=None, prior_precision=1,
        out_grad_f=None,
        initializer=None,
        minibatch_size=100, dev=mx.gpu()):
    if out_grad_f is None:
        label_key = list(set(data_inputs.keys()) - set(['data']))[0]
    exe, params, params_grad, _ = get_executor(sym, dev, data_inputs, initializer)
    optimizer = mx.optimizer.create('sgd', learning_rate=lr,
                                    rescale_grad=X.shape[0] / minibatch_size,
                                    lr_scheduler=lr_scheduler,
                                    wd=prior_precision,
                                    arg_names=params.keys())
    updater = mx.optimizer.get_updater(optimizer)
    start = time.time()
    for i in range(total_iter_num):
        indices = numpy.random.randint(X.shape[0], size=minibatch_size)
        X_batch = X[indices]
        Y_batch = Y[indices]
        exe.arg_dict['data'][:] = X_batch
        if out_grad_f is None:
            exe.arg_dict[label_key][:] = Y_batch
            exe.forward(is_train=True)
            exe.backward()
        else:
            exe.forward(is_train=True)
            exe.backward(out_grad_f(exe.outputs, nd.array(Y_batch, ctx=dev)))
        for k in params:
            updater(k, params_grad[k], params[k])
        if (i + 1) % 500 == 0:
            end = time.time()
            print("Current Iter Num: %d" % (i + 1), "Time Spent: %f" % (end - start))
            sample_test_acc(exe, X=X_test, Y=Y_test, label_num=10, minibatch_size=100)
            start = time.time()
    return exe, params, params_grad


def SGLD(sym, X, Y, X_test, Y_test, total_iter_num,
         data_inputs=None,
         learning_rate=None,
         lr_scheduler=None, prior_precision=1,
         out_grad_f=None,
         initializer=None,
         minibatch_size=100, thin_interval=100, burn_in_iter_num=1000, task='classification',
         dev=mx.gpu()):
    if out_grad_f is None:
        label_key = list(set(data_inputs.keys()) - set(['data']))[0]
    exe, params, params_grad, _ = get_executor(sym, dev, data_inputs, initializer)
    optimizer = mx.optimizer.create('sgld', learning_rate=learning_rate,
                                    rescale_grad=X.shape[0] / minibatch_size,
                                    lr_scheduler=lr_scheduler,
                                    wd=prior_precision)
    updater = mx.optimizer.get_updater(optimizer)
    sample_pool = []
    start = time.time()
    for i in range(total_iter_num):
        indices = numpy.random.randint(X.shape[0], size=minibatch_size)
        X_batch = X[indices]
        Y_batch = Y[indices]
        exe.arg_dict['data'][:] = X_batch
        if out_grad_f is None:
            exe.arg_dict[label_key][:] = Y_batch
            exe.forward(is_train=True)
            exe.backward()
        else:
            exe.forward(is_train=True)
            exe.backward(out_grad_f(exe.outputs, nd.array(Y_batch, ctx=dev)))
        for k in params:
            updater(k, params_grad[k], params[k])
        if i < burn_in_iter_num:
            continue
        else:
            if 0 == (i - burn_in_iter_num) % thin_interval:
                if optimizer.lr_scheduler is not None:
                    lr = optimizer.lr_scheduler(optimizer.num_update)
                else:
                    lr = learning_rate
                sample_pool.append([lr, copy_param(exe)])
        if (i + 1) % 100000 == 0:
            end = time.time()
            if task == 'classification':
                print("Current Iter Num: %d" % (i + 1), "Time Spent: %f" % (end - start))
                test_correct, test_total, test_acc = \
                    sample_test_acc(exe, sample_pool=sample_pool, X=X_test, Y=Y_test, label_num=10,
                                    minibatch_size=minibatch_size)
                print("Test %d/%d=%f" % (test_correct, test_total, test_acc))
            else:
                print("Current Iter Num: %d" % (i + 1), "Time Spent: %f" % (end - start), "MSE:",
                      sample_test_regression(exe=exe, sample_pool=sample_pool,
                                             X=X_test,
                                             Y=Y_test, minibatch_size=minibatch_size,
                                             save_path='regression_SGLD.txt'))
            start = time.time()
    return exe, sample_pool


def DistilledSGLD(teacher_sym, student_sym,
                  teacher_data_inputs, student_data_inputs,
                  X, Y, X_test, Y_test, total_iter_num,
                  teacher_learning_rate, student_learning_rate,
                  teacher_lr_scheduler=None, student_lr_scheduler=None,
                  student_optimizing_algorithm='sgd',
                  teacher_grad_f=None, student_grad_f=None,
                  teacher_prior_precision=1, student_prior_precision=0.001,
                  perturb_deviation=0.001,
                  student_initializer=None,
                  teacher_initializer=None,
                  minibatch_size=100,
                  task='classification',
                  dev=mx.gpu()):
    teacher_exe, teacher_params, teacher_params_grad, _ = \
        get_executor(teacher_sym, dev, teacher_data_inputs, teacher_initializer)
    student_exe, student_params, student_params_grad, _ = \
        get_executor(student_sym, dev, student_data_inputs, student_initializer)
    if teacher_grad_f is None:
        teacher_label_key = list(set(teacher_data_inputs.keys()) - set(['data']))[0]
    if student_grad_f is None:
        student_label_key = list(set(student_data_inputs.keys()) - set(['data']))[0]
    teacher_optimizer = mx.optimizer.create('sgld',
                                            learning_rate=teacher_learning_rate,
                                            rescale_grad=X.shape[0] / float(minibatch_size),
                                            lr_scheduler=teacher_lr_scheduler,
                                            wd=teacher_prior_precision)
    student_optimizer = mx.optimizer.create(student_optimizing_algorithm,
                                            learning_rate=student_learning_rate,
                                            rescale_grad=1.0 / float(minibatch_size),
                                            lr_scheduler=student_lr_scheduler,
                                            wd=student_prior_precision)
    teacher_updater = mx.optimizer.get_updater(teacher_optimizer)
    student_updater = mx.optimizer.get_updater(student_optimizer)
    start = time.time()
    for i in range(total_iter_num):
        # 1.1 Draw random minibatch
        indices = numpy.random.randint(X.shape[0], size=minibatch_size)
        X_batch = X[indices]
        Y_batch = Y[indices]

        # 1.2 Update teacher
        teacher_exe.arg_dict['data'][:] = X_batch
        if teacher_grad_f is None:
            teacher_exe.arg_dict[teacher_label_key][:] = Y_batch
            teacher_exe.forward(is_train=True)
            teacher_exe.backward()
        else:
            teacher_exe.forward(is_train=True)
            teacher_exe.backward(
                teacher_grad_f(teacher_exe.outputs, nd.array(Y_batch, ctx=dev)))

        for k in teacher_params:
            teacher_updater(k, teacher_params_grad[k], teacher_params[k])

        # 2.1 Draw random minibatch and do random perturbation
        if task == 'classification':
            indices = numpy.random.randint(X.shape[0], size=minibatch_size)
            X_student_batch = X[indices] + numpy.random.normal(0,
                                                               perturb_deviation,
                                                               X_batch.shape).astype('float32')
        else:
            X_student_batch = mx.random.uniform(-6, 6, X_batch.shape, mx.cpu())

        # 2.2 Get teacher predictions
        teacher_exe.arg_dict['data'][:] = X_student_batch
        teacher_exe.forward(is_train=False)
        teacher_pred = teacher_exe.outputs[0]
        teacher_pred.wait_to_read()

        # 2.3 Update student
        student_exe.arg_dict['data'][:] = X_student_batch
        if student_grad_f is None:
            student_exe.arg_dict[student_label_key][:] = teacher_pred
            student_exe.forward(is_train=True)
            student_exe.backward()
        else:
            student_exe.forward(is_train=True)
            student_exe.backward(student_grad_f(student_exe.outputs, teacher_pred))
        for k in student_params:
            student_updater(k, student_params_grad[k], student_params[k])

        if (i + 1) % 2000 == 0:
            end = time.time()
            if task == 'classification':
                print("Current Iter Num: %d" % (i + 1), "Time Spent: %f" % (end - start))
                test_correct, test_total, test_acc = \
                    sample_test_acc(student_exe, X=X_test, Y=Y_test, label_num=10,
                                    minibatch_size=minibatch_size)
                train_correct, train_total, train_acc = \
                    sample_test_acc(student_exe, X=X, Y=Y, label_num=10,
                                    minibatch_size=minibatch_size)
                teacher_test_correct, teacher_test_total, teacher_test_acc = \
                    sample_test_acc(teacher_exe, X=X_test, Y=Y_test, label_num=10,
                                    minibatch_size=minibatch_size)
                teacher_train_correct, teacher_train_total, teacher_train_acc = \
                    sample_test_acc(teacher_exe, X=X, Y=Y, label_num=10,
                                    minibatch_size=minibatch_size)
                print("Student: Test ACC %d/%d=%f, Train ACC %d/%d=%f" % (test_correct, test_total,
                                                    test_acc, train_correct, train_total, train_acc))
                print("Teacher: Test ACC %d/%d=%f, Train ACC %d/%d=%f" \
                      % (teacher_test_correct, teacher_test_total, teacher_test_acc,
                         teacher_train_correct, teacher_train_total, teacher_train_acc))
            else:
                print("Current Iter Num: %d" % (i + 1), "Time Spent: %f" % (end - start), "MSE:",
                       sample_test_regression(exe=student_exe, X=X_test, Y=Y_test,
                                             minibatch_size=minibatch_size,
                                             save_path='regression_DSGLD.txt'))
            start = time.time()

    return student_exe, student_params, student_params_grad