import random
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import pandas as pd
from dataset.TimeDataset import TimeSeriesDataset
#from model.Transformer_base import TimeSeriesTransformer
from model.LSTM import TimeSeriesTransformer
from tqdm import tqdm
from utils.Arg import Arg
from utils import ModeTest
import matplotlib.pyplot as plt
import training_model

arg = Arg()
# 超参数
input_dim = arg.input_dim
output_dim = arg.output_dim
input_seq_length = arg.input_seq_length
output_seq_length = arg.output_seq_length
d_model = arg.d_model
nhead = arg.nhead
num_layers = arg.num_layers
dropout = arg.dropout
batch_size = arg.batch_size
epochs = arg.epochs
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def setup_seed(seed):
    torch.manual_seed( seed )
    torch.cuda.manual_seed_all( seed )
    np.random.seed( seed )
    random.seed( seed )

def plot_fig(time,out,tag):
    plt.plot(time,out)
    plt.plot(time,tag)
    plt.show()

def train(model_use = False):
    #model = TimeSeriesTransformer(input_dim, output_dim, d_model, nhead, num_layers, dropout).to(device)
    model = TimeSeriesTransformer()
    #model = torch.compile(model, mode="reduce-overhead")
    if model_use:
        print("载入历史训练的模型")
        model.load_state_dict(torch.load('./save/best_loss_short_encoder.pt'))
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    criterion = nn.MSELoss()#nn.L1Loss()#nn.MSELoss()
    best_loss = float('inf')

    model.train()
    for epoch in range(epochs):
        epoch_loss = 0.0
        data_len = 0
        for i in tqdm(range(6), desc='Training progress:'):
            if i == 4 :
                continue
            file_inputs_2 = './data/Dataset_training/NWP/NWP_{}.csv'.format(i)
            file_inputs_3 = './data/Dataset_training/power/power_{}.csv'.format(i)
            dataset = TimeSeriesDataset(file_inputs_3, file_inputs_2)
            dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
            # 训练模型
            data_loss = 0.0
            for batch_idx, (input_seq, output_seq) in enumerate(dataloader):
                input_seq, output_seq = input_seq.to(device), output_seq.to(device)
                # 前向传播
                # input_seq = input_seq.permute(1, 0, 2)
                # tgt = input_seq[-1:]
                #predictions = model(input_seq,tgt)
                predictions = model(input_seq)
                # 计算损失
                loss = criterion(predictions, output_seq)
                # 反向传播
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                data_loss += loss.item()
            data_len += len(dataloader)
            epoch_loss += data_loss
            #print(f"Datasate is {i} ,Loss is {data_loss/data_len}")
        print(f"Epoch {epoch+1}, Loss: {epoch_loss / data_len}")
        # 保存 模型
        if epoch_loss < best_loss:
            best_loss = epoch_loss
            print("Best loss model is saved")
        torch.save(model.state_dict(), './save/best_loss_short_encoder.pt')




if __name__ == '__main__':
    setup_seed(50)
    model_use = True
    model = TimeSeriesTransformer()
    #model = torch.compile(model, mode="reduce-overhead")
    if model_use:
        print("载入历史训练的模型")
        model.load_state_dict(torch.load('save/lstm_base.pt'))

    #training_model.base_train(model)


    # re_train_for_data(model, 2022,11)
    # re_train_for_data(model, 2022,9)
    #train(model_use = True）
    #training_model.re_train_for_data(model, 2022, 5)
    #training_model.re_train_for_data(model, 2022, 10)

    # for i in [5]:
    #     #training_model.re_train_for_data(model, 2022,11)
    #     training_model.re_train_for_data(model,2023,4)
    #     ModeTest.test_model(2023,i,"lstm_base_pro.pt") #69
    #     break

    for i in range(4,5):
        #training_model.re_train_for_turbine_sum_power(model)
        #training_model.re_train_for_data(model, 2023,1)
        #training_model.re_train_for_data(model, 2022,1)
        #training_model.re_train_for_data(model, 2023, 2)
        ModeTest.test_model(2022, i, "lstm_base_pro.pt")