from functools import partial
from typing import Callable, Optional, Union
import torch
import torch.nn.functional as F
DEFAULT_ACTIVATION_FN = partial(F.gelu, approximate='tanh')

class _UniformExpertAssignment(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x: torch.Tensor, num_experts: int):
        out = torch.arange(x.numel(), dtype=x.dtype, device=x.device)
        out = torch.remainder(out, num_experts)
        return out.view(x.shape)

class LearnedRouter(torch.nn.Module):

    def __init__(self, hidden_size: int, moe_num_experts: int, moe_top_k: int, moe_jitter_eps: Optional[float], moe_normalize_expert_weights: Optional[Union[int, float]], uniform_expert_assignment: bool, device: Optional[torch.device]) -> None:
        super().__init__()
        self.hidden_size: int = hidden_size
        self.moe_num_experts: int = moe_num_experts
        self.moe_top_k: int = moe_top_k
        self.moe_jitter_eps: Optional[float] = moe_jitter_eps
        self.moe_normalize_expert_weights: Optional[Union[int, float]] = moe_normalize_expert_weights
        self.uniform_expert_assignment: bool = uniform_expert_assignment
        self.layer: torch.nn.Module = torch.nn.Linear(hidden_size, moe_num_experts, bias=False, device=device)

    def jitter(self, x: torch.Tensor) -> torch.Tensor:
        assert self.moe_jitter_eps is not None
        low: float = 1.0 - self.moe_jitter_eps
        high: float = 1.0 + self.moe_jitter_eps
        noise: torch.Tensor = torch.rand(x.size(), dtype=x.dtype, device=x.device)
        return low + noise * (high - low)

    def _top_k(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if self.moe_top_k == 1:
            values, indices = scores.max(dim=-1)
            return (values.unsqueeze(-1), indices.unsqueeze(-1))
        return torch.topk(scores, self.moe_top_k, dim=-1)

    def forward(self, x: torch.Tensor):
        if self.training and self.moe_jitter_eps is not None:
            x = x * self.jitter(x)
        scores = self.layer(x.view(-1, x.shape[-1])).softmax(dim=-1)
        expert_weights, top_experts = self._top_k(scores)
        if self.moe_normalize_expert_weights:
            expert_weights = expert_weights / torch.norm(expert_weights, p=self.moe_normalize_expert_weights, dim=-1, keepdim=True)
        top_experts = _UniformExpertAssignment.apply(top_experts, self.moe_num_experts) if self.uniform_expert_assignment else top_experts
        scores = scores.to(x.dtype)
        expert_weights = expert_weights.to(x.dtype)
        return (scores, expert_weights, top_experts)

class MLP(torch.nn.Module):

    def __init__(self, hidden_size: int, ffn_hidden_size: int, moe_num_experts: int, activation_fn: Callable, device: Optional[torch.device]) -> None:
        super().__init__()
        self.moe_num_experts: int = moe_num_experts
        self.ffn_hidden_size: int = ffn_hidden_size
        self.hidden_size: int = hidden_size
        self.activation_fn: Callable = activation_fn
        self.w1 = torch.nn.Parameter(torch.rand(moe_num_experts * ffn_hidden_size, hidden_size, device=device))
        self.w2 = torch.nn.Parameter(torch.rand(moe_num_experts * ffn_hidden_size, hidden_size, device=device))
        self.activation_fn = activation_fn

    def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor:
        expert_w1 = self.w1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size)[expert_idx]
        expert_w2 = self.w2.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size)[expert_idx]
        before_activation = x @ expert_w1.t()
        layer_1_output = self.activation_fn(before_activation)
        output = layer_1_output @ expert_w2
        return output

class GLU(torch.nn.Module):

    def __init__(self, hidden_size: int, ffn_hidden_size: int, moe_num_experts: int, activation_fn: Callable, device: Optional[torch.device]):
        super().__init__()
        self.hidden_size = hidden_size
        self.ffn_hidden_size = ffn_hidden_size
        self.moe_num_experts = moe_num_experts
        self.w1 = torch.nn.Parameter(torch.rand(moe_num_experts * ffn_hidden_size, hidden_size, device=device))
        self.v1 = torch.nn.Parameter(torch.rand(moe_num_experts * ffn_hidden_size, hidden_size, device=device))
        self.w2 = torch.nn.Parameter(torch.rand(moe_num_experts * ffn_hidden_size, hidden_size, device=device))
        self.activation_fn = activation_fn

    def forward(self, x: torch.Tensor, expert_idx: torch.Tensor):
        expert_w1 = self.w1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size)[expert_idx]
        expert_v1 = self.v1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size)[expert_idx]
        expert_w2 = self.w2.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size)[expert_idx]
        x1 = x.matmul(expert_w1.t())
        x2 = x.matmul(expert_v1.t())
        x1 = self.activation_fn(x1)
        x1 = x1 * x2
        x1 = x1.matmul(expert_w2)
        return x1

class DroplessMLP(torch.nn.Module):

    def __init__(self, hidden_size: int, ffn_hidden_size: int, mlp_type: str, moe_num_experts: int, activation_fn: Callable, bias: bool, device: Optional[torch.device]):
        super().__init__()
        self.moe_num_experts = moe_num_experts
        if mlp_type == 'mlp':
            self.mlp = MLP(hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size, moe_num_experts=moe_num_experts, activation_fn=activation_fn, device=device)
        elif mlp_type == 'glu':
            self.mlp = GLU(hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size, moe_num_experts=moe_num_experts, activation_fn=activation_fn, device=device)
        else:
            raise ValueError(f'Received unknown mlp_type={mlp_type!r}')

    def forward(self, x: torch.Tensor, scores: torch.Tensor, expert_weights: torch.Tensor, top_experts: torch.Tensor):
        in_shape = x.shape
        hidden_size = in_shape[-1]
        x = x.view(-1, hidden_size)
        out = torch.zeros_like(x)
        expert_mask = torch.nn.functional.one_hot(top_experts, num_classes=self.moe_num_experts).permute(2, 1, 0)
        for expert_idx in range(0, self.moe_num_experts):
            topk_idx, token_idx = torch.where(expert_mask[expert_idx])
            if token_idx.shape[0] == 0:
                continue
            token_list = token_idx.tolist()
            topk_list = topk_idx.tolist()
            expert_tokens = x[None, token_list].reshape(-1, hidden_size)
            mlp_output = self.mlp(expert_tokens, expert_idx)
            expert_weights = expert_weights.to(mlp_output.device)
            expert_out = mlp_output * expert_weights[token_list, topk_list, None]
            out = out.to(mlp_output.device)
            token_idx = token_idx.to(mlp_output.device)
            out.index_add_(0, token_idx, expert_out)
        out = out.view(in_shape)
        return out

class dMoE(torch.nn.Module):

    def __init__(self, device: Optional[torch.device], hidden_size: int=1024, ffn_hidden_size: int=4096, moe_num_experts: int=1, moe_top_k: int=1, mlp_type: str='mlp', activation_fn: Callable=DEFAULT_ACTIVATION_FN, moe_jitter_eps: Optional[float]=None, moe_normalize_expert_weights: Optional[Union[int, float]]=None, uniform_expert_assignment: bool=False, bias: bool=True):
        super().__init__()
        self.router = LearnedRouter(hidden_size, moe_num_experts=moe_num_experts, moe_top_k=moe_top_k, moe_jitter_eps=moe_jitter_eps, moe_normalize_expert_weights=moe_normalize_expert_weights, uniform_expert_assignment=uniform_expert_assignment, device=device)
        self.experts = DroplessMLP(hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size, mlp_type=mlp_type, moe_num_experts=moe_num_experts, activation_fn=activation_fn, bias=bias, device=device)

    def forward(self, x: torch.Tensor):
        scores, expert_weights, top_experts = self.router(x)
        return self.experts(x, scores, expert_weights, top_experts)