# Copyright 2022 Google LLC
#
# Licensed 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.
"""Configuration dataclasses."""

import abc
import dataclasses
from typing import Any, Callable, Optional, Union
from . import calibration
from . import int_numerics
from . import stochastic_rounding
import flax.struct
import jax
import jax.numpy as jnp

DType = Any
Context = Any  # TODO(lew): We could put Context in a separate file.

ClipAndRoundFn = Callable[[jnp.ndarray, Context], jnp.ndarray]


class Preprocess(abc.ABC):

    @abc.abstractmethod
    def __call__(self, inputs):
        pass


class NoNumerics(flax.struct.PyTreeNode):
    # TODO(lew): This is a workaround. We should separate Stochastic Rounding.
    # noise_fn has no effect in NoNumerics.
    noise_fn: Optional[stochastic_rounding.NoiseFn] = None

    """No quantization, use a native type such as bf16."""
    pass


Numerics = Union[NoNumerics, int_numerics.IntNumerics]


@dataclasses.dataclass
class Tensor:
    """Configuration of quantization of one tensor or one side of tensor op."""

    numerics: Numerics
    calib_shared_axes: Optional[list[int]]
    scale_stop_grad: bool
    # noise+clip+round
    # We apply gradient of clip_and_round in bwd pass.
    calibration: calibration.Calibration
    # Round up the calibration to power of 2 (po2).
    po2_scale: bool
    use_fake_quant: bool
    # Controls at what value of input tensor should be used.
    # Setting it to True, but not quantizing fwd pass will assert-fail.
    use_fwd_quant: Optional[bool]
    # Operation applied to the quantized inputs to the dot general
    preprocess_quant_cls: Optional[Callable[[], Preprocess]]
    # lax.dot_general output is multiplied by (int transposed) scales.
    # preprocess_scale will be applied to these scales before the multiplication.
    preprocess_scale_cls: Optional[Callable[[], Preprocess]]

    @classmethod
    def make(cls, *args, **kwargs) -> 'Tensor':
        return tensor_make(*args, **kwargs)


@dataclasses.dataclass
class LocalQ:
    contraction_axis_shard_count: int


@dataclasses.dataclass
class DotGeneralRaw:
    """Configuration of quantization of one dot_general without gradient."""

    lhs: Tensor
    rhs: Tensor
    dg_in_dtype: Optional[DType]
    dg_accumulator_dtype: Optional[DType]
    local_q: Optional[LocalQ]

    @classmethod
    def make(cls, *args, **kwargs) -> 'DotGeneralRaw':
        return dot_general_raw_make(*args, **kwargs)

    @classmethod
    def make_conv_general_dilated(cls, *args, **kwargs) -> 'DotGeneralRaw':
        return conv_general_dilated_make(*args, **kwargs)


@dataclasses.dataclass
class DotGeneral:
    """Configuration of quantization of dot_general and its gradients."""

    fwd: DotGeneralRaw
    dlhs: DotGeneralRaw
    drhs: DotGeneralRaw

    @classmethod
    def make(cls, *args, **kwargs) -> 'DotGeneral':
        return dot_general_make(*args, **kwargs)


def set_accumulator_dtype(
        cfg: DotGeneral,
        fwd_dtype: Optional[DType],
        dlhs_dtype: Optional[DType],
        drhs_dtype: Optional[DType],
):
    """
    The set_accumulator_dtype function sets the accumulator dtype for each of the three
    differentiable functions.  The accumulator dtype is used to store intermediate results
    during forward and backward passes.  It is also used to store gradients during backward pass.
    The default value for this parameter is None, which means that it will be set automatically by
    the library based on other parameters such as input data types and output data type.

    :param cfg: DotGeneral: Set the accumulator dtype for each of the three
    :param fwd_dtype: Optional[DType]: Set the accumulator dtype for forward mode
    :param dlhs_dtype: Optional[DType]: Set the dg_accumulator_dtype
    :param drhs_dtype: Optional[DType]: Set the data type for the
    :param : Set the accumulator dtype for the forward pass,
    :return: A tuple of three elements,
    
    """
    cfg.fwd.dg_accumulator_dtype = fwd_dtype
    cfg.dlhs.dg_accumulator_dtype = dlhs_dtype
    cfg.drhs.dg_accumulator_dtype = drhs_dtype


def set_stochastic_rounding(
        cfg: DotGeneral,
        # Typically we have (but it's a caller's responsibility to check):
        # - vjp_lhs_stochastic_rounding is referring to the gradient and
        # - vjp_rhs_stochastic_rounding is referring to the activations/weights.
        vjp_lhs_stochastic_rounding: bool,
        vjp_rhs_stochastic_rounding: bool,
        implementation: str,
):
    """Configure stochastic rounding implementation."""
    noise_implementations = {
        'jax.uniform': lambda shape, key: jax.random.uniform(key, shape) - 0.5,
        'custom-1': stochastic_rounding.random_centered_uniform,
    }
    msg = f'{implementation} not supported.'
    assert implementation in noise_implementations.keys(), msg
    noise_fn = noise_implementations[implementation]

    if vjp_lhs_stochastic_rounding:
        cfg.dlhs.lhs.numerics = cfg.dlhs.lhs.numerics.replace(noise_fn=noise_fn)
        cfg.drhs.lhs.numerics = cfg.drhs.lhs.numerics.replace(noise_fn=noise_fn)
    else:
        cfg.dlhs.lhs.numerics = cfg.dlhs.lhs.numerics.replace(noise_fn=None)
        cfg.drhs.lhs.numerics = cfg.drhs.lhs.numerics.replace(noise_fn=None)

    if vjp_rhs_stochastic_rounding:
        cfg.dlhs.rhs.numerics = cfg.dlhs.rhs.numerics.replace(noise_fn=noise_fn)
        cfg.drhs.rhs.numerics = cfg.drhs.rhs.numerics.replace(noise_fn=noise_fn)
    else:
        cfg.dlhs.rhs.numerics = cfg.dlhs.rhs.numerics.replace(noise_fn=None)
        cfg.drhs.rhs.numerics = cfg.drhs.rhs.numerics.replace(noise_fn=None)


def set_static_bound(cfg: DotGeneral, bound: float = 1.0):
    """
    The set_static_bound function sets the calibration of all the layers in a DotGeneral configuration to be constant.
    This is useful for when we want to use a static bound, as opposed to an adaptive one.


    :param cfg: DotGeneral: Set the static bound for all the dot-general's sub-functions
    :param bound: float: Set the bound of the calibration
    :return: The cfg which is the dot general object
    
    """
    cfg.fwd.lhs.calibration = calibration.ConstantCalibration(bound)
    cfg.fwd.rhs.calibration = calibration.ConstantCalibration(bound)
    cfg.drhs.lhs.calibration = calibration.ConstantCalibration(bound)
    cfg.drhs.rhs.calibration = calibration.ConstantCalibration(bound)
    cfg.dlhs.lhs.calibration = calibration.ConstantCalibration(bound)
    cfg.dlhs.rhs.calibration = calibration.ConstantCalibration(bound)


def tensor_make(bits: Optional[int]) -> 'Tensor':
    """Makes config.Tensor."""
    if bits is None:
        numerics = NoNumerics()
    else:
        pz = False if bits == 1 else True
        numerics = int_numerics.IntNumerics(
            bits=bits,
            preserve_zero=pz,
            preserve_max_val=False,
            clip=True,
            round=True,
            noise_fn=None,
            clip_gradient=False,  # This can be disabled when using abs-max scaling.
        )

    return Tensor(
        numerics=numerics,
        calib_shared_axes=None,
        scale_stop_grad=True,
        calibration=calibration.AbsMaxCalibration(),
        po2_scale=False,
        use_fake_quant=False,
        # dtype_x=dtype,
        use_fwd_quant=None,
        preprocess_quant_cls=None,
        preprocess_scale_cls=None,
    )


def dot_general_raw_make(
        lhs_bits=None,
        rhs_bits=None,
        local_q=None,
) -> 'DotGeneralRaw':
    """Create quantization configs for input matrices to a matmul."""
    lhs_cfg = tensor_make(lhs_bits)
    rhs_cfg = tensor_make(rhs_bits)

    # Binary uses 0.5 right now.
    if (
            lhs_bits is not None
            and rhs_bits is not None
            and 2 <= lhs_bits <= 8
            and 2 <= rhs_bits <= 8
    ):
        dg_in_dtype = jnp.int8
        dg_accumulator_dtype = jnp.int32
    else:
        # Use None to determine the dtype on the fly in q_dot_general
        dg_in_dtype = None
        dg_accumulator_dtype = None

    return DotGeneralRaw(
        lhs=lhs_cfg,
        rhs=rhs_cfg,
        dg_in_dtype=dg_in_dtype,
        dg_accumulator_dtype=dg_accumulator_dtype,
        local_q=local_q,
    )


def conv_general_dilated_make(
        spatial_dimensions=2,
        lhs_bits: Optional[int] = None,
        rhs_bits: Optional[int] = None,
) -> 'DotGeneralRaw':
    """Create quantization config conv_general_dilated."""
    config = dot_general_raw_make(lhs_bits, rhs_bits)
    # Hardcoding flax assumptions.
    if config.lhs:
        config.lhs.calib_shared_axes = list(range(1, spatial_dimensions + 2))
    if config.rhs:
        config.rhs.calib_shared_axes = list(range(0, spatial_dimensions + 2 - 1))
    return config


def dot_general_make(
        lhs_bits: Optional[int] = None,
        rhs_bits: Optional[int] = None,
        bwd_bits: Optional[int] = None,
        use_fwd_quant: bool = True,
        dlhs_local_q=None,
        drhs_local_q=None,
) -> 'DotGeneral':
    """Create quantization configs for input matrices to a matmul."""
    fwd = dot_general_raw_make(lhs_bits, rhs_bits)
    dlhs = dot_general_raw_make(bwd_bits, bwd_bits, local_q=dlhs_local_q)
    drhs = dot_general_raw_make(bwd_bits, bwd_bits, local_q=drhs_local_q)
    cfg = DotGeneral(fwd=fwd, dlhs=dlhs, drhs=drhs)

    # Surprising: lhs quantization determines what drhs can do.
    if lhs_bits is not None:
        # Only rhs is accepting MultiTensor.
        cfg.drhs.rhs.use_fwd_quant = use_fwd_quant
    if rhs_bits is not None:
        cfg.dlhs.rhs.use_fwd_quant = use_fwd_quant
    return cfg


def fully_quantized(
        *,
        fwd_bits: Union[int, None] = 8,
        bwd_bits: Union[int, None] = 8,
        use_fwd_quant: bool = True,
        use_stochastic_rounding: Optional[bool] = True,
        # Typically we have (but it's a caller's responsibility to check):
        # - vjp_lhs_stochastic_rounding is referring to the gradient and
        # - vjp_rhs_stochastic_rounding is referring to the activations/weights.
        vjp_lhs_stochastic_rounding: Optional[bool] = None,
        vjp_rhs_stochastic_rounding: Optional[bool] = None,
        # The dummy static bound flag is temporary, for performance benchmarking.
        use_dummy_static_bound: bool = False,
        dlhs_local_q: Optional[LocalQ] = None,
        drhs_local_q: Optional[LocalQ] = None,
) -> DotGeneral:
    """
    The fully_quantized function is a helper function that allows you to quickly
    configure the DotGeneral operator with quantization and stochastic rounding.
    It takes in the following arguments:

    :param *: Indicate that the function takes a variable number of arguments
    :param fwd_bits: Union[int, None]: Specify the number of bits used to represent the forward pass
    :param bwd_bits: Union[int, None]: Specify the number of bits used for backpropagation
    :param use_fwd_quant: bool: Control whether we quantize the forward pass
    :param use_stochastic_rounding: Optional[bool]: Enable stochastic rounding
    :param vjp_lhs_stochastic_rounding: Optional[bool]: Refer to the gradient
    :param vjp_rhs_stochastic_rounding: Optional[bool]: Enable stochastic rounding for the activations/weights
    :param dlhs_local_q: Optional[LocalQ]: Specify the quantization scheme for
    :param drhs_local_q: Optional[LocalQ]: Specify the quantization parameters for the right hand side of the matrix multiplication
    :param : Set the number of bits for forward and backward pass
    :return: A dotgeneral object
    
    """
    cfg = dot_general_make(
        lhs_bits=fwd_bits,
        rhs_bits=fwd_bits,
        bwd_bits=bwd_bits,
        use_fwd_quant=use_fwd_quant,
        dlhs_local_q=dlhs_local_q,
        drhs_local_q=drhs_local_q,
    )

    # Stochastic Rounding
    # These 3 variables are used to ensure we don't mix
    # old and new style of SR configuration.
    old_style_sr_config = use_stochastic_rounding is not None
    new_style_sr_config_lhs = vjp_lhs_stochastic_rounding is not None
    new_style_sr_config_rhs = vjp_rhs_stochastic_rounding is not None
    assert new_style_sr_config_lhs == new_style_sr_config_rhs, (
        'if you use new style SR config (vjp_xhs_stochastic_rounding), do pass'
        ' both lhs and rhs explicitely.'
    )
    assert new_style_sr_config_lhs != old_style_sr_config

    true = True  # A crude way to get around g-explicit-bool-comparison warning

    assert not (vjp_lhs_stochastic_rounding and vjp_rhs_stochastic_rounding), (
        'This config is buggy when you set both to True. Contact lew@ or use'
        ' config_v3'
    )

    # By default use jax.uniform for stochastic rounding
    if use_stochastic_rounding == true:
        set_stochastic_rounding(cfg, True, True, 'jax.uniform')

    if vjp_lhs_stochastic_rounding == true:
        set_stochastic_rounding(cfg, True, False, 'jax.uniform')

    if vjp_rhs_stochastic_rounding == true:
        set_stochastic_rounding(cfg, False, True, 'jax.uniform')

    if use_dummy_static_bound:
        set_static_bound(cfg, 1.0)

    return cfg


def config_v3(
        *,
        fwd_bits: Union[int, None],
        dlhs_bits: Union[int, None],
        drhs_bits: Union[int, None],
        use_dummy_static_bound: bool = False,
        rng_type: str = 'jax.uniform',  # 'custom-1'
        dlhs_local_q: Optional[LocalQ] = None,
        drhs_local_q: Optional[LocalQ] = None,
        fwd_accumulator_dtype: ... = jnp.int32,
        dlhs_accumulator_dtype: ... = jnp.int32,
        drhs_accumulator_dtype: ... = jnp.int32,
) -> DotGeneral:
    """Fully Quantized Training."""
    fwd = dot_general_raw_make(fwd_bits, fwd_bits)
    dlhs = dot_general_raw_make(dlhs_bits, dlhs_bits, local_q=dlhs_local_q)
    drhs = dot_general_raw_make(drhs_bits, drhs_bits, local_q=drhs_local_q)
    cfg = DotGeneral(fwd=fwd, dlhs=dlhs, drhs=drhs)

    cfg.dlhs.rhs.use_fwd_quant = False
    cfg.drhs.rhs.use_fwd_quant = False

    set_stochastic_rounding(
        cfg,
        vjp_lhs_stochastic_rounding=True,
        vjp_rhs_stochastic_rounding=False,
        implementation=rng_type,
    )

    if use_dummy_static_bound:
        set_static_bound(cfg, 1.0)

    set_accumulator_dtype(
        cfg,
        fwd_dtype=fwd_accumulator_dtype,
        dlhs_dtype=dlhs_accumulator_dtype,
        drhs_dtype=drhs_accumulator_dtype,
    )
    return cfg
