"""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_classifier")
class LightGBMClassifier(BaseModel):
    """LightGBM 分类模型

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

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

    name = "lightgbm_classifier"

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

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

        Examples:
            >>> model = LightGBMClassifier(params={
            ...     "objective": "binary",
            ...     "metric": "auc",
            ...     "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

    def fit(
        self,
        X: pl.DataFrame,
        y: pl.Series,
        eval_set: Optional[tuple] = None,
    ) -> "LightGBMClassifier":
        """训练模型

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

        Returns:
            self (支持链式调用)

        Raises:
            ImportError: 未安装 lightgbm
            RuntimeError: 训练失败
        """
        try:
            import lightgbm as lgb
        except ImportError:
            raise ImportError(
                "使用 LightGBMClassifier 需要安装 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 = [train_data]
        valid_names = ["train"]
        callbacks = []

        if eval_set is not None:
            X_val, y_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
            val_data = lgb.Dataset(X_val_np, label=y_val_np)
            valid_sets.append(val_data)
            valid_names.append("val")

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

        if len(valid_sets) > 1:
            callbacks.append(lgb.early_stopping(stopping_rounds=early_stopping_round))

        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 predict_class(self, X: pl.DataFrame) -> np.ndarray:
        """预测类别

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

        Returns:
            预测类别 (0 或 1)

        Raises:
            RuntimeError: 模型未训练时调用
        """
        proba = self.predict(X)
        return (proba >= 0.5).astype(int)

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

        Args:
            importance_type: 重要性类型，默认为 "gain"
                             必须使用 "gain"，"split" 会被噪音欺骗

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

        importance = self.model.feature_importance(importance_type=importance_type)
        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) -> "LightGBMClassifier":
        """加载模型

        从 LightGBM 原生格式加载模型。

        Args:
            path: 模型文件路径

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