from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import Ridge
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, SplineTransformer, StandardScaler


TARGET_FORBIDDEN = {"log10_discrepancy", "log10_g_obs", "g_obs", "g_obs_ms2"}


@dataclass(frozen=True)
class FeatureSpec:
    name: str
    local_gbar: list[str]
    radius: list[str]
    numeric: list[str]
    categorical: list[str]

    @property
    def all_features(self) -> list[str]:
        return self.local_gbar + self.radius + self.numeric + self.categorical


def feature_registry(config: dict[str, Any]) -> dict[str, FeatureSpec]:
    fg = config["feature_groups"]
    local_gbar = [config["columns"]["local_gbar"]]
    radius = [config["columns"]["radius"]]
    ordinary = fg["ordinary_numeric"]
    si_ak = fg["struc_i_Ak_numeric"]
    si_rho = fg["struc_i_rho_numeric"]
    si_cat = fg["struc_i_categorical"]
    sp_num = fg["struc_perc_numeric"]
    sp_cat = fg["struc_perc_categorical"]

    def make(name, include_gbar=True, include_radius=False, num=None, cat=None):
        spec = FeatureSpec(
            name=name,
            local_gbar=local_gbar if include_gbar else [],
            radius=radius if include_radius else [],
            numeric=list(num or []),
            categorical=list(cat or []),
        )
        forbidden = TARGET_FORBIDDEN & set(spec.all_features)
        if forbidden:
            raise RuntimeError(f"Target leakage in {name}: {sorted(forbidden)}")
        return spec

    registry = {
        "M0": make("M0", include_radius=False),
        "M1": make("M1", include_radius=True),
        "M2": make("M2", include_radius=True, num=ordinary),
        "M_I": make("M_I", include_radius=True, num=ordinary + si_ak + si_rho, cat=si_cat),
        "M_P": make("M_P", include_radius=True, num=ordinary + sp_num, cat=sp_cat),
        "M_IP": make("M_IP", include_radius=True, num=ordinary + si_ak + si_rho + sp_num, cat=si_cat + sp_cat),
        "M_IP_MINUS_AK": make("M_IP_MINUS_AK", include_radius=True, num=ordinary + si_rho + sp_num, cat=si_cat + sp_cat),
        "M_IP_MINUS_RHO": make("M_IP_MINUS_RHO", include_radius=True, num=ordinary + si_ak + sp_num, cat=si_cat + sp_cat),
        "M_IP_MINUS_P_NUMERIC": make("M_IP_MINUS_P_NUMERIC", include_radius=True, num=ordinary + si_ak + si_rho, cat=si_cat + sp_cat),
        "M_IP_MINUS_P_CATEGORICAL": make("M_IP_MINUS_P_CATEGORICAL", include_radius=True, num=ordinary + si_ak + si_rho + sp_num, cat=si_cat),
    }
    return registry


def build_pipeline(spec: FeatureSpec, config: dict[str, Any]) -> Pipeline:
    pp = config["preprocessing"]
    transformers = []
    if spec.local_gbar:
        gbar_pipe = Pipeline([
            ("imputer", SimpleImputer(strategy="median")),
            ("spline", SplineTransformer(
                n_knots=int(pp["gbar_spline_knots"]),
                degree=int(pp["spline_degree"]),
                include_bias=False,
                extrapolation="linear",
            )),
            ("scale", StandardScaler()),
        ])
        transformers.append(("gbar_spline", gbar_pipe, spec.local_gbar))
    if spec.radius:
        radius_pipe = Pipeline([
            ("imputer", SimpleImputer(strategy="median")),
            ("spline", SplineTransformer(
                n_knots=int(pp["radius_spline_knots"]),
                degree=int(pp["spline_degree"]),
                include_bias=False,
                extrapolation="linear",
            )),
            ("scale", StandardScaler()),
        ])
        transformers.append(("radius_spline", radius_pipe, spec.radius))
    if spec.numeric:
        numeric_pipe = Pipeline([
            ("imputer", SimpleImputer(strategy="median")),
            ("scale", StandardScaler()),
        ])
        transformers.append(("numeric", numeric_pipe, spec.numeric))
    if spec.categorical:
        categorical_pipe = Pipeline([
            ("imputer", SimpleImputer(strategy="most_frequent")),
            ("onehot", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
        ])
        transformers.append(("categorical", categorical_pipe, spec.categorical))

    if not transformers:
        raise RuntimeError(f"Model {spec.name} has no features.")
    preprocessor = ColumnTransformer(transformers=transformers, remainder="drop", sparse_threshold=0)
    estimator = Ridge(
        alpha=float(config["estimator"]["alpha"]),
        fit_intercept=bool(config["estimator"].get("fit_intercept", True)),
    )
    return Pipeline([("preprocess", preprocessor), ("model", estimator)])
