import torch
import matplotlib.pyplot as plt
import numpy as np
from torch import nn
from tqdm import tqdm


# モデルの定義
class MyModel(nn.Sequential):
    def __init__(self):
        super().__init__(
            nn.Linear(1, 3),
            nn.ReLU(),
            nn.Linear(3, 1),
        )
        
# 学習するパラメータの出力
def print_params(model: MyModel):
    p = []
    p_label = [
        'a1', 'a2', 'a3', 'b1', 'b2', 'b3',
        'a4', 'a5', 'a6', 'b4'
    ]
    for params in model.parameters():
        for para in params.flatten():
            p.append(para.item())
            
    for para, label in zip(p, p_label):
        print(f'{label}: {para:.3e}', end=', ')
    print()

# 学習
def train():
    torch.manual_seed(42)
    x_axis = np.arange(1, 100, 0.01)

    # データセットの準備
    x_axis_shuffled = np.random.permutation(x_axis)
    y_axis_shuffled = x_axis_shuffled * x_axis_shuffled

    x_train, x_test = x_axis_shuffled[:8000], x_axis_shuffled[8000:]
    y_train, y_test = y_axis_shuffled[:8000], y_axis_shuffled[8000:]

    x_dataset = torch.from_numpy(x_train).float()
    y_dataset = torch.from_numpy(y_train).float()
    
    # モデルの準備
    model = MyModel()

    # パラメータの初期値の出力
    print_params(model)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    criterion = nn.MSELoss()
    num_epochs = 10

    losses = []
    for epoch in tqdm(range(num_epochs)):
        losses_each_epoch = []
    
        # 1 epoch 分の学習
        model.train()
        for x, y in zip(x_dataset, y_dataset):
            out = model(x.unsqueeze(0))
            loss = criterion(out, y.unsqueeze(0))
            losses_each_epoch.append(loss.item())

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # 1 epoch 分の検証
        model.eval()
        with torch.no_grad():
            y_pred = model(torch.from_numpy(x_test).float().unsqueeze(0).T)
            y_pred = y_pred.numpy().flatten()
            s_x_test, s_y_test, s_y_pred = zip(*sorted(zip(x_test, y_test, y_pred)))
        
        plt.plot(s_x_test, s_y_pred, 'c-',
                 label=f'prediction (epoch {epoch + 1})',
                 alpha=(epoch + 1) / num_epochs)
        if epoch == num_epochs - 1:
            plt.plot(s_x_test, s_y_test, 'm-', label='ground truth')

        print_params(model)
        losses.append(np.mean(losses_each_epoch))
        
    plt.legend()
    plt.show()

    # Epoch ごとの Loss の出力
    plt.plot(range(num_epochs), losses, label='loss')
    plt.show()

    # 結果の表示
    model.eval()
    with torch.no_grad():
        y_pred = model(torch.from_numpy(x_test).float().unsqueeze(0).T)
        y_pred = y_pred.numpy().flatten()

        # sort x_test and y_test and y_pred
        x_test, y_test, y_pred = zip(*sorted(zip(x_test, y_test, y_pred)))

        print(list(zip(x_test, y_pred)))

        plt.plot(x_test, y_pred, label='prediction')
        plt.plot(x_test, y_test, label='ground truth')
        plt.legend()
        plt.show()


if __name__ == '__main__':
    train()