import json, math
from pathlib import Path
import numpy as np

RNG_SEED = 7

def mgda(G, tol=1e-10):
    """Exact active-set enumeration for small M: min alpha'G'G alpha on simplex."""
    M = G.shape[1]; K = G.T @ G
    best = None
    for mask in range(1, 1 << M):
        A = [i for i in range(M) if mask >> i & 1]
        KA = K[np.ix_(A, A)]
        # KKT: KA a + lambda 1 = 0, 1'a=1
        mat = np.block([[KA, np.ones((len(A),1))],
                        [np.ones((1,len(A))), np.zeros((1,1))]])
        rhs = np.r_[np.zeros(len(A)), 1.0]
        try: sol = np.linalg.solve(mat, rhs)[:-1]
        except np.linalg.LinAlgError: sol = np.linalg.lstsq(mat, rhs, rcond=None)[0][:-1]
        if np.min(sol) < -tol: continue
        a = np.zeros(M); a[A] = sol
        val = float(a @ K @ a)
        if best is None or val < best[0]: best = (val, a)
    if best is None:
        # numerical fallback
        a = np.ones(M)/M
    else: a = best[1]
    return a, -G @ a

def reduced_eig(G):
    M=G.shape[1]; P=np.eye(M)-np.ones((M,M))/M
    vals=np.linalg.eigvalsh(P @ G.T @ G @ P)
    return float(vals[1] if M > 1 else vals[0])

class GatedMGDA:
    def __init__(self, M, eps_eig=.02, eps_alpha=.03, eps_change=.18, beta=.8):
        self.M=M; self.ee=eps_eig; self.ea=eps_alpha; self.ec=eps_change; self.beta=beta
        self.prev_bar=None; self.prev_A=None; self.regular=0; self.total=0; self.ratios=[]
    def direction(self,G):
        a, d_m = mgda(G); A=tuple(np.flatnonzero(a > self.ea))
        bar = G if self.prev_bar is None else self.beta*self.prev_bar+(1-self.beta)*G
        r=0.0 if self.prev_bar is None else np.linalg.norm(bar-self.prev_bar)/(np.linalg.norm(self.prev_bar)+1e-8)
        # Active Gram proxy requested by the idea; first step cannot have unchanged A.
        eig = float(np.min(np.linalg.eigvalsh(G[:,A].T @ G[:,A]))) if A else 0.0
        ok = self.prev_bar is not None and eig >= self.ee and min(a[list(A)], default=0) >= self.ea and A == self.prev_A and r <= self.ec
        d = d_m if ok else -G @ (np.ones(self.M)/self.M)
        self.prev_bar=bar.copy(); self.prev_A=A; self.total+=1; self.regular+=int(ok); self.ratios.append(r)
        return d, ok, r, eig, a

def math_checks():
    rng=np.random.default_rng(RNG_SEED)
    # Prediction 1: on nondegenerate interior geometry, direction error scales linearly.
    base=np.array([[1.0,.15,.05],[.1,1.0,.2],[.2,.1,1.1]])
    deltas=np.logspace(-5,-1,8); errs=[]
    for q in deltas:
        H=base + q*rng.normal(size=base.shape)
        errs.append(np.linalg.norm(mgda(base)[1]-mgda(H)[1]))
    slope=float(np.polyfit(np.log(deltas),np.log(np.maximum(errs,1e-14)),1)[0])
    # Prediction 2: gate activation falls with temporal noise and rises at low noise.
    rates=[]
    for noise in [0.01,.05,.15,.4]:
        g=GatedMGDA(3,eps_eig=.02,eps_alpha=.03,eps_change=.18,beta=.8); G=base.copy()
        for _ in range(100): g.direction(G); G=base+noise*rng.normal(size=base.shape)
        rates.append(float(g.regular/g.total))
    # Prediction 3: near-degenerate geometry has zero reduced curvature and should trigger fallback.
    # Construct two nearly collinear columns and perturb the third; measure local log slope.
    deg=np.array([[1.,1.,0.001],[0.,0.,1.]])
    es=[]
    for q in deltas:
        H=deg + q*rng.normal(size=deg.shape)
        es.append(np.linalg.norm(mgda(deg)[1]-mgda(H)[1]))
    slope_deg=float(np.polyfit(np.log(deltas),np.log(np.maximum(es,1e-14)),1)[0])
    return {'regular_direction_loglog_slope':slope,'predicted_regular_slope':1.0,
            'gate_noise_levels':[.01,.05,.15,.4],'gate_regular_fraction':rates,
            'predicted_gate_trend':'decreases as noise increases',
            'degenerate_direction_loglog_slope':slope_deg,'predicted_worst_case_exponent':.5,
            'regular_reduced_eigenvalue':reduced_eig(base),'degenerate_reduced_eigenvalue':reduced_eig(deg),
            'degenerate_curvature_prediction':'near-zero curvature => fallback',
            'degenerate_curvature_observed':'near-zero curvature; sensitivity exponent was not sublinear'}

def run_optimizer(method, seed, steps=300):
    rng=np.random.default_rng(seed); M=3; p=2
    # Conflicting quadratic objectives, with noisy per-task gradient estimates.
    centers=np.array([[-2.,0.],[2.,0.],[0.,2.]])
    x=np.array([0.4,-1.0]); gate=GatedMGDA(M) if method=='gated' else None
    losses=[]; worst=[]; smooth=[]; oks=0
    prev=None
    for t in range(steps):
        G=np.column_stack([x-c + .18*rng.normal(size=p) for c in centers])
        if method=='uniform': d=-G.mean(axis=1)
        elif method=='mgda': d=mgda(G)[1]
        else:
            d,ok,_,_,_=gate.direction(G); oks+=ok
        # fixed-step SGD; direction is descent update
        x=x+0.045*d
        true=np.sum((x-centers)**2,axis=1)/2
        losses.append(float(true.mean())); worst.append(float(true.max()))
        if prev is not None: smooth.append(float(np.linalg.norm(d-prev)))
        prev=d.copy()
    return {'mean_final':losses[-1],'worst_final':worst[-1],
            'mean_best':min(losses),'worst_best':min(worst),
            'direction_change':float(np.mean(smooth)), 'regular_fraction':oks/steps}

def main():
    checks=math_checks(); results={}
    for method in ['uniform','mgda','gated']:
        vals=[run_optimizer(method,s) for s in [11,23,41,59,71]]
        results[method]={k:float(np.mean([v[k] for v in vals])) for k in vals[0]}
        results[method+'_std']={k:float(np.std([v[k] for v in vals])) for k in vals[0]}
    out={'seed':RNG_SEED,'math_checks':checks,'benchmark':results}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()