"""
TRI-CMU → Unified format converter (matches EgoVerse output schema).

Per episode, TRI-CMU provides:
  - video_1.mp4:         1920×1080 @ 30 fps egocentric video
  - hand_pose.json:      per-frame {hands: [LEFT, RIGHT] with 21 3D landmarks
                          in pelvis frame, hand_orientation quat,
                          camera_transform = camera pose in pelvis (pos+quat_xyzw)}
  - body_pose.json:      15 body landmarks in pelvis frame (idx 8, 12 wrists
                          are always null — supplied by hand_pose instead)
  - pelvis_pose.csv:     per-frame pelvis pose in world (pos + quat_xyzw)
  - intrinsics.json:     pinhole + radial/tangential distortion (k1,k2,p1,p2)
  - annotations.json:    dense temporal action segments {start_seconds,
                          end_seconds, label}; REQUIRED. Exported per-frame into
                          the lowdim `language.prompt` field (segments active at
                          frame_idx / FPS). An episode with no usable segments is
                          treated as a data bug and raises. The coarse
                          episode-level task_description is also kept in metadata.
  - metadata.json:       (written by download script) task_id, task_description,
                          environment_id, scene_id, user_id, duration

Unified output:
  rgb/stereo_left/NNNNNNNNNN.jpg       640×480 undistorted pinhole
  lowdim/stereo_left/NNNNNNNNNN.npz    intrinsics(3×3), extrinsics(4×4 = T_c2w),
                                       action.states = {
                                           hand_left_3d (21,3),
                                           hand_right_3d (21,3),
                                           wrist_left_3d (3,),     wrist position
                                           wrist_left_quat (4,),   wrist orientation (wxyz, camera frame)
                                           wrist_right_3d (3,),
                                           wrist_right_quat (4,),
                                           body_keypoints_cam (52,3) OpenCV cam frame, SMPL-H layout
                                           slam_pose (7,)
                                       }
                                       language = {prompt: [label, ...]}  (only on
                                           frames covered by an action segment)
  metadata.json (extrinsics declared as transform='cam2world')

  body_keypoints layout: SMPL-H 52-joint layout. TRI-CMU's 13-joint body
  skeleton maps to 12 SMPL-H body indices + 2 wrists sourced from hand[0]
  (= 14 of 52). Unpopulated SMPL-H slots (other spine joints, legs, feet,
  hand-finger joints) are filled with IGNORE_VAL (-1000). Hand fingers
  remain in the dedicated hand_left_3d / hand_right_3d (21,3) fields, same
  as EgoVerse / Xperience10M.

Conventions:
  - Pelvis frame is ROS-like (X fwd, Y left, Z up). Hand + body landmarks live there.
  - Camera frame is OpenCV (X right, Y down, Z fwd) — we apply M_PERM to convert.
  - extrinsics stored as T_camera_to_world (matches main's `transform='cam2world'`).
  - slam_pose is T_camera_to_world, format = [pos(3), quat_WXYZ(4)] (matches EgoVerse).
  - wrist_*_quat is in OpenCV-camera frame, wxyz order.
  - Missing/sentinel hands and unpopulated body slots are filled with IGNORE_VAL (-1000).

TODO when this branch is rebased onto main: port to the new converter API
(parse_sequence / process_sequence split, prepare_lowdim + write_lowdim +
write_labels helpers, --frames/--videos storage mode, args= kwarg in
fill_metadata, raise-instead-of-return error handling).

Usage:
    python -m anydata.converters.tri_cmu /data/cv_downloaded/TriCMU \
        --official --frames --num_procs 32
"""

import os
import json

import cv2
import numpy as np
from scipy.spatial.transform import Rotation

from anydata.converters.utils import (
    parse_dst_seq, frame_name, add_key_to_dict, fill_metadata,
    crawl, prepare_lowdim, run,
)
from anydata.utils.write import write_json, write_labels, write_lowdim
from anydata.converters.utils import extract_mp4

CAMERAS = ['stereo_left']
# Source video is 1920×1080 (16:9). Target preserves the same aspect ratio
# so undistort+resize stays a uniform scale (no squish, fx ≈ fy).
# Shortest dimension fixed at 512. Width is the closest multiple of 4 to
# 16/9 × 512 ≈ 910.2; this matches the framework's video/image writer
# constraint (dimensions trimmed to multiples of 4 for h264).
TARGET_W, TARGET_H = 908, 512
FPS = 30  # source video framerate; maps frame index → seconds for action segments
IGNORE_VAL = np.float32(-1000.0)
N_SMPLH = 52  # SMPL-H joint count (matches Xperience10M layout)

# Pelvis-ROS (X fwd, Y left, Z up) → OpenCV camera (X right, Y down, Z fwd)
M_PERM = np.array([[0, -1, 0],
                   [0,  0, -1],
                   [1,  0,  0]], dtype=np.float64)
M_PERM4 = np.eye(4)
M_PERM4[:3, :3] = M_PERM

# TRI-CMU body landmark idx → SMPL-H joint idx mapping (Xperience layout).
# TRI-CMU's Spine4 (idx 4) is dropped because SMPL-H only has 3 spine joints.
# Wrists (TRI-CMU 8, 12) are null in body data — populated from hand_*_3d[0].
SMPLH_FROM_TRICMU = {
    0:  0,    # Pelvis
    1:  3,    # Spine1
    2:  6,    # Spine2
    3:  9,    # Spine3
    # 4 (Spine4): no SMPL-H slot
    5:  13,   # L_Collar
    6:  16,   # L_Shoulder
    7:  18,   # L_Elbow
    # 8 (L_Wrist): filled from hand_left[0]
    9:  14,   # R_Collar
    10: 17,   # R_Shoulder
    11: 19,   # R_Elbow
    # 12 (R_Wrist): filled from hand_right[0]
    13: 12,   # Neck
    14: 15,   # Head
}
SMPLH_L_WRIST = 20
SMPLH_R_WRIST = 21


# ----------------------------- helpers -----------------------------

def make_T(pos, q_xyzw):
    """Build 4×4 SE(3) from (pos, quat_xyzw)."""
    T = np.eye(4)
    T[:3, :3] = Rotation.from_quat(q_xyzw).as_matrix()
    T[:3, 3] = pos
    return T


def quat_xyzw_to_wxyz(q):
    """scipy Rotation.as_quat() returns xyzw; EgoVerse expects wxyz."""
    return np.array([q[3], q[0], q[1], q[2]])


def read_pelvis_pose(path):
    """Parse pelvis_pose.csv. Returns list of (frame_no, pos[3], quat_xyzw[4])."""
    rows = []
    with open(path) as f:
        for line in f:
            if line.startswith('#') or line.startswith('frame_number'):
                continue
            parts = line.strip().split(',')
            if len(parts) < 9:
                continue
            try:
                frame_no = int(parts[0])
                pos = np.array([float(parts[2]), float(parts[3]), float(parts[4])])
                q = np.array([float(parts[5]), float(parts[6]),
                              float(parts[7]), float(parts[8])])
            except ValueError:
                continue
            rows.append((frame_no, pos, q))
    rows.sort(key=lambda r: r[0])
    return rows


def is_missing_hand(P_pelvis):
    """3D pelvis-frame landmarks are missing/degenerate if all zero.

    TRI-CMU's hand tracker always emits 21 3D landmarks (apparently temporally
    smoothed even through occlusion); the `landmarks_2d` field is a degenerate
    placeholder (all 21 entries share the same value) so it cannot serve as a
    per-landmark visibility flag. We therefore mirror EgoVerse's all-zero check.
    """
    return np.allclose(P_pelvis, 0, atol=1e-6)


def build_undistort_maps(intr, target_w, target_h):
    """Build a single remap that undistorts + resizes to target resolution.

    Uses cv2.getOptimalNewCameraMatrix(alpha=0) to pick a K such that the
    undistorted output has no black border (narrowed FOV instead). The K is
    then scaled from the valid ROI dimensions to (target_w, target_h).

    Returns (K_target, map1, map2) suitable for cv2.remap.
    """
    K_src = np.array([
        [intr['fl_x'], 0,           intr['cx']],
        [0,            intr['fl_y'], intr['cy']],
        [0,            0,            1.0],
    ])
    D = np.array([intr['k1'], intr['k2'], intr['p1'], intr['p2']])
    src_w, src_h = intr['w'], intr['h']

    K_src_undist, roi = cv2.getOptimalNewCameraMatrix(
        K_src, D, (src_w, src_h), alpha=0.0)
    rx, ry, rw, rh = roi
    if rw == 0 or rh == 0:
        K_src_undist = K_src.copy()
        rx, ry, rw, rh = 0, 0, src_w, src_h

    K_target = K_src_undist.copy()
    K_target[0, 2] -= rx
    K_target[1, 2] -= ry
    sx = target_w / rw
    sy = target_h / rh
    K_target[0, 0] *= sx
    K_target[0, 2] *= sx
    K_target[1, 1] *= sy
    K_target[1, 2] *= sy

    map1, map2 = cv2.initUndistortRectifyMap(
        K_src, D, None, K_target, (target_w, target_h), cv2.CV_16SC2)
    return K_target, map1, map2


def _build_body_smplh(body_landmarks, hand_left_pelvis, hand_right_pelvis,
                      hl_missing, hr_missing):
    """Pack TRI-CMU's 13-joint body skeleton into the 52-slot SMPL-H layout.

    Inputs in pelvis frame. Unpopulated slots filled with IGNORE_VAL.
    Wrists (idx 20, 21) come from hand_*_pelvis[0]; if the corresponding hand
    is missing/sentinel, those slots stay IGNORE.

    Returns (52, 3) array in pelvis frame.
    """
    out = np.full((N_SMPLH, 3), IGNORE_VAL, dtype=np.float32)
    for tri_idx, smpl_idx in SMPLH_FROM_TRICMU.items():
        lm = body_landmarks[tri_idx]
        if lm.get('x') is None:
            continue
        out[smpl_idx] = (lm['x'], lm['y'], lm['z'])
    if not hl_missing and hand_left_pelvis is not None:
        out[SMPLH_L_WRIST] = hand_left_pelvis[:, 0].astype(np.float32)
    if not hr_missing and hand_right_pelvis is not None:
        out[SMPLH_R_WRIST] = hand_right_pelvis[:, 0].astype(np.float32)
    return out


# ----------------------------- pipeline -----------------------------

def get_sequences(args):
    """An episode is a directory containing hand_pose.json."""
    if args.download:
        root = os.path.dirname(os.path.dirname(args.src))
        src = args.src.replace(root, args.s3_path)
        from anydata.converters.utils import list_s3_recursive
        seqs = list_s3_recursive(src, suffix='hand_pose.json')
        seqs = [f'{root}/{seq}' for seq in seqs]
    else:
        seqs = crawl(args.src, 'hand_pose.json')
    return [os.path.dirname(s) for s in seqs]


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args)


def process_sequence(i, seq, dst, args):
    return _process_inner(seq, args, dst)


def _process_inner(seq, args, dst):
    paths = {
        'video': f'{seq}/video_1.mp4',
        'hand':  f'{seq}/hand_pose.json',
        'body':  f'{seq}/body_pose.json',
        'pelvis': f'{seq}/pelvis_pose.csv',
        'intr':  f'{seq}/intrinsics.json',
        'meta':  f'{seq}/metadata.json',
        'annot': f'{seq}/annotations.json',
    }
    for key, p in paths.items():
        if not os.path.exists(p):
            raise ValueError(f'missing {key}: {p}')

    # Annotation tier filter: TRI-CMU annotations.json comes from two distinct
    # pipelines. Tier A (source='tierSnapshots.delivery.segments') is
    # consistently clean (~100% temporal coverage, 17 segments/ep avg).
    # Tier B (no `source` field) covers a different domain (mostly home
    # scenes) and ~5% of its episodes are severely under-annotated.
    # Per dataset policy we only convert Tier A episodes.
    ann = json.load(open(paths['annot']))
    if ann.get('source') != 'tierSnapshots.delivery.segments':
        raise ValueError(
            f'skip: annotation tier is not delivery '
            f'(source={ann.get("source")!r})')

    intr = json.load(open(paths['intr']))
    hp = json.load(open(paths['hand']))
    bp = json.load(open(paths['body']))
    pelvis = read_pelvis_pose(paths['pelvis'])
    meta = json.load(open(paths['meta']))

    # Dense temporal action segments: [(start_s, end_s, label), ...] from the
    # `ann` loaded above for the tier filter. Exported per-frame into the lowdim
    # `language` field (mirrors the epicfields / agibotworld convention) so each
    # frame carries the action label(s) active at its timestamp, in addition to
    # the coarse episode-level task_description prompt written into metadata
    # below. Tier-A episodes are required to carry usable segments; an empty set
    # is a data bug, so we raise rather than emit a sample without dense language.
    action_segments = [
        (s['start_seconds'], s['end_seconds'], s['label'])
        for s in ann.get('segments', [])
        if s.get('label')
    ]
    if not action_segments:
        raise ValueError(f'no language annotations (empty segments): {paths["annot"]}')

    # cap = cv2.VideoCapture(paths['video'])
    video = extract_mp4(paths['video'])

    n_video = int(video.shape[0])
    n_hand = len(hp['frames'])
    n_body = len(bp.get('poses', []))

    # Truncate at first invalid (zero-norm) pelvis quaternion
    pelvis_valid = []
    for r in pelvis:
        if np.linalg.norm(r[2]) > 1e-6:
            pelvis_valid.append(r)
        else:
            break
    pelvis = pelvis_valid
    n_pelvis = len(pelvis)

    n_frames = min(n_video, n_hand, n_pelvis, n_body)
    if n_frames == 0:
        raise ValueError('Zero usable frames')

    K_target, map1, map2 = build_undistort_maps(intr, TARGET_W, TARGET_H)
    K_target_out = K_target.astype(np.float64)  # match EgoVerse intrinsics dtype

    cam = 'stereo_left'
    num_frames = {cam: {}}
    resolution = {cam: {}}
    labels = []
    lowdim = {}
    dense = {'rgb': dict()}

    # Language (from manifest metadata.json)
    task_name = meta.get('task_id', os.path.basename(seq))
    task_description = meta.get('task_description', task_name)
    language = dict(task=task_name, prompt=[task_description])

    for frame_idx in range(n_frames):
        img = video[frame_idx]

        # RGB: undistort + resize in one remap. extract_mp4 already returns RGB
        # (pix_fmt=rgb24) and write_labels/write_image expect RGB, so no channel
        # flip here — flipping would store BGR and swap R/B in the output.
        img_warp = cv2.remap(img, map1, map2, cv2.INTER_LINEAR)
        frame = frame_name(frame_idx)
        dense['rgb'][frame] = img_warp

        # ---- Extrinsics chain ----
        cam_tr = hp['frames'][str(frame_idx)]['camera_transform']
        pos_c = np.array([cam_tr['position'][k] for k in 'xyz'])
        q_c = np.array([cam_tr['orientation'][k] for k in 'xyzw'])
        T_c2p = make_T(pos_c, q_c)            # camera-in-pelvis pose

        _, pos_p, q_p = pelvis[frame_idx]
        T_p2w = make_T(pos_p, q_p)            # pelvis-in-world pose
        T_c2w_ros = T_p2w @ T_c2p             # camera-in-world (ROS-axes camera)
        T_w2c_ros = np.linalg.inv(T_c2w_ros)
        T_w2c_cv = M_PERM4 @ T_w2c_ros        # world → OpenCV-axes camera
        T_c2w_cv = np.linalg.inv(T_w2c_cv)    # OpenCV-camera-in-world (stored)

        # slam_pose: T_c2w_cv as [pos(3) + quat_wxyz(4)] (matches EgoVerse format)
        slam_pos = T_c2w_cv[:3, 3]
        slam_q_wxyz = quat_xyzw_to_wxyz(
            Rotation.from_matrix(T_c2w_cv[:3, :3]).as_quat())
        slam_pose = np.concatenate([slam_pos, slam_q_wxyz]).astype(np.float32)

        # ---- Hand keypoints ----
        states = {}
        T_p2c = np.linalg.inv(T_c2p)
        R_p2c = T_p2c[:3, :3]
        t_p2c = T_p2c[:3, 3:4]
        hand_pelvis = {'LEFT': None, 'RIGHT': None}
        hand_missing = {'LEFT': True, 'RIGHT': True}

        for ht in ['LEFT', 'RIGHT']:
            ht_low = ht.lower()
            pos_key = f'wrist_{ht_low}_3d'
            quat_key = f'wrist_{ht_low}_quat'
            hand_key = f'hand_{ht_low}_3d'

            hand_record = next(
                (h for h in hp['frames'][str(frame_idx)]['hands']
                 if h['hand_type'] == ht), None)
            if hand_record is None:
                states[hand_key] = np.full((21, 3), IGNORE_VAL, np.float32)
                states[pos_key]  = np.full(3, IGNORE_VAL, np.float32)
                states[quat_key] = np.full(4, IGNORE_VAL, np.float32)
                continue

            P_pelvis = np.array(
                [(lm['x'], lm['y'], lm['z']) for lm in hand_record['landmarks']],
                dtype=np.float64).T  # 3 × 21
            if is_missing_hand(P_pelvis):
                states[hand_key] = np.full((21, 3), IGNORE_VAL, np.float32)
                states[pos_key]  = np.full(3, IGNORE_VAL, np.float32)
                states[quat_key] = np.full(4, IGNORE_VAL, np.float32)
                continue

            # Hand 21×3 in OpenCV-camera frame
            P_c_ros = R_p2c @ P_pelvis + t_p2c
            P_c_cv = (M_PERM @ P_c_ros).T  # 21 × 3
            states[hand_key] = P_c_cv.astype(np.float32)
            states[pos_key]  = P_c_cv[0].astype(np.float32)  # MANO wrist (idx 0)

            # Wrist orientation: hand_orientation is wrist-in-pelvis. Compose with
            # T_p2c, then permute to OpenCV-camera frame. Output as wxyz.
            ho = hand_record.get('hand_orientation')
            if (ho is not None and
                isinstance(ho, dict) and
                all(ho.get(k) is not None for k in 'xyzw')):
                q_xyzw_pelvis = np.array([ho['x'], ho['y'], ho['z'], ho['w']],
                                         dtype=np.float64)
                # SE3 for wrist-in-pelvis (position = MANO wrist landmark[0])
                T_wp = make_T(P_pelvis[:, 0], q_xyzw_pelvis)
                # → camera-ROS → camera-OpenCV
                T_wc_ros = T_p2c @ T_wp
                T_wc_cv = M_PERM4 @ T_wc_ros
                quat_xyzw_cam = Rotation.from_matrix(T_wc_cv[:3, :3]).as_quat()
                states[quat_key] = quat_xyzw_cam[[3, 0, 1, 2]].astype(np.float32)  # wxyz
            else:
                states[quat_key] = np.full(4, IGNORE_VAL, np.float32)

            hand_pelvis[ht] = P_pelvis
            hand_missing[ht] = False

        # ---- Body keypoints (SMPL-H layout, camera frame only) ----
        # We only emit body_keypoints_cam (camera-OpenCV frame). The world-frame
        # copy is recoverable from extrinsics + the camera-frame version if
        # ever needed, so storing both would just duplicate ~33 GB across the
        # full 5% sample after compression.
        body_landmarks = bp['poses'][frame_idx]['landmarks']
        body_smplh_pelvis = _build_body_smplh(
            body_landmarks,
            hand_pelvis['LEFT'], hand_pelvis['RIGHT'],
            hand_missing['LEFT'], hand_missing['RIGHT'])
        body_valid = ~np.isclose(body_smplh_pelvis, IGNORE_VAL).all(axis=1)

        bk_cam = np.full((N_SMPLH, 3), IGNORE_VAL, dtype=np.float32)
        if body_valid.any():
            P = body_smplh_pelvis[body_valid].T  # 3 × n_valid
            P_c_ros = R_p2c @ P + t_p2c
            P_c_cv = (M_PERM @ P_c_ros).T
            bk_cam[body_valid] = P_c_cv.astype(np.float32)

        states['body_keypoints_cam'] = bk_cam
        states['slam_pose'] = slam_pose

        filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
        prepare_lowdim(lowdim, dst, cam, frame)
        lowdim[filename_lowdim]['intrinsics'] = K_target_out
        lowdim[filename_lowdim]['extrinsics'] = T_c2w_cv.astype(np.float32)
        lowdim[filename_lowdim]['action'] = dict(states=states, actions={})

        # Dense per-frame language: labels of all action segments active at this
        # frame's timestamp (segments may overlap, so accumulate a list). Written
        # on *every* frame so videos-storage packing (stack_sample, which keys off
        # frame 0) sees a consistent field set across frames. TRI-CMU segments are
        # contiguous (gaps are labelled "no action"), so frame_labels is normally
        # non-empty for every frame.
        t_sec = frame_idx / FPS
        frame_labels = [lab for (s0, s1, lab) in action_segments
                        if s0 <= t_sec <= s1]
        lowdim[filename_lowdim]['language'] = dict(prompt=frame_labels)

    if not lowdim:
        raise ValueError('No frames produced')

    # Write RGB (and any future dense labels) + lowdim NPZs
    write_labels(dst, cam, args.storage, dense, labels, resolution, num_frames)
    write_lowdim(args, dst, labels, num_frames, lowdim)

    # Metadata
    tags = ['dynamic', 'egocentric', 'human', 'real', 'first_person']
    fps = FPS
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='TriCMU',
            tags=tags,
            raw_id=os.path.basename(seq),
        ),
        labels=labels,
        cameras=CAMERAS,
        resolution=resolution,
        num_frames=num_frames,
        framerate=int(fps),
        rgb=dict(extension='jpg'),
        intrinsics=dict(model='pinhole'),
        extrinsics=dict(transform='cam2world', metric=True),
        depth=None,
        semantic=None,
        action=dict(format='TriCMU'),
        language=language,
        specific=dict(
            domain='real',
            source='tricmu',
            user_id=meta.get('user_id'),
            environment_id=meta.get('environment_id'),
            scene_id=meta.get('scene_id'),
        ),
    )
    write_json(f'{dst}/metadata.json', seq_metadata)

    return dst


if __name__ == '__main__':
    converter = os.path.basename(__file__)
    run(converter, get_sequences, parse_sequence, process_sequence)