temp

import os
import random
import argparse
import numpy as np
import pandas as pd
from tqdm import tqdm
from datetime import datetime
import matplotlib.pyplot as plt
import torch
import torch.nn as nn

from src.models import PREDICTION_MODELS_MAP
from src.utils import set_logging, get_data_file_path
from src.dataset import get_dataloader_prediction


def parse_args():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Training Configuration')
    parser.add_argument('--data_train', type=str, default='./data/prediction_train_data.h5', help='Train Data directory')
    parser.add_argument('--data_test', type=str, default='./data/prediction_test_data.h5', help='Test Data directory')
    parser.add_argument('--results_dir', type=str, default='./results/save_run_001', help='Results directory')
    parser.add_argument('--models',
                        type=str,
                        nargs='+',
                        default=list(PREDICTION_MODELS_MAP.keys(),),
                        choices=list(PREDICTION_MODELS_MAP.keys()),)
    parser.add_argument('--batch_size', type=int, default=16, help='Batch size')
    parser.add_argument('--device', type=str, default='cuda:1', help='Device to use for training')
    args = parser.parse_args()
    
    config = vars(args)
    result_dir = config['results_dir']
    if not os.path.exists(result_dir):
        print(f"No such directory: {result_dir}!")
        exit()
    result_dir_eval = os.path.join(result_dir, 'Eval')
    os.makedirs(result_dir_eval, exist_ok=True)
    config['result_dir_eval'] = result_dir_eval
    global LOGGER
    LOGGER = set_logging(os.path.join(config['result_dir_eval'], 'eval.log'))
    LOGGER.info(datetime.now().strftime('%Y-%m-%d %H:%M'))
    for key, value in config.items():
        LOGGER.info(f'{key}: {value}')
    return config


def run(config):
    data_train, data_test = config['data_train'], config['data_test']
    batch_size = config['batch_size']
    models = config['models']
    device = config['device']
    results_dir = config['results_dir']
    result_dir_eval = config['result_dir_eval']

    dataloader_train, dataloader_test, dataset_train, dataset_test = get_dataloader_prediction(data_train, data_test, batch_size)
    model_results = []
    for model_name in models:
        LOGGER.info(f'\nStart evaluating model: {model_name}')
        result_dir_model = os.path.join(results_dir, model_name)
        model_class = PREDICTION_MODELS_MAP[model_name]
        model = model_class().to(device)
        model_path = os.path.join(result_dir_model, f'{model_name}_best_model.pth')
        model.load_state_dict(torch.load(model_path, map_location=device))
        LOGGER.info(f'Loaded model weights from {model_path}')
        train_metrics = model_evaluate(model, dataloader_train, dataset_train, device)
        val_metrics = model_evaluate(model, dataloader_test, dataset_test, device)
        LOGGER.info(f"Train Metrics for {model_name}:")
        for metric_name, metric_value in train_metrics.items():
            LOGGER.info(f"  {metric_name}: {metric_value:.8f}")
        LOGGER.info(f"Test Metrics for {model_name}:")
        for metric_name, metric_value in val_metrics.items():
            LOGGER.info(f"  {metric_name}: {metric_value:.8f}")
        train_metrics['model'] = model_name
        train_metrics['state'] = 'train'
        val_metrics['model'] = model_name
        val_metrics['state'] = 'test'
        model_results.append(train_metrics)
        model_results.append(val_metrics)
        LOGGER.info(f"Model {model_name} evaluating completed.\n")
    df = pd.DataFrame(model_results)
    results_csv_path = os.path.join(result_dir_eval, 'evaluation_results.csv')
    df.to_csv(results_csv_path, index=False, float_format='%.6f')


def model_evaluate(model, dataloader_test, dataset_test, device):
    model.eval()
    preds_list, targets_list = [], []
    val_metrics = {}
    with torch.no_grad():
        for X, Y, _ in tqdm(dataloader_test,
                         desc='Evaluating ',
                         total=len(dataloader_test),
                         dynamic_ncols=True):
            X, Y = X.to(device), Y.to(device)
            output = model(X)
            preds_list.append(output.cpu().numpy())
            targets_list.append(Y.cpu().numpy())
        ## MAE, MSE, MAPE
        preds = np.concatenate(preds_list, axis=0)
        targets = np.concatenate(targets_list, axis=0)
        preds = dataset_test.denormalize(preds)
        targets = dataset_test.denormalize(targets)
        bands = [7, 13, 14]
        for bi in range(3):
            band_mae = np.mean(np.abs(preds[:, bi] - targets[:, bi]))
            band_mse = np.mean((preds[:, bi] - targets[:, bi]) ** 2)
            band_mape = np.mean(np.abs((preds[:, bi] - targets[:, bi]) / (targets[:, bi] + 1e-6))) * 100.0
            val_metrics.update({
                f'mae_{bands[bi]}': band_mae,
                f'mse_{bands[bi]}': band_mse,
                f'mape_{bands[bi]}': band_mape,
            })
        for metric_name in ['mae', 'mse', 'mape']:
            val_metrics[f'{metric_name}_mean'] = np.mean([val_metrics[f'{metric_name}_{b}'] for b in bands])
    return val_metrics


if __name__ == '__main__':
    config = parse_args()
    run(config)
    print("ok")