#!/usr/bin/env python3
"""
Embedding Space Evaluation — Base vs All Fine-tuned ProtT5 XL Models.

Loads the base model + all fine-tuned variants (Here Just: ContraMLM) and produces exactly three publication-quality plots:

  Plot 1 — 2×2 t-SNE grid, one panel per model, coloured by PhrogCat.
  Plot 2 — 1×3 scatter of pairwise L2 distances (base x-axis, fine-tuned
            y-axis) for the same ~N_PAIRS protein pairs across all subplots.
  Plot 3 — Same as Plot 2 but using cosine similarity.
"""

import os
import sys
import json
import glob
import re
import argparse
from itertools import combinations
from typing import Any, Dict, List

import torch
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_score
from transformers import T5Tokenizer, T5ForConditionalGeneration
from peft import PeftModel

import warnings
warnings.filterwarnings("ignore")

# ============================================================================
# CONFIGURATION
# ============================================================================


VERSIONS: List[str] = ["ContraMLM_v1_1"]
ALL_MODEL_LABELS: List[str] = ["base"] + VERSIONS

DEFAULT_BASE_MODEL_NAME = "Rostlab/prot_t5_xl_uniref50"
DEFAULT_DATA_PATH       = "./data/envhog_phrog2/envhog_test_final_no_leakage.csv"
DEFAULT_MAX_LENGTH      = 512
DEFAULT_SAMPLE_SIZE     = 5000
DEFAULT_BATCH_SIZE      = 2
DEFAULT_RANDOM_STATE    = 42
DEFAULT_N_PAIRS         = 150   # protein pairs used in scatter plots 2 & 3


def parse_args():
    parser = argparse.ArgumentParser(
        description="Multi-model embedding space evaluation for ProtT5 XL"
    )
    parser.add_argument("--base-model",      type=str, default=DEFAULT_BASE_MODEL_NAME)
    parser.add_argument("--data-path",       type=str, default=DEFAULT_DATA_PATH)
    parser.add_argument("--output-dir",      type=str, default="./runs/evaluation_results_EmbeddingSpace")
    parser.add_argument("--max-length",      type=int, default=DEFAULT_MAX_LENGTH)
    parser.add_argument("--sample-size",     type=int, default=DEFAULT_SAMPLE_SIZE)
    parser.add_argument("--batch-size",      type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("--random-state",    type=int, default=DEFAULT_RANDOM_STATE)
    parser.add_argument(
        "--n-pairs",
        type=int,
        default=DEFAULT_N_PAIRS,
        help="Number of protein pairs for L2/cosine scatter plots (100–200 recommended)",
    )
    return parser.parse_args()


args = parse_args()

BASE_MODEL_NAME = args.base_model
DATA_PATH       = args.data_path
OUTPUT_DIR      = args.output_dir
IMAGES_DIR      = os.path.join(OUTPUT_DIR, "images")
TEXT_DIR        = os.path.join(OUTPUT_DIR, "text")
MAX_LENGTH      = args.max_length
SAMPLE_SIZE     = args.sample_size
BATCH_SIZE      = args.batch_size
RANDOM_STATE    = args.random_state
N_PAIRS_TARGET  = max(2, args.n_pairs)

os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(IMAGES_DIR, exist_ok=True)
os.makedirs(TEXT_DIR,   exist_ok=True)


# ============================================================================
# LOGGING
# ============================================================================

class Tee:
    def __init__(self, *streams):
        self.streams = streams
    def write(self, data):
        for s in self.streams:
            s.write(data)
            s.flush()
    def flush(self):
        for s in self.streams:
            s.flush()


run_log_path = os.path.join(TEXT_DIR, "run_log.txt")
log_file     = open(run_log_path, "w", encoding="utf-8")
sys.stdout   = Tee(sys.__stdout__, log_file)
sys.stderr   = Tee(sys.__stderr__, log_file)

print("=" * 80)
print("MULTI-MODEL EMBEDDING SPACE EVALUATION")
print("=" * 80)
print(f"\nBase model  : {BASE_MODEL_NAME}")
print(f"Versions    : {VERSIONS}")
print(f"Data        : {DATA_PATH}")
print(f"Sample size : {SAMPLE_SIZE}")
print(f"N pairs     : {N_PAIRS_TARGET}")
print(f"Output      : {OUTPUT_DIR}")

# ============================================================================
# DEVICE
# ============================================================================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"\nDevice: {device}")
if torch.cuda.is_available():
    print(f"GPU   : {torch.cuda.get_device_name(0)}")

# ============================================================================
# LOAD DATA
# ============================================================================

print("\n" + "=" * 80)
print("LOADING DATA")
print("=" * 80)

df = pd.read_csv(DATA_PATH)
print(f"Total proteins: {len(df)}")

if SAMPLE_SIZE > 0 and len(df) > SAMPLE_SIZE:
    df_sample = df.sample(n=SAMPLE_SIZE, random_state=RANDOM_STATE).reset_index(drop=True)
    print(f"Sampled {SAMPLE_SIZE} proteins")
else:
    df_sample = df.reset_index(drop=True)
    print(f"Using all {len(df)} proteins")

sequences   = df_sample["sequence"].tolist()
lengths     = df_sample["length"].tolist()
protein_ids = df_sample["id"].tolist()

# PhrogCat — used as colour label in t-SNE
if "PhrogCat" in df_sample.columns:
    phrog_cats = df_sample["PhrogCat"].fillna("unknown").tolist()
else:
    print("WARNING: 'PhrogCat' column not found — using 'unknown' for all proteins")
    phrog_cats = ["unknown"] * len(df_sample)


def prepare_t5_seq(seq: str) -> str:
    return " ".join(list(str(seq).replace(" ", "")))


sequences = [prepare_t5_seq(s) for s in sequences]

# ============================================================================
# TOKENIZER
# ============================================================================

print("\n" + "=" * 80)
print("LOADING TOKENIZER")
print("=" * 80)

tokenizer = T5Tokenizer.from_pretrained(
    BASE_MODEL_NAME, do_lower_case=False, legacy=True
)
print(f"Tokenizer loaded: {BASE_MODEL_NAME}")

# ============================================================================
# HELPERS
# ============================================================================

def get_encoder(model):
    """Return the encoder module regardless of wrapper type."""
    if hasattr(model, "encoder"):
        return model.encoder
    fn = getattr(model, "get_encoder", None)
    if callable(fn):
        return fn()
    for attr in ("base_model", "model"):
        inner = getattr(model, attr, None)
        if inner is not None:
            if hasattr(inner, "encoder"):
                return inner.encoder
            fn2 = getattr(inner, "get_encoder", None)
            if callable(fn2):
                return fn2()
    return model


def select_best_adapter_dir(version: str) -> str:
    """Return the first valid (NaN/Inf-free) adapter directory for *version*."""
    finetuned_path  = f"./runs/protrans_XL_Full_lora_envhog_{version}/lora_adapters"
    checkpoint_root = f"./runs/protrans_XL_Full_lora_envhog_{version}"

    candidates = []
    if os.path.isdir(finetuned_path):
        candidates.append(finetuned_path)

    ckpt_paths = sorted(
        glob.glob(os.path.join(checkpoint_root, "checkpoint-*")),
        key=lambda p: int(re.search(r"checkpoint-(\d+)", p).group(1))
        if re.search(r"checkpoint-(\d+)", p) else -1,
        reverse=True,
    )
    for cp in ckpt_paths:
        if os.path.isdir(cp):
            candidates.append(cp)

    def _resolve(candidate):
        for subdir in (candidate, os.path.join(candidate, "lora_adapters")):
            if (os.path.isfile(os.path.join(subdir, "adapter_model.safetensors"))
                    or os.path.isfile(os.path.join(subdir, "adapter_model.bin"))):
                return subdir
        return None

    for candidate in candidates:
        resolved = _resolve(candidate)
        if resolved is None:
            continue
        safe = os.path.join(resolved, "adapter_model.safetensors")
        bin_ = os.path.join(resolved, "adapter_model.bin")
        if os.path.isfile(safe):
            from safetensors.torch import load_file
            state_dict = load_file(safe, device="cpu")
        else:
            state_dict = torch.load(bin_, map_location="cpu")

        has_nan = any(torch.isnan(v).any().item() for v in state_dict.values())
        has_inf = any(torch.isinf(v).any().item() for v in state_dict.values())
        if not has_nan and not has_inf:
            print(f"  [{version}] Using adapter: {resolved}")
            return resolved

    raise RuntimeError(
        f"No valid (NaN/Inf-free) adapter found for version '{version}'. "
        f"Searched: {candidates}"
    )


def get_embeddings(model, seqs: List[str], batch_size: int = 16) -> np.ndarray:
    """Mean-pooled encoder embeddings for a list of pre-formatted sequences."""
    encoder = get_encoder(model)
    encoder.eval()
    all_embs = []

    with torch.no_grad():
        for i in range(0, len(seqs), batch_size):
            batch = seqs[i : i + batch_size]
            inputs = tokenizer(
                batch,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=MAX_LENGTH,
            )
            inputs = {k: v.to(device) for k, v in inputs.items()}
            out    = encoder(
                input_ids=inputs["input_ids"],
                attention_mask=inputs["attention_mask"],
            )
            hidden = out.last_hidden_state
            mask   = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
            pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
            all_embs.append(pooled.float().cpu().numpy())

            if (i // batch_size) % 10 == 0:
                print(f"  {i}/{len(seqs)} sequences processed...")

    return np.vstack(all_embs)


# ============================================================================
# LOAD BASE MODEL
# ============================================================================

print("\n" + "=" * 80)
print("LOADING BASE MODEL")
print("=" * 80)

model_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

base_model = T5ForConditionalGeneration.from_pretrained(
    BASE_MODEL_NAME, torch_dtype=model_dtype, low_cpu_mem_usage=True
)
base_model = base_model.to(device)
base_model.eval()
print("Base model loaded")

# ============================================================================
# GENERATE BASE EMBEDDINGS
# ============================================================================

print("\n" + "=" * 80)
print("GENERATING BASE EMBEDDINGS")
print("=" * 80)

base_embeddings = get_embeddings(base_model, sequences, BATCH_SIZE)
print(f"Base embeddings shape: {base_embeddings.shape}")

# Free base model GPU memory before loading fine-tuned models one-by-one
# (keep the numpy array — it is small)
del base_model
torch.cuda.empty_cache() if torch.cuda.is_available() else None

# ============================================================================
# LOAD FINE-TUNED MODELS AND GENERATE EMBEDDINGS
# ============================================================================

# Stores: { version_name: np.ndarray }
ft_embeddings: Dict[str, np.ndarray] = {}

for version in VERSIONS:
    print("\n" + "=" * 80)
    print(f"LOADING FINE-TUNED MODEL: {version}")
    print("=" * 80)

    adapter_dir = select_best_adapter_dir(version)

    ft_model = T5ForConditionalGeneration.from_pretrained(
        BASE_MODEL_NAME, torch_dtype=model_dtype, low_cpu_mem_usage=True
    )
    ft_model = PeftModel.from_pretrained(ft_model, adapter_dir)

    merge_fn = getattr(ft_model, "merge_and_unload", None)
    if callable(merge_fn):
        ft_model = merge_fn()
        print("  LoRA adapters merged")

    ft_model = ft_model.to(device)
    ft_model.eval()

    print(f"  Generating embeddings for {version}...")
    ft_embeddings[version] = get_embeddings(ft_model, sequences, BATCH_SIZE)
    print(f"  {version} embeddings shape: {ft_embeddings[version].shape}")

    del ft_model
    torch.cuda.empty_cache() if torch.cuda.is_available() else None

# ============================================================================
# SHARED t-SNE COLOUR MAP (PhrogCat)
# ============================================================================

print("\n" + "=" * 80)
print("PREPARING t-SNE COLOUR MAP")
print("=" * 80)

unique_cats   = sorted(set(phrog_cats))
n_cats        = len(unique_cats)
cmap_name     = "tab20" if n_cats > 10 else "tab10"
cmap          = plt.get_cmap(cmap_name, n_cats)
cat_to_idx    = {cat: i for i, cat in enumerate(unique_cats)}
colour_values = np.array([cat_to_idx[c] for c in phrog_cats])

print(f"Unique PhrogCat categories: {n_cats}")
print(f"Categories: {unique_cats}")

# ============================================================================
# t-SNE FOR ALL FOUR MODELS
# ============================================================================

print("\n" + "=" * 80)
print("RUNNING t-SNE (4 MODELS)")
print("=" * 80)

n_tsne     = min(2000, len(sequences))
rng_tsne   = np.random.default_rng(RANDOM_STATE)
tsne_idx   = rng_tsne.choice(len(sequences), n_tsne, replace=False)
perplexity = 30 if n_tsne > 30 else max(5, n_tsne - 1)

# Collect all embedding matrices for the 4 models
all_embeddings_ordered: Dict[str, np.ndarray] = {
    "base": base_embeddings,
    **ft_embeddings,
}

tsne_results: Dict[str, np.ndarray] = {}
tsne_silhouette: Dict[str, float] = {}

for label, emb in all_embeddings_ordered.items():
    print(f"  PCA → t-SNE for [{label}]...")
    pca   = PCA(n_components=50, random_state=RANDOM_STATE)
    emb_pca = pca.fit_transform(emb)
    tsne  = TSNE(n_components=2, random_state=RANDOM_STATE, perplexity=perplexity)
    tsne_results[label] = tsne.fit_transform(emb_pca[tsne_idx])

    # Silhouette index on the 2D t-SNE map using PhrogCat categories as labels.
    labels_tsne = np.array(phrog_cats, dtype=object)[tsne_idx]
    n_label_values = len(set(labels_tsne.tolist()))
    if 2 <= n_label_values < len(labels_tsne):
        try:
            tsne_silhouette[label] = float(silhouette_score(tsne_results[label], labels_tsne))
        except Exception:
            tsne_silhouette[label] = float("nan")
    else:
        tsne_silhouette[label] = float("nan")

    print(f"    Done.")
    if np.isnan(tsne_silhouette[label]):
        print("    Silhouette(PhrogCat): n/a")
    else:
        print(f"    Silhouette(PhrogCat): {tsne_silhouette[label]:.4f}")

colours_tsne = colour_values[tsne_idx]

# ============================================================================
# SHARED PROTEIN PAIRS FOR SCATTER PLOTS
# ============================================================================

print("\n" + "=" * 80)
print("BUILDING SHARED PROTEIN PAIRS")
print("=" * 80)

# We need N_PAIRS_TARGET pairs from a pool of proteins.
# Minimum proteins needed so combinations >= N_PAIRS_TARGET:
#   n*(n-1)/2 >= N_PAIRS_TARGET  →  n ≈ ceil((1 + sqrt(1+8k))/2)
import math
n_prot_needed = math.ceil((1 + math.sqrt(1 + 8 * N_PAIRS_TARGET)) / 2)
n_prot_needed = max(n_prot_needed, 2)
n_prot_needed = min(n_prot_needed, len(sequences))

rng_pairs     = np.random.default_rng(RANDOM_STATE + 1)
pair_indices  = rng_pairs.choice(len(sequences), n_prot_needed, replace=False)
pair_indices  = pair_indices.tolist()

all_pairs     = list(combinations(pair_indices, 2))

# Randomly subsample to exactly N_PAIRS_TARGET pairs if we have more
if len(all_pairs) > N_PAIRS_TARGET:
    rng_sub   = np.random.default_rng(RANDOM_STATE + 2)
    chosen    = rng_sub.choice(len(all_pairs), N_PAIRS_TARGET, replace=False)
    all_pairs = [all_pairs[i] for i in chosen]

n_pairs_actual = len(all_pairs)
print(f"Protein pool size : {n_prot_needed}")
print(f"Pairs generated   : {n_pairs_actual}")


def pairwise_l2(emb: np.ndarray, pairs: list) -> np.ndarray:
    return np.array([
        np.linalg.norm(emb[a] - emb[b])
        for a, b in pairs
    ])


def pairwise_cosine(emb: np.ndarray, pairs: list, eps: float = 1e-12) -> np.ndarray:
    sims = []
    for a, b in pairs:
        va, vb = emb[a], emb[b]
        denom  = np.linalg.norm(va) * np.linalg.norm(vb)
        sim    = np.dot(va, vb) / max(denom, eps)
        sims.append(float(np.clip(sim, -1.0, 1.0)))
    return np.array(sims)


# Compute for base
base_pair_l2  = pairwise_l2(base_embeddings,  all_pairs)
base_pair_cos = pairwise_cosine(base_embeddings, all_pairs)

# Compute for each fine-tuned version
ft_pair_l2:  Dict[str, np.ndarray] = {}
ft_pair_cos: Dict[str, np.ndarray] = {}

for version in VERSIONS:
    emb = ft_embeddings[version]
    ft_pair_l2[version]  = pairwise_l2(emb,  all_pairs)
    ft_pair_cos[version] = pairwise_cosine(emb, all_pairs)

# ============================================================================
# PLOT 1 — 2×2 t-SNE GRID  (coloured by PhrogCat)
# ============================================================================

print("\n" + "=" * 80)
print("PLOT 1: 2×2 t-SNE GRID")
print("=" * 80)

fig, axes = plt.subplots(2, 2, figsize=(16, 14))
axes_flat  = axes.flatten()

panel_order = ["base", "ContraMLM_v1_0", "Default_v2_1", "MLP_v0"]
panel_titles = {
    "base":           "Base Model",
    "ContraMLM_v1_0": "ContraMLM v1.0",
    "Default_v2_1":   "Default v2.1",
    "MLP_v0":         "MLP v0",
}

for ax, label in zip(axes_flat, panel_order):
    xy   = tsne_results[label]
    sc   = ax.scatter(
        xy[:, 0], xy[:, 1],
        c=colours_tsne,
        cmap=cmap_name,
        vmin=0, vmax=n_cats - 1,
        alpha=0.65,
        s=8,
        linewidths=0,
    )
    sil_txt = (
        f"Silhouette(PhrogCat): {tsne_silhouette[label]:.3f}"
        if not np.isnan(tsne_silhouette[label])
        else "Silhouette(PhrogCat): n/a"
    )
    ax.set_title(
        f"{panel_titles[label]}\n{sil_txt}",
        fontsize=14,
        fontweight="bold",
        pad=8,
    )
    ax.set_xlabel("t-SNE 1", fontsize=10)
    ax.set_ylabel("t-SNE 2", fontsize=10)
    ax.tick_params(labelsize=8)

# Shared legend for PhrogCat categories
legend_handles = [
    mpatches.Patch(color=cmap(cat_to_idx[cat] / max(n_cats - 1, 1)), label=cat)
    for cat in unique_cats
]
fig.legend(
    handles=legend_handles,
    title="PhrogCat",
    title_fontsize=14,
    fontsize=12,
    loc="lower center",
    ncol=min(n_cats, 6),
    bbox_to_anchor=(0.5, -0.02),
    frameon=True,
)

fig.suptitle(
    f"t-SNE Embedding Space — Base vs Fine-tuned Models\n"
    f"(n={n_tsne} proteins, coloured by PhrogCat)",
    fontsize=15,
    fontweight="bold",
    y=1.01,
)
plt.tight_layout()

plot1_path = os.path.join(IMAGES_DIR, "plot1_tsne_4models.png")
fig.savefig(plot1_path, dpi=300, bbox_inches="tight")
plt.close(fig)
print(f"Saved: {plot1_path}")

# ============================================================================
# PLOT 2 — PAIRWISE L2 SCATTER (base x-axis, fine-tuned y-axis)
# ============================================================================

print("\n" + "=" * 80)
print("PLOT 2: PAIRWISE L2 DISTANCE SCATTER")
print("=" * 80)

fig, axes = plt.subplots(1, 3, figsize=(18, 6))

version_titles = {
    "ContraMLM_v1_0": "ContraMLM v1.0",
    "Default_v2_1":   "Default v2.1",
    "MLP_v0":         "MLP v0",
}

for ax, version in zip(axes, VERSIONS):
    x = base_pair_l2
    y = ft_pair_l2[version]

    # Diagonal reference line
    lim_min = min(x.min(), y.min()) * 0.98
    lim_max = max(x.max(), y.max()) * 1.02
    ax.plot([lim_min, lim_max], [lim_min, lim_max],
            color="gray", linestyle="--", linewidth=1.0, alpha=0.7, label="y = x")

    ax.scatter(x, y, alpha=0.55, s=20, color="#2E86AB", linewidths=0)

    # Pearson r annotation
    r = float(np.corrcoef(x, y)[0, 1])
    ax.text(
        0.05, 0.93, f"r = {r:.3f}",
        transform=ax.transAxes,
        fontsize=10,
        verticalalignment="top",
        bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
    )

    ax.set_xlim(lim_min, lim_max)
    ax.set_ylim(lim_min, lim_max)
    ax.set_xlabel("Base model — L2 distance", fontsize=11)
    ax.set_ylabel(f"{version_titles[version]} — L2 distance", fontsize=11)
    ax.set_title(f"L2: Base vs {version_titles[version]}", fontsize=13, fontweight="bold")
    ax.set_aspect("equal", adjustable="box")
    ax.grid(True, alpha=0.25)
    ax.legend(fontsize=9)

fig.suptitle(
    f"Pairwise L2 Distance: Base vs Fine-tuned Models\n"
    f"({n_pairs_actual} protein pairs, same pairs across all subplots)",
    fontsize=14,
    fontweight="bold",
)
plt.tight_layout()

plot2_path = os.path.join(IMAGES_DIR, "plot2_pairwise_l2_scatter.png")
fig.savefig(plot2_path, dpi=300, bbox_inches="tight")
plt.close(fig)
print(f"Saved: {plot2_path}")

# ============================================================================
# PLOT 3 — PAIRWISE COSINE SIMILARITY SCATTER
# ============================================================================

print("\n" + "=" * 80)
print("PLOT 3: PAIRWISE COSINE SIMILARITY SCATTER")
print("=" * 80)

fig, axes = plt.subplots(1, 3, figsize=(18, 6))

for ax, version in zip(axes, VERSIONS):
    x = base_pair_cos
    y = ft_pair_cos[version]

    lim_min = min(x.min(), y.min()) - 0.02
    lim_max = max(x.max(), y.max()) + 0.02
    ax.plot([lim_min, lim_max], [lim_min, lim_max],
            color="gray", linestyle="--", linewidth=1.0, alpha=0.7, label="y = x")

    ax.scatter(x, y, alpha=0.55, s=20, color="#E84855", linewidths=0)

    r = float(np.corrcoef(x, y)[0, 1])
    ax.text(
        0.05, 0.93, f"r = {r:.3f}",
        transform=ax.transAxes,
        fontsize=10,
        verticalalignment="top",
        bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
    )

    ax.set_xlim(lim_min, lim_max)
    ax.set_ylim(lim_min, lim_max)
    ax.set_xlabel("Base model — cosine similarity", fontsize=11)
    ax.set_ylabel(f"{version_titles[version]} — cosine similarity", fontsize=11)
    ax.set_title(f"Cosine: Base vs {version_titles[version]}", fontsize=13, fontweight="bold")
    ax.set_aspect("equal", adjustable="box")
    ax.grid(True, alpha=0.25)
    ax.legend(fontsize=9)

fig.suptitle(
    f"Pairwise Cosine Similarity: Base vs Fine-tuned Models\n"
    f"({n_pairs_actual} protein pairs, same pairs across all subplots)",
    fontsize=14,
    fontweight="bold",
)
plt.tight_layout()

plot3_path = os.path.join(IMAGES_DIR, "plot3_pairwise_cosine_scatter.png")
fig.savefig(plot3_path, dpi=300, bbox_inches="tight")
plt.close(fig)
print(f"Saved: {plot3_path}")

# ============================================================================
# SAVE PAIR DATA AS CSV (reproducibility)
# ============================================================================

pair_records = []
for k, (a, b) in enumerate(all_pairs):
    row = {
        "pair_index":    k,
        "protein_id_a":  protein_ids[a],
        "protein_id_b":  protein_ids[b],
        "base_l2":       float(base_pair_l2[k]),
        "base_cosine":   float(base_pair_cos[k]),
    }
    for version in VERSIONS:
        row[f"{version}_l2"]     = float(ft_pair_l2[version][k])
        row[f"{version}_cosine"] = float(ft_pair_cos[version][k])
    pair_records.append(row)

pairs_csv_path = os.path.join(TEXT_DIR, "pairwise_distances_all_models.csv")
pd.DataFrame(pair_records).to_csv(pairs_csv_path, index=False)
print(f"\nPair data saved: {pairs_csv_path}")

# ============================================================================
# DONE
# ============================================================================

print("\n" + "=" * 80)
print("EVALUATION COMPLETE")
print("=" * 80)
print(f"\nOutputs written to : {OUTPUT_DIR}")
print(f"  Plot 1 (t-SNE)   : {plot1_path}")
print(f"  Plot 2 (L2)      : {plot2_path}")
print(f"  Plot 3 (cosine)  : {plot3_path}")
print(f"  Pair CSV         : {pairs_csv_path}")
print(f"  Run log          : {run_log_path}")