"""The gate at the deployed cell with the inputs section 25 left out, wallet MC 5,000 paths a run,
runs in sequence, common random numbers across variants (seeds 5 to 8).

  1. Jensen: the cost share of the variable stop is c E[1/l] (s04: x1.0326 on commission and
     spread), not c / E[l],
  2. the strict daily reading in the assessment: a losing assessment trade dies with probability
     q_l (s06: 1.15% at a 00:00 UTC day, 1.58% at 21:00 UTC); its draws come from a separate
     generator so the trade coins stay common,
  3. both,
  4. the 40-pip cell (section 26's richest) with its own Jensen factor (the beta knob, x1.0648)
     and the ceiling clock, 7 trades.
"""
import json
import time

import numpy as np

from common import PA, T3, SPLIT, G_COIN, G_SPREAD, SWAP_COIN, STOP0, pick

t0 = time.time()
N = 5000
SEEDS = (5, 6, 7, 8)
RF0 = 0.01575
G_J = G_COIN + G_SPREAD * 0.0326


def wallet_variant(g, wait, rf, h=0.8, q_eval_loss=0.0, n=N, seed=5, horizon=541, ra=0.03):
    """prime.wallet_mc's loop, plus an assessment loss that breaches the daily rule with probability
    q_eval_loss. With q_eval_loss = 0 it consumes the same random stream as wallet_mc."""
    rng = np.random.default_rng(seed)
    rng_b = np.random.default_rng(seed + 999)
    fee, B, min_gross = T3.fee, T3.bankroll, 100.0 / SPLIT / T3.bankroll
    w = np.full(n, 1000.0); ruined = np.zeros(n, bool); state = np.full(n, -1, np.int8)
    eq = np.zeros(n); since = np.zeros(n); traded = np.zeros(n); big = np.zeros(n)
    for d in range(horizon):
        idle = (state == -1) & ~ruined
        buy = idle & (w >= fee - PA.TOL)
        w = np.where(buy, w - fee, w); state = np.where(buy, 0, state)
        eq = np.where(buy, 0.0, eq); since = np.where(buy, 0, since); traded = np.where(buy, 0, traded); big = np.where(buy, 0.0, big)
        ruined |= idle & ~buy
        act = (state >= 0) & ~ruined
        prev = eq.copy()
        u = rng.random(n)
        st_a, win_a, a_a = PA._steps_clipped(rng, n, ra, g * ra, 1.0, eq, True, True, u=u)
        st_f, win_f, a_f = PA._steps_clipped(rng, n, rf, g * rf, 1.0, eq, False, True, u=u)
        fund = state == 1
        st = np.where(fund, st_f, st_a); win = np.where(fund, win_f, win_a); wsz = np.where(fund, a_f, a_a)
        eq = np.where(act, eq + st, eq); since += act; traded += act
        big = np.where(act & win & (wsz > big), wsz, big)
        breach = act & (state == 0) & ~win & (rng_b.random(n) < q_eval_loss)
        dead = act & (((prev - eq) > PA.DAILY * (1.0 + prev)) | (eq <= -PA.MAXDD) | breach)
        state = np.where(dead, -1, state)
        promoted = act & (state == 0) & (eq >= PA.TARGET)
        state = np.where(promoted, 1, state); eq = np.where(promoted, 0.0, eq)
        since = np.where(promoted, 0, since); traded = np.where(promoted, 0, traded); big = np.where(promoted, 0.0, big)
        can = (state == 1) & ~ruined & (since >= wait) & (traded >= PA.MIN_DAYS) & (eq >= min_gross) & (big <= PA.DIST_CAP * eq)
        honoured = can & (rng.random(n) < h)
        w = w + np.where(honoured, SPLIT * eq * B, 0.0)
        eq = np.where(can, 0.0, eq); since = np.where(can, 0, since); big = np.where(can, 0.0, big)
    return np.log(np.maximum(w, 1.0) / 1000.0)


def run(label, **kw):
    per = [wallet_variant(seed=s, **kw) for s in SEEDS]
    return label, per


# sanity: the variant at q = 0 equals wallet_mc on the same seed
ref = PA.wallet_mc(T3, 0.03, G_COIN * 0.03, SPLIT, 0.8, 1000, horizon_days=541, n=N, seed=5, rr=1.0, wait=7,
                   funded_risk=RF0, funded_cost=G_COIN * RF0, clip_up=True, clip_dn=True)["e_log_growth"]
print(f"check: wallet_mc {ref:+.5f} against the variant {wallet_variant(G_COIN, 7, RF0).mean():+.5f}")

variants = [
    run("deployed as published", g=G_COIN, wait=7, rf=RF0),
    run("Jensen cost share", g=G_J, wait=7, rf=RF0),
    run("strict daily, 00:00 UTC day", g=G_COIN, wait=7, rf=RF0, q_eval_loss=0.0115),
    run("strict daily, 21:00 UTC day", g=G_COIN, wait=7, rf=RF0, q_eval_loss=0.0158),
    run("Jensen and strict daily (00:00)", g=G_J, wait=7, rf=RF0, q_eval_loss=0.0115),
]
g40 = 1.0648 / 40.0 + SWAP_COIN * 40.0 / STOP0
rf40 = pick(g40, 7)[0]["risk"]
variants.append(run(f"40-pip beta knob, Jensen, clock 7 (pick {rf40:.3%})", g=g40, wait=7, rf=rf40))
g40p = 1.0 / 40.0 + SWAP_COIN * 40.0 / STOP0
variants.append(run("40-pip as published (clock 6, 1.575%)", g=g40p, wait=6, rf=0.01575))

base = np.concatenate(variants[0][1])
out = {}
for label, per in variants:
    allv = np.concatenate(per)
    diff = allv - base
    out[label] = dict(mean=float(allv.mean()), se=float(allv.std(ddof=1) / np.sqrt(len(allv))),
                      per_seed=[float(p.mean()) for p in per], d_mean=float(diff.mean()),
                      d_se=float(diff.std(ddof=1) / np.sqrt(len(diff))))
    o = out[label]
    print(f"{label:46}: E[log g] {o['mean']:+.4f} +/- {o['se']:.4f} over {len(allv)} paths "
          f"(seeds {' '.join(f'{x:+.4f}' for x in o['per_seed'])}); against deployed {o['d_mean']:+.4f} +/- {o['d_se']:.4f}", flush=True)
json.dump(out, open("/tmp/arena-prime-audit/candidate-1/out_s07.json", "w"), indent=1)
print(f"done {time.time() - t0:.0f}s")