"""
Cell 3: Trainer — Patch Cross-Attention Shape Classifier (8×16×16)
===================================================================
Run after Cell 1 (generator) and Cell 2 (model).
Everything from prior cells is already in kernel scope.
"""

import os, time, math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader


# === Augmentation (vectorized for 8×16×16) ====================================

def deform_grid(grid, p_dropout=0.05, p_add=0.05, p_shift=0.08):
    """Vectorized voxel augmentation for 8×16×16 grids."""
    B = grid.shape[0]
    device = grid.device
    r = torch.rand(B, 3, device=device)
    out = grid.clone()

    # Voxel dropout
    drop_sel = (r[:, 0] < p_dropout).view(B, 1, 1, 1)
    keep = torch.rand_like(out) > 0.10
    out = torch.where(drop_sel, out * keep.float(), out)

    # Boundary addition
    add_sel = (r[:, 1] < p_add).view(B, 1, 1, 1).float()
    dilated = F.max_pool3d(out.unsqueeze(1), kernel_size=3, stride=1, padding=1).squeeze(1)
    boundary = ((dilated > 0.5) & (out < 0.5)).float()
    add_noise = (torch.rand_like(out) < 0.2).float()
    out = (out + boundary * add_noise * add_sel).clamp(max=1.0)

    # Small translation (1 voxel — grid is small)
    shift_sel = (r[:, 2] < p_shift)
    axes = torch.randint(3, (B,), device=device)
    dirs = torch.randint(0, 2, (B,), device=device) * 2 - 1

    versions = []
    for ax in range(3):
        for d in [-1, 1]:
            s = torch.roll(out, shifts=d, dims=ax + 1)
            if d == 1:
                if ax == 0: s[:, 0, :, :] = 0
                elif ax == 1: s[:, :, 0, :] = 0
                else: s[:, :, :, 0] = 0
            else:
                if ax == 0: s[:, -1, :, :] = 0
                elif ax == 1: s[:, :, -1, :] = 0
                else: s[:, :, :, -1] = 0
            versions.append(s)
    versions.append(out)  # identity
    stacked = torch.stack(versions, dim=0)  # (7, B, 8, 16, 16)

    assign = torch.where(shift_sel, axes * 2 + (dirs == 1).long(), torch.full_like(axes, 6))
    out = stacked[assign, torch.arange(B, device=device)]

    return out


# === Training =================================================================

def train_vae_ca_classifier(model, train_loader, val_loader,
                            n_epochs=60, lr=3e-3, weight_decay=1e-4,
                            grad_clip=1.0, device='cuda',
                            checkpoint_dir='/content/checkpoints_vae_ca'):
    device = torch.device(device if torch.cuda.is_available() else 'cpu')
    model = model.to(device)

    amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=n_epochs, eta_min=lr * 0.01)

    ce_loss_fn = nn.CrossEntropyLoss()
    bce_loss_fn = nn.BCEWithLogitsLoss()

    # From Cell 1 globals
    dim_labels = torch.tensor([CLASS_META[n]["dim"] for n in CLASS_NAMES], dtype=torch.long, device=device)
    curved_labels = torch.tensor([1.0 if CLASS_META[n]["curved"] else 0.0 for n in CLASS_NAMES], device=device)
    curv_type_labels = torch.tensor([CURV_TO_IDX[CLASS_META[n]["curvature"]] for n in CLASS_NAMES], dtype=torch.long, device=device)

    os.makedirs(checkpoint_dir, exist_ok=True)

    best_acc = 0.0
    print("=" * 70)

    for epoch in range(1, n_epochs + 1):
        model.train()
        t0 = time.time()
        total_loss = 0
        correct = 0
        total = 0

        for grid, label in train_loader:
            grid = grid.to(device, non_blocking=True)
            label = label.to(device, non_blocking=True)

            # Augmentation
            grid = deform_grid(grid)

            with torch.amp.autocast('cuda', dtype=amp_dtype):
                out = model(grid, labels=label)

                loss_cls = ce_loss_fn(out["class_logits"], label)
                batch_dims = dim_labels[label]
                loss_dim = ce_loss_fn(out["dim_logits"], batch_dims)
                batch_curved = curved_labels[label].unsqueeze(-1)
                loss_curved = bce_loss_fn(out["is_curved_pred"], batch_curved)
                batch_curv_type = curv_type_labels[label]
                loss_curv = ce_loss_fn(out["curv_type_logits"], batch_curv_type)

                loss = loss_cls + 0.3 * loss_dim + 0.3 * loss_curved + 0.2 * loss_curv

            optimizer.zero_grad(set_to_none=True)
            if amp_dtype == torch.float16:
                scaler = torch.amp.GradScaler('cuda')
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                scaler.step(optimizer)
                scaler.update()
            else:
                loss.backward()
                nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                optimizer.step()

            total_loss += loss.item()
            pred = out["class_logits"].argmax(dim=-1)
            correct += (pred == label).sum().item()
            total += label.shape[0]

        scheduler.step()

        # === Validation ===
        model.eval()
        val_correct = 0
        val_total = 0
        curved_correct = 0
        curved_total = 0
        per_class_correct = torch.zeros(NUM_CLASSES, device=device)
        per_class_total = torch.zeros(NUM_CLASSES, device=device)

        with torch.no_grad():
            for grid, label in val_loader:
                grid = grid.to(device, non_blocking=True)
                label = label.to(device, non_blocking=True)

                with torch.amp.autocast('cuda', dtype=amp_dtype):
                    out = model(grid)

                pred = out["class_logits"].argmax(dim=-1)
                val_correct += (pred == label).sum().item()
                val_total += label.shape[0]

                curved_pred = (out["is_curved_pred"].squeeze(-1) > 0.0).float()
                curved_true = curved_labels[label]
                curved_correct += (curved_pred == curved_true).sum().item()
                curved_total += label.shape[0]

                for c in range(NUM_CLASSES):
                    mask = label == c
                    per_class_total[c] += mask.sum()
                    per_class_correct[c] += (pred[mask] == c).sum()

        val_acc = val_correct / val_total
        curved_acc = curved_correct / curved_total
        train_acc = correct / total
        train_loss = total_loss / len(train_loader)
        elapsed = time.time() - t0
        sps = total / elapsed

        per_class_acc = per_class_correct / per_class_total.clamp(min=1)
        worst = per_class_acc.argsort()[:10]

        print(f"\nEpoch {epoch}/{n_epochs} | {elapsed:.1f}s | {sps:.0f} samp/s")
        print(f"  Train: loss={train_loss:.4f} acc={train_acc:.4f}")
        print(f"  Val:   acc={val_acc:.4f} curved={curved_acc:.4f}")
        print(f"  LR: {scheduler.get_last_lr()[0]:.6f}")
        print(f"  Worst classes:")
        for idx in worst:
            idx = idx.item()
            acc = per_class_acc[idx].item() * 100
            print(f"    {CLASS_NAMES[idx]:20s} {acc:5.1f}%")

        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(model.state_dict(), os.path.join(checkpoint_dir, 'best.pt'))
            # Also save to flat path for Cell 5 compat
            torch.save(model.state_dict(), '/content/best_vae_ca_classifier.pt')
            print(f"  ★ New best: {best_acc:.2%}")

        # Always save latest
        torch.save(model.state_dict(), os.path.join(checkpoint_dir, 'latest.pt'))

    print(f"\n{'=' * 70}")
    print(f"Training complete. Best val acc: {best_acc:.2%}")
    return best_acc


# === Run ======================================================================

print("=" * 70)
print(f"Patch Cross-Attention Shape Classifier — {GZ}×{GY}×{GX}")
print(f"  Patches: {MACRO_Z}×{MACRO_Y}×{MACRO_X} = {MACRO_N} of {PATCH_Z}×{PATCH_Y}×{PATCH_X}")
print("=" * 70)

print("\nGenerating training data...")
gen = ShapeGenerator(seed=42)
train_data = gen.generate_dataset(2000, seed=42)
val_data = gen.generate_dataset(400, seed=999)

print(f"  Generated {len(train_data['grids'])} train + {len(val_data['grids'])} val")
print(f"  Classes: {NUM_CLASSES}")
occ = train_data['grids'].reshape(len(train_data['grids']), -1).sum(axis=1)
print(f"  Avg occupied voxels: {occ.mean():.1f} / {GZ*GY*GX} ({occ.mean()/(GZ*GY*GX)*100:.1f}%)")

batch_size = 1024
train_ds = TensorDataset(torch.from_numpy(train_data['grids']).float(),
                         torch.from_numpy(train_data['labels']).long())
val_ds = TensorDataset(torch.from_numpy(val_data['grids']).float(),
                       torch.from_numpy(val_data['labels']).long())

train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
                          num_workers=2, pin_memory=True, drop_last=True)
val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,
                        num_workers=2, pin_memory=True)

print(f"  Batch size: {batch_size}")
print(f"  Train batches: {len(train_loader)} | Val batches: {len(val_loader)}")

model = PatchCrossAttentionClassifier(n_classes=NUM_CLASSES)
n_params = sum(p.numel() for p in model.parameters())
print(f"  Model: {n_params:,} params")

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"  Device: {device}")

best_acc = train_vae_ca_classifier(
    model, train_loader, val_loader,
    n_epochs=60, lr=3e-3, weight_decay=1e-4, device=device)

print(f"\nDone. Best accuracy: {best_acc:.2%}")