@daft.func(return_dtype=EPISODE_DTYPE)
def process_egodex_episode(file_: daft.File) -> dict:
    h = file_.as_hdf5()
    transforms = h.read(list(dict.fromkeys(STATE_TRANSFORMS + SKELETON_TRANSFORMS + [CAMERA])))
 
    state = build_state(transforms)
    skeleton = build_skeleton(transforms)
    extrinsics = build_extrinsics(transforms)
    action = next_frame_action(state)
    task = resolve_task(h.attrs())  # native attrs() returns a dict
 
    frames = [
        {
            "frame_index": i,
            "observation.state": state[i],
            "observation.skeleton": skeleton[i],
            "observation.extrinsics": extrinsics[i],
            "action": action[i],
        }
        for i in range(len(state))
    ]
    return {"task": task, "frames": frames}

def read_egodex(hdf5_glob, with_video: bool = False):
    per_file = Window().order_by(col("file").file_path())
    episodes = (
        daft.from_files(hdf5_glob)
        .sort(col("file").file_path())
        .with_column("episode_index", row_number().over(per_file) - 1)
        # carry the HDF5 path so the video decoder can find each episode's sibling .mp4
        .with_column("_src", col("file").file_path())
        .into_batches(8)
        .with_column("_ep", process_egodex_episode(col("file")))
        .with_column("task", col("_ep")["task"])
        .with_column("frames", col("_ep")["frames"])
    )
 
    frames = (
        episodes.explode("frames")
        .select("episode_index", "task", "_src", col("frames").unnest())
        .with_column("timestamp", (col("frame_index") / FPS).cast(DataType.float32()))
    )
 
    if with_video:
        frames = frames.with_column("observation.image", _decode_sibling_mp4(col("_src"), col("timestamp")))
    return frames.exclude("_src")

@daft.func(return_dtype=DataType.float64())
def forearm_roll(rot6d, rot6d_next, forearm_axis):
    """Wrist roll (rad) about the forearm axis from one frame to the next (0 at an episode's last frame)."""
    if rot6d is None or rot6d_next is None:
        return 0.0
    delta = _rotation_matrix(rot6d_next) @ _rotation_matrix(rot6d).T
    angle = np.arccos(np.clip((np.trace(delta) - 1) / 2, -1, 1))
    axis = np.array([delta[2, 1] - delta[1, 2], delta[0, 2] - delta[2, 0], delta[1, 0] - delta[0, 1]])
    magnitude = np.linalg.norm(axis)
    if magnitude < 1e-9:
        return 0.0
    return float(abs(angle * np.dot(axis / magnitude, np.asarray(forearm_axis))))

def rotation_from_rot6d(rot6d):
    """rot6d (N, 6) -> (N, 3, 3) rotation matrices (columns = hand x, y axes + palm normal)."""
    first_column = rot6d[:, 0:3]
    first_column = first_column / (np.linalg.norm(first_column, axis=1, keepdims=True) + 1e-9)
    second_column = rot6d[:, 3:6]
    second_column = second_column - (first_column * second_column).sum(1, keepdims=True) * first_column
    second_column = second_column / (np.linalg.norm(second_column, axis=1, keepdims=True) + 1e-9)
    palm_normal = np.cross(first_column, second_column)
    return np.stack([first_column, second_column, palm_normal], axis=2)

def compute_raw_features(state: np.ndarray) -> dict:
    """Per-frame raw features for both hands, keyed by feature + hand tag ('L' / 'R')."""
    features = {}
    for side, tag in HAND_TAGS:
        wrist, rot6d, fingertips = _split_hand(state, HAND_BASE[side])
        tip_to_wrist = np.linalg.norm(fingertips - wrist[:, None, :], axis=2)  # (N, 5)
        features[f"fingerdist_{tag}"] = tip_to_wrist
        features[f"curl_{tag}"] = tip_to_wrist.mean(1)  # small = curled
        palm_normal = palm_normal_from_rot6d(rot6d)
        features[f"palmnormal_{tag}"] = palm_normal
        features[f"palm_up_{tag}"] = palm_normal[:, 1]  # +y component
        features[f"pinch_{tag}"] = np.linalg.norm(fingertips[:, 0] - fingertips[:, 1], axis=1)  # thumb-index
        tip_pairs = fingertips[:, :, None, :] - fingertips[:, None, :, :]
        features[f"aperture_{tag}"] = np.linalg.norm(tip_pairs, axis=3).max((1, 2))  # max tip-tip spread
        features[f"wrist_{tag}"] = wrist
    return features

def angle_between(first_bone, second_bone):
    """Row-wise angle (rad, [0, pi]) between two vectors. atan2 form is stable at 0/pi."""
    cross_magnitude = np.linalg.norm(np.cross(first_bone, second_bone), axis=-1)
    dot_product = (first_bone * second_bone).sum(-1)
    return np.arctan2(cross_magnitude, dot_product)

def finger_flexion(skeleton, side, finger):
    """(N, n_joints) per-joint flexion angles along a finger chain.
 
    Flexion = the angle between consecutive bone vectors (0 = straight, larger =
    more curled). Non-thumb fingers yield MCP / PIP / DIP (3); the thumb yields 2.
    """
    positions = [joint_position(skeleton, name) for name in finger_joint_names(side, finger)]
    bones = [positions[i + 1] - positions[i] for i in range(len(positions) - 1)]
    joint_angles = [angle_between(bones[i], bones[i + 1]) for i in range(len(bones) - 1)]
    return np.stack(joint_angles, axis=1)

def palm_normal(skeleton, side):
    """(N, 3) unit palm-plane normal via a best-fit plane through the knuckles."""
    wrist = joint_position(skeleton, side + "Hand")
    knuckles = [joint_position(skeleton, f"{side}{f}FingerKnuckle") for f in ("Index", "Middle", "Ring", "Little")]
    points = np.stack([wrist] + knuckles, axis=1)
    centered = points - points.mean(1, keepdims=True)
    covariance = np.einsum("nki,nkj->nij", centered, centered)
    _, eigenvectors = np.linalg.eigh(covariance)  # ascending; [:, :, 0] = smallest
    normal = eigenvectors[:, :, 0]
    # orient deterministically by handedness, then flip to the palm-facing convention
    index_knuckle = joint_position(skeleton, f"{side}IndexFingerKnuckle")
    little_knuckle = joint_position(skeleton, f"{side}LittleFingerKnuckle")
    reference = np.cross(index_knuckle - wrist, little_knuckle - wrist)
    orientation = np.sign((normal * reference).sum(1))
    orientation[orientation == 0] = 1.0
    return normal * orientation[:, None] * PALM_SIGN[side]

def arm_extension(skeleton, side):
    """(N,) wrist-to-shoulder distance / total arm length (0 = tucked, ~1 = straight)."""
    shoulder = joint_position(skeleton, side + "Shoulder")
    upper_arm = joint_position(skeleton, side + "Arm")
    forearm = joint_position(skeleton, side + "Forearm")
    wrist = joint_position(skeleton, side + "Hand")
    arm_length = (
        np.linalg.norm(upper_arm - shoulder, axis=1)
        + np.linalg.norm(forearm - upper_arm, axis=1)
        + np.linalg.norm(wrist - forearm, axis=1)
    )
    return np.linalg.norm(wrist - shoulder, axis=1) / (arm_length + 1e-9)

def hand_local_joints(skeleton, side):
    """(N, K, 3) the side's finger joints expressed in the hand's own frame.
 
    The frame is z = palm normal, x = across-palm, origin = wrist; expressing the
    joints in it removes gross hand translation/rotation, leaving only the
    fingers' own motion (used to detect in-hand manipulation).
    """
    wrist = joint_position(skeleton, side + "Hand")
    across_palm = joint_position(skeleton, f"{side}LittleFingerKnuckle") - joint_position(
        skeleton, f"{side}IndexFingerKnuckle"
    )
    z_axis = _unit(palm_normal(skeleton, side))
    x_axis = _unit(across_palm - (across_palm * z_axis).sum(1, keepdims=True) * z_axis)
    y_axis = np.cross(z_axis, x_axis)
    frame = np.stack([x_axis, y_axis, z_axis], axis=1)  # (N, 3, 3), rows = axes
    joint_names = [name for finger in FINGERS for name in finger_joint_names(side, finger)]
    positions = np.stack([joint_position(skeleton, name) for name in joint_names], axis=1)  # (N, K, 3)
    return np.einsum("nij,nkj->nki", frame, positions - wrist[:, None, :])

def add_angular_velocity(features, state, episode_index, frame_index, fps):
    """Add wrist angular speed (rad/s) from consecutive orientations, for the overlay table."""
    for side, tag in HAND_TAGS:
        rotations = rotation_from_rot6d(state[:, rot6d_slice(side)])
        angular_speed = np.zeros(len(rotations))
        for episode in np.unique(episode_index):
            rows = np.where(episode_index == episode)[0]
            rows = rows[np.argsort(frame_index[rows])]
            if len(rows) < 2:
                continue
            relative = np.einsum("nij,nkj->nik", rotations[rows][1:], rotations[rows][:-1])  # R_t @ R_{t-1}^T
            cosine = np.clip((np.trace(relative, axis1=1, axis2=2) - 1) / 2, -1, 1)
            angular_speed[rows[1:]] = np.arccos(cosine) * fps
        features[f"wrist_angvel_{tag}"] = angular_speed
    return features

# daft expression
def openness(hand, thr, open_lo=0.0, open_hi=1.0):
    """Openness band on `hand` (1 = fully open palm), mapped onto the closure column via the
    calibrated closure spread [closure_lo, closure_hi]: higher openness -> lower closure."""
    span = (thr["closure_hi"] - thr["closure_lo"]) or 1.0
    clo_lo = thr["closure_hi"] - open_hi * span
    clo_hi = thr["closure_hi"] - open_lo * span
    return (col(f"closure_{hand}") >= clo_lo) & (col(f"closure_{hand}") <= clo_hi)

 
# daft expression
def writing_grip(hand, thr):
    """Tripod/precision grip on `hand` from flexion + thumb-tip distance vs calibrated thresholds."""
    return _is_writing_grip(
        col(f"flex_nonthumb_{hand}"),
        col(f"thumb_min_tip_{hand}"),
        thr["curled_flexion"],
        thr["curl_gap"],
        thr["thumb_on_tip"],
    )
 
@daft.func(return_dtype=DataType.bool())
def _is_writing_grip(flex_nonthumb, thumb_min_tip, curled_flexion, curl_gap, thumb_on_tip):
    """Tripod: thumb on the index/middle tip, those two not fisted, ring+little more curled."""
    flex = np.asarray(flex_nonthumb)
    return bool(
        thumb_min_tip < thumb_on_tip
        and flex[0] < curled_flexion
        and flex[1] < curled_flexion
        and flex[2] > flex[0] + curl_gap
        and flex[3] > flex[1] + curl_gap
    )
 

 
# daft expression
def hammer_grip(hand, thr):
    """Power grip on `hand` from flexion + thumb-knuckle distance vs calibrated thresholds."""
    return _is_hammer_grip(
        col(f"flex_nonthumb_{hand}"), col(f"thumb_min_knuckle_{hand}"), thr["curled_flexion"], thr["thumb_on_knuckle"]
    )
 
# daft udf
@daft.func(return_dtype=DataType.bool())
def _is_hammer_grip(flex_nonthumb, thumb_min_knuckle, curled_flexion, thumb_on_knuckle):
    """Power: all four fingers curled and the thumb wrapped across a proximal knuckle."""
    flex = np.asarray(flex_nonthumb)
    return bool(bool((flex > curled_flexion).all()) and thumb_min_knuckle < thumb_on_knuckle)
 
 

 
def add_skeleton_features(df, fps=DEFAULT_FPS):
    """Per-episode continuous action rates via Daft window functions over the per-frame geometry.
 
    Differentiates the continuous quantities over time: curl / wrist-height / arm-extension by
    scalar diffs, wrist & finger-joint motion by native euclidean_distance, wrist rotation by
    forearm_roll (then smoothed). All native/in-DAG — no collect, no thresholds, no booleans.
    Keeps the per-frame columns the query-time scenarios read (closure, flex_nonthumb, thumb mins).
    Expects the output of add_state_features.
    """
    dt = 1.0 / fps
    per_episode = Window().partition_by("episode_index").order_by("frame_index")
    smooth = Window().partition_by("episode_index").order_by("frame_index").rows_between(-2, 2)
    for tag, _ in HANDS:
        # euclidean_distance needs fixed-size-list inputs; materialize the casts once per hand
        df = df.with_column(
            f"_wrist_v_{tag}", col(f"wrist_{tag}").cast(DataType.fixed_size_list(DataType.float64(), WRIST_DIM))
        )
        df = df.with_column(
            f"_joints_v_{tag}",
            col(f"local_joints_{tag}").cast(DataType.fixed_size_list(DataType.float64(), LOCAL_JOINTS_DIM)),
        )
        df = df.with_column(
            f"curl_rate_{tag}",
            ((col(f"curl_{tag}").lead(1).over(per_episode) - col(f"curl_{tag}")) / dt).fill_null(0.0),
        )
        df = df.with_column(
            f"wrist_vert_vel_{tag}",
            ((col(f"wrist_height_{tag}").lead(1).over(per_episode) - col(f"wrist_height_{tag}")) / dt).fill_null(0.0),
        )
        df = df.with_column(
            f"arm_ext_rate_{tag}",
            ((col(f"arm_extension_{tag}").lead(1).over(per_episode) - col(f"arm_extension_{tag}")) / dt).fill_null(0.0),
        )
        df = df.with_column(
            f"wrist_speed_{tag}",
            (
                euclidean_distance(col(f"_wrist_v_{tag}"), col(f"_wrist_v_{tag}").lead(1).over(per_episode)) / dt
            ).fill_null(0.0),
        )
        df = df.with_column(
            f"articulation_{tag}",
            (
                euclidean_distance(col(f"_joints_v_{tag}"), col(f"_joints_v_{tag}").lead(1).over(per_episode)) / dt
            ).fill_null(0.0),
        )
        df = df.with_column(
            f"roll_raw_{tag}",
            forearm_roll(
                col(f"wrist_rot6d_{tag}"),
                col(f"wrist_rot6d_{tag}").lead(1).over(per_episode),
                col(f"forearm_axis_{tag}"),
            )
            / dt,
        )
        df = df.with_column(f"roll_{tag}", col(f"roll_raw_{tag}").mean().over(smooth))
 
    keep = ["episode_index", "frame_index"]
    for tag, _ in HANDS:
        keep += [
            f"closure_{tag}",
            f"flex_nonthumb_{tag}",
            f"thumb_min_tip_{tag}",
            f"thumb_min_knuckle_{tag}",
            f"curl_rate_{tag}",
            f"wrist_vert_vel_{tag}",
            f"arm_ext_rate_{tag}",
            f"wrist_speed_{tag}",
            f"articulation_{tag}",
            f"roll_{tag}",
        ]
    return df.select(*keep)

# A custom UDF, fed by a window: lead(1) pulls the *next* frame's wrist
# rotation into this row, and the UDF does the rotation-axis math. (twisting)
@daft.func(return_dtype=DataType.float64())
def forearm_roll(rot6d, rot6d_next, forearm_axis):
    """Wrist roll (rad) about the forearm axis from one frame to the next (0 at an episode's last frame)."""
    if rot6d is None or rot6d_next is None:
        return 0.0
    delta = _rotation_matrix(rot6d_next) @ _rotation_matrix(rot6d).T
    angle = np.arccos(np.clip((np.trace(delta) - 1) / 2, -1, 1))
    rotation_axis = np.array([delta[2, 1] - delta[1, 2],
                              delta[0, 2] - delta[2, 0],
                              delta[1, 0] - delta[0, 1]])
    magnitude = np.linalg.norm(rotation_axis)
    if magnitude < 1e-9:
        return 0.0
    return float(abs(angle * np.dot(rotation_axis / magnitude, np.asarray(forearm_axis))))