import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import copy

# ============================================
# Weight Hub + Branch Architecture (CIFAR-10 MLP)
# ============================================

class WeightHub(nn.Module):
    """
    Central hub that learns shared representations
    and broadcasts weights to branches via protocol
    """
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.shared_layer = nn.Linear(input_size, hidden_size)
        self.broadcast_weight = nn.Parameter(torch.ones(1) * 0.5)
        
    def forward(self, x):
        return torch.relu(self.shared_layer(x))
    
    def get_weights(self):
        """Protocol: Export weights for sharing"""
        return {
            'shared_layer': copy.deepcopy(self.shared_layer.state_dict()),
            'broadcast_weight': self.broadcast_weight.data.clone()
        }
    
    def receive_weights(self, weight_dict):
        """Protocol: Import and apply shared weights"""
        self.shared_layer.load_state_dict(weight_dict['shared_layer'])


class Branch(nn.Module):
    """Branch that can receive and apply shared weights"""
    def __init__(self, hidden_size, output_size, branch_id):
        super().__init__()
        self.branch_id = branch_id
        self.layer1 = nn.Linear(hidden_size, hidden_size)
        self.layer2 = nn.Linear(hidden_size, output_size)
        
        # Weight quality tracker (momentum for this branch)
        self.quality_score = 0.0
        
    def forward(self, x):
        x = torch.relu(self.layer1(x))
        return self.layer2(x)
    
    def apply_shared_weights(self, hub_weights, alpha=0.5):
        """
        Protocol: Apply shared weights from hub
        alpha = how much to blend (0=all old, 1=all new)
        """
        shared_state = hub_weights['shared_layer']
        
        with torch.no_grad():
            self.layer1.weight.data = alpha * self.layer1.weight.data + (1-alpha) * shared_state['weight']
            self.layer1.bias.data = alpha * self.layer1.bias.data + (1-alpha) * shared_state['bias']
    
    def update_quality(self, accuracy):
        """Track branch performance"""
        self.quality_score = 0.9 * self.quality_score + 0.1 * accuracy


class HubBranchNetwork(nn.Module):
    """
    Combined architecture with Hub + multiple Branches
    Weight sharing protocol between components
    """
    def __init__(self, input_size=3072, hidden_size=256, num_classes=10, num_branches=3):
        super().__init__()
        
        self.hub = WeightHub(input_size, hidden_size, num_classes)
        self.branches = nn.ModuleList([
            Branch(hidden_size, num_classes, i) for i in range(num_branches)
        ])
        
        # Protocol parameters
        self.share_threshold = 0.7  # Quality threshold to trigger sharing
        self.share_interval = 5     # Share every N epochs
        
    def forward(self, x, branch_id=0):
        x = x.view(x.size(0), -1)  # Flatten
        hub_features = self.hub(x)
        return self.branches[branch_id](hub_features)
    
    def forward_all(self, x):
        """Forward through all branches for ensemble"""
        x = x.view(x.size(0), -1)
        hub_features = self.hub(x)
        outputs = [branch(hub_features) for branch in self.branches]
        return torch.stack(outputs)  # [num_branches, batch, classes]
    
    def run_share_protocol(self, epoch):
        """
        Protocol: Evaluate branch quality, broadcast good weights
        """
        if epoch % self.share_interval != 0:
            return
        
        print(f"\n  📡 SHARE PROTOCOL EPOCH {epoch}")
        
        # Find best branch(es)
        best_branch = max(self.branches, key=lambda b: b.quality_score)
        best_quality = best_branch.quality_score
        
        print(f"     Best branch: #{best_branch.branch_id} (quality: {best_quality:.3f})")
        
        if best_quality > self.share_threshold:
            hub_weights = self.hub.get_weights()
            
            # Broadcast to all branches
            for branch in self.branches:
                if branch.branch_id != best_branch.branch_id:
                    branch.apply_shared_weights(hub_weights, alpha=0.7)
                    print(f"     → Branch #{branch.branch_id} received shared weights")
        else:
            print(f"     ⚠ Quality too low ({best_quality:.3f} < {self.share_threshold}), skipping share")


def train_hub_branch_model():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}\n")
    
    # CIFAR-10 transforms
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    
    train_dataset = datasets.CIFAR10(root='./data', train=True, transform=transform, download=True)
    test_dataset = datasets.CIFAR10(root='./data', train=False, transform=transform, download=True)
    
    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)
    test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=2)
    
    # === Hub-Branch Model ===
    model = HubBranchNetwork(
        input_size=3072,
        hidden_size=512,
        num_classes=10,
        num_branches=3
    ).to(device)
    
    # Optimizers: hub learns slowly, branches learn faster
    hub_optimizer = optim.Adam(model.hub.parameters(), lr=0.0003)
    branch_optimizer = optim.Adam(model.branches.parameters(), lr=0.001)
    
    criterion = nn.CrossEntropyLoss()
    
    print("="*60)
    print("Hub-Branch Network with Weight Sharing Protocol on CIFAR-10")
    print("="*60)
    
    epochs = 20
    best_acc = 0
    
    for epoch in range(epochs):
        # Training phase
        model.train()
        total_loss = 0
        correct_by_branch = [0, 0, 0]
        total_by_branch = [0, 0, 0]
        
        for inputs, targets in train_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            
            branch_optimizer.zero_grad()
            hub_optimizer.zero_grad()
            
            # Train all branches on current batch
            outputs_all = model.forward_all(inputs)  # [3, batch, 10]
            loss = 0
            for i, outputs in enumerate(outputs_all):
                branch_loss = criterion(outputs, targets)
                loss += branch_loss
                _, predicted = outputs.max(1)
                correct_by_branch[i] += predicted.eq(targets).sum().item()
                total_by_branch[i] += targets.size(0)
            
            loss.backward()
            branch_optimizer.step()
            hub_optimizer.step()
            
            total_loss += loss.item()
        
        # Evaluate each branch
        model.eval()
        branch_acc = []
        for i, branch in enumerate(model.branches):
            acc = 100. * correct_by_branch[i] / total_by_branch[i]
            branch.update_quality(acc / 100)
            branch_acc.append(acc)
        
        # Test accuracy
        test_correct = 0
        test_total = 0
        with torch.no_grad():
            for inputs, targets in test_loader:
                inputs, targets = inputs.to(device), targets.to(device)
                outputs_all = model.forward_all(inputs)
                # Ensemble: average all branches
                ensemble_out = outputs_all.mean(0)
                _, predicted = ensemble_out.max(1)
                test_correct += predicted.eq(targets).sum().item()
                test_total += targets.size(0)
        
        test_acc = 100. * test_correct / test_total
        best_acc = max(best_acc, test_acc)
        
        print(f"Epoch {epoch+1:2d}/{epochs} | "
              f"Hub Acc: {branch_acc[0]:.1f}% | "
              f"B1: {branch_acc[0]:.1f}% B2: {branch_acc[1]:.1f}% B3: {branch_acc[2]:.1f}% | "
              f"Test: {test_acc:.1f}%")
        
        # Run weight sharing protocol
        model.run_share_protocol(epoch + 1)
    
    print(f"\n🏆 Best Test Accuracy: {best_acc:.2f}%")
    return model


# === Run ===
if __name__ == "__main__":
    model = train_hub_branch_model()