#!/usr/bin/env python3
"""
Generate publication-quality t-SNE figures for GlycanBERT V6 embeddings.
Uses FULL transformer forward pass (embeddings + 12 transformer layers).

Panels:
  (a) Taxonomy classification at domain level (Eukarya, Virus, Bacteria)
  (b) Glycosylation type (N-linked, O-linked, free)
  (c) Immunogenicity (non-immunogenic vs immunogenic)
  (d) Unsupervised cluster analysis with motif enrichment labels
"""

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

warnings.filterwarnings('ignore')

PROJECT_ROOT = Path(__file__).resolve().parents[2]
VOCAB_PATH = PROJECT_ROOT / 'bert_training_v4' / 'data' / 'vocabulary.json'
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',
}
BENCH_DIR = PROJECT_ROOT / 'bench' / 'GlycanML' / 'data'

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
from model.multimodal_glycan_bert_v3 import MultimodalGlycanBERT, MultimodalGlycanBERTConfig
from downstream_tasks.utils.tokenizer import WURCSTokenizer

# ═══════════════════════════════════════════════════════════════════
# Model loading & embedding
# ═══════════════════════════════════════════════════════════════════
def load_model(ckpt_path, device='cuda'):
    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)
    backbone_sd = {k: v for k, v in state_dict.items() if not k.startswith('proj_head.')}
    vocab_size = backbone_sd['seq_embeddings.token_embeddings.weight'].shape[0]
    ms_total_vocab = None
    if 'ms_embeddings.token_embeddings.weight' in backbone_sd:
        ms_total_vocab = backbone_sd['ms_embeddings.token_embeddings.weight'].shape[0]
    config_kwargs = dict(
        seq_vocab_size=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,
    )
    if ms_total_vocab is not None:
        config_kwargs['ms_vocab_size'] = ms_total_vocab - vocab_size
    config = MultimodalGlycanBERTConfig(**config_kwargs)
    model = MultimodalGlycanBERT(config)
    model.load_state_dict(backbone_sd, strict=False)
    model.to(device).eval()
    print(f"  Model loaded: {sum(p.numel() for p in model.parameters()):,} params")
    return model

def get_cls_embeddings(model, samples, device='cuda', batch_size=32, max_len=256):
    """Get CLS embeddings using FULL transformer forward pass."""
    tokenizer = WURCSTokenizer(str(VOCAB_PATH))
    all_embs, n_errors = [], 0
    for i in range(0, len(samples), batch_size):
        batch = samples[i:i+batch_size]
        for s in batch:
            try:
                result = tokenizer.tokenize(s['wurcs'], max_length=max_len)
                token_ids = torch.tensor(result['token_ids'], dtype=torch.long)
                branch_depths = torch.tensor(result.get('branch_depths', [0]*len(result['token_ids'])), dtype=torch.long)
                linkage_types = torch.tensor(result.get('linkage_types', [0]*len(result['token_ids'])), dtype=torch.long)
                min_l = min(len(token_ids), len(branch_depths), len(linkage_types))
                token_ids, branch_depths, linkage_types = token_ids[:min_l], branch_depths[:min_l], linkage_types[:min_l]
                if min_l > max_len:
                    token_ids, branch_depths, linkage_types = token_ids[:max_len], branch_depths[:max_len], linkage_types[:max_len]
                elif min_l < max_len:
                    pad = max_len - min_l
                    token_ids = F.pad(token_ids, (0, pad), value=0)
                    branch_depths = F.pad(branch_depths, (0, pad), value=0)
                    linkage_types = F.pad(linkage_types, (0, pad), value=0)
                with torch.no_grad():
                    hidden = model.seq_embeddings(
                        token_ids.unsqueeze(0).to(device),
                        branch_depths.unsqueeze(0).to(device),
                        linkage_types.unsqueeze(0).to(device)
                    )
                    for layer in model.seq_layers:
                        hidden = layer(hidden)
                    all_embs.append(hidden[0, 0, :].cpu().numpy())
            except Exception as e:
                n_errors += 1
                if n_errors <= 3: print(f"  ERROR: {e}")
                all_embs.append(np.zeros(768))
        if (i // batch_size) % 20 == 0 and i > 0:
            print(f"    Embedded {i}/{len(samples)}")
    if n_errors > 0: print(f"  WARNING: {n_errors}/{len(samples)} errors")
    return np.array(all_embs)

# ═══════════════════════════════════════════════════════════════════
# Data loading
# ═══════════════════════════════════════════════════════════════════
def load_domain_data():
    """Domain: Eukarya, Bacteria, Virus (no Archaea)."""
    csv_path = BENCH_DIR / 'glycan_classification_wurcs_subset.csv'
    samples, labels = [], []
    with open(csv_path) as f:
        for row in csv.DictReader(f):
            w, domain = row.get('wurcs',''), row.get('domain','')
            if w.startswith('WURCS') and domain in ('Eukarya', 'Bacteria', 'Virus'):
                samples.append({'wurcs': w})
                labels.append(domain)
    print(f"  Domain: {len(samples)}, {Counter(labels)}")
    return samples, labels

def load_glycosylation_data():
    """Glycosylation: N, O, free."""
    csv_path = BENCH_DIR / 'glycan_link_wurcs_subset.csv'
    samples, labels, iupacs = [], [], []
    with open(csv_path) as f:
        for row in csv.DictReader(f):
            w = row.get('wurcs','')
            link = row.get('link','')
            iupac = row.get('glycan','')
            if w.startswith('WURCS') and link in ('N', 'O', 'free'):
                samples.append({'wurcs': w})
                labels.append(link)
                iupacs.append(iupac)
    print(f"  Glycosylation: {len(samples)}, {Counter(labels)}")
    return samples, labels, iupacs

def load_immunogenicity_data():
    csv_path = BENCH_DIR / 'glycan_immunogenicity_wurcs_subset.csv'
    samples, labels = [], []
    with open(csv_path) as f:
        for row in csv.DictReader(f):
            w, imm = row.get('wurcs',''), row.get('immunogenicity','')
            if w.startswith('WURCS') and imm:
                samples.append({'wurcs': w})
                labels.append(float(imm))
    print(f"  Immunogenicity: {len(samples)}, {Counter(labels)}")
    return samples, labels

# ═══════════════════════════════════════════════════════════════════
# Motif enrichment using glycowork
# ═══════════════════════════════════════════════════════════════════
def run_motif_enrichment(cluster_labels, iupac_glycans, n_clusters=6):
    """Run glycowork motif enrichment per cluster using annotate_dataset."""
    try:
        from glycowork.motif.annotate import annotate_dataset
    except ImportError:
        print("  WARNING: glycowork not found, skipping motif enrichment")
        return {}
    
    # Step 1: Annotate all glycans with known motifs
    print("    Annotating glycans with known motifs...")
    try:
        motif_df = annotate_dataset(list(iupac_glycans), feature_set=['known'])
    except Exception as e:
        print(f"  Annotation error: {e}")
        return {ci: "—" for ci in range(n_clusters)}
    
    # Step 2: For each cluster, find the most enriched motif
    # (motif with highest ratio of cluster-mean / overall-mean)
    cluster_arr = np.array(cluster_labels)
    enriched = {}
    motif_details = {}
    
    # Get numeric columns only (motif presence/count columns)
    numeric_cols = [c for c in motif_df.columns if motif_df[c].dtype in ('float64','int64','float32','int32')]
    if len(numeric_cols) == 0:
        print("  WARNING: No numeric motif columns found")
        return {ci: "—" for ci in range(n_clusters)}
    
    overall_means = motif_df[numeric_cols].mean()
    
    for ci in range(n_clusters):
        mask = cluster_arr == ci
        if mask.sum() < 3:
            enriched[ci] = "N/A"
            continue
        
        cluster_means = motif_df.loc[mask, numeric_cols].mean()
        rest_means = motif_df.loc[~mask, numeric_cols].mean()
        
        # Enrichment ratio: cluster_mean / (rest_mean + eps) — higher = more specific
        eps = 0.001
        enrichment = (cluster_means + eps) / (rest_means + eps)
        
        # Filter: only motifs that actually appear in this cluster
        appears = cluster_means > 0.05
        if appears.sum() > 0:
            enrichment_filtered = enrichment[appears]
            top_motif = enrichment_filtered.idxmax()
            enriched[ci] = str(top_motif)
            motif_details[ci] = f"{top_motif} (ratio={enrichment[top_motif]:.2f})"
        else:
            enriched[ci] = "—"
            motif_details[ci] = "no enriched motifs"
    
    print("  Motif enrichment results:")
    for ci in range(n_clusters):
        n = int((cluster_arr == ci).sum())
        detail = motif_details.get(ci, enriched.get(ci, "—"))
        print(f"    Cluster {ci} (n={n}): {detail}")
    return enriched

# ═══════════════════════════════════════════════════════════════════
# Figure generation
# ═══════════════════════════════════════════════════════════════════
def generate_figure(embs_domain, labels_domain,
                    embs_glyco, labels_glyco, iupac_glyco,
                    embs_immuno, labels_immuno,
                    output_dir, model_name='V6'):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from sklearn.manifold import TSNE
    from sklearn.cluster import KMeans

    # Publication style
    plt.rcParams.update({
        'font.family': 'sans-serif',
        'font.sans-serif': ['Arial', 'Helvetica', 'DejaVu Sans'],
        'font.size': 15,
        'axes.titlesize': 18,
        'axes.labelsize': 16,
        'xtick.labelsize': 13,
        'ytick.labelsize': 13,
        'legend.fontsize': 14,
        'legend.title_fontsize': 15,
        'figure.dpi': 300,
        'savefig.dpi': 300,
        'savefig.bbox': 'tight',
        'axes.linewidth': 1.2,
        'axes.spines.top': False,
        'axes.spines.right': False,
    })

    # ─── t-SNE ───
    print("\n  Running t-SNE for domain (~12K, perplexity=50)...")
    tsne_domain = TSNE(n_components=2, perplexity=50, random_state=42, max_iter=1000,
                       learning_rate='auto', init='pca').fit_transform(embs_domain)
    print("  Running t-SNE for glycosylation...")
    tsne_glyco = TSNE(n_components=2, perplexity=30, random_state=42, max_iter=1000,
                       learning_rate='auto', init='pca').fit_transform(embs_glyco)
    print("  Running t-SNE for immunogenicity...")
    tsne_immuno = TSNE(n_components=2, perplexity=30, random_state=42, max_iter=1000,
                        learning_rate='auto', init='pca').fit_transform(embs_immuno)

    # Clip outliers (2nd/98th percentile)
    def clip_tsne(coords):
        for dim in range(2):
            lo, hi = np.percentile(coords[:, dim], [2, 98])
            coords[:, dim] = np.clip(coords[:, dim], lo, hi)
        return coords
    tsne_domain = clip_tsne(tsne_domain)
    tsne_glyco = clip_tsne(tsne_glyco)
    tsne_immuno = clip_tsne(tsne_immuno)

    # ─── K-Means + Motif Enrichment ───
    print("  Running K-Means (6 clusters)...")
    kmeans = KMeans(n_clusters=6, random_state=42, n_init=10)
    cluster_labels = kmeans.fit_predict(embs_glyco)

    print("  Running motif enrichment (glycowork)...")
    motif_labels = run_motif_enrichment(cluster_labels, iupac_glyco, n_clusters=6)

    # ─── Color palettes ───
    domain_colors = {'Eukarya': '#E8A838', 'Virus': '#2EC4B6', 'Bacteria': '#3D7DD8'}
    glyco_colors = {'N': '#56B870', 'O': '#5DADE2', 'free': '#E67E73'}
    immuno_colors = {0.0: '#4DAF7C', 1.0: '#6C7FE8'}
    cluster_cmap = ['#3D7DD8', '#F5A623', '#2ECC71', '#E8534E', '#9B59B6', '#D4A03C']

    # ─── Helper: draw panel ───
    def draw_panel(ax, tsne, labels, colors, order, legend_title, dot_size=20, alpha=0.65):
        for label in order:
            mask = np.array([l == label for l in labels])
            if mask.sum() > 0:
                ax.scatter(tsne[mask, 0], tsne[mask, 1],
                          c=colors[label], s=dot_size, alpha=alpha,
                          label=str(label), edgecolors='none', rasterized=True)
        ax.set_xlabel('t-SNE 1', fontsize=15, labelpad=8)
        ax.set_ylabel('t-SNE 2', fontsize=15, labelpad=8)
        leg = ax.legend(title=legend_title, frameon=True, fancybox=True, framealpha=0.92,
                        markerscale=3.5, loc='best', edgecolor='#bbb', fontsize=13)
        leg.get_title().set_fontsize(14)
        leg.get_frame().set_linewidth(0.6)

    # ═══ Combined 4-panel figure ═══
    fig, axes = plt.subplots(2, 2, figsize=(17, 15))
    plt.subplots_adjust(hspace=0.30, wspace=0.28)

    # (a) Domain
    draw_panel(axes[0,0], tsne_domain, labels_domain, domain_colors,
               ['Eukarya','Bacteria','Virus'], 'Domain', dot_size=8, alpha=0.5)
    axes[0,0].text(-0.08, 1.06, '(a)', transform=axes[0,0].transAxes, fontsize=20, fontweight='bold')

    # (b) Glycosylation (N, O, free)
    draw_panel(axes[0,1], tsne_glyco, labels_glyco, glyco_colors,
               ['N','O','free'], 'Glycosylation', dot_size=25, alpha=0.65)
    axes[0,1].text(-0.08, 1.06, '(b)', transform=axes[0,1].transAxes, fontsize=20, fontweight='bold')

    # (c) Immunogenicity
    draw_panel(axes[1,0], tsne_immuno, labels_immuno, immuno_colors,
               [0.0, 1.0], 'Immunogenicity', dot_size=25, alpha=0.65)
    axes[1,0].text(-0.08, 1.06, '(c)', transform=axes[1,0].transAxes, fontsize=20, fontweight='bold')

    # (d) Clusters with motif annotations
    ax = axes[1,1]
    for ci in range(6):
        mask = cluster_labels == ci
        if mask.sum() > 0:
            lbl = motif_labels.get(ci, str(ci))
            display = f"C{ci}: {lbl}" if lbl and lbl not in ('—','N/A') else f"C{ci}"
            ax.scatter(tsne_glyco[mask, 0], tsne_glyco[mask, 1],
                      c=cluster_cmap[ci], s=30, alpha=0.7,
                      label=display, edgecolors='none', rasterized=True)
    ax.set_xlabel('t-SNE 1', fontsize=15, labelpad=8)
    ax.set_ylabel('t-SNE 2', fontsize=15, labelpad=8)
    leg = ax.legend(title='Cluster (motif)', frameon=True, fancybox=True, framealpha=0.92,
                    markerscale=3, loc='best', edgecolor='#bbb', fontsize=11)
    leg.get_title().set_fontsize(13)
    leg.get_frame().set_linewidth(0.6)
    axes[1,1].text(-0.08, 1.06, '(d)', transform=axes[1,1].transAxes, fontsize=20, fontweight='bold')

    # Add cluster centroid annotations
    for ci in range(6):
        mask = cluster_labels == ci
        if mask.sum() > 5:
            cx, cy = np.median(tsne_glyco[mask, 0]), np.median(tsne_glyco[mask, 1])
            motif = motif_labels.get(ci, '')
            if motif and motif not in ('—', 'N/A'):
                ax.annotate(motif, (cx, cy), fontsize=8, fontweight='bold',
                           ha='center', va='center',
                           bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.8, edgecolor='#999'))

    out_path = Path(output_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    fig_png = out_path / f'tsne_figure_{model_name.lower()}.png'
    fig_pdf = out_path / f'tsne_figure_{model_name.lower()}.pdf'
    plt.savefig(fig_png, dpi=300, bbox_inches='tight', facecolor='white')
    plt.savefig(fig_pdf, dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"\n  Saved: {fig_png}")
    print(f"  Saved: {fig_pdf}")

    # ═══ Individual panels ═══
    panels = [
        ('domain', tsne_domain, labels_domain, domain_colors, ['Eukarya','Bacteria','Virus'], 'Domain', 8),
        ('glycosylation', tsne_glyco, labels_glyco, glyco_colors, ['N','O','free'], 'Glycosylation', 25),
        ('immunogenicity', tsne_immuno, labels_immuno, immuno_colors, [0.0, 1.0], 'Immunogenicity', 25),
    ]
    for pname, tdata, labs, cols, order, title, sz in panels:
        fig2, ax2 = plt.subplots(figsize=(8, 7))
        draw_panel(ax2, tdata, labs, cols, order, title, dot_size=sz)
        ax2.set_title(title, fontsize=18, fontweight='bold', pad=12)
        plt.tight_layout()
        plt.savefig(out_path / f'tsne_{pname}_{model_name.lower()}.png', dpi=300, bbox_inches='tight', facecolor='white')
        plt.savefig(out_path / f'tsne_{pname}_{model_name.lower()}.pdf', dpi=300, bbox_inches='tight', facecolor='white')
        plt.close()
        print(f"  Saved: tsne_{pname}_{model_name.lower()}.png")

    # Cluster standalone
    fig3, ax3 = plt.subplots(figsize=(9, 7))
    for ci in range(6):
        mask = cluster_labels == ci
        if mask.sum() > 0:
            lbl = motif_labels.get(ci, str(ci))
            display = f"C{ci}: {lbl}" if lbl and lbl not in ('—','N/A') else f"C{ci}"
            ax3.scatter(tsne_glyco[mask, 0], tsne_glyco[mask, 1],
                       c=cluster_cmap[ci], s=35, alpha=0.7,
                       label=display, edgecolors='none', rasterized=True)
    for ci in range(6):
        mask = cluster_labels == ci
        if mask.sum() > 5:
            cx, cy = np.median(tsne_glyco[mask, 0]), np.median(tsne_glyco[mask, 1])
            motif = motif_labels.get(ci, '')
            if motif and motif not in ('—', 'N/A'):
                ax3.annotate(motif, (cx, cy), fontsize=9, fontweight='bold',
                            ha='center', va='center',
                            bbox=dict(boxstyle='round,pad=0.3', facecolor='white', alpha=0.85, edgecolor='#999'))
    ax3.set_xlabel('t-SNE 1', fontsize=15, labelpad=8)
    ax3.set_ylabel('t-SNE 2', fontsize=15, labelpad=8)
    ax3.set_title('Cluster analysis (motif enrichment)', fontsize=18, fontweight='bold', pad=12)
    ax3.legend(title='Cluster', frameon=True, fancybox=True, framealpha=0.92,
               markerscale=3, edgecolor='#bbb', fontsize=12)
    plt.tight_layout()
    plt.savefig(out_path / f'tsne_clusters_{model_name.lower()}.png', dpi=300, bbox_inches='tight', facecolor='white')
    plt.savefig(out_path / f'tsne_clusters_{model_name.lower()}.pdf', dpi=300, bbox_inches='tight', facecolor='white')
    plt.close()
    print(f"  Saved: tsne_clusters_{model_name.lower()}.png")

    # ═══ Save CSV (for reproducibility) ═══
    # Domain
    df_domain = pd.DataFrame({
        'tsne_1': tsne_domain[:, 0],
        'tsne_2': tsne_domain[:, 1],
        'domain': labels_domain,
    })
    df_domain.to_csv(out_path / f'tsne_domain_{model_name.lower()}.csv', index=False)
    # Glycosylation
    df_glyco = pd.DataFrame({
        'tsne_1': tsne_glyco[:, 0],
        'tsne_2': tsne_glyco[:, 1],
        'glycosylation': labels_glyco,
        'cluster': cluster_labels,
        'iupac': iupac_glyco,
    })
    df_glyco.to_csv(out_path / f'tsne_glycosylation_{model_name.lower()}.csv', index=False)
    # Immunogenicity
    df_immuno = pd.DataFrame({
        'tsne_1': tsne_immuno[:, 0],
        'tsne_2': tsne_immuno[:, 1],
        'immunogenicity': labels_immuno,
    })
    df_immuno.to_csv(out_path / f'tsne_immunogenicity_{model_name.lower()}.csv', index=False)
    print(f"  Saved CSVs: tsne_domain/glycosylation/immunogenicity_{model_name.lower()}.csv")

    # Motif enrichment results
    motif_df = pd.DataFrame([
        {'cluster': ci, 'n_samples': sum(1 for c in cluster_labels if c == ci),
         'top_motif': motif_labels.get(ci, '')}
        for ci in range(6)
    ])
    motif_df.to_csv(out_path / f'motif_enrichment_{model_name.lower()}.csv', index=False)
    print(f"  Saved: motif_enrichment_{model_name.lower()}.csv")

    # Also save npz for backward compat
    np.savez(out_path / f'tsne_coordinates_{model_name.lower()}.npz',
             domain_tsne=tsne_domain, domain_labels=labels_domain,
             glyco_tsne=tsne_glyco, glyco_labels=labels_glyco,
             immuno_tsne=tsne_immuno, immuno_labels=labels_immuno,
             cluster_labels=cluster_labels)

# ═══════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════
def main():
    parser = argparse.ArgumentParser()
    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' / 'tsne_figures'))
    args = parser.parse_args()

    model_name = args.model.upper()
    print(f"\n{'='*60}")
    print(f"  t-SNE Figure — GlycanBERT {model_name} (Full Transformer)")
    print(f"{'='*60}")

    model = load_model(str(CHECKPOINTS[args.model]), device=args.device)

    print("\nLoading datasets...")
    domain_samples, domain_labels = load_domain_data()
    glyco_samples, glyco_labels, glyco_iupacs = load_glycosylation_data()
    immuno_samples, immuno_labels = load_immunogenicity_data()

    print(f"\nEmbedding with FULL transformer (12 layers)...")
    print(f"  Domain ({len(domain_samples)})...")
    embs_domain = get_cls_embeddings(model, domain_samples, device=args.device)
    print(f"  Glycosylation ({len(glyco_samples)})...")
    embs_glyco = get_cls_embeddings(model, glyco_samples, device=args.device)
    print(f"  Immunogenicity ({len(immuno_samples)})...")
    embs_immuno = get_cls_embeddings(model, immuno_samples, device=args.device)

    import gc; del model; torch.cuda.empty_cache(); gc.collect()

    print("\nGenerating t-SNE figures...")
    generate_figure(embs_domain, domain_labels,
                    embs_glyco, glyco_labels, glyco_iupacs,
                    embs_immuno, immuno_labels,
                    args.output_dir, model_name)

    print(f"\n{'='*60}")
    print(f"  COMPLETE")
    print(f"{'='*60}")

if __name__ == '__main__':
    main()