"""prime.wallet_mc's loop with an assessment daily-rule breach, shared by s07 and s10."""
import numpy as np

from common import PA, T3, SPLIT

N = 5000


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)


