"""
AXIOM Unified — The Definitive Memory-Efficient Optimizer
==========================================================

PyTorch wrapper for the unified AXIOM CUDA kernel.

Features:
    - 84% memory savings vs Adam (1.25-1.75 B/param vs 8 B/param)
    - Predictive Momentum: +75% faster convergence on Transformers
    - Adaptive Gamma: auto-tunes based on loss trends
    - Error Feedback: eliminates quantization bias
    - EFT precision: drift-free training

Usage:
    from quarterbit.torch.unified import AxiomUnified

    optimizer = AxiomUnified(model.parameters(), lr=1e-3)

    for batch in dataloader:
        loss = model(batch)
        loss.backward()
        optimizer.step(loss.item())  # Pass loss for adaptive gamma
        optimizer.zero_grad()

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

import os
import sys
import ctypes
from ctypes import c_void_p, c_float, c_int, c_size_t, POINTER, byref
from typing import Optional, Iterator

import torch
from torch.optim import Optimizer

# Find and load the DLL/SO
def _load_library():
    lib_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'lib')
    build_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'build')

    if sys.platform == 'win32':
        lib_names = ['axiom_unified.dll', 'libaxiom_unified.dll']
    else:
        lib_names = ['libaxiom_unified.so', 'axiom_unified.so']

    search_paths = [lib_dir, build_dir, '.']

    for path in search_paths:
        for name in lib_names:
            lib_path = os.path.join(path, name)
            if os.path.exists(lib_path):
                try:
                    return ctypes.CDLL(lib_path)
                except Exception:
                    continue

    raise RuntimeError(
        f"Could not load AXIOM Unified library. "
        f"Searched: {search_paths}. "
        f"Run build_all.bat (Windows) or build.sh (Linux) first."
    )


_lib = None


def _get_lib():
    global _lib
    if _lib is None:
        _lib = _load_library()
        _setup_functions(_lib)
    return _lib


def _setup_functions(lib):
    # axiom_unified_create
    lib.axiom_unified_create.argtypes = [
        c_int, c_int, c_int,           # n, total_steps, warmup_steps
        c_float, c_float, c_float,     # gamma_init, gamma_min, gamma_max
        c_float, c_int                 # weight_decay, use_error_feedback
    ]
    lib.axiom_unified_create.restype = c_void_p

    # axiom_unified_create_simple
    lib.axiom_unified_create_simple.argtypes = [c_int, c_int]
    lib.axiom_unified_create_simple.restype = c_void_p

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

    # axiom_unified_report_loss
    lib.axiom_unified_report_loss.argtypes = [c_void_p, c_float]
    lib.axiom_unified_report_loss.restype = None

    # axiom_unified_get_gamma
    lib.axiom_unified_get_gamma.argtypes = [c_void_p]
    lib.axiom_unified_get_gamma.restype = c_float

    # axiom_unified_set_gamma
    lib.axiom_unified_set_gamma.argtypes = [c_void_p, c_float]
    lib.axiom_unified_set_gamma.restype = None

    # axiom_unified_set_weight_decay
    lib.axiom_unified_set_weight_decay.argtypes = [c_void_p, c_float]
    lib.axiom_unified_set_weight_decay.restype = None

    # axiom_unified_destroy
    lib.axiom_unified_destroy.argtypes = [c_void_p]
    lib.axiom_unified_destroy.restype = None

    # axiom_unified_memory_info
    lib.axiom_unified_memory_info.argtypes = [
        c_void_p,
        POINTER(c_size_t),
        POINTER(c_float),
        POINTER(c_float)
    ]
    lib.axiom_unified_memory_info.restype = None

    # axiom_unified_reset
    lib.axiom_unified_reset.argtypes = [c_void_p]
    lib.axiom_unified_reset.restype = None

    # axiom_unified_version
    lib.axiom_unified_version.argtypes = []
    lib.axiom_unified_version.restype = ctypes.c_char_p


class AxiomUnified(Optimizer):
    """
    AXIOM Unified Optimizer — 84% memory savings with adaptive acceleration.

    Args:
        params: Model parameters
        lr: Learning rate (default: 1e-3)
        betas: Coefficients for momentum (default: (0.9, 0.999))
        eps: Numerical stability (default: 1e-8)
        weight_decay: Decoupled weight decay (default: 0.0)
        total_steps: Total training steps for LR schedule (default: 10000)
        warmup_steps: Warmup steps (default: 10% of total)
        gamma_init: Initial prediction strength (default: 1.0)
        gamma_range: (min, max) for adaptive gamma (default: (0.1, 5.0))
        use_error_feedback: Enable error feedback (default: True)
    """

    def __init__(
        self,
        params: Iterator[torch.nn.Parameter],
        lr: float = 1e-3,
        betas: tuple = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0.0,
        total_steps: int = 10000,
        warmup_steps: Optional[int] = None,
        gamma_init: float = 1.0,
        gamma_range: tuple = (0.1, 5.0),
        use_error_feedback: bool = True,
    ):
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta1: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta2: {betas[1]}")

        defaults = dict(
            lr=lr,
            betas=betas,
            eps=eps,
            weight_decay=weight_decay,
        )
        super().__init__(params, defaults)

        self.total_steps = total_steps
        self.warmup_steps = warmup_steps if warmup_steps is not None else total_steps // 10
        self.gamma_init = gamma_init
        self.gamma_min, self.gamma_max = gamma_range
        self.use_error_feedback = use_error_feedback

        self._lib = _get_lib()
        self._handles = {}
        self._step_count = 0
        self._initialized = False

    def _init_group(self, group_idx: int, params: list):
        """Initialize CUDA handle for a parameter group."""
        # Flatten all parameters in the group
        flat_params = []
        param_info = []

        for p in params:
            if p.requires_grad and p.grad is not None:
                flat_params.append(p.data.view(-1))
                param_info.append((p.data_ptr(), p.numel()))

        if not flat_params:
            return None

        total_n = sum(p.numel() for p in flat_params)

        handle = self._lib.axiom_unified_create(
            total_n,
            self.total_steps,
            self.warmup_steps,
            self.gamma_init,
            self.gamma_min,
            self.gamma_max,
            self.param_groups[group_idx]['weight_decay'],
            1 if self.use_error_feedback else 0
        )

        if handle is None:
            raise RuntimeError("Failed to create AXIOM Unified handle")

        return handle

    @torch.no_grad()
    def step(self, loss: Optional[float] = None, closure=None):
        """
        Perform a single optimization step.

        Args:
            loss: Current loss value for adaptive gamma (optional but recommended)
            closure: A closure that reevaluates the model and returns the loss
        """
        if closure is not None:
            with torch.enable_grad():
                loss_val = closure()
                if loss is None:
                    loss = loss_val.item() if isinstance(loss_val, torch.Tensor) else loss_val

        self._step_count += 1

        for group_idx, group in enumerate(self.param_groups):
            params_with_grad = [p for p in group['params'] if p.grad is not None]

            if not params_with_grad:
                continue

            # Process each parameter individually (simpler, works with any layout)
            for p in params_with_grad:
                if not p.is_cuda:
                    raise RuntimeError("AXIOM Unified requires CUDA tensors")

                n = p.numel()

                # Get or create handle for this parameter
                param_id = id(p)
                if param_id not in self._handles:
                    handle = self._lib.axiom_unified_create(
                        n,
                        self.total_steps,
                        self.warmup_steps,
                        self.gamma_init,
                        self.gamma_min,
                        self.gamma_max,
                        group['weight_decay'],
                        1 if self.use_error_feedback else 0
                    )
                    if handle is None:
                        raise RuntimeError("Failed to create AXIOM handle")
                    self._handles[param_id] = handle

                handle = self._handles[param_id]

                # Report loss for adaptive gamma
                if loss is not None:
                    self._lib.axiom_unified_report_loss(handle, loss)

                # Get data pointers
                params_ptr = p.data.data_ptr()
                grads_ptr = p.grad.data.data_ptr()

                # Call kernel
                beta1, beta2 = group['betas']
                ret = self._lib.axiom_unified_step(
                    handle,
                    params_ptr,
                    grads_ptr,
                    group['lr'],
                    beta1,
                    beta2,
                    group['eps'],
                    self._step_count
                )

                if ret != 0:
                    raise RuntimeError(f"AXIOM step failed with error {ret}")

        return loss

    def get_gamma(self) -> float:
        """Get current adaptive gamma value."""
        if self._handles:
            handle = next(iter(self._handles.values()))
            return self._lib.axiom_unified_get_gamma(handle)
        return self.gamma_init

    def set_gamma(self, gamma: float):
        """Manually set gamma for all parameters."""
        for handle in self._handles.values():
            self._lib.axiom_unified_set_gamma(handle, gamma)

    def memory_info(self) -> dict:
        """Get memory usage statistics."""
        total_bytes = 0
        total_params = 0

        for handle in self._handles.values():
            bytes_val = c_size_t()
            bpp = c_float()
            savings = c_float()
            self._lib.axiom_unified_memory_info(handle, byref(bytes_val), byref(bpp), byref(savings))
            total_bytes += bytes_val.value

        for group in self.param_groups:
            for p in group['params']:
                total_params += p.numel()

        adam_bytes = total_params * 8  # FP32 m + FP32 v

        return {
            'axiom_bytes': total_bytes,
            'adam_bytes': adam_bytes,
            'bytes_per_param': total_bytes / max(total_params, 1),
            'savings_percent': (1 - total_bytes / max(adam_bytes, 1)) * 100,
            'total_params': total_params,
        }

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

    def __del__(self):
        """Clean up CUDA handles."""
        if hasattr(self, '_handles'):
            for handle in self._handles.values():
                try:
                    self._lib.axiom_unified_destroy(handle)
                except Exception:
                    pass

    @staticmethod
    def version() -> str:
        """Get AXIOM version string."""
        lib = _get_lib()
        return lib.axiom_unified_version().decode('utf-8')


def get_version() -> str:
    """Get AXIOM Unified version."""
    return _get_lib().axiom_unified_version().decode('utf-8')
