"""模型评估与对比 — LSTM vs XGBoost 多时间尺度高温风险预测

功能:
    1. 加载测试数据 (与训练相同的 70/15/15 时间序划分)
    2. 加载已训练的 LSTM 模型，获取/计算预测结果
    3. 训练 XGBoost 基线并评估
    4. 生成混淆矩阵对比图和指标对比柱状图
    5. 打印格式化的模型对比表
"""

import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import accuracy_score, f1_score, confusion_matrix

from src.models.lstm_attention import HeatRiskPredictor
from src.models.xgboost_baseline import train_xgboost_baseline
from src.utils.config import DATA_PROCESSED, OUTPUT_MODELS, OUTPUT_FIGURES

# ============================================================================
# 全局常量
# ============================================================================

plt.rcParams["font.sans-serif"] = ["SimHei", "Microsoft YaHei", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False

RISK_LABELS = ["低", "中", "高", "严重"]
HORIZON_KEYS = ["short", "medium", "long"]
HORIZON_CN = ["短期", "中期", "长期"]


# ============================================================================
# 数据加载
# ============================================================================

def _load_all_data() -> tuple[np.ndarray, np.ndarray] | tuple[None, None]:
    """加载焦作和郑州的序列 NPZ 文件并沿样本轴拼接。

    Returns:
        (X, y) 或 (None, None) 当数据文件不存在时。
        X: (N, T, D) float32 特征数组
        y: (N, 3) int64 标签数组，列顺序: short, medium, long
    """
    cities = ["jiaozuo", "zhengzhou"]
    X_parts: list[np.ndarray] = []
    y_parts: list[np.ndarray] = []

    for city in cities:
        path = DATA_PROCESSED / f"{city}_sequences.npz"
        if path.exists():
            data = np.load(path)
            X_parts.append(data["X"])
            y_parts.append(data["y"])
            print(f"已加载 {city}: X {data['X'].shape}, y {data['y'].shape}")
        else:
            print(f"警告: {path} 不存在，跳过该城市")

    if not X_parts:
        print("错误: 未找到任何序列数据文件，无法进行评估。", file=sys.stderr)
        print(f"请确保 {DATA_PROCESSED / 'jiaozuo_sequences.npz'} "
              f"或 {DATA_PROCESSED / 'zhengzhou_sequences.npz'} 存在。", file=sys.stderr)
        return None, None

    X = np.concatenate(X_parts, axis=0)
    y = np.concatenate(y_parts, axis=0)
    print(f"合并后数据: X {X.shape}, y {y.shape}")
    return X, y


def load_test_data() -> tuple[np.ndarray, np.ndarray] | tuple[None, None]:
    """加载测试集数据 (与训练相同的 70/15/15 时间序划分)。

    Returns:
        (X_test, y_test) 或 (None, None) 当数据不可用时。
        X_test: (N_test, T, D) float32
        y_test: (N_test, 3) int64，列顺序: short, medium, long
    """
    X, y = _load_all_data()
    if X is None:
        return None, None

    n_total = len(X)
    n_train = int(n_total * 0.70)
    n_val = int(n_total * 0.15)

    X_test = X[n_train + n_val:]
    y_test = y[n_train + n_val:]

    print(f"测试集: X_test {X_test.shape}, y_test {y_test.shape}")
    return X_test, y_test


# ============================================================================
# LSTM 评估
# ============================================================================

def evaluate_lstm(
    X_test: np.ndarray,
    y_test: np.ndarray,
) -> tuple[dict[str, np.ndarray], np.ndarray] | tuple[None, None]:
    """加载已训练的 HeatRiskPredictor 并获取测试集预测。

    优先从 test_predictions.npz 加载已保存的预测结果；
    若文件不存在，则加载模型权重并重新推理。

    Args:
        X_test: (N, T, D) 测试集特征
        y_test: (N, 3) 测试集标签

    Returns:
        (predictions_dict, labels_array) 或 (None, None) 当模型不可用时。
        predictions_dict: {"short": (N,) int, "medium": (N,) int, "long": (N,) int}
        labels_array: (N, 3) int64，与 y_test 相同
    """
    pred_path = OUTPUT_MODELS / "test_predictions.npz"
    model_path = OUTPUT_MODELS / "best_model.pt"

    # ---- 路径 1: 从已保存的预测文件加载 ----
    if pred_path.exists():
        print(f"从 {pred_path} 加载 LSTM 预测结果...")
        data = np.load(pred_path)

        # 兼容 train.py 的 "y_true" 和任务约定的 "labels" 两种键名
        labels_key = "labels" if "labels" in data else "y_true"

        predictions = {
            "short": data["short"],
            "medium": data["medium"],
            "long": data["long"],
        }
        labels = data[labels_key]

        # 验证样本数匹配
        if len(predictions["short"]) != len(y_test):
            print(f"警告: 已保存预测样本数 ({len(predictions['short'])}) "
                  f"与测试集样本数 ({len(y_test)}) 不一致，将重新推理。")
        else:
            print(f"已加载 LSTM 预测: {len(predictions['short'])} 样本")
            return predictions, labels

    # ---- 路径 2: 重新推理 ----
    if not model_path.exists():
        print(f"错误: 模型文件 {model_path} 不存在，无法评估 LSTM。", file=sys.stderr)
        return None, None

    print(f"从 {model_path} 加载模型权重...")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")

    checkpoint = torch.load(model_path, map_location=device, weights_only=False)
    input_dim = checkpoint.get("input_dim", X_test.shape[2])

    model = HeatRiskPredictor(input_dim=input_dim).to(device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    # 批量推理
    X_tensor = torch.tensor(X_test, dtype=torch.float32).to(device)
    # 为避免 OOM，按 batch 处理
    batch_size = 128
    all_preds: dict[str, list[int]] = {"short": [], "medium": [], "long": []}

    with torch.no_grad():
        for i in range(0, len(X_tensor), batch_size):
            batch = X_tensor[i : i + batch_size]
            outputs = model(batch)
            for key in HORIZON_KEYS:
                preds = outputs[key].argmax(dim=1).cpu().numpy()
                all_preds[key].extend(preds.tolist())

    predictions = {k: np.array(v, dtype=np.int64) for k, v in all_preds.items()}
    print(f"LSTM 推理完成: {len(predictions['short'])} 样本")
    return predictions, y_test


# ============================================================================
# 辅助: 计算指标
# ============================================================================

def _compute_metrics(
    predictions: dict[str, np.ndarray],
    y_true: np.ndarray,
) -> dict[str, dict[str, float]]:
    """从预测和真实标签计算各时间尺度的 accuracy 和 macro F1。

    Args:
        predictions: {"short": (N,), "medium": (N,), "long": (N,)} 预测标签
        y_true: (N, 3) 真实标签，列顺序: short, medium, long

    Returns:
        {"short": {"accuracy": ..., "f1_macro": ...}, ...}
    """
    metrics: dict[str, dict[str, float]] = {}
    for i, key in enumerate(HORIZON_KEYS):
        y_pred = predictions[key]
        y_t = y_true[:, i]
        metrics[key] = {
            "accuracy": float(accuracy_score(y_t, y_pred)),
            "f1_macro": float(f1_score(y_t, y_pred, average="macro")),
        }
    return metrics


# ============================================================================
# 绘图: 混淆矩阵对比
# ============================================================================

def plot_confusion_matrices(
    lstm_preds: dict[str, np.ndarray],
    xgb_preds: dict[str, np.ndarray],
    y_true: np.ndarray,
    save_path: str | Path | None = None,
) -> None:
    """绘制 2x3 混淆矩阵对比图 (LSTM 行 / XGBoost 行; 短/中/长期 列)。

    Args:
        lstm_preds: LSTM 预测 {"short", "medium", "long"}
        xgb_preds: XGBoost 预测 {"short", "medium", "long"}
        y_true: (N, 3) 真实标签
        save_path: 保存路径，默认 OUTPUT_FIGURES / "confusion_matrix_comparison.png"
    """
    if save_path is None:
        save_path = OUTPUT_FIGURES / "confusion_matrix_comparison.png"
    OUTPUT_FIGURES.mkdir(parents=True, exist_ok=True)

    fig, axes = plt.subplots(2, 3, figsize=(16, 10))
    model_names = ["LSTM", "XGBoost"]
    preds_list = [lstm_preds, xgb_preds]

    for row_idx, (model_name, preds) in enumerate(zip(model_names, preds_list)):
        for col_idx, key in enumerate(HORIZON_KEYS):
            ax = axes[row_idx, col_idx]
            cm = confusion_matrix(
                y_true[:, col_idx], preds[key],
                labels=[0, 1, 2, 3],
            )
            im = ax.imshow(cm, cmap="Blues", interpolation="nearest")

            # 在每个单元格内标注数值
            for i in range(4):
                for j in range(4):
                    ax.text(
                        j, i, str(cm[i, j]),
                        ha="center", va="center",
                        fontsize=9,
                        color="white" if cm[i, j] > cm.max() / 2 else "black",
                    )

            ax.set_xticks(range(4))
            ax.set_yticks(range(4))
            ax.set_xticklabels(RISK_LABELS)
            ax.set_yticklabels(RISK_LABELS)

            title = f"{model_name} - {HORIZON_CN[col_idx]}"
            ax.set_title(title, fontsize=13, fontweight="bold")
            ax.set_xlabel("预测标签")
            ax.set_ylabel("真实标签")

    # 共享 colorbar
    fig.colorbar(
        axes[0, 2].get_images()[0],
        ax=axes[:, -1],
        shrink=0.8,
        label="样本数",
    )

    fig.suptitle("混淆矩阵对比: LSTM vs XGBoost", fontsize=15, fontweight="bold", y=1.01)
    plt.tight_layout()
    fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"混淆矩阵对比图已保存至 {save_path}")


# ============================================================================
# 绘图: 指标对比柱状图
# ============================================================================

def plot_metrics_comparison(
    lstm_metrics: dict[str, dict[str, float]],
    xgb_metrics: dict[str, dict[str, float]],
    save_path: str | Path | None = None,
) -> None:
    """绘制 1x3 指标对比柱状图 (每个时间尺度两柱: accuracy 和 F1)。

    Args:
        lstm_metrics: LSTM 指标 {"short": {"accuracy", "f1_macro"}, ...}
        xgb_metrics: XGBoost 指标 {"short": {"accuracy", "f1_macro"}, ...}
        save_path: 保存路径，默认 OUTPUT_FIGURES / "model_comparison.png"
    """
    if save_path is None:
        save_path = OUTPUT_FIGURES / "model_comparison.png"
    OUTPUT_FIGURES.mkdir(parents=True, exist_ok=True)

    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    metric_keys = ["accuracy", "f1_macro"]
    metric_cn = ["Accuracy", "F1 Macro"]
    bar_width = 0.35
    colors = ["#4C72B0", "#DD8452"]  # LSTM, XGBoost

    for col_idx, horizon_key in enumerate(HORIZON_KEYS):
        ax = axes[col_idx]

        lstm_vals = [lstm_metrics[horizon_key][m] for m in metric_keys]
        xgb_vals = [xgb_metrics[horizon_key][m] for m in metric_keys]

        x_pos = np.arange(len(metric_keys))
        bars1 = ax.bar(
            x_pos - bar_width / 2, lstm_vals, bar_width,
            label="LSTM", color=colors[0], edgecolor="white",
        )
        bars2 = ax.bar(
            x_pos + bar_width / 2, xgb_vals, bar_width,
            label="XGBoost", color=colors[1], edgecolor="white",
        )

        # 在柱顶标注数值
        for bar in bars1:
            height = bar.get_height()
            ax.text(
                bar.get_x() + bar.get_width() / 2, height + 0.01,
                f"{height:.3f}", ha="center", va="bottom", fontsize=9,
            )
        for bar in bars2:
            height = bar.get_height()
            ax.text(
                bar.get_x() + bar.get_width() / 2, height + 0.01,
                f"{height:.3f}", ha="center", va="bottom", fontsize=9,
            )

        ax.set_xticks(x_pos)
        ax.set_xticklabels(metric_cn)
        ax.set_ylim(0, 1.15)
        ax.set_title(HORIZON_CN[col_idx], fontsize=13, fontweight="bold")
        ax.set_ylabel("分数")
        ax.legend(loc="lower right")
        ax.grid(axis="y", alpha=0.3)

    fig.suptitle("模型指标对比: LSTM vs XGBoost", fontsize=15, fontweight="bold")
    plt.tight_layout()
    fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"模型对比图已保存至 {save_path}")


# ============================================================================
# 主评估函数
# ============================================================================

def evaluate() -> None:
    """执行完整的模型评估与对比流水线。

    流程:
        1. 加载数据并按 70/15/15 划分
        2. LSTM 评估 (加载最佳模型 + 预测)
        3. XGBoost 基线训练与评估
        4. 打印对比表
        5. 生成对比图 (混淆矩阵 + 指标柱状图)
    """
    print("=" * 60)
    print("模型评估与对比")
    print("=" * 60)

    # ---- 1. 加载数据 ----
    X, y = _load_all_data()
    if X is None:
        return

    n_total = len(X)
    n_train = int(n_total * 0.70)
    n_val = int(n_total * 0.15)

    X_train = X[:n_train]
    y_train = y[:n_train]
    X_test = X[n_train + n_val:]
    y_test = y[n_train + n_val:]

    print(f"数据划分: 训练 {len(X_train)}, 验证 {n_val} (保留), 测试 {len(X_test)}")

    # ---- 2. LSTM 评估 ----
    print("\n--- LSTM 评估 ---")
    lstm_result = evaluate_lstm(X_test, y_test)
    if lstm_result[0] is None:
        lstm_preds = None
        lstm_metrics = None
    else:
        lstm_preds, _ = lstm_result
        lstm_metrics = _compute_metrics(lstm_preds, y_test)
        for key in HORIZON_KEYS:
            m = lstm_metrics[key]
            print(f"  LSTM {HORIZON_CN[HORIZON_KEYS.index(key)]}: "
                  f"Accuracy={m['accuracy']:.4f}, F1 Macro={m['f1_macro']:.4f}")

    # ---- 3. XGBoost 基线 ----
    print("\n--- XGBoost 基线 ---")
    try:
        xgb_results = train_xgboost_baseline(X_train, y_train, X_test, y_test)
        xgb_preds = {
            key: xgb_results[key]["predictions"] for key in HORIZON_KEYS
        }
        xgb_metrics = {
            key: {
                "accuracy": xgb_results[key]["accuracy"],
                "f1_macro": xgb_results[key]["f1_macro"],
            }
            for key in HORIZON_KEYS
        }
    except Exception as e:
        print(f"警告: XGBoost 训练失败: {e}", file=sys.stderr)
        xgb_preds = None
        xgb_metrics = None

    # ---- 4. 打印对比表 ----
    if lstm_metrics is not None and xgb_metrics is not None:
        print("\n" + "=" * 56)
        print("=== 模型对比 ===")
        print(f"{'时间尺度':<10} {'指标':<12} {'LSTM':<12} {'XGBoost':<12}")
        print("-" * 56)
        for key, cn_name in zip(HORIZON_KEYS, HORIZON_CN):
            for metric_key, metric_cn_name in [
                ("accuracy", "accuracy"),
                ("f1_macro", "f1_macro"),
            ]:
                lstm_val = lstm_metrics[key][metric_key]
                xgb_val = xgb_metrics[key][metric_key]
                print(f"{cn_name:<10} {metric_cn_name:<12} "
                      f"{lstm_val:<12.4f} {xgb_val:<12.4f}")
        print("=" * 56)

    # ---- 5. 生成对比图 ----
    if lstm_preds is not None and xgb_preds is not None:
        print("\n生成对比图表...")
        plot_confusion_matrices(lstm_preds, xgb_preds, y_test)
        plot_metrics_comparison(lstm_metrics, xgb_metrics)
    else:
        if lstm_preds is None:
            print("跳过图表: LSTM 预测不可用", file=sys.stderr)
        if xgb_preds is None:
            print("跳过图表: XGBoost 预测不可用", file=sys.stderr)

    print("\n评估完成。")


# ============================================================================
# CLI 入口
# ============================================================================

if __name__ == "__main__":
    evaluate()
