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

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

# Plant: xdot = u, safe set h(x)=x >= 0. Thus r(x)=u(x)+alpha*x.
ALPHA = 1.0
EPS = 0.08

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))
    def forward(self, x):
        return 1.5 * torch.tanh(self.net(x))

def certification_sweep():
    # For each uniform delta-net, construct a smooth residual that equals EPS
    # at net points and dips between them.  r=EPS+A(cos(2*pi*x/s)-1),
    # s=2*delta and A=delta/pi.  Its true Lipschitz constant is 1 and
    # the dense gap is exactly 2*delta/pi: a direct O(delta) prediction.
    dense = np.linspace(0, 1, 200001)
    rows=[]
    for n in [3, 5, 9, 17, 33, 65, 129]:
        s = 1/(n-1)
        delta = s/2
        A = delta/np.pi
        r_dense = EPS + A*(np.cos(2*np.pi*dense/s)-1)
        grid = np.linspace(0, 1, n)
        r_grid = EPS + A*(np.cos(2*np.pi*grid/s)-1)
        L = 1.0
        actual_min = float(r_dense.min())
        gap = float(EPS - actual_min)
        bound = EPS - L*delta
        rows.append(dict(n=n, delta=delta, L=L, predicted_lower=bound,
                         actual_min=actual_min, observed_gap=max(0.0,gap),
                         positive_certificate=(bound > 0), actual_safe=(actual_min >= 0)))
    x = np.linspace(0,1,1001)
    # Check the claimed global Lipschitz inequality numerically.
    max_pair_excess = 0.0
    for n in [3, 5, 9, 17, 33]:
        s=1/(n-1); A=(s/2)/np.pi
        r=EPS+A*(np.cos(2*np.pi*x/s)-1)
        for i in range(0,1001,7):
            for j in range(0,1001,11):
                max_pair_excess=max(max_pair_excess, abs(r[i]-r[j])-abs(x[i]-x[j]))
    bound_ok = all(row['actual_min'] + 1e-7 >= row['predicted_lower'] for row in rows)
    ratios = [row['observed_gap']/row['delta'] for row in rows]
    scaling_ok = max(abs(q-2/np.pi) for q in ratios) < 2e-4
    ds=np.array([q['delta'] for q in rows]); gs=np.array([q['observed_gap'] for q in rows])
    slope=float(np.polyfit(np.log(ds), np.log(gs), 1)[0])
    predicted_boundary=EPS/1.0
    observed_boundary=max(q['delta'] for q in rows if q['actual_safe'])
    return dict(L=1.0, rows=rows, predicted_boundary_delta=predicted_boundary,
                observed_safe_boundary_delta=observed_boundary,
                pairwise_lipschitz_excess=max_pair_excess,
                prediction_checks={'lower_bound_holds':bound_ok,
                                   'gap_ratio_predicted':2/np.pi,
                                   'gap_ratio_max_error':max(abs(q-2/np.pi) for q in ratios),
                                   'observed_loglog_slope':slope,
                                   'predicted_slope':1.0,
                                   'boundary_prediction_conservative': observed_boundary >= predicted_boundary,
                                   'all_checks': bool(bound_ok and scaling_ok and max_pair_excess <= 1e-8) },
                observed_gap_over_delta=ratios)

def rollout(policy, x0, horizon=40, dt=0.05):
    x=x0
    xs=[]; us=[]; rs=[]
    for _ in range(horizon):
        x.requires_grad_(True)
        u=policy(x)
        # h=x, grad h=1, so residual is u+alpha*x.
        r=u + ALPHA*x
        xs.append(x); us.append(u); rs.append(r)
        x=x + dt*u
    return torch.cat(xs), torch.cat(us), torch.cat(rs)

def train(use_barrier, steps=700):
    torch.manual_seed(SEED + (1 if use_barrier else 0))
    p=Policy()
    opt=torch.optim.Adam(p.parameters(), lr=3e-3)
    x0=torch.tensor([[0.20]])
    for _ in range(steps):
        xs,us,rs=rollout(p,x0)
        # Deliberately unsafe target demonstrates safety/task tradeoff.
        task=(xs[-1] - (-0.45)).pow(2) + 0.01*us.pow(2).mean()
        barrier=torch.relu(EPS-rs).pow(2).mean()
        loss=task + (1000.0*barrier if use_barrier else 0.0)
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(p.parameters(), 10.0); opt.step()
    with torch.no_grad():
        xs,us,rs=rollout(p,x0)
        return {'final_x':float(xs[-1]), 'min_x':float(xs.min()), 'min_r':float(rs.min()),
                'task_terminal_error':float((xs[-1]+0.45).abs()),
                'barrier_violation_fraction':float((rs<EPS).float().mean()),
                'dense_unseen_min_r':float((p(torch.linspace(0,1,2001).reshape(-1,1)).flatten()+ALPHA*torch.linspace(0,1,2001)).min()),
                'policy':p}

def main():
    cert=certification_sweep()
    base=train(False); idea=train(True)
    # remove model object from serialization
    base.pop('policy'); idea.pop('policy')
    out={'seed':SEED,'system':'xdot=u, h=x, alpha=1', 'epsilon':EPS,
         'certificate':cert, 'training':{'baseline_task_only':base,'barrier_certified':idea}}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__=='__main__': main()