# 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 functools
from typing import Callable, Optional, Union
from fjformer.bit_quantization import config
from fjformer.bit_quantization import no_numerics
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 quantizing 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, no_numerics.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:
		# With floor the biggest value (we are using jnp.max) is in the range of
		# clipping and therefore have a correct gradinet.
		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)
	# We are passing quant_grad (and not more) ot the backward pass.
	# That is equivalent to having:
	# scale = stop_gradient(scale)
	#
	# This is not the only possible choice and we intend to allow experimentation.
	# However for today we hardcoded this choice.
	#
	# In order to achevie no-stop-gradiend solution, we should take vjp
	# of a larger piece of code like the whole _scale_quant.
	#
	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)
	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
	# The axis order in out is as follows: batch, lhs_ra, rhs_ra
	# - batch axes order is uniquely determined by either lhs_ba or rhs_ba
	# - contraction axes ca disappear from the output
	# - order of the remaining axes (ra) is preserved.
	(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(cfg: config.DotGeneralRaw):
	"""Makes quantized lax.dot_general replacement."""

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

	def dot_general_raw(
		lhs: jnp.ndarray,
		rhs: Union[jnp.ndarray, MultiTensor],
		# xhs_qt are used in serving.
		lhs_qt: Optional[QTensor],
		rhs_qt: Optional[QTensor],
		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.

		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_aqt is not None:

			def factor_reshape(x, ca, ba):
				factor = cfg.local_aqt.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
		if lhs_qt is not None:
			lhs_q, lhs_inv_scale_t = (lhs_qt.qvalue, lhs_qt.qvalue_scale_t)
			lhs_quant_grad = "Poison. Not needed in serving"
			lhs_inv_scale = "Poison. Fake quant not used in serving."
			lhs_qx = lhs_qt
		else:
			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)

		if rhs_qt is not None:
			rhs_q, rhs_inv_scale_t = (rhs_qt.qvalue, rhs_qt.qvalue_scale_t)
			rhs_quant_grad = "Poison. Not needed in serving"
			rhs_inv_scale = "Poison. Fake quant not used in serving."
			rhs_qx = rhs_qt
		else:
			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)
		lhs_cast_dtype = cfg.lhs.numerics.get_dtype()
		rhs_cast_dtype = cfg.rhs.numerics.get_dtype()
		if cfg.lhs.use_fake_quant:
			msg = "Can't cast dtype in fake_quant mode."
			assert lhs_cast_dtype is None and rhs_cast_dtype is None, msg
			lhs_q = _maybe_mul(lhs_q, lhs_inv_scale)
			rhs_q = _maybe_mul(rhs_q, rhs_inv_scale)
		else:
			if lhs_cast_dtype is not None:
				lhs_q = lhs_q.astype(lhs_cast_dtype)
			if rhs_cast_dtype is not None:
				rhs_q = rhs_q.astype(rhs_cast_dtype)

		dtype_ms = (
			f"Found {cfg.dg_accumulator_dtype=}, {lhs_cast_dtype=} and"
			f" {rhs_cast_dtype=}. Dot general accumulator dtype can only be"
			" jnp.int32 when both inputs are int8. Otherwise it is recommended to"
			" be None to let lax.dot_general automatically decide it."
		)
		if cfg.dg_accumulator_dtype == jnp.int32:
			assert lhs_cast_dtype == jnp.int8 and rhs_cast_dtype == jnp.int8, dtype_ms
		out = lax.dot_general(
			lhs_q,
			rhs_q,
			dimension_numbers=dimension_numbers,
			preferred_element_type=cfg.dg_accumulator_dtype,
			precision=lax.Precision.DEFAULT,
		).astype(jnp.promote_types(lhs, rhs))

		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_aqt 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,
):
	"""
	The _dot_general_raw_attach_gradient function is a helper function that attaches the gradient to the dot_general_raw
	   function. This allows for backpropagation through the dot general operation. The vjp (vector jacobian product)
	   calculates the gradients of each input tensor with respect to each output tensor, and returns them as a tuple in
	   order of inputs. In this case, it will return dlhs and drhs which are both TensorRes objects.

	:param fwd_dot_general_raw: Compute the forward pass of the dot general operation
	:param dlhs_dot_general_raw: Compute the gradient of dot_general with respect to lhs
	:param drhs_dot_general_raw: Calculate the gradient of rhs
	:param : Specify the function to be used for forward pass
	:return: A function
	"""

	def make_fwd(return_residual):
		def fwd(
			lhs,
			rhs,
			lhs_qt,
			rhs_qt,
			dimension_numbers,
			context,
		):
			assert lhs.dtype == rhs.dtype
			ret, res = fwd_dot_general_raw(
				lhs,
				rhs,
				lhs_qt,
				rhs_qt,
				dimension_numbers,
				context,
			)
			ret = ret.astype(lhs.dtype)
			# We return these values to allow for materialization.
			qret = (res.lhs.mt.qx, res.rhs.mt.qx)
			if return_residual:
				return ((ret, qret), res)
			else:
				return (ret, qret)

		return fwd

	def vjp_bwd(
		fwd_dimension_numbers,
		res: DotGeneralRes,
		g,
	):
		# g[1] is gradient with respect to qret which we are ignoring.
		g = g[0]

		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, None, None, 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,
		)
		# fwd_dimension_numbers are marked as nondiff_argnums instead of returning
		# None as grad to it. This is because it is a tuple of Python integers
		# that cannot be traced by Jax.
		return (dlhs, drhs, None, None, None)

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

	return vjp


def make_dot_general(cfg: Optional[config.DotGeneral]):
	"""
	The make_dot_general function is a wrapper around the dot_general function.
	It takes in two QTensors, lhs and rhs, and returns a QTensor out.
	The make_dot_general function also handles preprocessing of the inputs to dot_general (lhs and rhs)
	and postprocessing of the output from dot_general (out).  The pre-/post-processing steps are:

	:param cfg: Optional[config.DotGeneral]: Specify the configuration of the dot_general operation
	:return: A function that returns a function
	"""
	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),  # noqa :B008
		):
			"""
			The ret_lax_dg function is a wrapper for the jax.lax.dot_general function,
			which performs a general matrix multiplication of two arrays with batch dimensions
			and/or transpositions applied to either or both inputs. The ret_lax_dg function
			is used in the implementation of the dot operation in this module.

			:param lhs: Specify the left-hand side of the dot product
			:param rhs: Specify the right-hand side of the matrix multiplication
			:param dimension_numbers: Specify the dimensions of the operands
			:param precision: Specify the precision of the computation
			:param preferred_element_type: Specify the type of element that should be used to store the result
			:param *: Indicate that all the following parameters are keyword only
			:param context: Pass in the context of the computation
			:param : Specify the dimension numbers of the dot product
			:return: The result of the dot_general operation
			"""
			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), # noqa :B008
	):
		"""
		The ret_dg function is a wrapper around the dg function.
		It takes in two QTensors, lhs and rhs, and returns a QTensor out.
		The ret_dg function also handles preprocessing of the inputs to dg (lhs and rhs)
		and postprocessing of the output from dg (out).  The pre-/post-processing steps are:

		:param lhs: Pass the left hand side of the matrix multiplication
		:param rhs: Pass on the right hand side of the matrix multiplication
		:param dimension_numbers: Specify the contraction pattern
		:param precision: Specify the precision of the output
		:param preferred_element_type: Specify the dtype of the output
		:param *: Indicate that the argument is a keyword-only
		:param context: Pass the context to the dg function
		:return: A function that returns a function
		"""
		del preferred_element_type
		assert (
			precision is None
		), f"Precision {precision} requested together with quantization."

		msg = "AQT is not yet optimized to accept quantized types directly. "
		msg += f"lhs.dtype: {lhs.dtype}, rhs.dtype: {rhs.dtype}"
		assert lhs.dtype in [jnp.bfloat16, jnp.float32, jnp.float16], msg
		assert rhs.dtype in [jnp.bfloat16, jnp.float32, jnp.float16], msg
		if cfg.fwd.lhs.preprocess is not None:
			# lhs_q is quantized dtype.
			# we are breaking the invariant that QTensor has a float qvalue
			# But it will just be cast again to the same type.
			lhs_qt = cfg.fwd.lhs.preprocess(None)
		rhs_qt = None
		if cfg.fwd.rhs.preprocess is not None:
			rhs_qt = cfg.fwd.rhs.preprocess(None)

		out, (out_lhs_qt, out_rhs_qt) = dg(
			lhs=lhs,
			rhs=rhs,
			lhs_qt=lhs_qt,
			rhs_qt=rhs_qt,
			dimension_numbers=dimension_numbers,
			context=context,
		)

		if cfg.fwd.lhs.preprocess is not None:
			lhs_dtype = cfg.fwd.lhs.numerics.get_dtype()
			out_lhs_qt = QTensor(
				out_lhs_qt.qvalue.astype(lhs_dtype), out_lhs_qt.qvalue_scale_t
			)
			none = cfg.fwd.lhs.preprocess(out_lhs_qt)
			assert none is None
		if cfg.fwd.rhs.preprocess is not None:
			rhs_dtype = cfg.fwd.rhs.numerics.get_dtype()
			out_rhs_qt = QTensor(
				out_rhs_qt.qvalue.astype(rhs_dtype), out_rhs_qt.qvalue_scale_t
			)
			none = cfg.fwd.rhs.preprocess(out_rhs_qt)
			assert none is None

		return out

	return ret_dg
