import json, math, random
import numpy as np
import torch
from torch import nn

SEED=7
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)

def psd_check(SY,SX):
    d=SY-SX
    return float(np.linalg.eigvalsh((d+d.T)/2).min()), float(np.max(np.abs(d-d.T)))

def affine_prob(A,b,z,S,n=300000,seed=11):
    rng=np.random.default_rng(seed)
    L=np.linalg.cholesky(S+1e-12*np.eye(S.shape[0]))
    u=rng.standard_normal((n,S.shape[0]))@L.T + z
    ok=np.all(u@A.T <= b[None,:],axis=1)
    return float(ok.mean()), int(ok.sum())

def theorem_sanity():
    # A bounded convex rectangle, with deliberately inflated correlated covariance.
    SX=np.array([[.20,.08],[.08,.18]])
    SY=SX + .05*np.outer([1.,2.],[1.,2.])
    lo, asym=psd_check(SY,SX)
    A=np.array([[1,0],[-1,0],[0,1],[0,-1.]])
    b=np.array([1.55,1.55,1.55,1.55])
    # Centered rectangle has probability > 1/2; compare paired Monte Carlo estimates.
    py,ny=affine_prob(A,b,np.zeros(2),SY,300000,13)
    px,nx=affine_prob(A,b,np.zeros(2),SX,300000,14)
    # Also check an affine shifted set at several z values only where inflated p >= .5.
    rectangle_py, rectangle_px = py, px
    rows=[]
    for z in [np.array([0.,0.]),np.array([.25,-.15]),np.array([-.35,.10])]:
        py,_=affine_prob(A,b,z,SY,120000,30+len(rows))
        px,_=affine_prob(A,b,z,SX,120000,40+len(rows))
        rows.append((float(py),float(px)))
    return {'min_eigenvalue_SY_minus_SX':lo,'symmetry_error':asym,
            'rectangle_p_inflated':rectangle_py,'rectangle_p_smaller':rectangle_px,
            'shifted_pairs':rows,'all_observed_ordering':all(y>=.5 and x+0.012>=y for y,x in rows)}

class Net(nn.Module):
    def __init__(self):
        super().__init__(); self.body=nn.Sequential(nn.Linear(2,12),nn.Tanh(),nn.Linear(12,2))
    def forward(self,x): return self.body(x)

def train(mode, X, y, SY, steps=700, K=32):
    torch.manual_seed(SEED); model=Net(); opt=torch.optim.Adam(model.parameters(),lr=.025)
    L=torch.tensor(np.linalg.cholesky(SY),dtype=torch.float32)
    # acceptance is true-class logit minus false-class logit >= margin
    for step in range(steps):
        ix=torch.randint(0,len(X),(64,)); xb=X[ix]; yb=y[ix]
        z=model(xb); other=1-yb
        margin=z[torch.arange(64),yb]-z[torch.arange(64),other]
        # same inflated-noise Monte Carlo cost for both methods
        eps=torch.randn(K,64,2); noise=torch.einsum('ij,kbj->kbi',L,eps)
        zn=model(xb[None,:,:].expand(K,-1,-1).reshape(K*64,2)).reshape(K,64,2)
        zn=zn+noise
        noisy_margin=zn[:,:,0]-zn[:,:,1] # only valid for class order, signs below
        noisy_margin=torch.where(yb[None,:].bool(), noisy_margin, -noisy_margin)
        task=nn.functional.softplus(-margin).mean()
        if mode=='clean':
            loss=task
        elif mode=='ordinary':
            aug=nn.functional.softplus(-noisy_margin).mean()
            loss=task+0.7*aug
        else:
            # smooth estimate of P(margin >= 1), target p0=.80
            p_hat=torch.sigmoid(noisy_margin/.18).mean()
            loss=task+0.10*torch.relu(.55-p_hat)**2
        opt.zero_grad(); loss.backward(); opt.step()
    return model

def evaluate(model,X,y,Slist):
    with torch.no_grad():
        logits=model(X); acc=(logits.argmax(1)==y).float().mean().item()
    out={'clean_accuracy':acc}
    for name,S in Slist:
        L=torch.tensor(np.linalg.cholesky(S),dtype=torch.float32)
        n=12000; xx=X[:min(256,len(X))]; yy=y[:len(xx)]
        eps=torch.randn(n,len(xx),2); zz=model(xx[None].expand(n,-1,-1).reshape(n*len(xx),2)).reshape(n,len(xx),2)
        noise=torch.einsum('ij,kbj->kbi',L,eps)
        # representation perturbation is applied after model, matching the affine head.
        zz=zz+noise
        pred=zz.argmax(2); accept=(pred==yy[None,:]).float().mean().item()
        out[name]=accept
    return out

def mini_experiment():
    rng=np.random.default_rng(SEED); n=900
    X=np.r_[rng.multivariate_normal([-1.0,-.7],[[.45,.10],[.10,.35]],n//2),rng.multivariate_normal([1.0,.7],[[.45,.10],[.10,.35]],n//2)].astype('float32')
    y=np.r_[np.zeros(n//2,dtype=np.int64),np.ones(n//2,dtype=np.int64)]
    p=rng.permutation(n); X=X[p]; y=y[p]
    Xt=torch.tensor(X); yt=torch.tensor(y)
    SX=np.array([[.08,.025],[.025,.06]])
    SY=SX+.12*np.outer([1.,.5],[1.,.5])
    Slist=[('small_cov',SX),('inflated_cov',SY),('isotropic_small',np.eye(2)*.04)]
    results={}
    for mode in ['clean','ordinary','chance']:
        m=train(mode,Xt,yt,SY); results[mode]=evaluate(m,Xt,yt,Slist)
    return {'covariance_min_eigenvalue':float(np.linalg.eigvalsh(SY-SX).min()),'results':results}

if __name__=='__main__':
    print(json.dumps({'sanity':theorem_sanity(),'experiment':mini_experiment()},indent=2))