import math
from typing import Tuple

import torch
from torch.optim.optimizer import Optimizer

from typing import Any, Dict, Iterable, Union
# try:
#     from typing import TypeAlias, TypeVar
# except ImportError:
#     from typing_extensions import TypeAlias, TypeVar

import torch
from torch.optim.optimizer import Optimizer, required
from typing import Iterable, Optional, Callable, Dict, Any

import torch
import torch.optim

# try:
#     from torch.optim.optimizer import ParamsT
# except (ImportError, TypeError):
ParamsT = Union[Iterable[torch.Tensor], Iterable[Dict[str, Any]]]



class CAdamW(Optimizer):
    r"""Implements CAdamW algorithm.


    References:
        - Adam: A Method for Stochastic Optimization: https://arxiv.org/abs/1412.6980
        - Decoupled Weight Decay Regularization: https://arxiv.org/abs/1711.05101
        - On the Convergence of Adam and Beyond: https://openreview.net/forum?id=ryQu7f-RZ

    Args:
        params: iterable of parameters to optimize or dicts defining parameter groups
        lr: learning rate
        betas: coefficients used for computing running averages of gradient and its square
        eps: term added to the denominator to improve numerical stability
        weight_decay: weight decay coefficient
        amsgrad: whether to use the AMSGrad variant of this algorithm
            from the paper `On the Convergence of Adam and Beyond`
        caution: apply caution when using AdamW
    """

    def __init__(
            self,
            params: ParamsT,
            lr: float = 1e-3,
            betas: Tuple[float, float] = (0.9, 0.999),
            eps: float = 1e-8,
            weight_decay: float = 1e-2,
            amsgrad: bool = False,
            caution: bool = False,
    ):
        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
        defaults = dict(
            lr=lr,
            betas=betas,
            eps=eps,
            weight_decay=weight_decay,
            amsgrad=amsgrad,
            caution=caution,
        )
        super(CAdamW, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(CAdamW, self).__setstate__(state)
        for group in self.param_groups:
            group.setdefault('amsgrad', False)
            group.setdefault('caution', False)

    @torch.no_grad()
    def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

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

                # Perform stepweight decay
                p.data.mul_(1 - group['lr'] * group['weight_decay'])

                # Perform optimization step
                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
                amsgrad = group['amsgrad']

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state['step'] = 0
                    # Exponential moving average of gradient values
                    state['exp_avg'] = torch.zeros_like(p)
                    # Exponential moving average of squared gradient values
                    state['exp_avg_sq'] = torch.zeros_like(p)
                    if amsgrad:
                        # Maintains max of all exp. moving avg. of sq. grad. values
                        state['max_exp_avg_sq'] = torch.zeros_like(p)

                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                if amsgrad:
                    max_exp_avg_sq = state['max_exp_avg_sq']
                beta1, beta2 = group['betas']

                state['step'] += 1
                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']

                # 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:
                    # Maintains the maximum of all 2nd moment running avg. till now
                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
                    # Use the max. for normalizing running avg. of gradient
                    denom = (max_exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])
                else:
                    denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(group['eps'])

                step_size = group['lr'] / bias_correction1

                if group['caution']:
                    # Apply caution as per 'Cautious Optimizers' - https://arxiv.org/abs/2411.16085
                    mask = (exp_avg * grad > 0).to(grad.dtype)
                    mask.div_(mask.mean().clamp_(min=1e-3))
                    exp_avg = exp_avg * mask

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

        return loss



class AdamWHCO(Optimizer):
    r"""Implements AdamW algorithm.

    NOTE: This impl has been deprecated in favour of torch.optim.NAdam and remains as a reference

    References:
        - Adam: A Method for Stochastic Optimization: https://arxiv.org/abs/1412.6980
        - Decoupled Weight Decay Regularization: https://arxiv.org/abs/1711.05101
        - On the Convergence of Adam and Beyond: https://openreview.net/forum?id=ryQu7f-RZ

    Args:
        params: iterable of parameters to optimize or dicts defining parameter groups
        lr: learning rate
        betas: coefficients used for computing running averages of gradient and its square
        eps: term added to the denominator to improve numerical stability
        weight_decay: weight decay coefficient
        amsgrad: whether to use the AMSGrad variant of this algorithm
            from the paper `On the Convergence of Adam and Beyond`
        caution: apply caution when using AdamW
    """

    def __init__(
            self,
            params: ParamsT,
            lr: float = 1e-3,
            betas: Tuple[float, float] = (0.9, 0.999),
            eps: float = 1e-8,
            weight_decay: float = 1e-2,
            amsgrad: bool = False,
            caution: bool = False,
            hirearchical_caution: bool = True,
            unlabeled_momentum_update: bool = False,
    ):
        if not 0.0 <= lr:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if not 0.0 <= eps:
            raise ValueError("Invalid epsilon value: {}".format(eps))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
        defaults = dict(
            lr=lr,
            betas=betas,
            eps=eps,
            weight_decay=weight_decay,
            amsgrad=amsgrad,
            caution=caution,
            hirearchical_caution=hirearchical_caution,
            unlabeled_momentum_update=unlabeled_momentum_update,
        )
        super(AdamWHCO, self).__init__(params, defaults)
        self.accumulated_updates = {}
    def __setstate__(self, state):
        super(AdamWHCO, self).__setstate__(state)
        for group in self.param_groups:
            group.setdefault('amsgrad', False)
            group.setdefault('caution', False)
            group.setdefault('hirearchical_caution', True)
  
    
    @torch.no_grad()
    def step_labeled(self, closure=None):
        """Computes and stores parameter updates from labeled data without applying them."""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

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

                # Calculate weight decay factor
                # weight_decay_factor = 1 - group['lr'] * group['weight_decay']
                
                # Process gradients as before
                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError('Adam does not support sparse gradients')
                amsgrad = group['amsgrad']

                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state['step'] = 0
                    state['exp_avg'] = torch.zeros_like(p)
                    state['exp_avg_sq'] = torch.zeros_like(p)
                    if amsgrad:
                        state['max_exp_avg_sq'] = torch.zeros_like(p)

                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                if amsgrad:
                    max_exp_avg_sq = state['max_exp_avg_sq']
                beta1, beta2 = group['betas']

                state['step'] += 1
                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']

                # Update momentum state
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
                
                # Calculate update (but don't apply it)
                if amsgrad:
                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
                    denom = (max_exp_avg_sq.sqrt() / torch.sqrt(torch.tensor(bias_correction2, device=p.device))).add_(group['eps'])
                else:
                    denom = (exp_avg_sq.sqrt() / torch.sqrt(torch.tensor(bias_correction2, device=p.device))).add_(group['eps'])

                step_size = group['lr'] / bias_correction1
                
                if group['caution']:
                    mask = (exp_avg * grad > 0).to(grad.dtype)
                    mask.div_(mask.mean().clamp_(min=1e-3))
                    exp_avg_adjusted = exp_avg * mask
                else:
                    exp_avg_adjusted = exp_avg
                
                # Store updates instead of applying them
                update = -step_size * (exp_avg_adjusted / denom)
                self.accumulated_updates[p] = {
                    'update': update,
                    # 'weight_decay_factor': weight_decay_factor
                }

        return loss

    @torch.no_grad()
    def step_unlabeled(self, closure=None):
        """Computes updates from unlabeled data and applies both labeled and unlabeled updates."""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

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

                # Perform stepweight decay
                p.data.mul_(1 - group['lr'] * group['weight_decay'])
                
                # Process gradients
                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError('Adam does not support sparse gradients')
                amsgrad = group['amsgrad']

                state = self.state[p]

                # State initialization should already be done in step_labeled
                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                if amsgrad:
                    max_exp_avg_sq = state['max_exp_avg_sq']
                beta1, beta2 = group['betas']

                # Optionally update momentum for unlabeled data
                if group.get('unlabeled_momentum_update', False):
                    # Update only where the current momentum agrees with the gradient.
                    # Allow update if exp_avg is zero (no momentum yet)
                    agreement_mask = ((exp_avg == 0) | (exp_avg * grad > 0)).to(grad.dtype)
                    # Update exp_avg: scale current momentum then add the masked update
                    exp_avg.mul_(beta1)
                    update = (1 - beta1) * (grad * agreement_mask)
                    exp_avg.add_(update)
                    # Update exp_avg_sq similarly
                    exp_avg_sq.mul_(beta2)
                    update_sq = (1 - beta2) * (grad * grad * agreement_mask)
                    exp_avg_sq.add_(update_sq)
                # If unlabeled_momentum_update is False, leave exp_avg and exp_avg_sq unchanged

                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']

                # Calculate unlabeled update
                if amsgrad:
                    denom = (max_exp_avg_sq.sqrt() / torch.sqrt(torch.tensor(bias_correction2, device=p.device))).add_(group['eps'])
                else:
                    denom = (exp_avg_sq.sqrt() / torch.sqrt(torch.tensor(bias_correction2, device=p.device))).add_(group['eps'])

                step_size = group['lr'] / bias_correction1

                if group['hirearchical_caution']:
                    mask = (exp_avg * grad > 0).to(grad.dtype)
                    mask.div_(mask.mean().clamp_(min=1e-3))
                    exp_avg_adjusted = exp_avg * mask
                else:
                    exp_avg_adjusted = exp_avg

                # Calculate unlabeled update
                unlabeled_update = -step_size * (exp_avg_adjusted / denom)
                
                # Now apply both updates
                labeled_data = self.accumulated_updates.get(p, {'update': 0, 'weight_decay_factor': 1})
                        
                # Apply parameter updates
                p.data.add_(labeled_data['update'] + unlabeled_update)
                p.data = p.data.contiguous()
                
                # Clear accumulated updates
                if p in self.accumulated_updates:
                    del self.accumulated_updates[p]

        return loss

# generate class with weight decay 0
class AdamHCO(AdamWHCO):
    def __init__(self,
            params: ParamsT,
            lr: float = 1e-3,
            betas: Tuple[float, float] = (0.9, 0.999),
            eps: float = 1e-8,
            weight_decay: float = 0.0,
            amsgrad: bool = False,
            caution: bool = False,
            hirearchical_caution: bool = True,
            unlabeled_momentum_update: bool = False,
            ):
        super(AdamHCO, self).__init__(
            params,
            lr=lr,
            betas=betas,
            eps=eps,
            weight_decay=weight_decay,
            amsgrad=amsgrad,
            caution=caution,
            hirearchical_caution=hirearchical_caution,
            unlabeled_momentum_update=unlabeled_momentum_update,
        )




class HcoSGD(Optimizer):
    r"""
    Hierarchical-Cautious SGD with an *optional* “aligned-only” momentum
    update during the unlabeled phase.

    Typical loop:
        optimizer.zero_grad()
        loss_L.backward()
        optimizer.step_labeled()
        loss_U.backward()              # **do not** zero grads again
        optimizer.step_unlabeled()     # may update momentum if flag is on
    """

    def __init__(
        self,
        params: Iterable[torch.nn.Parameter],
        lr: float = required,
        lr_unlabeled: Optional[float] = None,
        momentum: float = 0.9,
        weight_decay: float = 0.0,
        hierarchical_caution: bool = True,
        unlabeled_momentum_update: bool = False,
        eps: float = 1e-8,
    ):
        if lr < 0 or (lr_unlabeled is not None and lr_unlabeled < 0):
            raise ValueError("Invalid learning rate")
        if not 0.0 <= momentum < 1.0:
            raise ValueError("Momentum must be in [0,1)")
        if weight_decay < 0:
            raise ValueError("Invalid weight_decay")

        defaults = dict(
            lr=lr,
            lr_unlabeled=lr_unlabeled if lr_unlabeled is not None else lr,
            momentum=momentum,
            weight_decay=weight_decay,
            hierarchical_caution=hierarchical_caution,
            unlabeled_momentum_update=unlabeled_momentum_update,
            eps=eps,
        )
        super().__init__(params, defaults)

    # -------------------------------------------------- #
    # 1) Labeled step – vanilla SGD + momentum + L2-WD   #
    # -------------------------------------------------- #
    @torch.no_grad()
    def step_labeled(self, closure: Optional[Callable[[], float]] = None):
        loss = closure() if closure is not None else None

        for group in self.param_groups:
            lr, beta, wd = group["lr"], group["momentum"], group["weight_decay"]

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

                # coupled weight-decay (like plain SGD)
                if wd:
                    grad = grad.add(p.data, alpha=wd)

                state: Dict[str, Any] = self.state[p]
                buf = state.setdefault("momentum_buffer", torch.zeros_like(p.data))

                buf.mul_(beta).add_(grad)            # m_t = β m + g^{(L)}
                p.data.add_(buf, alpha=-lr)          # SGD update

                state["trusted_direction"] = buf.clone()

        return loss

    # -------------------------------------------------- #
    # 2) Unlabeled step – cautious correction            #
    #    + (optional) aligned-only momentum update       #
    # -------------------------------------------------- #
    @torch.no_grad()
    def step_unlabeled(self, closure: Optional[Callable[[], float]] = None):
        loss = closure() if closure is not None else None

        for group in self.param_groups:
            lr_u   = group["lr_unlabeled"]
            hc     = group["hierarchical_caution"]
            updMom = group["unlabeled_momentum_update"]
            beta   = group["momentum"]
            eps    = group["eps"]

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

                m_trusted: torch.Tensor | None = state.get("trusted_direction")
                if m_trusted is None:
                    continue  # no labeled step yet

                # --- sign-agreement mask φ ---
                mask = (m_trusted * grad > 0).to(grad.dtype)
                if hc:
                    mask.div_(mask.mean().clamp_(min=eps))  # keep magnitude ≈ const

                # --- optional momentum update with aligned grads ---
                if updMom:
                    buf = state["momentum_buffer"]
                    # add only the aligned part of grad to momentum
                    aligned_grad = grad * mask
                    buf.mul_(beta).add_(aligned_grad)       # m ← β m + g_aligned
                    m_trusted = buf                         # use the updated momentum
                    state["trusted_direction"] = buf.clone()

                # --- cautious correction step ---
                correction = m_trusted * mask
                p.data.add_(correction, alpha=-lr_u)

        return loss


    # block ambiguous call
    def step(self, *_, **__):
        raise RuntimeError(
            "Call step_labeled() then step_unlabeled() explicitly; "
            "plain step() would mix the two."
        )


if __name__ == '__main__':
    # Example usage:
    # Create some parameters
    params = [torch.randn(2, 2, requires_grad=True)]
    # Create an optimizer
    optimizer = AdamWHCO(params)
    # Perform a step of optimization
    optimizer.step_labeled()
