"""Dataset audit, canonicalization and positive-graph statistics.

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

import hashlib

import pandas as pd

from .constants import AMINO_ACIDS, PLM_MAX_RESIDUES


def audit_and_canonicalize(proteins: pd.DataFrame, pairs: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame, dict]:
    """Returns (proteins_with_groups, canonical_pairs, audit_report).

    Canonicalization orders each undirected pair lexically (protein_a <
    protein_b) and removes duplicate rows. Fatal problems raise ValueError."""
    report: dict = {"input": {"n_proteins": int(len(proteins)), "n_pairs": int(len(pairs)),
                              "n_positive": int((pairs["label"] == 1).sum()),
                              "n_negative": int((pairs["label"] == 0).sum())}}
    issues = {}
    dup_ids = proteins["protein_id"][proteins["protein_id"].duplicated()].unique().tolist()
    if dup_ids:
        raise ValueError(f"Duplicate protein identifiers in protein table: {dup_ids[:10]}")
    empty = proteins.loc[proteins["sequence"].str.len() == 0, "protein_id"].tolist()
    if empty:
        raise ValueError(f"Proteins with missing sequences: {empty[:10]}")
    bad_labels = sorted(set(pairs["label"].unique()) - {0, 1})
    if bad_labels:
        raise ValueError(f"Labels must be 0 or 1; found {bad_labels}")
    known = set(proteins["protein_id"])
    missing = sorted((set(pairs["protein_a"]) | set(pairs["protein_b"])) - known)
    if missing:
        raise ValueError(f"{len(missing)} pair identifiers have no sequence, e.g. {missing[:10]}")

    nonstandard = proteins["sequence"].apply(lambda s: sum(c not in AMINO_ACIDS for c in s))
    issues["proteins_with_nonstandard_residues"] = int((nonstandard > 0).sum())

    self_pairs = pairs["protein_a"] == pairs["protein_b"]
    issues["self_pairs_removed"] = int(self_pairs.sum())
    work = pairs.loc[~self_pairs].copy()

    swap = work["protein_a"] > work["protein_b"]
    issues["rows_reordered_to_canonical"] = int(swap.sum())
    work.loc[swap, ["protein_a", "protein_b"]] = work.loc[swap, ["protein_b", "protein_a"]].values
    exact_dup = pairs.loc[~self_pairs].duplicated(["protein_a", "protein_b"]).sum()
    all_dup = work.duplicated(["protein_a", "protein_b"]).sum()
    issues["exact_duplicate_rows"] = int(exact_dup)
    issues["reversed_duplicate_rows"] = int(all_dup - exact_dup)

    n_labels = work.groupby(["protein_a", "protein_b"])["label"].nunique()
    conflicts = n_labels[n_labels > 1].index
    issues["conflicting_label_pairs_removed"] = int(len(conflicts))
    if len(conflicts):
        conflict_set = set(conflicts)
        work = work[[(a, b) not in conflict_set for a, b in zip(work["protein_a"], work["protein_b"])]]

    before = len(work)
    work = work.drop_duplicates(subset=["protein_a", "protein_b"], keep="first").reset_index(drop=True)
    issues["duplicate_pairs_removed"] = int(before - len(work))
    work["label"] = work["label"].astype(int)
    report["issues"] = issues

    prot = proteins.copy()
    prot["sequence_hash"] = prot["sequence"].apply(lambda s: hashlib.sha256(s.encode()).hexdigest())
    prot["sequence_group_id"] = prot.groupby("sequence_hash").ngroup()
    group_sizes = prot.groupby("sequence_group_id")["protein_id"].count()
    lengths = prot["sequence"].str.len()

    pos = work[work["label"] == 1]
    degree = pd.concat([pos["protein_a"], pos["protein_b"]]).value_counts().reindex(prot["protein_id"], fill_value=0)
    report["canonical"] = {
        "n_pairs": int(len(work)), "n_positive": int((work["label"] == 1).sum()),
        "n_negative": int((work["label"] == 0).sum()),
        "n_proteins_in_pairs": int(len(set(work["protein_a"]) | set(work["protein_b"]))),
        "identical_sequence_groups_size_gt1": int((group_sizes > 1).sum()),
        "proteins_in_identical_sequence_groups": int(group_sizes[group_sizes > 1].sum()),
        "sequence_length_median": float(lengths.median()),
        f"sequences_longer_than_{PLM_MAX_RESIDUES}": int((lengths > PLM_MAX_RESIDUES).sum()),
        "proteins_with_zero_positive_degree": int((degree == 0).sum()),
    }
    report["positive_graph"] = positive_graph_stats(pos)
    return prot, work, report


def connected_components(edges: pd.DataFrame) -> list[set[str]]:
    """Connected components of an undirected edge list (union-find)."""
    parent: dict[str, str] = {}

    def find(x):
        while parent.setdefault(x, x) != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    for a, b in zip(edges["protein_a"], edges["protein_b"]):
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb
    comps: dict[str, set[str]] = {}
    for node in list(parent):
        comps.setdefault(find(node), set()).add(node)
    return sorted(comps.values(), key=len, reverse=True)


def positive_graph_stats(pos: pd.DataFrame) -> dict:
    if pos.empty:
        return {}
    deg = pd.concat([pos["protein_a"], pos["protein_b"]]).value_counts()
    n, m = len(deg), len(pos)
    comps = connected_components(pos)
    return {"nodes": int(n), "edges": int(m), "density": round(2 * m / (n * (n - 1)), 5),
            "mean_degree": round(float(deg.mean()), 3), "max_degree": int(deg.max()),
            "connected_components": len(comps), "largest_component": len(comps[0])}


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