from __future__ import annotations
from dataclasses import dataclass
import numpy as np

@dataclass
class MotionClip:
    fps: float
    num_frames: int
    source: str                  # "smpl" | "mhr" | "bvh"
    joint_names: list[str]
    local_rotations: np.ndarray  # [F, J, 4] xyzw, parent-relative; [:,0] == root_rotation
    root_translation: np.ndarray # [F, 3] metres
    root_rotation: np.ndarray    # [F, 4] xyzw (copy of local_rotations[:,0])
    rest_offsets: np.ndarray     # [J, 3] at subject beta
    parents: list[int]           # parent index, root = -1
    source_height_m: float
    # Optional: when a backend already produces WORLD joint rotations (e.g. SAM 3D Body / MHR), it
    # can supply them directly + the source rest-pose world rotations, so retarget skips FK and uses
    # a proper rest-relative transfer. Both default None -> SMPL/BVH local-rotation path.
    world_rotations: np.ndarray = None   # [F, J, 4] xyzw, per-joint WORLD rotation
    rest_world: np.ndarray = None        # [J, 4] xyzw, per-joint WORLD rotation at rest
    preview_mesh: tuple = None           # optional (verts[N,3], faces[M,3]) raw source mesh, frame 0

    def save(self, path) -> None:
        np.savez(
            path, fps=self.fps, num_frames=self.num_frames, source=self.source,
            joint_names=np.array(self.joint_names, dtype=object),
            local_rotations=self.local_rotations, root_translation=self.root_translation,
            root_rotation=self.root_rotation, rest_offsets=self.rest_offsets,
            parents=np.array(self.parents), source_height_m=self.source_height_m,
        )

    @classmethod
    def load(cls, path) -> "MotionClip":
        z = np.load(path, allow_pickle=True)
        return cls(
            fps=float(z["fps"]), num_frames=int(z["num_frames"]), source=str(z["source"]),
            joint_names=list(z["joint_names"]), local_rotations=z["local_rotations"],
            root_translation=z["root_translation"], root_rotation=z["root_rotation"],
            rest_offsets=z["rest_offsets"], parents=[int(p) for p in z["parents"]],
            source_height_m=float(z["source_height_m"]),
        )