"""
QuarterBit AXIOM
================

Memory-efficient training components:
- AxiomUltra: Fused optimizer with gradient compression
- AxiomGradComp: Standalone gradient compression
- AxiomActCP: Activation checkpointing

Usage:
    from quarterbit import Axiom
    optimizer = Axiom(model.parameters(), lr=1e-3)

Copyright (c) 2026 Clouthier Simulation Labs. All rights reserved.
"""

import torch
from torch.optim import Optimizer
from typing import Optional, Callable, Iterable, List, Tuple
import ctypes
import os
import platform

__all__ = ['AxiomUltra', 'AxiomGradComp', 'AxiomActCP', 'load_ultra_libs']


def _ptr(tensor):
    """Get ctypes void pointer to tensor data."""
    if tensor is None:
        return None
    return ctypes.cast(tensor.data_ptr(), ctypes.c_void_p)


def _get_lib_paths():
    """Get library search paths based on platform."""
    if platform.system() == 'Windows':
        ext = '.dll'
        prefix = ''
    else:
        ext = '.so'
        prefix = 'lib'

    base_dir = os.path.dirname(__file__)

    return {
        'axiom': [
            os.path.join(base_dir, '..', 'lib', f'{prefix}quarterbit_axiom{ext}'),
            os.path.join(base_dir, '..', '..', 'build', f'quarterbit_axiom{ext}'),
            os.path.join('C:', 'QuarterBit', 'build', f'quarterbit_axiom{ext}'),
            f'/kaggle/working/build/{prefix}quarterbit_axiom{ext}',
        ],
        'gradcomp': [
            os.path.join(base_dir, '..', 'lib', f'{prefix}axiom_gradcomp{ext}'),
            os.path.join(base_dir, '..', '..', 'build', f'axiom_gradcomp{ext}'),
            os.path.join('C:', 'QuarterBit', 'build', f'axiom_gradcomp{ext}'),
            f'/kaggle/working/build/{prefix}axiom_gradcomp{ext}',
        ],
        'actcp': [
            os.path.join(base_dir, '..', 'lib', f'{prefix}axiom_actcp{ext}'),
            os.path.join(base_dir, '..', '..', 'build', f'axiom_actcp{ext}'),
            os.path.join('C:', 'QuarterBit', 'build', f'axiom_actcp{ext}'),
            f'/kaggle/working/build/{prefix}axiom_actcp{ext}',
        ],
        'ultra': [
            os.path.join(base_dir, '..', 'lib', f'{prefix}axiom_ultra{ext}'),
            os.path.join(base_dir, '..', '..', 'build', f'axiom_ultra{ext}'),
            os.path.join('C:', 'QuarterBit', 'build', f'axiom_ultra{ext}'),
            f'/kaggle/working/build/{prefix}axiom_ultra{ext}',
        ],
    }


def _load_lib(name: str, paths: List[str]):
    """Load a library from the first available path."""
    for path in paths:
        if os.path.exists(path):
            try:
                return ctypes.CDLL(path)
            except Exception:
                continue
    return None


# Global library cache
_libs = {}


def load_ultra_libs():
    """Load all AXIOM Ultra libraries. Returns dict of loaded libs."""
    global _libs
    if _libs:
        return _libs

    paths = _get_lib_paths()

    for name, search_paths in paths.items():
        lib = _load_lib(name, search_paths)
        if lib:
            _libs[name] = lib

    return _libs


def _setup_ultra_lib(lib):
    """Set up function signatures for axiom_ultra library."""
    lib.axiom_ultra_version.argtypes = []
    lib.axiom_ultra_version.restype = ctypes.c_char_p

    lib.axiom_ultra_create.argtypes = [
        ctypes.c_int,  # n
        ctypes.c_int,  # total_steps
        ctypes.c_int,  # warmup_steps
        ctypes.c_int   # use_gradcomp
    ]
    lib.axiom_ultra_create.restype = ctypes.c_void_p

    lib.axiom_ultra_step.argtypes = [
        ctypes.c_void_p,  # handle
        ctypes.c_void_p,  # params
        ctypes.c_void_p,  # grads
        ctypes.c_float,   # lr
        ctypes.c_float,   # beta1
        ctypes.c_float,   # beta2
        ctypes.c_float,   # eps
        ctypes.c_int      # step
    ]
    lib.axiom_ultra_step.restype = ctypes.c_int

    lib.axiom_ultra_bytes_per_param.argtypes = [ctypes.c_void_p]
    lib.axiom_ultra_bytes_per_param.restype = ctypes.c_float

    lib.axiom_ultra_reset.argtypes = [ctypes.c_void_p]
    lib.axiom_ultra_reset.restype = None

    lib.axiom_ultra_destroy.argtypes = [ctypes.c_void_p]
    lib.axiom_ultra_destroy.restype = None


def _setup_gradcomp_lib(lib):
    """Set up function signatures for axiom_gradcomp library."""
    lib.axiom_gradcomp_version.argtypes = []
    lib.axiom_gradcomp_version.restype = ctypes.c_char_p

    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_reset.argtypes = [ctypes.c_void_p]
    lib.axiom_gradcomp_reset.restype = None

    lib.axiom_gradcomp_destroy.argtypes = [ctypes.c_void_p]
    lib.axiom_gradcomp_destroy.restype = None


def _setup_actcp_lib(lib):
    """Set up function signatures for axiom_actcp library."""
    lib.axiom_actcp_version.argtypes = []
    lib.axiom_actcp_version.restype = ctypes.c_char_p

    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,  # handle
        ctypes.c_void_p,  # d_act
        ctypes.c_int,     # n
        ctypes.c_int      # slot
    ]
    lib.axiom_actcp_store.restype = ctypes.c_int

    lib.axiom_actcp_restore.argtypes = [
        ctypes.c_void_p,  # handle
        ctypes.c_void_p,  # d_out
        ctypes.c_int,     # n
        ctypes.c_int      # slot
    ]
    lib.axiom_actcp_restore.restype = ctypes.c_int

    lib.axiom_actcp_clear_slot.argtypes = [ctypes.c_void_p, ctypes.c_int]
    lib.axiom_actcp_clear_slot.restype = None

    lib.axiom_actcp_clear_all.argtypes = [ctypes.c_void_p]
    lib.axiom_actcp_clear_all.restype = None

    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

    lib.axiom_actcp_destroy.argtypes = [ctypes.c_void_p]
    lib.axiom_actcp_destroy.restype = None


class AxiomUltra(Optimizer):
    """
    AXIOM Optimizer with fused gradient compression.

    Args:
        params: Parameters to optimize
        lr: Learning rate (default: 1e-3)
        betas: Momentum coefficients (default: (0.9, 0.999))
        eps: Numerical stability term (default: 1e-8)
        weight_decay: Weight decay coefficient (default: 0)
        total_steps: Total training steps for scheduling (default: 0)
        warmup_steps: Warmup steps (default: 0)
        use_gradcomp: Enable gradient compression (default: True)

    Example:
        >>> optimizer = Axiom(model.parameters(), lr=1e-3)
        >>> loss.backward()
        >>> optimizer.step()
        >>> optimizer.zero_grad()
    """

    def __init__(
        self,
        params: Iterable[torch.nn.Parameter],
        lr: float = 1e-3,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0,
        total_steps: int = 0,
        warmup_steps: int = 0,
        use_gradcomp: bool = True
    ):
        defaults = dict(
            lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
            total_steps=total_steps, warmup_steps=warmup_steps
        )
        super().__init__(params, defaults)

        self._handles = {}
        self._step_count = 0
        self._total_steps = total_steps
        self._warmup_steps = warmup_steps
        self._use_gradcomp = use_gradcomp

        # Load library
        libs = load_ultra_libs()
        self._lib = libs.get('ultra')

        if self._lib is None:
            raise RuntimeError(
                "AXIOM library not found. "
                "Build with: cd src && build_all.bat (Windows) or ./build.sh (Linux)"
            )

        _setup_ultra_lib(self._lib)
        self._version = self._lib.axiom_ultra_version().decode()

    def _get_handle(self, p: torch.nn.Parameter):
        """Get or create handle for parameter tensor."""
        pid = id(p)
        if pid not in self._handles:
            if p.is_cuda and p.dtype == torch.float32:
                handle = self._lib.axiom_ultra_create(
                    p.numel(),
                    self._total_steps,
                    self._warmup_steps,
                    1 if self._use_gradcomp else 0
                )
                self._handles[pid] = handle if handle else None
            else:
                self._handles[pid] = None
        return self._handles[pid]

    @torch.no_grad()
    def step(self, closure: Optional[Callable] = None):
        """Perform a single optimization step."""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        self._step_count += 1

        for group in self.param_groups:
            beta1, beta2 = group['betas']
            lr = group['lr']
            eps = group['eps']
            weight_decay = group['weight_decay']

            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad.data
                handle = self._get_handle(p)

                if handle is not None:
                    # Apply weight decay
                    if weight_decay != 0:
                        grad = grad.add(p.data, alpha=weight_decay)

                    ret = self._lib.axiom_ultra_step(
                        handle,
                        _ptr(p.data),
                        _ptr(grad),
                        ctypes.c_float(lr),
                        ctypes.c_float(beta1),
                        ctypes.c_float(beta2),
                        ctypes.c_float(eps),
                        self._step_count
                    )

                    if ret != 0:
                        raise RuntimeError(f"AXIOM Ultra step failed with code {ret}")
                else:
                    # Fallback to CPU/non-float32
                    if 'fallback_state' not in self.state[p]:
                        self.state[p]['fallback_state'] = {
                            'm': torch.zeros_like(p.data),
                            'v': torch.zeros_like(p.data)
                        }
                    state = self.state[p]['fallback_state']

                    if weight_decay != 0:
                        grad = grad.add(p.data, alpha=weight_decay)

                    state['m'].mul_(beta1).add_(grad, alpha=1 - beta1)
                    state['v'].mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

                    bc1 = 1 - beta1 ** self._step_count
                    bc2 = 1 - beta2 ** self._step_count

                    step_size = lr / bc1
                    denom = (state['v'].sqrt() / (bc2 ** 0.5)).add_(eps)
                    p.data.addcdiv_(state['m'], denom, value=-step_size)

        return loss

    def reset(self):
        """Reset optimizer state."""
        for handle in self._handles.values():
            if handle:
                self._lib.axiom_ultra_reset(handle)
        self._step_count = 0

    @property
    def version(self) -> str:
        return self._version

    @property
    def bytes_per_param(self) -> float:
        """Get actual bytes per parameter from first handle."""
        for handle in self._handles.values():
            if handle:
                return self._lib.axiom_ultra_bytes_per_param(handle)
        return self.BYTES_PER_PARAM

    def __del__(self):
        if hasattr(self, '_handles') and hasattr(self, '_lib') and self._lib:
            for handle in self._handles.values():
                if handle:
                    try:
                        self._lib.axiom_ultra_destroy(handle)
                    except:
                        pass


class AxiomGradComp:
    """
    Gradient compression with error feedback.

    Args:
        n: Number of parameters to compress

    Example:
        >>> gradcomp = AxiomGradComp(total_params)
        >>> compressed = gradcomp.compress(gradients)
        >>> gradients = gradcomp.decompress(compressed)
    """

    def __init__(self, n: int):
        self.n = n

        libs = load_ultra_libs()
        self._lib = libs.get('gradcomp')

        if self._lib is None:
            raise RuntimeError("AXIOM GradComp library not found.")

        _setup_gradcomp_lib(self._lib)
        self._version = self._lib.axiom_gradcomp_version().decode()

        self._handle = self._lib.axiom_gradcomp_create(n)
        if not self._handle:
            raise RuntimeError("Failed to create GradComp handle")

        # Allocate compressed buffers
        self._data_bytes = (n + 1) // 2
        self._scale_count = (n + 15) // 16

        self._q_data = torch.zeros(self._data_bytes, dtype=torch.uint8, device='cuda')
        self._q_scales = torch.zeros(self._scale_count, dtype=torch.float32, device='cuda')
        self._r_data = torch.zeros(self._data_bytes, dtype=torch.uint8, device='cuda')
        self._r_scales = torch.zeros(self._scale_count, dtype=torch.float32, device='cuda')

    def compress(self, grads: torch.Tensor) -> Tuple[torch.Tensor, ...]:
        """Compress gradients with error feedback."""
        assert grads.numel() <= self.n, f"Gradient size {grads.numel()} exceeds capacity {self.n}"

        ret = self._lib.axiom_gradcomp_compress(
            self._handle,
            _ptr(grads),
            _ptr(self._q_data),
            _ptr(self._q_scales),
            _ptr(self._r_data),
            _ptr(self._r_scales),
            grads.numel()
        )

        if ret != 0:
            raise RuntimeError(f"GradComp compress failed with code {ret}")

        return (self._q_data, self._q_scales, self._r_data, self._r_scales)

    def decompress(self, compressed: Tuple[torch.Tensor, ...], out: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Decompress gradients."""
        q_data, q_scales, r_data, r_scales = compressed

        if out is None:
            out = torch.zeros(self.n, dtype=torch.float32, device='cuda')

        ret = self._lib.axiom_gradcomp_decompress(
            self._handle,
            _ptr(q_data),
            _ptr(q_scales),
            _ptr(r_data),
            _ptr(r_scales),
            _ptr(out),
            self.n
        )

        if ret != 0:
            raise RuntimeError(f"GradComp decompress failed with code {ret}")

        return out

    def reset(self):
        """Reset error feedback buffer."""
        self._lib.axiom_gradcomp_reset(self._handle)

    @property
    def version(self) -> str:
        return self._version

    def __del__(self):
        if hasattr(self, '_handle') and hasattr(self, '_lib') and self._lib and self._handle:
            try:
                self._lib.axiom_gradcomp_destroy(self._handle)
            except:
                pass


class AxiomActCP:
    """
    Activation checkpointing with compression.

    Args:
        max_slots: Number of checkpoint slots (e.g., num_layers)
        max_n: Maximum activation size per slot

    Example:
        >>> actcp = AxiomActCP(max_slots=12, max_n=65536)
        >>> # During forward:
        >>> actcp.store(activation, slot=layer_idx)
        >>> # During backward:
        >>> activation = actcp.restore(slot=layer_idx)
    """

    BYTES_PER_ELEM = 1.5

    def __init__(self, max_slots: int, max_n: int):
        self.max_slots = max_slots
        self.max_n = max_n

        libs = load_ultra_libs()
        self._lib = libs.get('actcp')

        if self._lib is None:
            raise RuntimeError("AXIOM ActCP library not found.")

        _setup_actcp_lib(self._lib)
        self._version = self._lib.axiom_actcp_version().decode()

        self._handle = self._lib.axiom_actcp_create(max_slots, max_n)
        if not self._handle:
            raise RuntimeError("Failed to create ActCP handle")

    def store(self, activation: torch.Tensor, slot: int):
        """Store activation to a slot."""
        assert 0 <= slot < self.max_slots, f"Slot {slot} out of range [0, {self.max_slots})"
        assert activation.numel() <= self.max_n, f"Activation size {activation.numel()} exceeds max {self.max_n}"

        ret = self._lib.axiom_actcp_store(
            self._handle,
            _ptr(activation.contiguous()),
            activation.numel(),
            slot
        )

        if ret != 0:
            raise RuntimeError(f"ActCP store failed with code {ret}")

    def restore(self, slot: int, size: Optional[int] = None, out: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Restore activation from a slot."""
        assert 0 <= slot < self.max_slots, f"Slot {slot} out of range [0, {self.max_slots})"

        n = size if size is not None else self.max_n

        if out is None:
            out = torch.zeros(n, dtype=torch.float32, device='cuda')

        ret = self._lib.axiom_actcp_restore(
            self._handle,
            _ptr(out),
            n,
            slot
        )

        if ret != 0:
            raise RuntimeError(f"ActCP restore failed with code {ret}")

        return out

    def clear_slot(self, slot: int):
        """Clear a specific slot."""
        self._lib.axiom_actcp_clear_slot(self._handle, slot)

    def clear_all(self):
        """Clear all slots."""
        self._lib.axiom_actcp_clear_all(self._handle)

    @property
    def version(self) -> str:
        return self._version

    @property
    def memory_bytes(self) -> int:
        """Get total memory used by checkpoint storage."""
        data_bytes = (self.max_n + 1) // 2
        scale_count = (self.max_n + 15) // 16
        per_slot = 2 * data_bytes + 2 * scale_count * 4
        return self.max_slots * per_slot

    def __del__(self):
        if hasattr(self, '_handle') and hasattr(self, '_lib') and self._lib and self._handle:
            try:
                self._lib.axiom_actcp_destroy(self._handle)
            except:
                pass
