#!/usr/bin/env python3
"""
Probe 15H: Phylogenetic Embedding from Milk Glycomes
=====================================================
Task H: Embed per-species milk glycomes via GlycanBERT V6,
compute species centroids, build UPGMA dendrogram, and test
whether embedding-derived phylogeny recapitulates taxonomy.

Data: Fig1a source data from Jin et al. (2025) — 2,524 glycans
      from 173 species (all mammals).

Usage:
  python probe_15h_phylogenetic.py --model v6
  python probe_15h_phylogenetic.py --model v6 --min-glycans 40
"""

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

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
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'
SOURCE_DATA  = (PROJECT_ROOT / 'bert_v6_contrastive'
                / '41467_2025_66075_MOESM6_ESM.xlsx')
OUTPUT_DIR   = (PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis'
                / 'probe_15_seal_milk')

# Reuse WURCS conversion infrastructure from probe_15
sys.path.insert(0, str(PROJECT_ROOT / 'bert_v6_contrastive' / 'scripts'))
from probe_15_seal_milk_ood import (
    load_model, get_cls_embeddings, batch_iupac_to_wurcs,
    iupac_to_wurcs
)

# --- Nature-style plotting ---
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,
})

ORDER_COLORS = {
    'Primates':       '#E64B35',
    'Carnivora':      '#4DBBD5',
    'Artiodactyla':   '#00A087',
    'Perissodactyla': '#3C5488',
    'Proboscidea':    '#F39B7F',
    'Rodentia':       '#8491B4',
    'Chiroptera':     '#91D1C2',
    'Diprotodontia':  '#DC9050',
    'Monotremata':    '#7E6148',
    'Pilosa':         '#B09C85',
}


# ====================================================================
# Data Loading
# ====================================================================
def load_cross_species_glycans(xlsx_path, min_glycans=10):
    """Load Fig1a from source data Excel."""
    print(f"  Loading cross-species data from {xlsx_path}")
    import openpyxl
    wb = openpyxl.load_workbook(str(xlsx_path), read_only=True)
    ws = wb['Fig1a']

    rows = list(ws.iter_rows(values_only=True))
    header_idx = None
    for i, row in enumerate(rows):
        if row and row[0] and str(row[0]).strip().lower() == 'glycan':
            header_idx = i
            break
    if header_idx is None:
        raise ValueError("Could not find header row in Fig1a")

    data_rows = rows[header_idx + 1:]

    records = []
    for row in data_rows:
        if not row or not row[0]:
            continue
        glycan = str(row[0]).strip()
        species = str(row[1]).strip() if row[1] else ''
        genus = str(row[2]).strip() if row[2] else ''
        family = str(row[3]).strip() if row[3] else ''
        order = str(row[4]).strip() if row[4] else ''
        class_ = str(row[5]).strip() if row[5] else ''
        if glycan and species:
            records.append({
                'glycan': glycan, 'species': species,
                'genus': genus, 'family': family,
                'order': order, 'class_': class_,
            })
    wb.close()

    species_glycans = defaultdict(list)
    species_info = {}
    for r in records:
        sp = r['species']
        species_glycans[sp].append(r['glycan'])
        if sp not in species_info:
            species_info[sp] = {
                'order': r['order'], 'family': r['family'],
                'genus': r['genus'], 'class_': r['class_'],
            }

    qualified = {sp: glycans for sp, glycans in species_glycans.items()
                 if len(glycans) >= min_glycans}

    total_glycans = sum(len(g) for g in species_glycans.values())
    qual_glycans = sum(len(g) for g in qualified.values())
    print(f"  Total: {len(species_glycans)} species, {total_glycans} glycans")
    print(f"  Qualified (>={min_glycans} glycans): {len(qualified)} species, "
          f"{qual_glycans} glycans")

    print(f"\n  {'Species':<35s} {'Order':<18s} {'N':>4s}")
    print(f"  {'-'*60}")
    for sp in sorted(qualified, key=lambda s: -len(qualified[s])):
        info = species_info[sp]
        print(f"  {sp:<35s} {info['order']:<18s} {len(qualified[sp]):>4d}")

    return qualified, species_info


# ====================================================================
# Phylogenetic Analysis
# ====================================================================
def compute_species_centroids(model, tokenizer, device,
                              species_glycans, species_info):
    """Embed all glycans, compute per-species centroid (mean embedding)."""
    all_iupacs = []
    glycan_to_species = {}
    for sp, glycans in species_glycans.items():
        for g in glycans:
            if g not in glycan_to_species:
                all_iupacs.append(g)
                glycan_to_species[g] = []
            glycan_to_species[g].append(sp)

    print(f"\n  Total unique glycans to embed: {len(all_iupacs)}")

    print("  Converting IUPAC -> WURCS...")
    wurcs_list, valid_iupacs, failed = batch_iupac_to_wurcs(all_iupacs)
    print(f"  WURCS conversion: {len(valid_iupacs)}/{len(all_iupacs)} "
          f"({len(failed)} failed)")

    print("  Embedding glycans...")
    embeddings, emb_valid_idx = get_cls_embeddings(
        model, tokenizer, wurcs_list, device
    )
    print(f"  Embedded: {embeddings.shape[0]}/{len(wurcs_list)}")

    final_iupacs = [valid_iupacs[i] for i in emb_valid_idx]

    centroids = {}
    centroid_n = {}
    for sp, glycans in species_glycans.items():
        sp_indices = []
        for i, iupac in enumerate(final_iupacs):
            if iupac in glycans:
                sp_indices.append(i)
        if len(sp_indices) >= 3:
            sp_embs = embeddings[sp_indices]
            centroids[sp] = sp_embs.mean(axis=0)
            centroid_n[sp] = len(sp_indices)
        else:
            print(f"  WARNING: {sp} only has {len(sp_indices)} embedded "
                  f"glycans, skipping")

    print(f"\n  Species with valid centroids: {len(centroids)}")
    for sp in sorted(centroids, key=lambda s: -centroid_n[s]):
        info = species_info[sp]
        print(f"    {sp:<35s} {info['order']:<18s} "
              f"n={centroid_n[sp]:>3d}")

    return centroids, centroid_n


def build_phylogenetic_tree(centroids, species_info, output_dir):
    """Build UPGMA dendrogram from species centroid cosine distances."""
    from scipy.cluster.hierarchy import linkage, dendrogram
    from scipy.spatial.distance import squareform
    from sklearn.metrics.pairwise import cosine_distances
    from scipy.stats import pearsonr, spearmanr

    species_list = sorted(centroids.keys())
    n = len(species_list)
    centroid_matrix = np.array([centroids[sp] for sp in species_list])

    dist_matrix = cosine_distances(centroid_matrix)

    dist_df = pd.DataFrame(dist_matrix, index=species_list, columns=species_list)
    csv_path = output_dir / 'task_h_distance_matrix.csv'
    dist_df.to_csv(csv_path)
    print(f"\n  Distance matrix saved to {csv_path}")

    upper = dist_matrix[np.triu_indices(n, k=1)]
    print(f"  Distance stats: min={upper.min():.4f}, max={upper.max():.4f}, "
          f"mean={upper.mean():.4f}, median={np.median(upper):.4f}")

    # Mantel test: embedding distance vs taxonomic distance
    print("\n  Computing Mantel test (embedding vs taxonomy)...")
    tax_dist = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            si = species_info[species_list[i]]
            sj = species_info[species_list[j]]
            if species_list[i] == species_list[j]:
                tax_dist[i, j] = 0
            elif si.get('genus') == sj.get('genus'):
                tax_dist[i, j] = 1
            elif si.get('family') == sj.get('family'):
                tax_dist[i, j] = 2
            elif si.get('order') == sj.get('order'):
                tax_dist[i, j] = 3
            elif si.get('class_') == sj.get('class_'):
                tax_dist[i, j] = 4
            else:
                tax_dist[i, j] = 5

    emb_upper = dist_matrix[np.triu_indices(n, k=1)]
    tax_upper = tax_dist[np.triu_indices(n, k=1)]

    r_pearson, p_pearson = pearsonr(emb_upper, tax_upper)
    r_spearman, p_spearman = spearmanr(emb_upper, tax_upper)
    print(f"  Mantel (Pearson):  r={r_pearson:.4f}, p={p_pearson:.4e}")
    print(f"  Mantel (Spearman): r={r_spearman:.4f}, p={p_spearman:.4e}")

    # Permutation test
    n_perm = 999
    perm_correlations = []
    for _ in range(n_perm):
        perm_idx = np.random.permutation(n)
        perm_tax = tax_dist[np.ix_(perm_idx, perm_idx)]
        perm_upper = perm_tax[np.triu_indices(n, k=1)]
        perm_r, _ = pearsonr(emb_upper, perm_upper)
        perm_correlations.append(perm_r)
    mantel_p = (np.sum(np.array(perm_correlations) >= r_pearson) + 1) / (n_perm + 1)
    print(f"  Mantel permutation p-value: {mantel_p:.4f} ({n_perm} permutations)")

    # UPGMA Dendrogram
    print("\n  Building UPGMA dendrogram...")
    condensed = squareform(dist_matrix)
    Z = linkage(condensed, method='average')

    COMMON_NAMES = {
        'Homo_sapiens': 'Human', 'Halichoerus_grypus': 'Grey seal',
        'Bos_taurus': 'Cow', 'Capra_hircus': 'Goat',
        'Sus_scrofa': 'Pig', 'Tursiops_truncatus': 'Dolphin',
        'Choeropsis_liberiensis': 'Pygmy hippo',
        'Delphinapterus_leucas': 'Beluga whale',
        'Ovis_aries': 'Sheep', 'Panthera_leo': 'Lion',
        'Camelus_dromedarius': 'Dromedary',
        'Equus_caballus': 'Horse',
        'Ailuropoda_melanoleuca': 'Giant panda',
        'Ornithorhynchus_anatinus': 'Platypus',
        'Tachyglossus_aculeatus': 'Echidna',
    }

    labels = []
    label_colors = {}
    for sp in species_list:
        common = COMMON_NAMES.get(sp, '')
        short_sp = sp.replace('_', ' ')
        label = f"{common} ({short_sp})" if common else short_sp
        labels.append(label)
        order = species_info[sp].get('order', 'Unknown')
        label_colors[label] = ORDER_COLORS.get(order, '#888888')

    fig_height = max(8, len(species_list) * 0.35)
    fig, ax = plt.subplots(1, 1, figsize=(10, fig_height))
    dendro = dendrogram(Z, labels=labels, orientation='left', ax=ax,
                        leaf_font_size=8, color_threshold=0,
                        above_threshold_color='#333333')

    ylbls = ax.get_yticklabels()
    for lbl in ylbls:
        txt = lbl.get_text()
        color = label_colors.get(txt, '#333333')
        lbl.set_color(color)
        lbl.set_fontweight('bold')

    ax.set_xlabel('Cosine Distance')
    ax.set_title(f'Task H: Phylogenetic Tree from V6 Milk Glycome Embeddings\n'
                 f'(Mantel r={r_pearson:.3f}, p={mantel_p:.4f}, '
                 f'n={len(species_list)} species)', fontsize=11)

    from matplotlib.patches import Patch
    orders_present = set(species_info[sp]['order'] for sp in species_list)
    legend_patches = [Patch(facecolor=ORDER_COLORS.get(o, '#888'), label=o)
                      for o in sorted(orders_present) if o]
    ax.legend(handles=legend_patches, loc='upper right', fontsize=7,
              frameon=True, framealpha=0.8, title='Taxonomic Order',
              title_fontsize=8)

    plt.tight_layout()
    tree_path = output_dir / 'task_h_dendrogram.png'
    plt.savefig(tree_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"  Dendrogram saved to {tree_path}")

    # Phylogenetic consistency checks
    print("\n  Phylogenetic consistency checks:")
    checks = {'Artiodactyla': [], 'Carnivora': [], 'Primates': []}
    for i, sp in enumerate(species_list):
        order = species_info[sp].get('order', '')
        if order in checks:
            checks[order].append(i)

    for group_name, indices in checks.items():
        if len(indices) >= 2:
            intra = [dist_matrix[i, j] for i in indices for j in indices if i < j]
            inter = [dist_matrix[i, j] for i in indices
                     for j in range(n) if j not in indices]
            if intra and inter:
                intra_mean = np.mean(intra)
                inter_mean = np.mean(inter)
                ratio = intra_mean / inter_mean if inter_mean > 0 else 0
                print(f"    {group_name}: intra={intra_mean:.4f}, "
                      f"inter={inter_mean:.4f}, ratio={ratio:.4f}")

    results = {
        'task': 'H_phylogenetic_embedding',
        'n_species': len(species_list),
        'mantel_pearson_r': float(r_pearson),
        'mantel_pearson_p': float(p_pearson),
        'mantel_spearman_r': float(r_spearman),
        'mantel_spearman_p': float(p_spearman),
        'mantel_permutation_p': float(mantel_p),
        'n_permutations': n_perm,
        'distance_mean': float(upper.mean()),
        'distance_min': float(upper.min()),
        'distance_max': float(upper.max()),
        'dendrogram_path': str(tree_path),
        'distance_matrix_path': str(csv_path),
        'species': {sp: {'order': species_info[sp].get('order', ''),
                         'family': species_info[sp].get('family', '')}
                    for sp in species_list},
    }
    return results


# ====================================================================
# Main
# ====================================================================
def main():
    parser = argparse.ArgumentParser(
        description='Probe 15H: Phylogenetic Embedding from Milk Glycomes'
    )
    parser.add_argument('--model', choices=['v5', 'v6'], default='v6')
    parser.add_argument('--device', default='cuda' if torch.cuda.is_available()
                        else 'cpu')
    parser.add_argument('--min-glycans', type=int, default=10,
                        help='Min glycans per species for centroid (default 10)')
    args = parser.parse_args()

    device = torch.device(args.device)
    print("=" * 60)
    print("Probe 15H: Phylogenetic Embedding from Milk Glycomes")
    print(f"Device: {device}")
    print(f"Model: {args.model}")
    print(f"Min glycans/species: {args.min_glycans}")
    print("=" * 60)

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print("\n[1] Loading cross-species glycan data...")
    species_glycans, species_info = load_cross_species_glycans(
        SOURCE_DATA, min_glycans=args.min_glycans
    )

    print(f"\n[2] Loading model {args.model}...")
    model, tokenizer = load_model(args.model, device)

    print("\n[3] Computing species centroids...")
    centroids, centroid_n = compute_species_centroids(
        model, tokenizer, device, species_glycans, species_info
    )

    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    print("\n[4] Building phylogenetic tree...")
    results = build_phylogenetic_tree(centroids, species_info, OUTPUT_DIR)

    results_path = OUTPUT_DIR / 'task_h_results.json'
    with open(results_path, 'w') as f:
        json.dump(results, f, indent=2, default=str)
    print(f"\nResults saved to {results_path}")

    print(f"\n{'=' * 60}")
    print("TASK H SUMMARY")
    print(f"{'=' * 60}")
    print(f"  Species:            {results['n_species']}")
    print(f"  Mantel r (Pearson): {results['mantel_pearson_r']:.4f} "
          f"(p={results['mantel_permutation_p']:.4f})")
    print(f"  Mantel r (Spearman): {results['mantel_spearman_r']:.4f}")
    print(f"  Distance range:     [{results['distance_min']:.4f}, "
          f"{results['distance_max']:.4f}]")
    print(f"\nDone. Dendrogram: {results['dendrogram_path']}")


if __name__ == '__main__':
    main()