"""Assign the 3 random companions to Hunyuan3D-2mv slots (side->right, back->back, oside->left)
by relative azimuth, identical cost to fb150/build_fb150.py:assign_slots. Builds a symlinked
inputs dir with renamed tags; writes per-object mapping."""
import json, itertools, os, sys, numpy as np
R='/lp-dev/jonghoon/mv-mesh'
def horiz(v):
    v=np.array(v,float).copy(); v[2]=0; return v/(np.linalg.norm(v)+1e-12)
def feats_of(exp,o,front,comps):
    C=lambda t: np.load(f'{exp}/renders/{o}_{t}.npz')['c2w_cv']
    F=horiz(C(front)[:3,3]); out={}
    for c in comps:
        M=C(c); d=horiz(M[:3,3]); r=horiz(M[:3,0])
        out[c]=dict(cos_to_front=round(float(F@d),4),F_dot_right=round(float(F@r),4),
                    rel_azimuth_deg=round(float(np.degrees(np.arctan2(np.cross(F,d)[2],F@d))),2))
    return out
def assign(feats):
    ideal={"side":(0.0,1.0),"back":(-1.0,0.0),"oside":(0.0,-1.0)}; best=None
    for perm in itertools.permutations(list(feats)):
        cost=sum((feats[c]["cos_to_front"]-ideal[t][0])**2+(feats[c]["F_dot_right"]-ideal[t][1])**2 for t,c in zip(["side","back","oside"],perm))
        if best is None or cost<best[0]: best=(cost,perm)
    return dict(zip(["side","back","oside"],best[1])),best[0]
if __name__=='__main__':
    if sys.argv[1]=='sanity':
        # canonical exp: side must be the 'side' (right) slot: cos~0, F.right~+1
        exp=f'{R}/exp_faithfulness/toys4k300'; o='airplane_000'
        print(feats_of(exp,o,'front',['side'])); exp=f'{R}/exp_faithfulness/omni3d300'; o='asparagus_001'; print(feats_of(exp,o,'front',['side'])); sys.exit()
    ds=sys.argv[1]; exp=f'{R}/exp_faithfulness/{ds}_rand'; hy=f'{R}/exp_faithfulness/{ds}_rand_hyslots'
    os.makedirs(f'{hy}/inputs',exist_ok=True)
    sel=json.load(open(f'{exp}/selection.json')); mp={}
    for s in sel['selections']:
        o=s['object']; f=feats_of(exp,o,'front',['side','back','oside']); a,cost=assign(f)
        mp[o]=dict(slot_to_exp_tag={"front":"front",**a},hunyuan_key={"front":"front","right":a["side"],"back":a["back"],"left":a["oside"]},cost=round(cost,4),feats=f)
        for slot,src in {"front":"front",**a}.items():
            dst=f'{hy}/inputs/{o}_{slot}.png'
            if os.path.lexists(dst): os.remove(dst)
            os.symlink(f'{exp}/inputs/{o}_{src}.png',dst)
    os.makedirs(f'{hy}',exist_ok=True)
    if not os.path.lexists(f'{hy}/renders'): os.symlink(f'{exp}/renders',f'{hy}/renders')
    json.dump(sel,open(f'{hy}/selection.json','w'),indent=1)
    json.dump(mp,open(f'/home/nvidia/jonghoon/mv-mesh/.debug/eval4v_300/hy_slot_mapping_{ds}.json','w'),indent=1)
    perm=sum(1 for v in mp.values() if v['slot_to_exp_tag']!={"front":"front","side":"side","back":"back","oside":"oside"})
    print(ds,len(mp),'objects; permuted (non-identity) =',perm,'; mean cost',np.mean([v['cost'] for v in mp.values()]))