build_cd_dataset.py

"""
Build the CD cross-dock test set.

Replicates the protocol described in DiffBindFR SI §1:
  1. ApoRef subset          – from Zhang et al. (2022)
  2. CASF2016 subset        – 57 proteins × 5 ligand-bound states
  3. Ensemble CDK2/EGFR/FXA – all Holo states from PDB by UniProt ID
  4. DUDE27-HoloEns         – SIENA-searched Holo ensembles for 27 DUD-E targets
  5. GPCR-AF2               – AlphaFold2 GPCR structures vs 66 PDB Holo states

Usage:
    python scripts/build_cd_testset.py \\
        --pdb_dir /path/to/pdb_mirror \\
        --output  data/cd_testset/
"""

import argparse
import os
import json
import csv
from typing import Optional


# ──────────────────────────────────────────────────────────────────────────────
# Constants – Table S1 summary
# ──────────────────────────────────────────────────────────────────────────────

SUBSET_STATS = {
    "Ensemble-CDK2": dict(n_pfam=1, n_apo=34,  n_holo=339, n_pairs=11317),
    "Ensemble-EGFR": dict(n_pfam=1, n_apo=1,   n_holo=72,  n_pairs=67),
    "Ensemble-FXA":  dict(n_pfam=1, n_apo=4,   n_holo=109, n_pairs=436),
    "ApoRef":        dict(n_pfam=32, n_apo=64,  n_holo=293, n_pairs=548),
    "CASF2016":      dict(n_pfam=57, n_apo=338, n_holo=285, n_pairs=1760),
    "DUDE27-HoloEns":dict(n_pfam=27, n_apo=0,   n_holo=93,  n_pairs=268),
    "GPCR-AF2":      dict(n_pfam=18, n_apo=66,  n_holo=66,  n_pairs=66),
}

# DUDE27 target info (Table S10 in paper)
DUDE27_TARGETS = {
    "aces": dict(pdb="1E66", chain="A", uniprot="P04058", apo_pdb=None),
    "akt2": dict(pdb="3D0E", chain="A", uniprot="P31751", apo_pdb=None),
    "bace1":dict(pdb="3L5D", chain="A", uniprot="P56817", apo_pdb=None),
    "hs90a":dict(pdb="1UYG", chain="A", uniprot="P07900", apo_pdb=None),
    "tgfr1":dict(pdb="3HMM", chain="A", uniprot="P36897", apo_pdb=None),
    "tryb1":dict(pdb="2ZEC", chain="A", uniprot="Q15661", apo_pdb=None),
    "try1": dict(pdb="2AYW", chain="A", uniprot="P00760", apo_pdb=None),
    "thrb": dict(pdb="1YPE", chain="H", uniprot="P00734", apo_pdb=None),
    "fabp4":dict(pdb="2NNQ", chain="A", uniprot="P15090", apo_pdb=None),
    "ppard":dict(pdb="2ZNP", chain="A", uniprot="Q03181", apo_pdb=None),
    "pparg":dict(pdb="2GTK", chain="A", uniprot="P37231", apo_pdb=None),
    "fa10": dict(pdb="3KL6", chain="A", uniprot="P00742", apo_pdb=None),
    "cdk2": dict(pdb="1H00", chain="A", uniprot="P24941", apo_pdb=None),
    "met":  dict(pdb="3LQ8", chain="A", uniprot="P08581", apo_pdb=None),
    "mk10": dict(pdb="2ZDT", chain="A", uniprot="P53779", apo_pdb=None),
    "rxra": dict(pdb="1MV9", chain="A", uniprot="P19793", apo_pdb=None),
    "mk14": dict(pdb="2QD9", chain="A", uniprot="Q16539", apo_pdb=None),
    "braf": dict(pdb="3D4Q", chain="A", uniprot="P15056", apo_pdb=None),
    "vgfr2":dict(pdb="2P2I", chain="A", uniprot="P35968", apo_pdb=None),
    "gria2":dict(pdb="3KGC", chain="B", uniprot="P19491", apo_pdb=None),
    "egfr": dict(pdb="2RGP", chain="A", uniprot="P00533", apo_pdb=None),
    "mapk2":dict(pdb="3M2W", chain="A", uniprot="P49137", apo_pdb=None),
    "ital": dict(pdb="2ICA", chain="A", uniprot="P20701", apo_pdb=None),
    "dpp4": dict(pdb="2I78", chain="B", uniprot="P27487", apo_pdb=None),
    "ptn1": dict(pdb="2AZR", chain="A", uniprot="P18031", apo_pdb=None),
    "igf1r":dict(pdb="2OJ9", chain="A", uniprot="P08069", apo_pdb=None),
    "ampc": dict(pdb="1L2S", chain="B", uniprot="P00811", apo_pdb=None),
}

# Ensemble targets (Table S1)
ENSEMBLE_TARGETS = {
    "CDK2": dict(apo_pdb="1FIN", uniprot="P24941"),
    "EGFR": dict(apo_pdb="7A2A", uniprot="P00533"),
    "FXA":  dict(apo_pdb="1EZQ", uniprot="P00742"),
}

# GPCR targets (18 receptors from Karelina et al.)
GPCR_UNIPROTS = [
    "P08908", "P21917", "P28223", "P34969", "P41143",
    "P41594", "P43657", "Q9Y5N1", "P35348", "P28335",
    "P14416", "P28223", "Q9Y5Y4", "P47869", "P30542",
    "Q8TDS5", "P49146", "P10826",
]


# ──────────────────────────────────────────────────────────────────────────────
# PDB download helper
# ──────────────────────────────────────────────────────────────────────────────

def download_pdb(pdb_id: str, out_dir: str) -> Optional[str]:
    """Download a PDB file from RCSB if not already present."""
    import urllib.request
    pdb_id = pdb_id.lower()
    path = os.path.join(out_dir, f"{pdb_id}.pdb")
    if os.path.exists(path):
        return path
    url = f"https://files.rcsb.org/download/{pdb_id.upper()}.pdb"
    try:
        urllib.request.urlretrieve(url, path)
        return path
    except Exception as e:
        print(f"  Failed to download {pdb_id}: {e}")
        return None


# ──────────────────────────────────────────────────────────────────────────────
# Pocket alignment (using Schrödinger align_binding_sites equivalent)
# ──────────────────────────────────────────────────────────────────────────────

def align_binding_sites(holo_path: str, target_path: str, cutoff: float = 5.0) -> Optional[str]:
    """
    Align target protein binding site onto holo structure.

    This wraps BioPython's Superimposer to overlay binding-site Cα atoms.
    In the paper, Schrödinger's align_binding_sites module was used with
    -cutoff 5 -dist 5 parameters.

    Returns path to the aligned target PDB.
    """
    try:
        from Bio.PDB import PDBParser, Superimposer, PDBIO
        parser = PDBParser(QUIET=True)
        holo   = parser.get_structure("holo",   holo_path)
        target = parser.get_structure("target", target_path)

        # Collect pocket Cα atoms (simplified: all Cα within 8 Å of any heteroatom)
        holo_cas   = [a for a in holo.get_atoms()   if a.name == "CA"]
        target_cas = [a for a in target.get_atoms() if a.name == "CA"]

        n = min(len(holo_cas), len(target_cas))
        if n < 3:
            return None

        sup = Superimposer()
        sup.set_atoms(holo_cas[:n], target_cas[:n])
        sup.apply(target.get_atoms())

        out_path = target_path.replace(".pdb", "_aligned.pdb")
        io = PDBIO(); io.set_structure(target); io.save(out_path)
        return out_path
    except Exception as e:
        print(f"  Alignment failed: {e}")
        return None


# ──────────────────────────────────────────────────────────────────────────────
# Per-subset builders
# ──────────────────────────────────────────────────────────────────────────────

def build_dude27_holoens_subset(pdb_dir: str, output_dir: str) -> list:
    """
    Build DUDE27-HoloEns subset: Holo-Holo cross-dock for 27 DUD-E targets.

    For each target, uses the Holo structures from Table S2 of the paper.
    """
    from models.evaluation.metrics import rmsd  # noqa – just for illustration

    # Searched Holo structures from Table S2 (excerpt)
    HOLO_STRUCTURES = {
        "dpp4":  ["2AJ8", "5LLS", "2BUC"],
        "ptn1":  ["8SKL", "7MM1", "7FQU"],
        "aces":  ["7AIS", "4TVK", "6H12", "5EHX", "1GQR"],
        "braf":  ["5ITA", "7M0X", "7P3V", "6P3D", "6N0Q"],
        "vgfr2": ["6GQO", "6XVK"],
        "akt2":  ["3E87", "1O6K", "2UW9", "2JDR"],
        "tgfr1": ["5FRI", "2WOT", "4X0M"],
        "mapk2": ["1NY3", "6T8X", "3KA0"],
        # ... (full list in Table S2)
    }

    pairs = []
    for target, holos in HOLO_STRUCTURES.items():
        info = DUDE27_TARGETS.get(target, {})
        ref_pdb = info.get("pdb", "")
        # Cross-dock: each holo against every other holo
        all_pdbs = [ref_pdb] + holos if ref_pdb else holos
        for i, rec_pdb in enumerate(all_pdbs):
            for j, lig_pdb in enumerate(all_pdbs):
                if i == j:
                    continue
                pairs.append({
                    "subset": "DUDE27-HoloEns",
                    "target": target,
                    "receptor_pdb": rec_pdb.lower(),
                    "ligand_pdb":   lig_pdb.lower(),
                    "type": "Holo-Holo",
                })

    return pairs


def build_casf2016_subset(pdb_dir: str, output_dir: str) -> list:
    """
    Build CASF2016 subset: 57 proteins × 5 Holo states each.
    Cross-dock all pairs (Apo-Holo and Holo-Holo).
    """
    # In practice, the CASF2016 PDB IDs and Apo structures are loaded from
    # the ApoBind database and AHoJ tool. Here we provide the framework.
    pairs = []
    # ... load CASF2016 core set, search ApoBind, generate pairs ...
    print("CASF2016 subset builder: implement with ApoBind/AHoJ API access")
    return pairs


def build_gpcr_af2_subset(pdb_dir: str, output_dir: str) -> list:
    """
    Build GPCR-AF2 subset: AlphaFold2 predicted GPCR structures vs 66 PDB Holo structures.
    AF2 structures predicted with templates ≤ April 30, 2018.
    """
    pairs = []
    # ... predict AF2 structures for 18 GPCRs, then pair with Holo structures ...
    print("GPCR-AF2 subset builder: requires AlphaFold2 access for structure prediction")
    return pairs


# ──────────────────────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────────────────────

def main():
    p = argparse.ArgumentParser("Build CD cross-dock test set")
    p.add_argument("--pdb_dir",  default="data/pdb_mirror")
    p.add_argument("--output",   default="data/cd_testset")
    p.add_argument("--subsets",  nargs="+",
                   default=["DUDE27-HoloEns", "CASF2016", "GPCR-AF2",
                            "Ensemble-CDK2", "Ensemble-EGFR", "Ensemble-FXA"],
                   help="Which subsets to build")
    args = p.parse_args()

    os.makedirs(args.pdb_dir, exist_ok=True)
    os.makedirs(args.output,  exist_ok=True)

    all_pairs = []

    if "DUDE27-HoloEns" in args.subsets:
        print("Building DUDE27-HoloEns subset...")
        pairs = build_dude27_holoens_subset(args.pdb_dir, args.output)
        all_pairs.extend(pairs)
        print(f"  Generated {len(pairs)} cross-dock pairs")

    if "CASF2016" in args.subsets:
        print("Building CASF2016 subset...")
        pairs = build_casf2016_subset(args.pdb_dir, args.output)
        all_pairs.extend(pairs)
        print(f"  Generated {len(pairs)} cross-dock pairs")

    if "GPCR-AF2" in args.subsets:
        print("Building GPCR-AF2 subset...")
        pairs = build_gpcr_af2_subset(args.pdb_dir, args.output)
        all_pairs.extend(pairs)

    # Save pairs manifest
    out_file = os.path.join(args.output, "cd_pairs.json")
    with open(out_file, "w") as f:
        json.dump(all_pairs, f, indent=2)
    print(f"\nSaved {len(all_pairs)} total pairs to {out_file}")

    # Print summary table
    print("\n=== CD Test Set Summary ===")
    print(f"{'Subset':<20} {'Expected':<10} {'Built':<10}")
    print("-" * 42)
    from collections import Counter
    built = Counter(p["subset"] for p in all_pairs)
    for name, stats in SUBSET_STATS.items():
        exp  = stats["n_pairs"]
        got  = built.get(name, 0)
        flag = "✓" if got == exp else "⚠"
        print(f"{flag} {name:<18} {exp:<10} {got:<10}")


if __name__ == "__main__":
    main()