"""Core definitions of Turbomole calculation documents."""

from typing import Any, Dict, List, Union

from atomate2.turbomole.schemas.calc_types import CalcType, RunType, TaskType
from jobflow.utils import ValueEnum
from pydantic import BaseModel, Field
from pymatgen.core.structure import Molecule, Structure


class Status(ValueEnum):
    """Turbomole calculation state."""

    SUCCESS = "successful"
    FAILED = "failed"


class TurbomoleObject(ValueEnum):
    """Types of Turbomole data objects."""

    ENERGY = "energy"


class CalculationInput(BaseModel):
    """Summary of inputs for a Turbomole calculation."""

    structure: Union[Structure, Molecule] = Field(
        None, description="The input structure/molecule object"
    )

    turbomole_input: Dict = Field(
        None, description="The Turbomole input used for this task"
    )

    dft: Dict = Field(
        None,
        description="DFT parameters used in the last calc of this task.",
    )

    turbomole_global: Dict = Field(
        None,
        description="Turbomole global parameters used in the last calc of this task.",
    )


class CalculationOutput(BaseModel):
    """Document defining Turbomole calculation outputs."""

    energy: float = Field(
        None, description="The final total DFT energy for the calculation"
    )
    energy_per_atom: float = Field(
        None, description="The final DFT energy per atom for the calculation"
    )
    structure: Union[Structure, Molecule] = Field(
        None, description="The final structure/molecule from the calculation"
    )
    efermi: float = Field(
        None, description="The Fermi level from the calculation in eV"
    )
    is_metal: bool = Field(None, description="Whether the system is metallic")
    bandgap: float = Field(None, description="The band gap from the calculation in eV")
    cbm: float = Field(
        None,
        description="The conduction band minimum in eV (if system is not metallic)",
    )
    vbm: float = Field(
        None, description="The valence band maximum in eV (if system is not metallic)"
    )
    ionic_steps: List[Dict[str, Any]] = Field(
        None, description="Energy, forces, and structure for each ionic step"
    )


class Calculation(BaseModel):
    """Full Turbomole calculation inputs and outputs."""

    dir_name: str = Field(
        None, description="The directory for this Turbomole calculation"
    )
    turbomole_version: str = Field(
        None, description="Turbomole version used to perform the calculation"
    )
    has_completed: Status = Field(
        None, description="Whether Turbomole completed the calculation successfully"
    )
    input: CalculationInput = Field(
        None, description="Turbomole input settings for the calculation"
    )
    output: CalculationOutput = Field(
        None, description="The Turbomole calculation output"
    )
    completed_at: str = Field(
        None, description="Timestamp for when the calculation was completed"
    )
    task_name: str = Field(
        None, description="Name of task given by custodian (e.g., relax1, relax2)"
    )
    output_file_paths: Dict[str, str] = Field(
        None,
        description="Paths (relative to dir_name) of the Turbomole output files "
        "associated with this calculation",
    )
    # bader: Dict = Field(None, description="Output from the bader software")
    run_type: RunType = Field(
        None, description="Calculation run type (e.g., HF, HSE06, PBE)"
    )
    task_type: TaskType = Field(
        None, description="Calculation task type (e.g., Structure Optimization)."
    )
    calc_type: CalcType = Field(
        None, description="Return calculation type (run type + task_type)."
    )
