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

        self.W_gate = nn.Linear(self.d_hidden, self.d_expert, bias=False)
        self.W_up = nn.Linear(self.d_hidden, 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:
        gate = self.act_fn(self.W_gate(x))
        out = self.W_down(gate * self.W_up(x))
        return out


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.num_experts: int = config["n_routed_experts"]
        self.top_k: int = config["n_experts_per_token"]

        # Create all experts
        self.experts = nn.ModuleList([Expert(config) for _ in range(self.num_experts)])
        self.gating_network = MoEGate(config)

        # A “shared” expert whose output is added for every token
        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:
        """
        x: [batch_size, seq_len, d_hidden]
        """
        # 1. Shared expert pass (no routing): shape → [B, S, d_hidden]
        shared_output = self.shared_expert(x)

        # 2. Compute routing for each token
        #    expert_indices: [B, S, top_k], expert_scores: [B, S, top_k]
        expert_indices, expert_scores = self.gating_network(x)

        batch_size, seq_len, hidden_dim = x.shape
        orig_shape = x.shape

        # 3. Flatten tokens → [B*S, d_hidden]
        x_flat = x.view(-1, hidden_dim)  # shape: [T, D], where T = B*S

        # 4. Also flatten the routing tensors → [T * top_k]
        flat_expert_indices = expert_indices.view(-1)  # [T * top_k]
        flat_expert_weights = expert_scores.view(-1, 1)  # [T * top_k, 1]

        # 5. Compute the “routed” output by grouping tokens per expert
        routed_output_flat = self.moe_infer(
            x_flat, flat_expert_indices, flat_expert_weights
        )  # returns [T, D]

        # 6. Reshape back to [B, S, D] and add the shared expert output
        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:
        """
        x:                  [T, D]       (T = batch_size*seq_len)
        flat_expert_indices: [T * top_k]  (each token replicated top_k times)
        flat_expert_weights: [T * top_k, 1]

        Returns:
            expert_cache: [T, D]  (each token’s final summed contribution across its top-k experts)
        """
        device = x.device
        T, D = x.shape
        top_k = self.config["n_experts_per_token"]

        # 1) Build a “token lookup” so that, for each entry in
        #    flat_expert_indices, we know which original token index it came from.
        #    E.g. for T=5, top_k=2, token_indices = [0,0, 1,1, 2,2, 3,3, 4,4]
        token_indices = (
            torch.arange(T, device=device).unsqueeze(1).expand(T, top_k).reshape(-1)
        )  # shape: [T * top_k]

        # 2) Prepare an empty “cache” to accumulate each token’s final outputs.
        #    It has shape [T, D].
        expert_cache = torch.zeros_like(x)  # same dtype/device as x

        # 3) Only loop over the *unique* experts that actually appear here
        unique_experts = torch.unique(flat_expert_indices)

        for expert_id in unique_experts.tolist():
            # 3a) Make a boolean mask identifying all assignments for this expert
            mask = flat_expert_indices == expert_id
            if not mask.any():
                continue

            # 3b) Which tokens (0..T-1) are assigned to this expert?
            tok_idxs = token_indices[mask]  # shape: [N_e]
            if tok_idxs.numel() == 0:
                continue

            # 3c) Gather those tokens from x
            expert_inputs = x[tok_idxs]  # [N_e, D]

            # 3d) Run them through this expert module
            expert_out = self.experts[expert_id](expert_inputs)  # [N_e, D]

            # 3e) Weight each output by the corresponding routing weight
            weights = flat_expert_weights[mask]  # [N_e, 1]
            weighted_out = expert_out * weights  # broadcast → [N_e, D]

            # 3f) Accumulate back into expert_cache via index_add_
            #    This does: expert_cache[tok_idxs[i]] += weighted_out[i] for each i
            expert_cache.index_add_(0, tok_idxs, weighted_out)

        return expert_cache


def custom_kernel(data: input_t) -> output_t:
    """
    Submission template for DeepSeek-style Mixture of Experts using PyTorch.

    Args:
        data: Tuple of (input: torch.Tensor, weights: Dict[str, torch.Tensor], config: Dict)
            - input: Input tensor of shape [batch_size, seq_len, hidden_size]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters

    Returns:
        Processed tensor [batch_size, seq_len, d_model]
    """
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    moe = MoE(config)

    # Load the router’s weight
    moe.gating_network.W_g.weight = nn.Parameter(weights["router.weight"])

    # Load each expert’s weights (transposed to match nn.Linear’s expectation)
    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"]

        moe.experts[i].W_gate.weight = nn.Parameter(gate_proj_weight.t())
        moe.experts[i].W_up.weight = nn.Parameter(up_proj_weight.t())
        moe.experts[i].W_down.weight = nn.Parameter(down_proj_weight.t())

    # Load shared expert’s weights
    moe.shared_expert.W_gate.weight = nn.Parameter(
        weights["shared_experts.0.weight"].t()
    )
    moe.shared_expert.W_up.weight = nn.Parameter(weights["shared_experts.1.weight"].t())
    moe.shared_expert.W_down.weight = nn.Parameter(
        weights["shared_experts.2.weight"].t()
    )

    # Run the forward pass
    output = moe(input_tensor)
    return output

python eval.py benchmark /tmp/tmp8pq0xlt3