from typing import Annotated, Self

from pydantic import AfterValidator, model_validator

from ..base import Base, round_float
from ..mode import Mode
from ..settings import Settings
from ..solvent import Solvent, SolventModel, SolventSettings
from ..types import UUID
from .conformer_search import ConformerClusteringSettings, ConformerProperties, ConformerSearchMixin, iMTDSettings
from .multistage_opt import MultiStageOptSettings
from .workflow import MoleculeWorkflow


class SolventDependentConformer(Base):
    """
    Stores a single conformer scored in many different solvents.

    :param calculation: the conformer (as a calculation)
    :param free_energy_by_solvent: the free energy of the solvent in every solvent (in Hartree)
    :param relative_free_energy_by_solvent: the relative free energy vs. the lowest-energy conformer in every solvent (in kcal/mol)
    :param population_by_solvent: the population of the conformer in every solvent (as a number between 0 and 1)
    """

    calculation: UUID
    free_energy_by_solvent: dict[Solvent, Annotated[float, AfterValidator(round_float(3))]]
    relative_free_energy_by_solvent: dict[Solvent, Annotated[float, AfterValidator(round_float(3))]]
    population_by_solvent: dict[Solvent, Annotated[float, AfterValidator(round_float(3))]]


class SolventDependentConformersWorkflow(ConformerSearchMixin, MoleculeWorkflow):
    """
    Conformers are generated through `conf_gen_settings` run in each solvent, and the final ensemble is generated by combining and deduplicating all conformers.

    Clustering is then performed, and optimization is conducted.
    The optimization of the conformers is governed by `multistage_opt_settings`.

    Final conformer scoring is done through a multi-level scheme:
        - A single-point gas-phase energy run through `multistage_opt_settings.sp_settings`.
        - A separate thermal free-energy correction computed via a single-point Hessian (GFN2-xTB).
        - A per-solvent CPCM-X calculation.

    Inherited:
    :param initial_molecule: Molecule of interest
    :param conf_gen_mode: Mode for calculations
    :param conf_gen_settings: settings for conformer generation
    :param mso_mode: Mode for MultiStageOptSettings
    :param multistage_opt_settings: set by mode unless mode=MANUAL (ignores additional settings if set)

    New:
    :param solvents: the solvents to study
    :param conformers: output conformers with per-solvent energies and weights
    :param per_solvent_properties: metrics for how overall distribution changes by solvent
    :param relative_free_energy_by_solvent: how free energy changes by solvent, for predicting ∆G_transfer
    """

    solvents: list[Solvent] = [
        Solvent.HEXANE,
        Solvent.OCTANOL,
        Solvent.CHLOROFORM,
        Solvent.DIMETHYLSULFOXIDE,
        Solvent.WATER,
    ]

    conf_gen_mode: Mode = Mode.MANUAL
    conf_gen_settings: iMTDSettings = iMTDSettings()

    conformer_clustering_settings: ConformerClusteringSettings | None = ConformerClusteringSettings()

    multistage_opt_settings: MultiStageOptSettings = MultiStageOptSettings(
        optimization_settings=[
            Settings(
                method="gfn2_xtb",
                solvent_settings=SolventSettings(
                    solvent=Solvent.WATER,
                    model=SolventModel.ALPB,
                ),
                tasks=["optimize"],
            )
        ],
        sp_settings=Settings(method="g_xtb", tasks=["energy"]),
        mode=Mode.MANUAL,
    )

    conformers: list[SolventDependentConformer] = []
    per_solvent_properties: dict[Solvent, ConformerProperties] = {}
    relative_free_energy_by_solvent: dict[Solvent, Annotated[float, AfterValidator(round_float(3))]] = {}

    @model_validator(mode="after")
    def validate_no_solvent(self) -> Self:
        """Ensure that no solvent is input."""
        if self.conf_gen_settings.solvent_settings:
            raise ValueError("Solvent settings cannot be specified for this workflow!")

        return self
