import numpy as np


def discounted(gaps, rho):
    r = 0.0
    out = []
    for g in gaps:
        r = rho * r + (1.0 - rho) * g
        out.append(r)
    return np.asarray(out)


def exact_gap(x, y, a, box=2.0):
    # f(x,y)=a*x*y, with x,y in [-box,box].
    return 2.0 * box * abs(a) * (abs(x) + abs(y))


def probe_gap(x, y, a, k=3, eta=0.08, box=2.0):
    yp, xp = float(y), float(x)
    for _ in range(k):
        yp = np.clip(yp + eta * a * x, -box, box)
    for _ in range(k):
        xp = np.clip(xp - eta * a * y, -box, box)
    return a * x * yp - a * xp * y


def verify_math():
    rng = np.random.default_rng(4)
    rho = 0.9
    gaps = rng.uniform(0, 3, 80)
    direct = discounted(gaps, rho)
    rec = np.zeros_like(gaps)
    for t, g in enumerate(gaps):
        rec[t] = rho * (rec[t - 1] if t else 0.0) + (1.0 - rho) * g
    recursion_err = float(np.max(np.abs(direct - rec)))
    impulse = discounted(np.r_[1.0, np.zeros(30)], rho)
    expected = impulse[0] * rho ** np.arange(1, 31)
    decay_err = float(np.max(np.abs(impulse[1:] - expected)))
    return recursion_err, decay_err, 1.0 / (1.0 - rho)


def run(controller, seed=7, steps=1800, change=650):
    rng = np.random.default_rng(seed)
    x, y = 0.9, -0.7
    ex = ey = 0.055
    rho, tau = 0.9, 0.025
    R = 0.0
    previous = None
    down_count = 0
    up_count = 0
    gaps, norms, rates, switches = [], [], [], []
    for t in range(steps):
        a = 1.0 if t < change else -1.0
        # Mild deterministic measurement noise models minibatch gap noise.
        noisy_a = a * (1.0 + 0.015 * rng.normal())
        x_old, y_old = x, y
        # Simultaneous descent/ascent on the current bilinear payoff.
        x = np.clip(x - ex * noisy_a * y_old, -2.0, 2.0)
        y = np.clip(y + ey * noisy_a * x_old, -2.0, 2.0)
        g = max(0.0, probe_gap(x, y, noisy_a))
        R = rho * R + (1.0 - rho) * g
        if controller and previous is not None:
            if R > previous * (1.0 + tau):
                ex *= 0.5; ey *= 0.5
                down_count += 1; up_count = 0
            elif R < previous * (1.0 - tau):
                up_count += 1
                if up_count >= 4:
                    ex = min(0.09, ex * 1.05); ey = min(0.09, ey * 1.05)
                    up_count = 0
            else:
                up_count = 0
        previous = R
        gaps.append(exact_gap(x, y, a))
        norms.append(abs(x) + abs(y))
        rates.append(ex)
        switches.append(down_count)
    arr = np.asarray(gaps)
    return {
        'mean_gap_all': float(arr.mean()),
        'mean_gap_after_change': float(arr[change:].mean()),
        'mean_gap_last_300': float(arr[-300:].mean()),
        'max_gap_after_change': float(arr[change:].max()),
        'final_norm': float(norms[-1]),
        'step_min': float(min(rates)),
        'step_final': float(rates[-1]),
        'down_events': int(down_count),
    }


if __name__ == '__main__':
    rec_err, decay_err, horizon = verify_math()
    print('MATH recursion_max_abs_error', rec_err)
    print('MATH impulse_decay_max_abs_error', decay_err)
    print('MATH effective_horizon', horizon)
    for controller in (False, True):
        vals = [run(controller, seed=s) for s in (7, 8, 9)]
        print('CONTROLLER', controller)
        for key in vals[0]:
            print(key, np.mean([v[key] for v in vals]), '+/-', np.std([v[key] for v in vals]))