# 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.
"""Quantized dot_general."""

# Lingo in this file:
#
# - lhs(rhs) - left(right) hand side of a binary operation
# - ca - contraction axes
# - ba - batch axes
# - ra - remaining axes

# pylint: disable=g-explicit-bool-comparison
# pylint: disable=g-explicit-length-test

import copy
import functools
from typing import Callable, Optional, Union
from . import config
import flax.struct
import jax
from jax import lax
import jax.numpy as jnp
import numpy as onp


@flax.struct.dataclass
class Context:
    key: Optional[jax.Array]
    train_step: Optional[int]


def _context_split(context: Context) -> tuple[Context, Context]:
    def mk_ctx(key):
        return Context(key=key, train_step=context.train_step)

    if context.key is not None:
        key1, key2 = jax.random.split(context.key)
        return mk_ctx(key1), mk_ctx(key2)
    return mk_ctx(None), mk_ctx(None)


def _scale_quant(x, *, cfg, ca, context):
    """The core quantifying function."""
    msg = (
        'use_fake_quant mode is used in tests and it is exactly equal when'
        ' po2_scale == True; Did you forget to set it?'
    )
    assert (not cfg.use_fake_quant) or cfg.po2_scale, msg

    if isinstance(cfg.numerics, config.NoNumerics):
        return x, None, None
    shared_axes = cfg.calib_shared_axes or ca
    bound = cfg.calibration.get_bound(x, shared_axes)
    abs_max_mapped_to = cfg.numerics.abs_val_mapped_to()
    scale = abs_max_mapped_to / bound

    if cfg.po2_scale:
        scale = 2 ** jnp.floor(jnp.log2(scale))
    if cfg.scale_stop_grad:
        scale = lax.stop_gradient(scale)

    x_s = _maybe_mul(x, scale)

    quant = jax.custom_vjp(cfg.numerics.fwd)
    quant.defvjp(cfg.numerics.vjp_fwd, cfg.numerics.vjp_bwd)
    quant = functools.partial(quant, context=context)

    x_q, quant_grad = jax.vjp(quant, x_s)
    inv_scale = _maybe_inv(scale)

    return x_q, inv_scale, quant_grad


def make_fake_quant(cfg: config.Tensor, ca=None):
    def fake_quant(x, context):
        x_q, inv_scale, _ = _scale_quant(x, cfg=cfg, ca=ca, context=context)
        return _maybe_mul(x_q, inv_scale)

    return fake_quant


@flax.struct.dataclass
# It is used only when use_fwd_quant = True
class QTensor:
    qvalue: jnp.ndarray
    qvalue_scale_t: jnp.ndarray


@flax.struct.dataclass
class MultiTensor:
    x: jnp.ndarray
    qx: Optional[QTensor]


@flax.struct.dataclass
class TensorRes:
    """All the things we pass from the forward pass to the backward pass."""
    mt: MultiTensor
    quant_grad: Union[Callable[[jnp.ndarray], tuple[jnp.ndarray]], None]


@flax.struct.dataclass
class DotGeneralRes:
    context_bwd: Context
    lhs: TensorRes
    rhs: TensorRes


def _scale_trans(x, ca, ba):
    """Transposes x to output dimension order."""
    ca = list(ca)
    ba = list(ba)
    for i in ca:
        assert x.shape[i] == 1
    ra = list(i for i in range(len(x.shape)) if i not in ba + ca)
    x = jnp.transpose(x, ba + ra + ca)
    # x = jnp.squeeze(x, axis=range(len(ba+ra): len(x.shape))
    shape_ba = x.shape[: len(ba)]
    shape_ra = x.shape[len(ba): len(x.shape) - len(ca)]
    # Will need to add additional axes (size 1) for the other shape_ra
    x = x.reshape(shape_ba + shape_ra)
    return x


def _lhs_scale_transpose(lhs_scale, dimension_numbers, lhs_shape, rhs_shape):
    """Transposes lhs_scale to output dimension order."""
    if lhs_scale is None:
        return None

    (lhs_ca, rhs_ca), (lhs_ba, rhs_ba) = dimension_numbers
    qlhs_scale_t = _scale_trans(lhs_scale, lhs_ca, lhs_ba)
    # inserting dummy axes for rhs_ra
    assert len(qlhs_scale_t.shape) == len(lhs_shape) - len(lhs_ca)
    start = len(qlhs_scale_t.shape)
    end = len(rhs_shape) - len(rhs_ca) - len(rhs_ba) + start
    lhs_dummy_axes = range(start, end)
    qlhs_scale_t = jnp.expand_dims(qlhs_scale_t, axis=lhs_dummy_axes)
    return qlhs_scale_t


def _rhs_scale_transpose(rhs_scale, dimension_numbers, lhs_shape, rhs_shape):
    if rhs_scale is None:
        return None
    del rhs_shape
    (lhs_ca, rhs_ca), (lhs_ba, rhs_ba) = dimension_numbers
    qrhs_scale_t = _scale_trans(rhs_scale, rhs_ca, rhs_ba)
    start = len(rhs_ba)
    end = len(lhs_shape) - len(lhs_ca) - len(lhs_ba) + start
    rhs_dummy_axes = range(start, end)
    qrhs_scale_t = jnp.expand_dims(qrhs_scale_t, axis=rhs_dummy_axes)
    return qrhs_scale_t


def _maybe_mul(x, scale):
    if scale is None:
        return x
    return x * scale


def _maybe_inv(x):
    if x is None:
        return None
    return 1.0 / x


def _make_dot_general_raw(gcfg: config.DotGeneralRaw):
    """Makes quantized lax.dot_general replacement."""

    msg = 'Custom calib_shared_axes not implemented for local AQT.'
    assert gcfg.lhs.calib_shared_axes is None, msg
    assert gcfg.rhs.calib_shared_axes is None, msg

    def dot_general_raw(
            lhs: jnp.ndarray,
            rhs: Union[jnp.ndarray, MultiTensor],
            dimension_numbers,
            context,
    ):
        """Creates a dot_general function without custom gradient."""
        (lhs_ca, rhs_ca), (lhs_ba, rhs_ba) = dimension_numbers
        # We need to copy because we modify cfg to populate some defaults.
        cfg = copy.deepcopy(gcfg)

        if isinstance(rhs, MultiTensor):
            # We are in gradient code.
            fwd_quantized = rhs.qx is not None
            expect_fwd_quantized = cfg.rhs.use_fwd_quant is not None
            msg = (
                'Misconfiguration: use_fwd_quant=True, but there is no fwd'
                ' quantization (but rhs.qx is None).'
            )
            assert fwd_quantized == expect_fwd_quantized, msg
            if cfg.rhs.use_fwd_quant:
                assert rhs.qx is not None, msg
                lhs = _maybe_mul(lhs, rhs.qx.qvalue_scale_t)
                rhs = rhs.qx.qvalue
            else:
                rhs = rhs.x
        else:
            assert cfg.rhs.use_fwd_quant is None, 'cannot set use_fwd_quant in fwd'

        if cfg.local_q is not None:

            def factor_reshape(x, ca, ba):
                factor = cfg.local_q.contraction_axis_shard_count
                assert factor is not None
                if len(ca) == 0:
                    return x, ca, ba
                shape = list(x.shape)
                ax = ca[0]
                orig_size = shape[ax]
                assert orig_size % factor == 0
                shape[ax] = factor
                shape.insert(ax + 1, orig_size // factor)
                new_ca = [(b + int(b >= ax)) for b in ca]
                assert new_ca[0] == ax + 1
                new_ba = [ax] + [(b + int(b > ax)) for b in ba]
                return x.reshape(shape), new_ca, new_ba

            lhs, lhs_ca, lhs_ba = factor_reshape(lhs, lhs_ca, lhs_ba)
            rhs, rhs_ca, rhs_ba = factor_reshape(rhs, rhs_ca, rhs_ba)

            dimension_numbers = (lhs_ca, rhs_ca), (lhs_ba, rhs_ba)

        assert isinstance(rhs, jnp.ndarray)

        context, context_bwd = _context_split(context)
        context_lhs, context_rhs = _context_split(context)
        del context

        lhs_q, lhs_inv_scale, lhs_quant_grad = _scale_quant(
            lhs, cfg=cfg.lhs, ca=lhs_ca, context=context_lhs
        )
        lhs_inv_scale_t = _lhs_scale_transpose(
            lhs_inv_scale, dimension_numbers, lhs.shape, rhs.shape
        )
        lhs_qx = (
            None
            if lhs_inv_scale_t is None
            else QTensor(qvalue=lhs_q, qvalue_scale_t=lhs_inv_scale_t)
        )
        lhs_mt = MultiTensor(x=lhs, qx=lhs_qx)
        lhs_res = TensorRes(mt=lhs_mt, quant_grad=lhs_quant_grad)

        rhs_q, rhs_inv_scale, rhs_quant_grad = _scale_quant(
            rhs, cfg=cfg.rhs, ca=rhs_ca, context=context_rhs
        )
        rhs_inv_scale_t = _rhs_scale_transpose(
            rhs_inv_scale, dimension_numbers, lhs.shape, rhs.shape
        )
        rhs_qx = (
            None
            if rhs_inv_scale_t is None
            else QTensor(qvalue=rhs_q, qvalue_scale_t=rhs_inv_scale_t)
        )
        rhs_mt = MultiTensor(x=rhs, qx=rhs_qx)
        rhs_res = TensorRes(mt=rhs_mt, quant_grad=rhs_quant_grad)

        assert lhs.dtype == rhs.dtype

        if cfg.lhs.use_fake_quant:
            msg = "Can't set dg_in_dtype in fake_quant mode."
            assert cfg.dg_in_dtype is None, msg
            lhs_q = _maybe_mul(lhs_q, lhs_inv_scale)
            rhs_q = _maybe_mul(rhs_q, rhs_inv_scale)
        else:
            if cfg.dg_in_dtype is not None:
                lhs_q = lhs_q.astype(cfg.dg_in_dtype)
                rhs_q = rhs_q.astype(cfg.dg_in_dtype)

        if cfg.lhs.preprocess_quant_cls is not None:
            preprocess_quant_lhs = cfg.lhs.preprocess_quant_cls()
            lhs_q = preprocess_quant_lhs(lhs_q)
        if cfg.lhs.preprocess_scale_cls is not None:
            preprocess_scale_lhs = cfg.lhs.preprocess_scale_cls()
            lhs_inv_scale_t = preprocess_scale_lhs(lhs_inv_scale_t)
        if cfg.rhs.preprocess_quant_cls is not None:
            preprocess_quant_rhs = cfg.rhs.preprocess_quant_cls()
            rhs_q = preprocess_quant_rhs(rhs_q)
        if cfg.rhs.preprocess_scale_cls is not None:
            preprocess_scale_rhs = cfg.rhs.preprocess_scale_cls()
            rhs_inv_scale_t = preprocess_scale_rhs(rhs_inv_scale_t)
        out = lax.dot_general(
            lhs_q,
            rhs_q,
            dimension_numbers=dimension_numbers,
            preferred_element_type=cfg.dg_accumulator_dtype,
            precision=lax.Precision.DEFAULT,
        ).astype(lhs.dtype)

        if not cfg.lhs.use_fake_quant:
            out = _maybe_mul(out, lhs_inv_scale_t)
            out = _maybe_mul(out, rhs_inv_scale_t)

        res = DotGeneralRes(
            context_bwd=context_bwd,
            lhs=lhs_res,
            rhs=rhs_res,
        )
        if cfg.local_q is not None:
            assert len(lhs_ca) == len(rhs_ca)
            if len(lhs_ca) > 0:
                out = jnp.sum(out, axis=0)
            # We are not supporting local AQT in fwd pass, so no res needed.
            res = None
        return out, res

    return dot_general_raw


def _dot_general_raw_attach_gradient(
        fwd_dot_general_raw,
        dlhs_dot_general_raw,
        drhs_dot_general_raw,
):
    """Makes quantized lax.dot_general replacement with attached gradients."""

    def make_fwd(return_residual):
        def fwd(
                lhs,
                rhs,
                dimension_numbers,
                context,
        ):
            assert lhs.dtype == rhs.dtype
            ret, res = fwd_dot_general_raw(
                lhs,
                rhs,
                dimension_numbers,
                context,
            )
            ret = ret.astype(lhs.dtype)
            return (ret, res) if return_residual else ret

        return fwd

    def vjp_bwd(
            fwd_dimension_numbers,
            res: DotGeneralRes,
            g,
    ):
        def ranges_like(*xs):
            start = 0
            for x in xs:
                yield tuple(range(start, start + len(x)))
                start += len(x)

        def grad_dot_general(
                y_res: TensorRes,
                quant_grad,
                dot_general,
                y_is_lhs,
                context,
        ):
            y_ndim = y_res.mt.x.ndim

            (x_ca, y_ca), (x_ba, y_ba) = fwd_dimension_numbers
            if y_is_lhs:
                (y_ca, x_ca) = (x_ca, y_ca)
                (y_ba, x_ba) = (x_ba, y_ba)
            g_ndim = g.ndim - y_ndim + len(x_ba) + 2 * len(x_ca)
            x_ra = tuple(i for i in range(g_ndim) if i not in x_ca and i not in x_ba)
            y_ra = tuple(i for i in range(y_ndim) if i not in y_ca and i not in y_ba)
            if y_is_lhs:
                g_ba, g_ca, _ = ranges_like(x_ba, y_ra, x_ra)
            else:
                g_ba, _, g_ca = ranges_like(x_ba, x_ra, y_ra)
            dims = ((g_ca, y_ra), (g_ba, y_ba))

            out, _ = dot_general(g, y_res.mt, dims, context)

            x_ca_sorted_by_y = tuple(onp.take(x_ca, onp.argsort(y_ca)))
            out_axes = tuple(onp.argsort(tuple(x_ba) + x_ra + x_ca_sorted_by_y))
            transposed_out = jax.lax.transpose(out, out_axes)
            if quant_grad is not None:
                transposed_out = quant_grad(transposed_out)[0]
            return transposed_out

        context1, context2 = _context_split(res.context_bwd)
        dlhs = grad_dot_general(
            res.rhs,
            res.lhs.quant_grad,
            dlhs_dot_general_raw,
            False,
            context1,
        )
        drhs = grad_dot_general(
            res.lhs,
            res.rhs.quant_grad,
            drhs_dot_general_raw,
            True,
            context2,
        )
        return (dlhs, drhs, None)

    vjp = jax.custom_vjp(make_fwd(False), nondiff_argnums=(2,))
    vjp.defvjp(make_fwd(True), vjp_bwd)
    return vjp


def make_dot_general(cfg: Optional[config.DotGeneral]):
    """Makes quantized lax.dot_general replacement with attached gradients."""
    if cfg is None:
        def ret_lax_dg(
                lhs,
                rhs,
                dimension_numbers,
                precision=None,
                preferred_element_type=None,
                *,
                context=Context(key=None, train_step=None),
        ):
            del context
            return jax.lax.dot_general(
                lhs, rhs, dimension_numbers, precision, preferred_element_type
            )

        return ret_lax_dg

    dg = _dot_general_raw_attach_gradient(
        fwd_dot_general_raw=_make_dot_general_raw(cfg.fwd),
        dlhs_dot_general_raw=_make_dot_general_raw(cfg.dlhs),
        drhs_dot_general_raw=_make_dot_general_raw(cfg.drhs),
    )

    def ret_dg(
            lhs,
            rhs,
            dimension_numbers,
            precision=None,
            preferred_element_type=None,
            *,
            context=Context(key=None, train_step=None),
    ):
        del preferred_element_type
        assert (
                precision is None
        ), f'Precision {precision} requested together with quantization.'
        assert lhs.dtype == rhs.dtype, (
            'The only reason we need that, is because we need to determine return'
            ' type.'
        )
        out = dg(
            lhs=lhs,
            rhs=rhs,
            dimension_numbers=dimension_numbers,
            context=context,
        )
        return out

    return ret_dg


def make_conv_general_dilated(cfg: config.DotGeneralRaw):
    """Makes quantized lax.make_conv_general_dilated replacement."""

    cfg = copy.deepcopy(cfg)
    if cfg is None:
        cfg = config.DotGeneralRaw.make()

    def my_conv_general_dilated(
            lhs,
            rhs,
            window_strides,
            padding,
            lhs_dilation=None,
            rhs_dilation=None,
            dimension_numbers=None,
            feature_group_count=1,
            batch_group_count=1,
            precision=None,
            preferred_element_type=None,
    ) -> jax.Array:
        msg1 = """
To simplify the code, we currently assume a Flax-particular layout of the data.
This makes sense, because this is the main use-case of this function.
However if there is any other use, we will drop that assumption.
"""
        rank = len(lhs.shape)
        assert len(rhs.shape) == rank
        assert dimension_numbers is not None, msg1
        assert dimension_numbers.lhs_spec[0:2] == (0, rank - 1), msg1
        assert dimension_numbers.rhs_spec[0:2] == (rank - 1, rank - 2), msg1
        assert dimension_numbers.out_spec[0:2] == (0, rank - 1), msg1
        # In Flax, lhs is the inputs, rhs is the kernel.
        # lhs layout is B, spatials..., Ci
        # rhs layout is: spatials..., Ci, Co
        # out layous it: B, spatials..., Co
        #
        # we need to share these axes: lhs[1:] , rhs[:-1]
        # we have a scale/invscale per: lhs[0] / out[0] and rhs[-1] / out[-1]

        # Flax assumptions.
        assert cfg.lhs.calib_shared_axes == list(range(1, rank))
        assert cfg.rhs.calib_shared_axes == list(range(0, rank - 1))

        lhs_q, lhs_inv_scale, _ = _scale_quant(
            lhs, cfg=cfg.lhs, ca=None, context=None
        )
        rhs_q, rhs_inv_scale, _ = _scale_quant(
            rhs, cfg=cfg.rhs, ca=None, context=None
        )

        out = lax.conv_general_dilated(
            lhs=lhs_q,
            rhs=rhs_q,
            window_strides=window_strides,
            padding=padding,
            lhs_dilation=lhs_dilation,
            rhs_dilation=rhs_dilation,
            dimension_numbers=dimension_numbers,
            feature_group_count=feature_group_count,
            batch_group_count=batch_group_count,
            precision=precision,
            preferred_element_type=preferred_element_type,
        )

        out = _maybe_mul(out, lhs_inv_scale)
        out = _maybe_mul(out, rhs_inv_scale)

        return out

    return my_conv_general_dilated
