trainer_v1.py

"""
Train DavidBeans: The Dynamic Duo
==================================

           ┌─────────────────┐
           │   BEANS         │  "I see the patches..."
           │   (ViT Backbone)│  
           │   🫘 → 🫘 → 🫘   │  Cantor-routed sparse attention
           └────────┬────────┘
                    │
                    ▼
           ┌─────────────────┐
           │   DAVID         │  "I know the crystals..."
           │   (Classifier)  │
           │   💎 → 💎 → 💎   │  Multi-scale projection
           └────────┬────────┘
                    │
                    ▼
              [Prediction]

Cross-contrast learning aligns patch features with crystal anchors.
Unified Cayley-Menger loss maintains geometric structure throughout.

Features:
- HuggingFace Hub integration for model upload
- Automatic model card generation
- Checkpoint management

Author: AbstractPhil
Date: November 28, 2025
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR
from tqdm.auto import tqdm
import time
import math
from pathlib import Path
from typing import Dict, Optional, Tuple, List
from dataclasses import dataclass, field
import json
import os
from datetime import datetime

# Import the model
from geofractal.model.david_beans.model import DavidBeans, DavidBeansConfig

# HuggingFace Hub integration
try:
    from huggingface_hub import HfApi, create_repo, upload_folder
    HF_HUB_AVAILABLE = True
except ImportError:
    HF_HUB_AVAILABLE = False
    print("  [!] huggingface_hub not installed. Run: pip install huggingface_hub")

# Safetensors support
try:
    from safetensors.torch import save_file as save_safetensors
    SAFETENSORS_AVAILABLE = True
except ImportError:
    SAFETENSORS_AVAILABLE = False

# TensorBoard support
try:
    from torch.utils.tensorboard import SummaryWriter
    TENSORBOARD_AVAILABLE = True
except ImportError:
    TENSORBOARD_AVAILABLE = False
    print("  [!] tensorboard not installed. Run: pip install tensorboard")


# ============================================================================
# TRAINING CONFIGURATION
# ============================================================================

@dataclass
class TrainingConfig:
    """Training hyperparameters."""
    
    # Run identification
    run_name: str = "default"  # Descriptive name for this run
    run_number: Optional[int] = None  # Auto-incremented if None
    
    # Data
    dataset: str = "cifar10"
    image_size: int = 32
    batch_size: int = 128
    num_workers: int = 4
    
    # Training schedule
    epochs: int = 100
    warmup_epochs: int = 5
    
    # Optimizer
    learning_rate: float = 1e-3
    weight_decay: float = 0.05
    betas: Tuple[float, float] = (0.9, 0.999)
    
    # Learning rate schedule
    scheduler: str = "cosine"
    min_lr: float = 1e-6
    
    # Loss weights
    ce_weight: float = 1.0
    cayley_weight: float = 0.01
    contrast_weight: float = 0.5
    scale_ce_weight: float = 0.1
    
    # Regularization
    gradient_clip: float = 1.0
    label_smoothing: float = 0.1
    
    # Augmentation
    use_augmentation: bool = True
    mixup_alpha: float = 0.2
    cutmix_alpha: float = 1.0
    
    # Checkpointing
    save_interval: int = 10
    output_dir: str = "./checkpoints"
    resume_from: Optional[str] = None  # Path to checkpoint or "latest"
    
    # TensorBoard
    use_tensorboard: bool = True
    log_interval: int = 50  # Log every N batches
    
    # HuggingFace Hub
    push_to_hub: bool = False
    hub_repo_id: Optional[str] = None
    hub_private: bool = False
    hub_append_run: bool = True  # Append run info to repo_id (e.g., repo-run001-baseline)
    
    # Device
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    
    def to_dict(self) -> Dict:
        return {k: v for k, v in self.__dict__.items()}


# ============================================================================
# HUGGINGFACE HUB INTEGRATION
# ============================================================================

def generate_model_card(
    model_config: DavidBeansConfig,
    train_config: TrainingConfig,
    best_acc: float,
    training_history: Optional[Dict] = None
) -> str:
    """Generate a model card for HuggingFace Hub."""
    
    scales_str = ", ".join([str(s) for s in model_config.scales])
    
    dataset_info = {
        "cifar10": ("CIFAR-10", 10, "Image classification on 32x32 images"),
        "cifar100": ("CIFAR-100", 100, "Fine-grained image classification on 32x32 images"),
    }.get(train_config.dataset, (train_config.dataset, model_config.num_classes, ""))
    
    card_content = f"""---
library_name: pytorch
license: apache-2.0
tags:
  - vision
  - image-classification
  - geometric-deep-learning
  - vit
  - cantor-routing
  - pentachoron
  - multi-scale
datasets:
  - {train_config.dataset}
metrics:
  - accuracy
model-index:
  - name: DavidBeans
    results:
      - task:
          type: image-classification
          name: Image Classification
        dataset:
          name: {dataset_info[0]}
          type: {train_config.dataset}
        metrics:
          - type: accuracy
            value: {best_acc:.2f}
            name: Top-1 Accuracy
---

# 🫘💎 DavidBeans: Unified Vision-to-Crystal Architecture

DavidBeans combines **ViT-Beans** (Cantor-routed sparse attention) with **David** (multi-scale crystal classification) into a unified geometric deep learning architecture.

## Model Description

This model implements several novel techniques:

- **Hybrid Cantor Routing**: Combines fractal Cantor set distances with positional proximity for sparse attention patterns
- **Pentachoron Experts**: 5-vertex simplex structure with Cayley-Menger geometric regularization
- **Multi-Scale Crystal Projection**: Projects features to multiple representation scales with learned fusion
- **Cross-Contrastive Learning**: Aligns patch-level features with crystal anchors

## Architecture

```
Image [B, 3, {model_config.image_size}, {model_config.image_size}]
       │
       ▼
┌─────────────────────────────────────────┐
│  BEANS BACKBONE                         │
│  ├─ Patch Embed → [{model_config.num_patches} patches, {model_config.dim}d]
│  ├─ Hybrid Cantor Router (α={model_config.cantor_weight})
│  ├─ {model_config.num_layers} × Attention Blocks ({model_config.num_heads} heads)
│  └─ {model_config.num_layers} × Pentachoron Expert Layers
└─────────────────────────────────────────┘
       │
       ▼
┌─────────────────────────────────────────┐
│  DAVID HEAD                             │
│  ├─ Multi-scale projection: [{scales_str}]
│  ├─ Per-scale Crystal Heads
│  └─ Geometric Fusion (learned weights)
└─────────────────────────────────────────┘
       │
       ▼
    [{model_config.num_classes} classes]
```

## Training Details

| Parameter | Value |
|-----------|-------|
| Dataset | {dataset_info[0]} |
| Classes | {model_config.num_classes} |
| Image Size | {model_config.image_size}×{model_config.image_size} |
| Patch Size | {model_config.patch_size}×{model_config.patch_size} |
| Embedding Dim | {model_config.dim} |
| Layers | {model_config.num_layers} |
| Attention Heads | {model_config.num_heads} |
| Experts | {model_config.num_experts} (pentachoron) |
| Sparse Neighbors | k={model_config.k_neighbors} |
| Scales | [{scales_str}] |
| Epochs | {train_config.epochs} |
| Batch Size | {train_config.batch_size} |
| Learning Rate | {train_config.learning_rate} |
| Weight Decay | {train_config.weight_decay} |
| Mixup α | {train_config.mixup_alpha} |
| CutMix α | {train_config.cutmix_alpha} |
| Label Smoothing | {train_config.label_smoothing} |

## Results

| Metric | Value |
|--------|-------|
| **Top-1 Accuracy** | **{best_acc:.2f}%** |

## TensorBoard Logs

Training logs are included in the `tensorboard/` directory. To view:

```bash
tensorboard --logdir tensorboard/
```

## Usage

```python
import torch
from safetensors.torch import load_file
from david_beans import DavidBeans, DavidBeansConfig

# Load config
config = DavidBeansConfig(
    image_size={model_config.image_size},
    patch_size={model_config.patch_size},
    dim={model_config.dim},
    num_layers={model_config.num_layers},
    num_heads={model_config.num_heads},
    num_experts={model_config.num_experts},
    k_neighbors={model_config.k_neighbors},
    cantor_weight={model_config.cantor_weight},
    scales={model_config.scales},
    num_classes={model_config.num_classes}
)

# Create model and load weights
model = DavidBeans(config)
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict)

# Inference
model.eval()
with torch.no_grad():
    output = model(images)
    predictions = output['logits'].argmax(dim=-1)
```

## Citation

```bibtex
@misc{{davidbeans2025,
  author = {{AbstractPhil}},
  title = {{DavidBeans: Unified Vision-to-Crystal Architecture}},
  year = {{2025}},
  publisher = {{HuggingFace}},
  url = {{https://huggingface.co/{train_config.hub_repo_id or 'AbstractPhil/david-beans'}}}
}}
```

## License

Apache 2.0
"""
    
    return card_content


def save_for_hub(
    model: DavidBeans,
    model_config: DavidBeansConfig,
    train_config: TrainingConfig,
    best_acc: float,
    output_dir: Path,
    training_history: Optional[Dict] = None
) -> Path:
    """Save model in HuggingFace Hub format."""
    
    hub_dir = output_dir / "hub"
    hub_dir.mkdir(parents=True, exist_ok=True)
    
    # 1. Save model weights - clone to avoid shared memory issues
    state_dict = {k: v.clone() for k, v in model.state_dict().items()}
    
    if SAFETENSORS_AVAILABLE:
        try:
            save_safetensors(state_dict, hub_dir / "model.safetensors")
            print(f"  ✓ Saved model.safetensors")
        except Exception as e:
            print(f"  [!] Safetensors failed ({e}), using pytorch format only")
    
    # Also save PyTorch format for compatibility
    torch.save(state_dict, hub_dir / "pytorch_model.bin")
    print(f"  ✓ Saved pytorch_model.bin")
    
    # 2. Save config
    config_dict = {
        "architecture": "DavidBeans",
        "model_type": "david_beans",
        **model_config.__dict__
    }
    with open(hub_dir / "config.json", "w") as f:
        json.dump(config_dict, f, indent=2, default=str)
    print(f"  ✓ Saved config.json")
    
    # 3. Save training config
    with open(hub_dir / "training_config.json", "w") as f:
        json.dump(train_config.to_dict(), f, indent=2, default=str)
    
    # 4. Generate and save model card
    model_card = generate_model_card(model_config, train_config, best_acc, training_history)
    with open(hub_dir / "README.md", "w") as f:
        f.write(model_card)
    print(f"  ✓ Generated README.md (model card)")
    
    # 5. Save training history if available
    if training_history:
        with open(hub_dir / "training_history.json", "w") as f:
            json.dump(training_history, f, indent=2)
    
    # 6. Copy TensorBoard logs if they exist
    tb_dir = output_dir / "tensorboard"
    if tb_dir.exists():
        import shutil
        hub_tb_dir = hub_dir / "tensorboard"
        if hub_tb_dir.exists():
            shutil.rmtree(hub_tb_dir)
        shutil.copytree(tb_dir, hub_tb_dir)
        print(f"  ✓ Copied TensorBoard logs")
    
    return hub_dir


def push_to_hub(
    hub_dir: Path,
    repo_id: str,
    private: bool = False,
    commit_message: Optional[str] = None
) -> str:
    """Push model to HuggingFace Hub."""
    
    if not HF_HUB_AVAILABLE:
        raise RuntimeError("huggingface_hub not installed. Run: pip install huggingface_hub")
    
    api = HfApi()
    
    # Create repo if it doesn't exist
    try:
        create_repo(repo_id, private=private, exist_ok=True)
        print(f"  ✓ Repository ready: {repo_id}")
    except Exception as e:
        print(f"  [!] Repo creation note: {e}")
    
    # Upload
    if commit_message is None:
        commit_message = f"Upload DavidBeans model - {datetime.now().strftime('%Y-%m-%d %H:%M')}"
    
    url = upload_folder(
        folder_path=str(hub_dir),
        repo_id=repo_id,
        commit_message=commit_message
    )
    
    print(f"  ✓ Uploaded to: https://huggingface.co/{repo_id}")
    
    return url


# ============================================================================
# DATA LOADING
# ============================================================================

def get_dataloaders(config: TrainingConfig) -> Tuple[DataLoader, DataLoader, int]:
    """Get train and test dataloaders."""
    
    try:
        import torchvision
        import torchvision.transforms as T
        
        if config.dataset == "cifar10":
            if config.use_augmentation:
                train_transform = T.Compose([
                    T.RandomCrop(32, padding=4),
                    T.RandomHorizontalFlip(),
                    T.AutoAugment(T.AutoAugmentPolicy.CIFAR10),
                    T.ToTensor(),
                    T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
                ])
            else:
                train_transform = T.Compose([
                    T.ToTensor(),
                    T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
                ])
            
            test_transform = T.Compose([
                T.ToTensor(),
                T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
            ])
            
            train_dataset = torchvision.datasets.CIFAR10(
                root='./data', train=True, download=True, transform=train_transform
            )
            test_dataset = torchvision.datasets.CIFAR10(
                root='./data', train=False, download=True, transform=test_transform
            )
            num_classes = 10
            
        elif config.dataset == "cifar100":
            if config.use_augmentation:
                train_transform = T.Compose([
                    T.RandomCrop(32, padding=4),
                    T.RandomHorizontalFlip(),
                    T.AutoAugment(T.AutoAugmentPolicy.CIFAR10),
                    T.ToTensor(),
                    T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
                ])
            else:
                train_transform = T.Compose([
                    T.ToTensor(),
                    T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
                ])
            
            test_transform = T.Compose([
                T.ToTensor(),
                T.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761))
            ])
            
            train_dataset = torchvision.datasets.CIFAR100(
                root='./data', train=True, download=True, transform=train_transform
            )
            test_dataset = torchvision.datasets.CIFAR100(
                root='./data', train=False, download=True, transform=test_transform
            )
            num_classes = 100
        else:
            raise ValueError(f"Unknown dataset: {config.dataset}")
        
        train_loader = DataLoader(
            train_dataset,
            batch_size=config.batch_size,
            shuffle=True,
            num_workers=config.num_workers,
            pin_memory=True,
            persistent_workers=config.num_workers > 0,
            drop_last=True
        )
        test_loader = DataLoader(
            test_dataset,
            batch_size=config.batch_size,
            shuffle=False,
            num_workers=config.num_workers,
            pin_memory=True,
            persistent_workers=config.num_workers > 0
        )
        
        return train_loader, test_loader, num_classes
        
    except ImportError:
        print("  [!] torchvision not available, using synthetic data")
        return get_synthetic_dataloaders(config)


def get_synthetic_dataloaders(config: TrainingConfig) -> Tuple[DataLoader, DataLoader, int]:
    """Fallback synthetic data for testing."""
    
    class SyntheticDataset(torch.utils.data.Dataset):
        def __init__(self, size: int, image_size: int, num_classes: int):
            self.size = size
            self.image_size = image_size
            self.num_classes = num_classes
            
        def __len__(self):
            return self.size
            
        def __getitem__(self, idx):
            x = torch.randn(3, self.image_size, self.image_size)
            y = idx % self.num_classes
            return x, y
    
    num_classes = 10
    train_dataset = SyntheticDataset(5000, config.image_size, num_classes)
    test_dataset = SyntheticDataset(1000, config.image_size, num_classes)
    
    train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False)
    
    return train_loader, test_loader, num_classes


# ============================================================================
# MIXUP / CUTMIX AUGMENTATION
# ============================================================================

def mixup_data(x: torch.Tensor, y: torch.Tensor, alpha: float = 0.2):
    """Mixup augmentation."""
    if alpha > 0:
        lam = torch.distributions.Beta(alpha, alpha).sample().item()
    else:
        lam = 1.0
    
    batch_size = x.size(0)
    index = torch.randperm(batch_size, device=x.device)
    
    mixed_x = lam * x + (1 - lam) * x[index]
    y_a, y_b = y, y[index]
    
    return mixed_x, y_a, y_b, lam


def cutmix_data(x: torch.Tensor, y: torch.Tensor, alpha: float = 1.0):
    """CutMix augmentation."""
    if alpha > 0:
        lam = torch.distributions.Beta(alpha, alpha).sample().item()
    else:
        lam = 1.0
    
    batch_size = x.size(0)
    index = torch.randperm(batch_size, device=x.device)
    
    _, _, H, W = x.shape
    
    cut_ratio = math.sqrt(1 - lam)
    cut_h = int(H * cut_ratio)
    cut_w = int(W * cut_ratio)
    
    cx = torch.randint(0, H, (1,)).item()
    cy = torch.randint(0, W, (1,)).item()
    
    x1 = max(0, cx - cut_h // 2)
    x2 = min(H, cx + cut_h // 2)
    y1 = max(0, cy - cut_w // 2)
    y2 = min(W, cy + cut_w // 2)
    
    mixed_x = x.clone()
    mixed_x[:, :, x1:x2, y1:y2] = x[index, :, x1:x2, y1:y2]
    
    lam = 1 - ((x2 - x1) * (y2 - y1)) / (H * W)
    
    y_a, y_b = y, y[index]
    
    return mixed_x, y_a, y_b, lam


# ============================================================================
# METRICS TRACKER
# ============================================================================

class MetricsTracker:
    """Track training metrics with EMA smoothing."""
    
    def __init__(self, ema_decay: float = 0.9):
        self.ema_decay = ema_decay
        self.metrics = {}
        self.ema_metrics = {}
        self.history = {}
    
    def update(self, **kwargs):
        for k, v in kwargs.items():
            if isinstance(v, torch.Tensor):
                v = v.item()
            
            if k not in self.metrics:
                self.metrics[k] = []
                self.ema_metrics[k] = v
                self.history[k] = []
            
            self.metrics[k].append(v)
            self.ema_metrics[k] = self.ema_decay * self.ema_metrics[k] + (1 - self.ema_decay) * v
    
    def get_ema(self, key: str) -> float:
        return self.ema_metrics.get(key, 0.0)
    
    def get_epoch_mean(self, key: str) -> float:
        values = self.metrics.get(key, [])
        return sum(values) / len(values) if values else 0.0
    
    def end_epoch(self):
        for k, v in self.metrics.items():
            if v:
                self.history[k].append(sum(v) / len(v))
        self.metrics = {k: [] for k in self.metrics}
    
    def get_history(self) -> Dict:
        return self.history


# ============================================================================
# CHECKPOINT UTILITIES
# ============================================================================

def find_latest_checkpoint(output_dir: Path) -> Optional[Path]:
    """Find the most recent checkpoint in output directory."""
    checkpoints = list(output_dir.glob("checkpoint_epoch_*.pt"))
    
    if not checkpoints:
        # Try best_model.pt as fallback
        best_model = output_dir / "best_model.pt"
        if best_model.exists():
            return best_model
        return None
    
    # Sort by epoch number
    def get_epoch(p):
        try:
            return int(p.stem.split("_")[-1])
        except:
            return 0
    
    checkpoints.sort(key=get_epoch, reverse=True)
    return checkpoints[0]


def get_next_run_number(base_dir: Path) -> int:
    """Get the next run number by scanning existing run directories."""
    if not base_dir.exists():
        return 1
    
    max_num = 0
    for d in base_dir.iterdir():
        if d.is_dir() and d.name.startswith("run_"):
            try:
                # Extract number from "run_XXX_name_timestamp"
                num = int(d.name.split("_")[1])
                max_num = max(max_num, num)
            except (IndexError, ValueError):
                continue
    
    return max_num + 1


def generate_run_dir_name(run_number: int, run_name: str) -> str:
    """Generate a run directory name with number, name, and timestamp."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    # Sanitize run_name: lowercase, replace spaces with underscores, remove special chars
    safe_name = "".join(c if c.isalnum() or c == "_" else "_" for c in run_name.lower())
    safe_name = "_".join(filter(None, safe_name.split("_")))  # Remove consecutive underscores
    return f"run_{run_number:03d}_{safe_name}_{timestamp}"


def find_latest_run_dir(base_dir: Path) -> Optional[Path]:
    """Find the most recent run directory."""
    if not base_dir.exists():
        return None
    
    run_dirs = [d for d in base_dir.iterdir() if d.is_dir() and d.name.startswith("run_")]
    
    if not run_dirs:
        return None
    
    # Sort by modification time (most recent first)
    run_dirs.sort(key=lambda d: d.stat().st_mtime, reverse=True)
    return run_dirs[0]


def find_checkpoint_in_runs(base_dir: Path, resume_from: str) -> Optional[Path]:
    """
    Find a checkpoint to resume from.
    
    Args:
        base_dir: Base checkpoint directory (e.g., ./checkpoints/cifar100)
        resume_from: Either "latest", a run directory name, or a full path
    
    Returns:
        Path to checkpoint file, or None
    """
    if resume_from == "latest":
        # Find most recent run directory
        run_dir = find_latest_run_dir(base_dir)
        if run_dir:
            return find_latest_checkpoint(run_dir)
        # Fallback: check base_dir itself (for old-style checkpoints)
        return find_latest_checkpoint(base_dir)
    
    # Check if it's a full path
    full_path = Path(resume_from)
    if full_path.exists():
        if full_path.is_file():
            return full_path
        elif full_path.is_dir():
            return find_latest_checkpoint(full_path)
    
    # Check if it's a run directory name within base_dir
    run_path = base_dir / resume_from
    if run_path.exists():
        return find_latest_checkpoint(run_path)
    
    return None


def load_checkpoint(
    checkpoint_path: Path,
    model: DavidBeans,
    optimizer: Optional[torch.optim.Optimizer] = None,
    device: str = "cuda"
) -> Tuple[int, float]:
    """
    Load checkpoint and return (start_epoch, best_acc).
    
    Returns:
        start_epoch: Epoch to resume from (checkpoint_epoch + 1)
        best_acc: Best accuracy so far
    """
    print(f"\n📂 Loading checkpoint: {checkpoint_path}")
    checkpoint = torch.load(checkpoint_path, map_location=device)
    
    model.load_state_dict(checkpoint['model_state_dict'])
    print(f"  ✓ Loaded model weights")
    
    if optimizer is not None and 'optimizer_state_dict' in checkpoint:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        print(f"  ✓ Loaded optimizer state")
    
    epoch = checkpoint.get('epoch', 0)
    best_acc = checkpoint.get('best_acc', 0.0)
    
    print(f"  ✓ Loaded checkpoint from epoch {epoch + 1}, best_acc={best_acc:.2f}%")
    print(f"  ✓ Will resume training from epoch {epoch + 2}")
    
    return epoch + 1, best_acc


def get_config_from_checkpoint(checkpoint_path: Path) -> Tuple[DavidBeansConfig, dict]:
    """
    Extract model and training configs from a checkpoint.
    
    Returns:
        (model_config, train_config_dict)
    """
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    
    model_config_dict = checkpoint.get('model_config', {})
    train_config_dict = checkpoint.get('train_config', {})
    
    # Handle tuple conversion for betas
    if 'betas' in train_config_dict and isinstance(train_config_dict['betas'], list):
        train_config_dict['betas'] = tuple(train_config_dict['betas'])
    
    model_config = DavidBeansConfig(**model_config_dict)
    
    return model_config, train_config_dict


# ============================================================================
# TRAINING LOOP
# ============================================================================

def train_epoch(
    model: DavidBeans,
    train_loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
    config: TrainingConfig,
    epoch: int,
    tracker: MetricsTracker,
    writer: Optional['SummaryWriter'] = None
) -> Dict[str, float]:
    """Train for one epoch."""
    
    model.train()
    device = config.device
    
    total_loss = 0.0
    total_correct = 0
    total_samples = 0
    global_step = epoch * len(train_loader)
    
    pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}", leave=True)
    
    for batch_idx, (images, targets) in enumerate(pbar):
        images = images.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)
        
        # Apply mixup/cutmix
        use_mixup = config.use_augmentation and config.mixup_alpha > 0
        use_cutmix = config.use_augmentation and config.cutmix_alpha > 0
        
        mixed = False
        if use_mixup or use_cutmix:
            r = torch.rand(1).item()
            if r < 0.5:
                pass
            elif r < 0.75 and use_mixup:
                images, targets_a, targets_b, lam = mixup_data(images, targets, config.mixup_alpha)
                mixed = True
            elif use_cutmix:
                images, targets_a, targets_b, lam = cutmix_data(images, targets, config.cutmix_alpha)
                mixed = True
        
        # Forward pass
        result = model(images, targets=targets, return_loss=True)
        losses = result['losses']
        
        if mixed:
            logits = result['logits']
            ce_loss = lam * F.cross_entropy(logits, targets_a, label_smoothing=config.label_smoothing) + \
                      (1 - lam) * F.cross_entropy(logits, targets_b, label_smoothing=config.label_smoothing)
            losses['ce'] = ce_loss
        
        # Compute total loss
        loss = (
            config.ce_weight * losses['ce'] +
            config.cayley_weight * losses.get('geometric', torch.tensor(0.0, device=device)) +
            config.contrast_weight * losses.get('contrast', torch.tensor(0.0, device=device))
        )
        
        for scale in model.config.scales:
            scale_ce = losses.get(f'ce_{scale}', 0.0)
            if isinstance(scale_ce, torch.Tensor):
                loss = loss + config.scale_ce_weight * scale_ce
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        
        if config.gradient_clip > 0:
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)
        else:
            grad_norm = 0.0
        
        optimizer.step()
        
        if scheduler is not None and config.scheduler == "onecycle":
            scheduler.step()
        
        # Compute accuracy
        with torch.no_grad():
            logits = result['logits']
            preds = logits.argmax(dim=-1)
            
            if mixed:
                correct = (lam * (preds == targets_a).float() + 
                          (1 - lam) * (preds == targets_b).float()).sum()
            else:
                correct = (preds == targets).sum()
            
            total_correct += correct.item()
            total_samples += targets.size(0)
            total_loss += loss.item()
        
        # Track metrics
        def to_float(v):
            return v.item() if isinstance(v, torch.Tensor) else float(v)
        
        geo_loss = to_float(losses.get('geometric', 0.0))
        contrast_loss = to_float(losses.get('contrast', 0.0))
        expert_vol = to_float(losses.get('expert_volume', 0.0))
        expert_collapse = to_float(losses.get('expert_collapse', 0.0))
        expert_edge = to_float(losses.get('expert_edge_dev', 0.0))
        current_lr = optimizer.param_groups[0]['lr']
        
        tracker.update(
            loss=loss.item(),
            ce=losses['ce'].item(),
            geo=geo_loss,
            contrast=contrast_loss,
            expert_vol=expert_vol,
            expert_collapse=expert_collapse,
            expert_edge=expert_edge,
            lr=current_lr
        )
        
        # TensorBoard logging (every log_interval batches)
        if writer is not None and (batch_idx + 1) % config.log_interval == 0:
            step = global_step + batch_idx
            
            # Loss components
            writer.add_scalar('train/loss_total', loss.item(), step)
            writer.add_scalar('train/loss_ce', losses['ce'].item(), step)
            writer.add_scalar('train/loss_geometric', geo_loss, step)
            writer.add_scalar('train/loss_contrast', contrast_loss, step)
            
            # Geometric metrics
            writer.add_scalar('train/expert_volume', expert_vol, step)
            writer.add_scalar('train/expert_collapse', expert_collapse, step)
            writer.add_scalar('train/expert_edge_dev', expert_edge, step)
            
            # Training dynamics
            writer.add_scalar('train/learning_rate', current_lr, step)
            writer.add_scalar('train/grad_norm', to_float(grad_norm), step)
            writer.add_scalar('train/batch_acc', 100.0 * correct.item() / targets.size(0), step)
        
        pbar.set_postfix({
            'loss': f"{tracker.get_ema('loss'):.3f}",
            'acc': f"{100.0 * total_correct / total_samples:.1f}%",
            'geo': f"{tracker.get_ema('geo'):.4f}",
            'vol': f"{tracker.get_ema('expert_vol'):.4f}"
        })
    
    if scheduler is not None and config.scheduler == "cosine":
        scheduler.step()
    
    return {
        'loss': total_loss / len(train_loader),
        'acc': 100.0 * total_correct / total_samples
    }


@torch.no_grad()
def evaluate(
    model: DavidBeans,
    test_loader: DataLoader,
    config: TrainingConfig
) -> Dict[str, float]:
    """Evaluate on test set."""
    
    model.eval()
    device = config.device
    
    total_loss = 0.0
    total_correct = 0
    total_samples = 0
    scale_correct = {s: 0 for s in model.config.scales}
    
    for images, targets in test_loader:
        images = images.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)
        
        result = model(images, targets=targets, return_loss=True)
        
        logits = result['logits']
        losses = result['losses']
        
        loss = losses['total']
        preds = logits.argmax(dim=-1)
        
        total_loss += loss.item() * targets.size(0)
        total_correct += (preds == targets).sum().item()
        total_samples += targets.size(0)
        
        for i, scale in enumerate(model.config.scales):
            scale_logits = result['scale_logits'][i]
            scale_preds = scale_logits.argmax(dim=-1)
            scale_correct[scale] += (scale_preds == targets).sum().item()
    
    metrics = {
        'loss': total_loss / total_samples,
        'acc': 100.0 * total_correct / total_samples
    }
    
    for scale, correct in scale_correct.items():
        metrics[f'acc_{scale}'] = 100.0 * correct / total_samples
    
    return metrics


# ============================================================================
# MAIN TRAINING FUNCTION
# ============================================================================

def train_david_beans(
    model_config: Optional[DavidBeansConfig] = None,
    train_config: Optional[TrainingConfig] = None
):
    """Main training function."""
    
    print("=" * 70)
    print("  DAVID-BEANS TRAINING: The Dynamic Duo")
    print("=" * 70)
    print()
    print("       🫘 BEANS (ViT)  +  💎 DAVID (Crystal)")
    print("       Sparse Attention     Multi-Scale Projection")
    print()
    print("=" * 70)
    
    if train_config is None:
        train_config = TrainingConfig()
    
    base_output_dir = Path(train_config.output_dir)
    base_output_dir.mkdir(parents=True, exist_ok=True)
    
    # Check for resume FIRST - load config from checkpoint if resuming
    checkpoint_path = None
    run_dir = None  # Will be set either from resume or new run
    
    if train_config.resume_from:
        # Find checkpoint using the new directory structure
        checkpoint_path = find_checkpoint_in_runs(base_output_dir, train_config.resume_from)
        
        if checkpoint_path and checkpoint_path.exists():
            print(f"\n📂 Found checkpoint: {checkpoint_path}")
            # The run directory is the parent of the checkpoint
            run_dir = checkpoint_path.parent
            print(f"  ✓ Resuming in run directory: {run_dir.name}")
            
            # Load config from checkpoint to ensure architecture matches
            loaded_model_config, loaded_train_config_dict = get_config_from_checkpoint(checkpoint_path)
            
            if model_config is None:
                model_config = loaded_model_config
                print(f"  ✓ Using model config from checkpoint")
            else:
                # Warn if configs differ
                if model_config.dim != loaded_model_config.dim or model_config.scales != loaded_model_config.scales:
                    print(f"  ⚠ WARNING: Provided config differs from checkpoint!")
                    print(f"    Checkpoint: dim={loaded_model_config.dim}, scales={loaded_model_config.scales}")
                    print(f"    Provided:   dim={model_config.dim}, scales={model_config.scales}")
                    print(f"  ✓ Using checkpoint config to ensure compatibility")
                    model_config = loaded_model_config
        else:
            print(f"  [!] Checkpoint not found: {train_config.resume_from}")
            checkpoint_path = None
    
    # If not resuming (or resume failed), create new run directory
    if run_dir is None:
        # Get run number
        if train_config.run_number is None:
            run_number = get_next_run_number(base_output_dir)
        else:
            run_number = train_config.run_number
        
        # Generate run directory name
        run_dir_name = generate_run_dir_name(run_number, train_config.run_name)
        run_dir = base_output_dir / run_dir_name
        run_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"\n📁 New run: {run_dir_name}")
        print(f"  Run #{run_number}: {train_config.run_name}")
    else:
        # Extract run number from existing directory name for hub repo
        try:
            run_number = int(run_dir.name.split("_")[1])
        except (IndexError, ValueError):
            run_number = 1
    
    # Update output_dir to point to the run directory
    output_dir = run_dir
    
    # Generate effective hub repo ID with run info
    effective_hub_repo_id = train_config.hub_repo_id
    if train_config.hub_repo_id and train_config.hub_append_run:
        # Extract run name from directory (run_XXX_name_timestamp -> name)
        parts = run_dir.name.split("_")
        if len(parts) >= 3:
            run_name_part = parts[2]  # Get the name part
        else:
            run_name_part = train_config.run_name
        effective_hub_repo_id = f"{train_config.hub_repo_id}-run{run_number:03d}-{run_name_part}"
        print(f"  Hub repo: {effective_hub_repo_id}")
    
    if model_config is None:
        model_config = DavidBeansConfig(
            image_size=train_config.image_size,
            patch_size=4,
            dim=256,
            num_layers=6,
            num_heads=8,
            num_experts=5,
            k_neighbors=16,
            cantor_weight=0.3,
            scales=[64, 128, 256],
            num_classes=10,
            contrast_weight=train_config.contrast_weight,
            cayley_weight=train_config.cayley_weight,
            dropout=0.1
        )
    
    device = train_config.device
    print(f"\nDevice: {device}")
    
    # Data
    print("\nLoading data...")
    train_loader, test_loader, num_classes = get_dataloaders(train_config)
    print(f"  Dataset: {train_config.dataset}")
    print(f"  Train: {len(train_loader.dataset)}, Test: {len(test_loader.dataset)}")
    print(f"  Classes: {num_classes}")
    
    model_config.num_classes = num_classes
    
    # Model
    print("\nBuilding model...")
    model = DavidBeans(model_config)
    model = model.to(device)
    
    print(f"\n{model}")
    
    num_params = sum(p.numel() for p in model.parameters())
    num_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"\nParameters: {num_params:,} ({num_trainable:,} trainable)")
    
    # Optimizer
    print("\nSetting up optimizer...")
    
    decay_params = []
    no_decay_params = []
    
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if 'bias' in name or 'norm' in name or 'embedding' in name:
            no_decay_params.append(param)
        else:
            decay_params.append(param)
    
    optimizer = AdamW([
        {'params': decay_params, 'weight_decay': train_config.weight_decay},
        {'params': no_decay_params, 'weight_decay': 0.0}
    ], lr=train_config.learning_rate, betas=train_config.betas)
    
    if train_config.scheduler == "cosine":
        scheduler = CosineAnnealingLR(
            optimizer,
            T_max=train_config.epochs - train_config.warmup_epochs,
            eta_min=train_config.min_lr
        )
    elif train_config.scheduler == "onecycle":
        scheduler = OneCycleLR(
            optimizer,
            max_lr=train_config.learning_rate,
            epochs=train_config.epochs,
            steps_per_epoch=len(train_loader),
            pct_start=train_config.warmup_epochs / train_config.epochs
        )
    else:
        scheduler = None
    
    print(f"  Optimizer: AdamW (lr={train_config.learning_rate}, wd={train_config.weight_decay})")
    print(f"  Scheduler: {train_config.scheduler}")
    print(f"  TensorBoard: {output_dir / 'tensorboard'}")
    
    tracker = MetricsTracker()
    best_acc = 0.0
    start_epoch = 0
    
    print(f"\nOutput directory: {output_dir}")
    
    # Load weights from checkpoint if we found one earlier
    if checkpoint_path and checkpoint_path.exists():
        start_epoch, best_acc = load_checkpoint(
            checkpoint_path, model, optimizer, device
        )
        
        # Adjust scheduler to correct position
        if scheduler is not None and train_config.scheduler == "cosine":
            for _ in range(start_epoch):
                scheduler.step()
    
    # TensorBoard setup
    writer = None
    if train_config.use_tensorboard and TENSORBOARD_AVAILABLE:
        tb_dir = output_dir / "tensorboard"
        tb_dir.mkdir(parents=True, exist_ok=True)
        writer = SummaryWriter(log_dir=str(tb_dir))
        print(f"  TensorBoard: {tb_dir}")
        
        # Log model config as text
        config_text = json.dumps(model_config.__dict__, indent=2, default=str)
        writer.add_text("config/model", f"```json\n{config_text}\n```", 0)
        
        train_text = json.dumps(train_config.to_dict(), indent=2, default=str)
        writer.add_text("config/training", f"```json\n{train_text}\n```", 0)
    elif train_config.use_tensorboard:
        print("  [!] TensorBoard requested but not available")
    
    with open(output_dir / "model_config.json", "w") as f:
        json.dump(model_config.__dict__, f, indent=2, default=str)
    with open(output_dir / "train_config.json", "w") as f:
        json.dump(train_config.to_dict(), f, indent=2, default=str)
    
    print(f"\nOutput directory: {output_dir}")
    
    # Training loop
    print("\n" + "=" * 70)
    print("  TRAINING")
    print("=" * 70)
    
    if start_epoch > 0:
        print(f"  Resuming from epoch {start_epoch + 1}/{train_config.epochs}")
    
    for epoch in range(start_epoch, train_config.epochs):
        epoch_start = time.time()
        
        if epoch < train_config.warmup_epochs and train_config.scheduler == "cosine":
            warmup_lr = train_config.learning_rate * (epoch + 1) / train_config.warmup_epochs
            for param_group in optimizer.param_groups:
                param_group['lr'] = warmup_lr
        
        train_metrics = train_epoch(
            model, train_loader, optimizer, scheduler,
            train_config, epoch, tracker, writer
        )
        
        test_metrics = evaluate(model, test_loader, train_config)
        
        epoch_time = time.time() - epoch_start
        
        # TensorBoard epoch logging
        if writer is not None:
            # Epoch-level metrics
            writer.add_scalar('epoch/train_loss', train_metrics['loss'], epoch)
            writer.add_scalar('epoch/train_acc', train_metrics['acc'], epoch)
            writer.add_scalar('epoch/test_loss', test_metrics['loss'], epoch)
            writer.add_scalar('epoch/test_acc', test_metrics['acc'], epoch)
            writer.add_scalar('epoch/learning_rate', optimizer.param_groups[0]['lr'], epoch)
            writer.add_scalar('epoch/time_seconds', epoch_time, epoch)
            
            # Per-scale accuracies
            for scale in model.config.scales:
                writer.add_scalar(f'scales/acc_{scale}', test_metrics[f'acc_{scale}'], epoch)
            
            # Generalization gap
            writer.add_scalar('epoch/generalization_gap', test_metrics['acc'] - train_metrics['acc'], epoch)
            
            # Flush periodically
            if (epoch + 1) % 5 == 0:
                writer.flush()
        
        scale_accs = " | ".join([f"{s}:{test_metrics[f'acc_{s}']:.1f}%" for s in model.config.scales])
        star = "★" if test_metrics['acc'] > best_acc else ""
        
        print(f"  → Train: {train_metrics['acc']:.1f}% | Test: {test_metrics['acc']:.1f}% | "
              f"Scales: [{scale_accs}] | {epoch_time:.0f}s {star}")
        
        if test_metrics['acc'] > best_acc:
            best_acc = test_metrics['acc']
            
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_acc': best_acc,
                'model_config': model_config.__dict__,
                'train_config': train_config.to_dict()
            }, output_dir / "best_model.pt")
        
        if (epoch + 1) % train_config.save_interval == 0:
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_acc': best_acc
            }, output_dir / f"checkpoint_epoch_{epoch + 1}.pt")
            
            # Periodic HuggingFace Hub upload
            if train_config.push_to_hub and HF_HUB_AVAILABLE and effective_hub_repo_id:
                try:
                    # Save current best for upload
                    checkpoint = torch.load(output_dir / "best_model.pt", map_location='cpu')
                    model_cpu = DavidBeans(model_config)
                    model_cpu.load_state_dict(checkpoint['model_state_dict'])
                    
                    hub_dir = save_for_hub(
                        model=model_cpu,
                        model_config=model_config,
                        train_config=train_config,
                        best_acc=best_acc,
                        output_dir=output_dir,
                        training_history=tracker.get_history()
                    )
                    
                    push_to_hub(
                        hub_dir=hub_dir,
                        repo_id=effective_hub_repo_id,
                        private=train_config.hub_private,
                        commit_message=f"Checkpoint epoch {epoch + 1} - {best_acc:.2f}% acc"
                    )
                    print(f"  📤 Uploaded to {effective_hub_repo_id}")
                except Exception as e:
                    print(f"  [!] Hub upload failed: {e}")
        
        tracker.end_epoch()
    
    # Final summary
    print("\n" + "=" * 70)
    print("  TRAINING COMPLETE")
    print("=" * 70)
    print(f"\n  Best Test Accuracy: {best_acc:.2f}%")
    print(f"  Model saved to: {output_dir / 'best_model.pt'}")
    
    # Save training history
    history = tracker.get_history()
    with open(output_dir / "training_history.json", "w") as f:
        json.dump(history, f, indent=2)
    
    # Final TensorBoard logging
    if writer is not None:
        # Log best accuracy as hparam metric
        hparams = {
            'dim': model_config.dim,
            'num_layers': model_config.num_layers,
            'num_heads': model_config.num_heads,
            'num_experts': model_config.num_experts,
            'k_neighbors': model_config.k_neighbors,
            'cantor_weight': model_config.cantor_weight,
            'learning_rate': train_config.learning_rate,
            'weight_decay': train_config.weight_decay,
            'batch_size': train_config.batch_size,
            'mixup_alpha': train_config.mixup_alpha,
            'cutmix_alpha': train_config.cutmix_alpha,
        }
        writer.add_hparams(hparams, {'hparam/best_acc': best_acc})
        writer.add_scalar('final/best_acc', best_acc, 0)
        writer.close()
        print(f"  TensorBoard logs: {output_dir / 'tensorboard'}")
    
    # HuggingFace Hub upload
    if train_config.push_to_hub:
        print("\n" + "=" * 70)
        print("  UPLOADING TO HUGGINGFACE HUB")
        print("=" * 70)
        
        if not HF_HUB_AVAILABLE:
            print("  [!] huggingface_hub not installed. Skipping upload.")
        elif not effective_hub_repo_id:
            print("  [!] hub_repo_id not set. Skipping upload.")
        else:
            checkpoint = torch.load(output_dir / "best_model.pt", map_location='cpu')
            model.load_state_dict(checkpoint['model_state_dict'])
            
            print(f"\n  Preparing model for upload...")
            hub_dir = save_for_hub(
                model=model,
                model_config=model_config,
                train_config=train_config,
                best_acc=best_acc,
                output_dir=output_dir,
                training_history=history
            )
            
            print(f"\n  Uploading to {effective_hub_repo_id}...")
            push_to_hub(
                hub_dir=hub_dir,
                repo_id=effective_hub_repo_id,
                private=train_config.hub_private
            )
            
            print(f"\n  🎉 Model uploaded to: https://huggingface.co/{effective_hub_repo_id}")
    
    return model, best_acc


# ============================================================================
# PRESETS
# ============================================================================

def train_cifar10_small(run_name: str = "cifar10_small"):
    """Small model for CIFAR-10."""
    model_config = DavidBeansConfig(
        image_size=32, patch_size=4, dim=256, num_layers=4,
        num_heads=4, num_experts=5, k_neighbors=16,
        cantor_weight=0.3, scales=[64, 128, 256, 512],
        num_classes=10, dropout=0.1
    )
    
    train_config = TrainingConfig(
        run_name=run_name,
        dataset="cifar10", epochs=50, batch_size=128,
        learning_rate=1e-3, weight_decay=0.05, warmup_epochs=5,
        cayley_weight=0.01, contrast_weight=0.3,
        output_dir="./checkpoints/cifar10"
    )
    
    return train_david_beans(model_config, train_config)


def train_cifar100(
    run_name: str = "cifar100_base",
    push_to_hub: bool = False, 
    hub_repo_id: Optional[str] = None, 
    resume: bool = False
):
    """Model for CIFAR-100 with optional HF Hub upload and resume."""
    model_config = DavidBeansConfig(
        image_size=32, patch_size=4, dim=512, num_layers=8,
        num_heads=8, num_experts=5, k_neighbors=32,
        cantor_weight=0.3, scales=[256, 512, 768],
        num_classes=100, dropout=0.15
    )
    
    train_config = TrainingConfig(
        run_name=run_name,
        dataset="cifar100", epochs=200, batch_size=128,
        learning_rate=5e-4, weight_decay=0.1, warmup_epochs=20,
        cayley_weight=0.01, contrast_weight=0.5,
        label_smoothing=0.1, mixup_alpha=0.3, cutmix_alpha=1.0,
        output_dir="./checkpoints/cifar100",
        resume_from="latest" if resume else None,
        push_to_hub=push_to_hub, hub_repo_id=hub_repo_id, hub_private=False
    )
    
    return train_david_beans(model_config, train_config)


def resume_training(
    checkpoint_dir: str = "./checkpoints/cifar100",
    push_to_hub: bool = False,
    hub_repo_id: Optional[str] = None
):
    """
    Resume training from the latest checkpoint in a directory.
    
    Usage:
        resume_training("./checkpoints/cifar100", push_to_hub=True, hub_repo_id="user/repo")
    """
    output_dir = Path(checkpoint_dir)
    
    # Load configs from checkpoint directory
    model_config_path = output_dir / "model_config.json"
    train_config_path = output_dir / "train_config.json"
    
    if not model_config_path.exists():
        raise FileNotFoundError(f"No model_config.json in {output_dir}")
    
    with open(model_config_path) as f:
        model_config_dict = json.load(f)
    
    with open(train_config_path) as f:
        train_config_dict = json.load(f)
    
    # Handle tuple conversion for betas
    if 'betas' in train_config_dict and isinstance(train_config_dict['betas'], list):
        train_config_dict['betas'] = tuple(train_config_dict['betas'])
    
    model_config = DavidBeansConfig(**model_config_dict)
    train_config = TrainingConfig(**train_config_dict)
    
    # Override with resume settings
    train_config.resume_from = "latest"
    train_config.push_to_hub = push_to_hub
    if hub_repo_id:
        train_config.hub_repo_id = hub_repo_id
    
    return train_david_beans(model_config, train_config)


# ============================================================================
# STANDALONE UPLOAD FUNCTION
# ============================================================================

def upload_checkpoint(
    checkpoint_path: str,
    repo_id: str,
    best_acc: Optional[float] = None,
    private: bool = False
):
    """
    Upload an existing checkpoint to HuggingFace Hub.
    
    Usage:
        upload_checkpoint(
            checkpoint_path="./checkpoints/cifar100/best_model.pt",
            repo_id="AbstractPhil/david-beans-cifar100",
            best_acc=70.0  # Optional, will read from checkpoint if available
        )
    """
    if not HF_HUB_AVAILABLE:
        raise RuntimeError("huggingface_hub not installed. Run: pip install huggingface_hub")
    
    print(f"\n📦 Loading checkpoint: {checkpoint_path}")
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    
    # Reconstruct configs
    model_config_dict = checkpoint.get('model_config', {})
    train_config_dict = checkpoint.get('train_config', {})
    
    model_config = DavidBeansConfig(**model_config_dict)
    train_config = TrainingConfig(**train_config_dict)
    train_config.hub_repo_id = repo_id
    
    # Build model and load weights
    model = DavidBeans(model_config)
    model.load_state_dict(checkpoint['model_state_dict'])
    
    actual_best_acc = best_acc or checkpoint.get('best_acc', 0.0)
    
    # Prepare and upload
    output_dir = Path(checkpoint_path).parent
    
    print(f"\n📝 Preparing files for upload...")
    hub_dir = save_for_hub(
        model=model,
        model_config=model_config,
        train_config=train_config,
        best_acc=actual_best_acc,
        output_dir=output_dir
    )
    
    print(f"\n🚀 Uploading to {repo_id}...")
    push_to_hub(hub_dir, repo_id, private=private)
    
    print(f"\n🎉 Done! https://huggingface.co/{repo_id}")


# ============================================================================
# MAIN
# ============================================================================

if __name__ == "__main__":
    # =====================================================
    # CONFIGURATION
    # =====================================================
    
    PRESET = "cifar100"  # "test", "small", "cifar100", "resume"
    RESUME = False       # Set True to resume from latest checkpoint
    RUN_NAME = "5expert_3scale"  # Descriptive name for this run
    
    # HuggingFace Hub settings
    PUSH_TO_HUB = False
    HUB_REPO_ID = "AbstractPhil/geovit-david-beans"
    
    # =====================================================
    # RUN
    # =====================================================
    
    if PRESET == "test":
        print("🧪 Quick test...")
        model_config = DavidBeansConfig(
            image_size=32, patch_size=4, dim=128, num_layers=2,
            num_heads=4, num_experts=5, k_neighbors=8,
            scales=[32, 64, 128], num_classes=10
        )
        train_config = TrainingConfig(
            run_name="test",
            epochs=2, batch_size=32,
            use_augmentation=False, mixup_alpha=0.0, cutmix_alpha=0.0
        )
        model, acc = train_david_beans(model_config, train_config)
        
    elif PRESET == "small":
        print("🫘💎 Training DavidBeans - Small (CIFAR-10)...")
        model, acc = train_cifar10_small()
        
    elif PRESET == "cifar100":
        print("🫘💎 Training DavidBeans - CIFAR-100...")
        model, acc = train_cifar100(
            run_name=RUN_NAME,
            push_to_hub=PUSH_TO_HUB, 
            hub_repo_id=HUB_REPO_ID,
            resume=RESUME
        )
        
    elif PRESET == "resume":
        print("🔄 Resuming training from latest checkpoint...")
        model, acc = resume_training(
            checkpoint_dir="./checkpoints/cifar100",
            push_to_hub=PUSH_TO_HUB,
            hub_repo_id=HUB_REPO_ID
        )
        
    else:
        raise ValueError(f"Unknown preset: {PRESET}")
    
    print(f"\n🎉 Done! Best accuracy: {acc:.2f}%")