"""TabM 排序模型实现 (TabM Rank)

基于 TabM (Tabular Multilayer Perceptron with Ensembles) 架构
引入 ListNet 列表级排序损失，实现类似 LambdaRank 的截面排序学习。
适用于股票未来收益率的截面排序预测。
"""

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

import numpy as np
import polars as pl
import scipy.stats as stats
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
from tabm import TabM

from src.training.components.base import BaseModel
from src.training.registry import register_model


class GroupSampler(Sampler):
    """排序学习专用的分组采样器

    确保每个 Batch 包含的是同一个 Query (如同一天) 的所有样本。
    这与 LightGBM 的 `group` 参数逻辑完全一致。
    """

    def __init__(self, group_counts: np.ndarray, shuffle_groups: bool = True):
        """初始化分组采样器

        Args:
            group_counts: 各组的样本数量数组，如 [100, 150, 120]
            shuffle_groups: 是否打乱组的训练顺序
        """
        self.group_counts = group_counts
        self.shuffle_groups = shuffle_groups
        # 计算每组在原数组中的起始和结束边界
        self.boundaries = np.insert(np.cumsum(group_counts), 0, 0)
        self.num_groups = len(group_counts)

    def __iter__(self):
        """迭代生成批次索引

        Yields:
            list: 同一组的所有样本索引
        """
        group_indices = list(range(self.num_groups))
        if self.shuffle_groups:
            np.random.shuffle(group_indices)

        for g_idx in group_indices:
            start = self.boundaries[g_idx]
            end = self.boundaries[g_idx + 1]
            # 返回该组（天）内所有样本的索引，作为一个完整的 batch
            yield list(range(start, end))

    def __len__(self):
        """返回组数量"""
        return self.num_groups


class EnsembleListNetLoss(nn.Module):
    """集成 ListNet 排序损失 (Listwise Ranking Loss)

    基于交叉熵将截面上的相关度转化为概率分布。
    支持 TabM 的 Ensemble 维度。
    """

    def __init__(self, topk_weight: float = 1.0):
        """初始化 ListNet 损失

        Args:
            topk_weight: 头部样本权重系数，>1.0 时强化对高分标签的关注
        """
        super().__init__()
        self.topk_weight = topk_weight

    def forward(self, preds: torch.Tensor, targets: torch.Tensor):
        """计算 ListNet 排序损失

        Args:
            preds: [BatchSize, EnsembleSize] - 当前组所有样本的预测 logits
            targets: [BatchSize] - 当前组所有样本的真实排序标签 (0, 1, 2...)

        Returns:
            标量损失值
        """
        # 如果该组内样本太少，无法排序，直接返回 0
        if preds.size(0) <= 1:
            return torch.tensor(0.0, requires_grad=True, device=preds.device)

        # [BatchSize] -> [BatchSize, EnsembleSize] 广播以对齐每个集成成员
        targets_expanded = targets.unsqueeze(1).expand_as(preds)

        # 1. 计算真实标签的分数概率分布 (Softmax 使得高分权重成倍提升)
        targets_prob = F.softmax(targets_expanded, dim=0)

        # 2. 【Top-K 优化】如果启用加权，给予高分标签更高权重
        if self.topk_weight > 1.0:
            # 基于标签值计算权重（标签越高，权重越大）
            # 归一化到 [1, topk_weight] 范围
            max_target = targets.max()
            min_target = targets.min()
            if max_target > min_target:
                # 线性插值：权重 = 1 + (target - min) / (max - min) * (topk_weight - 1)
                sample_weights = 1.0 + (targets - min_target) / (
                    max_target - min_target
                ) * (self.topk_weight - 1.0)
                sample_weights = sample_weights.unsqueeze(1).expand_as(preds)
                # 归一化权重使其和为样本数
                sample_weights = sample_weights * len(targets) / sample_weights.sum()
                # 应用权重到目标概率
                targets_prob = targets_prob * sample_weights

        # 3. 计算预测值的对数概率分布 (log_softmax 数值上比 log(softmax) 更稳定)
        preds_log_prob = F.log_softmax(preds, dim=0)

        # 4. 计算交叉熵损失 (在 Batch/组 维度求和)
        loss = -torch.sum(targets_prob * preds_log_prob, dim=0)  # [EnsembleSize]

        # 5. 对所有集成成员的 Loss 取平均
        return loss.mean()


class EnsembleLambdaLoss(nn.Module):
    """集成 LambdaLoss (支持 TabM 集成维度)

    基于 Pairwise 排序损失，引入 DeltaNDCG 权重。
    参考: "The LambdaLoss Framework for Ranking Metric Optimization" (Google Research)

    特点:
    - 计算每对样本交换位置后对 NDCG 的影响
    - 对头部样本（高 Gain 且排名靠前）给予更高权重
    - 更适合 Top-K 选股场景
    """

    def __init__(self, sigma: float = 1.0, ndcg_weight_power: float = 1.0):
        """初始化 LambdaLoss

        Args:
            sigma: Sigmoid 函数的陡峭程度，控制梯度大小
            ndcg_weight_power: DeltaNDCG 权重幂次，>1 时进一步放大头部效应
        """
        super().__init__()
        self.sigma = sigma
        self.ndcg_weight_power = ndcg_weight_power

    def forward(self, preds: torch.Tensor, targets: torch.Tensor):
        """计算 LambdaLoss

        Args:
            preds: [BatchSize, EnsembleSize] - 预测分
            targets: [BatchSize] - 相关性标签 (Gain)

        Returns:
            标量损失值
        """
        if preds.size(0) <= 1:
            return torch.tensor(0.0, requires_grad=True, device=preds.device)

        # 1. 计算两两对之间的差值
        preds_diff = preds.unsqueeze(1) - preds.unsqueeze(0)  # [B, B, E]

        # 【性能优化】: target_diff 不需要 E 维度！直接保持 [B, B]
        target_diff = targets.unsqueeze(1) - targets.unsqueeze(0)  # [B, B]

        # 2. 掩码矩阵: [B, B, 1] 方便后续广播
        mask = (target_diff > 0).float().unsqueeze(2)  # [B, B, 1]

        # 3. 计算 Delta NDCG
        # DeltaNDCG = |Gain_i - Gain_j| * |1/log(rank_i+1) - 1/log(rank_j+1)|
        with torch.no_grad():
            # 【性能核爆优化】: 使用两次 argsort 完全消灭 for 循环
            # 第1次 argsort: 获得从大到小的索引
            # 第2次 argsort: 直接将索引反转为排名 (加上 1 就是真实名次)
            ranks = preds.argsort(dim=0, descending=True).argsort(dim=0) + 1
            ranks = ranks.float()  # [B, E]

            # 计算位置惩罚项 (log2 排名)
            log_rank = torch.log2(ranks + 1)
            inv_log_rank_diff = torch.abs(
                1.0 / log_rank.unsqueeze(1) - 1.0 / log_rank.unsqueeze(0)
            )  # [B, B, E]

            # DeltaNDCG = |Gain 差| * |位置惩罚差|
            # target_diff 是 [B, B]，通过 unsqueeze 广播到 [B, B, 1]
            delta_ndcg = torch.abs(target_diff).unsqueeze(2) * inv_log_rank_diff

            # 应用幂次调整权重分布
            if self.ndcg_weight_power != 1.0:
                delta_ndcg = torch.pow(delta_ndcg, self.ndcg_weight_power)

        # 4. Pairwise Logistic Loss 并加权 Delta NDCG
        # loss = delta_ndcg * log(1 + exp(-sigma * (preds_i - preds_j)))
        pairwise_loss = F.binary_cross_entropy_with_logits(
            self.sigma * preds_diff,
            torch.ones_like(preds_diff),
            reduction="none",
        )

        # 5. 应用掩码和权重
        weighted_loss = pairwise_loss * mask * delta_ndcg

        # 避免全 0 导致的除零错误
        valid_pairs_count = mask.sum().clamp(min=1.0)
        mean_loss = weighted_loss.sum() / valid_pairs_count

        return mean_loss


@register_model("tabm_rank")
class TabMRankModel(BaseModel):
    """TabM 学习排序模型

    基于 TabM 架构的排序学习模型，支持 ListNet 损失。
    适用于股票截面排序任务，将未来收益率转换为分位数标签进行训练。

    特点:
    - 使用 ListNet 列表级排序损失
    - 支持 group 参数进行分组训练
    - 以 NDCG 作为验证指标
    - 与 LightGBMLambdaRank 接口兼容
    """

    name = "tabm_rank"

    def __init__(self, params: Optional[Dict[str, Any]] = None):
        """初始化 TabM Rank 模型

        Args:
            params: 模型参数字典，包含:
                - ensemble_size: 集成大小 (默认: 32)
                - n_blocks: MLP层数 (默认: 3)
                - d_block: 每层神经元数 (默认: 256)
                - dropout: Dropout率 (默认: 0.1)
                - batch_size: 批次大小 (默认: 2048，仅预测时使用)
                - learning_rate: 学习率 (默认: 1e-3)
                - weight_decay: 权重衰减 (默认: 1e-5)
                - epochs: 训练轮数 (默认: 50)
                - early_stopping_round: 早停轮数 (默认: 10)
                - max_grad_norm: 梯度裁剪阈值 (默认: 1.0)
                - ndcg_k: NDCG@k 的 k 值，None 表示全局 (默认: None)
                - loss_type: 损失函数类型 (默认: "listnet")
                    - "listnet": 标准 ListNet 损失
                    - "weighted_listnet": 加权 ListNet，通过 topk_weight 强化头部
                    - "lambda": LambdaLoss，基于 DeltaNDCG 加权
                - topk_weight: 头部样本权重系数，用于 weighted_listnet (默认: 5.0)
                - lambda_sigma: LambdaLoss 的 sigma 参数 (默认: 1.0)
                - ndcg_weight_power: DeltaNDCG 权重幂次 (默认: 1.0)
        """
        self.params = params or {}
        self.model = None
        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

        # 根据配置选择损失函数
        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:  # "listnet"
            self.criterion = EnsembleListNetLoss(topk_weight=1.0)

    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 截面打包)

        Args:
            X: 特征数组 [N, n_features]
            y: 标签数组 [N] 或 None
            group: 分组数组，表示每组样本数
            shuffle_groups: 是否打乱组的顺序

        Returns:
            DataLoader 实例
        """
        # 【性能核爆优化】: 显存管够时，直接把全量数据一把推进 GPU，避免训练时 PCIe 搬运
        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:
            # 训练和验证时使用 GroupSampler，每个 batch 就是一个 Query
            sampler = GroupSampler(group, shuffle_groups=shuffle_groups)
            return DataLoader(dataset, batch_sampler=sampler)
        else:
            # 预测时如果没有 group，则退化为普通批次预测
            batch_size = self.params.get("batch_size", 2048)
            return DataLoader(dataset, batch_size=batch_size, shuffle=False)

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

        Args:
            val_loader: 验证数据加载器
            k: NDCG@k 的 k 值，None 表示计算全局 NDCG

        Returns:
            平均 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
                bx = bx.to(self.device)
                by = by.cpu().numpy()

                if len(by) <= 1:
                    continue

                outputs = self.model(bx)  # [B, E, 1]
                preds = outputs.mean(dim=1).squeeze(-1).cpu().numpy()  # [B]

                try:
                    # ndcg_score 需要形状为 (1, n_samples) 的二维数组
                    score = ndcg_score([by], [preds], k=k)
                    ndcg_list.append(score)
                except ValueError:
                    pass

        return float(np.mean(ndcg_list)) if len(ndcg_list) > 0 else 0.0

    def fit(
        self,
        X: pl.DataFrame,
        y: pl.Series,
        group: Optional[np.ndarray] = None,
        eval_set: Optional[Tuple] = None,
    ) -> "TabMRankModel":
        """训练排序模型

        Args:
            X: 训练特征DataFrame
            y: 训练标签 (Polars Series)，应为分位数标签 (0, 1, 2, ...)
            group: 分组数组，表示每个 query 的样本数
            eval_set: 验证集元组 (X_val, y_val, group_val)，用于早停

        Returns:
            self (支持链式调用)

        Raises:
            ValueError: group 参数无效
        """
        self.feature_names_ = list(X.columns)

        X_np = X.to_numpy().astype(np.float32)
        y_np = y.to_numpy().astype(np.float32)

        # 检查和处理 group 参数
        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
            )

        ensemble_size = self.params.get("ensemble_size", 32)
        n_features = X_np.shape[1]

        # 初始化 TabM 模型
        self.model = TabM.make(
            n_num_features=n_features,
            cat_cardinalities=[],
            d_out=1,
            n_blocks=self.params.get("n_blocks", 3),
            d_block=self.params.get("d_block", 256),
            dropout=self.params.get("dropout", 0.1),
            k=ensemble_size,
        ).to(self.device)

        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)  # None 表示计算全局 NDCG

        print(f"[TabMRank] 开始训练... 设备: {self.device}, 集成大小: {ensemble_size}")

        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[0], batch[1]

                optimizer.zero_grad()
                outputs = self.model(bx)  # [B, E, 1]
                outputs_squeezed = outputs.squeeze(-1)  # [B, E]

                # 计算 ListNet 排序损失
                loss = self.criterion(outputs_squeezed, by)
                loss.backward()

                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(),
                    max_norm=self.params.get("max_grad_norm", 1.0),
                )
                optimizer.step()

                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)

            # 验证阶段 (基于 NDCG)
            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"[TabMRank] Epoch {epoch + 1}/{epochs} | "
                        f"Train Loss (ListNet): {avg_train_loss:.4f} | "
                        f"Val NDCG: {val_ndcg:.4f} (Best: {best_val_ndcg:.4f})"
                    )

                if patience_counter >= early_stopping_patience:
                    print(f"[TabMRank] 触发早停，停止于 epoch {epoch + 1}")
                    break
            else:
                if (epoch + 1) % 5 == 0 or epoch == 0:
                    print(
                        f"[TabMRank] Epoch {epoch + 1}/{epochs} | "
                        f"Train Loss: {avg_train_loss:.4f}"
                    )

            scheduler.step()

        # 恢复最佳权重
        if best_model_state is not None:
            self.model.load_state_dict(best_model_state)
            print(f"[TabMRank] 已恢复最佳模型权重 (Val NDCG: {best_val_ndcg:.4f})")

        return self

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

        Args:
            X: 特征矩阵 (Polars DataFrame)
            group: 分组数组，表示每个 query 的样本数。
                   如果提供，将使用 GroupSampler 确保预测顺序与分组一致。

        Returns:
            预测分数 (numpy ndarray)

        Raises:
            RuntimeError: 模型未训练时调用
            ValueError: 预测数据缺失特征
        """
        if self.model is None:
            raise RuntimeError("模型未训练，请先调用fit()")

        # 特征对齐检查
        if self.feature_names_:
            missing_cols = [c for c in self.feature_names_ if c not in X.columns]
            if missing_cols:
                raise ValueError(f"预测数据缺失特征: {missing_cols}")
            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].to(self.device)
                outputs = self.model(bx)  # [B, E, 1]
                # 排序模型预测时直接输出集成成员的均值作为最终分数
                preds = outputs.mean(dim=1).squeeze(-1)  # [B]
                all_preds.append(preds.cpu().numpy())

        return np.concatenate(all_preds)

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

        Returns:
            评估结果字典，包含 train_loss 和 val_ndcg
        """
        return self.training_history_

    def feature_importance(self) -> None:
        """获取特征重要性

        TabM没有内置特征重要性计算，返回None。
        """
        return None

    def save(self, path: str | Path) -> None:
        """保存模型

        Args:
            path: 保存路径

        Raises:
            RuntimeError: 模型未训练时调用
        """
        if self.model is None:
            raise RuntimeError("模型未训练，无法保存")

        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        # 保存模型权重
        model_path = path.with_suffix(".pt")
        torch.save(self.model.state_dict(), model_path)

        # 保存元数据
        meta_path = path.with_suffix(".meta")
        meta = {
            "params": self.params,
            "feature_names": self.feature_names_,
            "training_history": self.training_history_,
            "device": str(self.device),
        }
        with open(meta_path, "wb") as f:
            pickle.dump(meta, f)

        print(f"[TabMRank] 模型保存到: {path}")

    @classmethod
    def load(cls, path: str | Path) -> "TabMRankModel":
        """加载模型

        Args:
            path: 模型路径（不含扩展名）

        Returns:
            加载的 TabMRankModel 实例
        """
        path = Path(path)

        # 加载元数据
        meta_path = path.with_suffix(".meta")
        with open(meta_path, "rb") as f:
            meta = pickle.load(f)

        # 创建实例
        instance = cls(meta["params"])
        instance.feature_names_ = meta["feature_names"]
        instance.training_history_ = meta["training_history"]

        # 重建模型结构
        if instance.feature_names_ is not None:
            n_features = len(instance.feature_names_)
            ensemble_size = instance.params.get("ensemble_size", 32)

            instance.model = TabM.make(
                n_num_features=n_features,
                cat_cardinalities=[],
                d_out=1,
                n_blocks=instance.params.get("n_blocks", 3),
                d_block=instance.params.get("d_block", 256),
                dropout=instance.params.get("dropout", 0.1),
                k=ensemble_size,
            ).to(instance.device)

            # 加载权重
            model_path = path.with_suffix(".pt")
            instance.model.load_state_dict(
                torch.load(model_path, map_location=instance.device)
            )

        print(f"[TabMRank] 模型从 {path} 加载完成")
        return instance

    @staticmethod
    def prepare_group_from_dates(
        df: pl.DataFrame,
        date_col: str = "trade_date",
    ) -> np.ndarray:
        """从日期列生成 group 数组

        Args:
            df: 包含日期列的 DataFrame
            date_col: 日期列名，默认 "trade_date"

        Returns:
            group 数组
        """
        group_counts = df.group_by(date_col, maintain_order=True).agg(
            pl.count().alias("count")
        )
        return group_counts["count"].to_numpy()

    @staticmethod
    def convert_to_quantile_labels(
        df: pl.DataFrame,
        label_col: str,
        date_col: str = "trade_date",
        n_quantiles: int = 20,
        new_col_name: Optional[str] = None,
    ) -> pl.DataFrame:
        """将连续标签转换为分位数标签

        Args:
            df: 输入 DataFrame
            label_col: 原始标签列名
            date_col: 日期列名，默认 "trade_date"
            n_quantiles: 分位数数量，默认 20
            new_col_name: 新列名，默认为 {label_col}_rank

        Returns:
            添加了分位数标签列的 DataFrame
        """
        if new_col_name is None:
            new_col_name = f"{label_col}_rank"

        return df.with_columns(
            pl.col(label_col)
            .qcut(n_quantiles)
            .over(date_col)
            .to_physical()
            .cast(pl.Int64)
            .alias(new_col_name)
        )

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

        Args:
            X: 特征矩阵
            y: 真实标签
            group: 分组数组
            k: NDCG@k 的 k 值

        Returns:
            NDCG 分数
        """
        from sklearn.metrics import ndcg_score

        y_pred = self.predict(X)

        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.to_numpy()[start_idx:end_idx])
            y_score_list.append(y_pred[start_idx:end_idx])
            start_idx = end_idx

        ndcg_scores = []
        for y_true, y_score in zip(y_true_list, y_score_list):
            if len(y_true) > 1:
                try:
                    # 若标签存在负值（如 Rank-Gauss），平移到正值区间再计算 NDCG
                    if y_true.min() < 0:
                        y_true_shifted = y_true - y_true.min() + 1e-5
                        score = ndcg_score([y_true_shifted], [y_score], k=k)
                    else:
                        score = ndcg_score([y_true], [y_score], k=k)
                    ndcg_scores.append(score)
                except ValueError:
                    pass

        return float(np.mean(ndcg_scores)) if ndcg_scores else 0.0
