"""
EgoVerse → Unified format converter.

Converts Zarr v3 episodes into the unified dataset format, matching
Xperience10M's structure for downstream compatibility.

Per episode, EgoVerse provides:
  - images.front_1: (N,) JPEG frames, 640×480
  - left/right.obs_keypoints: (N, 63) = 21 MANO joints × 3D positions (world frame)
  - left/right.obs_wrist_pose: (N, 7) = pos(3) + quat_wxyz(4)
  - left/right.obs_ee_pose: (N, 7) end-effector pose
  - obs_head_pose: (N, 7) = pos(3) + quat_wxyz(4), T_camera_to_world
  - obs_eye_gaze: (N, 3) gaze direction
  - obs_rgb_timestamps_ns: (N,) timestamps in nanoseconds

Unified output matches Xperience10M:
  - rgb/stereo_left/NNNNNNNNNN.jpg
  - lowdim/stereo_left/NNNNNNNNNN.npz with extrinsics, intrinsics, action.states

Usage:
    python -m anydata.converters.egoverse /data/cv_downloaded/EgoVerse --official --num_procs 32
"""

import os
import struct
import json
import sys

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

from anydata.utils.read import read_json
from anydata.utils.write import write_json, write_lowdim, write_labels
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, crawl, prepare_lowdim
from anydata.utils.geometry import invert_extrinsics

#######################################################
# Per-embodiment camera intrinsics and joint conventions
#######################################################

# Aria glasses (aria_bimanual, aria_left_arm, aria_right_arm)
_ARIA_FL_SCALE = 133.25430222  # fl = _ARIA_FL_SCALE * (height / 240)

# Mecka (MECKA_BIMANUAL) — from 1920×1080 calibration scaled to 640×360
_MECKA_FX = 250.82
_MECKA_FY = 251.00
_MECKA_CX = 320.61
_MECKA_CY = 184.42

# Scale AI — from egomimicUtils.py SCALE_INTRINSICS (640×480)
_SCALE_FX = 214.134
_SCALE_FY = 256.968
_SCALE_CX = 324.593
_SCALE_CY = 260.146


def _get_intrinsics(embodiment, h, w):
    """Return 3×3 intrinsics matrix for the given embodiment and resolution."""
    emb = embodiment.lower()
    if emb.startswith('mecka'):
        return np.array([
            [_MECKA_FX, 0, _MECKA_CX],
            [0, _MECKA_FY, _MECKA_CY],
            [0, 0, 1],
        ], dtype=np.float64)
    elif emb.startswith('scale'):
        return np.array([
            [_SCALE_FX, 0, _SCALE_CX],
            [0, _SCALE_FY, _SCALE_CY],
            [0, 0, 1],
        ], dtype=np.float64)
    else:  # Aria (default)
        fl = _ARIA_FL_SCALE * (h / 240.0)
        return np.array([
            [fl, 0, w / 2.0],
            [0, fl, h / 2.0],
            [0, 0, 1],
        ], dtype=np.float64)


def _transform_keypoints_to_cam(kp_world, T_w2c, embodiment):
    """Transform 21×3 keypoints from world to camera frame.

    All embodiments store obs_keypoints in world frame in the zarr.
    Apply inv(head_pose) to transform to camera frame.
    """
    kp_h = np.hstack([kp_world, np.ones((21, 1))])
    return (T_w2c @ kp_h.T).T[:, :3].astype(np.float32)

from anydata.converters.utils import (
    parse_dst_seq, frame_name, add_key_to_dict, fill_metadata, run,
)
from anydata.utils.write import write_npz, write_json

#######################################################

CAMERAS = ['stereo_left']  # match Xperience10M camera naming


def read_zarr_array(zarr_dir):
    """Read a Zarr v3 array from a local directory (non-image arrays only).

    Handles sharded+zstd+vlen encoding for numeric arrays stored as a single
    shard chunk. Returns numpy array.
    """
    import zarr
    arr = zarr.open(zarr_dir, mode='r')
    return np.asarray(arr[:])


class ZarrImageReader:
    """Lazy reader for Zarr v3 sharded variable_length_bytes image arrays.

    Handles multi-chunk shards (episodes split across c/0, c/1, c/2, ...).
    Each chunk stores a fixed number of sub-chunks with its own shard index.
    Reads chunk data on demand and decodes images one at a time.
    """

    def __init__(self, zarr_dir, n_frames):
        self.n_frames = n_frames
        self._dctx = zstandard.ZstdDecompressor()
        self._zarr_dir = zarr_dir

        # Discover all chunk files
        chunk_dir = os.path.join(zarr_dir, 'c')
        if not os.path.isdir(chunk_dir):
            self._chunks = []
            return

        chunk_ids = sorted(int(f) for f in os.listdir(chunk_dir)
                           if f.isdigit())
        if not chunk_ids:
            self._chunks = []
            return

        # Read zarr.json to get chunk_shape (sub-chunks per shard chunk)
        import json as _json
        with open(os.path.join(zarr_dir, 'zarr.json')) as f:
            meta = _json.load(f)
        # chunk_shape from the sharding codec configuration
        shard_cfg = meta['codecs'][0]['configuration']
        inner_chunk = shard_cfg['chunk_shape'][0]  # sub-chunks per inner chunk
        outer_chunk = meta['chunk_grid']['configuration']['chunk_shape'][0]
        subs_per_chunk = outer_chunk // inner_chunk  # sub-chunks per shard file

        self._chunks = []
        for cid in chunk_ids:
            path = os.path.join(chunk_dir, str(cid))
            if not os.path.exists(path):
                continue
            data = open(path, 'rb').read()
            # Frames in this chunk
            start_frame = cid * subs_per_chunk
            end_frame = min(start_frame + subs_per_chunk, n_frames)
            # Shard index always has subs_per_chunk entries (padded with 0xFF for unused slots)
            index_size = subs_per_chunk * 16 + 4
            index_data = data[-index_size:-4]
            self._chunks.append({
                'data': data,
                'index': index_data,
                'start': start_frame,
                'end': end_frame,
                'n_subs': subs_per_chunk,
            })

    def __len__(self):
        return self.n_frames

    def read_frame(self, i):
        """Decode and return a single BGR image, or None on failure."""
        if i >= self.n_frames:
            return None
        # Find which chunk contains frame i
        for chunk in self._chunks:
            if chunk['start'] <= i < chunk['end']:
                local_i = i - chunk['start']
                off = struct.unpack('<Q', chunk['index'][local_i * 16:local_i * 16 + 8])[0]
                nb = struct.unpack('<Q', chunk['index'][local_i * 16 + 8:local_i * 16 + 16])[0]
                if off == 0xFFFFFFFFFFFFFFFF or nb > len(chunk['data']):
                    return None
                try:
                    dec = self._dctx.decompress(chunk['data'][off:off + nb])
                    return cv2.imdecode(np.frombuffer(dec[8:], dtype=np.uint8),
                                        cv2.IMREAD_COLOR)
                except Exception:
                    return None
        return None

    def close(self):
        """Release all chunk data."""
        self._chunks = []


def head_pose_to_matrices(head_pose):
    """Convert head_pose [x,y,z,qw,qx,qy,qz] → (T_c2w, T_w2c) as 4×4 float32.

    EgoVerse stores head_pose as T_camera_to_world.
    Format verified against EgoVerse source (_xyzwxyz_to_matrix in pose_utils.py).
    """
    pos = head_pose[:3]
    # head_pose[3:7] = [qw, qx, qy, qz]; scipy expects [qx, qy, qz, qw]
    quat_xyzw = head_pose[[4, 5, 6, 3]]
    R_mat = Rotation.from_quat(quat_xyzw).as_matrix()
    T_c2w = np.eye(4)
    T_c2w[:3, :3] = R_mat
    T_c2w[:3, 3] = pos
    T_w2c = np.linalg.inv(T_c2w)
    return T_c2w.astype(np.float32), T_w2c.astype(np.float32)


#######################################################

def get_sequences(args):
    """Find all episodes (each has a zarr.json at the root)."""
    seqs = crawl(args.src, 'zarr.json')
    seqs = [os.path.dirname(s) for s in seqs]
    return seqs


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


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

    return _process_sequence_inner(seq, args, dst)


def _process_sequence_inner(seq, args, dst):

    ### Read zarr.json metadata
    zarr_meta_path = os.path.join(seq, 'zarr.json')
    if not os.path.exists(zarr_meta_path):
        raise ValueError(f'Path {zarr_meta_path} does not exist')

    import json as _json
    with open(zarr_meta_path) as _f:
        _ep_meta = _json.load(_f)
    attributes = _ep_meta['attributes']
    embodiment = attributes.get('embodiment', 'aria_bimanual')

    fps = attributes['fps']
    task_name = attributes['task_name']
    task_description = attributes['task_description']
    language = dict(task=task_name, prompt=[task_description])

    ### Skip EVA episodes (robot data, no hand keypoints)
    if embodiment.lower().startswith('eva'):
        raise ValueError(f'Skipping {embodiment} embodiment')

    ### Check required streams exist (hands are optional — use ignore values if absent)
    required = ['images.front_1', 'obs_head_pose']
    for r in required:
        if not os.path.isdir(os.path.join(seq, r)):
            raise ValueError(f'required {r} does not exist')

    ### Load numeric arrays
    head_pose = read_zarr_array(os.path.join(seq, 'obs_head_pose'))
    left_kp = read_zarr_array(os.path.join(seq, 'left.obs_keypoints')) \
        if os.path.isdir(os.path.join(seq, 'left.obs_keypoints')) else None
    right_kp = read_zarr_array(os.path.join(seq, 'right.obs_keypoints')) \
        if os.path.isdir(os.path.join(seq, 'right.obs_keypoints')) else None

    # Truncate at first invalid (zero-norm) quaternion — tail padding
    quat_norms = np.linalg.norm(head_pose[:, 3:7], axis=1)
    valid_mask = quat_norms > 1e-6
    if not valid_mask.all():
        first_invalid = np.where(~valid_mask)[0][0]
        head_pose = head_pose[:first_invalid]
        if left_kp is not None:
            left_kp = left_kp[:first_invalid]
        if right_kp is not None:
            right_kp = right_kp[:first_invalid]

    n_frames = head_pose.shape[0]
    if n_frames == 0:
        raise ValueError('Zero frames')

    # Optional streams
    left_wrist = read_zarr_array(os.path.join(seq, 'left.obs_wrist_pose')) \
        if os.path.isdir(os.path.join(seq, 'left.obs_wrist_pose')) else None
    right_wrist = read_zarr_array(os.path.join(seq, 'right.obs_wrist_pose')) \
        if os.path.isdir(os.path.join(seq, 'right.obs_wrist_pose')) else None
    eye_gaze = read_zarr_array(os.path.join(seq, 'obs_eye_gaze')) \
        if os.path.isdir(os.path.join(seq, 'obs_eye_gaze')) else None
    timestamps = read_zarr_array(os.path.join(seq, 'obs_rgb_timestamps_ns')) \
        if os.path.isdir(os.path.join(seq, 'obs_rgb_timestamps_ns')) else None

    ### Open lazy image reader (use original zarr length for shard index)
    import json as _json
    with open(os.path.join(seq, 'images.front_1', 'zarr.json')) as _f:
        _img_meta = _json.load(_f)
    n_images_total = _img_meta['shape'][0]
    img_reader = ZarrImageReader(os.path.join(seq, 'images.front_1'), n_images_total)
    test_img = img_reader.read_frame(0)
    if test_img is None:
        img_reader.close()
        raise ValueError('Image error')

    ### Initialize
    cam = 'stereo_left'
    num_frames = {cam: {}}
    resolution = {cam: {}}
    labels, lowdim = [], {}
    dense_labels = ['rgb']

    ### Image resolution & intrinsics (per-embodiment)
    h, w = test_img.shape[:2]
    del test_img
    K_matrix = _get_intrinsics(embodiment, h, w)

    # ### Compute FPS
    # if timestamps is not None and len(timestamps) > 1:
    #     dt = abs(timestamps[-1] - timestamps[0]) / 1e9
    #     fps = round(n_frames / dt) if dt > 0 else 30
    # else:
    #     fps = 30

    # ### Language
    # ep_name = os.path.basename(seq)
    # language = dict(task=ep_name, prompt=[ep_name])

    ### Process frames
    dense = {label: dict() for label in dense_labels}
    for frame_idx in range(n_frames):
        frame = frame_name(frame_idx)
        img = img_reader.read_frame(frame_idx)
        if img is None:
            continue

        ######## RGB
        dense['rgb'][frame] = img[..., ::-1]

        ######## LOWDIM
        filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
        prepare_lowdim(lowdim, dst, cam, frame)

        ######## INTRINSICS
        lowdim[filename_lowdim]['intrinsics'] = K_matrix

        ######## EXTRINSICS (T_world_to_camera, matches Xperience10M)
        T_c2w, T_w2c = head_pose_to_matrices(head_pose[frame_idx])
        lowdim[filename_lowdim]['extrinsics'] = T_c2w

        ######## STATES (match Xperience10M format)
        states = {}

        # Hand keypoints: reshape to (21, 3), transform world → camera.
        # Missing hand streams or all-zero frames get ignore value (-1000),
        # matching Xperience10M's convention for out-of-view hands.
        IGNORE_VAL = np.float32(-1000.0)

        for hand_key, kp_arr, wrist_arr, pos_key, quat_key in [
            ('hand_left_3d', left_kp, left_wrist, 'wrist_left_3d', 'wrist_left_quat'),
            ('hand_right_3d', right_kp, right_wrist, 'wrist_right_3d','wrist_right_quat'),
        ]:
            if kp_arr is not None and frame_idx < len(kp_arr):
                kp_world = kp_arr[frame_idx].reshape(21, 3)
                if np.allclose(kp_world, 0, atol=1e-6):
                    states[hand_key] = np.full((21, 3), IGNORE_VAL, dtype=np.float32)
                    states[pos_key] = np.full(3, IGNORE_VAL, dtype=np.float32)
                    states[quat_key] = np.full(4, IGNORE_VAL, dtype=np.float32)
                else:
                    states[hand_key] = _transform_keypoints_to_cam(
                        kp_world, T_w2c, embodiment)
                    if wrist_arr is not None and frame_idx < len(wrist_arr):
                        w7_world = wrist_arr[frame_idx]                # [x,y,z,qw,qx,qy,qz]
                        quat_xyzw_world = w7_world[[4, 5, 6, 3]]       # wxyz → xyzw
                        T_ww = np.eye(4)
                        T_ww[:3, :3] = Rotation.from_quat(quat_xyzw_world).as_matrix()
                        T_ww[:3, 3]  = w7_world[:3]
                        T_wc = T_w2c @ T_ww
                        quat_xyzw_cam = Rotation.from_matrix(T_wc[:3, :3]).as_quat()
                        states[pos_key]  = T_wc[:3, 3].astype(np.float32)
                        states[quat_key] = quat_xyzw_cam[[3, 0, 1, 2]].astype(np.float32)  # wxyz
                    else:
                        states[pos_key]  = np.full(3, IGNORE_VAL, dtype=np.float32)
                        states[quat_key] = np.full(4, IGNORE_VAL, dtype=np.float32)
            else:
                states[hand_key] = np.full((21, 3), IGNORE_VAL, dtype=np.float32)
                states[pos_key]  = np.full(3, IGNORE_VAL, dtype=np.float32)
                states[quat_key] = np.full(4, IGNORE_VAL, dtype=np.float32)

        # Head pose as compact 7-vector (matches Xperience10M slam_pose)
        states['slam_pose'] = head_pose[frame_idx].astype(np.float32)

        lowdim[filename_lowdim]['action'] = dict(states=states, actions={})

    ######## WRITE LABELS
    write_labels(dst, cam, args.storage, dense, labels, resolution, num_frames)

    ######## WRITE LOWDIM
    write_lowdim(args, dst, labels, num_frames, lowdim)

    ######## CLEANUP
    img_reader.close()

    ######## METADATA
    tags = ['dynamic', 'egocentric', 'human', 'real', 'first_person']

    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='EgoVerse',
            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(metric=True),
        depth=None,
        semantic=None,
        action=dict(format='EgoVerse'),
        language=language,
        specific=dict(
            domain='real',
            source='egoverse',
            embodiment=embodiment,
        ),
    )
    write_json(filename, seq_metadata)

    return dst

#######################################################

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

#######################################################