"""Run-level inference. The unit of observation is one simulation run (cell × seed); trades and ticks
within a run are never treated as independent replicates.

* paired_effects : matched-seed differences treatment − control with t-based 95 % CI, Cohen's d (paired),
                   and a Wilcoxon signed-rank p-value.
* ols_seed_fe    : y = α_seed + Σ_k β_k 1[cell = k] + ε with cluster-robust (by seed) standard errors.
* dose_response  : effect at each treatment level vs control (paired), with Holm-adjusted p-values.
* local_projection: β_h from (gap_{t+h} − gap_{t-1}) regressed on treatment, h = 0..H, around a shock.
"""
from __future__ import annotations

import numpy as np
import pandas as pd
from scipy import stats


def holm(p: np.ndarray) -> np.ndarray:
    p = np.asarray(p, dtype=float)
    m = np.isfinite(p).sum()
    order = np.argsort(np.where(np.isfinite(p), p, np.inf))
    adj = np.full_like(p, np.nan)
    running = 0.0
    for rank, idx in enumerate(order[:m]):
        val = min(1.0, (m - rank) * p[idx])
        running = max(running, val)
        adj[idx] = running
    return adj


def paired_effects(df: pd.DataFrame, outcome: str, treat_col: str, treat_val, control_val, seed_col: str = "seed") -> dict:
    a = df[df[treat_col] == treat_val].set_index(seed_col)[outcome]
    b = df[df[treat_col] == control_val].set_index(seed_col)[outcome]
    idx = a.index.intersection(b.index)
    d = (a.loc[idx] - b.loc[idx]).astype(float).dropna()
    n = len(d)
    if n < 2:
        return {"outcome": outcome, "level": treat_val, "n": n, "diff": float(d.mean()) if n else np.nan}
    se = d.std(ddof=1) / np.sqrt(n)
    tcrit = stats.t.ppf(0.975, n - 1)
    t = d.mean() / se if se > 0 else np.nan
    p = 2 * stats.t.sf(abs(t), n - 1) if np.isfinite(t) else np.nan
    try:
        pw = stats.wilcoxon(d).pvalue if (d != 0).any() else 1.0
    except ValueError:
        pw = np.nan
    return {"outcome": outcome, "level": treat_val, "n": n, "control_mean": float(b.loc[idx].mean()), "treat_mean": float(a.loc[idx].mean()),
            "diff": float(d.mean()), "se": float(se), "ci_lo": float(d.mean() - tcrit * se), "ci_hi": float(d.mean() + tcrit * se),
            "t": float(t), "p": float(p), "p_wilcoxon": float(pw), "cohen_d": float(d.mean() / d.std(ddof=1)) if d.std(ddof=1) > 0 else np.nan,
            "pct_change": float(100 * d.mean() / abs(b.loc[idx].mean())) if b.loc[idx].mean() != 0 else np.nan}


def dose_response(df: pd.DataFrame, outcome: str, treat_col: str, control_val, levels: list | None = None, seed_col: str = "seed") -> pd.DataFrame:
    levels = levels or sorted(v for v in df[treat_col].unique() if v != control_val)
    rows = [paired_effects(df, outcome, treat_col, lv, control_val, seed_col) for lv in levels]
    out = pd.DataFrame(rows)
    if "p" in out:
        out["p_holm"] = holm(out["p"].to_numpy())
    return out


def ols_seed_fe(df: pd.DataFrame, outcome: str, cell_col: str, control_val, seed_col: str = "seed") -> pd.DataFrame:
    """Cluster-robust (by seed) OLS with seed fixed effects; returns one row per non-control cell."""
    import statsmodels.formula.api as smf
    d = df[[outcome, cell_col, seed_col]].dropna().copy()
    d["cell"] = d[cell_col].astype(str)
    ctrl = str(control_val)
    cats = [ctrl] + sorted(c for c in d["cell"].unique() if c != ctrl)
    d["cell"] = pd.Categorical(d["cell"], categories=cats)
    d["seed_f"] = d[seed_col].astype(str)
    d = d.rename(columns={outcome: "y"})
    m = smf.ols("y ~ C(cell) + C(seed_f)", data=d).fit(cov_type="cluster", cov_kwds={"groups": d[seed_col]})
    rows = []
    for c in cats[1:]:
        name = f"C(cell)[T.{c}]"
        if name in m.params:
            ci = m.conf_int().loc[name]
            rows.append({"outcome": outcome, "level": c, "beta": m.params[name], "se": m.bse[name], "ci_lo": ci[0], "ci_hi": ci[1],
                         "p": m.pvalues[name], "n": int(m.nobs), "r2": m.rsquared})
    out = pd.DataFrame(rows)
    if len(out):
        out["p_holm"] = holm(out["p"].to_numpy())
    return out


def local_projection(ticks_by_run: dict, treat: dict, shock_t: int, H: int = 40, control_val=0.0) -> pd.DataFrame:
    """ticks_by_run: {run_id: ticks DataFrame}; treat: {run_id: treatment value}. Gap = log(mid) - log(full_value).
    For each horizon h, regress (gap_{t+h} - gap_{t-1}) on 1[treated] across runs (paired by seed not required)."""
    rows = []
    gaps = {}
    for rid, tk in ticks_by_run.items():
        d = tk[tk.asset == 0].set_index("t")
        g = np.log(d["mid"]) - np.log(d["full_value"])
        gaps[rid] = g
    for h in range(0, H + 1):
        y, x = [], []
        for rid, g in gaps.items():
            if shock_t - 1 in g.index and shock_t + h in g.index:
                y.append(g.loc[shock_t + h] - g.loc[shock_t - 1])
                x.append(1.0 if treat[rid] != control_val else 0.0)
        y, x = np.array(y), np.array(x)
        if len(y) < 4 or x.var() == 0:
            continue
        X = np.column_stack([np.ones_like(x), x])
        beta, *_ = np.linalg.lstsq(X, y, rcond=None)
        resid = y - X @ beta
        # HC1 robust SE
        XtX_inv = np.linalg.inv(X.T @ X)
        V = XtX_inv @ (X.T * resid ** 2) @ X @ XtX_inv * len(y) / (len(y) - 2)
        rows.append({"h": h, "beta": beta[1], "se": np.sqrt(V[1, 1]), "control_path": float(y[x == 0].mean()), "treated_path": float(y[x == 1].mean())})
    return pd.DataFrame(rows)


def power_seeds(d: float, power: float = 0.8, alpha: float = 0.05) -> int:
    """Paired design: seeds needed to detect a standardized effect d (one-sample t approximation)."""
    z = stats.norm.ppf(1 - alpha / 2) + stats.norm.ppf(power)
    return int(np.ceil((z / d) ** 2))


def summarize_cells(df: pd.DataFrame, outcomes: list[str], cell_col: str = "name") -> pd.DataFrame:
    g = df.groupby(cell_col)[outcomes]
    m, s, n = g.mean(), g.std(ddof=1), g.size()
    out = m.copy()
    for c in outcomes:
        out[c + "_se"] = s[c] / np.sqrt(n)
    out["n_runs"] = n
    return out