# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.

from __future__ import annotations

from typing import Dict, Union, Iterable
from typing_extensions import Literal, Required, TypeAlias, TypedDict

from ..._types import SequenceNotStr

__all__ = [
    "StructureAndBindingStartParams",
    "Input",
    "InputEntity",
    "InputEntityProteinEntity",
    "InputEntityProteinEntityModification",
    "InputEntityProteinEntityModificationCcdModification",
    "InputEntityProteinEntityModificationSmilesModification",
    "InputEntityRnaEntity",
    "InputEntityRnaEntityModification",
    "InputEntityRnaEntityModificationCcdModification",
    "InputEntityRnaEntityModificationSmilesModification",
    "InputEntityDnaEntity",
    "InputEntityDnaEntityModification",
    "InputEntityDnaEntityModificationCcdModification",
    "InputEntityDnaEntityModificationSmilesModification",
    "InputEntityLigandCcdEntity",
    "InputEntityLigandSmilesEntity",
    "InputBinding",
    "InputBindingLigandBinding",
    "InputBindingPpiBinding",
    "InputBond",
    "InputBondAtom1",
    "InputBondAtom1LigandAtom",
    "InputBondAtom1PolymerAtom",
    "InputBondAtom2",
    "InputBondAtom2LigandAtom",
    "InputBondAtom2PolymerAtom",
    "InputConstraint",
    "InputConstraintPocketConstraint",
    "InputConstraintContactConstraint",
    "InputConstraintContactConstraintToken1",
    "InputConstraintContactConstraintToken1PolymerContactToken",
    "InputConstraintContactConstraintToken1LigandContactToken",
    "InputConstraintContactConstraintToken2",
    "InputConstraintContactConstraintToken2PolymerContactToken",
    "InputConstraintContactConstraintToken2LigandContactToken",
    "InputModelOptions",
]


class StructureAndBindingStartParams(TypedDict, total=False):
    input: Required[Input]

    model: Required[Literal["boltz-2.1"]]
    """Model to use for prediction"""

    idempotency_key: str
    """Client-provided key to prevent duplicate submissions on retries"""

    workspace_id: str
    """Target workspace ID (admin keys only; ignored for workspace keys)"""


class InputEntityProteinEntityModificationCcdModification(TypedDict, total=False):
    residue_index: Required[int]
    """0-based index of the residue to modify"""

    type: Required[Literal["ccd"]]

    value: Required[str]
    """CCD code from RCSB PDB (e.g.

    'MSE' for selenomethionine, 'SEP' for phosphoserine)
    """


class InputEntityProteinEntityModificationSmilesModification(TypedDict, total=False):
    residue_index: Required[int]
    """0-based index of the residue to modify"""

    type: Required[Literal["smiles"]]

    value: Required[str]
    """SMILES string for the modification"""


InputEntityProteinEntityModification: TypeAlias = Union[
    InputEntityProteinEntityModificationCcdModification, InputEntityProteinEntityModificationSmilesModification
]


class InputEntityProteinEntity(TypedDict, total=False):
    chain_ids: Required[SequenceNotStr[str]]
    """Chain IDs for this entity"""

    modifications: Required[Iterable[InputEntityProteinEntityModification]]
    """Post-translational modifications"""

    type: Required[Literal["protein"]]

    value: Required[str]
    """Amino acid sequence (one-letter codes)"""

    cyclic: bool
    """Whether the sequence is cyclic"""


class InputEntityRnaEntityModificationCcdModification(TypedDict, total=False):
    residue_index: Required[int]
    """0-based index of the residue to modify"""

    type: Required[Literal["ccd"]]

    value: Required[str]
    """CCD code from RCSB PDB (e.g.

    'MSE' for selenomethionine, 'SEP' for phosphoserine)
    """


class InputEntityRnaEntityModificationSmilesModification(TypedDict, total=False):
    residue_index: Required[int]
    """0-based index of the residue to modify"""

    type: Required[Literal["smiles"]]

    value: Required[str]
    """SMILES string for the modification"""


InputEntityRnaEntityModification: TypeAlias = Union[
    InputEntityRnaEntityModificationCcdModification, InputEntityRnaEntityModificationSmilesModification
]


class InputEntityRnaEntity(TypedDict, total=False):
    chain_ids: Required[SequenceNotStr[str]]
    """Chain IDs for this entity"""

    modifications: Required[Iterable[InputEntityRnaEntityModification]]
    """Chemical modifications"""

    type: Required[Literal["rna"]]

    value: Required[str]
    """RNA nucleotide sequence (A, C, G, U, N)"""

    cyclic: bool
    """Whether the sequence is cyclic"""


class InputEntityDnaEntityModificationCcdModification(TypedDict, total=False):
    residue_index: Required[int]
    """0-based index of the residue to modify"""

    type: Required[Literal["ccd"]]

    value: Required[str]
    """CCD code from RCSB PDB (e.g.

    'MSE' for selenomethionine, 'SEP' for phosphoserine)
    """


class InputEntityDnaEntityModificationSmilesModification(TypedDict, total=False):
    residue_index: Required[int]
    """0-based index of the residue to modify"""

    type: Required[Literal["smiles"]]

    value: Required[str]
    """SMILES string for the modification"""


InputEntityDnaEntityModification: TypeAlias = Union[
    InputEntityDnaEntityModificationCcdModification, InputEntityDnaEntityModificationSmilesModification
]


class InputEntityDnaEntity(TypedDict, total=False):
    chain_ids: Required[SequenceNotStr[str]]
    """Chain IDs for this entity"""

    modifications: Required[Iterable[InputEntityDnaEntityModification]]
    """Chemical modifications"""

    type: Required[Literal["dna"]]

    value: Required[str]
    """DNA nucleotide sequence (A, C, G, T, N)"""

    cyclic: bool
    """Whether the sequence is cyclic"""


class InputEntityLigandCcdEntity(TypedDict, total=False):
    chain_ids: Required[SequenceNotStr[str]]
    """Chain IDs for this ligand"""

    type: Required[Literal["ligand_ccd"]]

    value: Required[str]
    """CCD code (e.g., ATP, ADP)"""


class InputEntityLigandSmilesEntity(TypedDict, total=False):
    chain_ids: Required[SequenceNotStr[str]]
    """Chain IDs for this ligand"""

    type: Required[Literal["ligand_smiles"]]

    value: Required[str]
    """SMILES string representing the ligand"""


InputEntity: TypeAlias = Union[
    InputEntityProteinEntity,
    InputEntityRnaEntity,
    InputEntityDnaEntity,
    InputEntityLigandCcdEntity,
    InputEntityLigandSmilesEntity,
]


class InputBindingLigandBinding(TypedDict, total=False):
    binder_chain_id: Required[str]
    """
    Chain ID of the ligand binder (must have exactly 1 copy, <50 atoms, and only
    ligands+proteins in entities)
    """

    type: Required[Literal["ligand_binding"]]


class InputBindingPpiBinding(TypedDict, total=False):
    binder_chain_ids: Required[SequenceNotStr[str]]
    """Chain IDs of the protein binders"""

    type: Required[Literal["ppi_binding"]]


InputBinding: TypeAlias = Union[InputBindingLigandBinding, InputBindingPpiBinding]


class InputBondAtom1LigandAtom(TypedDict, total=False):
    atom_name: Required[str]
    """Standardized atom name (verifiable in CIF file on RCSB)"""

    chain_id: Required[str]
    """Chain ID containing the atom"""

    type: Required[Literal["ligand_atom"]]


class InputBondAtom1PolymerAtom(TypedDict, total=False):
    atom_name: Required[str]
    """Standardized atom name (verifiable in CIF file on RCSB)"""

    chain_id: Required[str]
    """Chain ID containing the atom"""

    residue_index: Required[int]
    """0-based residue index"""

    type: Required[Literal["polymer_atom"]]


InputBondAtom1: TypeAlias = Union[InputBondAtom1LigandAtom, InputBondAtom1PolymerAtom]


class InputBondAtom2LigandAtom(TypedDict, total=False):
    atom_name: Required[str]
    """Standardized atom name (verifiable in CIF file on RCSB)"""

    chain_id: Required[str]
    """Chain ID containing the atom"""

    type: Required[Literal["ligand_atom"]]


class InputBondAtom2PolymerAtom(TypedDict, total=False):
    atom_name: Required[str]
    """Standardized atom name (verifiable in CIF file on RCSB)"""

    chain_id: Required[str]
    """Chain ID containing the atom"""

    residue_index: Required[int]
    """0-based residue index"""

    type: Required[Literal["polymer_atom"]]


InputBondAtom2: TypeAlias = Union[InputBondAtom2LigandAtom, InputBondAtom2PolymerAtom]


class InputBond(TypedDict, total=False):
    atom1: Required[InputBondAtom1]

    atom2: Required[InputBondAtom2]


class InputConstraintPocketConstraint(TypedDict, total=False):
    """Constrains the binder to interact with specific pocket residues on the target."""

    binder_chain_id: Required[str]
    """Chain ID of the binder molecule"""

    contact_residues: Required[Dict[str, Iterable[int]]]
    """Binding pocket residues keyed by chain ID.

    Each key is a chain ID (e.g. "A") and the value is an array of 0-indexed residue
    indices that define the pocket on that chain.
    """

    max_distance_angstrom: Required[float]
    """Maximum allowed distance in Angstroms between binder and pocket residues.

    Typical range: 4-8 A.
    """

    type: Required[Literal["pocket"]]

    force: bool
    """Whether to force the constraint"""


class InputConstraintContactConstraintToken1PolymerContactToken(TypedDict, total=False):
    chain_id: Required[str]
    """Chain ID"""

    residue_index: Required[int]
    """0-based residue index"""

    type: Required[Literal["polymer_contact"]]


class InputConstraintContactConstraintToken1LigandContactToken(TypedDict, total=False):
    atom_name: Required[str]
    """Atom name"""

    chain_id: Required[str]
    """Chain ID"""

    type: Required[Literal["ligand_contact"]]


InputConstraintContactConstraintToken1: TypeAlias = Union[
    InputConstraintContactConstraintToken1PolymerContactToken, InputConstraintContactConstraintToken1LigandContactToken
]


class InputConstraintContactConstraintToken2PolymerContactToken(TypedDict, total=False):
    chain_id: Required[str]
    """Chain ID"""

    residue_index: Required[int]
    """0-based residue index"""

    type: Required[Literal["polymer_contact"]]


class InputConstraintContactConstraintToken2LigandContactToken(TypedDict, total=False):
    atom_name: Required[str]
    """Atom name"""

    chain_id: Required[str]
    """Chain ID"""

    type: Required[Literal["ligand_contact"]]


InputConstraintContactConstraintToken2: TypeAlias = Union[
    InputConstraintContactConstraintToken2PolymerContactToken, InputConstraintContactConstraintToken2LigandContactToken
]


class InputConstraintContactConstraint(TypedDict, total=False):
    max_distance_angstrom: Required[float]
    """Maximum distance in Angstroms"""

    token1: Required[InputConstraintContactConstraintToken1]

    token2: Required[InputConstraintContactConstraintToken2]

    type: Required[Literal["contact"]]

    force: bool
    """Whether to force the constraint"""


InputConstraint: TypeAlias = Union[InputConstraintPocketConstraint, InputConstraintContactConstraint]


class InputModelOptions(TypedDict, total=False):
    recycling_steps: int
    """The number of recycling steps to use for prediction. Default is 3."""

    sampling_steps: int
    """The number of sampling steps to use for prediction. Default is 200."""

    step_scale: float
    """Diffusion step scale (temperature).

    Controls sampling diversity — higher values produce more varied structures.
    Default is 1.638.
    """


class Input(TypedDict, total=False):
    entities: Required[Iterable[InputEntity]]
    """Entities (proteins, RNA, DNA, ligands) forming the complex to predict.

    Order determines chain assignment.
    """

    binding: InputBinding

    bonds: Iterable[InputBond]
    """Bond constraints between atoms"""

    constraints: Iterable[InputConstraint]
    """Structural constraints (pocket and contact)"""

    model_options: InputModelOptions

    num_samples: int
    """Number of structure samples to generate"""
