import torch
from typing import Tuple, Optional
from ..utils.math import align

# noinspection PyBroadException
try:
    # noinspection PyProtectedMember
    import torch.distributed._symmetric_memory as symm_mem
    import torch.distributed as dist
except Exception as exception:
    print(f'Failed to load mega kernels, please check your PyTorch version: {exception}')

from .. import _C


def _is_sm90() -> bool:
    return torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9


def _is_sm100() -> bool:
    return torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 10


class SymmBuffer:
    def __init__(self, group: dist.ProcessGroup,
                 # MoE arguments
                 num_experts: int,
                 num_max_tokens_per_rank: int, num_topk: int,
                 hidden: int, intermediate_hidden: int,
                 use_fp8_dispatch: bool = True,
                 activation: str = 'swiglu'):
        self.group = group
        self.num_experts = num_experts
        self.num_max_tokens_per_rank = num_max_tokens_per_rank
        self.num_topk = num_topk
        self.hidden = hidden
        self.intermediate_hidden = intermediate_hidden

        # Allocate a symmetric buffer (route by architecture)
        if _is_sm90():
            num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe(
                group.size(), num_experts,
                num_max_tokens_per_rank, num_topk,
                hidden, intermediate_hidden,
                use_fp8_dispatch, activation
            )
        elif _is_sm100():
            num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe(
                group.size(), num_experts,
                num_max_tokens_per_rank, num_topk,
                hidden, intermediate_hidden,
                use_fp8_dispatch, activation
            )
        else:
            raise RuntimeError('Unsupported architecture for MegaMoE')
        self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')
        self.handle = symm_mem.rendezvous(self.buffer, group=group)
        self.buffer.zero_()
        self.group.barrier()
        torch.cuda.synchronize()

        # Create input buffer views
        buffer_views = slice_input_buffers(self.buffer)
        if _is_sm90():
            (self.x, self.x_sf,
             self.topk_idx, self.topk_weights,
             self.l1_acts, self.l1_acts_sf, self.l1_topk_weights,
             self.l2_acts, self.l2_acts_sf,
             self.expert_recv_count_sum,
             self.l1_arrival_count,
             self.l2_arrival_mask,
             self.token_src_metadata,
             self.l1_accum_debug) = buffer_views
        else:
            (self.x, self.x_sf,
             self.topk_idx, self.topk_weights,
             self.l1_acts, self.l1_acts_sf,
             self.l2_acts, self.l2_acts_sf) = buffer_views
            self.l1_topk_weights = None
            self.expert_recv_count_sum = None
            self.l1_arrival_count = None
            self.l2_arrival_mask = None
            self.token_src_metadata = None
            self.l1_accum_debug = None

    def destroy(self):
        self.handle = None
        self.buffer = None
        self.group = None
        self.x = None
        self.x_sf = None
        self.topk_idx = None
        self.topk_weights = None
        self.l1_acts = None
        self.l1_acts_sf = None
        self.l1_topk_weights = None
        self.l2_acts = None
        self.l2_acts_sf = None
        self.expert_recv_count_sum = None
        self.l1_arrival_count = None
        self.l2_arrival_mask = None
        self.token_src_metadata = None
        self.l1_accum_debug = None


def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,
                                 num_experts: int,
                                 num_max_tokens_per_rank: int, num_topk: int,
                                 hidden: int, intermediate_hidden: int,
                                 use_fp8_dispatch: bool = True,
                                 activation: str = 'swiglu') -> SymmBuffer:
    # Token count must be aligned to block sizes
    if _is_sm90():
        alignment = _C.get_token_alignment_for_sm90_mega_moe()
    elif _is_sm100():
        alignment = _C.get_token_alignment_for_mega_moe()
    else:
        raise RuntimeError('Unsupported architecture for MegaMoE')
    num_max_tokens_per_rank = align(num_max_tokens_per_rank, alignment)

    return SymmBuffer(
        group, num_experts,
        num_max_tokens_per_rank, num_topk,
        hidden, intermediate_hidden,
        use_fp8_dispatch, activation
    )


def _interleave_l1_weight_tensor(t: torch.Tensor, gran: int = 8) -> torch.Tensor:
    # [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up]
    g, n, *rest = t.shape
    half = n // 2
    gate = t[:, :half].reshape(g, half // gran, gran, *rest)
    up = t[:, half:].reshape(g, half // gran, gran, *rest)
    return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest))


def _interleave_l1_weights(l1_weights: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
    return _interleave_l1_weight_tensor(l1_weights[0]), _interleave_l1_weight_tensor(l1_weights[1])


def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor:
    num_groups, mn, packed_sf_k = sf.shape
    assert sf.dtype == torch.int and mn % 128 == 0
    result = (sf.reshape(num_groups, -1, 4, 32, packed_sf_k)
                .transpose(2, 3)
                .reshape(num_groups, mn, packed_sf_k))
    return torch.empty_like(sf).copy_(result)


def transform_weights_for_mega_moe_sm90(
    l1_weights: Tuple[torch.Tensor, torch.Tensor],
    l2_weights: Tuple[torch.Tensor, torch.Tensor]
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
    # L1: interleave FP8 gate/up weights only; SM90 float weight SF stays natural MN-major.
    l1_weights = (_interleave_l1_weight_tensor(l1_weights[0]), l1_weights[1])
    # L2: no transform
    return l1_weights, l2_weights


def transform_weights_for_mega_moe(
    l1_weights: Tuple[torch.Tensor, torch.Tensor],
    l2_weights: Tuple[torch.Tensor, torch.Tensor]
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
    if _is_sm90():
        return transform_weights_for_mega_moe_sm90(l1_weights, l2_weights)
    # SM100: L1 interleave gate/up + UTCCP SF transpose, L2 UTCCP SF transpose
    l1_interleaved = _interleave_l1_weights(l1_weights)
    l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1]))
    l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1]))
    return l1_weights, l2_weights


def fp8_mega_moe(y: torch.Tensor,
                 l1_weights: Tuple[torch.Tensor, torch.Tensor],
                 l2_weights: Tuple[torch.Tensor, torch.Tensor],
                 sym_buffer: SymmBuffer,
                 cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
                 recipe: Optional[Tuple[int, int, int]] = None,
                 activation: str = 'swiglu',
                 activation_clamp: Optional[float] = None,
                 fast_math: bool = True):
    if _is_sm90():
        if recipe is None:
            recipe = (1, 128, 128)
        _C.fp8_mega_moe(
            y,
            l1_weights, l2_weights,
            cumulative_local_expert_recv_stats,
            sym_buffer.buffer,
            sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
            sym_buffer.num_max_tokens_per_rank,
            sym_buffer.num_experts, sym_buffer.num_topk,
            recipe,
            activation, activation_clamp,
            fast_math
        )
    elif _is_sm100():
        if recipe is None:
            recipe = (1, 1, 32)
        _C.fp8_fp4_mega_moe(
            y,
            l1_weights, l2_weights,
            cumulative_local_expert_recv_stats,
            sym_buffer.buffer,
            sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
            sym_buffer.num_max_tokens_per_rank,
            sym_buffer.num_experts, sym_buffer.num_topk,
            recipe,
            activation, activation_clamp,
            fast_math
        )
    else:
        raise RuntimeError('Unsupported architecture for MegaMoE')


# Backward-compatible alias
def fp8_fp4_mega_moe(y: torch.Tensor,
                     l1_weights: Tuple[torch.Tensor, torch.Tensor],
                     l2_weights: Tuple[torch.Tensor, torch.Tensor],
                     sym_buffer: SymmBuffer,
                     cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
                     recipe: Optional[Tuple[int, int, int]] = None,
                     activation: str = 'swiglu',
                     activation_clamp: Optional[float] = None,
                     fast_math: bool = True):
    fp8_mega_moe(y, l1_weights, l2_weights, sym_buffer,
                 cumulative_local_expert_recv_stats, recipe,
                 activation, activation_clamp, fast_math)
