src/eval_PhrogRetrieval.py

#!/usr/bin/env python3

from __future__ import annotations

# ============================================================================
# USER CONFIGURATION
# ============================================================================

FT_SUFFIX   = "ContraMLM_v1_1"   # used for legacy single-FT outputs
FT40_SUFFIX = "base"
ENABLE_FT40 = False

# All four model labels and their pkl basenames
ALL_VERSIONS = ["base", "ContraMLM_v1_1"]

# The 17 functional groups to show in new_outputs_all clustermaps.
# These must match the top-level keys in functional_groups.json exactly.
CLUSTERMAP_GROUPS = [
    "lysis",             # was "lysis"             ✓
    "anti-restriction",  # was "anti_restriction"   fixed
    "super_infection",   # was "super infection"    fixed
    "toxin",             # was "toxin"              ✓
    "crispr",            # was "crispr"             ✓
    "sir2",              # was "sir2"               ✓
    "pvp",               # was "PVP"                fixed
    "packaging_assembly",# was "packaging and assembly" fixed
    "DNA-associated",    # was "DNA-associated"     ✓
    "RNA-associated",    # was "RNA-associated"     ✓
    "nucleotide_metabolism",  # was "nucleotide metabolism"  fixed
    "cell_wall_depolymerase", # was "cell wall depolymerase" fixed
    "transferase",       # was "transferase"        ✓
    "reductase",         # was "reductase"          ✓
    "adsorption-related",# was "adsorption-related" ✓
    "phosphorylation",   # was "phosphorylation"    ✓
    "ejection",          # was "internal/ejection"  fixed
]

# ============================================================================

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(line_buffering=True)
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(line_buffering=True)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

K_VALUES         = [5, 10, 50]
MIN_FAMILY_SIZE  = 5

HNSW_M               = 16
HNSW_EF_CONSTRUCTION = 512
HNSW_EF_SEARCH       = 512

# Colour palette for all models
MODEL_COLORS: Dict[str, str] = {
    "base":           "#9fc2e6",
    "ContraMLM_v1_1": "#2d80c4",
    # legacy keys kept for backward compat
    "FT40":           "#2d80c4",
    "FT500":          "#1a5276",
}
MODEL_LINESTYLES: Dict[str, str] = {
    "base":           "--",
    "ContraMLM_v1_1": "-",
    "FT40":           "-.",
    "FT500":          "-",
}

# Human-readable display names for plot labels and titles.
# Internal keys (used for file paths, dicts, CSVs) stay unchanged;
# only what the reader sees in figures is translated.
MODEL_DISPLAY_NAMES: Dict[str, str] = {
    "base":           "Base",
    "ContraMLM_v1_1": "ContraMLM",
    "FT40":           "FT40",
    "FT500":          "FT500",
}

def disp(model_key: str) -> str:
    """Return the display name for a model key, falling back to the key itself."""
    return MODEL_DISPLAY_NAMES.get(model_key, model_key)

FAMILY_SIZE_BINS = [
    (5,    20,  "rare\n(5–19)"),
    (20,  100,  "medium\n(20–99)"),
    (100, None, "common\n(≥100)"),
]


# ---------------------------------------------------------------------------
# Model configuration
# ---------------------------------------------------------------------------

def build_model_configs() -> Dict[str, str]:
    """Maps model label → embedding pkl basename."""
    configs: Dict[str, str] = {}
    for v in ALL_VERSIONS:
        configs[v] = f"test_embeddings_protrans_lora_{v}.pkl"
    return configs


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    # Dynamically locate the repository root (the parent of the 'src' folder)
    ROOT_DIR = Path(__file__).resolve().parent.parent

    parser = argparse.ArgumentParser(
        description="Zero-shot PHROG family retrieval benchmark — all models.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--protein-csv",
        default=str(ROOT_DIR / "data" / "envhog_test_final_no_leakage.csv"),
    )
    parser.add_argument(
        "--emb-dir",    
        default=str(ROOT_DIR / "data") 
    )
    parser.add_argument(
        "--output-dir", 
        default=str(ROOT_DIR / "runs" / "phrog_retrieval_results")
    )
    parser.add_argument("--min-family-size", default=MIN_FAMILY_SIZE, type=int)
    parser.add_argument("--hnsw-m",               default=HNSW_M,               type=int)
    parser.add_argument("--hnsw-ef-construction",  default=HNSW_EF_CONSTRUCTION, type=int)
    parser.add_argument("--hnsw-ef-search",        default=HNSW_EF_SEARCH,       type=int)
    parser.add_argument("--ft-suffix",  default=FT_SUFFIX, type=str)
    parser.add_argument(
        "--functional-groups",
        # Correctly maps to the 'data' folder
        default=str(ROOT_DIR / "data" / "functional_groups.json"),
        type=str,
    )
    return parser.parse_args()


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------

def load_embeddings(pkl_path: Path) -> pd.DataFrame:
    print(f"Loading embeddings: {pkl_path} …")
    t0   = time.time()
    embs = pd.read_pickle(pkl_path)
    embs.index = embs.index.astype(str)

    if embs.index.duplicated().any():
        n_dup = int(embs.index.duplicated().sum())
        print(f"  Warning: {n_dup} duplicate IDs — keeping first occurrence.")
        embs = embs.loc[~embs.index.duplicated(keep="first")]

    int_cols = sorted([c for c in embs.columns if isinstance(c, int)])
    embs = embs[int_cols]
    print(f"  Shape: {embs.shape}  ({time.time() - t0:.1f} s)")
    return embs


def load_phrog_labels(csv_path: Path, min_family_size: int) -> pd.Series:
    print(f"Loading PHROG labels: {csv_path} …")
    df = pd.read_csv(csv_path, usecols=["id", "bestPhrog"])
    df["id"] = df["id"].astype(str)
    df = df.set_index("id")
    df = df.rename(columns={"bestPhrog": "phrog"})

    if df.index.duplicated().any():
        n_dup = int(df.index.duplicated().sum())
        print(f"  Warning: {n_dup} duplicate protein IDs in metadata — keeping first.")
        df = df.loc[~df.index.duplicated(keep="first")]

    invalid = (
        df["phrog"].isna()
        | df["phrog"].astype(str).str.lower().isin(
            ["no_phrog", "nan", "", "unknown", "none", "na"]
        )
    )
    df = df.loc[~invalid].copy()
    df["phrog"] = df["phrog"].astype(str)

    family_sizes = df["phrog"].value_counts()
    valid_families = family_sizes[family_sizes >= min_family_size].index
    df = df.loc[df["phrog"].isin(valid_families)]

    n_fam = df["phrog"].nunique()
    fs    = family_sizes[valid_families]
    print(f"  Proteins with valid PHROG (family ≥ {min_family_size}): {len(df)}")
    print(f"  PHROG families : {n_fam}")
    print(
        f"  Family size — mean={fs.mean():.1f}  "
        f"median={fs.median():.1f}  max={fs.max()}"
    )
    return df["phrog"]


# ---------------------------------------------------------------------------
# Precision@k via HNSW
# ---------------------------------------------------------------------------

def compute_precision_at_k(
    X:               np.ndarray,
    labels:          np.ndarray,
    k_values:        List[int],
    M:               int = HNSW_M,
    ef_construction: int = HNSW_EF_CONSTRUCTION,
    ef_search:       int = HNSW_EF_SEARCH,
) -> Tuple[Dict[int, np.ndarray], np.ndarray]:
    try:
        import hnswlib
    except ImportError:
        raise ImportError("hnswlib is required.  Install with:  pip install hnswlib")

    n, d   = X.shape
    max_k  = max(k_values)
    labels = np.asarray(labels, dtype=str)

    norms  = np.linalg.norm(X, axis=1, keepdims=True)
    norms  = np.where(norms == 0, 1.0, norms)
    X_norm = (X / norms).astype(np.float32)

    print(f"  Building HNSW index (n={n}, d={d}, M={M}, ef_construction={ef_construction}) …")
    t0    = time.time()
    index = hnswlib.Index(space="ip", dim=d)
    index.init_index(max_elements=n, ef_construction=ef_construction, M=M, random_seed=42)
    index.add_items(X_norm, np.arange(n, dtype=np.int32))
    index.set_ef(max(ef_search, max_k * 2))
    print(f"  Index built in {time.time() - t0:.1f} s")

    print(f"  Querying {n} proteins (k={max_k}+1) …")
    t1 = time.time()
    indices, _ = index.knn_query(X_norm, k=max_k + 1)
    print(f"  Query done in {time.time() - t1:.1f} s")

    row_idx   = np.arange(n)[:, None]
    is_self   = (indices == row_idx)
    sort_ord  = np.argsort(is_self.astype(np.uint8), axis=1, kind="stable")
    neighbors = indices[row_idx, sort_ord][:, :max_k]

    prec_at_k: Dict[int, np.ndarray] = {}
    for k in k_values:
        k_idx    = neighbors[:, :k]
        k_labels = labels[k_idx]
        matches  = k_labels == labels[:, None]
        prec_at_k[k] = matches.mean(axis=1).astype(np.float32)

    return prec_at_k, neighbors


# ---------------------------------------------------------------------------
# Aggregation
# ---------------------------------------------------------------------------

def aggregate_results(
    prec_at_k:    Dict[int, np.ndarray],
    labels:       np.ndarray,
    family_sizes: pd.Series,
) -> Dict:
    labels = np.asarray(labels, dtype=str)
    overall: Dict[int, float] = {k: float(arr.mean()) for k, arr in prec_at_k.items()}

    per_family: Dict[int, pd.Series] = {}
    for k, arr in prec_at_k.items():
        df_pf = pd.DataFrame({"phrog": labels, "prec": arr})
        per_family[k] = df_pf.groupby("phrog")["prec"].mean()

    by_bin: Dict[str, Dict[int, float]] = {}
    for lo, hi, bin_label in FAMILY_SIZE_BINS:
        mask = (
            (family_sizes >= lo) if hi is None
            else (family_sizes >= lo) & (family_sizes < hi)
        )
        valid_fams = family_sizes.index[mask]
        prot_mask  = np.isin(labels, valid_fams)
        by_bin[bin_label] = {
            k: (float(arr[prot_mask].mean()) if prot_mask.any() else float("nan"))
            for k, arr in prec_at_k.items()
        }

    return {"overall": overall, "per_family": per_family, "by_bin": by_bin}


# ---------------------------------------------------------------------------
# Confusion analysis helpers (unchanged from original)
# ---------------------------------------------------------------------------

def build_confusion_table(
    family:    str,
    neighbors: np.ndarray,
    labels:    np.ndarray,
    df_ann:    pd.DataFrame,
    k:         int,
    topn:      int = 15,
) -> pd.DataFrame:
    fam_mask = labels == family
    fam_idx  = np.where(fam_mask)[0]
    if len(fam_idx) == 0:
        return pd.DataFrame(columns=["wrong_phrog", "annotation", "count", "pct"])

    k       = min(k, neighbors.shape[1])
    nb_labs = labels[neighbors[fam_idx, :k]]
    wrong   = nb_labs[nb_labs != family]

    if len(wrong) == 0:
        return pd.DataFrame(columns=["wrong_phrog", "annotation", "count", "pct"])

    counts      = pd.Series(wrong).value_counts()
    total_wrong = int(counts.sum())
    rows = []
    for phrog, cnt in counts.head(topn).items():
        ann = (
            str(df_ann.loc[phrog, "Annotation"])
            if phrog in df_ann.index else "unknown"
        )
        rows.append({
            "wrong_phrog": phrog,
            "annotation":  ann,
            "count":       int(cnt),
            "pct":         100.0 * cnt / total_wrong,
        })
    return pd.DataFrame(rows)


def write_confusion_analysis(
    bins:          Dict[str, pd.DataFrame],
    all_neighbors: Dict[str, np.ndarray],
    labels:        np.ndarray,
    df_ann:        pd.DataFrame,
    ft_label:      Optional[str],
    ref:           str,
    k_ref:         int,
    out_dir:       Path,
    confusion_k:   int = 50,
    topn:          int = 15,
) -> None:
    if ft_label is None or ft_label not in all_neighbors or ref not in all_neighbors:
        return

    def get_ann(phrog: str) -> str:
        return (
            str(df_ann.loc[phrog, "Annotation"]).strip()
            if phrog in df_ann.index else "unknown function"
        )

    lines: List[str] = [
        "",
        "=" * 70,
        "CONFUSION ANALYSIS — MOST DEGRADED KNOWN FAMILIES  (one per size bin)",
        f"  Selected : worst-delta family per bin whose annotation ≠ 'unknown function'",
        f"  Neighbors: top-{confusion_k} retrieved per protein",
        f"  Table    : top-{topn} wrong PHROG families ranked by retrieval count",
        "=" * 70,
    ]

    ft_col  = f"{ft_label}_P@{k_ref}"
    ref_col = f"{ref}_P@{k_ref}"

    for bin_name, df_bin in bins.items():
        if df_bin.empty or ft_col not in df_bin.columns or ref_col not in df_bin.columns:
            lines.append(f"\n  --- {bin_name.upper()} ---  (insufficient data)")
            continue

        df_bin           = df_bin.copy()
        df_bin["delta"]  = df_bin[ft_col] - df_bin[ref_col]
        df_bin["_ann"]   = df_bin["phrog"].apply(get_ann)
        df_known = df_bin[~df_bin["_ann"].str.lower().str.contains("unknown", na=True)]

        if df_known.empty:
            lines.append(f"\n  --- {bin_name.upper()} ---  (no known-annotation families with negative delta)")
            continue

        worst       = df_known.nsmallest(1, "delta").iloc[0]
        family      = str(worst["phrog"])
        family_ann  = get_ann(family)
        family_size = int(worst["n_proteins"])
        base_p5     = float(worst[ref_col])
        ft_p5       = float(worst[ft_col])
        delta       = float(worst["delta"])

        lines += [
            "",
            f"  --- {bin_name.upper()} ---",
            f"  Family     : {family}  (size={family_size})",
            f"  Annotation : {family_ann}",
            f"  {ref} P@{k_ref}={base_p5:.4f}  |  {ft_label} P@{k_ref}={ft_p5:.4f}  |  Delta={delta:+.4f}",
        ]

        k_use     = min(confusion_k, all_neighbors[ref].shape[1])
        fam_mask  = labels == family
        fam_idx   = np.where(fam_mask)[0]
        total_slots = len(fam_idx) * k_use

        for model_lbl in [ft_label, ref]:
            nb  = all_neighbors[model_lbl]
            tbl = build_confusion_table(family, nb, labels, df_ann, k=k_use, topn=topn)

            n_correct = int(np.sum(labels[nb[fam_idx, :k_use]] == family))
            n_wrong   = total_slots - n_correct

            lines += [
                "",
                f"  Wrong neighbours by {model_lbl}  "
                f"(top-{k_use} per protein; {n_wrong}/{total_slots} slots are wrong):",
                f"  {'Wrong PHROG':<14}  {'Annotation':<40}  {'Count':>6}  {'% wrong':>8}",
                "  " + "-" * 74,
            ]
            if tbl.empty:
                lines.append("  (all neighbours correct — no confusion)")
            else:
                for _, row in tbl.iterrows():
                    ann_trunc = str(row["annotation"])[:38]
                    lines.append(
                        f"  {row['wrong_phrog']:<14}  {ann_trunc:<40}  "
                        f"{row['count']:>6}  {row['pct']:>7.1f}%"
                    )
        lines.append("")

    with open(out_dir / "summary.txt", "a", encoding="utf-8") as fh:
        fh.write("\n".join(lines) + "\n")
    print("  Saved: summary.txt (confusion analysis)")


# ---------------------------------------------------------------------------
# Original plotting helpers (unchanged logic, extended to N models)
# ---------------------------------------------------------------------------

def _bar_chart(
    ax:            "plt.Axes",
    x:             np.ndarray,
    model_labels:  List[str],
    get_val:       "callable",
    label_fmt:     str = "{:.4f}",
) -> None:
    n_mod   = len(model_labels)
    width   = 0.7 / n_mod
    offsets = np.linspace(
        -(0.7 / 2) + width / 2,
         (0.7 / 2) - width / 2,
        n_mod,
    )
    for ml, offset in zip(model_labels, offsets):
        color = MODEL_COLORS.get(ml, "#888888")
        vals  = [get_val(ml, i) for i in range(len(x))]
        bars  = ax.bar(x + offset, vals, width, label=disp(ml), color=color)
        for bar, val in zip(bars, vals):
            if not np.isnan(val):
                ax.text(
                    bar.get_x() + bar.get_width() / 2,
                    bar.get_height() + 0.005,
                    label_fmt.format(val),
                    ha="center", va="bottom", fontsize=7,
                )


def plot_precision_at_k_bar(
    all_results:  Dict[str, Dict],
    out_dir:      Path,
    model_labels: List[str],
) -> None:
    k_vals = sorted(K_VALUES)
    x      = np.arange(len(k_vals))

    fig, ax = plt.subplots(figsize=(10, 5))
    _bar_chart(
        ax, x, model_labels,
        get_val=lambda ml, i: all_results[ml]["overall"][k_vals[i]],
    )
    ax.set_xticks(x)
    ax.set_xticklabels([f"Precision@{k}" for k in k_vals])
    ax.set_ylabel("Mean Precision@k")
    ax.set_ylim(0.0, 1.05)
    ax.set_title("PHROG Family Retrieval — Overall Precision@k (All Models)")
    ax.legend()
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_dir / "precision_at_k_overall.png", dpi=150)
    plt.close(fig)
    print("  Saved: precision_at_k_overall.png")


def plot_precision_by_family_size(
    all_results:  Dict[str, Dict],
    out_dir:      Path,
    model_labels: List[str],
    k:            int,
) -> None:
    bins = [b[2] for b in FAMILY_SIZE_BINS]
    x    = np.arange(len(bins))

    fig, ax = plt.subplots(figsize=(10, 5))
    _bar_chart(
        ax, x, model_labels,
        get_val=lambda ml, i: all_results[ml]["by_bin"].get(bins[i], {}).get(k, float("nan")),
    )
    ax.set_xticks(x)
    ax.set_xticklabels(bins)
    ax.set_ylabel(f"Mean Precision@{k}")
    ax.set_ylim(0.0, 1.05)
    ax.set_title(f"PHROG Retrieval — Precision@{k} by Family Size (All Models)")
    ax.legend()
    ax.grid(axis="y", alpha=0.3)
    fig.tight_layout()
    fname = f"precision_at_{k}_by_family_size.png"
    fig.savefig(out_dir / fname, dpi=150)
    plt.close(fig)
    print(f"  Saved: {fname}")


def plot_per_family_scatter(
    results_base: Dict,
    results_ft:   Dict,
    label_ft:     str,
    out_dir:      Path,
    k:            int,
) -> None:
    pf_base = results_base["per_family"][k]
    pf_ft   = results_ft["per_family"][k]
    common  = pf_base.index.intersection(pf_ft.index)

    vals_base = pf_base.loc[common].values
    vals_ft   = pf_ft.loc[common].values

    n_ft_better   = int((vals_ft > vals_base).sum())
    n_base_better = int((vals_base > vals_ft).sum())
    n_equal       = len(common) - n_ft_better - n_base_better

    fig, ax = plt.subplots(figsize=(6, 6))
    ax.scatter(vals_base, vals_ft, s=10, alpha=0.4,
               color=MODEL_COLORS.get(label_ft, "#1a5276"))
    ax.plot([0, 1], [0, 1], "k--", lw=0.8, label="no change")
    ax.set_xlabel(f"Precision@{k} — Base")
    ax.set_ylabel(f"Precision@{k} — {disp(label_ft)}")
    ax.set_xlim(-0.02, 1.05)
    ax.set_ylim(-0.02, 1.05)
    ax.set_title(
        f"Per-family Precision@{k}  ({disp(label_ft)} vs Base)\n"
        f"{disp(label_ft)} better: {n_ft_better}  |  "
        f"Base better: {n_base_better}  |  "
        f"equal: {n_equal}"
    )
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)
    fig.tight_layout()
    fname = f"scatter_per_family_at_{k}_{label_ft}_vs_base.png"
    fig.savefig(out_dir / fname, dpi=150)
    plt.close(fig)
    print(f"  Saved: {fname}")


def plot_delta_histogram(
    prec_base: np.ndarray,
    prec_ft:   np.ndarray,
    label_ft:  str,
    out_dir:   Path,
    k:         int,
) -> None:
    delta      = prec_ft - prec_base
    n_improved = int((delta > 0).sum())
    n_worse    = int((delta < 0).sum())
    n_same     = int((delta == 0).sum())
    mean_delta = float(delta.mean())

    fig, ax = plt.subplots(figsize=(7, 4))
    ax.hist(delta, bins=60, color=MODEL_COLORS.get(label_ft, "#1a5276"),
            alpha=0.75, edgecolor="white", linewidth=0.3)
    ax.axvline(0,          color="red",    lw=1.2, linestyle="--", label="no change")
    ax.axvline(mean_delta, color="orange", lw=1.2, linestyle="-",
               label=f"mean Δ = {mean_delta:+.5f}")
    ax.set_xlabel(f"Δ Precision@{k}  ({disp(label_ft)} − Base)")
    ax.set_ylabel("Number of proteins")
    ax.set_title(
        f"Per-protein Δ Precision@{k}  ({disp(label_ft)} vs Base)\n"
        f"improved: {n_improved}  |  worse: {n_worse}  |  unchanged: {n_same}"
    )
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)
    fig.tight_layout()
    fname = f"delta_histogram_at_{k}_{label_ft}.png"
    fig.savefig(out_dir / fname, dpi=150)
    plt.close(fig)
    print(f"  Saved: {fname}")


# ---------------------------------------------------------------------------
# Text outputs (unchanged)
# ---------------------------------------------------------------------------

def write_summary(
    all_results:  Dict[str, Dict],
    model_labels: List[str],
    out_dir:      Path,
) -> None:
    k_vals = sorted(K_VALUES)
    lines = [
        "=" * 70,
        "PHROG FAMILY RETRIEVAL SUMMARY",
        f"ALL_VERSIONS = {ALL_VERSIONS}",
        f"MIN_FAMILY_SIZE = {MIN_FAMILY_SIZE}",
        "=" * 70,
        "",
        "Overall Precision@k",
        "-" * 50,
        "  {:<22}  {}".format(
            "Model", "  ".join(f"P@{k:>2}" for k in k_vals)
        ),
    ]
    for ml in model_labels:
        row = "  ".join(f"{all_results[ml]['overall'][k]:.5f}" for k in k_vals)
        lines.append(f"  {ml:<22}  {row}")

    lines.append("")
    for _, _, bin_label in FAMILY_SIZE_BINS:
        clean = bin_label.replace("\n", " ")
        lines.append(f"Family size: {clean}")
        lines.append("-" * 50)
        lines.append(
            "  {:<22}  {}".format(
                "Model", "  ".join(f"P@{k:>2}" for k in k_vals)
            )
        )
        for ml in model_labels:
            row = "  ".join(
                f"{all_results[ml]['by_bin'].get(bin_label, {}).get(k, float('nan')):.5f}"
                for k in k_vals
            )
            lines.append(f"  {ml:<22}  {row}")
        lines.append("")

    with open(out_dir / "summary.txt", "w") as fh:
        fh.write("\n".join(lines))
    print("  Saved: summary.txt")


def write_per_family_csv(
    all_results:  Dict[str, Dict],
    model_labels: List[str],
    out_dir:      Path,
) -> None:
    k_vals   = sorted(K_VALUES)
    ref      = model_labels[0]
    all_fams = all_results[ref]["per_family"][k_vals[0]].index

    rows = []
    for fam in all_fams:
        row: Dict = {"phrog": fam}
        for ml in model_labels:
            for k in k_vals:
                v = all_results[ml]["per_family"][k].get(fam, float("nan"))
                row[f"{ml}_P@{k}"] = round(float(v), 6)
        rows.append(row)

    pd.DataFrame(rows).to_csv(out_dir / "per_family_precision.csv", index=False)
    print("  Saved: per_family_precision.csv")


def write_summary_appendix(
    all_results:   Dict[str, Dict],
    model_labels:  List[str],
    family_sizes:  pd.Series,
    protein_csv:   Path,
    out_dir:       Path,
    all_neighbors: Optional[Dict[str, np.ndarray]] = None,
    labels:        Optional[np.ndarray] = None,
) -> None:
    k_ref = sorted(K_VALUES)[0]

    df_ann = pd.read_csv(protein_csv, usecols=["bestPhrog", "bestPhrogAn", "PhrogCat"])
    df_ann = df_ann.rename(columns={
        "bestPhrog": "phrog", "bestPhrogAn": "Annotation", "PhrogCat": "Category"
    })
    df_ann = df_ann.drop_duplicates("phrog").set_index("phrog")

    ref      = model_labels[0]
    all_fams = all_results[ref]["per_family"][k_ref].index
    pf_rows  = []
    for fam in all_fams:
        row: Dict = {"phrog": fam}
        for ml in model_labels:
            row[f"{ml}_P@{k_ref}"] = float(
                all_results[ml]["per_family"][k_ref].get(fam, float("nan"))
            )
        pf_rows.append(row)
    pf = pd.DataFrame(pf_rows)

    sizes_df = family_sizes.reset_index()
    sizes_df.columns = ["phrog", "n_proteins"]
    pf = pf.merge(sizes_df, on="phrog", how="left")

    bins = {
        "rare   (5-19)":  pf[pf["n_proteins"].between(5, 19)],
        "medium (20-99)": pf[pf["n_proteins"].between(20, 99)],
        "common (>=100)": pf[pf["n_proteins"] >= 100],
    }

    all_sizes      = family_sizes[family_sizes >= MIN_FAMILY_SIZE]
    bin_labels_order = ["rare   (5-19)", "medium (20-99)", "common (>=100)"]
    bin_edges = [5, 20, 100, int(all_sizes.max()) + 1]
    binned    = pd.cut(all_sizes, bins=bin_edges, labels=bin_labels_order, right=False)
    bin_counts = binned.value_counts().sort_index()
    total_fam  = len(all_sizes)

    # Use first non-base model as ft_label for confusion analysis
    ft_label = next((ml for ml in model_labels if ml != ref), None)

    lines: List[str] = [""]

    lines += [
        "=" * 70,
        "DATA PROVENANCE",
        "=" * 70,
        f"  Source file : protein CSV",
        f"  PHROG label : column 'bestPhrog'",
        f"  Filter      : families with < {MIN_FAMILY_SIZE} proteins excluded",
        "",
    ]

    lines += [
        "=" * 70,
        "FAMILY SIZE DISTRIBUTION",
        "=" * 70,
        f"  Total unique PHROG families in dataset  : {family_sizes.shape[0]:>8,}",
        f"  Families passing filter (>= {MIN_FAMILY_SIZE} proteins): {len(all_sizes):>8,}",
        f"  Total proteins in filtered set          : {int(all_sizes.sum()):>8,}",
        "",
        f"  {'Bin':<17} | {'# Families':>10} | {'% of total':>10}",
        f"  {'-'*17}-+-{'-'*10}-+-{'-'*10}",
    ]
    for lbl in bin_labels_order:
        cnt = bin_counts[lbl]
        lines.append(f"  {lbl:<17} | {cnt:>10,} | {100*cnt/total_fam:>9.1f}%")
    lines.append("")

    col_w  = max(22, max(len(ml) for ml in model_labels) + 2)
    ref_col = f"{ref}_P@{k_ref}"
    hdr_cols = f"  {'PHROG':<12}  {'Size':>5}  {ref+' P@'+str(k_ref):>{col_w}}"
    for ml in model_labels:
        if ml != ref:
            hdr_cols += f"  {ml+' P@'+str(k_ref):>{col_w}}  {'Delta':>7}"
    hdr_cols += "  Annotation"
    sep = "  " + "-" * 120

    def example_rows(df_bin: pd.DataFrame, topn: int = 5) -> List[str]:
        rows_out: List[str] = []
        for _, r in df_bin.nlargest(topn, "n_proteins").iterrows():
            ann  = df_ann.loc[r["phrog"], "Annotation"] if r["phrog"] in df_ann.index else "unknown"
            line = f"  {r['phrog']:<12}  {int(r['n_proteins']):>5}  {r[ref_col]:>{col_w}.4f}"
            for ml in model_labels:
                if ml != ref:
                    ml_col = f"{ml}_P@{k_ref}"
                    if ml_col in r:
                        delta = r[ml_col] - r[ref_col]
                        line += f"  {r[ml_col]:>{col_w}.4f}  {delta:>+7.4f}"
            line += f"  {ann}"
            rows_out.append(line)
        return rows_out

    lines += [
        "=" * 70,
        f"EXAMPLE FAMILIES PER SIZE BIN  (top 5 largest; P@{k_ref})",
        "=" * 70,
    ]
    for bin_name, df_bin in bins.items():
        lines += [f"  --- {bin_name.upper()} ---", hdr_cols, sep]
        lines += example_rows(df_bin)
        lines.append("")

    with open(out_dir / "summary.txt", "a", encoding="utf-8") as fh:
        fh.write("\n".join(lines) + "\n")
    print("  Saved: summary.txt (appendix)")

    if all_neighbors is not None and labels is not None and ft_label is not None:
        write_confusion_analysis(
            bins=bins,
            all_neighbors=all_neighbors,
            labels=labels,
            df_ann=df_ann,
            ft_label=ft_label,
            ref=ref,
            k_ref=k_ref,
            out_dir=out_dir,
        )


# ---------------------------------------------------------------------------
# Shared helpers for new_outputs and new_outputs_all
# ---------------------------------------------------------------------------

def load_full_metadata(csv_path: Path) -> pd.DataFrame:
    df = pd.read_csv(csv_path, usecols=["id", "bestPhrogAn", "PhrogCat"])
    df["id"] = df["id"].astype(str)
    if df["id"].duplicated().any():
        df = df.drop_duplicates(subset="id", keep="first")
    df = df.set_index("id")
    df["bestPhrogAn"] = df["bestPhrogAn"].fillna("unknown").astype(str)
    df["PhrogCat"]    = df["PhrogCat"].fillna("unknown").astype(str)
    return df


def load_functional_groups(fg_path: Path) -> Tuple[Dict[str, set], Dict[str, set]]:
    with fg_path.open("r", encoding="utf-8") as fh:
        raw = json.load(fh)

    exclude_neg_map: Dict[str, set] = {}
    group_members:   Dict[str, set] = {}
    for key, members in raw.items():
        if key.startswith("exclude_neg_"):
            exclude_neg_map[key[len("exclude_neg_"):]] = set(members)
        else:
            group_members[key] = set(members)

    print(f"  Functional groups: {len(group_members)} groups, "
          f"{len(exclude_neg_map)} exclude_neg entries")
    return group_members, exclude_neg_map


def compute_precision_from_neighbors(
    neighbors: np.ndarray,
    labels:    np.ndarray,
    k_values:  List[int],
) -> Dict[int, np.ndarray]:
    prec: Dict[int, np.ndarray] = {}
    for k in k_values:
        k_nb    = neighbors[:, :k]
        k_labs  = labels[k_nb]
        matches = k_labs == labels[:, None]
        prec[k] = matches.mean(axis=1).astype(np.float32)
    return prec


def _macro_avg_by_annotation(labels: np.ndarray, prec_arr: np.ndarray) -> float:
    valid = ~np.isnan(prec_arr.astype(np.float64))
    if not valid.any():
        return float("nan")
    df = pd.DataFrame({"ann": labels[valid], "prec": prec_arr[valid]})
    return float(df.groupby("ann")["prec"].mean().mean())


def compute_group_precision_with_excl(
    neighbors:       np.ndarray,
    ann_labels:      np.ndarray,
    group_members:   Dict[str, set],
    exclude_neg_map: Dict[str, set],
    k_values:        List[int],
) -> Dict[str, Dict[int, np.ndarray]]:
    max_k  = max(k_values)
    nb_max = neighbors[:, :max_k]
    nb_ann = ann_labels[nb_max]

    result: Dict[str, Dict[int, np.ndarray]] = {}
    for group, members in group_members.items():
        excl         = exclude_neg_map.get(group, set())
        members_list = list(members)

        group_mask = np.isin(ann_labels, members_list)
        group_idx  = np.where(group_mask)[0]
        if len(group_idx) == 0:
            result[group] = {k: np.array([], dtype=np.float32) for k in k_values}
            continue

        g_nb_ann = nb_ann[group_idx, :]
        is_excl  = (
            np.isin(g_nb_ann, list(excl))
            if excl else np.zeros(g_nb_ann.shape, dtype=bool)
        )

        result[group] = {}
        for k in k_values:
            g_nb_k = g_nb_ann[:, :k]
            excl_k = is_excl[:, :k]
            is_tp  = np.isin(g_nb_k, members_list) & ~excl_k
            eff_k  = k - excl_k.sum(axis=1)
            tps    = is_tp.sum(axis=1)
            with np.errstate(invalid="ignore", divide="ignore"):
                prec = np.where(eff_k > 0, tps / eff_k, np.nan)
            result[group][k] = prec.astype(np.float32)

    return result


def macro_avg_group_precision(
    group_prec:    Dict[str, Dict[int, np.ndarray]],
    ann_labels:    np.ndarray,
    group_members: Dict[str, set],
    k_values:      List[int],
) -> Dict[str, Dict[int, float]]:
    result: Dict[str, Dict[int, float]] = {}
    for group, members in group_members.items():
        group_mask = np.isin(ann_labels, list(members))
        group_ann  = ann_labels[group_mask]
        result[group] = {}
        for k in k_values:
            prec_arr = group_prec[group][k]
            if len(prec_arr) == 0:
                result[group][k] = float("nan")
            else:
                result[group][k] = _macro_avg_by_annotation(group_ann, prec_arr)
    return result


def compute_functional_confusion_matrix(
    neighbors:       np.ndarray,
    ann_labels:      np.ndarray,
    group_members:   Dict[str, set],
    exclude_neg_map: Dict[str, set],
    k:               int,
) -> Tuple[np.ndarray, List[str]]:
    group_names = list(group_members.keys())
    G           = len(group_names)
    matrix      = np.full((G, G), np.nan, dtype=np.float64)

    nb_k   = neighbors[:, :k]
    nb_ann = ann_labels[nb_k]

    for i, g_row in enumerate(group_names):
        members_row = group_members[g_row]
        excl_row    = exclude_neg_map.get(g_row, set())

        group_mask = np.isin(ann_labels, list(members_row))
        group_idx  = np.where(group_mask)[0]
        if len(group_idx) == 0:
            matrix[i, :] = np.nan
            continue

        g_nb_ann = nb_ann[group_idx, :]
        is_excl  = (
            np.isin(g_nb_ann, list(excl_row))
            if excl_row else np.zeros(g_nb_ann.shape, dtype=bool)
        )
        eff_k = k - is_excl.sum(axis=1)

        for j, g_col in enumerate(group_names):
            in_col     = np.isin(g_nb_ann, list(group_members[g_col])) & ~is_excl
            col_counts = in_col.sum(axis=1)
            with np.errstate(invalid="ignore", divide="ignore"):
                fracs = np.where(eff_k > 0, col_counts / eff_k, np.nan)
            matrix[i, j] = float(np.nanmean(fracs))

    return matrix, group_names


# ---------------------------------------------------------------------------
# Original new_outputs (heatmaps, 2-model bar charts) — unchanged
# ---------------------------------------------------------------------------

def plot_functional_confusion_heatmap(
    matrix:      np.ndarray,
    group_names: List[str],
    k:           int,
    model_label: str,
    out_dir:     Path,
) -> None:
    G   = len(group_names)
    fig, ax = plt.subplots(figsize=(max(10, G * 0.65), max(8, G * 0.6)))
    disp = np.nan_to_num(matrix, nan=0.0)
    vmax = max(float(disp.max()), 0.01)
    im   = ax.imshow(disp, aspect="auto", cmap="Blues", vmin=0, vmax=vmax)
    plt.colorbar(im, ax=ax, shrink=0.8, label="Mean fraction of neighbors")
    ax.set_xticks(range(G))
    ax.set_yticks(range(G))
    ax.set_xticklabels(group_names, rotation=45, ha="right", fontsize=7)
    ax.set_yticklabels(group_names, fontsize=7)
    ax.set_xlabel("Neighbor group")
    ax.set_ylabel("Query group")
    ax.set_title(f"Functional Group Confusion — {model_label}  P@{k}")
    for ii in range(G):
        for jj in range(G):
            v = disp[ii, jj]
            if v > 0.001:
                ax.text(jj, ii, f"{v:.2f}", ha="center", va="center",
                        fontsize=5, color="white" if v > 0.5 * vmax else "black")
    fig.tight_layout()
    fname = out_dir / f"confusion_k{k}_{model_label}.png"
    fig.savefig(fname, dpi=150, bbox_inches="tight")
    plt.close(fig)
    pd.DataFrame(disp, index=group_names, columns=group_names).to_csv(
        out_dir / f"confusion_k{k}_{model_label}.csv"
    )
    print(f"  Saved: {fname.name}")


def plot_group_bar_chart_2models(
    macro_base: Dict[str, Dict[int, float]],
    macro_ft:   Dict[str, Dict[int, float]],
    ft_label:   str,
    k:          int,
    out_dir:    Path,
) -> None:
    groups     = list(macro_ft.keys())
    ft_scores  = np.array([macro_ft[g].get(k, float("nan")) for g in groups])
    base_scores = np.array([macro_base[g].get(k, float("nan")) for g in groups])

    valid = ~np.isnan(ft_scores)
    order = np.argsort(ft_scores[valid])[::-1]
    v_groups = np.array(groups)[valid][order][:10]
    v_ft     = ft_scores[valid][order][:10]
    v_base_  = base_scores[valid][order][:10]

    x     = np.arange(len(v_groups))
    width = 0.35
    fig, ax = plt.subplots(figsize=(max(10, len(v_groups) * 1.2), 5))
    bars_b = ax.bar(x - width / 2, v_base_, width, label="base",
                    color=MODEL_COLORS.get("base", "#9fc2e6"))
    bars_f = ax.bar(x + width / 2, v_ft,    width, label=ft_label,
                    color=MODEL_COLORS.get(ft_label, "#1a5276"))
    ax.set_xticks(x)
    ax.set_xticklabels(v_groups, rotation=35, ha="right", fontsize=9)
    ax.set_ylabel(f"P@{k} (macro-avg by annotation)")
    ax.set_title(f"Top-10 Functional Groups — P@{k}")
    ax.legend()
    ax.set_ylim(0, 1.05)
    ax.yaxis.grid(True, linestyle="--", alpha=0.7)
    ax.set_axisbelow(True)
    for bar in (*bars_b, *bars_f):
        h = bar.get_height()
        if not np.isnan(h):
            ax.text(bar.get_x() + bar.get_width() / 2, h + 0.01,
                    f"{h:.3f}", ha="center", va="bottom", fontsize=6.5)
    fig.tight_layout()
    fname = out_dir / f"group_bar_k{k}.png"
    fig.savefig(fname, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {fname.name}")


def write_new_summary(
    prec_ann_base: Dict[int, np.ndarray],
    prec_ann_ft:   Dict[int, np.ndarray],
    prec_cat_base: Dict[int, np.ndarray],
    prec_cat_ft:   Dict[int, np.ndarray],
    ann_labels:    np.ndarray,
    cat_labels:    np.ndarray,
    ft_label:      str,
    k_values:      List[int],
    out_dir:       Path,
) -> None:
    k_vals = sorted(k_values)
    col_w  = max(22, len(ft_label) + 2)
    hdr    = "  ".join(f"P@{k:>2}" for k in k_vals)

    lines = [
        "=" * 70,
        "ANNOTATION & CATEGORY RETRIEVAL SUMMARY",
        f"FT model : {ft_label}",
        "=" * 70,
    ]

    for label_name, labs, prec_b, prec_f in [
        ("bestPhrogAn", ann_labels, prec_ann_base, prec_ann_ft),
        ("PhrogCat",    cat_labels, prec_cat_base, prec_cat_ft),
    ]:
        pm_b  = [f"{float(np.nanmean(prec_b[k])):.5f}" for k in k_vals]
        pm_f  = [f"{float(np.nanmean(prec_f[k])):.5f}" for k in k_vals]
        dpm   = [
            f"{float(np.nanmean(prec_f[k])) - float(np.nanmean(prec_b[k])):+.5f}"
            for k in k_vals
        ]
        mac_b = [f"{_macro_avg_by_annotation(labs, prec_b[k]):.5f}" for k in k_vals]
        mac_f = [f"{_macro_avg_by_annotation(labs, prec_f[k]):.5f}" for k in k_vals]
        dmac  = [
            f"{_macro_avg_by_annotation(labs, prec_f[k]) - _macro_avg_by_annotation(labs, prec_b[k]):+.5f}"
            for k in k_vals
        ]

        lines += [
            "",
            f"Label set: {label_name}",
            "-" * 50,
            f"  Per-protein mean:",
            f"  {'Model':<{col_w}}  {hdr}",
            f"  {'base':<{col_w}}  {'  '.join(pm_b)}",
            f"  {ft_label:<{col_w}}  {'  '.join(pm_f)}",
            f"  Deltas vs base:",
            f"  {ft_label:<{col_w}}  {'  '.join(dpm)}",
            "",
            f"  Macro-average by annotation:",
            f"  {'base':<{col_w}}  {'  '.join(mac_b)}",
            f"  {ft_label:<{col_w}}  {'  '.join(mac_f)}",
            f"  Deltas vs base (macro):",
            f"  {ft_label:<{col_w}}  {'  '.join(dmac)}",
        ]

    with open(out_dir / "summary.txt", "w", encoding="utf-8") as fh:
        fh.write("\n".join(lines) + "\n")
    print("  Saved: new_outputs/summary.txt")


def run_new_outputs(
    all_neighbors: Dict[str, np.ndarray],
    common:        "pd.Index",
    protein_csv:   Path,
    fg_path:       Optional[Path],
    out_dir:       Path,
    k_values:      List[int],
) -> None:
    """
    Original new_outputs/ folder: 2-model comparison (base vs ContraMLM_v1_1),
    heatmaps, bar charts. Logic unchanged from original script.
    """
    if fg_path is None or not fg_path.exists():
        print(f"  WARNING: functional_groups.json not found — skipping new_outputs.")
        return

    active_models = list(all_neighbors.keys())
    ft_candidates = [m for m in active_models if m != "base"]
    if "base" not in active_models or not ft_candidates:
        print("  WARNING: need 'base' and at least one FT model — skipping new_outputs.")
        return
    ft_label = "ContraMLM_v1_1" if "ContraMLM_v1_1" in ft_candidates else ft_candidates[0]

    new_dir = out_dir / "new_outputs"
    new_dir.mkdir(parents=True, exist_ok=True)
    print(f"\n{'=' * 60}")
    print(f"Generating new_outputs (base vs {ft_label}) → {new_dir}")
    print(f"{'=' * 60}")

    df_meta    = load_full_metadata(protein_csv)
    common_ids = list(common)
    ann_labels = np.array(
        [str(df_meta.loc[pid, "bestPhrogAn"]) if pid in df_meta.index else "unknown"
         for pid in common_ids], dtype=str
    )
    cat_labels = np.array(
        [str(df_meta.loc[pid, "PhrogCat"]) if pid in df_meta.index else "unknown"
         for pid in common_ids], dtype=str
    )

    nb_base = all_neighbors["base"]
    nb_ft   = all_neighbors[ft_label]

    prec_ann_base = compute_precision_from_neighbors(nb_base, ann_labels, k_values)
    prec_ann_ft   = compute_precision_from_neighbors(nb_ft,   ann_labels, k_values)
    prec_cat_base = compute_precision_from_neighbors(nb_base, cat_labels, k_values)
    prec_cat_ft   = compute_precision_from_neighbors(nb_ft,   cat_labels, k_values)

    group_members, exclude_neg_map = load_functional_groups(fg_path)

    gp_base = compute_group_precision_with_excl(
        nb_base, ann_labels, group_members, exclude_neg_map, k_values
    )
    gp_ft = compute_group_precision_with_excl(
        nb_ft, ann_labels, group_members, exclude_neg_map, k_values
    )
    macro_base = macro_avg_group_precision(gp_base, ann_labels, group_members, k_values)
    macro_ft   = macro_avg_group_precision(gp_ft,   ann_labels, group_members, k_values)

    gp_rows = []
    for group in group_members:
        row: Dict = {
            "group":      group,
            "n_proteins": int(np.isin(ann_labels, list(group_members[group])).sum()),
        }
        for k in sorted(k_values):
            pb = gp_base[group][k]
            pf = gp_ft[group][k]
            row[f"base_mean_P@{k}"]        = round(float(np.nanmean(pb)) if len(pb) > 0 else float("nan"), 6)
            row[f"{ft_label}_mean_P@{k}"]  = round(float(np.nanmean(pf)) if len(pf) > 0 else float("nan"), 6)
            row[f"base_macro_P@{k}"]       = round(macro_base[group][k], 6)
            row[f"{ft_label}_macro_P@{k}"] = round(macro_ft[group][k], 6)
        gp_rows.append(row)
    pd.DataFrame(gp_rows).to_csv(new_dir / "group_precision.csv", index=False)
    print("  Saved: group_precision.csv")

    for k in sorted(k_values):
        plot_group_bar_chart_2models(macro_base, macro_ft, ft_label, k, new_dir)

    for model_label, nb in [("base", nb_base), (ft_label, nb_ft)]:
        for k in sorted(k_values):
            print(f"  Heatmap: {model_label} P@{k} …")
            matrix, gnames = compute_functional_confusion_matrix(
                nb, ann_labels, group_members, exclude_neg_map, k
            )
            plot_functional_confusion_heatmap(matrix, gnames, k, model_label, new_dir)

    write_new_summary(
        prec_ann_base=prec_ann_base,
        prec_ann_ft=prec_ann_ft,
        prec_cat_base=prec_cat_base,
        prec_cat_ft=prec_cat_ft,
        ann_labels=ann_labels,
        cat_labels=cat_labels,
        ft_label=ft_label,
        k_values=k_values,
        out_dir=new_dir,
    )
    print(f"  new_outputs complete → {new_dir}")


# ---------------------------------------------------------------------------
# NEW: new_outputs_all  (all models, clustermaps, 4-model bar charts)
# ---------------------------------------------------------------------------

def _filter_confusion_matrix_to_groups(
    matrix:      np.ndarray,
    group_names: List[str],
    keep_groups: List[str],
) -> Tuple[np.ndarray, List[str]]:
    """
    Slice a confusion matrix to keep only the rows/columns whose names
    appear in keep_groups (preserving the order of keep_groups).
    Groups not present in group_names are silently skipped.
    """
    valid_keep = [g for g in keep_groups if g in group_names]
    idx        = [group_names.index(g) for g in valid_keep]
    sub_matrix = matrix[np.ix_(idx, idx)]
    return sub_matrix, valid_keep


def plot_functional_clustermap(
    matrix:      np.ndarray,
    group_names: List[str],
    k:           int,
    model_label: str,
    out_dir:     Path,
    suffix:      str = "",
) -> None:
    """
    Seaborn clustermap of a confusion matrix.
    Rows and columns are clustered by hierarchical clustering.
    NaN values are replaced with 0 before clustering.
    """
    disp = np.nan_to_num(matrix, nan=0.0)
    df   = pd.DataFrame(disp, index=group_names, columns=group_names)

    vmax = max(float(disp.max()), 0.01)
    g = sns.clustermap(
        df,
        cmap="Blues",
        vmin=0,
        vmax=vmax,
        annot=True,
        fmt=".2f",
        annot_kws={"size": 6},
        linewidths=0.3,
        linecolor="white",
        figsize=(max(12, len(group_names) * 0.75), max(10, len(group_names) * 0.7)),
        cbar_kws={"label": "Mean fraction of neighbors", "shrink": 0.6},
        xticklabels=True,
        yticklabels=True,
    )
    g.ax_heatmap.set_xlabel("Neighbor group", fontsize=9)
    g.ax_heatmap.set_ylabel("Query group",    fontsize=9)
    g.ax_heatmap.set_xticklabels(
        g.ax_heatmap.get_xticklabels(), rotation=45, ha="right", fontsize=7
    )
    g.ax_heatmap.set_yticklabels(
        g.ax_heatmap.get_yticklabels(), rotation=0, fontsize=7
    )
    title = f"Functional Group Clustermap — {model_label}  P@{k}"
    if suffix:
        title += f"  [{suffix}]"
    g.fig.suptitle(title, y=1.01, fontsize=10, fontweight="bold")

    fname = out_dir / f"clustermap_k{k}_{model_label}{('_' + suffix) if suffix else ''}.png"
    g.fig.savefig(fname, dpi=150, bbox_inches="tight")
    plt.close(g.fig)

    # Also save the reordered CSV (clustered order)
    row_order = g.dendrogram_row.reordered_ind
    col_order = g.dendrogram_col.reordered_ind
    reordered = df.iloc[row_order, col_order]
    reordered.to_csv(
        out_dir / f"clustermap_k{k}_{model_label}{('_' + suffix) if suffix else ''}.csv"
    )
    print(f"  Saved: {fname.name}")


def _collect_group_sizes(
    ann_labels:    np.ndarray,
    group_members: Dict[str, set],
    target_groups: List[str],
) -> Dict[str, int]:
    """Count proteins per group (only groups in target_groups)."""
    sizes = {}
    for g in target_groups:
        if g in group_members:
            sizes[g] = int(np.isin(ann_labels, list(group_members[g])).sum())
        else:
            sizes[g] = 0
    return sizes


def plot_4model_group_barchart(
    macro_all:     Dict[str, Dict[str, Dict[int, float]]],  # model → group → k → value
    model_labels:  List[str],
    groups:        List[str],
    k:             int,
    title:         str,
    fname:         Path,
) -> None:
    """
    Grouped bar chart: for each group on x-axis, one bar per model side by side.
    Uses macro-average P@k values with the same exclude_neg logic as original.
    """
    n_groups = len(groups)
    n_models = len(model_labels)
    width    = 0.7 / n_models
    offsets  = np.linspace(
        -(0.7 / 2) + width / 2,
         (0.7 / 2) - width / 2,
        n_models,
    )
    x = np.arange(n_groups)

    fig, ax = plt.subplots(figsize=(max(14, n_groups * 1.1), 5))

    for ml, offset in zip(model_labels, offsets):
        vals = [
            macro_all[ml].get(g, {}).get(k, float("nan"))
            for g in groups
        ]
        color = MODEL_COLORS.get(ml, "#888888")
        bars  = ax.bar(x + offset, vals, width, label=ml, color=color, alpha=0.85)
        for bar, val in zip(bars, vals):
            if not np.isnan(val) and val > 0:
                ax.text(
                    bar.get_x() + bar.get_width() / 2,
                    bar.get_height() + 0.01,
                    f"{val:.2f}",
                    ha="center", va="bottom", fontsize=5.5, rotation=90,
                )

    ax.set_xticks(x)
    ax.set_xticklabels(groups, rotation=40, ha="right", fontsize=8)
    ax.set_ylabel(f"Macro P@{k}", fontsize=10)
    ax.set_ylim(0, 1.12)
    ax.set_title(title, fontsize=11, fontweight="bold")
    ax.legend(fontsize=8, loc="upper right")
    ax.yaxis.grid(True, linestyle="--", alpha=0.5)
    ax.set_axisbelow(True)
    fig.tight_layout()
    fig.savefig(fname, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"  Saved: {fname.name}")


def run_new_outputs_all(
    all_neighbors: Dict[str, np.ndarray],
    common:        "pd.Index",
    protein_csv:   Path,
    fg_path:       Optional[Path],
    out_dir:       Path,
    k_values:      List[int],
    model_labels:  List[str],
) -> None:
    """
    new_outputs_all/ folder:
      - Clustermaps (instead of heatmaps) for CLUSTERMAP_GROUPS only,
        for each of the all models × each k value.
      - Two 4-model bar charts per k:
          * top-10 largest functional groups (by protein count)
          * top-10 smallest functional groups (by protein count, min 1 protein)
    All computation uses the same exclude_neg logic as the rest of the script.
    """
    if fg_path is None or not fg_path.exists():
        print(f"  WARNING: functional_groups.json not found — skipping new_outputs_all.")
        return

    new_dir = out_dir / "new_outputs_all"
    new_dir.mkdir(parents=True, exist_ok=True)
    print(f"\n{'=' * 60}")
    print(f"Generating new_outputs_all (all models) → {new_dir}")
    print(f"{'=' * 60}")

    # ---- metadata ----
    df_meta    = load_full_metadata(protein_csv)
    common_ids = list(common)
    ann_labels = np.array(
        [str(df_meta.loc[pid, "bestPhrogAn"]) if pid in df_meta.index else "unknown"
         for pid in common_ids], dtype=str
    )

    # ---- functional groups ----
    group_members, exclude_neg_map = load_functional_groups(fg_path)

    # Validate clustermap groups against what's in functional_groups.json
    valid_clustermap_groups = [g for g in CLUSTERMAP_GROUPS if g in group_members]
    missing = [g for g in CLUSTERMAP_GROUPS if g not in group_members]
    if missing:
        print(f"  WARNING: These CLUSTERMAP_GROUPS were not found in functional_groups.json "
              f"and will be skipped: {missing}")
    print(f"  Clustermap groups to use: {valid_clustermap_groups}")

    # Sub-dict of group_members containing only clustermap groups
    clustermap_group_members = {
        g: group_members[g] for g in valid_clustermap_groups
    }

    # ---- compute macro P@k for ALL groups, for ALL models ----
    # macro_all[model][group][k] = float
    macro_all: Dict[str, Dict[str, Dict[int, float]]] = {}

    for ml in model_labels:
        print(f"\n  Computing group precision for [{ml}] …")
        nb = all_neighbors[ml]
        gp = compute_group_precision_with_excl(
            nb, ann_labels, group_members, exclude_neg_map, k_values
        )
        macro_all[ml] = macro_avg_group_precision(gp, ann_labels, group_members, k_values)

    # ---- clustermaps for each model × each k (only CLUSTERMAP_GROUPS) ----
    print("\n  Generating clustermaps …")
    for ml in model_labels:
        nb = all_neighbors[ml]
        for k in sorted(k_values):
            print(f"    Clustermap: [{ml}] P@{k} …")
            # Compute confusion matrix for ALL groups, then slice to clustermap groups
            full_matrix, full_names = compute_functional_confusion_matrix(
                nb, ann_labels, clustermap_group_members, exclude_neg_map, k
            )
            # full_matrix is already restricted to clustermap_group_members,
            # so no further slicing needed
            plot_functional_clustermap(
                matrix=full_matrix,
                group_names=full_names,
                k=k,
                model_label=ml,
                out_dir=new_dir,
            )

    # ---- collect group sizes (from clustermap groups + all groups for bar charts) ----
    all_group_sizes = _collect_group_sizes(ann_labels, group_members, list(group_members.keys()))

    # Filter to groups that have at least 1 protein and a valid macro score for base
    scored_groups = [
        g for g in group_members
        if all_group_sizes.get(g, 0) > 0
        and not np.isnan(macro_all["base"].get(g, {}).get(sorted(k_values)[0], float("nan")))
    ]
    scored_sizes = {g: all_group_sizes[g] for g in scored_groups}

    sorted_by_size_desc = sorted(scored_sizes.keys(), key=lambda g: scored_sizes[g], reverse=True)
    sorted_by_size_asc  = sorted(scored_sizes.keys(), key=lambda g: scored_sizes[g])

    top10_largest  = sorted_by_size_desc[:10]
    top10_smallest = [g for g in sorted_by_size_asc if scored_sizes[g] > 0][:10]

    # ---- 4-model bar charts per k ----
    print("\n  Generating 4-model bar charts …")
    for k in sorted(k_values):
        # Top-10 largest
        plot_4model_group_barchart(
            macro_all    = macro_all,
            model_labels = model_labels,
            groups       = top10_largest,
            k            = k,
            title        = (
                f"Top-10 Largest Functional Groups — Macro P@{k}\n"
                f"(ranked by protein count; all models)"
            ),
            fname        = new_dir / f"barchart_top10_largest_k{k}.png",
        )

        # Top-10 smallest
        plot_4model_group_barchart(
            macro_all    = macro_all,
            model_labels = model_labels,
            groups       = top10_smallest,
            k            = k,
            title        = (
                f"Top-10 Smallest Functional Groups — Macro P@{k}\n"
                f"(ranked by protein count ascending; all models)"
            ),
            fname        = new_dir / f"barchart_top10_smallest_k{k}.png",
        )

    # ---- save group sizes CSV for reference ----
    size_rows = [
        {"group": g, "n_proteins": all_group_sizes.get(g, 0)}
        for g in group_members
    ]
    pd.DataFrame(size_rows).sort_values("n_proteins", ascending=False).to_csv(
        new_dir / "group_sizes.csv", index=False
    )
    print("  Saved: group_sizes.csv")

    # ---- save macro precision CSV for all models and groups ----
    prec_rows = []
    for g in group_members:
        row: Dict = {
            "group":      g,
            "n_proteins": all_group_sizes.get(g, 0),
        }
        for ml in model_labels:
            for k in sorted(k_values):
                row[f"{ml}_macro_P@{k}"] = round(
                    macro_all[ml].get(g, {}).get(k, float("nan")), 6
                )
        prec_rows.append(row)
    pd.DataFrame(prec_rows).sort_values("n_proteins", ascending=False).to_csv(
        new_dir / "group_macro_precision_all_models.csv", index=False
    )
    print("  Saved: group_macro_precision_all_models.csv")

    print(f"\n  new_outputs_all complete → {new_dir}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    global FT_SUFFIX
    t_start  = time.time()
    args     = parse_args()
    FT_SUFFIX = args.ft_suffix
    emb_dir  = Path(args.emb_dir)
    out_dir  = Path(args.output_dir)
    figs_dir   = out_dir / "figs"
    texts_dir  = out_dir / "texts"
    for d in (figs_dir, texts_dir):
        d.mkdir(parents=True, exist_ok=True)

    print(f"ALL_VERSIONS : {ALL_VERSIONS}")
    print(f"FT_SUFFIX    : {FT_SUFFIX}")

    # --- PHROG labels ---
    phrog_series = load_phrog_labels(Path(args.protein_csv), args.min_family_size)
    family_sizes = phrog_series.value_counts()

    # --- Load embeddings ---
    MODEL_CONFIGS = build_model_configs()
    model_labels  = list(MODEL_CONFIGS.keys())

    embs_dict: Dict[str, pd.DataFrame] = {}
    for ml, pkl_name in MODEL_CONFIGS.items():
        pkl_path = emb_dir / pkl_name
        if not pkl_path.exists():
            print(f"WARNING: pkl not found: {pkl_path} — skipping {ml}.")
            continue
        embs_dict[ml] = load_embeddings(pkl_path)

    if not embs_dict:
        print("No embedding files found. Check --emb-dir.")
        return

    # --- Common protein set (all models ∩ valid PHROG) ---
    common = phrog_series.index
    for ml, embs in embs_dict.items():
        common = common.intersection(embs.index)
    print(f"\nProteins in evaluation (all models ∩ valid PHROG): {len(common)}")

    labels        = np.asarray(phrog_series.loc[common].values, dtype=str)
    active_labels = list(embs_dict.keys())

    # --- Compute Precision@k for each model ---
    all_prec_at_k: Dict[str, Dict[int, np.ndarray]] = {}
    all_neighbors: Dict[str, np.ndarray]             = {}
    all_results:   Dict[str, Dict]                   = {}

    for ml in active_labels:
        print(f"\n{'=' * 60}")
        print(f"Model: {ml}  ({MODEL_CONFIGS[ml]})")
        print(f"{'=' * 60}")

        X                    = embs_dict[ml].loc[common].values.astype(np.float32)
        prec_at_k, neighbors = compute_precision_at_k(
            X, labels, K_VALUES,
            M=args.hnsw_m,
            ef_construction=args.hnsw_ef_construction,
            ef_search=args.hnsw_ef_search,
        )
        all_prec_at_k[ml] = prec_at_k
        all_neighbors[ml] = neighbors

        for k in sorted(K_VALUES):
            print(f"  Mean Precision@{k}: {prec_at_k[k].mean():.5f}")

        all_results[ml] = aggregate_results(prec_at_k, labels, family_sizes)

        df_prot = pd.DataFrame(
            {f"P@{k}": prec_at_k[k] for k in K_VALUES},
            index=common,
        )
        df_prot.to_csv(texts_dir / f"per_protein_precision_{ml}.csv")

    # --- Standard plots (all models) ---
    print("\n--- Generating standard plots ---")
    plot_precision_at_k_bar(all_results, figs_dir, active_labels)
    for k in K_VALUES:
        plot_precision_by_family_size(all_results, figs_dir, active_labels, k=k)

    if "base" in all_prec_at_k:
        for ml in active_labels:
            if ml == "base":
                continue
            for k in K_VALUES:
                plot_per_family_scatter(
                    all_results["base"], all_results[ml], ml, figs_dir, k=k
                )
                plot_delta_histogram(
                    all_prec_at_k["base"][k],
                    all_prec_at_k[ml][k],
                    ml, figs_dir, k=k,
                )

    # --- Text outputs ---
    write_summary(all_results, active_labels, texts_dir)
    write_per_family_csv(all_results, active_labels, texts_dir)
    write_summary_appendix(
        all_results, active_labels, family_sizes,
        Path(args.protein_csv), texts_dir,
        all_neighbors=all_neighbors,
        labels=labels,
    )

    # --- new_outputs/ (original: base vs ContraMLM_v1_1, heatmaps) ---
    run_new_outputs(
        all_neighbors=all_neighbors,
        common=common,
        protein_csv=Path(args.protein_csv),
        fg_path=Path(args.functional_groups),
        out_dir=out_dir,
        k_values=K_VALUES,
    )

    # --- new_outputs_all/ (all models, clustermaps, 4-model bar charts) ---
    run_new_outputs_all(
        all_neighbors=all_neighbors,
        common=common,
        protein_csv=Path(args.protein_csv),
        fg_path=Path(args.functional_groups),
        out_dir=out_dir,
        k_values=K_VALUES,
        model_labels=active_labels,
    )

    elapsed = time.time() - t_start
    print(f"\nDone. Total time: {elapsed / 60:.1f} min")
    print(f"Results in: {out_dir}")


if __name__ == "__main__":
    main()