"""Check SPIDER imports, dependency versions and a two-world GPU contact simulation."""
import os
from pathlib import Path
ROOT=Path(__file__).resolve().parents[1]
os.environ['WARP_CACHE_PATH']=str(ROOT/'.spider_cache/warp')
os.environ['TORCHINDUCTOR_CACHE_DIR']=str(ROOT/'.spider_cache/torchinductor')
os.environ.setdefault('MUJOCO_GL','osmesa')
import json,sys,importlib,importlib.metadata as metadata,tomllib
import numpy as np
from packaging.requirements import Requirement
import torch,mujoco,mujoco_warp as mw,warp as wp

report={'python':sys.version,'compatibility_mode':'Python 3.11 with reused Torch; upstream declares Python >=3.12','versions':{},'imports':{},'requirements':{}}
reqs=tomllib.loads((ROOT/'third_party/spider/pyproject.toml').read_text())['project']['dependencies']
for text in reqs:
    r=Requirement(text);v=metadata.version(r.name);report['versions'][r.name]=v;report['requirements'][r.name]=v in r.specifier
for name in ['spider.config','spider.simulators.mjwp','spider.optimizers.sampling_fast','spider.viewers','spider.preprocess.ik','open3d','pymeshlab','mink']:
    importlib.import_module(name);report['imports'][name]=True
report['torch_path']=torch.__file__;report['cuda_available']=torch.cuda.is_available()
assert report['cuda_available'],'GPU test requires access outside the device-restricted sandbox'
report['gpu']=torch.cuda.get_device_name(0)
assert float((torch.ones(8,device='cuda')*2).sum())==16
wp.init()
xml='<mujoco><option timestep="0.002" iterations="5"/><worldbody><geom type="plane" size="1 1 .1"/><body pos="0 0 .105"><freejoint/><geom type="sphere" size=".1" mass=".1"/></body></worldbody></mujoco>'
m=mujoco.MjModel.from_xml_string(xml);d=mujoco.MjData(m);mujoco.mj_forward(m,d)
with wp.ScopedDevice('cuda:0'):
    wm=mw.put_model(m);wd=mw.put_data(m,d,nworld=2,nconmax=32,njmax=128)
    for _ in range(60):mw.step(wm,wd)
    wp.synchronize()
    q=wd.qpos.numpy();time=wd.time.numpy();contacts=wd.nacon.numpy()
assert np.isfinite(q).all() and np.all(q[:,2]>.085) and np.all(q[:,2]<.105),q
assert np.allclose(time,.12,atol=1e-5),time
report['gpu_physics']={'worlds':2,'steps':60,'time_s':time.tolist(),'sphere_height_m':q[:,2].tolist(),'contacts':contacts.tolist(),'passed':True}
assert all(report['requirements'].values())
out=ROOT/'output/spider_install';out.mkdir(exist_ok=True)
(out/'validation.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2),flush=True)