# =============================================================================
# CELL 4: Qwen × Geometric Classifier Cross-Contrast Training
# Requires: Cell 1 (constants), Cell 2 (model classes), Cell 3 (trained checkpoint)
# Outputs: crosscontrast/ and qwen_embeddings/ on HF
#
# Features:
#   - Loads geo classifier from Cell 3 checkpoint (no notebook scope dependency)
#   - Uses model.forward()["features"] (no duplicated internals)
#   - Dataset + Qwen embeddings cached to disk
#   - CC model checkpointed with resume
# =============================================================================

import math, time, json, os
from pathlib import Path
from tqdm.auto import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F

HF_REPO = "AbstractPhil/grid-geometric-classifier-proto"
CKPT_DIR = Path("./checkpoints")
CC_CKPT_DIR = Path("./cc_checkpoints")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True

use_amp = device.type == "cuda"
amp_dtype = torch.bfloat16 if (device.type == "cuda" and
            torch.cuda.is_bf16_supported()) else torch.float16


# =============================================================================
# Shape Descriptions
# =============================================================================

SHAPE_DESCRIPTIONS = {
    "point": "A zero-dimensional geometric primitive occupying a single discrete location in three-dimensional space with no extent along any axis.",
    "line_x": "A one-dimensional line segment extending along the horizontal x-axis, connecting two endpoints with uniform spacing between occupied voxels.",
    "line_y": "A one-dimensional line segment extending along the vertical y-axis, a straight structure rising upward through the grid.",
    "line_z": "A one-dimensional line segment extending along the depth z-axis, projecting straight backward through the voxel grid.",
    "line_diag": "A one-dimensional diagonal line segment cutting across multiple axes simultaneously, connecting opposite corners of the grid.",
    "cross": "Two perpendicular line segments intersecting at their midpoints forming a plus-shaped cross pattern in a single plane.",
    "l_shape": "Two connected line segments meeting at a right angle to form an L-shaped corner, like two edges of a rectangle.",
    "collinear": "Three or more points arranged along a single straight line with equal spacing, demonstrating perfect linear alignment.",
    "triangle_xy": "A flat triangular outline formed by three connected edges lying in the horizontal xy-plane, the simplest polygon.",
    "triangle_xz": "A flat triangular outline formed by three connected edges lying in the vertical xz-plane, a triangle standing upright.",
    "triangle_3d": "A triangular outline with vertices at different heights, forming a non-planar triangle tilted in three-dimensional space.",
    "square_xy": "A square outline formed by four equal edges in the xy-plane, a regular quadrilateral with right angles at each corner.",
    "square_xz": "A square outline formed by four equal edges in the xz-plane, a square standing vertically like a window frame.",
    "rectangle": "A rectangular outline with two pairs of parallel edges of different lengths, wider than it is tall.",
    "coplanar": "A set of points all lying in the same plane but not forming a regular polygon, a scattered planar arrangement.",
    "plane": "A solid flat surface filling an entire plane with occupied voxels, a two-dimensional sheet extending across the grid.",
    "tetrahedron": "A three-dimensional simplex with four triangular faces meeting at four vertices and six edges, the simplest polyhedron.",
    "pyramid": "A solid with a square base and four triangular faces converging to a single apex point above the base center.",
    "pentachoron": "A four-dimensional simplex projected into three dimensions, consisting of five tetrahedral cells sharing faces.",
    "cube": "A regular hexahedron with six identical square faces, twelve edges, and eight vertices forming a perfect box shape.",
    "cuboid": "A rectangular box with six rectangular faces, similar to a cube but with at least one pair of edges longer than the others.",
    "triangular_prism": "A solid with two parallel triangular faces connected by three rectangular faces, like a tent or Toblerone shape.",
    "octahedron": "A regular polyhedron with eight equilateral triangular faces, twelve edges, and six vertices, like two pyramids base-to-base.",
    "arc": "A curved one-dimensional segment forming part of a circle, a smooth bend connecting two endpoints along a circular path.",
    "helix": "A three-dimensional spiral curve that winds around a central axis while advancing along it, like a corkscrew or spring.",
    "circle": "A closed curved outline where every point is equidistant from the center, forming a perfect round ring in a plane.",
    "ellipse": "A closed curved outline forming an elongated circle, an oval shape with two focal points and varying curvature.",
    "disc": "A solid filled circular region, a flat round plate occupying all voxels within a circular boundary in a plane.",
    "sphere": "A perfectly round three-dimensional solid where every surface point is equidistant from the center, fully filled inside.",
    "hemisphere": "Half of a sphere cut along a great circle, a dome shape with a flat base and a convex curved upper surface.",
    "cylinder": "A solid with two parallel circular faces connected by a curved rectangular surface, like a can or pillar.",
    "cone": "A solid tapering smoothly from a circular base to a single apex point, with a curved surface of decreasing radius.",
    "capsule": "A cylinder capped with hemispheres at both ends, a smooth elongated pill shape with no sharp edges.",
    "torus": "A donut-shaped solid formed by revolving a circle around an external axis, with a hole through the center.",
    "shell": "A hollow spherical surface with no interior fill, an empty ball where only the outer boundary layer is occupied.",
    "tube": "A hollow cylindrical surface with no interior fill, an empty pipe where only the curved wall is occupied.",
    "bowl": "A concave open surface curving inward like a dish, the bottom half of a hollow sphere with the opening facing up.",
    "saddle": "A hyperbolic surface that curves upward along one axis and downward along the perpendicular axis, like a horse saddle.",
}
assert set(SHAPE_DESCRIPTIONS.keys()) == set(CLASS_NAMES), "Description/class mismatch!"


# =============================================================================
# Qwen Embedding Extractor
# =============================================================================

class QwenEmbeddingExtractor:
    MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
    HIDDEN_DIM = 1536

    def __init__(self, device="cuda"):
        self.device = device
        self.model = None
        self.tokenizer = None

    def load_model(self):
        from transformers import AutoModelForCausalLM, AutoTokenizer
        print(f"Loading {self.MODEL_NAME}...")
        self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_NAME, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.MODEL_NAME, dtype=torch.float16,
            device_map=self.device, trust_remote_code=True)
        self.model.eval()
        print(f"Qwen loaded: {self.HIDDEN_DIM}-dim hidden states")

    def _build_encode_prompt(self, description):
        messages = [
            {"role": "system", "content": "You are a geometric shape analyst."},
            {"role": "user", "content": f"Analyze this shape: {description}"},
        ]
        return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    @torch.no_grad()
    def extract_embedding(self, text):
        prompt = self._build_encode_prompt(text)
        inputs = self.tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(self.device)
        outputs = self.model(**inputs, output_hidden_states=True)
        hidden = outputs.hidden_states[-1]
        return hidden.mean(dim=1).squeeze(0).float()

    def cache_all_embeddings(self, class_names):
        print(f"Extracting embeddings for {len(class_names)} classes...")
        embeddings = {}
        for name in class_names:
            embeddings[name] = self.extract_embedding(SHAPE_DESCRIPTIONS[name])
        emb_tensor = torch.stack([embeddings[n] for n in class_names])
        normed = F.normalize(emb_tensor, dim=-1)
        sim = normed @ normed.T
        mean_sim = (sim.sum() - sim.trace()) / (len(class_names) * (len(class_names) - 1))
        print(f"Cached: {emb_tensor.shape} | mean cross-class sim: {mean_sim:.4f}")
        return emb_tensor

    def unload(self):
        del self.model, self.tokenizer
        self.model = self.tokenizer = None
        torch.cuda.empty_cache()
        print("Qwen unloaded")


# =============================================================================
# Projection Heads + Cross-Contrast Model
# =============================================================================

class TextProjection(nn.Module):
    def __init__(self, text_dim=1536, latent_dim=256):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(text_dim, latent_dim * 2), nn.GELU(),
            nn.Linear(latent_dim * 2, latent_dim), nn.GELU(),
            nn.Linear(latent_dim, latent_dim))
        self.norm = nn.LayerNorm(latent_dim)
    def forward(self, x): return self.norm(self.proj(x))

class VoxelProjection(nn.Module):
    def __init__(self, voxel_dim=645, latent_dim=256):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(voxel_dim, latent_dim * 2), nn.GELU(),
            nn.Linear(latent_dim * 2, latent_dim), nn.GELU(),
            nn.Linear(latent_dim, latent_dim))
        self.norm = nn.LayerNorm(latent_dim)
    def forward(self, x): return self.norm(self.proj(x))


class CrossContrastModel(nn.Module):
    def __init__(self, text_dim=1536, voxel_dim=645, latent_dim=256,
                 n_classes=38, temperature=0.07):
        super().__init__()
        self.text_proj = TextProjection(text_dim, latent_dim)
        self.voxel_proj = VoxelProjection(voxel_dim, latent_dim)
        self.log_temperature = nn.Parameter(torch.tensor(math.log(1.0 / temperature)))

    @property
    def temperature(self):
        return torch.exp(-self.log_temperature)

    def forward(self, voxel_features, class_labels, text_embeddings_table):
        text_emb = text_embeddings_table[class_labels]
        z_text = F.normalize(self.text_proj(text_emb), dim=-1)
        z_voxel = F.normalize(self.voxel_proj(voxel_features), dim=-1)

        temp = self.temperature
        logits_v2t = z_voxel @ z_text.T / temp
        logits_t2v = z_text @ z_voxel.T / temp

        labels_matrix = (class_labels.unsqueeze(0) == class_labels.unsqueeze(1)).float()
        labels_matrix = labels_matrix / labels_matrix.sum(dim=1, keepdim=True).clamp(min=1)

        loss_v2t = (-labels_matrix * F.log_softmax(logits_v2t, dim=1)).sum(dim=1).mean()
        loss_t2v = (-labels_matrix * F.log_softmax(logits_t2v, dim=1)).sum(dim=1).mean()
        loss = (loss_v2t + loss_t2v) / 2.0

        with torch.no_grad():
            v2t_preds = logits_v2t.argmax(dim=1)
            pred_classes = class_labels[v2t_preds]
            acc = (pred_classes == class_labels).float().mean()
            pos_sim = (z_voxel * z_text).sum(dim=-1).mean()
            neg_mask = ~(class_labels.unsqueeze(0) == class_labels.unsqueeze(1))
            neg_sim = (z_voxel @ z_text.T)[neg_mask].mean() if neg_mask.any() else torch.tensor(0.0)

        return loss, {"acc": acc.item(), "pos_sim": pos_sim.item(),
                      "neg_sim": neg_sim.item(), "temperature": temp.item()}


# =============================================================================
# Helpers
# =============================================================================

def get_hf_token():
    try:
        from google.colab import userdata
        return userdata.get('HF_TOKEN')
    except Exception:
        return os.environ.get('HF_TOKEN')


def load_geo_model(ckpt_dir=CKPT_DIR):
    """Load trained GeometricShapeClassifier from Cell 3 checkpoint."""
    latest = ckpt_dir / "latest.pt"
    if not latest.exists():
        raise FileNotFoundError(
            f"No checkpoint at {latest}. Run Cell 3 first to train the classifier.")
    print(f"Loading geo classifier from {latest}...")
    ckpt = torch.load(latest, weights_only=False, map_location=device)
    geo = GeometricShapeClassifier().to(device)
    geo.load_state_dict(ckpt["model_state_dict"])
    geo.eval()
    for p in geo.parameters():
        p.requires_grad = False
    print(f"Loaded: epoch {ckpt['epoch']}, val_acc={ckpt['best_val_acc']:.4f}, "
          f"{sum(p.numel() for p in geo.parameters()):,} params (frozen)")
    return geo


@torch.no_grad()
def extract_voxel_features(geo_model, grid):
    """Extract pre-classifier features using model.forward()['features']."""
    with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
        out = geo_model(grid)
    return out["features"].float()


def save_cc_checkpoint(cc_model, cc_opt, cc_sched, epoch, best_acc, ckpt_dir=CC_CKPT_DIR):
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    ckpt = {
        "epoch": epoch,
        "best_cc_acc": best_acc,
        "cc_model_state_dict": cc_model.state_dict(),
        "cc_optimizer_state_dict": cc_opt.state_dict(),
        "cc_scheduler_state_dict": cc_sched.state_dict(),
    }
    torch.save(ckpt, ckpt_dir / "latest.pt")


def load_cc_checkpoint(cc_model, cc_opt, cc_sched, ckpt_dir=CC_CKPT_DIR):
    latest = ckpt_dir / "latest.pt"
    if not latest.exists():
        return 0, 0.0
    print(f"Resuming CC from {latest}...")
    ckpt = torch.load(latest, weights_only=False)
    cc_model.load_state_dict(ckpt["cc_model_state_dict"])
    cc_opt.load_state_dict(ckpt["cc_optimizer_state_dict"])
    cc_sched.load_state_dict(ckpt["cc_scheduler_state_dict"])
    start = ckpt["epoch"] + 1
    best = ckpt["best_cc_acc"]
    print(f"Resumed: epoch {start}, best_cc_acc={best:.4f}")
    return start, best


def upload_cc_to_hf(cc_model, best_acc, epoch, token, reason="periodic",
                    text_dim_=None, voxel_dim_=None, latent_dim_=None):
    """Upload crosscontrast weights + config to HF. Called mid-training and at end."""
    if not token:
        return
    try:
        from huggingface_hub import HfApi, create_repo
        from safetensors.torch import save_file as st_save

        staging = Path("./hf_staging/crosscontrast")
        staging.mkdir(parents=True, exist_ok=True)

        st_save(cc_model.text_proj.state_dict(), str(staging / "text_proj.safetensors"))
        st_save(cc_model.voxel_proj.state_dict(), str(staging / "voxel_proj.safetensors"))
        st_save({"log_temperature": cc_model.log_temperature.data.unsqueeze(0)},
                str(staging / "temperature.safetensors"))

        # Write config so uploads are self-contained
        if text_dim_ and voxel_dim_ and latent_dim_:
            cfg = {
                "model_type": "CrossContrastModel",
                "text_dim": text_dim_, "voxel_dim": voxel_dim_,
                "latent_dim": latent_dim_,
                "best_val_accuracy": best_acc,
                "epoch": epoch + 1, "upload_reason": reason,
                "temperature": cc_model.temperature.item(),
            }
            with open(staging / "config.json", "w") as f:
                json.dump(cfg, f, indent=2)

        api = HfApi(token=token)
        create_repo(HF_REPO, token=token, exist_ok=True)
        api.upload_folder(
            folder_path=str(staging), repo_id=HF_REPO,
            path_in_repo="crosscontrast", token=token,
            commit_message=f"crosscontrast ep{epoch+1} | acc={best_acc:.4f} | {reason}")
        tqdm.write(f"  ✓ HF upload ({reason}): ep{epoch+1} acc={best_acc:.4f}")
    except Exception as e:
        tqdm.write(f"  ⚠ HF upload failed: {e}")


# =============================================================================
# PHASE 1: Qwen Embeddings (cached)
# =============================================================================

print("=" * 70)
print("Phase 1: Qwen Embeddings")
print("=" * 70)

cache_path = Path("qwen_geo_cache.pt")
if cache_path.exists():
    _cache = torch.load(cache_path, map_location="cpu", weights_only=True)
    text_embeddings = _cache["embeddings"]
    print(f"Loaded cached: {text_embeddings.shape}")
else:
    extractor = QwenEmbeddingExtractor(device=str(device))
    extractor.load_model()
    text_embeddings = extractor.cache_all_embeddings(CLASS_NAMES)
    torch.save({"embeddings": text_embeddings.cpu(), "class_names": CLASS_NAMES}, cache_path)
    extractor.unload()

text_embeddings = text_embeddings.to(device)
text_dim = text_embeddings.shape[1]

# Upload qwen_embeddings/ to HF
hf_token = get_hf_token()
qwen_staging = Path("./hf_staging/qwen_embeddings")
qwen_staging.mkdir(parents=True, exist_ok=True)

qwen_config = {
    "model_name": QwenEmbeddingExtractor.MODEL_NAME,
    "hidden_dim": QwenEmbeddingExtractor.HIDDEN_DIM,
    "extraction_method": "mean_pool_last_layer",
    "prompt_style": "2shot_geometric",
    "num_classes": NUM_CLASSES,
    "class_names": CLASS_NAMES,
    "embedding_shape": list(text_embeddings.shape),
}
with open(qwen_staging / "config.json", "w") as f:
    json.dump(qwen_config, f, indent=2)
with open(qwen_staging / "descriptions.json", "w") as f:
    json.dump(SHAPE_DESCRIPTIONS, f, indent=2)

try:
    from safetensors.torch import save_file as st_save
    st_save({"embeddings": text_embeddings.cpu()}, str(qwen_staging / "embeddings.safetensors"))
except ImportError:
    torch.save({"embeddings": text_embeddings.cpu()}, qwen_staging / "embeddings.pt")

if hf_token:
    try:
        from huggingface_hub import HfApi, create_repo
        api = HfApi(token=hf_token)
        create_repo(HF_REPO, token=hf_token, exist_ok=True)
        api.upload_folder(
            folder_path=str(qwen_staging), repo_id=HF_REPO,
            path_in_repo="qwen_embeddings", token=hf_token,
            commit_message=f"qwen_embeddings | {QwenEmbeddingExtractor.MODEL_NAME} | {NUM_CLASSES} classes")
        print(f"Uploaded: https://huggingface.co/{HF_REPO}/tree/main/qwen_embeddings")
    except Exception as e:
        print(f"Upload failed: {e}")
else:
    print("No HF_TOKEN — saved locally")


# =============================================================================
# PHASE 2: Load Geo Model + Dataset
# =============================================================================

print("\n" + "=" * 70)
print("Phase 2: Geo Model + Voxel Data")
print("=" * 70)

geo_model = load_geo_model()

CC_SAMPLES = 500000
CC_BATCH = 4096
CC_EPOCHS = 40
CC_LR = 2e-3
CC_LATENT = 256

# Reuse Cell 3's cached dataset
DATASET_PATH = Path("./cached_dataset.pt")
if not DATASET_PATH.exists():
    raise FileNotFoundError("No cached_dataset.pt — run Cell 3 first.")

print(f"Loading dataset from {DATASET_PATH}...")
_cached = torch.load(DATASET_PATH, weights_only=True)
cc_train_ds = ShapeDataset.__new__(ShapeDataset)
cc_val_ds = ShapeDataset.__new__(ShapeDataset)
for k in ["grids", "labels", "dim_conf", "peak_dim", "volume", "cm_det", "is_curved", "curvature"]:
    setattr(cc_train_ds, k, _cached["train"][k])
    setattr(cc_val_ds, k, _cached["val"][k])
print(f"Loaded {len(cc_train_ds)} train + {len(cc_val_ds)} val (from Cell 3 cache)")

cc_train_loader = torch.utils.data.DataLoader(
    cc_train_ds, batch_size=CC_BATCH, shuffle=True,
    num_workers=4, pin_memory=True, persistent_workers=True)
cc_val_loader = torch.utils.data.DataLoader(
    cc_val_ds, batch_size=CC_BATCH, shuffle=False,
    num_workers=4, pin_memory=True, persistent_workers=True)

# Probe voxel feature dim
with torch.no_grad():
    _dummy = torch.zeros(1, GS, GS, GS, device=device)
    voxel_dim = extract_voxel_features(geo_model, _dummy).shape[1]
print(f"Voxel dim: {voxel_dim} | Text dim: {text_dim}")


# =============================================================================
# PHASE 3: Cross-Contrast Training
# =============================================================================

print("\n" + "=" * 70)
print("Phase 3: Cross-Contrast Training")
print("=" * 70)

cc_model = CrossContrastModel(
    text_dim=text_dim, voxel_dim=voxel_dim,
    latent_dim=CC_LATENT, n_classes=NUM_CLASSES
).to(device)
cc_params = sum(p.numel() for p in cc_model.parameters())
print(f"CrossContrast: {cc_params:,} params | latent={CC_LATENT}")

cc_opt = torch.optim.AdamW(cc_model.parameters(), lr=CC_LR, weight_decay=1e-4)
_warmup = 3
def _cc_lr(ep):
    if ep < _warmup: return (ep + 1) / _warmup
    return 0.5 * (1 + math.cos(math.pi * (ep - _warmup) / (CC_EPOCHS - _warmup)))
cc_sched = torch.optim.lr_scheduler.LambdaLR(cc_opt, _cc_lr)

# Resume CC training if checkpoint exists
cc_start, best_cc_acc = load_cc_checkpoint(cc_model, cc_opt, cc_sched)

t_start = time.time()
epoch_bar = tqdm(range(cc_start, CC_EPOCHS), desc="Training", unit="ep")

for epoch in epoch_bar:
    t0 = time.time()
    cc_model.train()
    tot_loss, tot_acc, nb = 0, 0, 0

    batch_bar = tqdm(cc_train_loader, desc=f"Ep {epoch+1}/{CC_EPOCHS} train",
                     leave=False, unit="batch")
    for grid, label, *_ in batch_bar:
        grid = grid.to(device, non_blocking=True)
        label = label.to(device, non_blocking=True)
        vf = extract_voxel_features(geo_model, grid)

        cc_opt.zero_grad(set_to_none=True)
        with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
            loss, metrics = cc_model(vf, label, text_embeddings)

        loss.backward()
        torch.nn.utils.clip_grad_norm_(cc_model.parameters(), 1.0)
        cc_opt.step()
        tot_loss += loss.item(); tot_acc += metrics["acc"]; nb += 1
        batch_bar.set_postfix(loss=f"{loss.item():.4f}", acc=f"{metrics['acc']:.3f}")

    cc_sched.step()

    cc_model.eval()
    vl, va, vps, vns, vnb = 0, 0, 0, 0, 0
    with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
        for grid, label, *_ in tqdm(cc_val_loader, desc=f"Ep {epoch+1}/{CC_EPOCHS} val",
                                     leave=False, unit="batch"):
            grid = grid.to(device, non_blocking=True)
            label = label.to(device, non_blocking=True)
            vf = extract_voxel_features(geo_model, grid)
            loss, m = cc_model(vf, label, text_embeddings)
            vl += loss.item(); va += m["acc"]; vps += m["pos_sim"]; vns += m["neg_sim"]; vnb += 1

    tl = tot_loss/max(nb,1); ta = tot_acc/max(nb,1)
    vl_ = vl/max(vnb,1); va_ = va/max(vnb,1)
    ps = vps/max(vnb,1); ns = vns/max(vnb,1)
    temp = cc_model.temperature.item()
    dt = time.time() - t0
    mk = " *" if va_ > best_cc_acc else ""
    if va_ > best_cc_acc: best_cc_acc = va_

    save_cc_checkpoint(cc_model, cc_opt, cc_sched, epoch, best_cc_acc)

    # Upload to HF on new best or every 10 epochs
    is_new_best = mk == " *"
    periodic = (epoch + 1) % 10 == 0
    if is_new_best:
        upload_cc_to_hf(cc_model, best_cc_acc, epoch, hf_token, reason="new_best",
                        text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT)
    elif periodic:
        upload_cc_to_hf(cc_model, best_cc_acc, epoch, hf_token, reason="periodic",
                        text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT)

    epoch_bar.set_postfix(loss=f"{vl_:.4f}", acc=f"{va_:.3f}", best=f"{best_cc_acc:.3f}",
                          tau=f"{temp:.4f}")

    if (epoch+1) % 5 == 0 or epoch == cc_start or mk:
        tqdm.write(f"Ep {epoch+1:3d}/{CC_EPOCHS} [{dt:.1f}s] | "
                   f"loss {tl:.4f}/{vl_:.4f} | acc {ta:.3f}/{va_:.3f} | "
                   f"pos {ps:.3f} neg {ns:.3f} | τ {temp:.4f}{mk}")

    if epoch == cc_start and device.type == "cuda":
        tqdm.write(f"VRAM peak: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")

tt = time.time() - t_start
print(f"\nDone in {tt:.0f}s ({tt/60:.1f}min) | Best acc: {best_cc_acc:.4f}")


# =============================================================================
# Per-Class Alignment Analysis
# =============================================================================

print("\n" + "=" * 70)
print("Per-Class Alignment")
print("=" * 70)

cc_model.eval()
cls_vox = {n: [] for n in CLASS_NAMES}
with torch.no_grad():
    text_proj_all = F.normalize(cc_model.text_proj(text_embeddings), dim=-1)

with torch.no_grad(), torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
    for grid, label, *_ in cc_val_loader:
        grid = grid.to(device, non_blocking=True)
        label = label.to(device, non_blocking=True)
        vf = extract_voxel_features(geo_model, grid)
        zv = F.normalize(cc_model.voxel_proj(vf), dim=-1)
        for k in range(len(label)):
            cls_vox[CLASS_NAMES[label[k].item()]].append(zv[k].cpu())

print(f"\n{'Class':22s} | {'Align':>6s} | {'N':>5s} | {'Nearest':22s} | {'OK':>3s}")
print("-" * 70)
correct = 0; total_c = 0
for name in CLASS_NAMES:
    if not cls_vox[name]: continue
    mv = F.normalize(torch.stack(cls_vox[name]).mean(dim=0), dim=-1)
    own = text_proj_all[CLASS_NAMES.index(name)].cpu()
    align = (mv * own).sum().item()
    sims = (mv.unsqueeze(0) @ text_proj_all.cpu().T).squeeze(0)
    ni = sims.argmax().item(); nn_ = CLASS_NAMES[ni]
    ok = "Y" if nn_ == name else f"X->{nn_}"
    if nn_ == name: correct += 1
    total_c += 1
    print(f"  {name:20s} | {align:.4f} | {len(cls_vox[name]):5d} | {nn_:22s} | {ok}")
print(f"\nNearest-text accuracy: {correct}/{total_c} = {correct/max(total_c,1):.1%}")

print("\nTop 10 Text-Space Confusions:")
with torch.no_grad():
    sim = (text_proj_all @ text_proj_all.T).cpu().numpy()
import numpy as np
confusions = []
for i in range(len(CLASS_NAMES)):
    for j in range(i+1, len(CLASS_NAMES)):
        confusions.append((CLASS_NAMES[i], CLASS_NAMES[j], sim[i,j]))
confusions.sort(key=lambda x: x[2], reverse=True)
for a, b, s in confusions[:10]:
    print(f"  {a:20s} <-> {b:20s} | {s:.4f}")


# =============================================================================
# Upload crosscontrast/ to HuggingFace (final — with full configs)
# =============================================================================

print("\n" + "=" * 70)
print("Saving crosscontrast/ to HuggingFace")
print("=" * 70)

cc_staging = Path("./hf_staging/crosscontrast")
cc_staging.mkdir(parents=True, exist_ok=True)

# Write detailed architecture config (overwrites minimal mid-training config)
cc_arch = {
    "model_type": "CrossContrastModel",
    "text_dim": text_dim,
    "voxel_dim": voxel_dim,
    "latent_dim": CC_LATENT,
    "n_classes": NUM_CLASSES,
    "text_proj_layers": [text_dim, CC_LATENT * 2, CC_LATENT, CC_LATENT],
    "voxel_proj_layers": [voxel_dim, CC_LATENT * 2, CC_LATENT, CC_LATENT],
    "activation": "GELU",
    "normalization": "LayerNorm",
    "total_params": cc_params,
    "class_names": CLASS_NAMES,
    "text_encoder": QwenEmbeddingExtractor.MODEL_NAME,
    "voxel_encoder": "GeometricShapeClassifier_v8",
}
with open(cc_staging / "config.json", "w") as f:
    json.dump(cc_arch, f, indent=2)

cc_train_cfg = {
    "n_samples": CC_SAMPLES, "epochs": CC_EPOCHS, "batch_size": CC_BATCH,
    "lr": CC_LR, "weight_decay": 1e-4, "optimizer": "AdamW",
    "scheduler": "cosine_with_warmup", "warmup_epochs": _warmup,
    "loss": "symmetric_InfoNCE", "initial_temperature": 0.07,
    "final_temperature": cc_model.temperature.item(),
    "amp_dtype": str(amp_dtype),
    "best_val_accuracy": best_cc_acc,
    "nearest_text_accuracy": correct / max(total_c, 1),
    "total_training_time_seconds": tt,
}
with open(cc_staging / "training_config.json", "w") as f:
    json.dump(cc_train_cfg, f, indent=2)

# Final upload with full configs + weights
upload_cc_to_hf(cc_model, best_cc_acc, CC_EPOCHS - 1, hf_token, reason="final",
                text_dim_=text_dim, voxel_dim_=voxel_dim, latent_dim_=CC_LATENT)
if not hf_token:
    print("No HF_TOKEN — saved locally at ./hf_staging/crosscontrast/")

print(f"\nAll three subdirectories staged:")
print(f"  geometric_classifier/ — from Cell 3")
print(f"  qwen_embeddings/      — {text_embeddings.shape}")
print(f"  crosscontrast/        — latent={CC_LATENT}, acc={best_cc_acc:.4f}")
print(f"  Repo: https://huggingface.co/{HF_REPO}")