"""
Utilities and parameterization for Grouped GEMM dSReLU backward tests.
Contains test configuration fixtures, tensor creation, and reference implementations.
"""

import torch
import pytest
from typing import Optional, Tuple, List, Dict, Any
from test_fe_api_utils import (
    ceil_div,
    compute_reference_amax,
    create_and_permute_tensor,
    create_scale_factor_tensor,
    create_sf_layout_tensor,
    cvt_sf_MKL_to_M32x4xrm_K4xrk_L,
)
from test_grouped_gemm_swiglu_utils import (
    allocate_grouped_gemm_input_tensors as allocate_grouped_gemm_input_tensors_base,
    get_dtype_rcp_limits as get_grouped_gemm_dtype_rcp_limits,
    grouped_gemm_swiglu_init as grouped_gemm_dsrelu_init,
)

# =============================================================================
# Parameterization Marks
# =============================================================================

GROUPED_GEMM_DSRELU_COMMON_MARKS = [
    pytest.mark.parametrize("cd_major", ["n"]),
    pytest.mark.parametrize("acc_dtype", [torch.float32]),
    pytest.mark.parametrize("mma_tiler_mn", [(256, 256), (128, 256)]),
    pytest.mark.parametrize("cluster_shape_mn", [(2, 1), (1, 1)]),
    pytest.mark.parametrize("vector_f32", [True, False]),
]

GROUPED_GEMM_DSRELU_FP8_TYPE_MARKS = [
    pytest.mark.parametrize(
        "ab_dtype",
        [
            torch.float8_e4m3fn,
        ],
    ),
    pytest.mark.parametrize("c_dtype", [torch.bfloat16]),
    pytest.mark.parametrize(
        "d_dtype",
        [
            torch.float8_e4m3fn,
        ],
    ),
    pytest.mark.parametrize("b_major", ["k"]),
]

GROUPED_GEMM_DSRELU_FP4_TYPE_MARKS = [
    pytest.mark.parametrize(
        "ab_dtype",
        [
            torch.uint8,
        ],
    ),
    pytest.mark.parametrize(
        "c_dtype",
        [
            # torch.float16,
            torch.bfloat16,
            # torch.float32,
        ],
    ),
    pytest.mark.parametrize(
        "d_dtype",
        [
            # torch.float16,
            torch.bfloat16,
            torch.float32,
        ],
    ),
    pytest.mark.parametrize("b_major", ["k"]),
]

GROUPED_GEMM_DSRELU_PARAM_MARKS_FP8 = GROUPED_GEMM_DSRELU_FP8_TYPE_MARKS + [
    pytest.mark.parametrize("cd_major", ["n"]),
    pytest.mark.parametrize("acc_dtype", [torch.float32]),
    pytest.mark.parametrize("mma_tiler_mn", [(256, 256)]),
    pytest.mark.parametrize("cluster_shape_mn", [(2, 1)]),
    pytest.mark.parametrize("vector_f32", [False]),
    pytest.mark.parametrize("sf_vec_size,sf_dtype", [(32, torch.float8_e8m0fnu)]),
    pytest.mark.parametrize("discrete_col_sfd", [True]),
]

GROUPED_GEMM_DSRELU_PARAM_MARKS_FP4 = (
    GROUPED_GEMM_DSRELU_FP4_TYPE_MARKS
    + GROUPED_GEMM_DSRELU_COMMON_MARKS
    + [
        pytest.mark.parametrize(
            "sf_vec_size,sf_dtype",
            [
                (16, torch.float8_e8m0fnu),
                (16, torch.float8_e4m3fn),
                (32, torch.float8_e8m0fnu),
                (32, torch.float8_e4m3fn),
            ],
        ),
        pytest.mark.parametrize("discrete_col_sfd", [False]),
    ]
)

GROUPED_GEMM_DSRELU_PARAM_MARKS_DBIAS_FP4 = (
    GROUPED_GEMM_DSRELU_FP4_TYPE_MARKS
    + GROUPED_GEMM_DSRELU_COMMON_MARKS
    + [
        pytest.mark.parametrize(
            "sf_vec_size,sf_dtype",
            [
                (16, torch.float8_e8m0fnu),
                (16, torch.float8_e4m3fn),
                (32, torch.float8_e8m0fnu),
            ],
        ),
        pytest.mark.parametrize("discrete_col_sfd", [False]),
    ]
)


def with_grouped_gemm_dsrelu_params_fp4(func):
    """Decorator to apply grouped GEMM dSReLU FP4 test parameters."""
    for mark in reversed(GROUPED_GEMM_DSRELU_PARAM_MARKS_FP4):
        func = mark(func)
    return func


def with_grouped_gemm_dsrelu_params_fp8(func):
    """Decorator to apply grouped GEMM dSReLU FP8 test parameters."""
    for mark in reversed(GROUPED_GEMM_DSRELU_PARAM_MARKS_FP8):
        func = mark(func)
    return func


def with_grouped_gemm_dsrelu_params_dbias_fp4(func):
    """Decorator to apply grouped GEMM dSReLU dense dbias FP4 test parameters."""
    for mark in reversed(GROUPED_GEMM_DSRELU_PARAM_MARKS_DBIAS_FP4):
        func = mark(func)
    return func


# =============================================================================
# Tensor Allocation
# =============================================================================
def allocate_grouped_gemm_input_tensors(*args, **kwargs) -> Dict[str, Any]:
    """Allocate grouped GEMM inputs and restore the dsReLU-specific alpha/prob range.

    The shared grouped GEMM utility initializes alpha/prob from ``(-2, 2)`` to
    match the forward kernel inputs. The upstream grouped dsReLU kernel uses
    ``torch.randint(1, 2, ...)`` instead, which effectively produces constant ones.
    Keep that behavior for the backward tests here.
    """

    tensors = allocate_grouped_gemm_input_tensors_base(*args, **kwargs)

    alpha_tensor = tensors["alpha_tensor"]
    tensors["alpha_tensor"] = torch.randint(1, 2, alpha_tensor.shape, dtype=torch.float32, device=alpha_tensor.device)

    prob_tensor = tensors["prob_tensor"]
    tensors["prob_tensor"] = torch.randint(1, 2, prob_tensor.shape, dtype=torch.float32, device=prob_tensor.device)

    return tensors


def allocate_grouped_gemm_dsrelu_tensors(
    tensor_m: int,
    n: int,
    l: int,
    ab_dtype: torch.dtype,
    c_dtype: torch.dtype,
    d_dtype: torch.dtype,
    cd_major: str,
    sf_dtype: torch.dtype,
    sf_vec_size: int = 16,
    generate_dbias: bool = False,
    input_tensors: Optional[Dict] = None,
    output_tensors: Optional[Dict] = None,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """Allocate backward tensors for grouped GEMM dSReLU backward.

    :return: Newly allocated input and output tensors. Modifies input_tensors and output_tensors dictionaries in place if provided.
    """

    # D has same shape as C - contains interleaved ab and dsrelu in 32-column blocks
    c_ref, c_tensor = create_and_permute_tensor(
        1, tensor_m, n, cd_major == "m", c_dtype
    )  # Note: c_tensor is an input tensor rather than an output tensor but is being kept as an output tensor to eventually merge with the forward allocation
    _, d_row_tensor = create_and_permute_tensor(1, tensor_m, n, cd_major == "m", d_dtype)
    _, d_col_tensor = create_and_permute_tensor(1, tensor_m, n, cd_major == "m", d_dtype)
    _, d_srelu_tensor = create_and_permute_tensor(1, tensor_m, n, cd_major == "m", d_dtype)
    dprob_tensor = torch.zeros((tensor_m, 1, 1), dtype=torch.float32).cuda()

    _input_tensors = {
        "c_ref": c_ref,
        "c_tensor": c_tensor,
    }
    _output_tensors = {
        "d_row_tensor": d_row_tensor,
        "d_col_tensor": d_col_tensor,
        "d_srelu_tensor": d_srelu_tensor,
        "dprob_tensor": dprob_tensor,
        "dbias_tensor": None,
        "sfd_row_tensor": None,
        "sfd_col_tensor": None,
        "sfd_col_d_srelu_tensor": None,
        "amax_tensor": None,
    }

    if d_dtype in [torch.bfloat16, torch.float16]:
        _output_tensors["amax_tensor"] = torch.full((l, 1), float("-inf"), dtype=torch.float32).cuda()

    if generate_dbias:
        _output_tensors["dbias_tensor"] = torch.zeros((l, n, 1), dtype=torch.bfloat16).cuda()

    if ab_dtype in [torch.float8_e4m3fn, torch.float8_e5m2] and sf_dtype in [
        torch.float8_e8m0fnu,
        torch.float8_e4m3fn,
    ]:  # generate_sfd
        sfd_row_ref, sfd_row_tensor = create_scale_factor_tensor(1, tensor_m, n, sf_vec_size, sf_dtype)
        _output_tensors["sfd_row_tensor"] = sfd_row_tensor
        _output_tensors["sfd_row_ref"] = sfd_row_ref

        sfd_col_ref, sfd_col_tensor = create_scale_factor_tensor(1, n, tensor_m, sf_vec_size, sf_dtype)
        _output_tensors["sfd_col_tensor"] = sfd_col_tensor
        _output_tensors["sfd_col_ref"] = sfd_col_ref

        sfd_col_d_srelu_ref, sfd_col_d_srelu_tensor = create_scale_factor_tensor(1, n, tensor_m, sf_vec_size, sf_dtype)
        _output_tensors["sfd_col_d_srelu_tensor"] = sfd_col_d_srelu_tensor
        _output_tensors["sfd_col_d_srelu_ref"] = sfd_col_d_srelu_ref

    if input_tensors is not None:
        input_tensors.update(_input_tensors)
        _input_tensors = input_tensors
    if output_tensors is not None:
        output_tensors.update(_output_tensors)
        _output_tensors = output_tensors
    return _input_tensors, _output_tensors


# =============================================================================
# Reference Implementations
# =============================================================================


def compute_reference_row_quant(src, d_dtype, sf_dtype, vec_size, norm_const) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Compute reference quantized value on CPU.

    Args:
        src: torch.Tensor, source tensor
        dst_type: Type[cutlass.Numeric], destination type
        vec_size: int, vector size

    Returns:
        torch.Tensor: quantized reference tensor
        torch.Tensor: scale factor tensor
    """

    try:
        from cutlass.cute.runtime import from_dlpack
        import cutlass.cute as cute
        from cudnn.datatypes import _convert_to_cutlass_data_type
    except ImportError:
        pytest.skip("CUTLASS not available for scale factor conversion")

    m = src.shape[0]
    n = src.shape[1]
    l = src.shape[2]

    # 1. Compute reference SFD (m, sfn, l) in fp32
    sfn = ceil_div(n, vec_size)
    n_aligned = ceil_div(n, 128) * 128
    sfn_aligned = n_aligned // vec_size
    sfm = ceil_div(m, 128) * 128
    if sfn_aligned != sfn:
        zeros = torch.zeros(
            src.shape[0],
            n_aligned - n,
            src.shape[2],
            dtype=src.dtype,
            device=src.device,
        )
        src_sf = torch.cat([src, zeros], dim=1)
        src_reshaped = src_sf.permute(2, 0, 1).contiguous()
    else:
        src_reshaped = src.permute(2, 0, 1).contiguous()
    src_reshaped = src_reshaped.view(l, sfm, sfn_aligned, vec_size)
    # Take abs max over vec_size dimension
    src_reshaped, _ = torch.abs(src_reshaped).max(dim=3)  # (l, m, sfn)
    # Multiply by norm_const and rcp_limits
    src_sfd_f32 = src_reshaped * norm_const * get_grouped_gemm_dtype_rcp_limits(d_dtype)
    # Permute to (m, sfn, l)
    src_sfd_f32 = src_sfd_f32.permute(1, 2, 0)
    # Convert fp32 -> f8 -> fp32 for src_sfd_f32
    src_sfd_f8_torch = torch.empty(*(l, sfm, sfn_aligned), dtype=torch.uint8, device=src.device).permute(1, 2, 0)
    src_sfd_f8 = from_dlpack(src_sfd_f8_torch, assumed_align=16).mark_layout_dynamic(leading_dim=1)
    src_sfd_f8.element_type = _convert_to_cutlass_data_type(sf_dtype)
    src_sfd_f32_device = src_sfd_f32.to(src.device)
    ref_sfd_f32_tensor = from_dlpack(src_sfd_f32_device, assumed_align=16).mark_layout_dynamic(leading_dim=1)

    # 2. Convert sfd from fp32 to scale factor
    cute.testing.convert(ref_sfd_f32_tensor, src_sfd_f8)
    cute.testing.convert(src_sfd_f8, ref_sfd_f32_tensor)
    src_sfd_f32 = src_sfd_f32_device.cpu()
    # ref_sfd_f32 for fp32 reference check
    ref_sfd_f32, _ = create_sf_layout_tensor(l, sfm, n, vec_size)
    cvt_sf_MKL_to_M32x4xrm_K4xrk_L(
        from_dlpack(src_sfd_f32),
        from_dlpack(ref_sfd_f32),
    )

    # 3. Quantized output with scale factor
    # Compute reciprocal of src_sfd_f32 and multiply by norm_const
    src_sfd_f32_rcp = norm_const * src_sfd_f32.to(src.device).reciprocal()
    # Expand the sfn dimension by repeating each value sf_vec_size times
    # src_sfd_f32_rcp: (m, sfn, l) -> (m, sfn, sf_vec_size, l) -> (m, n, l)
    src_sfd_f32_rcp_expanded = src_sfd_f32_rcp.unsqueeze(2).expand(sfm, sfn_aligned, vec_size, l)
    src_sfd_f32_rcp_expanded = src_sfd_f32_rcp_expanded.reshape(sfm, sfn_aligned * vec_size, l)
    # Trim to exact n dimension if needed
    src_sfd_f32_rcp_expanded = src_sfd_f32_rcp_expanded[:, :n, :]
    # Apply scale to reference output: ref = ref * src_sfd_f32_rcp
    src_d_f32_torch = torch.einsum("mnl,mnl->mnl", src, src_sfd_f32_rcp_expanded)
    # Convert to d_dtype, then convert back to fp32 for reference check
    src_d_f8_torch = torch.empty(*(l, m, n), dtype=torch.uint8, device=src.device).permute(1, 2, 0)
    src_d_f8 = from_dlpack(src_d_f8_torch, assumed_align=16).mark_layout_dynamic(leading_dim=1)
    src_d_f8.element_type = _convert_to_cutlass_data_type(d_dtype)
    src_d_f32_torch = src_d_f32_torch.to(src.device)
    src_d_f32 = from_dlpack(src_d_f32_torch, assumed_align=16).mark_layout_dynamic(leading_dim=1)
    cute.testing.convert(src_d_f32, src_d_f8)
    cute.testing.convert(src_d_f8, src_d_f32)

    return (ref_sfd_f32, src_d_f32_torch)


def run_grouped_gemm_dsrelu_ref(
    a_ref: torch.Tensor,
    b_ref: torch.Tensor,
    c_ref: torch.Tensor,  # C tensor (intermediate from forward pass) in float32
    sfa_ref: torch.Tensor,
    sfb_ref: torch.Tensor,
    alpha_tensor: torch.Tensor,
    prob_tensor: torch.Tensor,
    aligned_group_m_list: List[int],
    valid_m: int,
    generate_dbias: bool = False,
    generate_amax: bool = False,
    generate_sfd: bool = False,
    norm_const_tensor: Optional[torch.Tensor] = None,
    c_dtype: torch.dtype = torch.bfloat16,
    d_dtype: torch.dtype = torch.float32,
    sf_vec_size: int = 16,
    sf_dtype: torch.dtype = torch.float8_e8m0fnu,
) -> Dict[str, torch.Tensor]:
    """Run reference implementation for grouped GEMM dSReLU backward.

    Based on the reference in contiguous_blockscaled_grouped_gemm_dsrelu_quant_fusion.py

    The dSReLU backward pass computes:
    1. GEMM: ref = alpha * (SFA * A) @ (SFB * B)^T per group
    2. dprob = sum over N of (relu(C)^2 * ref)
    3. D = 2 * relu(C) * ref * prob

    :param a_ref: A tensor (tensor_m, k, 1) in float32
    :param b_ref: B tensor (n, k, l) in float32
    :param c_ref: C tensor (tensor_m, n, 1) from forward pass in float32
    :param sfa_ref: Scale factor A tensor (tensor_m, k, 1) in float32
    :param sfb_ref: Scale factor B tensor (n, k, l) in float32
    :param alpha_tensor: Per-group alpha scaling (l,)
    :param prob_tensor: Per-row probability scaling (tensor_m, 1, 1)
    :param aligned_group_m_list: Aligned M values per group
    :param valid_m: Total valid M dimension
    :param generate_amax: Generate AMAX tensor
    :param generate_sfd: Generate SFD tensor
    :param norm_const_tensor: Normalization constant tensor (1,)
    :param c_dtype: Intermediate C tensor dtype
    :param d_dtype: Output D tensor dtype
    :param sf_vec_size: Scale factor vector size
    :param sf_dtype: Scale factor dtype
    :return: Dictionary of reference tensors
    """
    n, k, l = b_ref.shape
    ref_tensors = {}

    # Step 1: Compute GEMM per group with scale factors
    ref = torch.empty((1, valid_m, n), dtype=torch.float32, device=a_ref.device)
    start = 0
    for i, group_m in enumerate(aligned_group_m_list):
        end = start + group_m
        res_a = torch.einsum(
            "mk,mk->mk",
            a_ref[start:end, :, 0],
            sfa_ref[start:end, :, 0],
        )
        res_b = torch.einsum("nk,nk->nk", b_ref[:, :, i], sfb_ref[:, :, i])
        ref[0, start:end, :] = torch.einsum("mk,nk->mn", res_a, res_b) * alpha_tensor[i].item()
        start = end
    ref = ref.permute((1, 2, 0))  # shape [M, N, 1]

    # Step 2: Apply dsquared-ReLU backward elementwise
    c_full = c_ref.clone()
    c_relu = torch.relu(c_full)
    ref_dprob = c_relu**2 * ref
    chunk_sums = [torch.sum(chunk, dim=1, keepdim=True) for chunk in torch.split(ref_dprob, 32, dim=1)]
    ref_dprob = torch.sum(torch.cat(chunk_sums, dim=1), dim=1, keepdim=True)
    ref_tensors["dprob_ref"] = ref_dprob

    # Step 3: Compute dSReLU formula
    prob = prob_tensor.expand(-1, n, -1)
    ref_d = 2 * c_relu * ref * prob

    ref_tensors["d_ref"] = ref_d.clone()
    ref_d_srelu = (c_relu**2 * prob).to(d_dtype).to(torch.float32)
    ref_tensors["d_srelu_ref"] = ref_d_srelu.clone()

    if generate_dbias:
        ref_dbias = torch.zeros((l, n, 1), dtype=torch.bfloat16, device=a_ref.device)
        start = 0
        for i, group_m in enumerate(aligned_group_m_list):
            end = start + group_m
            ref_dbias[i, :, 0] = ref_d[start:end, :, 0].sum(dim=0).to(torch.bfloat16)
            start = end
        ref_tensors["dbias_ref"] = ref_dbias

    # Step 6: Generate amax for FP4/BF16 output
    if generate_amax:
        ref_amax = torch.empty((l, 1), dtype=torch.float32, device=a_ref.device)
        start = 0
        for i, group_m in enumerate(aligned_group_m_list):
            end = start + group_m
            ref_amax[i, 0] = torch.tensor(compute_reference_amax(ref_d[start:end, :, 0].clone()))
            start = end
        ref_tensors["amax_ref"] = ref_amax

    # Step 7: Generate SFD for FP8 output
    if generate_sfd:
        norm_const = norm_const_tensor[0].item()
        sfd_row_ref_f32, d_ref_f32 = compute_reference_row_quant(ref_d, d_dtype, sf_dtype, sf_vec_size, norm_const)
        ref_tensors["sfd_row_ref"] = sfd_row_ref_f32.clone()
        ref_tensors["d_ref"] = d_ref_f32.clone()

        ref_d_col = ref_d.permute(2, 1, 0).contiguous().permute(1, 2, 0)
        sfd_col_ref_f32, d_col_ref_f32 = compute_reference_row_quant(ref_d_col, d_dtype, sf_dtype, sf_vec_size, norm_const)
        ref_tensors["sfd_col_ref"] = sfd_col_ref_f32.clone()
        ref_tensors["d_col_ref"] = d_col_ref_f32.clone()

        ref_d_srelu_col = (c_relu**2 * prob).permute(2, 1, 0).contiguous().permute(1, 2, 0)
        sfd_col_d_srelu_ref_f32, d_srelu_col_ref_f32 = compute_reference_row_quant(ref_d_srelu_col, d_dtype, sf_dtype, sf_vec_size, norm_const)
        ref_tensors["sfd_col_d_srelu_ref"] = sfd_col_d_srelu_ref_f32.clone()
        ref_tensors["d_srelu_col_ref"] = d_srelu_col_ref_f32.clone()

    return ref_tensors


# =============================================================================
# Reference Checking
# =============================================================================


def check_ref_grouped_gemm_dsrelu(
    inputs: Dict[str, Any],
    outputs: Dict[str, Any],
    cfg: Dict[str, Any],
    atol: float = 1e-1,
    rtol: float = 1e-2,
    skip_ref: bool = False,
) -> None:
    if skip_ref:
        return

    torch.cuda.synchronize()
    ref_tensors = run_grouped_gemm_dsrelu_ref(
        a_ref=inputs["a_ref"],
        b_ref=inputs["b_ref"],
        c_ref=inputs["c_ref"],
        sfa_ref=inputs["sfa_ref"],
        sfb_ref=inputs["sfb_ref"],
        alpha_tensor=inputs["alpha_tensor"],
        prob_tensor=inputs["prob_tensor"],
        aligned_group_m_list=inputs["aligned_group_m_list"],
        valid_m=inputs["valid_m"],
        generate_dbias=outputs.get("dbias_tensor") is not None,
        generate_amax=outputs.get("amax_tensor") is not None,
        generate_sfd=outputs.get("sfd_row_tensor") is not None and outputs.get("sfd_col_tensor") is not None,
        norm_const_tensor=inputs.get("norm_const_tensor"),
        c_dtype=cfg["c_dtype"],
        d_dtype=cfg["d_dtype"],
        sf_vec_size=cfg["sf_vec_size"],
        sf_dtype=cfg["sf_dtype"],
    )

    torch.testing.assert_close(outputs["dprob_tensor"].float(), ref_tensors["dprob_ref"].float(), atol=atol, rtol=rtol)
    torch.testing.assert_close(outputs["d_row_tensor"].float(), ref_tensors["d_ref"].float(), atol=atol, rtol=rtol)

    if "d_col_ref" in ref_tensors:
        torch.testing.assert_close(
            outputs["d_col_tensor"].float().permute(1, 0, 2),
            ref_tensors["d_col_ref"].float(),
            atol=atol,
            rtol=rtol,
        )

    if outputs.get("d_srelu_tensor") is not None:
        if "d_srelu_col_ref" in ref_tensors:
            torch.testing.assert_close(
                outputs["d_srelu_tensor"].float().permute(1, 0, 2),
                ref_tensors["d_srelu_col_ref"].float(),
                atol=atol,
                rtol=rtol,
            )
        else:
            torch.testing.assert_close(outputs["d_srelu_tensor"].float(), ref_tensors["d_srelu_ref"].float(), atol=atol, rtol=rtol)

    if outputs.get("amax_tensor") is not None and "amax_ref" in ref_tensors:
        torch.testing.assert_close(outputs["amax_tensor"].float(), ref_tensors["amax_ref"].float(), atol=atol, rtol=rtol)

    if outputs.get("sfd_row_tensor") is not None and "sfd_row_ref" in ref_tensors:
        torch.testing.assert_close(
            outputs["sfd_row_tensor"].float(),
            ref_tensors["sfd_row_ref"].to(outputs["sfd_row_tensor"].device).float(),
            atol=atol,
            rtol=rtol,
        )

    if outputs.get("sfd_col_tensor") is not None and "sfd_col_ref" in ref_tensors:
        torch.testing.assert_close(
            outputs["sfd_col_tensor"].float(),
            ref_tensors["sfd_col_ref"].to(outputs["sfd_col_tensor"].device).float(),
            atol=atol,
            rtol=rtol,
        )

    if outputs.get("sfd_col_d_srelu_tensor") is not None and "sfd_col_d_srelu_ref" in ref_tensors:
        torch.testing.assert_close(
            outputs["sfd_col_d_srelu_tensor"].float(),
            ref_tensors["sfd_col_d_srelu_ref"].to(outputs["sfd_col_d_srelu_tensor"].device).float(),
            atol=atol,
            rtol=rtol,
        )
