import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, Optional
from task import input_t, output_t


class Expert(nn.Module):
    def __init__(self, config: Dict, d_expert: Optional[int] = None):
        super().__init__()
        self.config = config
        self.act_fn = nn.SiLU()
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"] if d_expert is None else d_expert

        # Fuse gate and up projections into a single weight
        self.W_gate_up = nn.Linear(self.d_hidden, 2 * self.d_expert, bias=False)
        self.W_down = nn.Linear(self.d_expert, self.d_hidden, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Single matmul for gate and up, then split
        gate_up = self.W_gate_up(x)
        gate, up = gate_up.chunk(2, dim=-1)
        gate = self.act_fn(gate)
        return self.W_down(gate * up)


class MoEGate(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.top_k: int = config["n_experts_per_token"]
        self.num_experts: int = config["n_routed_experts"]
        self.d_hidden: int = config["d_hidden"]

        self.W_g = nn.Linear(self.d_hidden, self.num_experts, bias=False)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        logits = self.W_g(x)
        scores = logits.softmax(dim=-1)
        topk_scores, topk_indices = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

        return topk_indices, topk_scores


class MoE(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([Expert(config) for _ in range(config["n_routed_experts"])])
        self.gating_network = MoEGate(config)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert(config=config, d_expert=shared_expert_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        shared_output = self.shared_expert(x)
        expert_indices, expert_scores = self.gating_network(x)
        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape
        x_flat = x.view(-1, hidden_dim)
        flat_expert_indices = expert_indices.view(-1)
        flat_expert_weights = expert_scores.view(-1, 1)
        routed_output_flat = self.moe_infer(x_flat, flat_expert_indices, flat_expert_weights)

        routed_output = routed_output_flat.view(*orig_shape)
        return routed_output + shared_output

    @torch.no_grad()
    def moe_infer(
        self, x: torch.Tensor, flat_expert_indices: torch.Tensor, flat_expert_weights: torch.Tensor
    ) -> torch.Tensor:
        # Prepare mapping from gating entries to token indices
        top_k = self.config["n_experts_per_token"]
        total_entries = flat_expert_indices.size(0)
        # Map each gating entry to its original token position
        token_idxs = torch.arange(total_entries, device=x.device) // top_k  # [M]
        # Expand input tokens per gating entry
        x_expanded = x[token_idxs]  # [M, H]

        # Stack expert weights once
        gate_up_w = torch.stack([e.W_gate_up.weight for e in self.experts], dim=0).permute(
            0, 2, 1
        )  # [E, H, 2D_e]
        down_w = torch.stack([e.W_down.weight for e in self.experts], dim=0).permute(
            0, 2, 1
        )  # [E, D_e, H]

        # Buffer for outputs
        expert_cache = torch.zeros_like(x)
        # Process in chunks to save memory
        chunk_size = 32  # reduced chunk size to fit GPU memory  # adjust as needed
        for start in range(0, total_entries, chunk_size):
            end = min(start + chunk_size, total_entries)
            idxs_chunk = flat_expert_indices[start:end]  # [C]
            weights_chunk = flat_expert_weights[start:end]  # [C,1]
            token_chunk = token_idxs[start:end]  # [C]
            x_chunk = x_expanded[start:end]  # [C, H]

            # Gather per-entry expert weights
            W_gu = gate_up_w[idxs_chunk]  # [C, H, 2D_e]
            W_do = down_w[idxs_chunk]  # [C, D_e, H]

            # Gate and up projections via batch matmul
            gu = torch.bmm(x_chunk.unsqueeze(1), W_gu).squeeze(1)  # [C,2D_e]
            gate, up = gu.chunk(2, dim=-1)  # [C,D_e]
            gate = self.shared_expert.act_fn(gate)
            gated = gate * up  # [C, D_e]

            # Down projection
            out = torch.bmm(gated.unsqueeze(1), W_do).squeeze(1)  # [C, H]
            out = out * weights_chunk  # apply routing score

            # Scatter-add into cache by original token indices
            expert_cache.index_add_(0, token_chunk, out)

        return expert_cache


def custom_kernel(data: input_t) -> output_t:
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    moe = MoE(config)

    moe.gating_network.W_g.weight = nn.Parameter(weights["router.weight"])

    for i in range(num_experts):
        gate_proj_weight = weights[f"experts.{i}.0.weight"]
        up_proj_weight = weights[f"experts.{i}.1.weight"]
        down_proj_weight = weights[f"experts.{i}.2.weight"]

        # Combine gate and up weights for fused projection
        fused = torch.cat([gate_proj_weight.t(), up_proj_weight.t()], dim=0)
        moe.experts[i].W_gate_up.weight = nn.Parameter(fused)
        moe.experts[i].W_down.weight = nn.Parameter(down_proj_weight.t())

    # Shared expert weight assignment
    sg = weights["shared_experts.0.weight"]
    sup = weights["shared_experts.1.weight"]
    sdown = weights["shared_experts.2.weight"]
    fused_shared = torch.cat([sg.t(), sup.t()], dim=0)
    moe.shared_expert.W_gate_up.weight = nn.Parameter(fused_shared)
    moe.shared_expert.W_down.weight = nn.Parameter(sdown.t())

    output = moe(input_tensor)
    return output

python eval.py benchmark /tmp/tmp6fglj8xy