import torch
import numpy as np
import matplotlib.pyplot as plt
from config import *
from dataset import get_dataloaders
from model import BuildingTCN

def evaluate_model():
    _, _, test_loader, scaler_X, scaler_Y = get_dataloaders()
    if test_loader is None:
        raise ValueError("当前数据配置下没有可用的测试集。")
    
    model = BuildingTCN(input_size=1, output_size=5, num_channels=CHANNELS, 
                        kernel_size=KERNEL_SIZE, dropout=DROPOUT).to(DEVICE)
    model.load_state_dict(torch.load('best_model.pth', map_location=DEVICE))
    model.eval()
    
    all_preds = []
    all_targets = []
    
    with torch.no_grad():
        for inputs, targets, masks in test_loader:
            inputs = inputs.to(DEVICE)
            outputs = model(inputs)
            
            all_preds.append(outputs.cpu().numpy())
            all_targets.append(targets.numpy())
            
    # Concatenate results
    all_preds = np.concatenate(all_preds, axis=0)
    all_targets = np.concatenate(all_targets, axis=0)

    # Inverse transform to physical scale for metrics
    all_preds_inv = scaler_Y.inverse_transform(all_preds.reshape(-1, len(OUTPUT_SENSORS))).reshape(all_preds.shape)
    all_targets_inv = scaler_Y.inverse_transform(all_targets.reshape(-1, len(OUTPUT_SENSORS))).reshape(all_targets.shape)

    print("\n=== Test Metrics (All Windows) ===")
    for i, sens in enumerate(OUTPUT_SENSORS):
        pred_i = all_preds_inv[:, :, i].reshape(-1)
        target_i = all_targets_inv[:, :, i].reshape(-1)
        mae = float(np.mean(np.abs(pred_i - target_i)))
        pred_std = float(np.std(pred_i))
        target_std = float(np.std(target_i))
        amp_ratio = pred_std / (target_std + 1e-12)
        pred_peak = float(np.max(np.abs(pred_i)))
        target_peak = float(np.max(np.abs(target_i)))
        peak_ratio = pred_peak / (target_peak + 1e-12)
        corr = float(np.corrcoef(pred_i, target_i)[0, 1]) if len(pred_i) > 1 else float("nan")

        print(
            f"{sens}: MAE={mae:.4f}, Corr={corr:.4f}, "
            f"AmpRatio(std_pred/std_true)={amp_ratio:.4f}, "
            f"PeakRatio(max|pred|/max|true|)={peak_ratio:.4f}"
        )
    
    # 取一个 batch 的第一条序列进行可视化
    sample_pred = all_preds_inv[0]
    sample_target = all_targets_inv[0]
    
    plt.figure(figsize=(15, 10))
    for i, sens in enumerate(OUTPUT_SENSORS):
        plt.subplot(5, 1, i+1)
        plt.plot(sample_target[:, i], label='True', alpha=0.7)
        plt.plot(sample_pred[:, i], label='Pred', alpha=0.7, linestyle='--')
        plt.title(f'Sensor {sens} Z-axis Response (Test: Earthquake)')
        plt.legend()
        
    plt.tight_layout()
    plt.savefig('test_results.png')
    print("Evaluation done. Result saved to test_results.png")

if __name__ == '__main__':
    evaluate_model()