"""Audit check 3: how a refused payout is modelled, and what the gate reads under each reading.

`prime.wallet_mc` applies the honour rate per request as a Bernoulli and, on a refusal, still
sets the account's equity to zero (the equity leaves the account and the wallet gets nothing).
Two other readings are at least as plausible under the firm's documents: the equity stays in
the account and the trader asks again after the next clock (refusal = delay), or the firm
terminates the account (General Terms 23.1 and clause 16 forfeit future payouts). This copy of
wallet_mc adds a `refusal` argument with those three readings and leaves the module untouched.

    cd /root/Projects/trading-framework && python3 /tmp/arena-prime-audit/candidate-2/c2_refusal.py
"""
import json
import sys

import numpy as np

sys.path.insert(0, "/root/Projects/trading-framework")
from src.notebooks.reports.topics.prime_academy.modules import prime as PA  # noqa: E402

OUT = "/tmp/arena-prime-audit/candidate-2/c2_refusal.json"
DEC = json.load(open("/root/Projects/trading-framework/src/notebooks/reports/topics/prime_academy/paper/decision.json"))
T3 = PA.Tier(3000.0, 38.0, 51.0)
G = DEC["scenarios"]["random side, swap charged"]["cost"]
RF = DEC["scenarios"]["random side, swap charged"]["pick"]["risk"]
WAIT, H = DEC["wait"], DEC["horizon"]


def wallet_mc_refusal(tier, risk, cost, split, p_paid, wallet0, horizon_days, n, seed, rr, wait,
                      funded_risk, funded_cost, clip_up, clip_dn, refusal="vanish"):
    """prime.wallet_mc, lines 600-702 of modules/prime.py, with one change: what happens to
    a refused request. 'vanish' is the module's behaviour; 'stay' keeps the equity and the
    clock restarts; 'terminate' ends the account."""
    rng = np.random.default_rng(seed)
    fee, B = tier.fee, tier.bankroll
    min_gross = 100.0 / split / B
    w = np.full(n, float(wallet0)); peak = w.copy(); maxdd = np.zeros(n)
    ruined = np.zeros(n, bool); first_cash = np.full(n, np.nan)
    state = np.full(n, -1, np.int8); eq = np.zeros(n); since = np.zeros(n); traded = np.zeros(n); big = np.zeros(n)
    refused = 0; requests = 0
    for d in range(horizon_days):
        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
        if not act.any() and ruined.all():
            break
        prev = eq.copy()
        u = rng.random(n)
        st_a, win_a, a_a = PA._steps_clipped(rng, n, risk, cost, rr, eq, clip_up, clip_dn, u=u)
        st_f, win_f, a_f = PA._steps_clipped(rng, n, funded_risk, funded_cost, rr, eq, False, clip_dn, 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)
        dead = act & (((prev - eq) > PA.DAILY * (1.0 + prev)) | (eq <= -PA.MAXDD))
        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) < p_paid)
        denied = can & ~honoured
        requests += int(can.sum()); refused += int(denied.sum())
        cash = np.where(honoured, split * eq * B, 0.0)
        w = w + cash
        first_cash = np.where(np.isnan(first_cash) & (cash > 0), d, first_cash)
        if refusal == "vanish":
            reset = can
        elif refusal == "stay":
            reset = honoured
        elif refusal == "terminate":
            reset = honoured
            state = np.where(denied, -1, state)
        else:
            raise ValueError(refusal)
        eq = np.where(reset, 0.0, eq); since = np.where(can, 0, since); big = np.where(reset, 0.0, big)
        peak = np.maximum(peak, w); maxdd = np.maximum(maxdd, peak - w)
    return {"p_ruin": float(ruined.mean()), "e_log_growth": float(np.mean(np.log(np.maximum(w, 1.0) / wallet0))),
            "p_any_cash": float(np.mean(~np.isnan(first_cash))), "mean_final": float(w.mean()),
            "refused_share": refused / max(requests, 1)}


res = {}
# sanity: 'vanish' must reproduce prime.wallet_mc exactly at the same seed
ref = PA.wallet_mc(T3, 0.03, G * 0.03, 0.7, 0.8, 1000, horizon_days=H, n=5000, seed=1, rr=1.0, wait=WAIT,
                   funded_risk=RF, funded_cost=G * RF, clip_up=True, clip_dn=True)
mine = wallet_mc_refusal(T3, 0.03, G * 0.03, 0.7, 0.8, 1000, H, 5000, 1, 1.0, WAIT, RF, G * RF, True, True, "vanish")
print(f"[0] copy reproduces prime.wallet_mc at seed 1: {ref['e_log_growth']:+.5f} against {mine['e_log_growth']:+.5f}, "
      f"ruin {ref['p_ruin']:.4f} against {mine['p_ruin']:.4f}")
res["reproduce"] = dict(module=ref["e_log_growth"], copy=mine["e_log_growth"])

print("[1] the gate at 80% honoured under three readings of a refusal, 4 seeds x 5,000 paths:")
for mode in ("vanish", "stay", "terminate"):
    rows = [wallet_mc_refusal(T3, 0.03, G * 0.03, 0.7, 0.8, 1000, H, 5000, s, 1.0, WAIT, RF, G * RF, True, True, mode) for s in (1, 2, 3, 4)]
    lg = [r["e_log_growth"] for r in rows]
    print(f"    {mode:9}: g_80 mean {np.mean(lg):+.4f} sd {np.std(lg, ddof=1):.4f}  ruin {np.mean([r['p_ruin'] for r in rows]):.4f}  "
          f"mean final ${np.mean([r['mean_final'] for r in rows]):,.0f}  refused share of requests {np.mean([r['refused_share'] for r in rows]):.3f}")
    res[mode] = dict(g80=float(np.mean(lg)), sd=float(np.std(lg, ddof=1)), ruin=float(np.mean([r["p_ruin"] for r in rows])),
                     mean_final=float(np.mean([r["mean_final"] for r in rows])))

print("[2] the honour rate at which growth is zero under each reading (bisection on the 4-seed mean):")
for mode in ("vanish", "stay", "terminate"):
    lo, hi = 0.6, 1.0
    for _ in range(7):
        mid = 0.5 * (lo + hi)
        m = np.mean([wallet_mc_refusal(T3, 0.03, G * 0.03, 0.7, mid, 1000, H, 5000, s, 1.0, WAIT, RF, G * RF, True, True, mode)["e_log_growth"] for s in (1, 2, 3, 4)])
        lo, hi = (mid, hi) if m < 0 else (lo, mid)
    print(f"    {mode:9}: breakeven honour about {0.5 * (lo + hi):.3f}")
    res[mode]["honour_root"] = 0.5 * (lo + hi)

json.dump(res, open(OUT, "w"), indent=1)
print("done ->", OUT)