import json
import numpy as np
from pathlib import Path

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


def entropy_point(p, g, lam):
    """Simplex entropy mirror map: normalize p*exp(-lam*g), stably."""
    z = np.log(p) - lam * g
    z -= z.max()
    w = np.exp(z)
    return w / w.sum()


def entropy_root(p, g, delta, max_expand=80, bisect=80):
    """Solve g.(p-x(lam))=delta, or return failure if outside reachable range."""
    if delta <= 0:
        return 0.0, p.copy(), True
    gmin = g.min()
    reachable = float(np.dot(g, p) - gmin)
    if delta > reachable + 1e-12:
        return np.nan, p.copy(), False
    def phi(lam):
        x = entropy_point(p, g, lam)
        return float(np.dot(g, p-x) - delta)
    lo, hi = 0.0, 1.0
    while phi(hi) < 0 and hi < 1e12:
        hi *= 2
    if phi(hi) < 0:
        return np.nan, p.copy(), False
    for _ in range(bisect):
        mid = (lo + hi) / 2
        if phi(mid) >= 0:
            hi = mid
        else:
            lo = mid
    lam = (lo + hi) / 2
    x = entropy_point(p, g, lam)
    return lam, x, abs(float(np.dot(g, p-x)-delta)) < 1e-10


def simplex_project(v):
    u = np.sort(v)[::-1]
    cssv = np.cumsum(u) - 1
    ind = np.arange(1, len(v)+1)
    rho = np.nonzero(u - cssv / ind > 0)[0][-1]
    theta = cssv[rho] / (rho + 1)
    return np.maximum(v-theta, 0)


def core_sweeps():
    # Fixed non-symmetric point and gradient makes the predictions identifiable.
    p = np.array([0.55, 0.30, 0.15])
    g = np.array([0.8, -0.3, 1.4])
    mean = np.dot(p, g)
    variance = np.dot(p, (g-mean)**2)
    reachable = np.dot(p, g) - g.min()

    # Prediction 1: as delta -> 0, lambda/delta -> 1/Var_p(g).
    deltas = np.logspace(-7, -2, 8)
    small = []
    for d in deltas:
        lam, x, ok = entropy_root(p, g, d)
        small.append((float(d), float(lam), float(lam/d), bool(ok)))
    asymptotic = 1.0 / variance
    observed_small = small[0][2]

    # Prediction 2: simultaneous g, delta scaling leaves x unchanged and lambda -> lambda/a.
    # This is the exact homogeneity law; fixed absolute delta is not invariant.
    d0 = 0.23 * reachable
    l1, x1, ok1 = entropy_root(p, g, d0)
    scaling = []
    for a in [0.25, 0.5, 1.0, 2.0, 4.0]:
        la, xa, oka = entropy_root(p, a*g, a*d0)
        scaling.append({"a": a, "lambda": float(la), "predicted_lambda": float(l1/a),
                        "lambda_ratio": float(la/(l1/a)), "x_error": float(np.max(abs(xa-x1))), "ok": bool(oka)})

    # A local fixed-delta prediction also follows from the Taylor expansion:
    # lambda/delta ~ 1/(a^2 Var_p(g)) as delta -> 0.
    fixed_delta_scaling = []
    tiny_delta = 1e-7
    for a in [0.25, 0.5, 1.0, 2.0, 4.0]:
        la, xa, oka = entropy_root(p, a*g, tiny_delta)
        predicted = tiny_delta / (a*a*variance)
        fixed_delta_scaling.append({"a": a, "lambda": float(la),
                                    "predicted_lambda": float(predicted),
                                    "ratio": float(la/predicted), "ok": bool(oka)})
    # Prediction 3: root exists iff delta <= limiting displacement, and fails above it.
    boundary = []
    for frac in [0.0, 0.25, 0.5, 0.9, 0.999999, 1.001, 1.5]:
        d = frac * reachable
        lam, x, ok = entropy_root(p, g, d)
        boundary.append({"fraction": frac, "delta": float(d), "success": bool(ok),
                         "lambda": None if not ok else float(lam)})
    return {"p": p.tolist(), "g": g.tolist(), "weighted_variance": float(variance),
            "reachable_delta": float(reachable), "small_gap": small,
            "small_gap_prediction_lambda_over_delta": float(asymptotic),
            "small_gap_observed_first": float(observed_small),
            "scaling": scaling, "fixed_delta_local_scaling": fixed_delta_scaling, "boundary": boundary}


def optimization_comparison(seed=7, n=8, steps=120):
    rng = np.random.default_rng(seed)
    q = rng.dirichlet(np.ones(n)*1.5)
    p0 = np.ones(n)/n
    # f=1/2||p-q||², f*=0, so the exact Polyak gap is known.
    pe, pu = p0.copy(), p0.copy()
    ent_losses, euc_losses = [], []
    ent_resid, euc_resid = [], []
    for _ in range(steps):
        for p, kind, losses, residuals in [(pe, 'entropy', ent_losses, ent_resid), (pu, 'euclidean', euc_losses, euc_resid)]:
            g = p-q
            f = 0.5*np.dot(g,g)
            delta = f
            if kind == 'entropy':
                lam, x, ok = entropy_root(p, g, delta)
                if not ok: x = p
            else:
                # Euclidean Polyak step, with simplex projection for a fair feasible baseline.
                alpha = delta / max(np.dot(g,g), 1e-30)
                x = simplex_project(p-alpha*g)
            residuals.append(abs(np.dot(g, p-x)-delta))
            losses.append(f)
            if kind == 'entropy': pe = x
            else: pu = x
    return {"target_q": q.tolist(), "steps": steps,
            "entropy_loss_start": ent_losses[0], "entropy_loss_final": ent_losses[-1],
            "euclidean_loss_start": euc_losses[0], "euclidean_loss_final": euc_losses[-1],
            "entropy_loss_at_10": ent_losses[9], "euclidean_loss_at_10": euc_losses[9],
            "entropy_max_halfspace_residual": max(ent_resid),
            "euclidean_max_halfspace_residual_after_projection": max(euc_resid)}


def main():
    out = {"core": core_sweeps(), "optimization": optimization_comparison()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()