# 2D 绕柱玩具实验：比较不同动作头在"左右(上下)双模态绕行"下的闭环表现
import torch, torch.nn as nn, numpy as np, math, time, sys, json
torch.manual_seed(0); np.random.seed(0); torch.set_num_threads(2)
R=1.0; H=16; STEP=0.05
def demo(s,g,mode):
    via=np.array([0.0, 1.35*mode])
    pts=[s, np.array([-0.9,0.95*mode]) if s[0]<-0.9 else s, via, np.array([0.9,0.95*mode]), g]
    P=np.array(pts); L=np.r_[0,np.cumsum(np.linalg.norm(np.diff(P,axis=0),axis=1))]
    n=int(L[-1]/STEP)+1; u=np.linspace(0,L[-1],n)
    x=np.interp(u,L,P[:,0]); y=np.interp(u,L,P[:,1]); Q=np.c_[x,y]
    # 平滑
    for _ in range(20): Q[1:-1]=0.5*Q[1:-1]+0.25*(Q[:-2]+Q[2:])
    return Q
def sample_task(hard=False):
    sy=np.random.uniform(-0.3,0.3) if hard else np.random.uniform(-1.5,1.5)
    gy=np.random.uniform(-0.3,0.3) if hard else np.random.uniform(-1.5,1.5)
    return np.array([np.random.uniform(-2.5,-1.6),sy]), np.array([np.random.uniform(1.6,2.5),gy])
def collide_path(Q):
    return np.any(np.linalg.norm(Q,axis=1)<R)
def build(n=3000):
    O=[];A=[]
    for _ in range(n):
        s,g=sample_task(hard=np.random.rand()<0.4)
        d=s[1]+g[1]
        mode = np.sign(d) if abs(d)>0.3 else (1 if np.random.rand()<0.5 else -1)
        Q=demo(s,g,mode)
        if collide_path(Q): continue
        D=np.diff(Q,axis=0); D=np.r_[D,np.zeros((H,2))]
        for i in range(len(Q)-1):
            O.append(np.r_[Q[i],g]); A.append(D[i:i+H].reshape(-1)/STEP)
    return torch.tensor(np.array(O),dtype=torch.float32), torch.tensor(np.array(A),dtype=torch.float32)
O,A=build(); print('samples',len(O),flush=True)
AD=2*H
def mlp(i,o,h=256): return nn.Sequential(nn.Linear(i,h),nn.SiLU(),nn.Linear(h,h),nn.SiLU(),nn.Linear(h,h),nn.SiLU(),nn.Linear(h,o))
class MSE(nn.Module):
    def __init__(s): super().__init__(); s.f=mlp(4,AD)
    def loss(s,o,a): return ((s.f(o)-a)**2).mean()
    def sample(s,o,k=1): return s.f(o).unsqueeze(1).expand(-1,k,-1)
class WTA(nn.Module):
    K=4
    def __init__(s): super().__init__(); s.f=mlp(4,AD*s.K+s.K)
    def out(s,o):
        y=s.f(o); return y[:,:AD*s.K].view(-1,s.K,AD), y[:,AD*s.K:]
    def loss(s,o,a):
        p,l=s.out(o); e=(p-a[:,None]).abs().mean(-1); k=e.argmin(1)
        return e.gather(1,k[:,None]).mean()+0.5*nn.functional.cross_entropy(l,k)
    def sample(s,o,k=1):
        p,l=s.out(o)
        if k==1: return p.gather(1,l.argmax(1)[:,None,None].expand(-1,1,AD))
        o_=l.argsort(1,descending=True); return p.gather(1,o_[:,:,None].expand(-1,-1,AD))  # 按概率排序
class FM(nn.Module):
    def __init__(s,steps): super().__init__(); s.f=mlp(4+AD+1,AD); s.steps=steps
    def loss(s,o,a):
        e=torch.randn_like(a); t=torch.rand(len(a),1); z=(1-t)*a+t*e
        return ((s.f(torch.cat([o,z,t],1))-(e-a))**2).mean()
    def sample(s,o,k=1):
        B=len(o); oo=o.repeat_interleave(k,0); z=torch.randn(B*k,AD); dt=1/s.steps
        for i in range(s.steps):
            t=torch.full((B*k,1),1-i*dt); z=z-dt*s.f(torch.cat([oo,z,t],1))
        return z.view(B,k,AD)
class MF(nn.Module):  # MeanFlow (Geng et al. 2025)
    def __init__(s): super().__init__(); s.f=mlp(4+AD+2,AD)
    def u(s,o,z,r,t): return s.f(torch.cat([o,z,r,t],1))
    def loss(s,o,a):
        B=len(a); e=torch.randn_like(a)
        t=torch.rand(B,1); r=torch.rand(B,1); t,r=torch.maximum(t,r),torch.minimum(t,r)
        eq=(torch.rand(B,1)<0.75); r=torch.where(eq,t,r)
        z=(1-t)*a+t*e; v=e-a
        fn=lambda z_,r_,t_: s.u(o,z_,r_,t_)
        u,dudt=torch.func.jvp(fn,(z,r,t),(v,torch.zeros_like(r),torch.ones_like(t)))
        tgt=(v-(t-r)*dudt).detach()
        return ((u-tgt)**2).mean()
    def sample(s,o,k=1):
        B=len(o); oo=o.repeat_interleave(k,0); e=torch.randn(B*k,AD)
        x=e-s.u(oo,e,torch.zeros(B*k,1),torch.ones(B*k,1)); return x.view(B,k,AD)
def train(m,iters=6000,bs=512):
    opt=torch.optim.AdamW(m.parameters(),1e-3,weight_decay=1e-4)
    sch=torch.optim.lr_scheduler.CosineAnnealingLR(opt,iters)
    for it in range(iters):
        idx=torch.randint(0,len(O),(bs,)); l=m.loss(O[idx],A[idx])
        opt.zero_grad(); l.backward(); opt.step(); sch.step()
    return m
def seg_collide(p,q):
    for a in np.linspace(0,1,6):
        if np.linalg.norm(p+(q-p)*a)<R: return True
    return False
@torch.no_grad()
def rollout(m,s,g,k=1,safety=False,commit=3,T=500):
    p=s.copy(); hits=0; ups=[]; brakes=0
    for t in range(0,T,commit):
        o=torch.tensor(np.r_[p,g],dtype=torch.float32)[None]
        C=m.sample(o,k)[0].numpy().reshape(-1,H,2)*STEP
        chosen=None
        for c in C:
            path=p+np.cumsum(c[:commit],0)
            if not safety: chosen=c; break
            pts=np.r_[p[None],path]
            if not any(seg_collide(pts[i],pts[i+1]) for i in range(commit)): chosen=c; break
        if chosen is None: brakes+=1; continue
        for d in chosen[:commit]:
            q=p+d
            if seg_collide(p,q): hits+=1
            p=q
            if abs(p[0])<0.5: ups.append(np.sign(p[1]))
            if np.linalg.norm(p-g)<0.1: return dict(ok=1,hits=hits,brakes=brakes,steps=t,mode=np.mean(ups) if ups else 0)
    return dict(ok=0,hits=hits,brakes=brakes,steps=T,mode=np.mean(ups) if ups else 0)
def evaluate(m,k,safety,n=150,hard=True):
    np.random.seed(123); res=[rollout(m,*sample_task(hard),k=k,safety=safety) for _ in range(n)]
    ok=np.mean([r['ok'] for r in res]); col=np.mean([r['hits']>0 for r in res])
    br=np.mean([r['brakes'] for r in res]); up=np.mean([r['mode']>0 for r in res if r['ok']]) if ok>0 else float('nan')
    return dict(reach=round(ok,3),collide_ep=round(col,3),brakes=round(br,1),frac_up=round(up,2))
methods={'MSE':lambda:MSE(),'WTA4':lambda:WTA(),'FM10':lambda:FM(10),'FM1':lambda:FM(1),'MeanFlow1':lambda:MF()}
which=sys.argv[1:] or list(methods)
out={}
for name in which:
    torch.manual_seed(0); t0=time.time(); m=train(methods[name]()); tt=time.time()-t0
    o=torch.tensor(np.r_[[-2,0],[2,0]],dtype=torch.float32)[None]
    t1=time.time(); [m.sample(o,8) for _ in range(50)]; lat=(time.time()-t1)/50*1000
    r={'train_s':round(tt),'infer_ms_k8':round(lat,2)}
    r['hard_raw_k1']=evaluate(m,1,False); r['hard_safe_k8']=evaluate(m,4 if name=='WTA4' else (1 if name=='MSE' else 8),True)
    r['easy_raw_k1']=evaluate(m,1,False,hard=False)
    out[name]=r; print(name,json.dumps(r,ensure_ascii=False),flush=True)
json.dump(out,open(f'/workspace/work/toy/res_{"_".join(which)}.json','w'),indent=1)