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

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

# Primitive values are signed distances for two target/constraint primitives.
def primitives(x):
    # Two smooth primitive value functions; <=0 is the desired set.
    z1=(x[:,0]-0.35)**2+(x[:,1]+0.15)**2-0.42
    z2=(x[:,0]+0.40)**2+(x[:,1]-0.20)**2-0.36
    return torch.stack([z1,z2],1)

def positive_linear_fit(z,y,steps=1200):
    # A monotone aggregator: b + softplus(a_i) z_i, with strictly nonnegative derivatives.
    raw=nn.Parameter(torch.zeros(2)); b=nn.Parameter(torch.zeros(()))
    opt=torch.optim.Adam([raw,b],lr=.04)
    for _ in range(steps):
        opt.zero_grad(); w=torch.nn.functional.softplus(raw)
        loss=((b+z@w-y)**2).mean(); loss.backward(); opt.step()
    return b.detach(), torch.nn.functional.softplus(raw).detach()

def unconstrained_fit(z,y,steps=1200):
    w=nn.Parameter(torch.zeros(2)); b=nn.Parameter(torch.zeros(()))
    opt=torch.optim.Adam([w,b],lr=.04)
    for _ in range(steps):
        opt.zero_grad(); loss=((b+z@w-y)**2).mean(); loss.backward(); opt.step()
    return b.detach(),w.detach()

def mask_iou(a,b):
    inter=(a&b).sum(); union=(a|b).sum()
    return (inter/union).item() if union else 1.0

def main():
    # ---------------- core math checks ----------------
    n=50000
    z=torch.randn(n,2)*1.3
    # Positive aggregator with known Lipschitz coordinate coefficients.
    w=torch.tensor([0.7,1.4]); b=0.2
    e=torch.randn(n,2)*0.03
    err=(b+(z+e)@w-(b+z@w)).abs()
    bound=(e.abs()@w)
    bound_ratio=(err/(bound+1e-12)).max().item()
    # Coordinatewise ordering preservation (finite differences, all coordinates nonnegative).
    delta=torch.rand(n,2)*2
    diff=(b+(z+delta)@w-(b+z@w))
    order_fraction=(diff>=-1e-10).all().item() if diff.ndim==0 else (diff>=-1e-10).float().mean().item()
    # In the linear case error scales exactly linearly with primitive error amplitude.
    scales=torch.tensor([0.25,0.5,1.,2.,4.])
    scaling=[]
    for s in scales:
        ee=torch.randn(n,2)*.02*s
        scaling.append((ee@w).abs().mean().item())
    slopes=np.polyfit(np.log(scales.numpy()),np.log(np.array(scaling)),1)[0]

    # ---------------- learned critic mini experiment ----------------
    # Train on left half of state square and test on right half: primitive critics are fixed,
    # while the aggregator must preserve their task ordering under shift.
    N=5000
    x=torch.rand(N,2)*2-1
    z0=primitives(x)
    noise=torch.randn_like(z0)*0.025
    zhat=z0+noise
    # Ground truth composite is monotone and has a known zero-sublevel set.
    y=z0@torch.tensor([1.,0.8])+0.03
    train=x[:,0]<0.05; test=~train
    bp,wp=positive_linear_fit(zhat[train],y[train])
    bu,wu=unconstrained_fit(zhat[train],y[train])
    yp=bp+zhat@wp; yu=bu+zhat@wu
    rmse_p=torch.sqrt(((yp[test]-y[test])**2).mean()).item()
    rmse_u=torch.sqrt(((yu[test]-y[test])**2).mean()).item()
    iou_p=mask_iou(yp[test]<=0,y[test]<=0); iou_u=mask_iou(yu[test]<=0,y[test]<=0)

    # Negative derivative sweep: compare c2*z2 with the true c2=.8 composite.
    # The ordering theorem predicts failure precisely after c2 crosses zero.
    cs=np.array([1.0,.5,.2,.05,0.,-.02,-.05,-.1,-.2,-.5,-1.0])
    sweep=[]
    for c in cs:
        pred=z0[:,0]+float(c)*z0[:,1]+.03
        # set disagreement and order reversals relative to increasing z2 at fixed z1
        set_err=(pred<=0).ne(y<=0).float().mean().item()
        # derivative sign is c; for c<0, every positive z2 increment reverses contribution
        # Direct monotonicity test: increasing z2 must not decrease the output.
        dz=torch.rand(N)*2.0
        reversal=(float(c)*dz < -1e-12).float().mean().item()
        sweep.append({'c2':float(c),'derivative':float(c),'set_disagreement':set_err,
                      'order_reversal_fraction':reversal})

    result={
      'predictions':{
        'P1_monotone_derivatives_nonnegative':{'predicted':'100% samples','observed_fraction':order_fraction,'tolerance':'exact up to numerical precision'},
        'P2_error_bound':{'predicted':'|delta V| <= 0.7|e1|+1.4|e2| and linear scaling','max_ratio_observed':bound_ratio,'loglog_scaling_exponent':float(slopes),'expected_exponent':1.0},
        'P3_negative_crossing':{'predicted':'set/order failure begins at c2<0','sweep':sweep}
      },
      'mini_experiment':{
        'positive_aggregator_weights':wp.tolist(),'unconstrained_weights':wu.tolist(),
        'heldout_rmse':{'monotone':rmse_p,'unconstrained':rmse_u},
        'heldout_set_iou':{'monotone':iou_p,'unconstrained':iou_u}
      }
    }
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()