#!/usr/bin/env python3
"""
Tier 2: Core Biology Embedding Analysis
Analyzes whether GlycanBERT understands glycan biology at a structural level.

Analyses:
  1. Glycan structural type clustering (N-glycan, O-glycan, GAG, etc.)
  2. Monosaccharide composition vs embedding distance correlation
  3. Branching complexity analysis
  4. Linkage chemistry sensitivity (alpha vs beta)
  5. Decorations (fucosylation, sialylation)
  6. Plantae / Phytomining investigation

Usage:
  python3 analyze_core_biology.py --model v6
  python3 analyze_core_biology.py --model v5
"""

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

try:
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
except ImportError:
    print("ERROR: matplotlib required"); sys.exit(1)

try:
    import umap
    def compute_umap(data, n=2):
        return umap.UMAP(n_components=n, n_neighbors=30, min_dist=0.3,
                         random_state=42, n_jobs=1).fit_transform(data)
except ImportError:
    from sklearn.manifold import TSNE
    def compute_umap(data, n=2):
        return TSNE(n_components=n, perplexity=30, random_state=42).fit_transform(data)

from scipy.stats import spearmanr

# ── Glycan Type Classification ───────────────────────────────────────

def classify_glycan_type(iupac):
    """Classify a glycan by structural type from its IUPAC name."""
    if not iupac or iupac == 'nan':
        return 'Unknown'
    iupac_upper = iupac.upper()
    if 'GLCNAC(B1-4)GLCNAC' in iupac_upper and 'MAN' in iupac_upper:
        return 'N-glycan'
    if iupac_upper.count('MAN') >= 5 and 'GLCNAC' in iupac_upper:
        return 'N-glycan (high-mannose)'
    if 'GALNAC(A1-' in iupac_upper:
        return 'O-glycan'
    if any(x in iupac_upper for x in ['GLCA', 'IDOA', 'HEPARAN', 'CHONDROITIN']):
        return 'GAG'
    if iupac_upper.startswith('GLC(B1-') or iupac_upper.startswith('GLCCER'):
        return 'Glycolipid'
    if 'LAC' in iupac_upper and len(iupac) < 60:
        return 'Milk OS'
    return 'Other'

# ── Monosaccharide Composition ───────────────────────────────────────

MONOSACCHARIDES = ['Glc', 'GlcNAc', 'Man', 'Gal', 'GalNAc', 'Fuc',
                   'Neu5Ac', 'Xyl', 'Rha', 'GlcA', 'IdoA', 'Ara']

def get_composition(iupac):
    """Count monosaccharides in an IUPAC string."""
    if not iupac:
        return np.zeros(len(MONOSACCHARIDES))
    vec = []
    for mono in MONOSACCHARIDES:
        count = len(re.findall(re.escape(mono) + r'[\(\[]', iupac, re.IGNORECASE))
        if count == 0:
            count = iupac.upper().count(mono.upper())
        vec.append(count)
    return np.array(vec, dtype=float)

# ── Analysis Functions ───────────────────────────────────────────────

def analysis_1_structural_types(embs, iupacs, output_dir, name):
    """UMAP colored by glycan structural type."""
    print(f"\n=== Tier 2 Analysis 1: Structural Type Clustering ({name}) ===")
    types = [classify_glycan_type(s) for s in iupacs]
    type_counts = Counter(types)
    print(f"  Type distribution: {dict(type_counts)}")
    valid_types = {t for t, c in type_counts.items() if c >= 5 and t != 'Unknown'}
    mask = np.array([t in valid_types for t in types])
    if mask.sum() < 20:
        print(f"  Only {mask.sum()} samples with known type, skipping UMAP.")
        return {'structural_type_counts': dict(type_counts)}
    embs_f = embs[mask]
    types_f = [t for t, m in zip(types, mask) if m]
    proj = compute_umap(embs_f)
    colors = ['#E53935', '#1E88E5', '#43A047', '#FB8C00', '#8E24AA',
              '#00ACC1', '#6D4C41', '#546E7A']
    unique_types = sorted(set(types_f))
    fig, ax = plt.subplots(figsize=(14, 10))
    for i, gtype in enumerate(unique_types):
        m = np.array([t == gtype for t in types_f])
        ax.scatter(proj[m, 0], proj[m, 1], c=colors[i % len(colors)],
                   s=15, alpha=0.6, label=f'{gtype} (n={m.sum()})', rasterized=True)
    ax.set_title(f'{name}: Glycan Structural Types in Embedding Space', fontsize=16, fontweight='bold')
    ax.legend(fontsize=10, markerscale=3, loc='best')
    ax.set_xlabel('UMAP-1', fontsize=12); ax.set_ylabel('UMAP-2', fontsize=12)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'structural_types_{name}.png'), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  Saved structural_types_{name}.png")
    try:
        from sklearn.metrics import silhouette_score
        labels_num = [unique_types.index(t) for t in types_f]
        if len(set(labels_num)) >= 2:
            sil = silhouette_score(embs_f, labels_num, metric='cosine', sample_size=min(5000, len(embs_f)))
            print(f"  Silhouette (structural type): {sil:.4f}")
            return {'structural_type_counts': dict(type_counts), 'silhouette_structural_type': round(sil, 4)}
    except: pass
    return {'structural_type_counts': dict(type_counts)}


def analysis_2_composition(embs, iupacs, output_dir, name, n_sample=2000):
    """Monosaccharide composition vs embedding distance correlation."""
    print(f"\n=== Tier 2 Analysis 2: Composition Correlation ({name}) ===")
    n = min(n_sample, len(embs))
    idx = np.random.RandomState(42).choice(len(embs), n, replace=False)
    comp_vecs = np.array([get_composition(iupacs[i]) for i in idx])
    emb_subset = embs[idx]
    norms = np.linalg.norm(comp_vecs, axis=1)
    valid = norms > 0
    if valid.sum() < 50:
        print(f"  Only {valid.sum()} samples with parseable composition, skipping.")
        return {}
    comp_vecs = comp_vecs[valid]
    emb_subset = emb_subset[valid]
    print(f"  {len(comp_vecs)} samples with valid composition")
    emb_n = emb_subset / (np.linalg.norm(emb_subset, axis=1, keepdims=True) + 1e-8)
    comp_n = comp_vecs / (np.linalg.norm(comp_vecs, axis=1, keepdims=True) + 1e-8)
    emb_sim = emb_n @ emb_n.T
    comp_sim = comp_n @ comp_n.T
    mask = np.triu(np.ones_like(emb_sim, dtype=bool), k=1)
    emb_flat = emb_sim[mask]
    comp_flat = comp_sim[mask]
    if len(emb_flat) > 50000:
        sample_idx = np.random.RandomState(42).choice(len(emb_flat), 50000, replace=False)
        emb_flat = emb_flat[sample_idx]
        comp_flat = comp_flat[sample_idx]
    rho, pval = spearmanr(comp_flat, emb_flat)
    print(f"  Spearman rho = {rho:.4f} (p = {pval:.2e})")
    fig, ax = plt.subplots(figsize=(10, 8))
    sample = np.random.RandomState(42).choice(len(emb_flat), min(5000, len(emb_flat)), replace=False)
    ax.scatter(comp_flat[sample], emb_flat[sample], s=1, alpha=0.15, c='#1565C0', rasterized=True)
    ax.set_xlabel('Compositional Similarity (cosine)', fontsize=14)
    ax.set_ylabel('Embedding Similarity (cosine)', fontsize=14)
    ax.set_title(f'{name}: Composition vs Embedding Similarity\nSpearman rho = {rho:.4f}', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'composition_correlation_{name}.png'), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  Saved composition_correlation_{name}.png")
    return {'composition_spearman_rho': round(rho, 4), 'composition_p_value': float(f'{pval:.2e}'),
            'n_composition_samples': int(len(comp_vecs))}


def analysis_3_branching(embs, branch_depths_list, output_dir, name):
    """Branching complexity analysis."""
    print(f"\n=== Tier 2 Analysis 3: Branching Complexity ({name}) ===")
    max_depths = []
    for bd in branch_depths_list:
        if isinstance(bd, (list, np.ndarray)) and len(bd) > 0:
            max_depths.append(int(max(bd)))
        else:
            max_depths.append(0)
    max_depths = np.array(max_depths)
    depth_counts = Counter(max_depths)
    print(f"  Depth distribution: {dict(sorted(depth_counts.items()))}")
    def depth_category(d):
        if d <= 1: return 'Linear (depth<=1)'
        elif d == 2: return 'Branched (depth=2)'
        elif d == 3: return 'Complex (depth=3)'
        else: return 'Hyper-branched (depth>=4)'
    categories = [depth_category(d) for d in max_depths]
    cat_counts = Counter(categories)
    valid_cats = {c for c, n in cat_counts.items() if n >= 10}
    mask = np.array([c in valid_cats for c in categories])
    if mask.sum() < 30:
        print(f"  Insufficient category diversity ({mask.sum()} valid), skipping UMAP.")
        return {'branching_distribution': dict(depth_counts)}
    embs_f = embs[mask]
    cats_f = [c for c, m in zip(categories, mask) if m]
    proj = compute_umap(embs_f)
    colors = {'Linear (depth<=1)': '#42A5F5', 'Branched (depth=2)': '#66BB6A',
              'Complex (depth=3)': '#FFA726', 'Hyper-branched (depth>=4)': '#EF5350'}
    unique_cats = sorted(set(cats_f))
    fig, ax = plt.subplots(figsize=(14, 10))
    for cat in unique_cats:
        m = np.array([c == cat for c in cats_f])
        ax.scatter(proj[m, 0], proj[m, 1], c=colors.get(cat, '#999'),
                   s=15, alpha=0.5, label=f'{cat} (n={m.sum()})', rasterized=True)
    ax.set_title(f'{name}: Branching Complexity in Embedding Space', fontsize=16, fontweight='bold')
    ax.legend(fontsize=11, markerscale=3)
    ax.set_xlabel('UMAP-1', fontsize=12); ax.set_ylabel('UMAP-2', fontsize=12)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'branching_complexity_{name}.png'), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  Saved branching_complexity_{name}.png")
    try:
        from sklearn.metrics import silhouette_score
        labels_num = [unique_cats.index(c) for c in cats_f]
        if len(set(labels_num)) >= 2:
            sil = silhouette_score(embs_f, labels_num, metric='cosine', sample_size=min(5000, len(embs_f)))
            print(f"  Silhouette (branching): {sil:.4f}")
            return {'branching_distribution': dict(depth_counts), 'silhouette_branching': round(sil, 4)}
    except: pass
    return {'branching_distribution': dict(depth_counts)}


def analysis_4_linkage(embs, linkage_types_list, output_dir, name):
    """Linkage chemistry analysis: alpha vs beta distribution in embedding space."""
    print(f"\n=== Tier 2 Analysis 4: Linkage Chemistry ({name}) ===")
    alpha_fracs = []
    for lt in linkage_types_list:
        if isinstance(lt, (list, np.ndarray)) and len(lt) > 0:
            arr = np.array(lt)
            n_valid = np.sum(arr > 0)
            if n_valid > 0:
                n_alpha = np.sum(arr == 1)
                alpha_fracs.append(n_alpha / n_valid)
            else:
                alpha_fracs.append(0.5)
        else:
            alpha_fracs.append(0.5)
    alpha_fracs = np.array(alpha_fracs)
    def linkage_category(frac):
        if frac >= 0.8: return 'Mostly alpha'
        elif frac <= 0.2: return 'Mostly beta'
        else: return 'Mixed alpha/beta'
    categories = [linkage_category(f) for f in alpha_fracs]
    cat_counts = Counter(categories)
    print(f"  Linkage distribution: {dict(cat_counts)}")
    valid_cats = {c for c, n in cat_counts.items() if n >= 10}
    mask = np.array([c in valid_cats for c in categories])
    if mask.sum() < 30:
        print(f"  Insufficient diversity, skipping UMAP.")
        return {'linkage_distribution': dict(cat_counts)}
    embs_f = embs[mask]
    cats_f = [c for c, m in zip(categories, mask) if m]
    proj = compute_umap(embs_f)
    colors = {'Mostly alpha': '#E53935', 'Mixed alpha/beta': '#FDD835', 'Mostly beta': '#1E88E5'}
    unique_cats = sorted(set(cats_f))
    fig, ax = plt.subplots(figsize=(14, 10))
    for cat in unique_cats:
        m = np.array([c == cat for c in cats_f])
        ax.scatter(proj[m, 0], proj[m, 1], c=colors.get(cat, '#999'),
                   s=15, alpha=0.5, label=f'{cat} (n={m.sum()})', rasterized=True)
    ax.set_title(f'{name}: Linkage Chemistry (alpha vs beta) in Embedding Space', fontsize=16, fontweight='bold')
    ax.legend(fontsize=12, markerscale=3)
    ax.set_xlabel('UMAP-1', fontsize=12); ax.set_ylabel('UMAP-2', fontsize=12)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'linkage_chemistry_{name}.png'), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  Saved linkage_chemistry_{name}.png")
    alpha_mask = np.array([c == 'Mostly alpha' for c in categories])
    beta_mask = np.array([c == 'Mostly beta' for c in categories])
    if alpha_mask.sum() >= 10 and beta_mask.sum() >= 10:
        a_embs = embs[alpha_mask][:500]
        b_embs = embs[beta_mask][:500]
        a_n = a_embs / (np.linalg.norm(a_embs, axis=1, keepdims=True) + 1e-8)
        b_n = b_embs / (np.linalg.norm(b_embs, axis=1, keepdims=True) + 1e-8)
        cross_sim = np.mean(a_n @ b_n.T)
        within_a = np.mean(a_n @ a_n.T)
        within_b = np.mean(b_n @ b_n.T)
        print(f"  a-a sim: {within_a:.4f}, b-b sim: {within_b:.4f}, a-b sim: {cross_sim:.4f}")
        return {'linkage_distribution': dict(cat_counts),
                'alpha_alpha_sim': round(float(within_a), 4),
                'beta_beta_sim': round(float(within_b), 4),
                'alpha_beta_cross_sim': round(float(cross_sim), 4)}
    return {'linkage_distribution': dict(cat_counts)}


def analysis_5_decorations(embs, iupacs, output_dir, name):
    """Fucosylation / sialylation decoration analysis."""
    print(f"\n=== Tier 2 Analysis 5: Decorations ({name}) ===")
    fuc_mask = np.array(['FUC' in (s or '').upper() for s in iupacs])
    sia_mask = np.array([any(x in (s or '').upper() for x in ['NEU5AC','NEU5GC','SIA']) for s in iupacs])
    print(f"  Fucosylated: {fuc_mask.sum()}/{len(fuc_mask)}")
    print(f"  Sialylated:  {sia_mask.sum()}/{len(sia_mask)}")
    n = min(5000, len(embs))
    idx = np.random.RandomState(42).choice(len(embs), n, replace=False)
    proj = compute_umap(embs[idx])
    fig, axes = plt.subplots(1, 2, figsize=(22, 9))
    for ax, (dec_name, mask_full, color) in zip(axes, [
        ('Fucosylated', fuc_mask, '#E53935'),
        ('Sialylated', sia_mask, '#7B1FA2')
    ]):
        m = mask_full[idx]
        ax.scatter(proj[~m, 0], proj[~m, 1], c='#BDBDBD', s=5, alpha=0.2,
                   label=f'No (n={int((~m).sum())})', rasterized=True)
        ax.scatter(proj[m, 0], proj[m, 1], c=color, s=15, alpha=0.6,
                   label=f'{dec_name} (n={int(m.sum())})', rasterized=True)
        ax.set_title(f'{name}: {dec_name}', fontsize=16, fontweight='bold')
        ax.legend(fontsize=12, markerscale=3)
        ax.set_xlabel('UMAP-1'); ax.set_ylabel('UMAP-2')
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'decorations_{name}.png'), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  Saved decorations_{name}.png")
    metrics = {
        'n_fucosylated': int(fuc_mask.sum()),
        'n_sialylated': int(sia_mask.sum()),
        'n_total': int(len(iupacs))
    }
    if fuc_mask.sum() >= 20 and (~fuc_mask).sum() >= 20:
        f_embs = embs[fuc_mask][:500]
        nf_embs = embs[~fuc_mask][:500]
        f_n = f_embs / (np.linalg.norm(f_embs, axis=1, keepdims=True) + 1e-8)
        nf_n = nf_embs / (np.linalg.norm(nf_embs, axis=1, keepdims=True) + 1e-8)
        cross = np.mean(f_n @ nf_n.T)
        within_f = np.mean(f_n @ f_n.T)
        print(f"  Fuc-Fuc sim: {within_f:.4f}, Fuc-NonFuc sim: {cross:.4f}")
        metrics['fuc_fuc_sim'] = round(float(within_f), 4)
        metrics['fuc_nonfuc_sim'] = round(float(cross), 4)
    return metrics


def analysis_6_phytomining_plantae(benchmark_embs, benchmark_kingdoms, output_dir, name):
    """Plantae glycan cluster analysis (phytomining proxy).
    Literature basis: Plant polysaccharides (cellulose, EPS, sulfated sugars)
    bind metals via hydroxyl groups and 1,2-diolato coordination. Plantae glycans
    with GlcA, Xyl, Rha, Ara monosaccharides are characteristic of plant cell wall
    components relevant to phytomining/phytoremediation.
    """
    print(f"\n=== Tier 2 Analysis 6: Plantae / Phytomining ({name}) ===")
    plant_mask = np.array([k == 'Plantae' for k in benchmark_kingdoms])
    animal_mask = np.array([k == 'Animalia' for k in benchmark_kingdoms])
    bacteria_mask = np.array([k == 'Bacteria' for k in benchmark_kingdoms])
    fungi_mask = np.array([k == 'Fungi' for k in benchmark_kingdoms])
    print(f"  Plantae: {plant_mask.sum()}, Animalia: {animal_mask.sum()}, Bacteria: {bacteria_mask.sum()}, Fungi: {fungi_mask.sum()}")
    if plant_mask.sum() < 10:
        print("  Not enough Plantae samples.")
        return {}
    proj = compute_umap(benchmark_embs)
    fig, ax = plt.subplots(figsize=(14, 10))
    other = ~(plant_mask | animal_mask | bacteria_mask | fungi_mask)
    if other.sum() > 0:
        ax.scatter(proj[other, 0], proj[other, 1], c='#BDBDBD', s=8, alpha=0.3,
                   label=f'Other (n={other.sum()})', rasterized=True)
    if fungi_mask.sum() > 0:
        ax.scatter(proj[fungi_mask, 0], proj[fungi_mask, 1], c='#AB47BC', s=20, alpha=0.6,
                   label=f'Fungi (n={fungi_mask.sum()})', rasterized=True)
    if bacteria_mask.sum() > 0:
        ax.scatter(proj[bacteria_mask, 0], proj[bacteria_mask, 1], c='#42A5F5', s=20, alpha=0.6,
                   label=f'Bacteria (n={bacteria_mask.sum()})', rasterized=True)
    if animal_mask.sum() > 0:
        ax.scatter(proj[animal_mask, 0], proj[animal_mask, 1], c='#FFA726', s=20, alpha=0.6,
                   label=f'Animalia (n={animal_mask.sum()})', rasterized=True)
    ax.scatter(proj[plant_mask, 0], proj[plant_mask, 1], c='#2E7D32', s=30, alpha=0.8,
               edgecolors='black', linewidths=0.5,
               label=f'Plantae (n={plant_mask.sum()})', rasterized=True)
    ax.set_title(f'{name}: Plantae Glycans in Embedding Space (Phytomining Proxy)',
                 fontsize=16, fontweight='bold')
    ax.legend(fontsize=12, markerscale=2.5)
    ax.set_xlabel('UMAP-1', fontsize=12); ax.set_ylabel('UMAP-2', fontsize=12)
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f'plantae_phytomining_{name}.png'), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  Saved plantae_phytomining_{name}.png")
    metrics = {}
    for k_name, k_mask in [('Plantae', plant_mask), ('Animalia', animal_mask), ('Bacteria', bacteria_mask), ('Fungi', fungi_mask)]:
        if k_mask.sum() >= 10:
            e = benchmark_embs[k_mask][:300]
            e_n = e / (np.linalg.norm(e, axis=1, keepdims=True) + 1e-8)
            within = np.triu(e_n @ e_n.T, k=1)
            mask_tri = np.triu(np.ones_like(within, dtype=bool), k=1)
            if mask_tri.sum() > 0:
                metrics[f'{k_name.lower()}_within_sim'] = round(float(np.mean(within[mask_tri])), 4)
    if plant_mask.sum() >= 10 and animal_mask.sum() >= 10:
        p_n = benchmark_embs[plant_mask][:200]
        p_n = p_n / (np.linalg.norm(p_n, axis=1, keepdims=True) + 1e-8)
        a_n = benchmark_embs[animal_mask][:200]
        a_n = a_n / (np.linalg.norm(a_n, axis=1, keepdims=True) + 1e-8)
        metrics['plant_animal_cross_sim'] = round(float(np.mean(p_n @ a_n.T)), 4)
    if plant_mask.sum() >= 10 and bacteria_mask.sum() >= 10:
        p_n = benchmark_embs[plant_mask][:200]
        p_n = p_n / (np.linalg.norm(p_n, axis=1, keepdims=True) + 1e-8)
        b_n = benchmark_embs[bacteria_mask][:200]
        b_n = b_n / (np.linalg.norm(b_n, axis=1, keepdims=True) + 1e-8)
        metrics['plant_bacteria_cross_sim'] = round(float(np.mean(p_n @ b_n.T)), 4)
    for k, v in metrics.items():
        print(f"  {k}: {v}")
    return metrics


# ── Main ─────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', required=True, choices=['v5', 'v6'])
    parser.add_argument('--output-dir', default='bert_v6_contrastive/analysis')
    args = parser.parse_args()

    root = Path('.')
    output_dir = args.output_dir
    os.makedirs(output_dir, exist_ok=True)

    # Load embeddings
    emb_path = os.path.join(output_dir, f'embeddings_{args.model}.npz')
    print(f"Loading embeddings from {emb_path}...")
    data = np.load(emb_path, allow_pickle=True)
    train_embs = data['train_embs']
    benchmark_embs = data['benchmark_embs']
    benchmark_kingdoms = data['benchmark_kingdom']

    # Load original training data for IUPAC names, branch_depths, linkage_types
    data_path = root / 'bert_v5_bpe_topo' / 'data' / 'sequences_bpe_expanded.pkl'
    print(f"Loading training data from {data_path}...")
    with open(data_path, 'rb') as f:
        all_samples = pickle.load(f)

    # We need IUPAC names for the same 10K training samples
    # The extraction script used random.sample with seed 42
    import random
    random.seed(42)
    n_train = min(10000, len(all_samples))
    sampled = random.sample(all_samples, n_train)

    iupacs = [s.get('iupac_name', '') or '' for s in sampled]
    branch_depths = [s.get('branch_depths', []) for s in sampled]
    linkage_types = [s.get('linkage_types', []) for s in sampled]

    print(f"  {len(iupacs)} training samples with metadata")
    print(f"  IUPAC example: {iupacs[0][:80]}")

    # Run analyses
    metrics = {'model': args.model, 'tier': 2}

    m = analysis_1_structural_types(train_embs, iupacs, output_dir, args.model)
    metrics.update(m)

    m = analysis_2_composition(train_embs, iupacs, output_dir, args.model)
    metrics.update(m)

    m = analysis_3_branching(train_embs, branch_depths, output_dir, args.model)
    metrics.update(m)

    m = analysis_4_linkage(train_embs, linkage_types, output_dir, args.model)
    metrics.update(m)

    m = analysis_5_decorations(train_embs, iupacs, output_dir, args.model)
    metrics.update(m)

    m = analysis_6_phytomining_plantae(benchmark_embs, benchmark_kingdoms, output_dir, args.model)
    metrics.update(m)

    # Save metrics
    out_path = os.path.join(output_dir, f'tier2_metrics_{args.model}.json')
    def convert(o):
        if isinstance(o, (np.integer,)): return int(o)
        if isinstance(o, (np.floating,)): return float(o)
        if isinstance(o, np.ndarray): return o.tolist()
        return o
    with open(out_path, 'w') as f:
        json.dump(metrics, f, indent=2, default=convert)
    print(f"\nAll Tier 2 metrics saved to {out_path}")

if __name__ == '__main__':
    main()