"""TabM 排序学习任务实现

实现基于 TabM 的排序学习训练流程：
- Label 转换为分位数标签
- 生成 group 数组
- 使用 TabMRankModel（基于 ListNet Loss）
- 支持 NDCG@k 评估
"""

from typing import Any, Dict, List, Optional, Type
import numpy as np
import polars as pl

from src.training.tasks.base import BaseTask
from src.training.components.models.tabm_rank_model import TabMRankModel


class TabMRankTask(BaseTask):
    """TabM 排序学习任务

    使用 TabMRankModel（或兼容模型）进行排序学习训练。
    将连续收益率转换为分位数标签进行训练。
    支持指数化增益标签以增强 Top-K 关注。
    """

    def __init__(
        self,
        model_params: Dict[str, Any],
        label_name: str = "future_return_5",
        n_quantiles: int = 20,
        label_transform: Optional[str] = None,
        label_scale: float = 20.0,
        model_class: Optional[Type] = None,
    ):
        """初始化排序学习任务

        Args:
            model_params: TabM 参数字典
            label_name: Label 列名
            n_quantiles: 分位数数量
            label_transform: 标签变换类型，可选:
                - None: 标准分位数标签 (0, 1, ..., n_quantiles-1)
                - "exponential": 指数化增益: 2^(rank/scale) - 1
            label_scale: 指数变换的缩放因子，用于控制增益幅度
            model_class: 模型类，默认 TabMRankModel。可传入兼容接口的其他模型，
                         如 TabMSetRankModel。
        """
        super().__init__(model_params, label_name)
        self.n_quantiles = n_quantiles
        self.label_transform = label_transform
        self.label_scale = label_scale
        self.model_class = model_class or TabMRankModel

    @staticmethod
    def _rank_gauss(arr: np.ndarray) -> np.ndarray:
        """对一维数组做 Rank-Gauss 变换

        将数组值转换为基于排名的标准正态分布值，结果近似服从 N(0, 1)。

        Args:
            arr: 输入一维数组

        Returns:
            Rank-Gauss 变换后的数组
        """
        from scipy.special import erfinv
        from scipy.stats import rankdata

        n = len(arr)
        if n <= 1:
            return np.zeros(n, dtype=np.float64)

        # 1. 获取平均排名 (1-based)
        ranks = rankdata(arr, method="average")
        # 2. 映射到 (0, 1)，使用 (ranks - 0.5) / n 避免边界
        p = (ranks - 0.5) / n
        # 3. 映射到 (-1, 1)
        p = 2.0 * p - 1.0
        # 4. 裁剪边界避免 erfinv 输入越界
        eps = 1e-6
        p = np.clip(p, -1.0 + eps, 1.0 - eps)
        # 5. Rank-Gauss: erfinv(p) * sqrt(2)
        return erfinv(p) * np.sqrt(2.0)

    def prepare_labels(self, data: Dict[str, Dict]) -> Dict[str, Dict]:
        """准备标签（转换为分位数标签或 Rank-Gauss 连续标签）

        支持三种模式：
        - 标准分位数标签 (0, 1, ..., n_quantiles-1)
        - 指数化增益变换 (rank^2)
        - Rank-Gauss 连续标签（截面正态化，不依赖分桶数）

        Args:
            data: 数据字典

        Returns:
            处理后的数据字典（添加了 y_rank 和 groups）
        """
        for split in ["train", "val", "test"]:
            if split not in data:
                continue

            df = data[split]["raw_data"]
            rank_col = f"{self.label_name}_rank"

            if self.label_transform == "rank_gauss":
                # Rank-Gauss 连续化标签：先截面高斯化，再缩放到 (0, 1) 全局正区间
                df_ranked = df.with_columns(
                    pl.col(self.label_name)
                    .map_batches(
                        lambda s: pl.Series(self._rank_gauss(s.to_numpy())),
                        is_elementwise=False,
                    )
                    .map_batches(
                        lambda s: pl.Series((s.to_numpy() + 4.0) / 8.0).clip(0.0, 1.0),
                        is_elementwise=False,
                    )
                    .over("trade_date")
                    .alias(rank_col)
                )
            else:
                # 1. 基础分位数标签 (0 到 n_quantiles-1)
                df_ranked = df.with_columns(
                    pl.col(self.label_name)
                    .rank(method="min")
                    .over("trade_date")
                    .alias("_rank")
                ).with_columns(
                    (
                        (pl.col("_rank") - 1)
                        / pl.len().over("trade_date")
                        * self.n_quantiles
                    )
                    .floor()
                    .cast(pl.Int64)
                    .clip(0, self.n_quantiles - 1)
                    .alias("_base_rank")
                )

                # 2. 【Top-K 优化】可选指数化增益变换
                if self.label_transform == "exponential":
                    # 平方变换: rank^2
                    # 例如 rank=0 -> 0, rank=10 -> 100, rank=19 -> 361
                    # 效果：高分样本与低分样本的差距被平方级拉大
                    df_ranked = df_ranked.with_columns(
                        (pl.col("_base_rank").cast(pl.Float64) ** 2).alias(rank_col)
                    )
                else:
                    # 标准分位数标签
                    df_ranked = df_ranked.with_columns(
                        pl.col("_base_rank").cast(pl.Float64).alias(rank_col)
                    )

                # 清理临时列
                df_ranked = df_ranked.drop(["_rank", "_base_rank"])

            # 更新数据
            data[split]["raw_data"] = df_ranked
            data[split]["y"] = df_ranked[rank_col]
            data[split]["y_raw"] = df_ranked[self.label_name]  # 保留原始值

            # 生成 group 数组
            data[split]["groups"] = self._compute_group_array(df_ranked, "trade_date")

        return data

    def _compute_group_array(
        self,
        df: pl.DataFrame,
        date_col: str = "trade_date",
    ) -> np.ndarray:
        """计算 group 数组

        Args:
            df: 数据框
            date_col: 日期列名

        Returns:
            group 数组（每个日期的样本数）
        """
        group_counts = df.group_by(date_col, maintain_order=True).agg(
            pl.count().alias("count")
        )
        return group_counts["count"].to_numpy()

    def fit(self, train_data: Dict, val_data: Dict) -> None:
        """训练排序模型

        Args:
            train_data: 训练数据
            val_data: 验证数据
        """
        self.model_params["n_features"] = train_data["X"].shape[1]
        self.model = self.model_class(params=self.model_params)

        self.model.fit(
            train_data["X"],
            train_data["y"],
            group=train_data["groups"],
            eval_set=(val_data["X"], val_data["y"], val_data["groups"])
            if val_data
            else None,
        )

    def predict(self, test_data: Dict) -> np.ndarray:
        """生成预测

        Args:
            test_data: 测试数据

        Returns:
            预测结果
        """
        # 传入 groups 参数，确保预测顺序与分组一致，与验证逻辑保持一致
        return self.model.predict(test_data["X"], group=test_data.get("groups"))

    def evaluate_ndcg(
        self,
        test_data: Dict,
        k_list: List[int] = None,
    ) -> Dict[str, float]:
        """评估 NDCG@k

        Args:
            test_data: 测试数据
            k_list: k 值列表，默认 [1, 5, 10, 20]

        Returns:
            NDCG 分数字典 {"ndcg@1": score, ...}
        """
        if k_list is None:
            k_list = [1, 5, 10, 20]

        y_pred = self.predict(test_data)
        groups = test_data["groups"]

        from sklearn.metrics import ndcg_score

        results = {}

        # 按 group 拆分
        start_idx = 0
        y_true_groups = []
        y_pred_groups = []

        # 使用分位数标签 y (0-19) 作为真实相关性分数，而非原始收益率 y_raw
        # 这样与模型学习目标一致，避免原始收益率中负值的影响
        y_true_array = test_data["y"].to_numpy()

        for group_size in groups:
            end_idx = start_idx + group_size
            y_true_groups.append(y_true_array[start_idx:end_idx])
            y_pred_groups.append(y_pred[start_idx:end_idx])
            start_idx = end_idx

        # 计算每个 k 的 NDCG
        for k in k_list:
            ndcg_scores = []
            for yt, yp in zip(y_true_groups, y_pred_groups):
                if len(yt) > 1:
                    try:
                        score = ndcg_score([yt], [yp], k=k)
                        ndcg_scores.append(score)
                    except ValueError:
                        pass

            results[f"ndcg@{k}"] = float(np.mean(ndcg_scores)) if ndcg_scores else 0.0

        return results

    def plot_training_metrics(self) -> None:
        """绘制训练指标曲线（NDCG）"""
        if self.model and hasattr(self.model, "get_evals_result"):
            try:
                import matplotlib.pyplot as plt

                evals_result = self.model.get_evals_result()
                if not evals_result:
                    print("[警告] 没有训练指标数据可供绘制")
                    return

                fig, ax = plt.subplots(1, 2, figsize=(12, 4))

                # 绘制训练损失
                if "train_loss" in evals_result:
                    ax[0].plot(evals_result["train_loss"], label="Train Loss")
                    ax[0].set_xlabel("Epoch")
                    ax[0].set_ylabel("ListNet Loss")
                    ax[0].set_title("Training Loss")
                    ax[0].legend()
                    ax[0].grid(True)

                # 绘制验证 NDCG
                if "val_ndcg" in evals_result:
                    ax[1].plot(evals_result["val_ndcg"], label="Val NDCG")
                    ax[1].set_xlabel("Epoch")
                    ax[1].set_ylabel("NDCG")
                    ax[1].set_title("Validation NDCG")
                    ax[1].legend()
                    ax[1].grid(True)

                plt.tight_layout()
                plt.show()
            except Exception as e:
                print(f"[警告] 无法绘制训练曲线: {e}")
