"""MotionClone backend adapter.

MotionClone (Bujiazi et al., ICLR 2025) extracts a motion descriptor from one or
more reference videos via sparse temporal attention — training-free, no inversion.
Repo: https://github.com/Bujiazi/MotionClone (514 stars, last push 2025-06-17).

This file is a contract adapter. The cursor and op-set are wired correctly; the
last-mile model call is stubbed pending vendoring. To activate:
  1. Confirm upstream license is compatible (no LICENSE declared at last check).
  2. git submodule add https://github.com/Bujiazi/MotionClone vendor/motionclone
  3. Implement _extract_motion_descriptor and _apply_motion below against the
     vendored module.
  4. Call MotionCloneAdapter().register() at studio startup.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Sequence

import numpy as np

from pixel_cursor import PixelCursor, Image, IdentityLock, FrameStack, ops
from pixel_cursor.artifact import _new_framestack


@dataclass
class MotionDescriptor:
    sparse_attention: np.ndarray
    source_paths: tuple[str, ...]


class MotionCloneAdapter:
    """Adapter exposing MotionClone primitives to the PixelCursor op registry."""

    SOURCE_TAG = "motionclone"

    def register(self) -> None:
        ops.register_backend("bind_motion_exemplar", "motionclone", self.bind_motion_exemplar)
        ops.register_backend("write_motion", "motionclone", self.write_motion)

    def bind_motion_exemplar(
        self, cursor: PixelCursor, exemplar_paths: Sequence[str]
    ) -> PixelCursor:
        descriptor = self._extract_motion_descriptor(exemplar_paths)
        lock = IdentityLock(
            embedding=descriptor.sparse_attention.flatten(),
            source=f"{self.SOURCE_TAG}:{','.join(exemplar_paths)}",
        )
        return cursor.lock_identity(lock)

    def write_motion(self, cursor: PixelCursor, motion_spec) -> FrameStack:
        if cursor.identity_lock is None or not cursor.identity_lock.source.startswith(
            f"{self.SOURCE_TAG}:"
        ):
            raise ValueError(
                "write_motion(motionclone) requires bind_motion_exemplar(..., backend='motionclone') first"
            )
        if not isinstance(cursor.artifact, Image):
            raise TypeError(
                f"MotionClone write_motion expects an Image artifact; got {type(cursor.artifact).__name__}"
            )
        frames = self._apply_motion(cursor.artifact.pixels, cursor.identity_lock, motion_spec)
        return _new_framestack(frames, fps=24, name=f"{self.SOURCE_TAG}-generated")

    # ── pending vendoring ──

    def _extract_motion_descriptor(self, paths: Sequence[str]) -> MotionDescriptor:
        return MotionDescriptor(
            sparse_attention=np.zeros((1, 1, 1), dtype=np.float32),
            source_paths=tuple(paths),
        )

    def _apply_motion(
        self, image_pixels: np.ndarray, lock: IdentityLock, motion_spec
    ) -> np.ndarray:
        return np.repeat(image_pixels[None, ...], 24, axis=0)