import json, math, random, time
import numpy as np
import torch

SEED = 17
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)

# Diagonal specialization: G=diag(exp(g)), H=diag(exp(h)).
def plan_metric(g0, hhat, K=2.0, T=6, alpha=0.15, beta=12.0, iters=30):
    g0t = torch.tensor(g0, dtype=torch.float64)
    ht = torch.tensor(hhat, dtype=torch.float64)
    z = g0t.repeat(T, 1).clone().detach().requires_grad_(True)
    opt = torch.optim.Adam([z], lr=0.08)
    for _ in range(iters):
        opt.zero_grad()
        spread = (ht-z[-1]).max() - (ht-z[-1]).min()
        viol = torch.relu(spread - math.log(K))
        kinetic = ((z[1:]-z[:-1])**2).sum()/(2*T)
        loss = beta*viol**2 + alpha*kinetic
        loss.backward(); opt.step()
    return z.detach().numpy()

def affine_dist(A, B):
    w,V=np.linalg.eigh(A); Ai=V@np.diag(w**-.5)@V.T
    q=np.linalg.eigvalsh(Ai@B@Ai)
    return float(np.linalg.norm(np.log(q)))

def math_check():
    g=np.array([-.7,.2,1.1]); g2=np.array([.4,-.3,.8])
    exact=np.linalg.norm(g2-g)
    got=affine_dist(np.diag(np.exp(g)),np.diag(np.exp(g2)))
    A=np.array([[2.0,.4],[.4,1.0]]); B=np.array([[1.3,.2],[.2,2.4]])
    Q=np.linalg.qr(np.random.randn(2,2))[0]; scale=3.7
    inv_err=abs(affine_dist(A,B)-affine_dist(scale*Q@A@Q.T,scale*Q@B@Q.T))
    T=6; path=np.linspace(g,g2,T+1); ds=[np.linalg.norm(path[i+1]-path[i]) for i in range(T)]
    kinetic=sum(x*x for x in ds)/(2*T); length=sum(ds)
    identity_err=abs(kinetic-length*length/(2*T*T))
    # Controller should reduce terminal condition number toward K from a bad metric.
    h=np.log(np.array([1., 10., 100., 1000.]))
    g0=np.zeros(4); planned=plan_metric(g0,h,K=3,T=6)
    before=math.exp((h-g0).max()-(h-g0).min())
    after=math.exp((h-planned[-1]).max()-(h-planned[-1]).min())
    return dict(diagonal_distance_error=abs(exact-got), affine_invariance_error=inv_err,
                kinetic_identity_error=identity_err, condition_before=before,
                condition_after=after, target=3.0, math_pass=(abs(exact-got)<1e-10 and inv_err<1e-9 and identity_err<1e-10 and after<before))

def run(kind, n=10, steps=500, lr=.07, noise=.35):
    # Fixed ill-conditioned quadratic with noisy gradients; metric updates every R steps.
    eig=np.geomspace(1.,100.,n); h=np.log(eig); x=np.ones(n)*2.; gmetric=np.zeros(n)
    ema=np.zeros(n); prev_metric=gmetric.copy(); jumps=[]; conds=[]; losses=[]; gradvars=[]
    t0=time.perf_counter(); window=[]
    for step in range(steps):
        grad=eig*x
        noisy=grad + noise*np.sqrt(eig)*np.random.randn(n)
        window.append(noisy.copy())
        if kind=='idea' and step%20==0:
            # curvature estimate from noisy diagonal gradient/parameter statistics
            hhat=h + .10*np.random.randn(n)
            plan=plan_metric(gmetric,hhat,K=8,T=4,alpha=.25,beta=8,iters=18)
            newg=plan[0]
            jumps.append(float(np.linalg.norm(newg-gmetric)))
            gmetric=newg
        elif kind=='ema' and step%20==0:
            # standard EMA diagonal second moment, normalized only for scale
            if len(window):
                sq=np.mean(np.array(window[-20:])**2,axis=0)
                ema=.95*ema+.05*sq
            newg=.5*np.log(ema+1e-4); newg-=newg.mean()
            jumps.append(float(np.linalg.norm(newg-gmetric)))
            gmetric=newg
        # Use normalized metrics so the comparison is about shape, not global scale.
        inv=np.exp(-gmetric); inv=inv/np.mean(inv)
        x -= lr*inv*noisy/np.sqrt(eig.mean())
        losses.append(float(.5*np.sum(eig*x*x)))
        M=eig*np.exp(-gmetric); conds.append(float(M.max()/M.min()))
        gradvars.append(float(np.var(noisy)))
    return dict(final_loss=losses[-1], best_loss=min(losses), final_condition=conds[-1],
                mean_condition=float(np.mean(conds[-100:])), mean_jump=float(np.mean(jumps)),
                gradient_variance=float(np.mean(gradvars[-100:])), seconds=time.perf_counter()-t0,
                losses=losses, conditions=conds)

def main():
    check=math_check()
    # Reset randomness before fair paired runs.
    np.random.seed(SEED+1); base=run('ema')
    np.random.seed(SEED+1); idea=run('idea')
    repeats=[]
    for seed in [3, 11, 29, 47]:
        np.random.seed(seed); b=run('ema')
        np.random.seed(seed); q=run('idea')
        repeats.append({'seed':seed, 'baseline':{k:b[k] for k in ('final_loss','best_loss','mean_condition','mean_jump','gradient_variance','seconds')}, 'idea':{k:q[k] for k in ('final_loss','best_loss','mean_condition','mean_jump','gradient_variance','seconds')}})
    out={'math_check':check,'baseline':base,'idea':idea,'repeats':repeats}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps({'math_check':check,'baseline':{k:v for k,v in base.items() if k not in ('losses','conditions')},'idea':{k:v for k,v in idea.items() if k not in ('losses','conditions')},'repeats':repeats},indent=2))

if __name__=='__main__': main()