"""Dataset normalisation statistics and delta-base inversion for online execution."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from .data_utils import matrix_to_rotation_6d as _matrix_to_rotation_6d
from .data_utils import rotation_6d_to_matrix


def delta_base_to_absolute(
    reference_state: np.ndarray, delta_base_action: np.ndarray
) -> np.ndarray:
    """Invert T-Rex delta-base arm poses; hand targets are already absolute."""

    reference = np.asarray(reference_state, dtype=np.float64)
    delta = np.asarray(delta_base_action, dtype=np.float64)
    if reference.shape != (62,) or delta.shape[-1] != 62:
        raise ValueError("expected reference [62] and delta action [...,62]")
    output = np.empty_like(delta)
    for pose_slice, hand_slice in (
        (slice(0, 9), slice(9, 31)),
        (slice(31, 40), slice(40, 62)),
    ):
        reference_pose = reference[pose_slice]
        delta_pose = delta[..., pose_slice]
        reference_rotation = rotation_6d_to_matrix(reference_pose[3:9])
        delta_rotation = rotation_6d_to_matrix(delta_pose[..., 3:9])
        target_translation = (
            reference_pose[:3]
            + np.einsum("ij,...j->...i", reference_rotation, delta_pose[..., :3])
        )
        target_rotation = np.einsum(
            "ij,...jk->...ik", reference_rotation, delta_rotation
        )
        output[..., pose_slice] = np.concatenate(
            (target_translation, _matrix_to_rotation_6d(target_rotation)), axis=-1
        )
        output[..., hand_slice] = delta[..., hand_slice]
    return output.astype(np.float32)


@dataclass(frozen=True)
class TrexRuntimeStatistics:
    action_q01: np.ndarray
    action_q99: np.ndarray
    state_q01: np.ndarray
    state_q99: np.ndarray
    force_q01: np.ndarray
    force_q99: np.ndarray

    @classmethod
    def from_dataset(cls, dataset_root: str | Path) -> "TrexRuntimeStatistics":
        root = Path(dataset_root)
        with (root / "meta" / "relative_stats_dreamzero.json").open() as handle:
            action = json.load(handle)["eef62"]
        with (root / "meta" / "stats.json").open() as handle:
            stats = json.load(handle)
        force = stats["observation.tactile_force"]
        state = stats["observation.state_eef62"]
        return cls(
            action_q01=np.asarray(action["q01"], dtype=np.float32),
            action_q99=np.asarray(action["q99"], dtype=np.float32),
            state_q01=np.asarray(state["q01"], dtype=np.float32),
            state_q99=np.asarray(state["q99"], dtype=np.float32),
            force_q01=np.asarray(force["q01"], dtype=np.float32).reshape(10, 6),
            force_q99=np.asarray(force["q99"], dtype=np.float32).reshape(10, 6),
        )

    def normalize_force(self, force: np.ndarray) -> np.ndarray:
        values = np.asarray(force, dtype=np.float32)
        scale = self.force_q99 - self.force_q01
        normalized = np.where(
            scale != 0,
            2.0 * (values - self.force_q01) / np.where(scale == 0, 1.0, scale)
            - 1.0,
            values,
        )
        return np.clip(normalized, -1.0, 1.0).astype(np.float32)

    def normalize_state(self, state: np.ndarray) -> np.ndarray:
        values = np.asarray(state, dtype=np.float32)
        scale = self.state_q99 - self.state_q01
        normalized = np.where(
            scale != 0,
            2.0 * (values - self.state_q01) / np.where(scale == 0, 1.0, scale)
            - 1.0,
            values,
        )
        return np.clip(normalized, -1.0, 1.0).astype(np.float32)

    def denormalize_action(self, normalized_action: np.ndarray) -> np.ndarray:
        values = np.asarray(normalized_action, dtype=np.float32)[..., :62]
        return (
            (values + 1.0) * 0.5 * (self.action_q99 - self.action_q01)
            + self.action_q01
        ).astype(np.float32)



__all__ = ["TrexRuntimeStatistics", "delta_base_to_absolute"]