"""排序学习任务实现

实现排序学习任务的训练流程：
- Label 转换为分位数标签
- 生成 group 数组
- 使用 LightGBM LambdaRank
- 支持 NDCG@k 评估
"""

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

from src.training.tasks.base import BaseTask
from src.training.components.models.lightgbm_lambdarank import LightGBMLambdaRankModel


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

    使用 LightGBM LambdaRank 进行排序学习训练。
    将连续收益率转换为分位数标签进行训练。
    """

    def __init__(
        self,
        model_params: Dict[str, Any],
        label_name: str = "future_return_5",
        n_quantiles: int = 20,
    ):
        """初始化排序学习任务

        Args:
            model_params: LightGBM 参数字典
            label_name: Label 列名
            n_quantiles: 分位数数量
        """
        super().__init__(model_params, label_name)
        self.n_quantiles = n_quantiles

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

        将连续收益率转换为分位数标签，并生成 group 数组。

        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"
            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(rank_col)
                )
                .drop("_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 = LightGBMLambdaRankModel(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:
            预测结果
        """
        return self.model.predict(test_data["X"])

    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 lightgbm as lgb

                evals_result = self.model.get_evals_result()
                if evals_result:
                    lgb.plot_metric(evals_result)
                else:
                    print("[警告] 没有训练指标数据可供绘制")
            except Exception as e:
                print(f"[警告] 无法绘制训练曲线: {e}")
