"""Synthetic demonstration dataset with a planted degree bias and a planted sequence rule.

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

from pathlib import Path

import numpy as np
import pandas as pd

from .constants import AMINO_ACIDS


def make_synthetic_dataset(out_dir: Path, n_proteins: int = 600, n_positive: int = 1800, seed: int = 7):
    """Toy benchmark with two properties the framework should detect:

    * a genuine, sequence-encoded interaction rule - proteins belong to hidden
      types; basic (K/R-rich) proteins preferentially bind acidic (D/E-rich)
      ones - so sequence models retain signal after degree matching; and
    * a planted construction bias - the supplied negatives are drawn mainly from
      low-degree proteins, so a degree-only classifier separates them easily.
    """
    rng = np.random.default_rng(seed)
    out_dir.mkdir(parents=True, exist_ok=True)
    base = np.ones(20) / 20
    comp = {0: base.copy(), 1: base.copy(), 2: base.copy()}
    for aa in "KR":
        comp[1][AMINO_ACIDS.index(aa)] *= 3.0
    for aa in "DE":
        comp[2][AMINO_ACIDS.index(aa)] *= 3.0
    comp = {k: v / v.sum() for k, v in comp.items()}
    types = rng.choice([0, 1, 2], size=n_proteins, p=[0.4, 0.3, 0.3])
    lengths = np.clip(rng.lognormal(np.log(380), 0.55, n_proteins).astype(int), 60, 2500)
    ids = [f"SYN{i:04d}" for i in range(n_proteins)]
    seqs = ["".join(rng.choice(list(AMINO_ACIDS), size=L, p=comp[t])) for t, L in zip(types, lengths)]
    proteins = pd.DataFrame({"protein_id": ids, "sequence": seqs})

    activity = rng.pareto(2.5, n_proteins) + 0.3          # heavy-tailed propensity -> degree heterogeneity
    affinity = np.ones((3, 3))
    affinity[1, 2] = affinity[2, 1] = 6.0                   # basic-acidic complementarity
    positives: set = set()
    w = activity / activity.sum()
    while len(positives) < n_positive:
        u = rng.choice(n_proteins, p=w)
        pv = activity * affinity[types[u], types]
        pv[u] = 0
        v = rng.choice(n_proteins, p=pv / pv.sum())
        positives.add((min(ids[u], ids[v]), max(ids[u], ids[v])))
    pos = pd.DataFrame(sorted(positives), columns=["protein_a", "protein_b"]).assign(label=1)

    degree = pd.concat([pos.protein_a, pos.protein_b]).value_counts().reindex(ids, fill_value=0).to_numpy()
    wn = 1.0 / (1.0 + degree) ** 1.5                         # biased negatives: low-degree endpoints
    wn /= wn.sum()
    negatives: set = set()
    while len(negatives) < n_positive:
        a, b = rng.choice(n_proteins, size=2, p=wn)
        key = (min(ids[a], ids[b]), max(ids[a], ids[b]))
        if a != b and key not in positives:
            negatives.add(key)
    neg = pd.DataFrame(sorted(negatives), columns=["protein_a", "protein_b"]).assign(label=0)
    pairs = pd.concat([pos, neg], ignore_index=True).sample(frac=1.0, random_state=seed)
    # present half of the rows in reversed order and add a few duplicates, to exercise the audit
    flip = rng.random(len(pairs)) < 0.5
    pairs.loc[flip, ["protein_a", "protein_b"]] = pairs.loc[flip, ["protein_b", "protein_a"]].values
    pairs = pd.concat([pairs, pairs.sample(12, random_state=seed)], ignore_index=True)

    proteins.to_csv(out_dir / "proteins.tsv", sep="\t", index=False)
    pairs.to_csv(out_dir / "pairs.tsv", sep="\t", index=False)
    pd.DataFrame({"protein_id": ids, "hidden_type": types}).to_csv(out_dir / "hidden_types.tsv", sep="\t", index=False)
    return out_dir / "proteins.tsv", out_dir / "pairs.tsv"