# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/optimizer.lion.ipynb.

# %% ../../nbs/optimizer.lion.ipynb 1
# Lion implementation based on the paper's code release
# https://github.com/google/automl/tree/master/lion - Apache License 2.0 - Copyright 2020 Google Research

# %% ../../nbs/optimizer.lion.ipynb 4
from __future__ import annotations

import numpy as np

from fastai.optimizer import Optimizer

from .foreach import ForEachOptimizer

try:
    from fastxtend.optimizer.eightbit import Lion8bitOptimizer
    EIGHTBIT = True
except ImportError:
    EIGHTBIT = False

from ..imports import *

# %% auto 0
__all__ = ['Lion', 'lion']

# %% ../../nbs/optimizer.lion.ipynb 6
def lion_step(p:Tensor, lr:float, wd:float, beta1:float, beta2:float,
              grad_avg:Tensor|None=None, do_wd:bool=True, **kwargs):

    if grad_avg is None:
        grad_avg = torch.zeros_like(p, memory_format=torch.preserve_format)

    # weight decay
    if do_wd and wd != 0:
        p.data.mul_(1-lr*wd)

    # lion step
    update = grad_avg.mul(beta1) + p.grad.data.mul(1 - beta1)
    p.data.add_(torch.sign(update), alpha=-lr)

    # Update m_k
    grad_avg.mul_(beta2).add_(p.grad.data, alpha=1-beta2)

    return {'grad_avg': grad_avg}

lion_step.defaults = dict(beta1=0.9, beta2=0.99)

# %% ../../nbs/optimizer.lion.ipynb 8
def lion_foreach_step(p:list[Tensor], g:list[Tensor], grad_avg:list[Tensor],
                      do_wd:np.ndarray[Any, bool], lr:float, wd:float, beta1:float,
                      beta2:float, **kwargs):

    # weight_decay
    if wd != 0:
        wd = np.where(do_wd, 1-lr*wd, 1.)
        torch._foreach_mul_(p, scalars=wd.tolist())

    # lion update step
    update = torch._foreach_mul(grad_avg, scalar=beta1)
    torch._foreach_add_(update, g, alpha=1-beta1)
    for u in update: u.sign_()
    torch._foreach_add_(p, update, alpha=-lr)

    # update m_k
    torch._foreach_mul_(grad_avg, scalar=beta2)
    torch._foreach_add_(grad_avg, g, alpha=1-beta2)

# %% ../../nbs/optimizer.lion.ipynb 9
class LionForEachOptimizer(ForEachOptimizer):
    "An `Optimizer` with a modified step for Lion ForEach"
    def __init__(self,
        params:Listified[Tensor], # Model parameters
        opt_step:Callable, # `ForEachOptimizer` optimizer step
        **defaults # Optimizer specific hyper parameters
    ):
        super().__init__(params, opt_step, **defaults)

    @torch.no_grad()
    def step(self, closure=None):
        if closure is not None: raise NotImplementedError("fastai optimizers currently do not support closure")
        for pg, hyper in zip(self.param_lists, self.hypers):
            pl, gl, grad_avg, do_wd = [], [], [], []

            for p in pg:
                if p.grad is not None:
                    state = self.state[p]

                    if 'grad_avg' not in state:
                        state['grad_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)

                    pl.append(p)
                    gl.append(p.grad)
                    grad_avg.append(state['grad_avg'])
                    do_wd.append(state.get('do_wd', True))

            self.opt_step(p=pl, g=gl, grad_avg=grad_avg, do_wd=np.array(do_wd, dtype=bool), **hyper)

# %% ../../nbs/optimizer.lion.ipynb 11
def Lion(
    params:Listified[Tensor], # Model parameters or parameter groups
    lr:float, # Default learning rate
    beta1:float=0.9, # Update gradient moving average (β1) coefficient
    beta2:float=0.99, # Gradient moving average (β2) coefficient
    wd:float=0.1, # True weight decay
    foreach:bool=False, # Use fused ForEach implementation
    eightbit:bool=False, # Use fused 8-bit implementation
    **eightbitargs
) -> Optimizer|LionForEachOptimizer|Lion8bitOptimizer:
    "A fastai Lion optimizer with fused ForEach and 8-bit implementations"
    if foreach:
        return LionForEachOptimizer(params, lion_foreach_step, lr=lr,
                                    beta1=beta1, beta2=beta2, wd=wd)
    elif eightbit:
        if EIGHTBIT:
            return Lion8bitOptimizer(params, lr=lr, beta1=beta1,
                                     beta2=beta2, wd=wd, **eightbitargs)
        else:
            raise ImportError(f'{eightbit=}. bitsandbytes package not found. Run `pip install bitsandbytes`')
    else:
        return Optimizer(params, [lion_step], lr=lr, beta1=beta1, beta2=beta2, wd=wd)

# %% ../../nbs/optimizer.lion.ipynb 12
def lion(
    beta1:float=0.9, # Update gradient moving average (β1) coefficient
    beta2:float=0.99, # Gradient moving average (β2) coefficient
    wd:float=0.1, # True weight decay
    foreach:bool=False, # Use fused ForEach implementation
    eightbit:bool=False, # Use fused 8-bit implementation
    **eightbitargs
) -> Optimizer|LionForEachOptimizer|Lion8bitOptimizer:
    "Partial function for the Lion optimizer with fused ForEach and 8-bit implementations"
    return partialler(Lion, beta1=beta1, beta2=beta2, wd=wd, foreach=foreach,
                      eightbit=eightbit, **eightbitargs)
