"""
AXIOM GradComp — Drift-Free Gradient Compression
Copyright (c) 2026 Clouthier Simulation Labs

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

Features:
    - Error feedback for drift-free compression
    - FP4 L1 format (2.67x compression)
    - GPU-resident buffers

Usage:
    from quarterbit.axiom import GradientCompressor

    gradcomp = GradientCompressor(max_elements=model_params)

    # After backward:
    for p in model.parameters():
        if p.grad is not None:
            compressed = gradcomp.compress(p.grad)
            p.grad = gradcomp.decompress(compressed)
"""

import ctypes
import torch
from pathlib import Path
from dataclasses import dataclass
from typing import Tuple

_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_gradcomp.dll",
        pkg_root / "lib" / "libaxiom_gradcomp.so",
        Path(__file__).parent.parent / "build" / "axiom_gradcomp.dll",
        Path(__file__).parent.parent / "build" / "libaxiom_gradcomp.so",
        Path("/kaggle/working/build/libaxiom_gradcomp.so"),
        Path("/kaggle/working/libaxiom_gradcomp.so"),
    ]

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

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

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

    _lib.axiom_gradcomp_compress.argtypes = [
        ctypes.c_void_p,  # handle
        ctypes.c_void_p,  # d_grads
        ctypes.c_void_p,  # d_q
        ctypes.c_void_p,  # d_qs
        ctypes.c_void_p,  # d_r
        ctypes.c_void_p,  # d_rs
        ctypes.c_int      # n
    ]
    _lib.axiom_gradcomp_compress.restype = ctypes.c_int

    _lib.axiom_gradcomp_decompress.argtypes = [
        ctypes.c_void_p,  # handle
        ctypes.c_void_p,  # d_q
        ctypes.c_void_p,  # d_qs
        ctypes.c_void_p,  # d_r
        ctypes.c_void_p,  # d_rs
        ctypes.c_void_p,  # d_out
        ctypes.c_int      # n
    ]
    _lib.axiom_gradcomp_decompress.restype = ctypes.c_int

    _lib.axiom_gradcomp_destroy.argtypes = [ctypes.c_void_p]
    _lib.axiom_gradcomp_reset.argtypes = [ctypes.c_void_p]

    return _lib


@dataclass
class CompressedGradient:
    """Compressed gradient in FP4 L1 format."""
    q_data: torch.Tensor
    q_scales: torch.Tensor
    r_data: torch.Tensor
    r_scales: torch.Tensor
    shape: Tuple[int, ...]
    n: int


class GradientCompressor:
    """
    Drift-Free Gradient Compression

    Compresses gradients to FP4 L1 format (1.5 bytes/element) with error
    feedback to prevent drift over training.

    Args:
        max_elements: Maximum gradient size to compress

    Example:
        gradcomp = GradientCompressor(max_elements=10_000_000)

        # After loss.backward():
        for p in model.parameters():
            if p.grad is not None:
                compressed = gradcomp.compress(p.grad)
                # Send compressed over network for DDP
                p.grad = gradcomp.decompress(compressed)
    """

    def __init__(self, max_elements: int):
        lib = _load_lib()
        self._handle = lib.axiom_gradcomp_create(max_elements)
        if not self._handle:
            raise RuntimeError("Failed to create GradientCompressor handle")
        self._lib = lib
        self._max_elements = max_elements

        # Pre-allocate output buffers
        data_bytes = (max_elements + 1) // 2
        scale_count = (max_elements + 15) // 16

        self._q_data = torch.empty(data_bytes, dtype=torch.uint8, device='cuda')
        self._q_scales = torch.empty(scale_count, dtype=torch.float32, device='cuda')
        self._r_data = torch.empty(data_bytes, dtype=torch.uint8, device='cuda')
        self._r_scales = torch.empty(scale_count, dtype=torch.float32, device='cuda')

    def compress(self, grad: torch.Tensor) -> CompressedGradient:
        """Compress gradient tensor with error feedback."""
        grad = grad.contiguous().float()
        n = grad.numel()

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

        ret = self._lib.axiom_gradcomp_compress(
            self._handle,
            ctypes.c_void_p(grad.data_ptr()),
            ctypes.c_void_p(self._q_data.data_ptr()),
            ctypes.c_void_p(self._q_scales.data_ptr()),
            ctypes.c_void_p(self._r_data.data_ptr()),
            ctypes.c_void_p(self._r_scales.data_ptr()),
            n
        )
        if ret != 0:
            raise RuntimeError(f"gradcomp_compress failed with code {ret}")

        # Return copies of the compressed data
        data_bytes = (n + 1) // 2
        scale_count = (n + 15) // 16

        return CompressedGradient(
            q_data=self._q_data[:data_bytes].clone(),
            q_scales=self._q_scales[:scale_count].clone(),
            r_data=self._r_data[:data_bytes].clone(),
            r_scales=self._r_scales[:scale_count].clone(),
            shape=grad.shape,
            n=n
        )

    def decompress(self, compressed: CompressedGradient) -> torch.Tensor:
        """Decompress gradient tensor."""
        out = torch.empty(compressed.n, dtype=torch.float32, device='cuda')

        ret = self._lib.axiom_gradcomp_decompress(
            self._handle,
            ctypes.c_void_p(compressed.q_data.data_ptr()),
            ctypes.c_void_p(compressed.q_scales.data_ptr()),
            ctypes.c_void_p(compressed.r_data.data_ptr()),
            ctypes.c_void_p(compressed.r_scales.data_ptr()),
            ctypes.c_void_p(out.data_ptr()),
            compressed.n
        )
        if ret != 0:
            raise RuntimeError(f"gradcomp_decompress failed with code {ret}")

        return out.view(compressed.shape)

    def reset(self) -> None:
        """Reset error feedback state."""
        self._lib.axiom_gradcomp_reset(self._handle)

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

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