"""Tautomer prediction workflow."""

from typing import Annotated

from pydantic import AfterValidator

from ..base import Base, round_float, round_optional_float
from ..conformers import ConformerGenSettingsUnion, ETKDGSettings
from ..mode import Mode
from ..settings import Settings
from .multistage_opt import MultiStageOptSettings
from .workflow import DBCalculation, MoleculeWorkflow


class Tautomer(Base):
    """
    A tautomer.

    :param energy: energy of the tautomer
    :param weight: statistical weight of the tautomer
    :param predicted_relative_energy: relative energy of the tautomer
    :param structures: UUIDs of the structures
    """

    energy: Annotated[float, AfterValidator(round_float(6))]
    weight: Annotated[float | None, AfterValidator(round_optional_float(6))] = None
    predicted_relative_energy: Annotated[float | None, AfterValidator(round_optional_float(6))] = None

    structures: list[DBCalculation] = []


class TautomerWorkflow(MoleculeWorkflow):
    """
    A workflow to calculate tautomers.

    Inherited:
    :param initial_molecule: Molecule of interest
    :param mode: Mode for workflow (deprecated)

    New:
    :param conf_gen_settings: conformer-search settings, defaults to old "careful"
    :param multistage_opt_settings: optimization settings

    Results:
    :param tautomers: resulting Tautomers
    """

    conf_gen_settings: ConformerGenSettingsUnion = ETKDGSettings(
        num_initial_confs=250,
        max_confs=20,
        max_mmff_energy=15,
    )

    multistage_opt_settings: MultiStageOptSettings = MultiStageOptSettings(
        mode=Mode.MANUAL,
        optimization_settings=[Settings(method="aimnet2_wb97md3", tasks=["optimize"])],
        singlepoint_settings=Settings(method="aimnet2_wb97md3", tasks=["energy"], solvent_settings={"solvent": "water", "model": "cpcmx"}),
    )

    mode: Mode = Mode.CAREFUL  # deprecated
    tautomers: list[Tautomer] = []
