"""
QuarterBit Optimizers
=====================

Precision optimizers for PyTorch.
Copyright (c) 2026 Clouthier Simulation Labs. All rights reserved.
"""

import torch
from torch.optim import Optimizer
from typing import List, Optional, Callable, Iterable
import ctypes
from .utils import get_lib, is_available

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))


class SGD(Optimizer):
    """
    EFT-powered Stochastic Gradient Descent.

    Drop-in replacement for torch.optim.SGD with precise gradient accumulation.

    Args:
        params: Iterable of parameters to optimize
        lr: Learning rate
        momentum: Momentum factor (default: 0)
        dampening: Dampening for momentum (default: 0)
        weight_decay: Weight decay (L2 penalty) (default: 0)
        nesterov: Enables Nesterov momentum (default: False)

    Example:
        >>> from quarterbit.torch import SGD
        >>> optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)
    """

    def __init__(
        self,
        params: Iterable[torch.nn.Parameter],
        lr: float = 1e-3,
        momentum: float = 0,
        dampening: float = 0,
        weight_decay: float = 0,
        nesterov: bool = False
    ):
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if momentum < 0.0:
            raise ValueError(f"Invalid momentum value: {momentum}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
        if nesterov and (momentum <= 0 or dampening != 0):
            raise ValueError("Nesterov momentum requires a momentum and zero dampening")

        defaults = dict(
            lr=lr, momentum=momentum, dampening=dampening,
            weight_decay=weight_decay, nesterov=nesterov
        )
        super().__init__(params, defaults)

        self._use_eft = is_available()
        if self._use_eft:
            self._lib = get_lib()

        # Initialize compensation buffers
        for group in self.param_groups:
            for p in group['params']:
                state = self.state[p]
                state['weight_comp'] = torch.zeros_like(p.data)
                if momentum != 0:
                    state['momentum_buffer'] = torch.zeros_like(p.data)
                    state['momentum_comp'] = torch.zeros_like(p.data)

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

        for group in self.param_groups:
            lr = group['lr']
            momentum = group['momentum']
            dampening = group['dampening']
            weight_decay = group['weight_decay']
            nesterov = group['nesterov']

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

                grad = p.grad.data
                state = self.state[p]

                if self._use_eft and p.is_cuda and p.dtype == torch.float32:
                    # Use EFT kernel
                    n = p.numel()

                    if momentum != 0:
                        # TODO: Implement momentum with EFT
                        # For now, fall back to standard for momentum
                        self._standard_step(p, grad, state, group)
                    else:
                        # Pure SGD with EFT
                        self._lib.eft_sgd_step(
                            _ptr(p.data),
                            _ptr(state['weight_comp']),
                            _ptr(grad),
                            ctypes.c_float(lr),
                            ctypes.c_float(weight_decay),
                            ctypes.c_int(n)
                        )
                else:
                    # Fallback to standard PyTorch
                    self._standard_step(p, grad, state, group)

        return loss

    def _standard_step(self, p, grad, state, group):
        """Standard PyTorch SGD step (fallback)."""
        weight_decay = group['weight_decay']
        momentum = group['momentum']
        dampening = group['dampening']
        nesterov = group['nesterov']
        lr = group['lr']

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

        if momentum != 0:
            buf = state.get('momentum_buffer')
            if buf is None:
                buf = torch.clone(grad).detach()
                state['momentum_buffer'] = buf
            else:
                buf.mul_(momentum).add_(grad, alpha=1 - dampening)

            if nesterov:
                grad = grad.add(buf, alpha=momentum)
            else:
                grad = buf

        p.data.add_(grad, alpha=-lr)


class Adam(Optimizer):
    """
    EFT-powered Adam optimizer.

    Drop-in replacement for torch.optim.Adam with precise moment accumulation.

    The key insight: Adam's moment estimates (m and v) accumulate small updates
    over millions of steps. Standard FP32 loses precision. EFT tracks the
    exact error, ensuring stable training.

    Args:
        params: Iterable of parameters to optimize
        lr: Learning rate (default: 1e-3)
        betas: Coefficients for computing running averages (default: (0.9, 0.999))
        eps: Term added to denominator for numerical stability (default: 1e-8)
        weight_decay: Weight decay (L2 penalty) (default: 0)
        amsgrad: Whether to use AMSGrad variant (default: False)

    Example:
        >>> from quarterbit.torch import Adam
        >>> optimizer = Adam(model.parameters(), lr=1e-3)
    """

    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,
        amsgrad: bool = False
    ):
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if eps < 0.0:
            raise ValueError(f"Invalid epsilon value: {eps}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

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

        self._use_eft = is_available()
        if self._use_eft:
            self._lib = get_lib()

    def _init_state(self, p):
        """Initialize state for a parameter."""
        state = self.state[p]
        state['step'] = 0
        state['exp_avg'] = torch.zeros_like(p.data)
        state['exp_avg_sq'] = torch.zeros_like(p.data)

        # EFT compensation buffers
        state['weight_comp'] = torch.zeros_like(p.data)
        state['exp_avg_comp'] = torch.zeros_like(p.data)
        state['exp_avg_sq_comp'] = torch.zeros_like(p.data)

        return state

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

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

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

                grad = p.grad.data
                if grad.is_sparse:
                    raise RuntimeError("Adam does not support sparse gradients")

                # Initialize state if needed
                state = self.state[p]
                if len(state) == 0:
                    state = self._init_state(p)

                state['step'] += 1
                step = state['step']

                if self._use_eft and p.is_cuda and p.dtype == torch.float32 and not amsgrad:
                    # Use EFT kernel
                    n = p.numel()

                    self._lib.eft_adam_step(
                        _ptr(p.data),
                        _ptr(state['weight_comp']),
                        _ptr(grad),
                        _ptr(state['exp_avg']),
                        _ptr(state['exp_avg_sq']),
                        _ptr(state['exp_avg_comp']),
                        _ptr(state['exp_avg_sq_comp']),
                        ctypes.c_float(lr),
                        ctypes.c_float(beta1),
                        ctypes.c_float(beta2),
                        ctypes.c_float(eps),
                        ctypes.c_float(weight_decay),
                        ctypes.c_int(step),
                        ctypes.c_int(n)
                    )
                else:
                    # Fallback to standard PyTorch Adam
                    self._standard_step(p, grad, state, group)

        return loss

    def _standard_step(self, p, grad, state, group):
        """Standard PyTorch Adam step (fallback)."""
        beta1, beta2 = group['betas']
        lr = group['lr']
        eps = group['eps']
        weight_decay = group['weight_decay']
        amsgrad = group['amsgrad']
        step = state['step']

        exp_avg = state['exp_avg']
        exp_avg_sq = state['exp_avg_sq']

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

        # Decay the first and second moment running average coefficient
        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

        if amsgrad:
            max_exp_avg_sq = state.get('max_exp_avg_sq')
            if max_exp_avg_sq is None:
                max_exp_avg_sq = torch.zeros_like(p.data)
                state['max_exp_avg_sq'] = max_exp_avg_sq
            torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
            denom = (max_exp_avg_sq.sqrt() / (1 - beta2 ** step) ** 0.5).add_(eps)
        else:
            bias_correction1 = 1 - beta1 ** step
            bias_correction2 = 1 - beta2 ** step
            denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)

        step_size = lr / (1 - beta1 ** step) if not amsgrad else lr / (1 - beta1 ** step)
        p.data.addcdiv_(exp_avg, denom, value=-step_size)


class AdamW(Adam):
    """
    EFT-powered AdamW optimizer (Adam with decoupled weight decay).

    Same as Adam but with proper weight decay (not L2 regularization).

    Example:
        >>> from quarterbit.torch import AdamW
        >>> optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
    """

    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 = 1e-2,
        amsgrad: bool = False
    ):
        super().__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)

    def _standard_step(self, p, grad, state, group):
        """AdamW step - weight decay applied directly to weights."""
        beta1, beta2 = group['betas']
        lr = group['lr']
        eps = group['eps']
        weight_decay = group['weight_decay']
        amsgrad = group['amsgrad']
        step = state['step']

        exp_avg = state['exp_avg']
        exp_avg_sq = state['exp_avg_sq']

        # Decoupled weight decay (applied to weights, not gradients)
        if weight_decay != 0:
            p.data.mul_(1 - lr * weight_decay)

        # Standard Adam update
        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

        if amsgrad:
            max_exp_avg_sq = state.get('max_exp_avg_sq')
            if max_exp_avg_sq is None:
                max_exp_avg_sq = torch.zeros_like(p.data)
                state['max_exp_avg_sq'] = max_exp_avg_sq
            torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
            denom = (max_exp_avg_sq.sqrt() / (1 - beta2 ** step) ** 0.5).add_(eps)
        else:
            bias_correction1 = 1 - beta1 ** step
            bias_correction2 = 1 - beta2 ** step
            denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)

        step_size = lr / (1 - beta1 ** step)
        p.data.addcdiv_(exp_avg, denom, value=-step_size)


class CompactAdam(Optimizer):
    """
    Memory-efficient Adam optimizer with compressed state storage.

    Uses FP16+FP4 split weights and FP4+INT2 compressed optimizer states
    for ~1.7x memory savings vs standard FP32 Adam, while maintaining
    training quality through error feedback.

    Memory budget per parameter:
    - Standard FP32 Adam: 16 bytes/param
    - CompactAdam (no EF): 9.25 bytes/param (1.7x savings)
    - CompactAdam (with EF): 10.75 bytes/param (1.5x savings)

    Args:
        params: Iterable of parameters to optimize
        lr: Learning rate (default: 1e-3)
        betas: Coefficients for computing running averages (default: (0.9, 0.999))
        eps: Term added to denominator for numerical stability (default: 1e-8)
        weight_decay: Weight decay (L2 penalty) (default: 0)
        use_error_feedback: Enable error feedback for better precision (default: True)

    Example:
        >>> from quarterbit.torch import CompactAdam
        >>> optimizer = CompactAdam(model.parameters(), lr=1e-3)
        >>> # For maximum memory savings (slight precision tradeoff):
        >>> optimizer = CompactAdam(model.parameters(), lr=1e-3, use_error_feedback=False)
    """

    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,
        use_error_feedback: bool = True
    ):
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if eps < 0.0:
            raise ValueError(f"Invalid epsilon value: {eps}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

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

        self._use_ef = 1 if use_error_feedback else 0
        self._handles = {}  # Per-parameter handles
        self._step_count = 0

        # Try to load compact library
        self._lib = None
        self._use_compact = False
        try:
            from .utils import get_lib
            lib = get_lib()
            # Check if compact functions exist
            if hasattr(lib, 'compact_adam_create'):
                self._lib = lib
                self._use_compact = True
        except Exception:
            pass

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

    @torch.no_grad()
    def step(self, closure: Optional[Callable] = None):
        """Performs 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("CompactAdam does not support sparse gradients")

                handle = self._get_handle(p)

                if handle is not None and self._use_compact:
                    # Use compact CUDA kernel
                    n = p.numel()

                    # Apply weight decay to gradients (L2 regularization)
                    if weight_decay != 0:
                        grad = grad.add(p.data, alpha=weight_decay)

                    self._lib.compact_adam_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:
                    # Fallback to standard Adam
                    self._standard_step(p, grad, group)

        return loss

    def _standard_step(self, p, grad, group):
        """Standard PyTorch Adam step (fallback for CPU or non-float32)."""
        beta1, beta2 = group['betas']
        lr = group['lr']
        eps = group['eps']
        weight_decay = group['weight_decay']

        state = self.state[p]
        if len(state) == 0:
            state['step'] = 0
            state['exp_avg'] = torch.zeros_like(p.data)
            state['exp_avg_sq'] = torch.zeros_like(p.data)

        state['step'] += 1
        step = state['step']

        exp_avg = state['exp_avg']
        exp_avg_sq = state['exp_avg_sq']

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

        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

        bias_correction1 = 1 - beta1 ** step
        bias_correction2 = 1 - beta2 ** step
        denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)

        step_size = lr / bias_correction1
        p.data.addcdiv_(exp_avg, denom, value=-step_size)

    def __del__(self):
        """Clean up handles on deletion."""
        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.compact_adam_destroy(handle)
                    except Exception:
                        pass

    def memory_savings(self):
        """Report memory savings vs standard FP32 Adam."""
        total_params = sum(p.numel() for group in self.param_groups for p in group['params'])
        fp32_bytes = total_params * 16  # 4 bytes each for param, grad, m, v
        compact_bytes = total_params * (10.75 if self._use_ef else 9.25)
        savings = 1 - compact_bytes / fp32_bytes
        return {
            'total_params': total_params,
            'fp32_bytes': fp32_bytes,
            'compact_bytes': compact_bytes,
            'savings_ratio': fp32_bytes / compact_bytes,
            'savings_percent': savings * 100
        }


class CompactEFTAdam(Optimizer):
    """
    Production Adam optimizer: Memory-efficient + EFT precision.

    Combines the best of both worlds:
    - Compressed storage: FP16+FP4 weights, FP4+INT2 states (27% memory savings)
    - EFT precision: TwoSum/TwoProduct for exact arithmetic in Adam updates

    Memory budget per parameter:
    - Standard FP32 Adam: 16 bytes/param
    - CompactEFTAdam: ~13.25 bytes/param (17% savings with full precision)

    This is the RECOMMENDED optimizer for production training.

    Args:
        params: Iterable of parameters to optimize
        lr: Learning rate (default: 1e-3)
        betas: Coefficients for computing running averages (default: (0.9, 0.999))
        eps: Term added to denominator for numerical stability (default: 1e-8)
        weight_decay: Weight decay (L2 penalty) (default: 0)

    Example:
        >>> from quarterbit.torch import CompactEFTAdam
        >>> optimizer = CompactEFTAdam(model.parameters(), lr=1e-3)
    """

    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,
    ):
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr}")
        if eps < 0.0:
            raise ValueError(f"Invalid epsilon value: {eps}")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
        if weight_decay < 0.0:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

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

        self._handles = {}
        self._step_count = 0

        # Try to load CompactEFT library
        self._lib = None
        self._use_compact_eft = False
        try:
            from .utils import _load_compact_eft_lib
            self._lib = _load_compact_eft_lib()
            if self._lib is not None:
                self._use_compact_eft = True
        except Exception:
            pass

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

    @torch.no_grad()
    def step(self, closure: Optional[Callable] = None):
        """Performs 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("CompactEFTAdam does not support sparse gradients")

                handle = self._get_handle(p)

                if handle is not None and self._use_compact_eft:
                    # Apply weight decay to gradients (L2 regularization)
                    if weight_decay != 0:
                        grad = grad.add(p.data, alpha=weight_decay)

                    self._lib.compact_eft_adam_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:
                    # Fallback to standard Adam
                    self._standard_step(p, grad, group)

        return loss

    def _standard_step(self, p, grad, group):
        """Standard PyTorch Adam step (fallback for CPU or non-float32)."""
        beta1, beta2 = group['betas']
        lr = group['lr']
        eps = group['eps']
        weight_decay = group['weight_decay']

        state = self.state[p]
        if len(state) == 0:
            state['step'] = 0
            state['exp_avg'] = torch.zeros_like(p.data)
            state['exp_avg_sq'] = torch.zeros_like(p.data)

        state['step'] += 1
        step = state['step']

        exp_avg = state['exp_avg']
        exp_avg_sq = state['exp_avg_sq']

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

        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

        bias_correction1 = 1 - beta1 ** step
        bias_correction2 = 1 - beta2 ** step
        denom = (exp_avg_sq.sqrt() / (bias_correction2 ** 0.5)).add_(eps)

        step_size = lr / bias_correction1
        p.data.addcdiv_(exp_avg, denom, value=-step_size)

    def __del__(self):
        """Clean up handles on deletion."""
        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.compact_eft_adam_destroy(handle)
                    except Exception:
                        pass

    def memory_savings(self):
        """Report memory savings vs standard FP32 Adam."""
        total_params = sum(p.numel() for group in self.param_groups for p in group['params'])
        fp32_bytes = total_params * 16
        # Weights: 2.75, states: 2.5, EFT comp: 4, grads: 4 = 13.25 B/param
        compact_eft_bytes = total_params * 13.25
        savings = 1 - compact_eft_bytes / fp32_bytes
        return {
            'total_params': total_params,
            'fp32_bytes': fp32_bytes,
            'compact_eft_bytes': compact_eft_bytes,
            'savings_ratio': fp32_bytes / compact_eft_bytes,
            'savings_percent': savings * 100
        }
