pCoMole/peptide/constraints.py

from dataclasses import dataclass
from typing import List, Dict, Any, Tuple, Optional
from rdkit import Chem


class DisulfideBuildError(Exception):
    pass

def _clean_peptide_sequence(seq: str) -> str:
    seq = seq.strip().upper().replace(" ", "")
    if not seq:
        raise DisulfideBuildError("Empty sequence.")
    return seq


def _validate_disulfide_pair(seq: str, disulfide_pair: Tuple[int, int]) -> Tuple[int, int]:
    if not isinstance(disulfide_pair, tuple) or len(disulfide_pair) != 2:
        raise DisulfideBuildError("disulfide_pair must be a tuple of two 0-based sequence indices.")

    i, j = disulfide_pair
    n = len(seq)

    if not (0 <= i < n and 0 <= j < n):
        raise DisulfideBuildError(
            f"disulfide_pair {disulfide_pair} is out of range for sequence length {n}."
        )
    if i == j:
        raise DisulfideBuildError("The two cysteine indices must be different.")
    if seq[i] != "C" or seq[j] != "C":
        raise DisulfideBuildError(
            f"disulfide_pair {disulfide_pair} does not point to two cysteines."
        )

    return tuple(sorted((i, j)))


def _get_residue_number_to_sg_atom_idx(mol: Chem.Mol) -> dict:
    residue_to_sg = {}

    for atom in mol.GetAtoms():
        info = atom.GetPDBResidueInfo()
        if info is None:
            continue

        atom_name = info.GetName().strip()
        residue_name = info.GetResidueName().strip()
        residue_number = info.GetResidueNumber()

        if residue_name == "CYS" and atom_name == "SG":
            if residue_number in residue_to_sg:
                raise DisulfideBuildError(
                    f"Found multiple SG atoms for residue number {residue_number}."
                )
            residue_to_sg[residue_number] = atom.GetIdx()

    return residue_to_sg


def _find_ss_bond_idx(mol: Chem.Mol) -> Optional[int]:
    for bond in mol.GetBonds():
        a1 = bond.GetBeginAtom()
        a2 = bond.GetEndAtom()
        if a1.GetSymbol() == "S" and a2.GetSymbol() == "S":
            return bond.GetIdx()
    return None


def _bond_is_in_ring(mol: Chem.Mol, bond_idx: int) -> bool:
    for ring_bond_indices in mol.GetRingInfo().BondRings():
        if bond_idx in ring_bond_indices:
            return True
    return False


def build_peptide_with_disulfide(
    seq: str,
    disulfide_pair: Tuple[int, int],
    require_ring_closure: bool = True,
) -> Chem.Mol:
    """
    Build an oxidized peptide with an explicit disulfide bond between two cysteines.

    Parameters
    ----------
    seq : str
        One-letter amino-acid sequence.
    disulfide_pair : tuple[int, int]
        0-based sequence indices of the two cysteines to connect.
    require_ring_closure : bool
        If True, raise an error unless the constructed S-S bond is part of a ring.

    Returns
    -------
    mol : rdkit.Chem.Mol
        Sanitized RDKit molecule with explicit S-S bond.

    Notes
    -----
    - This uses RDKit's MolFromFASTA, so it is intended for standard amino acids.
    - The function does not generate 3D coordinates; it builds the molecular graph.
    """

    seq = _clean_peptide_sequence(seq)
    i, j = _validate_disulfide_pair(seq, disulfide_pair)

    mol = Chem.MolFromFASTA(seq)
    if mol is None:
        raise DisulfideBuildError("RDKit failed to build peptide from sequence.")

    residue_to_sg = _get_residue_number_to_sg_atom_idx(mol)

    # RDKit residue numbering from MolFromFASTA starts at 1
    res_i = i + 1
    res_j = j + 1

    if res_i not in residue_to_sg:
        raise DisulfideBuildError(f"Could not find SG atom for cysteine at sequence index {i}.")
    if res_j not in residue_to_sg:
        raise DisulfideBuildError(f"Could not find SG atom for cysteine at sequence index {j}.")

    sg_i = residue_to_sg[res_i]
    sg_j = residue_to_sg[res_j]

    existing_bond = mol.GetBondBetweenAtoms(sg_i, sg_j)
    if existing_bond is not None:
        ss_bond_idx = existing_bond.GetIdx()
        if require_ring_closure and not _bond_is_in_ring(mol, ss_bond_idx):
            raise DisulfideBuildError("An S-S bond already exists, but it is not part of a ring.")
        return mol

    rw = Chem.RWMol(mol)
    rw.AddBond(sg_i, sg_j, Chem.BondType.SINGLE)
    out = rw.GetMol()

    try:
        Chem.SanitizeMol(out)
    except Exception as e:
        raise DisulfideBuildError(f"Sanitization failed after adding disulfide bond: {e}")

    ss_bond_idx = _find_ss_bond_idx(out)
    if ss_bond_idx is None:
        raise DisulfideBuildError("Failed to find S-S bond after construction.")

    if require_ring_closure and not _bond_is_in_ring(out, ss_bond_idx):
        raise DisulfideBuildError("Constructed S-S bond exists but does not close a loop.")

    return out

@dataclass
class ConstraintResult:
    feasible: bool
    details: Dict[str, Any]

class DisulfideLoopConstraint:
    def __init__(self, min_separation: int = 1):
        """
        builder_fn(seq, disulfide_pair) -> rdkit.Chem.Mol

        The builder_fn must:
        1. build the peptide from the sequence
        2. explicitly add a disulfide bond between the two cysteines
        3. return the oxidized molecule
        """
        self.builder_fn = build_peptide_with_disulfide
        self.min_separation = min_separation

    def _clean_seq(self, seq: str) -> str:
        return seq.strip().upper().replace(" ", "")

    def _get_cys_positions(self, seq: str) -> List[int]:
        return [i for i, aa in enumerate(seq) if aa == "C"]

    def _find_ss_bond_idx(self, mol):
        for bond in mol.GetBonds():
            a1 = bond.GetBeginAtom()
            a2 = bond.GetEndAtom()
            if a1.GetSymbol() == "S" and a2.GetSymbol() == "S":
                return bond.GetIdx()
        return None

    def _ss_bond_closes_loop(self, mol, ss_bond_idx: int) -> bool:
        ring_info = mol.GetRingInfo()
        bond_rings = ring_info.BondRings()
        for ring in bond_rings:
            if ss_bond_idx in ring:
                return True
        return False

    def evaluate_one(self, seq: str) -> ConstraintResult:
        seq = self._clean_seq(seq)
        cys_positions = self._get_cys_positions(seq)

        if seq[-1] != 'C':
            return ConstraintResult(
                feasible=False,
                details={
                    "reason": "terminal residue is not cysteine",
                    "cys_positions_0based": cys_positions,
                },
            )

        if len(cys_positions) != 2:
            return ConstraintResult(
                feasible=False,
                details={
                    "reason": "sequence must contain exactly two cysteines",
                    "num_cys": len(cys_positions),
                    "cys_positions_0based": cys_positions,
                },
            )

        if cys_positions[1] - cys_positions[0] < self.min_separation:
            return ConstraintResult(
                feasible=False,
                details={
                    "reason": "two cysteines are too close in sequence",
                    "num_cys": 2,
                    "cys_positions_0based": cys_positions,
                },
            )

        try:
            mol = self.builder_fn(seq, disulfide_pair=(cys_positions[0], cys_positions[1]), require_ring_closure=True)
        except Exception as e:
            return ConstraintResult(
                feasible=False,
                details={
                    "reason": "oxidized molecule build failed",
                    "error": str(e),
                    "cys_positions_0based": cys_positions,
                },
            )

        if mol is None:
            return ConstraintResult(
                feasible=False,
                details={
                    "reason": "builder returned None",
                    "cys_positions_0based": cys_positions,
                },
            )

        ss_bond_idx = self._find_ss_bond_idx(mol)
        if ss_bond_idx is None:
            return ConstraintResult(
                feasible=False,
                details={
                    "reason": "no S-S bond found in built molecule",
                    "cys_positions_0based": cys_positions,
                },
            )

        closes_loop = self._ss_bond_closes_loop(mol, ss_bond_idx)
        if not closes_loop:
            return ConstraintResult(
                feasible=False,
                details={
                    "reason": "S-S bond exists but is not part of a ring",
                    "cys_positions_0based": cys_positions,
                    "ss_bond_idx": ss_bond_idx,
                },
            )

        return ConstraintResult(
            feasible=True,
            details={
                "reason": "valid disulfide-closed loop",
                "cys_positions_0based": cys_positions,
                "ss_bond_idx": ss_bond_idx,
            },
        )

    def __call__(self, input_seqs: List[str]) -> List[int]:
        return [1 if self.evaluate_one(seq).feasible else 0 for seq in input_seqs]


class Length:
    def __init__(self, orig_seq):
        self.L0 = len(orig_seq)
        print("Initial Length: ", self.L0)

    def __call__(self, seqs):
        scores = [1 if (len(seq) < self.L0 - 3) else 0 for seq in seqs]

        return scores