"""Protein-novelty and sequence-cluster splits with leakage gates.

Algorithms copied verbatim from ppi_benchmark_controls.py v1.0.0 (lines 341-456).
Do not edit the algorithms without re-running the golden tests."""
from __future__ import annotations

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split


def partition_protein_pool(proteins: pd.DataFrame, test_fraction: float, seed: int,
                           group_col: str = "sequence_group_id", sort_groups: bool = False):
    """Split proteins into (train_pool, test_pool) at the group level so that no
    identical-sequence group (or sequence cluster) straddles the boundary."""
    group_ids = proteins[group_col].unique()
    if sort_groups:
        group_ids = np.asarray(sorted(group_ids))
    shuffled = np.random.default_rng(seed).permutation(group_ids)
    n_test = int(round(len(shuffled) * test_fraction))
    test_groups, train_groups = set(shuffled[:n_test]), set(shuffled[n_test:])
    train_pool = set(proteins.loc[proteins[group_col].isin(train_groups), "protein_id"])
    test_pool = set(proteins.loc[proteins[group_col].isin(test_groups), "protein_id"])
    return train_pool, test_pool


def _pool_membership(pairs: pd.DataFrame, train_pool: set, test_pool: set) -> pd.Series:
    a_tr, b_tr = pairs["protein_a"].isin(train_pool), pairs["protein_b"].isin(train_pool)
    a_te, b_te = pairs["protein_a"].isin(test_pool), pairs["protein_b"].isin(test_pool)
    cat = pd.Series("unassigned", index=pairs.index)
    cat[a_tr & b_tr] = "both_train"
    cat[a_te & b_te] = "both_test"
    cat[(a_tr & b_te) | (a_te & b_tr)] = "mixed"
    return cat


def _degree_safe_holdout(both_train: pd.DataFrame, fraction: float, seed: int):
    """Hold out R1 pairs without removing the last training occurrence of any protein."""
    degree = pd.concat([both_train["protein_a"], both_train["protein_b"]]).value_counts().to_dict()
    target = int(round(len(both_train) * fraction))
    held = []
    for idx in np.random.default_rng(seed).permutation(both_train.index):
        if len(held) >= target:
            break
        a, b = both_train.at[idx, "protein_a"], both_train.at[idx, "protein_b"]
        if degree[a] > 1 and degree[b] > 1:
            degree[a] -= 1
            degree[b] -= 1
            held.append(idx)
    held = set(held)
    mask = both_train.index.isin(held)
    return both_train.loc[~mask].reset_index(drop=True), both_train.loc[mask].reset_index(drop=True)


def make_splits(pairs: pd.DataFrame, proteins: pd.DataFrame, seed: int, test_fraction: float = 0.2,
                r1_fraction: float = 0.2, clusters: dict | None = None) -> dict:
    """R0 random pairs; R1/R2/R3 from one protein-pool partition sharing one
    training set; optional R3-C from a sequence-cluster partition."""
    splits = {}
    strat = pairs["label"] if pairs["label"].nunique() > 1 else None
    tr, te = train_test_split(pairs, test_size=test_fraction, stratify=strat, random_state=seed)
    splits["R0"] = (tr.reset_index(drop=True), te.reset_index(drop=True))

    train_pool, test_pool = partition_protein_pool(proteins, test_fraction, seed)
    cat = _pool_membership(pairs, train_pool, test_pool)
    both_train = pairs[cat == "both_train"]
    r2 = pairs[cat == "mixed"].reset_index(drop=True)
    r3 = pairs[cat == "both_test"].reset_index(drop=True)
    shared_train, r1 = _degree_safe_holdout(both_train, r1_fraction, seed)
    seen = set(shared_train["protein_a"]) | set(shared_train["protein_b"])
    train_side = r2["protein_a"].where(r2["protein_a"].isin(train_pool), r2["protein_b"])
    r2 = r2[train_side.isin(seen)].reset_index(drop=True)
    splits.update({"R1": (shared_train, r1), "R2": (shared_train, r2), "R3": (shared_train, r3)})

    if clusters:
        prot_c = proteins[["protein_id"]].copy()
        prot_c["cluster_id"] = prot_c["protein_id"].map(clusters).fillna(prot_c["protein_id"])
        tp, sp = partition_protein_pool(prot_c, test_fraction, seed, group_col="cluster_id", sort_groups=True)
        cat_c = _pool_membership(pairs, tp, sp)
        splits["R3-C"] = (pairs[cat_c == "both_train"].reset_index(drop=True),
                          pairs[cat_c == "both_test"].reset_index(drop=True))
    return splits


def _canon(df: pd.DataFrame) -> set:
    return {(a, b) if a < b else (b, a) for a, b in zip(df["protein_a"], df["protein_b"])}


def _proteins_of(df: pd.DataFrame) -> set:
    return set(df["protein_a"]) | set(df["protein_b"])


def leakage_checks(name: str, train: pd.DataFrame, test: pd.DataFrame, id_to_group: dict,
                   id_to_cluster: dict | None = None) -> dict:
    """Pair, reversed-pair, protein, sequence-group (and cluster) overlap gates.
    Raises RuntimeError if a required check fails."""
    need_disjoint = name in {"R3", "R3-C"}
    tr_p, te_p = _proteins_of(train), _proteins_of(test)
    res = {
        "exact_pair_overlap": len(set(zip(train.protein_a, train.protein_b)) & set(zip(test.protein_a, test.protein_b))),
        "reversed_pair_overlap": len(_canon(train) & _canon(test)),
        "protein_overlap": len(tr_p & te_p),
        "sequence_group_overlap": len({id_to_group[p] for p in tr_p} & {id_to_group[p] for p in te_p}),
        "labels_valid": bool(set(train.label) | set(test.label) <= {0, 1}),
    }
    if name == "R3-C" and id_to_cluster is not None:
        res["cluster_overlap"] = len({id_to_cluster.get(p, p) for p in tr_p} & {id_to_cluster.get(p, p) for p in te_p})
    failed = [k for k in ("exact_pair_overlap", "reversed_pair_overlap") if res[k]]
    if need_disjoint:
        failed += [k for k in ("protein_overlap", "sequence_group_overlap", "cluster_overlap") if res.get(k)]
    if name == "R1" and not te_p <= tr_p:
        failed.append("R1_test_protein_not_seen")
    if name == "R2":
        n_seen = test["protein_a"].isin(tr_p).astype(int) + test["protein_b"].isin(tr_p).astype(int)
        if not (n_seen == 1).all():
            failed.append("R2_not_exactly_one_seen")
    if not res["labels_valid"]:
        failed.append("labels_valid")
    res["passed"] = not failed
    if failed:
        raise RuntimeError(f"Split {name} failed leakage checks: {failed}")
    return res


# =============================================================================
# 4. Putative-negative construction and balance diagnostics
# =============================================================================