"""Putative-negative construction and endpoint balance diagnostics.

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

import numpy as np
import pandas as pd
from scipy import stats


def compute_protein_stats(proteins: pd.DataFrame, positive_pairs: pd.DataFrame, n_deciles: int = 10) -> pd.DataFrame:
    """protein_id -> length, positive-network degree, and their decile strata."""
    degree = pd.concat([positive_pairs["protein_a"], positive_pairs["protein_b"]]).value_counts()
    st = proteins[["protein_id", "sequence"]].copy()
    st["length"] = st["sequence"].str.len()
    st["positive_degree"] = st["protein_id"].map(degree).fillna(0).astype(int)
    st["length_decile"] = pd.qcut(st["length"], n_deciles, labels=False, duplicates="drop")
    st["degree_decile"] = pd.qcut(st["positive_degree"].rank(method="first"), n_deciles, labels=False, duplicates="drop")
    return st.drop(columns="sequence").set_index("protein_id")


def known_pair_set(pairs: pd.DataFrame) -> set:
    return set(zip(pairs["protein_a"], pairs["protein_b"])) | set(zip(pairs["protein_b"], pairs["protein_a"]))


def sample_random_negatives(protein_ids: list[str], known: set, n: int, seed: int) -> pd.DataFrame:
    """N1: uniformly random unobserved pairs (no self-pairs, no known positives)."""
    rng = np.random.default_rng(seed)
    ids = np.array(sorted(protein_ids))
    max_pairs = len(ids) * (len(ids) - 1) // 2 - len(known) // 2
    n = min(n, max_pairs)
    chosen: list[tuple[str, str]] = []
    seen: set = set()
    while len(chosen) < n:
        for a, b in rng.choice(ids, size=(max(2 * n, 100), 2)):
            if a == b:
                continue
            key = (a, b) if a < b else (b, a)
            if key in known or key in seen:
                continue
            seen.add(key)
            chosen.append(key)
            if len(chosen) >= n:
                break
    return pd.DataFrame(chosen, columns=["protein_a", "protein_b"]).assign(label=0)


def _substitute(positive_pairs, protein_stats, known, seed, candidate_fn):
    """Shared one-sided substitution loop: for each positive (u, v) keep one
    endpoint, replace the other with a candidate from ``candidate_fn``. A
    positive with no valid candidate yields no negative (never relaxed)."""
    rng = np.random.default_rng(seed)
    negatives, seen, n_skipped = [], set(), 0
    for _, row in positive_pairs.sample(frac=1.0, random_state=seed).iterrows():
        u, v = row["protein_a"], row["protein_b"]
        fixed, to_replace = (v, u) if rng.random() < 0.5 else (u, v)
        decoy = None
        for cand in candidate_fn(to_replace, rng):
            if cand == fixed or cand == to_replace:
                continue
            key = (cand, fixed) if cand < fixed else (fixed, cand)
            if key in known or key in seen:
                continue
            decoy = cand
            break
        if decoy is None:
            n_skipped += 1
            continue
        pair = (decoy, fixed) if decoy < fixed else (fixed, decoy)
        seen.add(pair)
        negatives.append(pair)
    df = pd.DataFrame(negatives, columns=["protein_a", "protein_b"]).assign(label=0)
    df.attrs["n_skipped"] = n_skipped
    return df


def sample_matched_negatives(positive_pairs, protein_stats, known, match_cols, seed, candidate_pool=None):
    """N2 (degree decile), N3 (length decile), N5 (both): one-sided matched substitution."""
    pool = protein_stats.loc[protein_stats.index.isin(candidate_pool)] if candidate_pool is not None else protein_stats
    buckets = pool.groupby(match_cols).apply(lambda g: list(g.index), include_groups=False).to_dict()

    def candidates(p, rng):
        key = tuple(protein_stats.loc[p, match_cols]) if len(match_cols) > 1 else protein_stats.loc[p, match_cols[0]]
        c = list(buckets.get(key, []))
        rng.shuffle(c)
        return c
    return _substitute(positive_pairs, protein_stats, known, seed, candidates)


def sample_exact_degree_negatives(positive_pairs, protein_stats, known, seed, candidate_pool=None):
    """N2-exact: replacement must have exactly the same positive degree."""
    pool = protein_stats.loc[protein_stats.index.isin(candidate_pool)] if candidate_pool is not None else protein_stats
    buckets = pool.groupby("positive_degree").apply(lambda g: list(g.index), include_groups=False).to_dict()

    def candidates(p, rng):
        c = list(buckets.get(protein_stats.loc[p, "positive_degree"], []))
        rng.shuffle(c)
        return c
    return _substitute(positive_pairs, protein_stats, known, seed, candidates)


def sample_nearest_degree_negatives(positive_pairs, protein_stats, known, seed, candidate_pool=None):
    """N2-nearest: replacement with the smallest absolute degree difference
    (random order within each distance tier)."""
    pool = protein_stats.loc[protein_stats.index.isin(candidate_pool)] if candidate_pool is not None else protein_stats
    groups = pool.groupby("positive_degree").apply(lambda g: g.index.to_numpy(), include_groups=False)
    degrees, arrays = groups.index.to_numpy(), groups.to_numpy()

    def candidates(p, rng):
        dist = np.abs(degrees - protein_stats.loc[p, "positive_degree"])
        for gi in np.argsort(dist, kind="stable"):
            grp = arrays[gi]
            for j in rng.permutation(len(grp)):
                yield grp[j]
    return _substitute(positive_pairs, protein_stats, known, seed, candidates)


NEGATIVE_BUILDERS = {
    "N2": lambda pos, st, kn, seed, pool=None: sample_matched_negatives(pos, st, kn, ["degree_decile"], seed, pool),
    "N3": lambda pos, st, kn, seed, pool=None: sample_matched_negatives(pos, st, kn, ["length_decile"], seed, pool),
    "N5": lambda pos, st, kn, seed, pool=None: sample_matched_negatives(pos, st, kn, ["degree_decile", "length_decile"], seed, pool),
    "N2-exact": lambda pos, st, kn, seed, pool=None: sample_exact_degree_negatives(pos, st, kn, seed, pool),
    "N2-nearest": lambda pos, st, kn, seed, pool=None: sample_nearest_degree_negatives(pos, st, kn, seed, pool),
}


def endpoint_values(pairs: pd.DataFrame, scalar: pd.Series) -> np.ndarray:
    """Endpoint-occurrence values (each pair contributes both endpoints)."""
    return np.concatenate([pairs["protein_a"].map(scalar).to_numpy(float), pairs["protein_b"].map(scalar).to_numpy(float)])


def balance_row(pos_vals: np.ndarray, neg_vals: np.ndarray) -> dict:
    na, nb = len(pos_vals), len(neg_vals)
    pooled = np.sqrt(((na - 1) * pos_vals.var(ddof=1) + (nb - 1) * neg_vals.var(ddof=1)) / (na + nb - 2))
    return {"positive_mean": pos_vals.mean(), "negative_mean": neg_vals.mean(),
            "positive_median": np.median(pos_vals), "negative_median": np.median(neg_vals),
            "smd": (pos_vals.mean() - neg_vals.mean()) / pooled if pooled > 0 else 0.0,
            "ks_statistic": stats.ks_2samp(pos_vals, neg_vals).statistic,
            "wasserstein": stats.wasserstein_distance(pos_vals, neg_vals)}


# =============================================================================
# 5. Protein representations and symmetric pair features
# =============================================================================