# copy dependencies from transformers/optimization.py
import math
import warnings
from typing import Callable, Iterable, Tuple

import torch
from torch import nn
from torch.optim import Optimizer

from transformers.utils.versions import require_version

from .galore_projector import GaLoreProjector

class Adam_mini(Optimizer):
    """
    Implements Adam algorithm with weight decay fix as introduced in [Decoupled Weight Decay
    Regularization](https://arxiv.org/abs/1711.05101).

    Parameters:
        params (`Iterable[nn.parameter.Parameter]`):
            Iterable of parameters to optimize or dictionaries defining parameter groups.
        lr (`float`, *optional*, defaults to 0.001):
            The learning rate to use.
        betas (`Tuple[float,float]`, *optional*, defaults to `(0.9, 0.999)`):
            Adam's betas parameters (b1, b2).
        eps (`float`, *optional*, defaults to 1e-06):
            Adam's epsilon for numerical stability.
        weight_decay (`float`, *optional*, defaults to 0.0):
            Decoupled weight decay to apply.
        correct_bias (`bool`, *optional*, defaults to `True`):
            Whether or not to correct bias in Adam (for instance, in Bert TF repository they use `False`).
        no_deprecation_warning (`bool`, *optional*, defaults to `False`):
            A flag used to disable the deprecation warning (set to `True` to disable the warning).
    """

    def __init__(
        self,
        params: Iterable[nn.parameter.Parameter],
        lr: float = 1e-3,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-6,
        weight_decay: float = 0.0,
        correct_bias: bool = True,
        no_deprecation_warning: bool = True,
    ):
        if not no_deprecation_warning:
            warnings.warn(
                "This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch"
                " implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this"
                " warning",
                FutureWarning,
            )
        require_version("torch>=1.5.0")  # add_ with alpha
        if lr < 0.0:
            raise ValueError(f"Invalid learning rate: {lr} - should be >= 0.0")
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError(f"Invalid beta parameter: {betas[0]} - should be in [0.0, 1.0)")
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError(f"Invalid beta parameter: {betas[1]} - should be in [0.0, 1.0)")
        if not 0.0 <= eps:
            raise ValueError(f"Invalid epsilon value: {eps} - should be >= 0.0")
        defaults = {"lr": lr, "betas": betas, "eps": eps, "weight_decay": weight_decay, "correct_bias": correct_bias}
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure: Callable = 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:
            loss = closure()

        for group in self.param_groups:
            for p in group["params"]:
                if p.grad is None:
                    continue
                grad = p.grad
                if grad.is_sparse:
                    raise RuntimeError("Adam does not support sparse gradients, please consider SparseAdam instead")

                state = self.state[p]
                
                if "step" not in state:
                    state["step"] = 0





                if 'dim' not in group:
                    group['dim'] = 2 # will only be used for those who has 'rank' in group

                beta1, beta2 = group["betas"]

                # State initialization of vmean
                if "rank" not in group:
                    #embedding and layernorm
                    if "exp_avg_sq" not in state:
                        # Exponential moving average of squared gradient values
                        state["exp_avg_sq"] = torch.zeros_like(grad)
                else:
                    # MLP
                    if "exp_avg_sq" not in state:
                        # Exponential moving average of squared gradient values
                        # state["exp_avg_sq"] = 0  #torch.zeros_like(grad)
                        state['exp_avg_sq'] = torch.zeros(grad.size(0),1).cuda()

                # update exp_avg_sq
                if "rank" not in group: # layernorm  and embedding use Adam
                    state["exp_avg_sq"] = state["exp_avg_sq"] * beta2 + (1 - beta2) * (grad * grad)

                else: 
                    # state["exp_avg_sq"] = state["exp_avg_sq"] * beta2 + (1 - beta2) * torch.mean(grad * grad)
                    #grad_view = grad.view(grad.size(1), -1)
                    #print('post shape, ', grad.size(), y.size(), z.size())
                    # tmp_lr = torch.mean(grad_view * grad_view, dim=1, keepdim=True)
                    tmp_lr = torch.mean(grad * grad, dim=1, keepdim=True)
                        #print('tmp lr', tmp_lr.size(), exp_avg_sq.size())
                    state["exp_avg_sq"].mul_(beta2).add_(tmp_lr, alpha=1-beta2)


                state["step"] += 1    

                # GaLore Projection
                if "rank" in group:
                    if "projector" not in state:
                        if group['dim'] <=2:
                            state["projector"] = GaLoreProjector(group["rank"], update_proj_gap=group["update_proj_gap"], scale=group["scale"], proj_type=group["proj_type"])
                        else:
                            state["projector"] = GaLoreProjectorTensor(group["rank"], update_proj_gap=group["update_proj_gap"], scale=group["scale"], proj_type=group["proj_type"])
                    
                    grad = state["projector"].project(grad, state["step"])

                # initialize m
                if "exp_avg" not in state:
                    # Exponential moving average of gradient values
                    state["exp_avg"] = torch.zeros_like(grad)
    
                # update m
                state["exp_avg"] = state["exp_avg"] * beta1 + (1 - beta1) * grad
                # exp_avg.mul_(beta1).add_(grad.view(exp_avg.size()), alpha=(1.0 - beta1))
                #exp_avg_view = exp_avg.view(grad.size())
                # exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                if "rank" in group:
                    exp_avg_proj_back = state["projector"].project_back(state["exp_avg"])
                else: 
                    exp_avg_proj_back = state["exp_avg"]

                #exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
                denom = state["exp_avg_sq"].sqrt().add_(group["eps"])

                step_size = group["lr"]
                if group["correct_bias"]:  # No bias correction for Bert
                    bias_correction1 = 1.0 - beta1 ** state["step"]
                    bias_correction2 = 1.0 - beta2 ** state["step"]
                    step_size = step_size * math.sqrt(bias_correction2) / bias_correction1

                # compute norm gradient
                # norm_grad = exp_avg / denom

                # if grad.dim()>1 and grad.size(0) == 768 and grad.size(1) == 768:
                #     exp_avg_proj_back = exp_avg_proj_back.view(12,-1)
                

                # if "rank" not in group:
                #     exp_avg_proj_back = exp_avg_proj_back 
                # else:
                #     exp_avg_proj_back = exp_avg_proj_back.view(grad.size(1),-1)
                
                # print('dimension check', grad.size(), exp_avg_proj_back.size(), denom.size())
                
                norm_grad = exp_avg_proj_back / denom
                # GaLore Projection Back
                # if "rank" in group:
                #     norm_grad = state["projector"].project_back(norm_grad)
                
                norm_grad = norm_grad.view(p.size())
                p.add_(norm_grad, alpha=-step_size)

                # Just adding the square of the weights to the loss function is *not*
                # the correct way of using L2 regularization/weight decay with Adam,
                # since that will interact with the m and v parameters in strange ways.
                #
                # Instead we want to decay the weights in a manner that doesn't interact
                # with the m/v parameters. This is equivalent to adding the square
                # of the weights to the loss with plain (non-momentum) SGD.
                # Add weight decay at the end (fixed version)
                if group["weight_decay"] > 0.0:
                    p.add_(p, alpha=(-group["lr"] * group["weight_decay"]))

        return loss
