import json
import math
import numpy as np


def weighted_box_projection(z, lo, hi, m):
    """Exact argmin over a box of 0.5*(y-z)^T M^{-1}(y-z)."""
    z = np.asarray(z, dtype=float)
    return np.minimum(hi, np.maximum(lo, z))


class StateMetricProjected:
    def __init__(self, dim, alpha=0.1, eta=1.0, beta2=0.9,
                 metric_beta=0.5, m_min=1e-3, m_max=10.0,
                 lo=-2.0, hi=2.0, eps=1e-12):
        self.alpha, self.eta, self.beta2 = alpha, eta, beta2
        self.metric_beta = metric_beta
        self.m_min, self.m_max = m_min, m_max
        self.v = np.zeros(dim)
        self.lo, self.hi, self.eps = lo, hi, eps

    def step(self, theta, grad):
        self.v = self.beta2 * self.v + (1.0 - self.beta2) * grad * grad
        t = self.v / max(1.0 - self.beta2, self.eps)  # bias correction
        logm = self.metric_beta * np.log(t + self.eps)
        m = np.exp(np.clip(logm, np.log(self.m_min), np.log(self.m_max)))
        z = theta - self.alpha * m * grad
        p = weighted_box_projection(z, self.lo, self.hi, m)
        new_theta = theta + self.eta * (p - theta)
        return new_theta, m


def projected_sgd_step(theta, grad, lr, lo=-2.0, hi=2.0):
    return np.clip(theta - lr * grad, lo, hi)


def projected_adam_step(theta, grad, state, lr=0.1, beta1=.9, beta2=.999,
                        lo=-2.0, hi=2.0, eps=1e-8):
    state[0] = beta1 * state[0] + (1-beta1) * grad
    state[1] = beta2 * state[1] + (1-beta2) * grad * grad
    state[2] += 1
    mh = state[0] / (1-beta1**state[2])
    vh = state[1] / (1-beta2**state[2])
    return np.clip(theta - lr * mh / (np.sqrt(vh)+eps), lo, hi)


def verify_projection():
    rng = np.random.default_rng(7)
    z = rng.normal(size=8)
    lo, hi = -0.4, 0.6
    m = np.exp(rng.normal(size=8))
    p = weighted_box_projection(z, lo, hi, m)
    # KKT sign check for the separable weighted projection.
    stationarity = (p-z)/m
    ok = True
    for i in range(len(p)):
        if lo < p[i] < hi:
            ok &= abs(stationarity[i]) < 1e-12
        elif p[i] <= lo:
            ok &= stationarity[i] >= -1e-12
        elif p[i] >= hi:
            ok &= stationarity[i] <= 1e-12
    return {"max_box_violation": float(max(np.max(lo-p), np.max(p-hi), 0)),
            "kkt_sign_check": bool(ok)}


def contraction_sweep():
    # For x' = x + eta(P(x-alpha*m*h*x)-x), interior factor is 1-eta*alpha*m*h.
    h, eta = 37.0, 1.0
    alphas = np.array([.001, .01, .02, .03, .04, .05, .06, .07, .08])
    observed, predicted = [], []
    for alpha in alphas:
        x = 0.73
        vals = []
        for _ in range(30):
            vals.append(abs(x))
            x = x + eta * ((x-alpha*h*x)-x)
        # geometric ratio over early nonzero iterates
        ratios = np.abs(np.asarray(vals[1:]) / np.asarray(vals[:-1]))
        observed.append(float(np.median(ratios)))
        predicted.append(float(abs(1-eta*alpha*h)))
    # Stability transition predicted at alpha*h=2; classify by a 100-step bounded run.
    boundary_pred = 2.0 / (eta*h)
    boundary_obs = None
    for alpha in np.linspace(.045, .06, 301):
        x = .73
        for _ in range(100):
            x = x + eta*((x-alpha*h*x)-x)
        if abs(x) > 1e3:
            boundary_obs = float(alpha)
            break
    return {"h": h, "eta": eta, "alphas": alphas.tolist(),
            "predicted_contraction": predicted, "observed_contraction": observed,
            "stability_boundary_predicted": boundary_pred,
            "stability_boundary_observed_first_divergence": boundary_obs,
            "contraction_max_abs_error": float(np.max(np.abs(np.array(observed)-np.array(predicted))))}


def metric_scaling_sweep():
    # Prediction: for an interior coordinate, the per-step fractional decrease is
    # alpha*m*h (until the sign-flip boundary), hence gain is linear in m.
    h, alpha, eta = 23.0, .01, 1.0
    ms = np.array([.1, .2, .4, .8, 1.2])
    observed, predicted = [], []
    for m in ms:
        x = .91
        x1 = x + eta*((x-alpha*m*h*x)-x)
        observed.append(float((x-x1)/x))
        predicted.append(float(eta*alpha*m*h))
    return {"h": h, "alpha": alpha, "metrics": ms.tolist(),
            "predicted_fractional_decrease": predicted,
            "observed_fractional_decrease": observed,
            "linearity_max_abs_error": float(np.max(np.abs(np.array(observed)-np.array(predicted))))}


def quadratic_comparison():
    rng = np.random.default_rng(123)
    H = np.array([1.0, 10000.0])
    theta0 = np.array([1.0, 1.0])
    lo, hi = -2.0, 2.0
    steps = 180
    def grad(x): return H*x
    def loss(x): return .5*float(np.sum(H*x*x))
    records = {}
    # Stable but deliberately conservative SGD for the ill-conditioned direction.
    x = theta0.copy(); ls=[]; viol=0
    for _ in range(steps):
        g=grad(x); x=projected_sgd_step(x,g,1.5e-4,lo,hi); ls.append(loss(x)); viol=max(viol,float(max(np.max(lo-x),np.max(x-hi),0)))
    records['projected_sgd']={'loss_curve':ls,'final_loss':ls[-1],'violations':viol}
    x=theta0.copy(); ls=[]; viol=0; st=[np.zeros(2),np.zeros(2),0]
    for _ in range(steps):
        x=projected_adam_step(x,grad(x),st,lr=.1,lo=lo,hi=hi); ls.append(loss(x)); viol=max(viol,float(max(np.max(lo-x),np.max(x-hi),0)))
    records['projected_adam']={'loss_curve':ls,'final_loss':ls[-1],'violations':viol}
    x=theta0.copy(); ls=[]; viol=0; opt=StateMetricProjected(2,alpha=.03,eta=1,beta2=.9,metric_beta=.5,m_min=.01,m_max=10,lo=lo,hi=hi)
    metric_trace=[]
    for _ in range(steps):
        g=grad(x); x,m=opt.step(x,g); metric_trace.append(m.copy()); ls.append(loss(x)); viol=max(viol,float(max(np.max(lo-x),np.max(x-hi),0)))
    records['state_metric_projected']={'loss_curve':ls,'final_loss':ls[-1],'violations':viol,'final_metric':metric_trace[-1].tolist()}
    return {k:{kk:vv for kk,vv in v.items() if kk!='loss_curve'} for k,v in records.items()}, records


def main():
    proj=verify_projection()
    contraction=contraction_sweep()
    scaling=metric_scaling_sweep()
    summary, curves=quadratic_comparison()
    out={'projection_check':proj,'prediction_sweeps':{'contraction':contraction,'metric_scaling':scaling},'quadratic_summary':summary,
         'loss_curves':{k:v['loss_curve'] for k,v in curves.items()}}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps({'projection_check':proj,'prediction_sweeps':out['prediction_sweeps'],'quadratic_summary':summary},indent=2))

if __name__=='__main__': main()