# =============================================================================
# CELL 3: Train Geometric Classifier + Upload to HuggingFace
# Requires: Cell 1 (generator/constants), Cell 2 (model classes)
# Outputs: `model` in notebook scope + geometric_classifier/ on HF
#
# Features:
#   - Dataset cached to disk (skip regeneration on resume)
#   - Checkpoint saved every epoch (model, optimizer, scheduler, epoch, best_acc)
#   - Auto-resume from latest checkpoint
# =============================================================================

import json, time, os, shutil
from pathlib import Path

HF_REPO = "AbstractPhil/grid-geometric-classifier-proto"
CKPT_DIR = Path("./checkpoints")
DATASET_PATH = Path("./cached_dataset.pt")

# --- Loss Functions ---

def _safe_bce(inp, tgt):
    """BCE that forces fp32 and clamps to prevent log(0) from BF16 sigmoid saturation."""
    with torch.amp.autocast('cuda', enabled=False):
        return F.binary_cross_entropy(
            inp.float().clamp(1e-7, 1 - 1e-7),
            tgt.float())

def capacity_fill_loss(fr, dt): return _safe_bce(fr, dt)
def overflow_reg(on, dt):
    """Vectorized overflow penalty — no Python loops, no .item() calls."""
    pk = dt.sum(dim=-1).long().clamp(min=0)  # (B,) peak dim index
    n_caps = on.shape[1]
    arange = torch.arange(n_caps, device=on.device).unsqueeze(0)  # (1, n_caps)
    mask = (arange >= pk.unsqueeze(1)).float()  # (B, n_caps)
    return (on * mask).sum() / (on.shape[0] + 1e-8)
def cap_diversity(c): return -c.var()
def peak_loss(l, t): return F.cross_entropy(l, t)
def cm_loss(p, t): return F.mse_loss(p, torch.sign(t))
def curved_bce(p, t): return _safe_bce(p.squeeze(-1), t)
def ctype_loss(l, t): return F.cross_entropy(l, t)

# --- Dataset Cache ---

def get_or_generate_dataset(n_samples, seed, path=DATASET_PATH):
    """Load cached dataset from disk, or generate + cache it."""
    if path.exists():
        print(f"Loading cached dataset from {path}...")
        t0 = time.time()
        cached = torch.load(path, weights_only=True)
        if cached["n_samples"] == n_samples and cached["seed"] == seed:
            train_ds = ShapeDataset.__new__(ShapeDataset)
            val_ds = ShapeDataset.__new__(ShapeDataset)
            for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]:
                setattr(train_ds, k, cached["train"][k])
                setattr(val_ds, k, cached["val"][k])
            dt = time.time() - t0
            print(f"Loaded {len(train_ds)} train + {len(val_ds)} val in {dt:.1f}s (cached)")
            return train_ds, val_ds
        else:
            print(f"Cache mismatch (n={cached['n_samples']}, seed={cached['seed']}) — regenerating")

    all_samples = generate_parallel(n_samples, seed=seed, n_workers=8)
    n_train = int(len(all_samples) * 0.8)
    train_ds = ShapeDataset(all_samples[:n_train])
    val_ds = ShapeDataset(all_samples[n_train:])

    print(f"Caching dataset to {path}...")
    cache_data = {
        "n_samples": n_samples, "seed": seed,
        "train": {k: getattr(train_ds, k) for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]},
        "val": {k: getattr(val_ds, k) for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]},
    }
    torch.save(cache_data, path)
    size_mb = path.stat().st_size / 1e6
    print(f"Cached: {size_mb:.0f}MB")
    return train_ds, val_ds

# --- Checkpoint helpers ---

def save_checkpoint(model, optimizer, scheduler, epoch, best_val_acc, ckpt_dir=CKPT_DIR):
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    raw = model._orig_mod if hasattr(model, '_orig_mod') else model
    ckpt = {
        "epoch": epoch,
        "best_val_acc": best_val_acc,
        "model_state_dict": raw.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict(),
    }
    path = ckpt_dir / f"epoch_{epoch:03d}.pt"
    torch.save(ckpt, path)
    latest = ckpt_dir / "latest.pt"
    torch.save(ckpt, latest)
    return path


def load_checkpoint(model, optimizer, scheduler, ckpt_dir=CKPT_DIR):
    latest = ckpt_dir / "latest.pt"
    if not latest.exists():
        return 0, 0.0
    print(f"Resuming from {latest}...")
    ckpt = torch.load(latest, weights_only=False)
    raw = model._orig_mod if hasattr(model, '_orig_mod') else model
    raw.load_state_dict(ckpt["model_state_dict"])
    optimizer.load_state_dict(ckpt["optimizer_state_dict"])
    scheduler.load_state_dict(ckpt["scheduler_state_dict"])
    start_epoch = ckpt["epoch"] + 1
    best_val_acc = ckpt["best_val_acc"]
    print(f"Resumed: epoch {start_epoch}, best_val_acc={best_val_acc:.4f}")
    return start_epoch, best_val_acc


# --- Training ---

def train(n_samples=500000, epochs=80, batch_size=4096, lr=3e-3, seed=42):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    if device.type == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        if hasattr(torch.backends.cuda.matmul, 'fp32_precision'):
            torch.backends.cuda.matmul.fp32_precision = 'tf32'
        if hasattr(torch.backends.cudnn, 'conv') and hasattr(torch.backends.cudnn.conv, 'fp32_precision'):
            torch.backends.cudnn.conv.fp32_precision = 'tf32'
        torch.backends.cudnn.benchmark = True
        props = torch.cuda.get_device_properties(0)
        print(f"GPU: {props.name} | {props.total_memory / 1e9:.1f}GB | SM {props.major}.{props.minor}")
        print(f"TF32: enabled | cuDNN benchmark: enabled | batch: {batch_size}")

    train_ds, val_ds = get_or_generate_dataset(n_samples, seed)
    print(f"Train: {len(train_ds)} | Val: {len(val_ds)} | {NUM_CLASSES} classes | pre-tensored")

    train_loader = torch.utils.data.DataLoader(
        train_ds, batch_size=batch_size, shuffle=True,
        num_workers=4, pin_memory=True, persistent_workers=True)
    val_loader = torch.utils.data.DataLoader(
        val_ds, batch_size=batch_size, shuffle=False,
        num_workers=4, pin_memory=True, persistent_workers=True)

    model = GeometricShapeClassifier().to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Model: {n_params:,} parameters")
    if device.type == "cuda":
        print(f"VRAM after model load: {torch.cuda.memory_allocated()/1e9:.2f}GB / "
              f"{torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB")

    use_amp = device.type == "cuda"
    amp_dtype = torch.bfloat16 if (device.type == "cuda" and
                torch.cuda.is_bf16_supported()) else torch.float16
    use_scaler = use_amp and amp_dtype == torch.float16
    scaler = torch.amp.GradScaler('cuda', enabled=use_scaler)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    warmup_epochs = 5
    def lr_lambda(epoch):
        if epoch < warmup_epochs:
            return (epoch + 1) / warmup_epochs
        return 0.5 * (1 + math.cos(math.pi * (epoch - warmup_epochs) / (epochs - warmup_epochs)))
    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    # Resume from checkpoint (loads into model BEFORE compile)
    start_epoch, best_val_acc = load_checkpoint(model, optimizer, scheduler)

    # Compile AFTER loading checkpoint weights
    if device.type == "cuda" and hasattr(torch, 'compile'):
        try:
            model = torch.compile(model, mode="default")
            print("torch.compile: enabled (default mode)")
        except Exception as e:
            print(f"torch.compile: skipped ({e})")

    print(f"AMP: {'bf16' if amp_dtype == torch.bfloat16 else 'fp16'}" +
          (f" (scaler: {'on' if use_scaler else 'off'})" if use_amp else " disabled"))

    w = {"cls": 1.0, "fill": 0.3, "peak": 0.3, "ovf": 0.05,
         "div": 0.02, "vol": 0.1, "cm": 0.1, "curved": 0.2, "ctype": 0.2,
         "arb_cls": 0.8, "arb_traj": 0.2, "arb_conf": 0.1, "flow": 0.5}

    epoch_start = time.time()

    for epoch in range(start_epoch, epochs):
        t0 = time.time()
        model.train()
        correct, total = 0, 0
        correct_init, correct_ref = 0, 0

        for batch_idx, (grid, label, dc, pd, vol, cm, ic, ct) in enumerate(train_loader):
            grid, label = grid.to(device, non_blocking=True), label.to(device, non_blocking=True)
            dc, pd = dc.to(device, non_blocking=True), pd.to(device, non_blocking=True)
            vol, cm = vol.to(device, non_blocking=True), cm.to(device, non_blocking=True)
            ic, ct = ic.to(device, non_blocking=True), ct.to(device, non_blocking=True)

            grid = deform_grid(grid, p_dropout=0.05, p_add=0.05, p_shift=0.08)
            optimizer.zero_grad(set_to_none=True)

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

                    loss_first = (w["cls"] * F.cross_entropy(out["initial_logits"], label) +
                                  w["fill"] * capacity_fill_loss(out["fill_ratios"], dc) +
                                  w["peak"] * peak_loss(out["peak_logits"], pd) +
                                  w["ovf"] * overflow_reg(out["overflows"], dc) +
                                  w["div"] * cap_diversity(out["capacities"]) +
                                  w["vol"] * F.mse_loss(out["volume_pred"], torch.log1p(vol)) +
                                  w["cm"] * cm_loss(out["cm_pred"], cm) +
                                  w["curved"] * curved_bce(out["is_curved_pred"], ic) +
                                  w["ctype"] * ctype_loss(out["curv_type_logits"], ct))

                    loss_arb = w["arb_cls"] * F.cross_entropy(out["refined_logits"], label)
                    traj_loss = 0
                    for step_i, step_logits in enumerate(out["trajectory_logits"]):
                        step_weight = (step_i + 1) / len(out["trajectory_logits"])
                        traj_loss += step_weight * F.cross_entropy(step_logits, label)
                    traj_loss /= len(out["trajectory_logits"])
                    loss_arb += w["arb_traj"] * traj_loss
                    loss_arb += w["flow"] * out["flow_loss"]

                    with torch.no_grad():
                        is_correct = (out["refined_logits"].argmax(1) == label).float()
                    loss_arb += w["arb_conf"] * _safe_bce(
                        out["refined_confidence"].squeeze(-1), is_correct)

                    with torch.no_grad():
                        init_correct = (out["initial_logits"].argmax(1) == label).float()
                        ref_correct = (out["refined_logits"].argmax(1) == label).float()
                        blend_target = torch.where(init_correct >= ref_correct,
                            torch.ones_like(init_correct) * 0.8,
                            torch.ones_like(init_correct) * 0.2)
                    loss_arb += 0.1 * _safe_bce(out["blend_weight"], blend_target)

                    loss_blend = w["cls"] * F.cross_entropy(out["class_logits"], label)
                    loss = loss_first + loss_arb + loss_blend

                # NaN guard: skip batch if loss is non-finite
                if not torch.isfinite(loss).item():
                    optimizer.zero_grad(set_to_none=True)
                    total += grid.size(0)
                    continue
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                scaler.step(optimizer)
                scaler.update()

            except RuntimeError as e:
                if "CUDA" in str(e) or "device-side" in str(e):
                    print(f"\n!!! CUDA error at epoch {epoch+1}, batch {batch_idx} !!!")
                    print(f"    Error: {e}")
                    print(f"    label range: [{label.min().item()}, {label.max().item()}]")
                    print(f"    pd range: [{pd.min().item()}, {pd.max().item()}]")
                    print(f"    ct range: [{ct.min().item()}, {ct.max().item()}]")
                    print(f"    Checkpoint saved at epoch {epoch-1}")
                    print(f"    To diagnose: add os.environ['CUDA_LAUNCH_BLOCKING']='1' before training")
                raise

            correct += (out["class_logits"].argmax(1) == label).sum().item()
            correct_init += (out["initial_logits"].argmax(1) == label).sum().item()
            correct_ref += (out["refined_logits"].argmax(1) == label).sum().item()
            total += grid.size(0)

        scheduler.step()
        train_acc = correct / total

        if epoch == start_epoch and device.type == "cuda":
            peak = torch.cuda.max_memory_allocated() / 1e9
            print(f"VRAM peak: {peak:.2f}GB | throughput: {total/(time.time()-t0):.0f} samples/s")

        model.eval()
        vc, vt, vcc, vct = 0, 0, 0, 0
        vc_init, vc_ref = 0, 0
        val_fills, val_alts, val_confs, val_blends = [], [], [], []

        with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
            for grid, label, dc, pd, vol, cm, ic, ct in val_loader:
                grid, label = grid.to(device, non_blocking=True), label.to(device, non_blocking=True)
                ic = ic.to(device, non_blocking=True)
                out = model(grid)
                vc += (out["class_logits"].argmax(1) == label).sum().item()
                vc_init += (out["initial_logits"].argmax(1) == label).sum().item()
                vc_ref += (out["refined_logits"].argmax(1) == label).sum().item()
                vt += grid.size(0)
                vcc += ((out["is_curved_pred"].squeeze(-1) > 0.5).float() == ic).sum().item()
                vct += grid.size(0)
                val_fills.append(out["fill_ratios"].cpu())
                val_alts.append(out["alternation"].cpu())
                val_confs.append(out["confidence"].cpu())
                val_blends.append(out["blend_weight"].cpu())

        val_acc = vc / vt; val_init = vc_init / vt; val_ref = vc_ref / vt
        curved_acc = vcc / vct
        mf = torch.cat(val_fills).mean(dim=0)
        mc = torch.cat(val_confs).mean().item()
        mb = torch.cat(val_blends).mean().item()
        marker = " *" if val_acc > best_val_acc else ""
        if val_acc > best_val_acc: best_val_acc = val_acc

        raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
        with torch.no_grad():
            caps = [F.softplus(getattr(raw_model, f"dim{d}")._raw_capacity).item() for d in range(4)]

        dt = time.time() - t0

        # Save checkpoint every epoch
        save_checkpoint(model, optimizer, scheduler, epoch, best_val_acc)

        if (epoch + 1) % 10 == 0 or epoch == start_epoch or marker:
            if (epoch + 1) % 10 == 0 or epoch == start_epoch:
                print(f"Epoch {epoch+1:3d}/{epochs} [{dt:.1f}s {total/dt:.0f} s/s] | "
                      f"blend {val_acc:.3f} init {val_init:.3f} arb {val_ref:.3f} | "
                      f"conf {mc:.3f} blend_w {mb:.2f} | curved {curved_acc:.3f} | "
                      f"fill [{mf[0]:.2f} {mf[1]:.2f} {mf[2]:.2f} {mf[3]:.2f}] | "
                      f"cap [{caps[0]:.2f} {caps[1]:.2f} {caps[2]:.2f} {caps[3]:.2f}]{marker}")
            elif marker:
                print(f"Epoch {epoch+1:3d}/{epochs} [{dt:.1f}s] | "
                      f"blend {val_acc:.3f} init {val_init:.3f} arb {val_ref:.3f} | conf {mc:.3f}{marker}")

    total_time = time.time() - epoch_start
    print(f"\nTraining complete in {total_time:.0f}s ({total_time/60:.1f}min)")
    print(f"Best val accuracy: {best_val_acc:.4f}")
    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model

    # --- Per-class breakdown ---
    cc_b, cc_i, cc_r = {n: 0 for n in CLASS_NAMES}, {n: 0 for n in CLASS_NAMES}, {n: 0 for n in CLASS_NAMES}
    ct_c = {n: 0 for n in CLASS_NAMES}
    cf = {n: [] for n in CLASS_NAMES}
    cconf = {n: [] for n in CLASS_NAMES}
    cblend = {n: [] for n in CLASS_NAMES}

    with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
        for grid, label, *_ in val_loader:
            grid, label = grid.to(device, non_blocking=True), label.to(device, non_blocking=True)
            out = model(grid)
            pb = out["class_logits"].argmax(1)
            pi = out["initial_logits"].argmax(1)
            pr = out["refined_logits"].argmax(1)
            for k in range(len(label)):
                name = CLASS_NAMES[label[k].item()]
                cc_b[name] += (pb[k] == label[k]).item()
                cc_i[name] += (pi[k] == label[k]).item()
                cc_r[name] += (pr[k] == label[k]).item()
                ct_c[name] += 1
                cf[name].append(out["fill_ratios"][k].cpu().numpy())
                cconf[name].append(out["confidence"][k].item())
                cblend[name].append(out["blend_weight"][k].item())

    print(f"\n{'Class':22s} | {'Blend':>5s} {'Init':>5s} {'Arb':>5s} | "
          f"{'Conf':>5s} {'Bld':>4s} | {'Corr':>4s}/{'Tot':>4s} | "
          f"{'Fill Ratios':22s} | {'Type':8s} Curvature")
    print("-" * 110)
    for name in CLASS_NAMES:
        if ct_c[name] == 0: continue
        ab = cc_b[name]/ct_c[name]; ai = cc_i[name]/ct_c[name]; ar = cc_r[name]/ct_c[name]
        mfv = np.mean(cf[name], axis=0)
        mconf = np.mean(cconf[name]); mblend = np.mean(cblend[name])
        info = SHAPE_CATALOG[name]
        arb_flag = f" +{ar-ai:+.3f}" if ar-ai > 0.01 else ""
        print(f"  {name:20s} | {ab:.3f} {ai:.3f} {ar:.3f} | "
              f"{mconf:.3f} {mblend:.2f} | {cc_b[name]:4d}/{ct_c[name]:4d} | "
              f"[{mfv[0]:.2f} {mfv[1]:.2f} {mfv[2]:.2f} {mfv[3]:.2f}] | "
              f"{'CURVED' if info['curved'] else 'rigid':8s} {info['curvature']}{arb_flag}")

    print(f"\n--- Arbiter Impact Summary ---")
    imps = [(n, cc_i[n]/ct_c[n], cc_r[n]/ct_c[n], cc_b[n]/ct_c[n], cc_r[n]/ct_c[n]-cc_i[n]/ct_c[n])
            for n in CLASS_NAMES if ct_c[n] > 0]
    imps.sort(key=lambda x: x[4], reverse=True)
    print(f"  {'Class':20s} | {'Init':>5s} {'Arb':>5s} {'Blend':>5s} | {'Δ':>6s}")
    for name, ai, ar, ab, delta in imps[:10]:
        print(f"  {name:20s} | {ai:.3f} {ar:.3f} {ab:.3f} | {delta:+.3f}")

    # =========================================================================
    # Upload geometric_classifier/ to HuggingFace
    # =========================================================================
    print("\n" + "=" * 70)
    print("Saving geometric_classifier/ to HuggingFace")
    print("=" * 70)

    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
    staging = Path("./hf_staging/geometric_classifier")
    staging.mkdir(parents=True, exist_ok=True)

    arch_config = {
        "model_type": "GeometricShapeClassifier",
        "version": "v8",
        "grid_size": GS,
        "num_classes": NUM_CLASSES,
        "class_names": CLASS_NAMES,
        "curvature_types": CURVATURE_TYPES,
        "embed_dim": 128,
        "n_tracers": 5,
        "capacity_dims": [64, 64, 64, 64],
        "curvature_embed_dim": 128,
        "arbiter_latent_dim": 128,
        "arbiter_flow_steps": 4,
        "total_params": sum(p.numel() for p in raw_model.parameters()),
        "shape_catalog": {k: v for k, v in SHAPE_CATALOG.items()},
    }
    with open(staging / "config.json", "w") as f:
        json.dump(arch_config, f, indent=2)

    train_cfg = {
        "n_samples": n_samples, "epochs": epochs, "batch_size": batch_size,
        "lr": lr, "seed": seed, "optimizer": "AdamW", "weight_decay": 1e-4,
        "scheduler": "cosine_with_warmup", "warmup_epochs": warmup_epochs,
        "amp_dtype": str(amp_dtype), "loss_weights": w,
        "best_val_accuracy": best_val_acc, "learned_capacities": caps,
        "total_training_time_seconds": total_time,
    }
    with open(staging / "training_config.json", "w") as f:
        json.dump(train_cfg, f, indent=2)

    try:
        from safetensors.torch import save_file as st_save
        st_save(raw_model.state_dict(), str(staging / "model.safetensors"))
        print(f"  Saved: model.safetensors")
    except ImportError:
        torch.save(raw_model.state_dict(), staging / "model.pt")
        print(f"  Saved: model.pt (install safetensors for .safetensors)")

    try:
        from huggingface_hub import HfApi, create_repo
        token = None
        try:
            from google.colab import userdata
            token = userdata.get('HF_TOKEN')
        except Exception:
            token = os.environ.get('HF_TOKEN')

        if token:
            api = HfApi(token=token)
            create_repo(HF_REPO, token=token, exist_ok=True)
            readme = Path("./hf_staging/README.md")
            readme.write_text(f"""---
license: mit
tags:
  - geometric-deep-learning
  - voxel-classifier
  - cross-contrast
  - pentachoron
---

# Grid Geometric Classifier Proto

Geometric primitive classifier using 5x5x5 binary voxel grids with capacity cascade,
curvature analysis, differentiation gates, and rectified flow arbiter.

## Structure

```
geometric_classifier/   # Voxel classifier (v8, ~1.85M params)
crosscontrast/           # Text-Voxel alignment heads
qwen_embeddings/         # Cached Qwen 2.5-1.5B embeddings
```

## 38 Shape Classes

Rigid 0D-3D: point, lines, triangles, quads, tetrahedra, cubes, prisms, octahedra, pentachoron
Curved 1D-3D: arcs, helices, circles, ellipses, discs, spheres, hemispheres, cylinders, cones, capsules, tori, shells, tubes, bowls, saddles
""")
            api.upload_file(path_or_fileobj=str(readme), path_in_repo="README.md",
                           repo_id=HF_REPO, token=token, commit_message="README")
            api.upload_folder(
                folder_path=str(staging), repo_id=HF_REPO,
                path_in_repo="geometric_classifier", token=token,
                commit_message=f"geometric_classifier v8 | acc={best_val_acc:.4f} | {sum(p.numel() for p in raw_model.parameters()):,} params")
            print(f"Uploaded: https://huggingface.co/{HF_REPO}/tree/main/geometric_classifier")
        else:
            print("No HF_TOKEN — saved locally at ./hf_staging/geometric_classifier/")
    except Exception as e:
        print(f"HF upload failed: {e}\n  Weights at ./hf_staging/geometric_classifier/")

    return model


model = train()