colab_training_enhanced_v3.py

# ARC Prize 2025 - ENHANCED Training Script V3 for Google Colab
# Major improvements for exact match accuracy

# Install packages
import subprocess
import sys
subprocess.check_call([sys.executable, "-m", "pip", "install", "torch", "torchvision", "matplotlib", "numpy", "pandas", "tqdm", "onnx", "onnxruntime", "plotly", "scikit-learn", "albumentations", "-q"])
print("✓ Packages installed")

# Imports
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, random_split
import numpy as np
import json
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from tqdm import tqdm
import os
import shutil
from typing import Dict, List, Tuple, Optional
import plotly.graph_objects as go
import plotly.express as px
import pandas as pd
from datetime import datetime
import time
import gc
import random
from torchvision.transforms import v2 as transforms_v2
import albumentations as A
from albumentations.pytorch import ToTensorV2

# Check GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'\nUsing device: {device}')
if torch.cuda.is_available():
    print(f'GPU: {torch.cuda.get_device_name(0)}')
    print(f'Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB')
    print(f'\n🚀 A100 80GB DETECTED! Using optimized settings for maximum performance!')

# Clone repository
print("\n📥 Cloning ARC Prize 2025 repository...")
if os.path.exists('Arc2025'):
    shutil.rmtree('Arc2025')
os.system("git clone https://github.com/AutomataControls/Arc2025.git")
os.system("cd Arc2025 && git pull")
print("✓ Repository cloned and updated")

# Setup paths and imports
sys.path.append('/content/Arc2025')
sys.path.append('/content')

# Import with fallbacks
try:
    from Arc2025.models.arc_models_enhanced import create_enhanced_models
except:
    try:
        from Arc2025.arc_models_enhanced import create_enhanced_models
    except:
        sys.path.append('/content/Arc2025/models')
        from arc_models_enhanced import create_enhanced_models

try:
    from Arc2025.colab_monitor_integration import setup_colab_monitor
except:
    from colab_monitor_integration import setup_colab_monitor

# Enable mixed precision training
from torch.amp import GradScaler, autocast

# Verify models
print("\n🔍 Verifying enhanced models...")
try:
    test_models = create_enhanced_models()
    print(f"✓ Successfully loaded {len(test_models)} models: {list(test_models.keys())}")
    del test_models
except Exception as e:
    print(f"❌ Error loading models: {e}")
    print("Downloading models directly...")
    os.system("wget -q https://raw.githubusercontent.com/AutomataControls/Arc2025/main/models/arc_models_enhanced.py -O /content/arc_models_enhanced.py")
    from arc_models_enhanced import create_enhanced_models

# IMPROVED HYPERPARAMETERS FOR V3
BATCH_SIZE = 32  # Much smaller for more frequent updates
GRADIENT_ACCUMULATION_STEPS = 1  # No accumulation - update every batch
LEARNING_RATE = 0.001  # Increased 20x - we're stuck and need to force learning
NUM_EPOCHS = 200  # Increased from 100
MAX_GRID_SIZE = 30
NUM_COLORS = 10
DEVICE = device

# IMPROVED LOSS WEIGHTS - FIXED FOR EXACT MATCH
RECONSTRUCTION_WEIGHT = 1.0
PATTERN_WEIGHT = 0.0  # Removed - not helping
CONSISTENCY_WEIGHT = 0.01  # Reduced even further
EDGE_WEIGHT = 0.5  # Reduced: Too much edge focus hurts active regions
COLOR_BALANCE_WEIGHT = 0.3  # Balanced: Must get colors right but not dominate
STRUCTURE_WEIGHT = 0.4  # Moderate: Important but not overwhelming
TRANSFORMATION_PENALTY = 1.5  # INCREASED: Much stronger penalty for copying input

print("\n⚙️ V3 Configuration:")
print(f"  Batch size: {BATCH_SIZE} (effective: {BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS})")
print(f"  Learning rate: {LEARNING_RATE}")
print(f"  Epochs: {NUM_EPOCHS}")
print(f"  New features: Edge-aware loss, Color balance loss, Focal loss for hard pixels")

# Data setup
print("\n📊 Setting up dataset...")
if os.path.exists('/content/Arc2025/data'):
    DATA_DIR = '/content/Arc2025/data'
elif os.path.exists('/content/data'):
    DATA_DIR = '/content/data'
else:
    print("Downloading ARC data...")
    os.makedirs('/content/data', exist_ok=True)
    os.system("wget -q https://github.com/fchollet/ARC-AGI/raw/master/data/training/arc-agi_training_challenges.json -O /content/data/arc-agi_training_challenges.json")
    os.system("wget -q https://github.com/fchollet/ARC-AGI/raw/master/data/training/arc-agi_training_solutions.json -O /content/data/arc-agi_training_solutions.json")
    DATA_DIR = '/content/data'
print(f"✓ Dataset location: {DATA_DIR}")


class ImprovedReconstructionLoss(nn.Module):
    """Enhanced loss with edge awareness and focal loss for hard pixels"""
    
    def __init__(self):
        super().__init__()
        self.ce_loss = nn.CrossEntropyLoss(reduction='none')
        
    def forward(self, pred: torch.Tensor, target: torch.Tensor, input_grid: torch.Tensor = None) -> Dict[str, torch.Tensor]:
        B, C, H, W = pred.shape
        
        # Convert target to indices
        target_indices = target.argmax(dim=1)  # (B, H, W)
        
        # 1. Focal loss for hard pixels
        pred_flat = pred.permute(0, 2, 3, 1).reshape(-1, C)
        target_flat = target_indices.reshape(-1)
        
        ce_loss = self.ce_loss(pred_flat, target_flat)
        
        # Focal loss: focus on hard examples with moderate gamma
        pt = torch.exp(-ce_loss)  # probability of correct class
        focal_loss = (1 - pt) ** 1.5 * ce_loss  # gamma=1.5 for gentler focus
        focal_loss = focal_loss.reshape(B, H, W)
        
        # 2. Edge-aware loss
        # Detect edges in target
        target_edges = self._detect_edges(target_indices)
        
        # Weight edge pixels more - but not too much
        edge_weight = 1.0 + target_edges * 1.0  # 2x weight on edges (reduced from 3x)
        weighted_loss = focal_loss * edge_weight
        
        reconstruction_loss = weighted_loss.mean()
        
        # 3. Color balance loss
        pred_colors = pred.argmax(dim=1)
        color_balance_loss = self._color_balance_loss(pred_colors, target_indices)
        
        # 4. Consistency loss
        probs = F.softmax(pred, dim=1)
        entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=1)
        consistency_loss = entropy.mean()
        
        # 5. Structure preservation loss
        structure_loss = self._structure_loss(pred_colors, target_indices)
        
        # 6. Transformation penalty - penalize if prediction is too similar to input
        if input_grid is not None:
            input_indices = input_grid.argmax(dim=1)  # Get actual input colors
            similarity_to_input = (pred_colors == input_indices).float().mean()
            transformation_penalty = similarity_to_input  # High similarity = high penalty
        else:
            transformation_penalty = 0.0
        
        total_loss = (
            RECONSTRUCTION_WEIGHT * reconstruction_loss +
            COLOR_BALANCE_WEIGHT * color_balance_loss +
            CONSISTENCY_WEIGHT * consistency_loss +
            STRUCTURE_WEIGHT * structure_loss +
            TRANSFORMATION_PENALTY * transformation_penalty  # Encourage transformation
        )
        
        return {
            'reconstruction': reconstruction_loss,
            'color_balance': color_balance_loss,
            'consistency': consistency_loss,
            'structure': structure_loss,
            'transformation': transformation_penalty,
            'total': total_loss
        }
    
    def _detect_edges(self, grid: torch.Tensor) -> torch.Tensor:
        """Detect edges in grid"""
        # Sobel-like edge detection
        dx = torch.abs(grid[:, 1:, :] - grid[:, :-1, :])
        dy = torch.abs(grid[:, :, 1:] - grid[:, :, :-1])
        
        # Pad to original size
        dx = F.pad(dx, (0, 0, 0, 1), value=0)
        dy = F.pad(dy, (0, 1, 0, 0), value=0)
        
        edges = ((dx + dy) > 0).float()
        return edges
    
    def _color_balance_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Encourage correct color distribution"""
        B = pred.shape[0]
        loss = 0
        
        for b in range(B):
            # Get color histograms
            pred_hist = torch.histc(pred[b].float(), bins=10, min=0, max=9)
            target_hist = torch.histc(target[b].float(), bins=10, min=0, max=9)
            
            # Normalize
            pred_hist = pred_hist / (pred_hist.sum() + 1e-8)
            target_hist = target_hist / (target_hist.sum() + 1e-8)
            
            # KL divergence
            kl_div = F.kl_div(torch.log(pred_hist + 1e-8), target_hist, reduction='sum')
            loss += kl_div
        
        return loss / B
    
    def _structure_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Preserve structural patterns"""
        # Simple connected component preservation
        # Check if objects maintain their connectivity
        
        # For each color, check if regions are preserved
        B = pred.shape[0]
        loss = 0
        
        for b in range(B):
            for color in range(1, 10):  # Skip background
                pred_mask = (pred[b] == color).float()
                target_mask = (target[b] == color).float()
                
                # IoU for this color
                intersection = (pred_mask * target_mask).sum()
                union = pred_mask.sum() + target_mask.sum() - intersection
                
                if union > 0:
                    iou = intersection / union
                    loss += 1.0 - iou
        
        return loss / (B * 9)


class CurriculumARCDataset(Dataset):
    """Dataset with curriculum learning - start with easier samples"""
    
    def __init__(self, data_dir: str, split: str = 'train', curriculum_stage: int = 0):
        self.data_dir = data_dir
        self.split = split
        self.curriculum_stage = curriculum_stage  # 0: easy, 1: medium, 2: hard
        self.samples = []
        self.pattern_labels = []
        self._load_data()
        
    def _load_data(self):
        """Load data with difficulty assessment"""
        print(f"Loading {self.split} data (curriculum stage {self.curriculum_stage})...")
        
        with open(f'{self.data_dir}/arc-agi_training_challenges.json', 'r') as f:
            challenges = json.load(f)
        
        with open(f'{self.data_dir}/arc-agi_training_solutions.json', 'r') as f:
            solutions = json.load(f)
        
        # Process and categorize by difficulty
        easy_samples = []
        medium_samples = []
        hard_samples = []
        
        for task_id, task_data in challenges.items():
            train_examples = task_data['train']
            test_examples = task_data['test']
            task_solutions = solutions[task_id]
            
            # Process examples
            for example in train_examples:
                input_grid = np.array(example['input'])
                output_grid = np.array(example['output'])
                
                difficulty = self._assess_difficulty(input_grid, output_grid)
                sample = {
                    'task_id': task_id,
                    'input': input_grid,
                    'output': output_grid,
                    'type': 'train_example',
                    'difficulty': difficulty
                }
                
                if difficulty == 0:
                    easy_samples.append(sample)
                elif difficulty == 1:
                    medium_samples.append(sample)
                else:
                    hard_samples.append(sample)
        
        # Select samples based on curriculum stage
        if self.curriculum_stage == 0:
            self.samples = easy_samples
        elif self.curriculum_stage == 1:
            self.samples = easy_samples + medium_samples
        else:
            self.samples = easy_samples + medium_samples + hard_samples
        
        # Apply augmentation
        self._augment_data_v3()
        print(f"Loaded {len(self.samples)} samples for stage {self.curriculum_stage}")
    
    def _assess_difficulty(self, input_grid: np.ndarray, output_grid: np.ndarray) -> int:
        """Assess task difficulty (0: easy, 1: medium, 2: hard)"""
        # Easy: same size, simple transformations
        if input_grid.shape == output_grid.shape:
            # Check for simple patterns
            if np.array_equal(input_grid, np.rot90(output_grid, k=1)):
                return 0
            if np.array_equal(input_grid, np.fliplr(output_grid)):
                return 0
            if np.array_equal(input_grid, np.flipud(output_grid)):
                return 0
            
            # Check if only color mapping
            if input_grid.shape == output_grid.shape:
                unique_in = len(np.unique(input_grid))
                unique_out = len(np.unique(output_grid))
                if unique_in <= 3 and unique_out <= 3:
                    return 0
        
        # Medium: size changes or moderate complexity
        if input_grid.size < 100 or output_grid.size < 100:
            return 1
        
        # Hard: everything else
        return 2
    
    def _augment_data_v3(self):
        """Improved augmentation focusing on exact patterns"""
        augmented = []
        
        for sample in self.samples:
            input_grid = sample['input']
            output_grid = sample['output']
            
            # Only augment easy samples to create more training data
            if sample['difficulty'] == 0:
                # Rotation augmentation (all 4 rotations)
                for k in range(1, 4):
                    aug_input = np.rot90(input_grid, k)
                    aug_output = np.rot90(output_grid, k)
                    
                    augmented.append({
                        'task_id': sample['task_id'] + f'_rot{k*90}',
                        'input': aug_input,
                        'output': aug_output,
                        'type': 'augmented',
                        'difficulty': 0
                    })
                
                # Reflection augmentation
                for axis in [0, 1]:
                    aug_input = np.flip(input_grid, axis=axis)
                    aug_output = np.flip(output_grid, axis=axis)
                    
                    augmented.append({
                        'task_id': sample['task_id'] + f'_flip{axis}',
                        'input': aug_input,
                        'output': aug_output,
                        'type': 'augmented',
                        'difficulty': 0
                    })
        
        self.samples.extend(augmented)
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        
        input_grid = self._to_one_hot(sample['input'])
        output_grid = self._to_one_hot(sample['output'])
        
        input_grid = self._pad_grid(input_grid)
        output_grid = self._pad_grid(output_grid)
        
        return {
            'input': torch.FloatTensor(input_grid),
            'output': torch.FloatTensor(output_grid),
            'difficulty': sample['difficulty'],
            'task_id': sample['task_id']
        }
    
    def _to_one_hot(self, grid: np.ndarray) -> np.ndarray:
        h, w = grid.shape
        one_hot = np.zeros((NUM_COLORS, h, w))
        
        for i in range(h):
            for j in range(w):
                color = int(grid[i, j])
                if 0 <= color < NUM_COLORS:
                    one_hot[color, i, j] = 1
        
        return one_hot
    
    def _pad_grid(self, grid: np.ndarray) -> np.ndarray:
        c, h, w = grid.shape
        if h >= MAX_GRID_SIZE or w >= MAX_GRID_SIZE:
            grid = grid[:, :MAX_GRID_SIZE, :MAX_GRID_SIZE]
            h = min(h, MAX_GRID_SIZE)
            w = min(w, MAX_GRID_SIZE)
        
        padded = np.zeros((c, MAX_GRID_SIZE, MAX_GRID_SIZE))
        padded[:, :h, :w] = grid[:, :h, :w]
        
        return padded


class ModelWithDropoutSchedule(nn.Module):
    """Wrapper with scheduled dropout for better convergence"""
    
    def __init__(self, base_model: nn.Module, initial_dropout: float = 0.2):
        super().__init__()
        self.base_model = base_model
        self.current_dropout = initial_dropout
        # CRITICAL FIX: NEVER apply dropout to output predictions!
        
    def set_dropout(self, rate: float):
        self.current_dropout = rate
        # Dropout should be inside the model layers, not on final output
        
    def forward(self, *args, **kwargs):
        outputs = self.base_model(*args, **kwargs)
        # REMOVED DROPOUT ON OUTPUT - this was killing exact match!
        return outputs


def train_enhanced_models_v3():
    """V3 training with curriculum learning and improved losses"""
    print("\n🚀 Starting Enhanced Model Training V3")
    print("="*60)
    print("New features: Curriculum learning, Edge-aware loss, Focal loss")
    print("="*60)
    
    monitor = setup_colab_monitor()
    
    # Start with curriculum stage 0 (easy samples)
    current_stage = 0
    dataset = CurriculumARCDataset(DATA_DIR, split='train', curriculum_stage=current_stage)
    
    train_size = int(0.9 * len(dataset))
    val_size = len(dataset) - train_size
    
    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
    
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, 
                            num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, 
                          num_workers=4, pin_memory=True)
    
    print(f"\n📊 Dataset Statistics (Stage {current_stage}):")
    print(f"  Total samples: {len(dataset)}")
    print(f"  Train samples: {len(train_dataset)}")
    print(f"  Val samples: {len(val_dataset)}")
    
    os.makedirs('/content/arc_models', exist_ok=True)
    os.makedirs('/content/results', exist_ok=True)
    
    # Create models with dropout schedule
    base_models = create_enhanced_models()
    models = {name: ModelWithDropoutSchedule(model, initial_dropout=0.2) 
              for name, model in base_models.items()}
    
    # Use improved loss
    loss_fn = ImprovedReconstructionLoss()
    
    training_history = {}
    
    # Train each model
    for model_name, model in models.items():
        print(f"\n{'='*60}")
        print(f"🧠 Training {model_name.upper()} with V3 improvements")
        print(f"{'='*60}")
        
        model = model.to(DEVICE)
        
        # Use different optimizers for different models
        if model_name in ['minerva', 'atlas']:
            optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, 
                                  weight_decay=0.01, betas=(0.9, 0.999))
        else:
            # Use SGD for others - sometimes works better
            optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE * 2,  # Reduced multiplier
                                momentum=0.9, weight_decay=0.01, nesterov=True)
        
        # Use ReduceLROnPlateau for adaptive learning rate
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.2, patience=5)
        
        scaler = GradScaler('cuda')
        
        history = {
            'train_loss': [],
            'val_loss': [],
            'val_accuracy': [],
            'val_pixel_acc': [],
            'val_active_acc': [],
            'val_structure': []
        }
        
        best_val_loss = float('inf')
        best_val_acc = 0
        patience = 30
        patience_counter = 0
        # Set a schedule to switch stages
        stage_switch_epoch_1 = 30  # Switch to medium at epoch 30
        stage_switch_epoch_2 = 80  # Switch to hard at epoch 80
        
        # Training loop
        for epoch in range(NUM_EPOCHS):
            # Curriculum learning: increase difficulty
            if epoch == stage_switch_epoch_1 and current_stage == 0:
                current_stage = 1
                print(f"\n🎯 Switching to curriculum stage {current_stage}")
            elif epoch == stage_switch_epoch_2 and current_stage == 1:
                current_stage = 2
                print(f"\n🎯 Switching to curriculum stage {current_stage}")
                
                # Reload dataset with new difficulty
                dataset = CurriculumARCDataset(DATA_DIR, split='train', 
                                             curriculum_stage=current_stage)
                train_size = int(0.9 * len(dataset))
                val_size = len(dataset) - train_size
                
                train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
                
                train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, 
                                        shuffle=True, num_workers=4, pin_memory=True)
                val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, 
                                      shuffle=False, num_workers=4, pin_memory=True)
                
                stage_switch_epoch = epoch + 50  # Next switch
            
            # Adjust dropout schedule
            dropout_rate = 0.2 * (1.0 - epoch / NUM_EPOCHS)  # Decay dropout
            model.set_dropout(dropout_rate)
            
            # Training phase
            model.train()
            train_loss = 0.0
            train_steps = 0
            
            pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{NUM_EPOCHS}")
            for batch_idx, batch in enumerate(pbar):
                input_grids = batch['input'].to(DEVICE)
                output_grids = batch['output'].to(DEVICE)
                
                # Accumulate gradients
                with autocast('cuda'):
                    if model_name == 'chronos':
                        outputs = model.base_model([input_grids])
                    else:
                        outputs = model(input_grids, output_grids, mode='train')
                    
                    pred_output = outputs['predicted_output']
                    
                    # CRITICAL: Check output shape and add debugging
                    if pred_output.dim() != 4 or pred_output.shape[1] != 10:
                        print(f"ERROR: Invalid output shape {pred_output.shape} for {model_name}")
                        pred_output = torch.zeros(input_grids.shape).to(DEVICE)
                    
                    # Debug: check if outputs are reasonable
                    if epoch == 0 and train_steps == 0:
                        print(f"\nDEBUG {model_name}: output range [{pred_output.min():.3f}, {pred_output.max():.3f}]")
                        print(f"Output shape: {pred_output.shape}, Input shape: {input_grids.shape}")
                    
                    losses = loss_fn(pred_output, output_grids, input_grids)
                    loss = losses['total'] / GRADIENT_ACCUMULATION_STEPS
                
                scaler.scale(loss).backward()
                
                if (train_steps + 1) % GRADIENT_ACCUMULATION_STEPS == 0:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)  # Increased to allow stronger updates
                    scaler.step(optimizer)
                    scaler.update()
                    optimizer.zero_grad()
                
                train_loss += losses['total'].item()
                train_steps += 1
                
                pbar.set_postfix({
                    'loss': f'{losses["total"].item():.4f}',
                    'recon': f'{losses["reconstruction"].item():.4f}',
                    'struct': f'{losses["structure"].item():.4f}'
                })
            
            # Step scheduler with validation accuracy for ReduceLROnPlateau
            # scheduler.step()  # Removed - will be done after validation
            
            # Validation phase
            model.eval()
            val_loss = 0.0
            val_correct = 0
            val_total = 0
            val_pixel_acc_sum = 0
            val_active_acc_sum = 0
            val_structure_sum = 0
            val_batches_count = 0
            
            with torch.no_grad():
                for batch in val_loader:
                    input_grids = batch['input'].to(DEVICE)
                    output_grids = batch['output'].to(DEVICE)
                    
                    with autocast('cuda'):
                        if model_name == 'chronos':
                            outputs = model.base_model([input_grids])
                        else:
                            outputs = model(input_grids)
                        
                        pred_output = outputs['predicted_output']
                        
                        # CRITICAL FIX: Ensure proper shape and no activation
                        # The models output raw logits, loss function expects them
                        if pred_output.dim() != 4 or pred_output.shape[1] != 10:
                            print(f"WARNING: Invalid output shape {pred_output.shape}")
                            pred_output = torch.zeros(input_grids.shape).to(DEVICE)
                        
                        losses = loss_fn(pred_output, output_grids, input_grids)
                    
                    val_loss += losses['total'].item()
                    val_structure_sum += losses['structure'].item()
                    
                    # Calculate metrics
                    pred_colors = pred_output.argmax(dim=1)
                    target_colors = output_grids.argmax(dim=1)
                    
                    # Exact match
                    matches = (pred_colors == target_colors).all(dim=[1,2])
                    val_correct += matches.sum().item()
                    val_total += input_grids.size(0)
                    
                    # DIAGNOSTIC: How close are we to exact matches?
                    if epoch % 10 == 0 and val_batches_count == 0:
                        per_sample_accuracy = (pred_colors == target_colors).float().mean(dim=[1,2])
                        best_accuracy = per_sample_accuracy.max().item()
                        avg_accuracy = per_sample_accuracy.mean().item()
                        print(f"  Best sample accuracy: {best_accuracy*100:.1f}%")
                        print(f"  Average sample accuracy: {avg_accuracy*100:.1f}%")
                        print(f"  Samples >90% accurate: {(per_sample_accuracy > 0.9).sum().item()}/{len(per_sample_accuracy)}")
                    
                    # Pixel accuracy
                    pixel_correct = (pred_colors == target_colors).float()
                    val_pixel_acc_sum += pixel_correct.mean().item() * 100
                    
                    # DIAGNOSTIC: Check what colors are being predicted
                    if epoch % 10 == 0 and val_batches_count == 0:
                        unique_pred = torch.unique(pred_colors)
                        unique_target = torch.unique(target_colors)
                        print(f"\nDIAGNOSTIC Epoch {epoch+1}:")
                        print(f"  Predicted colors: {unique_pred.tolist()}")
                        print(f"  Target colors: {unique_target.tolist()}")
                        print(f"  Most common predicted: {pred_colors.flatten().mode().values.item()}")
                        
                        # Check if model is just copying input
                        input_colors = input_grids.argmax(dim=1)
                        copying_accuracy = (pred_colors == input_colors).float().mean().item()
                        print(f"  Copying input accuracy: {copying_accuracy*100:.1f}%")
                    
                    # Active region accuracy
                    active_mask = (target_colors != 0) | (pred_colors != 0)
                    if active_mask.any():
                        active_correct = pixel_correct[active_mask]
                        val_active_acc_sum += active_correct.mean().item() * 100
                    else:
                        val_active_acc_sum += 100.0
                    
                    val_batches_count += 1
            
            # Visualize predictions every 10 epochs
            if epoch % 10 == 0:
                # Get a batch for visualization
                vis_batch = next(iter(val_loader))
                vis_input = vis_batch['input'].to(DEVICE)
                vis_output = vis_batch['output'].to(DEVICE)
                
                with torch.no_grad():
                    with autocast('cuda'):
                        if model_name == 'chronos':
                            vis_outputs = model.base_model([vis_input])
                        else:
                            vis_outputs = model(vis_input)
                        vis_pred = vis_outputs['predicted_output']
                
                visualize_predictions(vis_input, vis_output, vis_pred, epoch+1, model_name, num_samples=3)
                print(f"  📸 Saved visualization for epoch {epoch+1}")
            
            # Calculate metrics
            avg_train_loss = train_loss / train_steps
            avg_val_loss = val_loss / len(val_loader)
            val_accuracy = val_correct / val_total * 100
            avg_pixel_acc = val_pixel_acc_sum / val_batches_count
            avg_active_acc = val_active_acc_sum / val_batches_count
            avg_structure = val_structure_sum / val_batches_count
            
            history['train_loss'].append(avg_train_loss)
            history['val_loss'].append(avg_val_loss)
            history['val_accuracy'].append(val_accuracy)
            history['val_pixel_acc'].append(avg_pixel_acc)
            history['val_active_acc'].append(avg_active_acc)
            history['val_structure'].append(avg_structure)
            
            print(f"Epoch {epoch+1}: Train Loss: {avg_train_loss:.4f}, "
                  f"Val Loss: {avg_val_loss:.4f}, Exact: {val_accuracy:.2f}%, "
                  f"Pixel: {avg_pixel_acc:.2f}%, Active: {avg_active_acc:.2f}%, "
                  f"Structure: {avg_structure:.4f}")
            
            # Step ReduceLROnPlateau with validation accuracy
            scheduler.step(val_accuracy)
            
            # Update monitor
            monitor.update(
                model_name=model_name,
                epoch=epoch + 1,
                metrics={
                    'train_loss': avg_train_loss,
                    'val_loss': avg_val_loss,
                    'train_acc': 0.0,
                    'val_acc': val_accuracy / 100.0
                }
            )
            
            if (epoch + 1) % 20 == 0:
                monitor.show_dashboard()
            
            # Save best model
            if val_accuracy > best_val_acc:
                best_val_acc = val_accuracy
                patience_counter = 0
                
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'val_loss': avg_val_loss,
                    'val_accuracy': val_accuracy,
                    'val_pixel_acc': avg_pixel_acc
                }, f'/content/arc_models/{model_name}_v3_best.pt')
                
                print(f"✅ New best model! Exact: {val_accuracy:.2f}%, Pixel: {avg_pixel_acc:.2f}%")
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    print(f"Early stopping at epoch {epoch+1}")
                    break
        
        training_history[model_name] = history
        
        # Clear memory
        del model
        torch.cuda.empty_cache()
        gc.collect()
    
    # Final visualization and reporting
    create_v3_report(training_history)
    monitor.complete()
    print("\n🎉 V3 Training complete!")


def visualize_predictions(input_grids, output_grids, predicted_grids, epoch, model_name, num_samples=5):
    """Visualize predictions vs ground truth"""
    # Create color map for ARC (0-9 colors)
    colors = ['#000000', '#0074D9', '#FF4136', '#2ECC40', '#FFDC00',
              '#AAAAAA', '#F012BE', '#FF851B', '#7FDBFF', '#870C25']
    cmap = plt.matplotlib.colors.ListedColormap(colors)
    
    fig, axes = plt.subplots(num_samples, 3, figsize=(12, 4*num_samples))
    if num_samples == 1:
        axes = axes.reshape(1, -1)
    
    for i in range(min(num_samples, input_grids.shape[0])):
        # Convert from one-hot to color indices
        input_colors = input_grids[i].argmax(dim=0).cpu().numpy()
        target_colors = output_grids[i].argmax(dim=0).cpu().numpy()
        pred_colors = predicted_grids[i].argmax(dim=0).cpu().numpy()
        
        # Plot input
        axes[i, 0].imshow(input_colors, cmap=cmap, vmin=0, vmax=9)
        axes[i, 0].set_title('Input')
        axes[i, 0].axis('off')
        
        # Plot target
        axes[i, 1].imshow(target_colors, cmap=cmap, vmin=0, vmax=9)
        axes[i, 1].set_title('Target')
        axes[i, 1].axis('off')
        
        # Plot prediction
        axes[i, 2].imshow(pred_colors, cmap=cmap, vmin=0, vmax=9)
        # Check if exact match
        exact_match = np.array_equal(pred_colors, target_colors)
        title_color = 'green' if exact_match else 'red'
        axes[i, 2].set_title('Prediction', color=title_color)
        axes[i, 2].axis('off')
    
    plt.suptitle(f'{model_name} - Epoch {epoch}', fontsize=16)
    plt.tight_layout()
    
    # Save figure
    os.makedirs('/content/visualizations', exist_ok=True)
    plt.savefig(f'/content/visualizations/{model_name}_epoch_{epoch}.png', dpi=150, bbox_inches='tight')
    plt.close()


def create_v3_report(history: Dict):
    """Create V3 training report"""
    print("\n📝 Creating V3 training report...")
    
    report = f"""# ARC Prize 2025 - V3 Training Report

Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

## V3 Improvements

1. **Curriculum Learning**: Start with easy samples, gradually increase difficulty
2. **Improved Loss Function**:
   - Focal loss for hard pixels
   - Edge-aware weighting (3x on edges)
   - Color balance loss
   - Structure preservation loss
3. **Dropout Scheduling**: Decay dropout during training
4. **Mixed Optimizers**: AdamW for some models, SGD for others
5. **OneCycle Learning Rate**: Better convergence

## Results

| Model | Best Exact Match | Best Pixel Acc | Final Structure Loss |
|-------|------------------|----------------|---------------------|
"""
    
    for model_name, h in history.items():
        best_exact = max(h['val_accuracy']) if h['val_accuracy'] else 0
        best_pixel = max(h['val_pixel_acc']) if h['val_pixel_acc'] else 0
        final_structure = h['val_structure'][-1] if h['val_structure'] else 999
        
        report += f"| {model_name.upper()} | {best_exact:.2f}% | {best_pixel:.2f}% | {final_structure:.4f} |\n"
    
    report += "\n## Key Findings\n\n"
    report += "- Curriculum learning helps models learn basic patterns first\n"
    report += "- Edge-aware loss improves boundary precision\n"
    report += "- Structure preservation loss maintains object integrity\n"
    
    with open('/content/results/v3_training_report.md', 'w') as f:
        f.write(report)
    
    print("✅ Report saved!")


if __name__ == "__main__":
    print("="*80)
    print("ARC PRIZE 2025 - ENHANCED MODEL TRAINING V3")
    print("="*80)
    print("Major improvements for exact match accuracy")
    print("Target: 85% accuracy for $700,000 grand prize")
    print("="*80)
    print()
    
    train_enhanced_models_v3()