from __future__ import annotations

from dataclasses import dataclass
import pandas as pd


@dataclass(frozen=True)
class ScalerState:
    means: dict[str, float]
    stds: dict[str, float]


def fit_standard_scaler(df: pd.DataFrame, cols: list[str]) -> ScalerState:
    """Fit z-score stats on train data only."""
    means = {c: float(df[c].mean()) for c in cols}
    stds = {c: float(df[c].std(ddof=0) or 1.0) for c in cols}
    return ScalerState(means=means, stds=stds)


def apply_standard_scaler(df: pd.DataFrame, state: ScalerState) -> pd.DataFrame:
    """Apply pre-fitted z-score stats (no refit)."""
    out = df.copy()
    for c, m in state.means.items():
        out[c] = (out[c] - m) / state.stds[c]
    return out


def fit_one_hot(df: pd.DataFrame, cols: list[str]) -> dict[str, list[str]]:
    """Capture category vocabulary per column."""
    return {c: sorted(df[c].dropna().astype(str).unique().tolist()) for c in cols}


def apply_one_hot(df: pd.DataFrame, vocab: dict[str, list[str]]) -> pd.DataFrame:
    """One-hot encode using fixed vocab to avoid leakage/drift surprises."""
    out = df.copy()
    for col, cats in vocab.items():
        series = out[col].astype(str)
        for cat in cats:
            out[f"{col}__{cat}"] = (series == cat).astype(int)
        out = out.drop(columns=[col])
    return out
