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

SEED=2100
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device='cuda' if torch.cuda.is_available() else 'cpu'
try:
    if device=='cuda': torch.cuda.init()
except Exception:
    device='cpu'

# Scalar robust interval: actuator [-1,1] intersect two opposing constraints.
# For the tested region x<2.8, M(x)=2*width-0.5*x.
def interval_np(x,width=.7):
    x=np.asarray(x)
    return np.maximum(-1.,-width+.25*x), np.minimum(1.,width-.25*x)
def interval_t(x,width=.7):
    z=torch.zeros_like(x)
    return torch.maximum(z-1.,-width+.25*x), torch.minimum(z+1.,width-.25*x)
def project_np(u,lo,hi): return np.minimum(np.maximum(u,lo),hi)

def math_checks():
    xs=np.linspace(0,3.6,3601); lo,hi=interval_np(xs); M=hi-lo
    feasible=lo<=hi
    # Tolerance is needed because M at the exact floating grid boundary is roundoff-negative.
    boundary=float(xs[np.argmin(np.abs(M))])
    iff_err=int(np.sum(feasible!=(M>=-1e-10)))
    # Prediction 1: feasibility boundary x*=2*width/.5 = 2.8 at width=.7.
    # Prediction 2: hinge derivative is -2(m0-M) below m0 and zero above.
    m0=.7; g=np.linspace(0,1.4,141); e=1e-5
    f=lambda q: np.maximum(0,m0-q)**2
    fd=(f(g+e)-f(g-e))/(2*e); pred=np.where(g<m0,-2*(m0-g),0.)
    # Prediction 3: projection correction outside an endpoint is exactly endpoint-u.
    us=np.linspace(.7,1.8,100); endpoint=.4
    correction=project_np(us,-1,endpoint)-us
    return {'boundary_prediction':{'predicted_x':2.8,'observed_x':boundary,'abs_error':abs(boundary-2.8),'feasibility_iff_errors':iff_err},
            'hinge_prediction':{'finite_difference_rmse':float(np.sqrt(np.mean((fd-pred)**2))),
                                'max_derivative_above_m0':float(np.max(np.abs(fd[g>=m0])))},
            'projection_prediction':{'max_error_vs_endpoint_minus_action':float(np.max(np.abs(correction+(us-endpoint))))}}

class Policy(nn.Module):
    def __init__(self):
        super().__init__()
        self.net=nn.Sequential(nn.Linear(1,24),nn.Tanh(),nn.Linear(24,24),nn.Tanh(),nn.Linear(24,1),nn.Tanh())
    def forward(self,x): return self.net(x).squeeze(-1)

def rollout_metrics(p,width=.7,project=True):
    with torch.no_grad():
        x=torch.full((128,),1.45,device=device)
        mins=[]; actions=[]; corrections=[]
        for _ in range(18):
            u=p(x[:,None]); lo,hi=interval_t(x,width); M=hi-lo
            ue=torch.maximum(torch.minimum(u,hi),lo) if project else u
            mins.append(M); actions.append(u); corrections.append(ue-u)
            x=x+.075*ue+.018
        M=torch.cat(mins); C=torch.cat(corrections); U=torch.cat(actions)
        return {'p05_margin':float(torch.quantile(M,.05).cpu()),'min_margin':float(M.min().cpu()),
                'negative_margin_rate':float((M<0).float().mean().cpu()),
                'intervention_rate':float((C.abs()>1e-5).float().mean().cpu()),
                'correction_rms':float(torch.sqrt((C*C).mean()).cpu()),
                'intervention_energy':float((C*C).mean().cpu()),
                'mean_action':float(U.mean().cpu())}

def train(use_margin=False,use_baseline=False,width=.7,seed=SEED):
    torch.manual_seed(seed); p=Policy().to(device); opt=torch.optim.Adam(p.parameters(),lr=4e-3)
    m0=.30
    for _ in range(180):
        x=torch.full((128,),1.45,device=device)
        task=0.; ml=0.; bl=0.
        for _ in range(18):
            u=p(x[:,None]); lo,hi=interval_t(x,width); M=hi-lo
            # Training uses the exact projection, with endpoint gradients retained for this smooth affine toy.
            ue=torch.maximum(torch.minimum(u,hi),lo)
            # Aggressive task drives x toward the boundary; margin regularization opposes that drift.
            task=task+(u-.82).pow(2).mean()
            ml=ml+torch.relu(m0-M).pow(2).mean()
            baseline=torch.clamp(-.18*x,-1,1)
            bl=bl+(ue-baseline).pow(2).mean()
            x=x+.075*ue+.018
        loss=task/28 + (3.0*ml/28 if use_margin else 0.) + (.12*bl/28 if use_baseline else 0.)
        opt.zero_grad(); loss.backward(); opt.step()
    r=rollout_metrics(p,width)
    r['task_action_mse']=float((p(torch.full((128,1),1.45,device=device))-.82).pow(2).mean().detach().cpu())
    return r

def run():
    widths=[.55,.70,.85]
    sweep=[]
    for w in widths:
        b=train(False,False,w); m=train(True,False,w); mb=train(True,True,w)
        # Analytic width prediction at the initial state: Delta p05 M=2 Delta width.
        sweep.append({'width':w,'baseline':b,'margin':m,'margin_baseline':mb,
                      'predicted_initial_margin':2*w-.5*1.45})
    return {'device':device,'math_checks':math_checks(),'training_width_sweep':sweep}

if __name__=='__main__':
    out=run(); Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))