from typing import Any

import torch

from fastvideo.layers.quantization.nvfp4_config import nvfp4_quantize_fenced

try:
    import flashinfer
except ImportError:
    flashinfer = None


def _require_flashinfer() -> Any:
    if flashinfer is None:
        raise ImportError("flashinfer is required for FP4 linear layers. "
                          "Please install flashinfer to use this path.")
    return flashinfer


def _global_sf(t: torch.Tensor) -> torch.Tensor:
    # Eager amax: the previous torch.compile'd reduction lowered to a single-block
    # Triton kernel (~38 ms for a 21k x 2048 activation on GB200). Max is exact
    # in any order, so finite input gives the same scale bit for bit; NaN and inf
    # keep the elementwise reduction's behavior. No host sync.
    maxabs = torch.nan_to_num(t.abs()).amax().float()
    # The reduction above runs in the input dtype, so an infinite input becomes the
    # fp16/bf16 maximum. The reference reduction replaces it in float32.
    # new_full fills on the device; torch.tensor would copy from the host and sync.
    maxabs = torch.where(torch.isinf(t).any(), maxabs.new_full((), torch.finfo(torch.float32).max), maxabs)
    return (448.0 * 6.0) / maxabs.clamp(min=1e-12)


class _LinearFWD4BWD16Fn(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, weight, bias, backend="cutlass", block_size=16, use_128x4_sf_layout=True):
        flashinfer_mod = _require_flashinfer()

        # assert activation dtype
        if x.dtype not in (torch.float16, torch.bfloat16):
            x = x.to(dtype=torch.bfloat16)

        # cast params (can be fp32) to activation dtype for quantization
        weight_cast = weight.to(dtype=x.dtype)
        bias_cast = bias.to(dtype=x.dtype) if bias is not None else None

        # shapes
        orig_shape = x.shape
        k = weight_cast.shape[1]
        n = weight_cast.shape[0]
        x2d = x.reshape(-1, k).contiguous()
        M = x2d.shape[0]

        out2d = torch.empty((M, n), device=x.device, dtype=x.dtype)

        a_sf_layout = (flashinfer_mod.SfLayout.layout_128x4
                       if use_128x4_sf_layout else flashinfer_mod.SfLayout.layout_8x4)
        global_sf_a = _global_sf(x2d)
        global_sf_b = _global_sf(weight_cast)

        # Shared helper: carries the DGX Spark (GB10) ordering fence.
        a_fp4, a_inv_s = nvfp4_quantize_fenced(x2d, global_sf_a, a_sf_layout.value)
        b_fp4, b_inv_s = nvfp4_quantize_fenced(weight_cast, global_sf_b, flashinfer_mod.SfLayout.layout_128x4.value)

        alpha = 1.0 / (global_sf_a * global_sf_b)

        flashinfer_mod.mm_fp4(
            a_fp4,
            b_fp4.T,
            a_inv_s,
            b_inv_s.T,
            alpha,
            x.dtype,
            out2d,
            block_size=block_size,
            use_8x4_sf_layout=(not use_128x4_sf_layout),
            backend=backend,
        )

        if bias_cast is not None:
            out2d.add_(bias_cast)

        # save tensors for backward (keep original dtypes)
        ctx.save_for_backward(x2d, weight, bias)
        ctx.k = k
        ctx.n = n
        ctx.orig_shape = orig_shape
        return out2d.reshape(*orig_shape[:-1], n)

    @staticmethod
    def backward(ctx, grad_out):
        x2d, weight, bias = ctx.saved_tensors
        M = x2d.shape[0]
        n = ctx.n

        grad_out_2d = grad_out.reshape(M, n).contiguous()

        # cast to grad dtype for matmuls
        weight_cast = weight.to(dtype=grad_out.dtype)
        x_cast = x2d.to(dtype=grad_out.dtype)

        grad_x = grad_out_2d.matmul(weight_cast).reshape(*ctx.orig_shape)
        grad_w = grad_out_2d.t().matmul(x_cast)
        grad_b = grad_out_2d.sum(dim=0) if bias is not None else None

        # None for the three extra forward args
        return grad_x, grad_w, grad_b, None, None, None


def fp4_linear_forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
    # pass config **positionally**; autograd.Function.apply ignores kwargs
    bias = self.bias if not self.skip_bias_add else None
    output = _LinearFWD4BWD16Fn.apply(x, self.weight, bias, "cutlass", 16, True)
    output_bias = self.bias if self.skip_bias_add else None
    return output, output_bias