#!/usr/bin/env python3
"""Decisive offline check: feed the model the CORRECT-ORDER state (from the
lerobot parquet) + the demo images, and compare predicted action to the recorded
ground-truth action (also lerobot order). If pred ~= gt, the model + preprocessing
+ normalization are correct and the ONLY bug was joint ordering in our live scripts.
"""
import io, json, sys
import numpy as np
import go1_env  # 路径来自 go1_env(GO1_* 环境变量, 仓库相对默认)
import types as _t
if "decord" not in sys.modules:
    _f=_t.ModuleType("decord"); _f.VideoReader=type("VideoReader",(),{}); _f.cpu=lambda *a,**k:None
    _f.bridge=_t.ModuleType("decord.bridge"); _f.bridge.set_bridge=lambda *a,**k:None
    sys.modules["decord"]=_f; sys.modules["decord.bridge"]=_f.bridge
import torch
from PIL import Image

from go1_env import CKPT, STATS
PROMPT="What action should the robot take to reach into the bin, attach both suction cups to the part and lift it out level?"
CAM=["cam_head_color","cam_hand_l_color","cam_hand_r_color"]
DEV="cuda"
with open(STATS) as f: st=json.load(f)
sm=np.array(st["state"]["mean"],np.float32); ss=np.array(st["state"]["std"],np.float32)
am=np.array(st["action"]["mean"],np.float32); asd=np.array(st["action"]["std"],np.float32)

print("loading model...")
from go1.internvl.model.go1 import GO1Model, GO1ModelConfig
from go1.internvl.train.dataset import build_transform, dynamic_preprocess, preprocess_internvl2_5
from transformers import AutoTokenizer
cfg=GO1ModelConfig.from_pretrained(CKPT,torch_dtype=torch.bfloat16,low_cpu_mem_usage=True)
model=GO1Model.from_pretrained(CKPT,config=cfg).to(torch.bfloat16).to(DEV).eval()
tok=AutoTokenizer.from_pretrained(CKPT,trust_remote_code=True,use_fast=False,add_eos_token=False)
tf=build_transform(is_train=False,input_size=cfg.force_image_size,pad2square=cfg.pad2square)
nit=int((cfg.force_image_size//cfg.vision_config.patch_size)**2*(cfg.downsample_ratio**2))
torch.cuda.synchronize(); print("ready.\n")

def decode(a): return Image.open(io.BytesIO(bytes(a.tobytes()))).convert("RGB")

def predict(state, images):
    pp,nt=[],[]
    for im in images:
        tiles=[im] if not cfg.dynamic_image_size else dynamic_preprocess(im,min_num=cfg.min_dynamic_patch,
            max_num=cfg.max_dynamic_patch,image_size=cfg.force_image_size,use_thumbnail=cfg.use_thumbnail)
        pp+=tiles; nt.append(len(tiles))
    pv=torch.stack([tf(x) for x in pp]); npatch=pv.size(0)
    conv=[{"from":"human","value":("<image>"*len(images))+PROMPT},{"from":"gpt","value":""}]
    rp=preprocess_internvl2_5("internvl2_5",[conv],tok,[nit*n for n in nt],num_image=len(images),group_by_length=True)
    pos=rp["attention_mask"].long().cumsum(-1)-1; pos.masked_fill_(rp["attention_mask"]==0,1)
    sn=(state-sm)/(ss+1e-6)
    inp=dict(input_ids=rp["input_ids"][0].cuda().unsqueeze(0),attention_mask=rp["attention_mask"][0].cuda().unsqueeze(0),
        position_ids=pos[0].cuda().unsqueeze(0),pixel_values=pv.to(torch.bfloat16).cuda(),
        image_flags=torch.tensor([1]*npatch,dtype=torch.long).cuda(),
        state=torch.from_numpy(sn).to(torch.bfloat16).cuda().unsqueeze(0).unsqueeze(0),
        ctrl_freqs=torch.tensor([30.0],dtype=torch.bfloat16,device="cuda").unsqueeze(0))
    with torch.no_grad(): out=model(**inp)
    torch.cuda.synchronize()
    return out[1][0].float().cpu().numpy()*asd+am

d=np.load("/tmp/demo_frames2.npz")
LBL=["L1","L2","L3","L4","L5","L6","L7","R1","R2","R3","R4","R5","R6","R7","Lsuc","Rsuc","b1","b2","b3","b4","b5","hpit"]
print("%-6s %-10s %-11s %-10s %s"%("frame","|p0-gt0|","|p0-state|","chunkMAE","worst dim"))
for fr in d["frames"].tolist():
    state=d["state_%d"%fr].astype(np.float32)
    images=[decode(d["img_%d_%s"%(fr,k)]) for k in CAM]
    gt=d["act_%d"%fr].astype(np.float32)
    pred=predict(state,images)[:gt.shape[0]]
    d0=np.abs(pred[0]-gt[0]); d0s=np.abs(pred[0]-state)
    mae=np.abs(pred-gt).mean(); wj=int(d0.argmax())
    print("%-6d %-10.3f %-11.3f %-10.3f %s(%.2f)"%(fr,d0.max(),d0s.max(),mae,LBL[wj],d0[wj]))
    if fr==0:
        print("   pred[0]=",[round(x,2) for x in pred[0].tolist()])
        print("   gt[0]  =",[round(x,2) for x in gt[0].tolist()])
print("\n|p0-gt0| small => model CORRECT with right order. |p0-state| small => action[0]~current (no jump).")