import json
import math
import numpy as np

SEED = 7
np.random.seed(SEED)

# Bounded diagonal quadratic manifold. The barrier is smooth in (-5,5)^d.
LO, HI, LAM, GAMMA = -5.0, 5.0, 0.015, 0.03
H = np.linspace(0.7, 1.8, 8)


def potential(x):
    if np.any(x <= LO) or np.any(x >= HI):
        return np.inf
    return 0.5 * np.sum(H*x*x) - LAM*np.log(x-LO).sum() - LAM*np.log(HI-x).sum()


def grad_potential(x):
    return H*x + LAM*(-1.0/(x-LO) + 1.0/(HI-x))


def metric(x):
    return 1.0 + GAMMA*x*x


def state(x):
    G = metric(x)
    g = grad_potential(x)/G
    r = math.sqrt(float(np.sum(G*g*g)))
    return g, r


def step_fixed(x, eta=0.035, a=0.75, b=0.10, p=0.5, q=2.0, eps=1e-10):
    g, r0 = state(x)
    r = math.sqrt(r0*r0 + eps*eps)
    return x - eta*(a*r**(p-1.0) + b*r**(q-1.0))*g


def step_gd(x, eta=0.035):
    return x - eta*grad_potential(x)/metric(x)


def run(step, x0, n=1200, tol=1e-5):
    x = x0.copy(); vstar = potential(np.zeros_like(x)); vals=[]
    feasible = True; settle = None
    for t in range(n+1):
        v = potential(x)-vstar
        vals.append(v)
        feasible = feasible and np.all(np.isfinite(x)) and np.all(x > LO) and np.all(x < HI)
        if settle is None and v <= tol: settle=t
        if t < n: x=step(x)
    return {"settle": settle if settle is not None else n+1,
            "residual": float(vals[-1]), "min_residual": float(np.min(vals)),
            "feasible": bool(feasible)}


def math_check():
    # Check the exact claimed inequality using the metric gradient. Estimate
    # the largest valid mu on the same central neighborhood, then use 90% of it.
    p,q,a,b=.5,2.,.75,.10
    alpha,beta=(p+1)/2,(q+1)/2
    ratios=[]; xs=np.linspace(-3.8,3.8,2001)
    v0=potential(np.zeros(1))
    for z in xs:
        x=np.array([z]); V=potential(x)-v0
        if V > 1e-12:
            _,r=state(x); ratios.append(r*r/(2*V))
    mu=.9*min(ratios)
    A=a*(2*mu)**alpha; BB=b*(2*mu)**beta
    worst=-np.inf; violations=0
    for z in xs:
        x=np.array([z]); V=potential(x)-v0
        if V <= 1e-12: continue
        _,r=state(x)
        dV=-a*r**(p+1)-b*r**(q+1)
        rhs=-A*V**alpha-BB*V**beta
        worst=max(worst,dV-rhs)
        violations += int(dV > rhs+1e-10)
    bound=1/(A*(1-alpha))+1/(BB*(beta-1))
    return {"metric_mu_estimate":float(mu),"alpha":alpha,"beta":beta,
            "settling_bound":float(bound),"max_dV_minus_rhs":float(worst),
            "violations":violations,"samples":len(xs)-1}


def epsilon_check():
    # For the smoothed norm, measure residual potential at a fixed long horizon.
    # This is a small numerical check of the stated O(epsilon) neighborhood claim.
    out=[]; x0=np.full(8,3.0); vstar=potential(np.zeros(8))
    for eps in [1e-4, 3e-4, 1e-3, 3e-3, 1e-2]:
        x=x0.copy()
        for _ in range(3000): x=step_fixed(x,eps=eps)
        out.append([eps,float(potential(x)-vstar)])
    logs=np.log(np.array(out)); slope=float(np.polyfit(logs[:,0],logs[:,1],1)[0])
    return {"epsilon_residual_pairs":out,"loglog_slope":slope}


def main():
    results={"math_check":math_check(),"epsilon_check":epsilon_check(),"runs":[]}
    for scale in [0.1,0.5,1.0,2.0,3.5,4.5]:
        x0=np.full(8,scale)
        results["runs"].append({"scale":scale,
          "fixed_time":run(step_fixed,x0),"gradient_descent":run(step_gd,x0)})
    results["summary"]={
      "fixed_time_settles":[r["fixed_time"]["settle"] for r in results["runs"]],
      "gd_settles":[r["gradient_descent"]["settle"] for r in results["runs"]],
      "fixed_time_residuals":[r["fixed_time"]["residual"] for r in results["runs"]],
      "gd_residuals":[r["gradient_descent"]["residual"] for r in results["runs"]]}
    with open("results.json","w") as f: json.dump(results,f,indent=2)
    print(json.dumps(results,indent=2))

if __name__ == "__main__": main()