"""Species-aware diagnostics for multi-species PPI benchmarks.

Pooling several organisms adds a shortcut the single-species framework cannot see.
Positives are (almost always) within-species, while constructed negatives that
substitute a protein from the pooled set often join two species. A model can then
score well just by recognising taxonomy. This module quantifies that and repeats
the core controls within each species.

Within-species replication follows the protocol of
``github_release/scripts/49_species_stratified_replication.py``:
subset to pairs whose endpoints are both in the species, recompute degree
statistics within the species, build N0/N2/N2-exact/N2-nearest, split positives
and each negative set 80/20 with ``random_state=seed``, then fit degree-only LR and
sequence-representation LR/MLP.
"""
from __future__ import annotations

import warnings

import numpy as np
import pandas as pd
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

from .features import pair_matrix, scalar_pair_matrix
from .models import build_model, classification_metrics
from .negatives import (balance_row, compute_protein_stats, endpoint_values, known_pair_set, sample_exact_degree_negatives,
                        sample_matched_negatives, sample_nearest_degree_negatives, sample_random_negatives)

UNASSIGNED = "unassigned"
WITHIN_VARIANTS = ("N0", "N2", "N2-exact", "N2-nearest")


def species_series(proteins: pd.DataFrame, species_map: dict[str, str]) -> pd.Series:
    return proteins["protein_id"].map(species_map).fillna(UNASSIGNED).set_axis(proteins["protein_id"].to_numpy())


def pair_species(pairs: pd.DataFrame, sp: pd.Series) -> tuple[np.ndarray, np.ndarray]:
    return pairs["protein_a"].map(sp).to_numpy(), pairs["protein_b"].map(sp).to_numpy()


def cross_species_row(name: str, pairs: pd.DataFrame, sp: pd.Series) -> dict:
    a, b = pair_species(pairs, sp)
    assigned = (a != UNASSIGNED) & (b != UNASSIGNED)
    cross = assigned & (a != b)
    return {"set": name, "n_pairs": int(len(pairs)), "n_assigned": int(assigned.sum()),
            "cross_species": int(cross.sum()),
            "cross_species_fraction": float(cross.sum() / assigned.sum()) if assigned.any() else float("nan")}


def composition(proteins: pd.DataFrame, pos: pd.DataFrame, neg0: pd.DataFrame, sp: pd.Series) -> list[dict]:
    rows = []
    counts = sp.value_counts()
    pa, pb = pair_species(pos, sp)
    na, nb = pair_species(neg0, sp)
    for name in counts.index:
        rows.append({"species": name, "proteins": int(counts[name]),
                     "positive_within": int(((pa == name) & (pb == name)).sum()),
                     "negative_within": int(((na == name) & (nb == name)).sum()),
                     "positive_cross": int(((pa == name) ^ (pb == name)).sum()),
                     "negative_cross": int(((na == name) ^ (nb == name)).sum())})
    return sorted(rows, key=lambda r: -r["proteins"])


def same_species_shortcut(test_sets: dict[str, pd.DataFrame], sp: pd.Series) -> list[dict]:
    """ROC-AUC of a score that is 1 when both endpoints share a species, else 0.
    0.5 means taxonomy carries no label information; well above 0.5 means
    negatives can be recognised by their species composition alone."""
    rows = []
    for name, te in test_sets.items():
        a, b = pair_species(te, sp)
        ok = (a != UNASSIGNED) & (b != UNASSIGNED)
        y, s = te["label"].to_numpy()[ok], (a[ok] == b[ok]).astype(float)
        auc = roc_auc_score(y, s) if len(np.unique(y)) == 2 else float("nan")
        rows.append({"evaluated_on": name, "n_evaluated": int(ok.sum()), "n_excluded": int((~ok).sum()), "roc_auc": float(auc)})
    return rows


def species_matched_negatives(pos: pd.DataFrame, stats_tbl: pd.DataFrame, known: set, seed: int, sp: pd.Series,
                              candidate_pool=None) -> pd.DataFrame:
    """N2-species: one-sided degree-decile substitution restricted to the species of the
    replaced protein, so within-species positives yield within-species negatives."""
    st = stats_tbl.copy()
    st["species"] = sp.reindex(st.index).fillna(UNASSIGNED).to_numpy()
    return sample_matched_negatives(pos, st, known, ["species", "degree_decile"], seed, candidate_pool)


def within_species(proteins: pd.DataFrame, pairs: pd.DataFrame, sp: pd.Series, reps: dict[str, pd.DataFrame],
                   seed: int, test_fraction: float, min_positive: int = 50, models=("lr", "mlp"), progress=None) -> dict:
    """Repeat degree-only and sequence controls inside each species."""
    metrics, balances, skipped = [], [], []
    species_list = [s for s in sp.value_counts().index if s != UNASSIGNED]
    for k, name in enumerate(species_list):
        ids = set(sp.index[sp == name])
        prot = proteins[proteins["protein_id"].isin(ids)].reset_index(drop=True)
        sub = pairs[pairs["protein_a"].isin(ids) & pairs["protein_b"].isin(ids)].reset_index(drop=True)
        pos = sub[sub["label"] == 1].reset_index(drop=True)
        neg0 = sub[sub["label"] == 0].reset_index(drop=True)
        if len(pos) < min_positive:
            skipped.append({"species": name, "reason": f"only {len(pos)} within-species positives (< {min_positive})"})
            continue
        if progress:
            progress(k / max(len(species_list), 1), f"Within-species controls: {name}")
        st = compute_protein_stats(prot, pos)
        degree = st["positive_degree"]
        known = known_pair_set(pos)
        supplied = len(neg0) > 0
        variants = {"N0": neg0 if supplied else sample_random_negatives(list(prot["protein_id"]), known, len(pos), seed),
                    "N2": sample_matched_negatives(pos, st, known, ["degree_decile"], seed),
                    "N2-exact": sample_exact_degree_negatives(pos, st, known, seed),
                    "N2-nearest": sample_nearest_degree_negatives(pos, st, known, seed)}
        pt, pv = train_test_split(pos, test_size=test_fraction, random_state=seed)
        pos_deg = endpoint_values(pos, degree)
        for variant, neg in variants.items():
            if len(neg) < 5:
                continue
            balances.append({"species": name, "negative_set": variant, "n_negatives": len(neg),
                             "supplied": supplied or variant != "N0",
                             **balance_row(pos_deg, endpoint_values(neg, degree))})
            nt, nv = train_test_split(neg, test_size=test_fraction, random_state=seed)
            train = pd.concat([pt, nt], ignore_index=True)
            test = pd.concat([pv, nv], ignore_index=True)
            designs = [("Degree only", "lr", lambda df: scalar_pair_matrix(df, degree))]
            for rep_name, feats in reps.items():
                f = feats[feats["protein_id"].isin(ids)]
                designs += [(rep_name, m, (lambda df, f=f: pair_matrix(df, f))) for m in models]
            for rep_name, model, matrix in designs:
                x, y = matrix(train)
                xt, yt = matrix(test)
                scaler = StandardScaler().fit(x)
                clf = build_model(model, seed)
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    clf.fit(scaler.transform(x), y)
                p = clf.predict_proba(scaler.transform(xt))[:, 1]
                metrics.append({"species": name, "negative_set": variant, "representation": rep_name, "model": model,
                                "train_pairs": len(train), "test_pairs": len(test), **classification_metrics(yt, p)})
    return {"metrics": pd.DataFrame(metrics), "balance": pd.DataFrame(balances), "skipped": skipped}