"""Exercise compiled extensions, GPU rendering, official weights and red-box initialization."""
import sys,json
from pathlib import Path
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/'third_party/FoundationPose'))
import numpy as np
import torch,trimesh
import nvdiffrast.torch as dr
from estimater import FoundationPose,ScorePredictor,PoseRefinePredictor
import Utils
assert torch.cuda.is_available()
assert Utils.mycpp is not None and hasattr(Utils.mycpp,'cluster_poses')
ctx=dr.RasterizeCudaContext()
pos=torch.tensor([[[-.5,-.5,0,1],[.5,-.5,0,1],[0,.5,0,1]]],dtype=torch.float32,device='cuda')
tri=torch.tensor([[0,1,2]],dtype=torch.int32,device='cuda')
rast,_=dr.rasterize(ctx,pos,tri,resolution=[64,64])
assert torch.isfinite(rast).all() and (rast[...,3]>0).sum()>0
score=ScorePredictor();refine=PoseRefinePredictor()
outputs={}
with torch.inference_mode():
 for name,predictor in [('score',score),('refine',refine)]:
  h,w=predictor.cfg.input_resize
  a=torch.zeros((1,predictor.cfg.c_in,h,w),device='cuda')
  out=predictor.model(a,a,L=1) if name=='score' else predictor.model(a,a)
  assert all(torch.isfinite(v).all() for v in out.values())
  outputs[name]={k:list(v.shape) for k,v in out.items()}
mesh=trimesh.load(ROOT/'docs/上半.stl')
est=FoundationPose(model_pts=mesh.vertices,model_normals=mesh.vertex_normals,mesh=mesh,scorer=score,refiner=refine,glctx=ctx,debug=0,debug_dir=str(ROOT/'output/foundationpose_setup/smoke_debug'))
torch.cuda.synchronize()
report={'status':'passed','gpu':torch.cuda.get_device_name(),'torch':torch.__version__,'cuda':torch.version.cuda,'mycpp':Utils.mycpp.__file__,'rasterized_pixels':int((rast[...,3]>0).sum()),'network_outputs':outputs,'red_box_rotation_hypotheses':list(est.rot_grid.shape),'scope':'GPU rendering, network forward passes and STL estimator initialization; no recorded-video pose accuracy validation','optional_bundlesdf_mycuda':'not built; not required for model-based FoundationPose'}
(ROOT/'output/foundationpose_setup/runtime_verification.json').write_text(json.dumps(report,indent=2))
print(json.dumps(report,indent=2))