import random

class SimpleExpert:
    """One expert student in our class"""
    def __init__(self, name, specialty):
        self.name = name
        self.specialty = specialty
    
    def answer(self, question):
        # Each expert answers in their own style
        return f"{self.name} says: Here's my {self.specialty} take on '{question}'"

class SimpleMoE:
    """Our classroom of experts"""
    def __init__(self):
        # Create our expert students
        self.experts = [
            SimpleExpert("Rimi", "science"),
            SimpleExpert("Hasan", "literature"),
            SimpleExpert("Hamim", "math"),
            SimpleExpert("Richi", "history"),
            SimpleExpert("Emma", "art")
        ]
    
    def router(self, question):
        """Decide which experts should answer"""
        scores = []
        
        # Give each expert a score based on the question
        for expert in self.experts:
            if expert.specialty in question.lower():
                score = 0.9  # High score if specialty matches!
            else:
                score = random.uniform(0.1, 0.3)  # Low random score
            scores.append(score)
        
        return scores
    
    def answer_question(self, question):
        """Get answer from the best 2 experts"""
        # Get scores from router
        scores = self.router(question)
        
        # Find top 2 experts
        expert_scores = list(zip(self.experts, scores))
        expert_scores.sort(key=lambda x: x[1], reverse=True)
        top_experts = expert_scores[:2]
        
        print(f"\n📚 Question: '{question}'")
        print(f"🎯 Choosing experts...")
        
        # Show the selection process
        for expert, score in expert_scores:
            status = "✅ SELECTED" if (expert, score) in top_experts else "❌"
            print(f"   {expert.name} ({expert.specialty}): {score:.2f} {status}")
        
        # Get answers from top experts
        print(f"\n💭 Expert answers:")
        total_score = sum(score for _, score in top_experts)
        final_answer = ""
        
        for expert, score in top_experts:
            weight = score / total_score
            answer = expert.answer(question)
            print(f"   {answer} (weight: {weight:.2f})")
            final_answer += f"{answer} "
        
        return final_answer

# Let's try it!
moe = SimpleMoE()

# Ask different questions
questions = [
    "Tell me about science experiments",
    "What's the best math formula?",
    "Explain art techniques",
    "Write a literature story"
]

for q in questions:
    moe.answer_question(q)
    print("-" * 50)

import numpy as np

class NeuralExpert:
    """A tiny neural network expert"""
    def __init__(self, name, input_size=10, hidden_size=5):
        self.name = name
        # Random weights (like random knowledge!)
        self.weights1 = np.random.randn(input_size, hidden_size) * 0.1
        self.weights2 = np.random.randn(hidden_size, 1) * 0.1
    
    def forward(self, x):
        """Process input through the neural network"""
        # Simple neural network: input -> hidden -> output
        hidden = np.maximum(0, np.dot(x, self.weights1))  # ReLU activation
        output = np.dot(hidden, self.weights2)
        return output[0]

class NeuralMoE:
    """MoE with actual neural networks"""
    def __init__(self, num_experts=4, input_size=10):
        self.num_experts = num_experts
        self.input_size = input_size
        
        # Create expert networks
        self.experts = [
            NeuralExpert(f"Expert_{i}", input_size) 
            for i in range(num_experts)
        ]
        
        # Router network (decides which experts to use)
        self.router_weights = np.random.randn(input_size, num_experts) * 0.1
    
    def forward(self, x):
        """Process input through MoE"""
        # Step 1: Router scores each expert
        router_scores = np.dot(x, self.router_weights)
        
        # Step 2: Pick top 2 experts (using softmax for probabilities)
        exp_scores = np.exp(router_scores - np.max(router_scores))
        probabilities = exp_scores / np.sum(exp_scores)
        
        top_2_indices = np.argsort(probabilities)[-2:]
        
        # Step 3: Get outputs from top experts
        final_output = 0
        print(f"\n🧠 Expert Selection:")
        
        for i in range(self.num_experts):
            if i in top_2_indices:
                expert_output = self.experts[i].forward(x)
                weight = probabilities[i]
                final_output += weight * expert_output
                print(f"   Expert_{i}: {probabilities[i]:.3f} ✅ (output: {expert_output:.3f})")
            else:
                print(f"   Expert_{i}: {probabilities[i]:.3f} ❌")
        
        return final_output

# Demo time!
moe = NeuralMoE(num_experts=4, input_size=10)

# Create some sample inputs
print("🎮 Testing our Neural MoE:\n")

for i in range(3):
    # Random input (like a random question)
    test_input = np.random.randn(10)
    output = moe.forward(test_input)
    print(f"\n📊 Final output: {output:.3f}")
    print("-" * 40)

def demonstrate_sparsity():
    """Show how MoE saves computation"""
    
    # Traditional approach: Everyone works
    traditional_work = 8 * 100  # 8 experts × 100 units of work each
    print(f"🏃 Traditional NN: {traditional_work} units of work")
    
    # MoE approach: Only top 2 work
    moe_work = 2 * 100  # Only 2 experts × 100 units of work
    print(f"🚀 MoE: {moe_work} units of work")
    
    savings = (traditional_work - moe_work) / traditional_work * 100
    print(f"💰 Savings: {savings:.0f}% less work!")
    
    # Visual representation
    print("\n📊 Work Distribution:")
    print("Traditional: ████████ (all 8 experts)")
    print("MoE:         ██       (only 2 experts)")

demonstrate_sparsity()

class DIYMoE:
    """Build your own MoE step by step!"""
    
    def __init__(self):
        print("🏗️ Building your MoE...\n")
        
        # Step 1: Create experts
        print("Step 1: Creating experts 👥")
        self.experts = self.create_experts()
        
        # Step 2: Create router
        print("\nStep 2: Creating router 🎯")
        self.router = self.create_router()
        
        # Step 3: Ready to go!
        print("\n✅ Your MoE is ready!")
    
    def create_experts(self):
        """Make different expert types"""
        experts = {
            "math": lambda x: f"Math says: {x} × 2 = {x*2}",
            "science": lambda x: f"Science says: {x} atoms make a molecule",
            "art": lambda x: f"Art says: {x} is a beautiful number",
            "music": lambda x: f"Music says: {x} notes make a chord"
        }
        for name in experts:
            print(f"   ✓ Created {name} expert")
        return experts
    
    def create_router(self):
        """Decide which experts to use based on input type"""
        def router(question_type):
            if "calculate" in question_type:
                return ["math", "science"]
            elif "create" in question_type:
                return ["art", "music"]
            else:
                # Random selection
                import random
                return random.sample(list(self.experts.keys()), 2)
        
        print("   ✓ Router ready to select experts")
        return router
    
    def process(self, question_type, value):
        """Use the MoE to process input"""
        print(f"\n🎤 Processing: '{question_type}' with value {value}")
        
        # Router selects experts
        selected = self.router(question_type)
        print(f"🎯 Router selected: {selected}")
        
        # Get answers from selected experts
        print("💭 Expert responses:")
        for expert_name in selected:
            response = self.experts[expert_name](value)
            print(f"   - {response}")

# Try it out!
my_moe = DIYMoE()
my_moe.process("calculate something", 5)
my_moe.process("create something", 3)

def expert_popularity_contest(num_questions=100):
    """See which experts get picked most often"""
    expert_counts = {"math": 0, "science": 0, "art": 0, "history": 0}
    
    for _ in range(num_questions):
        # Randomly pick 2 experts
        import random
        chosen = random.sample(list(expert_counts.keys()), 2)
        for expert in chosen:
            expert_counts[expert] += 1
    
    print("🏆 Expert Popularity Contest Results:")
    for expert, count in sorted(expert_counts.items(), key=lambda x: x[1], reverse=True):
        bar = "█" * (count // 5)
        print(f"{expert:8} {bar} {count} times")

expert_popularity_contest()

def show_expert_combinations():
    """Show all possible expert teams"""
    experts = ["Math", "Science", "Art", "Music"]
    
    print("🤝 Possible Expert Teams (picking 2):")
    teams = []
    for i in range(len(experts)):
        for j in range(i+1, len(experts)):
            team = f"{experts[i]} + {experts[j]}"
            teams.append(team)
            print(f"   • {team}")
    
    print(f"\n📊 Total combinations: {len(teams)}")
    print(f"💡 With 256 experts (like big AI), there would be {256*255//2:,} combinations!")

show_expert_combinations()

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
from collections import Counter

# Step 1: Create a simple dataset
class SimpleTextDataset:
    """Our study materials for the expert students"""
    
    def __init__(self):
        # Each entry: (text, category)
        self.data = [
            # Tech examples
            ("Python is a programming language", "tech"),
            ("The new iPhone has amazing features", "tech"),
            ("Debugging code can be challenging", "tech"),
            ("Machine learning uses neural networks", "tech"),
            ("The laptop has 16GB of RAM", "tech"),
            ("JavaScript runs in web browsers", "tech"),
            ("The GPU processes graphics quickly", "tech"),
            ("Linux is an operating system", "tech"),
            
            # Sports examples
            ("The basketball game was exciting", "sports"),
            ("She scored three goals in soccer", "sports"),
            ("The Olympics happen every four years", "sports"),
            ("Tennis requires good hand-eye coordination", "sports"),
            ("The marathon runner finished first", "sports"),
            ("Baseball is America's pastime", "sports"),
            ("Swimming is great exercise", "sports"),
            ("The football team won the championship", "sports"),
            
            # Food examples
            ("Pizza is my favorite food", "food"),
            ("The recipe needs two cups of flour", "food"),
            ("Chocolate cake tastes delicious", "food"),
            ("Vegetables are healthy to eat", "food"),
            ("The restaurant serves Italian cuisine", "food"),
            ("Cooking pasta is easy and quick", "food"),
            ("Fresh fruits contain vitamins", "food"),
            ("The chef prepared a gourmet meal", "food"),
            
            # Science examples
            ("Water boils at 100 degrees Celsius", "science"),
            ("Photosynthesis produces oxygen", "science"),
            ("Gravity pulls objects toward Earth", "science"),
            ("DNA contains genetic information", "science"),
            ("Chemical reactions can produce heat", "science"),
            ("The microscope magnifies small objects", "science"),
            ("Planets orbit around the sun", "science"),
            ("Atoms are building blocks of matter", "science"),
        ]
        
        # Create vocabulary (word to number mapping)
        self.build_vocabulary()
        
    def build_vocabulary(self):
        """Create a dictionary of all words"""
        # Collect all words
        all_words = []
        for text, _ in self.data:
            words = text.lower().split()
            all_words.extend(words)
        
        # Create vocabulary
        unique_words = sorted(set(all_words))
        self.vocab = {word: idx for idx, word in enumerate(unique_words)}
        self.vocab['<PAD>'] = len(self.vocab)  # For padding
        self.vocab_size = len(self.vocab)
        
        # Category mapping
        self.categories = {"tech": 0, "sports": 1, "food": 2, "science": 3}
        self.num_categories = len(self.categories)
        
        print(f"📚 Vocabulary size: {self.vocab_size} words")
        print(f"📂 Categories: {list(self.categories.keys())}")
    
    def text_to_numbers(self, text, max_length=10):
        """Convert text to numbers for the neural network"""
        words = text.lower().split()
        numbers = [self.vocab.get(word, 0) for word in words]
        
        # Pad or truncate to fixed length
        if len(numbers) < max_length:
            numbers.extend([self.vocab['<PAD>']] * (max_length - len(numbers)))
        else:
            numbers = numbers[:max_length]
            
        return numbers
    
    def prepare_data(self):
        """Prepare data for training"""
        X = []  # Input texts as numbers
        y = []  # Categories
        
        for text, category in self.data:
            X.append(self.text_to_numbers(text))
            y.append(self.categories[category])
        
        return torch.tensor(X), torch.tensor(y)

# Create and explore our dataset
dataset = SimpleTextDataset()

# Show some examples
print("\n📝 Sample data:")
for i in range(4):
    text, category = dataset.data[i]
    print(f"   '{text}' → {category}")

# Prepare training data
X_train, y_train = dataset.prepare_data()
print(f"\n🔢 Data shape: {X_train.shape}")
print(f"🏷️ Labels shape: {y_train.shape}")

class ExpertNetwork(nn.Module):
    """One expert student with their own knowledge"""
    
    def __init__(self, input_size, hidden_size, output_size, expert_name):
        super().__init__()
        self.expert_name = expert_name
        self.network = nn.Sequential(
            nn.Embedding(input_size, 32),  # Word embeddings
            nn.Flatten(),
            nn.Linear(32 * 10, hidden_size),  # 10 is max sequence length
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.ReLU(),
            nn.Linear(hidden_size // 2, output_size)
        )
    
    def forward(self, x):
        return self.network(x)

class RouterNetwork(nn.Module):
    """The teacher who assigns experts to questions"""
    
    def __init__(self, input_size, num_experts):
        super().__init__()
        self.embedding = nn.Embedding(input_size, 32)
        self.router = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 10, 64),
            nn.ReLU(),
            nn.Linear(64, num_experts)
        )
    
    def forward(self, x):
        embedded = self.embedding(x)
        router_logits = self.router(embedded)
        return F.softmax(router_logits, dim=-1)

class MoEClassifier(nn.Module):
    """Our complete MoE model"""
    
    def __init__(self, vocab_size, num_categories, num_experts=4, experts_per_input=2):
        super().__init__()
        self.num_experts = num_experts
        self.experts_per_input = experts_per_input
        self.num_categories = num_categories
        
        # Create router
        self.router = RouterNetwork(vocab_size, num_experts)
        
        # Create experts
        self.experts = nn.ModuleList([
            ExpertNetwork(vocab_size, 64, num_categories, f"Expert_{i}")
            for i in range(num_experts)
        ])
        
        # Track expert usage
        self.expert_usage = torch.zeros(num_experts)
        
    def forward(self, x, return_expert_info=False):
        batch_size = x.shape[0]
        
        # Step 1: Router decides which experts to use
        router_probs = self.router(x)  # Shape: (batch_size, num_experts)
        
        # Step 2: Select top k experts
        top_k_probs, top_k_indices = torch.topk(
            router_probs, self.experts_per_input, dim=1
        )
        
        # Step 3: Get predictions from selected experts
        final_output = torch.zeros(batch_size, self.num_categories).to(x.device)
        
        # Process each sample in the batch
        for i in range(batch_size):
            expert_outputs = []
            
            # Get outputs from selected experts
            for j in range(self.experts_per_input):
                expert_idx = top_k_indices[i, j].item()
                expert_prob = top_k_probs[i, j]
                
                # Get expert's prediction
                expert_output = self.experts[expert_idx](x[i:i+1])
                
                # Weight by router probability
                final_output[i] += expert_prob * expert_output.squeeze()
                
                # Track usage
                self.expert_usage[expert_idx] += 1
        
        if return_expert_info:
            return final_output, router_probs, top_k_indices
        return final_output

# Create our model
model = MoEClassifier(
    vocab_size=dataset.vocab_size,
    num_categories=dataset.num_categories,
    num_experts=4,
    experts_per_input=2
)

print("🤖 Model created!")
print(f"   Total parameters: {sum(p.numel() for p in model.parameters()):,}")

def train_moe_model(model, X_train, y_train, epochs=50):
    """Train our MoE model"""
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    
    # Training history
    losses = []
    accuracies = []
    
    print("🏋️ Starting training...")
    
    for epoch in range(epochs):
        # Reset expert usage tracking
        model.expert_usage.zero_()
        
        # Forward pass
        outputs = model(X_train)
        loss = criterion(outputs, y_train)
        
        # Calculate accuracy
        _, predicted = torch.max(outputs, 1)
        accuracy = (predicted == y_train).float().mean().item()
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # Store history
        losses.append(loss.item())
        accuracies.append(accuracy)
        
        # Print progress every 10 epochs
        if (epoch + 1) % 10 == 0:
            print(f"   Epoch {epoch+1}/{epochs}: Loss={loss.item():.4f}, Accuracy={accuracy:.2%}")
            
            # Show expert usage
            usage_percent = model.expert_usage / model.expert_usage.sum() * 100
            print(f"   Expert usage: ", end="")
            for i, usage in enumerate(usage_percent):
                print(f"Expert_{i}: {usage:.1f}%  ", end="")
            print()
    
    return losses, accuracies

# Train the model
losses, accuracies = train_moe_model(model, X_train, y_train, epochs=100)

# Plot training progress
plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(losses)
plt.title('Training Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')

plt.subplot(1, 2, 2)
plt.plot(accuracies)
plt.title('Training Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.ylim([0, 1])

plt.tight_layout()
plt.show()

def test_model(model, dataset):
    """Test the model with new examples"""
    model.eval()
    
    test_examples = [
        "The smartphone has a fast processor",
        "The soccer player scored a goal",
        "I love eating chocolate ice cream",
        "Oxygen is essential for breathing",
        "Python code runs on computers",
        "Basketball players are very tall",
        "Pizza with cheese is delicious",
        "Molecules are made of atoms"
    ]
    
    print("🧪 Testing our MoE model:\n")
    
    category_names = {v: k for k, v in dataset.categories.items()}
    
    with torch.no_grad():
        for text in test_examples:
            # Convert text to numbers
            input_tensor = torch.tensor([dataset.text_to_numbers(text)])
            
            # Get prediction with expert info
            output, router_probs, expert_indices = model(input_tensor, return_expert_info=True)
            
            # Get predicted category
            _, predicted = torch.max(output, 1)
            predicted_category = category_names[predicted.item()]
            
            # Get confidence
            probs = F.softmax(output, dim=1)
            confidence = probs[0, predicted].item()
            
            print(f"📝 Text: '{text}'")
            print(f"🎯 Prediction: {predicted_category} (confidence: {confidence:.2%})")
            
            # Show which experts were used
            print(f"👥 Experts used: ", end="")
            for i in range(model.experts_per_input):
                expert_idx = expert_indices[0, i].item()
                expert_prob = router_probs[0, expert_idx].item()
                print(f"Expert_{expert_idx} ({expert_prob:.2%})  ", end="")
            print("\n" + "-"*50)

# Test our model
test_model(model, dataset)

class AdvancedMoEClassifier(nn.Module):
    """Enhanced MoE with load balancing and better routing"""
    
    def __init__(self, vocab_size, num_categories, num_experts=4, experts_per_input=2):
        super().__init__()
        self.num_experts = num_experts
        self.experts_per_input = experts_per_input
        self.num_categories = num_categories
        
        # Enhanced router with attention mechanism
        self.router = nn.Sequential(
            nn.Embedding(vocab_size, 64),
            nn.TransformerEncoderLayer(d_model=64, nhead=4, dim_feedforward=128, batch_first=True),
            nn.Flatten(),
            nn.Linear(64 * 10, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, num_experts)
        )
        
        # Specialized experts with different architectures
        self.experts = nn.ModuleList()
        expert_configs = [
            ("Tech Expert", 128),     # Larger network for tech
            ("Sports Expert", 64),    # Medium network
            ("Food Expert", 64),      # Medium network
            ("Science Expert", 128),  # Larger network for science
        ]
        
        for i, (name, hidden_size) in enumerate(expert_configs):
            self.experts.append(
                ExpertNetwork(vocab_size, hidden_size, num_categories, name)
            )
        
        # Load balancing loss weight
        self.load_balance_loss_weight = 0.01
        
        # Expert usage tracking
        self.register_buffer('expert_usage', torch.zeros(num_experts))
        
    def compute_load_balancing_loss(self, router_probs):
        """Encourage balanced usage of experts"""
        # Calculate the fraction of routing probability per expert
        expert_scores = router_probs.mean(dim=0)
        
        # Ideal uniform distribution
        uniform_distribution = 1.0 / self.num_experts
        
        # MSE loss from uniform distribution
        load_balance_loss = torch.mean((expert_scores - uniform_distribution) ** 2)
        
        return self.load_balance_loss_weight * load_balance_loss
    
    def forward(self, x, return_expert_info=False):
        batch_size = x.shape[0]
        
        # Get routing probabilities
        router_logits = self.router(x)
        router_probs = F.softmax(router_logits, dim=-1)
        
        # Add noise during training for exploration
        if self.training:
            noise = torch.randn_like(router_logits) * 0.1
            router_logits += noise
            router_probs = F.softmax(router_logits, dim=-1)
        
        # Select top k experts
        top_k_probs, top_k_indices = torch.topk(
            router_probs, self.experts_per_input, dim=1
        )
        
        # Normalize top k probabilities
        top_k_probs = top_k_probs / top_k_probs.sum(dim=1, keepdim=True)
        
        # Get predictions from experts
        final_output = torch.zeros(batch_size, self.num_categories).to(x.device)
        
        for i in range(batch_size):
            for j in range(self.experts_per_input):
                expert_idx = top_k_indices[i, j].item()
                expert_prob = top_k_probs[i, j]
                
                # Get expert's prediction
                expert_output = self.experts[expert_idx](x[i:i+1])
                final_output[i] += expert_prob * expert_output.squeeze()
                
                # Update usage statistics
                if not self.training:
                    self.expert_usage[expert_idx] += 1
        
        # Compute load balancing loss during training
        aux_loss = None
        if self.training:
            aux_loss = self.compute_load_balancing_loss(router_probs)
        
        if return_expert_info:
            return final_output, router_probs, top_k_indices, aux_loss
        return final_output, aux_loss

def train_advanced_moe(model, X_train, y_train, epochs=100):
    """Train with load balancing"""
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs)
    criterion = nn.CrossEntropyLoss()
    
    losses = []
    accuracies = []
    
    print("🚀 Training Advanced MoE...")
    
    for epoch in range(epochs):
        model.train()
        
        # Forward pass
        outputs, aux_loss = model(X_train)
        
        # Calculate main loss
        main_loss = criterion(outputs, y_train)
        
        # Total loss includes load balancing
        total_loss = main_loss
        if aux_loss is not None:
            total_loss = main_loss + aux_loss
        
        # Calculate accuracy
        _, predicted = torch.max(outputs, 1)
        accuracy = (predicted == y_train).float().mean().item()
        
        # Backward pass
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
        scheduler.step()
        
        losses.append(total_loss.item())
        accuracies.append(accuracy)
        
        if (epoch + 1) % 20 == 0:
            print(f"   Epoch {epoch+1}: Loss={total_loss.item():.4f}, Accuracy={accuracy:.2%}")
    
    return losses, accuracies

# Create and train advanced model
advanced_model = AdvancedMoEClassifier(
    vocab_size=dataset.vocab_size,
    num_categories=dataset.num_categories,
    num_experts=4,
    experts_per_input=2
)

adv_losses, adv_accuracies = train_advanced_moe(advanced_model, X_train, y_train)

# Compare with basic model
plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(losses, label='Basic MoE', alpha=0.7)
plt.plot(adv_losses, label='Advanced MoE', alpha=0.7)
plt.title('Training Loss Comparison')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(accuracies, label='Basic MoE', alpha=0.7)
plt.plot(adv_accuracies, label='Advanced MoE', alpha=0.7)
plt.title('Training Accuracy Comparison')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.legend()
plt.ylim([0, 1])

plt.tight_layout()
plt.show()

def analyze_expert_specialization(model, dataset, num_samples=50):
    """See which experts specialize in which categories"""
    model.eval()
    
    expert_category_counts = torch.zeros(model.num_experts, model.num_categories)
    
    with torch.no_grad():
        for text, category in dataset.data:
            # Convert to tensor
            input_tensor = torch.tensor([dataset.text_to_numbers(text)])
            
            # Get expert selections
            _, router_probs, expert_indices, _ = model(input_tensor, return_expert_info=True)
            
            # Get category index
            cat_idx = dataset.categories[category]
            
            # Count which experts were selected for this category
            for i in range(model.experts_per_input):
                expert_idx = expert_indices[0, i].item()
                expert_category_counts[expert_idx, cat_idx] += 1
    
    # Visualize specialization
    plt.figure(figsize=(10, 6))
    
    categories = list(dataset.categories.keys())
    expert_names = [f"Expert {i}" for i in range(model.num_experts)]
    
    # Create heatmap
    plt.imshow(expert_category_counts.numpy(), cmap='YlOrRd', aspect='auto')
    plt.colorbar(label='Selection Count')
    
    # Add labels
    plt.xticks(range(len(categories)), categories)
    plt.yticks(range(len(expert_names)), expert_names)
    
    # Add text annotations
    for i in range(model.num_experts):
        for j in range(model.num_categories):
            count = int(expert_category_counts[i, j].item())
            plt.text(j, i, str(count), ha='center', va='center')
    
    plt.title('Expert Specialization Heatmap')
    plt.xlabel('Categories')
    plt.ylabel('Experts')
    plt.tight_layout()
    plt.show()
    
    # Print analysis
    print("📊 Expert Specialization Analysis:")
    for i in range(model.num_experts):
        specialties = expert_category_counts[i]
        best_category = categories[torch.argmax(specialties).item()]
        total_selections = specialties.sum().item()
        
        print(f"\nExpert {i}:")
        print(f"  Most used for: {best_category}")
        print(f"  Total selections: {int(total_selections)}")
        print(f"  Distribution: ", end="")
        for j, cat in enumerate(categories):
            percentage = (specialties[j] / total_selections * 100) if total_selections > 0 else 0
            print(f"{cat}: {percentage:.1f}%  ", end="")

# Analyze our trained model
analyze_expert_specialization(advanced_model, dataset)

def create_custom_expert_model():
    """Challenge: Modify this to create your own unique MoE!"""
    
    # Ideas to try:
    # 1. Add more experts (8 or 16 instead of 4)
    # 2. Change how many experts are selected (top 3 instead of top 2)
    # 3. Add a "generalist" expert that always gets selected
    # 4. Create experts with different sizes (some big, some small)
    # 5. Add attention mechanisms between experts
    
    print("🏆 Your turn! Some ideas to try:")
    print("1. Add more training data categories (music, history, etc.)")
    print("2. Create specialized expert architectures")
    print("3. Implement expert dropout for robustness")
    print("4. Add a confidence threshold for routing")
    print("5. Create hierarchical experts (experts that call sub-experts)")
    
    # Your code here!
    pass

# Save your trained model
torch.save(advanced_model.state_dict(), 'my_first_moe_model.pth')
print("\n💾 Model saved! You've built your first real MoE AI!")

# Summary statistics
total_params = sum(p.numel() for p in advanced_model.parameters())
router_params = sum(p.numel() for p in advanced_model.router.parameters())
expert_params = sum(p.numel() for name, p in advanced_model.named_parameters() if 'expert' in name)

print(f"\n📊 Model Statistics:")
print(f"   Total parameters: {total_params:,}")
print(f"   Router parameters: {router_params:,} ({router_params/total_params*100:.1f}%)")
print(f"   Expert parameters: {expert_params:,} ({expert_params/total_params*100:.1f}%)")
print(f"   Parameters per expert: ~{expert_params//4:,}")