import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn.functional as F

from models import MLP

TRAIN_RANGE = [-5, 5]
TEST_RANGE = [-20, 20]
LEARNING_RATE = 1e-2
NUM_ITERS = int(1e4)
NON_LINEARITIES = [
    'hardtanh', 'sigmoid',
    'relu6', 'tanh',
    'tanhshrink', 'hardshrink',
    'leakyrelu', 'softshrink',
    'softsign', 'relu',
    'prelu', 'softplus',
    'elu', 'selu',
]


def train(model, optimizer, data, num_iters):
    for i in range(num_iters):
        out = model(data)
        loss = F.mse_loss(out, data)
        mea = torch.mean(torch.abs(data - out))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if i % 1000 == 0:
            print("\t{}/{}: loss: {:.3f} - mea: {:.3f}".format(
                i+1, num_iters, loss.item(), mea.item())
            )


def test(model, data):
    with torch.no_grad():
        out = model(data)
        return torch.abs(data - out)


def main():
    save_dir = './imgs/'

    TRAIN_RANGE[-1] += 1
    TEST_RANGE[-1] += 1

    # datasets
    train_data = torch.arange(*TRAIN_RANGE).unsqueeze_(1).float()
    test_data = torch.arange(*TEST_RANGE).unsqueeze_(1).float()

    # train
    all_mses = []
    for non_lin in NON_LINEARITIES:
        print("Working with {}...".format(non_lin))
        mses = []
        for i in range(100):
            net = MLP(4, 1, 8, 1, non_lin)
            optim = torch.optim.RMSprop(net.parameters(), lr=LEARNING_RATE)
            train(net, optim, train_data, NUM_ITERS)
            mses.append(test(net, test_data))
        all_mses.append(torch.cat(mses, dim=1).mean(dim=1))
    all_mses = [x.numpy().flatten() for x in all_mses]

    # plot
    fig, ax = plt.subplots(figsize=(8, 7))
    x_axis = np.arange(-20, 21)
    for i, non_lin in enumerate(NON_LINEARITIES):
        ax.plot(x_axis, all_mses[i], label=non_lin)
    plt.grid()
    plt.legend(loc='best')
    plt.ylabel('Mean Absolute Error')
    plt.savefig(save_dir + 'extrapolation.png', format='png', dpi=300)
    plt.show()


if __name__ == '__main__':
    main()