"""TabM + SetRank 排序模型实现 (TabM_SetRank)

基于 TabM 特征提取 + SetRank 组内注意力建模。
引入 ListNet/LambdaLoss 列表级排序损失，支持 AMP 混合精度与显存优化。
适用于股票未来收益率的截面排序预测。
"""

from typing import Dict, Any, List, Optional, Tuple
from pathlib import Path
import pickle

import numpy as np
import polars as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset, Sampler

# 注：torch.cuda.amp.autocast 已被废弃，请统一使用 torch.amp.autocast('cuda', ...)
# 同时从子模块显式导入以兼容各类 IDE 的类型提示
from torch.amp.grad_scaler import GradScaler
from torch.amp.autocast_mode import autocast
from tabm import TabM

from src.training.components.base import BaseModel
from src.training.registry import register_model
from src.training.components.models.tabm_rank_model import (
    EnsembleListNetLoss,
    EnsembleLambdaLoss,
    GroupSampler,
)


class SetRankHead(nn.Module):
    """标准的 SetRank 组内上下文注意力头

    输入: [N, E] (TabM集成预测)
    输出: [N, E] (经过截面交互后的集成预测)
    """

    def __init__(
        self,
        d_in: int,  # TabM 的 Ensemble 数量 (例如 32)
        n_heads: int = 4,
        d_ff: int = 128,  # SetRank 的隐藏特征维度
        dropout: float = 0.1,
    ):
        super().__init__()
        # 1. 特征升维: 将 Logits 映射为高维语义特征
        self.feature_proj = nn.Linear(d_in, d_ff)

        # 2. 多头注意力层 (SetRank 核心)
        # 注意：不要加入位置编码 (Positional Encoding)，因为股票截面是一个无序集合(Set)，必须保持排列等变性
        self.attn = nn.MultiheadAttention(
            embed_dim=d_ff, num_heads=n_heads, dropout=dropout, batch_first=True
        )
        self.norm1 = nn.LayerNorm(d_ff)

        # 3. 前馈神经网络 (标准 Transformer FFN 结构)
        self.ffn = nn.Sequential(
            nn.Linear(d_ff, d_ff * 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff * 2, d_ff),
        )
        self.norm2 = nn.LayerNorm(d_ff)

        # 4. 恢复集成维度: [N, d_ff] -> [N, E]
        self.output_proj = nn.Linear(d_ff, d_in)

        # 5. 可学习的残差缩放因子，初始化为 0，训练初期退化为纯 TabM
        self.res_scale = nn.Parameter(torch.zeros(1))

        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.LayerNorm):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
        # Zero-init 输出层，确保训练启动时为恒等映射
        nn.init.zeros_(self.output_proj.weight)
        nn.init.zeros_(self.output_proj.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x 初始维度: [N, E]

        # 升维: [N, d_ff]
        h = self.feature_proj(x)

        # 增加 Batch 维度以适配 MultiheadAttention: [1, N, d_ff]
        # B=1 (代表一天的截面), L=N (股票数量), D=d_ff
        h = h.unsqueeze(0)

        # Pre-Norm Attention 交互
        attn_out, _ = self.attn(self.norm1(h), self.norm1(h), self.norm1(h))
        h = h + attn_out

        # Pre-Norm FFN
        ffn_out = self.ffn(self.norm2(h))
        h = h + ffn_out

        # 降维并去除 Batch 维度: [1, N, d_ff] -> [N, d_ff]
        h = h.squeeze(0)

        # 输出回 Ensemble 维度: [N, E]
        out_logits = self.output_proj(h)

        # 带有可学习缩放因子的残差连接，初始阶段退化为纯 TabM
        return x + self.res_scale * out_logits


class TabMSetRankNet(nn.Module):
    def __init__(
        self,
        n_features: int,
        ensemble_size: int,
        n_blocks: int,
        d_block: int,
        dropout: float,
        setrank_heads: int,
        setrank_hidden: int,
        setrank_dropout: float,
        mlp_hidden: Optional[int] = None,
    ):
        super().__init__()
        mlp_hidden = mlp_hidden or d_block

        # TabM 特征提取器 (作为隐藏层，输出维度为 d_block)
        self.tabm = TabM.make(
            n_num_features=n_features,
            cat_cardinalities=[],
            d_out=d_block,
            n_blocks=n_blocks,
            d_block=d_block,
            dropout=dropout,
            k=ensemble_size,
        )
        # 截面注意力修正器 (处理 TabM 隐藏特征)
        self.setrank_head = SetRankHead(
            d_in=d_block,
            n_heads=setrank_heads,
            d_ff=setrank_hidden,
            dropout=setrank_dropout,
        )

        # 最终两次 MLP 输出 Ensemble 排序分
        self.final_mlp = nn.Sequential(
            nn.Linear(d_block, mlp_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_hidden, ensemble_size),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. TabM 独立特征提取 -> [N, E, d_block]
        tabm_out = self.tabm(x)

        # 2. 对 ensemble 维度取平均，得到原始隐藏特征 -> [N, d_block]
        tabm_hidden = tabm_out.mean(dim=1)

        # 3. 截面上下文交互修正 -> [N, d_block]
        # SetRankHead 内部实现残差: x + res_scale * transform(x)
        setrank_hidden = self.setrank_head(tabm_hidden)

        # 4. 两次 MLP 输出最终 Ensemble 排序分 -> [N, E]
        scores = self.final_mlp(setrank_hidden)
        return scores


@register_model("tabm_setrank_rank")
class TabMSetRankModel(BaseModel):
    """TabM + SetRank 学习排序模型

    使用统一的 nn.Module (TabMSetRankNet) 封装 TabM 与 SetRankHead。
    支持 AMP 混合精度与显存优化。
    """

    name = "tabm_setrank_rank"

    def __init__(self, params: Optional[Dict[str, Any]] = None):
        self.params = params or {}
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.training_history_: Dict[str, List[float]] = {
            "train_loss": [],
            "val_ndcg": [],
        }
        self.feature_names_: Optional[List[str]] = None

        self.use_amp = self.params.get("use_amp", True)
        self.scaler = GradScaler(enabled=self.use_amp)

        loss_type = self.params.get("loss_type", "listnet")
        if loss_type == "lambda":
            self.criterion = EnsembleLambdaLoss(
                sigma=self.params.get("lambda_sigma", 1.0),
                ndcg_weight_power=self.params.get("ndcg_weight_power", 1.0),
            )
        elif loss_type == "weighted_listnet":
            self.criterion = EnsembleListNetLoss(
                topk_weight=self.params.get("topk_weight", 5.0)
            )
        else:
            self.criterion = EnsembleListNetLoss(topk_weight=1.0)

        n_features = self.params.get("n_features")
        if n_features is None:
            raise ValueError("TabMSetRankModel 初始化需要在 params 中传入 'n_features'")
        ensemble_size = self.params.get("ensemble_size", 32)
        self.model = TabMSetRankNet(
            n_features=n_features,
            ensemble_size=ensemble_size,
            n_blocks=self.params.get("n_blocks", 3),
            d_block=self.params.get("d_block", 256),
            dropout=self.params.get("dropout", 0.1),
            setrank_heads=self.params.get("setrank_heads", 4),
            setrank_hidden=self.params.get("setrank_hidden", 128),
            setrank_dropout=self.params.get("setrank_dropout", 0.1),
            mlp_hidden=self.params.get("mlp_hidden", None),
        ).to(self.device)

    def _make_loader(
        self,
        X: np.ndarray,
        y: Optional[np.ndarray] = None,
        group: Optional[np.ndarray] = None,
        shuffle_groups: bool = False,
    ) -> DataLoader:
        """创建 DataLoader (支持 Query/Group 截面打包)

        数据在一开始就直接放到 GPU，避免训练/预测循环中反复拷贝。
        """
        X_tensor = torch.from_numpy(X).to(self.device)
        if y is not None:
            y_tensor = torch.from_numpy(y).to(self.device)
            dataset = TensorDataset(X_tensor, y_tensor)
        else:
            dataset = TensorDataset(X_tensor)

        if group is not None:
            sampler = GroupSampler(group, shuffle_groups=shuffle_groups)
            return DataLoader(
                dataset,
                batch_sampler=sampler,
                num_workers=self.params.get("num_workers", 0),
                pin_memory=self.params.get("pin_memory", False),
            )
        else:
            batch_size = self.params.get("batch_size", 2048)
            return DataLoader(
                dataset,
                batch_size=batch_size,
                shuffle=False,
                num_workers=self.params.get("num_workers", 0),
                pin_memory=self.params.get("pin_memory", False),
            )

    def _validate_ndcg(self, val_loader: DataLoader, k: Optional[int] = None) -> float:
        """验证模型 (使用 NDCG 排序指标)"""
        from sklearn.metrics import ndcg_score

        assert self.model is not None, "模型未训练，无法验证"
        self.model.eval()
        ndcg_list = []

        with torch.no_grad():
            for batch in val_loader:
                if len(batch) != 2:
                    continue
                bx, by = batch
                by = by.cpu().numpy()
                if len(by) <= 1:
                    continue

                with autocast("cuda", enabled=self.use_amp):
                    scores = self.model(bx)  # [N, E]

                # 对 Ensemble 维度求平均，得到最终排序分 [N]
                preds = scores.mean(dim=1).cpu().numpy()
                try:
                    ndcg_list.append(ndcg_score([by], [preds], k=k))
                except ValueError:
                    pass
        return float(np.mean(ndcg_list)) if ndcg_list else 0.0

    def fit(
        self,
        X: pl.DataFrame,
        y: pl.Series,
        group: Optional[np.ndarray] = None,
        eval_set: Optional[Tuple] = None,
    ) -> "TabMSetRankModel":
        """训练排序模型"""
        self.feature_names_ = list(X.columns)
        X_np = X.to_numpy().astype(np.float32)
        y_np = y.to_numpy().astype(np.float32)

        if group is None:
            group = np.array([len(y_np)])
        if group.sum() != len(y_np):
            raise ValueError(
                f"group 数组的和 ({group.sum()}) 必须等于样本数 ({len(y_np)})"
            )

        train_loader = self._make_loader(X_np, y_np, group=group, shuffle_groups=True)
        val_loader = None
        if eval_set is not None:
            X_val, y_val, group_val = eval_set
            X_val_np = (
                X_val.to_numpy().astype(np.float32)
                if isinstance(X_val, pl.DataFrame)
                else X_val
            )
            y_val_np = (
                y_val.to_numpy().astype(np.float32)
                if isinstance(y_val, pl.Series)
                else y_val
            )
            if group_val is None:
                group_val = np.array([len(y_val_np)])
            val_loader = self._make_loader(
                X_val_np, y_val_np, group=group_val, shuffle_groups=False
            )

        n_features = X_np.shape[1]
        expected_n_features = self.params.get("n_features")
        if expected_n_features is not None and n_features != expected_n_features:
            raise ValueError(
                f"输入特征维度 ({n_features}) 与初始化时指定的 n_features ({expected_n_features}) 不一致"
            )

        optimizer = optim.AdamW(
            self.model.parameters(),
            lr=self.params.get("learning_rate", 1e-3),
            weight_decay=self.params.get("weight_decay", 1e-5),
        )
        epochs = self.params.get("epochs", 50)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=epochs, eta_min=1e-6
        )
        early_stopping_patience = self.params.get(
            "early_stopping_patience",
            self.params.get("early_stopping_round", 10),
        )
        best_val_ndcg = -float("inf")
        patience_counter = 0
        best_model_state = None
        ndcg_k = self.params.get("ndcg_k", None)

        print(f"[TabM+SetRank] 开始训练... 设备: {self.device}, AMP: {self.use_amp}")

        for epoch in range(epochs):
            self.model.train()
            train_loss = 0.0
            n_train_batches = 0

            for batch in train_loader:
                if len(batch) != 2:
                    continue
                bx, by = batch

                optimizer.zero_grad()
                with autocast("cuda", enabled=self.use_amp):
                    scores = self.model(bx)  # [N, E]
                    loss = self.criterion(scores, by)

                self.scaler.scale(loss).backward()
                self.scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(),
                    max_norm=self.params.get("max_grad_norm", 1.0),
                )
                self.scaler.step(optimizer)
                self.scaler.update()

                train_loss += loss.item()
                n_train_batches += 1

            avg_train_loss = train_loss / max(n_train_batches, 1)
            self.training_history_["train_loss"].append(avg_train_loss)

            if val_loader is not None:
                val_ndcg = self._validate_ndcg(val_loader, k=ndcg_k)
                self.training_history_["val_ndcg"].append(val_ndcg)
                if val_ndcg > best_val_ndcg:
                    best_val_ndcg = val_ndcg
                    patience_counter = 0
                    best_model_state = {
                        k: v.cpu().clone() for k, v in self.model.state_dict().items()
                    }
                else:
                    patience_counter += 1

                if (epoch + 1) % 5 == 0 or epoch == 0:
                    print(
                        f"[TabM+SetRank] Epoch {epoch + 1}/{epochs} | "
                        f"Loss: {avg_train_loss:.4f} | "
                        f"Val NDCG: {val_ndcg:.4f} (Best: {best_val_ndcg:.4f})"
                    )
                if patience_counter >= early_stopping_patience:
                    print(f"[TabM+SetRank] 早停于 epoch {epoch + 1}")
                    break
            else:
                if (epoch + 1) % 5 == 0 or epoch == 0:
                    print(
                        f"[TabM+SetRank] Epoch {epoch + 1}/{epochs} | "
                        f"Loss: {avg_train_loss:.4f}"
                    )
            scheduler.step()

        if best_model_state:
            self.model.load_state_dict(best_model_state)
            print(f"[TabM+SetRank] 恢复最佳权重 (Val NDCG: {best_val_ndcg:.4f})")
        return self

    def predict(
        self, X: pl.DataFrame, group: Optional[np.ndarray] = None
    ) -> np.ndarray:
        """预测排序分数

        对于 TabM + SetRank 模型，group 参数是必需的。
        SetRankHead 中的截面注意力机制假设每个 batch 就是一天的所有股票，
        若缺失 group 退化为普通 batch_size=2048 预测，会导致跨天样本错误交互。
        """
        if self.model is None:
            raise RuntimeError("模型未训练，请先调用fit()")
        if group is None:
            raise ValueError(
                "TabMSetRankModel.predict() 需要提供 group 数组，"
                "确保每个 batch 为完整的一天截面。请通过 test_data['groups'] 传入。"
            )
        if self.feature_names_:
            missing = [c for c in self.feature_names_ if c not in X.columns]
            if missing:
                raise ValueError(f"缺失特征: {missing}")
            X = X.select(self.feature_names_)

        X_np = X.to_numpy().astype(np.float32)
        loader = self._make_loader(X_np, group=group, shuffle_groups=False)
        self.model.eval()
        all_preds = []

        with torch.no_grad():
            for batch in loader:
                bx = batch[0]
                with autocast("cuda", enabled=self.use_amp):
                    scores = self.model(bx)  # [N, E]

                # 对 Ensemble 维度求平均，得到最终排序分 [N]
                preds = scores.mean(dim=1)
                all_preds.append(preds.cpu().numpy())
        return np.concatenate(all_preds)

    def get_evals_result(self) -> Optional[Dict[str, List[float]]]:
        """获取训练评估结果"""
        return self.training_history_

    def feature_importance(self) -> None:
        """TabM 没有内置特征重要性计算，返回 None。"""
        return None

    def save(self, path: str | Path) -> None:
        """保存模型"""
        if self.model is None:
            raise RuntimeError("模型未训练，无法保存")
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save(self.model.state_dict(), path.with_suffix(".pt"))
        with open(path.with_suffix(".meta"), "wb") as f:
            pickle.dump(
                {
                    "params": self.params,
                    "feature_names": self.feature_names_,
                    "training_history": self.training_history_,
                    "device": str(self.device),
                },
                f,
            )

    @classmethod
    def load(cls, path: str | Path) -> "TabMSetRankModel":
        """加载模型"""
        path = Path(path)
        with open(path.with_suffix(".meta"), "rb") as f:
            meta = pickle.load(f)
        meta_params = meta["params"]
        feature_names = meta["feature_names"]
        assert feature_names is not None, "加载的模型缺少 feature_names，无法重建网络"
        meta_params["n_features"] = len(feature_names)
        instance = cls(meta_params)
        instance.feature_names_ = feature_names
        instance.training_history_ = meta["training_history"]
        ckpt = torch.load(path.with_suffix(".pt"), map_location=instance.device)
        instance.model.load_state_dict(ckpt)
        return instance

    def evaluate_ndcg(
        self,
        X: pl.DataFrame,
        y: pl.Series,
        group: np.ndarray,
        k: Optional[int] = None,
    ) -> float:
        """评估 NDCG 指标"""
        from sklearn.metrics import ndcg_score

        y_pred = self.predict(X, group=group)
        y_true_array = y.to_numpy() if isinstance(y, pl.Series) else y
        y_true_list = []
        y_score_list = []
        start_idx = 0
        for group_size in group:
            end_idx = start_idx + group_size
            y_true_list.append(y_true_array[start_idx:end_idx])
            y_score_list.append(y_pred[start_idx:end_idx])
            start_idx = end_idx

        ndcg_scores = []
        for yt, yp in zip(y_true_list, y_score_list):
            if len(yt) > 1:
                try:
                    ndcg_scores.append(ndcg_score([yt], [yp], k=k))
                except ValueError:
                    pass
        return float(np.mean(ndcg_scores)) if ndcg_scores else 0.0
