import itertools
import numpy as np
import torch
from torch import optim
from torch import nn
import torch.random as rnd

iteration_count = 3
n = 10 # number of features
l = 5 # number of data points (aka batch size)

X_ = torch.randn(l, n)
Y_ = torch.randn(l, 1)

# An example implementation of ONNX Momentum when "nesterov" attribute is false.
def apply_momentum(t, r, x, g, v, norm_coefficient, alpha_, beta_):  # type: ignore
    # Add gradient of regularization term.
    g_regularized = norm_coefficient * x + g
    # Coefficient of gradient should be 1 at the first iteration.
    beta_adjusted = beta_ if t > 0 else 1
    # Update momentum.
    v_new = alpha_ * v + beta_adjusted * g_regularized
    # Apply SG with momentum update rule.
    x_new = x - r * v_new
    return x_new, v_new

# An example implementation of ONNX Momentum when "nesterov" attribute is true.
def apply_nesterov(t, r, x, g, v, norm_coefficient, alpha_, beta_):  # type: ignore
    # Add gradient of regularization term.
    g_regularized = norm_coefficient * x + g
    # Coefficient of gradient should be 1 at the first iteration.
    beta_adjusted = beta_ if t > 0 else 1
    v_new = alpha_ * v + beta_adjusted * g_regularized
    # Apply Nesterov with momentum update rule.
    x_new = x - r * (g_regularized + alpha_ * v_new)
    return x_new, v_new

def show_pytorch(lr, alpha, beta, weight_decay, nesterov):
    X = X_.clone()
    Y = Y_.clone()

    rnd.manual_seed(0)
    model = nn.Sequential(
        nn.Linear(n, 1, bias=False)
    )

    loss_fn = nn.MSELoss(reduction='sum')

    solver = optim.SGD(model.parameters(), lr=lr, momentum=alpha, dampening=1-beta,
            weight_decay=weight_decay, nesterov=nesterov)

    results = []
    for t in range(iteration_count):
        Y_pred = model(X)
        loss = loss_fn(Y_pred, Y)
        results.append(float(loss))
        model.zero_grad()
        loss.backward()
        solver.step()
    return results

def show_tensorflow(lr, alpha, nesterov):
    import tensorflow as tf
    rnd.manual_seed(0)
    layer = nn.Linear(n, 1, bias=False)

    X = tf.placeholder('float', shape=[l, n])
    W = tf.Variable(torch.Tensor.numpy(layer.weight.detach()))
    Y = tf.placeholder('float', shape=[l, 1])
    Y_pred = tf.matmul(X, W, transpose_b=True)
    loss = tf.reduce_sum(tf.square(Y - Y_pred))
    optimizer = tf.train.MomentumOptimizer(learning_rate=lr, momentum=alpha, use_nesterov=nesterov)
    minimizer = optimizer.minimize(loss)

    sess = tf.Session()
    init = tf.global_variables_initializer()
    sess.run(init)

    results = []
    for t in range(iteration_count):
        loss_value = sess.run(loss, {X: X_, Y: Y_})
        results.append(float(loss_value ))
        sess.run([minimizer], {X: X_, Y: Y_})

    return results


def show_mine(lr, alpha, beta, weight_decay, nesterov):
    X = X_.clone()
    Y = Y_.clone()

    rnd.manual_seed(0)
    model = nn.Sequential(
        nn.Linear(n, 1, bias=False)
    )

    loss_fn = nn.MSELoss(reduction='sum')

    solver = optim.SGD(model.parameters(), lr=lr, momentum=alpha, dampening=1-beta,
            weight_decay=weight_decay, nesterov=nesterov)

    results = []
    for t in range(iteration_count):
        Y_pred = model(X)
        loss = loss_fn(Y_pred, Y)
        results.append(float(loss))
        model.zero_grad()
        loss.backward()

        with torch.no_grad():
            for param in model.parameters():
                if 'momentum_buffer' not in solver.state[param]:
                    solver.state[param]['momentum_buffer'] = torch.zeros_like(param)
                
                if (nesterov):
                    new_tensor, new_state = apply_nesterov(t=t,
                            r=lr, x=param.data, g=param.grad.data,
                            v=solver.state[param]['momentum_buffer'].data,
                            norm_coefficient=weight_decay,
                            alpha_=alpha, beta_=beta)
                else:
                    new_tensor, new_state = apply_momentum(t=t,
                            r=lr, x=param.data, g=param.grad.data,
                            v=solver.state[param]['momentum_buffer'].data,
                            norm_coefficient=weight_decay,
                            alpha_=alpha, beta_=beta)
                solver.state[param]['momentum_buffer'].data = new_state.data
                param.data = new_tensor.data

    return results


# Compare TF SG and its ONNX counterpart.
lr_pool = [0.1, 0.001]
alpha_pool = [1.2, 0.8, 0.01]
beta_pool = [1]
weight_decay_pool = [0]
nesterov_pool = [True, False]

for lr_, alpha_, beta_, weight_decay_, nesterov_ in itertools.product(lr_pool, alpha_pool, beta_pool,
                                                                      weight_decay_pool, nesterov_pool):
    tf_result = show_tensorflow(lr=lr_, alpha=alpha_, nesterov=nesterov_)
    mine_result = show_mine(lr=lr_, alpha=alpha_, beta=beta_, weight_decay=weight_decay_, nesterov=nesterov_)
    assert np.allclose(tf_result, mine_result)


# Compare Pytorch SG with momentum and its ONNX counterpart.
lr_pool = [0.1, 0.001]
alpha_pool = [1.2, 0.8, 0.01]
beta_pool = [1, 0.2, 0.5]
weight_decay_pool = [0, 0.01, 0.1, 1]
nesterov_pool = [False]
for lr_, alpha_, beta_, weight_decay_, nesterov_ in itertools.product(lr_pool, alpha_pool, beta_pool,
                                                                      weight_decay_pool, nesterov_pool):
    torch_result = show_pytorch(lr=lr_, alpha=alpha_, beta=beta_, weight_decay=weight_decay_, nesterov=nesterov_)
    mine_result = show_mine(lr=lr_, alpha=alpha_, beta=beta_, weight_decay=weight_decay_, nesterov=nesterov_)
    assert np.allclose(torch_result, mine_result)


# Compare Pytorch SG with Nesterov momentum and its ONNX counterpart.
lr_pool = [0.1, 0.001]
alpha_pool = [1.2, 0.8, 0.01]
beta_pool = [1]
weight_decay_pool = [0, 0.01, 0.1, 1]
nesterov_pool = [True]
for lr_, alpha_, beta_, weight_decay_, nesterov_ in itertools.product(lr_pool, alpha_pool, beta_pool,
                                                                      weight_decay_pool, nesterov_pool):
    torch_result = show_pytorch(lr=lr_, alpha=alpha_, beta=beta_, weight_decay=weight_decay_, nesterov=nesterov_)
    mine_result = show_mine(lr=lr_, alpha=alpha_, beta=beta_, weight_decay=weight_decay_, nesterov=nesterov_)
    assert np.allclose(torch_result, mine_result)