import torch
import triton
import triton.language as tl
from dataclasses import dataclass
from .base import Layout, LayoutTransformation
from triton_kernels.numerics_details.mxfp_details._downcast_to_mxfp import MXFP_BLOCK_SIZE
from triton_kernels.target_info import cuda_capability_geq
from .torch_utils import repack


# ------------------- Hopper MX Value Layout -------------------
@dataclass(frozen=True)
class HopperMXValueLayout(Layout):
    mx_axis: int
    mma_version: int

    def __post_init__(self):
        assert self.mx_axis in [-1, -2]
        assert self.mma_version in [2, 3]

    @property
    def name(self):
        return "HOPPER_VALUE"

    def swizzle_block_shape(self, block_shape):
        if self.mx_axis == -1:
            *head, N, K = block_shape
            assert N % 4 == 0
            return [*head, N // 4, K * 4]
        else:
            assert self.mx_axis == -2
            *head, K, N = block_shape
            assert N % 4 == 0
            return [*head, K * 4, N // 4]

    def make_transformation(self, shape: list[int], is_fp4) -> LayoutTransformation:
        return HopperMXValueLayoutTransformation(shape, is_fp4, self.mx_axis, self.mma_version)


# ------------------- Hopper MX Value Layout Transformation -------------------


@dataclass(frozen=True)
class HopperMXValueLayoutTransformation(LayoutTransformation):
    mx_axis: int
    mma_version: int

    def __post_init__(self):
        *leading_shape, K, N, = self.shape
        if self.mx_axis < 0:
            object.__setattr__(self, "mx_axis", self.mx_axis + len(self.shape))
        object.__setattr__(self, "leading_shape", leading_shape)
        object.__setattr__(self, "K", K)
        object.__setattr__(self, "N", N)

    def _padded_shape(self, shape) -> list[int]:
        *leading_shape, M, K = shape
        if self.mx_axis == len(leading_shape):
            align_m, align_k = 64, 256
        else:
            align_m, align_k = 256, 64
        M = (M + align_m - 1) // align_m * align_m
        K = (K + align_k - 1) // align_k * align_k
        return [*leading_shape, M, K]

    @property
    def storage_shape(self) -> list[int]:
        *leading_shape, M, K = self.shape
        if self.is_fp4:
            K //= 2
            if self.mx_axis == len(leading_shape):
                M //= 2
                K *= 2
        *leading_shape, M, K = self._padded_shape((*leading_shape, M, K))
        if self.mx_axis == len(leading_shape):
            return [*leading_shape, M * 4, K // 4]
        return [*leading_shape, M // 4, K * 4]

    def _maybe_mT(self, data):
        if self.mx_axis == len(self.leading_shape):
            return data.mT
        return data

    def swizzle_data(self, data):
        """
        Given a uint8 tensor of shape (*, M, K), returns a tensor of shape
        (*, M // 4, K * 4) such that:

        1) Groups contiguously all the elements owned by the same thread of 4
        mma tiles along the K axis. The following animation shows a similar
        grouping for 2 tiles along M and 2 tiles along K rather than 4 along K
        as done here:
        https://neuralmagic.com/wp-content/uploads/2024/10/animation_4.gif

        2) Moves the elements belonging to thread 4-7 to be contiguous with those
        from thread 0-3. This is done to get a full cache line when loading them
        from HBM.

        mx_axis selects the lhs or rhs of the matmul.

        WARNING: Assumes that the matmul will be done in bf16 or fp16!
        Implementing it for fp8 is as easy as making the tile size (8, 8)
        """
        # re-pack as column-major
        data = repack(data, -1, self.mx_axis, self.is_fp4)
        batch = data.ndim - 2
        assert batch >= 0
        assert self.mma_version in (2, 3)
        # Align the dimension packed by four to a 64-byte load extent.
        *_, M_in, K_in = data.shape
        *_, M, K = self._padded_shape(data.shape)
        pad_m = M - M_in
        pad_k = K - K_in
        data = torch.nn.functional.pad(data, (0, pad_k, 0, pad_m))

        data = self._maybe_mT(data)
        init_shape = data.shape

        # We are loading 8 bf16 elements per thread to use ld.global.v4
        # Every u8 represents 2 mxfp4 elements
        u8_kwidth = 8 // 2 if self.mma_version == 2 else 1

        # Pack the 4 // u8_kwidth subtiles of an mma into a u4x8
        contig = (1, u8_kwidth)
        scott_trick = (2, 1)
        threads = (4, 4)
        warp_tile = (2, 2)
        k_tile = (1, 4 // u8_kwidth)

        sizes = list(data.shape[:-2])
        # [rest, K, tile, threads] per dimension
        for i, (a, b, c, s, d) in enumerate(zip(k_tile, warp_tile, threads, scott_trick, contig)):
            packed = a * b * c * s * d
            size = data.shape[batch + i]
            sizes.append(size // packed)
            sizes += [a, b, c, s, d]

        # 0: rest[0]
        # 1: k_tile[0]
        # 2: warp_tile[0]
        # 3: threads[0]
        # 4: scott_trick[0]
        # 5: contig[0]
        # 6: rest[1]
        # 7: k_tile[1]
        # 8: warp_tile[1]
        # 9: threads[1]
        # 10: scott_trick[1]
        # 11: contig[1]
        data = data.view(*sizes)
        # Want [rest[0], threads[0], rest[1], scott_trick[0], scott_trick[0], threads[1], contig[1], contig[0], k_tile[1], k_tile[0], warp_tile[1], warp_tile[0]]
        perm = [0, 3, 6, 10, 4, 9, 7, 1, 8, 2, 5, 11]
        perm = list(range(batch)) + [batch + p for p in perm]
        data = data.permute(*perm).contiguous()
        # These are views
        data = data.flatten(-10, -1)
        data = data.flatten(-3, -2)
        assert data.is_contiguous()
        assert data.shape[-2] == init_shape[-2] // 4
        assert data.shape[-1] == init_shape[-1] * 4
        # twiddle the bits
        data = _pack_bits(data, self.mx_axis)
        data = self._maybe_mT(data)
        return self._validate_storage_shape(data)

    def unswizzle_data(self, data):
        data = self._maybe_mT(data)
        data = _unpack_bits(data, self.mx_axis)
        *batch, M, K = data.shape
        # We have two times the elements if we already upcasted to bfloat16
        mult = 2 if data.dtype == torch.bfloat16 else 1
        assert M % 4 == 0, "M must be divisible by 4"
        assert K % (4 * 8 * 2 * 2 * mult) == 0, f"K must be divisible by {4 * 8 * 2 * 2 * mult}"
        # We are loading 8 bf16 elements per thread to use ld.global.v4
        # Every u8 represents 2 mxfp4 elements
        u8_kwidth = 8 // 2 if self.mma_version == 2 else 1
        data = data.reshape(*batch, M // 4, 4, K // (4 * 8 * 2 * 2 * mult), 2, 4, 8 // u8_kwidth, 2, u8_kwidth * mult)
        b = len(batch)
        perm = [0, 6, 1, 3, 2, 5, 4, 7]
        perm = list(range(b)) + [b + p for p in perm]
        data = data.permute(*perm)
        data = data.reshape(*batch, M * 4, K // 4)
        data = self._maybe_mT(data)
        data = repack(data, self.mx_axis, -1, self.is_fp4)
        data = data[..., :self.K, :self.N // 2]
        data = data.contiguous()
        return data


def right_shift_unsigned(x, shift):
    return (x >> shift) & ((1 << (32 - shift)) - 1)


# -----------------------------------------------------------------------
# Interleave the bits of four consecutive fp4 values (i.e. 16-bits) as:
#     1000000111000000         (first fp4)
#        1000000111000000      (second fp4)
#           1000000111000000   (third fp4)
#     0110110000000000         (fourth fp4)
# This is done so that dequantization can be done in 14 SASS instructions
# -----------------------------------------------------------------------


def _compress_fp4(x):
    x = x.to(torch.int32)
    return ((x & 0x8) << 12) | ((x & 0x7) << 6)


def _compress_fourth(x):
    x = x.to(torch.int32)
    return ((x & 0x8) << 11) | ((x & 0x6) << 9) | ((x & 0x1) << 13)


def _pack_bits(x: torch.Tensor, mx_axis: int):
    x = x.contiguous()
    x = x.reshape(x.shape[:-1] + (x.shape[-1] // 4, 4))
    ret = _compress_fp4(x[..., 0]) | (_compress_fp4(x[..., 0] >> 4) << 16)
    ret |= right_shift_unsigned(_compress_fp4(x[..., 1]) | (_compress_fp4(x[..., 1] >> 4) << 16), 3)
    ret |= right_shift_unsigned(_compress_fp4(x[..., 2]) | (_compress_fp4(x[..., 2] >> 4) << 16), 6)
    ret |= _compress_fourth(x[..., 3]) | (_compress_fourth(x[..., 3] >> 4) << 16)
    assert ret.is_contiguous()
    ret = ret.view(torch.uint8)
    return ret


# -----------------------------------------------------------------------
# inverse operation of _pack_bits
# -----------------------------------------------------------------------


def _bf16_to_fp4e2m1(x):
    # 0bAxxxxxxBCDxxxxxx (int16) -> 0b0000ABCD (uint8)
    assert x.dtype == torch.int16
    s = (right_shift_unsigned(x, 15) & 0x1) << 3
    em = right_shift_unsigned(x, 6) & 0x7
    return (s | em).to(torch.uint8)


def _bf16x2_to_fp4e2m1x2(x):
    # 0bAxxxxxxBCDxxxxxx_0bExxxxxxFGHxxxxxx  (int32) -> 0bABCD_EFGH (uint8)
    assert x.dtype == torch.int32
    lo = (x & 0xFFFF).to(torch.int16)
    hi = (right_shift_unsigned(x, 16) & 0xFFFF).to(torch.int16)
    ret_lo = _bf16_to_fp4e2m1(lo)
    ret_hi = _bf16_to_fp4e2m1(hi)
    return ret_lo | (ret_hi << 4)


def _unpack_bits(x, mx_axis: int):
    x = x.view(torch.int32)
    m = 0b10000001110000001000000111000000
    a = (x << 1) & 0b10000000000000001000000000000000
    b = right_shift_unsigned(x, 3) & 0b00000001100000000000000110000000
    c = right_shift_unsigned(x, 7) & 0b00000000010000000000000001000000
    unpacked = [x & m, (x << 3) & m, (x << 6) & m, (a | b) | c]
    x = torch.stack(unpacked, dim=-1)
    x = x.flatten(-2, -1)
    x = _bf16x2_to_fp4e2m1x2(x)
    return x


# -----------------------------------------------------------------------


@triton.jit
def _unshuffle_triton(x, mma_version: tl.constexpr):
    """
    Triton inverse of swizzle_mxfp4_value_hopper
    """
    tl.static_assert(mma_version == 2 or mma_version == 3, "mma_version must be 2 or 3")
    # if mx_axis == 0:
    #     x = x.trans()

    # We have two times the elements if we already upcasted to bfloat16
    mult: tl.constexpr = 2 if x.dtype == tl.bfloat16 else 1
    M: tl.constexpr = x.shape[0]
    K: tl.constexpr = x.shape[1]
    tl.static_assert(M % 4 == 0, "M must be divisible by 4")
    tl.static_assert(K % (4 * 8 * 2 * 2 * mult) == 0, f"K must be divisible by {4 * 8 * 2 * 2 * mult}")

    # We are loading 8 bf16 elements per thread to use ld.global.v4
    # Every u8 represents 2 mxfp4 elements
    u8_kwidth: tl.constexpr = 8 // 2 if mma_version == 2 else 1
    x = x.reshape(M // 4, 4, K // (4 * 8 * 2 * 2 * mult), 2, 4, 8 // u8_kwidth, 2, u8_kwidth * mult)
    x = x.trans(0, 6, 1, 3, 2, 5, 4, 7)
    x = x.reshape(M * 4, K // 4)
    # if mx_axis == 0:
    #     x = x.trans()
    return x


@triton.jit
def _unpack_fp4_to_bf16_triton(x):
    # Use fma on a100 as there is no mul.bf16x2.
    use_mul: tl.constexpr = cuda_capability_geq(9)
    op_instr: tl.constexpr = "mul.bf16x2" if use_mul else "fma.rn.bf16x2"
    op_suffix: tl.constexpr = "" if use_mul else ", z"
    r0, r1 = tl.inline_asm_elementwise(
        asm=f"""{{
            .reg .b32 b, c, z, d<7>, scale;
            .reg .b32 bias;
            mov.b32 z, 0;
            mov.b32 bias, 0x7e807e80; // 2 ** 126 == 2 ** (bias_bf16 - bias_fp2)
            // We add the missing bias to the scale directly
            and.b32 $0, $4, 0b10000001110000001000000111000000;
            {op_instr} $0, $0, bias{op_suffix};
            shl.b32 b, $4, 3;
            and.b32 $1, b,  0b10000001110000001000000111000000;
            {op_instr} $1, $1, bias{op_suffix};
            shl.b32 c, $4, 6;
            and.b32 $2, c,  0b10000001110000001000000111000000;
            {op_instr} $2, $2, bias{op_suffix};
            // Unpack last two elements
            shl.b32 d0, $4, 1;
            and.b32 d1, d0, 0b10000000000000001000000000000000;
            shr.b32 d2, $4, 3;
            and.b32 d3, d2, 0b00000001100000000000000110000000;
            or.b32 d4, d1, d3;
            shr.b32 d5, $4, 7;
            and.b32 d6, d5, 0b00000000010000000000000001000000;
            or.b32 $3, d4, d6;
            {op_instr} $3, $3, bias{op_suffix};
        }}""",
        constraints="=r,=r,=r,=r,r",
        args=[x],
        dtype=(tl.bfloat16, tl.bfloat16),
        is_pure=True,
        pack=4,
    )
    # Concat each pack of 4
    x = tl.join(r0, r1)
    x = x.reshape(x.shape[0], x.shape[1] // 4, 4, x.shape[2])
    x = x.trans(0, 1, 3, 2)
    x = x.reshape(x.shape[0], x.shape[1] * x.shape[2] * x.shape[3])
    return x


@triton.jit
def mul_bf16x2(a, b):
    use_mul: tl.constexpr = cuda_capability_geq(9)
    op_instr: tl.constexpr = "mul.bf16x2" if use_mul else "fma.rn.bf16x2"
    op_suffix: tl.constexpr = "" if use_mul else ", z"

    return tl.inline_asm_elementwise(
        asm=f"{op_instr} $0, $1, $2{op_suffix};",
        constraints="=r,r,r",
        args=[a, b],
        dtype=tl.bfloat16,
        is_pure=True,
        pack=2,
    )


@triton.jit
def mxfp4_to_bf16_triton(x, scale, mx_axis: tl.constexpr):
    """
    Implements the bit-untwiddling of a 32-bit integer (8 mxfp4 elements):
    (x << 0) & 0b1000000111000000
    (x << 3) & 0b1000000111000000
    (x << 6) & 0b1000000111000000
    ((x << 1) & 0b1000000000000000) | ((x >> 3) & 0b0000000110000000) | ((x >> 7) & 0b0000000001000000)
    """
    # upcast values to bfloat16
    tl.static_assert(len(x.shape) == 2)
    tl.static_assert(mx_axis == 0 or mx_axis == 1, "mx_axis must be 0 or 1")
    tl.static_assert(x.shape[1] % 4 == 0)
    tl.static_assert(x.dtype == tl.uint8)
    if mx_axis == 0:
        x = x.trans()
    x = _unpack_fp4_to_bf16_triton(x)
    x = _unshuffle_triton(x, mma_version=3)
    if mx_axis == 0:
        x = x.trans()

    # upcast scale to bfloat16
    # Add bias missing from the bf16 upcasting sequence
    # triton / LLVM generates terrible code for this sequence
    # scale = scale.to(tl.uint16)
    # scale = scale << 7
    # scale = scale.to(tl.bfloat16, bitcast=True)
    scale = tl.inline_asm_elementwise(
        r"""
        {
            prmt.b32 $0, $2, 0, 0x5140;
            shl.b32 $0, $0, 7;
            prmt.b32 $1, $2, 0, 0x7362;
            shl.b32 $1, $1, 7;
        }
        """,
        constraints="=r,=r,r",
        args=[scale],
        dtype=tl.bfloat16,
        is_pure=True,
        pack=4,
    )
    # Sanity check shape
    for axis in tl.static_range(len(x.shape)):
        if axis == mx_axis:
            tl.static_assert(x.shape[axis] == MXFP_BLOCK_SIZE * scale.shape[axis])
        else:
            tl.static_assert(x.shape[axis] == scale.shape[axis])
    # Broadcast scale
    scale = scale.expand_dims(mx_axis + 1)
    scale = scale.broadcast_to(scale.shape[:mx_axis + 1] + [MXFP_BLOCK_SIZE] + scale.shape[mx_axis + 2:])
    scale = scale.reshape(x.shape)

    # Combine scale and x
    x = mul_bf16x2(x, scale)
    return x
