staplebridge/data/catalog.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Iterable

from staplebridge.data.schemas import BuildingBlock


def default_catalog() -> list[BuildingBlock]:
    """Single CP-Composer-style stapled block: K ↔ D/E lactam at i,i+3 or i,i+4.

    Motif and Cα window come straight from CP-Composer's success criterion
    (`evaluate_utils/success_utils.ipynb`):
        - peptide[i] == 'K' and peptide[i+3] in {'D','E'} or peptide[i+4] in {'D','E'}
        - 4.0 ≤ Cα(i)-Cα(j) ≤ 6.5 Å
    """
    return [
        BuildingBlock(
            block_id="STAPLE_LACTAM",
            name="K-(D/E) lactam staple",
            chemistry_class="stapled",
            synthetic_accessibility_score=0.85,
            cost_score=0.5,
            spps_score=0.7,
            motif={"i_aa": ["K"], "j_aa": ["D", "E"], "spacings": [3, 4]},
            ca_window=(4.0, 6.5),
        ),
    ]


def save_catalog(blocks: Iterable[BuildingBlock], out_path: str | Path) -> None:
    path = Path(out_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump([b.to_dict() for b in blocks], f, indent=2)


def load_catalog(path: str | Path) -> list[BuildingBlock]:
    with Path(path).open("r", encoding="utf-8") as f:
        data = json.load(f)
    return [BuildingBlock.from_dict(x) for x in data]


def catalog_index(blocks: Iterable[BuildingBlock]) -> dict[str, BuildingBlock]:
    return {b.block_id: b for b in blocks}