from __future__ import annotations
import sys, json, math, itertools
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import make_model, train_model, sweep_baseline, make_report
from bench.protocol import evaluate

META = {"name":"tangent_evasion_dynamics", "domain":"dynamics", "description":"2D evader policy with pursuer exclusion disk and tangent-branch safety structure."}
RADIUS = .8
CENTER = np.array([0., 0.], dtype=np.float32)
GOAL = np.array([3., 0.], dtype=np.float32)
DT = .08
VMAX = 1.0


def rot(theta):
    c,s=np.cos(theta),np.sin(theta)
    return np.array([[c,-s],[s,c]], dtype=np.float32)

def tangent(center, target, radius, side):
    z=center-target; d=np.linalg.norm(z)
    if d <= radius: return None
    a=np.arcsin(np.clip(radius/d,-1,1))
    return target + np.sqrt(d*d-radius*radius)/d * (rot(-side*a) @ z)

def label_velocity(x, y, heading):
    """Analytic branch reference, used only to generate task observations/labels."""
    p=np.array([x,y], dtype=np.float32); h=np.array([np.cos(heading),np.sin(heading)], dtype=np.float32)
    candidates=[]
    for side in (-1,1):
        s=tangent(CENTER, GOAL, RADIUS, side)
        if s is None: continue
        vec=s-p; ang=abs(np.arctan2(np.cross(h,vec),np.dot(h,vec)))
        candidates.append((ang+np.linalg.norm(GOAL-s),s))
    if np.linalg.norm(p-CENTER) < 1.35 and candidates:
        s=min(candidates,key=lambda z:z[0])[1]
        v=s-p
    else: v=GOAL-p
    n=np.linalg.norm(v)
    return (v/max(n,1e-6)*VMAX).astype(np.float32)

def get_dataset(seed, n_train, n_test):
    rng=np.random.RandomState(seed)
    def make(n):
        X=np.empty((n,24),np.float32); Y=np.empty((n,2),np.float32)
        for i in range(n):
            # Include safe and near-boundary states, with heading as the third GRU channel.
            a=rng.uniform(-np.pi,np.pi); rad=rng.uniform(.82,2.0)
            x,y=rad*np.cos(a),rad*np.sin(a)
            heading=rng.uniform(-np.pi,np.pi)
            state=np.tile(np.array([x,y,heading],np.float32),(8,1))
            state[:,2]+=rng.normal(0,.025,8).astype(np.float32)
            X[i]=state.reshape(-1)
            Y[i]=label_velocity(x,y,heading)
        return X,Y
    xtr,ytr=make(n_train); xte,yte=make(n_test)
    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,"task":"regression","metric":"mse","input_shape":(24,),"out_dim":2}


def project_np(pos, vel):
    p=np.asarray(pos,float); v=np.asarray(vel,float).copy(); d=np.linalg.norm(p-CENTER)
    if d>1e-9:
        n=(p-CENTER)/d
        if d <= RADIUS + DT*np.linalg.norm(v) and np.dot(v,n)<0: v-=np.dot(v,n)*n
    z=np.linalg.norm(v)
    return v if z<=VMAX else v/z*VMAX

def tangent_project_torch(x, out):
    # Network input is (x,y,heading) repeated; compute active-disk half-space projection.
    p=x[:, -1, :2]; v=out
    d=torch.linalg.vector_norm(p,dim=1,keepdim=True).clamp_min(1e-6)
    n=p/d
    inward=(v*n).sum(1,keepdim=True)
    active=(d <= RADIUS + DT*torch.linalg.vector_norm(v,dim=1,keepdim=True)).float()
    v=v-active*torch.minimum(inward,torch.zeros_like(inward))*n
    speed=torch.linalg.vector_norm(v,dim=1,keepdim=True).clamp_min(1e-6)
    return v*torch.minimum(torch.ones_like(speed), VMAX/speed)

def run(seed, lr, idea, epochs=15):
    torch.manual_seed(10000+int(seed)); np.random.seed(10000+int(seed))
    ds=get_dataset(int(seed),400,200)
    net=make_model("rnn_small", ds["input_shape"], ds["out_dim"])
    device="cuda" if torch.cuda.is_available() else "cpu"
    try:
        net=net.to(device); opt=torch.optim.Adam(net.parameters(),lr=lr)
        xtr=torch.as_tensor(ds["xtr"],device=device); ytr=torch.as_tensor(ds["ytr"],device=device)
        lossf=nn.MSELoss()
        for _ in range(epochs):
            net.train(); perm=torch.randperm(len(xtr),device=device)
            for j in range(0,len(xtr),128):
                ix=perm[j:j+128]; raw=net(xtr[ix])
                pred=tangent_project_torch(xtr[ix].view(-1,8,3),raw) if idea else raw
                loss=lossf(pred,ytr[ix]); opt.zero_grad(); loss.backward(); opt.step()
        net.eval(); xt=torch.as_tensor(ds["xte"],device=device)
        with torch.no_grad():
            raw=net(xt); pred=tangent_project_torch(xt.view(-1,8,3),raw) if idea else raw
            metric=float(lossf(pred,torch.as_tensor(ds["yte"],device=device)))
        return metric
    except Exception:
        # Explicit CPU fallback, matching the bench's robust-device requirement.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        torch.manual_seed(10000+int(seed)); net=make_model("rnn_small",ds["input_shape"],2)
        opt=torch.optim.Adam(net.parameters(),lr=lr); xtr=torch.as_tensor(ds["xtr"]); ytr=torch.as_tensor(ds["ytr"])
        for _ in range(epochs):
            for j in range(0,400,128):
                raw=net(xtr[j:j+128]); pred=tangent_project_torch(xtr[j:j+128].view(-1,8,3),raw) if idea else raw
                loss=nn.functional.mse_loss(pred,ytr[j:j+128]); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            raw=net(torch.as_tensor(ds["xte"])); pred=tangent_project_torch(torch.as_tensor(ds["xte"]).view(-1,8,3),raw) if idea else raw
            return float(nn.functional.mse_loss(pred,torch.as_tensor(ds["yte"])))

def main():
    # Equal union search space: every idea LR is also evaluated by baseline.
    grid=[{"lr":1e-3},{"lr":3e-3},{"lr":1e-2}]
    base=sweep_baseline(lambda c: lambda s: run(s,c["lr"],False),grid)
    # Idea is run at best baseline LR and two nearby grid settings; all are already baseline-swept.
    idea_cfgs=[{"lr":c["lr"]} for c in grid]
    vals=[]
    for c in idea_cfgs:
        r=evaluate(lambda s,c=c: run(s,c["lr"],True),tuple(range(8)))
        vals.append((r["mean"],c,r))
    _,best_cfg,idea=min(vals,key=lambda z:z[0])
    # Trained-model behavior signature: one-step disk penetration on held-out states.
    def penetration(seed, isidea):
        ds=get_dataset(seed,400,200); x=torch.as_tensor(ds["xte"]); torch.manual_seed(10000+seed); net=make_model("rnn_small",(24,),2); opt=torch.optim.Adam(net.parameters(),lr=best_cfg["lr"])
        for _ in range(15):
            for j in range(0,400,128):
                raw=net(torch.as_tensor(ds["xtr"])[j:j+128]); p=tangent_project_torch(torch.as_tensor(ds["xtr"])[j:j+128].view(-1,8,3),raw) if isidea else raw
                l=nn.functional.mse_loss(p,torch.as_tensor(ds["ytr"])[j:j+128]); opt.zero_grad();l.backward();opt.step()
        with torch.no_grad(): out=net(x); out=tangent_project_torch(x.view(-1,8,3),out) if isidea else out
        pos=x[:,-3:-1].numpy(); nxt=pos+DT*out.numpy(); return float(np.maximum(0,RADIUS-np.linalg.norm(nxt-CENTER,axis=1)).max())
    bp=float(np.mean([penetration(s,False) for s in range(8)])); ip=float(np.mean([penetration(s,True) for s in range(8)]))
    sig={"predicted":"projection keeps one-step disk penetration at numerical zero","observed_baseline_max_penetration":bp,"observed_idea_max_penetration":ip,"confirmed":ip < 1e-6 and ip <= bp+1e-8}
    rep=make_report("tangent_evasion_dynamics","rnn_small",base,idea,{**sig,"custom_track":{"name":"tangent_evasion_dynamics","file":"bench_tangent_stage2.py","domain":"dynamics"},"idea_sweep":[{"cfg":c,"mean":m} for m,c,_ in vals]})
    Path("bench_report.json").write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=="__main__": main()