# experiments/train_pcam_cv_rtx4070_laptop.py
import torch
from sklearn.model_selection import StratifiedKFold

def train_with_cv(config, n_folds=5):
    # Load full dataset
    full_dataset = PatchCamelyonDataset(config.data_root)
    
    # Get labels for stratification
    labels = [full_dataset[i][1] for i in range(len(full_dataset))]
    
    skf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42)
    
    results = []
    for fold, (train_idx, val_idx) in enumerate(skf.split(range(len(full_dataset)), labels)):
        print(f"\n{'='*50}")
        print(f"Fold {fold+1}/{n_folds}")
        print(f"{'='*50}\n")
        
        # Create fold-specific loaders
        train_subset = torch.utils.data.Subset(full_dataset, train_idx)
        val_subset = torch.utils.data.Subset(full_dataset, val_idx)
        
        train_loader = DataLoader(
            train_subset,
            batch_size=64,  # RTX 4070 Laptop optimized
            shuffle=True,
            num_workers=12,  # Utilize 16-core CPU
            pin_memory=True
        )
        
        val_loader = DataLoader(
            val_subset,
            batch_size=96,  # Larger for inference
            shuffle=False,
            num_workers=12,
            pin_memory=True
        )
        
        # Train fold
        model = create_model(config)
        metrics = train_fold(model, train_loader, val_loader, config)
        
        results.append(metrics)
        
        # Clear GPU memory
        del model
        torch.cuda.empty_cache()
    
    # Aggregate results
    print_cv_results(results)
    return results

def print_cv_results(results):
    import numpy as np
    
    accs = [r['accuracy'] for r in results]
    aucs = [r['auc'] for r in results]
    
    print(f"\n{'='*50}")
    print("Cross-Validation Results")
    print(f"{'='*50}")
    print(f"Accuracy: {np.mean(accs):.4f} ± {np.std(accs):.4f}")
    print(f"AUC:      {np.mean(aucs):.4f} ± {np.std(aucs):.4f}")
    print(f"\nPer-fold results:")
    for i, r in enumerate(results):
        print(f"  Fold {i+1}: Acc={r['accuracy']:.4f}, AUC={r['auc']:.4f}")

# Add to training loop
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for epoch in range(num_epochs):
    for batch in train_loader:
        images, labels = batch
        images = images.cuda()
        labels = labels.cuda()
        
        optimizer.zero_grad()
        
        # Mixed precision forward pass
        with autocast():
            outputs = model(images)
            loss = criterion(outputs, labels)
        
        # Mixed precision backward pass
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

# src/utils/gradcam_rtx4070.py
import torch
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
import matplotlib.pyplot as plt

def generate_gradcam_batch(model, images, target_layer, batch_size=32):
    """
    Generate Grad-CAM for batch of images
    Configuration for RTX 4070 Laptop
    """
    model.eval()
    cam = GradCAM(model=model, target_layers=[target_layer])
    
    all_cams = []
    
    # Process in batches to avoid OOM
    for i in range(0, len(images), batch_size):
        batch = images[i:i+batch_size].cuda()
        
        with torch.no_grad():
            grayscale_cams = cam(input_tensor=batch)
        
        all_cams.append(grayscale_cams)
        
        # Clear cache
        torch.cuda.empty_cache()
    
    return torch.cat(all_cams)

# Usage example
model = load_model('checkpoints/best_model.pth').cuda()
target_layer = model.layer4[-1]  # Last conv layer

# Generate for test set
test_images = load_test_images()
heatmaps = generate_gradcam_batch(model, test_images, target_layer)

# Save visualizations
save_gradcam_grid(test_images, heatmaps, 'results/gradcam_examples.png')

import torch
torch.backends.cudnn.benchmark = True  # 10-20% speedup

train_loader = DataLoader(
    dataset,
    batch_size=64,
    pin_memory=True,  # Faster CPU->GPU transfer
    num_workers=12  # Utilize 16-core CPU (start with 0, increase gradually)
)

# PCamDataset now opens h5 files lazily in each worker process
# This avoids pickling errors with multiprocessing
def __getitem__(self, idx):
    # Open h5 files lazily if not already open
    if self._images_h5 is None or self._labels_h5 is None:
        self._open_h5_files()
    
    # Load data...

# Add to training loop
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for epoch in range(num_epochs):
    for batch in train_loader:
        images, labels = batch
        images = images.cuda()
        labels = labels.cuda()
        
        optimizer.zero_grad()
        
        # Mixed precision forward pass
        with autocast():
            outputs = model(images)
            loss = criterion(outputs, labels)
        
        # Mixed precision backward pass
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

accumulation_steps = 2  # Effective batch size = 64 * 2 = 128

for i, (images, labels) in enumerate(train_loader):
    outputs = model(images)
    loss = criterion(outputs, labels) / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# Clear cache between experiments
torch.cuda.empty_cache()

# Delete unused tensors
del model, optimizer
torch.cuda.empty_cache()

# Cache clearing in training loop
if batch_idx % 100 == 0:
    torch.cuda.empty_cache()

# Profiling for bottleneck detection
with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU,
        torch.profiler.ProfilerActivity.CUDA,
    ]
) as prof:
    train_one_epoch(model, train_loader)

print(prof.key_averages().table(sort_by="cuda_time_total"))