"""Retry residual collision failures from the preceding valid pose."""
import json
import numpy as np
from l20_bottle_sequence import BASE,OUT
from l20_contact import ContactRefiner

src=np.load(BASE/'motion.npz');m=dict(np.load(OUT/'registered_motion.npz'))
s=ContactRefiner(json.loads((OUT/'object_spec.json').read_text()),OUT)
names=list(m['joint_names']);reports=[]
for t in np.where(m['penetration_m']>.0015)[0]:
    active=np.array([src['qpos'][t,names.index(n)] for n in s.cal.active])
    x,_,r=s.solve(active,src['wrist_pos'][t],src['wrist_quat_wxyz'][t],m['bottle_pos'][t],m['bottle_quat_wxyz'][t],initial=m['optimization_variables'][t-1])
    if r['max_penetration_after_mm']<m['penetration_m'][t]*1000:
        m['optimization_variables'][t]=x
        m['qpos'][t]=[s.expand(x[:16])[str(n)] for n in names]
        m['wrist_pos'][t]=s.data.qpos[s.wrist:s.wrist+3]
        m['wrist_quat_wxyz'][t]=s.data.qpos[s.wrist+3:s.wrist+7]
        m['penetration_m'][t]=r['max_penetration_after_mm']/1000
        m['command_u8'][t]=s.cal.command(dict(zip(names,m['qpos'][t])))
    reports.append(dict(frame=int(t),**r))
np.savez_compressed(OUT/'registered_motion.npz',**m)
header=(OUT/'registered_trajectory.csv').open().readline().strip()
np.savetxt(OUT/'registered_trajectory.csv',np.c_[np.arange(len(m['qpos']))/30,m['wrist_pos'],m['wrist_quat_wxyz'],m['qpos'],m['bottle_pos'],m['bottle_quat_wxyz']],delimiter=',',header=header,comments='')
r=json.loads((OUT/'sequence_validation.json').read_text());r['neighbor_retry_reports']=reports
r['max_all_frame_penetration_mm']=float(m['penetration_m'].max()*1000)
r['max_holding_penetration_mm']=float(m['penetration_m'][160:].max()*1000)
(OUT/'sequence_validation.json').write_text(json.dumps(r,indent=2))
print({'retries':len(reports),'max_penetration_mm':r['max_all_frame_penetration_mm'],'remaining_bad_frames':np.where(m['penetration_m']>.0015)[0].tolist()})