"""
LightGBM 模型封装。
"""

import numpy as np
import pandas as pd
import lightgbm as lgb

from models.base import BaseModel
from sklearn.model_selection import TimeSeriesSplit

_DEFAULT_PARAMS = {
    "objective": "regression",
    "metric": "rmse",
    "boosting_type": "gbdt",
    "num_leaves": 15,
    "learning_rate": 0.03,
    "feature_fraction": 0.7,
    "bagging_fraction": 0.7,
    "bagging_freq": 5,
    "verbose": -1,
    "n_estimators": 1000,
    "random_state": 42,
    "min_data_in_leaf": 20,
}


class LightGBMModel(BaseModel):
    """LightGBM 回归模型。"""

    name = "lightgbm"

    def __init__(
        self,
        params: dict | None = None,
        early_stopping: int = 50,
        eval_ratio: float = 0.2,
        random_seed: int = 42,
    ):
        self.params = params or _DEFAULT_PARAMS.copy()
        self.params["random_state"] = random_seed
        self.early_stopping = early_stopping
        self.eval_ratio = eval_ratio
        self._model: lgb.Booster | None = None
        self._feature_names: list[str] = []

    def fit(self, X: pd.DataFrame, y: pd.Series) -> "LightGBMModel":
        self._feature_names = list(X.columns)

        n = len(X)
        val_size = int(n * self.eval_ratio)

        callbacks = []
        eval_set = None
        X_train, y_train = X, y

        # 验证集足够大时才启用早停（至少 50 条）
        if val_size >= 50:
            split_idx = n - val_size
            X_train, X_val = X.iloc[:split_idx], X.iloc[split_idx:]
            y_train, y_val = y.iloc[:split_idx], y.iloc[split_idx:]
            eval_set = [(X_val, y_val)]
            callbacks = [
                lgb.early_stopping(stopping_rounds=self.early_stopping, verbose=False),
                lgb.log_evaluation(0),
            ]

        self._model = lgb.LGBMRegressor(**self.params)
        self._model.fit(
            X_train, y_train,
            eval_set=eval_set,
            callbacks=callbacks if callbacks else None,
        )
        return self

    def predict(self, X: pd.DataFrame) -> pd.Series:
        if self._model is None:
            raise RuntimeError("模型尚未训练")
        preds = self._model.predict(X[self._feature_names])
        return pd.Series(preds, index=X.index, name="pred")

    def get_feature_importance(self, importance_type: str = "gain") -> pd.DataFrame:
        """特征重要性。importance_type: 'gain' | 'split'"""
        if self._model is None:
            return pd.DataFrame()
        imp = self._model.booster_.feature_importance(importance_type=importance_type)
        names = self._model.booster_.feature_name()
        df = pd.DataFrame({"feature": names, "importance": imp})
        df["importance_pct"] = df["importance"] / df["importance"].sum() * 100
        return df.sort_values("importance", ascending=False)

    def cv_evaluate(
        self, X: pd.DataFrame, y: pd.Series, n_folds: int = 5
    ) -> pd.DataFrame:
        """
        时间序列交叉验证评估（不 shuffle）。

        返回每折的 IC (相关系数) 和 MSE。
        """
        tscv = TimeSeriesSplit(n_splits=n_folds)
        results = []
        for fold, (train_idx, test_idx) in enumerate(tscv.split(X)):
            X_train, X_test = X.iloc[train_idx], X.iloc[test_idx]
            y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]

            model = LightGBMModel(
                params=self.params,
                early_stopping=self.early_stopping,
                eval_ratio=0.0,  # 不使用内部验证，直接全量训练
            )
            model.fit(X_train, y_train)
            preds = model.predict(X_test)
            ic = preds.corr(y_test)
            mse = ((preds - y_test) ** 2).mean()
            results.append({"fold": fold, "ic": round(ic, 4), "mse": round(mse, 4)})

        df = pd.DataFrame(results)
        df.loc["mean"] = df.mean()
        return df

    @property
    def n_estimators_used(self) -> int | None:
        """实际使用的树数量（早停后可能 < n_estimators）。"""
        if self._model is None:
            return None
        return self._model.booster_.current_iteration()