"""Full-length contact registration; object poses are inferred targets, not physics."""
import json
from pathlib import Path
import numpy as np
from scipy.spatial.transform import Rotation,Slerp
from l20_contact import ContactRefiner
from l20_calibrated import ROOT

BASE=ROOT/'output/l20_2047635068';OUT=BASE/'bottle_grasp'
def main():
    motion=np.load(BASE/'motion.npz');spec=json.loads((OUT/'object_spec.json').read_text())
    solver=ContactRefiner(spec,OUT)
    count=len(motion['qpos']);times=np.arange(count)
    rotation=Rotation.from_quat(motion['wrist_quat_wxyz'][:,[1,2,3,0]])
    reference=666
    # Infer a constant bottle-to-palm registration over the holding phase.
    rel_pos=rotation[reference].inv().apply(np.array(spec['position'])-motion['wrist_pos'][reference])
    rel_rot=rotation[reference].inv()*Rotation.from_quat(np.array(spec['quat_wxyz'])[[1,2,3,0]])
    bottle_pos=rotation.apply(np.tile(rel_pos,(count,1)))+motion['wrist_pos']
    bottle_rot=rotation*rel_rot
    # Contact onset bracket comes from visual inspection: open at 100, grasped by 200.
    onset=160
    bottle_pos[:onset]=bottle_pos[onset]
    bottle_quat=bottle_rot.as_quat()[:,[3,0,1,2]]
    bottle_quat[:onset]=bottle_quat[onset]
    samples=np.unique(np.r_[np.arange(onset,count,8),count-1])
    fits=[];reports=[]
    for t in samples:
        active=np.array([motion['qpos'][t,list(motion['joint_names']).index(n)] for n in solver.cal.active])
        x,q,report=solver.solve(active,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
        fits.append(x);reports.append(report)
        if len(fits)%20==0:print(f'contact keyframes {len(fits)}/{len(samples)}',flush=True)
    fits=np.array(fits)
    interp=np.stack([np.interp(times,samples,fits[:,j]) for j in range(22)],axis=1)
    source_active=np.stack([motion['qpos'][:,list(motion['joint_names']).index(n)] for n in solver.cal.active],axis=1)
    # Blend acquisition only; preserve pre-contact source motion.
    weight=np.clip((times-(onset-30))/30,0,1);weight=weight*weight*(3-2*weight)
    interp[:,:16]=source_active*(1-weight[:,None])+interp[:,:16]*weight[:,None]
    interp[:,16:]*=weight[:,None]
    repairs=[]
    for t in range(onset,count):
        solver.set_pose(interp[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
        if -solver.distances().min()>.0015:
            x,_,r=solver.solve(source_active[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t],initial=interp[t])
            if r['max_penetration_after_mm']>1.5 and t>onset:
                alternative,_,retry=solver.solve(source_active[t],motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t],initial=interp[t-1])
                if retry['max_penetration_after_mm']<r['max_penetration_after_mm']:x,r=alternative,retry
            interp[t]=x;repairs.append(dict(frame=t,**r))
        if t%100==0:print(f'frame collision repair {t}/{count}: {len(repairs)} repaired',flush=True)
    qpos=[];wp=[];wq=[];penetration=[]
    for t,x in enumerate(interp):
        solver.set_pose(x,motion['wrist_pos'][t],motion['wrist_quat_wxyz'][t],bottle_pos[t],bottle_quat[t])
        values=solver.expand(x[:16]);qpos.append([values[str(n)] for n in motion['joint_names']])
        wp.append(solver.data.qpos[solver.wrist:solver.wrist+3].copy());wq.append(solver.data.qpos[solver.wrist+3:solver.wrist+7].copy())
        penetration.append(max(0,-solver.distances().min()))
    qpos=np.array(qpos);wp=np.array(wp);wq=np.array(wq)
    commands=np.array([solver.cal.command(dict(zip(motion['joint_names'],q))) for q in qpos])
    np.savez_compressed(OUT/'registered_motion.npz',qpos=qpos,joint_names=motion['joint_names'],wrist_pos=wp,wrist_quat_wxyz=wq,
           bottle_pos=bottle_pos,bottle_quat_wxyz=bottle_quat,command_u8=commands,fps=30,
           keyframes=samples,optimization_variables=interp,penetration_m=penetration,contact_onset_frame=onset)
    np.savetxt(OUT/'registered_trajectory.csv',np.c_[times/30,wp,wq,qpos,bottle_pos,bottle_quat],delimiter=',',
           header=','.join(['time_s','wrist_x','wrist_y','wrist_z','wrist_qw','wrist_qx','wrist_qy','wrist_qz']+list(motion['joint_names'])+['bottle_x','bottle_y','bottle_z','bottle_qw','bottle_qx','bottle_qy','bottle_qz']),comments='')
    report=dict(frames=count,fps=30,contact_keyframes=len(samples),converged_keyframes=sum(r['converged'] for r in reports),
      mean_keyframe_gap_before_mm=float(np.mean([r['mean_selected_surface_gap_before_mm'] for r in reports])),
      mean_keyframe_gap_after_mm=float(np.mean([r['mean_selected_surface_gap_after_mm'] for r in reports])),
      max_holding_penetration_mm=float(np.max(penetration[onset:])*1000),max_all_frame_penetration_mm=float(np.max(penetration)*1000),
      contact_onset_frame=onset,source='current 2047635068 video + supplied tri-camera bottle mesh',
      repaired_frames=len(repairs),repair_reports=repairs,
      alignment='Inferred fixed bottle-to-palm transform during holding, with contact pose corrections; no measured object trajectory or camera extrinsic alignment',
      boundary='Object poses in this file are prescribed registration targets. Dynamic grasp tests are separate and have not passed.',keyframe_reports=reports)
    (OUT/'sequence_validation.json').write_text(json.dumps(report,indent=2));print({k:v for k,v in report.items() if not k.endswith('_reports')},flush=True)

if __name__=='__main__':main()