import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import copy
import matplotlib.pyplot as plt


class SimpleMLP(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleMLP, self).__init__()
        self.fc1 = nn.Linear(32*32*3, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        x = x.view(-1, 32*32*3)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# -------------------------------
# Simple CNN Model
# -------------------------------
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 64 * 8 * 8)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


def train_cct_transfer_reset(bend_model, ground_model, trainloader, testloader,
                             cycles=50, batches_per_cycle=10,
                             lr=0.001, device='cpu',
                             entropy_threshold=0.27,
                             transfer_strength=0.9,   # how much bend -> ground on overfit 0.1
                             entropy_penalty_weight=0.5):
    """
    Cycles: number of meta-iterations (each cycle = batches_per_cycle batches)
    """
    bend_model.to(device)
    ground_model.to(device)
    ground_model.train()
    #ground_model.eval()
    
    optimizer = optim.Adam(bend_model.parameters(), lr=lr)
    criterion_ce = nn.CrossEntropyLoss()
    
    entropy_history = []
    test_acc_history = []
    transfer_cycles = []
    
    total_batches = 0
    for cycle in range(cycles):
        bend_model.train()
        cycle_entropy_sum = 0.0
        cycle_batches = 0
        
        # Run a small number of batches
        for batch_idx, (inputs, labels) in enumerate(trainloader):
            if batch_idx >= batches_per_cycle:
                break
            for _ in range(batch_idx*2+1):
                inputs, labels = inputs.to(device), labels.to(device)
                
                outputs = bend_model(inputs)
                loss_ce = criterion_ce(outputs, labels)
                
                # Entropy = divergence from ground
                with torch.no_grad():
                    ground_out = ground_model(inputs)
                    entropy = F.mse_loss(outputs, ground_out).item()
                
                loss = loss_ce + entropy_penalty_weight * entropy
                
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            
            cycle_entropy_sum += entropy
            cycle_batches += 1
            total_batches += 1
        
        avg_entropy = cycle_entropy_sum / cycle_batches
        entropy_history.append(avg_entropy)
        
        # ---- CCT Transfer & Reset ----
        if avg_entropy > entropy_threshold:
            # Transfer useful learning: ground ← (1 - α)*ground + α*bend
            with torch.no_grad():
                for g_param, b_param in zip(ground_model.parameters(), bend_model.parameters()):
                    g_param.data.mul_(1 - transfer_strength).add_(b_param.data, alpha=transfer_strength)
            # Reset bend to the improved ground
            bend_model.load_state_dict(copy.deepcopy(ground_model.state_dict()))
            # Also reset optimizer state (optional, but good to avoid momentum from bad trajectory)
            optimizer = optim.Adam(bend_model.parameters(), lr=lr)
            print(f"Cycle {cycle+1}: Entropy {avg_entropy:.3f} > {entropy_threshold} → TRANSFERRED (α={transfer_strength}) and RESET")
            transfer_cycles.append(cycle+1)
        else:
            # Optional: even when low entropy, slowly transfer (optional)
            # to keep ground improving gradually
            if avg_entropy < 0.15:  # relaxed state
                with torch.no_grad():
                    for g_param, b_param in zip(ground_model.parameters(), bend_model.parameters()):
                        g_param.data.mul_(0.99).add_(b_param.data, alpha=0.01)  # very slow consolidation
                print(f"Cycle {cycle+1}: Entropy {avg_entropy:.3f} < 0.15 → SLOW CONSOLIDATION")
        
        # Evaluate test accuracy after each cycle
        bend_model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in testloader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = bend_model(inputs)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        acc = 100 * correct / total
        test_acc_history.append(acc)
        print(f"         Test Acc: {acc:.2f}%")
    
    return bend_model, ground_model, entropy_history, test_acc_history, transfer_cycles
    
# -------------------------------
# CCT Relaxation & Consolidation Trainer
# -------------------------------
def train_cct_relax_consolidate(bend_model, ground_model, trainloader, testloader,
                                epochs=50, lr=0.001, device='cuda',
                                entropy_threshold=0.27,
                                relax_threshold=0.15,
                                ground_update_decay=0.99,   # how fast ground absorbs bend
                                reset_on_overfit=True,
                                entropy_penalty_weight=0.5):
    """
    - bend_model: fast learner (trainable)
    - ground_model: slow anchor (initially copy of bend)
    - ground_update_decay: high = slow consolidation
    """
    bend_model.to(device)
    ground_model.to(device)
    ground_model.eval()
    
    optimizer = optim.Adam(bend_model.parameters(), lr=lr)
    criterion_ce = nn.CrossEntropyLoss()
    
    entropy_history = []
    test_acc_history = []
    reset_epochs = []
    consolidate_epochs = []
    
    for epoch in range(epochs):
        bend_model.train()
        running_ce_loss = 0.0
        running_entropy = 0.0
        batch_count = 0
        
        for inputs, labels in trainloader:
            inputs, labels = inputs.to(device), labels.to(device)
            
            outputs = bend_model(inputs)
            loss_ce = criterion_ce(outputs, labels)
            
            # Entropy = divergence from ground model
            with torch.no_grad():
                ground_out = ground_model(inputs)
                entropy = F.mse_loss(outputs, ground_out).item()
            
            # Total loss = cross‑entropy + small entropy penalty
            # (keeps bend from drifting too fast)
            loss = loss_ce + entropy_penalty_weight * entropy
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            running_ce_loss += loss_ce.item()
            running_entropy += entropy
            batch_count += 1
        
        avg_entropy = running_entropy / batch_count
        entropy_history.append(avg_entropy)
        
        # ---- CCT Logic ----
        if avg_entropy > entropy_threshold:
            # Overfitting → reset bend to ground
            if reset_on_overfit:
                bend_model.load_state_dict(copy.deepcopy(ground_model.state_dict()))
                print(f"⚠️ Epoch {epoch+1}: Entropy {avg_entropy:.3f} > {entropy_threshold} → RESET bend to ground")
                reset_epochs.append(epoch+1)
            # Optionally reinitialize optimizer state?
        elif avg_entropy < relax_threshold:
            # Relaxation achieved → consolidate bend into ground
            # Slow update: ground = decay*ground + (1-decay)*bend
            with torch.no_grad():
                for g_param, b_param in zip(ground_model.parameters(), bend_model.parameters()):
                    g_param.data.mul_(ground_update_decay).add_(b_param.data, alpha=1 - ground_update_decay)
            print(f"✓ Epoch {epoch+1}: Entropy {avg_entropy:.3f} < {relax_threshold} → CONSOLIDATE ground ← bend")
            consolidate_epochs.append(epoch+1)
            # Optional: also reset bend to ground after consolidation?
            # If you reset, bend starts fresh from the improved ground.
            # Uncomment next line to reset bend after consolidation:
            # bend_model.load_state_dict(copy.deepcopy(ground_model.state_dict()))
        else:
            print(f"  Epoch {epoch+1}: Entropy {avg_entropy:.3f} (normal range)")
        
        # Evaluate test accuracy
        bend_model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for inputs, labels in testloader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = bend_model(inputs)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        acc = 100 * correct / total
        test_acc_history.append(acc)
        
        print(f"         Test Acc: {acc:.2f}%")
    
    return bend_model, ground_model, entropy_history, test_acc_history, reset_epochs, consolidate_epochs

# -------------------------------
# Main
# -------------------------------
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using {device}")
    
    # Data
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
    ])
    trainset = torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform)
    testset = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)
    testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=2)
    
    # Initialize models
    bend = SimpleMLP(num_classes=10)
    ground = SimpleMLP(num_classes=10)
    ground.load_state_dict(bend.state_dict())  # start identical
    
    # Train with CCT relaxation & consolidation
    bend_final, ground_final, entropy_hist, acc_hist, transfer_cycles = train_cct_transfer_reset(
        bend_model=bend,
        ground_model=ground,
        trainloader=trainloader,
        testloader=testloader,
        cycles=100,
        batches_per_cycle=2
    )
    
    # Plot
    plt.figure(figsize=(12,4))
    plt.subplot(1,2,1)
    plt.plot(entropy_hist, label='Entropy (bend vs ground)')
    plt.axhline(y=0.27, color='r', linestyle='--', label='Overfit threshold')
    plt.axhline(y=0.15, color='g', linestyle='--', label='Relax threshold')
    plt.xlabel('Epoch')
    plt.ylabel('Entropy')
    plt.title('CCT Entropy – Resets prevent overfitting')
    plt.legend()
    
    plt.subplot(1,2,2)
    plt.plot(acc_hist, label='Test Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title('No overfitting – steady improvement')
    plt.legend()
    plt.tight_layout()
    plt.savefig('cct_relax_consolidate.png')
    plt.show()
    
    print("\n=== Final ===")
    print(f"Final test accuracy: {acc_hist[-1]:.2f}%")
    
if __name__ == "__main__":
    main()