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

SEED=7
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
try:
    if DEVICE=='cuda': torch.cuda.get_device_properties(0)
except Exception:
    DEVICE='cpu'

# Scalar reconstructed barrier: psi = safe_action(x)-u.  psi>=0 means nominally safe.
# Positive r=kappa*epsilon-psi is the proposed margin violation.
def softplus_sq(r,tau):
    return np.logaddexp(0.0,r/tau)**2

def verify_update():
    delta=.1; rho=.25; e=0.32; eps=0.02; vals=[]
    for _ in range(100):
        eps=np.clip(eps+rho*(e-(1-delta)*eps), .001, 2.)
        vals.append(eps)
    fixed=e/(1-delta)
    # exact linear convergence: error after n = (1-rho*(1-delta))^n error0
    pred_factor=1-rho*(1-delta)
    observed_factor=np.mean(np.abs(np.diff(vals[-30:])[1:])/np.abs(np.diff(vals[-30:])[:-1]))
    # coverage transition for Gaussian error magnitudes: eps at 90th percentile
    rng=np.random.default_rng(SEED)
    cover=[]
    for sigma in [.05,.10,.20,.40]:
        err=np.abs(rng.normal(0,sigma,200000))
        q=float(np.quantile(err,.9)); rm_eps=float(np.mean(err)/(1-delta)); rm_cov=float(np.mean(err<=rm_eps)); pred_rm_cov=math.erf((math.sqrt(2/math.pi)/(1-delta))/math.sqrt(2)); cover.append({'sigma':sigma,'q90':q,'coverage_at_q90':float(np.mean(err<=q)), 'rm_epsilon_pred':rm_eps, 'rm_coverage_pred':pred_rm_cov, 'rm_coverage_obs':rm_cov})
    return {'fixed_point_pred':fixed,'fixed_point_obs':vals[-1],
            'convergence_factor_pred':pred_factor,'convergence_factor_obs':float(observed_factor),
            'coverage_sweep':cover}

class Policy(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,1),nn.Tanh())
    def forward(self,x): return 1.5*self.net(x)

def train(method, sigma, seed=SEED, steps=700):
    torch.manual_seed(seed); rng=np.random.default_rng(seed)
    p=Policy().to(DEVICE); opt=torch.optim.Adam(p.parameters(),lr=.012)
    eps=.03; delta=.1; rho=.08; kappa=1.0; tau=.08; beta=2.0
    x=torch.linspace(-1,1,96,device=DEVICE).reshape(-1,1)
    target=torch.full_like(x,.75)
    for step in range(steps):
        # deterministic known observer bias plus random true disturbance; detach observer quantities
        d=torch.tensor(rng.normal(0,sigma,96),dtype=torch.float32,device=DEVICE).reshape(-1,1)
        dh=torch.zeros_like(d)
        err=float(torch.mean(torch.abs(d-dh)).item())
        if method=='adaptive': eps=float(np.clip(eps+rho*(err-(1-delta)*eps),.001,1.5))
        elif method=='fixed': eps=.22
        else: eps=0.
        u=p(x); safe=-.15*x + .05 # nominal action allowed below this
        psi=safe-u
        r=kappa*eps-psi
        task=torch.mean((u-target)**2)
        barrier=beta*torch.mean(torch.nn.functional.softplus(r/tau)**2)
        loss=task + (barrier if method!='baseline' else 0.)
        opt.zero_grad(); loss.backward(); opt.step()
    # Evaluation over states/disturbances. Filter projects u+d_hat <= safe - k eps.
    xx=torch.linspace(-1,1,256,device=DEVICE).reshape(-1,1)
    with torch.no_grad(): raw=p(xx).cpu().numpy().ravel()
    safe=(-.15*xx+.05).cpu().numpy().ravel()
    # fresh disturbances, observer d_hat=0; projection is the inference-time barrier filter
    erng=np.random.default_rng(seed+1000)
    d=erng.normal(0,sigma,(40,256)); dh=0.
    threshold=safe-kappa*eps
    filt=np.minimum(raw,threshold)
    intervention=np.abs(raw-filt)
    actual=filt[None,:]+d
    violations=float(np.mean(actual>safe[None,:]))
    nominal_viol=float(np.mean(raw>safe))
    track=float(np.mean((filt-.75)**2))
    return {'eps':eps,'violations_after_filter':violations,'nominal_violations':nominal_viol,
            'intervention_norm':float(np.mean(intervention)),'tracking_mse':track,
            'barrier_fraction':float(np.mean(raw>threshold))}

def main():
    checks=verify_update(); results={'device':DEVICE,'verification':checks,'experiments':{}}
    for sigma in [.05,.10,.20,.40]:
        results['experiments'][str(sigma)]={m:train(m,sigma) for m in ['baseline','fixed','adaptive']}
    Path('results.json').write_text(json.dumps(results,indent=2))
    print(json.dumps(results,indent=2))
if __name__=='__main__': main()