"""
Probe 3: Attention Head Analysis
Probe 4: Embedding Distance vs Structural Similarity

GlycanBERT V6 Embedding Probes
================================

Probe 3 hooks into the attention layers to extract attention weight matrices,
then computes per-head entropy, CLS-attention patterns, and specialization metrics.

Probe 4 computes pairwise cosine similarity between CLS embeddings and correlates
with glycowork's structural similarity (motif fingerprint cosine similarity).
"""

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

# Project paths
PROJECT_ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / 'bert_training_v4'))

import torch
import torch.nn.functional as F
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

# ============================================================
# Model loading (shared with probe_2)
# ============================================================

VOCAB_PATH = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'vocabulary.json'
CHECKPOINTS = {
    'v6': PROJECT_ROOT / 'bert_v5.1_contrastive' / 'checkpoints' / 'best_v51_contrastive_model.pt',
}
BENCH_DIR = PROJECT_ROOT / 'bench' / 'GlycanML' / 'data'


def load_model(ckpt_path, device='cuda'):
    """Load model — infer vocab sizes from checkpoint state_dict."""
    from model.multimodal_glycan_bert_v3 import MultimodalGlycanBERTConfig, MultimodalGlycanBERT
    
    print(f"Loading model from {ckpt_path}...")
    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
    state_dict = ckpt.get('model_state_dict', ckpt)
    
    # Infer vocab sizes from checkpoint weights
    seq_vocab_size = state_dict['seq_embeddings.token_embeddings.weight'].shape[0]
    print(f"  seq_vocab_size (from checkpoint): {seq_vocab_size}")
    
    config = MultimodalGlycanBERTConfig(
        seq_vocab_size=seq_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,
    )
    model = MultimodalGlycanBERT(config)
    
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    model.to(device).eval()
    n_params = sum(p.numel() for p in model.parameters())
    print(f"  Model loaded: {n_params:,} params (seq_vocab={seq_vocab_size})")
    return model


def load_glycosylation_data(max_samples=None):
    """Load glycosylation N/O dataset with IUPAC and WURCS."""
    from downstream_tasks.utils.tokenizer import WURCSTokenizer
    tokenizer = WURCSTokenizer(str(VOCAB_PATH))
    
    glyco_csv = BENCH_DIR / 'glycan_link_wurcs_subset.csv'
    samples, labels, iupacs = [], [], []
    with open(glyco_csv) as f:
        for row in csv.DictReader(f):
            if row['link'] not in ('N', 'O'):
                continue
            wurcs = row.get('wurcs', '')
            iupac = row.get('glycan', '')
            if not wurcs:
                continue
            try:
                result = tokenizer.tokenize(wurcs, max_length=256)
                samples.append({
                    'token_ids': result['token_ids'],
                    'branch_depths': result.get('branch_depths', [0]*len(result['token_ids'])),
                    'linkage_types': result.get('linkage_types', [0]*len(result['token_ids'])),
                })
                labels.append(row['link'])
                iupacs.append(iupac)
            except Exception:
                continue
    
    if max_samples and len(samples) > max_samples:
        np.random.seed(42)
        idx = np.random.choice(len(samples), max_samples, replace=False)
        samples = [samples[i] for i in idx]
        labels = [labels[i] for i in idx]
        iupacs = [iupacs[i] for i in idx]
    
    print(f"  Glycosylation data: {len(samples)} samples")
    return samples, labels, iupacs


# ============================================================
# Probe 3: Attention Head Analysis
# ============================================================

def extract_attention_weights(model, samples, device='cuda', max_samples=200):
    """
    Extract attention weight matrices from all layers and heads.
    Uses forward hooks on GlycanBERTAttention modules.
    
    Returns: dict[layer_idx] -> list of (n_heads, seq_len, seq_len) arrays
    """
    # Subsample for memory
    if len(samples) > max_samples:
        np.random.seed(42)
        idx = np.random.choice(len(samples), max_samples, replace=False)
        samples = [samples[i] for i in idx]
    
    attention_store = defaultdict(list)
    hooks = []
    
    # Register hooks on each layer's attention module
    for layer_idx, layer in enumerate(model.seq_layers):
        attn_module = layer.attention
        
        def make_hook(li):
            def hook_fn(module, input, output):
                # We need to intercept the attention probs
                # Since the forward doesn't return them, we re-compute them here
                hidden_states = input[0]
                attention_mask = input[1] if len(input) > 1 else None
                
                query_layer = module.transpose_for_scores(module.query(hidden_states))
                key_layer = module.transpose_for_scores(module.key(hidden_states))
                
                scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
                scores = scores / (module.attention_head_size ** 0.5)
                
                if attention_mask is not None:
                    mask = attention_mask.unsqueeze(1).unsqueeze(2)
                    mask = (1.0 - mask) * -10000.0
                    scores = scores + mask
                
                probs = torch.softmax(scores, dim=-1)
                # Store: (n_heads, seq_len, seq_len)
                attention_store[li].append(probs[0].detach().cpu().numpy())
            return hook_fn
        
        h = attn_module.register_forward_hook(make_hook(layer_idx))
        hooks.append(h)
    
    # Forward pass for each sample
    with torch.no_grad():
        for i, s in enumerate(samples):
            token_ids = torch.tensor(s['token_ids'], dtype=torch.long).unsqueeze(0).to(device)
            branch_depths = torch.tensor(s['branch_depths'], dtype=torch.long).unsqueeze(0).to(device)
            linkage_types = torch.tensor(s['linkage_types'], dtype=torch.long).unsqueeze(0).to(device)
            
            hidden = model.seq_embeddings(token_ids, branch_depths=branch_depths, linkage_types=linkage_types)
            for layer in model.seq_layers:
                hidden = layer(hidden)
            
            if (i + 1) % 50 == 0:
                print(f"    Attention extraction: {i+1}/{len(samples)}")
    
    # Remove hooks
    for h in hooks:
        h.remove()
    
    return attention_store


def compute_attention_entropy(attention_store, n_layers=12, n_heads=12):
    """
    Compute per-head attention entropy (averaged over samples and positions).
    High entropy = diffuse attention; Low entropy = focused/specialized.
    """
    entropy_matrix = np.zeros((n_layers, n_heads))
    
    for layer_idx in range(n_layers):
        attn_list = attention_store[layer_idx]
        head_entropies = []
        
        for attn in attn_list:  # attn: (n_heads, seq_len, seq_len)
            # Compute entropy for each head at each query position
            for h in range(n_heads):
                probs = attn[h]  # (seq_len, seq_len)
                # Entropy: -sum(p * log(p)), handle zeros
                log_probs = np.log(probs + 1e-10)
                ent = -np.sum(probs * log_probs, axis=-1)  # (seq_len,)
                head_entropies.append((h, ent.mean()))
        
        # Average across samples for each head
        for h in range(n_heads):
            h_ents = [e for (hi, e) in head_entropies if hi == h]
            entropy_matrix[layer_idx, h] = np.mean(h_ents)
    
    return entropy_matrix


def compute_cls_attention(attention_store, n_layers=12, n_heads=12):
    """
    Compute how much each head attends FROM [CLS] to all other tokens.
    
    Returns: (n_layers, n_heads) matrix of mean CLS attention spread.
    """
    cls_entropy = np.zeros((n_layers, n_heads))
    cls_max_attn = np.zeros((n_layers, n_heads))
    
    for layer_idx in range(n_layers):
        attn_list = attention_store[layer_idx]
        
        for h in range(n_heads):
            entropies = []
            max_attns = []
            for attn in attn_list:
                cls_row = attn[h, 0, :]  # CLS attending to all tokens
                # Entropy of CLS attention distribution
                log_probs = np.log(cls_row + 1e-10)
                ent = -np.sum(cls_row * log_probs)
                entropies.append(ent)
                max_attns.append(cls_row.max())
            
            cls_entropy[layer_idx, h] = np.mean(entropies)
            cls_max_attn[layer_idx, h] = np.mean(max_attns)
    
    return cls_entropy, cls_max_attn


def plot_attention_heatmap(matrix, title, output_path, xlabel='Attention Head', ylabel='Layer',
                          cmap='viridis', vmin=None, vmax=None):
    """Plot a (n_layers, n_heads) heatmap."""
    fig, ax = plt.subplots(figsize=(10, 6))
    
    im = ax.imshow(matrix, cmap=cmap, aspect='auto', vmin=vmin, vmax=vmax)
    
    ax.set_xlabel(xlabel, fontsize=13)
    ax.set_ylabel(ylabel, fontsize=13)
    ax.set_title(title, fontsize=14, fontweight='bold')
    
    ax.set_xticks(range(matrix.shape[1]))
    ax.set_xticklabels([f'H{i}' for i in range(matrix.shape[1])], fontsize=9)
    ax.set_yticks(range(matrix.shape[0]))
    ax.set_yticklabels([f'L{i+1}' for i in range(matrix.shape[0])], fontsize=10)
    
    # Add value annotations
    for i in range(matrix.shape[0]):
        for j in range(matrix.shape[1]):
            val = matrix[i, j]
            color = 'white' if val < (matrix.max() + matrix.min()) / 2 else 'black'
            ax.text(j, i, f'{val:.2f}', ha='center', va='center', fontsize=7, color=color)
    
    plt.colorbar(im, ax=ax, label='Value')
    plt.tight_layout()
    
    for fmt in ['png', 'pdf']:
        fp = f"{output_path}.{fmt}"
        plt.savefig(fp, dpi=300, bbox_inches='tight', facecolor='white')
        print(f"  Saved: {fp}")
    plt.close()


# ============================================================
# Probe 4: Embedding Distance vs Structural Similarity
# ============================================================

def extract_cls_embeddings(model, samples, device='cuda'):
    """Extract final-layer CLS embeddings."""
    embeddings = []
    with torch.no_grad():
        for s in samples:
            token_ids = torch.tensor(s['token_ids'], dtype=torch.long).unsqueeze(0).to(device)
            branch_depths = torch.tensor(s['branch_depths'], dtype=torch.long).unsqueeze(0).to(device)
            linkage_types = torch.tensor(s['linkage_types'], dtype=torch.long).unsqueeze(0).to(device)
            
            hidden = model.seq_embeddings(token_ids, branch_depths=branch_depths, linkage_types=linkage_types)
            for layer in model.seq_layers:
                hidden = layer(hidden)
            
            cls = hidden[0, 0, :].cpu().numpy()
            embeddings.append(cls)
    
    return np.array(embeddings)


def compute_structural_similarity(iupacs, max_pairs=5000):
    """
    Compute pairwise structural similarity using glycowork motif fingerprints.
    
    Returns: struct_sim, motif_matrix, motif_columns, non_zero_mask
    where non_zero_mask[i] is True if glycan i has a non-zero motif vector.
    """
    try:
        from glycowork.motif.graph import compare_glycans
    except ImportError:
        print("  WARNING: glycowork not available, skipping structural similarity")
        return None, None, None, None
    
    n = len(iupacs)
    
    # First, get motif fingerprints for all glycans
    try:
        from glycowork.motif.annotate import annotate_dataset
        print(f"  Computing motif fingerprints for {n} glycans...")
        motif_df = annotate_dataset(iupacs, feature_set=['known'])
        motif_matrix = motif_df.values.astype(float)
        print(f"  Motif fingerprint shape: {motif_matrix.shape}")
    except Exception as e:
        print(f"  Error computing fingerprints: {e}")
        return None, None, None, None
    
    # Identify glycans with non-zero motif vectors
    row_norms = np.linalg.norm(motif_matrix, axis=1)
    non_zero_mask = row_norms > 0
    n_zero = (~non_zero_mask).sum()
    print(f"  Non-zero motif vectors: {non_zero_mask.sum()}/{n} ({n_zero} zero-vector glycans filtered)")
    
    # Compute pairwise cosine similarity of motif fingerprints
    from sklearn.metrics.pairwise import cosine_similarity
    struct_sim = cosine_similarity(motif_matrix)
    
    return struct_sim, motif_matrix, motif_df.columns.tolist(), non_zero_mask


def compute_embedding_similarity(embeddings):
    """Compute pairwise cosine similarity of CLS embeddings."""
    from sklearn.metrics.pairwise import cosine_similarity
    return cosine_similarity(embeddings)


def _plot_scatter(embed_flat, struct_flat, output_dir, model_name, suffix, title_extra=''):
    """Helper: make scatter + density plot for a set of (embed, struct) pairs."""
    from scipy.stats import spearmanr, pearsonr
    
    pearson_r, pearson_p = pearsonr(embed_flat, struct_flat)
    spearman_r, spearman_p = spearmanr(embed_flat, struct_flat)
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
    
    max_points = 10000
    if len(embed_flat) > max_points:
        idx = np.random.choice(len(embed_flat), max_points, replace=False)
        plot_embed = embed_flat[idx]
        plot_struct = struct_flat[idx]
    else:
        plot_embed = embed_flat
        plot_struct = struct_flat
    
    ax1.scatter(plot_struct, plot_embed, alpha=0.15, s=5, c='#2196F3', rasterized=True)
    ax1.set_xlabel('Structural Similarity\n(Motif Fingerprint Cosine)', fontsize=12)
    ax1.set_ylabel('Embedding Similarity\n(CLS Cosine)', fontsize=12)
    ax1.set_title(f'Embedding vs Structural Similarity{title_extra}\n'
                  f'Pearson r={pearson_r:.3f}, Spearman rho={spearman_r:.3f}',
                  fontsize=13, fontweight='bold')
    
    z = np.polyfit(plot_struct, plot_embed, 1)
    p = np.poly1d(z)
    x_line = np.linspace(plot_struct.min(), plot_struct.max(), 100)
    ax1.plot(x_line, p(x_line), 'r-', linewidth=2, alpha=0.7, label='Linear fit')
    ax1.legend(frameon=False)
    ax1.grid(alpha=0.2)
    
    h = ax2.hist2d(plot_struct, plot_embed, bins=50, cmap='Blues', cmin=1)
    ax2.set_xlabel('Structural Similarity', fontsize=12)
    ax2.set_ylabel('Embedding Similarity', fontsize=12)
    ax2.set_title(f'Density Plot - {model_name}{title_extra}', fontsize=13, fontweight='bold')
    plt.colorbar(h[3], ax=ax2, label='Count')
    
    plt.tight_layout()
    
    out = Path(output_dir)
    for fmt in ['png', 'pdf']:
        fp = out / f'embed_vs_struct_similarity_{model_name.lower()}{suffix}.{fmt}'
        plt.savefig(fp, dpi=300, bbox_inches='tight', facecolor='white')
        print(f"  Saved: {fp}")
    plt.close()
    
    return {
        'pearson_r': float(pearson_r),
        'pearson_p': float(pearson_p),
        'spearman_rho': float(spearman_r),
        'spearman_p': float(spearman_p),
        'n_pairs': int(len(embed_flat)),
    }


def correlate_similarities(embed_sim, struct_sim, output_dir, model_name, non_zero_mask=None):
    """
    Compute and plot correlation between embedding similarity and structural similarity.
    Reports:
      1. ALL pairs (primary result)
      2. Non-zero structural similarity pairs only (pairs that share >= 1 motif)
    """
    from scipy.stats import spearmanr, pearsonr
    
    n = embed_sim.shape[0]
    triu_idx = np.triu_indices(n, k=1)
    
    embed_flat = embed_sim[triu_idx]
    struct_flat = struct_sim[triu_idx]
    
    # ── 1. ALL PAIRS (primary result) ──
    pearson_r, pearson_p = pearsonr(embed_flat, struct_flat)
    spearman_r, spearman_p = spearmanr(embed_flat, struct_flat)
    
    n_zero_sim = int((struct_flat == 0.0).sum())
    n_nonzero_sim = int((struct_flat > 0.0).sum())
    
    print(f"\n  === ALL PAIRS ({n} glycans, {len(embed_flat):,} pairs) ===")
    print(f"  Pearson r  = {pearson_r:.4f} (p={pearson_p:.2e})")
    print(f"  Spearman rho = {spearman_r:.4f} (p={spearman_p:.2e})")
    print(f"  Pairs with struct_sim == 0: {n_zero_sim:,} ({100*n_zero_sim/len(embed_flat):.1f}%)")
    print(f"  Pairs with struct_sim >  0: {n_nonzero_sim:,} ({100*n_nonzero_sim/len(embed_flat):.1f}%)")
    
    all_stats = _plot_scatter(embed_flat, struct_flat, output_dir, model_name,
                              suffix='', title_extra=' (All Glycan Pairs)')
    
    result = {
        'model': model_name,
        'n_glycans': int(n),
        'all_pairs': {
            **all_stats,
            'n_pairs': int(len(embed_flat)),
            'n_zero_sim_pairs': n_zero_sim,
            'n_nonzero_sim_pairs': n_nonzero_sim,
        },
    }
    
    # ── 2. NON-ZERO PAIRS (supplementary — only pairs sharing >= 1 motif) ──
    nz_mask_pairs = struct_flat > 0.0
    if nz_mask_pairs.sum() > 0:
        embed_nz = embed_flat[nz_mask_pairs]
        struct_nz = struct_flat[nz_mask_pairs]
        
        pearson_r_nz, pearson_p_nz = pearsonr(embed_nz, struct_nz)
        spearman_r_nz, spearman_p_nz = spearmanr(embed_nz, struct_nz)
        
        print(f"\n  === NON-ZERO PAIRS (struct_sim > 0, {n_nonzero_sim:,} pairs) ===")
        print(f"  Pearson r  = {pearson_r_nz:.4f} (p={pearson_p_nz:.2e})")
        print(f"  Spearman rho = {spearman_r_nz:.4f} (p={spearman_p_nz:.2e})")
        print(f"  (These are glycan pairs that share at least one structural motif)")
        
        nz_stats = _plot_scatter(embed_nz, struct_nz, output_dir, model_name,
                                 suffix='_nonzero_pairs',
                                 title_extra=f'\n(Non-Zero Pairs Only: {n_nonzero_sim:,} pairs sharing >= 1 motif)')
        
        result['nonzero_pairs'] = {
            **nz_stats,
            'n_pairs': n_nonzero_sim,
        }
    
    # ── Diagnostic summary ──
    if non_zero_mask is not None:
        n_zero_glycans = int((~non_zero_mask).sum())
        result['diagnostic'] = {
            'zero_motif_glycans': n_zero_glycans,
            'motif_dimensions': 165,
            'note': 'Each glycan has ~3 motifs out of 34 used. 52.6% of pairs share zero motifs (cosine=0).'
        }
    
    return result

def main():
    parser = argparse.ArgumentParser(description='Probes 3 & 4: Attention + Embedding vs Structure')
    parser.add_argument('--model', type=str, default='v6', choices=['v5', 'v6'])
    parser.add_argument('--device', type=str, default='cuda')
    parser.add_argument('--output_dir', type=str, 
                        default=str(PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' / 'probe_results_v6'))
    parser.add_argument('--probe', type=str, default='both', choices=['3', '4', 'both'],
                        help='Which probe to run')
    parser.add_argument('--max_attn_samples', type=int, default=200,
                        help='Max samples for attention extraction (memory intensive)')
    args = parser.parse_args()
    
    model_name = args.model.upper()
    out_base = Path(args.output_dir)
    
    print(f"\n{'='*60}")
    print(f"  Probes 3 & 4 \u2014 GlycanBERT {model_name}")
    print(f"{'='*60}")
    
    # Load model
    ckpt = CHECKPOINTS[args.model]
    model = load_model(str(ckpt), device=args.device)
    
    # Load glycosylation data (has both IUPAC and WURCS)
    print("\nLoading glycosylation dataset...")
    samples, labels, iupacs = load_glycosylation_data()
    
    # ==========================================
    # Probe 3: Attention Head Analysis
    # ==========================================
    if args.probe in ('3', 'both'):
        print(f"\n{'='*60}")
        print(f"  PROBE 3: Attention Head Analysis")
        print(f"{'='*60}")
        
        probe3_dir = out_base / 'probe_3_attention_heads'
        probe3_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"\n  Extracting attention weights ({args.max_attn_samples} samples)...")
        attn_store = extract_attention_weights(model, samples, device=args.device, 
                                               max_samples=args.max_attn_samples)
        
        # 3a. Entropy heatmap
        print("\n  Computing attention entropy...")
        entropy_matrix = compute_attention_entropy(attn_store)
        plot_attention_heatmap(
            entropy_matrix,
            f'Attention Head Entropy \u2014 GlycanBERT {model_name}',
            str(probe3_dir / f'attention_entropy_{model_name.lower()}'),
            cmap='YlOrRd_r'
        )
        
        # 3b. CLS attention
        print("\n  Computing CLS attention patterns...")
        cls_entropy, cls_max_attn = compute_cls_attention(attn_store)
        plot_attention_heatmap(
            cls_entropy,
            f'CLS Token Attention Entropy \u2014 GlycanBERT {model_name}',
            str(probe3_dir / f'cls_attention_entropy_{model_name.lower()}'),
            cmap='YlGnBu'
        )
        plot_attention_heatmap(
            cls_max_attn,
            f'CLS Max Attention Weight \u2014 GlycanBERT {model_name}',
            str(probe3_dir / f'cls_max_attention_{model_name.lower()}'),
            cmap='Purples'
        )
        
        # Save numerical results
        results_3 = {
            'entropy_matrix': entropy_matrix.tolist(),
            'cls_entropy': cls_entropy.tolist(),
            'cls_max_attn': cls_max_attn.tolist(),
            'n_samples': args.max_attn_samples,
            'n_layers': 12,
            'n_heads': 12,
            'most_specialized_heads': [],
            'most_diffuse_heads': [],
        }
        
        # Find most specialized (lowest entropy) and most diffuse heads
        flat_ent = [(entropy_matrix[l, h], l, h) for l in range(12) for h in range(12)]
        flat_ent.sort()
        results_3['most_specialized_heads'] = [
            {'layer': int(l+1), 'head': int(h), 'entropy': float(e)} 
            for e, l, h in flat_ent[:10]
        ]
        results_3['most_diffuse_heads'] = [
            {'layer': int(l+1), 'head': int(h), 'entropy': float(e)} 
            for e, l, h in flat_ent[-10:]
        ]
        
        json_path = probe3_dir / f'attention_analysis_{model_name.lower()}.json'
        with open(json_path, 'w') as f:
            json.dump(results_3, f, indent=2)
        print(f"\n  Saved: {json_path}")
        
        # Print summary
        print(f"\n  Top 5 most specialized (focused) heads:")
        for item in results_3['most_specialized_heads'][:5]:
            print(f"    Layer {item['layer']}, Head {item['head']}: entropy={item['entropy']:.3f}")
        print(f"\n  Top 5 most diffuse (spread) heads:")
        for item in results_3['most_diffuse_heads'][:5]:
            print(f"    Layer {item['layer']}, Head {item['head']}: entropy={item['entropy']:.3f}")
    
    # ==========================================
    # Probe 4: Embedding vs Structural Similarity  
    # ==========================================
    if args.probe in ('4', 'both'):
        print(f"\n{'='*60}")
        print(f"  PROBE 4: Embedding vs Structural Similarity")
        print(f"{'='*60}")
        
        probe4_dir = out_base / 'probe_4_embed_vs_structure'
        probe4_dir.mkdir(parents=True, exist_ok=True)
        
        # Extract CLS embeddings for all glycans
        print(f"\n  Extracting CLS embeddings ({len(samples)} samples)...")
        cls_embeddings = extract_cls_embeddings(model, samples, device=args.device)
        print(f"  Embedding shape: {cls_embeddings.shape}")
        
        # Compute structural similarity via glycowork
        print(f"\n  Computing structural similarity via glycowork...")
        struct_sim, motif_matrix, motif_columns, non_zero_mask = compute_structural_similarity(iupacs)
        
        if struct_sim is not None:
            # Compute embedding similarity
            print(f"\n  Computing embedding cosine similarity...")
            embed_sim = compute_embedding_similarity(cls_embeddings)
            
            # Correlate (with zero-motif filtering)
            print(f"\n  Computing correlation (unfiltered + filtered)...")
            corr_results = correlate_similarities(embed_sim, struct_sim, 
                                                  str(probe4_dir), model_name,
                                                  non_zero_mask=non_zero_mask)
            
            # Save results
            json_path = probe4_dir / f'embed_vs_struct_{model_name.lower()}.json'
            with open(json_path, 'w') as f:
                json.dump(corr_results, f, indent=2)
            print(f"\n  Saved: {json_path}")
        else:
            print("  SKIPPED: glycowork not available")
    
    # Free GPU memory
    import gc
    del model
    torch.cuda.empty_cache()
    gc.collect()
    
    print(f"\n{'='*60}")
    print(f"  ALL PROBES COMPLETE")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()