"""Export the existing right-hand replay without changing its model or motion.

Uses the dataset layout of HDF5_REQUIREMENTS.md, with explicit right-hand
model/coordinate exceptions. This is not a left-hand l20_tracking_v1 PASS.
"""
import argparse
import hashlib
import json
from pathlib import Path
import xml.etree.ElementTree as ET

import h5py
import numpy as np

ROOT = Path(__file__).resolve().parents[1]
NAMES = [f'{f}_{j}' for f in ('index', 'middle', 'pinky', 'ring')
         for j in ('dip', 'mcp_pitch', 'mcp_roll', 'pip')] + [
             'thumb_cmc_pitch', 'thumb_cmc_roll', 'thumb_cmc_yaw', 'thumb_ip', 'thumb_mcp']


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--source', type=Path, default=ROOT/'output/l20_15886123_stable/motion.npz')
    parser.add_argument('--output', type=Path, default=ROOT/'output/hdf5_delivery/right')
    parser.add_argument('--scene', type=Path, help='MuJoCo scene for a full-replay qpos source')
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    a = np.load(args.source, allow_pickle=False)
    scene_description = None
    if args.scene:
        import mujoco
        model = mujoco.MjModel.from_xml_path(str(args.scene.resolve()))
        data = mujoco.MjData(model)
        rows = a['qpos']
        assert rows.ndim == 2 and rows.shape[1] == model.nq and np.isfinite(rows).all()
        addresses = [model.jnt_qposadr[model.joint(n).id] for n in NAMES]
        body_id = model.body('hand_base_link').id
        positions, rotations = [], []
        for row in rows:
            data.qpos[:] = row
            mujoco.mj_forward(model, data)
            positions.append(data.xpos[body_id].copy())
            rotations.append(data.xquat[body_id].copy())
        positions = np.asarray(positions)
        assert np.allclose(positions, a['wrist_pos'], atol=1e-9, rtol=0)
        a = dict(qpos=rows[:,addresses], joint_names=np.asarray(NAMES), wrist_pos=positions,
                 wrist_quat_wxyz=np.asarray(rotations), time=a['time'],
                 detection_valid=np.ones(len(rows), dtype=bool),
                 scene_rotation=np.eye(3), scene_translation=np.zeros(3))
        scene_description = (
            'Video 2047635068, full 1333-frame right-hand replay: Dyn-HaMR -> '
            'run_dex_l20.py -> temporal smoothing -> build_l20_full_replay.py contact registration. '
            'Actual source is the final replay qpos, not the earlier raw retargeting. '
            'hand_base_link world position and active local-to-world WXYZ rotation extracted '
            'with MuJoCo mj_forward for every frame; finger columns resolved by joint name. '
            'Source world is explicitly the FINAL replay scene world; world_from_source is identity '
            'because exported FK is already in that same world, not because camera calibration is known. '
            'Axes/origin are the existing fixed right-handed replay scene with display Z up: upstream '
            'scene was aligned using estimated bottle orientation at frame 600, shifted to the display '
            'placement, then raised for floor clearance; physical gravity alignment is unverified. '
            'Original human wrist/palm mapping, temporal smoothing and per-frame contact-registration '
            'corrections are inherited unchanged. Root is robot link origin, not its center of mass. '
            'valid=true means finite kinematic reference state satisfying the actual right URDF '
            'limits and mimic; original detection validity was not supplied, so this is NOT an '
            'observed-frame mask or measurement confidence. Motion and object registration are '
            'estimated, not ground truth; no full dynamic rollout or hardware execution claim. '
            'Source scene: ' + str(args.scene.resolve()))
    asset = ROOT/'third_party/l20_assets/L20/RIGHT'
    urdf = asset/'linkerhand_g20_right.urdf'
    tree = ET.parse(urdf).getroot()
    files = {urdf}
    for mesh in tree.iter('mesh'):
        files.add((asset/mesh.attrib['filename']).resolve())
    entries = [{'path': p.relative_to(asset.resolve()).as_posix(),
                'sha256': hashlib.sha256(p.read_bytes()).hexdigest()}
               for p in sorted(files, key=lambda p: p.relative_to(asset.resolve()).as_posix())]
    # Explicit local algorithm, not the unavailable receiver's asset hash tool.
    payload = ''.join(f"{e['path']}\t{e['sha256']}\n" for e in entries).encode('utf-8')
    digest = hashlib.sha256(payload).hexdigest()
    js = {j.attrib['name']: j for j in tree.findall('joint') if j.attrib['type'] != 'fixed'}
    manifest = dict(hand_side='right', root_link='hand_base_link', asset_sha256=digest,
                    hash_algorithm='sha256(UTF-8 concatenation of sorted relative_path TAB file_sha256 LF); URDF plus referenced meshes only; local export convention',
                    source_urdf=str(urdf), files=entries, joints=[])
    for name in NAMES:
        j = js[name]
        item = dict(name=name, lower_rad=float(j.find('limit').get('lower')),
                    upper_rad=float(j.find('limit').get('upper')))
        m = j.find('mimic')
        if m is not None:
            item['mimic'] = dict(joint=m.get('joint'), multiplier=float(m.get('multiplier', '1')),
                                 offset=float(m.get('offset', '0')))
        manifest['joints'].append(item)
    (args.output/'right_model_manifest.json').write_text(json.dumps(manifest, indent=2)+'\n')
    source_names = a['joint_names'].tolist()
    assert len(source_names) == len(set(source_names)) == len(NAMES)
    assert set(source_names) == set(NAMES)
    q = a['qpos'][:, [source_names.index(n) for n in NAMES]].astype(np.float32)
    pos = a['wrist_pos'].astype(np.float32)
    quat = a['wrist_quat_wxyz'].copy()
    quat /= np.linalg.norm(quat, axis=1, keepdims=True)
    for i in range(1, len(quat)):
        if quat[i] @ quat[i-1] < 0:
            quat[i] *= -1
    quat = quat.astype(np.float32)
    valid = a['detection_valid'].astype(bool)
    transform = np.eye(4, dtype=np.float64)
    transform[:3, :3] = a['scene_rotation']
    transform[:3, 3] = a['scene_translation']
    description = (
        'Video 15886123; HandFlow/manopth + ViPE camera-to-world; '
        'scripts/export_handflow_dex.py -> retarget_l20_video.py -> stabilize_l20_temporal.py. '
        'Right L20 URDF linear mimic, 16 independent and 21 total joints. '
        'hand_base_link origin assigned to estimated MANO wrist; orientation uses human palm_basis '
        'and robot neutral palm_basis, not a measured anatomical mount. '
        'Positions already use the stored fixed scene transform, followed by offline temporal smoothing. '
        'Scene origin and XYZ axes follow the original replay: first unsmoothed wrist at (0,0,0.2), '
        'first unsmoothed root orientation identity. Right-handed fixed scene with display Z up; '
        'physical gravity alignment is unknown. No new pose recentering applied by exporter. '
        'world_from_source maps the upstream estimated world to this existing display world; '
        'smoothing additionally changes individual poses. '
        'valid is upstream usable detection mask, not ground-truth tracking accuracy: '
        'first 24 frames are held/interpolated finite estimates and remain false. '
        'All frames have monocular scale/depth uncertainty, possible drift/occlusion errors, '
        'future-frame temporal smoothing; no contact, dynamics or hardware validation.')
    attrs = dict(schema_version='l20_tracking_v1', embodiment='L20', hand_side='right',
                 asset_sha256=digest, root_link='hand_base_link', provenance='expert_retargeted',
                 source_description=description,
                 metric_scale_provenance='Upstream MANO FK millimeters divided by 1000; camera translations used in upstream meter convention. Monocular/model-estimated metric scale, no measured reference length or gravity calibration; absolute scale error unknown. Input motion.npz positions already in estimated meters, so export scale_to_meters=1.0.',
                 scale_to_meters=1.0,
                 contract_status='RIGHT_HAND_VARIANT_NOT_STRICT_LEFT_V1; actual URDF mimic; local asset hash algorithm; display-world gravity unverified',
                 source_sha256=hashlib.sha256(args.source.read_bytes()).hexdigest(),
                 asset_hash_algorithm=manifest['hash_algorithm'])
    if scene_description:
        attrs['source_description'] = scene_description
        attrs['metric_scale_provenance'] = (
            'Final replay positions already in estimated meters, inherited from Dyn-HaMR MANO '
            'world reconstruction and robot geometry; export scale_to_meters=1.0. '
            'No measured absolute scale reference is established; error unknown. '
            'Scene alignment and contact fitting are engineering estimates, not scale measurements.')
        attrs['valid_semantics'] = 'kinematic_reference_valid; source_detection_validity_unavailable'
        attrs['source_scene_sha256'] = hashlib.sha256(args.scene.read_bytes()).hexdigest()
    good = np.flatnonzero(valid)
    runs = [r for r in np.split(good, np.flatnonzero(np.diff(good) > 1)+1) if len(r) >= 2]
    assert runs
    run = max(runs, key=len)
    sample = slice(int(run[0]), int(min(run[-1]+1, run[0]+90)))
    reports = []
    for filename, selection in [('sample_right.hdf5', sample), ('demonstrations_right.hdf5', slice(0, len(q)))]:
        path = args.output/filename
        times = a['time'][selection].astype(np.float64)
        times -= times[0]
        with h5py.File(path, 'w') as f:
            for k, v in attrs.items():
                f.attrs[k] = v
            meta = f.create_group('metadata')
            meta.create_dataset('joint_names', data=NAMES, dtype=h5py.string_dtype('utf-8'))
            meta.create_dataset('world_from_source', data=transform)
            ep = f.create_group('episodes/demo_000000')
            ep.attrs['source_frame_start'] = selection.start
            ep.attrs['source_frame_end_exclusive'] = selection.stop
            for name, value in dict(time=times, wrist_position=pos[selection],
                                    wrist_quaternion=quat[selection], joint_position=q[selection],
                                    valid=valid[selection]).items():
                ep.create_dataset(name, data=value)
        # Reopen the actual serialized file and validate against the actual right URDF.
        with h5py.File(path, 'r') as f:
            ep = f['episodes/demo_000000']
            t, p, r, angles, mask = [ep[n][:] for n in
                                    ('time','wrist_position','wrist_quaternion','joint_position','valid')]
            count = len(t)
            assert [(x.shape, x.dtype) for x in (t,p,r,angles,mask)] == [
                ((count,),np.dtype('float64')),((count,3),np.dtype('float32')),
                ((count,4),np.dtype('float32')),((count,21),np.dtype('float32')),((count,),np.dtype('bool'))]
            assert count >= 2 and t[0] == 0 and np.all(np.diff(t)>0) and mask.any()
            assert all(np.isfinite(x).all() for x in (t,p,r,angles))
            assert f['metadata/joint_names'].asstr()[:].tolist() == NAMES
            assert h5py.check_string_dtype(f['metadata/joint_names'].dtype).encoding == 'utf-8'
            assert np.max(np.abs(np.linalg.norm(r[mask],axis=1)-1)) < 1e-4
            assert np.all(np.sum(r[1:]*r[:-1],axis=1)[mask[1:] & mask[:-1]] >= 0)
            assert np.allclose(transform[3], [0,0,0,1])
            assert np.allclose(transform[:3,:3].T@transform[:3,:3], np.eye(3))
            assert abs(np.linalg.det(transform[:3,:3])-1)<1e-8
            residuals = {}
            for i, joint in enumerate(manifest['joints']):
                assert angles[mask,i].min() >= joint['lower_rad']-1e-6
                assert angles[mask,i].max() <= joint['upper_rad']+1e-6
                if 'mimic' in joint:
                    m = joint['mimic']
                    err = np.abs(angles[mask,i]-m['multiplier']*angles[mask,NAMES.index(m['joint'])]-m['offset'])
                    residuals[joint['name']] = float(err.max())
                    assert err.max() <= 1e-3
            reports.append(dict(file=filename, right_model_validation='PASS', strict_left_v1='NOT_COMPATIBLE',
                                frames=count, valid_frames=int(mask.sum()), source_frame_start=selection.start,
                                source_frame_end_exclusive=selection.stop, time_last_s=float(t[-1]),
                                mimic_residual_rad=residuals, sha256=hashlib.sha256(path.read_bytes()).hexdigest()))
    (args.output/'validation.json').write_text(json.dumps(reports,indent=2)+'\n')
    print(json.dumps(reports,indent=2))


if __name__ == '__main__':
    main()