import functools
from typing import Optional, Union, NamedTuple

import jax
from jax import lax
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu
import jax.numpy as jnp
import numpy as np

DEFAULT_MASK_VALUE = -0.7 * float(np.finfo(np.dtype("float32")).max)

P = jax.sharding.PartitionSpec
MAX_INT8 = 127.5


class QuantizedTensor(NamedTuple):
    """A tensor which has been quantized to int8 and its scales.

    Attributes:
      weight: Weight
      scales: Scales
    """

    weight: jnp.ndarray
    scales: jnp.ndarray


def to_int8(x: jnp.ndarray, h: jnp.ndarray) -> jnp.ndarray:
    """Converts a float array to an int8 array with a scale.

    Args:
      x: Float array.
      h: Quantization scale.

    Returns:
      Int8 array.
    """
    return jnp.int8(jnp.rint(x * (MAX_INT8 / h)))


def from_int8(
        x: jnp.ndarray, h: jnp.ndarray, dtype: jnp.dtype = jnp.bfloat16
) -> jnp.ndarray:
    """Converts an int8 array to a float array with a scale.

    Args:
      x: Int8 array.
      h: Quantization scale.
      dtype: Float dtype to convert to.

    Returns:
      Float array.
    """
    return x.astype(dtype) * h / MAX_INT8


def get_quantization_scales(x: jnp.ndarray) -> jnp.ndarray:
    """Computes the quantization scales for a float array.

    These are the maximum values of the trailing dimension.

    Args:
      x: Float array to quantize.

    Returns:
      Array of the same shape as input but with the trailing dimension reduced to
      a size 1 absolute max value.
    """
    return jnp.max(jnp.abs(x), axis=-1, keepdims=True)


def quantize_to_int8(
        x: jnp.ndarray,
) -> QuantizedTensor:
    """Quantizes a float array to an int8 QuantizedTensor.

    Args:
      x: Float array to quantize.

    Returns:
      Int8 QuantizedTensor.
    """
    x_scales = get_quantization_scales(x)
    return QuantizedTensor(weight=to_int8(x, x_scales), scales=x_scales)


def unquantize_from_int8(
        x: QuantizedTensor,
        dtype: jnp.dtype = jnp.bfloat16,
) -> jnp.ndarray:
    """Unquantizes an int8 QuantizedTensor to a float array.

    Args:
      x: Int8 QuantizedTensor to unquantize.
      dtype: Float dtype to unquantize to.

    Returns:
      Float array.
    """
    return from_int8(x.weight, x.scales, dtype)


class MultiPageAsyncCopyDescriptor:
    """Descriptor for async copy of multiple K/V pages from HBM."""

    def __init__(
            self,
            pages_hbm_ref,
            scales_pages_hbm_ref,
            vmem_buffer,
            scales_vmem_buffer,
            sem,
            page_indices,
            page_indices_start_offset,
            num_pages_to_load,
            head_index,
    ):
        self._vmem_buffer = vmem_buffer
        self._scales_vmem_buffer = scales_vmem_buffer
        self._num_pages_to_load = num_pages_to_load
        if head_index is not None:
            self._pages_hbm_ref = pages_hbm_ref.at[head_index]
            if scales_pages_hbm_ref is not None:
                self._scales_pages_hbm_ref = scales_pages_hbm_ref.at[head_index]
            else:
                self._scales_pages_hbm_ref = None
        else:
            self._pages_hbm_ref = pages_hbm_ref
            self._scales_pages_hbm_ref = scales_pages_hbm_ref
        self._sem = sem
        self._page_indices = page_indices
        self._page_indices_start_offset = page_indices_start_offset
        self._async_copies = [
            self._make_async_copy(i) for i in range(self._num_pages_to_load)
        ]
        if (
                self._scales_pages_hbm_ref is not None
                and self._scales_vmem_buffer is not None
        ):
            self._async_copies += [
                self._make_scales_async_copy(i)
                for i in range(self._num_pages_to_load)
            ]

    def _make_async_copy(self, i):
        page_index = self._page_indices[self._page_indices_start_offset + i]
        return pltpu.make_async_copy(
            self._pages_hbm_ref.at[page_index], self._vmem_buffer.at[i], self._sem
        )

    def _make_scales_async_copy(self, i):
        page_index = self._page_indices[self._page_indices_start_offset + i]
        return pltpu.make_async_copy(
            self._scales_pages_hbm_ref.at[page_index],  # pytype: disable=attribute-error
            self._scales_vmem_buffer.at[i],  # pytype: disable=attribute-error
            self._sem,
        )

    def start(self):
        """Starts the async copies."""
        for async_copy in self._async_copies:
            async_copy.start()

    def _maybe_dequantize(self, x, x_scale, dtype=jnp.bfloat16):
        if x_scale is None:
            return x.astype(dtype)
        return from_int8(x, x_scale, dtype=dtype)

    def wait_and_get_loaded(self) -> jax.Array:
        """Wait async copies and gets the loaded buffer as a jax.Array."""
        for async_copy in self._async_copies:
            async_copy.wait()
        head_dim = self._vmem_buffer.shape[-1]
        jax_array = self._vmem_buffer[...].astype(jnp.float32)
        if self._scales_vmem_buffer is not None:
            scales_jax_array = self._scales_vmem_buffer[...].astype(jnp.float32)
        else:
            scales_jax_array = None
        jax_array = self._maybe_dequantize(jax_array, scales_jax_array)
        return jax_array.reshape(-1, head_dim)


def paged_flash_attention_kernel(
        lengths_ref,
        page_indices_ref,
        buffer_index_ref,
        step_ref,
        q_ref,
        k_pages_hbm_ref,
        k_scales_pages_hbm_ref,
        v_pages_hbm_ref,
        v_scales_pages_hbm_ref,
        o_ref,
        m_ref,
        l_ref,
        k_vmem_buffer,
        k_scales_vmem_buffer,
        v_vmem_buffer,
        v_scales_vmem_buffer,
        sem,
        *,
        batch_size: int,
        pages_per_compute_block: int,
        pages_per_sequence: int,
        mask_value: float,
        megacore_mode: str,
        program_ids=(),
):
    """Pallas kernel for paged attention."""
    if program_ids:
        core_index, b, h, i = program_ids  # noqa
    else:
        core_index, b, h, i = (
            pl.program_id(0),
            pl.program_id(1),
            pl.program_id(2),
            pl.program_id(3),
        )
    num_kv_heads, _, page_size, _ = k_pages_hbm_ref.shape
    bk = page_size * pages_per_compute_block
    num_cores = pl.num_programs(0)

    b_step = num_cores if megacore_mode == "batch" else 1
    b_start = core_index if megacore_mode == "batch" else 0
    h_step = num_cores if megacore_mode == "kv_head" else 1
    h_start = core_index if megacore_mode == "kv_head" else 0

    h = h * h_step + h_start
    b = b * b_step + b_start
    length = lengths_ref[b]

    def compute_block_indices(b, h, i):

        def advance_b():
            next_b = b + b_step

            def advance_to_next_non_zero_length():
                next_next_b = next_b + b_step
                return lax.fori_loop(
                    lax.div(next_next_b, b_step),
                    lax.div(batch_size, b_step),
                    lambda _, b: jnp.where(lengths_ref[b] == 0, b + b_step, b),
                    next_next_b,
                )

            return (
                lax.cond(
                    jnp.logical_and(next_b < batch_size, lengths_ref[next_b] == 0),
                    advance_to_next_non_zero_length,
                    lambda: next_b,
                ),
                h_start,
                0,
            )

        def advance_h():
            next_h = h + h_step
            return lax.cond(next_h < num_kv_heads, lambda: (b, next_h, 0), advance_b)

        return lax.cond(i * bk < lengths_ref[b], lambda: (b, h, i), advance_h)

    def create_kv_async_copy_descriptors(b, h, i, buffer_index):
        page_offset = b * pages_per_sequence + i * pages_per_compute_block
        pages_to_load = pages_per_compute_block
        async_copy_k = MultiPageAsyncCopyDescriptor(
            k_pages_hbm_ref,
            k_scales_pages_hbm_ref,
            k_vmem_buffer.at[buffer_index],
            k_scales_vmem_buffer.at[buffer_index]
            if k_scales_vmem_buffer is not None
            else None,
            sem,
            page_indices_ref,
            page_offset,
            pages_to_load,
            h,
        )
        async_copy_v = MultiPageAsyncCopyDescriptor(
            v_pages_hbm_ref,
            v_scales_pages_hbm_ref,
            v_vmem_buffer.at[buffer_index],
            v_scales_vmem_buffer.at[buffer_index]
            if v_scales_vmem_buffer is not None
            else None,
            sem,
            page_indices_ref,
            page_offset,
            pages_to_load,
            h,
        )
        return async_copy_k, async_copy_v

    @pl.when(i * bk < length)
    def flash_attention():
        step = step_ref[0]
        buffer_index = buffer_index_ref[0]

        @pl.when(i == 0)
        def init():
            m_ref[...] = jnp.full_like(m_ref, -jnp.inf)
            l_ref[...] = jnp.zeros_like(l_ref)
            o_ref[...] = jnp.zeros_like(o_ref)

        @pl.when(step == 0)
        def prefetch_first_block():
            async_copy_k, async_copy_v = create_kv_async_copy_descriptors(
                b, h, i, buffer_index
            )
            async_copy_k.start()
            async_copy_v.start()

        next_b, next_h, next_i = compute_block_indices(b, h, i + 1)

        @pl.when(next_b < batch_size)
        def prefetch_next_block():
            next_buffer_index = jnp.where(buffer_index == 0, 1, 0)
            async_copy_next_k, async_copy_next_v = create_kv_async_copy_descriptors(
                next_b, next_h, next_i, next_buffer_index
            )
            async_copy_next_k.start()
            async_copy_next_v.start()
            buffer_index_ref[0] = next_buffer_index

        async_copy_k, async_copy_v = create_kv_async_copy_descriptors(
            b, h, i, buffer_index
        )
        q = q_ref[...].astype(jnp.float32)
        k = async_copy_k.wait_and_get_loaded()
        qk = jnp.einsum('hd,td->ht', q, k, preferred_element_type=jnp.float32)
        mask = i * bk + jax.lax.broadcasted_iota(jnp.int32, qk.shape, 1) < length
        qk = qk + jnp.where(mask, 0.0, mask_value)
        m_curr = qk.max(axis=-1)

        s_curr = jnp.exp(qk - m_curr[..., None])
        m_prev, l_prev = m_ref[...], l_ref[...]
        l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,))
        m_curr = jax.lax.broadcast_in_dim(m_curr, m_prev.shape, (0,))
        m_next = jnp.maximum(m_prev, m_curr)
        alpha = jnp.exp(m_prev - m_next)
        beta = jnp.exp(m_curr - m_next)
        l_next = alpha * l_prev + beta * l_curr
        l_next_safe = jnp.where(l_next == 0.0, 1.0, l_next)

        v = async_copy_v.wait_and_get_loaded()
        o_curr_times_l_curr = jnp.dot(s_curr, v)

        m_ref[...], l_ref[...] = m_next, l_next_safe
        o_ref[...] = (
                (l_prev * alpha * o_ref[...] + beta * o_curr_times_l_curr) / l_next_safe
        ).astype(o_ref.dtype)

        step_ref[0] = step + 1


def paged_flash_attention_kernel_inline_seq_dim(
        lengths_ref,
        page_indices_ref,
        buffer_index_ref,
        step_ref,
        q_ref,
        k_pages_hbm_ref,
        k_scales_pages_hbm_ref,
        v_pages_hbm_ref,
        v_scales_pages_hbm_ref,
        o_ref,
        m_ref,
        l_ref,
        k_vmem_buffer,
        k_scales_vmem_buffer,
        v_vmem_buffer,
        v_scales_vmem_buffer,
        sem,
        *,
        batch_size: int,
        pages_per_compute_block: int,
        pages_per_sequence: int,
        mask_value: float,
        megacore_mode: str,
):
    core_index, b, h = pl.program_id(0), pl.program_id(1), pl.program_id(2)

    # Initialize the output HBM buffers to avoid accessing garbage memory inside
    # the kernel body below.
    m_ref[...] = jnp.full_like(m_ref, -jnp.inf)
    l_ref[...] = jnp.zeros_like(l_ref)
    o_ref[...] = jnp.zeros_like(o_ref)

    def body(i, _):
        paged_flash_attention_kernel(
            lengths_ref,
            page_indices_ref,
            buffer_index_ref,
            step_ref,
            q_ref,
            k_pages_hbm_ref,
            k_scales_pages_hbm_ref,
            v_pages_hbm_ref,
            v_scales_pages_hbm_ref,
            o_ref,
            m_ref,
            l_ref,
            k_vmem_buffer,
            k_scales_vmem_buffer,
            v_vmem_buffer,
            v_scales_vmem_buffer,
            sem,
            batch_size=batch_size,
            pages_per_compute_block=pages_per_compute_block,
            pages_per_sequence=pages_per_sequence,
            mask_value=mask_value,
            megacore_mode=megacore_mode,
            program_ids=(core_index, b, h, i),
        )
        return ()

    bk = pages_per_compute_block * k_pages_hbm_ref.shape[-2]

    if megacore_mode == "batch":
        num_cores = pl.num_programs(0)
        length = lengths_ref[b * num_cores + core_index]
    else:
        length = lengths_ref[b]

    lax.fori_loop(0, lax.div(length + bk - 1, bk), body, ())


@functools.partial(
    jax.jit,
    static_argnames=[
        "pages_per_compute_block",
        "mask_value",
        "megacore_mode",
        "inline_seq_dim",
    ],
)
def paged_attention(
        q: jax.Array,
        k_pages: Union[jax.Array, QuantizedTensor],
        v_pages: Union[jax.Array, QuantizedTensor],
        lengths: jax.Array,
        page_indices: jax.Array,
        *,
        mask_value: float = DEFAULT_MASK_VALUE,
        pages_per_compute_block: int,
        megacore_mode: Optional[str] = None,
        inline_seq_dim: bool = True,
) -> jax.Array:
    """Paged grouped query attention.

    Args:
      q: A [batch_size, num_heads, head_dim] jax.Array.
      k_pages: A [num_kv_heads, total_num_pages, page_size, head_dim] jax.Array.
      v_pages: A [num_kv_heads, total_num_pages, page_size, head_dim] jax.Array.
      lengths: A i32[batch_size] jax.Array the length of each example.
      page_indices: A i32[batch_size, pages_per_sequence] jax.Array. Each entry
        should be in the range of [0, total_num_pages), indicating where to locate
        the page in `k_pages` or `v_pages`.
      mask_value: The value used for padding in attention. By default it is a very
        negative floating point number.
      pages_per_compute_block: how many pages to be processed in one flash
        attention block in the pallas kernel.
      megacore_mode: if set, enable megacore to parallelize the computation. Must
        be one of ['kv_head', 'batch', None]. Caveat: set this only if megacore is
        enabled, otherwise the kernel may hang. If you are not sure, leave it to
        None.
        * None: disable megacore parallelism.
        * kv_head: megacore parallelism on KV heads; requires number of KV heads
          divisible by 2.
        * batch: megacore parallelism on batch dimension; requires batch divisible
          by 2.
      inline_seq_dim: whether to fuse kernel instances along the sequence dim into
        one kernel.

    Returns:
      The output of attention([batch_size, num_heads, head_dim]).
    """
    if isinstance(k_pages, QuantizedTensor):
        k_pages, k_scales_pages = k_pages.weight, k_pages.scales
        assert isinstance(k_scales_pages, jax.Array)  # For typing.
        k_scales_pages = jnp.broadcast_to(
            k_scales_pages, (*k_scales_pages.shape[:-1], k_pages.shape[-1])
        )
    else:
        k_scales_pages = None
    if isinstance(v_pages, QuantizedTensor):
        v_pages, v_scales_pages = v_pages.weight, v_pages.scales
        assert isinstance(v_scales_pages, jax.Array)  # For typing.
        v_scales_pages = jnp.broadcast_to(
            v_scales_pages, (*v_scales_pages.shape[:-1], v_pages.shape[-1])
        )
    else:
        v_scales_pages = None

    batch_size, num_heads, head_dim = q.shape
    num_kv_heads, _, page_size, head_dim_k = k_pages.shape
    batch_size_paged_indices, pages_per_sequence = page_indices.shape

    if k_pages.shape != v_pages.shape:
        raise ValueError(
            f"k_pages and v_pages must have the same shape. Got {k_pages.shape} and"
            f" {v_pages.shape}"  # pytype: disable=attribute-error
        )
    if num_heads % num_kv_heads != 0:
        raise ValueError(
            "Number of Q heads must be divisible by number of KV heads. Got"
            f" {num_heads} and {num_kv_heads}."
        )
    if head_dim_k != head_dim:
        raise ValueError(
            "head_dim of Q must be the same as that of K/V. Got"
            f" {head_dim} and {head_dim_k}."
        )
    if pages_per_sequence % pages_per_compute_block != 0:
        raise ValueError(
            "pages_per_compute_block must be divisible by pages per sequence. Got"
            f" {pages_per_compute_block} and {pages_per_sequence}."
        )
    if lengths.shape != (batch_size,):
        raise ValueError("`lengths` and `q` must have the same batch size")
    if batch_size_paged_indices != batch_size:
        raise ValueError("`page_indices` and `q` must have the same batch size")
    if lengths.dtype != jnp.int32:
        raise ValueError(
            "The dtype of `lengths` must be int32. Got {lengths.dtype}"
        )

    if megacore_mode == "kv_head":
        if num_kv_heads % 2 != 0:
            raise ValueError(
                "number of KV heads must be even when megacore_mode is 'kv_head'"
            )
        num_cores = 2
    elif megacore_mode == "batch":
        if batch_size % 2 != 0:
            raise ValueError("batch size must be even when megacore_mode is 'batch'")
        num_cores = 2
    elif megacore_mode is None:
        num_cores = 1
    else:
        raise ValueError("megacore_mode must be one of ['kv_head', 'batch', None]")

    if (num_heads // num_kv_heads) % 8 != 0:
        # Reshape q to hint XLA to pick a <1x128> layout otherwise it will pick a
        # <8x128> layout for a <1x128> memref inside the kernel and error out.
        q = q.reshape(batch_size, num_heads, 1, head_dim)
        if megacore_mode == "kv_head":
            q_block_spec = pl.BlockSpec(
                lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0, 0),
                (None, num_heads // num_kv_heads, None, head_dim),
            )
        elif megacore_mode == "batch":
            q_block_spec = pl.BlockSpec(
                lambda core_index, b, h, *_: (b * num_cores + core_index, h, 0, 0),
                (None, num_heads // num_kv_heads, None, head_dim),
            )
        else:
            q_block_spec = pl.BlockSpec(
                lambda core_index, b, h, *_: (b, h, 0, 0),
                (None, num_heads // num_kv_heads, None, head_dim),
            )
        q_dtype_for_kernel_launch = jnp.float32
    else:
        if megacore_mode == "kv_head":
            q_block_spec = pl.BlockSpec(
                lambda core_index, b, h, *_: (b, h * num_cores + core_index, 0),
                (None, num_heads // num_kv_heads, head_dim),
            )
        elif megacore_mode == "batch":
            q_block_spec = pl.BlockSpec(
                lambda core_index, b, h, *_: (b * num_cores + core_index, h, 0),
                (None, num_heads // num_kv_heads, head_dim),
            )
        else:
            q_block_spec = pl.BlockSpec(
                lambda core_index, b, h, *_: (b, h, 0),
                (None, num_heads // num_kv_heads, head_dim),
            )
        q_dtype_for_kernel_launch = q.dtype

    if inline_seq_dim:
        kernel = paged_flash_attention_kernel_inline_seq_dim
        grid = (
            num_cores,
            batch_size // num_cores if megacore_mode == "batch" else batch_size,
            num_kv_heads // num_cores
            if megacore_mode == "kv_head"
            else num_kv_heads,
        )
        dimension_sematics = ("parallel", "arbitrary", "arbitrary")
    else:
        kernel = paged_flash_attention_kernel
        grid = (
            num_cores,
            batch_size // num_cores if megacore_mode == "batch" else batch_size,
            num_kv_heads // num_cores
            if megacore_mode == "kv_head"
            else num_kv_heads,
            pages_per_sequence // pages_per_compute_block,
        )  # type: ignore
        dimension_sematics = ("parallel", "arbitrary", "arbitrary", "arbitrary")  # type: ignore

    if k_scales_pages is not None and v_scales_pages is not None:
        in_specs = [
            q_block_spec,
            pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
            pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
            pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
            pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
        ]
        scratch_shapes = (
            pltpu.VMEM(
                (
                    2,  # For double buffering during DMA copies.
                    pages_per_compute_block,
                    page_size,
                    head_dim,
                ),
                k_pages.dtype,
            ),  # k_pages buffer
            pltpu.VMEM(
                (
                    2,  # For double buffering during DMA copies.
                    pages_per_compute_block,
                    page_size,
                    head_dim,
                ),
                k_scales_pages.dtype,  # pytype: disable=attribute-error
            ),  # k_scales_pages buffer
            pltpu.VMEM(
                (
                    2,  # For double buffering during DMA copies.
                    pages_per_compute_block,
                    page_size,
                    head_dim,
                ),
                v_pages.dtype,
            ),  # v_pages buffer
            pltpu.VMEM(
                (
                    2,  # For double buffering during DMA copies.
                    pages_per_compute_block,
                    page_size,
                    head_dim,
                ),
                v_scales_pages.dtype,  # pytype: disable=attribute-error
            ),  # v_scales_pages buffer
            pltpu.SemaphoreType.DMA,
        )
    else:
        in_specs = [
            q_block_spec,
            pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
            None,  # type: ignore[list-item]
            pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
            None,  # type: ignore[list-item]
        ]
        scratch_shapes = (
            pltpu.VMEM(
                (
                    2,  # For double buffering during DMA copies.
                    pages_per_compute_block,
                    page_size,
                    head_dim,
                ),
                k_pages.dtype,
            ),  # k_pages buffer
            None,
            pltpu.VMEM(
                (
                    2,  # For double buffering during DMA copies.
                    pages_per_compute_block,
                    page_size,
                    head_dim,
                ),
                v_pages.dtype,
            ),  # v_pages buffer
            None,
            pltpu.SemaphoreType.DMA,
        )

    out, _, _ = pl.pallas_call(
        functools.partial(
            kernel,
            pages_per_sequence=pages_per_sequence,
            batch_size=batch_size,
            pages_per_compute_block=pages_per_compute_block,
            mask_value=mask_value,
            megacore_mode=megacore_mode,
        ),
        grid_spec=pltpu.PrefetchScalarGridSpec(
            # There are 4 scalars prefetched per kernel call: `lengths_ref`,
            # `page_indices_ref`, `buffer_index_ref`, `step_ref`
            num_scalar_prefetch=4,
            in_specs=in_specs,
            out_specs=[
                q_block_spec,
                q_block_spec,
                q_block_spec,
            ],
            grid=grid,
            scratch_shapes=scratch_shapes,
        ),
        compiler_params=dict(mosaic=dict(dimension_semantics=dimension_sematics)),
        out_shape=[
            jax.ShapeDtypeStruct(q.shape, q_dtype_for_kernel_launch),
            jax.ShapeDtypeStruct((*q.shape[:-1], 1), jnp.float32),
            jax.ShapeDtypeStruct((*q.shape[:-1], 1), jnp.float32),
        ],
    )(
        lengths,
        page_indices.reshape(-1),
        jnp.zeros((1,), jnp.int32),  # buffer index
        jnp.zeros((1,), jnp.int32),  # step
        q.astype(q_dtype_for_kernel_launch),
        k_pages,
        k_scales_pages,
        v_pages,
        v_scales_pages,
    )
    return out.reshape(batch_size, num_heads, head_dim).astype(q.dtype)
