#!/usr/bin/env python3
"""
Probe 8: Disease/Tissue Embedding Analysis - GlycanBERT
========================================================
Uses glycowork metadata (disease_association, tissue_sample, glycan_type)
to probe what biological knowledge GlycanBERT embeddings encode.

Sub-probes:
  8a. Glycan Type Classification (N / O / free / repeat / lipid)  [16,556 samples]
  8b. Tissue Source Classification (top-10 tissues)                [~2,000 samples]
  8c. Disease Association (cancer vs non-cancer vs healthy)        [~200 samples]
  8d. t-SNE Visualization colored by glycan_type + tissue          [all 39,873]

Data: bert_v6_contrastive/analysis/probe8_disease_tissue_data.csv

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

import sys
import os
import json
import csv
import argparse
import ast
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_PATH = PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' / 'probe8_disease_tissue_data.csv'

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

GTYPE_COLORS = {
    'N': '#0072B2', 'O': '#E69F00', 'free': '#009E73',
    'repeat': '#D55E00', 'lipid': '#CC79A7',
}
TISSUE_COLORS = [
    '#0072B2', '#E69F00', '#009E73', '#D55E00', '#CC79A7',
    '#56B4E9', '#F0E442', '#999999', '#882255', '#44AA99',
]

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'):
    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()
    print(f"  Loaded: {sum(p.numel() for p in model.parameters()):,} params, vocab={vocab_size}, hidden={hidden}")
    return model


def batch_cls_embeddings(model, wurcs_list, tokenizer, device='cuda',
                         batch_size=128, max_len=256):
    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)
        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)
        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)
        if (i // batch_size) % 20 == 0:
            print(f"    Embedded {i+len(batch)}/{len(wurcs_list)} ({errors} errors)")
    print(f"  Total 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():
    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


def probe_8a_glycan_type(embs, labels, output_dir, model_name):
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import StratifiedKFold
    from sklearn.metrics import accuracy_score, f1_score
    from sklearn.preprocessing import LabelEncoder
    print(f"\n{'='*60}")
    print(f"PROBE 8a: Glycan Type Classification ({model_name})")
    print(f"{'='*60}")
    le = LabelEncoder()
    y = le.fit_transform(labels)
    classes = le.classes_
    print(f"  Classes: {list(classes)}")
    print(f"  Distribution: {dict(zip(classes, np.bincount(y)))}")
    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    accs, f1s = [], []
    for fold, (train_idx, test_idx) in enumerate(skf.split(embs, y)):
        clf = LogisticRegression(max_iter=1000, C=1.0, solver='lbfgs',
                                 multi_class='multinomial')
        clf.fit(embs[train_idx], y[train_idx])
        preds = clf.predict(embs[test_idx])
        acc = accuracy_score(y[test_idx], preds)
        f1 = f1_score(y[test_idx], preds, average='macro')
        accs.append(acc)
        f1s.append(f1)
        print(f"  Fold {fold+1}: acc={acc:.4f}, F1={f1:.4f}")
    mean_acc, std_acc = np.mean(accs), np.std(accs)
    mean_f1, std_f1 = np.mean(f1s), np.std(f1s)
    print(f"\n  RESULT: Accuracy = {mean_acc:.4f} +/- {std_acc:.4f}")
    print(f"          F1 Macro = {mean_f1:.4f} +/- {std_f1:.4f}")
    results = {
        'task': 'glycan_type_classification', 'model': model_name,
        'n_samples': len(labels), 'n_classes': len(classes),
        'classes': list(classes),
        'distribution': {str(c): int(n) for c, n in zip(classes, np.bincount(y))},
        'accuracy_mean': float(mean_acc), 'accuracy_std': float(std_acc),
        'f1_macro_mean': float(mean_f1), 'f1_macro_std': float(std_f1),
        'per_fold_accuracy': [float(a) for a in accs],
    }
    with open(os.path.join(output_dir, f'probe8a_glycan_type_{model_name}.json'), 'w') as f:
        json.dump(results, f, indent=2)
    return results


def probe_8b_tissue(embs, labels, output_dir, model_name):
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import StratifiedKFold
    from sklearn.metrics import accuracy_score, f1_score
    from sklearn.preprocessing import LabelEncoder
    print(f"\n{'='*60}")
    print(f"PROBE 8b: Tissue Classification ({model_name})")
    print(f"{'='*60}")
    counts = Counter(labels)
    top_tissues = [t for t, c in counts.most_common(10) if c >= 30]
    mask = np.array([l in top_tissues for l in labels])
    embs_f = embs[mask]
    labels_f = np.array(labels)[mask]
    if len(labels_f) < 50:
        print("  SKIP: Too few samples for tissue classification")
        return {'task': 'tissue_classification', 'model': model_name, 'skipped': True}
    le = LabelEncoder()
    y = le.fit_transform(labels_f)
    classes = le.classes_
    print(f"  Classes ({len(classes)}): {list(classes)}")
    print(f"  Samples: {len(labels_f)}")
    n_splits = min(5, min(Counter(y).values()))
    if n_splits < 2:
        print("  SKIP: Not enough samples per class for stratified CV")
        return {'task': 'tissue_classification', 'model': model_name, 'skipped': True}
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
    accs, f1s = [], []
    for fold, (train_idx, test_idx) in enumerate(skf.split(embs_f, y)):
        clf = LogisticRegression(max_iter=1000, C=1.0, solver='lbfgs',
                                 multi_class='multinomial')
        clf.fit(embs_f[train_idx], y[train_idx])
        preds = clf.predict(embs_f[test_idx])
        acc = accuracy_score(y[test_idx], preds)
        f1 = f1_score(y[test_idx], preds, average='macro')
        accs.append(acc)
        f1s.append(f1)
        print(f"  Fold {fold+1}: acc={acc:.4f}, F1={f1:.4f}")
    mean_acc, std_acc = np.mean(accs), np.std(accs)
    mean_f1, std_f1 = np.mean(f1s), np.std(f1s)
    print(f"\n  RESULT: Accuracy = {mean_acc:.4f} +/- {std_acc:.4f}")
    print(f"          F1 Macro = {mean_f1:.4f} +/- {std_f1:.4f}")
    results = {
        'task': 'tissue_classification', 'model': model_name,
        'n_samples': int(len(labels_f)), 'n_classes': int(len(classes)),
        'classes': list(classes),
        'accuracy_mean': float(mean_acc), 'accuracy_std': float(std_acc),
        'f1_macro_mean': float(mean_f1), 'f1_macro_std': float(std_f1),
        'per_fold_accuracy': [float(a) for a in accs],
    }
    with open(os.path.join(output_dir, f'probe8b_tissue_{model_name}.json'), 'w') as f:
        json.dump(results, f, indent=2)
    return results


def probe_8c_disease(embs_all, disease_labels, wurcs_all, output_dir, model_name):
    print(f"\n{'='*60}")
    print(f"PROBE 8c: Disease Association ({model_name})")
    print(f"{'='*60}")
    cancer_keywords = ['cancer', 'carcinoma', 'tumor', 'glioblastoma', 'melanoma',
                       'leukemia', 'lymphoma', 'sarcoma', 'adenoma', 'myeloma',
                       'neuroblastoma', 'hepatoblastoma']
    labels = []
    for d in disease_labels:
        if not d:
            labels.append('healthy')
        elif any(kw in d.lower() for kw in cancer_keywords):
            labels.append('cancer')
        else:
            labels.append('non_cancer_disease')
    labels = np.array(labels)
    counts = Counter(labels)
    print(f"  Distribution: {dict(counts)}")
    has_disease = labels != 'healthy'
    n_disease = has_disease.sum()
    print(f"  Disease glycans: {n_disease}")
    if n_disease < 50:
        print("  SKIP: Too few disease glycans for classification")
        results = {
            'task': 'disease_association', 'model': model_name,
            'n_total': int(len(labels)), 'distribution': {str(k): int(v) for k, v in counts.items()},
            'skipped': True, 'reason': f'Only {n_disease} disease glycans'
        }
        with open(os.path.join(output_dir, f'probe8c_disease_{model_name}.json'), 'w') as f:
            json.dump(results, f, indent=2)
        return results
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import StratifiedKFold
    from sklearn.metrics import accuracy_score, f1_score, roc_auc_score
    disease_idx = np.where(has_disease)[0]
    healthy_idx = np.where(~has_disease)[0]
    np.random.seed(42)
    healthy_sub = np.random.choice(healthy_idx, size=min(len(healthy_idx), len(disease_idx)*3), replace=False)
    idx = np.concatenate([disease_idx, healthy_sub])
    np.random.shuffle(idx)
    X = embs_all[idx]
    y = (labels[idx] != 'healthy').astype(int)
    print(f"  Binary probe: disease={y.sum()}, healthy={len(y)-y.sum()}")
    n_splits = min(5, min(Counter(y).values()))
    if n_splits < 2:
        results = {'task': 'disease_association', 'model': model_name, 'skipped': True}
        with open(os.path.join(output_dir, f'probe8c_disease_{model_name}.json'), 'w') as f:
            json.dump(results, f, indent=2)
        return results
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
    accs, f1s, aucs = [], [], []
    for fold, (tr, te) in enumerate(skf.split(X, y)):
        clf = LogisticRegression(max_iter=1000, C=1.0)
        clf.fit(X[tr], y[tr])
        preds = clf.predict(X[te])
        proba = clf.predict_proba(X[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:
            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}+/-{np.std(accs):.4f}, AUC={np.mean(aucs):.4f}+/-{np.std(aucs):.4f}")
    results = {
        'task': 'disease_association', 'model': model_name,
        'n_disease': int(n_disease), 'n_healthy_subsample': int(len(healthy_sub)),
        'distribution': {str(k): int(v) for k, v in counts.items()},
        'accuracy_mean': float(np.mean(accs)), 'accuracy_std': float(np.std(accs)),
        'f1_mean': float(np.mean(f1s)), 'auc_mean': float(np.mean(aucs)),
    }
    with open(os.path.join(output_dir, f'probe8c_disease_{model_name}.json'), 'w') as f:
        json.dump(results, f, indent=2)
    return results


def probe_8d_visualization(embs, gtype_labels, tissue_labels, disease_labels,
                           output_dir, model_name):
    plt = setup_nature_style()
    from sklearn.manifold import TSNE
    print(f"\n{'='*60}")
    print(f"PROBE 8d: t-SNE Visualization ({model_name})")
    print(f"{'='*60}")
    n = len(embs)
    if n > 10000:
        np.random.seed(42)
        idx = np.random.choice(n, 10000, replace=False)
        embs_sub = embs[idx]
        gtype_sub = np.array(gtype_labels)[idx]
        tissue_sub = np.array(tissue_labels)[idx]
        disease_sub = np.array(disease_labels)[idx]
    else:
        embs_sub = embs
        gtype_sub = np.array(gtype_labels)
        tissue_sub = np.array(tissue_labels)
        disease_sub = np.array(disease_labels)
    print(f"  Running t-SNE on {len(embs_sub)} samples (perplexity=50)...")
    tsne = TSNE(n_components=2, perplexity=50, random_state=42, max_iter=1000,
                learning_rate='auto', init='pca')
    coords = tsne.fit_transform(embs_sub)
    fig, axes = plt.subplots(1, 3, figsize=(21, 6))
    ax = axes[0]
    for gtype in ['N', 'O', 'free', 'repeat', 'lipid']:
        mask = gtype_sub == gtype
        if mask.sum() > 0:
            ax.scatter(coords[mask, 0], coords[mask, 1],
                      c=GTYPE_COLORS.get(gtype, 'gray'), s=8, alpha=0.5,
                      label=f'{gtype} ({mask.sum()})', edgecolors='none', rasterized=True)
    no_type = gtype_sub == ''
    if no_type.sum() > 0:
        ax.scatter(coords[no_type, 0], coords[no_type, 1],
                  c='#DDDDDD', s=4, alpha=0.2, label=f'unlabeled ({no_type.sum()})',
                  edgecolors='none', rasterized=True)
    ax.set_xlabel('t-SNE 1'); ax.set_ylabel('t-SNE 2')
    ax.set_title(f'Glycan Type - {model_name}')
    ax.legend(frameon=False, markerscale=2, fontsize=8)
    ax = axes[1]
    tissue_counts = Counter([t for t in tissue_sub if t])
    top_tissues = [t for t, _ in tissue_counts.most_common(10)]
    for i, tissue in enumerate(top_tissues):
        mask = tissue_sub == tissue
        if mask.sum() > 0:
            ax.scatter(coords[mask, 0], coords[mask, 1],
                      c=TISSUE_COLORS[i % len(TISSUE_COLORS)], s=12, alpha=0.6,
                      label=f'{tissue} ({mask.sum()})', edgecolors='none', rasterized=True)
    other = np.array([t not in top_tissues for t in tissue_sub])
    if other.sum() > 0:
        ax.scatter(coords[other, 0], coords[other, 1],
                  c='#EEEEEE', s=3, alpha=0.15, edgecolors='none', rasterized=True)
    ax.set_xlabel('t-SNE 1'); ax.set_ylabel('t-SNE 2')
    ax.set_title(f'Tissue Source (top 10) - {model_name}')
    ax.legend(frameon=False, markerscale=2, fontsize=7, loc='center left',
              bbox_to_anchor=(1.0, 0.5))
    ax = axes[2]
    cancer_kw = ['cancer', 'carcinoma', 'tumor', 'glioblastoma', 'melanoma',
                 'leukemia', 'lymphoma']
    disease_cats = []
    for d in disease_sub:
        if not d:
            disease_cats.append('healthy')
        elif any(kw in d.lower() for kw in cancer_kw):
            disease_cats.append('cancer')
        else:
            disease_cats.append('other_disease')
    disease_cats = np.array(disease_cats)
    for cat, color, s, alpha in [('healthy', '#DDDDDD', 3, 0.15),
                                  ('other_disease', '#E69F00', 25, 0.8),
                                  ('cancer', '#D55E00', 30, 0.9)]:
        mask = disease_cats == cat
        if mask.sum() > 0:
            ax.scatter(coords[mask, 0], coords[mask, 1],
                      c=color, s=s, alpha=alpha,
                      label=f'{cat} ({mask.sum()})', edgecolors='none', rasterized=True)
    ax.set_xlabel('t-SNE 1'); ax.set_ylabel('t-SNE 2')
    ax.set_title(f'Disease Association - {model_name}')
    ax.legend(frameon=False, markerscale=2, fontsize=8)
    plt.suptitle(f'Probe 8: Glycowork Metadata Embedding Analysis - GlycanBERT {model_name}',
                 fontsize=13, fontweight='bold')
    plt.tight_layout()
    fig_path = os.path.join(output_dir, f'probe8_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'probe8_tsne_coords_{model_name.lower()}.npz'),
             coords=coords, gtype=gtype_sub, tissue=tissue_sub, disease=disease_sub)


def main():
    parser = argparse.ArgumentParser()
    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('--max_samples', type=int, default=None)
    args = parser.parse_args()
    model_name = {'v5': 'V5-A', 'v6': 'V6'}[args.model]
    if args.output_dir is None:
        args.output_dir = str(PROJECT_ROOT / 'bert_v6_contrastive' / 'analysis' /
                              f'probe8_disease_tissue_{args.model}')
    os.makedirs(args.output_dir, exist_ok=True)
    print(f"\nLoading probe data from {DATA_PATH}...")
    wurcs_list, iupac_list = [], []
    gtype_labels, tissue_labels, disease_labels = [], [], []
    with open(DATA_PATH) as f:
        for row in csv.DictReader(f):
            wurcs_list.append(row['wurcs'])
            iupac_list.append(row.get('iupac', ''))
            gtype_labels.append(row.get('glycan_type', ''))
            tissue_labels.append(row.get('tissue', '').split(';')[0] if row.get('tissue') else '')
            disease_labels.append(row.get('disease', ''))
    total = len(wurcs_list)
    if args.max_samples and args.max_samples < total:
        np.random.seed(42)
        idx = np.random.choice(total, args.max_samples, replace=False)
        wurcs_list = [wurcs_list[i] for i in idx]
        gtype_labels = [gtype_labels[i] for i in idx]
        tissue_labels = [tissue_labels[i] for i in idx]
        disease_labels = [disease_labels[i] for i in idx]
        total = args.max_samples
    print(f"  Total: {total:,}")
    print(f"  glycan_type: {sum(1 for g in gtype_labels if g):,}")
    print(f"  tissue: {sum(1 for t in tissue_labels if t):,}")
    print(f"  disease: {sum(1 for d in disease_labels if d):,}")
    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)
    print(f"\nEmbedding {total:,} glycans...")
    embs = batch_cls_embeddings(model, wurcs_list, tokenizer, device=args.device)
    import gc, torch
    del model
    torch.cuda.empty_cache()
    gc.collect()
    print(f"\nEmbeddings shape: {embs.shape}")
    all_results = {}
    gtype_mask = np.array([bool(g) for g in gtype_labels])
    if gtype_mask.sum() > 100:
        all_results['8a_glycan_type'] = probe_8a_glycan_type(
            embs[gtype_mask], np.array(gtype_labels)[gtype_mask],
            args.output_dir, model_name)
    tissue_mask = np.array([bool(t) for t in tissue_labels])
    if tissue_mask.sum() > 50:
        all_results['8b_tissue'] = probe_8b_tissue(
            embs[tissue_mask], np.array(tissue_labels)[tissue_mask],
            args.output_dir, model_name)
    all_results['8c_disease'] = probe_8c_disease(
        embs, disease_labels, wurcs_list, args.output_dir, model_name)
    probe_8d_visualization(embs, gtype_labels, tissue_labels, disease_labels,
                           args.output_dir, model_name)
    with open(os.path.join(args.output_dir, f'probe8_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 8 COMPLETE - {model_name}")
    print(f"Results: {args.output_dir}")
    print(f"{'='*60}")


if __name__ == '__main__':
    main()