"""Debug hooks for capturing intermediate tensors during DiT forward pass."""

import mlx.core as mx
import numpy as np


class IntermediateCapture:
    """Captures named intermediate tensors for debugging.

    Usage::

        cap = IntermediateCapture()
        # Pass to model via debug_capture parameter
        model(x, debug_capture=cap)
        # Save for comparison
        cap.save("mlx_intermediates.npz")
    """

    def __init__(self, enabled: bool = True):
        self._data: dict[str, mx.array] = {}
        self.enabled = enabled

    def record(self, name: str, tensor: mx.array) -> None:
        """Record a tensor if enabled."""
        if self.enabled:
            self._data[name] = tensor

    def save(self, path: str) -> None:
        """Save all captured tensors to a .npz file."""
        np_data = {k: np.array(v) for k, v in self._data.items()}
        np.savez(path, **np_data)

    @classmethod
    def load(cls, path: str) -> "IntermediateCapture":
        """Load captured tensors from a .npz file."""
        cap = cls()
        data = np.load(path)
        for k in data.files:
            cap._data[k] = mx.array(data[k])
        return cap

    def __contains__(self, name: str) -> bool:
        return name in self._data

    def __getitem__(self, name: str) -> mx.array:
        return self._data[name]

    def __len__(self) -> int:
        return len(self._data)

    def keys(self) -> list[str]:
        return list(self._data.keys())
