"""Protein-binder-design workflow."""

from enum import Enum
from typing import Annotated

from pydantic import AfterValidator

from ..base import Base, LowercaseStrEnum, round_optional_float
from ..protein import ProteinSequence
from ..types import ProteinUUID
from .workflow import Workflow


class BoltzGenSecondaryStructure(Base):
    """
    Represents the secondary structure assignments for a protein.

    :param id: Optional identifier for this secondary structure annotation.
    :param sheet: String encoding the residue indices comprising β-sheet structures
                  (e.g., "1,3..11" for residues 1, and 3 through 11).
    :param helix: String encoding residue indices comprising helices.
    :param loop: String encoding residue indices comprising loop or coil regions.
    """

    id: str | None = None
    sheet: str | None = None
    helix: str | None = None
    loop: str | None = None


class BoltzGenProteinEntity(Base):
    """
    Represents a protein entity, either a designed or natural sequence.

    :param id: Unique identifier for the protein.
    :param sequence: Protein sequence, may contain amino acids and numbers for designed regions.
    :param secondary_structure: Optional assigned secondary structure.
    :param cyclic: Whether the protein is cyclic (True/False). Optional.
    """

    id: str
    sequence: str  # can include amino acids as well as numbers for designed regions
    secondary_structure: BoltzGenSecondaryStructure | None = None
    # binding_types: BindingType | None = None - we may want to add this later but not used in examples.
    cyclic: bool | None = None


class BoltzGenLigandEntity(Base):
    """
    Represents a ligand entity (non-protein), such as a small molecule.

    :param id: Unique identifier for the ligand.
    :param smiles: SMILES string representation of the ligand.
    """

    id: str
    smiles: str
    # binding_types: str | None = None - we may want to add this later but not used in examples.


class BoltzGenRegionSelection(Base):
    """
    Defines a region of a protein chain by specifying its chain identifier and (optionally) residue indices.

    :param chain_id: Identifier for the protein chain (e.g., 'A', 'B', etc.).
    :param residue_indices: Residues to select, specified as a string in the format "5..7,13" or "5..15,50..".
    """

    chain_id: str | None = None
    residue_indices: str | None = None


class BoltzGenProximityRegionSelection(BoltzGenRegionSelection):
    """
    Defines a region of a protein chain based on spatial proximity to a selection of residues.

    Inherits:
        BoltzGenRegionSelection

    :param radius: Radius in angstroms (Å) used to select all residues within proximity to the specified region.
    """

    radius: int | None = None


class BoltzGenBindingType(Base):
    """
    Represents the binding interface specification for a given protein chain.

    :param chain_id: Identifier for the protein chain (e.g., 'A', 'B', etc.).
    :param binding: Residue indices or regions that are required to participate in binding
        (e.g., "5..7,13" or "all" for the whole chain).
    :param not_binding: Residue indices or regions that should explicitly not participate in binding
        (e.g., "5..7,13" or "all" for excluding the entire chain).
    """

    chain_id: str | None = None
    binding: str | None = None
    not_binding: str | None = None


class BoltzGenSecondaryStructureOptions(str, Enum):
    UNSPECIFIED = "UNSPECIFIED"
    LOOP = "LOOP"
    HELIX = "HELIX"
    SHEET = "SHEET"


class BoltzGenDesignInsertion(Base):
    """
    Represents an insertion site for protein design in a specific chain.

    :param chain_id: Identifier of the chain where the insertion occurs.
    :param residue_index: Position in the chain after which the insertion is to be made.
    :param number_of_residues: Number of residues to insert at the specified site (can be a string pattern).
    :param secondary_structure: Desired secondary structure type for the inserted residues
        ("UNSPECIFIED", "LOOP", "HELIX", or "SHEET"). Optional.
    """

    chain_id: str
    residue_index: int
    number_of_residues: str
    secondary_structure: BoltzGenSecondaryStructureOptions | None = None


class BoltzGenFileEntity(Base):
    """
    Represents a protein structure input and its associated region selection and design specifications
    for the BoltzGen binder design workflow.

    :param uuid: Unique identifier for the protein structure.
    :param include: List of regions to include in the design or analysis context.
    :param exclude: List of regions to explicitly exclude from consideration (e.g., for ignoring noisy/irrelevant regions).
    :param include_proximity: List of regions defined by spatial proximity (e.g., residues within a given radius).
    :param binding_types: List of binding type constraints or permitted interface regions.
    :param design: List of regions that are being subject to design (mutable, allowed to change).
    :param secondary_structure: List of desired or annotated secondary structure definitions for selected regions.
    :param design_insertions: List of new regions to be inserted with specified properties (e.g., insertion sites, structure preferences).
    """

    uuid: ProteinUUID
    include: list[BoltzGenRegionSelection] = []
    exclude: list[BoltzGenRegionSelection] = []
    # fuse: None  - we may want to add this later but not used in examples.
    include_proximity: list[BoltzGenProximityRegionSelection] = []
    binding_types: list[BoltzGenBindingType] = []
    # structure_groups: None - we may want to add this later but not used in examples.
    design: list[BoltzGenRegionSelection] = []
    secondary_structure: list[BoltzGenSecondaryStructure] = []
    design_insertions: list[BoltzGenDesignInsertion] = []


class BoltzGenAtomSpecification(Base):
    """
    Atom specification for a protein chain, used for applying constraints or referencing atoms.

    :param chain_id: Identifier for the protein chain (e.g., "A", "B").
    :param index: Residue index the atom belongs to (integer, 1-based).
    :param atom_name: Name of the atom (e.g., "CA", "N", "O", etc.).
    """

    chain_id: str
    index: int
    atom_name: str


class BoltzGenConstraint(Base):
    """
    Describes a covalent or spatial constraint between two specified atoms in the context of protein design.

    :param atom1: First atom in the constraint.
    :param atom2: Second atom in the constraint.
    """

    atom1: BoltzGenAtomSpecification
    atom2: BoltzGenAtomSpecification


class BoltzGenInput(Base):
    """
    Represents the primary input schema for the boltzgen application.

    :param protein_entities: Protein chains that are designed or targeted for binding.
    :param ligand_entities: Small molecules or other non-protein ligands relevant to the design.
    :param file_entities: 3d protein structures and input settings related to them.
    :param constraints: Covalent bond constraints
    """

    protein_entities: list[BoltzGenProteinEntity] = []
    ligand_entities: list[BoltzGenLigandEntity] = []
    file_entities: list[BoltzGenFileEntity] = []
    constraints: list[BoltzGenConstraint] = []


class BoltzGenScores(Base):
    """
    Compact, interpretable metrics for a designed binder.
    ↑ higher is better, ↓ lower is better

    :param quality_score: aggregate model quality (↑)
    :param num_filters_passed: number of QC/heuristic filters passed (↑)
    :param iptm: inter-chain pTM confidence, 0–1 (↑)
    :param design_ptm: design pTM confidence, 0–1 (↑)
    :param min_interaction_pae: minimum interface PAE, in Å (↓)
    :param bb_rmsd: backbone RMSD, in Å (↓)
    :param delta_sasa_refolded: ΔSASA of interface after refolding, Å² (↑ typically indicates better burial)
    :param plip_hbonds_refolded: count of hydrogen bonds at the interface (↑)
    :param plip_saltbridge_refolded: count of salt bridges at the interface (↑)
    :param liability_score: composite liabilities score (↓)
    :param liability_high_severity_violations: count of high-severity liabilities (↓)
    :param liability_num_violations: total liability count (↓)
    :param helix: fraction helical content, 0–1
    :param sheet: fraction β-sheet content, 0–1
    :param loop: fraction loop/coil content, 0–1
    :param design_largest_hydrophobic_patch_refolded: largest hydrophobic patch area after refolding, Å²
    :param design_hydrophobicity: overall design hydrophobicity score (unitless)
    :param num_tokens: sequence length / token count
    """

    quality_score: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    num_filters_passed: int | None = None

    iptm: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    design_ptm: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    min_design_to_target_pae: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    design_to_target_iptm: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    min_interaction_pae: Annotated[float | None, AfterValidator(round_optional_float(3))] = None

    bb_rmsd: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    delta_sasa_refolded: Annotated[float | None, AfterValidator(round_optional_float(3))] = None

    plip_hbonds_refolded: int | None = None
    plip_saltbridge_refolded: int | None = None

    liability_score: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    liability_high_severity_violations: int | None = None
    liability_num_violations: int | None = None

    helix: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    sheet: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    loop: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    design_largest_hydrophobic_patch_refolded: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    design_hydrophobicity: Annotated[float | None, AfterValidator(round_optional_float(3))] = None
    num_tokens: int | None = None


class ProteinBinderDesignResult(Base):
    """
    The output; a designed binder.

    :param sequence: sequence
    :param bound_structure: PDB of structure bound to target
    :param scores: scores for generated structure
    """

    binder_sequence: str | None = None  # deprecated
    binder_sequences: list[ProteinSequence] | None = None
    bound_structure: ProteinUUID | None = None
    scores: BoltzGenScores | None = None


class BoltzGenProtocol(LowercaseStrEnum):
    """
    The predefined protocol used for generation + filtering.
    """

    PROTEIN_ANYTHING = "protein-anything"
    PEPTIDE_ANYTHING = "peptide-anything"
    PROTEIN_SMALL_MOLECULE = "protein-small_molecule"
    NANOBODY_ANYTHING = "nanobody-anything"


class BoltzGenSettings(Base):
    """
    The settings for running BoltzGen.

    :param num_designs: how many designs to generate
    :param protocol: which protocol to use
    :param binding_residue: dict mapping chain ID to which residues should bind;
        the string follows the BoltzGen format of specifying ranges of residue indices (refer to their documentation).
        examples include "5..7,13" or "5..15,50..".
    """

    protocol: BoltzGenProtocol = BoltzGenProtocol.PROTEIN_ANYTHING
    num_designs: int = 100
    budget: int = 20


class ProteinBinderDesignWorkflow(Workflow):
    """
    A workflow for generating proteins or peptides that bind to something.


    New:
    :param binder_design_input: input to protein binder design workflow
    :param binder_design_settings: settings for protein generation method employed
    :param generated_binders: output structures
    """

    binder_design_input: BoltzGenInput = BoltzGenInput()
    binder_design_settings: BoltzGenSettings = BoltzGenSettings()

    generated_binders: list[ProteinBinderDesignResult] = []
