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)

def apply_adam(t, r, x, g, v, h, norm_coefficient, alpha, beta, epsilon):  # type: ignore
    # Add gradient of regularization term.
    g_regularized = norm_coefficient * x + g
    # Update momentum.
    v_new = alpha * v + (1 - alpha) * g_regularized
    # Update second-order momentum.
    h_new = beta * h + (1 - beta) * (g_regularized * g_regularized)
    # Compute element-wise square root.
    h_sqrt = np.sqrt(h_new) + epsilon
    # Adjust learning rate.
    t = t + 1
    r_adjusted = r * np.sqrt(1 - beta**t) / (1 - alpha**t)
    # Apply Adam update rule.
    x_new = x - r_adjusted * (v_new / h_sqrt)
    return x_new, v_new, h_new

def show_pytorch(lr, alpha, beta, weight_decay, epsilon):
    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.Adam(model.parameters(), lr=lr, betas=[alpha, beta],
            weight_decay=weight_decay, eps=epsilon)

    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, beta, epsilon):
    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.AdamOptimizer(learning_rate=lr, beta1=alpha, beta2=beta, epsilon=epsilon)
    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, epsilon):
    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.Adam(model.parameters(), lr=lr, betas=[alpha, beta],
            weight_decay=weight_decay, eps=epsilon)

    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 'exp_avg' not in solver.state[param]:
                    solver.state[param]['exp_avg'] = torch.zeros_like(param)
                if 'exp_avg_sq' not in solver.state[param]:
                    solver.state[param]['exp_avg_sq'] = torch.zeros_like(param)

                new_tensor, new_v, new_h = apply_adam(t=t,
                        r=lr, x=param.data, g=param.grad.data,
                        v=solver.state[param]['exp_avg'].data,
                        h=solver.state[param]['exp_avg_sq'].data,
                        norm_coefficient=weight_decay,
                        alpha=alpha, beta=beta, epsilon=epsilon)

                solver.state[param]['exp_avg'].data = new_v.data
                solver.state[param]['exp_avg_sq'].data = new_h.data
                param.data = new_tensor.data

    return results


# Compare TF Adam and its ONNX counterpart.
lr_pool = [0.1, 0.001]
alpha_pool = [0.89, 0.8, 0.01]
beta_pool = [0.99]
weight_decay_pool = [0]
epsilon_pool = [1e-7, 1e-1]
for lr_, alpha_, beta_, weight_decay_, epsilon_ in itertools.product(lr_pool, alpha_pool, beta_pool, weight_decay_pool, epsilon_pool):
    tf_result = show_tensorflow(lr=lr_, alpha=alpha_, beta=beta_, epsilon=epsilon_)
    mine_result = show_mine(lr=lr_, alpha=alpha_, beta=beta_, weight_decay=weight_decay_, epsilon=epsilon_)
    assert np.allclose(tf_result, mine_result)

# Compare Pytorch Adam and its ONNX counterpart.
lr_pool = [0.1, 0.001]
alpha_pool = [0.89, 0.8, 0.01]
beta_pool = [0.99, 0.2, 0.5]
weight_decay_pool = [0, 0.01, 0.1, 1]
epsilon_pool = [1e-7, 1e-1]
for lr_, alpha_, beta_, weight_decay_, epsilon_ in itertools.product(lr_pool, alpha_pool, beta_pool, weight_decay_pool, epsilon_pool):
    torch_result = show_pytorch(lr=lr_, alpha=alpha_, beta=beta_, weight_decay=weight_decay_, epsilon=epsilon_)
    mine_result = show_mine(lr=lr_, alpha=alpha_, beta=beta_, weight_decay=weight_decay_, epsilon=epsilon_)
    assert np.allclose(torch_result, mine_result)