"""LightGBM LambdaRank 排序学习模型

提供 LightGBM LambdaRank 模型的实现，支持学习排序任务。
适用于将股票按未来收益率进行排序的场景。
"""

from typing import Any, Optional, List

import numpy as np
import pandas as pd
import polars as pl

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


@register_model("lightgbm_lambdarank")
class LightGBMLambdaRankModel(BaseModel):
    """LightGBM LambdaRank 排序学习模型

    使用 LightGBM 的 LambdaRank 实现学习排序。
    适用于股票排序任务，将未来收益率转换为分位数标签进行训练。

    Attributes:
        name: 模型名称 "lightgbm_lambdarank"
        params: LightGBM 参数字典
        model: 训练后的 LightGBM Booster 对象
        feature_names_: 特征名称列表
    """

    name = "lightgbm_lambdarank"

    def __init__(self, params: Optional[dict] = None):
        """初始化 LambdaRank 模型

        Args:
            params: LightGBM 参数字典，直接传递给 lgb.train()。
                   包含所有模型参数和训练控制参数。

        Examples:
            >>> model = LightGBMLambdaRankModel(params={
            ...     "objective": "lambdarank",
            ...     "metric": "ndcg",
            ...     "ndcg_at": [1, 5, 10, 20],
            ...     "num_leaves": 31,
            ...     "learning_rate": 0.05,
            ...     "n_estimators": 1000,
            ...     "early_stopping_round": 50,
            ... })
        """
        self.params = dict(params) if params is not None else {}
        self.model = None
        self.feature_names_: Optional[list] = None
        self.evals_result_: Optional[dict] = None

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

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

        Returns:
            self (支持链式调用)

        Raises:
            ImportError: 未安装 lightgbm
            RuntimeError: 训练失败
            ValueError: group 参数无效
        """
        try:
            import lightgbm as lgb
        except ImportError:
            raise ImportError(
                "使用 LightGBMLambdaRankModel 需要安装 lightgbm: pip install lightgbm"
            )

        self.feature_names_ = X.columns
        X_np = X.to_numpy()
        y_np = y.to_numpy()

        if group is None:
            group = np.array([len(y_np)])

        if not isinstance(group, np.ndarray):
            group = np.array(group)
        if group.sum() != len(y_np):
            raise ValueError(
                f"group 数组的和 ({group.sum()}) 必须等于样本数 ({len(y_np)})"
            )

        train_data = lgb.Dataset(X_np, label=y_np, group=group)

        valid_sets = [train_data]
        valid_names = ["train"]
        if eval_set is not None:
            X_val, y_val, group_val = eval_set
            X_val_np = X_val.to_numpy() if isinstance(X_val, pl.DataFrame) else X_val
            y_val_np = y_val.to_numpy() if isinstance(y_val, pl.Series) else y_val

            if group_val is None:
                group_val = np.array([len(y_val_np)])
            elif not isinstance(group_val, np.ndarray):
                group_val = np.array(group_val)

            val_data = lgb.Dataset(X_val_np, label=y_val_np, group=group_val)
            valid_sets.append(val_data)
            valid_names.append("val")

        self.evals_result_ = {}

        # 从 params 提取训练控制参数
        params_copy = dict(self.params)
        num_boost_round = params_copy.pop("n_estimators", 100)
        early_stopping_round = params_copy.pop("early_stopping_round", 50)

        callbacks = [
            lgb.early_stopping(stopping_rounds=early_stopping_round),
            lgb.record_evaluation(self.evals_result_),
        ]

        self.model = lgb.train(
            params_copy,
            train_data,
            num_boost_round=num_boost_round,
            valid_sets=valid_sets,
            valid_names=valid_names,
            callbacks=callbacks,
        )

        return self

    def predict(self, X: pl.DataFrame) -> np.ndarray:
        """预测排序分数

        Args:
            X: 特征矩阵 (Polars DataFrame)

        Returns:
            预测分数 (numpy ndarray)

        Raises:
            RuntimeError: 模型未训练时调用
        """
        if self.model is None:
            raise RuntimeError("模型尚未训练，请先调用 fit()")

        X_np = X.to_numpy()
        result = self.model.predict(X_np)
        return np.asarray(result)

    def get_evals_result(self) -> Optional[dict]:
        """获取训练评估结果

        Returns:
            评估结果字典，如果模型尚未训练返回 None
        """
        return self.evals_result_

    def get_best_iteration(self) -> Optional[int]:
        """获取最佳迭代轮数（考虑早停）

        Returns:
            最佳迭代轮数，如果模型未训练返回 None
        """
        if self.model is None:
            return None
        return self.model.best_iteration

    def get_best_score(self) -> Optional[dict]:
        """获取最佳评分

        Returns:
            最佳评分字典，如果模型未训练返回 None
        """
        if self.model is None:
            return None
        return self.model.best_score

    def feature_importance(self) -> Optional[pd.Series]:
        """返回特征重要性

        Returns:
            特征重要性序列，如果模型未训练则返回 None
        """
        if self.model is None or self.feature_names_ is None:
            return None

        importance = self.model.feature_importance(importance_type="gain")
        return pd.Series(importance, index=self.feature_names_)

    def save(self, path: str) -> None:
        """保存模型（使用 LightGBM 原生格式）

        Args:
            path: 保存路径

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

        self.model.save_model(path)

        import json

        meta_path = path + ".meta.json"
        with open(meta_path, "w") as f:
            json.dump(
                {
                    "feature_names": self.feature_names_,
                    "params": self.params,
                },
                f,
            )

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

        Args:
            path: 模型文件路径

        Returns:
            加载的 LightGBMLambdaRankModel 实例
        """
        import lightgbm as lgb
        import json

        instance = cls()
        instance.model = lgb.Booster(model_file=path)

        meta_path = path + ".meta.json"
        try:
            with open(meta_path, "r") as f:
                meta = json.load(f)
                instance.feature_names_ = meta.get("feature_names")
                instance.params = meta.get("params", {})
        except FileNotFoundError:
            pass

        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: 新列名

        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:
                    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
