"""Run-level outcome metrics. One dict per run; every number is traceable to ticks/trades/agents/llm frames.

Primary outcomes (pre-registered in docs/identification.md):
  efficiency  : msd_log (mean squared log deviation of mid from the latent value), mad_pct, time_to_absorb,
                post-shock response ratio, variance ratio VR(5) of log mid changes.
  liquidity   : quoted spread (bp), effective spread (bp), depth (lots within 5 levels), Kyle lambda
                (bp per 100 lots signed volume), fill rate, resiliency (half-life of spread after a widening).
  instability : bubble amplitude/duration (P/V - 1 above `bubble_thr` for >= `bubble_min_len` ticks),
                max drawdown, crash count (drawdown >= crash_thr within crash_window), volatility clustering
                (sum of ACF(|r|) lags 1..10), excess kurtosis of tick returns, halts.
  contagion   : (multi-asset) excess return correlation over fundamental correlation, Diebold-Yilmaz total
                spillover from a VAR(2) FEVD, lead-lag cross-correlation.
Secondary: volume, turnover, cancel ratio, P&L by policy, Gini, bankruptcies, MM inventory, LLM validity/cost.
"""
from __future__ import annotations

import numpy as np
import pandas as pd


def _acf(x: np.ndarray, lags: int) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    x = x - x.mean()
    n = len(x)
    denom = (x ** 2).sum()
    if denom <= 0 or n <= lags + 1:
        return np.full(lags, np.nan)
    return np.array([(x[:-k] * x[k:]).sum() / denom for k in range(1, lags + 1)])


def variance_ratio(logp: np.ndarray, q: int = 5) -> float:
    r = np.diff(logp)
    if len(r) < 5 * q or r.var() <= 0:
        return np.nan
    rq = logp[q:] - logp[:-q]
    return float(rq.var() / (q * r.var()))


def gini(w: np.ndarray) -> float:
    w = np.sort(np.clip(np.asarray(w, dtype=float), 0, None))
    n = len(w)
    if n == 0 or w.sum() <= 0:
        return np.nan
    return float((2 * np.arange(1, n + 1) - n - 1).dot(w) / (n * w.sum()))


def bubbles(ratio: np.ndarray, thr: float, min_len: int) -> tuple[int, float, float]:
    """Episodes where ratio = P/V - 1 stays above thr (or below -thr) for >= min_len consecutive ticks."""
    above = np.abs(ratio) > thr
    n, amp, dur = 0, 0.0, 0.0
    i = 0
    while i < len(above):
        if above[i]:
            j = i
            while j < len(above) and above[j]:
                j += 1
            if j - i >= min_len:
                n += 1
                amp = max(amp, float(np.abs(ratio[i:j]).max()))
                dur += j - i
            i = j
        else:
            i += 1
    return n, amp, dur


def drawdowns(p: np.ndarray, crash_thr: float, window: int) -> tuple[float, int]:
    p = np.asarray(p, dtype=float)
    run_max = np.maximum.accumulate(p)
    dd = 1.0 - p / run_max
    max_dd = float(dd.max()) if len(dd) else np.nan
    # crash: drop of >= crash_thr from the max of the trailing `window` ticks
    n_crash = 0
    i = window
    while i < len(p):
        local_max = p[i - window:i].max()
        if p[i] <= (1 - crash_thr) * local_max:
            n_crash += 1
            i += window           # one crash per window
        else:
            i += 1
    return max_dd, n_crash


def kyle_lambda(ticks: pd.DataFrame) -> float:
    """OLS of mid change (bp) on signed volume (per 100 lots), tick by tick."""
    d = ticks.dropna(subset=["mid"])
    if len(d) < 30:
        return np.nan
    y = 1e4 * np.diff(np.log(d["mid"].to_numpy()))
    x = d["signed_volume"].to_numpy()[1:] / 100.0
    if x.var() <= 0:
        return np.nan
    return float(np.cov(x, y)[0, 1] / x.var())


def time_to_absorb(ticks: pd.DataFrame, shock_t: int, tol: float = 0.01, hold: int = 5, max_h: int = 100) -> float:
    """First horizon h after shock_t such that |P/V - 1| < tol for `hold` consecutive ticks."""
    d = ticks[(ticks.t >= shock_t) & (ticks.t < shock_t + max_h)]
    dev = np.abs(d["mid"].to_numpy() / d["full_value"].to_numpy() - 1.0)
    ok = dev < tol
    for h in range(len(ok) - hold + 1):
        if ok[h:h + hold].all():
            return float(h)
    return float(max_h)


def response_ratio(ticks: pd.DataFrame, shock_t: int, h: int = 20) -> float:
    """(P_{t+h} - P_{t-1}) / (V_{t+h} - V_{t-1}) : 1 = full absorption; <1 under-reaction; >1 over-reaction."""
    d = ticks.set_index("t")
    try:
        dp = d.loc[shock_t + h, "mid"] - d.loc[shock_t - 1, "mid"]
        dv = d.loc[shock_t + h, "full_value"] - d.loc[shock_t - 1, "full_value"]
    except KeyError:
        return np.nan
    return float(dp / dv) if abs(dv) > 1e-9 else np.nan


def resiliency(spread: np.ndarray) -> float:
    """Half-life (ticks) of spread deviations from their median: from AR(1) of log spread."""
    s = np.log(np.clip(spread[~np.isnan(spread)], 1e-6, None))
    if len(s) < 30:
        return np.nan
    s = s - np.median(s)
    phi = (s[:-1] * s[1:]).sum() / max((s[:-1] ** 2).sum(), 1e-12)
    phi = min(max(phi, 1e-6), 0.999)
    return float(np.log(0.5) / np.log(phi))


def series_outcomes(ticks: pd.DataFrame, warmup: int = 50, bubble_thr: float = 0.05, bubble_min_len: int = 10,
                    crash_thr: float = 0.10, crash_window: int = 20, shocks: list[dict] | None = None) -> dict:
    out: dict = {}
    for a, d in ticks.groupby("asset"):
        d = d[d.t >= warmup].sort_values("t")
        P, V = d["mid"].to_numpy(), d["full_value"].to_numpy()
        logp = np.log(P)
        r = np.diff(logp)
        sfx = "" if a == 0 else f"_a{a}"
        dev = logp - np.log(V)
        out["msd_log" + sfx] = float(np.mean(dev ** 2))
        out["mad_pct" + sfx] = float(100 * np.mean(np.abs(dev)))
        out["bias_pct" + sfx] = float(100 * np.mean(dev))
        out["vr5" + sfx] = variance_ratio(logp, 5)
        out["vr10" + sfx] = variance_ratio(logp, 10)
        out["vol_bp" + sfx] = float(1e4 * r.std()) if len(r) > 2 else np.nan
        out["kurtosis" + sfx] = float(pd.Series(r).kurt()) if len(r) > 10 else np.nan
        out["acf_absr_1_10" + sfx] = float(np.nansum(_acf(np.abs(r), 10))) if len(r) > 30 else np.nan
        out["acf_r_1" + sfx] = float(_acf(r, 1)[0]) if len(r) > 30 else np.nan
        with np.errstate(all="ignore"):
            import warnings
            warnings.simplefilter("ignore", RuntimeWarning)
            out["spread_bp" + sfx] = float(np.nanmean(1e4 * d["spread"] / d["mid"])) if d["spread"].notna().any() else np.nan
        out["eff_spread_bp" + sfx] = float(1e4 * np.nanmean(d["eff_spread"] / d["mid"]))
        out["depth5" + sfx] = float(np.mean(d["depth_bid5"] + d["depth_ask5"]))
        out["lambda_bp_100" + sfx] = kyle_lambda(d)
        out["resiliency" + sfx] = resiliency(d["spread"].to_numpy())
        nb, amp, dur = bubbles(P / V - 1.0, bubble_thr, bubble_min_len)
        out["n_bubbles" + sfx], out["bubble_amp" + sfx], out["bubble_dur" + sfx] = nb, amp, dur
        out["max_dd" + sfx], out["n_crashes" + sfx] = drawdowns(P, crash_thr, crash_window)
        out["volume" + sfx] = float(d["volume"].sum())
        out["n_halts" + sfx] = int(d["n_halts"].max()) if "n_halts" in d else 0
        out["share_ticks_halted" + sfx] = float(d["halted"].mean()) if "halted" in d else 0.0
        if shocks:
            tta, rr = [], []
            for s in shocks:
                if s.get("shock") in ("fundamental_jump", "news_burst") and s["t"] >= warmup:
                    tta.append(time_to_absorb(ticks[ticks.asset == a], s["t"]))
                    rr.append(response_ratio(ticks[ticks.asset == a], s["t"]))
            if tta:
                out["time_to_absorb" + sfx] = float(np.mean(tta))
                out["response_ratio" + sfx] = float(np.nanmean(rr))
    assets = sorted(ticks["asset"].unique())
    if len(assets) >= 2:
        piv = ticks[ticks.t >= warmup].pivot(index="t", columns="asset", values="mid")
        fv = ticks[ticks.t >= warmup].pivot(index="t", columns="asset", values="full_value")
        rp, rv = np.log(piv).diff().dropna(), np.log(fv).diff().dropna()
        cp, cv = rp.corr().to_numpy(), rv.corr().to_numpy()
        iu = np.triu_indices(len(assets), 1)
        out["ret_corr"] = float(np.nanmean(cp[iu]))
        out["fund_corr"] = float(np.nanmean(cv[iu]))
        out["excess_comovement"] = out["ret_corr"] - out["fund_corr"]
        out["spillover_dy"] = diebold_yilmaz(rp.to_numpy())
        out["leadlag_01"] = float(np.corrcoef(rp.iloc[:-1, 0], rp.iloc[1:, 1])[0, 1]) if len(rp) > 10 else np.nan
    return out


def diebold_yilmaz(R: np.ndarray, p: int = 2, h: int = 10) -> float:
    """Total spillover index from the generalized FEVD of a VAR(p) (Diebold–Yilmaz 2012)."""
    T, K = R.shape
    if T < 10 * K * p + 10:
        return np.nan
    X = np.hstack([R[p - j - 1:T - j - 1] for j in range(p)])
    Y = R[p:]
    X = np.hstack([np.ones((len(Y), 1)), X])
    B, *_ = np.linalg.lstsq(X, Y, rcond=None)
    E = Y - X @ B
    Sigma = np.cov(E.T)
    A = [np.eye(K)]
    coefs = [B[1 + K * j:1 + K * (j + 1)].T for j in range(p)]
    for i in range(1, h):
        A.append(sum(coefs[j] @ A[i - j - 1] for j in range(p) if i - j - 1 >= 0))
    sig = np.diag(Sigma)
    theta = np.zeros((K, K))
    for i in range(K):
        for j in range(K):
            num = sum((A[l][i] @ Sigma[:, j]) ** 2 for l in range(h)) / sig[j]
            den = sum(A[l][i] @ Sigma @ A[l][i] for l in range(h))
            theta[i, j] = num / den
    theta = theta / theta.sum(axis=1, keepdims=True)
    return float(100 * (theta.sum() - np.trace(theta)) / K)


def agent_outcomes(agents: pd.DataFrame, trades: pd.DataFrame, llm: pd.DataFrame, n_ticks: int) -> dict:
    out: dict = {}
    g = agents.groupby("policy")["pnl_pct"].mean()
    for k, v in g.items():
        out[f"pnl_{k}"] = float(v)
    out["gini_final"] = gini(agents["final_wealth"].to_numpy())
    out["gini_initial"] = gini(agents["initial_wealth"].to_numpy())
    out["n_bankrupt"] = int(agents["frozen"].sum())
    out["cancel_ratio"] = float(agents["n_cancels"].sum() / max(agents["n_orders"].sum(), 1))
    out["reject_ratio"] = float(agents["n_rejects"].sum() / max(agents["n_orders"].sum(), 1))
    out["fill_ratio"] = float(agents["n_fills"].sum() / max(agents["n_orders"].sum(), 1))
    out["turnover"] = float(trades["qty"].sum() / max(agents["position"].abs().sum() + 1, 1)) if len(trades) else 0.0
    mm = agents[agents.policy == "market_maker"]
    if len(mm):
        out["mm_pnl"] = float(mm["pnl_pct"].mean())
        out["mm_inventory_abs"] = float(mm["position"].abs().mean())
    if len(llm):
        out["llm_valid_rate"] = float(llm["valid"].mean())
        out["llm_hold_rate"] = float((llm["action"] == "hold").mean())
        out["llm_market_share"] = float((llm["order_type"] == "market").mean())
        out["llm_usd"] = float(llm["usd"].sum())
        out["llm_tokens"] = float(llm["input_tokens"].sum() + llm["output_tokens"].sum())
        out["llm_latency_s"] = float(llm["latency_s"].mean())
        out["llm_cached_hits"] = float(llm["cached_hit"].mean())
        # synchrony and order-flow correlation among LLM agents: decisions grouped by decision round
        d = llm[llm.valid & (llm.action == "place")].copy()
        if len(d) > 5:
            d["sign"] = np.where(d["side"] == "buy", 1, -1)
            rounds = d.groupby("t")["sign"]
            sync = rounds.apply(lambda s: abs(s.mean()) if len(s) > 1 else np.nan).dropna()
            out["llm_synchrony"] = float(sync.mean()) if len(sync) else np.nan
            out["llm_buy_share"] = float((d["sign"] > 0).mean())
            out["llm_pairwise_corr"] = pairwise_sign_corr(d)
    return out


def pairwise_sign_corr(d: pd.DataFrame, window: int = 20) -> float:
    """Average pairwise correlation of agents' net order sign aggregated per `window` ticks."""
    d = d.copy()
    d["w"] = d["t"] // window
    piv = d.pivot_table(index="w", columns="agent_id", values="sign", aggfunc="sum").fillna(0.0)
    if piv.shape[1] < 2 or piv.shape[0] < 3:
        return np.nan
    c = piv.corr().to_numpy()
    iu = np.triu_indices(piv.shape[1], 1)
    vals = c[iu]
    return float(np.nanmean(vals)) if np.isfinite(vals).any() else np.nan


def flow_corr_by_policy(trades: pd.DataFrame, agents: pd.DataFrame, window: int = 20) -> dict:
    """Average pairwise correlation of agents' net signed traded volume per `window` ticks, by policy group
    (a classical analogue of llm_pairwise_corr that needs no decision log)."""
    if len(trades) < 50:
        return {}
    t = trades.copy()
    t["w"] = (t["time"] // window).astype(int)
    buy = t.groupby(["w", "buyer"])["qty"].sum().rename("q").reset_index().rename(columns={"buyer": "agent_id"})
    sell = t.groupby(["w", "seller"])["qty"].sum().rename("q").reset_index().rename(columns={"seller": "agent_id"})
    sell["q"] = -sell["q"]
    flow = pd.concat([buy, sell]).groupby(["w", "agent_id"])["q"].sum().reset_index()
    piv = flow.pivot(index="w", columns="agent_id", values="q").fillna(0.0)
    pol = agents.set_index("agent_id")["policy"]
    out = {}
    for name, grp in pol.groupby(pol):
        ids = [i for i in grp.index if i in piv.columns]
        if len(ids) < 3 or piv.shape[0] < 4:
            continue
        c = piv[ids].corr().to_numpy()
        iu = np.triu_indices(len(ids), 1)
        vals = c[iu]
        if np.isfinite(vals).any():
            out[f"flow_corr_{name}"] = float(np.nanmean(vals))
    return out


def run_outcomes(res, bubble_thr: float = 0.05, bubble_min_len: int = 10, crash_thr: float = 0.10, crash_window: int = 20) -> dict:
    cfg = res.config
    out = {"name": cfg.name, "seed": cfg.seed, "config_hash": cfg.config_hash(), "n_traders": cfg.population.n_traders,
           "llm_share": cfg.population.llm_share, **{f"tr_{k}": v for k, v in cfg.treatment.items()}}
    out.update({k: res.manifest[k] for k in ("llm_share_traders", "llm_share_orders", "llm_share_volume", "n_llm_decisions", "n_trades", "wall_s")})
    out.update(series_outcomes(res.ticks, cfg.warmup, bubble_thr, bubble_min_len, crash_thr, crash_window, res.manifest.get("shocks")))
    out.update(agent_outcomes(res.agents, res.trades, res.llm, cfg.n_ticks))
    out.update(flow_corr_by_policy(res.trades, res.agents))
    return out