# main.py
import os
import torch
import numpy as np
import random
from gat_lstm import GAT_LSTM
from data_trainer import Trainer
from args import lstm_args_parser
from torch.nn import MSELoss
from data_preprocessor import DataPreprocessor

def set_seed(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def main():
    args = lstm_args_parser()
    set_seed(args.random_seed)
    
    device = torch.device(f"cuda:{args.device}" if torch.cuda.is_available() else "cpu")
    args.device = device  # 将device存入args，方便后续使用

    # 数据预处理
    data = DataPreprocessor.read_and_combine_csv_files(args)
    train_loader, val_loader, test_loader, _ = DataPreprocessor.load_and_process_data(args, data)
    
    # 初始化包含16个子模型的整体模型
    model = GAT_LSTM(args).to(device)

    # 初始化训练器和MSE损失函数
    trainer = Trainer(model, args, data)
    criterion = MSELoss()

    # 优化器：优化所有子模型的参数（联合训练）
    optimizer = torch.optim.Adam(
        model.parameters(),  # 整体模型参数
        lr=args.lr
    )
    scheduler = torch.optim.lr_scheduler.StepLR(
        optimizer,
        step_size=args.scheduler_step_size,
        gamma=args.scheduler_gamma
    )

    # 整体训练大模型（包含所有16个子模型）
    print("=== 开始训练包含16个子模型的整体模型 ===")
    trainer.train_full_model(
        train_loader,
        val_loader,
        optimizer,
        criterion,
        scheduler
    )

    # 保存包含所有16个子模型参数的整体模型
    trainer.save_model()
    print("\n=== 模型训练完成，已保存整体模型 ===")

    # 评估模型
    trainer.evaluate_model(test_loader, MSELoss())

if __name__ == "__main__":
    main()