trainer.py

"""
ViT-Beatrix V5 - Contrarian Tower Collective
============================================

Architecture using geofractal router infrastructure with pos/neg tower pairs.

Key insights from V4 200-epoch run:
- λ converged to 0.217 ≈ 1/5 (structure ~15% of routing)
- patch_weight → -0.575 (emergent contrastive readout)
- Model naturally learned to subtract common-mode signal

V5 Design:
- Explicit pos/neg tower pairs (what V4 learned implicitly)
- WideRouter for parallel tower execution
- Contrastive fusion: pos_output - α * neg_output
- Cantor routing within each tower

Geofractal infrastructure:
- BaseTower: stages as nn.ModuleList
- WideRouter: discover_towers(), wide_forward()
- TorchComponent: for attention blocks
- FusionComponent pattern for contrastive fusion

COLAB SETUP:
------------
# Install geofractal first:
try:
    !pip uninstall -qy geofractal geometricvocab
except:
    pass
!pip install -q git+https://github.com/AbstractEyes/geofractal.git

Copyright 2025 AbstractPhil
Licensed under the Apache License, Version 2.0
"""

import math
from typing import Optional, Dict, List, Tuple
from dataclasses import dataclass
from datetime import datetime
import os

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.utils.tensorboard import SummaryWriter
from tqdm.auto import tqdm
from huggingface_hub import HfApi, upload_folder

# Geofractal imports
from geofractal.router.base_tower import BaseTower
from geofractal.router.wide_router import WideRouter
from geofractal.router.components.torch_component import TorchComponent


# =============================================================================
# CONFIGURATION
# =============================================================================

@dataclass
class BeatrixV5Config:
    image_size: int = 32
    patch_size: int = 4
    in_channels: int = 3
    embed_dim: int = 384
    depth: int = 6                    # Layers per tower
    num_heads: int = 6
    mlp_ratio: float = 4.0
    
    # Tower configuration
    num_tower_pairs: int = 2          # pos/neg pairs
    
    # Cantor routing (inherited from V4)
    cantor_levels: int = 5
    cantor_tau: float = 0.25
    routing_weight_init: float = 0.22  # Start near V4's converged value
    learnable_routing_weight: bool = True
    num_wormholes: int = 8
    wormhole_temperature: float = 0.1
    
    # Contrastive fusion
    contrastive_alpha_init: float = 0.5  # learnable neg contribution
    
    dropout: float = 0.1
    drop_path: float = 0.1
    num_classes: int = 100
    
    @property
    def num_patches(self) -> int:
        return (self.image_size // self.patch_size) ** 2
    
    @property
    def head_dim(self) -> int:
        return self.embed_dim // self.num_heads
    
    @property
    def num_towers(self) -> int:
        return self.num_tower_pairs * 2  # pos + neg for each pair
    
    def to_dict(self) -> dict:
        """Serialize config to dict for checkpoint saving."""
        return {
            'image_size': self.image_size,
            'patch_size': self.patch_size,
            'in_channels': self.in_channels,
            'embed_dim': self.embed_dim,
            'depth': self.depth,
            'num_heads': self.num_heads,
            'mlp_ratio': self.mlp_ratio,
            'num_tower_pairs': self.num_tower_pairs,
            'cantor_levels': self.cantor_levels,
            'cantor_tau': self.cantor_tau,
            'routing_weight_init': self.routing_weight_init,
            'learnable_routing_weight': self.learnable_routing_weight,
            'num_wormholes': self.num_wormholes,
            'wormhole_temperature': self.wormhole_temperature,
            'contrastive_alpha_init': self.contrastive_alpha_init,
            'dropout': self.dropout,
            'drop_path': self.drop_path,
            'num_classes': self.num_classes,
            'num_patches': self.num_patches,
            'num_towers': self.num_towers,
        }


# =============================================================================
# CANTOR STAIRCASE (from V4)
# =============================================================================

class BeatrixStaircase(nn.Module):
    """Cantor-based branch path encoding."""
    
    def __init__(self, levels: int = 5, tau: float = 0.25, alpha: float = 0.5):
        super().__init__()
        self.levels = levels
        self.tau = tau
        
        centers = torch.tensor([0.5, 1.5, 2.5], dtype=torch.float32)
        self.register_buffer('centers', centers)
        self.register_buffer('_alpha', torch.tensor(alpha))
        
        scales = 3.0 ** torch.arange(1, levels + 1, dtype=torch.float32)
        self.register_buffer('scales', scales)
        
        level_weights = 0.5 ** torch.arange(1, levels + 1, dtype=torch.float32)
        self.register_buffer('level_weights', level_weights)
    
    def forward(self, x):
        original_shape = x.shape
        x = x.clamp(1e-6, 1.0 - 1e-6)
        x_flat = x.reshape(-1)
        
        y = (x_flat.unsqueeze(-1) * self.scales) % 3
        d2 = (y.unsqueeze(-1) - self.centers) ** 2
        logits = -d2 / (self.tau + 1e-8)
        branch_path = logits.argmax(dim=-1)
        
        return branch_path.reshape(*original_shape, self.levels)


class HierarchicalRoutingBias(nn.Module):
    """Cantor-based routing bias for attention."""
    
    def __init__(
        self,
        num_positions: int,
        levels: int = 5,
        tau: float = 0.25,
        learnable_weight: bool = True,
        init_weight: float = 0.22,
    ):
        super().__init__()
        self.num_positions = num_positions
        self.levels = levels
        
        self.staircase = BeatrixStaircase(levels=levels, tau=tau)
        
        positions = torch.linspace(0, 1, num_positions)
        with torch.no_grad():
            branch_paths = self.staircase(positions)
        self.register_buffer('branch_paths', branch_paths)
        
        alignment = self._compute_alignment_matrix(branch_paths)
        self.register_buffer('alignment_matrix', alignment)
        
        if learnable_weight:
            self.routing_weight = nn.Parameter(torch.tensor(init_weight))
        else:
            self.register_buffer('routing_weight', torch.tensor(init_weight))
    
    def _compute_alignment_matrix(self, paths):
        P, L = paths.shape
        level_weights = 0.5 ** torch.arange(1, L + 1, device=paths.device)
        matches = (paths.unsqueeze(0) == paths.unsqueeze(1)).float()
        alignment = (matches * level_weights).sum(dim=-1)
        alignment.fill_diagonal_(0)
        return alignment
    
    def forward(self, content_scores):
        return content_scores + self.routing_weight * self.alignment_matrix
    
    def get_structure_only_scores(self, batch_size: int, device: torch.device):
        return self.alignment_matrix.unsqueeze(0).expand(batch_size, -1, -1)


# =============================================================================
# DROP PATH
# =============================================================================

class DropPath(nn.Module):
    def __init__(self, drop_prob: float = 0.0):
        super().__init__()
        self.drop_prob = drop_prob
    
    def forward(self, x):
        if self.drop_prob == 0.0 or not self.training:
            return x
        keep_prob = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
        random_tensor.floor_()
        return x.div(keep_prob) * random_tensor


# =============================================================================
# WORMHOLE ATTENTION
# =============================================================================

class WormholeAttention(nn.Module):
    """Attention with Cantor-based routing."""
    
    def __init__(
        self,
        dim: int,
        num_heads: int,
        num_patches: int,
        num_wormholes: int = 8,
        temperature: float = 0.1,
        routing_bias: Optional[HierarchicalRoutingBias] = None,
        dropout: float = 0.0,
        layer_idx: int = 0,
        num_layers: int = 6,
        inverted: bool = False,  # NEW: contrarian mode
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.num_patches = num_patches
        self.num_wormholes = min(num_wormholes, num_patches - 1)
        self.temperature = temperature
        self.routing_bias = routing_bias
        self.layer_idx = layer_idx
        self.is_final_layer = (layer_idx == num_layers - 1)
        self.inverted = inverted  # Contrarian tower uses inverted routing
        
        self.qkv = nn.Linear(dim, dim * 3)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(dropout)
        self.attn_drop = nn.Dropout(dropout)
        
        if not self.is_final_layer:
            self.route_q = nn.Linear(dim, dim)
            self.route_k = nn.Linear(dim, dim)
    
    def _compute_routes(self, x):
        B, S, D = x.shape
        P = self.num_patches
        K = self.num_wormholes
        
        x_patches = x[:, 1:, :]
        
        if self.is_final_layer:
            scores = self.routing_bias.get_structure_only_scores(B, x.device)
        else:
            q = F.normalize(self.route_q(x_patches), dim=-1)
            k = F.normalize(self.route_k(x_patches), dim=-1)
            content_scores = torch.bmm(q, k.transpose(1, 2))
            
            if self.routing_bias is not None:
                scores = self.routing_bias(content_scores)
            else:
                scores = content_scores
        
        # CONTRARIAN: invert routing scores
        if self.inverted:
            scores = -scores
        
        mask = torch.eye(P, device=x.device, dtype=torch.bool)
        scores = scores.masked_fill(mask.unsqueeze(0), -1e9)
        
        scores_scaled = scores / self.temperature
        topk_scores, routes = torch.topk(scores_scaled, K, dim=-1)
        weights = F.softmax(topk_scores, dim=-1)
        
        return routes, weights
    
    def _gather_wormhole(self, x, routes):
        B, H, P, D = x.shape
        K = routes.shape[-1]
        
        x_flat = x.reshape(B * H, P, D)
        routes_exp = routes.unsqueeze(1).expand(-1, H, -1, -1).reshape(B * H, P * K)
        routes_exp = routes_exp.unsqueeze(-1).expand(-1, -1, D)
        
        gathered = torch.gather(x_flat, 1, routes_exp)
        return gathered.view(B, H, P, K, D)
    
    def forward(self, x):
        B, S, D = x.shape
        H = self.num_heads
        P = self.num_patches
        head_dim = self.head_dim
        
        routes, route_weights = self._compute_routes(x)
        
        qkv = self.qkv(x).reshape(B, S, 3, H, head_dim).permute(2, 0, 3, 1, 4)
        Q, K_full, V = qkv.unbind(0)
        
        # CLS attention
        Q_cls = Q[:, :, :1, :]
        attn_cls = F.softmax(
            torch.einsum('bhqd,bhkd->bhqk', Q_cls, K_full) * self.scale,
            dim=-1
        )
        attn_cls = self.attn_drop(attn_cls)
        out_cls = torch.einsum('bhqk,bhkd->bhqd', attn_cls, V)
        
        # Patch attention with wormholes
        Q_patches = Q[:, :, 1:, :]
        K_patches = K_full[:, :, 1:, :]
        V_patches = V[:, :, 1:, :]
        
        K_gathered = self._gather_wormhole(K_patches, routes)
        V_gathered = self._gather_wormhole(V_patches, routes)
        
        scores_patches = torch.einsum('bhpd,bhpkd->bhpk', Q_patches, K_gathered) * self.scale
        scores_patches = scores_patches + route_weights.unsqueeze(1).log().clamp(min=-10)
        
        attn_patches = F.softmax(scores_patches, dim=-1)
        attn_patches = self.attn_drop(attn_patches)
        
        out_patches = torch.einsum('bhpk,bhpkd->bhpd', attn_patches, V_gathered)
        
        out = torch.cat([out_cls, out_patches], dim=2)
        out = out.transpose(1, 2).reshape(B, S, D)
        
        return self.proj_drop(self.proj(out))


# =============================================================================
# TRANSFORMER BLOCK (TorchComponent)
# =============================================================================

class BeatrixBlock(TorchComponent):
    """Transformer block as TorchComponent for proper stage registration."""
    
    def __init__(
        self,
        name: str,
        dim: int,
        num_heads: int,
        num_patches: int,
        num_wormholes: int = 8,
        mlp_ratio: float = 4.0,
        routing_bias: Optional[HierarchicalRoutingBias] = None,
        dropout: float = 0.0,
        drop_path: float = 0.0,
        layer_idx: int = 0,
        num_layers: int = 6,
        inverted: bool = False,
    ):
        super().__init__(name)
        
        self.norm1 = nn.LayerNorm(dim)
        self.attn = WormholeAttention(
            dim=dim, num_heads=num_heads, num_patches=num_patches,
            num_wormholes=num_wormholes, routing_bias=routing_bias,
            dropout=dropout, layer_idx=layer_idx, num_layers=num_layers,
            inverted=inverted,
        )
        
        self.norm2 = nn.LayerNorm(dim)
        mlp_hidden = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_hidden, dim),
            nn.Dropout(dropout),
        )
        
        self.drop_path = DropPath(drop_path) if drop_path > 0 else nn.Identity()
    
    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x)))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


# =============================================================================
# BEATRIX TOWER (BaseTower)
# =============================================================================

class BeatrixTower(BaseTower):
    """
    Single tower using geofractal BaseTower infrastructure.
    
    Uses:
    - self.append() to add stages
    - self.attach() for named components  
    - self.stages for iteration
    - self['name'] for component access
    
    Can be positive (normal) or negative (contrarian/inverted routing).
    """
    
    def __init__(
        self,
        name: str,
        config: BeatrixV5Config,
        inverted: bool = False,
    ):
        super().__init__(name, strict=False)
        self.inverted = inverted
        self._config = config
        
        # Shared routing bias (Cantor alignment matrix)
        self.attach('routing_bias', HierarchicalRoutingBias(
            num_positions=config.num_patches,
            levels=config.cantor_levels,
            tau=config.cantor_tau,
            learnable_weight=config.learnable_routing_weight,
            init_weight=config.routing_weight_init,
        ))
        
        # Stages via append() - geofractal pattern
        dpr = torch.linspace(0, config.drop_path, config.depth).tolist()
        for i in range(config.depth):
            self.append(BeatrixBlock(
                name=f'{name}_block_{i}',
                dim=config.embed_dim,
                num_heads=config.num_heads,
                num_patches=config.num_patches,
                num_wormholes=config.num_wormholes,
                mlp_ratio=config.mlp_ratio,
                routing_bias=self['routing_bias'],
                dropout=config.dropout,
                drop_path=dpr[i],
                layer_idx=i,
                num_layers=config.depth,
                inverted=inverted,
            ))
        
        # Named component via attach()
        self.attach('norm', nn.LayerNorm(config.embed_dim))
    
    def forward(self, x: Tensor) -> Tensor:
        """Process input and return opinion (CLS token)."""
        for stage in self.stages:
            x = stage(x)
        x = self['norm'](x)
        return x[:, 0]  # Return CLS token as opinion
    
    def get_routing_weight(self) -> float:
        return self['routing_bias'].routing_weight.item()


# =============================================================================
# CONTRASTIVE FUSION (TorchComponent)
# =============================================================================

class ContrastiveFusion(TorchComponent):
    """
    Fuses pos/neg tower pairs via learned contrastive combination.
    
    For each pair: output = pos + α * neg
    Where α is learnable and typically becomes negative (subtracting common-mode).
    
    This makes explicit what V4 learned implicitly with patch_weight.
    """
    
    def __init__(
        self,
        name: str,
        num_pairs: int,
        dim: int,
        alpha_init: float = 0.5,
    ):
        super().__init__(name)
        self.num_pairs = num_pairs
        
        # Per-pair learnable alpha (expect to go negative)
        self.alphas = nn.Parameter(torch.full((num_pairs,), alpha_init))
        
        # Final projection if multiple pairs
        if num_pairs > 1:
            self.pair_fusion = nn.Linear(dim * num_pairs, dim)
        else:
            self.pair_fusion = None
    
    def forward(self, pos_opinions: List[Tensor], neg_opinions: List[Tensor]) -> Tensor:
        """
        Args:
            pos_opinions: List of [B, D] tensors from positive towers
            neg_opinions: List of [B, D] tensors from negative towers
        Returns:
            Fused output [B, D]
        """
        assert len(pos_opinions) == len(neg_opinions) == self.num_pairs
        
        # Contrastive combination per pair
        fused_pairs = []
        for i, (pos, neg) in enumerate(zip(pos_opinions, neg_opinions)):
            # pos + α*neg where α learns to be negative
            fused = pos + self.alphas[i] * neg
            fused_pairs.append(fused)
        
        if self.pair_fusion is not None:
            # Concatenate and project
            combined = torch.cat(fused_pairs, dim=-1)
            return self.pair_fusion(combined)
        else:
            return fused_pairs[0]
    
    def get_alphas(self) -> List[float]:
        return self.alphas.tolist()


# =============================================================================
# WIDE ROUTER COLLECTIVE (WideRouter)
# =============================================================================

class EmbeddingParams(TorchComponent):
    """Wrapper for learnable embedding parameters."""
    def __init__(self, name: str, num_patches: int, embed_dim: int):
        super().__init__(name)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + num_patches, embed_dim))
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
    
    def forward(self, x: Tensor) -> Tensor:
        B = x.shape[0]
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)
        return x + self.pos_embed


class BeatrixCollective(WideRouter):
    """
    WideRouter collective managing pos/neg tower pairs.
    
    Follows geofractal WideRouter pattern:
    1. super().__init__(name, auto_discover=True)
    2. attach towers with self.attach(name, tower)
    3. call self.discover_towers() AFTER attaching
    4. wide_forward(x) returns Dict[tower_name, output]
    
    "Individual towers don't need to be accurate.
     They need to see differently.
     The routing fabric triangulates truth from divergent viewpoints."
    """
    
    def __init__(self, config: BeatrixV5Config):
        # auto_discover=True enables tower discovery
        super().__init__(name='beatrix_collective', auto_discover=True)
        self.config = config
        
        # Patch embedding (attached as component)
        self.attach('patch_embed', nn.Conv2d(
            config.in_channels, config.embed_dim,
            kernel_size=config.patch_size, stride=config.patch_size
        ))
        
        # Position/CLS embedding as TorchComponent (moves with .to())
        self.attach('embeddings', EmbeddingParams(
            'embeddings', config.num_patches, config.embed_dim
        ))
        self.attach('pos_drop', nn.Dropout(config.dropout))
        
        # Create tower pairs via attach() - towers inherit BaseTower
        for i in range(config.num_tower_pairs):
            pos_name = f'pos_{i}'
            neg_name = f'neg_{i}'
            self.attach(pos_name, BeatrixTower(pos_name, config, inverted=False))
            self.attach(neg_name, BeatrixTower(neg_name, config, inverted=True))
        
        # IMPORTANT: Call discover_towers() AFTER attaching all towers
        self.discover_towers()
        
        # Contrastive fusion
        self.attach('fusion', ContrastiveFusion(
            name='contrastive_fusion',
            num_pairs=config.num_tower_pairs,
            dim=config.embed_dim,
            alpha_init=config.contrastive_alpha_init,
        ))
        
        # Classification head
        self.attach('head', nn.Linear(config.embed_dim, config.num_classes))
        
        self._init_weights()
    
    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.LayerNorm):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
    
    def _prepare_input(self, images: Tensor) -> Tensor:
        """Shared input preparation: patch embed + pos embed."""
        # Patch embedding
        x = self['patch_embed'](images)
        x = x.flatten(2).transpose(1, 2)
        
        # Add CLS token and position embedding via TorchComponent
        x = self['embeddings'](x)
        x = self['pos_drop'](x)
        
        return x
    
    def forward(self, images: Tensor) -> Tensor:
        # Prepare shared input
        x = self._prepare_input(images)
        
        # wide_forward returns Dict[tower_name, output]
        opinions = self.wide_forward(x)
        
        # Separate pos/neg opinions using tower_names
        pos_opinions = []
        neg_opinions = []
        for i in range(self.config.num_tower_pairs):
            pos_opinions.append(opinions[f'pos_{i}'])
            neg_opinions.append(opinions[f'neg_{i}'])
        
        # Contrastive fusion
        fused = self['fusion'](pos_opinions, neg_opinions)
        
        # Classification
        return self['head'](fused)
    
    def get_diagnostics(self) -> Dict:
        """Get diagnostic info about tower states."""
        diag = {
            'fusion_alphas': self['fusion'].get_alphas(),
            'tower_lambdas': {},
        }
        for name in self.tower_names:
            diag['tower_lambdas'][name] = self[name].get_routing_weight()
        return diag


# =============================================================================
# MODEL FACTORY
# =============================================================================

def create_beatrix_v5_small(num_classes=100, **kwargs) -> BeatrixCollective:
    """Small model: 2 tower pairs, 384 dim, 6 depth."""
    config = BeatrixV5Config(
        embed_dim=384,
        depth=6,
        num_heads=6,
        num_tower_pairs=2,
        num_wormholes=8,
        num_classes=num_classes,
        **kwargs
    )
    return BeatrixCollective(config)


def create_beatrix_v5_base(num_classes=100, **kwargs) -> BeatrixCollective:
    """Base model: 2 tower pairs, 512 dim, 8 depth."""
    config = BeatrixV5Config(
        embed_dim=512,
        depth=8,
        num_heads=8,
        num_tower_pairs=2,
        num_wormholes=12,
        num_classes=num_classes,
        **kwargs
    )
    return BeatrixCollective(config)


def create_beatrix_v5_wide(num_classes=100, **kwargs) -> BeatrixCollective:
    """Wide model: 4 tower pairs, 384 dim, 4 depth."""
    config = BeatrixV5Config(
        embed_dim=512,
        depth=2,
        num_heads=8,
        num_tower_pairs=8,
        num_wormholes=32,
        num_classes=num_classes,
        patch_size=4,
        **kwargs
    )
    return BeatrixCollective(config)


# =============================================================================
# TRAINING UTILITIES
# =============================================================================

class CosineWarmupScheduler:
    def __init__(self, optimizer, warmup_epochs, total_epochs, min_lr=1e-6, base_lr=1e-3):
        self.optimizer = optimizer
        self.warmup_epochs = warmup_epochs
        self.total_epochs = total_epochs
        self.min_lr = min_lr
        self.base_lr = base_lr
    
    def step(self, epoch):
        if epoch < self.warmup_epochs:
            lr = self.base_lr * (epoch + 1) / self.warmup_epochs
        else:
            progress = (epoch - self.warmup_epochs) / (self.total_epochs - self.warmup_epochs)
            lr = self.min_lr + 0.5 * (self.base_lr - self.min_lr) * (1 + math.cos(math.pi * progress))
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
        return lr


def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss, correct, total = 0, 0, 0
    
    pbar = tqdm(loader, desc='Train', leave=False)
    for inputs, targets in pbar:
        inputs, targets = inputs.to(device), targets.to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
        # Update progress bar
        pbar.set_postfix({
            'loss': f'{loss.item():.3f}',
            'acc': f'{100.*correct/total:.1f}%'
        })
    
    return total_loss / len(loader), 100. * correct / total


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0, 0, 0
    
    pbar = tqdm(loader, desc='Eval', leave=False)
    for inputs, targets in pbar:
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
        pbar.set_postfix({'acc': f'{100.*correct/total:.1f}%'})
    
    return total_loss / len(loader), 100. * correct / total


def print_diagnostics(epoch: int, model: BeatrixCollective):
    diag = model.get_diagnostics()
    
    print(f"\n  ┌─ DIAGNOSTICS (Epoch {epoch}) ─────────────────────────────────────")
    print(f"  │ Fusion alphas (expect negative): {diag['fusion_alphas']}")
    print(f"  │ Tower routing weights (λ):")
    for name, lam in diag['tower_lambdas'].items():
        tower_type = "POS" if name.startswith('pos') else "NEG"
        print(f"  │   {name} ({tower_type}): {lam:.4f}")
    print(f"  └───────────────────────────────────────────────────────────────")


# =============================================================================
# MAIN - TRAINING
# =============================================================================

def main():
    import torchvision
    import torchvision.transforms as transforms
    
    # =========================================================================
    # CONFIGURATION
    # =========================================================================
    MODEL_TYPE = 'wide'  # 'small', 'base', or 'wide'
    EPOCHS = 100
    BASE_LR = 1e-3
    WARMUP_EPOCHS = 10
    BATCH_SIZE = 128
    # =========================================================================
    
    print("=" * 70)
    print("ViT-Beatrix V5 - CONTRARIAN TOWER COLLECTIVE")
    print("=" * 70)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"\nDevice: {device}")
    print(f"Model type: {MODEL_TYPE}")
    
    # Create model
    if MODEL_TYPE == 'small':
        model = create_beatrix_v5_small()
    elif MODEL_TYPE == 'base':
        model = create_beatrix_v5_base()
    elif MODEL_TYPE == 'wide':
        model = create_beatrix_v5_wide()
    else:
        raise ValueError(f"Unknown model type: {MODEL_TYPE}")
    
    # Move to device
    model = model.to(device)
    
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params:,}")
    print(f"Towers: {model.tower_names}")
    
    # Compile for performance (geofractal pattern)
    print("\nPreparing and compiling model...")
    torch.set_float32_matmul_precision('high')
    model_raw = model  # Keep reference for diagnostics
    model = model.prepare_and_compile()
    print("✓ Model compiled")
    
    # Data
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.AutoAugment(transforms.AutoAugmentPolicy.CIFAR10),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])
    
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])
    
    print("\nLoading CIFAR-100...")
    trainset = torchvision.datasets.CIFAR100(
        root='./data', train=True, download=True, transform=transform_train
    )
    testset = torchvision.datasets.CIFAR100(
        root='./data', train=False, download=True, transform=transform_test
    )
    
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True
    )
    testloader = torch.utils.data.DataLoader(
        testset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True
    )
    
    # Training setup
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=BASE_LR, weight_decay=0.05, betas=(0.9, 0.999)
    )
    scheduler = CosineWarmupScheduler(
        optimizer, warmup_epochs=WARMUP_EPOCHS, total_epochs=EPOCHS,
        min_lr=1e-6, base_lr=BASE_LR
    )
    
    # =========================================================================
    # HuggingFace and TensorBoard Setup
    # =========================================================================
    HF_REPO = "AbstractPhil/vit-beatrix-contrarian"
    CHECKPOINT_INTERVAL = 10  # Upload every N epochs
    
    # Create timestamp for this run
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_name = f"v5_{MODEL_TYPE}_{timestamp}"
    
    # Checkpoint directory
    checkpoint_dir = f"checkpoints/{run_name}"
    os.makedirs(checkpoint_dir, exist_ok=True)
    
    # TensorBoard writer
    tb_dir = f"{checkpoint_dir}/tensorboard"
    writer = SummaryWriter(tb_dir)
    
    # Log model config
    writer.add_text("config/model_type", MODEL_TYPE)
    writer.add_text("config/num_towers", str(len(model_raw.tower_names)))
    writer.add_text("config/total_params", f"{total_params:,}")
    
    # HuggingFace API - check/create repo
    hf_api = HfApi()
    try:
        hf_api.repo_info(repo_id=HF_REPO, repo_type="model")
        print(f"✓ HF repo exists: {HF_REPO}")
    except Exception:
        print(f"Creating HF repo: {HF_REPO}")
        try:
            hf_api.create_repo(repo_id=HF_REPO, repo_type="model", exist_ok=True)
            print(f"✓ Created HF repo: {HF_REPO}")
        except Exception as e:
            print(f"⚠️ Could not create repo: {e}")
    
    def save_best_locally(epoch, model_raw, history, diag, test_acc):
        """Save best checkpoint locally (no upload)."""
        ckpt_path = f"{checkpoint_dir}/{run_name}_best.pth"
        torch.save({
            'epoch': epoch,
            'model_state_dict': model_raw.state_dict(),
            'config': model_raw.config.to_dict(),
            'test_acc': test_acc,
            'history': history,
            'diagnostics': diag,
            'run_name': run_name,
            'timestamp': timestamp,
        }, ckpt_path)
        print(f"  💾 Saved best locally: {run_name}_best.pth")
    
    def save_interval_and_upload(epoch, model_raw, history, diag, test_acc):
        """Save interval checkpoint and upload everything to HuggingFace."""
        # Save interval checkpoint
        ckpt_name = f"{run_name}_e{epoch+1}.pth"
        ckpt_path = f"{checkpoint_dir}/{ckpt_name}"
        torch.save({
            'epoch': epoch,
            'model_state_dict': model_raw.state_dict(),
            'config': model_raw.config.to_dict(),
            'test_acc': test_acc,
            'history': history,
            'diagnostics': diag,
            'run_name': run_name,
            'timestamp': timestamp,
        }, ckpt_path)
        print(f"  💾 Saved interval: {ckpt_name}")
        
        # Update README
        readme_content = f"""# ViT-Beatrix V5 Contrarian Tower Collective

## Run: {run_name}

### Model Configuration
- **Type**: {MODEL_TYPE}
- **Total Parameters**: {total_params:,}
- **Towers**: {len(model_raw.tower_names)} ({model_raw.config.num_tower_pairs} pos/neg pairs)
- **Embed Dim**: {model_raw.config.embed_dim}
- **Depth**: {model_raw.config.depth} layers per tower

### Training Progress (Epoch {epoch+1})
- **Test Accuracy**: {test_acc:.2f}%
- **Best Accuracy**: {best_acc:.2f}%

### Files
- `{run_name}_best.pth` - Best checkpoint
- `{run_name}_e*.pth` - Interval checkpoints  
- `tensorboard/` - Training metrics

### Usage
```python
import torch
from vit_beatrix_v5_contrarian import BeatrixCollective, BeatrixV5Config

ckpt = torch.load("{run_name}_best.pth")
config = BeatrixV5Config(**ckpt['config'])
model = BeatrixCollective(config)
model.load_state_dict(ckpt['model_state_dict'])
```
"""
        with open(f"{checkpoint_dir}/README.md", 'w') as f:
            f.write(readme_content)
        
        # Upload entire folder (includes best, interval, tensorboard, readme)
        try:
            upload_folder(
                folder_path=checkpoint_dir,
                repo_id=HF_REPO,
                path_in_repo=run_name,
                repo_type="model",
            )
            print(f"  ☁️ Uploaded to {HF_REPO}/{run_name}")
        except Exception as e:
            print(f"  ⚠️ Upload failed: {e}")
    
    # =========================================================================
    
    # History tracking
    history = {
        'train_loss': [], 'train_acc': [], 'test_loss': [], 'test_acc': [],
        'fusion_alphas': [], 'tower_lambdas': [],
    }
    
    print("\n" + "=" * 70)
    print(f"Starting Training ({EPOCHS} epochs)")
    print(f"Run: {run_name}")
    print(f"Checkpoints: {checkpoint_dir}")
    print(f"HuggingFace: {HF_REPO}")
    print("=" * 70)
    
    best_acc = 0
    
    epoch_pbar = tqdm(range(EPOCHS), desc='Training')
    for epoch in epoch_pbar:
        lr = scheduler.step(epoch)
        
        train_loss, train_acc = train_epoch(model, trainloader, criterion, optimizer, device)
        test_loss, test_acc = evaluate(model, testloader, criterion, device)
        
        # Get diagnostics from raw model (compiled model may not expose methods)
        diag = model_raw.get_diagnostics()
        gap = train_acc - test_acc
        
        epoch_pbar.set_postfix({
            'test': f'{test_acc:.1f}%',
            'gap': f'{gap:.1f}%',
            'α': f'{diag["fusion_alphas"][0]:.2f}'
        })
        
        # TensorBoard logging
        writer.add_scalar('Loss/train', train_loss, epoch)
        writer.add_scalar('Loss/test', test_loss, epoch)
        writer.add_scalar('Accuracy/train', train_acc, epoch)
        writer.add_scalar('Accuracy/test', test_acc, epoch)
        writer.add_scalar('Accuracy/gap', gap, epoch)
        writer.add_scalar('LR', lr, epoch)
        
        # Log fusion alphas
        for i, alpha in enumerate(diag['fusion_alphas']):
            writer.add_scalar(f'Fusion/alpha_{i}', alpha, epoch)
        
        # Log tower lambdas
        for name, lam in diag['tower_lambdas'].items():
            writer.add_scalar(f'Lambda/{name}', lam, epoch)
        
        print(f"\nEpoch {epoch+1}/{EPOCHS} | LR: {lr:.6f}")
        print(f"  Train: {train_acc:.2f}% (loss={train_loss:.4f})")
        print(f"  Test:  {test_acc:.2f}% (loss={test_loss:.4f}) | Gap: {gap:.2f}%")
        print(f"  Fusion α: {diag['fusion_alphas'][:4]}{'...' if len(diag['fusion_alphas']) > 4 else ''}")
        
        # Diagnostics every 10 epochs or new best
        if (epoch + 1) % 10 == 0 or test_acc > best_acc:
            print_diagnostics(epoch + 1, model_raw)
        
        # Track history
        history['train_loss'].append(train_loss)
        history['train_acc'].append(train_acc)
        history['test_loss'].append(test_loss)
        history['test_acc'].append(test_acc)
        history['fusion_alphas'].append(diag['fusion_alphas'])
        history['tower_lambdas'].append(diag['tower_lambdas'])
        
        # Save best locally (no upload) every time we beat best
        if test_acc > best_acc:
            best_acc = test_acc
            print(f"  ★ New best: {best_acc:.2f}%")
            save_best_locally(epoch, model_raw, history, diag, test_acc)
        
        # Upload at intervals only (includes best + interval + tensorboard)
        if (epoch + 1) % CHECKPOINT_INTERVAL == 0:
            save_interval_and_upload(epoch, model_raw, history, diag, test_acc)
    
    # Final upload
    save_interval_and_upload(EPOCHS-1, model_raw, history, diag, test_acc)
    
    # Close TensorBoard writer
    writer.close()
    
    # Final summary
    print("\n" + "=" * 70)
    print(f"Training Complete!")
    print(f"Best accuracy: {best_acc:.2f}%")
    print(f"Checkpoints: {checkpoint_dir}")
    print("=" * 70)
    
    # Final diagnostics
    print_diagnostics(EPOCHS, model_raw)
    
    return model_raw, history


if __name__ == "__main__":
    main()