code/probes/probe_exploratory_biology.py

#!/usr/bin/env python3
"""
Exploratory Biology Probes — GlycanBERT Embedding Analysis
============================================================
Generates t-SNE and UMAP visualizations colored by biologically meaningful
properties, with quantitative clustering metrics (k-NN purity, silhouette).

Themes covered:
  A) Lectin-binding motifs (galectin-3, mannose receptor, DC-SIGN, core fucose,
     sialic acid linkage)
  B) Cancer glycan signatures (Tn antigen, core fucose, poly-LacNAc, branching)
  C) Glycan type (N-linked, O-linked, etc.)
  D) Tissue groupings (organ systems from glycowork metadata)

No logistic regression — purely exploratory visualization + clustering metrics.
"""

import sys, os, json, argparse, warnings
import numpy as np
import pandas as pd
from pathlib import Path
from collections import Counter

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import torch

warnings.filterwarnings('ignore', category=FutureWarning)

# ─── 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

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


def load_model(model_version, device):
    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")
    tokenizer = WURCSTokenizer(str(VOCAB_PATH))
    return model, tokenizer


def get_cls_embeddings(model, tokenizer, wurcs_list, device,
                       batch_size=128, max_len=256):
    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 = tokenizer.tokenize(w)
                token_ids_list.append(tok['token_ids'][:max_len])
                bd_list.append(tok['branch_depths'][:max_len])
                lt_list.append(tok['linkage_types'][:max_len])
            except Exception:
                errors += 1
                continue
        if not token_ids_list:
            continue
        ml = max(len(x) for x in token_ids_list)
        ids_t = torch.zeros(len(token_ids_list), ml, dtype=torch.long)
        bd_t  = torch.zeros_like(ids_t)
        lt_t  = torch.zeros_like(ids_t)
        for j, (ids, bd, lt) in enumerate(zip(token_ids_list, bd_list, lt_list)):
            ids_t[j, :len(ids)] = torch.tensor(ids, dtype=torch.long)
            bd_t[j, :len(bd)]  = torch.tensor(bd, dtype=torch.long)
            lt_t[j, :len(lt)]  = torch.tensor(lt, dtype=torch.long)
        ids_t, bd_t, lt_t = ids_t.to(device), bd_t.to(device), lt_t.to(device)
        with torch.no_grad():
            seq_out = model.seq_embeddings(ids_t, branch_depths=bd_t,
                                           linkage_types=lt_t)
            all_embs.append(seq_out[:, 0, :].cpu().numpy())
        if (i // batch_size) % 10 == 0:
            print(f"  Embedded {i+len(batch)}/{len(wurcs_list)} ({errors} errors)")
    print(f"  Total: {sum(e.shape[0] for e in all_embs):,} ({errors} errors)")
    return np.vstack(all_embs) if all_embs else np.zeros((0, 768))


# ═══ Clustering metrics ══════════════════════════════════════════════════════

def knn_label_purity(embeddings, labels, k_values=[10, 20, 50]):
    from sklearn.neighbors import NearestNeighbors
    unique_labels = sorted(set(labels))
    if len(unique_labels) < 2:
        return {k: (1.0, 1.0) for k in k_values}
    label_arr = np.array(labels)
    results = {}
    for k in k_values:
        k_actual = min(k, len(embeddings) - 1)
        nn = NearestNeighbors(n_neighbors=k_actual + 1, metric='cosine')
        nn.fit(embeddings)
        _, indices = nn.kneighbors(embeddings)
        purities = []
        for i in range(len(embeddings)):
            neighbors = indices[i, 1:]
            same_label = np.sum(label_arr[neighbors] == label_arr[i])
            purities.append(same_label / len(neighbors))
        avg_purity = np.mean(purities)
        label_counts = Counter(labels)
        baseline = sum((c / len(labels)) ** 2 for c in label_counts.values())
        results[k] = (float(avg_purity), float(baseline))
        print(f"    k-NN purity (k={k}): {avg_purity:.4f} (baseline={baseline:.4f}, lift={avg_purity/baseline:.2f}x)")
    return results


def compute_silhouette(embeddings, labels, sample_size=5000):
    from sklearn.metrics import silhouette_score
    unique = set(labels)
    if len(unique) < 2:
        print("    Silhouette: N/A (only 1 class)")
        return None
    label_arr = np.array(labels)
    if len(embeddings) > sample_size:
        np.random.seed(42)
        idx = np.random.choice(len(embeddings), sample_size, replace=False)
        embs_sub, labels_sub = embeddings[idx], label_arr[idx]
    else:
        embs_sub, labels_sub = embeddings, label_arr
    if len(set(labels_sub)) < 2:
        print("    Silhouette: N/A (sampling left only 1 class)")
        return None
    score = silhouette_score(embs_sub, labels_sub, metric='cosine',
                             sample_size=min(2000, len(embs_sub)))
    print(f"    Silhouette score: {score:.4f}")
    return float(score)


# ═══ Dimensionality reduction ═════════════════════════════════════════════════

def compute_dim_reduction(embeddings, max_points=10000):
    from sklearn.manifold import TSNE
    n = len(embeddings)
    if n > max_points:
        np.random.seed(42)
        idx = np.random.choice(n, max_points, replace=False)
        embs_sub = embeddings[idx]
    else:
        idx = np.arange(n)
        embs_sub = embeddings
    print(f"\n  Computing t-SNE on {len(embs_sub)} samples...")
    tsne = TSNE(n_components=2, perplexity=40, random_state=42,
                max_iter=1000, learning_rate='auto', init='pca')
    tsne_coords = tsne.fit_transform(embs_sub)
    try:
        import umap
        print(f"  Computing UMAP on {len(embs_sub)} samples...")
        reducer = umap.UMAP(n_components=2, n_neighbors=30, min_dist=0.3,
                           random_state=42, metric='cosine')
        umap_coords = reducer.fit_transform(embs_sub)
    except ImportError:
        print("  UMAP not available, skipping...")
        umap_coords = None
    return tsne_coords, umap_coords, idx


# ═══ Visualization ════════════════════════════════════════════════════════════

def plot_multi_panel_binary(coords, label_dict, title, output_path, method='t-SNE'):
    n_panels = len(label_dict)
    if n_panels == 0:
        return
    cols = min(3, n_panels)
    rows = (n_panels + cols - 1) // cols
    fig, axes = plt.subplots(rows, cols, figsize=(6*cols, 5*rows))
    if n_panels == 1:
        axes = np.array([axes])
    axes = axes.flatten()
    for i, (name, (labels, color)) in enumerate(label_dict.items()):
        ax = axes[i]
        neg_mask = np.array(labels) == 0
        pos_mask = np.array(labels) == 1
        ax.scatter(coords[neg_mask, 0], coords[neg_mask, 1],
                   c='#E8E8E8', s=3, alpha=0.15, edgecolors='none', rasterized=True)
        ax.scatter(coords[pos_mask, 0], coords[pos_mask, 1],
                   c=color, s=12, alpha=0.6, edgecolors='none', rasterized=True,
                   label=f'{name} ({pos_mask.sum()})')
        ax.legend(frameon=False, fontsize=8, loc='upper right')
        ax.set_xlabel(f'{method} 1')
        ax.set_ylabel(f'{method} 2')
        ax.set_title(name, fontsize=10)
    for j in range(n_panels, len(axes)):
        axes[j].set_visible(False)
    plt.suptitle(title, fontsize=13, fontweight='bold')
    plt.tight_layout()
    plt.savefig(str(output_path) + '.png', dpi=300, facecolor='white')
    plt.savefig(str(output_path) + '.pdf', facecolor='white')
    plt.close()
    print(f"    Saved: {output_path}.png")


def plot_multi_class_overlay(coords, labels, title, color_map, output_path, method='t-SNE'):
    fig, ax = plt.subplots(figsize=(9, 7))
    unique_labels = sorted(set(labels))
    for label in unique_labels:
        mask = np.array(labels) == label
        color = color_map.get(label, '#999999')
        ax.scatter(coords[mask, 0], coords[mask, 1],
                   c=color, s=8, alpha=0.5, edgecolors='none', rasterized=True,
                   label=f'{label} ({mask.sum()})')
    ax.legend(frameon=False, fontsize=8, loc='upper right', markerscale=2)
    ax.set_xlabel(f'{method} 1')
    ax.set_ylabel(f'{method} 2')
    ax.set_title(title, fontsize=12, fontweight='bold')
    plt.tight_layout()
    plt.savefig(str(output_path) + '.png', dpi=300, facecolor='white')
    plt.savefig(str(output_path) + '.pdf', facecolor='white')
    plt.close()
    print(f"    Saved: {output_path}.png")


def plot_continuous_overlay(coords, values, title, output_path, cmap='viridis',
                           method='t-SNE', label='value'):
    fig, ax = plt.subplots(figsize=(9, 7))
    sc = ax.scatter(coords[:, 0], coords[:, 1], c=values, cmap=cmap,
                    s=6, alpha=0.5, edgecolors='none', rasterized=True)
    cbar = plt.colorbar(sc, ax=ax, shrink=0.8)
    cbar.set_label(label)
    ax.set_xlabel(f'{method} 1')
    ax.set_ylabel(f'{method} 2')
    ax.set_title(title, fontsize=12, fontweight='bold')
    plt.tight_layout()
    plt.savefig(str(output_path) + '.png', dpi=300, facecolor='white')
    plt.savefig(str(output_path) + '.pdf', facecolor='white')
    plt.close()
    print(f"    Saved: {output_path}.png")


# ═══ Biological label assignment ══════════════════════════════════════════════

def assign_lectin_labels(annotated_df, motif_cols, iupac_list):
    lectin_classes = {}

    lacnac_cols = [c for c in motif_cols if 'lacnac' in c.lower() or 'polylacnac' in c.lower()]
    if lacnac_cols:
        labels = (annotated_df[lacnac_cols].sum(axis=1) > 0).astype(int).values
        lectin_classes['Galectin-3 (LacNAc)'] = (labels, '#0072B2')

    hm_cols = [c for c in motif_cols if 'high_mannose' in c.lower()]
    if hm_cols:
        labels = (annotated_df[hm_cols].sum(axis=1) > 0).astype(int).values
        lectin_classes['Mannose receptor (high-man)'] = (labels, '#009E73')

    lewis_cols = [c for c in motif_cols if 'lewis' in c.lower()]
    if lewis_cols:
        labels = (annotated_df[lewis_cols].sum(axis=1) > 0).astype(int).values
        lectin_classes['DC-SIGN (Lewis)'] = (labels, '#E69F00')

    fuc_cols = [c for c in motif_cols if 'core_fucose' in c.lower()]
    if fuc_cols:
        labels = (annotated_df[fuc_cols].sum(axis=1) > 0).astype(int).values
        lectin_classes['Core fucose'] = (labels, '#CC79A7')

    bis_cols = [c for c in motif_cols if 'bisecting' in c.lower()]
    if bis_cols:
        labels = (annotated_df[bis_cols].sum(axis=1) > 0).astype(int).values
        lectin_classes['Bisecting GlcNAc'] = (labels, '#D55E00')

    sia_labels = []
    for iupac in iupac_list:
        s = str(iupac) if iupac else ''
        has_a23 = 'Neu5Ac(a2-3)' in s or 'Sia(a2-3)' in s or 'NeuAc(a2-3)' in s
        has_a26 = 'Neu5Ac(a2-6)' in s or 'Sia(a2-6)' in s or 'NeuAc(a2-6)' in s
        if has_a23 and not has_a26:
            sia_labels.append('a2-3')
        elif has_a26 and not has_a23:
            sia_labels.append('a2-6')
        elif has_a23 and has_a26:
            sia_labels.append('both')
        else:
            sia_labels.append('none')
    lectin_classes['_sia_linkage'] = sia_labels
    return lectin_classes


def assign_cancer_labels(annotated_df, motif_cols, iupac_list):
    cancer_motifs = {}

    fuc_cols = [c for c in motif_cols if 'core_fucose' in c.lower()]
    if fuc_cols:
        labels = (annotated_df[fuc_cols].sum(axis=1) > 0).astype(int).values
        cancer_motifs['Core fucose (AFP-L3)'] = (labels, '#D55E00')

    oglycan_cols = [c for c in motif_cols if 'oglycan' in c.lower() or 'mucin' in c.lower()]
    if oglycan_cols:
        labels = (annotated_df[oglycan_cols].sum(axis=1) > 0).astype(int).values
        cancer_motifs['O-glycan cores (Tn context)'] = (labels, '#CC79A7')

    hm_cols = [c for c in motif_cols if 'high_mannose' in c.lower()]
    if hm_cols:
        labels = (annotated_df[hm_cols].sum(axis=1) > 0).astype(int).values
        cancer_motifs['High mannose (immature)'] = (labels, '#009E73')

    plac_cols = [c for c in motif_cols if 'polylacnac' in c.lower() or 'PolyLacNAc' in c]
    if plac_cols:
        labels = (annotated_df[plac_cols].sum(axis=1) > 0).astype(int).values
        cancer_motifs['Poly-LacNAc (GnT-V)'] = (labels, '#0072B2')

    bis_cols = [c for c in motif_cols if 'bisecting' in c.lower()]
    if bis_cols:
        labels = (annotated_df[bis_cols].sum(axis=1) > 0).astype(int).values
        cancer_motifs['Bisecting GlcNAc (anti-met)'] = (labels, '#56B4E9')

    sia_count = []
    for iupac in iupac_list:
        s = str(iupac) if iupac else ''
        count = s.count('Neu5Ac') + s.count('NeuAc') + s.count('Sia(')
        sia_count.append(count)
    cancer_motifs['_sia_count'] = sia_count

    cancer_score = np.zeros(len(iupac_list))
    for name, data in cancer_motifs.items():
        if name.startswith('_') or 'anti-met' in name:
            continue
        if not isinstance(data, tuple):
            continue
        labels, _ = data
        cancer_score += labels
    cancer_motifs['_cancer_score'] = cancer_score.astype(int)
    return cancer_motifs


def assign_glycan_type_labels(df):
    col = 'glycan_type'
    if col not in df.columns:
        return None
    labels = df[col].fillna('Unknown').values
    type_counts = Counter(labels)
    return [l if type_counts[l] >= 50 else 'Other' for l in labels]


def assign_tissue_organ_labels(df):
    col = 'tissue_sample'
    if col not in df.columns or df[col].isna().all():
        return None, None
    tissue_to_organ = {
        'brain': 'Nervous', 'cerebral_cortex': 'Nervous', 'cerebellum': 'Nervous',
        'hippocampus': 'Nervous', 'hypothalamus': 'Nervous', 'medulla': 'Nervous',
        'white_matter': 'Nervous', 'cortex': 'Nervous',
        'serum': 'Immune/Blood', 'blood': 'Immune/Blood', 'plasma': 'Immune/Blood',
        'lymph_node': 'Immune/Blood', 'spleen': 'Immune/Blood', 'thymus': 'Immune/Blood',
        'stomach_mucosa': 'Digestive', 'stomach': 'Digestive', 'intestine': 'Digestive',
        'colon': 'Digestive', 'liver': 'Digestive', 'pancreas': 'Digestive',
        'kidney': 'Urogenital', 'urine': 'Urogenital', 'prostate': 'Urogenital',
        'seminal_fluid': 'Urogenital', 'bladder': 'Urogenital', 'testis': 'Urogenital',
        'uterus': 'Urogenital', 'ovary': 'Urogenital',
        'milk': 'Secretory', 'saliva': 'Secretory', 'tears': 'Secretory',
        'skin': 'Secretory', 'lung': 'Secretory',
    }
    organ_labels, tissue_labels = [], []
    for _, row in df.iterrows():
        tissue = str(row.get(col, '')).strip().lower() if pd.notna(row.get(col)) else ''
        if tissue and tissue != 'nan':
            organ_labels.append(tissue_to_organ.get(tissue, 'Other'))
            tissue_labels.append(tissue)
        else:
            organ_labels.append(None)
            tissue_labels.append(None)
    return organ_labels, tissue_labels


# ═══ Main analysis runner ═════════════════════════════════════════════════════

def run_exploratory(model_version, device, max_glycans=15000):
    model_name = {'v5': 'V5-A', 'v6': 'V6'}[model_version]
    output_dir = (PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' /
                  'probing_analysis' / f'exploratory_bio_{model_version}')
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n{'='*70}")
    print(f"Exploratory Biology Analysis - GlycanBERT {model_name}")
    print(f"{'='*70}")

    # 1. Load data
    print(f"\n1. Loading data from {DATA_PATH}")
    df = pd.read_csv(DATA_PATH)
    mask = df['glycan'].notna() & df['wurcs'].notna()
    df = df[mask].head(max_glycans).reset_index(drop=True)
    print(f"   {len(df)} glycans with IUPAC + WURCS")

    # 2. Annotate motifs
    print(f"\n2. Annotating glycowork motifs...")
    from glycowork.motif.annotate import annotate_dataset
    iupac_list = df['glycan'].tolist()
    try:
        annotated = annotate_dataset(iupac_list, feature_set=['known'], condense=True)
    except TypeError:
        annotated = annotate_dataset(iupac_list)
    motif_cols = [c for c in annotated.columns if c != 'glycan']
    print(f"   Found {len(motif_cols)} motif columns with hits")
    for c in sorted(motif_cols):
        print(f"   - {c}: {int(annotated[c].sum())} ({annotated[c].sum()/len(df)*100:.1f}%)")

    # 3. Assign biological labels
    print(f"\n3. Assigning biological labels...")
    lectin_labels = assign_lectin_labels(annotated, motif_cols, iupac_list)
    cancer_labels = assign_cancer_labels(annotated, motif_cols, iupac_list)
    glycan_types  = assign_glycan_type_labels(df)
    organ_labels, tissue_labels = assign_tissue_organ_labels(df)

    for name, data in lectin_labels.items():
        if name.startswith('_'): continue
        labels, _ = data
        print(f"   Lectin: {name}: {labels.sum()} pos ({labels.sum()/len(labels)*100:.1f}%)")
    for name, data in cancer_labels.items():
        if name.startswith('_'): continue
        labels, _ = data
        print(f"   Cancer: {name}: {labels.sum()} pos ({labels.sum()/len(labels)*100:.1f}%)")
    if glycan_types:
        print(f"   Glycan types: {dict(Counter(glycan_types))}")
    if organ_labels:
        print(f"   Organs: {dict(Counter(o for o in organ_labels if o))}")

    # 4. Load model & extract embeddings
    print(f"\n4. Loading model and extracting CLS embeddings...")
    model, tokenizer = load_model(model_version, device)
    embeddings = get_cls_embeddings(model, tokenizer, df['wurcs'].tolist(), device)
    print(f"   Embeddings: {embeddings.shape}")
    import gc; del model; torch.cuda.empty_cache(); gc.collect()

    # 5. Dimensionality reduction
    print(f"\n5. Dimensionality reduction...")
    tsne_coords, umap_coords, sub_idx = compute_dim_reduction(embeddings)

    def subset(labels):
        if isinstance(labels, np.ndarray): return labels[sub_idx]
        elif isinstance(labels, list): return [labels[i] for i in sub_idx]
        return labels

    all_metrics = {'model': model_name, 'themes': {}}

    # 6. Theme A: Lectin motifs
    print(f"\n{'~'*50}\nTheme A: Lectin-Binding Motifs\n{'~'*50}")
    lectin_binary = {k: (subset(v[0]), v[1]) for k, v in lectin_labels.items() if not k.startswith('_') and isinstance(v, tuple)}
    plot_multi_panel_binary(tsne_coords, lectin_binary,
                           f'Lectin Motif Clustering - {model_name} (t-SNE)',
                           output_dir / f'lectin_tsne_{model_name.lower()}')
    if umap_coords is not None:
        plot_multi_panel_binary(umap_coords, lectin_binary,
                               f'Lectin Motif Clustering - {model_name} (UMAP)',
                               output_dir / f'lectin_umap_{model_name.lower()}')

    lectin_metrics = {}
    for name, data in lectin_labels.items():
        if name.startswith('_'): continue
        if not isinstance(data, tuple): continue
        labels_full, _ = data
        print(f"\n  {name}:")
        knn = knn_label_purity(embeddings, labels_full)
        sil = compute_silhouette(embeddings, labels_full)
        lectin_metrics[name] = {'knn_purity': knn, 'silhouette': sil,
                                'n_positive': int(labels_full.sum()), 'n_total': len(labels_full)}

    sia_labels = lectin_labels.get('_sia_linkage', [])
    sia_nonzero = [(l, i) for i, l in enumerate(sia_labels) if l != 'none']
    if len(sia_nonzero) > 50:
        sia_vals, sia_idx = zip(*sia_nonzero)
        sia_idx = np.array(sia_idx)
        sia_embs = embeddings[sia_idx]
        print(f"\n  Sialic acid linkage ({len(sia_nonzero)} glycans):")
        knn = knn_label_purity(sia_embs, list(sia_vals))
        sil = compute_silhouette(sia_embs, list(sia_vals))
        lectin_metrics['Sialic acid linkage'] = {'knn_purity': knn, 'silhouette': sil,
                                                  'distribution': dict(Counter(sia_vals))}
        sia_in_sub = [i for i, idx in enumerate(sub_idx) if idx in set(sia_idx)]
        if len(sia_in_sub) > 10:
            sia_sub_labels = [sia_labels[sub_idx[i]] for i in sia_in_sub]
            color_map = {'a2-3': '#D55E00', 'a2-6': '#0072B2', 'both': '#009E73'}
            plot_multi_class_overlay(tsne_coords[sia_in_sub], sia_sub_labels,
                f'Sialic Acid Linkage - {model_name} (t-SNE)',
                color_map, output_dir / f'sia_linkage_tsne_{model_name.lower()}')
    all_metrics['themes']['lectin'] = lectin_metrics

    # 7. Theme B: Cancer signatures
    print(f"\n{'~'*50}\nTheme B: Cancer Glycan Signatures\n{'~'*50}")
    cancer_binary = {k: (subset(v[0]), v[1]) for k, v in cancer_labels.items() if not k.startswith('_') and isinstance(v, tuple)}
    plot_multi_panel_binary(tsne_coords, cancer_binary,
                           f'Cancer Glycan Signatures - {model_name} (t-SNE)',
                           output_dir / f'cancer_tsne_{model_name.lower()}')
    if umap_coords is not None:
        plot_multi_panel_binary(umap_coords, cancer_binary,
                               f'Cancer Glycan Signatures - {model_name} (UMAP)',
                               output_dir / f'cancer_umap_{model_name.lower()}')

    cancer_metrics = {}
    for name, data in cancer_labels.items():
        if name.startswith('_'): continue
        if not isinstance(data, tuple): continue
        labels_full, _ = data
        print(f"\n  {name}:")
        knn = knn_label_purity(embeddings, labels_full)
        sil = compute_silhouette(embeddings, labels_full)
        cancer_metrics[name] = {'knn_purity': knn, 'silhouette': sil,
                                'n_positive': int(labels_full.sum()), 'n_total': len(labels_full)}

    cancer_score = cancer_labels.get('_cancer_score')
    if cancer_score is not None:
        plot_continuous_overlay(tsne_coords, subset(cancer_score),
            f'Cancer Motif Score - {model_name} (t-SNE)',
            output_dir / f'cancer_score_tsne_{model_name.lower()}', cmap='YlOrRd', label='# cancer motifs')
    sia_count = cancer_labels.get('_sia_count')
    if sia_count is not None:
        plot_continuous_overlay(tsne_coords, subset(sia_count),
            f'Sialylation Count - {model_name} (t-SNE)',
            output_dir / f'sialylation_tsne_{model_name.lower()}', cmap='PuBu', label='# Sia residues')
    all_metrics['themes']['cancer'] = cancer_metrics

    # 8. Theme C: Glycan type
    print(f"\n{'~'*50}\nTheme C: Glycan Type\n{'~'*50}")
    if glycan_types:
        sub_types = subset(glycan_types)
        type_colors = {'N': '#0072B2', 'O': '#D55E00', 'lipid': '#009E73',
                       'free': '#CC79A7', 'GAG': '#E69F00', 'Other': '#999999'}
        plot_multi_class_overlay(tsne_coords, sub_types,
            f'Glycan Type - {model_name} (t-SNE)',
            type_colors, output_dir / f'glycan_type_tsne_{model_name.lower()}')
        if umap_coords is not None:
            plot_multi_class_overlay(umap_coords, sub_types,
                f'Glycan Type - {model_name} (UMAP)',
                type_colors, output_dir / f'glycan_type_umap_{model_name.lower()}')
        print(f"\n  Glycan type clustering:")
        knn = knn_label_purity(embeddings, glycan_types)
        sil = compute_silhouette(embeddings, glycan_types)
        all_metrics['themes']['glycan_type'] = {'knn_purity': knn, 'silhouette': sil,
                                                 'distribution': dict(Counter(glycan_types))}

    # 9. Theme D: Tissue/organ system
    print(f"\n{'~'*50}\nTheme D: Tissue & Organ System\n{'~'*50}")
    if organ_labels:
        has_organ = [i for i, o in enumerate(organ_labels) if o is not None]
        if len(has_organ) > 100:
            organ_embs = embeddings[has_organ]
            organ_labs = [organ_labels[i] for i in has_organ]
            print(f"\n  {len(has_organ)} glycans with organ annotations:")
            knn = knn_label_purity(organ_embs, organ_labs)
            sil = compute_silhouette(organ_embs, organ_labs)
            organ_colors = {'Nervous': '#0072B2', 'Immune/Blood': '#D55E00',
                           'Digestive': '#009E73', 'Urogenital': '#CC79A7',
                           'Secretory': '#E69F00', 'Other': '#999999'}
            sub_organ = subset(organ_labels)
            has_organ_sub = [i for i, o in enumerate(sub_organ) if o is not None]
            if len(has_organ_sub) > 50:
                plot_multi_class_overlay(tsne_coords[has_organ_sub],
                    [sub_organ[i] for i in has_organ_sub],
                    f'Organ System - {model_name} (t-SNE)',
                    organ_colors, output_dir / f'organ_system_tsne_{model_name.lower()}')
            all_metrics['themes']['organ_system'] = {'knn_purity': knn, 'silhouette': sil,
                                                      'distribution': dict(Counter(organ_labs))}
        else:
            print(f"  Only {len(has_organ)} glycans with organ annotations - skipping")

    # 10. Save results
    def make_serializable(obj):
        if isinstance(obj, (np.integer,)): return int(obj)
        if isinstance(obj, (np.floating,)): return float(obj)
        if isinstance(obj, np.ndarray): return obj.tolist()
        if isinstance(obj, dict): return {str(k): make_serializable(v) for k, v in obj.items()}
        if isinstance(obj, (list, tuple)): return [make_serializable(x) for x in obj]
        return obj

    results_path = output_dir / f'exploratory_results_{model_name.lower()}.json'
    with open(results_path, 'w') as f:
        json.dump(make_serializable(all_metrics), f, indent=2)
    print(f"\n  Saved results: {results_path}")

    # Summary
    print(f"\n{'='*70}\nEXPLORATORY SUMMARY ({model_name})\n{'='*70}")
    for theme, metrics in all_metrics['themes'].items():
        print(f"\n  {theme.upper()}:")
        if isinstance(metrics, dict) and 'knn_purity' in metrics:
            knn = metrics['knn_purity']
            for k, (p, b) in knn.items():
                print(f"    k={k}: purity={p:.4f} (baseline={b:.4f}, lift={p/b:.2f}x)")
        else:
            for sub_name, sub_m in metrics.items():
                if isinstance(sub_m, dict) and 'knn_purity' in sub_m:
                    knn = sub_m['knn_purity']
                    if 10 in knn:
                        p, b = knn[10]
                        sil = sub_m.get('silhouette')
                        sil_str = f", sil={sil:.4f}" if sil else ""
                        print(f"    {sub_name}: k10={p:.4f} (x{p/b:.2f}){sil_str}")
    print(f"\n  Output: {output_dir}")
    print(f"{'='*70}")
    return all_metrics


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--max-glycans', type=int, default=15000)
    parser.add_argument('--model', choices=['v5', 'v6', 'both'], default='both')
    args = parser.parse_args()
    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    print("Python:", sys.executable)
    print(f"PyTorch: {torch.__version__} CUDA: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"{torch.cuda.get_device_name()}, "
              f"{torch.cuda.get_device_properties(0).total_memory // 1024**2} MiB")
    models = ['v5', 'v6'] if args.model == 'both' else [args.model]
    for m in models:
        run_exploratory(m, device, args.max_glycans)
    print(f"\nAll done: {pd.Timestamp.now()}")

if __name__ == '__main__':
    main()