import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple, List

# ==============================================================
# 1. Expert Module
# ==============================================================

class Expert(nn.Module):
    """
    GMoE Expert Module.
    A standard FFN with two linear layers and activation.
    """
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        dropout: float = 0.1,
        activation: str = "gelu"
    ):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)
        
        if activation == "gelu":
            self.act = nn.GELU()
        elif activation == "relu":
            self.act = nn.ReLU()
        elif activation == "silu":
            self.act = nn.SiLU()
        else:
            raise ValueError(f"Unknown activation: {activation}")
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (batch_size, seq_len, d_model)
        Returns:
            (batch_size, seq_len, d_model)
        """
        return self.w2(self.dropout(self.act(self.w1(x))))


# ==============================================================
# 2. GRU-based Global Router
# ==============================================================

class GRURouter(nn.Module):
    """
    GMoE's Global Router with GRU-based recurrent routing.
    Core innovation: propagates routing logits across layers 
    to solve the path collapse problem.
    """
    def __init__(
        self,
        d_model: int,
        num_global_experts: int,
        gru_hidden_size: int = 128,
        num_experts_per_token: int = 2,
    ):
        super().__init__()
        self.num_global_experts = num_global_experts
        self.num_experts_per_token = num_experts_per_token
        
        # Input projection: map input to GRU hidden space
        self.input_proj = nn.Linear(d_model, gru_hidden_size, bias=False)
        
        # GRU cell: passes routing state across layers
        self.gru_cell = nn.GRUCell(gru_hidden_size, gru_hidden_size)
        
        # Routing head: generates expert selection scores for each token
        self.routing_head = nn.Linear(gru_hidden_size, num_global_experts, bias=False)
        
        self._init_weights()
    
    def _init_weights(self):
        for name, param in self.named_parameters():
            if 'weight' in name:
                nn.init.xavier_uniform_(param)
            elif 'bias' in name:
                nn.init.zeros_(param)
    
    def forward(
        self,
        x: torch.Tensor,
        prev_hidden: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Args:
            x: (batch_size, seq_len, d_model) current layer input
            prev_hidden: previous layer GRU hidden state
        Returns:
            routing_weights: (batch, seq, num_experts) routing weights
            selected_experts: (batch, seq, K) selected expert indices
            hidden_state: current GRU state for next layer
        """
        batch_size, seq_len, _ = x.shape
        
        # 1. Input projection
        projected = self.input_proj(x)
        
        # 2. GRU state update
        if prev_hidden is None:
            prev_hidden = torch.zeros(
                batch_size, seq_len, self.gru_cell.hidden_size,
                device=x.device, dtype=x.dtype
            )
        
        # Flatten sequence dimension for GRUCell (requires 2D input)
        flat_projected = projected.view(-1, projected.size(-1))
        flat_prev = prev_hidden.view(-1, prev_hidden.size(-1))
        
        flat_hidden = self.gru_cell(flat_projected, flat_prev)
        hidden_state = flat_hidden.view(batch_size, seq_len, -1)
        
        # 3. Generate routing logits
        logits = self.routing_head(hidden_state)
        
        # 4. Top-K selection
        routing_weights, selected_experts = self._top_k_routing(logits)
        
        return routing_weights, selected_experts, hidden_state
    
    def _top_k_routing(
        self, logits: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Top-K routing selection per token."""
        scores = F.softmax(logits, dim=-1)
        routing_weights, selected_experts = torch.topk(
            scores, self.num_experts_per_token, dim=-1
        )
        routing_weights = routing_weights / (
            routing_weights.sum(dim=-1, keepdim=True) + 1e-8
        )
        return routing_weights, selected_experts


# ==============================================================
# 3. GMoE Layer
# ==============================================================

class GMoELayer(nn.Module):
    """
    Core GMoE Layer.
    Contains: 1 Local Expert + Global Expert Pool + Global Router.
    """
    def __init__(
        self,
        d_model: int,
        d_ff: int,
        num_global_experts: int,
        num_experts_per_token: int = 2,
        dropout: float = 0.1,
        gru_hidden_size: int = 128,
        layer_id: int = 0,
    ):
        super().__init__()
        self.layer_id = layer_id
        self.d_model = d_model
        self.num_global_experts = num_global_experts
        self.num_experts_per_token = num_experts_per_token
        
        # Local expert (unique per layer)
        self.local_expert = Expert(d_model, d_ff, dropout)
        
        # Global expert reference (set by GMoEModel)
        self.global_experts: Optional[nn.ModuleList] = None
        
        # Global router with GRU
        self.router = GRURouter(
            d_model=d_model,
            num_global_experts=num_global_experts,
            gru_hidden_size=gru_hidden_size,
            num_experts_per_token=num_experts_per_token,
        )
        
        self.norm = nn.LayerNorm(d_model)
    
    def set_global_experts(self, global_experts: nn.ModuleList):
        """Set reference to the global expert pool."""
        self.global_experts = global_experts
    
    def forward(
        self,
        x: torch.Tensor,
        prev_routing_state: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        residual = x
        x = self.norm(x)
        
        # 1. Local expert output
        local_output = self.local_expert(x)
        
        # 2. Routing
        routing_weights, selected_experts, routing_state = self.router(x, prev_routing_state)
        
        # 3. Sparse global expert forward
        global_output = self._sparse_global_forward(x, routing_weights, selected_experts)
        
        # 4. Combine outputs
        output = residual + local_output + global_output
        
        return output, routing_state
    
    def _sparse_global_forward(
        self,
        x: torch.Tensor,
        routing_weights: torch.Tensor,
        selected_experts: torch.Tensor
    ) -> torch.Tensor:
        """
        Sparse activation of global experts.
        Each token activates only K experts.
        """
        assert self.global_experts is not None, "Global experts not set!"
        
        batch_size, seq_len, d_model = x.shape
        K = self.num_experts_per_token
        num_experts = len(self.global_experts)
        
        final_output = torch.zeros_like(x)
        
        flat_x = x.view(-1, d_model)
        flat_weights = routing_weights.view(-1, K)
        flat_experts = selected_experts.view(-1, K)
        
        total_tokens = flat_x.size(0)
        
        for expert_idx in range(num_experts):
            mask = (flat_experts == expert_idx)
            
            if not mask.any():
                continue
            
            token_indices, expert_positions = torch.where(mask)
            selected_x = flat_x[token_indices]
            selected_weights = flat_weights[token_indices, expert_positions]
            
            expert_output = self.global_experts[expert_idx](selected_x)
            weighted_output = expert_output * selected_weights.unsqueeze(-1)
            
            final_output.view(-1, d_model).index_add_(
                0, token_indices, weighted_output
            )
        
        return final_output


# ==============================================================
# 4. Complete GMoE Model
# ==============================================================

class GMoEModel(nn.Module):
    """
    Complete GMoE Transformer model.
    Embedding + N GMoE layers + Global Expert Pool + Output layer.
    """
    def __init__(
        self,
        vocab_size: int = 50257,
        d_model: int = 768,
        d_ff: int = 3072,
        num_layers: int = 12,
        num_global_experts: int = 16,
        num_experts_per_token: int = 2,
        num_heads: int = 12,
        dropout: float = 0.1,
        max_seq_len: int = 1024,
        gru_hidden_size: int = 128,
    ):
        super().__init__()
        self.d_model = d_model
        self.num_layers = num_layers
        self.num_global_experts = num_global_experts
        
        # Token embedding
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, max_seq_len, dropout)
        
        # Global expert pool (shared across all layers)
        self.global_experts = nn.ModuleList([
            Expert(d_model, d_ff, dropout)
            for _ in range(num_global_experts)
        ])
        
        # GMoE layers
        self.layers = nn.ModuleList([
            GMoELayer(
                d_model=d_model, d_ff=d_ff,
                num_global_experts=num_global_experts,
                num_experts_per_token=num_experts_per_token,
                dropout=dropout, gru_hidden_size=gru_hidden_size,
                layer_id=i,
            )
            for i in range(num_layers)
        ])
        
        # Set global expert references
        for layer in self.layers:
            layer.set_global_experts(self.global_experts)
        
        # Output layer
        self.final_norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.embedding.weight
        
        self._init_weights()
    
    def _init_weights(self):
        for name, param in self.named_parameters():
            if 'weight' in name and param.dim() >= 2:
                nn.init.xavier_uniform_(param)
            elif 'bias' in name:
                nn.init.zeros_(param)
    
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        x = self.embedding(input_ids)
        x = self.pos_encoding(x)
        
        routing_state = None
        for layer in self.layers:
            x, routing_state = layer(x, routing_state)
        
        x = self.final_norm(x)
        logits = self.lm_head(x)
        return logits
    
    def count_active_parameters(self) -> dict:
        """Count active parameters during inference."""
        total_params = sum(p.numel() for p in self.parameters())
        
        must_params = 0
        for name, p in self.named_parameters():
            if 'global_experts' not in name:
                must_params += p.numel()
        
        expert_params = sum(p.numel() for p in self.global_experts.parameters())
        active_expert_params = expert_params * 2 // self.num_global_experts  
        
        return {
            "total_params": total_params,
            "must_params": must_params,
            "expert_params": expert_params,
            "active_params_per_token": must_params + active_expert_params,
        }


class PositionalEncoding(nn.Module):
    """Sinusoidal positional encoding."""
    def __init__(self, d_model: int, max_len: int, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)

# ==============================================================
# 5. Auxiliary Load Balancing Loss
# ==============================================================

class GMoELoss(nn.Module):
    """
    GMoE auxiliary loss function.
    Includes load balancing loss and router Z-loss.
    Reference: https://arxiv.org/abs/2101.03961 (Switch Transformer)
    """
    def __init__(self, alpha: float = 0.01):
        super().__init__()
        self.alpha = alpha
    
    def forward(
        self,
        logits: torch.Tensor,
        labels: torch.Tensor,
        routing_weights_list: List[torch.Tensor],
        selected_experts_list: List[torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        total_loss = cross_entropy_loss + alpha * load_balancing_loss
        """
        # 1. Cross-entropy loss
        vocab_size = logits.size(-1)
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()
        
        ce_loss = F.cross_entropy(
            shift_logits.view(-1, vocab_size),
            shift_labels.view(-1),
            ignore_index=-100,
        )
        
        # 2. Load balancing loss
        balance_loss = 0.0
        for routing_weights, selected_experts in zip(
            routing_weights_list, selected_experts_list
        ):
            balance_loss += self._load_balancing_loss(
                routing_weights, selected_experts
            )
        balance_loss = balance_loss / len(routing_weights_list)
        
        # 3. Total loss
        total_loss = ce_loss + self.alpha * balance_loss
        
        return total_loss, ce_loss, balance_loss
    
    def _load_balancing_loss(
        self,
        routing_weights: torch.Tensor,
        selected_experts: torch.Tensor,
    ) -> torch.Tensor:
        """
        Load balancing loss encourages uniform expert utilization.
        loss = N * sum(f_i * P_i) where f_i is frequency and P_i is average weight.
        """
        batch_size, seq_len, num_experts = routing_weights.shape
        K = selected_experts.size(-1)
        total_tokens = batch_size * seq_len
        
        flat_experts = selected_experts.view(-1, K)
        
        expert_freq = torch.zeros(num_experts, device=routing_weights.device)
        for k in range(K):
            expert_freq += torch.bincount(
                flat_experts[:, k].long(),
                minlength=num_experts,
            ).float()
        expert_freq = expert_freq / total_tokens
        
        flat_weights = routing_weights.view(-1, num_experts)
        expert_weight = flat_weights.mean(dim=0)
        
        balance_loss = num_experts * (expert_freq * expert_weight).sum()
        return balance_loss


# ==============================================================
# 6. Training Loop
# ==============================================================

def train_step(
    model: GMoEModel,
    loss_fn: GMoELoss,
    optimizer: torch.optim.Optimizer,
    batch: Tuple[torch.Tensor, torch.Tensor],
    device: torch.device,
) -> dict:
    """Single training step."""
    input_ids, labels = batch
    input_ids = input_ids.to(device)
    labels = labels.to(device)
    
    optimizer.zero_grad()
    logits = model(input_ids)
    
    # In a complete implementation, we'd collect routing info from the model
    # For demonstration, we use dummy routing data
    routing_weights_list = []
    selected_experts_list = []
    for layer in model.layers:
        dummy_weights = torch.zeros(
            input_ids.size(0), input_ids.size(1), model.num_global_experts,
            device=device
        )
        dummy_experts = torch.zeros(
            input_ids.size(0), input_ids.size(1), model.num_global_experts,
            device=device, dtype=torch.long
        )
        routing_weights_list.append(dummy_weights)
        selected_experts_list.append(dummy_experts)
    
    total_loss, ce_loss, balance_loss = loss_fn(
        logits, labels, routing_weights_list, selected_experts_list
    )
    
    total_loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    
    return {
        "total_loss": total_loss.item(),
        "ce_loss": ce_loss.item(),
        "balance_loss": balance_loss.item(),
    }


# ==============================================================
# 7. Parameter Efficiency Comparison
# ==============================================================

def parameter_efficiency_comparison():
    """
    Compare parameter efficiency of traditional MoE vs GMoE.
    """
    print("=" * 70)
    print("GMoE vs Traditional MoE: Parameter Efficiency Comparison")
    print("=" * 70)
    
    d_model = 768
    d_ff = 3072
    num_layers = 12
    traditional_experts_per_layer = 64
    num_global_experts = 16
    num_experts_per_token = 2
    vocab_size = 50257
    gru_hidden_size = 128
    
    expert_params = 2 * d_model * d_ff
    
    traditional_total_expert_params = (
        num_layers * traditional_experts_per_layer * expert_params
    )
    
    gmoe_global_expert_params = num_global_experts * expert_params
    gmoe_local_expert_params = num_layers * 1 * expert_params
    gmoe_total_expert_params = gmoe_global_expert_params + gmoe_local_expert_params
    
    # Non-expert parameters
    non_expert_params = (
        vocab_size * d_model +                                    # embedding
        num_layers * 4 * d_model * d_model +                      # attention
        num_layers * 2 * d_model +                                # LayerNorm
        d_model * vocab_size +                                    # output
        num_layers * gru_hidden_size * (d_model + gru_hidden_size + 
                                        d_model + num_global_experts)  # router
    )
    
    traditional_total = traditional_total_expert_params + non_expert_params
    gmoe_total = gmoe_total_expert_params + non_expert_params
    
    print(f"\nModel Configuration:")
    print(f"  d_model={d_model}, d_ff={d_ff}, num_layers={num_layers}")
    print(f"  Traditional MoE: {traditional_experts_per_layer} experts/layer")
    print(f"  GMoE: {num_global_experts} global experts + 1 local expert/layer")
    print(f"  Experts activated per token: K={num_experts_per_token}")
    
    print(f"\n--- Total Parameters ---")
    print(f"  Traditional MoE: {traditional_total:,}")
    print(f"  GMoE:            {gmoe_total:,}")
    print(f"  Compression:     {gmoe_total / traditional_total * 100:.1f}%")
    print(f"  Reduction:       {(1 - gmoe_total / traditional_total) * 100:.1f}%")
    
    print(f"\n--- Expert Parameters ---")
    print(f"  Traditional MoE: {traditional_total_expert_params:,}")
    print(f"  GMoE:            {gmoe_total_expert_params:,}")
    print(f"  Compression:     {gmoe_total_expert_params / traditional_total_expert_params * 100:.1f}%")
    
    return {
        "traditional_total": traditional_total,
        "gmoe_total": gmoe_total,
        "compression_ratio": gmoe_total / traditional_total,
    }


if __name__ == "__main__":
    results = parameter_efficiency_comparison()
    
    print("\n" + "=" * 70)
    print("Creating GMoE model instance...")
    print("=" * 70)
    
    model = GMoEModel(
        vocab_size=50257,
        d_model=256,
        d_ff=1024,
        num_layers=4,
        num_global_experts=8,
        num_experts_per_token=2,
        num_heads=4,
        dropout=0.1,
        max_seq_len=512,
        gru_hidden_size=64,
    )
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"  Total parameters: {total_params:,}")
    print(f"  GPU memory: {total_params * 4 / 1024 / 1024:.2f} MB (FP32)")
    
    batch_size, seq_len = 2, 128
    dummy_input = torch.randint(0, 50257, (batch_size, seq_len))
    
    with torch.no_grad():
        output = model(dummy_input)
    
    print(f"  Input shape: {dummy_input.shape}")
    print(f"  Output shape: {output.shape}")
    print(f"  Forward pass successful!")

// ==============================================================
// GMoE Inference Engine (Go Implementation)
// Focused on high-performance inference for edge deployment
// ==============================================================

package main

import (
	"encoding/binary"
	"fmt"
	"math"
	"os"
)

// ==============================================================
// Data Type Definitions
// ==============================================================

// Matrix represents a 2D matrix in row-major order
type Matrix struct {
	Rows int
	Cols int
	Data []float32
}

func NewMatrix(rows, cols int) *Matrix {
	return &Matrix{
		Rows: rows,
		Cols: cols,
		Data: make([]float32, rows*cols),
	}
}

func (m *Matrix) At(r, c int) float32 {
	return m.Data[r*m.Cols+c]
}

func (m *Matrix) Set(r, c int, v float32) {
	m.Data[r*m.Cols+c] = v
}

// ==============================================================
// Basic Math Operations
// ==============================================================

// MatMul computes C = A * B where A:(M,K), B:(K,N), C:(M,N)
func MatMul(A, B, C *Matrix) {
	if A.Cols != B.Rows {
		panic(fmt.Sprintf("Dimension mismatch: A(%d,%d) B(%d,%d)",
			A.Rows, A.Cols, B.Rows, B.Cols))
	}
	if C.Rows != A.Rows || C.Cols != B.Cols {
		panic("C dimension mismatch")
	}

	M, K, N := A.Rows, A.Cols, B.Cols

	for i := 0; i < M; i++ {
		for j := 0; j < N; j++ {
			var sum float32
			for k := 0; k < K; k++ {
				sum += A.At(i, k) * B.At(k, j)
			}
			C.Set(i, j, sum)
		}
	}
}

// Softmax normalizes along the last dimension
func Softmax(m *Matrix) {
	for i := 0; i < m.Rows; i++ {
		maxVal := float32(math.Inf(-1))
		for j := 0; j < m.Cols; j++ {
			if m.At(i, j) > maxVal {
				maxVal = m.At(i, j)
			}
		}

		var sum float32
		for j := 0; j < m.Cols; j++ {
			val := float32(math.Exp(float64(m.At(i, j) - maxVal)))
			m.Set(i, j, val)
			sum += val
		}

		for j := 0; j < m.Cols; j++ {
			m.Set(i, j, m.At(i, j)/sum)
		}
	}
}

// GELU activation function
func GELU(x float32) float32 {
	return float32(0.5 * float64(x) *
		(1 + math.Erf(float64(x)/math.Sqrt2)))
}

// ==============================================================
// GMoE Expert Module (Go version)
// ==============================================================

type ExpertWeights struct {
	W1 *Matrix // (d_model, d_ff)
	W2 *Matrix // (d_ff, d_model)
}

func NewExpertWeights(dModel, dFF int) *ExpertWeights {
	return &ExpertWeights{
		W1: NewMatrix(dModel, dFF),
		W2: NewMatrix(dFF, dModel),
	}
}

func (e *ExpertWeights) Forward(input *Matrix) *Matrix {
	// input: (batch, d_model)
	// hidden = GELU(input * W1): (batch, d_ff)
	hidden := NewMatrix(input.Rows, e.W1.Cols)
	MatMul(input, e.W1, hidden)

	for i := 0; i < hidden.Rows; i++ {
		for j := 0; j < hidden.Cols; j++ {
			hidden.Set(i, j, GELU(hidden.At(i, j)))
		}
	}

	// output = hidden * W2: (batch, d_model)
	output := NewMatrix(hidden.Rows, e.W2.Cols)
	MatMul(hidden, e.W2, output)

	return output
}

// ==============================================================
// GRU Router (Go version)
// ==============================================================

type GRURouterWeights struct {
	InputProj *Matrix // (d_model, gru_hidden)

	// GRU cell parameters
	Wz *Matrix // (gru_hidden, gru_hidden) update gate
	Wr *Matrix // (gru_hidden, gru_hidden) reset gate
	Wh *Matrix // (gru_hidden, gru_hidden) candidate hidden

	Uz *Matrix // (gru_hidden, gru_hidden)
	Ur *Matrix // (gru_hidden, gru_hidden)
	Uh *Matrix // (gru_hidden, gru_hidden)

	Bz []float32 // (gru_hidden,)
	Br []float32 // (gru_hidden,)
	Bh []float32 // (gru_hidden,)

	RoutingHead *Matrix // (gru_hidden, num_experts)
}

func NewGRURouterWeights(dModel, gruHidden, numExperts int) *GRURouterWeights {
	return &GRURouterWeights{
		InputProj:   NewMatrix(dModel, gruHidden),
		Wz:          NewMatrix(gruHidden, gruHidden),
		Wr:          NewMatrix(gruHidden, gruHidden),
		Wh:          NewMatrix(gruHidden, gruHidden),
		Uz:          NewMatrix(gruHidden, gruHidden),
		Ur:          NewMatrix(gruHidden, gruHidden),
		Uh:          NewMatrix(gruHidden, gruHidden),
		Bz:          make([]float32, gruHidden),
		Br:          make([]float32, gruHidden),
		Bh:          make([]float32, gruHidden),
		RoutingHead: NewMatrix(gruHidden, numExperts),
	}
}

func (r *GRURouterWeights) Forward(
	x *Matrix,
	prevHidden *Matrix,
) (*Matrix, *Matrix) {
	// x: (batch, d_model)
	// prevHidden: (batch, gru_hidden)
	// Returns: routingWeights (batch, num_experts), newHidden (batch, gru_hidden)

	batch := x.Rows
	gruHidden := r.Wz.Rows
	numExperts := r.RoutingHead.Cols

	// 1. Input projection
	projected := NewMatrix(batch, gruHidden)
	MatMul(x, r.InputProj, projected)

	// 2. GRU forward pass
	newHidden := NewMatrix(batch, gruHidden)

	for i := 0; i < batch; i++ {
		for j := 0; j < gruHidden; j++ {
			var zSum, rSum, hSum float32

			// W_z * x_t + b_z
			for k := 0; k < projected.Cols; k++ {
				zSum += projected.At(i, k) * r.Wz.At(k, j)
				rSum += projected.At(i, k) * r.Wr.At(k, j)
				hSum += projected.At(i, k) * r.Wh.At(k, j)
			}
			zSum += r.Bz[j]
			rSum += r.Br[j]
			hSum += r.Bh[j]

			// U_z * h_{t-1}
			if prevHidden != nil {
				for k := 0; k < gruHidden; k++ {
					zSum += prevHidden.At(i, k) * r.Uz.At(k, j)
					rSum += prevHidden.At(i, k) * r.Ur.At(k, j)
				}
			}

			z := sigmoid(zSum)
			r := sigmoid(rSum)

			if prevHidden != nil {
				for k := 0; k < gruHidden; k++ {
					hSum += (r * prevHidden.At(i, k)) * r.Uh.At(k, j)
				}
			}

			hCandidate := float32(math.Tanh(float64(hSum)))
			newHidden.Set(i, j, (1-z)*hCandidate+z*0)
		}
	}

	// 3. Routing computation
	logits := NewMatrix(batch, numExperts)
	MatMul(newHidden, r.RoutingHead, logits)

	// 4. Softmax
	Softmax(logits)

	return logits, newHidden
}

func sigmoid(x float32) float32 {
	return 1.0 / (1.0 + float32(math.Exp(float64(-x))))
}

// ==============================================================
// GMoE Layer (Go version)
// ==============================================================

type GMoELayerWeights struct {
	LocalExpert *ExpertWeights
	Router      *GRURouterWeights
	LayerNorm   *LayerNormWeights
}

type LayerNormWeights struct {
	Gamma []float32 // (d_model,)
	Beta  []float32 // (d_model,)
}

func NewGMoELayerWeights(dModel, dFF, gruHidden, numExperts int) *GMoELayerWeights {
	return &GMoELayerWeights{
		LocalExpert: NewExpertWeights(dModel, dFF),
		Router:      NewGRURouterWeights(dModel, gruHidden, numExperts),
		LayerNorm: &LayerNormWeights{
			Gamma: make([]float32, dModel),
			Beta:  make([]float32, dModel),
		},
	}
}

// ==============================================================
// GMoE Inference Engine (Go version)
// ==============================================================

type GMoEInferenceEngine struct {
	Config ModelConfig

	// Global expert pool
	GlobalExperts []*ExpertWeights

	// Per-layer parameters
	Layers []*GMoELayerWeights

	// Embedding layer
	Embedding *Matrix // (vocab_size, d_model)

	// Output layer
	LMHead *Matrix // (d_model, vocab_size)

	// Final LayerNorm
	FinalNorm *LayerNormWeights
}

type ModelConfig struct {
	VocabSize          int
	DModel             int
	DFF                int
	NumLayers          int
	NumGlobalExperts   int
	NumExpertsPerToken int
	GRUHiddenSize      int
	MaxSeqLen          int
}

func NewGMoEInferenceEngine(config ModelConfig) *GMoEInferenceEngine {
	engine := &GMoEInferenceEngine{
		Config:        config,
		GlobalExperts: make([]*ExpertWeights, config.NumGlobalExperts),
		Layers:        make([]*GMoELayerWeights, config.NumLayers),
		Embedding:     NewMatrix(config.VocabSize, config.DModel),
		LMHead:        NewMatrix(config.DModel, config.VocabSize),
		FinalNorm: &LayerNormWeights{
			Gamma: make([]float32, config.DModel),
			Beta:  make([]float32, config.DModel),
		},
	}

	for i := 0; i < config.NumGlobalExperts; i++ {
		engine.GlobalExperts[i] = NewExpertWeights(config.DModel, config.DFF)
	}

	for i := 0; i < config.NumLayers; i++ {
		engine.Layers[i] = NewGMoELayerWeights(
			config.DModel, config.DFF,
			config.GRUHiddenSize, config.NumGlobalExperts,
		)
	}

	return engine
}

// Forward performs inference forward pass
func (e *GMoEInferenceEngine) Forward(inputIDs []int) []float32 {
	seqLen := len(inputIDs)
	batchSize := 1

	// 1. Embedding
	hidden := NewMatrix(seqLen, e.Config.DModel)
	for i := 0; i < seqLen; i++ {
		tokenID := inputIDs[i]
		if tokenID >= e.Config.VocabSize {
			tokenID = 0
		}
		for j := 0; j < e.Config.DModel; j++ {
			hidden.Set(i, j, e.Embedding.At(tokenID, j))
		}
	}

	// 2. Layer-by-layer forward pass
	var routingState *Matrix
	for layerIdx := 0; layerIdx < e.Config.NumLayers; layerIdx++ {
		layer := e.Layers[layerIdx]

		// LayerNorm
		hidden = e.applyLayerNorm(hidden, layer.LayerNorm)

		// Local expert
		localOutput := layer.LocalExpert.Forward(hidden)

		// Routing
		routingWeights, routingState := layer.Router.Forward(hidden, routingState)

		// Global experts (sparse activation)
		globalOutput := e.sparseGlobalForward(hidden, routingWeights)

		// Combine: output = input + local + global
		for i := 0; i < hidden.Rows; i++ {
			for j := 0; j < hidden.Cols; j++ {
				hidden.Set(i, j,
					hidden.At(i, j)+localOutput.At(i, j)+globalOutput.At(i, j))
			}
		}
	}

	// 3. Final LayerNorm
	hidden = e.applyLayerNorm(hidden, e.FinalNorm)

	// 4. Output projection
	logits := NewMatrix(batchSize, e.Config.VocabSize)
	MatMul(hidden, e.LMHead, logits)

	// Return last token's logits
	result := make([]float32, e.Config.VocabSize)
	for i := 0; i < e.Config.VocabSize; i++ {
		result[i] = logits.At(seqLen-1, i)
	}

	return result
}

func (e *GMoEInferenceEngine) applyLayerNorm(
	x *Matrix, norm *LayerNormWeights,
) *Matrix {
	output := NewMatrix(x.Rows, x.Cols)

	for i := 0; i < x.Rows; i++ {
		var mean, variance float32
		for j := 0; j < x.Cols; j++ {
			mean += x.At(i, j)
		}
		mean /= float32(x.Cols)

		for j := 0; j < x.Cols; j++ {
			diff := x.At(i, j) - mean
			variance += diff * diff
		}
		variance /= float32(x.Cols)
		std := float32(math.Sqrt(float64(variance + 1e-5)))

		for j := 0; j < x.Cols; j++ {
			normalized := (x.At(i, j) - mean) / std
			output.Set(i, j, normalized*norm.Gamma[j]+norm.Beta[j])
		}
	}

	return output
}

func (e *GMoEInferenceEngine) sparseGlobalForward(
	x *Matrix, routingWeights *Matrix,
) *Matrix {
	numExperts := e.Config.NumGlobalExperts
	K := e.Config.NumExpertsPerToken
	batchSize := x.Rows

	output := NewMatrix(batchSize, e.Config.DModel)

	for b := 0; b < batchSize; b++ {
		// Find Top-K experts
		type expertScore struct {
			idx   int
			score float32
		}
		scores := make([]expertScore, numExperts)
		for i := 0; i < numExperts; i++ {
			scores[i] = expertScore{idx: i, score: routingWeights.At(b, i)}
		}

		// Selection sort for Top-K
		for i := 0; i < K; i++ {
			maxIdx := i
			for j := i + 1; j < numExperts; j++ {
				if scores[j].score > scores[maxIdx].score {
					maxIdx = j
				}
			}
			scores[i], scores[maxIdx] = scores[maxIdx], scores[i]
		}

		// Normalize Top-K weights
		var weightSum float32
		for i := 0; i < K; i++ {
			weightSum += scores[i].score
		}

		// Sparse activation
		for i := 0; i < K; i++ {
			expertIdx := scores[i].idx
			weight := scores[i].score / weightSum

			tokenInput := NewMatrix(1, e.Config.DModel)
			for j := 0; j < e.Config.DModel; j++ {
				tokenInput.Set(0, j, x.At(b, j))
			}

			expertOutput := e.GlobalExperts[expertIdx].Forward(tokenInput)

			for j := 0; j < e.Config.DModel; j++ {
				output.Set(b, j,
					output.At(b, j)+weight*expertOutput.At(0, j))
			}
		}
	}

	return output
}

// ==============================================================
// Model Export
// ==============================================================

func (e *GMoEInferenceEngine) SaveWeights(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()

	binary.Write(f, binary.LittleEndian, int32(e.Config.VocabSize))
	binary.Write(f, binary.LittleEndian, int32(e.Config.DModel))
	binary.Write(f, binary.LittleEndian, int32(e.Config.DFF))
	binary.Write(f, binary.LittleEndian, int32(e.Config.NumLayers))
	binary.Write(f, binary.LittleEndian, int32(e.Config.NumGlobalExperts))
	binary.Write(f, binary.LittleEndian, int32(e.Config.NumExpertsPerToken))
	binary.Write(f, binary.LittleEndian, int32(e.Config.GRUHiddenSize))
	binary.Write(f, binary.LittleEndian, int32(e.Config.MaxSeqLen))

	fmt.Printf("Model weights saved to: %s\n", path)
	return nil
}

// ==============================================================
// Main: Inference Demo
// ==============================================================

func main() {
	fmt.Println("=" + strings.Repeat("=", 69))
	fmt.Println("  GMoE Inference Engine (Go Implementation)")
	fmt.Println("=" + strings.Repeat("=", 69))

	config := ModelConfig{
		VocabSize:         50257,
		DModel:            256,
		DFF:               1024,
		NumLayers:         4,
		NumGlobalExperts:  8,
		NumExpertsPerToken: 2,
		GRUHiddenSize:     64,
		MaxSeqLen:         512,
	}

	engine := NewGMoEInferenceEngine(config)

	fmt.Printf("\nModel Configuration:\n")
	fmt.Printf("  VocabSize:        %d\n", config.VocabSize)
	fmt.Printf("  DModel:           %d\n", config.DModel)
	fmt.Printf("  DFF:              %d\n", config.DFF)
	fmt.Printf("  NumLayers:        %d\n", config.NumLayers)
	fmt.Printf("  NumGlobalExperts: %d\n", config.NumGlobalExperts)
	fmt.Printf("  NumExpertsPerToken: %d\n", config.NumExpertsPerToken)

	expertParams := 2 * config.DModel * config.DFF
	globalExpertParams := config.NumGlobalExperts * expertParams
	localExpertParams := config.NumLayers * expertParams
	routerParams := config.DModel*config.GRUHiddenSize +
		6*config.GRUHiddenSize*config.GRUHiddenSize +
		3*config.GRUHiddenSize +
		config.GRUHiddenSize*config.NumGlobalExperts
	embedParams := config.VocabSize * config.DModel
	outputParams := config.DModel * config.VocabSize
	attentionParams := config.NumLayers * 4 * config.DModel * config.DModel

	totalParams := globalExpertParams + localExpertParams +
		routerParams + embedParams + outputParams + attentionParams

	fmt.Printf("\nParameter Statistics:\n")
	fmt.Printf("  Global Expert Params:      %d\n", globalExpertParams)
	fmt.Printf("  Local Expert Params:       %d\n", localExpertParams)
	fmt.Printf("  Router Params:             %d\n", routerParams)
	fmt.Printf("  Embedding Params:          %d\n", embedParams)
	fmt.Printf("  Total Parameters:          %d (%.2fM)\n",
		totalParams, float64(totalParams)/1e6)

	fmt.Printf("\nRunning inference...\n")
	dummyInput := []int{101, 202, 303, 404, 505}
	logits := engine.Forward(dummyInput)
	fmt.Printf("  Input sequence length: %d\n", len(dummyInput))
	fmt.Printf("  Output logits dimension: %d\n", len(logits))

	fmt.Printf("  Top-5 predictions:\n")
	type pred struct {
		id    int
		score float32
	}
	top5 := make([]pred, 5)
	for i := 0; i < 5; i++ {
		top5[i] = pred{id: -1, score: float32(math.Inf(-1))}
	}
	for i, score := range logits {
		if score > top5[4].score {
			top5[4] = pred{id: i, score: score}
			for j := 4; j > 0; j-- {
				if top5[j].score > top5[j-1].score {
					top5[j], top5[j-1] = top5[j-1], top5[j]
				}
			}
		}
	}
	for i, p := range top5 {
		fmt.Printf("    %d. token_id=%d, score=%.4f\n", i+1, p.id, p.score)
	}

	fmt.Printf("\n✅ GMoE inference engine running successfully!\n")
}

# ==============================================================
# GMoE INT8 Quantization for Edge Deployment
# ==============================================================

import torch
import torch.nn as nn
import numpy as np

class GMoEInt8Quantizer:
    """
    INT8 quantizer for GMoE models.
    Leverages the shared expert pool for efficient quantization.
    """
    
    def __init__(self, model: GMoEModel, calibration_data: torch.Tensor):
        self.model = model
        self.calibration_data = calibration_data
        self.scales = {}
        self.zero_points = {}
    
    def calibrate(self):
        """Calibrate quantization parameters using calibration dataset."""
        self.model.eval()
        
        with torch.no_grad():
            x = self.model.embedding(self.calibration_data)
            x = self.model.pos_encoding(x)
            
            routing_state = None
            for layer_idx, layer in enumerate(self.model.layers):
                x = layer.norm(x)
                
                local_out = layer.local_expert(x)
                self._update_scale(f"layer_{layer_idx}_local", local_out)
                
                weights, experts, routing_state = layer.router(x, routing_state)
                self._update_scale(f"layer_{layer_idx}_router", weights)
                
                for expert_idx in range(len(self.model.global_experts)):
                    mask = (experts == expert_idx)
                    if mask.any():
                        selected_x = x[mask]
                        expert_out = self.model.global_experts[expert_idx](selected_x)
                        self._update_scale(
                            f"global_expert_{expert_idx}", expert_out
                        )
                
                x = x + local_out
                x = x + self._sparse_forward_demo(x, weights, experts)
        
        print(f"Calibration complete. Collected {len(self.scales)} scale values")
    
    def _update_scale(self, name: str, tensor: torch.Tensor):
        if name not in self.scales:
            self.scales[name] = tensor.abs().max().item()
            self.zero_points[name] = 0
        else:
            self.scales[name] = max(
                self.scales[name], tensor.abs().max().item()
            )
    
    def _sparse_forward_demo(self, x, weights, experts):
        """Demo sparse forward pass (simplified)."""
        K = weights.size(-1)
        batch_size, seq_len, d_model = x.shape
        output = torch.zeros_like(x)
        
        for k in range(K):
            expert_indices = experts[:, :, k]
            expert_weights = weights[:, :, k]
            
            for expert_idx in range(len(self.model.global_experts)):
                mask = (expert_indices == expert_idx)
                if not mask.any():
                    continue
                
                selected_x = x[mask]
                selected_w = expert_weights[mask].unsqueeze(-1)
                expert_out = self.model.global_experts[expert_idx](selected_x)
                output[mask] += selected_w * expert_out
        
        return output
    
    def quantize_weights(self) -> dict:
        """Quantize model weights to INT8."""
        quantized = {}
        
        for name, param in self.model.named_parameters():
            if 'global_experts' in name or 'local_expert' in name:
                scale = param.abs().max(dim=-1, keepdim=True)[0] / 127.0
                quantized_data = torch.round(param / scale).to(torch.int8)
                quantized[f"{name}_quant"] = quantized_data
                quantized[f"{name}_scale"] = scale
            else:
                scale = param.abs().max().item() / 127.0
                quantized_data = torch.round(param / scale).to(torch.int8)
                quantized[f"{name}_quant"] = quantized_data
                quantized[f"{name}_scale"] = torch.tensor(scale)
        
        total_bytes = sum(
            p.numel() for p in quantized.values() if p.dtype == torch.int8
        )
        print(f"Quantized model size: {total_bytes / 1024 / 1024:.2f} MB")
        
        return quantized
    
    def save_for_deployment(self, path: str):
        """Save quantized model for edge deployment."""
        quantized = self.quantize_weights()
        torch.save(quantized, path)
        print(f"Quantized model saved to: {path}")


def edge_deployment_example():
    """
    Complete edge deployment workflow example.
    """
    print("=" * 70)
    print("GMoE Edge Deployment Workflow")
    print("=" * 70)
    
    model = GMoEModel(
        vocab_size=50257,
        d_model=256,
        d_ff=1024,
        num_layers=4,
        num_global_experts=8,
        num_experts_per_token=2,
        num_heads=4,
        dropout=0.0,
        max_seq_len=512,
        gru_hidden_size=64,
    )
    model.eval()
    
    total_params = sum(p.numel() for p in model.parameters())
    fp32_size = total_params * 4 / 1024 / 1024
    int8_size = total_params / 1024 / 1024
    
    d_model, d_ff = 256, 1024
    expert_params = 2 * d_model * d_ff
    active_params_per_layer = (2 + 1) * expert_params
    active_params = active_params_per_layer * 4
    
    print(f"\nDeployment Resource Requirements:")
    print(f"  Total parameters: {total_params:,}")
    print(f"  FP32 model size: {fp32_size:.2f} MB")
    print(f"  INT8 model size: {int8_size:.2f} MB")
    print(f"  Active params per token: ~{active_params:,}")
    print(f"  Recommended device: Phone/Tablet/Edge Gateway")
    
    flops_per_token = 2 * d_model * d_ff * (2 + 1) * 4
    estimated_latency = flops_per_token / (10 * 1e9) * 1000
    print(f"  Compute per token: ~{flops_per_token/1e6:.1f} MFLOPs")
    print(f"  Estimated latency: ~{estimated_latency:.2f} ms/token")
    print(f"  Generation speed: ~{1000/estimated_latency:.0f} tokens/s")
    
    print(f"\n✅ Edge deployment feasible!")


if __name__ == "__main__":
    edge_deployment_example()