models.py

"""
ML Predictor — 模型定义与训练
Logistic Regression / Random Forest / XGBoost / LightGBM / Stacking Ensemble
+ 概率校准（Isotonic Regression / Platt Scaling）
"""

import os
import pickle
import logging
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.calibration import CalibratedClassifierCV
from sklearn.isotonic import IsotonicRegression

import config

logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════
# 单模型包装器
# ═══════════════════════════════════════════

class BaseModel:
    """模型基类"""

    def __init__(self, name: str):
        self.name = name
        self.model = None
        self.scaler = None      # 特征标准化器
        self.calibrator = None  # 概率校准器
        self.feature_names: list[str] = []
        self.is_fitted = False

    def fit(self, X: pd.DataFrame, y: np.ndarray,
            X_cal: pd.DataFrame = None, y_cal: np.ndarray = None):
        """
        训练模型。
        X_cal, y_cal: 校准集（如果提供，用于训练概率校准器）
        """
        self.feature_names = list(X.columns)

        # 标准化
        self.scaler = StandardScaler()
        X_scaled = self.scaler.fit_transform(X)

        # 训练
        self._fit_impl(X_scaled, y)
        self.is_fitted = True

        # 概率校准
        if X_cal is not None and y_cal is not None:
            X_cal_scaled = self.scaler.transform(X_cal)
            raw_probs = self._predict_proba_impl(X_cal_scaled)
            self._fit_calibrator(raw_probs, y_cal)

        logger.info(f"[{self.name}] Trained on {len(X)} samples, "
                    f"{len(self.feature_names)} features")

    def predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        """预测 P(UP) — 返回校准后的概率"""
        if not self.is_fitted:
            raise ValueError(f"{self.name} not fitted")

        X_scaled = self.scaler.transform(X[self.feature_names])
        raw = self._predict_proba_impl(X_scaled)

        if self.calibrator is not None:
            calibrated = self.calibrator.predict(raw)
            return np.clip(calibrated, 0.01, 0.99)

        return np.clip(raw, 0.01, 0.99)

    def _fit_impl(self, X: np.ndarray, y: np.ndarray):
        raise NotImplementedError

    def _predict_proba_impl(self, X: np.ndarray) -> np.ndarray:
        raise NotImplementedError

    def _fit_calibrator(self, raw_probs: np.ndarray, y: np.ndarray):
        """训练概率校准器"""
        method = config.CALIBRATION_METHOD
        if method == "isotonic":
            self.calibrator = IsotonicRegression(
                y_min=0.01, y_max=0.99, out_of_bounds="clip"
            )
            self.calibrator.fit(raw_probs, y)
        elif method == "sigmoid":
            # Platt Scaling — 用 Logistic Regression 做校准
            from sklearn.linear_model import LogisticRegression as LR
            self.calibrator = _PlattScaler()
            self.calibrator.fit(raw_probs, y)
        logger.info(f"[{self.name}] Calibrator fitted ({method})")

    def feature_importance(self) -> Optional[pd.Series]:
        """返回特征重要性（如果模型支持）"""
        return None

    def save(self, path: str):
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump(self, f)
        logger.info(f"[{self.name}] Saved to {path}")

    @staticmethod
    def load(path: str) -> "BaseModel":
        with open(path, "rb") as f:
            return pickle.load(f)


class _PlattScaler:
    """Platt Scaling 校准器"""
    def __init__(self):
        self.lr = LogisticRegression(C=1e10, max_iter=1000)

    def fit(self, raw_probs, y):
        self.lr.fit(raw_probs.reshape(-1, 1), y)

    def predict(self, raw_probs):
        return self.lr.predict_proba(raw_probs.reshape(-1, 1))[:, 1]


# ═══════════════════════════════════════════
# 具体模型实现
# ═══════════════════════════════════════════

class LogisticModel(BaseModel):
    """Logistic Regression — 基线模型"""

    def __init__(self):
        super().__init__("LogisticRegression")

    def _fit_impl(self, X, y):
        self.model = LogisticRegression(
            C=config.LR_C,
            penalty=config.LR_PENALTY,
            max_iter=config.LR_MAX_ITER,
            solver="lbfgs",
            random_state=42,
        )
        self.model.fit(X, y)

    def _predict_proba_impl(self, X):
        return self.model.predict_proba(X)[:, 1]

    def feature_importance(self):
        if self.model is None:
            return None
        coefs = np.abs(self.model.coef_[0])
        return pd.Series(coefs, index=self.feature_names).sort_values(ascending=False)


class RandomForestModel(BaseModel):
    """Random Forest — 非线性特征交互"""

    def __init__(self):
        super().__init__("RandomForest")

    def _fit_impl(self, X, y):
        self.model = RandomForestClassifier(
            n_estimators=config.RF_N_ESTIMATORS,
            max_depth=config.RF_MAX_DEPTH,
            min_samples_leaf=config.RF_MIN_SAMPLES_LEAF,
            max_features=config.RF_MAX_FEATURES,
            random_state=42,
            n_jobs=-1,
        )
        self.model.fit(X, y)

    def _predict_proba_impl(self, X):
        return self.model.predict_proba(X)[:, 1]

    def feature_importance(self):
        if self.model is None:
            return None
        return pd.Series(
            self.model.feature_importances_, index=self.feature_names
        ).sort_values(ascending=False)


class XGBoostModel(BaseModel):
    """XGBoost — 梯度提升"""

    def __init__(self):
        super().__init__("XGBoost")

    def _fit_impl(self, X, y):
        import xgboost as xgb
        self.model = xgb.XGBClassifier(
            n_estimators=config.XGB_N_ESTIMATORS,
            max_depth=config.XGB_MAX_DEPTH,
            learning_rate=config.XGB_LEARNING_RATE,
            subsample=config.XGB_SUBSAMPLE,
            colsample_bytree=config.XGB_COLSAMPLE_BYTREE,
            reg_alpha=config.XGB_REG_ALPHA,
            reg_lambda=config.XGB_REG_LAMBDA,
            eval_metric="logloss",
            random_state=42,
            n_jobs=-1,
            verbosity=0,
        )
        self.model.fit(X, y)

    def _predict_proba_impl(self, X):
        return self.model.predict_proba(X)[:, 1]

    def feature_importance(self):
        if self.model is None:
            return None
        return pd.Series(
            self.model.feature_importances_, index=self.feature_names
        ).sort_values(ascending=False)


class LightGBMModel(BaseModel):
    """LightGBM — 轻量梯度提升"""

    def __init__(self):
        super().__init__("LightGBM")

    def _fit_impl(self, X, y):
        import lightgbm as lgb
        self.model = lgb.LGBMClassifier(
            n_estimators=config.LGB_N_ESTIMATORS,
            max_depth=config.LGB_MAX_DEPTH,
            num_leaves=config.LGB_NUM_LEAVES,
            learning_rate=config.LGB_LEARNING_RATE,
            min_child_samples=config.LGB_MIN_CHILD_SAMPLES,
            subsample=config.LGB_SUBSAMPLE,
            colsample_bytree=config.LGB_COLSAMPLE_BYTREE,
            reg_alpha=config.LGB_REG_ALPHA,
            reg_lambda=config.LGB_REG_LAMBDA,
            random_state=42,
            n_jobs=-1,
            verbose=-1,
        )
        self.model.fit(X, y)

    def _predict_proba_impl(self, X):
        return self.model.predict_proba(X)[:, 1]

    def feature_importance(self):
        if self.model is None:
            return None
        return pd.Series(
            self.model.feature_importances_, index=self.feature_names
        ).sort_values(ascending=False)


# ═══════════════════════════════════════════
# Stacking Ensemble
# ═══════════════════════════════════════════

class StackingEnsemble(BaseModel):
    """
    Stacking 集成模型。

    Layer 1: Logistic, RF, XGBoost, LightGBM（基模型）
    Layer 2: Logistic Regression（元模型）

    基模型的 out-of-fold 预测作为元模型的输入特征。
    """

    def __init__(self, base_models: list[BaseModel] = None):
        super().__init__("StackingEnsemble")
        if base_models is None:
            self.base_models = [
                LogisticModel(),
                RandomForestModel(),
                XGBoostModel(),
                LightGBMModel(),
            ]
        else:
            self.base_models = base_models
        self.meta_model = LogisticRegression(C=1.0, max_iter=1000)

    def fit(self, X: pd.DataFrame, y: np.ndarray,
            X_cal: pd.DataFrame = None, y_cal: np.ndarray = None):
        """
        训练 Stacking Ensemble。

        使用时间序列分割生成 out-of-fold 预测：
        - 保持时间顺序
        - 训练集 < 验证集（无前视偏差）
        """
        from sklearn.model_selection import TimeSeriesSplit

        self.feature_names = list(X.columns)
        self.scaler = StandardScaler()
        X_scaled_df = pd.DataFrame(
            self.scaler.fit_transform(X), columns=X.columns, index=X.index
        )

        n_splits = min(config.CV_N_SPLITS, max(2, len(X) // 500))
        tscv = TimeSeriesSplit(n_splits=n_splits, gap=config.CV_GAP)

        # 生成 out-of-fold 预测
        oof_preds = np.zeros((len(X), len(self.base_models)))

        for fold, (train_idx, val_idx) in enumerate(tscv.split(X_scaled_df)):
            X_tr = X_scaled_df.iloc[train_idx]
            y_tr = y[train_idx]
            X_val = X_scaled_df.iloc[val_idx]

            for j, model in enumerate(self.base_models):
                # 每个 base model 独立训练
                model_copy = self._clone_base_model(model)
                model_copy.scaler = None  # 已经标准化
                model_copy._scaler_bypass = True
                # 直接训练内部模型
                model_copy._fit_impl(X_tr.values, y_tr)
                model_copy.is_fitted = True
                oof_preds[val_idx, j] = model_copy._predict_proba_impl(X_val.values)

        # 用有 OOF 预测的样本训练元模型
        valid_mask = oof_preds.sum(axis=1) > 0
        meta_X = oof_preds[valid_mask]
        meta_y = y[valid_mask]

        self.meta_model.fit(meta_X, meta_y)

        # 用全量数据训练所有 base models
        for model in self.base_models:
            model._fit_impl(X_scaled_df.values, y)
            model.feature_names = self.feature_names
            model.is_fitted = True

        self.is_fitted = True

        # 校准
        if X_cal is not None and y_cal is not None:
            raw = self._predict_raw(X_cal)
            self._fit_calibrator(raw, y_cal)

        logger.info(f"[Ensemble] Stacking trained: {len(self.base_models)} base models")

    def predict_proba(self, X: pd.DataFrame) -> np.ndarray:
        if not self.is_fitted:
            raise ValueError("Ensemble not fitted")

        raw = self._predict_raw(X)

        if self.calibrator is not None:
            calibrated = self.calibrator.predict(raw)
            return np.clip(calibrated, 0.01, 0.99)

        return np.clip(raw, 0.01, 0.99)

    def _predict_raw(self, X: pd.DataFrame) -> np.ndarray:
        X_scaled = self.scaler.transform(X[self.feature_names])
        base_preds = np.column_stack([
            model._predict_proba_impl(X_scaled) for model in self.base_models
        ])
        return self.meta_model.predict_proba(base_preds)[:, 1]

    def _clone_base_model(self, model: BaseModel) -> BaseModel:
        """创建 base model 的浅拷贝"""
        cls = type(model)
        new_model = cls.__new__(cls)
        new_model.name = model.name
        new_model.model = None
        new_model.scaler = None
        new_model.calibrator = None
        new_model.feature_names = model.feature_names
        new_model.is_fitted = False
        return new_model

    def feature_importance(self):
        """聚合所有 base model 的特征重要性"""
        importances = []
        for model in self.base_models:
            imp = model.feature_importance()
            if imp is not None:
                imp_norm = imp / (imp.sum() + 1e-10)
                importances.append(imp_norm)

        if not importances:
            return None

        combined = pd.DataFrame(importances).mean()
        return combined.sort_values(ascending=False)

    def get_base_model_predictions(self, X: pd.DataFrame) -> dict:
        """获取各 base model 的独立预测"""
        X_scaled = self.scaler.transform(X[self.feature_names])
        return {
            model.name: float(model._predict_proba_impl(X_scaled)[0])
            for model in self.base_models
        }


# ═══════════════════════════════════════════
# 模型工厂
# ═══════════════════════════════════════════

def create_model(name: str) -> BaseModel:
    """根据名称创建模型"""
    models = {
        "logistic": LogisticModel,
        "random_forest": RandomForestModel,
        "xgboost": XGBoostModel,
        "lightgbm": LightGBMModel,
        "ensemble": StackingEnsemble,
    }
    if name not in models:
        raise ValueError(f"Unknown model: {name}. Choose from {list(models.keys())}")
    return models[name]()


def save_model(model: BaseModel, name: str = None):
    """保存模型到磁盘"""
    name = name or model.name
    path = os.path.join(config.MODEL_DIR, f"{name}.pkl")
    model.save(path)


def load_model(name: str) -> BaseModel:
    """从磁盘加载模型"""
    path = os.path.join(config.MODEL_DIR, f"{name}.pkl")
    return BaseModel.load(path)