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

Memory-efficient Adam variant with enhanced numerical precision.

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
import ctypes
import os


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


def _load_axiom_lib():
    """Load the AXIOM CUDA library."""
    if os.name == 'nt':
        lib_name = 'quarterbit_axiom.dll'
    else:
        lib_name = 'libquarterbit_axiom.so'

    search_paths = [
        os.path.join(os.path.dirname(__file__), '..', 'lib', lib_name),
        os.path.join(os.path.dirname(__file__), '..', '..', 'build', lib_name),
        os.path.join(os.path.dirname(__file__), '..', '..', '..', 'build', lib_name),
        os.path.join('C:', 'QuarterBit', 'build', lib_name),
        '/kaggle/working/build/' + lib_name,
    ]

    for path in search_paths:
        if os.path.exists(path):
            try:
                lib = ctypes.CDLL(path)

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

                lib.axiom_step.argtypes = [
                    ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p,
                    ctypes.c_float, ctypes.c_float, ctypes.c_float, ctypes.c_float,
                    ctypes.c_int
                ]
                lib.axiom_step.restype = ctypes.c_int

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

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

                return lib
            except Exception:
                continue
    return None


class Axiom(Optimizer):
    """
    QuarterBit AXIOM Optimizer.

    World's most memory-efficient drift-free optimizer.
    90%+ memory savings vs Adam with full numerical precision.

    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 LR scheduling (default: 0)
        warmup_steps: Warmup steps (default: 0)

    Example:
        >>> optimizer = Axiom(model.parameters(), lr=1e-3)
        >>> optimizer = Axiom(model.parameters(), lr=1e-3, total_steps=10000)
    """

    def __init__(
        self,
        params: Iterable[torch.nn.Parameter],
        lr: float = 1e-3,
        betas: tuple = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0,
        total_steps: int = 0,
        warmup_steps: int = 0,
    ):
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if eps < 0.0:
            raise ValueError(f"Invalid epsilon: {eps}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta[0]: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta[1]: {betas[1]}")

        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._lib = _load_axiom_lib()
        if self._lib is None:
            raise RuntimeError(
                "QuarterBit AXIOM CUDA library not found. "
                "Ensure CUDA is available and the library is installed."
            )

    def _get_handle(self, p):
        """Get or create handle for parameter."""
        if id(p) not in self._handles:
            if p.is_cuda and p.dtype == torch.float32:
                handle = self._lib.axiom_create(
                    ctypes.c_int(p.numel()),
                    ctypes.c_int(self._total_steps),
                    ctypes.c_int(self._warmup_steps)
                )
                self._handles[id(p)] = handle if handle else None
            else:
                self._handles[id(p)] = None
        return self._handles[id(p)]

    def set_schedule(self, total_steps: int, warmup_steps: int = 0):
        """
        Update training schedule.

        Args:
            total_steps: Total training steps
            warmup_steps: Warmup steps (default: 0)
        """
        self._total_steps = total_steps
        self._warmup_steps = warmup_steps

        for group in self.param_groups:
            group['total_steps'] = total_steps
            group['warmup_steps'] = warmup_steps

        for handle in self._handles.values():
            if handle is not None:
                self._lib.axiom_set_schedule(
                    handle,
                    ctypes.c_int(total_steps),
                    ctypes.c_int(warmup_steps)
                )

    @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
                if grad.is_sparse:
                    raise RuntimeError("Axiom does not support sparse gradients")

                handle = self._get_handle(p)

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

                    self._lib.axiom_step(
                        handle,
                        _ptr(p.data),
                        _ptr(grad),
                        ctypes.c_float(lr),
                        ctypes.c_float(beta1),
                        ctypes.c_float(beta2),
                        ctypes.c_float(eps),
                        ctypes.c_int(self._step_count)
                    )
                else:
                    raise RuntimeError(
                        f"Parameter not on CUDA or not float32. "
                        f"Got device={p.device}, dtype={p.dtype}"
                    )

        return loss

    def __del__(self):
        """Clean up handles."""
        if hasattr(self, '_handles') and hasattr(self, '_lib') and self._lib is not None:
            for handle in self._handles.values():
                if handle is not None:
                    try:
                        self._lib.axiom_destroy(handle)
                    except Exception:
                        pass
