"""Stage-0 only comparison: for every test pair, build the swapped model with BOTH the
baseline runner and the PRIO runner, place the object upright, settle 60 steps and record
up_cos / bottom_err / verticality verdict. No policy, no rollout."""
import os, sys, json, math
os.environ.setdefault("MUJOCO_GL","egl"); os.environ.setdefault("PYOPENGL_PLATFORM","egl")
W="/lp-dev/jonghoon/actaug/robocasa_coacd"
sys.path.insert(0,W); sys.path.insert(0,"/lp-dev/jonghoon/actaug"); sys.path.insert(0,"/lp-dev/jonghoon/mimicgen_augment")
import numpy as np
import actaug_runner_coacd as B, actaug_runner_coacd_prio as P
from swap_and_render import swap_object_in_xml
from constrained import inject_weld
from actaug_core import load_episode, make_env, reset_episode

pairs=[]
for k in "ABCD": pairs += json.load(open(f"{W}/priotest/pairs_prio_{k}.json"))
env=make_env(load_episode(B.DS,int(pairs[0]["ep"]))["env_args"])
out=[]
for p in pairs:
    e,t,s0=int(p["ep"]),p["target"],int(p["seed0"])
    row={"ep":e,"target":t,"seed0":s0}
    ep=load_episode(B.DS,e); base=B.remap(ep["xml"])
    full=swap_object_in_xml(base,*t.split("/"))
    sjit=float(np.random.default_rng(s0).uniform(0.96,1.04))
    row["s_jit"]=round(sjit,4)
    for tag,R in (("baseline",B),("prio",P)):
        vis=R.swap_full_to_sam3d(R.remap(full),t,sjit)
        xml=inject_weld(R.remap(vis),"0.005 1"); ep_a=dict(ep,xml=xml)
        reset_episode(env,ep_a)
        ok,up,err=R.place_upright(env,e,quiet=True)
        row[tag]={"vertical":int(ok),"up_cos":round(float(up),4),"bottom_err_cm":round(float(err)*100,2)}
        print(f"  ep{e} {t} {tag}: vert={int(ok)} up_cos={up:.4f} bottom_err={err*100:.2f}cm",flush=True)
    out.append(row); json.dump(out,open(f"{W}/priotest/settle_compare.json","w"),indent=1)
print("SETTLE_DONE",flush=True)