import json
import numpy as np

SEED = 2293

def projK(x, identity=False): return x.copy() if identity else np.maximum(x, 0.0)

def projector(A, b):
    A = np.asarray(A, float); b = np.asarray(b, float)
    G = np.linalg.pinv(A @ A.T)
    return lambda w: w - A.T @ (G @ (A @ w - b))

def run(z0, A, b, c, gamma, beta, lam, T, identity=False):
    P = projector(A, b); z = z0.copy(); out=[]
    for _ in range(T):
        u=projK(z, identity); v=P(2*u-z-gamma*beta*c)
        out.append({'z':z.copy(),'u':u.copy(),'v':v.copy(),
                    'r':float(np.linalg.norm(v-u)),
                    'p':float(np.linalg.norm(A@v-b))})
        z=z+lam*(v-u)
    return out

def controller(z0,A,b,c,gamma,T):
    P=projector(A,b); z=z0.copy(); lam=1.; beta=1.; oldr=np.inf; oldp=np.inf; out=[]
    for t in range(T):
        u=projK(z); v=P(2*u-z-gamma*beta*c)
        r=float(np.linalg.norm(v-u)); p=float(np.linalg.norm(A@v-b))
        if np.isfinite(oldr):
            decay=.96**t
            lam=float(np.clip(lam*np.exp(.18*decay*np.clip(np.log((oldr+1e-12)/(r+1e-12)),-1,1)),.15,1.85))
            if r < oldr and p <= oldp*(1+1e-6): beta=min(2.,beta*(1+.08*decay))
            elif r > 1.1*oldr or p > 1.1*oldp:
                beta=max(.1,beta*(1-.12*decay)); lam=max(.15,.8*lam)
        out.append({'u':u,'v':v,'r':r,'p':p,'lam':lam,'beta':beta})
        z=z+lam*(v-u); oldr=r; oldp=p
    return out

def affine_check():
    rng=np.random.default_rng(SEED); A=rng.normal(size=(3,5)); b=rng.normal(size=3); P=projector(A,b)
    w=rng.normal(size=5); x=P(w); _,_,vh=np.linalg.svd(A); null=vh[3:].T
    return {'equality_error':float(np.linalg.norm(A@x-b)),
            'idempotence_error':float(np.linalg.norm(P(x)-x)),
            'nullspace_orthogonality':float(np.linalg.norm(null.T@(w-x)))}

def exact_toy():
    # K=R^n and A=I,b=0: u=z,v=0, hence z_next=(1-lambda)z.
    n=3; A=np.eye(n); b=np.zeros(n); c=np.zeros(n); z=np.array([1.,-2.,.5]); T=12; ans=[]
    for lam in [.25,.5,1.,1.5,1.75,2.,2.1,2.5]:
        q=run(z,A,b,c,1.,1.,lam,T,identity=True)
        rate=(q[-1]['r']/q[0]['r'])**(1/(T-1))
        ans.append({'lambda':lam,'predicted_rate_abs_1_minus_lambda':abs(1-lam),
                    'observed_rate':float(rate),'final_over_initial':float(q[-1]['r']/q[0]['r']),
                    'contractive':bool(abs(1-lam)<1)})
    return ans

def finite_bound():
    z=np.array([1.,-2.,.5]); D=float(z@z); N=40; ans=[]
    for lam in [.25,.75,1.25,1.75]:
        q=run(z,np.eye(3),np.zeros(3),np.zeros(3),1,1,lam,N,identity=True)
        vals=np.array([x['r'] for x in q]); alpha=lam/2; bound=alpha/(1-alpha)*D/N
        ans.append({'lambda':lam,'N_min_r2':float(N*np.min(vals**2)),
                    'bound':float(bound),'ratio':float(N*np.min(vals**2)/bound)})
    return ans

def mini():
    rng=np.random.default_rng(SEED); B=[]; C=[]; lams=[]; betas=[]
    for _ in range(100):
        n=8; A=np.ones((1,n)); b=np.ones(1); c=rng.uniform(.1,2.,n); z=rng.normal(size=n); opt=float(c.min())
        f=run(z,A,b,c,.35,1.,1.,24)[-1]; cc=controller(z,A,b,c,.35,24); g=cc[-1]
        def m(q):
            u=q['u']; eq=abs(float((A@u-b)[0])); cone=float(np.linalg.norm(np.minimum(u,0)))
            return [float(c@u-opt+10*eq),eq,cone,q['r']]
        B.append(m(f)); C.append(m(g)); lams.append(g['lam']); betas.append(g['beta'])
    B=np.array(B); C=np.array(C)
    return {'trials':100,'T':24,'fixed_median_merit':float(np.median(B[:,0])),
      'controller_median_merit':float(np.median(C[:,0])),'fixed_median_eq_violation':float(np.median(B[:,1])),
      'controller_median_eq_violation':float(np.median(C[:,1])),'fixed_median_residual':float(np.median(B[:,3])),
      'controller_median_residual':float(np.median(C[:,3])),'controller_final_lambda_median':float(np.median(lams)),
      'controller_final_beta_median':float(np.median(betas)),'residual_win_fraction':float(np.mean(C[:,3]<B[:,3]))}

def main():
    print(json.dumps({'affine_projection':affine_check(),'exact_toy':exact_toy(),
                      'finite_step_bound':finite_bound(),'mini_experiment':mini()},indent=2))
if __name__=='__main__': main()