dataset.py

"""
PDBbind / CD cross-dock dataset loader.

Handles:
  - PDBbind time-split (train/val/test)
  - CD cross-dock test set (Apo-Holo and Holo-Holo pairs)

Each item returns a featurised protein-ligand complex ready for the
DiffBindFR score network.
"""

import os
import math
import json
import random
import numpy as np
import torch
from torch import Tensor
from torch.utils.data import Dataset
from typing import Optional

try:
    from rdkit import Chem
    from rdkit.Chem import AllChem
    RDKIT_AVAILABLE = True
except ImportError:
    RDKIT_AVAILABLE = False


# ──────────────────────────────────────────────────────────────────────────────
# Protein featurisation helpers
# ──────────────────────────────────────────────────────────────────────────────

AA_3TO1 = {
    "ALA": "A", "ARG": "R", "ASN": "N", "ASP": "D", "CYS": "C",
    "GLN": "Q", "GLU": "E", "GLY": "G", "HIS": "H", "ILE": "I",
    "LEU": "L", "LYS": "K", "MET": "M", "PHE": "F", "PRO": "P",
    "SER": "S", "THR": "T", "TRP": "W", "TYR": "Y", "VAL": "V",
}
AA_TYPES = sorted(set(AA_3TO1.values())) + ["X"]  # 21 types
AA_TO_IDX = {aa: i for i, aa in enumerate(AA_TYPES)}


def rbf_encode(d: Tensor, d_min: float = 0., d_max: float = 20., n: int = 16) -> Tensor:
    centers = torch.linspace(d_min, d_max, n, device=d.device)
    sigma = (d_max - d_min) / n
    return torch.exp(-((d.unsqueeze(-1) - centers) ** 2) / (2 * sigma ** 2))


def parse_pocket_from_pdb(pdb_path: str, radius: float = 10.0, ligand_center=None):
    """
    Parse protein pocket residues within `radius` Å of `ligand_center`.

    Returns:
        residues: list of dicts with keys: name, ca_pos, n_pos, c_pos, cb_pos, seq_idx
    """
    from Bio.PDB import PDBParser
    parser = PDBParser(QUIET=True)
    try:
        structure = parser.get_structure("prot", pdb_path)
    except Exception as e:
        return []

    all_residues = []
    for model in structure:
        for chain in model:
            for res in chain:
                if res.id[0] != " ":   # skip HETATM
                    continue
                resname = res.resname.strip()
                aa = AA_3TO1.get(resname, "X")
                try:
                    ca = torch.tensor(res["CA"].get_vector().get_array(), dtype=torch.float)
                except KeyError:
                    continue
                try:
                    n_pos = torch.tensor(res["N"].get_vector().get_array(), dtype=torch.float)
                    c_pos = torch.tensor(res["C"].get_vector().get_array(), dtype=torch.float)
                except KeyError:
                    n_pos = ca.clone()
                    c_pos = ca.clone()
                try:
                    cb = torch.tensor(res["CB"].get_vector().get_array(), dtype=torch.float)
                except KeyError:
                    cb = ca.clone()

                all_residues.append({
                    "aa": aa,
                    "ca": ca,
                    "n":  n_pos,
                    "c":  c_pos,
                    "cb": cb,
                    "chi1": None,
                })

    if ligand_center is None or len(all_residues) == 0:
        return all_residues

    # Filter to pocket
    lc = torch.tensor(ligand_center, dtype=torch.float) if not isinstance(ligand_center, Tensor) else ligand_center
    pocket = [r for r in all_residues if (r["ca"] - lc).norm() < radius]
    return pocket


def featurise_pocket(residues: list) -> dict:
    """
    Convert a list of pocket residues into GVP-ready tensors.

    node_s: [N, 27]   scalar node features (aa-type, backbone angles)
    node_v: [N, 3, 3] vector node features (N→CA, CA→C, CB direction)
    edge_s: [E, 20]   scalar edge features (RBF distance)
    edge_v: [E, 1, 3] vector edge features (unit displacement)
    edge_index: [2, E]
    ca_pos: [N, 3]
    """
    N = len(residues)
    if N == 0:
        return None

    # Node scalar features: one-hot AA (21) + backbone angles sines/cosines (6)
    aa_idx = torch.tensor([AA_TO_IDX.get(r["aa"], 20) for r in residues], dtype=torch.long)
    aa_onehot = torch.nn.functional.one_hot(aa_idx, num_classes=21).float()

    # Backbone dihedral angles (simplified: zeros if not computed)
    dihedrals = torch.zeros(N, 6)  # sin/cos of phi, psi, omega
    node_s = torch.cat([aa_onehot, dihedrals], dim=-1)  # [N, 27]

    # Node vector features: normalised backbone bond directions
    ca = torch.stack([r["ca"] for r in residues])   # [N, 3]
    n  = torch.stack([r["n"]  for r in residues])
    c  = torch.stack([r["c"]  for r in residues])
    cb = torch.stack([r["cb"] for r in residues])

    def safe_norm(v):
        return v / v.norm(dim=-1, keepdim=True).clamp(min=1e-8)

    n_to_ca = safe_norm(ca - n)
    ca_to_c = safe_norm(c - ca)
    ca_to_cb = safe_norm(cb - ca)
    node_v = torch.stack([n_to_ca, ca_to_c, ca_to_cb], dim=1)  # [N, 3, 3]

    # Build k-NN edges (k=30)
    k = min(30, N - 1) if N > 1 else 0
    dist_mat = torch.cdist(ca, ca)   # [N, N]
    topk = dist_mat.topk(k + 1, dim=1, largest=False)
    src = torch.arange(N).unsqueeze(1).expand(-1, k).reshape(-1)
    dst = topk.indices[:, 1:].reshape(-1)  # exclude self
    edge_index = torch.stack([src, dst])  # [2, E]

    edge_dist = dist_mat[src, dst]  # [E]
    edge_s = rbf_encode(edge_dist)   # [E, 16]

    # Positional encoding extra features
    rel_pos = ca[dst] - ca[src]           # [E, 3]
    edge_s_extra = rbf_encode(edge_dist, n=4)
    edge_s = torch.cat([edge_s, edge_s_extra], dim=-1)  # [E, 20]

    disp = rel_pos / (edge_dist.unsqueeze(-1) + 1e-8)
    edge_v = disp.unsqueeze(1)  # [E, 1, 3]

    return dict(
        node_s=node_s,
        node_v=node_v,
        edge_s=edge_s,
        edge_v=edge_v,
        edge_index=edge_index,
        ca_pos=ca,
        n_residues=N,
    )


# ──────────────────────────────────────────────────────────────────────────────
# Ligand torsion graph
# ──────────────────────────────────────────────────────────────────────────────

def get_torsion_bonds(mol):
    """
    Return list of (i, j, k, l) atom indices defining rotatable bonds.
    A bond (j, k) is rotatable if it is a single bond, not in a ring,
    and neither j nor k is terminal.
    """
    from rdkit.Chem import rdMolTransforms
    rot_bonds = []
    for bond in mol.GetBonds():
        if bond.IsInRing():
            continue
        if bond.GetBondTypeAsDouble() != 1.0:
            continue
        j = bond.GetBeginAtomIdx()
        k = bond.GetEndAtomIdx()
        # Need at least one neighbour on each side (excluding j↔k)
        j_nbrs = [n.GetIdx() for n in mol.GetAtomWithIdx(j).GetNeighbors() if n.GetIdx() != k]
        k_nbrs = [n.GetIdx() for n in mol.GetAtomWithIdx(k).GetNeighbors() if n.GetIdx() != j]
        if j_nbrs and k_nbrs:
            rot_bonds.append((j_nbrs[0], j, k, k_nbrs[0]))
    return rot_bonds


# ──────────────────────────────────────────────────────────────────────────────
# Dataset
# ──────────────────────────────────────────────────────────────────────────────

class PDBBindDataset(Dataset):
    """
    PDBbind v2020 dataset using the time-split convention.

    Each item corresponds to one protein-ligand complex.
    The protein structure may be the crystal Holo state (for redocking)
    or an Apo/cross-dock structure.
    """

    def __init__(
        self,
        data_dir: str,
        split: str = "train",          # train / val / test
        split_file: Optional[str] = None,
        pocket_radius: float = 10.0,
        max_residues: int = 80,
        max_atoms: int = 100,
        augment: bool = True,
    ):
        self.data_dir = data_dir
        self.pocket_radius = pocket_radius
        self.max_residues = max_residues
        self.max_atoms = max_atoms
        self.augment = augment

        # Load split indices
        if split_file is not None:
            with open(split_file) as f:
                self.ids = json.load(f)[split]
        else:
            # Fallback: list all subdirectories
            self.ids = [
                d for d in os.listdir(data_dir)
                if os.path.isdir(os.path.join(data_dir, d))
            ]
        self.ids = sorted(self.ids)

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, idx: int):
        pdb_id = self.ids[idx]
        base = os.path.join(self.data_dir, pdb_id)

        prot_path = os.path.join(base, f"{pdb_id}_protein.pdb")
        lig_path  = os.path.join(base, f"{pdb_id}_ligand.sdf")

        if not os.path.exists(prot_path) or not os.path.exists(lig_path):
            return None

        # Load ligand
        mol = Chem.SDMolSupplier(lig_path, removeHs=False, sanitize=True)[0]
        if mol is None:
            return None

        try:
            conf = mol.GetConformer()
            lig_pos = torch.tensor(conf.GetPositions(), dtype=torch.float)
        except Exception:
            return None

        ligand_center = lig_pos.mean(dim=0).numpy()

        # Load protein pocket
        residues = parse_pocket_from_pdb(prot_path, self.pocket_radius, ligand_center)
        if len(residues) == 0:
            return None
        if len(residues) > self.max_residues:
            # Keep closest residues
            ca_stack = torch.stack([r["ca"] for r in residues])
            lc = torch.tensor(ligand_center)
            dists = (ca_stack - lc).norm(dim=-1)
            keep = dists.topk(self.max_residues, largest=False).indices.tolist()
            residues = [residues[i] for i in sorted(keep)]

        pocket_feats = featurise_pocket(residues)
        if pocket_feats is None:
            return None

        # Featurise ligand
        from models.ligand_encoder import mol_to_graph
        lig_graph = mol_to_graph(mol)

        # Get torsion bonds
        rot_bonds = get_torsion_bonds(mol)
        if rot_bonds:
            tor_i = torch.tensor([b[1] for b in rot_bonds], dtype=torch.long)
            tor_j = torch.tensor([b[2] for b in rot_bonds], dtype=torch.long)
            tor_edge_index = torch.stack([tor_i, tor_j])  # [2, d]
            # Initial torsion angles from coords
            lig_torsions = torch.zeros(len(rot_bonds))
        else:
            tor_edge_index = torch.zeros(2, 0, dtype=torch.long)
            lig_torsions   = torch.zeros(0)

        return dict(
            pdb_id=pdb_id,
            # Protein
            **{f"prot_{k}": v for k, v in pocket_feats.items()},
            # Ligand
            lig_atom_feats=lig_graph["atom_feats"],
            lig_bond_feats=lig_graph["bond_feats"],
            lig_edge_index=lig_graph["edge_index"],
            lig_pos_crystal=lig_pos,
            lig_pos=lig_pos.clone(),
            # Torsions
            tor_edge_index=tor_edge_index,
            lig_torsions=lig_torsions,
            # Placeholders for sc torsions (populated per-residue)
            sc_torsions=torch.zeros(0),
            sc_residue_idx=torch.zeros(0, dtype=torch.long),
        )


class CDCrossDockDataset(Dataset):
    """
    CD cross-dock test set: pairs of (apo/holo receptor, holo ligand).

    Expects a CSV with columns:
        receptor_pdb, receptor_chain, ligand_pdb, ligand_chain
    or a pre-built JSON file listing all cross-dock pairs.
    """

    def __init__(
        self,
        pairs_file: str,
        pdb_dir: str,
        pocket_radius: float = 10.0,
        max_residues: int = 80,
    ):
        import pandas as pd
        self.pdb_dir = pdb_dir
        self.pocket_radius = pocket_radius
        self.max_residues = max_residues

        if pairs_file.endswith(".json"):
            with open(pairs_file) as f:
                self.pairs = json.load(f)
        else:
            df = pd.read_csv(pairs_file)
            self.pairs = df.to_dict("records")

    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, idx: int):
        pair = self.pairs[idx]
        # Load receptor (apo or other-holo) + crystal ligand
        # Implementation mirrors PDBBindDataset but with explicit pair paths
        # (omitted for brevity – follows same featurization pipeline)
        raise NotImplementedError("Implement path resolution for your CD test set layout")