import json, random
import numpy as np
import torch

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

def support(q, H):
    return float(np.sum(np.abs(q @ H)))

def margin(q, c0, H0, cj, Hj):
    return float(q @ (c0 - cj) - support(q, H0) - support(q, Hj))

def math_check():
    # Exact numerical checks of Minkowski addition and linear-map identities.
    c1=np.array([.3,-.2]); H1=np.array([[.4,.1],[.0,.25]])
    c2=np.array([-.1,.5]); H2=np.array([[.2],[.3]])
    K=np.array([[1.2,-.4],[.5,.8]])
    z1=np.random.uniform(-1,1,(20000,2)); z2=np.random.uniform(-1,1,(20000,1))
    lhs=(c1+z1@H1.T)+(c2+z2@H2.T)
    rhs=(c1+c2)+np.concatenate([z1,z2],1)@np.concatenate([H1,H2],1).T
    mink_err=float(np.max(np.abs(lhs-rhs)))
    lhs2=(c1+z1@H1.T)@K.T
    rhs2=(K@c1)+z1@(K@H1).T
    map_err=float(np.max(np.abs(lhs2-rhs2)))

    # Choose q toward the trusted center, so the certificate is positive.
    c0=np.array([0.,0.]); H0=np.diag([.1,.08])
    cj=np.array([.55,0.]); Hj=np.diag([.08,.06]); q=np.array([-1.,0.])
    m=margin(q,c0,H0,cj,Hj)
    trusted_right=q@c0-support(q,H0)
    attack_left=q@cj+support(q,Hj)
    return dict(minkowski_max_error=mink_err, linear_map_max_error=map_err,
                positive_certificate_margin=m,
                sampled_projection_gap=float(trusted_right-attack_left))

class Fusion(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net=torch.nn.Sequential(torch.nn.Linear(2,12),torch.nn.Tanh(),torch.nn.Linear(12,1))
    def forward(self,x): return self.net(x)

def jacobian(model, x):
    x=x.detach().requires_grad_(True); y=model(x); rows=[]
    for k in range(y.shape[1]):
        rows.append(torch.autograd.grad(y[:,k].sum(),x,retain_graph=True)[0])
    return torch.stack(rows,1)  # B,out,in

def probe_test():
    # Two-sensor nonlinear fusion regression; sensor 0 has a fixed attack shift.
    torch.manual_seed(SEED)
    n=256; x=torch.rand(n,2)*2-1
    y=x[:,0:1]+.6*x[:,1:2]+.2*torch.sin(3*x[:,0:1])
    model=Fusion(); opt=torch.optim.Adam(model.parameters(),lr=.03)
    for _ in range(250):
        opt.zero_grad(); loss=((model(x)-y)**2).mean(); loss.backward(); opt.step()

    x0=x[:1].clone(); J=jacobian(model,x0)[0,0].detach().numpy()[None,:]
    He=np.diag([.08,.08]); attack=np.array([[.45],[0.]])
    cbase=float(model(x0)); c0=np.array([cbase]); cj=np.array([cbase+float((J@attack).item())])
    H0=J@He; Hj=J@He; eps=.25
    ds=np.linspace(-eps,eps,101)
    qs=np.array([[-1.],[1.]])
    local_margins=[]
    exact_sep=[]
    for d in ds:
        # Under the supplied local model, both hypotheses receive the same J*d.
        local_margins.append(max(margin(q,c0+J[:,0]*d,H0,cj+J[:,0]*d,Hj) for q in qs))
        with torch.no_grad():
            a=model(x0+torch.tensor([[d,0.]],dtype=x0.dtype)).item()
            b=model(x0+torch.tensor([[d+.45,0.]],dtype=x0.dtype)).item()
        exact_sep.append(abs(b-a))
    local_margins=np.asarray(local_margins); exact_sep=np.asarray(exact_sep)
    zero_i=len(ds)//2; best_i=int(np.argmax(exact_sep))
    # Monte Carlo confirms the support bound for the trusted linearized zonotope.
    z=np.random.uniform(-1,1,(10000,2)); proj=(c0+(z@(H0.T))).ravel()
    bound=float(abs(c0[0])+support(np.array([1.]),H0))
    return dict(train_mse=float(((model(x)-y)**2).mean()), jacobian=J.tolist(),
        local_margin_at_zero=float(local_margins[zero_i]),
        local_margin_best=float(local_margins.max()),
        local_margin_range=float(local_margins.max()-local_margins.min()),
        exact_separation_at_zero=float(exact_sep[zero_i]),
        exact_separation_best=float(exact_sep.max()),
        exact_separation_gain=float(exact_sep[best_i]-exact_sep[zero_i]),
        best_probe=float(ds[best_i]),
        trusted_projection_max_abs=float(np.max(np.abs(proj))),
        trusted_support_bound=bound)

def main():
    out={'seed':SEED,'math_check':math_check(),'probe_test':probe_test()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()