"""Protein representations (AAC+CTD) and symmetric pair features.

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

import numpy as np
import pandas as pd

from .constants import AMINO_ACIDS


CTD_ATTRIBUTES = {
    "hydrophobicity": ("RKEDQN", "GASTPHY", "CLVIMFW"),
    "normalized_vdw_volume": ("GASTPDC", "NVEQIL", "MHKFRYW"),
    "polarity": ("LIFWCMVY", "PATGS", "HQRKNED"),
    "charge": ("KR", "ANCQGHILMFPSTWYV", "DE"),
    "secondary_structure": ("EALMQKRH", "VIYCWFT", "GNPSD"),
    "solvent_accessibility": ("ALFCGIVW", "RKQEND", "MSPTHY"),
    "polarizability": ("GASDT", "CPNVEQIL", "KMHFRYW"),
}
_CTD_INDEX = {attr: {aa: g for g, grp in enumerate(groups, 1) for aa in grp} for attr, groups in CTD_ATTRIBUTES.items()}


def compute_aac(seq: str) -> np.ndarray:
    """20-dim amino-acid composition."""
    return np.array([seq.count(aa) / len(seq) for aa in AMINO_ACIDS], dtype=np.float64)


def compute_ctd(seq: str) -> np.ndarray:
    """147-dim composition/transition/distribution descriptors (7 attributes x 21)."""
    n = len(seq)
    feats = []
    for attr in CTD_ATTRIBUTES:
        coded = np.array([_CTD_INDEX[attr][aa] for aa in seq])
        feats.extend([(coded == g).sum() / n for g in (1, 2, 3)])
        if n < 2:
            feats.extend([0.0, 0.0, 0.0])
        else:
            a, b = coded[:-1], coded[1:]
            for x, y in ((1, 2), (1, 3), (2, 3)):
                feats.append((((a == x) & (b == y)) | ((a == y) & (b == x))).sum() / (n - 1))
        for g in (1, 2, 3):
            pos = np.where(coded == g)[0]
            if len(pos) == 0:
                feats.extend([0.0] * 5)
                continue
            c = len(pos)
            marks = [pos[0], pos[max(int(c * 0.25) - 1, 0)], pos[max(int(c * 0.50) - 1, 0)],
                     pos[max(int(c * 0.75) - 1, 0)], pos[-1]]
            feats.extend([(p + 1) / n for p in marks])
    return np.array(feats, dtype=np.float64)


def aac_ctd_features(proteins: pd.DataFrame) -> pd.DataFrame:
    """AAC+CTD (167 dims). Non-standard residues are removed before computing."""
    rows = []
    for seq in proteins["sequence"]:
        clean = "".join(c for c in seq if c in AMINO_ACIDS) or "A"
        rows.append(np.concatenate([compute_aac(clean), compute_ctd(clean)]))
    out = pd.DataFrame(np.vstack(rows), columns=[f"aacctd_{i}" for i in range(167)])
    out.insert(0, "protein_id", proteins["protein_id"].to_numpy())
    return out


def pair_matrix(pairs: pd.DataFrame, feats: pd.DataFrame) -> tuple[np.ndarray, np.ndarray]:
    """Symmetric fusion [hA+hB, |hA-hB|, hA*hB] so f(A,B) == f(B,A)."""
    table = feats.set_index("protein_id")
    ha = table.loc[pairs["protein_a"]].to_numpy(np.float64)
    hb = table.loc[pairs["protein_b"]].to_numpy(np.float64)
    return np.concatenate([ha + hb, np.abs(ha - hb), ha * hb], axis=1), pairs["label"].to_numpy()


def scalar_pair_matrix(pairs: pd.DataFrame, scalar: pd.Series) -> tuple[np.ndarray, np.ndarray]:
    """Degree-only or length-only control: sum, |diff|, min, max of endpoint values."""
    a = pairs["protein_a"].map(scalar).to_numpy(float)
    b = pairs["protein_b"].map(scalar).to_numpy(float)
    return np.column_stack([a + b, np.abs(a - b), np.minimum(a, b), np.maximum(a, b)]), pairs["label"].to_numpy()


# =============================================================================
# 6. Models and metrics