"""Fused SM120 NVFP4 experts; original checkpoint weights and calibrated scales."""
import torch
from torch import nn
from flashinfer import nvfp4_block_scale_interleave
from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout
from flashinfer.fused_moe import B12xMoEWrapper

_SHARED_WORKSPACES = {}


class FusedNVFP4Experts(nn.Module):
    def __init__(self, reader, prefix, num_experts, hidden, intermediate, device, top_k=8, max_tokens=256):
        super().__init__()
        self.num_experts, self.hidden, self.intermediate = num_experts, hidden, intermediate
        w1 = torch.empty((num_experts, 2 * intermediate, hidden // 2), dtype=torch.uint8, device=device)
        w2 = torch.empty((num_experts, hidden, intermediate // 2), dtype=torch.uint8, device=device)
        sf1 = torch.empty((num_experts, 2 * intermediate, hidden // 16), dtype=torch.uint8, device=device)
        sf2 = torch.empty((num_experts, hidden, intermediate // 16), dtype=torch.uint8, device=device)
        globals1, globals2, inputs1, inputs2 = [], [], [], []
        for expert in range(num_experts):
            name = f'{prefix}{expert}.'
            values = {p + '.' + k: reader.get_tensor(name + p + '.' + k)
                for p in ('gate_proj', 'up_proj', 'down_proj')
                for k in ('weight', 'weight_scale', 'weight_scale_2', 'input_scale')}
            for kind in ('weight_scale_2', 'input_scale'):
                assert torch.equal(values['gate_proj.' + kind], values['up_proj.' + kind])
            # b12x uses [up, gate]; the reference linear uses [gate, up].
            w1[expert, :intermediate].copy_(values['up_proj.weight'])
            w1[expert, intermediate:].copy_(values['gate_proj.weight'])
            w2[expert].copy_(values['down_proj.weight'])
            logical = torch.cat([values['up_proj.weight_scale'], values['gate_proj.weight_scale']])
            sf1[expert].copy_(nvfp4_block_scale_interleave(logical.view(torch.uint8)).reshape_as(sf1[expert]))
            sf2[expert].copy_(nvfp4_block_scale_interleave(values['down_proj.weight_scale'].view(torch.uint8)).reshape_as(sf2[expert]))
            globals1.append(values['gate_proj.weight_scale_2'].reshape(()))
            globals2.append(values['down_proj.weight_scale_2'].reshape(()))
            inputs1.append(values['gate_proj.input_scale'].reshape(()))
            inputs2.append(values['down_proj.input_scale'].reshape(()))
        for key, value in dict(w1=w1, w2=w2,
                sf1=convert_sf_to_mma_layout(sf1.view(torch.float8_e4m3fn), m=2*intermediate, k=hidden, num_groups=num_experts),
                sf2=convert_sf_to_mma_layout(sf2.view(torch.float8_e4m3fn), m=hidden, k=intermediate, num_groups=num_experts),
                alpha1=torch.stack(globals1).float(),
                alpha2=torch.stack(globals2).float() * torch.stack(inputs2).float(),
                input1=torch.stack(inputs1).float(), input2=torch.stack(inputs2).float()).items():
            self.register_buffer(key, value)
        # One set of B12x workspaces per (shape, device), shared by every layer.
        # Layers run strictly in sequence, so per-layer copies (0.665 GiB each,
        # ~31 GiB per GPU across 47 MoE layers) only ever had one in use. Each
        # layer keeps its own thin wrapper so FlashInfer's per-weight caches
        # (_weight_views etc.) stay keyed to that layer's tensors.
        key = (num_experts, top_k, hidden, intermediate, max_tokens, str(device))
        shared = _SHARED_WORKSPACES.get(key)
        if shared is None:
            self.wrapper = B12xMoEWrapper(num_experts, top_k, hidden, intermediate,
                use_cuda_graph=True, max_num_tokens=max_tokens, device=str(device), quant_mode='nvfp4')
            _SHARED_WORKSPACES[key] = (self.wrapper._static_workspace,
                                       self.wrapper._dynamic_workspace, self.wrapper._moe_output)
        else:
            static, dynamic, output = shared
            self.wrapper = B12xMoEWrapper(num_experts, top_k, hidden, intermediate,
                use_cuda_graph=True, max_num_tokens=max_tokens, device=str(device), quant_mode='nvfp4',
                shared_static_workspace=static, shared_dynamic_workspace=dynamic, shared_output=output)

    def forward(self, hidden_states, top_k_index, top_k_weights):
        with torch.cuda.device(hidden_states.device):
            return self.wrapper.run(hidden_states, self.w1, self.sf1, self.w2, self.sf2,
                top_k_index, top_k_weights, w1_alpha=self.alpha1, w2_alpha=self.alpha2,
                fc2_input_scale=self.input2, input_global_scale=self.input1)