from typing import Callable, Iterable, Optional, Tuple, Union

import torch
from torch.optim.optimizer import Optimizer

from ..kernels.adamw_kernel import TritonAdamWKernel


class FastAdamW(Optimizer):
    """Fast AdamW optimizer with automatic Triton/PyTorch backend selection.

    Implements AdamW algorithm with decoupled weight decay. Unlike Adam, AdamW applies
    weight decay directly to parameters before the gradient update, which has been shown
    to improve generalization. Automatically uses Triton kernels for CUDA tensors when
    available, falling back to PyTorch implementation otherwise.

    Args:
        params (iterable): Iterable of parameters to optimize or dicts defining parameter groups.
        lr (float, optional): Learning rate. Defaults to 1e-3.
        betas (Tuple[float, float], optional): Coefficients for computing running averages
            of gradient and its square. Defaults to (0.9, 0.999).
        eps (float, optional): Term added to denominator for numerical stability. Defaults to 1e-8.
        weight_decay (float, optional): Decoupled weight decay coefficient. Defaults to 1e-2.
        use_triton (bool, optional): Whether to enable Triton kernels for CUDA tensors.
            Defaults to True.

    Examples:
        >>> optimizer = FastAdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
        >>> optimizer.zero_grad()
        >>> loss.backward()
        >>> optimizer.step()

        >>> # AdamW is commonly used for training transformers
        >>> optimizer = FastAdamW(model.parameters(), lr=5e-5, weight_decay=0.01)
    """

    def __init__(
        self,
        params: Union[Iterable[torch.Tensor], Iterable[dict]],
        lr: float = 1e-3,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 1e-2,
        use_triton: bool = True,
    ):
        if not 0.0 <= lr:
            raise ValueError(f"Invalid learning rate: {lr}")
        if not 0.0 <= eps:
            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 not 0.0 <= weight_decay:
            raise ValueError(f"Invalid weight_decay value: {weight_decay}")

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

    def _check_triton_availability(self, param: torch.Tensor) -> bool:
        if not self.use_triton:
            return False
        if not TritonAdamWKernel.is_available():
            return False
        if not param.is_cuda:
            return False
        return True

    def _triton_step(
        self,
        param: torch.Tensor,
        grad: torch.Tensor,
        exp_avg: torch.Tensor,
        exp_avg_sq: torch.Tensor,
        lr: float,
        beta1: float,
        beta2: float,
        eps: float,
        weight_decay: float,
        step: int,
        grad_scale: float = 1.0,
    ):
        p_flat = param.view(-1)
        grad_flat = grad.view(-1)
        exp_avg_flat = exp_avg.view(-1)
        exp_avg_sq_flat = exp_avg_sq.view(-1)

        TritonAdamWKernel.apply(
            p_flat,
            grad_flat,
            exp_avg_flat,
            exp_avg_sq_flat,
            lr,
            beta1,
            beta2,
            eps,
            weight_decay,
            step,
            grad_scale,
        )

    def _pytorch_step(
        self,
        param: torch.Tensor,
        grad: torch.Tensor,
        exp_avg: torch.Tensor,
        exp_avg_sq: torch.Tensor,
        lr: float,
        beta1: float,
        beta2: float,
        eps: float,
        weight_decay: float,
        step: int,
        grad_scale: float = 1.0,
    ):
        if grad_scale != 1.0:
            grad = grad / grad_scale

        param.mul_(1 - lr * 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
        step_size = lr / bias_correction1

        denom = exp_avg_sq.sqrt().div_(bias_correction2**0.5).add_(eps)

        param.addcdiv_(exp_avg, denom, value=-step_size)

    @torch.no_grad()
    def step(
        self, closure: Optional[Callable] = None, grad_scaler: Optional[float] = None
    ):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        grad_scale = grad_scaler if grad_scaler is not None else 1.0

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

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

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

                state = self.state[p]

                if len(state) == 0:
                    state["step"] = 0
                    state["exp_avg"] = torch.zeros_like(
                        p, memory_format=torch.preserve_format
                    )
                    state["exp_avg_sq"] = torch.zeros_like(
                        p, memory_format=torch.preserve_format
                    )

                exp_avg = state["exp_avg"]
                exp_avg_sq = state["exp_avg_sq"]
                state["step"] += 1
                step = state["step"]

                if self._check_triton_availability(p):
                    self._triton_step(
                        p,
                        grad,
                        exp_avg,
                        exp_avg_sq,
                        lr,
                        beta1,
                        beta2,
                        eps,
                        weight_decay,
                        step,
                        grad_scale,
                    )
                else:
                    self._pytorch_step(
                        p,
                        grad,
                        exp_avg,
                        exp_avg_sq,
                        lr,
                        beta1,
                        beta2,
                        eps,
                        weight_decay,
                        step,
                        grad_scale,
                    )

        return loss
