import torch
import torch.nn as nn
import torch.nn.functional as F
from model.config import ModelConfig


class SwiGLU(nn.Module):
    """SwiGLU activation: gate(x) * value(x) with SiLU gating."""
    def __init__(self, n_embd, intermediate_size):
        super().__init__()
        self.w1 = nn.Linear(n_embd, intermediate_size, bias=False)  # Gate
        self.w2 = nn.Linear(intermediate_size, n_embd, bias=False)  # Output
        self.w3 = nn.Linear(n_embd, intermediate_size, bias=False)  # Value

    def forward(self, x):
        # Try fused CUDA SwiGLU if available
        if x.is_cuda:
            try:
                from model.custom_ops import fused_swiglu
                gate = self.w1(x)
                val = self.w3(x)
                return self.w2(fused_swiglu(gate, val))
            except Exception:
                pass
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class SparseMoE(nn.Module):
    """
    DeepSeek-inspired Sparse Mixture of Experts with:
      • Shared expert(s) always active for common patterns
      • Top-K routing with load-balancing auxiliary loss
      • Expert capacity factor to prevent token dropping
      • CPU offloading for idle experts (VRAM savings)
    """
    def __init__(self, config: ModelConfig):
        super().__init__()
        n_embd = config.n_embd
        intermediate_size = config.moe_intermediate_size
        n_experts = config.n_experts
        n_shared = getattr(config, 'n_shared_experts', 1)

        self.n_experts = n_experts
        self.n_shared = n_shared
        self.top_k = min(config.top_k, n_experts)
        self.capacity_factor = getattr(config, 'expert_capacity_factor', 1.25)

        # Gating network (routes tokens to routed experts only)
        self.gate = nn.Linear(n_embd, n_experts, bias=False)

        # Shared expert(s) — always active, no routing needed
        self.shared_experts = nn.ModuleList([
            SwiGLU(n_embd, intermediate_size) for _ in range(n_shared)
        ])

        # Routed expert pool
        self.experts = nn.ModuleList([
            SwiGLU(n_embd, intermediate_size) for _ in range(n_experts)
        ])

        # Expert offloading state
        self.register_buffer('experts_on_device', torch.ones(n_experts, dtype=torch.bool))

    def offload_experts(self, expert_indices):
        """Moves specified experts to CPU RAM to save VRAM."""
        for idx in expert_indices:
            self.experts[idx].to('cpu')
            self.experts_on_device[idx] = False

    def fetch_expert(self, expert_idx, device):
        """Brings an expert back to GPU when needed."""
        if not self.experts_on_device[expert_idx]:
            self.experts[expert_idx].to(device)
            self.experts_on_device[expert_idx] = True

    def forward(self, x):
        B, T, C = x.shape
        x_flat = x.view(-1, C)  # (N, C)
        N = x_flat.size(0)

        # ── Shared expert output (always computed) ──
        shared_out = torch.zeros_like(x_flat)
        for se in self.shared_experts:
            shared_out = shared_out + se(x_flat)
        if self.n_shared > 1:
            shared_out = shared_out / self.n_shared

        # ── Routing for sparse experts ──
        gate_logits = self.gate(x_flat)  # (N, n_experts)

        # Noise for exploration during training
        if self.training:
            noise = torch.randn_like(gate_logits) * 0.1
            gate_logits = gate_logits + noise

        # Top-K selection
        weights, indices = torch.topk(gate_logits, self.top_k, dim=-1)
        weights = F.softmax(weights, dim=-1)

        # ── Load-Balancing Auxiliary Loss ──
        # Probability distribution per expert
        probs = F.softmax(gate_logits, dim=-1)  # (N, n_experts)
        p_avg = probs.mean(dim=0)  # Mean probability per expert

        # Fraction of tokens routed to each expert (top-1)
        f = torch.zeros(self.n_experts, device=x.device)
        top1_indices = indices[:, 0]
        f.scatter_add_(0, top1_indices, torch.ones_like(top1_indices, dtype=x.dtype))
        f = f / max(N, 1)

        # Standard MoE load-balancing loss
        balance_loss = self.n_experts * torch.sum(p_avg * f)

        # Entropy regularization — encourage uniform expert usage
        entropy = -(probs * (probs + 1e-10).log()).sum(dim=-1).mean()
        aux_loss = balance_loss - 0.01 * entropy

        # ── Route tokens to experts ──
        routed_out = torch.zeros_like(x_flat)
        for i in range(self.n_experts):
            expert_mask = (indices == i).any(dim=-1)
            if not expert_mask.any():
                continue

            # Dynamic expert fetching (offload support)
            self.fetch_expert(i, x.device)

            for k in range(self.top_k):
                k_mask = (indices[:, k] == i)
                if not k_mask.any():
                    continue

                token_indices = k_mask.nonzero().squeeze(-1)
                expert_output = self.experts[i](x_flat[token_indices])
                routed_out[token_indices] += weights[token_indices, k:k+1] * expert_output

        # ── Combine shared + routed ──
        output = shared_out + routed_out
        return output.view(B, T, C), aux_loss