import json, math
from pathlib import Path
import numpy as np

SEED = 2078
np.set_printoptions(precision=6, suppress=True)


def prox(p, c, eta):
    """KL prox: argmin_q <c,q> + KL(q||p)/eta."""
    logw = np.log(np.maximum(p, 1e-300)) - eta * c
    logw -= np.max(logw)
    q = np.exp(logw)
    return q / q.sum()


def make_operator(lam):
    # On the tangent plane, A is skew-symmetric with eigenvalues +/- i*lam.
    return lam * np.array([[0., -1., 1.], [1., 0., -1.], [-1., 1., 0.]]) / math.sqrt(3.)


def md_step(p, A, eta):
    return prox(p, A @ p, eta)


def mp_step(p, A, eta):
    z = prox(p, A @ p, eta)
    return prox(p, A @ z, eta)


def radius(p):
    return float(np.linalg.norm(p - np.ones(3) / 3))


def jacobian_factor(step, u, A, eta, eps=1e-7):
    # Two tangent directions; both have the same theoretical factor here.
    basis = np.array([[1., -1., 0.], [1., 1., -2.]])
    basis /= np.linalg.norm(basis, axis=1)[:, None]
    vals = []
    for v in basis:
        p = u + eps * v
        vals.append(radius(step(p, A, eta)) / radius(p))
    return float(np.mean(vals))


def trajectory(step, p0, A, eta, steps=60):
    p = p0.copy(); rs = [radius(p)]
    for _ in range(steps):
        p = step(p, A, eta); rs.append(radius(p))
    return np.asarray(rs)


def main():
    u = np.ones(3) / 3
    p0 = u + np.array([0.035, -0.021, -0.014])
    rows = []

    # Prediction 1: zero coupling gives exactly identity updates.
    A = make_operator(0.)
    zero_effect = max(float(np.max(np.abs(md_step(p0, A, .7) - p0))),
                      float(np.max(np.abs(mp_step(p0, A, .7) - p0))))

    # At uniform p, linearized KL prox is I-(eta/3)A. Thus x=eta*Lambda/3,
    # MD factor=sqrt(1+x^2), MP factor=sqrt(1-x^2+x^4), neutral at x=1.
    etas = [0.30, 1.00, 2.00]
    lambdas = [0.50, 1.00, 2.00, 3.00, 5.00]
    for eta in etas:
        for lam in lambdas:
            x = eta * lam / 3.0
            A = make_operator(lam)
            measured_md = jacobian_factor(md_step, u, A, eta)
            measured_mp = jacobian_factor(mp_step, u, A, eta)
            rows.append({
                'eta': eta, 'Lambda': lam, 'x': x,
                'md_measured': measured_md, 'md_predicted': math.sqrt(1 + x*x),
                'mp_measured': measured_mp, 'mp_predicted': math.sqrt(1 - x*x + x**4)
            })

    # Equal-step coupled-routing proxy at x=2/3: MP should damp the cycle while
    # standard one-stage entropic mirror descent amplifies it.
    eta, lam = 1.00, 2.00
    A = make_operator(lam)
    md_r = trajectory(md_step, p0, A, eta)
    mp_r = trajectory(mp_step, p0, A, eta)

    def vi_residual(p):
        c = A @ p
        return float(np.max(c - c @ p))

    pmd, pmp = p0.copy(), p0.copy(); md_vi, mp_vi = [], []
    for _ in range(60):
        md_vi.append(vi_residual(pmd)); mp_vi.append(vi_residual(pmp))
        pmd = md_step(pmd, A, eta); pmp = mp_step(pmp, A, eta)

    tr = [r for r in rows if abs(r['eta'] - 2.0) < 1e-12]
    below = [r['x'] for r in tr if r['mp_measured'] < 1.0]
    above = [r['x'] for r in tr if r['mp_measured'] > 1.0]
    md_err = max(abs(r['md_measured'] - r['md_predicted']) for r in rows)
    mp_err = max(abs(r['mp_measured'] - r['mp_predicted']) for r in rows)
    result = {
        'seed': SEED,
        'zero_lambda_max_change': zero_effect,
        'zero_lambda_prediction_confirmed': bool(zero_effect < 1e-14),
        'local_sweep': rows,
        'max_abs_local_prediction_error_md': float(md_err),
        'max_abs_local_prediction_error_mp': float(mp_err),
        'transition': {
            'predicted_x': 1.0,
            'observed_below_max_x': float(max(below)) if below else None,
            'observed_above_min_x': float(min(above)) if above else None,
            'note': 'MP factor is below 1 for 0<x<1 and above 1 for x>1.'
        },
        'mini_experiment': {
            'eta': eta, 'Lambda': lam, 'x': eta * lam / 3.0,
            'md_radius_initial': float(md_r[0]), 'md_radius_final': float(md_r[-1]),
            'mp_radius_initial': float(mp_r[0]), 'mp_radius_final': float(mp_r[-1]),
            'md_vi_initial': float(md_vi[0]), 'md_vi_final': float(md_vi[-1]),
            'mp_vi_initial': float(mp_vi[0]), 'mp_vi_final': float(mp_vi[-1]),
            'md_radius_max': float(md_r.max()), 'mp_radius_max': float(mp_r.max())
        }
    }
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()