"""Model C factorial experiments with the Bayesian mean-variance control (no LLM): market design × delegation share ×
model diversity × leverage, 20 seeds (common random numbers across cells), plus stress shocks.

Outputs: results/model_c/factorial.csv (one row per cell × seed), results/model_c/stress.csv,
results/model_c/summary.csv (means and bootstrap 95% CIs by cell).
Run: python3 scripts/30_model_c.py [--quick]
"""
from __future__ import annotations

import itertools
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'src'))
from age.markets.simulation import MarketConfig, run_market  # noqa: E402

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / 'results' / 'model_c'
OUT.mkdir(parents=True, exist_ok=True)
QUICK = '--quick' in sys.argv
SEEDS = range(1, 6) if QUICK else range(1, 21)
T = 80 if QUICK else 160
KEYS = ['vol', 'of_corr', 'of_corr_deleg', 'of_corr_human', 'pressure', 'r2_common', 'kyle_lambda', 'spread', 'depth',
        'dispersion', 'defaults', 'forced', 'losses', 'leverage', 'n_trades', 'mm_inventory']


def cell(design, D, M, lev, seed, **kw):
    out = run_market(MarketConfig(design=design, delegation_share=D, M=M, max_leverage=lev, seed=seed, T=T, **kw))
    row = dict(design=design, D=D, M=M, leverage_cap=lev, seed=seed)
    row.update({k: out.get(k, np.nan) for k in KEYS})
    for k in ('shock_return', 'shock_forced', 'shock_defaults', 'shock_min_return', 'shock_trough', 'shock_losses'):
        row[k] = out.get(k, np.nan)
    return row


STRESS_ONLY = '--stress-only' in sys.argv
rows = []
for design, D, M, lev in ([] if STRESS_ONLY else itertools.product(['walrasian', 'cda', 'bilateral'], [0.0, 0.25, 0.5, 0.75, 1.0], [1, 2, 4, 8], [0.0, 0.5])):
    if D == 0.0 and M > 1:
        continue
    for s in SEEDS:
        rows.append(cell(design, D, M, lev, s))
df = pd.DataFrame(rows) if rows else pd.read_csv(OUT / 'factorial.csv')
df.to_csv(OUT / 'factorial.csv', index=False)

# stress: a 3-sd common forecast shock at period T/2, with and without leverage, by design and M
srows = []
STRESS = dict(sig_A=0.02, rho_d=0.95, liquidation_impact=1.0, mm_tick=0.0025, mm_levels=40, mm_depth=0.1,
              mm_inventory_max=10, auto_repay=False)
for design, M, persist, margin, pre in itertools.product(['walrasian', 'cda'], [1, 4], [1, 3], [0.5, 0.2], [0.0, 0.95]):
    for s in SEEDS:
        r = cell(design, 1.0, M, 0.0, s, shock_period=T // 2, shock_size=-10.0, shock_persist=persist, margin=margin,
                 preload_leverage=pre, **STRESS)
        r.update(margin=margin, persist=persist, preload=pre, control=False)
        srows.append(r)
    for s in SEEDS:   # control: same seeds, no shock
        r = cell(design, 1.0, M, 0.0, s, margin=margin, preload_leverage=pre, **STRESS)
        r.update(margin=margin, persist=persist, preload=pre, control=True, shock_return=np.nan)
        srows.append(r)
sdf = pd.DataFrame(srows)
sdf['control'] = sdf.get('control', False).fillna(False)
sdf.to_csv(OUT / 'stress.csv', index=False)


def boot_ci(x, n=1000, seed=0):
    x = np.asarray(x, float)
    x = x[~np.isnan(x)]
    if len(x) == 0:
        return np.nan, np.nan, np.nan
    rng = np.random.default_rng(seed)
    means = [rng.choice(x, len(x)).mean() for _ in range(n)]
    return x.mean(), np.percentile(means, 2.5), np.percentile(means, 97.5)


summ = []
for keys, g in df.groupby(['design', 'D', 'M', 'leverage_cap']):
    row = dict(zip(['design', 'D', 'M', 'leverage_cap'], keys))
    for k in KEYS:
        mu, lo, hi = boot_ci(g[k])
        row[k], row[k + '_lo'], row[k + '_hi'] = mu, lo, hi
    summ.append(row)
pd.DataFrame(summ).to_csv(OUT / 'summary.csv', index=False)
ssumm = []
for keys, g in sdf.groupby(['design', 'M', 'persist', 'margin', 'preload', 'control']):
    row = dict(zip(['design', 'M', 'persist', 'margin', 'preload', 'control'], keys))
    for k in ['shock_return', 'shock_min_return', 'shock_trough', 'shock_forced', 'shock_defaults', 'shock_losses', 'vol', 'defaults', 'forced', 'losses', 'kyle_lambda', 'mm_inventory', 'leverage']:
        mu, lo, hi = boot_ci(g[k])
        row[k], row[k + '_lo'], row[k + '_hi'] = mu, lo, hi
    ssumm.append(row)
pd.DataFrame(ssumm).to_csv(OUT / 'stress_summary.csv', index=False)
print(pd.DataFrame(summ)[['design', 'D', 'M', 'leverage_cap', 'vol', 'of_corr', 'of_corr_deleg', 'r2_common', 'defaults']].to_string())
print(pd.DataFrame(ssumm)[['design', 'M', 'persist', 'margin', 'preload', 'control', 'shock_return', 'shock_trough', 'shock_forced', 'shock_defaults', 'leverage']].to_string())