"""
AXIOM ActCP — Compressed Activation Checkpointing
Copyright (c) 2026 Clouthier Simulation Labs

Memory: 1.5 bytes/element (62% reduction vs FP32)

Features:
    - Slot-based storage for transformer layers
    - FP4 L1 compression (2.67x compression)
    - GPU-resident buffers

Usage:
    from quarterbit.axiom import ActivationCheckpoint

    actcp = ActivationCheckpoint(max_slots=12, max_elements=batch*seq*dim)

    # In forward pass:
    actcp.store(activation, slot=layer_idx)

    # In backward pass:
    activation = actcp.restore(slot=layer_idx)
"""

import ctypes
import torch
from pathlib import Path

_lib = None

def _load_lib():
    global _lib
    if _lib is not None:
        return _lib

    pkg_root = Path(__file__).parent.parent.parent

    candidates = [
        pkg_root / "lib" / "axiom_actcp.dll",
        pkg_root / "lib" / "libaxiom_actcp.so",
        Path(__file__).parent.parent / "build" / "axiom_actcp.dll",
        Path(__file__).parent.parent / "build" / "libaxiom_actcp.so",
        Path("/kaggle/working/build/libaxiom_actcp.so"),
        Path("/kaggle/working/libaxiom_actcp.so"),
    ]

    for path in candidates:
        if path.exists():
            _lib = ctypes.CDLL(str(path))
            break

    if _lib is None:
        raise RuntimeError(
            f"AXIOM ActCP library not found. Searched: {[str(p) for p in candidates]}"
        )

    # API bindings
    _lib.axiom_actcp_create.argtypes = [ctypes.c_int, ctypes.c_int]
    _lib.axiom_actcp_create.restype = ctypes.c_void_p

    _lib.axiom_actcp_store.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int]
    _lib.axiom_actcp_store.restype = ctypes.c_int

    _lib.axiom_actcp_restore.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_int]
    _lib.axiom_actcp_restore.restype = ctypes.c_int

    _lib.axiom_actcp_destroy.argtypes = [ctypes.c_void_p]

    _lib.axiom_actcp_clear_slot.argtypes = [ctypes.c_void_p, ctypes.c_int]
    _lib.axiom_actcp_clear_all.argtypes = [ctypes.c_void_p]

    _lib.axiom_actcp_get_slots.argtypes = [ctypes.c_void_p]
    _lib.axiom_actcp_get_slots.restype = ctypes.c_int

    _lib.axiom_actcp_get_max_n.argtypes = [ctypes.c_void_p]
    _lib.axiom_actcp_get_max_n.restype = ctypes.c_int

    return _lib


class ActivationCheckpoint:
    """
    Compressed Activation Checkpointing

    Stores activations in FP4 L1 format (1.5 bytes/element) for gradient
    checkpointing with 62% memory reduction.

    Args:
        max_slots: Number of slots (e.g., transformer layers)
        max_elements: Maximum activation size per slot

    Example:
        actcp = ActivationCheckpoint(max_slots=12, max_elements=4*128*768)

        # During forward:
        for i, layer in enumerate(layers):
            actcp.store(x.detach(), slot=i)
            x = layer(x)

        # Activations auto-restored during backward via hooks
    """

    def __init__(self, max_slots: int, max_elements: int):
        lib = _load_lib()
        self._handle = lib.axiom_actcp_create(max_slots, max_elements)
        if not self._handle:
            raise RuntimeError("Failed to create ActivationCheckpoint handle")
        self._lib = lib
        self._max_slots = max_slots
        self._max_elements = max_elements
        self._stored_shapes = {}

    def store(self, tensor: torch.Tensor, slot: int) -> None:
        """Store activation tensor in compressed format."""
        if slot < 0 or slot >= self._max_slots:
            raise ValueError(f"Slot {slot} out of range [0, {self._max_slots})")

        tensor = tensor.contiguous().float()
        n = tensor.numel()

        if n > self._max_elements:
            raise ValueError(f"Tensor size {n} exceeds max_elements {self._max_elements}")

        self._stored_shapes[slot] = tensor.shape

        ret = self._lib.axiom_actcp_store(
            self._handle,
            ctypes.c_void_p(tensor.data_ptr()),
            n,
            slot
        )
        if ret != 0:
            raise RuntimeError(f"actcp_store failed with code {ret}")

    def restore(self, slot: int, device: str = 'cuda') -> torch.Tensor:
        """Restore activation tensor from compressed format."""
        if slot not in self._stored_shapes:
            raise ValueError(f"Slot {slot} has no stored activation")

        shape = self._stored_shapes[slot]
        n = 1
        for s in shape:
            n *= s

        out = torch.empty(n, dtype=torch.float32, device=device)

        ret = self._lib.axiom_actcp_restore(
            self._handle,
            ctypes.c_void_p(out.data_ptr()),
            n,
            slot
        )
        if ret != 0:
            raise RuntimeError(f"actcp_restore failed with code {ret}")

        return out.view(shape)

    def clear_slot(self, slot: int) -> None:
        """Clear a single slot."""
        self._lib.axiom_actcp_clear_slot(self._handle, slot)
        if slot in self._stored_shapes:
            del self._stored_shapes[slot]

    def clear_all(self) -> None:
        """Clear all slots."""
        self._lib.axiom_actcp_clear_all(self._handle)
        self._stored_shapes.clear()

    @property
    def max_slots(self) -> int:
        return self._max_slots

    @property
    def max_elements(self) -> int:
        return self._max_elements

    def __del__(self):
        if hasattr(self, '_handle') and self._handle:
            self._lib.axiom_actcp_destroy(self._handle)
            self._handle = None
