import torch
import torch._functorch.config
import torch._inductor
from common.diff_engine import DiffCase

torch._functorch.config.donated_buffer = False

from activation.grouped_poly_norm import (fused_mul_grouped_poly_norm,
                                          fused_mul_grouped_poly_norm_ref)

# 384 / 8 (EP) = 48 experts per rank
# total_tokens = bs * sl, which equals per-rank tokens
# since top_k=8 and EP=8, each rank sees all tokens once
NUM_EXPERTS = 48


class GroupedRefModule(torch.nn.Module):
    """Wraps the PyTorch reference for grouped FusedMulPolyNorm."""

    def __init__(self, weight, bias, offsets, eps, expert_offset=0):
        super().__init__()
        self.weight = torch.nn.Parameter(weight)
        self.bias = torch.nn.Parameter(bias)
        self.offsets = offsets
        self.eps = eps
        self.expert_offset = expert_offset

    def forward(self, x, mul):
        return fused_mul_grouped_poly_norm_ref(
            x,
            mul,
            self.weight,
            self.bias,
            self.offsets,
            self.eps,
            expert_offset=self.expert_offset)


class GroupedCUDAModule(torch.nn.Module):
    """Wraps the CUDA kernel for grouped FusedMulPolyNorm."""

    def __init__(self, weight, bias, offsets, eps, expert_offset=0):
        super().__init__()
        self.weight = torch.nn.Parameter(weight)
        self.bias = torch.nn.Parameter(bias)
        self.offsets = offsets
        self.eps = eps
        self.expert_offset = expert_offset

    def forward(self, x, mul):
        return fused_mul_grouped_poly_norm(x,
                                           mul,
                                           self.weight,
                                           self.bias,
                                           self.offsets,
                                           self.eps,
                                           expert_offset=self.expert_offset)


class GroupedMulPoly(DiffCase):
    """Benchmark case for Grouped FusedMulPolyNorm (MoE).

    Maps the framework's (bs, sl, hidden) to grouped polynorm's
    (total_tokens, D) where total_tokens = bs * sl.
    Uses a fixed number of experts with uniform token distribution.
    """

    def build_inputs(self, bs, sl, hidden, dtype, eps):
        total_tokens = bs * sl
        num_experts = min(NUM_EXPERTS, total_tokens)

        torch.manual_seed(42)
        probs = torch.ones(num_experts) / num_experts
        assignments = torch.multinomial(probs, total_tokens, replacement=True)
        counts = torch.bincount(assignments, minlength=num_experts).tolist()
        offsets = torch.cumsum(torch.tensor(counts, dtype=torch.int32),
                               dim=0).to(torch.int32)

        return {
            "x":
            torch.randn(total_tokens, hidden, dtype=dtype,
                        requires_grad=True) * 0.5,
            "mul":
            torch.randn(total_tokens, hidden, dtype=dtype,
                        requires_grad=True) * 0.5,
            "weight":
            torch.ones(num_experts, 3, dtype=dtype) / 3 +
            torch.randn(num_experts, 3, dtype=dtype) * 0.01,
            "bias":
            torch.randn(num_experts, 1, dtype=dtype) * 0.01,
            "offsets":
            offsets,
            "dim":
            hidden,
            "eps":
            eps,
            "dtype":
            dtype,
        }

    def make_naive(self, I):
        return GroupedRefModule(
            I["weight"].detach().clone(),
            I["bias"].detach().clone(),
            I["offsets"],
            I["eps"],
        )

    def make_compiled(self, I):
        m = GroupedRefModule(
            I["weight"].detach().clone(),
            I["bias"].detach().clone(),
            I["offsets"],
            I["eps"],
        )
        return torch.compile(m)

    def make_cuda(self, I):
        return GroupedCUDAModule(
            I["weight"].detach().clone(),
            I["bias"].detach().clone(),
            I["offsets"],
            I["eps"],
        )

    def make_compiled_cuda(self, I):
        m = GroupedCUDAModule(
            I["weight"].detach().clone(),
            I["bias"].detach().clone(),
            I["offsets"],
            I["eps"],
        )
        return torch.compile(m)

    def forward(self, obj, I):
        return obj(I["x"], I["mul"])

    def grad_inputs(self, I):
        return [I["x"], I["mul"]]


CASE = GroupedMulPoly()