src/utils/models.py

"""Model fitting helpers used across initiative pages.

All routines return small, consistent dataclasses so the page-side code
can render results uniformly.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Iterable, Optional

import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf


# --------------------------------------------------------------------------
# Result containers
# --------------------------------------------------------------------------
@dataclass
class CoefRow:
    term: str
    coef: float
    se: float
    t: float
    p: float
    ci_low: float
    ci_high: float

    def as_dict(self) -> dict:
        return self.__dict__


@dataclass
class FitResult:
    """Lightweight wrapper around a fitted statsmodels regression."""
    coefs: pd.DataFrame
    n_obs: int
    n_clusters: Optional[int]
    r2: Optional[float]
    formula: str
    raw: object = field(repr=False, default=None)

    @classmethod
    def from_results(cls, res, formula: str, *, raw_terms: Iterable[str] | None = None,
                     n_clusters: Optional[int] = None) -> "FitResult":
        ci = res.conf_int()
        rows = []
        for term in res.params.index:
            if raw_terms is not None and term not in raw_terms:
                continue
            rows.append(CoefRow(
                term=term,
                coef=float(res.params[term]),
                se=float(res.bse[term]),
                t=float(res.tvalues[term]),
                p=float(res.pvalues[term]),
                ci_low=float(ci.loc[term, 0]),
                ci_high=float(ci.loc[term, 1]),
            ).as_dict())
        coefs = pd.DataFrame(rows)
        try:
            r2 = float(res.rsquared)
        except AttributeError:
            r2 = None
        return cls(coefs=coefs, n_obs=int(res.nobs), n_clusters=n_clusters,
                   r2=r2, formula=formula, raw=res)


# --------------------------------------------------------------------------
# Two-way fixed effects DiD / TWFE
# --------------------------------------------------------------------------
def fit_twfe(
    df: pd.DataFrame, *,
    outcome: str,
    unit: str = "fips",
    period: str = "year",
    treatment_terms: list[str],
    controls: Optional[list[str]] = None,
    cluster_col: Optional[str] = None,
) -> FitResult:
    """Two-way fixed effects regression.

    Specification:
        y_{it} = alpha_i + lambda_t + sum_k beta_k * Treat_k_{it} + gamma X_{it} + e

    Implemented via OLS with explicit dummies (small panels, fast enough).
    """
    controls = controls or []
    needed = [outcome, unit, period] + treatment_terms + controls
    needed = [c for c in needed if c is not None]
    work = df.dropna(subset=needed).copy()
    formula_parts = [outcome, "~"]
    formula_parts.append(" + ".join(treatment_terms + controls))
    formula_parts.append(f" + C({unit}) + C({period})")
    formula = " ".join(formula_parts)

    model = smf.ols(formula, data=work)
    if cluster_col is not None:
        groups = work[cluster_col]
        res = model.fit(cov_type="cluster", cov_kwds={"groups": groups})
        n_clusters = int(work[cluster_col].nunique())
    else:
        res = model.fit()
        n_clusters = None

    return FitResult.from_results(
        res, formula,
        raw_terms=treatment_terms + controls,
        n_clusters=n_clusters,
    )


# --------------------------------------------------------------------------
# Event study around a treatment time
# --------------------------------------------------------------------------
def build_event_time(
    df: pd.DataFrame, *,
    unit: str,
    period: str,
    treat_unit_col: str,
    event_period_col: str,
    leads: int = 4, lags: int = 6,
    reference_lead: int = -1,
) -> pd.DataFrame:
    """Construct event-time indicators for a staggered-adoption event study.

    Returns the original df augmented with columns ev_m4, ..., ev_p6 plus
    a binned ev_minus / ev_plus for periods outside the window.
    Reference period (reference_lead) is omitted so coefficients are
    interpreted relative to that period.
    """
    out = df.copy()
    et = out[period] - out[event_period_col]
    et = et.where(out[treat_unit_col] == 1)  # NaN for never-treated

    # Bin tails
    et_binned = et.copy()
    et_binned = et_binned.where(et_binned >= -leads, other=-(leads + 99))
    et_binned = et_binned.where(et_binned <= lags, other=(lags + 99))

    cols = []
    for k in range(-leads, lags + 1):
        if k == reference_lead:
            continue
        col = f"ev_{('m' if k < 0 else 'p')}{abs(k)}"
        out[col] = ((et_binned == k) & (out[treat_unit_col] == 1)).astype(int)
        cols.append(col)
    # binned tails
    out["ev_pre"] = ((et_binned == -(leads + 99)) & (out[treat_unit_col] == 1)).astype(int)
    out["ev_post"] = ((et_binned == (lags + 99)) & (out[treat_unit_col] == 1)).astype(int)
    return out, cols + ["ev_pre", "ev_post"]


def event_study_table(
    fit: FitResult, *, leads: int = 4, lags: int = 6,
    reference_lead: int = -1,
) -> pd.DataFrame:
    """Reshape an event-study fit into a long table for plotting."""
    rows = []
    rows.append({
        "event_time": reference_lead,
        "coef": 0.0, "se": 0.0,
        "ci_low": 0.0, "ci_high": 0.0, "p": 1.0,
    })
    for k in range(-leads, lags + 1):
        if k == reference_lead:
            continue
        col = f"ev_{('m' if k < 0 else 'p')}{abs(k)}"
        m = fit.coefs[fit.coefs["term"] == col]
        if m.empty:
            continue
        m = m.iloc[0]
        rows.append({
            "event_time": k,
            "coef": m["coef"], "se": m["se"],
            "ci_low": m["ci_low"], "ci_high": m["ci_high"], "p": m["p"],
        })
    return pd.DataFrame(rows).sort_values("event_time").reset_index(drop=True)


# --------------------------------------------------------------------------
# Logistic regression with marginal effects
# --------------------------------------------------------------------------
def fit_logit_marginal(
    df: pd.DataFrame, *,
    outcome: str, predictors: list[str],
) -> tuple[FitResult, pd.DataFrame]:
    """Logit + average marginal effects (AME)."""
    work = df.dropna(subset=[outcome] + predictors).copy()
    formula = f"{outcome} ~ {' + '.join(predictors)}"
    res = smf.logit(formula, data=work).fit(disp=False)
    fit = FitResult.from_results(res, formula, raw_terms=predictors)
    me = res.get_margeff(at="overall", method="dydx")
    me_df = pd.DataFrame({
        "term": me.results_table_data[0][1:],
        "marginal_effect": me.margeff,
        "se": me.margeff_se,
        "p": me.pvalues,
    }) if False else None  # The API differs across versions; build manually:
    margeff = res.get_margeff()
    me_summary = margeff.summary_frame()
    me_summary = me_summary.reset_index().rename(columns={
        "index": "term", "dy/dx": "marginal_effect",
        "Std. Err.": "se", "Pr(>|z|)": "p",
    })
    return fit, me_summary


# --------------------------------------------------------------------------
# Parallel-trends pre-period test
# --------------------------------------------------------------------------
def parallel_trends_pvalue(
    df: pd.DataFrame, *,
    outcome: str, unit: str, period: str,
    treat_unit_col: str, pre_periods: list[int],
) -> dict:
    """Joint F-test on pre-period treated × period interactions."""
    pre = df[df[period].isin(pre_periods)].copy()
    pre["t"] = pre[period] - min(pre_periods)
    pre["interact"] = pre[treat_unit_col] * pre["t"]
    formula = f"{outcome} ~ {treat_unit_col} + C({period}) + interact + C({unit})"
    res = smf.ols(formula, data=pre).fit()
    coef = float(res.params.get("interact", np.nan))
    se = float(res.bse.get("interact", np.nan))
    p = float(res.pvalues.get("interact", np.nan))
    return {"coef": coef, "se": se, "p": p, "n": int(res.nobs)}


# --------------------------------------------------------------------------
# Convenience: pretty p-value
# --------------------------------------------------------------------------
def stars(p: float) -> str:
    if pd.isna(p):
        return ""
    if p < 0.001:
        return "***"
    if p < 0.01:
        return "**"
    if p < 0.05:
        return "*"
    if p < 0.1:
        return "·"
    return ""