from __future__ import annotations

from typing import Any, Dict, Tuple
import warnings

from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.linear_model import ElasticNet, Lasso, LinearRegression, Ridge
from sklearn.neural_network import MLPRegressor
from sklearn.ensemble import RandomForestRegressor, ExtraTreesRegressor, GradientBoostingRegressor
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler


class RawModel(BaseEstimator, RegressorMixin):
    """Placeholder for RAW; prediction is handled outside because it uses original model field."""

    def fit(self, X, y):
        return self

    def predict(self, X):
        raise RuntimeError("RawModel prediction is handled outside train/predict pipeline.")


def _as_tuple(x: Any) -> Tuple[int, ...]:
    if isinstance(x, tuple):
        return x
    if isinstance(x, list):
        return tuple(int(v) for v in x)
    if isinstance(x, int):
        return (int(x),)
    raise ValueError(f"Cannot convert {x!r} to hidden_layer_sizes tuple")


def _xgb_params_for_method(method: str, cfg: Dict[str, Any], random_seed: int) -> Dict[str, Any]:
    base = {
        "objective": "reg:squarederror",
        "n_jobs": 1,
        "random_state": random_seed,
        "tree_method": "hist",
        "verbosity": 0,
    }
    base.update(dict(cfg.get("xgboost_base_params", {})))
    variants = dict(cfg.get("xgboost_variants", {}))
    if method in variants:
        base.update(dict(variants[method]))
        return base
    builtins = {
        "xgb_d2": dict(n_estimators=500, max_depth=2, learning_rate=0.03, subsample=0.90, colsample_bytree=0.90, reg_lambda=5.0, reg_alpha=0.0, min_child_weight=3.0),
        "xgb_d3": dict(n_estimators=500, max_depth=3, learning_rate=0.03, subsample=0.90, colsample_bytree=0.90, reg_lambda=5.0, reg_alpha=0.0, min_child_weight=3.0),
        "xgb_d4": dict(n_estimators=500, max_depth=4, learning_rate=0.03, subsample=0.85, colsample_bytree=0.85, reg_lambda=8.0, reg_alpha=0.0, min_child_weight=4.0),
        "xgb_slow": dict(n_estimators=900, max_depth=3, learning_rate=0.015, subsample=0.90, colsample_bytree=0.90, reg_lambda=8.0, reg_alpha=0.0, min_child_weight=3.0),
        "xgb_regularized": dict(n_estimators=700, max_depth=2, learning_rate=0.02, subsample=0.75, colsample_bytree=0.75, reg_lambda=15.0, reg_alpha=0.5, min_child_weight=5.0),
    }
    if method not in builtins:
        raise ValueError(f"Unknown XGBoost method variant: {method}")
    base.update(builtins[method])
    return base


def _mlp_params_for_method(method: str, cfg: Dict[str, Any], random_seed: int) -> Dict[str, Any]:
    base = {
        "random_state": random_seed,
        "max_iter": 700,
        "early_stopping": True,
        "validation_fraction": 0.15,
        "n_iter_no_change": 40,
        "learning_rate_init": 0.001,
        "alpha": 0.001,
        "solver": "adam",
        "activation": "relu",
    }
    base.update(dict(cfg.get("mlp_base_params", {})))
    variants = dict(cfg.get("mlp_variants", {}))
    if method in variants:
        base.update(dict(variants[method]))
    else:
        builtins = {
            "mlp_small": dict(hidden_layer_sizes=(32,)),
            "mlp_medium": dict(hidden_layer_sizes=(64, 32)),
            "mlp_deep": dict(hidden_layer_sizes=(128, 64, 32), alpha=0.003),
            "ann_tanh": dict(hidden_layer_sizes=(64, 32), activation="tanh", alpha=0.001),
            "ann_lbfgs": dict(hidden_layer_sizes=(32,), activation="tanh", solver="lbfgs", alpha=0.01, max_iter=700),
            "mlp_1layer_32": dict(hidden_layer_sizes=(32,), activation="relu", alpha=0.001),
            "mlp_1layer_64": dict(hidden_layer_sizes=(64,), activation="relu", alpha=0.001),
            "mlp_2layer_64_32": dict(hidden_layer_sizes=(64, 32), activation="relu", alpha=0.001),
            "mlp_2layer_128_64": dict(hidden_layer_sizes=(128, 64), activation="relu", alpha=0.003),
            "mlp_tanh_64": dict(hidden_layer_sizes=(64,), activation="tanh", alpha=0.001),
        }
        if method not in builtins:
            raise ValueError(f"Unknown MLP/ANN method variant: {method}")
        base.update(builtins[method])
    if "hidden_layer_sizes" in base:
        base["hidden_layer_sizes"] = _as_tuple(base["hidden_layer_sizes"])
    if base.get("solver") == "lbfgs":
        base.pop("early_stopping", None)
        base.pop("validation_fraction", None)
        base.pop("n_iter_no_change", None)
        base.pop("learning_rate_init", None)
    return base


def get_model_params(method: str, cfg: Dict[str, Any], random_seed: int = 42) -> Dict[str, Any]:
    """Return human-readable hyperparameters used by a method variant."""
    method = str(method).lower()
    if method == "raw":
        return {"type": "raw", "description": "uncorrected direct model field"}
    if method == "ridge":
        return {"type": "Ridge", "alpha": float(cfg.get("ridge_alpha", 1.0)), "scaler": "StandardScaler"}
    if method == "elasticnet":
        return {"type": "ElasticNet", "alpha": float(cfg.get("elasticnet_alpha", 0.001)), "l1_ratio": float(cfg.get("elasticnet_l1_ratio", 0.25)), "scaler": "StandardScaler"}
    if method == "lasso":
        return {"type": "Lasso", "alpha": float(cfg.get("lasso_alpha", 0.0005)), "scaler": "StandardScaler"}
    if method == "random_forest":
        return {"type": "RandomForestRegressor", **dict(cfg.get("random_forest_params", {"n_estimators": 300, "max_depth": 12, "min_samples_leaf": 5, "n_jobs": 1})), "scaler": "none"}
    if method in {"gbm", "gradient_boosting"}:
        return {"type": "GradientBoostingRegressor", **dict(cfg.get("gbm_params", {"n_estimators": 300, "max_depth": 3, "learning_rate": 0.03, "subsample": 0.8, "min_samples_leaf": 5})), "scaler": "none"}
    if method == "extratrees":
        return {"type": "ExtraTreesRegressor", **dict(cfg.get("extratrees_params", {"n_estimators": 300, "max_depth": 12, "min_samples_leaf": 5, "n_jobs": 1})), "scaler": "none"}
    if method in {"mlp", "mlp_small", "mlp_medium", "mlp_deep", "ann_tanh", "ann_lbfgs", "mlp_1layer_32", "mlp_1layer_64", "mlp_2layer_64_32", "mlp_2layer_128_64", "mlp_tanh_64"}:
        m = "mlp_medium" if method == "mlp" else method
        p = _mlp_params_for_method(m, cfg, random_seed)
        p = {k: (list(v) if isinstance(v, tuple) else v) for k, v in p.items()}
        p.update({"type": "MLPRegressor", "scaler": "StandardScaler"})
        return p
    if method in {"xgboost", "xgb_d2", "xgb_d3", "xgb_d4", "xgb_slow", "xgb_regularized"}:
        m = "xgb_d3" if method == "xgboost" else method
        p = _xgb_params_for_method(m, cfg, random_seed)
        p.update({"type": "XGBRegressor", "scaler": "none"})
        return p
    if method == "cnn1d_small":
        return {"type": "TorchCNN1DRegressor", **dict(cfg.get("cnn1d_small_params", {})), "scaler": "StandardScaler"}
    if method in {"stack_mean", "stack_ridge", "stack_elasticnet", "stack_nnls"}:
        return get_stack_meta_params(method, cfg)
    return {"type": "unknown"}


def make_model(method: str, cfg: Dict[str, Any], random_seed: int = 42):
    method = method.lower()
    if method == "ridge":
        alpha = float(cfg.get("ridge_alpha", 1.0))
        return Pipeline([("scaler", StandardScaler()), ("model", Ridge(alpha=alpha, random_state=random_seed))])
    if method == "elasticnet":
        return Pipeline([("scaler", StandardScaler()), ("model", ElasticNet(alpha=float(cfg.get("elasticnet_alpha", 0.001)), l1_ratio=float(cfg.get("elasticnet_l1_ratio", 0.25)), random_state=random_seed, max_iter=10000))])
    if method == "lasso":
        return Pipeline([("scaler", StandardScaler()), ("model", Lasso(alpha=float(cfg.get("lasso_alpha", 0.0005)), random_state=random_seed, max_iter=10000))])
    if method == "random_forest":
        params = {"n_estimators": 300, "max_depth": 12, "min_samples_leaf": 5, "n_jobs": 1, "random_state": random_seed}
        params.update(dict(cfg.get("random_forest_params", {})))
        return RandomForestRegressor(**params)
    if method in {"gbm", "gradient_boosting"}:
        params = {"n_estimators": 300, "max_depth": 3, "learning_rate": 0.03, "subsample": 0.8, "min_samples_leaf": 5, "random_state": random_seed}
        params.update(dict(cfg.get("gbm_params", {})))
        return GradientBoostingRegressor(**params)
    if method == "extratrees":
        params = {"n_estimators": 300, "max_depth": 12, "min_samples_leaf": 5, "n_jobs": 1, "random_state": random_seed}
        params.update(dict(cfg.get("extratrees_params", {})))
        return ExtraTreesRegressor(**params)
    if method in {"mlp", "mlp_small", "mlp_medium", "mlp_deep", "ann_tanh", "ann_lbfgs", "mlp_1layer_32", "mlp_1layer_64", "mlp_2layer_64_32", "mlp_2layer_128_64", "mlp_tanh_64"}:
        if method == "mlp" and "mlp_params" in cfg:
            params = dict(cfg.get("mlp_params", {}))
            if "hidden_layer_sizes" in params:
                params["hidden_layer_sizes"] = _as_tuple(params["hidden_layer_sizes"])
            params.setdefault("random_state", random_seed)
        else:
            params = _mlp_params_for_method("mlp_medium" if method == "mlp" else method, cfg, random_seed)
        return Pipeline([("scaler", StandardScaler()), ("model", MLPRegressor(**params))])
    if method in {"xgboost", "xgb_d2", "xgb_d3", "xgb_d4", "xgb_slow", "xgb_regularized"}:
        try:
            from xgboost import XGBRegressor
        except Exception as exc:  # pragma: no cover
            raise ImportError("xgboost is not installed. Remove XGBoost variants or install xgboost.") from exc
        params = dict(cfg.get("xgboost_params", {})) if (method == "xgboost" and "xgboost_params" in cfg) else _xgb_params_for_method("xgb_d3" if method == "xgboost" else method, cfg, random_seed)
        params.setdefault("random_state", random_seed)
        params.setdefault("n_jobs", 1)
        return XGBRegressor(**params)
    if method == "cnn1d_small":
        try:
            import torch  # noqa: F401
            from .torch_models import TorchCNN1DRegressor
            return Pipeline([("scaler", StandardScaler()), ("model", TorchCNN1DRegressor(random_seed=random_seed, **dict(cfg.get("cnn1d_small_params", {}))))])
        except Exception as exc:  # pragma: no cover - depends on HPC torch availability
            # Updated Dr_SFC.docx requests a CNN candidate among the 48 members.
            # PyTorch is not always available on the HPC login environment.  To
            # keep the large experiment from crashing, keep the member in the
            # registry but fall back to a compact MLP benchmark and make the
            # fallback visible in stderr/logs.
            warnings.warn("cnn1d_small requested but PyTorch is unavailable; using MLPRegressor fallback for this member.", RuntimeWarning)
            params = _mlp_params_for_method("mlp_1layer_64", cfg, random_seed)
            return Pipeline([("scaler", StandardScaler()), ("model", MLPRegressor(**params))])
    raise ValueError(f"Unknown model method: {method}")


def get_stack_meta_params(method: str, cfg: Dict[str, Any]) -> Dict[str, Any]:
    st = dict(cfg.get("stacking", {}))
    method = method.lower()
    if method == "stack_mean":
        return {"type": "simple_mean", "fit_intercept": False, "weights": "equal"}
    if method == "stack_ridge":
        return {"type": "Ridge", "alpha": float(st.get("stack_ridge_alpha", 1.0)), "scaler": "StandardScaler"}
    if method == "stack_elasticnet":
        return {"type": "ElasticNet", "alpha": float(st.get("stack_elasticnet_alpha", 0.001)), "l1_ratio": float(st.get("stack_elasticnet_l1_ratio", 0.25)), "scaler": "StandardScaler"}
    if method == "stack_nnls":
        return {"type": "LinearRegression_positive", "positive": True, "fit_intercept": True, "scaler": "none"}
    raise ValueError(f"Unknown stack meta method: {method}")


def make_stack_meta_model(method: str, cfg: Dict[str, Any], random_seed: int = 42):
    st = dict(cfg.get("stacking", {}))
    method = method.lower()
    if method == "stack_mean":
        return None
    if method == "stack_ridge":
        return Pipeline([("scaler", StandardScaler()), ("model", Ridge(alpha=float(st.get("stack_ridge_alpha", 1.0)), random_state=random_seed))])
    if method == "stack_elasticnet":
        return Pipeline([("scaler", StandardScaler()), ("model", ElasticNet(alpha=float(st.get("stack_elasticnet_alpha", 0.001)), l1_ratio=float(st.get("stack_elasticnet_l1_ratio", 0.25)), random_state=random_seed, max_iter=10000))])
    if method == "stack_nnls":
        return LinearRegression(positive=True, fit_intercept=True)
    raise ValueError(f"Unknown stack meta method: {method}")


def extract_linear_weights(model, feature_names):
    """Return intercept/coefs for linear stack meta-models when available."""
    if model is None:
        return {"intercept": 0.0, "weights": {name: 1.0 / max(1, len(feature_names)) for name in feature_names}}
    m = model
    if hasattr(model, "named_steps") and "model" in model.named_steps:
        m = model.named_steps["model"]
    if hasattr(m, "coef_"):
        coef = m.coef_.ravel().tolist()
        intercept = float(getattr(m, "intercept_", 0.0))
        return {"intercept": intercept, "weights": {name: float(coef[i]) for i, name in enumerate(feature_names)}}
    return {"intercept": None, "weights": {}}