"""CoreX fallback for FP8 W8A8 MoE on Iluvatar BI-V150.

The CoreX ixformer MoE kernels are W16A16 only and torch._scaled_mm supports
only tensor-wise scales, so per-token/per-channel FP8 MoE cannot use them.
This module dequantizes one expert at a time (each expert shard is only a few
hundred MB on a TP rank), runs BF16 matmuls, and scatters the result.
"""
from __future__ import annotations

import torch
import torch.nn.functional as F

from vllm.logger import init_logger

logger = init_logger(__name__)
_FP8_DTYPES = (
    torch.float8_e4m3fn,
    torch.float8_e5m2,
    torch.float8_e4m3fnuz,
    torch.float8_e5m2fnuz,
)
_LOGGED = False


def _dequant_weight(w: torch.Tensor, scale: torch.Tensor | None) -> torch.Tensor:
    """w: [rows, cols] FP8, scale: per-output-row (len rows) or scalar."""
    if scale is None:
        return w.to(torch.bfloat16)
    s = scale.reshape(-1)
    if s.numel() == 1:
        return (w.float() * s[0]).to(torch.bfloat16)
    if s.numel() != w.shape[0]:
        # try to broadcast a [1, rows] style scale
        s = s.reshape(w.shape[0], -1)[:, 0]
    return (w.float() * s.reshape(-1, 1)).to(torch.bfloat16)


def corexx_fp8_moe(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    activation,
    apply_router_weight_on_input: bool,
    global_num_experts: int,
    expert_map: torch.Tensor | None,
    w1_scale: torch.Tensor | None,
    w2_scale: torch.Tensor | None,
    a1_scale: torch.Tensor | None,
    a2_scale: torch.Tensor | None,
    block_shape,
    w1_bias: torch.Tensor | None,
    w2_bias: torch.Tensor | None,
) -> torch.Tensor:
    global _LOGGED
    if expert_map is not None:
        raise NotImplementedError("CoreX FP8 MoE fallback does not support expert_map/EP")
    if block_shape is not None:
        raise NotImplementedError("CoreX FP8 MoE fallback does not support block FP8")
    if not _LOGGED:
        logger.info(
            "CoreX FP8 MoE fallback: hidden=%s %s a1_scale=%s w1=%s w1_scale=%s w2=%s w2_scale=%s topk=%s",
            tuple(hidden_states.shape), hidden_states.dtype,
            None if a1_scale is None else tuple(a1_scale.shape),
            tuple(w1.shape), None if w1_scale is None else tuple(w1_scale.shape),
            tuple(w2.shape), None if w2_scale is None else tuple(w2_scale.shape),
            tuple(topk_ids.shape),
        )
        _LOGGED = True

    if hidden_states.dtype in _FP8_DTYPES:
        if a1_scale is None:
            raise ValueError("FP8 MoE input requires a1_scale")
        A = (hidden_states.float() * a1_scale.reshape(-1, 1)).to(torch.bfloat16)
    else:
        A = hidden_states

    E = w1.shape[0]
    N = w1.shape[1]
    K = w1.shape[2]
    I = w2.shape[2]
    assert N == 2 * I, f"expected w1 N=2*I, got {N} vs {I}"

    act_name = getattr(activation, "value", activation)

    M = A.shape[0]
    out = torch.zeros((M, K), dtype=torch.bfloat16, device=A.device)
    topk = topk_ids.shape[1]

    # COREXX_FP8_MOE_PATCH: only touch experts that actually received tokens
    active_experts = torch.unique(topk_ids).tolist()
    for e in active_experts:
        if e < 0 or e >= E:
            continue
        # rows of A routed to expert e (global id == local id when expert_map is None)
        for k in range(topk):
            idx = (topk_ids[:, k] == e).nonzero(as_tuple=False).flatten()
            if idx.numel() == 0:
                continue
            x = A.index_select(0, idx)
            if apply_router_weight_on_input:
                x = x * topk_weights[idx, k].to(x.dtype).unsqueeze(1)
            w1_e = _dequant_weight(w1[e], None if w1_scale is None else w1_scale[e])
            h = x @ w1_e.t()
            if w1_bias is not None:
                h = h + w1_bias[e]
            gate = h[:, :I]
            up = h[:, I:N]
            if act_name == "silu":
                act = F.silu(gate) * up
            elif act_name == "gelu":
                act = F.gelu(gate) * up
            else:
                raise NotImplementedError(f"CoreX FP8 MoE activation {act_name}")
            w2_e = _dequant_weight(w2[e], None if w2_scale is None else w2_scale[e])
            y = act @ w2_e.t()
            if w2_bias is not None:
                y = y + w2_bias[e]
            if not apply_router_weight_on_input:
                y = y * topk_weights[idx, k].to(y.dtype).unsqueeze(1)
            out.index_add_(0, idx, y)
            del w1_e, w2_e, h, act, y

    return out