# src/models.py
from __future__ import annotations

from sklearn.ensemble import (
    ExtraTreesRegressor,
    RandomForestRegressor,
    GradientBoostingRegressor,
    StackingRegressor,
)

# Optional XGBoost (paper uses XGBR in base layer)
try:
    from xgboost import XGBRegressor
    HAS_XGB = True
except Exception:
    HAS_XGB = False


def build_sm2_stacking(random_state: int = 42) -> StackingRegressor:
    """
    SM2-style stacking:
      Base: ETR, RFR, GBR, (XGBR if available)
      Meta: GBR
    """
    base_estimators = [
        ("etr", ExtraTreesRegressor(n_estimators=1000, random_state=random_state, n_jobs=-1)),
        ("rfr", RandomForestRegressor(n_estimators=1000, random_state=random_state, n_jobs=-1)),
        ("gbr", GradientBoostingRegressor(random_state=random_state)),
    ]

    if HAS_XGB:
        base_estimators.append(
            ("xgbr", XGBRegressor(
                n_estimators=100,
                max_depth=6,
                learning_rate=0.1,
                subsample=0.9,
                reg_lambda=1.0,
                random_state=random_state,
                n_jobs=-1,
                tree_method="hist",
            ))
        )

    meta = GradientBoostingRegressor(random_state=random_state)

    return StackingRegressor(
        estimators=base_estimators,
        final_estimator=meta,
        passthrough=True,
        n_jobs=-1,
    )