"""CVMM with FP8-activations (safe wrapper over existing kernel).

This module provides a minimal FP8-like activations API by quantizing inputs to
int8 with per-row scale (E4M3-style emulation) on the host, then calling the
existing cvmm kernel to compute results using dequantized fp16/bf16 inputs.

It does NOT modify any Triton kernels. This keeps the build stable while
providing the new cvmm_fp8() API to iterate on.

Next steps:
- Add an in-kernel dequant path (guarded by a constexpr flag) to avoid host-side dequant
  and leverage SM90-friendly tiling/pipelining.
- Add minimal tests/benchmarks to compare FP8-activations vs baseline on H100.
- Optional: extend calibration to support e5m2 if you need it.

"""

# ruff: noqa: ANN001
# type: ignore[reportMissingParameterType]
# pyright: ignore[reportMissingParameterType]

from typing import Literal, cast

import torch

from mt_dao.models.cvmm_triton_kernel import (
    CVMMSel,
    cvmm_prepare_sel,
)
from mt_dao.models.cvmm_triton_kernel import (
    cvmm as _cvmm,
)
from mt_dao.models.cvmm_triton_kernel import (
    cvmm_triton_backward as _cvmm_bw,
)


def _quantize_per_row_int8(
    x: torch.Tensor, eps: float = 1e-8,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize rows of x to int8 with per-row scale (FP8-like emulation).

    Parameters
    ----------
    x : torch.Tensor
        Input of shape [..., K]. Last dim is treated as feature.
    eps : float, optional
        Clamp minimum for scale to avoid division by zero.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        q: int8 tensor with same shape as x; s: float32 scales of shape [...].

    """
    x_f = x.flatten(end_dim=-2)
    # per-row amax
    amax = x_f.abs().amax(dim=-1).to(torch.float32)
    scale = (amax / 127.0).clamp_min(eps)
    inv = torch.reciprocal(scale)
    q = torch.round(x_f.to(torch.float32) * inv.unsqueeze(-1)).clamp_(-128, 127)
    q = q.to(torch.int8)
    return q.view_as(x), scale


class CVMMFp8(torch.autograd.Function):
    """Autograd Function for cvmm with FP8-like activations (host dequant).

    Forward: quantize activations to int8 + per-row scale, dequantize back to
    float (bf16/fp16) on host, call existing _cvmm kernel path.
    Backward: uses the same cvmm backward path via CVMM in cvmm_triton_kernel.
    """

    @staticmethod
    def forward(  # noqa: PLR0913, PLR0917
        ctx,  # pyright: ignore[reportMissingParameterType]
        x: torch.Tensor,
        sel_index: torch.Tensor,
        sel: torch.Tensor,
        keys: torch.Tensor,
        _fp8_format: Literal["e4m3", "e5m2"] = "e4m3",
        out_index: torch.Tensor | None = None,
        reduction_weight: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Forward pass with host dequantization and existing kernel path.

        Returns
        -------
        torch.Tensor
            Output tensor produced by cvmm.

        """
        # Quantize activations to int8 + per-row scale (emulating FP8)
        x_q, s = _quantize_per_row_int8(x)
        # Host dequant to compute dtype to reuse existing kernel safely
        compute_dtype = (
            torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else x.dtype
        )
        if compute_dtype == torch.bfloat16:
            x_dq = (
                x_q.to(torch.float32) * s.view(*x_q.shape[:-1], 1)
            ).to(torch.bfloat16)
        elif compute_dtype == torch.float16:
            x_dq = (
                x_q.to(torch.float32) * s.view(*x_q.shape[:-1], 1)
            ).to(torch.float16)
        else:
            x_dq = (
                x_q.to(torch.float32) * s.view(*x_q.shape[:-1], 1)
            ).to(torch.float32)
        ctx.save_for_backward(x_dq, keys, sel, sel_index, out_index, reduction_weight)
        out_type = compute_dtype
        if out_index is None:
            out_index = torch.tensor(-1, device=x.device)
        res = cast(
            "torch.Tensor",
            torch.ops.mylib.cvmm_triton(  # pyright: ignore[reportCallIssue]
                x_dq,
                sel_index,
                sel,
                keys,
                out_type,
                out_index,
            ),
        )
        if reduction_weight is not None:
            res = res.view(*reduction_weight.shape, res.shape[-1])
            res = (reduction_weight.unsqueeze(-2).type_as(res) @ res).squeeze(-2)
        ctx.op_type = out_type
        ctx.keys_type = keys.dtype
        ctx.is_autocast = torch.is_autocast_enabled()
        return res

    @staticmethod
    def backward(  # pyright: ignore[reportIncompatibleMethodOverride]  # noqa: ANN205
        ctx,  # pyright: ignore[reportMissingParameterType]
        grad_output,  # pyright: ignore[reportMissingParameterType]
    ):
        """Backward pass computing grads for input and weights.

        Returns
        -------
        tuple
            Gradients for (x, sel_index, sel, keys, fp8_format, out_index,
            reduction_weight), matching forward args. None for non-tensor args.

        """
        x_dq, keys, sel, sel_index, out_index, reduction_weight = ctx.saved_tensors
        keys_dt = keys
        if reduction_weight is not None:
            grad_output_w = reduction_weight.unsqueeze(-1).type_as(grad_output) @ (
                grad_output.unsqueeze(-2)
            )
        else:
            grad_output_w = grad_output
        out_index_is_none = False
        if out_index is None:
            out_index_is_none = True
            out_index = torch.tensor(-1, device=grad_output.device)
        # Weight grad via existing backward
        grad_w = _cvmm_bw(
            x_dq,
            sel_index,
            sel,
            grad_output_w,
            keys_dt.shape[0],
            ctx.keys_type,
            ctx.is_autocast,
            out_index=out_index,
        )
        # Input grad via cvmm with transposed keys
        bw_index = sel_index if out_index_is_none else out_index
        bw_index_out = torch.tensor(-1, device=grad_output.device)
        if reduction_weight is not None:
            bw_index_out = bw_index
            # Avoid in-place modification on a saved tensor (sel_index/out_index)
            # which breaks autograd versioning. Use out-of-place integer division.
            bw_index = bw_index // reduction_weight.shape[-1]  # noqa: PLR6104
        grad_x_full = torch.ops.mylib.cvmm_triton(  # pyright: ignore[reportCallIssue]
            grad_output,
            bw_index,
            sel,
            keys_dt.transpose(1, 2),
            ctx.op_type,
            bw_index_out,
        )
        grad_x_full = grad_x_full.view(*x_dq.shape[:-1], -1, x_dq.shape[-1])
        grad_w_off = None
        if reduction_weight is not None:
            grad_x = (
                reduction_weight.view(*grad_x_full.shape[:-1]).unsqueeze(-2).type_as(grad_x_full)
                @ grad_x_full
            ).squeeze(-2)
            grad_w_off = (
                (
                    grad_x_full.type_as(reduction_weight)
                    @ x_dq.unsqueeze(-1).type_as(reduction_weight)
                )
                .squeeze(-1)
                .view_as(reduction_weight)
            )
        elif grad_x_full.shape[-2] != 1:
            grad_x = grad_x_full.sum(-2)
        else:
            grad_x = grad_x_full
        grad_x = grad_x.view_as(x_dq)
        return grad_x, None, None, grad_w, None, None, grad_w_off


def cvmm_fp8(
    x: torch.Tensor,
    sel: torch.Tensor | CVMMSel,
    keys: torch.Tensor,
    fp8_format: Literal["e4m3", "e5m2"] = "e4m3",
) -> torch.Tensor:
    """FP8-like cvmm using per-row int8 quantization + host dequant.

    Stable entrypoint that quantizes x to int8 + scale, then dequantizes to the
    compute dtype and calls the existing Triton kernel via
    torch.ops.mylib.cvmm_triton.

    Returns
    -------
    torch.Tensor
        Output of shape sel.shape + (N,).

    """
    if not isinstance(sel, CVMMSel):
        sel = cvmm_prepare_sel(sel)
    return cast(
        "torch.Tensor",
        CVMMFp8.apply(
            x,
            sel.sel_index,
            sel.sel,
            keys,
            fp8_format,
            sel.out_index,
            sel.reduction_weight,
        ),
    )


# Convenience alias keeping parity with non-FP8 API naming
cvmm = _cvmm
