"""LSTM-Attention 模型训练脚本

完整的训练流水线：数据加载、训练、验证、早停、测试评估。
"""

import json
import sys
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, f1_score
from torch.optim import AdamW
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader, TensorDataset

from src.models.lstm_attention import HeatRiskPredictor
from src.utils.config import (
    BATCH_SIZE,
    DATA_PROCESSED,
    EARLY_STOP_PATIENCE,
    LEARNING_RATE,
    MAX_EPOCHS,
    OUTPUT_LOGS,
    OUTPUT_MODELS,
)


class FocalLoss(nn.Module):
    """Focal Loss — 聚焦困难样本"""

    def __init__(self, alpha: float = 0.5, gamma: float = 2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        ce = F.cross_entropy(logits, targets, reduction="none")
        pt = torch.exp(-ce)
        focal = self.alpha * (1 - pt) ** self.gamma * ce
        return focal.mean()


def load_data() -> tuple[np.ndarray, np.ndarray]:
    """加载焦作和郑州的序列数据，拼接后返回。

    Returns:
        X: (N, 14, F) 特征数组
        y: (N, 3) 标签数组，列顺序: short, medium, long (0-3 等级)

    Raises:
        FileNotFoundError: 两个城市的 npz 文件均不存在时
    """
    cities = ["jiaozuo", "zhengzhou"]
    X_parts, y_parts = [], []

    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:
        msg = (
            "未找到任何序列数据文件。\n"
            f"请确保 {DATA_PROCESSED / 'jiaozuo_sequences.npz'} 或 "
            f"{DATA_PROCESSED / 'zhengzhou_sequences.npz'} 存在。\n"
            "运行数据预处理流水线以生成这些文件。"
        )
        print(msg, file=sys.stderr)
        sys.exit(1)

    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 _compute_metrics(
    outputs: dict[str, torch.Tensor], labels: torch.Tensor
) -> dict[str, dict[str, float]]:
    """计算每个预测头的 loss、accuracy 和 macro F1。

    Args:
        outputs: 模型前向输出，{"short": (B,4), "medium": (B,4), "long": (B,4)}
        labels: (B, 3) 标签，列顺序: short, medium, long

    Returns:
        {"short": {"loss": ..., "acc": ..., "f1": ...}, ...}
    """
    loss_fn = FocalLoss()
    horizons = ["short", "medium", "long"]
    metrics = {}

    for i, name in enumerate(horizons):
        logits = outputs[name]
        targets = labels[:, i]
        loss_val = loss_fn(logits, targets).item()
        preds = logits.argmax(dim=1).cpu().numpy()
        trues = targets.cpu().numpy()
        acc = accuracy_score(trues, preds)
        f1 = f1_score(trues, preds, average="macro")
        metrics[name] = {"loss": loss_val, "acc": acc, "f1": f1}

    return metrics


def train() -> HeatRiskPredictor:
    """执行完整的训练流水线。

    Returns:
        训练好的 HeatRiskPredictor 模型 (best checkpoint)
    """
    # -------------------- 设备 --------------------
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"使用设备: {device}")

    # -------------------- 加载数据 --------------------
    X, y = load_data()
    print(f"数据加载完成: X {X.shape}, y {y.shape}")

    # -------------------- 时间序划分 (不打乱) --------------------
    n_total = len(X)
    n_train = int(n_total * 0.70)
    n_val = int(n_total * 0.15)

    X_train_np = X[:n_train]
    y_train_np = y[:n_train]
    X_val_np = X[n_train : n_train + n_val]
    y_val_np = y[n_train : n_train + n_val]
    X_test_np = X[n_train + n_val :]
    y_test_np = y[n_train + n_val :]

    print(f"划分: 训练 {len(X_train_np)}, 验证 {len(X_val_np)}, 测试 {len(X_test_np)}")

    # -------------------- Tensor 转换 --------------------
    X_train_t = torch.tensor(X_train_np, dtype=torch.float32)
    y_train_t = torch.tensor(y_train_np, dtype=torch.long)
    X_val_t = torch.tensor(X_val_np, dtype=torch.float32)
    y_val_t = torch.tensor(y_val_np, dtype=torch.long)
    X_test_t = torch.tensor(X_test_np, dtype=torch.float32)
    y_test_t = torch.tensor(y_test_np, dtype=torch.long)

    # -------------------- DataLoader (加权采样, 温和过采样) --------------------
    from torch.utils.data import WeightedRandomSampler
    y_short = y_train_np[:, 0]
    class_counts = np.bincount(y_short)
    # 温和权重: 比 sqrt 更弱, 避免 overcorrect
    weights = 1.0 / np.sqrt(class_counts)
    sample_w = weights[y_short]
    sampler = WeightedRandomSampler(
        torch.from_numpy(sample_w).float(),
        num_samples=len(y_short),
        replacement=True,
    )
    train_loader = DataLoader(
        TensorDataset(X_train_t, y_train_t),
        batch_size=BATCH_SIZE,
        sampler=sampler,
    )
    val_loader = DataLoader(
        TensorDataset(X_val_t, y_val_t),
        batch_size=BATCH_SIZE,
        shuffle=False,
    )

    # -------------------- 模型 --------------------
    input_dim = X.shape[2]
    model = HeatRiskPredictor(input_dim=input_dim).to(device)
    print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")

    # -------------------- 损失、优化器、调度器 --------------------
    focal_loss = FocalLoss()
    optimizer = AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
    scheduler = ReduceLROnPlateau(optimizer, factor=0.5, patience=5)

    # -------------------- 训练状态 --------------------
    best_val_loss = float("inf")
    best_epoch = 0
    patience_counter = 0
    history: dict[str, list] = {
        "train_loss": [],
        "val_loss": [],
        "val_acc_short": [],
        "val_acc_medium": [],
        "val_acc_long": [],
    }

    OUTPUT_MODELS.mkdir(parents=True, exist_ok=True)
    best_model_path = OUTPUT_MODELS / "best_model.pt"

    # -------------------- 训练循环 --------------------
    for epoch in range(1, MAX_EPOCHS + 1):
        # ---- 训练阶段 ----
        model.train()
        train_losses: list[float] = []

        for batch_X, batch_y in train_loader:
            batch_X = batch_X.to(device)
            batch_y = batch_y.to(device)

            outputs = model(batch_X)
            loss_short = focal_loss(outputs["short"], batch_y[:, 0])
            loss_medium = focal_loss(outputs["medium"], batch_y[:, 1])
            loss_long = focal_loss(outputs["long"], batch_y[:, 2])
            loss = (loss_short + loss_medium + loss_long) / 3.0

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            train_losses.append(loss.item())

        avg_train_loss = float(np.mean(train_losses))

        # ---- 验证阶段 ----
        model.eval()
        val_outputs_all: list[dict] = []
        val_labels_all: list[torch.Tensor] = []

        with torch.no_grad():
            for batch_X, batch_y in val_loader:
                batch_X = batch_X.to(device)
                batch_y = batch_y.to(device)
                outputs = model(batch_X)
                val_outputs_all.append(
                    {k: v.cpu() for k, v in outputs.items()}
                )
                val_labels_all.append(batch_y.cpu())

        # 合并所有 batch
        merged_outputs = {
            name: torch.cat([o[name] for o in val_outputs_all], dim=0)
            for name in ["short", "medium", "long"]
        }
        merged_labels = torch.cat(val_labels_all, dim=0)
        val_metrics = _compute_metrics(merged_outputs, merged_labels)

        avg_val_loss = np.mean(
            [val_metrics[h]["loss"] for h in ["short", "medium", "long"]]
        )

        # ---- 记录历史 ----
        history["train_loss"].append(avg_train_loss)
        history["val_loss"].append(avg_val_loss)
        history["val_acc_short"].append(val_metrics["short"]["acc"])
        history["val_acc_medium"].append(val_metrics["medium"]["acc"])
        history["val_acc_long"].append(val_metrics["long"]["acc"])

        # ---- 学习率调度 ----
        scheduler.step(avg_val_loss)

        # ---- 打印进度 ----
        if True:  # print every epoch
            lr_now = optimizer.param_groups[0]["lr"]
            print(
                f"Epoch {epoch:3d}/{MAX_EPOCHS} | "
                f"Train Loss: {avg_train_loss:.4f} | "
                f"Val Loss: {avg_val_loss:.4f} | "
                f"Acc S/M/L: "
                f"{val_metrics['short']['acc']:.3f}/"
                f"{val_metrics['medium']['acc']:.3f}/"
                f"{val_metrics['long']['acc']:.3f} | "
                f"LR: {lr_now:.2e}"
            )

        # ---- 保存最佳模型 ----
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            best_epoch = epoch
            patience_counter = 0
            torch.save(
                {
                    "epoch": epoch,
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "val_loss": avg_val_loss,
                    "val_metrics": val_metrics,
                    "input_dim": input_dim,
                },
                best_model_path,
            )
        else:
            patience_counter += 1

        # ---- 早停 ----
        if patience_counter >= EARLY_STOP_PATIENCE:
            print(
                f"早停触发: {EARLY_STOP_PATIENCE} 轮未改善 "
                f"(最佳 epoch: {best_epoch}, Val Loss: {best_val_loss:.4f})"
            )
            break

    # -------------------- 保存训练历史 --------------------
    OUTPUT_LOGS.mkdir(parents=True, exist_ok=True)
    history_path = OUTPUT_LOGS / "training_history.json"
    with open(history_path, "w", encoding="utf-8") as f:
        json.dump(
            {
                "best_epoch": best_epoch,
                "best_val_loss": best_val_loss,
                **history,
            },
            f,
            indent=2,
            ensure_ascii=False,
        )
    print(f"训练历史已保存至 {history_path}")

    # -------------------- 加载最佳模型，测试评估 --------------------
    checkpoint = torch.load(best_model_path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    print("\n========== 测试集评估 ==========")
    test_dataset = TensorDataset(X_test_t, y_test_t)
    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

    test_outputs_all: list[dict] = []
    test_labels_all: list[torch.Tensor] = []

    with torch.no_grad():
        for batch_X, batch_y in test_loader:
            batch_X = batch_X.to(device)
            batch_y = batch_y.to(device)
            outputs = model(batch_X)
            test_outputs_all.append(
                {k: v.cpu() for k, v in outputs.items()}
            )
            test_labels_all.append(batch_y.cpu())

    merged_test_outputs = {
        name: torch.cat([o[name] for o in test_outputs_all], dim=0)
        for name in ["short", "medium", "long"]
    }
    merged_test_labels = torch.cat(test_labels_all, dim=0)
    test_metrics = _compute_metrics(merged_test_outputs, merged_test_labels)

    # 打印结果
    for name in ["short", "medium", "long"]:
        m = test_metrics[name]
        print(
            f"  {name:>6}: Accuracy={m['acc']:.4f}, "
            f"F1 Macro={m['f1']:.4f}, Loss={m['loss']:.4f}"
        )

    avg_acc = np.mean([test_metrics[h]["acc"] for h in ["short", "medium", "long"]])
    avg_f1 = np.mean([test_metrics[h]["f1"] for h in ["short", "medium", "long"]])
    print(f"\n平均 Accuracy: {avg_acc:.4f}")
    print(f"平均 F1 Macro: {avg_f1:.4f}")

    # -------------------- 保存测试预测 --------------------
    test_predictions: dict[str, np.ndarray] = {}
    for name in ["short", "medium", "long"]:
        test_predictions[name] = merged_test_outputs[name].argmax(dim=1).numpy()
    test_predictions["y_true"] = merged_test_labels.numpy()
    test_predictions["logits_short"] = merged_test_outputs["short"].numpy()
    test_predictions["logits_medium"] = merged_test_outputs["medium"].numpy()
    test_predictions["logits_long"] = merged_test_outputs["long"].numpy()

    pred_path = OUTPUT_MODELS / "test_predictions.npz"
    np.savez_compressed(pred_path, **test_predictions)
    print(f"测试预测已保存至 {pred_path}")

    return model


if __name__ == "__main__":
    import logging
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=[
            logging.FileHandler("outputs/logs/training.log"),
            logging.StreamHandler(),
        ],
    )
    train()
