from __future__ import annotations

from typing import Any, Dict

import numpy as np
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.linear_model import Ridge
from sklearn.neural_network import MLPRegressor
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 make_model(method: str, cfg: Dict[str, Any], random_seed: int = 42):
    method = method.lower()
    if method == "ridge":
        return Pipeline([
            ("scaler", StandardScaler()),
            ("model", Ridge(alpha=1.0, random_state=random_seed)),
        ])
    if method == "mlp":
        params = dict(cfg.get("mlp_params", {}))
        if "hidden_layer_sizes" in params and isinstance(params["hidden_layer_sizes"], list):
            params["hidden_layer_sizes"] = tuple(params["hidden_layer_sizes"])
        params.setdefault("random_state", random_seed)
        return Pipeline([
            ("scaler", StandardScaler()),
            ("model", MLPRegressor(**params)),
        ])
    if method == "xgboost":
        try:
            from xgboost import XGBRegressor
        except Exception as exc:  # pragma: no cover
            raise ImportError(
                "xgboost is not installed. Install it or remove 'xgboost' from methods in the config."
            ) from exc
        params = dict(cfg.get("xgboost_params", {}))
        params.setdefault("random_state", random_seed)
        return XGBRegressor(**params)
    raise ValueError(f"Unknown model method: {method}")