code/probes/novel_probes_v2.py

#!/usr/bin/env python3
"""
Novel Embedding Probes v2 — GlycanBERT V5
==========================================
Uses the FULL pretraining data (254k WURCS) instead of benchmark subsets.
All plots use Nature BGP color palette, 300 DPI, publication-ready.

Probes:
  1. Ambiguity (? marks)     — 98k ambiguous vs 156k clean WURCS
  2. Composition              — monosaccharide fingerprint from [CLS]
  3. KNN Purity (expanded)   — domain, kingdom, link (N vs O), immunogenicity
  4. Polymerization           — chain length / branch depth regression
  5. Size Prediction          — small/med/large/xlarge from frozen [CLS]
  6. N-vs-O Link (binary)    — only N and O linkages embedded
  7. MLM Zero-Shot (fixed)   — random token replacement instead of [MASK]
  8. Token Importance (fixed) — leave-one-out CLS shift analysis

Usage:
  python novel_probes_v2.py --model v5 --probe all --max_samples 5000
"""

import os, sys, json, argparse, csv
import numpy as np
from pathlib import Path
from collections import Counter

# ─── Paths ───────────────────────────────────────────────────────────────
PROJECT_ROOT = Path(__file__).resolve().parents[2]
VOCAB_PATH = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'vocabulary.json'
CHECKPOINTS = {
    'v5': PROJECT_ROOT / 'checkpoints_v5_bpe_topo' / 'best_v5_bpe_topo_model.pt',
    'v6': PROJECT_ROOT / 'bert_v5.1_contrastive' / 'checkpoints' / 'best_v51_contrastive_model.pt',
}
PRETRAIN_CSV = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'multimodal_index.csv'
BENCH_DIR    = PROJECT_ROOT / 'bench' / 'GlycanML' / 'data'

# ─── Nature BGP Color Palette ──────────────────────────────────────────
# From: https://www.nature.com/documents/natrev-artworkguide.pdf
NATURE_COLORS = {
    'blue':    '#0072B2',
    'orange':  '#E69F00',
    'green':   '#009E73',
    'red':     '#D55E00',
    'purple':  '#CC79A7',
    'cyan':    '#56B4E9',
    'yellow':  '#F0E442',
    'black':   '#000000',
    'grey':    '#999999',
}
# Categorical palettes
PALETTE_2 = ['#0072B2', '#D55E00']
PALETTE_3 = ['#0072B2', '#E69F00', '#009E73']
PALETTE_4 = ['#0072B2', '#E69F00', '#009E73', '#D55E00']
PALETTE_5 = ['#0072B2', '#E69F00', '#009E73', '#D55E00', '#CC79A7']
PALETTE_8 = ['#0072B2', '#E69F00', '#009E73', '#D55E00', '#CC79A7',
             '#56B4E9', '#F0E442', '#999999']
PALETTE_11 = PALETTE_8 + ['#000000', '#882255', '#44AA99']

def get_palette(n):
    if n <= 2: return PALETTE_2[:n]
    if n <= 3: return PALETTE_3[:n]
    if n <= 4: return PALETTE_4[:n]
    if n <= 5: return PALETTE_5[:n]
    if n <= 8: return PALETTE_8[:n]
    return (PALETTE_11 * ((n // 11) + 1))[:n]

# ─── Plot setup ─────────────────────────────────────────────────────────
def setup_nature_style():
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    plt.rcParams.update({
        'font.family': 'sans-serif',
        'font.sans-serif': ['Arial', 'Helvetica', 'DejaVu Sans'],
        'font.size': 10,
        'axes.titlesize': 12,
        'axes.labelsize': 11,
        'xtick.labelsize': 9,
        'ytick.labelsize': 9,
        'legend.fontsize': 9,
        'figure.dpi': 300,
        'savefig.dpi': 300,
        'savefig.bbox': 'tight',
        'axes.linewidth': 0.8,
        'axes.spines.top': False,
        'axes.spines.right': False,
    })
    return plt

# ─── Model loading ──────────────────────────────────────────────────────
# Matches the working pattern from embed_benchmark_tasks.py and extract_embeddings.py
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / 'bert_training_v4'))
from model.multimodal_glycan_bert_v3 import MultimodalGlycanBERT, MultimodalGlycanBERTConfig
from downstream_tasks.utils.tokenizer import WURCSTokenizer

def load_model(ckpt_path, device='cuda'):
    """Load MultimodalGlycanBERT from checkpoint (matches embed_benchmark_tasks.py)."""
    import torch
    print(f"Loading model from {ckpt_path}...")
    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)

    if 'model_state_dict' in ckpt:
        state_dict = ckpt['model_state_dict']
    else:
        state_dict = ckpt

    # Strip projection head keys (V6 only)
    backbone_sd = {k: v for k, v in state_dict.items() if not k.startswith('proj_head.')}
    n_stripped = len(state_dict) - len(backbone_sd)
    if n_stripped > 0:
        print(f"  Stripped {n_stripped} projection head keys")

    # Infer vocab size from state dict
    vocab_size = backbone_sd['seq_embeddings.token_embeddings.weight'].shape[0]

    # Check for MS embeddings
    ms_total_vocab = None
    if 'ms_embeddings.token_embeddings.weight' in backbone_sd:
        ms_total_vocab = backbone_sd['ms_embeddings.token_embeddings.weight'].shape[0]

    config_kwargs = dict(
        seq_vocab_size=vocab_size,
        seq_hidden_size=768,
        seq_num_layers=12,
        seq_num_heads=12,
        seq_max_length=256,
        use_cnn_frontend=True,
        cnn_kernel_size=3,
    )
    if ms_total_vocab is not None:
        config_kwargs['ms_vocab_size'] = ms_total_vocab - vocab_size

    config = MultimodalGlycanBERTConfig(**config_kwargs)
    model = MultimodalGlycanBERT(config)
    model.load_state_dict(backbone_sd, strict=False)
    model.to(device)
    model.eval()

    n_params = sum(p.numel() for p in model.parameters())
    print(f"  Model loaded: {n_params:,} params, vocab_size={vocab_size}")
    return model

# ─── Data loading ───────────────────────────────────────────────────────
def load_pretrain_wurcs(tokenizer, max_n=None):
    """Load ALL WURCS from multimodal_index.csv + metadata."""
    samples = []
    with open(PRETRAIN_CSV) as f:
        reader = csv.DictReader(f)
        for row in reader:
            w = row['wurcs']
            if not w.startswith('WURCS'): continue
            try:
                n_res = int(w.split('/')[1].split(',')[1]) if '/' in w else 0
            except:
                n_res = 0
            has_q = '?' in w
            q_count = w.count('?')
            samples.append({
                'wurcs': w,
                'accession': row.get('accession', ''),
                'n_residues': n_res,
                'has_ambiguity': has_q,
                'ambiguity_count': q_count,
                'monosaccharide_names': row.get('monosaccharide_names', ''),
            })
            if max_n and len(samples) >= max_n:
                break
    print(f"  Loaded {len(samples)} WURCS from pretraining data")
    print(f"  Ambiguous (has ?): {sum(1 for s in samples if s['has_ambiguity'])}")
    return samples

def load_benchmark_glycans(tokenizer, csv_name, max_n=None):
    """Load glycans from a benchmark CSV."""
    csv_path = BENCH_DIR / csv_name
    if not csv_path.exists():
        print(f"  WARNING: {csv_path} not found")
        return []
    samples = []
    with open(csv_path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            w = row.get('wurcs', '')
            if not w.startswith('WURCS'): continue
            samples.append(row)
            if max_n and len(samples) >= max_n:
                break
    return samples

# ─── Embedding ──────────────────────────────────────────────────────────
# Matches extract_embeddings.py pattern: use model.seq_embeddings() with
# branch_depths and linkage_types from the WURCSTokenizer.
def batch_cls_embeddings(model, samples, device='cuda', batch_size=64, max_len=256):
    """Extract [CLS] embeddings for a list of samples.

    Uses WURCSTokenizer.tokenize() to get token_ids, branch_depths, and
    linkage_types, then runs model.seq_embeddings() — the working forward
    pass pattern from extract_embeddings.py.
    """
    import torch
    import torch.nn.functional as F
    tokenizer = WURCSTokenizer(str(VOCAB_PATH))

    if not samples:
        return np.zeros((0, 768))

    all_embs = []
    n_errors = 0
    for i in range(0, len(samples), batch_size):
        batch = samples[i:i+batch_size]
        batch_embs = []
        for s in batch:
            try:
                result = tokenizer.tokenize(s['wurcs'], max_length=max_len)
                token_ids = torch.tensor(result['token_ids'], dtype=torch.long)
                branch_depths = torch.tensor(result.get('branch_depths', [0]*len(result['token_ids'])), dtype=torch.long)
                linkage_types = torch.tensor(result.get('linkage_types', [0]*len(result['token_ids'])), dtype=torch.long)

                # Ensure same length
                min_l = min(len(token_ids), len(branch_depths), len(linkage_types))
                token_ids = token_ids[:min_l]
                branch_depths = branch_depths[:min_l]
                linkage_types = linkage_types[:min_l]

                # Truncate / pad to max_len
                if min_l > max_len:
                    token_ids = token_ids[:max_len]
                    branch_depths = branch_depths[:max_len]
                    linkage_types = linkage_types[:max_len]
                elif min_l < max_len:
                    pad_len = max_len - min_l
                    token_ids = F.pad(token_ids, (0, pad_len), value=0)
                    branch_depths = F.pad(branch_depths, (0, pad_len), value=0)
                    linkage_types = F.pad(linkage_types, (0, pad_len), value=0)

                # Forward through seq encoder
                token_ids = token_ids.unsqueeze(0).to(device)
                branch_depths = branch_depths.unsqueeze(0).to(device)
                linkage_types = linkage_types.unsqueeze(0).to(device)

                with torch.no_grad():
                    seq_out = model.seq_embeddings(token_ids, branch_depths=branch_depths, linkage_types=linkage_types)
                    cls_emb = seq_out[0, 0, :].cpu().numpy()
                batch_embs.append(cls_emb)
            except Exception as e:
                n_errors += 1
                if n_errors <= 5:
                    import traceback as tb
                    print(f"  ERROR (sample {i}): {e}")
                    tb.print_exc()
                batch_embs.append(np.zeros(768))

        all_embs.extend(batch_embs)
        if (i // batch_size) % 20 == 0 and i > 0:
            print(f"    Embedded {i}/{len(samples)} ({n_errors} errors)")

    if n_errors > 0:
        print(f"  WARNING: {n_errors}/{len(samples)} tokenization errors")
    print(f"  Embedded {len(all_embs)} total samples", flush=True)
    return np.array(all_embs) if all_embs else np.zeros((0, 768))

# ═══════════════════════════════════════════════════════════════════════
# PROBE 1: Ambiguity (? marks) — FULL DATA
# ═══════════════════════════════════════════════════════════════════════

def save_publication_plots(X, labels, label_name, out_dir, title_prefix=""):
    """Generate PCA publication-quality plot (UMAP disabled to save memory)."""
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from sklearn.decomposition import PCA
    
    out_dir = Path(out_dir)
    safe_name = re.sub(r'[^a-zA-Z0-9_-]', '_', label_name.lower().replace(' ', '_'))
    
    # Subsample for plotting
    max_plot = min(len(X), 10000)
    if len(X) > max_plot:
        idx = np.random.RandomState(42).choice(len(X), max_plot, replace=False)
        X_sub = X[idx]
        labels_sub = [labels[i] for i in idx]
    else:
        X_sub = X
        labels_sub = list(labels)
    
    unique_labels = sorted(set(labels_sub))
    cmap = plt.cm.get_cmap('tab20', len(unique_labels))
    color_map = {lbl: cmap(i) for i, lbl in enumerate(unique_labels)}
    
    fig, ax = plt.subplots(1, 1, figsize=(8, 6))
    
    pca = PCA(n_components=2)
    X_pca = pca.fit_transform(X_sub)
    
    for lbl in unique_labels:
        mask = [l == lbl for l in labels_sub]
        ax.scatter(X_pca[np.array(mask), 0], X_pca[np.array(mask), 1], c=[color_map[lbl]],
                   label=lbl, s=8, alpha=0.5)
    ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})')
    ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})')
    ax.set_title(f'{title_prefix} — PCA by {label_name}')
    ax.legend(fontsize=6, markerscale=2, loc='best')
    
    plt.tight_layout()
    plt.savefig(out_dir / f'{safe_name}_pca.png', dpi=150)
    plt.close(fig)
    del X_sub, X_pca, fig
    print(f"  Saved: {safe_name}_pca.png")



def probe_ambiguity(model, tokenizer, device, output_dir, max_samples=10000, **kwargs):
    print("\n" + "="*60)
    print("PROBE 1: Ambiguity Analysis (? marks in WURCS)")
    print("="*60)

    samples = load_pretrain_wurcs(tokenizer, max_n=max_samples)
    ambig = [s for s in samples if s['has_ambiguity']]
    clean = [s for s in samples if not s['has_ambiguity']]
    print(f"  Ambiguous: {len(ambig)}, Clean: {len(clean)}")

    # Bin ambiguity into levels
    for s in samples:
        qc = s['ambiguity_count']
        if qc == 0: s['amb_level'] = 'none'
        elif qc <= 2: s['amb_level'] = 'low (1-2)'
        elif qc <= 5: s['amb_level'] = 'medium (3-5)'
        else: s['amb_level'] = 'high (6+)'

    level_counts = Counter(s['amb_level'] for s in samples)
    print(f"  Levels: {dict(level_counts)}")

    # Subsample for balance if needed
    min_group = min(len(ambig), len(clean), 2000)
    np.random.seed(42)
    if len(ambig) > min_group:
        ambig = [ambig[i] for i in np.random.choice(len(ambig), min_group, replace=False)]
    if len(clean) > min_group:
        clean = [clean[i] for i in np.random.choice(len(clean), min_group, replace=False)]

    all_samp = ambig + clean
    print(f"  Embedding {len(all_samp)} samples (balanced)...")
    embeddings = batch_cls_embeddings(model, all_samp, device=device)

    if embeddings.shape[0] == 0:
        print("  SKIPPING probe_ambiguity — no valid embeddings")
        return {'error': 'no_embeddings', 'n_ambig': len(ambig), 'n_clean': len(clean)}

    labels = ['ambiguous']*len(ambig) + ['clean']*len(clean)

    # Metrics
    from sklearn.metrics import silhouette_score
    from sklearn.neighbors import KNeighborsClassifier
    from sklearn.model_selection import cross_val_score
    from sklearn.preprocessing import StandardScaler
    int_labels = np.array([0 if l == 'ambiguous' else 1 for l in labels])
    sil = float(silhouette_score(embeddings, int_labels))
    X = StandardScaler().fit_transform(embeddings)
    knn_acc = float(cross_val_score(
        KNeighborsClassifier(n_neighbors=10), X, int_labels, cv=5, scoring='accuracy'
    ).mean())
    print(f"  Silhouette (ambig vs clean): {sil:.4f}")
    print(f"  KNN classification accuracy: {knn_acc:.4f}")

    # Cosine similarity analysis
    from sklearn.metrics.pairwise import cosine_similarity
    emb_ambig = embeddings[:len(ambig)]
    emb_clean = embeddings[len(ambig):]
    within_ambig = float(np.mean(cosine_similarity(emb_ambig)))
    within_clean = float(np.mean(cosine_similarity(emb_clean)))
    between = float(np.mean(cosine_similarity(emb_ambig, emb_clean)))
    print(f"  Within-ambig sim: {within_ambig:.4f}")
    print(f"  Within-clean sim: {within_clean:.4f}")
    print(f"  Between sim: {between:.4f}")

    # t-SNE plot
    plt = setup_nature_style()
    from sklearn.manifold import TSNE
    perp = min(30, len(embeddings) - 1)
    coords = TSNE(n_components=2, perplexity=perp, max_iter=1000,
                  init='pca', random_state=42, learning_rate='auto').fit_transform(embeddings)

    fig, ax = plt.subplots(figsize=(8, 6))
    colors = {'clean': NATURE_COLORS['blue'], 'ambiguous': NATURE_COLORS['orange']}
    for label in ['clean', 'ambiguous']:
        mask = np.array(labels) == label
        ax.scatter(coords[mask, 0], coords[mask, 1], c=colors[label],
                   label=f'{label} (n={mask.sum()})', s=8, alpha=0.5, edgecolors='none')
    ax.set_title(f'Ambiguity Probe: WURCS with ? marks vs Clean\n'
                 f'Silhouette={sil:.4f} | KNN Acc={knn_acc:.4f}')
    ax.set_xlabel('t-SNE 1')
    ax.set_ylabel('t-SNE 2')
    ax.legend(loc='best', framealpha=0.8)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'ambiguity_probe.png'), dpi=300, bbox_inches='tight')
    plt.close()

    results = {
        'silhouette': sil, 'knn_accuracy': knn_acc,
        'within_ambig_sim': within_ambig, 'within_clean_sim': within_clean,
        'between_sim': between,
        'n_ambiguous': len(ambig), 'n_clean': len(clean),
        'level_counts': dict(level_counts),
    }
    with open(os.path.join(output_dir, 'ambiguity_probe.json'), 'w') as f:
        json.dump(results, f, indent=2, default=str)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 2: Monosaccharide Composition — FULL DATA
# ═══════════════════════════════════════════════════════════════════════
def probe_composition(model, tokenizer, device, output_dir, max_samples=5000, **kwargs):
    _cached_embs = kwargs.get("_cached_embs")
    _cached_samples = kwargs.get("_cached_samples")
    print("\n" + "="*60)
    print("PROBE 2: Monosaccharide Composition")
    print("="*60)

    if _cached_samples is not None:


        samples = _cached_samples[:max_samples]


    else:


        samples = load_pretrain_wurcs(tokenizer, max_n=max_samples)
    # Parse monosaccharide names
    for s in samples:
        names = s.get('monosaccharide_names', '')
        s['monos'] = [m.strip() for m in names.split(',') if m.strip()] if names else []

    # Find top-20 most common monosaccharides
    all_monos = []
    for s in samples:
        all_monos.extend(s['monos'])
    mono_counts = Counter(all_monos)
    top_k = [m for m, _ in mono_counts.most_common(20)]
    print(f"  Top-20 monos: {top_k[:5]}...")

    embeddings = batch_cls_embeddings(model, samples, device=device)

    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import cross_val_score
    from sklearn.preprocessing import StandardScaler

    X = StandardScaler().fit_transform(embeddings)
    per_mono_results = {}
    for mono in top_k:
        y = np.array([1 if mono in s['monos'] else 0 for s in samples])
        n_pos = int(y.sum())
        if n_pos < 20 or n_pos > len(y) - 20: continue
        scores = cross_val_score(
            LogisticRegression(max_iter=500, class_weight='balanced'),
            X, y, cv=5, scoring='roc_auc'
        )
        per_mono_results[mono] = {'auc': float(scores.mean()), 'std': float(scores.std()), 'n_pos': n_pos}
        print(f"    {mono:35s}: AUC={scores.mean():.4f} ± {scores.std():.4f} (n+={n_pos})")

    # Bar chart
    plt = setup_nature_style()
    monos_sorted = sorted(per_mono_results.keys(), key=lambda m: per_mono_results[m]['auc'], reverse=True)
    fig, ax = plt.subplots(figsize=(12, 6))
    x = range(len(monos_sorted))
    aucs = [per_mono_results[m]['auc'] for m in monos_sorted]
    stds = [per_mono_results[m]['std'] for m in monos_sorted]
    bars = ax.bar(x, aucs, yerr=stds, color=NATURE_COLORS['blue'], alpha=0.8,
                  edgecolor='white', linewidth=0.5, capsize=3)
    ax.axhline(0.5, color=NATURE_COLORS['grey'], linestyle='--', linewidth=0.8, label='Random baseline')
    ax.set_xticks(x)
    ax.set_xticklabels(monos_sorted, rotation=45, ha='right', fontsize=7)
    ax.set_ylabel('ROC AUC')
    ax.set_title(f'Monosaccharide Detection from Frozen [CLS] Embedding\n(n={len(samples)}, {len(monos_sorted)} monosaccharides)')
    ax.set_ylim(0.4, 1.05)
    ax.legend(loc='lower right')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'composition_probe.png'), dpi=300, bbox_inches='tight')
    plt.close()

    results = {'per_mono_auc': per_mono_results, 'n_samples': len(samples), 'top_k': top_k}
    with open(os.path.join(output_dir, 'composition_probe.json'), 'w') as f:
        json.dump(results, f, indent=2, default=str)
    # Publication UMAP — color by top-3 monosaccharides
    try:
        from pathlib import Path
        top3 = top_k[:3]
        labels_mono = []
        for s in samples[:len(embs)]:
            monos = set(m.strip() for m in s.get('monosaccharide_names', '').split(','))
            found = [m for m in top3 if m in monos]
            labels_mono.append(found[0] if len(found) == 1 else ('Multi' if len(found) > 1 else 'None'))
        save_publication_plots(np.array(embs), labels_mono,
                             'Top-3 Monosaccharides', Path(args.output_dir),
                             title_prefix='Probe 2')
    except Exception as e:
        print(f"  Pub plot error: {e}")
    
    return results

def probe_knn_purity(model, tokenizer, device, output_dir, max_samples=15000, **kwargs):
    print("\n" + "="*60)
    print("PROBE 3: KNN Purity (Expanded)")
    print("="*60)

    # Load classification data (domain + kingdom)
    cls_samples = load_benchmark_glycans(tokenizer, 'glycan_classification_wurcs_subset.csv', max_n=max_samples)
    # Load link data
    link_samples = load_benchmark_glycans(tokenizer, 'glycan_link_wurcs_subset.csv', max_n=max_samples)
    # Load immunogenicity
    immuno_samples = load_benchmark_glycans(tokenizer, 'glycan_immunogenicity_wurcs_subset.csv', max_n=max_samples)

    results = {}

    # Domain + Kingdom KNN
    if cls_samples:
        print(f"  Classification samples: {len(cls_samples)}")
        cls_embs = batch_cls_embeddings(model, cls_samples, device=device)

        for task_col in ['domain', 'kingdom']:
            labels = []
            valid_mask = []
            for i, s in enumerate(cls_samples):
                label = s.get(task_col, '')
                if label:
                    labels.append(label)
                    valid_mask.append(i)
            if not labels: continue

            embs = cls_embs[valid_mask]
            label_arr = np.array(labels)
            n_classes = len(set(labels))
            class_counts = Counter(labels)
            print(f"  {task_col}: {len(labels)} samples, {n_classes} classes")
            print(f"    Distribution: {dict(class_counts)}")

            for k in [5, 10, 20, 50]:
                from sklearn.neighbors import NearestNeighbors
                nn = NearestNeighbors(n_neighbors=k+1, metric='cosine')
                nn.fit(embs)
                _, indices = nn.kneighbors(embs)
                purities = []
                for i in range(len(embs)):
                    neighbors = indices[i, 1:]  # exclude self
                    same_class = np.sum(label_arr[neighbors] == label_arr[i])
                    purities.append(same_class / k)
                purity = float(np.mean(purities))
                results[f'{task_col}_k{k}'] = purity
                print(f"    KNN Purity (k={k:2d}): {purity:.4f}")

            # Per-class purity at k=10
            for cls_name in sorted(set(labels)):
                cls_mask = label_arr == cls_name
                cls_purity = float(np.mean([purities[i] for i in range(len(purities)) if cls_mask[i]]))
                results[f'{task_col}_{cls_name}_k10'] = cls_purity

    # Link KNN (N vs O only — binary)
    if link_samples:
        no_samples = [s for s in link_samples if s.get('target', '') in ('N', 'O')]
        print(f"  Link N-vs-O samples: {len(no_samples)}")
        if len(no_samples) > 50:
            link_embs = batch_cls_embeddings(model, no_samples, device=device)
            link_labels = np.array([s['target'] for s in no_samples])
            for k in [5, 10, 20]:
                from sklearn.neighbors import NearestNeighbors
                nn = NearestNeighbors(n_neighbors=k+1, metric='cosine')
                nn.fit(link_embs)
                _, indices = nn.kneighbors(link_embs)
                purities = []
                for i in range(len(link_embs)):
                    neighbors = indices[i, 1:]
                    same_class = np.sum(link_labels[neighbors] == link_labels[i])
                    purities.append(same_class / k)
                purity = float(np.mean(purities))
                results[f'link_NO_k{k}'] = purity
                print(f"    Link N-vs-O KNN (k={k:2d}): {purity:.4f}")

    # Immunogenicity KNN
    if immuno_samples:
        print(f"  Immunogenicity samples: {len(immuno_samples)}")
        i_embs = batch_cls_embeddings(model, immuno_samples, device=device)
        i_labels = np.array([s.get('target', s.get('immunogenicity', '')) for s in immuno_samples])
        for k in [5, 10, 20]:
            from sklearn.neighbors import NearestNeighbors
            nn = NearestNeighbors(n_neighbors=k+1, metric='cosine')
            nn.fit(i_embs)
            _, indices = nn.kneighbors(i_embs)
            purities = []
            for i in range(len(i_embs)):
                neighbors = indices[i, 1:]
                same = np.sum(i_labels[neighbors] == i_labels[i])
                purities.append(same / k)
            purity = float(np.mean(purities))
            results[f'immunogenicity_k{k}'] = purity
            print(f"    Immunogenicity KNN (k={k:2d}): {purity:.4f}")

    with open(os.path.join(output_dir, 'knn_purity.json'), 'w') as f:
        json.dump(results, f, indent=2, default=str)
    # Publication UMAP — color by domain
    try:
        from pathlib import Path
        domain_labels_all = [s.get('domain','') for s in samples[:len(embs)]]
        if len(set(l for l in domain_labels_all if l)) >= 2:
            save_publication_plots(np.array(embs), domain_labels_all,
                                 'Taxonomy Domain', Path(args.output_dir),
                                 title_prefix='Probe 3a')
    except Exception as e:
        print(f"  Pub plot error: {e}")
    
    return results

def probe_polymerization(model, tokenizer, device, output_dir, max_samples=5000, **kwargs):
    _cached_embs = kwargs.get("_cached_embs")
    _cached_samples = kwargs.get("_cached_samples")
    print("\n" + "="*60)
    print("PROBE 4: Polymerization / Complexity Probe")
    print("="*60)

    if _cached_samples is not None:


        samples = _cached_samples[:max_samples]


    else:


        samples = load_pretrain_wurcs(tokenizer, max_n=max_samples)
    # Parse complexity features from WURCS
    for s in samples:
        w = s['wurcs']
        try:
            parts = w.split('/')
            counts = parts[1].split(',')
            s['n_unique_res'] = int(counts[0])
            s['n_total_res'] = int(counts[1])
            s['n_linkages'] = int(counts[2]) if len(counts) > 2 else s['n_total_res'] - 1
        except:
            s['n_unique_res'] = s['n_residues']
            s['n_total_res'] = s['n_residues']
            s['n_linkages'] = max(0, s['n_residues'] - 1)
        # Branch depth
        try:
            link_str = w.split('/')[-1] if '/' in w else ''
            depth = link_str.count('-') - (s['n_total_res'] - 1) if link_str else 0
            s['branch_depth'] = max(0, depth)
        except:
            s['branch_depth'] = 0

    print(f"  Samples: {len(samples)}, Residues: {min(s['n_total_res'] for s in samples)}-{max(s['n_total_res'] for s in samples)}")

    embeddings = batch_cls_embeddings(model, samples, device=device)

    from sklearn.linear_model import Ridge
    from sklearn.model_selection import cross_val_score
    from sklearn.preprocessing import StandardScaler
    from scipy.stats import spearmanr

    X = StandardScaler().fit_transform(embeddings)
    features = {
        'n_total_residues': [s['n_total_res'] for s in samples],
        'n_unique_residues': [s['n_unique_res'] for s in samples],
        'n_linkages': [s['n_linkages'] for s in samples],
        'branch_depth': [s['branch_depth'] for s in samples],
    }

    results = {'linear_probe_r2': {}, 'spearman_correlations': {}, 'n_samples': len(samples)}
    for fname, values in features.items():
        y = np.array(values, dtype=float)
        if np.std(y) < 1e-6: continue
        r2 = cross_val_score(Ridge(alpha=1.0), X, y, cv=5, scoring='r2')
        results['linear_probe_r2'][fname] = {'r2_mean': float(r2.mean()), 'r2_std': float(r2.std())}
        print(f"  {fname:25s}: R²={r2.mean():.4f} ± {r2.std():.4f}")

    # Pairwise Spearman
    from sklearn.metrics.pairwise import euclidean_distances
    dists = euclidean_distances(embeddings)
    upper_idx = np.triu_indices(len(embeddings), k=1)
    emb_dists = dists[upper_idx]
    for fname, values in features.items():
        y = np.array(values, dtype=float)
        feat_diffs = np.abs(y[upper_idx[0]] - y[upper_idx[1]])
        # Subsample for speed
        if len(emb_dists) > 500000:
            idx = np.random.choice(len(emb_dists), 500000, replace=False)
            rho, p = spearmanr(emb_dists[idx], feat_diffs[idx])
        else:
            rho, p = spearmanr(emb_dists, feat_diffs)
        results['spearman_correlations'][fname] = {'rho': float(rho), 'p': float(p)}
        print(f"  Spearman ρ ({fname}): {rho:.4f} (p={p:.2e})")

    # Scatter plot: n_residues vs CLS PCA1
    plt = setup_nature_style()
    from sklearn.decomposition import PCA
    pca = PCA(n_components=2)
    pca_coords = pca.fit_transform(embeddings)

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    for ax, (fname, values) in zip(axes, [('n_total_residues', features['n_total_residues']),
                                            ('n_unique_residues', features['n_unique_residues'])]):
        sc = ax.scatter(pca_coords[:, 0], pca_coords[:, 1], c=values,
                        cmap='viridis', s=5, alpha=0.5, edgecolors='none')
        plt.colorbar(sc, ax=ax, label=fname)
        ax.set_xlabel('PCA 1')
        ax.set_ylabel('PCA 2')
        ax.set_title(f'{fname}\nR²={results["linear_probe_r2"].get(fname, {}).get("r2_mean", 0):.4f}')
    plt.suptitle(f'Polymerization Probe (n={len(samples)})', fontsize=13)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'polymerization_probe.png'), dpi=300, bbox_inches='tight')
    plt.close()

    with open(os.path.join(output_dir, 'polymerization_probe.json'), 'w') as f:
        json.dump(results, f, indent=2, default=str)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 5: Size Category Prediction — FULL DATA
# ═══════════════════════════════════════════════════════════════════════
def probe_size(model, tokenizer, device, output_dir, max_samples=5000, **kwargs):
    _cached_embs = kwargs.get("_cached_embs")
    _cached_samples = kwargs.get("_cached_samples")
    print("\n" + "="*60)
    print("PROBE 5: Size Category Prediction")
    print("="*60)

    if _cached_samples is not None:


        samples = _cached_samples[:max_samples]


    else:


        samples = load_pretrain_wurcs(tokenizer, max_n=max_samples)
    for s in samples:
        n = s['n_residues']
        s['size'] = 'small' if n <= 3 else ('medium' if n <= 6 else ('large' if n <= 10 else 'very_large'))

    size_dist = Counter(s['size'] for s in samples)
    print(f"  Sizes: {dict(size_dist)}")

    embeddings = batch_cls_embeddings(model, samples, device=device)
    labels = [s['size'] for s in samples]

    from sklearn.metrics import silhouette_score
    from sklearn.neighbors import KNeighborsClassifier
    from sklearn.model_selection import cross_val_score
    from sklearn.preprocessing import StandardScaler

    unique = sorted(set(labels))
    l2i = {l: i for i, l in enumerate(unique)}
    int_labels = np.array([l2i[l] for l in labels])
    sil = float(silhouette_score(embeddings, int_labels))
    X = StandardScaler().fit_transform(embeddings)
    knn_acc = float(cross_val_score(
        KNeighborsClassifier(n_neighbors=10), X, int_labels, cv=5, scoring='accuracy'
    ).mean())
    print(f"  Silhouette: {sil:.4f}, KNN Acc: {knn_acc:.4f}")

    # t-SNE plot
    plt = setup_nature_style()
    from sklearn.manifold import TSNE
    perp = min(30, len(embeddings) - 1)
    coords = TSNE(n_components=2, perplexity=perp, max_iter=1000,
                  init='pca', random_state=42, learning_rate='auto').fit_transform(embeddings)
    fig, ax = plt.subplots(figsize=(8, 6))
    colors = {'small': NATURE_COLORS['green'], 'medium': NATURE_COLORS['blue'],
              'large': NATURE_COLORS['orange'], 'very_large': NATURE_COLORS['red']}
    for cat in ['small', 'medium', 'large', 'very_large']:
        mask = np.array(labels) == cat
        if mask.any():
            ax.scatter(coords[mask, 0], coords[mask, 1], c=colors[cat],
                       label=f'{cat} (n={mask.sum()})', s=8, alpha=0.5, edgecolors='none')
    ax.set_title(f'Size Category Prediction\nSilhouette={sil:.4f} | KNN Acc={knn_acc:.4f}')
    ax.set_xlabel('t-SNE 1')
    ax.set_ylabel('t-SNE 2')
    ax.legend(loc='best', framealpha=0.8)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'size_probe.png'), dpi=300, bbox_inches='tight')
    plt.close()

    results = {'silhouette': sil, 'knn_accuracy': knn_acc, 'sizes': dict(size_dist)}
    with open(os.path.join(output_dir, 'size_probe.json'), 'w') as f:
        json.dump(results, f, indent=2, default=str)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 6: N-vs-O Link (Binary Embedding)
# ═══════════════════════════════════════════════════════════════════════
def probe_link_binary(model, tokenizer, device, output_dir, max_samples=5000, **kwargs):
    print("\n" + "="*60)
    print("PROBE 6: N-linked vs O-linked (Binary)")
    print("="*60)

    link_samples = load_benchmark_glycans(tokenizer, 'glycan_link_wurcs_subset.csv', max_n=max_samples)
    # Filter to N and O only
    no_samples = [s for s in link_samples if s.get('target', '') in ('N', 'O')]
    print(f"  N+O samples: {len(no_samples)}")
    label_dist = Counter(s['target'] for s in no_samples)
    print(f"  Distribution: {dict(label_dist)}")

    if len(no_samples) < 50:
        print("  Too few samples, skipping")
        return {'error': 'Too few N/O samples'}

    embeddings = batch_cls_embeddings(model, no_samples, device=device)
    labels = [s['target'] for s in no_samples]

    from sklearn.metrics import silhouette_score
    from sklearn.neighbors import KNeighborsClassifier
    from sklearn.model_selection import cross_val_score
    from sklearn.preprocessing import StandardScaler
    from sklearn.linear_model import LogisticRegression

    int_labels = np.array([0 if l == 'N' else 1 for l in labels])
    sil = float(silhouette_score(embeddings, int_labels))
    X = StandardScaler().fit_transform(embeddings)
    knn_acc = float(cross_val_score(
        KNeighborsClassifier(n_neighbors=10), X, int_labels, cv=5, scoring='accuracy'
    ).mean())
    lr_auc = float(cross_val_score(
        LogisticRegression(max_iter=500), X, int_labels, cv=5, scoring='roc_auc'
    ).mean())
    print(f"  Silhouette: {sil:.4f}, KNN Acc: {knn_acc:.4f}, LR AUC: {lr_auc:.4f}")

    # t-SNE
    plt = setup_nature_style()
    from sklearn.manifold import TSNE
    perp = min(30, len(embeddings) - 1)
    coords = TSNE(n_components=2, perplexity=perp, max_iter=1000,
                  init='pca', random_state=42, learning_rate='auto').fit_transform(embeddings)
    fig, ax = plt.subplots(figsize=(8, 6))
    for label, color, name in [('N', NATURE_COLORS['blue'], 'N-linked'),
                                 ('O', NATURE_COLORS['orange'], 'O-linked')]:
        mask = np.array(labels) == label
        ax.scatter(coords[mask, 0], coords[mask, 1], c=color,
                   label=f'{name} (n={mask.sum()})', s=15, alpha=0.6, edgecolors='none')
    ax.set_title(f'N-linked vs O-linked Glycans\nSilhouette={sil:.4f} | KNN={knn_acc:.4f} | AUC={lr_auc:.4f}')
    ax.set_xlabel('t-SNE 1')
    ax.set_ylabel('t-SNE 2')
    ax.legend(loc='best', framealpha=0.8)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'link_binary_probe.png'), dpi=300, bbox_inches='tight')
    plt.close()

    results = {'silhouette': sil, 'knn_accuracy': knn_acc, 'lr_auc': lr_auc,
               'distribution': dict(label_dist)}
    with open(os.path.join(output_dir, 'link_binary_probe.json'), 'w') as f:
        json.dump(results, f, indent=2, default=str)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 7: MLM Zero-Shot (Fixed — random token replacement)
# ═══════════════════════════════════════════════════════════════════════
def probe_mlm_zeroshot(model, tokenizer, device, output_dir, max_samples=500, **kwargs):
    _cached_embs = kwargs.get("_cached_embs")
    _cached_samples = kwargs.get("_cached_samples")
    print("\n" + "="*60)
    print("PROBE 7: MLM Zero-Shot (Token Replacement)")
    print("="*60)
    import torch
    import torch.nn.functional as F

    if _cached_samples is not None:


        samples = _cached_samples[:max_samples]


    else:


        samples = load_pretrain_wurcs(tokenizer, max_n=max_samples)
    tok = WURCSTokenizer(str(VOCAB_PATH))
    MAX_LEN = 256

    def _get_cls(token_ids, branch_depths, linkage_types):
        """Helper: pad/truncate and run model.seq_embeddings(), return CLS numpy."""
        min_l = min(len(token_ids), len(branch_depths), len(linkage_types))
        token_ids = token_ids[:min_l]
        branch_depths = branch_depths[:min_l]
        linkage_types = linkage_types[:min_l]
        if min_l > MAX_LEN:
            token_ids = token_ids[:MAX_LEN]
            branch_depths = branch_depths[:MAX_LEN]
            linkage_types = linkage_types[:MAX_LEN]
        elif min_l < MAX_LEN:
            p = MAX_LEN - min_l
            token_ids = F.pad(token_ids, (0, p), value=0)
            branch_depths = F.pad(branch_depths, (0, p), value=0)
            linkage_types = F.pad(linkage_types, (0, p), value=0)
        with torch.no_grad():
            out = model.seq_embeddings(
                token_ids.unsqueeze(0).to(device),
                branch_depths=branch_depths.unsqueeze(0).to(device),
                linkage_types=linkage_types.unsqueeze(0).to(device),
            )
            return out[0, 0, :].cpu().numpy()

    correct_predictions = 0
    total_predictions = 0
    per_position_shifts = []

    for idx, s in enumerate(samples):
        if idx > 200: break  # cap for speed
        try:
            result = tok.tokenize(s['wurcs'], max_length=MAX_LEN)
            ids = torch.tensor(result['token_ids'], dtype=torch.long)
            bd = torch.tensor(result.get('branch_depths', [0]*len(result['token_ids'])), dtype=torch.long)
            lt = torch.tensor(result.get('linkage_types', [0]*len(result['token_ids'])), dtype=torch.long)
            real_len = result.get('length', len(ids))
            if real_len < 3: continue

            # Get original CLS
            cls_orig = _get_cls(ids.clone(), bd.clone(), lt.clone())

            # For each non-special token position, replace with UNK
            for pos in range(1, min(real_len - 1, 20)):
                original_token = ids[pos].item()
                ids_modified = ids.clone()
                ids_modified[pos] = 1  # UNK token
                cls_modified = _get_cls(ids_modified, bd.clone(), lt.clone())

                shift = float(np.linalg.norm(cls_orig - cls_modified))
                per_position_shifts.append({
                    'sample_idx': idx,
                    'position': pos,
                    'original_token': original_token,
                    'cls_shift': shift,
                })
                total_predictions += 1

        except Exception as e:
            continue

        if idx % 50 == 0:
            print(f"  Processed {idx}/{min(len(samples), 200)}")

    if not per_position_shifts:
        return {'error': 'No predictions could be made'}

    shifts = [p['cls_shift'] for p in per_position_shifts]
    mean_shift = float(np.mean(shifts))
    std_shift = float(np.std(shifts))
    print(f"  Total token replacements: {total_predictions}")
    print(f"  Mean CLS shift: {mean_shift:.4f} ± {std_shift:.4f}")

    # Plot shift distribution
    plt = setup_nature_style()
    fig, ax = plt.subplots(figsize=(8, 5))
    ax.hist(shifts, bins=50, color=NATURE_COLORS['blue'], alpha=0.8, edgecolor='white')
    ax.axvline(mean_shift, color=NATURE_COLORS['red'], linestyle='--', linewidth=1.5,
               label=f'Mean = {mean_shift:.3f}')
    ax.set_xlabel('CLS Embedding Shift (L2 norm)')
    ax.set_ylabel('Count')
    ax.set_title(f'Token Replacement → CLS Shift Distribution\n(n={total_predictions} replacements)')
    ax.legend()
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'mlm_zeroshot_probe.png'), dpi=300, bbox_inches='tight')
    plt.close()

    results = {
        'total_replacements': total_predictions,
        'mean_cls_shift': mean_shift,
        'std_cls_shift': std_shift,
        'median_cls_shift': float(np.median(shifts)),
    }
    with open(os.path.join(output_dir, 'mlm_zeroshot_probe.json'), 'w') as f:
        json.dump(results, f, indent=2, default=str)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 8: Token Importance (Leave-one-out CLS shift)
# ═══════════════════════════════════════════════════════════════════════
def probe_token_importance(model, tokenizer, device, output_dir, max_samples=200, **kwargs):
    _cached_embs = kwargs.get("_cached_embs")
    _cached_samples = kwargs.get("_cached_samples")
    print("\n" + "="*60)
    print("PROBE 8: Token Importance (Leave-One-Out)")
    print("="*60)
    import torch
    import torch.nn.functional as F

    if _cached_samples is not None:


        samples = _cached_samples[:max_samples]


    else:


        samples = load_pretrain_wurcs(tokenizer, max_n=max_samples)
    tok = WURCSTokenizer(str(VOCAB_PATH))
    MAX_LEN = 256

    def _get_cls2(token_ids, branch_depths, linkage_types):
        """Helper: pad/truncate and run model.seq_embeddings(), return CLS numpy."""
        min_l = min(len(token_ids), len(branch_depths), len(linkage_types))
        token_ids = token_ids[:min_l]
        branch_depths = branch_depths[:min_l]
        linkage_types = linkage_types[:min_l]
        if min_l > MAX_LEN:
            token_ids = token_ids[:MAX_LEN]
            branch_depths = branch_depths[:MAX_LEN]
            linkage_types = linkage_types[:MAX_LEN]
        elif min_l < MAX_LEN:
            p = MAX_LEN - min_l
            token_ids = F.pad(token_ids, (0, p), value=0)
            branch_depths = F.pad(branch_depths, (0, p), value=0)
            linkage_types = F.pad(linkage_types, (0, p), value=0)
        with torch.no_grad():
            out = model.seq_embeddings(
                token_ids.unsqueeze(0).to(device),
                branch_depths=branch_depths.unsqueeze(0).to(device),
                linkage_types=linkage_types.unsqueeze(0).to(device),
            )
            return out[0, 0, :].cpu().numpy()

    # For each sample, drop one token at a time and measure CLS shift
    all_importance_by_position = {}  # position -> list of shifts
    token_importance_map = {}  # token_id -> list of shifts

    for idx, s in enumerate(samples):
        if idx > 100: break
        try:
            result = tok.tokenize(s['wurcs'], max_length=MAX_LEN)
            ids = torch.tensor(result['token_ids'], dtype=torch.long)
            bd = torch.tensor(result.get('branch_depths', [0]*len(result['token_ids'])), dtype=torch.long)
            lt = torch.tensor(result.get('linkage_types', [0]*len(result['token_ids'])), dtype=torch.long)
            real_len = result.get('length', len(ids))
            if real_len < 4: continue

            # Original CLS
            cls_orig = _get_cls2(ids.clone(), bd.clone(), lt.clone())

            seq_len = real_len
            for pos in range(1, min(seq_len - 1, 30)):
                # Remove token at position pos (drop from all 3 tensors)
                ids_dropped = torch.cat([ids[:pos], ids[pos+1:]])
                bd_dropped = torch.cat([bd[:pos], bd[pos+1:]])
                lt_dropped = torch.cat([lt[:pos], lt[pos+1:]])
                cls_dropped = _get_cls2(ids_dropped, bd_dropped, lt_dropped)

                shift = float(np.linalg.norm(cls_orig - cls_dropped))
                rel_pos = pos / seq_len  # relative position

                # Track by relative position bin
                bin_key = f'{int(rel_pos * 10) / 10:.1f}'
                if bin_key not in all_importance_by_position:
                    all_importance_by_position[bin_key] = []
                all_importance_by_position[bin_key].append(shift)

                # Track by token id
                tid = ids[pos].item()
                if tid not in token_importance_map:
                    token_importance_map[tid] = []
                token_importance_map[tid].append(shift)

        except Exception as e:
            continue

        if idx % 25 == 0:
            print(f"  Processed {idx}/100")

    if not all_importance_by_position:
        return {'error': 'No importance data'}

    # Average importance by position
    pos_importance = {k: float(np.mean(v)) for k, v in sorted(all_importance_by_position.items())}
    print(f"  Position importance: {pos_importance}")

    # Top-10 most important tokens
    token_avg = {k: float(np.mean(v)) for k, v in token_importance_map.items() if len(v) >= 3}
    top_tokens = sorted(token_avg.items(), key=lambda x: x[1], reverse=True)[:20]
    print(f"  Top-10 most important tokens (by CLS shift):")
    for tid, shift in top_tokens[:10]:
        print(f"    Token {tid}: shift={shift:.4f}")

    # Plot: importance by position
    plt = setup_nature_style()
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

    positions = sorted(pos_importance.keys())
    imp_vals = [pos_importance[p] for p in positions]
    ax1.bar(range(len(positions)), imp_vals, color=NATURE_COLORS['blue'], alpha=0.8, edgecolor='white')
    ax1.set_xticks(range(len(positions)))
    ax1.set_xticklabels(positions, fontsize=8)
    ax1.set_xlabel('Relative Position in Sequence')
    ax1.set_ylabel('Mean CLS Shift')
    ax1.set_title('Token Importance by Position')

    # Plot: top token importance
    top_ids = [str(t[0]) for t in top_tokens[:15]]
    top_shifts = [t[1] for t in top_tokens[:15]]
    ax2.barh(range(len(top_ids)), top_shifts, color=NATURE_COLORS['orange'], alpha=0.8, edgecolor='white')
    ax2.set_yticks(range(len(top_ids)))
    ax2.set_yticklabels(top_ids, fontsize=8)
    ax2.set_xlabel('Mean CLS Shift')
    ax2.set_title('Top-15 Most Important Tokens')
    ax2.invert_yaxis()

    plt.suptitle('Token Importance Analysis (Leave-One-Out)', fontsize=13)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, 'token_importance_probe.png'), dpi=300, bbox_inches='tight')
    plt.close()

    results = {
        'position_importance': pos_importance,
        'top_tokens': {str(k): v for k, v in top_tokens},
        'n_samples_processed': min(len(samples), 100),
    }
    with open(os.path.join(output_dir, 'token_importance_probe.json'), 'w') as f:
        json.dump(results, f, indent=2, default=str)
    return results


# ═══════════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════════


# ============================================================
# PROBE 9: Cancer Glycan Marker Signatures
# ============================================================
def probe_cancer_markers(model, tokenizer, device, output_dir, max_samples=5000):
    # Create args-like namespace for compatibility
    import types
    args = types.SimpleNamespace(output_dir=output_dir, max_samples=max_samples, batch_size=64)
    """Probe whether embeddings separate cancer-associated glycan signatures.
    
    Cancer cells have aberrant glycosylation: hyper-sialylation (Neu5Ac),
    hyper-fucosylation (Fuc), and truncated O-glycans (Tn antigen = single GalNAc).
    We classify glycans as "cancer-associated" if they have >=2 sialylation markers
    OR specific truncation patterns.
    """
    print("\n" + "="*60)
    print("PROBE 9: Cancer Glycan Marker Signatures")
    print("="*60)
    
    import csv, json
    from pathlib import Path
    
    root = Path(__file__).resolve().parent.parent.parent
    csv_path = root / 'bert_training_v4' / 'data' / 'multimodal_index.csv'
    
    samples = []
    with open(csv_path) as fh:
        reader = csv.DictReader(fh)
        for i, row in enumerate(reader):
            if i >= args.max_samples:
                break
            w = row.get('wurcs', '')
            monos = row.get('monosaccharide_names', '')
            if not w:
                continue
            
            mono_list = [m.strip() for m in monos.split(',') if m.strip()]
            
            # Cancer-associated markers
            n_sialic = sum(1 for m in mono_list if m in ('Neu5Ac', 'Neu5Gc', 'KDN'))
            n_fuc = sum(1 for m in mono_list if m == 'Fuc')
            n_galnac = sum(1 for m in mono_list if m == 'GalNAc')
            total_monos = len(mono_list)
            
            # Cancer score: hyper-sialylation OR hyper-fucosylation OR truncated
            sialylation_ratio = n_sialic / max(total_monos, 1)
            fucosylation_ratio = n_fuc / max(total_monos, 1)
            is_truncated = (total_monos <= 2 and n_galnac >= 1)  # Tn-like
            
            # Binary: cancer-associated if high sialylation/fucosylation or truncated
            cancer_assoc = (sialylation_ratio >= 0.3 or fucosylation_ratio >= 0.3 
                          or is_truncated)
            
            label = 'cancer_associated' if cancer_assoc else 'normal'
            samples.append({'wurcs': w, 'label': label, 
                           'n_sialic': n_sialic, 'n_fuc': n_fuc,
                           'sialylation_ratio': sialylation_ratio})
    
    labels = [s['label'] for s in samples]
    from collections import Counter
    dist = Counter(labels)
    print(f"  Total: {len(samples)}, Distribution: {dict(dist)}")
    
    if dist['cancer_associated'] < 20 or dist['normal'] < 20:
        print("  Too few samples in one class, skipping")
        return {}
    
    # Balance classes
    min_n = min(dist.values())
    balanced = []
    counts = {'cancer_associated': 0, 'normal': 0}
    for s in samples:
        if counts[s['label']] < min_n:
            balanced.append(s)
            counts[s['label']] += 1
    
    print(f"  Balanced: {len(balanced)} ({min_n} per class)")
    
    embs = batch_cls_embeddings(model, balanced, device=device, 
                                 batch_size=args.batch_size if hasattr(args, 'batch_size') else 64)
    if embs is None or len(embs) == 0:
        print("  SKIPPING — no valid embeddings")
        return {}
    
    from sklearn.metrics import roc_auc_score, silhouette_score
    from sklearn.model_selection import cross_val_score
    from sklearn.linear_model import LogisticRegression
    
    X = np.array(embs)
    y = np.array([1 if s['label'] == 'cancer_associated' else 0 for s in balanced[:len(embs)]])
    
    # Linear probe
    lr = LogisticRegression(max_iter=1000, random_state=42)
    scores = cross_val_score(lr, X, y, cv=5, scoring='roc_auc')
    mean_auc = scores.mean()
    std_auc = scores.std()
    
    # KNN
    from sklearn.neighbors import KNeighborsClassifier
    knn = KNeighborsClassifier(n_neighbors=5)
    knn_scores = cross_val_score(knn, X, y, cv=5, scoring='accuracy')
    
    # Silhouette
    try:
        sil = silhouette_score(X, y)
    except:
        sil = float('nan')
    
    results = {
        'n_samples': len(embs),
        'n_cancer': int(y.sum()),
        'n_normal': int((1-y).sum()),
        'linear_probe_auc': float(mean_auc),
        'linear_probe_auc_std': float(std_auc),
        'knn_accuracy': float(knn_scores.mean()),
        'silhouette': float(sil),
    }
    
    print(f"  Linear Probe AUC: {mean_auc:.4f} ± {std_auc:.4f}")
    print(f"  KNN Accuracy:     {knn_scores.mean():.4f}")
    print(f"  Silhouette:       {sil:.4f}")
    
    # Save results
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    with open(out_dir / 'cancer_markers_probe.json', 'w') as fh:
        json.dump(results, fh, indent=2)
    
    # Plot
    try:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        from sklearn.decomposition import PCA
        
        pca = PCA(n_components=2)
        X2 = pca.fit_transform(X)
        
        fig, ax = plt.subplots(1, 1, figsize=(8, 6))
        for label_val, label_name, color in [(1, 'Cancer-associated', 'red'), 
                                              (0, 'Normal', 'blue')]:
            mask = y == label_val
            ax.scatter(X2[mask, 0], X2[mask, 1], c=color, alpha=0.3, s=10, label=label_name)
        ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})')
        ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})')
        ax.set_title(f'Cancer Glycan Markers (AUC={mean_auc:.3f})')
        ax.legend()
        plt.tight_layout()
        plt.savefig(out_dir / 'cancer_markers_probe.png', dpi=150)
        plt.close()
    except Exception as e:
        print(f"  Plot error: {e}")
    
    # Publication plots
    try:
        save_publication_plots(X, [s['label'] for s in balanced[:len(embs)]], 
                             'Cancer Glycan Markers', out_dir,
                             title_prefix='Probe 9')
    except Exception as e:
        print(f"  Pub plot error: {e}")
    
    return results

def probe_glycosylation_type(model, tokenizer, device, output_dir, max_samples=5000):
    # Create args-like namespace for compatibility
    import types
    args = types.SimpleNamespace(output_dir=output_dir, max_samples=max_samples, batch_size=64)
    """Probe whether embeddings separate N-linked vs O-linked vs free glycans.
    Uses curated GlycanML benchmark link data.
    """
    print("\n" + "="*60)
    print("PROBE 10: Glycosylation Type (N vs O)")
    print("="*60)
    
    import csv, json
    from pathlib import Path
    
    root = Path(__file__).resolve().parent.parent.parent
    link_csv = root / 'bench' / 'GlycanML' / 'data' / 'glycan_link_wurcs_subset.csv'
    
    if not link_csv.exists():
        print(f"  Link data not found: {link_csv}")
        return {}
    
    samples = []
    with open(link_csv) as fh:
        reader = csv.DictReader(fh)
        for row in reader:
            w = row.get('wurcs', '')
            link = row.get('link', '')
            if w and link in ('N', 'O'):
                samples.append({'wurcs': w, 'label': link})
    
    from collections import Counter
    dist = Counter(s['label'] for s in samples)
    print(f"  Samples: {len(samples)}, Distribution: {dict(dist)}")
    
    if len(samples) < 50:
        print("  Too few samples, skipping")
        return {}
    
    embs = batch_cls_embeddings(model, samples, device=device,
                                 batch_size=args.batch_size if hasattr(args, 'batch_size') else 64)
    if embs is None or len(embs) == 0:
        print("  SKIPPING — no valid embeddings")
        return {}
    
    from sklearn.metrics import silhouette_score
    from sklearn.neighbors import KNeighborsClassifier
    from sklearn.model_selection import cross_val_score
    from sklearn.preprocessing import LabelEncoder
    
    X = np.array(embs)
    le = LabelEncoder()
    y = le.fit_transform([s['label'] for s in samples[:len(embs)]])
    classes = list(le.classes_)
    
    # KNN at multiple k
    results = {'n_samples': len(embs), 'distribution': dict(dist), 'classes': classes}
    for k in [5, 10, 20]:
        if len(embs) > k:
            knn = KNeighborsClassifier(n_neighbors=k)
            scores = cross_val_score(knn, X, y, cv=5, scoring='accuracy')
            results[f'knn_k{k}_accuracy'] = float(scores.mean())
            print(f"  KNN (k={k:2d}): {scores.mean():.4f}")
    
    # Silhouette
    try:
        sil = silhouette_score(X, y)
        results['silhouette'] = float(sil)
        print(f"  Silhouette: {sil:.4f}")
    except:
        pass
    
    # Linear probe (N vs O only, binary)
    n_o_mask = np.array([s['label'] in ('N', 'O') for s in samples[:len(embs)]])
    if n_o_mask.sum() > 50:
        from sklearn.linear_model import LogisticRegression
        X_no = X[n_o_mask]
        y_no = np.array([1 if s['label'] == 'N' else 0 
                         for s in samples[:len(embs)]])[n_o_mask]
        lr = LogisticRegression(max_iter=1000, random_state=42)
        auc_scores = cross_val_score(lr, X_no, y_no, cv=5, scoring='roc_auc')
        results['n_vs_o_auc'] = float(auc_scores.mean())
        results['n_vs_o_auc_std'] = float(auc_scores.std())
        print(f"  N-vs-O AUC: {auc_scores.mean():.4f} ± {auc_scores.std():.4f}")
    
    # Save
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    with open(out_dir / 'glycosylation_type_probe.json', 'w') as fh:
        json.dump(results, fh, indent=2)
    
    # Plot
    try:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        from sklearn.decomposition import PCA
        
        pca = PCA(n_components=2)
        X2 = pca.fit_transform(X)
        
        fig, ax = plt.subplots(1, 1, figsize=(8, 6))
        colors = {'N': 'blue', 'O': 'red', 'free': 'green'}
        for c in classes:
            mask = np.array([s['label'] == c for s in samples[:len(embs)]])
            ax.scatter(X2[mask, 0], X2[mask, 1], c=colors.get(c, 'gray'), 
                      alpha=0.4, s=15, label=f'{c} (n={mask.sum()})')
        ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})')
        ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})')
        ax.set_title('Glycosylation Type: N vs O')
        ax.legend()
        plt.tight_layout()
        plt.savefig(out_dir / 'glycosylation_type_probe.png', dpi=150)
        plt.close()
    except Exception as e:
        print(f"  Plot error: {e}")
    
    # Publication plots
    try:
        save_publication_plots(X, [s['label'] for s in samples[:len(embs)]],
                             'Glycosylation Type', out_dir,
                             title_prefix='Probe 10')
    except Exception as e:
        print(f"  Pub plot error: {e}")
    
    return results

def probe_taxonomic_class(model, tokenizer, device, output_dir, max_samples=5000):
    # Create args-like namespace for compatibility
    import types
    args = types.SimpleNamespace(output_dir=output_dir, max_samples=max_samples, batch_size=64)
    """Probe whether embeddings separate glycans by biological class.
    Uses GlycanML classification data with 90+ taxonomic classes.
    """
    print("\n" + "="*60)
    print("PROBE 11: Taxonomic Classification (GlycanML)")
    print("="*60)
    
    import csv, json
    from pathlib import Path
    from collections import Counter
    
    root = Path(__file__).resolve().parent.parent.parent
    cls_csv = root / 'bench' / 'GlycanML' / 'data' / 'glycan_classification_wurcs_subset.csv'
    
    if not cls_csv.exists():
        print(f"  Classification data not found: {cls_csv}")
        return {}
    
    samples = []
    with open(cls_csv) as fh:
        reader = csv.DictReader(fh)
        for row in reader:
            w = row.get('wurcs', '')
            cls_label = row.get('class', '').strip()
            domain = row.get('domain', '').strip()
            kingdom = row.get('kingdom', '').strip()
            phylum = row.get('phylum', '').strip()
            if w and cls_label:
                samples.append({
                    'wurcs': w, 'class': cls_label,
                    'domain': domain, 'kingdom': kingdom, 'phylum': phylum
                })
    
    # Filter to classes with >= 20 samples for meaningful evaluation
    class_dist = Counter(s['class'] for s in samples)
    valid_classes = {c for c, n in class_dist.items() if n >= 20}
    samples = [s for s in samples if s['class'] in valid_classes]
    
    class_dist = Counter(s['class'] for s in samples)
    print(f"  Samples: {len(samples)}, Classes (n>=20): {len(valid_classes)}")
    print(f"  Top-10: {class_dist.most_common(10)}")
    
    if len(samples) < 100 or len(valid_classes) < 3:
        print("  Too few samples/classes, skipping")
        return {}
    
    # Cap at max_samples
    if len(samples) > args.max_samples:
        samples = samples[:args.max_samples]
    
    embs = batch_cls_embeddings(model, samples, device=device,
                                 batch_size=args.batch_size if hasattr(args, 'batch_size') else 64)
    if embs is None or len(embs) == 0:
        print("  SKIPPING — no valid embeddings")
        return {}
    
    from sklearn.metrics import silhouette_score
    from sklearn.neighbors import KNeighborsClassifier
    from sklearn.model_selection import cross_val_score
    from sklearn.preprocessing import LabelEncoder
    
    X = np.array(embs)
    
    results = {'n_samples': len(embs), 'n_classes': len(valid_classes)}
    
    # Evaluate at class level
    le = LabelEncoder()
    y_class = le.fit_transform([s['class'] for s in samples[:len(embs)]])
    
    for k in [5, 10, 20]:
        if len(embs) > k:
            knn = KNeighborsClassifier(n_neighbors=k)
            scores = cross_val_score(knn, X, y_class, cv=5, scoring='accuracy')
            results[f'class_knn_k{k}'] = float(scores.mean())
            print(f"  Class KNN (k={k:2d}): {scores.mean():.4f}")
    
    # Silhouette at class level
    try:
        sil = silhouette_score(X, y_class)
        results['class_silhouette'] = float(sil)
        print(f"  Class Silhouette: {sil:.4f}")
    except:
        pass
    
    # Also evaluate at domain level (coarser, fewer classes)
    domain_labels = [s.get('domain', '') for s in samples[:len(embs)]]
    domain_dist = Counter(domain_labels)
    valid_domains = {d for d, n in domain_dist.items() if n >= 10 and d}
    if len(valid_domains) >= 2:
        domain_mask = np.array([s.get('domain', '') in valid_domains 
                               for s in samples[:len(embs)]])
        le_d = LabelEncoder()
        y_domain = le_d.fit_transform([s.get('domain', '') 
                                       for s in samples[:len(embs)] 
                                       if s.get('domain', '') in valid_domains])
        X_d = X[domain_mask]
        if len(X_d) > 20:
            knn_d = KNeighborsClassifier(n_neighbors=10)
            d_scores = cross_val_score(knn_d, X_d, y_domain, cv=5, scoring='accuracy')
            results['domain_knn_k10'] = float(d_scores.mean())
            results['domain_distribution'] = dict(Counter(
                s.get('domain', '') for s in samples[:len(embs)] 
                if s.get('domain', '') in valid_domains))
            print(f"  Domain KNN (k=10): {d_scores.mean():.4f}")
    
    # Save
    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    with open(out_dir / 'taxonomic_class_probe.json', 'w') as fh:
        json.dump(results, fh, indent=2)
    
    # Plot — PCA colored by top-5 classes
    try:
        import matplotlib
        matplotlib.use('Agg')
        import matplotlib.pyplot as plt
        from sklearn.decomposition import PCA
        
        pca = PCA(n_components=2)
        X2 = pca.fit_transform(X)
        
        top5 = [c for c, _ in class_dist.most_common(5)]
        fig, ax = plt.subplots(1, 1, figsize=(10, 7))
        colors_list = ['red', 'blue', 'green', 'orange', 'purple']
        for idx, cls_name in enumerate(top5):
            mask = np.array([s['class'] == cls_name for s in samples[:len(embs)]])
            ax.scatter(X2[mask, 0], X2[mask, 1], c=colors_list[idx], 
                      alpha=0.3, s=10, label=f'{cls_name} (n={mask.sum()})')
        # Plot rest in gray
        rest_mask = np.array([s['class'] not in top5 for s in samples[:len(embs)]])
        ax.scatter(X2[rest_mask, 0], X2[rest_mask, 1], c='lightgray', 
                  alpha=0.1, s=5, label=f'Other ({rest_mask.sum()})')
        ax.set_xlabel(f'PC1 ({pca.explained_variance_ratio_[0]:.1%})')
        ax.set_ylabel(f'PC2 ({pca.explained_variance_ratio_[1]:.1%})')
        ax.set_title(f'Taxonomic Classification ({len(valid_classes)} classes)')
        ax.legend(fontsize=8)
        plt.tight_layout()
        plt.savefig(out_dir / 'taxonomic_class_probe.png', dpi=150)
        plt.close()
    except Exception as e:
        print(f"  Plot error: {e}")
    
    # Publication plots — by domain (cleaner than 90 classes)
    try:
        domain_labels = [s.get('domain', 'unknown') for s in samples[:len(embs)]]
        save_publication_plots(X, domain_labels,
                             'Taxonomic Domain', out_dir,
                             title_prefix='Probe 11a')
        # Also by top-5 classes
        from collections import Counter as Ctr
        top5cls = [c for c, _ in Ctr(s['class'] for s in samples[:len(embs)]).most_common(5)]
        labels_top5 = [s['class'] if s['class'] in top5cls else 'Other' 
                       for s in samples[:len(embs)]]
        save_publication_plots(X, labels_top5,
                             'Top-5 Taxonomic Classes', out_dir,
                             title_prefix='Probe 11b')
    except Exception as e:
        print(f"  Pub plot error: {e}")
    
    return results


# ============================================================
# Probe Registry
# ============================================================
PROBES = {
    # 'ambiguity': probe_ambiguity,  # Deprecated: ? marks are annotation artifacts, not biology
    'composition': probe_composition,
    'knn_purity': probe_knn_purity,
    'polymerization': probe_polymerization,
    'size_prediction': probe_size,
    'link_binary': probe_link_binary,
    'mlm_zeroshot': probe_mlm_zeroshot,
    'token_importance': probe_token_importance,
    'cancer_markers': probe_cancer_markers,
    'glycosylation_type': probe_glycosylation_type,
    'taxonomic_class': probe_taxonomic_class,
}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', choices=['v5', 'v6'], required=True)
    parser.add_argument('--probe', nargs='+', default=['all'],
                        choices=['all'] + list(PROBES.keys()))
    parser.add_argument('--output_dir', default=None)
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--resolved_only', action='store_true',
                        help='Filter out ambiguous WURCS (containing ?) before probing')
    parser.add_argument('--max_samples', type=int, default=5000)
    args = parser.parse_args()

    if args.output_dir is None:
        args.output_dir = str(PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' / f'novel_probes_v2_{args.model}')
    os.makedirs(args.output_dir, exist_ok=True)

    print(f"Loading tokenizer from {VOCAB_PATH}...")
    tokenizer = WURCSTokenizer(str(VOCAB_PATH))
    print(f"  Vocab size: {tokenizer.vocab_size}")

    ckpt = CHECKPOINTS[args.model]
    if not ckpt.exists():
        print(f"ERROR: Checkpoint not found: {ckpt}")
        sys.exit(1)

    model = load_model(str(ckpt), device=args.device)

    # ============================================================
    # PHASE 1: Embed pretraining samples ONCE, save to disk
    # ============================================================
    cache_dir = Path(args.output_dir) / 'embedding_cache'
    cache_dir.mkdir(parents=True, exist_ok=True)
    emb_npy = cache_dir / 'cls_embeddings.npy'
    samples_pkl = cache_dir / 'samples.pkl'

    if emb_npy.exists() and samples_pkl.exists():
        print(f"\n  Loading cached embeddings from {cache_dir}...")
        pretrain_embs = np.load(str(emb_npy))
        import pickle
        with open(samples_pkl, 'rb') as pf:
            pretrain_samples = pickle.load(pf)
        print(f"  Loaded: {pretrain_embs.shape[0]} embeddings, {len(pretrain_samples)} samples")
    else:
        print(f"\n  Embedding pretraining samples (embed once, save to disk)...")
        pretrain_samples = load_pretrain_wurcs(tokenizer, max_n=args.max_samples)
        
        if args.resolved_only:
            before = len(pretrain_samples)
            pretrain_samples = [s for s in pretrain_samples if '?' not in s.get('wurcs', '')]
            print(f"  Filtered resolved: {before} -> {len(pretrain_samples)} (removed {before-len(pretrain_samples)} ambiguous)")
        
        print(f"  Embedding {len(pretrain_samples)} samples...")
        pretrain_embs_list = batch_cls_embeddings(model, pretrain_samples, device=args.device)
        pretrain_embs = np.array(pretrain_embs_list)
        del pretrain_embs_list
        
        # Save to disk
        np.save(str(emb_npy), pretrain_embs)
        import pickle
        with open(samples_pkl, 'wb') as pf:
            pickle.dump(pretrain_samples, pf)
        print(f"  Saved: {emb_npy} ({pretrain_embs.nbytes / 1e9:.2f} GB)")
    
    import gc
    gc.collect()
    print(f"  Pretrain embeddings: shape={pretrain_embs.shape}")

    # ============================================================
    # PHASE 2: Run probes
    # ============================================================
    # Probes that use pretraining data get pre-computed embeddings
    # Probes that use external data (GlycanML) still call batch_cls_embeddings
    
    PRETRAIN_PROBES = {'composition', 'polymerization', 
                       'size_prediction', 'mlm_zeroshot',
                       'token_importance', 'ambiguity'}
    EXTERNAL_PROBES = {'cancer_markers', 'glycosylation_type', 'taxonomic_class',
                       'knn_purity', 'link_binary'}
    
    probes_to_run = list(PROBES.keys()) if 'all' in args.probe else args.probe
    all_results = {}

    for pn in probes_to_run:
        try:
            if pn in PRETRAIN_PROBES:
                # Pass pre-computed embeddings — probe won't re-embed
                all_results[pn] = PROBES[pn](
                    model, tokenizer, args.device, args.output_dir, args.max_samples,
                    _cached_embs=pretrain_embs, _cached_samples=pretrain_samples
                )
            else:
                # External probes embed their own (small) datasets
                all_results[pn] = PROBES[pn](
                    model, tokenizer, args.device, args.output_dir, args.max_samples
                )
        except Exception as e:
            print(f"\n  ERROR in '{pn}': {e}")
            import traceback; traceback.print_exc()
            all_results[pn] = {'error': str(e)}
        
        gc.collect()  # Free memory between probes
    
    # Save combined results
    with open(os.path.join(args.output_dir, 'all_probe_results_v2.json'), 'w') as f:
        json.dump(all_results, f, indent=2, default=str)
    
    print(f"\nALL PROBES COMPLETE — V2 ({args.model.upper()})")
    print(f"Results: {args.output_dir}")

if __name__ == '__main__':
    main()