#!/usr/bin/env python3
"""
Probe 10: Tissue-Specific Glycan Presence — GlycanBERT
=======================================================
Tests whether GlycanBERT embeddings encode tissue-specific glycan identity
using mass-spectrometry-validated ground truth from:

  "A comprehensive N-glycoproteome atlas reveals tissue-specific glycan
   remodeling but non-random structural microheterogeneities"
   (Nature Communications, 2025)

Data: 189 N-glycan structures × 23 mouse tissues (binary presence/absence)
      sourced from Figure 4d + Supplement Figure 1a of the paper.
      Glycans mapped: StrucGP → GlyTouCan → WURCS (100% resolution).

Sub-probes:
  10a. Multi-label tissue prediction (23-class sigmoid probe)
  10b. Organ-system classification (7 systems: Cardio, Digestive, etc.)
  10c. Brain vs Non-brain binary classification
  10d. Tissue-specificity vs Ubiquity prediction (regression)
  10e. t-SNE visualization colored by organ system and tissue count

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

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

PROJECT_ROOT = Path(__file__).resolve().parents[2]
VOCAB_PATH = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'vocabulary.json'
DATA_PATHS = {
    'v1': PROJECT_ROOT / 'bert_v6_contrastive' / 'additional_probes' / 'probe_10_tissue_presence_data.csv',
    'v2': PROJECT_ROOT / 'bert_v6_contrastive' / 'additional_probes' / 'probe_10_tissue_presence_data_v2.csv',
}

# Tissues with >=30 positives in v2 dataset (well-powered for classification)
WELL_POWERED_TISSUES = [
    'Colon', 'Intestine', 'Liver', 'Pancreas', 'Stomach',
    'Kidney', 'Spleen', 'Cerebellum', 'Hippocampus', 'Hypothalamus',
    'Medulla', 'Olfactory bulb', 'White matter', 'Testis', 'Uterus', 'Lung'
]

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

TISSUE_NAMES = [
    'Heart', 'Colon', 'Gallbladder', 'Intestine', 'Liver', 'Pancreas', 'Stomach',
    'Bladder', 'Kidney', 'Spleen', 'Muscle',
    'Cerebellum', 'Cortex', 'Hippocampus', 'Hypothalamus', 'Medulla', 'Olfactory bulb', 'White matter',
    'Ovary', 'Seminal vesicle', 'Testis', 'Uterus', 'Lung'
]

# Organ system groupings
ORGAN_SYSTEMS = {
    'Cardiovascular': ['Heart'],
    'Digestive': ['Colon', 'Gallbladder', 'Intestine', 'Liver', 'Pancreas', 'Stomach'],
    'Excretory': ['Bladder', 'Kidney'],
    'Immune': ['Spleen'],
    'Musculoskeletal': ['Muscle'],
    'Nervous': ['Cerebellum', 'Cortex', 'Hippocampus', 'Hypothalamus', 'Medulla',
                'Olfactory bulb', 'White matter'],
    'Reproductive': ['Ovary', 'Seminal vesicle', 'Testis', 'Uterus'],
    'Respiratory': ['Lung'],
}

BRAIN_TISSUES = ['Cerebellum', 'Cortex', 'Hippocampus', 'Hypothalamus',
                 'Medulla', 'Olfactory bulb', 'White matter']

# Nature-style colors
SYSTEM_COLORS = {
    'Cardiovascular': '#D55E00',
    'Digestive': '#E69F00',
    'Excretory': '#56B4E9',
    'Immune': '#009E73',
    'Musculoskeletal': '#CC79A7',
    'Nervous': '#0072B2',
    'Reproductive': '#F0E442',
    'Respiratory': '#882255',
}

# ─── Model loading (matches probe_8) ────────────────────────────────
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."""
    import torch
    print(f"Loading model 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()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"  Loaded: {n_params:,} params, vocab={vocab_size}, hidden={hidden}")
    return model


def batch_cls_embeddings(model, wurcs_list, tokenizer, device='cuda',
                         batch_size=64, max_len=256):
    """Extract [CLS] embeddings for a list of WURCS strings."""
    import torch
    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)
        mask = torch.zeros(len(token_ids_list), max_l, dtype=torch.long)

        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)
            mask[j, :len(ids)] = 1

        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)

    print(f"  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))


def setup_nature_style():
    """Configure matplotlib for Nature-style publication plots."""
    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


# ═══════════════════════════════════════════════════════════════════════
# PROBE 10a: Multi-label Tissue Prediction
# ═══════════════════════════════════════════════════════════════════════

def probe_10a_multilabel(embs, tissue_matrix, output_dir, model_name):
    """Per-tissue binary classifiers (one-vs-rest) using linear probes."""
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import StratifiedKFold, LeaveOneOut
    from sklearn.metrics import accuracy_score, f1_score, roc_auc_score

    print(f"\n{'='*60}")
    print(f"PROBE 10a: Multi-label Tissue Prediction ({model_name})")
    print(f"{'='*60}")

    per_tissue_results = {}
    valid_tissues = []

    for t_idx, tissue in enumerate(TISSUE_NAMES):
        y = tissue_matrix[:, t_idx]
        n_pos = y.sum()
        n_neg = len(y) - n_pos

        if n_pos < 3 or n_neg < 3:
            print(f"  {tissue:<20s}: SKIP (pos={n_pos}, neg={n_neg})")
            continue

        valid_tissues.append(tissue)

        # Use stratified 5-fold or LOO for very small classes
        n_splits = min(5, n_pos, n_neg)
        if n_splits < 2:
            print(f"  {tissue:<20s}: SKIP (can't split, pos={n_pos})")
            continue

        skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
        accs, f1s, aucs = [], [], []

        for train_idx, test_idx in skf.split(embs, y):
            clf = LogisticRegression(max_iter=1000, C=1.0, solver='lbfgs')
            clf.fit(embs[train_idx], y[train_idx])
            preds = clf.predict(embs[test_idx])
            proba = clf.predict_proba(embs[test_idx])

            acc = accuracy_score(y[test_idx], preds)
            f1 = f1_score(y[test_idx], preds, average='binary', zero_division=0)
            try:
                auc = roc_auc_score(y[test_idx], proba[:, 1])
            except (ValueError, IndexError):
                auc = 0.5
            accs.append(acc)
            f1s.append(f1)
            aucs.append(auc)

        result = {
            'tissue': tissue,
            'n_positive': int(n_pos),
            'n_negative': int(n_neg),
            'accuracy_mean': float(np.mean(accs)),
            'accuracy_std': float(np.std(accs)),
            'f1_mean': float(np.mean(f1s)),
            'f1_std': float(np.std(f1s)),
            'auc_mean': float(np.mean(aucs)),
            'auc_std': float(np.std(aucs)),
        }
        per_tissue_results[tissue] = result
        print(f"  {tissue:<20s}: AUC={np.mean(aucs):.3f}±{np.std(aucs):.3f}, "
              f"F1={np.mean(f1s):.3f}, pos={n_pos}")

    # Macro averages
    if per_tissue_results:
        macro_auc = np.mean([r['auc_mean'] for r in per_tissue_results.values()])
        macro_f1 = np.mean([r['f1_mean'] for r in per_tissue_results.values()])
        macro_acc = np.mean([r['accuracy_mean'] for r in per_tissue_results.values()])
        print(f"\n  MACRO AVERAGE across {len(per_tissue_results)} tissues:")
        print(f"    AUC  = {macro_auc:.4f}")
        print(f"    F1   = {macro_f1:.4f}")
        print(f"    Acc  = {macro_acc:.4f}")

    results = {
        'task': 'multilabel_tissue_prediction',
        'model': model_name,
        'n_glycans': int(len(embs)),
        'n_tissues_probed': len(per_tissue_results),
        'macro_auc': float(macro_auc) if per_tissue_results else 0.0,
        'macro_f1': float(macro_f1) if per_tissue_results else 0.0,
        'macro_accuracy': float(macro_acc) if per_tissue_results else 0.0,
        'per_tissue': per_tissue_results,
    }
    with open(os.path.join(output_dir, f'probe10a_multilabel_{model_name}.json'), 'w') as f:
        json.dump(results, f, indent=2)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 10b: Organ System Classification
# ═══════════════════════════════════════════════════════════════════════

def probe_10b_organ_system(embs, tissue_matrix, output_dir, model_name):
    """Classify glycans by dominant organ system."""
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import StratifiedKFold
    from sklearn.metrics import accuracy_score, f1_score

    print(f"\n{'='*60}")
    print(f"PROBE 10b: Organ System Classification ({model_name})")
    print(f"{'='*60}")

    # Assign each glycan to the organ system where it appears most
    # (for glycans in multiple systems, assign to MOST SPECIFIC one)
    system_labels = []
    for i in range(len(embs)):
        system_counts = {}
        for sys_name, tissues in ORGAN_SYSTEMS.items():
            count = sum(tissue_matrix[i, TISSUE_NAMES.index(t)]
                       for t in tissues if t in TISSUE_NAMES)
            if count > 0:
                system_counts[sys_name] = count

        if not system_counts:
            system_labels.append('Unknown')
        elif len(system_counts) == 1:
            system_labels.append(list(system_counts.keys())[0])
        else:
            # Assign to the system with highest relative presence
            # (count / number of tissues in that system)
            best_sys = max(system_counts.keys(),
                          key=lambda s: system_counts[s] / len(ORGAN_SYSTEMS[s]))
            system_labels.append(best_sys)

    system_labels = np.array(system_labels)
    counts = Counter(system_labels)
    print(f"  Distribution: {dict(sorted(counts.items()))}")

    # Filter to systems with ≥5 samples
    valid_systems = [s for s, c in counts.items() if c >= 5 and s != 'Unknown']
    mask = np.array([l in valid_systems for l in system_labels])

    if mask.sum() < 20:
        print("  SKIP: Too few valid samples")
        return {'task': 'organ_system_classification', 'model': model_name, 'skipped': True}

    embs_f = embs[mask]
    labels_f = system_labels[mask]

    from sklearn.preprocessing import LabelEncoder
    le = LabelEncoder()
    y = le.fit_transform(labels_f)
    classes = le.classes_
    print(f"  Classes ({len(classes)}): {list(classes)}")

    n_splits = min(5, min(Counter(y).values()))
    if n_splits < 2:
        print("  SKIP: Not enough per-class for CV")
        return {'task': 'organ_system_classification', 'model': model_name, 'skipped': True}

    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
    accs, f1s = [], []
    for fold, (tr, te) in enumerate(skf.split(embs_f, y)):
        clf = LogisticRegression(max_iter=1000, C=1.0, solver='lbfgs',
                                 multi_class='multinomial')
        clf.fit(embs_f[tr], y[tr])
        preds = clf.predict(embs_f[te])
        acc = accuracy_score(y[te], preds)
        f1 = f1_score(y[te], preds, average='macro')
        accs.append(acc)
        f1s.append(f1)
        print(f"  Fold {fold+1}: acc={acc:.4f}, F1={f1:.4f}")

    print(f"\n  RESULT: Acc={np.mean(accs):.4f}±{np.std(accs):.4f}, "
          f"F1={np.mean(f1s):.4f}±{np.std(f1s):.4f}")

    results = {
        'task': 'organ_system_classification', 'model': model_name,
        'n_samples': int(mask.sum()),
        'n_classes': int(len(classes)),
        'classes': list(classes),
        'distribution': {str(k): int(v) for k, v in counts.items()},
        'accuracy_mean': float(np.mean(accs)),
        'accuracy_std': float(np.std(accs)),
        'f1_macro_mean': float(np.mean(f1s)),
        'f1_macro_std': float(np.std(f1s)),
    }
    with open(os.path.join(output_dir, f'probe10b_organ_system_{model_name}.json'), 'w') as f:
        json.dump(results, f, indent=2)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 10c: Brain vs Non-brain
# ═══════════════════════════════════════════════════════════════════════

def probe_10c_brain(embs, tissue_matrix, output_dir, model_name):
    """Binary: does this glycan appear in any brain tissue?"""
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import StratifiedKFold
    from sklearn.metrics import accuracy_score, f1_score, roc_auc_score

    print(f"\n{'='*60}")
    print(f"PROBE 10c: Brain vs Non-brain ({model_name})")
    print(f"{'='*60}")

    brain_indices = [TISSUE_NAMES.index(t) for t in BRAIN_TISSUES if t in TISSUE_NAMES]
    y = (tissue_matrix[:, brain_indices].sum(axis=1) > 0).astype(int)

    n_brain = y.sum()
    n_nonbrain = len(y) - n_brain
    print(f"  Brain-present: {n_brain}, Non-brain-only: {n_nonbrain}")

    if n_brain < 5 or n_nonbrain < 5:
        print("  SKIP: Too few samples")
        return {'task': 'brain_classification', 'model': model_name, 'skipped': True}

    n_splits = min(5, n_brain, n_nonbrain)
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
    accs, f1s, aucs = [], [], []

    for fold, (tr, te) in enumerate(skf.split(embs, y)):
        clf = LogisticRegression(max_iter=1000, C=1.0)
        clf.fit(embs[tr], y[tr])
        preds = clf.predict(embs[te])
        proba = clf.predict_proba(embs[te])[:, 1]
        acc = accuracy_score(y[te], preds)
        f1 = f1_score(y[te], preds, average='binary')
        try:
            auc = roc_auc_score(y[te], proba)
        except ValueError:
            auc = 0.5
        accs.append(acc); f1s.append(f1); aucs.append(auc)
        print(f"  Fold {fold+1}: acc={acc:.4f}, F1={f1:.4f}, AUC={auc:.4f}")

    print(f"\n  RESULT: Acc={np.mean(accs):.4f}, AUC={np.mean(aucs):.4f}")

    results = {
        'task': 'brain_classification', 'model': model_name,
        'n_brain': int(n_brain), 'n_nonbrain': int(n_nonbrain),
        'accuracy_mean': float(np.mean(accs)),
        'accuracy_std': float(np.std(accs)),
        'f1_mean': float(np.mean(f1s)),
        'auc_mean': float(np.mean(aucs)),
        'auc_std': float(np.std(aucs)),
    }
    with open(os.path.join(output_dir, f'probe10c_brain_{model_name}.json'), 'w') as f:
        json.dump(results, f, indent=2)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 10d: Tissue Breadth Prediction (regression)
# ═══════════════════════════════════════════════════════════════════════

def probe_10d_breadth(embs, tissue_matrix, output_dir, model_name):
    """Predict number of tissues a glycan appears in (1–23)."""
    from sklearn.linear_model import Ridge
    from sklearn.model_selection import KFold
    from scipy.stats import spearmanr

    print(f"\n{'='*60}")
    print(f"PROBE 10d: Tissue Breadth Prediction ({model_name})")
    print(f"{'='*60}")

    y = tissue_matrix.sum(axis=1)
    print(f"  N={len(y)}, range=[{y.min()}, {y.max()}], mean={y.mean():.2f}")
    print(f"  Unique=1: {(y==1).sum()}, ≤3: {(y<=3).sum()}, ≥10: {(y>=10).sum()}")

    kf = KFold(n_splits=5, shuffle=True, random_state=42)
    r2s, spearmans, maes = [], [], []

    for fold, (tr, te) in enumerate(kf.split(embs)):
        clf = Ridge(alpha=1.0)
        clf.fit(embs[tr], y[tr])
        preds = clf.predict(embs[te])

        r2 = clf.score(embs[te], y[te])
        rho, p = spearmanr(y[te], preds)
        mae = np.mean(np.abs(y[te] - preds))
        r2s.append(r2); spearmans.append(rho); maes.append(mae)
        print(f"  Fold {fold+1}: R²={r2:.4f}, ρ={rho:.4f}, MAE={mae:.2f}")

    print(f"\n  RESULT: R²={np.mean(r2s):.4f}, ρ={np.mean(spearmans):.4f}, "
          f"MAE={np.mean(maes):.2f}")

    results = {
        'task': 'tissue_breadth_regression', 'model': model_name,
        'n_samples': int(len(y)),
        'tissue_count_range': [int(y.min()), int(y.max())],
        'r2_mean': float(np.mean(r2s)),
        'r2_std': float(np.std(r2s)),
        'spearman_mean': float(np.mean(spearmans)),
        'mae_mean': float(np.mean(maes)),
    }
    with open(os.path.join(output_dir, f'probe10d_breadth_{model_name}.json'), 'w') as f:
        json.dump(results, f, indent=2)
    return results


# ═══════════════════════════════════════════════════════════════════════
# PROBE 10e: t-SNE Visualization
# ═══════════════════════════════════════════════════════════════════════

def probe_10e_visualization(embs, tissue_matrix, compositions, output_dir, model_name):
    """t-SNE colored by organ system and tissue breadth."""
    plt = setup_nature_style()
    from sklearn.manifold import TSNE

    print(f"\n{'='*60}")
    print(f"PROBE 10e: t-SNE Visualization ({model_name})")
    print(f"{'='*60}")

    n_tissues = tissue_matrix.sum(axis=1)

    # Assign primary organ system for coloring
    system_labels = []
    for i in range(len(embs)):
        system_counts = {}
        for sys_name, tissues in ORGAN_SYSTEMS.items():
            count = sum(tissue_matrix[i, TISSUE_NAMES.index(t)]
                       for t in tissues if t in TISSUE_NAMES)
            if count > 0:
                system_counts[sys_name] = count
        if system_counts:
            best = max(system_counts.keys(),
                      key=lambda s: system_counts[s] / len(ORGAN_SYSTEMS[s]))
            system_labels.append(best)
        else:
            system_labels.append('Unknown')
    system_labels = np.array(system_labels)

    # Run t-SNE
    print(f"  Running t-SNE on {len(embs)} samples...")
    tsne = TSNE(n_components=2, perplexity=min(30, len(embs)-1),
                random_state=42, max_iter=1000, learning_rate='auto', init='pca')
    coords = tsne.fit_transform(embs)

    fig, axes = plt.subplots(1, 3, figsize=(21, 6))

    # ─── Panel 1: Organ System ───
    ax = axes[0]
    for sys_name in sorted(ORGAN_SYSTEMS.keys()):
        mask = system_labels == sys_name
        if mask.sum() > 0:
            ax.scatter(coords[mask, 0], coords[mask, 1],
                      c=SYSTEM_COLORS.get(sys_name, 'gray'), s=40, alpha=0.7,
                      label=f'{sys_name} ({mask.sum()})', edgecolors='white',
                      linewidths=0.3, rasterized=True)
    ax.set_xlabel('t-SNE 1')
    ax.set_ylabel('t-SNE 2')
    ax.set_title(f'Dominant Organ System — {model_name}')
    ax.legend(frameon=False, markerscale=1.2, fontsize=7,
              loc='center left', bbox_to_anchor=(1.0, 0.5))

    # ─── Panel 2: Tissue Breadth (colormap) ───
    ax = axes[1]
    scatter = ax.scatter(coords[:, 0], coords[:, 1],
                        c=n_tissues, cmap='YlOrRd', s=40, alpha=0.7,
                        edgecolors='white', linewidths=0.3, rasterized=True,
                        vmin=1, vmax=max(n_tissues))
    cbar = plt.colorbar(scatter, ax=ax, shrink=0.8)
    cbar.set_label('Number of tissues')
    ax.set_xlabel('t-SNE 1')
    ax.set_ylabel('t-SNE 2')
    ax.set_title(f'Tissue Breadth — {model_name}')

    # ─── Panel 3: Brain vs Non-brain ───
    ax = axes[2]
    brain_indices = [TISSUE_NAMES.index(t) for t in BRAIN_TISSUES if t in TISSUE_NAMES]
    is_brain = (tissue_matrix[:, brain_indices].sum(axis=1) > 0)
    brain_only = is_brain & (tissue_matrix[:, [i for i in range(len(TISSUE_NAMES))
                                               if TISSUE_NAMES[i] not in BRAIN_TISSUES]].sum(axis=1) == 0)

    non_brain = ~is_brain
    both = is_brain & ~brain_only

    for mask, color, label, s, alpha in [
        (non_brain, '#E69F00', 'Non-brain only', 35, 0.7),
        (both, '#009E73', 'Brain + peripheral', 45, 0.8),
        (brain_only, '#0072B2', 'Brain only', 50, 0.9),
    ]:
        if mask.sum() > 0:
            ax.scatter(coords[mask, 0], coords[mask, 1],
                      c=color, s=s, alpha=alpha,
                      label=f'{label} ({mask.sum()})', edgecolors='white',
                      linewidths=0.3, rasterized=True)
    ax.set_xlabel('t-SNE 1')
    ax.set_ylabel('t-SNE 2')
    ax.set_title(f'Brain vs Non-brain — {model_name}')
    ax.legend(frameon=False, markerscale=1.5, fontsize=8)

    plt.suptitle(f'Probe 10: MS-validated Tissue-Specific Glycan Presence — '
                 f'GlycanBERT {model_name}\n'
                 f'({len(embs)} N-glycans × {len(TISSUE_NAMES)} mouse tissues, '
                 f'Nature Commun. 2025)',
                 fontsize=12, fontweight='bold', y=1.02)
    plt.tight_layout()

    fig_path = os.path.join(output_dir, f'probe10_tsne_{model_name.lower()}.png')
    plt.savefig(fig_path, dpi=300, bbox_inches='tight', facecolor='white')
    plt.savefig(fig_path.replace('.png', '.pdf'), bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"  Saved: {fig_path}")

    np.savez(os.path.join(output_dir, f'probe10_tsne_coords_{model_name.lower()}.npz'),
             coords=coords, systems=system_labels, n_tissues=n_tissues)


# ═══════════════════════════════════════════════════════════════════════
# PROBE 10f: Per-tissue AUC bar chart
# ═══════════════════════════════════════════════════════════════════════

def probe_10f_auc_barplot(results_10a, output_dir, model_name):
    """Bar chart of per-tissue AUC from Probe 10a."""
    plt = setup_nature_style()

    print(f"\n{'='*60}")
    print(f"PROBE 10f: Per-tissue AUC Bar Chart ({model_name})")
    print(f"{'='*60}")

    per_tissue = results_10a.get('per_tissue', {})
    if not per_tissue:
        print("  SKIP: No per-tissue results")
        return

    # Sort by AUC descending
    tissues_sorted = sorted(per_tissue.keys(),
                           key=lambda t: per_tissue[t]['auc_mean'], reverse=True)

    fig, ax = plt.subplots(figsize=(12, 6))

    x = np.arange(len(tissues_sorted))
    aucs = [per_tissue[t]['auc_mean'] for t in tissues_sorted]
    auc_stds = [per_tissue[t]['auc_std'] for t in tissues_sorted]
    n_pos = [per_tissue[t]['n_positive'] for t in tissues_sorted]

    # Color by organ system
    colors = []
    for t in tissues_sorted:
        sys_found = 'Unknown'
        for sys_name, members in ORGAN_SYSTEMS.items():
            if t in members:
                sys_found = sys_name
                break
        colors.append(SYSTEM_COLORS.get(sys_found, '#999999'))

    bars = ax.bar(x, aucs, yerr=auc_stds, capsize=3, color=colors,
                  edgecolor='white', linewidth=0.5, alpha=0.85)

    # Add sample count annotations
    for i, (bar, n) in enumerate(zip(bars, n_pos)):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + auc_stds[i] + 0.02,
                f'n={n}', ha='center', va='bottom', fontsize=7, color='#555555')

    ax.axhline(y=0.5, color='#CCCCCC', linestyle='--', linewidth=1, label='Chance')
    ax.set_xticks(x)
    ax.set_xticklabels(tissues_sorted, rotation=45, ha='right', fontsize=8)
    ax.set_ylabel('AUC-ROC')
    ax.set_ylim(0, 1.15)
    ax.set_title(f'Probe 10a: Per-Tissue Presence Prediction AUC — {model_name}')

    # Legend for organ systems
    from matplotlib.patches import Patch
    legend_elements = [Patch(facecolor=c, label=s) for s, c in SYSTEM_COLORS.items()]
    ax.legend(handles=legend_elements, frameon=False, fontsize=7,
              loc='upper right', ncol=2)

    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'probe10a_auc_barplot_{model_name.lower()}.png')
    plt.savefig(fig_path, dpi=300, bbox_inches='tight', facecolor='white')
    plt.savefig(fig_path.replace('.png', '.pdf'), bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"  Saved: {fig_path}")


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

def main():
    parser = argparse.ArgumentParser(description='Probe 10: Tissue-Specific Glycan Presence')
    parser.add_argument('--model', choices=['v5', 'v6'], required=True)
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--output_dir', default=None)
    parser.add_argument('--data-version', choices=['v1', 'v2'], default='v2',
                        help='Dataset version: v1=189 glycans, v2=243 glycans')
    parser.add_argument('--well-powered', action='store_true',
                        help='Filter to well-powered tissues only (>=30 positives)')
    args = parser.parse_args()

    model_name = {'v5': 'V5-A', 'v6': 'V6'}[args.model]
    suffix = f'_{args.data_version}' if args.data_version != 'v1' else ''
    if args.well_powered:
        suffix += '_wp'
    if args.output_dir is None:
        args.output_dir = str(PROJECT_ROOT / 'bert_v6_contrastive' / 'additional_probes' /
                              f'probe10_results_{args.model}{suffix}')
    os.makedirs(args.output_dir, exist_ok=True)
    DATA_PATH = DATA_PATHS[args.data_version]

    # -- Determine active tissues --
    global TISSUE_NAMES
    if args.well_powered:
        active_tissues = [t for t in TISSUE_NAMES if t in WELL_POWERED_TISSUES]
        print(f"\nUsing WELL-POWERED tissues only: {len(active_tissues)} tissues")
    else:
        active_tissues = TISSUE_NAMES[:]
    TISSUE_NAMES = active_tissues

    # ── Load data ──
    print(f"\nLoading probe data from {DATA_PATH} (data={args.data_version})...")
    wurcs_list = []
    compositions = []
    tissue_rows = []

    with open(DATA_PATH) as f:
        for row in csv.DictReader(f):
            wurcs_list.append(row['wurcs'])
            compositions.append(row.get('composition', ''))
            tissue_rows.append([int(row.get(t, 0)) for t in TISSUE_NAMES])

    tissue_matrix = np.array(tissue_rows)
    total = len(wurcs_list)
    print(f"  Glycans: {total}")
    print(f"  Tissues: {len(TISSUE_NAMES)}")
    print(f"  Positive labels: {tissue_matrix.sum()} / {tissue_matrix.size} "
          f"({tissue_matrix.sum()/tissue_matrix.size*100:.1f}% fill rate)")
    print(f"  Tissues per glycan: mean={tissue_matrix.sum(axis=1).mean():.1f}, "
          f"median={np.median(tissue_matrix.sum(axis=1)):.0f}")

    # ── Load model + tokenizer ──
    print(f"\nLoading 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)

    # ── Embed ──
    print(f"\nEmbedding {total} glycans...")
    embs = batch_cls_embeddings(model, wurcs_list, tokenizer, device=args.device)

    # Free GPU
    import gc, torch
    del model
    torch.cuda.empty_cache()
    gc.collect()

    print(f"\nEmbeddings shape: {embs.shape}")

    # Check for tokenization failures (rows that didn't get embedded)
    if embs.shape[0] < total:
        print(f"  WARNING: {total - embs.shape[0]} glycans failed tokenization")
        # Truncate data to match
        tissue_matrix = tissue_matrix[:embs.shape[0]]
        compositions = compositions[:embs.shape[0]]

    # ── Run all sub-probes ──
    all_results = {}

    # 10a: Multi-label
    results_10a = probe_10a_multilabel(embs, tissue_matrix, args.output_dir, model_name)
    all_results['10a_multilabel'] = results_10a

    # 10b: Organ system
    all_results['10b_organ_system'] = probe_10b_organ_system(
        embs, tissue_matrix, args.output_dir, model_name)

    # 10c: Brain vs non-brain
    all_results['10c_brain'] = probe_10c_brain(
        embs, tissue_matrix, args.output_dir, model_name)

    # 10d: Tissue breadth regression
    all_results['10d_breadth'] = probe_10d_breadth(
        embs, tissue_matrix, args.output_dir, model_name)

    # 10e: t-SNE
    probe_10e_visualization(embs, tissue_matrix, compositions,
                           args.output_dir, model_name)

    # 10f: AUC bar plot
    probe_10f_auc_barplot(results_10a, args.output_dir, model_name)

    # Save combined results
    with open(os.path.join(args.output_dir, f'probe10_all_results_{model_name.lower()}.json'), 'w') as f:
        json.dump(all_results, f, indent=2, default=str)

    print(f"\n{'='*60}")
    print(f"PROBE 10 COMPLETE — {model_name}")
    print(f"Results: {args.output_dir}")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()