"""Independently validate saved video/robot outputs, including full video decode."""
import argparse
import json
from pathlib import Path
import xml.etree.ElementTree as ET

import cv2
import mujoco
import numpy as np

ROOT = Path(__file__).resolve().parents[1]


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--inference', type=Path, required=True)
    p.add_argument('--retarget', type=Path, required=True)
    a = p.parse_args()
    h = np.load(a.inference / 'handflow_results.npz')
    m = np.load(a.retarget / 'motion.npz')
    count, fps = len(h['pose']), float(h['fps'])
    assert len(m['qpos']) == count and float(m['fps']) == fps
    names = list(m['joint_names'])
    model = mujoco.MjModel.from_xml_path(str(a.retarget / 'l20_moving.xml'))
    data = mujoco.MjData(model)
    addr = [model.jnt_qposadr[model.joint(n).id] for n in names]
    wrist = model.jnt_qposadr[model.joint('wrist_free').id]
    sites = [model.site(f'landmark_{i:02d}').id for i in range(21)]
    limits = []
    mimic_error = 0.
    urdf = ET.parse(ROOT / 'third_party/l20_assets/L20/RIGHT/linkerhand_g20_right.urdf').getroot()
    for n in names:
        j = urdf.find(f"joint[@name='{n}']")
        limits.append([float(j.find('limit').get(k)) for k in ['lower','upper']])
        mi = j.find('mimic')
        if mi is not None:
            expected = m['qpos'][:,names.index(mi.get('joint'))] * float(mi.get('multiplier','1')) + float(mi.get('offset','0'))
            mimic_error = max(mimic_error, float(np.abs(expected-m['qpos'][:,names.index(n)]).max()))
    limits = np.array(limits)
    violation = float(max(np.maximum(limits[:,0]-m['qpos'],0).max(),np.maximum(m['qpos']-limits[:,1],0).max()))
    fk_error = 0.
    from scipy.spatial.transform import Rotation
    rotations = Rotation.from_quat(m['wrist_quat_wxyz'][:,[1,2,3,0]]).as_matrix()
    for t in range(count):
        data.qpos[addr] = m['qpos'][t]
        data.qpos[wrist:wrist+3] = m['wrist_pos'][t]
        data.qpos[wrist+3:wrist+7] = m['wrist_quat_wxyz'][t]
        mujoco.mj_forward(model,data)
        expected = m['actual'][t] @ rotations[t].T + m['wrist_pos'][t]
        fk_error = max(fk_error,float(np.abs(data.site_xpos[sites]-expected).max()))
    csv = np.loadtxt(a.retarget/'trajectory.csv',delimiter=',',skiprows=1)
    expected_csv = np.c_[m['time'],m['wrist_pos'],m['wrist_quat_wxyz'],m['qpos']]
    assert np.allclose(csv,expected_csv,atol=1e-12)
    assert np.isfinite(expected_csv).all() and violation < 1e-8 and mimic_error < 1e-10 and fk_error < 1e-5
    video_paths = [a.inference/'overlay.mp4',a.inference/'ortho_third_person.mp4',a.retarget/'l20_moving_replay.mp4']
    if (a.retarget/'comparison.mp4').exists():
        video_paths.append(a.retarget/'comparison.mp4')
    videos = []
    for path in video_paths:
        cap = cv2.VideoCapture(str(path))
        vf = cap.get(cv2.CAP_PROP_FPS)
        decoded = 0
        while True:
            ok, image = cap.read()
            if not ok:
                break
            assert image is not None
            decoded += 1
        cap.release()
        assert decoded == count and abs(vf-fps)<1e-3, (path,decoded,vf)
        videos.append(dict(path=str(path.resolve()),decoded_frames=decoded,fps=vf))
    report = dict(status='COMPLETE',frames=count,fps=fps,duration_seconds=count/fps,
                  finite=True,urdf_limit_violation_rad=violation,linear_mimic_error_rad=mimic_error,
                  independent_world_fk_max_error_m=fk_error,checked_frames=count,csv_matches_npz=True,
                  videos=videos,boundary='Kinematic hand motion only; monocular estimates; no object or hardware validation')
    (a.retarget/'delivery_validation.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))


if __name__ == '__main__':
    main()