"""Shared estimation utilities for Phase 4 (statsmodels-based, deterministic)."""
from __future__ import annotations
import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf
from sklearn.metrics import roc_auc_score

SEED = 20261101


def fit_formula(formula: str, data: pd.DataFrame, family, groups: str | None = None):
    """GLM (or OLS if family is None) with cluster-robust SE by `groups`. Returns the fitted result."""
    d = data.dropna(subset=_vars(formula, data) + ([groups] if groups else []))
    if family is None:
        mod = smf.ols(formula, d)
    else:
        mod = smf.glm(formula, d, family=family)
    if groups:
        g = pd.factorize(d[groups])[0]
        res = mod.fit(cov_type="cluster", cov_kwds={"groups": g})
    else:
        res = mod.fit(cov_type="HC1")
    res._n_used = len(d); res._data_used = d
    return res


def _vars(formula: str, data: pd.DataFrame) -> list[str]:
    import re
    toks = set(re.findall(r"[A-Za-z_][A-Za-z0-9_]*", formula))
    return [t for t in toks if t in data.columns]


def coef_row(res, name: str, label: str, extra: dict | None = None) -> dict:
    b = res.params.get(name, np.nan); se = res.bse.get(name, np.nan); p = res.pvalues.get(name, np.nan)
    row = {"model": label, "term": name, "coef": b, "se": se, "ci_lo": b - 1.96 * se, "ci_hi": b + 1.96 * se, "p": p,
           "n": getattr(res, "_n_used", np.nan), "llf": getattr(res, "llf", np.nan)}
    if extra:
        row.update(extra)
    return row


def oos_loglik_binomial(y, p, trials=None):
    """Out-of-sample log-likelihood for binary or fractional outcomes (per observation mean)."""
    p = np.clip(np.asarray(p, float), 1e-6, 1 - 1e-6); y = np.asarray(y, float)
    ll = y * np.log(p) + (1 - y) * np.log(1 - p)
    return float(np.mean(ll))


def binary_metrics(y, p) -> dict:
    y = np.asarray(y, float); p = np.clip(np.asarray(p, float), 1e-6, 1 - 1e-6)
    out = {"oos_loglik": oos_loglik_binomial(y, p), "brier": float(np.mean((y - p) ** 2)), "n": int(len(y)), "events": int(y.sum())}
    if 0 < y.sum() < len(y):
        out["auc"] = float(roc_auc_score(y, p))
        # calibration slope: logit(y) on logit(p)
        lp = np.log(p / (1 - p))
        try:
            cal = sm.Logit(y, sm.add_constant(lp)).fit(disp=0)
            out["calibration_slope"] = float(cal.params[1]); out["calibration_intercept"] = float(cal.params[0])
        except Exception:
            out["calibration_slope"] = np.nan
    else:
        out["auc"] = np.nan; out["calibration_slope"] = np.nan
    return out


def _cat_cols(formula: str) -> list[str]:
    import re
    return re.findall(r"C\(([A-Za-z_][A-Za-z0-9_]*)\)", formula)


def map_unseen_levels(train: pd.DataFrame, test: pd.DataFrame, formula: str) -> pd.DataFrame:
    """Test rows whose categorical fixed-effect level is absent from the training fold get the training mode (A-004);
    the alternative (dropping them) would change the test sample across models."""
    test = test.copy()
    for col in _cat_cols(formula):
        if col in test.columns:
            seen = set(train[col].dropna().unique()); mode = train[col].mode().iloc[0]
            test.loc[~test[col].isin(seen), col] = mode
    return test


def cv_compare(data: pd.DataFrame, formulas: dict[str, str], y: str, family, fold_col: str | None = None,
               n_repeats: int = 10, k: int = 5, groups: str | None = None, fractional: bool = False, block_col: str | None = None) -> pd.DataFrame:
    """Out-of-sample comparison. If fold_col is given → leave-one-group-out on that column; else repeated k-fold.
    Returns one row per (model, fold) with OOS log-likelihood and, for binary y, AUC/Brier."""
    rows = []
    allvars = sorted(set(sum([_vars(f, data) for f in formulas.values()], [])) | {y} | ({groups} if groups else set()))
    d = data.dropna(subset=allvars).reset_index(drop=True)
    if fold_col:
        folds = [(f"loo_{g}", d[fold_col] != g, d[fold_col] == g) for g in sorted(d[fold_col].unique())]
    else:
        rng = np.random.default_rng(SEED); folds = []
        units = d[block_col].values if block_col else np.arange(len(d))          # polity-blocked folds when block_col is given
        uniq = np.unique(units)
        for r in range(n_repeats):
            perm = rng.permutation(len(uniq)); parts = np.array_split(perm, k)
            for i, te in enumerate(parts):
                mask = np.isin(units, uniq[te])
                folds.append((f"rep{r}_fold{i}", ~mask, mask))
    for name, formula in formulas.items():
        for fid, tr, te in folds:
            dtr, dte = d[tr], d[te]
            if dte[y].notna().sum() == 0 or len(dtr) < 20:
                continue
            try:
                mod = smf.glm(formula, dtr, family=family).fit()
                p = mod.predict(map_unseen_levels(dtr, dte, formula))
                ok = np.isfinite(p)
                m = binary_metrics(dte[y][ok], p[ok]) if not fractional else {"oos_loglik": oos_loglik_binomial(dte[y][ok], p[ok]), "n": int(ok.sum())}
                m.update({"model": name, "fold": fid, "n_train": int(len(dtr))}); rows.append(m)
            except Exception as e:
                rows.append({"model": name, "fold": fid, "error": str(e)[:80]})
    return pd.DataFrame(rows)


def fold_win_share(cv: pd.DataFrame, model: str, baseline: str, metric: str = "oos_loglik") -> dict:
    a = cv[cv.model == model].set_index("fold")[metric]; b = cv[cv.model == baseline].set_index("fold")[metric]
    j = a.to_frame("m").join(b.to_frame("b"), how="inner").dropna()
    if len(j) == 0:
        return {"model": model, "baseline": baseline, "folds": 0, "win_share": np.nan, "mean_diff": np.nan}
    return {"model": model, "baseline": baseline, "folds": int(len(j)), "win_share": float((j.m > j.b).mean()), "mean_diff": float((j.m - j.b).mean())}


def holm(pvals: dict[str, float]) -> dict[str, float]:
    items = sorted(pvals.items(), key=lambda kv: kv[1]); m = len(items); out = {}; running = 0
    for i, (k, p) in enumerate(items):
        adj = min(1.0, (m - i) * p); running = max(running, adj); out[k] = running
    return out


def conley_ols(formula: str, data: pd.DataFrame, lat: str, lon: str, cutoff_km: float):
    """OLS with Conley (1999) spatial HAC standard errors, Bartlett kernel with distance cutoff."""
    d = data.dropna(subset=_vars(formula, data) + [lat, lon]).reset_index(drop=True)
    mod = smf.ols(formula, d); res = mod.fit()
    X = mod.exog; u = res.resid.values; n, kx = X.shape
    la, lo = np.radians(d[lat].values), np.radians(d[lon].values)
    dlat = la[:, None] - la[None, :]; dlon = lo[:, None] - lo[None, :]
    h = np.sin(dlat / 2) ** 2 + np.cos(la)[:, None] * np.cos(la)[None, :] * np.sin(dlon / 2) ** 2
    dist = 2 * 6371.0 * np.arcsin(np.sqrt(np.clip(h, 0, 1)))
    w = np.clip(1 - dist / cutoff_km, 0, None)
    Xu = X * u[:, None]
    meat = Xu.T @ w @ Xu
    bread = np.linalg.inv(X.T @ X)
    V = bread @ meat @ bread
    se = np.sqrt(np.diag(V))
    return res, pd.Series(se, index=res.params.index), len(d)


def hinge_fit(data: pd.DataFrame, y: str, base_formula: str, pvar: str, family, grid=None):
    """Profile-likelihood hinge: y ~ base + P + max(P-c,0), c chosen on the given data (call on the TRAINING fold)."""
    d = data.dropna(subset=_vars(base_formula, data) + [y, pvar]).copy()
    grid = np.quantile(d[pvar], np.linspace(0.15, 0.85, 15)) if grid is None else grid
    best = None
    for c in grid:
        d["_hinge"] = np.clip(d[pvar] - c, 0, None)
        try:
            m = smf.glm(f"{base_formula} + _hinge", d, family=family).fit()
            if best is None or m.llf > best[0]:
                best = (m.llf, c, m)
        except Exception:
            continue
    return best