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

from typing import Dict, List, Union, Optional
from datetime import datetime
from typing_extensions import Literal, TypeAlias

from pydantic import Field as FieldInfo

from ..._models import BaseModel

__all__ = [
    "StructureAndBindingRetrieveResponse",
    "Error",
    "Input",
    "InputEntity",
    "InputEntityProteinEntityResponse",
    "InputEntityProteinEntityResponseModification",
    "InputEntityProteinEntityResponseModificationCcdModificationResponse",
    "InputEntityProteinEntityResponseModificationSmilesModificationResponse",
    "InputEntityRnaEntityResponse",
    "InputEntityRnaEntityResponseModification",
    "InputEntityRnaEntityResponseModificationCcdModificationResponse",
    "InputEntityRnaEntityResponseModificationSmilesModificationResponse",
    "InputEntityDnaEntityResponse",
    "InputEntityDnaEntityResponseModification",
    "InputEntityDnaEntityResponseModificationCcdModificationResponse",
    "InputEntityDnaEntityResponseModificationSmilesModificationResponse",
    "InputEntityLigandCcdEntityResponse",
    "InputEntityLigandSmilesEntityResponse",
    "InputBinding",
    "InputBindingLigandBindingResponse",
    "InputBindingPpiBindingResponse",
    "InputBond",
    "InputBondAtom1",
    "InputBondAtom1LigandAtomResponse",
    "InputBondAtom1PolymerAtomResponse",
    "InputBondAtom2",
    "InputBondAtom2LigandAtomResponse",
    "InputBondAtom2PolymerAtomResponse",
    "InputConstraint",
    "InputConstraintPocketConstraintResponse",
    "InputConstraintContactConstraintResponse",
    "InputConstraintContactConstraintResponseToken1",
    "InputConstraintContactConstraintResponseToken1PolymerContactTokenResponse",
    "InputConstraintContactConstraintResponseToken1LigandContactTokenResponse",
    "InputConstraintContactConstraintResponseToken2",
    "InputConstraintContactConstraintResponseToken2PolymerContactTokenResponse",
    "InputConstraintContactConstraintResponseToken2LigandContactTokenResponse",
    "InputModelOptions",
    "Output",
    "OutputAllSampleResult",
    "OutputAllSampleResultMetrics",
    "OutputAllSampleResultStructure",
    "OutputBestSample",
    "OutputBestSampleMetrics",
    "OutputBestSampleStructure",
    "OutputArchive",
    "OutputBindingMetrics",
    "OutputBindingMetricsLigandBindingMetrics",
    "OutputBindingMetricsPpiBindingMetrics",
]


class Error(BaseModel):
    """Error details when failed"""

    code: str
    """Machine-readable error code"""

    message: str
    """Human-readable error message"""

    details: Optional[object] = None
    """Additional error details"""


class InputEntityProteinEntityResponseModificationCcdModificationResponse(BaseModel):
    residue_index: int
    """0-based index of the residue to modify"""

    type: Literal["ccd"]

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

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


class InputEntityProteinEntityResponseModificationSmilesModificationResponse(BaseModel):
    residue_index: int
    """0-based index of the residue to modify"""

    type: Literal["smiles"]

    value: str
    """SMILES string for the modification"""


InputEntityProteinEntityResponseModification: TypeAlias = Union[
    InputEntityProteinEntityResponseModificationCcdModificationResponse,
    InputEntityProteinEntityResponseModificationSmilesModificationResponse,
]


class InputEntityProteinEntityResponse(BaseModel):
    chain_ids: List[str]
    """Chain IDs for this entity"""

    modifications: List[InputEntityProteinEntityResponseModification]
    """Post-translational modifications"""

    type: Literal["protein"]

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

    cyclic: Optional[bool] = None
    """Whether the sequence is cyclic"""


class InputEntityRnaEntityResponseModificationCcdModificationResponse(BaseModel):
    residue_index: int
    """0-based index of the residue to modify"""

    type: Literal["ccd"]

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

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


class InputEntityRnaEntityResponseModificationSmilesModificationResponse(BaseModel):
    residue_index: int
    """0-based index of the residue to modify"""

    type: Literal["smiles"]

    value: str
    """SMILES string for the modification"""


InputEntityRnaEntityResponseModification: TypeAlias = Union[
    InputEntityRnaEntityResponseModificationCcdModificationResponse,
    InputEntityRnaEntityResponseModificationSmilesModificationResponse,
]


class InputEntityRnaEntityResponse(BaseModel):
    chain_ids: List[str]
    """Chain IDs for this entity"""

    modifications: List[InputEntityRnaEntityResponseModification]
    """Chemical modifications"""

    type: Literal["rna"]

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

    cyclic: Optional[bool] = None
    """Whether the sequence is cyclic"""


class InputEntityDnaEntityResponseModificationCcdModificationResponse(BaseModel):
    residue_index: int
    """0-based index of the residue to modify"""

    type: Literal["ccd"]

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

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


class InputEntityDnaEntityResponseModificationSmilesModificationResponse(BaseModel):
    residue_index: int
    """0-based index of the residue to modify"""

    type: Literal["smiles"]

    value: str
    """SMILES string for the modification"""


InputEntityDnaEntityResponseModification: TypeAlias = Union[
    InputEntityDnaEntityResponseModificationCcdModificationResponse,
    InputEntityDnaEntityResponseModificationSmilesModificationResponse,
]


class InputEntityDnaEntityResponse(BaseModel):
    chain_ids: List[str]
    """Chain IDs for this entity"""

    modifications: List[InputEntityDnaEntityResponseModification]
    """Chemical modifications"""

    type: Literal["dna"]

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

    cyclic: Optional[bool] = None
    """Whether the sequence is cyclic"""


class InputEntityLigandCcdEntityResponse(BaseModel):
    chain_ids: List[str]
    """Chain IDs for this ligand"""

    type: Literal["ligand_ccd"]

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


class InputEntityLigandSmilesEntityResponse(BaseModel):
    chain_ids: List[str]
    """Chain IDs for this ligand"""

    type: Literal["ligand_smiles"]

    value: str
    """SMILES string representing the ligand"""


InputEntity: TypeAlias = Union[
    InputEntityProteinEntityResponse,
    InputEntityRnaEntityResponse,
    InputEntityDnaEntityResponse,
    InputEntityLigandCcdEntityResponse,
    InputEntityLigandSmilesEntityResponse,
]


class InputBindingLigandBindingResponse(BaseModel):
    binder_chain_id: str
    """
    Chain ID of the ligand binder (must have exactly 1 copy, <50 atoms, and only
    ligands+proteins in entities)
    """

    type: Literal["ligand_binding"]


class InputBindingPpiBindingResponse(BaseModel):
    binder_chain_ids: List[str]
    """Chain IDs of the protein binders"""

    type: Literal["ppi_binding"]


InputBinding: TypeAlias = Union[InputBindingLigandBindingResponse, InputBindingPpiBindingResponse]


class InputBondAtom1LigandAtomResponse(BaseModel):
    atom_name: str
    """Standardized atom name (verifiable in CIF file on RCSB)"""

    chain_id: str
    """Chain ID containing the atom"""

    type: Literal["ligand_atom"]


class InputBondAtom1PolymerAtomResponse(BaseModel):
    atom_name: str
    """Standardized atom name (verifiable in CIF file on RCSB)"""

    chain_id: str
    """Chain ID containing the atom"""

    residue_index: int
    """0-based residue index"""

    type: Literal["polymer_atom"]


InputBondAtom1: TypeAlias = Union[InputBondAtom1LigandAtomResponse, InputBondAtom1PolymerAtomResponse]


class InputBondAtom2LigandAtomResponse(BaseModel):
    atom_name: str
    """Standardized atom name (verifiable in CIF file on RCSB)"""

    chain_id: str
    """Chain ID containing the atom"""

    type: Literal["ligand_atom"]


class InputBondAtom2PolymerAtomResponse(BaseModel):
    atom_name: str
    """Standardized atom name (verifiable in CIF file on RCSB)"""

    chain_id: str
    """Chain ID containing the atom"""

    residue_index: int
    """0-based residue index"""

    type: Literal["polymer_atom"]


InputBondAtom2: TypeAlias = Union[InputBondAtom2LigandAtomResponse, InputBondAtom2PolymerAtomResponse]


class InputBond(BaseModel):
    atom1: InputBondAtom1

    atom2: InputBondAtom2


class InputConstraintPocketConstraintResponse(BaseModel):
    """Constrains the binder to interact with specific pocket residues on the target."""

    binder_chain_id: str
    """Chain ID of the binder molecule"""

    contact_residues: Dict[str, List[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: float
    """Maximum allowed distance in Angstroms between binder and pocket residues.

    Typical range: 4-8 A.
    """

    type: Literal["pocket"]

    force: Optional[bool] = None
    """Whether to force the constraint"""


class InputConstraintContactConstraintResponseToken1PolymerContactTokenResponse(BaseModel):
    chain_id: str
    """Chain ID"""

    residue_index: int
    """0-based residue index"""

    type: Literal["polymer_contact"]


class InputConstraintContactConstraintResponseToken1LigandContactTokenResponse(BaseModel):
    atom_name: str
    """Atom name"""

    chain_id: str
    """Chain ID"""

    type: Literal["ligand_contact"]


InputConstraintContactConstraintResponseToken1: TypeAlias = Union[
    InputConstraintContactConstraintResponseToken1PolymerContactTokenResponse,
    InputConstraintContactConstraintResponseToken1LigandContactTokenResponse,
]


class InputConstraintContactConstraintResponseToken2PolymerContactTokenResponse(BaseModel):
    chain_id: str
    """Chain ID"""

    residue_index: int
    """0-based residue index"""

    type: Literal["polymer_contact"]


class InputConstraintContactConstraintResponseToken2LigandContactTokenResponse(BaseModel):
    atom_name: str
    """Atom name"""

    chain_id: str
    """Chain ID"""

    type: Literal["ligand_contact"]


InputConstraintContactConstraintResponseToken2: TypeAlias = Union[
    InputConstraintContactConstraintResponseToken2PolymerContactTokenResponse,
    InputConstraintContactConstraintResponseToken2LigandContactTokenResponse,
]


class InputConstraintContactConstraintResponse(BaseModel):
    max_distance_angstrom: float
    """Maximum distance in Angstroms"""

    token1: InputConstraintContactConstraintResponseToken1

    token2: InputConstraintContactConstraintResponseToken2

    type: Literal["contact"]

    force: Optional[bool] = None
    """Whether to force the constraint"""


InputConstraint: TypeAlias = Union[InputConstraintPocketConstraintResponse, InputConstraintContactConstraintResponse]


class InputModelOptions(BaseModel):
    recycling_steps: Optional[int] = None
    """The number of recycling steps to use for prediction. Default is 3."""

    sampling_steps: Optional[int] = None
    """The number of sampling steps to use for prediction. Default is 200."""

    step_scale: Optional[float] = None
    """Diffusion step scale (temperature).

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


class Input(BaseModel):
    entities: List[InputEntity]
    """Entities (proteins, RNA, DNA, ligands) forming the complex to predict.

    Order determines chain assignment.
    """

    binding: Optional[InputBinding] = None

    bonds: Optional[List[InputBond]] = None
    """Bond constraints between atoms"""

    constraints: Optional[List[InputConstraint]] = None
    """Structural constraints (pocket and contact)"""

    api_model_options: Optional[InputModelOptions] = FieldInfo(alias="model_options", default=None)

    num_samples: Optional[int] = None
    """Number of structure samples to generate"""


class OutputAllSampleResultMetrics(BaseModel):
    complex_ipde: float
    """Complex interface predicted distance error. Lower is better."""

    complex_iplddt: float
    """Complex interface pLDDT (0-100). Confidence at inter-chain interfaces."""

    complex_pde: float
    """Complex predicted distance error. Lower is better."""

    complex_plddt: float
    """Complex pLDDT (0-100). Per-residue confidence averaged over the complex."""

    iptm: float
    """Interface predicted TM score (0-1). Confidence in domain interfaces."""

    ligand_iptm: float
    """Ligand interface pTM (0-1). Only present when ligands are included."""

    protein_iptm: float
    """Protein-protein interface pTM (0-1). Only present for multi-protein complexes."""

    ptm: float
    """Predicted TM score (0-1). Global structure quality."""

    structure_confidence: float
    """Overall structure confidence (0-1)."""


class OutputAllSampleResultStructure(BaseModel):
    url: str
    """URL to download the file"""

    url_expires_at: datetime
    """When the presigned URL expires"""


class OutputAllSampleResult(BaseModel):
    metrics: OutputAllSampleResultMetrics

    structure: OutputAllSampleResultStructure


class OutputBestSampleMetrics(BaseModel):
    complex_ipde: float
    """Complex interface predicted distance error. Lower is better."""

    complex_iplddt: float
    """Complex interface pLDDT (0-100). Confidence at inter-chain interfaces."""

    complex_pde: float
    """Complex predicted distance error. Lower is better."""

    complex_plddt: float
    """Complex pLDDT (0-100). Per-residue confidence averaged over the complex."""

    iptm: float
    """Interface predicted TM score (0-1). Confidence in domain interfaces."""

    ligand_iptm: float
    """Ligand interface pTM (0-1). Only present when ligands are included."""

    protein_iptm: float
    """Protein-protein interface pTM (0-1). Only present for multi-protein complexes."""

    ptm: float
    """Predicted TM score (0-1). Global structure quality."""

    structure_confidence: float
    """Overall structure confidence (0-1)."""


class OutputBestSampleStructure(BaseModel):
    url: str
    """URL to download the file"""

    url_expires_at: datetime
    """When the presigned URL expires"""


class OutputBestSample(BaseModel):
    metrics: OutputBestSampleMetrics

    structure: OutputBestSampleStructure


class OutputArchive(BaseModel):
    url: str
    """URL to download the file"""

    url_expires_at: datetime
    """When the presigned URL expires"""


class OutputBindingMetricsLigandBindingMetrics(BaseModel):
    binding_confidence: float
    """Confidence that binding occurs (0-1). Primary metric for hit discovery."""

    optimization_score: float
    """Binding strength ranking score for lead optimization.

    Higher values indicate stronger predicted binding.
    """

    type: Literal["ligand_binding_metrics"]


class OutputBindingMetricsPpiBindingMetrics(BaseModel):
    binding_confidence: float
    """Confidence that binding occurs (0-1). Primary metric for hit discovery."""

    type: Literal["ppi_binding_metrics"]


OutputBindingMetrics: TypeAlias = Union[OutputBindingMetricsLigandBindingMetrics, OutputBindingMetricsPpiBindingMetrics]


class Output(BaseModel):
    """Prediction output when succeeded"""

    all_sample_results: List[OutputAllSampleResult]
    """Per-sample structure results"""

    best_sample: OutputBestSample

    archive: Optional[OutputArchive] = None

    binding_metrics: Optional[OutputBindingMetrics] = None


class StructureAndBindingRetrieveResponse(BaseModel):
    id: str
    """Unique prediction identifier"""

    completed_at: Optional[datetime] = None

    created_at: datetime

    data_deleted_at: Optional[datetime] = None
    """When the input/output data was deleted, or null if still available"""

    error: Optional[Error] = None
    """Error details when failed"""

    expires_at: Optional[datetime] = None
    """When this resource and its associated data will be permanently deleted.

    Null while still in progress.
    """

    input: Input

    livemode: bool
    """Whether this resource was created with a live API key."""

    model: Literal["boltz-2.1"]
    """Model used for prediction"""

    output: Optional[Output] = None
    """Prediction output when succeeded"""

    started_at: Optional[datetime] = None

    status: Literal["pending", "running", "succeeded", "failed"]

    version: str
    """Model version used for prediction"""

    workspace_id: str
    """Workspace ID"""

    idempotency_key: Optional[str] = None
    """Client-provided idempotency key"""
