import json
import math
import numpy as np

SEED = 123
EPS = 1e-12


def bregman(y, x):
    y = np.asarray(y, dtype=np.float64)
    x = np.asarray(x, dtype=np.float64)
    return float(np.sum(-np.log(y / x) + (y - x) / x))


def barrier_step(p, g, alpha, tol=1e-13, max_iter=200):
    """Solve y_i=(1/p_i+alpha*(g_i+lambda))^-1, sum(y)=1."""
    p = np.asarray(p, dtype=np.float64)
    g = np.asarray(g, dtype=np.float64)
    if np.any(p <= 0) or abs(p.sum() - 1) > 1e-9:
        raise ValueError("p must be strictly positive and sum to one")
    invp = 1.0 / p
    # Domain is lambda > max_i(-g_i-invp_i/alpha).
    lo = float(np.max(-g - invp / alpha)) + 1e-14
    def F(lam):
        den = invp + alpha * (g + lam)
        if np.any(den <= 0):
            return np.inf
        return float(np.sum(1.0 / den) - 1.0)
    # F decreases continuously from +infinity to -1.
    left = lo
    right = max(1.0, float(np.max(np.abs(g)) + np.max(invp) / alpha + 1.0))
    while F(right) > 0:
        right = 2.0 * right + 1.0
        if right > 1e14:
            raise FloatingPointError("could not bracket lambda")
    for _ in range(max_iter):
        mid = (left + right) / 2.0
        if F(mid) > 0:
            left = mid
        else:
            right = mid
        if right - left <= tol * max(1.0, abs(mid)):
            break
    lam = (left + right) / 2.0
    den = invp + alpha * (g + lam)
    y = 1.0 / den
    y = y / y.sum()
    return y, lam, float(abs(y.sum() - 1.0)), float(np.min(den))


def math_checks():
    rng = np.random.default_rng(SEED)
    bvals = []
    residuals = []
    mins = []
    for _ in range(300):
        x = rng.dirichlet(np.ones(7) * .7)
        y = rng.dirichlet(np.ones(7) * .7)
        bvals.append(bregman(y, x))
        g = rng.normal(size=7)
        q, _, r, md = barrier_step(x, g, 0.37)
        residuals.append(r)
        mins.append((q > 0).all() and md > 0)
    return {"min_bregman": float(min(bvals)),
            "max_simplex_residual": float(max(residuals)),
            "all_denominators_positive": bool(all(mins))}


def prediction_sweeps():
    # Prediction 1: every finite alpha gives strictly positive coordinates and exact root.
    p = np.array([.55, .25, .15, .05])
    g = np.array([0., .4, 1.2, 2.0])
    alphas = np.logspace(-3, 3, 13)
    rows = []
    for a in alphas:
        q, lam, res, md = barrier_step(p, g, float(a))
        rows.append({"alpha": float(a), "min_q": float(q.min()),
                     "root_residual": res, "min_denominator": md})

    # Prediction 2: for a boundary-seeking linear objective, 1/q_bad grows
    # approximately linearly in steps, with slope alpha (asymptotically).
    p0 = np.array([.5, .5])
    slopes = []
    for a in [0.5, 1., 2., 5., 10., 20.]:
        pcur = p0.copy(); inv_history = []
        for _ in range(80):
            pcur, _, _, _ = barrier_step(pcur, np.array([0., 1.]), a)
            inv_history.append(1.0 / pcur[1])
        t = np.arange(1, 81)
        slope = float(np.polyfit(t[30:], np.asarray(inv_history)[30:], 1)[0])
        slopes.append({"alpha": a, "observed_reciprocal_slope": slope,
                       "predicted_slope": a, "ratio": slope / a})

    # Prediction 3: one step from p_bad has q_bad approximately p_bad/(1+alpha*p_bad)
    # when alpha*p_bad is large; report relative error across initial masses.
    scaling = []
    for pb in [.01, .03, .1, .3, .5]:
        pp = np.array([1-pb, pb]); aa = 100.
        q, _, _, _ = barrier_step(pp, np.array([0., 1.]), aa)
        pred = pb / (1 + aa * pb)
        scaling.append({"p_bad": pb, "observed_q_bad": float(q[1]),
                        "approx_predicted_q_bad": float(pred),
                        "relative_error": float(abs(q[1]-pred)/q[1])})
    return {"positivity_and_root_sweep": rows,
            "reciprocal_scaling_sweep": slopes,
            "boundary_one_step_scaling": scaling}


def adam_logits_step(z, grad, m, v, t, lr=.08):
    b1, b2 = .9, .999
    m = b1*m + (1-b1)*grad
    v = b2*v + (1-b2)*grad*grad
    mh = m/(1-b1**t); vh = v/(1-b2**t)
    return z - lr*mh/(np.sqrt(vh)+1e-8), m, v


def routing_benchmark():
    # Convex cross entropy to an imbalanced useful routing prior. Same 160 steps.
    target = np.array([.70, .20, .08, .019, .001])
    target /= target.sum()
    nsteps = 160
    start = np.ones(5)/5
    results = {}
    # Barrier operates directly on p with exact gradient d(-r log p)/dp.
    for name in ["barrier", "eg", "adam_logits"]:
        p = start.copy(); z = np.log(start); m = np.zeros(5); v = np.zeros(5)
        losses = []
        for t in range(1, nsteps+1):
            if name == "barrier":
                g = -target / p
                p, _, _, _ = barrier_step(p, g, .08)
            elif name == "eg":
                g = -target / p
                p = p * np.exp(-.08*g); p /= p.sum()
            else:
                ex = np.exp(z-z.max()); p = ex/ex.sum()
                gradz = p-target
                z, m, v = adam_logits_step(z, gradz, m, v, t)
                p = np.exp(z-z.max()); p /= p.sum()
            loss = float(-np.sum(target*np.log(np.maximum(p, 1e-300))))
            losses.append(loss)
        entropy = float(-np.sum(p*np.log(np.maximum(p, 1e-300))))
        results[name] = {"final_loss": losses[-1], "loss_step_20": losses[19],
                         "final_entropy": entropy, "min_probability": float(p.min()),
                         "losses": losses}
    return {"target": target.tolist(), "steps": nsteps, "methods": results}


def main():
    out = {"seed": SEED, "math_checks": math_checks(),
           "predictions": prediction_sweeps(), "routing_benchmark": routing_benchmark()}
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()