import time
from typing import Dict, Tuple, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from task import input_t, output_t


class Expert_ref(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_ref(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
        # return topk_indices, topk_scores, logits, scores


class MoE_ref(nn.Module):
    def __init__(self, config: Dict):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            Expert_ref(config)
            for _ in range(config["n_routed_experts"])
        ])
        self.gating_network = MoEGate_ref(config)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert_ref(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:
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        counts = flat_expert_indices.bincount().cpu().numpy()
        tokens_per_expert = counts.cumsum()
        num_per_tok = self.config["n_experts_per_token"]
        token_idxs = idxs // num_per_tok
        for expert_id, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
            if start_idx == end_idx:
                continue

            expert = self.experts[expert_id]
            exp_token_idxs = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idxs]
            expert_out    = expert(expert_tokens)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            expert_cache.scatter_reduce_(
                0,
                exp_token_idxs.view(-1, 1).repeat(1, x.shape[-1]),
                expert_out,
                reduce='sum'
            )

        return expert_cache
    

# class Expert(nn.Module):
#     def __init__(self, config: Dict, weights: Dict, shared: bool, i: int):
#         super().__init__()
#         self.config = config
#         self.d_hidden: int = config["d_hidden"]
#         self.d_expert: int = config["d_expert"]

#         if shared:
#             self.W_gate = weights[f'shared_experts.0.weight'].t().clone()
#             self.W_up   = weights[f'shared_experts.1.weight'].t().clone()
#             self.W_down = weights[f'shared_experts.2.weight'].t().clone()
#         else:
#             self.W_gate = weights[f'experts.{i}.0.weight'].t().clone()
#             self.W_up   = weights[f'experts.{i}.1.weight'].t().clone()
#             self.W_down = weights[f'experts.{i}.2.weight'].t().clone()

#     def forward(self, x: torch.Tensor) -> torch.Tensor:
#         x_gate = F.linear(x, self.W_gate)
#         x_gate = F.silu(x_gate)
#         x_up = F.linear(x, self.W_up)
#         out = x_gate * x_up
#         out = F.linear(out, self.W_down)

#         return out

class Expert(nn.Module):
    def __init__(self, config: Dict, weights: Dict, shared: bool, i: int):
        super().__init__()
        self.config = config
        self.d_hidden: int = config["d_hidden"]
        self.d_expert: int = config["d_expert"]

        if shared:
            self.W_gate = weights[f'shared_experts.0.weight'].clone()
            self.W_up   = weights[f'shared_experts.1.weight'].clone()
            self.W_down = weights[f'shared_experts.2.weight'].clone()
        else:
            self.W_gate = weights[f'experts.{i}.0.weight'].clone()
            self.W_up   = weights[f'experts.{i}.1.weight'].clone()
            self.W_down = weights[f'experts.{i}.2.weight'].clone()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_gate = x @ self.W_gate
        x_gate = F.silu(x_gate)
        x_up = x @ self.W_up
        out = x_gate * x_up
        out = out @ self.W_down

        return out


class MoEGate(nn.Module):
    def __init__(self, config: Dict, weights: 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 = weights['router.weight'].clone()

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        # logits = x @ self.W_g
        logits = F.linear(x, self.W_g)
        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
        # return topk_indices, topk_scores, logits, scores


class MoE(nn.Module):
    def __init__(self, config: Dict, weights: Dict):
        super().__init__()
        self.config = config
        self.experts = nn.ModuleList([
            # Expert_ref(config)
            Expert(config, weights, False, i)
            for i in range(config["n_routed_experts"])
        ])
        self.gating_network = MoEGate_ref(config)
        # self.gating_network = MoEGate(config, weights)
        shared_expert_dim = config["d_expert"] * config["n_shared_experts"]
        self.shared_expert = Expert_ref(config, d_expert=shared_expert_dim)
        # self.shared_expert = Expert(config, weights, True, 0)

    @torch.no_grad()
    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:
        expert_cache = torch.zeros_like(x)
        idxs = flat_expert_indices.argsort()
        counts = flat_expert_indices.bincount().cpu().numpy()
        tokens_per_expert = counts.cumsum()
        num_per_tok = self.config["n_experts_per_token"]
        token_idxs = idxs // num_per_tok
        for expert_id, end_idx in enumerate(tokens_per_expert):
            start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
            if start_idx == end_idx:
                continue

            expert = self.experts[expert_id]
            exp_token_idxs = token_idxs[start_idx:end_idx]
            expert_tokens = x[exp_token_idxs]
            expert_out    = expert(expert_tokens)
            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
            expert_cache.scatter_reduce_(
                0,
                exp_token_idxs.view(-1, 1).repeat(1, x.shape[-1]),
                expert_out,
                reduce='sum'
            )

        return expert_cache


# @torch.no_grad()
def moe_gate_kernel(x: torch.Tensor, w_g: torch.Tensor, topk: int):
    logits = F.linear(x, w_g)
    scores = logits.softmax(dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=topk, dim=-1, sorted=False)

    return topk_indices, topk_scores

# @torch.no_grad()
def moe_expert_kernel(x: torch.Tensor, w_gate: torch.Tensor, w_up: torch.Tensor, w_down: torch.Tensor):
    x_gate = x @ w_gate
    x_gate = F.silu(x_gate)
    x_up = x @ w_up
    out = x_gate * x_up
    out = out @ w_down
    return out

# @torch.no_grad()
def moe_kernel(x: torch.Tensor, weights: Dict, config: Dict):
    shared_output = moe_expert_kernel(
        x,
        weights['shared_experts.0.weight'],
        weights['shared_experts.1.weight'],
        weights['shared_experts.2.weight'])
    expert_indices, expert_scores = moe_gate_kernel(
        x,
        weights['router.weight'],
        config["n_experts_per_token"])
    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)

    expert_cache = torch.zeros_like(x_flat)
    idxs = flat_expert_indices.argsort()
    counts = flat_expert_indices.bincount().cpu().numpy()
    tokens_per_expert = counts.cumsum()
    num_per_tok = config["n_experts_per_token"]
    token_idxs = idxs // num_per_tok
    for expert_id, end_idx in enumerate(tokens_per_expert):
        start_idx = 0 if expert_id == 0 else tokens_per_expert[expert_id - 1]
        if start_idx == end_idx:
            continue
        exp_token_idxs = token_idxs[start_idx:end_idx]
        expert_tokens = x_flat[exp_token_idxs]
        expert_out = moe_expert_kernel(
            expert_tokens,
            weights[f'experts.{expert_id}.0.weight'],
            weights[f'experts.{expert_id}.1.weight'],
            weights[f'experts.{expert_id}.2.weight'])
        expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
        expert_cache.scatter_reduce_(
            0,
            exp_token_idxs.view(-1, 1).repeat(1, x_flat.shape[-1]),
            expert_out,
            reduce='sum'
        )

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


def custom_kernel(data: input_t) -> output_t:
    """
    Reference implementation of 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_dim]
            - weights: Dictionary containing model weights
            - config: Dictionary containing model configuration parameters
            
    Returns:
        Tuple containing:
            - output: Processed tensor [batch_size, seq_len, d_model]
            - aux_data: Dictionary with auxiliary data
    """
    # print(f"begin kernel {time.time_ns()}")
    input_tensor, weights, config = data
    # moe = MoE(config, weights)

    # print(f"begin copy {time.time_ns()}")

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

    # for i in range(config["n_routed_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']

    #     # Transpose weights to match expected shape for nn.Linear
    #     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())

    # 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())

    # print(f"begin compute {time.time_ns()}")
    # output = moe(input_tensor)
    # print(f"end compute {time.time_ns()}")

    output = moe_kernel(input_tensor, weights, config)
    return output


# if __name__ == "__main__":
#     data = generate_input(7168, 2048, 128, 1, 8, 1, 1024, 17)
#     custom_kernel(data)

    # input_tensor, weights, config = data

    # moe_gate = MoEGate(config, weights)
    # moe_gate_ref = MoEGate_ref(config)
    # moe_gate_ref.W_g.weight = nn.Parameter(weights['router.weight'])

    # output = moe_gate(input_tensor)
    # output_ref = moe_gate_ref(input_tensor)

    # print(verbose_allclose(output[0], output_ref[0], rtol=1e-2, atol=1e-2))
    # print(verbose_allclose(output[1], output_ref[1], rtol=1e-2, atol=1e-2))
    # print(verbose_allclose(output[2], output_ref[2], rtol=1e-2, atol=1e-2))
    # print(verbose_allclose(output[3], output_ref[3], rtol=1e-2, atol=1e-2))

    # for i in range(config["n_routed_experts"]):
    #     expert = Expert(config, weights, False, i)
    #     expert_ref = Expert_ref(config)
    #     expert_ref.W_gate.weight = nn.Parameter(weights[f'experts.{i}.0.weight'].t())
    #     expert_ref.W_up.weight = nn.Parameter(weights[f'experts.{i}.1.weight'].t())
    #     expert_ref.W_down.weight = nn.Parameter(weights[f'experts.{i}.2.weight'].t())

    #     output = expert(input_tensor)
    #     output_ref = expert_ref(input_tensor)

    #     print(f"expert {i}")
    #     print(verbose_allclose(output, output_ref, rtol=1e-2, atol=1e-2))

    # moe = MoE(config, weights)
    # moe_ref = MoE_ref(config)

    # moe_ref.gating_network.W_g.weight = nn.Parameter(weights['router.weight'])

    # for i in range(config["n_routed_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']

    #     # Transpose weights to match expected shape for nn.Linear
    #     moe_ref.experts[i].W_gate.weight = nn.Parameter(gate_proj_weight.t())
    #     moe_ref.experts[i].W_up.weight = nn.Parameter(up_proj_weight.t())
    #     moe_ref.experts[i].W_down.weight = nn.Parameter(down_proj_weight.t())

    # moe_ref.shared_expert.W_gate.weight = nn.Parameter(weights['shared_experts.0.weight'].t())
    # moe_ref.shared_expert.W_up.weight = nn.Parameter(weights['shared_experts.1.weight'].t())
    # moe_ref.shared_expert.W_down.weight = nn.Parameter(weights['shared_experts.2.weight'].t())

    # output = moe(input_tensor)
    # output_ref = moe_ref(input_tensor)

    # print(verbose_allclose(output, output_ref, rtol=1e-2, atol=1e-2))

python eval.py benchmark /tmp/tmphu83oq1h