code/probes/probe_9_biosynthesis_order.py

#!/usr/bin/env python3
"""
Probe 9: N-Glycan Biosynthesis Pathway Ordering
=================================================
Tests whether GlycanBERT embeddings recapitulate the N-glycan maturation
pathway — a strict enzymatic ordering from ER to trans-Golgi:

  Stage 1: High-Mannose (Man5–Man9)  — ER / cis-Golgi
  Stage 2: Hybrid                    — medial-Golgi (GnT-I acted)
  Stage 3: Complex (basic)           — medial-Golgi (GnT-II acted)
  Stage 4: Decorated Complex         — trans-Golgi (core fucose, bisecting GlcNAc)
  Stage 5: Capped/Terminal           — trans-Golgi (sialylation, Lewis epitopes)

Biological hypothesis: If GlycanBERT learned glycan biology, glycans at
adjacent biosynthesis stages (|i−j|=1) should be CLOSER in embedding space
than glycans at distant stages (|i−j|≥3).

Metrics:
  - Spearman ρ between |stage_i − stage_j| and cos_distance(emb_i, emb_j)
  - Within-stage vs between-stage distance ratio (silhouette-like)
  - t-SNE colored by biosynthesis stage

Data source: glycowork_iupac_wurcs_unified.csv (32,428 glycans with IUPAC+WURCS)
Motif source: glycowork.motif.annotate — 165 curated motifs

Usage:
  python probe_9_biosynthesis_order.py --model v6 --device cuda
  python probe_9_biosynthesis_order.py --model v5 --device cuda
"""

import sys
import os
import json
import argparse
import numpy as np
import pandas as pd
from pathlib import Path
from collections import Counter
from scipy.stats import spearmanr
from scipy.spatial.distance import pdist, squareform
from sklearn.manifold import TSNE
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import torch

# ─── Project paths ────────────────────────────────────────────────────────────
PROJECT_ROOT = Path(__file__).resolve().parents[2]
VOCAB_PATH = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'vocabulary.json'
DATA_PATH = PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' / 'glycowork_iupac_wurcs_unified.csv'

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',
}

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

# ─── Nature BGP color palette ────────────────────────────────────────────────
STAGE_COLORS = {
    1: '#0072B2',  # Blue  — High-mannose
    2: '#009E73',  # Green — Hybrid
    3: '#E69F00',  # Gold  — Complex
    4: '#D55E00',  # Red   — Decorated
    5: '#CC79A7',  # Pink  — Capped
}
STAGE_LABELS = {
    1: 'Stage 1: High-Mannose\n(ER / cis-Golgi)',
    2: 'Stage 2: Hybrid\n(medial-Golgi, GnT-I)',
    3: 'Stage 3: Complex\n(medial-Golgi, GnT-II)',
    4: 'Stage 4: Decorated\n(trans-Golgi, Fuc/bisect)',
    5: 'Stage 5: Capped\n(trans-Golgi, Sia/Lewis)',
}

# ─── Biosynthesis stage assignment ────────────────────────────────────────────
def assign_biosynthesis_stage(motif_row):
    """
    Assign a glycan to one of 5 N-glycan biosynthesis stages based on
    glycowork motif annotations. Uses a hierarchical rule system that
    mirrors the actual enzymatic pathway:

    The key insight: the N-glycan pathway is ORDERED. A glycan can only
    reach stage N if it has passed through stages 1...N-1. So we check
    from most mature (stage 5) backwards.

    Returns: stage (1-5) or 0 if not an N-glycan
    """
    # Must have the chitobiose core (GlcNAc-β1-4-GlcNAc) to be an N-glycan
    if motif_row.get('Chitobiose', 0) == 0 and motif_row.get('Trimannosylcore', 0) == 0:
        return 0  # Not an N-glycan

    # Stage 5: Capped — terminal sialylation, Lewis epitopes, poly-LacNAc
    sialyl_markers = ['SialylLewisX', 'SialylLewisA', 'GM3', 'GM2', 'GM1',
                       'GD3', 'GD1a', 'GD2', 'GD1b', 'GT1b', 'polySia',
                       'DisialylLewisA', 'DisialylLewisC']
    if any(motif_row.get(m, 0) > 0 for m in sialyl_markers):
        return 5

    # Stage 4: Decorated — core fucose, bisecting GlcNAc, tetra-antennary
    decoration_markers = ['core_fucose', 'bisectingGlcNAc', 'Tetraantennary_Nglycan',
                          'Difucosylated_core', 'GalFuc_core']
    if any(motif_row.get(m, 0) > 0 for m in decoration_markers):
        return 4

    # Stage 3: Complex — has complex N-glycan motifs (LacNAc extensions)
    complex_markers = ['Nglycan_complex', 'Nglycan_complex2',
                       'Internal_LacNAc_type2', 'Terminal_LacNAc_type2',
                       'Internal_LacNAc_type1', 'Terminal_LacNAc_type1',
                       'PolyLacNAc', 'I_antigen', 'i_antigen',
                       'Terminal_LewisX', 'Internal_LewisX', 'LewisY',
                       'Terminal_LewisA', 'Internal_LewisA', 'LewisB',
                       'H_antigen_type2', 'H_antigen_type1',
                       'A_antigen', 'B_antigen', 'Galili_antigen']
    if any(motif_row.get(m, 0) > 0 for m in complex_markers):
        return 3

    # Stage 2: Hybrid — has hybrid N-glycan motif
    if motif_row.get('Nglycan_hybrid', 0) > 0:
        return 2

    # Stage 1: High-mannose — has high_mannose motif or just trimannosylcore
    if motif_row.get('high_mannose', 0) > 0 or motif_row.get('Trimannosylcore', 0) > 0:
        return 1

    # Has chitobiose but no clear stage — likely unusual N-glycan
    if motif_row.get('Chitobiose', 0) > 0:
        return 1  # Default to high-mannose precursor

    return 0  # Not assignable


def annotate_glycans_with_stages(iupac_list):
    """
    Annotate a list of IUPAC glycan strings with biosynthesis stages
    using glycowork's motif annotation.
    """
    from glycowork.motif.annotate import annotate_glycan

    stages = []
    motif_dfs = []
    valid_indices = []

    for i, iupac in enumerate(iupac_list):
        try:
            df = annotate_glycan(iupac)
            if df is not None and len(df) > 0:
                row = df.iloc[0].to_dict()
                stage = assign_biosynthesis_stage(row)
                stages.append(stage)
                motif_dfs.append(row)
                valid_indices.append(i)
            else:
                stages.append(0)
                valid_indices.append(i)
                motif_dfs.append({})
        except Exception:
            stages.append(0)
            valid_indices.append(i)
            motif_dfs.append({})
        if (i + 1) % 1000 == 0:
            print(f"  Annotated {i+1}/{len(iupac_list)}...")

    return stages, motif_dfs, valid_indices


# ─── Model loading (same pattern as probe_8) ─────────────────────────────────
def load_model(model_version, device):
    """Load GlycanBERT model using MultimodalGlycanBERT (same as probe_8)."""
    ckpt_path = CHECKPOINTS[model_version]
    print(f"Loading {model_version} from {ckpt_path}")

    state = torch.load(ckpt_path, map_location='cpu', weights_only=False)
    sd = state.get('model_state_dict', state)
    if 'proj_head_state_dict' in state:
        sd = {k: v for k, v in sd.items() if not k.startswith('proj_head')}

    emb_weight = sd.get('seq_embeddings.token_embeddings.weight',
                        sd.get('token_embeddings.weight'))
    vocab_size = emb_weight.shape[0] if emb_weight is not None else 2200
    hidden = emb_weight.shape[1] if emb_weight is not None else 768

    config = MultimodalGlycanBERTConfig(
        seq_vocab_size=vocab_size, seq_hidden_size=hidden,
        seq_num_layers=12, seq_num_heads=12, seq_max_length=256,
        use_cnn_frontend=True, cnn_kernel_size=3,
    )
    model = MultimodalGlycanBERT(config)
    model.load_state_dict(sd, strict=False)
    model = model.to(device).eval()
    print(f"  Loaded: {sum(p.numel() for p in model.parameters()):,} params, "
          f"vocab={vocab_size}, hidden={hidden}")

    tokenizer = WURCSTokenizer(str(VOCAB_PATH))
    return model, tokenizer


def get_cls_embeddings(model, tokenizer, wurcs_list, device, batch_size=128, max_len=256):
    """Extract CLS embeddings using WURCSTokenizer (same as probe_8)."""
    all_embs = []
    errors = 0

    for i in range(0, len(wurcs_list), batch_size):
        batch = wurcs_list[i:i+batch_size]
        token_ids_list, bd_list, lt_list = [], [], []
        for w in batch:
            try:
                tok_out = tokenizer.tokenize(w)
                ids = tok_out['token_ids'][:max_len]
                bd = tok_out['branch_depths'][:max_len]
                lt = tok_out['linkage_types'][:max_len]
                token_ids_list.append(ids)
                bd_list.append(bd)
                lt_list.append(lt)
            except Exception:
                errors += 1
                continue
        if not token_ids_list:
            continue

        max_l = max(len(x) for x in token_ids_list)
        padded_ids = torch.zeros(len(token_ids_list), max_l, dtype=torch.long)
        padded_bd = torch.zeros_like(padded_ids)
        padded_lt = torch.zeros_like(padded_ids)
        for j, (ids, bd, lt) in enumerate(zip(token_ids_list, bd_list, lt_list)):
            padded_ids[j, :len(ids)] = torch.tensor(ids, dtype=torch.long)
            padded_bd[j, :len(bd)] = torch.tensor(bd, dtype=torch.long)
            padded_lt[j, :len(lt)] = torch.tensor(lt, dtype=torch.long)

        padded_ids = padded_ids.to(device)
        padded_bd = padded_bd.to(device)
        padded_lt = padded_lt.to(device)

        with torch.no_grad():
            seq_out = model.seq_embeddings(padded_ids, branch_depths=padded_bd,
                                           linkage_types=padded_lt)
            cls_emb = seq_out[:, 0, :].cpu().numpy()
            all_embs.append(cls_emb)

        if (i // batch_size) % 10 == 0:
            print(f"  Embedded {i+len(batch)}/{len(wurcs_list)} ({errors} errors)")

    print(f"  Total embedded: {sum(e.shape[0] for e in all_embs):,} ({errors} errors)")
    return np.vstack(all_embs) if all_embs else np.zeros((0, 768))


# ─── Probe analyses ──────────────────────────────────────────────────────────
def compute_pathway_correlation(embeddings, stages, output_dir):
    """
    Core metric: Spearman ρ between |stage difference| and cosine distance.
    If the model learned biosynthesis ordering, this should be positive.
    """
    print("\n=== Pathway Correlation Analysis ===")

    # Subsample for tractability (pairwise = O(n²))
    n = len(embeddings)
    if n > 2000:
        idx = np.random.RandomState(42).choice(n, 2000, replace=False)
        emb_sub = embeddings[idx]
        stg_sub = np.array(stages)[idx]
    else:
        emb_sub = embeddings
        stg_sub = np.array(stages)

    # Cosine distances
    cos_dists = pdist(emb_sub, metric='cosine')

    # Stage differences
    n_sub = len(emb_sub)
    stage_diffs = []
    for i in range(n_sub):
        for j in range(i+1, n_sub):
            stage_diffs.append(abs(stg_sub[i] - stg_sub[j]))
    stage_diffs = np.array(stage_diffs)

    # Spearman correlation
    rho, pval = spearmanr(stage_diffs, cos_dists)
    print(f"  Spearman ρ (|Δstage| vs cos_dist): {rho:.4f}  (p={pval:.2e})")
    print(f"    Interpretation: {'Model captures maturation ordering' if rho > 0.15 else 'Weak/no pathway awareness'}")

    # Within-stage vs between-stage distances
    stages_arr = np.array(stages)
    cos_matrix = squareform(cos_dists) if len(emb_sub) == len(embeddings) else None

    # Compute per-stage-pair mean distances
    stage_pairs = {}
    pair_idx = 0
    for i in range(n_sub):
        for j in range(i+1, n_sub):
            key = (min(stg_sub[i], stg_sub[j]), max(stg_sub[i], stg_sub[j]))
            if key not in stage_pairs:
                stage_pairs[key] = []
            stage_pairs[key].append(cos_dists[pair_idx])
            pair_idx += 1

    print(f"\n  Mean cosine distance by stage pair:")
    within_dists = []
    between_dists = []
    pair_results = {}
    for (s1, s2) in sorted(stage_pairs.keys()):
        mean_d = np.mean(stage_pairs[(s1, s2)])
        std_d = np.std(stage_pairs[(s1, s2)])
        n_pairs = len(stage_pairs[(s1, s2)])
        pair_results[(s1, s2)] = {'mean': mean_d, 'std': std_d, 'n': n_pairs}
        label = f"({s1},{s2})"
        print(f"    {label:>8s}: {mean_d:.4f} ± {std_d:.4f}  (n={n_pairs})")
        if s1 == s2:
            within_dists.extend(stage_pairs[(s1, s2)])
        else:
            between_dists.extend(stage_pairs[(s1, s2)])

    within_mean = np.mean(within_dists) if within_dists else 0
    between_mean = np.mean(between_dists) if between_dists else 0
    ratio = within_mean / between_mean if between_mean > 0 else float('inf')
    print(f"\n  Within-stage mean distance:  {within_mean:.4f}")
    print(f"  Between-stage mean distance: {between_mean:.4f}")
    print(f"  Ratio (within/between):      {ratio:.4f}")
    print(f"    Interpretation: {'Good clustering' if ratio < 0.85 else 'Moderate' if ratio < 0.95 else 'Weak clustering'}")

    results = {
        'spearman_rho': float(rho),
        'spearman_pval': float(pval),
        'within_stage_mean_dist': float(within_mean),
        'between_stage_mean_dist': float(between_mean),
        'within_between_ratio': float(ratio),
        'pair_distances': {f"({k[0]},{k[1]})": {'mean': float(v['mean']), 'n': int(v['n'])}
                          for k, v in pair_results.items()},
        'n_samples': int(n_sub),
    }

    with open(output_dir / 'pathway_correlation.json', 'w') as f:
        json.dump(results, f, indent=2)

    return results


def plot_stage_distance_heatmap(results, output_dir):
    """Heatmap of mean cosine distance between biosynthesis stages."""
    fig, ax = plt.subplots(figsize=(7, 6))

    # Build 5x5 matrix
    matrix = np.zeros((5, 5))
    for key_str, vals in results['pair_distances'].items():
        s1, s2 = int(key_str[1]), int(key_str[3])
        matrix[s1-1, s2-1] = vals['mean']
        matrix[s2-1, s1-1] = vals['mean']

    labels = ['High-\nMannose', 'Hybrid', 'Complex', 'Decorated', 'Capped/\nSialylated']
    im = ax.imshow(matrix, cmap='RdYlBu_r', aspect='equal')
    ax.set_xticks(range(5))
    ax.set_yticks(range(5))
    ax.set_xticklabels(labels, fontsize=9)
    ax.set_yticklabels(labels, fontsize=9)

    # Add text annotations
    for i in range(5):
        for j in range(5):
            ax.text(j, i, f'{matrix[i,j]:.3f}', ha='center', va='center',
                    fontsize=10, fontweight='bold',
                    color='white' if matrix[i,j] > 0.5 * matrix.max() else 'black')

    ax.set_title(f'Mean Cosine Distance Between Biosynthesis Stages\n'
                 f'Spearman ρ = {results["spearman_rho"]:.3f} (p = {results["spearman_pval"]:.1e})',
                 fontsize=12, fontweight='bold')
    plt.colorbar(im, ax=ax, label='Cosine Distance', shrink=0.8)
    plt.tight_layout()
    fig.savefig(output_dir / 'stage_distance_heatmap.png', dpi=300, bbox_inches='tight')
    fig.savefig(output_dir / 'stage_distance_heatmap.pdf', bbox_inches='tight')
    plt.close()
    print(f"  Saved heatmap to {output_dir / 'stage_distance_heatmap.png'}")


def plot_tsne_by_stage(embeddings, stages, output_dir, max_points=5000):
    """t-SNE visualization colored by biosynthesis stage."""
    n = len(embeddings)
    if n > max_points:
        idx = np.random.RandomState(42).choice(n, max_points, replace=False)
        emb = embeddings[idx]
        stg = np.array(stages)[idx]
    else:
        emb = embeddings
        stg = np.array(stages)

    print(f"\n=== t-SNE Visualization ({len(emb)} points) ===")
    tsne = TSNE(n_components=2, perplexity=30, random_state=42, max_iter=1000)
    coords = tsne.fit_transform(emb)

    fig, ax = plt.subplots(figsize=(10, 8))
    for stage in sorted(STAGE_COLORS.keys()):
        mask = stg == stage
        if mask.sum() > 0:
            ax.scatter(coords[mask, 0], coords[mask, 1],
                      c=STAGE_COLORS[stage], label=STAGE_LABELS[stage],
                      alpha=0.5, s=15, edgecolors='none')

    ax.set_xlabel('t-SNE 1', fontsize=12)
    ax.set_ylabel('t-SNE 2', fontsize=12)
    ax.set_title('GlycanBERT Embeddings Colored by N-Glycan Biosynthesis Stage',
                 fontsize=13, fontweight='bold')
    ax.legend(fontsize=8, loc='best', framealpha=0.9, markerscale=2)
    ax.set_xticks([])
    ax.set_yticks([])

    plt.tight_layout()
    fig.savefig(output_dir / 'tsne_biosynthesis_stages.png', dpi=300, bbox_inches='tight')
    fig.savefig(output_dir / 'tsne_biosynthesis_stages.pdf', bbox_inches='tight')
    plt.close()
    print(f"  Saved t-SNE to {output_dir / 'tsne_biosynthesis_stages.png'}")


def plot_stage_distribution(stages, output_dir):
    """Bar chart of glycan counts per biosynthesis stage."""
    counts = Counter(stages)
    fig, ax = plt.subplots(figsize=(8, 5))

    stage_nums = sorted([s for s in counts if s > 0])
    bars = ax.bar([str(s) for s in stage_nums],
                  [counts[s] for s in stage_nums],
                  color=[STAGE_COLORS[s] for s in stage_nums],
                  edgecolor='white', linewidth=0.5)

    # Labels inside bars
    for bar, s in zip(bars, stage_nums):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() - 20,
                f'n={counts[s]}', ha='center', va='top',
                fontsize=10, fontweight='bold', color='white')

    short_labels = ['High-\nMannose', 'Hybrid', 'Complex', 'Decorated', 'Capped']
    ax.set_xticks(range(len(stage_nums)))
    ax.set_xticklabels([short_labels[s-1] for s in stage_nums], fontsize=10)
    ax.set_ylabel('Number of Glycans', fontsize=12)
    ax.set_title('Distribution of N-Glycans Across Biosynthesis Stages', fontsize=13, fontweight='bold')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    not_nglycan = counts.get(0, 0)
    ax.text(0.98, 0.95, f'Not N-glycan: {not_nglycan}',
            transform=ax.transAxes, ha='right', va='top', fontsize=9, color='gray')

    plt.tight_layout()
    fig.savefig(output_dir / 'stage_distribution.png', dpi=300, bbox_inches='tight')
    plt.close()
    print(f"  Saved distribution to {output_dir / 'stage_distribution.png'}")


def plot_adjacent_vs_distant(results, output_dir):
    """Compare adjacent (|Δ|=1) vs distant (|Δ|≥3) stage distances."""
    pair_dists = results['pair_distances']

    adjacent = []   # |Δ|=1
    moderate = []   # |Δ|=2
    distant = []    # |Δ|≥3

    for key_str, vals in pair_dists.items():
        s1, s2 = int(key_str[1]), int(key_str[3])
        delta = abs(s1 - s2)
        if delta == 0:
            continue
        elif delta == 1:
            adjacent.append(vals['mean'])
        elif delta == 2:
            moderate.append(vals['mean'])
        else:
            distant.append(vals['mean'])

    fig, ax = plt.subplots(figsize=(6, 5))
    categories = ['Adjacent\n(|Δ|=1)', 'Moderate\n(|Δ|=2)', 'Distant\n(|Δ|≥3)']
    means = [np.mean(adjacent) if adjacent else 0,
             np.mean(moderate) if moderate else 0,
             np.mean(distant) if distant else 0]
    colors = ['#009E73', '#E69F00', '#D55E00']

    bars = ax.bar(categories, means, color=colors, edgecolor='white', width=0.6)
    for bar, m in zip(bars, means):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.005,
                f'{m:.3f}', ha='center', va='bottom', fontsize=11, fontweight='bold')

    ax.set_ylabel('Mean Cosine Distance', fontsize=12)
    ax.set_title('Embedding Distance vs Biosynthesis Stage Separation\n'
                 '(Adjacent stages should be closer than distant)',
                 fontsize=12, fontweight='bold')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    plt.tight_layout()
    fig.savefig(output_dir / 'adjacent_vs_distant.png', dpi=300, bbox_inches='tight')
    fig.savefig(output_dir / 'adjacent_vs_distant.pdf', bbox_inches='tight')
    plt.close()
    print(f"  Saved adjacent vs distant to {output_dir / 'adjacent_vs_distant.png'}")


# ─── Main ─────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description='Probe 9: N-Glycan Biosynthesis Pathway Ordering')
    parser.add_argument('--model', choices=['v5', 'v6'], required=True)
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--max-glycans', type=int, default=15000,
                        help='Max glycans to process (for speed)')
    args = parser.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    output_dir = PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' / 'probing_analysis' / f'09_biosynthesis_pathway_{args.model}'
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"{'='*60}")
    print(f"Probe 9: N-Glycan Biosynthesis Pathway Ordering ({args.model})")
    print(f"{'='*60}")

    # ── Load data ─────────────────────────────────────────────────────────
    print(f"\n1. Loading data from {DATA_PATH}")
    df = pd.read_csv(DATA_PATH)
    print(f"   Total glycans: {len(df)}")

    # Column names: glycan (IUPAC), wurcs (WURCS), glytoucan_id, disease_association, tissue_sample, glycan_type
    iupac_col = 'glycan'   # IUPAC condensed notation
    wurcs_col = 'wurcs'    # WURCS notation for model
    print(f"   Columns: {list(df.columns)}")

    # Filter to rows with both
    mask = df[iupac_col].notna() & df[wurcs_col].notna()
    df = df[mask].head(args.max_glycans).reset_index(drop=True)
    print(f"   After filtering (IUPAC+WURCS present): {len(df)}")

    # ── Annotate with glycowork ───────────────────────────────────────────
    print(f"\n2. Annotating {len(df)} glycans with glycowork motifs...")
    stages, motif_data, valid_indices = annotate_glycans_with_stages(df[iupac_col].tolist())

    stage_counts = Counter(stages)
    print(f"   Stage distribution:")
    for s in sorted(stage_counts.keys()):
        label = STAGE_LABELS.get(s, 'Not N-glycan') if s > 0 else 'Not N-glycan'
        print(f"     Stage {s}: {stage_counts[s]:5d}  ({label.split(chr(10))[0]})")

    # Filter to only N-glycans (stage > 0)
    nglycan_mask = [s > 0 for s in stages]
    nglycan_df = df[nglycan_mask].reset_index(drop=True)
    nglycan_stages = [s for s in stages if s > 0]
    print(f"   N-glycans for analysis: {len(nglycan_stages)}")

    if len(nglycan_stages) < 100:
        print("ERROR: Too few N-glycans found. Check data/annotations.")
        return

    # Plot stage distribution
    plot_stage_distribution(stages, output_dir)

    # ── Load model & extract embeddings ───────────────────────────────────
    print(f"\n3. Loading model and extracting CLS embeddings...")
    model, tokenizer = load_model(args.model, device)
    embeddings = get_cls_embeddings(model, tokenizer, nglycan_df[wurcs_col].tolist(), device)
    print(f"   Embeddings shape: {embeddings.shape}")

    # ── Pathway correlation ───────────────────────────────────────────────
    print(f"\n4. Computing pathway correlation...")
    results = compute_pathway_correlation(embeddings, nglycan_stages, output_dir)

    # ── Visualizations ────────────────────────────────────────────────────
    print(f"\n5. Generating visualizations...")
    plot_stage_distance_heatmap(results, output_dir)
    plot_tsne_by_stage(embeddings, nglycan_stages, output_dir)
    plot_adjacent_vs_distant(results, output_dir)

    # ── Summary ──────────────────────────────────────────────────────────
    print(f"\n{'='*60}")
    print(f"PROBE 9 SUMMARY ({args.model})")
    print(f"{'='*60}")
    print(f"  N-glycans analyzed:    {len(nglycan_stages)}")
    print(f"  Spearman ρ:            {results['spearman_rho']:.4f}  (p={results['spearman_pval']:.2e})")
    print(f"  Within/between ratio:  {results['within_between_ratio']:.4f}")
    print(f"  Output directory:      {output_dir}")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()