"""Timestamp alignment and delta-base EEF62 pose math shared by the FACT data path.

These helpers define the 20 Hz logical control grid (nearest-timestamp
gathering from the native-rate LeRobot parquet) and the T-Rex delta-base
action parameterisation: arm poses are expressed in the frame of the block
reference state, hand joint targets stay absolute.
"""

from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass

import numpy as np


def matrix_to_rotation_6d(matrix: np.ndarray) -> np.ndarray:
    """First two columns of a rotation matrix, concatenated (inverse of ``rotation_6d_to_matrix``)."""
    return np.concatenate((matrix[..., :, 0], matrix[..., :, 1]), axis=-1)


@dataclass(frozen=True)
class TimestampSelection:
    indices: np.ndarray
    target_timestamps: np.ndarray
    padding_mask: np.ndarray
    alignment_error: np.ndarray


def nearest_timestamp_indices(
    source_timestamps: np.ndarray,
    target_timestamps: np.ndarray,
    *,
    tie_tolerance: float = 1e-9,
) -> TimestampSelection:
    """Map target times to source rows, choosing the earlier row on a tie."""

    source = np.asarray(source_timestamps, dtype=np.float64)
    target = np.asarray(target_timestamps, dtype=np.float64)
    if source.ndim != 1 or source.size < 2:
        raise ValueError("source timestamps must be one-dimensional with at least two rows")
    if target.ndim != 1 or target.size == 0:
        raise ValueError("target timestamps must be a non-empty one-dimensional array")
    if not np.isfinite(source).all() or not np.isfinite(target).all():
        raise ValueError("timestamps must be finite")
    if not np.all(np.diff(source) > 0):
        raise ValueError("source timestamps must be strictly increasing")
    if not np.all(np.diff(target) > 0):
        raise ValueError("target timestamps must be strictly increasing")

    insertion = np.searchsorted(source, target, side="left")
    lower = np.clip(insertion - 1, 0, source.size - 1)
    upper = np.clip(insertion, 0, source.size - 1)
    lower_error = np.abs(target - source[lower])
    upper_error = np.abs(source[upper] - target)
    choose_upper = upper_error < lower_error - tie_tolerance
    indices = np.where(choose_upper, upper, lower).astype(np.int64)
    padding = (target < source[0]) | (target > source[-1])
    indices[target < source[0]] = 0
    indices[target > source[-1]] = source.size - 1
    return TimestampSelection(
        indices=indices,
        target_timestamps=target,
        padding_mask=padding,
        alignment_error=np.abs(source[indices] - target),
    )


def uniform_target_times(
    anchor_timestamp: float,
    offsets: Iterable[int],
    rate_hz: float,
) -> np.ndarray:
    offsets_array = np.asarray(tuple(offsets), dtype=np.float64)
    if offsets_array.ndim != 1 or offsets_array.size == 0:
        raise ValueError("offsets must be a non-empty sequence")
    if not np.all(np.diff(offsets_array) > 0):
        raise ValueError("offsets must be strictly increasing")
    if not np.isfinite(rate_hz) or rate_hz <= 0:
        raise ValueError("rate_hz must be positive")
    return float(anchor_timestamp) + offsets_array / float(rate_hz)


def rotation_6d_to_matrix(rotation_6d: np.ndarray) -> np.ndarray:
    """Convert matrix-column 6D rotations to proper rotation matrices."""

    values = np.asarray(rotation_6d, dtype=np.float64)
    if values.shape[-1] != 6:
        raise ValueError("rotation_6d must end in six values")
    first = values[..., :3]
    first /= np.linalg.norm(first, axis=-1, keepdims=True).clip(min=1e-8)
    second = values[..., 3:6]
    second = second - np.sum(first * second, axis=-1, keepdims=True) * first
    second /= np.linalg.norm(second, axis=-1, keepdims=True).clip(min=1e-8)
    third = np.cross(first, second)
    return np.stack((first, second, third), axis=-1)


def eef62_delta_base(reference_state: np.ndarray, absolute_targets: np.ndarray) -> np.ndarray:
    """Apply T-Rex delta-base pose math; hand targets remain absolute."""

    reference = np.asarray(reference_state, dtype=np.float64)
    targets = np.asarray(absolute_targets, dtype=np.float64)
    if reference.shape != (62,) or targets.shape[-1] != 62:
        raise ValueError("expected reference [62] and targets [...,62]")

    output = np.empty_like(targets, dtype=np.float64)
    for pose_slice, hand_slice in ((slice(0, 9), slice(9, 31)), (slice(31, 40), slice(40, 62))):
        reference_pose = reference[pose_slice]
        target_pose = targets[..., pose_slice]
        reference_rotation = rotation_6d_to_matrix(reference_pose[3:9])
        target_rotation = rotation_6d_to_matrix(target_pose[..., 3:9])
        translation_delta = np.einsum(
            "ji,...j->...i",
            reference_rotation,
            target_pose[..., :3] - reference_pose[:3],
        )
        rotation_delta = np.einsum(
            "ji,...jk->...ik", reference_rotation, target_rotation
        )
        output[..., pose_slice] = np.concatenate(
            (
                translation_delta,
                rotation_delta[..., :, 0],
                rotation_delta[..., :, 1],
            ),
            axis=-1,
        )
        output[..., hand_slice] = targets[..., hand_slice]
    return output.astype(np.float32)



__all__ = [
    "TimestampSelection",
    "eef62_delta_base",
    "matrix_to_rotation_6d",
    "nearest_timestamp_indices",
    "rotation_6d_to_matrix",
    "uniform_target_times",
]