#!/usr/bin/env python3
"""
MMseqs2 cluster-based train/test split from CSV (highFRET / noFRET).

Convention (Bushuiev & Bushuiev-style): cluster all sequences at a fixed identity
threshold, rank clusters by how distant they are from the rest of the dataset, then
hold out entire clusters for test. No sequence in a test cluster appears in train,
which controls leakage across both classes.

Pipeline:
1. Read CSV (variant, Protein, label).
2. MMseqs2 easy-cluster on all sequences (--min-seq-id / coverage).
3. All-against-all MMseqs2 search for pairwise identities.
4. Per cluster: mean sequence identity to sequences *outside* that cluster (global).
5. Assign the most distant clusters to test (whole clusters only).
   Default: hold out --test-fraction of all clusters by count (e.g. 20% of clusters).
   Optional --holdout-by sequences: add distant clusters until ~that fraction of sequences.

Also reports cross-class similarity (test vs train, other label) for rebuttal text.

Primary outputs: <prefix>_train.csv and <prefix>_test.csv in FRET selection layout
(Unnamed: 0, variant, DNA, Protein, TPM columns, label — no split column).
Analysis/debug: <prefix>_with_split.csv, clustering stats, cluster summary.

Outputs include clustering diagnostics (*_clustering_stats.txt / .csv) to verify
whether clusters are mostly singletons (degenerate cluster split).

Usage:
    python mmseqs_cluster_split.py <input.csv> \\
        [--cluster-identity 0.8] [--test-fraction 0.2]
"""

import argparse
import os
import sys
import tempfile
import shutil
from pathlib import Path

import pandas as pd
import subprocess


def run_mmseqs(cmd: list[str]) -> None:
    """Run an MMseqs2 command; print stderr/stdout and exit on failure."""
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"ERROR: MMseqs2 failed: {' '.join(cmd)}", file=sys.stderr)
        if result.stderr:
            print(result.stderr, file=sys.stderr)
        if result.stdout:
            print(result.stdout, file=sys.stderr)
        sys.exit(result.returncode)


def check_mmseqs_available() -> str:
    """Ensure `mmseqs` is on PATH; return version string for logging."""
    if shutil.which("mmseqs") is None:
        print(
            "ERROR: MMseqs2 not found on PATH.\n"
            "  conda:  conda install -c bioconda mmseqs2\n"
            "  HPC:    module load <mmseqs2-module>  (site-specific)\n"
            "  verify: which mmseqs && mmseqs version",
            file=sys.stderr,
        )
        sys.exit(1)
    result = subprocess.run(
        ["mmseqs", "version"],
        capture_output=True,
        text=True,
    )
    version = (result.stdout or result.stderr).strip().splitlines()[0] if result.returncode == 0 else "unknown"
    return version


LABEL_HIGH = "highFRET"
LABEL_NO = "noFRET"

# Train/test deliverables: same column layout as the FRET selection CSV (no split column).
FRET_EXPORT_COLUMNS = [
    "variant",
    "DNA",
    "Protein",
    "R0_TPM",
    "R1F_TPM",
    "R2F_TPM",
    "R3F_TPM",
    "R1N_TPM",
    "R2N_TPM",
    "label",
]

_LABEL_ALIASES = {
    "highfret": LABEL_HIGH,
    "high fret": LABEL_HIGH,
    "no fret": LABEL_NO,
    "nofret": LABEL_NO,
    "no_fret": LABEL_NO,
    "lowfret": LABEL_NO,
    "low fret": LABEL_NO,
    "low_fret": LABEL_NO,
}


def _normalize_label(raw: str) -> str | None:
    key = " ".join(str(raw).strip().lower().split())
    return _LABEL_ALIASES.get(key)


def load_csv(csv_path: str) -> pd.DataFrame:
    df = pd.read_csv(csv_path)
    for col in ("variant", "Protein", "label"):
        if col not in df.columns:
            raise ValueError(f"CSV must have column '{col}'. Found: {list(df.columns)}")
    df["variant"] = df["variant"].astype(str).str.strip()
    df["Protein"] = df["Protein"].astype(str).str.strip()
    df["label"] = df["label"].astype(str).str.strip()
    return df


def filter_labels(df: pd.DataFrame) -> pd.DataFrame:
    canonical = []
    for raw in df["label"]:
        c = _normalize_label(str(raw))
        canonical.append(c)
    df = df.copy()
    df["label_original"] = df["label"].astype(str).str.strip()
    df["label"] = canonical
    out = df[df["label"].notna()].copy()
    out["label"] = out["label"].astype(str)
    dropped = len(df) - len(out)
    if dropped:
        print(
            f"  Dropped {dropped} rows with unrecognized labels "
            "(accepted: highFRET, noFRET, no FRET, lowFRET, etc.)"
        )
    return out.reset_index(drop=True)


def csv_to_fasta(df: pd.DataFrame, fasta_path: str, id_col: str = "variant", seq_col: str = "Protein") -> None:
    with open(fasta_path, "w") as f:
        for _, row in df.iterrows():
            seq_id = str(row[id_col]).replace(" ", "_")
            seq = str(row[seq_col])
            if not seq:
                continue
            f.write(f">{seq_id}\n{seq}\n")
    print(f"  Wrote {len(df)} sequences to {fasta_path}")


def perform_mmseqs_easy_cluster(
    fasta_file: str,
    tmp_dir: str,
    identity: float,
    coverage: float,
    threads: int,
    evalue: float,
) -> str:
    """Run MMseqs2 easy-cluster; return path to *_cluster.tsv."""
    print(
        f"Running MMseqs2 easy-cluster (min-seq-id={identity}, coverage={coverage})..."
    )
    cluster_prefix = os.path.join(tmp_dir, "cluster")
    mmseqs_tmp = os.path.join(tmp_dir, "mmseqs_cluster_tmp")
    os.makedirs(mmseqs_tmp, exist_ok=True)

    cmd = [
        "mmseqs",
        "easy-cluster",
        fasta_file,
        cluster_prefix,
        mmseqs_tmp,
        "--min-seq-id",
        str(identity),
        "-c",
        str(coverage),
        "--cov-mode",
        "0",
        "-e",
        str(evalue),
        "--threads",
        str(threads),
        "--compressed",
        "1",
        "--remove-tmp-files",
        "1",
        "--split",
        "0",
        "--db-load-mode",
        "1",
    ]
    run_mmseqs(cmd)

    cluster_file = f"{cluster_prefix}_cluster.tsv"
    if not os.path.exists(cluster_file):
        raise FileNotFoundError(f"Cluster file not found: {cluster_file}")
    print(f"  Clustering finished: {cluster_file}")
    return cluster_file


def parse_cluster_assignments(cluster_file: str) -> tuple[dict[str, str], dict[str, set[str]]]:
    """
    Parse MMseqs2 cluster TSV (rep_id, member_id).

    Returns:
        member_to_rep: member_id -> representative cluster id
        cluster_members: rep_id -> set of member ids
    """
    member_to_rep: dict[str, str] = {}
    cluster_members: dict[str, set[str]] = {}

    with open(cluster_file) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split("\t")
            if len(parts) < 2:
                continue
            rep_id, member_id = parts[0], parts[1]
            member_to_rep[member_id] = rep_id
            cluster_members.setdefault(rep_id, set()).add(member_id)

    print(f"  Parsed {len(cluster_members)} clusters, {len(member_to_rep)} sequence assignments")
    return member_to_rep, cluster_members


def compute_clustering_stats(
    cluster_members: dict[str, set[str]],
    df: pd.DataFrame,
    cluster_identity: float,
    cluster_coverage: float,
) -> tuple[dict[str, object], list[str]]:
    """
    Summarize cluster size distribution (detect all-singleton clustering).

    Returns (stats_dict, report_lines for printing / text file).
    """
    sizes = [len(m) for m in cluster_members.values()]
    n_clusters = len(sizes)
    n_sequences = int(sum(sizes))
    n_singleton_clusters = sum(1 for s in sizes if s == 1)
    seq_in_singletons = sum(s for s in sizes if s == 1)
    multi_sizes = [s for s in sizes if s > 1]
    n_multi_clusters = len(multi_sizes)
    seq_in_multi = sum(multi_sizes)

    def _pct(n: int, d: int) -> float:
        return 100.0 * n / d if d else 0.0

    size_series = pd.Series(sizes) if sizes else pd.Series(dtype=float)
    buckets = {
        "size_1": sum(1 for s in sizes if s == 1),
        "size_2": sum(1 for s in sizes if s == 2),
        "size_3_5": sum(1 for s in sizes if 3 <= s <= 5),
        "size_6_10": sum(1 for s in sizes if 6 <= s <= 10),
        "size_11_plus": sum(1 for s in sizes if s >= 11),
    }

    largest_rep = max(cluster_members, key=lambda r: len(cluster_members[r])) if cluster_members else ""
    largest_size = len(cluster_members[largest_rep]) if largest_rep else 0

    # Label composition inside multi-member clusters only
    multi_mixed_label = 0
    multi_pure_high = 0
    multi_pure_no = 0
    if "variant_key" in df.columns and "label" in df.columns:
        for rep_id, members in cluster_members.items():
            if len(members) < 2:
                continue
            labels = df.loc[df["variant_key"].isin(members), "label"].unique()
            if len(labels) > 1:
                multi_mixed_label += 1
            elif LABEL_HIGH in labels:
                multi_pure_high += 1
            elif LABEL_NO in labels:
                multi_pure_no += 1

    all_singleton = n_singleton_clusters == n_clusters and n_clusters > 0
    degenerate_split = all_singleton or n_multi_clusters == 0

    stats: dict[str, object] = {
        "n_sequences": n_sequences,
        "n_clusters": n_clusters,
        "compression_ratio_seq_per_cluster": n_sequences / n_clusters if n_clusters else 0.0,
        "cluster_identity_threshold": cluster_identity,
        "cluster_coverage_threshold": cluster_coverage,
        "n_singleton_clusters": n_singleton_clusters,
        "pct_singleton_clusters": _pct(n_singleton_clusters, n_clusters),
        "n_sequences_in_singleton_clusters": seq_in_singletons,
        "pct_sequences_in_singleton_clusters": _pct(seq_in_singletons, n_sequences),
        "n_multi_member_clusters": n_multi_clusters,
        "pct_multi_member_clusters": _pct(n_multi_clusters, n_clusters),
        "n_sequences_in_multi_member_clusters": seq_in_multi,
        "pct_sequences_in_multi_member_clusters": _pct(seq_in_multi, n_sequences),
        "cluster_size_min": int(size_series.min()) if len(size_series) else 0,
        "cluster_size_max": largest_size,
        "cluster_size_mean": float(size_series.mean()) if len(size_series) else 0.0,
        "cluster_size_median": float(size_series.median()) if len(size_series) else 0.0,
        "cluster_size_std": float(size_series.std()) if len(size_series) > 1 else 0.0,
        "largest_cluster_id": largest_rep,
        "largest_cluster_n_members": largest_size,
        "bucket_clusters_size_1": buckets["size_1"],
        "bucket_clusters_size_2": buckets["size_2"],
        "bucket_clusters_size_3_5": buckets["size_3_5"],
        "bucket_clusters_size_6_10": buckets["size_6_10"],
        "bucket_clusters_size_11_plus": buckets["size_11_plus"],
        "multi_member_clusters_mixed_label": multi_mixed_label,
        "multi_member_clusters_pure_highFRET": multi_pure_high,
        "multi_member_clusters_pure_noFRET": multi_pure_no,
        "all_clusters_singleton": all_singleton,
        "cluster_split_degenerate": degenerate_split,
    }

    lines = [
        "=" * 60,
        "MMseqs2 clustering statistics",
        "=" * 60,
        f"Parameters: min-seq-id={cluster_identity}, coverage={cluster_coverage}",
        "",
        f"Sequences: {n_sequences}",
        f"Clusters:  {n_clusters}",
        f"Compression (seq/cluster): {stats['compression_ratio_seq_per_cluster']:.3f}",
        "",
        "Cluster size distribution",
        f"  min / median / mean / max: "
        f"{stats['cluster_size_min']} / {stats['cluster_size_median']:.2f} / "
        f"{stats['cluster_size_mean']:.2f} / {stats['cluster_size_max']}",
        f"  std dev: {stats['cluster_size_std']:.2f}",
        "",
        "Singleton clusters (size = 1)",
        f"  clusters:   {n_singleton_clusters} / {n_clusters} "
        f"({stats['pct_singleton_clusters']:.1f}%)",
        f"  sequences:  {seq_in_singletons} / {n_sequences} "
        f"({stats['pct_sequences_in_singleton_clusters']:.1f}%)",
        "",
        "Multi-member clusters (size > 1)",
        f"  clusters:   {n_multi_clusters} / {n_clusters} "
        f"({stats['pct_multi_member_clusters']:.1f}%)",
        f"  sequences:  {seq_in_multi} / {n_sequences} "
        f"({stats['pct_sequences_in_multi_member_clusters']:.1f}%)",
        f"  mixed label (highFRET + noFRET): {multi_mixed_label}",
        f"  pure highFRET only: {multi_pure_high}",
        f"  pure noFRET only: {multi_pure_no}",
        "",
        "Cluster count by size bucket",
        f"  size 1:     {buckets['size_1']}",
        f"  size 2:     {buckets['size_2']}",
        f"  size 3-5:   {buckets['size_3_5']}",
        f"  size 6-10:  {buckets['size_6_10']}",
        f"  size 11+:   {buckets['size_11_plus']}",
        "",
        f"Largest cluster: id={largest_rep!r}, n_members={largest_size}",
    ]

    if all_singleton:
        lines.extend(
            [
                "",
                "WARNING: Every cluster has exactly one sequence.",
                "  Cluster-based hold-out is equivalent to holding out individual sequences.",
                "  Consider lowering --cluster-identity / --cluster-coverage, or check",
                "  sequence diversity and FASTA IDs (duplicates / parsing).",
            ]
        )
    elif degenerate_split:
        lines.extend(
            [
                "",
                "WARNING: No multi-member clusters; cluster split has no within-cluster groups.",
            ]
        )

    lines.append("=" * 60)
    return stats, lines


def save_clustering_stats(
    stats: dict[str, object],
    report_lines: list[str],
    txt_path: str,
    csv_path: str,
) -> None:
    """Write human-readable report and one-row-per-metric CSV."""
    with open(txt_path, "w") as f:
        f.write("\n".join(report_lines) + "\n")
    pd.DataFrame([stats]).to_csv(csv_path, index=False)


def perform_all_vs_all_mmseqs_search(
    fasta_file: str,
    tmp_dir: str,
    sensitivity: float = 7.5,
    threads: int = 1,
    max_seqs: int = 100,
    evalue: float = 0.001,
) -> str:
    print("Running all-against-all MMseqs2 search...")
    query_db = os.path.join(tmp_dir, "query_db")
    target_db = query_db
    results_db = os.path.join(tmp_dir, "results_db")
    results_file = os.path.join(tmp_dir, "mmseqs_all_vs_all.tsv")
    mmseqs_tmp = os.path.join(tmp_dir, "mmseqs_search_tmp")
    os.makedirs(mmseqs_tmp, exist_ok=True)

    run_mmseqs(["mmseqs", "createdb", fasta_file, query_db])
    run_mmseqs(
        [
            "mmseqs",
            "search",
            query_db,
            target_db,
            results_db,
            mmseqs_tmp,
            "--threads",
            str(threads),
            "-s",
            str(sensitivity),
            "--max-seqs",
            str(max_seqs),
            "-e",
            str(evalue),
            "--compressed",
            "1",
            "--remove-tmp-files",
            "1",
        ]
    )
    run_mmseqs(
        ["mmseqs", "convertalis", query_db, target_db, results_db, results_file]
    )
    print("  MMseqs2 search finished.")
    return results_file


def parse_mmseqs_results(results_file: str) -> pd.DataFrame:
    return pd.read_csv(
        results_file,
        sep="\t",
        header=None,
        names=[
            "query",
            "target",
            "seqid",
            "alnlen",
            "mismatch",
            "gapopen",
            "qstart",
            "qend",
            "tstart",
            "tend",
            "evalue",
            "bits",
            "qcov",
            "tcov",
        ],
    )


def compute_cluster_external_similarity(
    search_results: pd.DataFrame,
    cluster_members: dict[str, set[str]],
    all_ids: set[str],
) -> dict[str, float]:
    """
    For each cluster, mean seqid from cluster members to sequences outside the cluster.
    Lower = more distant cluster (better candidate for test hold-out).
    """
    member_to_rep: dict[str, str] = {}
    for rep_id, members in cluster_members.items():
        for m in members:
            member_to_rep[m] = rep_id

    external_hits: dict[str, list[float]] = {rep: [] for rep in cluster_members}

    for _, row in search_results.iterrows():
        q, t, seqid = row["query"], row["target"], float(row["seqid"])
        if q == t:
            continue
        if q not in all_ids or t not in all_ids:
            continue
        q_cluster = member_to_rep.get(q)
        t_cluster = member_to_rep.get(t)
        if q_cluster is None or t_cluster is None:
            continue
        if q_cluster == t_cluster:
            continue
        external_hits[q_cluster].append(seqid)

    scores: dict[str, float] = {}
    for rep_id in cluster_members:
        hits = external_hits.get(rep_id, [])
        scores[rep_id] = sum(hits) / len(hits) if hits else 0.0
    return scores


def select_test_clusters(
    cluster_members: dict[str, set[str]],
    cluster_scores: dict[str, float],
    test_fraction: float,
    holdout_by: str,
) -> set[str]:
    """
    Select whole clusters for test (most distant = lowest external mean seqid first).

    holdout_by:
      clusters   — top test_fraction of cluster count (default, Bushuiev-style)
      sequences  — greedy add clusters until ~test_fraction of sequences are in test
    """
    ranked = sorted(cluster_scores.items(), key=lambda x: x[1])
    n_clusters = len(cluster_members)
    n_seq = sum(len(m) for m in cluster_members.values())

    if holdout_by == "clusters":
        n_take = max(1, int(round(test_fraction * n_clusters)))
        if n_take >= n_clusters and n_clusters > 1:
            n_take = n_clusters - 1
        return {rep_id for rep_id, _ in ranked[:n_take]}

    target_n = max(1, int(round(test_fraction * n_seq)))
    test_clusters: set[str] = set()
    test_count = 0
    for rep_id, _score in ranked:
        size = len(cluster_members[rep_id])
        if test_count >= target_n:
            break
        remaining_train = n_seq - test_count - size
        if remaining_train < 1 and test_clusters:
            break
        test_clusters.add(rep_id)
        test_count += size
    if not test_clusters and ranked:
        test_clusters.add(ranked[0][0])
    return test_clusters


def compute_cross_class_similarity_report(
    search_results: pd.DataFrame,
    df: pd.DataFrame,
    test_ids: set[str],
) -> pd.DataFrame:
    """
    Per test sequence: max/mean seqid to train sequences in the *other* label.
    """
    id_to_label = {
        str(row["variant_key"]): row["label"]
        for _, row in df.iterrows()
    }
    train_ids = {k for k in id_to_label if k not in test_ids}

    other_label_hits: dict[str, list[float]] = {tid: [] for tid in test_ids}
    same_label_train_hits: dict[str, list[float]] = {tid: [] for tid in test_ids}
    any_train_hits: dict[str, list[float]] = {tid: [] for tid in test_ids}

    for _, row in search_results.iterrows():
        q, t, seqid = row["query"], row["target"], float(row["seqid"])
        if q == t:
            continue
        if q not in test_ids or t not in train_ids:
            continue
        any_train_hits[q].append(seqid)
        q_label = id_to_label.get(q)
        t_label = id_to_label.get(t)
        if q_label is None or t_label is None:
            continue
        if q_label != t_label:
            other_label_hits[q].append(seqid)
        else:
            same_label_train_hits[q].append(seqid)

    rows = []
    for tid in sorted(test_ids):
        other = other_label_hits.get(tid, [])
        same = same_label_train_hits.get(tid, [])
        any_tr = any_train_hits.get(tid, [])
        rows.append(
            {
                "variant_key": tid,
                "label": id_to_label.get(tid, ""),
                "max_seqid_to_other_label_train": max(other) if other else 0.0,
                "mean_seqid_to_other_label_train": sum(other) / len(other) if other else 0.0,
                "max_seqid_to_same_label_train": max(same) if same else 0.0,
                "mean_seqid_to_same_label_train": sum(same) / len(same) if same else 0.0,
                "max_seqid_to_any_train": max(any_tr) if any_tr else 0.0,
                "mean_seqid_to_any_train": sum(any_tr) / len(any_tr) if any_tr else 0.0,
                "n_hits_other_label_train": len(other),
                "n_hits_same_label_train": len(same),
            }
        )
    return pd.DataFrame(rows)


def prepare_fret_export_df(df: pd.DataFrame) -> pd.DataFrame:
    """Subset and order columns for FRET-style train/test CSVs (label = original input)."""
    out = df.copy()
    if "label_original" in out.columns:
        out["label"] = out["label_original"]
    cols = [c for c in FRET_EXPORT_COLUMNS if c in out.columns]
    missing = [c for c in FRET_EXPORT_COLUMNS if c not in out.columns]
    if missing:
        print(f"  WARNING: FRET export missing columns: {missing}")
    return out[cols].reset_index(drop=True)


def write_fret_format_csv(df: pd.DataFrame, path: str) -> None:
    """Write CSV with pandas index as 'Unnamed: 0' (matches FRET selection export)."""
    prepare_fret_export_df(df).to_csv(path, index=True)


def main():
    parser = argparse.ArgumentParser(
        description="MMseqs2 cluster-based train/test split (hold out distant whole clusters)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("input_csv", type=str, help="Input CSV: variant, Protein, label")
    parser.add_argument(
        "--output-prefix",
        type=str,
        default=None,
        help="Prefix for output files (default: <input_stem>_cluster_split)",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="Directory for output CSVs (created if missing)",
    )
    parser.add_argument(
        "--test-fraction",
        type=float,
        default=0.2,
        help="Fraction for hold-out (default 0.2): clusters or sequences per --holdout-by",
    )
    parser.add_argument(
        "--holdout-by",
        choices=("clusters", "sequences"),
        default="clusters",
        help="Hold out test_fraction of clusters (default) or of sequences via whole clusters",
    )
    parser.add_argument(
        "--cluster-identity",
        type=float,
        default=0.8,
        help="MMseqs2 easy-cluster --min-seq-id (default 0.8)",
    )
    parser.add_argument(
        "--cluster-coverage",
        type=float,
        default=0.8,
        help="MMseqs2 easy-cluster -c coverage (default 0.8)",
    )
    parser.add_argument("--sensitivity", "-s", type=float, default=7.5, help="Search sensitivity")
    parser.add_argument("--threads", "-t", type=int, default=1, help="MMseqs threads (0 = all CPUs)")
    parser.add_argument("--max-seqs", type=int, default=100, help="Search --max-seqs")
    parser.add_argument("--evalue", "-e", type=float, default=0.001, help="E-value for cluster and search")
    parser.add_argument("--keep-temp", action="store_true", help="Keep temporary directory")
    parser.add_argument("--tmp-dir", type=str, default=None, help="Temp directory for MMseqs")
    args = parser.parse_args()

    mmseqs_version = check_mmseqs_available()
    print(f"MMseqs2 version: {mmseqs_version}")

    if not os.path.exists(args.input_csv):
        print(f"ERROR: Input CSV not found: {args.input_csv}", file=sys.stderr)
        sys.exit(1)
    if not 0 < args.test_fraction < 1:
        print("ERROR: --test-fraction must be between 0 and 1", file=sys.stderr)
        sys.exit(1)
    if not 0 < args.cluster_identity <= 1:
        print("ERROR: --cluster-identity must be in (0, 1]", file=sys.stderr)
        sys.exit(1)

    prefix = args.output_prefix or (Path(args.input_csv).stem + "_cluster_split")

    if args.threads == 0:
        import multiprocessing

        args.threads = multiprocessing.cpu_count()
        print(f"Using {args.threads} threads")

    tmp_dir = args.tmp_dir
    if tmp_dir:
        tmp_dir = os.path.abspath(tmp_dir)
        os.makedirs(tmp_dir, exist_ok=True)
    else:
        tmp_dir = tempfile.mkdtemp(prefix="mmseqs_csv_cluster_split_")
    print(f"Temporary directory: {tmp_dir}")

    try:
        df = load_csv(args.input_csv)
        df = filter_labels(df)
        if df.empty:
            print("ERROR: No rows with label highFRET or noFRET", file=sys.stderr)
            sys.exit(1)
        print(f"Loaded {len(df)} sequences ({df['label'].value_counts().to_dict()})")

        df["variant_key"] = df["variant"].astype(str).str.replace(" ", "_", regex=False)
        all_ids = set(df["variant_key"].astype(str))

        fasta_path = os.path.join(tmp_dir, "sequences.fasta")
        csv_to_fasta(df, fasta_path)

        cluster_file = perform_mmseqs_easy_cluster(
            fasta_path,
            tmp_dir,
            identity=args.cluster_identity,
            coverage=args.cluster_coverage,
            threads=args.threads,
            evalue=args.evalue,
        )
        member_to_rep, cluster_members = parse_cluster_assignments(cluster_file)
        df["cluster_id"] = df["variant_key"].map(member_to_rep)

        cluster_stats, cluster_report_lines = compute_clustering_stats(
            cluster_members,
            df,
            args.cluster_identity,
            args.cluster_coverage,
        )
        cluster_stats["mmseqs2_version"] = mmseqs_version
        cluster_report_lines.insert(3, f"MMseqs2 version: {mmseqs_version}")
        print("")
        for line in cluster_report_lines:
            print(line)

        missing = df["cluster_id"].isna().sum()
        if missing:
            print(f"  WARNING: {missing} sequences missing from cluster TSV", file=sys.stderr)

        results_file = perform_all_vs_all_mmseqs_search(
            fasta_path,
            tmp_dir,
            sensitivity=args.sensitivity,
            threads=args.threads,
            max_seqs=args.max_seqs,
            evalue=args.evalue,
        )
        search_results = parse_mmseqs_results(results_file)
        print(f"  Parsed {len(search_results)} MMseqs2 hits")

        cluster_scores = compute_cluster_external_similarity(
            search_results, cluster_members, all_ids
        )
        test_cluster_ids = select_test_clusters(
            cluster_members,
            cluster_scores,
            args.test_fraction,
            args.holdout_by,
        )

        test_member_ids: set[str] = set()
        for rep_id in test_cluster_ids:
            test_member_ids.update(cluster_members[rep_id])

        test_variants = set(
            df.loc[df["variant_key"].isin(test_member_ids), "variant"].astype(str)
        )
        df["split"] = df["variant"].apply(lambda v: "test" if v in test_variants else "train")
        df["cluster_external_mean_seqid"] = df["cluster_id"].map(cluster_scores)

        n_clusters = len(cluster_members)
        n_test_clusters = len(test_cluster_ids)
        print(
            f"  Test clusters: {n_test_clusters}/{n_clusters} "
            f"({100 * n_test_clusters / n_clusters:.1f}% of clusters)"
        )
        print(
            f"  Test sequences: {len(test_member_ids)}/{len(df)} "
            f"({100 * len(test_member_ids) / len(df):.1f}% of sequences)"
        )
        print(f"  Hold-out mode: --holdout-by {args.holdout_by}")

        if args.output_dir:
            os.makedirs(args.output_dir, exist_ok=True)
            base = Path(prefix).name
            train_csv = os.path.join(args.output_dir, f"{base}_train.csv")
            test_csv = os.path.join(args.output_dir, f"{base}_test.csv")
            out_csv = os.path.join(args.output_dir, f"{base}_with_split.csv")
            cross_csv = os.path.join(args.output_dir, f"{base}_cross_class_similarity.csv")
            cluster_summary_csv = os.path.join(args.output_dir, f"{base}_cluster_summary.csv")
            cluster_stats_txt = os.path.join(args.output_dir, f"{base}_clustering_stats.txt")
            cluster_stats_csv = os.path.join(args.output_dir, f"{base}_clustering_stats.csv")
        else:
            train_csv = f"{prefix}_train.csv"
            test_csv = f"{prefix}_test.csv"
            out_csv = f"{prefix}_with_split.csv"
            cross_csv = f"{prefix}_cross_class_similarity.csv"
            cluster_summary_csv = f"{prefix}_cluster_summary.csv"
            cluster_stats_txt = f"{prefix}_clustering_stats.txt"
            cluster_stats_csv = f"{prefix}_clustering_stats.csv"

        save_clustering_stats(
            cluster_stats, cluster_report_lines, cluster_stats_txt, cluster_stats_csv
        )
        print(f"Saved clustering statistics to: {cluster_stats_txt}")
        print(f"Saved clustering statistics (CSV) to: {cluster_stats_csv}")

        analysis_cols = [
            c
            for c in df.columns
            if c not in ("variant_key", "label_original")
        ]
        df[analysis_cols].to_csv(out_csv, index=False)
        print(f"Saved analysis table (split, cluster_id, …) to: {out_csv}")

        train_df = df[df["split"] == "train"]
        test_df = df[df["split"] == "test"]
        write_fret_format_csv(train_df, train_csv)
        print(f"Saved train set ({len(train_df)} rows, FRET format) to: {train_csv}")
        write_fret_format_csv(test_df, test_csv)
        print(f"Saved test set ({len(test_df)} rows, FRET format) to: {test_csv}")

        cross_df = compute_cross_class_similarity_report(
            search_results, df, test_member_ids
        )
        cross_df.to_csv(cross_csv, index=False)
        print(f"Saved per-test cross-class similarity to: {cross_csv}")

        cluster_rows = []
        for rep_id, members in cluster_members.items():
            labels = df.loc[df["variant_key"].isin(members), "label"]
            n_mem = len(members)
            cluster_rows.append(
                {
                    "cluster_id": rep_id,
                    "n_members": n_mem,
                    "is_singleton": n_mem == 1,
                    "external_mean_seqid": cluster_scores.get(rep_id, 0.0),
                    "split": "test" if rep_id in test_cluster_ids else "train",
                    "n_highFRET": int((labels == LABEL_HIGH).sum()),
                    "n_noFRET": int((labels == LABEL_NO).sum()),
                    "has_mixed_labels": int(labels.nunique() > 1) if n_mem > 1 else 0,
                }
            )
        cluster_summary = pd.DataFrame(cluster_rows).sort_values(
            "external_mean_seqid", ascending=True
        )
        cluster_summary.to_csv(cluster_summary_csv, index=False)
        print(f"Saved cluster summary to: {cluster_summary_csv}")

        print("\nSplit summary:")
        for label in (LABEL_HIGH, LABEL_NO):
            sub = df[df["label"] == label]
            if len(sub) == 0:
                continue
            n_test = (sub["split"] == "test").sum()
            n_train = (sub["split"] == "train").sum()
            print(f"  {label}: train={n_train}, test={n_test} ({100 * n_test / len(sub):.1f}% test)")

        if not cross_df.empty:
            print("\nCross-class similarity (test sequences vs train, other label):")
            print(
                f"  mean of per-seq max seqid: "
                f"{cross_df['max_seqid_to_other_label_train'].mean():.4f}"
            )
            print(
                f"  mean of per-seq mean seqid: "
                f"{cross_df['mean_seqid_to_other_label_train'].mean():.4f}"
            )
            print(
                f"  fraction with max other-label train seqid >= {args.cluster_identity}: "
                f"{(cross_df['max_seqid_to_other_label_train'] >= args.cluster_identity).mean():.2%}"
            )
            print(
                "  (Low values / low fraction above threshold => stronger cross-class separation in test.)"
            )

        print(
            f"\nClustering: MMseqs2 easy-cluster at min-seq-id={args.cluster_identity}, "
            f"coverage={args.cluster_coverage}. "
            "Test = whole clusters with lowest external mean seqid."
        )
    finally:
        if not args.keep_temp and os.path.exists(tmp_dir):
            shutil.rmtree(tmp_dir)
        elif args.keep_temp:
            print(f"Temporary files kept: {tmp_dir}")


if __name__ == "__main__":
    main()