#!/usr/bin/env python3
"""Create a sequence-similarity-disjoint protein split using MMseqs2."""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
from collections import Counter
from pathlib import Path

from mitointeract_recovery.splits import assign_grouped_splits, assert_group_disjoint


def read_rows(path: Path) -> list[dict]:
    with path.open() as handle:
        return [json.loads(line) for line in handle if line.strip()]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", type=Path, default=Path("artifacts/dev-10k"))
    parser.add_argument("--min-seq-id", type=float, default=0.5)
    parser.add_argument("--coverage", type=float, default=0.8)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--threads", type=int, default=8)
    args = parser.parse_args()

    if not shutil.which("mmseqs"):
        raise RuntimeError("MMseqs2 is required; install the mmseqs2 package")
    rows = read_rows(args.data_dir / "sample.jsonl")
    audit_path = args.data_dir / "audit.json"
    audit = json.loads(audit_path.read_text())
    split_ratios = audit.get(
        "split_ratios", {"train": 0.8, "validation": 0.1, "test": 0.1}
    )
    proteins = sorted({row["protein_id"]: row["sequence"] for row in rows}.items())
    suffix = (
        f"identity-{int(args.min_seq_id * 100)}-coverage-{int(args.coverage * 100)}"
    )
    work = args.data_dir / f"mmseqs-{suffix}"
    work.mkdir(parents=True, exist_ok=True)
    fasta = work / "proteins.fasta"
    with fasta.open("w") as handle:
        for protein_id, sequence in proteins:
            handle.write(f">{protein_id}\n{sequence}\n")

    prefix = work / "clusters"
    tmp = work / "tmp"
    command = [
        "mmseqs",
        "easy-cluster",
        str(fasta),
        str(prefix),
        str(tmp),
        "--min-seq-id",
        str(args.min_seq_id),
        "-c",
        str(args.coverage),
        "--cov-mode",
        "0",
        "--alignment-mode",
        "3",
        "--threads",
        str(args.threads),
    ]
    cluster_tsv = prefix.with_name(prefix.name + "_cluster.tsv")
    if not cluster_tsv.exists():
        subprocess.run(command, check=True)
    protein_to_cluster: dict[str, str] = {}
    with cluster_tsv.open() as handle:
        for line in handle:
            representative, member = line.rstrip("\n").split("\t")
            protein_to_cluster[member] = representative
    missing = {protein_id for protein_id, _ in proteins} - protein_to_cluster.keys()
    if missing:
        raise AssertionError(f"MMseqs2 omitted {len(missing)} proteins")

    enriched = [
        dict(row, protein_similarity_cluster=protein_to_cluster[row["protein_id"]])
        for row in rows
    ]
    assignments = assign_grouped_splits(
        enriched,
        group_key="protein_similarity_cluster",
        seed=args.seed,
        ratios=split_ratios,
    )
    assert_group_disjoint(enriched, assignments, group_key="protein_similarity_cluster")
    manifest_name = f"split-protein_similarity_{int(args.min_seq_id * 100)}.jsonl"
    manifest = args.data_dir / manifest_name
    with manifest.open("w") as handle:
        for row in rows:
            handle.write(
                json.dumps(
                    {"pair_id": row["pair_id"], "split": assignments[row["pair_id"]]}
                )
                + "\n"
            )

    cluster_sizes = Counter(protein_to_cluster.values())
    row_counts = Counter(assignments.values())
    report = {
        "tool": "MMseqs2",
        "command": command,
        "min_sequence_identity": args.min_seq_id,
        "minimum_alignment_coverage": args.coverage,
        "coverage_mode": "both query and target",
        "alignment_mode": "actual identical residues / aligned columns",
        "unique_proteins": len(proteins),
        "clusters": len(cluster_sizes),
        "largest_cluster_proteins": max(cluster_sizes.values()),
        "split_rows": dict(row_counts),
        "manifest": manifest_name,
        "protein_cluster_overlap": 0,
    }
    (args.data_dir / f"protein-clusters-{suffix}.json").write_text(
        json.dumps(report, indent=2) + "\n"
    )
    audit.setdefault("split_manifests", {})[
        f"protein_similarity_{int(args.min_seq_id * 100)}"
    ] = manifest_name
    audit.setdefault("external_split_audits", {})[
        f"protein_similarity_{int(args.min_seq_id * 100)}"
    ] = report
    audit_path.write_text(json.dumps(audit, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()