# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F401
"""Provide abstraction for defining optimizers and a set of common optimizers."""

from decimal import Decimal
from typing import Optional, Union

import numpy as np  # type: ignore
import tvm_ffi

import tvm

from ..block_builder import BlockBuilder
from ..expr import Function, TupleGetItem, Var, const
from ..expr import Tuple as RxTuple
from ..op import add, divide, multiply, sqrt, subtract
from ..type import TensorType, TupleType


# TODO(chaofan, yixin): Migrate key logics to C++
class Optimizer:
    """Relax training optimizer. This class could generate relax Functions for optimizing specified
    parameters, and store the states used in the optimization process, such as momentum.

    Parameters
    ----------
    name : str
        The name of the optimizer function. This parameter is provided by subclasses.

    Attributes
    ----------
    dtype : str
        The only dtype of the optimizer. It will be used as the dtype of the optimizer states,
        and the dtype of necessary constants, such as the learning rate. Will be set in `init()`.

    name : str
        The name of the optimizer.

    param_list : List[Var]
        The list of variables to optimize. Will be set in `init()`.

    state : tvm_ffi.Array
        `state` is an runtime Array representing the state of the optimizer. Will be set in
        `init()`.

        The states of the optimizer can store necessary information in the optimization process at
        runtime, such as the number of steps, the momentum in momentum SGD, etc.

        `opt.state` should be used as the last argument of the function that is got through
        `get_function()`, and its new value is returned as the last return value of that function.

        See examples for more details.

    Examples
    --------
    The usage of optimizers should resemble the following pattern. We will take SGD as an example.
    For detailed examples, please see the tutorial.

    .. code-block:: python

        # Construct the optimizer
        opt = relax.optimizer.SGD(0.1)

        # Initialize the parameter list, the dtype and the optimizer state
        # x is the relax Var we want to optimize
        opt.init(x)

        # The above two lines is equivalent to one line:
        opt = relax.optimizer.SGD(0.1).init(x)

        # Get the optimizer function
        # mod is an IRModule constructed earlier
        mod["SGD"] = opt.get_function()

        # Legalize and build mod
        lowered_mod = LegalizeOps()(mod)
        ex = build(lowered_mod, target="llvm")
        vm = VirtualMachine(ex, tvm.cpu())

        # Optimization process
        # param_tuple is a runtime tuple of parameters
        # param_gradient is a runtime tuple of the gradient of the parameters in param_tuple,
        # respectively
        # param_gradient can be gained by the automatic differentiation pass. Please see
        # `relax.transform.Gradient`
        param_tuple, opt.state = vm["SGD"](param_tuple, param_gradient, opt.state)
    """

    dtype: str
    name: str
    param_list: list[Var]
    state: tvm_ffi.Array

    def __init__(self, name: str) -> None:
        self.name = name
        self.param_list = None
        self.state = None
        self.dtype = None

    def init(self, params: Var | list[Var]) -> "Optimizer":
        """Set the parameters, determine the dtype, and construct the initial state for the
        optimizer.

        Parameters
        ----------
        params : Union[Var, List[Var]]
            The parameter or the list of parameters to optimize.

            Parameters should all be Vars of floating point Tensors, including float32, float64,
            float16, etc. Currently, all parameters should have the same dtype, and that dtype
            will be used as the dtype of the optimizer states.

        Returns
        -------
        self : Optimizer
            The optimizer itself.
        """
        if not isinstance(params, list):
            params = [params]
        self._set_params_and_dtype(params)
        # State should be initialized in any implementation of optimizer.
        self.state = None
        return self

    def _set_params_and_dtype(self, params: list[Var]) -> None:
        """Check params is legal and set the param_list and dtype of the optimizer."""
        params_set = set()
        dtype = None
        for x in params:
            if not isinstance(x, Var):
                raise ValueError(f"Parameter {x} is not a Var")
            if not isinstance(x.ty, TensorType):
                raise ValueError(
                    f"Optimizers only support Tensor parameters, but parameter {x.name_hint} has "
                    f"type {x.ty}"
                )
            data_type = tvm.DataType(x.ty.dtype.dtype)
            if data_type.type_code not in (tvm.DataTypeCode.BFLOAT, tvm.DataTypeCode.FLOAT):
                raise ValueError(
                    f"Optimizers only support Tensor parameters of floating point dtype, but dtype "
                    f"of {x.name_hint} is {x.ty.dtype}"
                )
            if dtype is None:
                dtype = x.ty.dtype
            else:
                if dtype != x.ty.dtype:
                    raise ValueError(
                        f"All parameters should have the same dtype, but parameter {x.name_hint} "
                        f"has dtype {x.ty.dtype}, which differs from the previous dtype "
                        f"{dtype}"
                    )
            if x in params_set:
                raise ValueError(f"Parameter {x.name_hint} appears more than once")
            params_set.add(x)
        self.param_list = params
        self.dtype = dtype

    def _check_init(self):
        """Check that the optimizer is initialized. This method should be called at the start of
        get_function().
        """
        if self.param_list is None or self.state is None or self.dtype is None:
            raise RuntimeError("Please call init() for the optimizer before calling get_function()")

    def get_function(self) -> Function:
        """Use blockbuilder to construct an optimizer function that executes updates of the
        parameters and the optimizer state.

        The optimizer function will take in a tuple of parameters, a tuple of gradients of
        parameters, and a tuple of optimizer states. It will return a tuple of updated parameters,
        and a tuple of optimizer states.

        Returns
        -------
        func : Function
            The optimizer function.

        Examples
        --------
        An example of the returned optimizer function. This function executes the stochastic
        gradient descent method with lr = 0.1.

        .. code-block:: python

            @R.function
            def SGD(
                params: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
                gradients: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
                optim_states: R.Tuple(R.Tensor((), "int64")),
            ) -> R.Tuple(
                R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
                R.Tuple(R.Tensor((), "int64")),
            ):
                with R.dataflow():
                    num_steps: R.Tensor((), "int64") = optim_states[0]
                    num_steps_new: R.Tensor((), "int64") = R.add(num_steps, R.const(1, "int64"))
                    x: R.Tensor((3, 3), "float32") = params[0]
                    x_grad: R.Tensor((3, 3), "float32") = gradients[0]
                    lv: R.Tensor((3, 3), "float32") = R.multiply(R.const(0.01, "float32"), x_grad)
                    x_new: R.Tensor((3, 3), "float32") = R.subtract(x, lv)
                    y: R.Tensor((3,), "float32") = params[1]
                    y_grad: R.Tensor((3,), "float32") = gradients[1]
                    lv1: R.Tensor((3,), "float32") = R.multiply(R.const(0.01, "float32"), y_grad)
                    y_new: R.Tensor((3,), "float32") = R.subtract(y, lv1)
                    params_new: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")) = (
                        x_new,
                        y_new,
                    )
                    optim_states_new: R.Tuple(R.Tensor((), "int64")) = (num_steps_new,)
                    R.output(params_new, optim_states_new)
                return (params_new, optim_states_new)
        """
        self._check_init()
        raise NotImplementedError()


# TODO(chaofan, yixin): Support symbolic shapes
def _get_shape_as_int_list(var: Var) -> list[int]:
    return [int(val) for val in var.ty.shape]


# We need to subtract on hyperparameters, but do not want to introduce floating point error.
# Floating point error would lead to a few problems, such as making assert_structural_equal not
# pass in unit tests
def _high_precision_subtract(lhs: float, rhs: float) -> float:
    return float(Decimal(str(lhs)) - Decimal(str(rhs)))


class SGD(Optimizer):
    """Implements stochastic gradient descent.

    The returned function of `get_function()` is equivalent to the following numpy code:

    .. code-block:: python

        def SGD(param_tuple, grad_tuple, state_tuple):
            num_steps = state_tuple[0]
            param_tuple_new, state_tuple_new = [], []
            state_tuple_new.append(num_steps + 1)
            for i in range(len(param_tuple)):
                param = param_tuple[i]
                grad = grad_tuple[i]
                param_tuple_new.append(param - lr * (grad + weight_decay * param))
            return param_tuple_new, state_tuple_new

    Parameters
    ----------
    lr : float
        learning rate

    weight_decay : float
        weight decay (L2 penalty) (default: 0)
    """

    def __init__(self, lr: float, weight_decay: float = 0) -> None:
        super().__init__("SGD")
        self.lr = float(lr)
        self.weight_decay = float(weight_decay)

    def init(self, params: Var | list[Var]) -> "SGD":
        """Set the parameters, determine the dtype, and construct the initial state for the
        optimizer.

        The state of SGD is `(num_steps,)`.

        Parameters
        ----------
        params : Union[Var, List[Var]]
            The parameter or the list of parameters to optimize.

            Parameters should all be Vars of floating point Tensors, including float32, float64,
            float16, etc. Currently, all parameters should have the same dtype, and that dtype
            will be used as the dtype of the optimizer states.

        Returns
        -------
        self : SGD
            The SGD optimizer itself.
        """
        if not isinstance(params, list):
            params = [params]
        self._set_params_and_dtype(params)
        self.state = (
            # num_steps = 0
            tvm.runtime.tensor(np.zeros((), "int64")),
        )
        return self

    def get_function(self) -> Function:
        """Use blockbuilder to construct an optimizer function that executes updates of the
        parameters and the optimizer state. `init()` should be called before `get_function()`.

        Returns
        -------
        func : Function
            The optimizer function.
        """
        self._check_init()

        plist = self.param_list
        len_param = len(plist)
        dtype = self.dtype

        # input variables
        param_var = Var("params", TupleType([p.ty for p in plist]))
        grad_var = Var("gradients", TupleType([p.ty for p in plist]))
        state_var = Var("optim_states", TupleType([TensorType((), "int64")]))

        # constants
        lr = const(self.lr, dtype)
        weight_decay = const(self.weight_decay, dtype)
        one = const(1, "int64")

        builder = BlockBuilder()
        with builder.function(self.name, [param_var, grad_var, state_var]):
            with builder.dataflow():
                param_list_new, state_list_new = [], []

                # handle num_steps
                num_steps = builder.emit(TupleGetItem(state_var, 0), "num_steps")
                num_steps_new = builder.emit(add(num_steps, one), "num_steps_new")
                state_list_new.append(num_steps_new)

                # computation logics
                for i in range(len_param):
                    name = self.param_list[i].name_hint
                    p = builder.emit(TupleGetItem(param_var, i), name)
                    g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
                    if self.weight_decay:
                        g = builder.emit(add(multiply(weight_decay, p), g), name + "_grad_new")
                    p_new = builder.emit(subtract(p, multiply(lr, g)), name + "_new")
                    param_list_new.append(p_new)

                # handle return values
                params_new = builder.emit_output(RxTuple(param_list_new), "params_new")
                optim_states_new = builder.emit_output(RxTuple(state_list_new), "optim_states_new")
            builder.emit_func_output((params_new, optim_states_new))
        return builder.get()[self.name]


class MomentumSGD(Optimizer):
    """Implements stochastic gradient descent with momentum. Optionally supports Nesterov
    momentum.

    The returned function of `get_function()` is equivalent to the following numpy code:

    .. code-block:: python

        def MomentumSGD(param_tuple, grad_tuple, state_tuple):
            num_steps = state_tuple[0]
            param_tuple_new, state_tuple_new = [], []
            state_tuple_new.append(num_steps + 1)

            for i in range(len(param_tuple)):
                param = param_tuple[i]
                grad = grad_tuple[i]
                velocity = state_tuple[i + 1]
                grad = param * weight_decay + grad
                velocity = momentum * velocity + grad * (1 - dampening)
                if nesterov:
                    param = param - (grad + momentum * velocity) * lr
                else:
                    param = param - velocity * lr
                param_tuple_new.append(param)
                state_tuple_new.append(velocity)

            return param_tuple_new, state_tuple_new

    Parameters
    ----------
    lr : float
        learning rate

    momentum : float
        momentum factor (default: 0)

    weight_decay : float
        weight decay (L2 penalty) (default: 0)

    dampening : float
        dampening for momentum (default: 0)

    nesterov : bool
        enables Nesterov momentum (default: False)
    """

    def __init__(
        self,
        lr: float,
        momentum: float,
        dampening: float = 0,
        weight_decay: float = 0,
        nesterov: bool = False,
    ) -> None:
        super().__init__("MomentumSGD")
        self.lr = float(lr)
        self.momentum = float(momentum)
        self.weight_decay = float(weight_decay)
        self.dampening = float(dampening)
        self.nesterov = nesterov

    def init(self, params: Var | list[Var]) -> "MomentumSGD":
        """Set the parameters, determine the dtype, and construct the initial state for the
        optimizer.

        The state of MomentumSGD is
        `(num_steps, velocity_of_param_0, ..., velocity_of_param_n-1)`.

        Parameters
        ----------
        params : Union[Var, List[Var]]
            The parameter or the list of parameters to optimize.

            Parameters should all be Vars of floating point Tensors, including float32, float64,
            float16, etc. Currently, all parameters should have the same dtype, and that dtype
            will be used as the dtype of the optimizer states.

        Returns
        -------
        self : MomentumSGD
            The MomentumSGD optimizer itself.
        """
        if not isinstance(params, list):
            params = [params]
        self._set_params_and_dtype(params)
        self.state = (
            # num_steps = 0
            tvm.runtime.tensor(np.zeros((), "int64")),
            # v_{param} is initialized to all zeros
            *(
                tvm.runtime.tensor(np.zeros(_get_shape_as_int_list(p), p.ty.dtype.dtype))
                for p in self.param_list
            ),
        )
        return self

    def get_function(self) -> Function:
        """Use blockbuilder to construct an optimizer function that executes updates of the
        parameters and the optimizer state. `init()` should be called before `get_function()`.

        Returns
        -------
        func : Function
            The optimizer function.
        """
        self._check_init()
        plist = self.param_list
        len_param = len(plist)
        dtype = self.dtype

        # input variables
        param_var = Var("params", TupleType([p.ty for p in plist]))
        grad_var = Var("gradients", TupleType([p.ty for p in plist]))
        state_var = Var(
            "optim_states",
            TupleType([TensorType((), "int64"), *(p.ty for p in plist)]),
        )

        # constants
        lr = const(self.lr, dtype)
        momentum = const(self.momentum, dtype)
        weight_decay = const(self.weight_decay, dtype)
        dampening_inv = const(_high_precision_subtract(1, self.dampening), dtype)
        one = const(1, "int64")

        builder = BlockBuilder()
        with builder.function(self.name, [param_var, grad_var, state_var]):
            with builder.dataflow():
                param_list_new, state_list_new = [], []

                # handle num_steps
                num_steps = builder.emit(TupleGetItem(state_var, 0), "num_steps")
                num_steps_new = builder.emit(add(num_steps, one), "num_steps_new")
                state_list_new.append(num_steps_new)

                # computation logics
                for i in range(len_param):
                    name = self.param_list[i].name_hint
                    p = builder.emit(TupleGetItem(param_var, i), name)
                    g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
                    v = builder.emit(TupleGetItem(state_var, i + 1), name + "_v")
                    if self.weight_decay:
                        g = builder.emit(add(multiply(weight_decay, p), g), name + "_grad_new")
                    damp_g = multiply(dampening_inv, g) if self.dampening else g
                    v_new = builder.emit(add(multiply(momentum, v), damp_g), name + "_v_new")
                    g_new = (
                        builder.emit(add(g, multiply(momentum, v_new)), name + "_g_nest")
                        if self.nesterov
                        else v_new
                    )
                    p_new = builder.emit(subtract(p, multiply(lr, g_new)), name + "_new")
                    param_list_new.append(p_new)
                    state_list_new.append(v_new)

                # handle return values
                params_new = builder.emit_output(RxTuple(param_list_new), "params_new")
                optim_states_new = builder.emit_output(RxTuple(state_list_new), "optim_states_new")
            builder.emit_func_output((params_new, optim_states_new))
        return builder.get()[self.name]


class Adam(Optimizer):
    """Implements Adam optimization algorithm.

    The returned function of `get_function()` is equivalent to the following numpy code:

    .. code-block:: python

        def Adam(param_tuple, grad_tuple, state_tuple):
            num_steps = state_tuple[0]
            num_steps_new = num_steps + 1

            param_tuple_new = []
            state_tuple_new = [None] * len(state_tuple)
            state_tuple_new[0] = num_steps_new
            state_tuple_new[1] = state_tuple[1] * betas[0]
            state_tuple_new[2] = state_tuple[2] * betas[1]

            for i in range(len(param_tuple)):
                param = param_tuple[i]
                grad = grad_tuple[i]
                m = state_tuple[i + 3]
                v = state_tuple[i + 3 + len(param_tuple)]
                grad = grad + weight_decay * param
                m = betas[0] * m + (1 - betas[0]) * grad
                v = betas[1] * v + (1 - betas[1]) * grad * grad
                m_hat = m / (1 - betas[0] ** num_steps_new)
                v_hat = v / (1 - betas[1] ** num_steps_new)
                param = param - lr * m_hat / (np.sqrt(v_hat) + eps)
                param_tuple_new.append(param)
                state_tuple_new[i + 3] = m
                state_tuple_new[i + 3 + len(param_tuple)] = v

            return param_tuple_new, state_tuple_new

    Parameters
    ----------
    lr : float
        learning rate

    betas : Tuple[float, float]
        coefficients used for computing running averages of gradient and its square
        (default: (0.9, 0.999))

    eps : float
        term added to the denominator to improve numerical stability (default: 1e-8)

    weight_decay : float
        weight decay (L2 penalty) (default: 0)
    """

    def __init__(
        self,
        lr: float,
        betas: tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-08,
        weight_decay: float = 0,
    ) -> None:
        super().__init__("Adam")
        self.lr = float(lr)
        self.beta1 = float(betas[0])
        self.beta2 = float(betas[1])
        self.eps = float(eps)
        self.weight_decay = float(weight_decay)

    def init(self, params: Var | list[Var]) -> "Adam":
        """Set the parameters, determine the dtype, and construct the initial state for the
        optimizer.

        The state of Adam is

        .. code-block:: python

            (
                num_steps,
                beta_0_prod, # beta0 ** num_steps
                beta_1_prod, # beta1 ** num_steps
                first_momentum_of_param_0, ..., first_momentum_of_param_n-1,
                second_momentum_of_param_0, ..., second_momentum_of_param_n-1
            )

        Parameters
        ----------
        params : Union[Var, List[Var]]
            The parameter or the list of parameters to optimize.

            Parameters should all be Vars of floating point Tensors, including float32, float64,
            float16, etc. Currently, all parameters should have the same dtype, and that dtype
            will be used as the dtype of the optimizer states.

        Returns
        -------
        self : Adam
            The Adam optimizer itself.
        """
        if not isinstance(params, list):
            params = [params]
        self._set_params_and_dtype(params)
        self.state = (
            # num_steps, beta_0_prod, beta_1_prod
            tvm.runtime.tensor(np.zeros((), "int64")),
            tvm.runtime.tensor(np.ones((), self.dtype)),
            tvm.runtime.tensor(np.ones((), self.dtype)),
            # first_momentum
            *(
                tvm.runtime.tensor(np.zeros(_get_shape_as_int_list(p), p.ty.dtype.dtype))
                for p in self.param_list
            ),
            # second_momentum
            *(
                tvm.runtime.tensor(np.zeros(_get_shape_as_int_list(p), p.ty.dtype.dtype))
                for p in self.param_list
            ),
        )
        return self

    def get_function(self) -> Function:
        """Use blockbuilder to construct an optimizer function that executes updates of the
        parameters and the optimizer state. `init()` should be called before `get_function()`.

        Returns
        -------
        func : Function
            The optimizer function.
        """
        self._check_init()
        plist = self.param_list
        len_param = len(plist)
        dtype = self.dtype

        # input variables
        param_var = Var("params", TupleType([p.ty for p in plist]))
        grad_var = Var("gradients", TupleType([p.ty for p in plist]))
        state_var = Var(
            "optim_states",
            TupleType(
                [
                    TensorType((), "int64"),
                    TensorType((), dtype),
                    TensorType((), dtype),
                    *(p.ty for p in plist),
                    *(p.ty for p in plist),
                ]
            ),
        )

        # constants
        lr = const(self.lr, dtype)
        beta1 = const(self.beta1, dtype)
        beta2 = const(self.beta2, dtype)
        beta1_inv = const(_high_precision_subtract(1, self.beta1), dtype)
        beta2_inv = const(_high_precision_subtract(1, self.beta2), dtype)
        eps = const(self.eps, dtype)
        weight_decay = const(self.weight_decay, dtype)
        one_int = const(1, "int64")
        one_float = const(1, dtype)

        builder = BlockBuilder()
        with builder.function(self.name, [param_var, grad_var, state_var]):
            with builder.dataflow():
                param_list_new = []
                state_list_new = [None] * (len_param * 2 + 3)  # type: List[Optional[Var]]

                # handle num_steps
                num_steps = builder.emit(TupleGetItem(state_var, 0), "num_steps")
                num_steps_new = builder.emit(add(num_steps, one_int), "num_steps_new")
                state_list_new[0] = num_steps_new
                beta1_prod = builder.emit(multiply(TupleGetItem(state_var, 1), beta1), "beta1_prod")
                beta2_prod = builder.emit(multiply(TupleGetItem(state_var, 2), beta2), "beta2_prod")
                state_list_new[1] = beta1_prod
                state_list_new[2] = beta2_prod

                # computation logics
                for i in range(len_param):
                    name = self.param_list[i].name_hint
                    p = builder.emit(TupleGetItem(param_var, i), name)
                    g = builder.emit(TupleGetItem(grad_var, i), name + "_grad")
                    m = builder.emit(TupleGetItem(state_var, i + 3), name + "_m")
                    v = builder.emit(TupleGetItem(state_var, i + 3 + len_param), name + "_v")
                    if self.weight_decay:
                        g = builder.emit(add(multiply(weight_decay, p), g), name + "_grad_new")
                    m_new = builder.emit(
                        add(multiply(beta1, m), multiply(beta1_inv, g)), name + "_m_new"
                    )
                    v_new = builder.emit(
                        add(multiply(beta2, v), multiply(beta2_inv, multiply(g, g))),
                        name + "_v_new",
                    )
                    m_hat = builder.emit(
                        divide(m_new, subtract(one_float, state_list_new[1])), name + "_m_hat"
                    )
                    v_hat = builder.emit(
                        divide(v_new, subtract(one_float, state_list_new[2])), name + "_v_hat"
                    )
                    p_new = builder.emit(
                        subtract(p, multiply(lr, divide(m_hat, add(sqrt(v_hat), eps)))),
                        name + "_new",
                    )
                    param_list_new.append(p_new)
                    state_list_new[i + 3] = m_new
                    state_list_new[i + 3 + len_param] = v_new

                # handle return values
                params_new = builder.emit_output(RxTuple(param_list_new), "params_new")
                optim_states_new = builder.emit_output(RxTuple(state_list_new), "optim_states_new")
            builder.emit_func_output((params_new, optim_states_new))
        return builder.get()[self.name]