"""LightGBM 模型实现

提供 LightGBM 回归模型的实现，支持特征重要性和原生模型保存。
"""

from typing import Any, Optional

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")
class LightGBMModel(BaseModel):
    """LightGBM 回归模型

    使用 LightGBM 库实现梯度提升回归树。
    支持自定义参数、特征重要性提取和原生模型格式保存。

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

    name = "lightgbm"

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

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

        Examples:
            >>> model = LightGBMModel(params={
            ...     "objective": "regression",
            ...     "metric": "rmse",
            ...     "num_leaves": 31,
            ...     "learning_rate": 0.05,
            ...     "n_estimators": 100,
            ... })
        """
        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,
        eval_set: Optional[tuple] = None,
    ) -> "LightGBMModel":
        """训练模型

        Args:
            X: 特征矩阵 (Polars DataFrame)
            y: 目标变量 (Polars Series)
            eval_set: 验证集元组 (X_val, y_val)，用于早停

        Returns:
            self (支持链式调用)

        Raises:
            ImportError: 未安装 lightgbm
            RuntimeError: 训练失败
        """
        try:
            import lightgbm as lgb
        except ImportError:
            raise ImportError(
                "使用 LightGBMModel 需要安装 lightgbm: pip install lightgbm"
            )

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

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

        # 准备验证集
        valid_sets = None
        if eval_set is not None:
            X_val, y_val = eval_set
            X_val_np = X_val.to_numpy()
            y_val_np = y_val.to_numpy()
            valid_sets = lgb.Dataset(X_val_np, label=y_val_np, reference=train_data)

        # 从 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)

        self.evals_result_ = {}
        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] if valid_sets else None,
            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 原生格式）

        使用 LightGBM 的原生格式保存，不依赖 pickle，
        可以在不同环境中加载。

        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) -> "LightGBMModel":
        """加载模型

        从 LightGBM 原生格式加载模型。

        Args:
            path: 模型文件路径

        Returns:
            加载的 LightGBMModel 实例
        """
        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
