config.py

"""
Contains the configuration Class for the program.
"""

from __future__ import annotations

from pathlib import Path
from abc import ABC, abstractmethod
from dataclasses import dataclass
import warnings
import multiprocessing as mp
from typing import Mapping, Sequence, Any

import numpy as np
import toml

from mindlessgen.data.constants import PSE_SYMBOLS, PSE_NUMBERS, PSE


# abstract base class for configuration
class BaseConfig(ABC):
    """
    Abstract base class for configuration settings.
    """

    @abstractmethod
    def __init__(self):
        pass

    @abstractmethod
    def get_identifier(self) -> str:
        """
        Get the identifier of the configuration.
        """


def _symbol_to_atomic_number(symbol: str) -> int:
    """
    Convert an element symbol to its atomic number.
    """
    if not isinstance(symbol, str):
        raise TypeError("Element symbol must be a string.")
    normalized = symbol.strip().lower()
    atomic_number = PSE_NUMBERS.get(normalized)
    if atomic_number is None:
        raise ValueError(f"Element '{symbol}' not found in the periodic table.")
    return atomic_number


def _parse_distance(value: Any) -> float:
    """
    Convert a user-provided distance into a validated float.
    """
    if isinstance(value, (float, int)):
        distance = float(value)
    elif isinstance(value, str):
        try:
            distance = float(value)
        except ValueError as exc:
            raise ValueError(
                f"Distance '{value}' could not be parsed as float."
            ) from exc
    else:
        raise TypeError("Distance must be a float-compatible value.")

    if distance <= 0:
        raise ValueError("Distance must be greater than 0.")
    return distance


@dataclass
class DistanceConstraint:
    """
    Representation of an atom-type distance constraint.
    """

    atom_a: int
    atom_b: int
    distance: float

    def __post_init__(self) -> None:
        for attr in ("atom_a", "atom_b"):
            value = getattr(self, attr)
            if not isinstance(value, int):
                raise TypeError("Atomic numbers must be integers.")
            if value <= 0 or value not in PSE:
                raise ValueError("Atomic numbers must reference known elements.")
        if not isinstance(self.distance, (float, int)):
            raise TypeError("Distance must be a float.")
        self.distance = float(self.distance)
        if self.distance <= 0:
            raise ValueError("Distance must be greater than 0.")

    @property
    def element_a(self) -> str:
        """
        Canonical symbol of the first atom type.
        """
        return PSE[self.atom_a]

    @property
    def element_b(self) -> str:
        """
        Canonical symbol of the second atom type.
        """
        return PSE[self.atom_b]

    @property
    def atomic_numbers(self) -> tuple[int, int]:
        """
        Atomic numbers of the constrained atom pair.
        """
        return (self.atom_a, self.atom_b)

    def required_counts(self) -> dict[int, int]:
        """
        Minimum number of atoms required in the molecule for this constraint.
        """
        if self.atom_a == self.atom_b:
            return {self.atom_a: 2}
        return {self.atom_a: 1, self.atom_b: 1}

    def symbol_for(self, atomic_number: int) -> str:
        """
        Return the canonical element symbol for a stored atomic number.
        """
        return PSE[atomic_number]

    @classmethod
    def from_mapping(cls, data: Mapping[str, Any]) -> DistanceConstraint:
        """
        Build a constraint from TOML data.
        """
        if "pair" in data:
            pair = data["pair"]
        elif "elements" in data:
            pair = data["elements"]
        elif all(key in data for key in ("element_a", "element_b")):
            pair = [data["element_a"], data["element_b"]]
        else:
            raise KeyError(
                "Distance constraint requires either 'pair', 'elements', or "
                + "both 'element_a'/'element_b'."
            )
        if not isinstance(pair, Sequence) or len(pair) != 2:
            raise TypeError("Constraint pair must contain two entries.")
        distance = data.get("distance")
        if distance is None:
            raise KeyError("Distance constraint requires a 'distance' value.")
        atom_a = _symbol_to_atomic_number(str(pair[0]))
        atom_b = _symbol_to_atomic_number(str(pair[1]))
        return cls(atom_a, atom_b, _parse_distance(distance))

    @classmethod
    def from_cli_string(cls, spec: str) -> DistanceConstraint:
        """
        Parse a CLI constraint string formatted as 'A,B,distance'.
        """
        if not isinstance(spec, str):
            raise TypeError("Constraint specification must be a string.")
        parts = [part.strip() for part in spec.split(",")]
        if len(parts) != 3:
            raise ValueError(
                "Constraint specification must be formatted as 'ElementA,ElementB,distance'."
            )
        atom_a = _symbol_to_atomic_number(parts[0])
        atom_b = _symbol_to_atomic_number(parts[1])
        return cls(atom_a, atom_b, _parse_distance(parts[2]))

    def __str__(self) -> str:
        return f"{self.element_a}-{self.element_b}: {self.distance:g} Å"


class GeneralConfig(BaseConfig):
    """
    Configuration class for general settings.
    """

    def __init__(self: GeneralConfig) -> None:
        self._verbosity: int = 1
        self._max_cycles: int = 200
        self._print_config: bool = False
        self._parallel: int = 4
        self._num_molecules: int = 1
        self._postprocess: bool = False
        self._write_xyz: bool = True
        self._symmetrization: bool = False
        self._tmp_dir: None | str | Path = None

    def get_identifier(self) -> str:
        return "general"

    @property
    def verbosity(self):
        """
        Get the verbosity level.
        """
        return self._verbosity

    @verbosity.setter
    def verbosity(self, verbosity: int):
        """
        Set the verbosity level.
        """
        if not isinstance(verbosity, int):
            raise TypeError("Verbosity should be an integer.")
        if verbosity not in [-1, 0, 1, 2, 3]:
            raise ValueError("Verbosity can only be -1, 0, 1, 2, or 3.")
        self._verbosity = verbosity

    @property
    def max_cycles(self):
        """
        Get the maximum number of cycles.
        """
        return self._max_cycles

    @max_cycles.setter
    def max_cycles(self, max_cycles: int):
        """
        Set the maximum number of cycles.
        """
        if not isinstance(max_cycles, int):
            raise TypeError("Max cycles should be an integer.")
        if max_cycles < 1:
            raise ValueError("Max cycles should be greater than 0.")
        self._max_cycles = max_cycles

    @property
    def print_config(self):
        """
        Get the print config flag.
        """
        return self._print_config

    @print_config.setter
    def print_config(self, print_config: bool):
        """
        Set the print config flag.
        """
        if not isinstance(print_config, bool):
            raise TypeError("Print config should be a boolean.")
        self._print_config = print_config

    @property
    def parallel(self):
        """
        Get the parallel flag.
        """
        return self._parallel

    @parallel.setter
    def parallel(self, parallel: int):
        """
        Set the parallel flag.
        """
        if not isinstance(parallel, int):
            raise TypeError("Parallel should be an integer.")
        if parallel < 1:
            raise ValueError("Parallel should be greater than 0.")
        self._parallel = parallel

    @property
    def num_molecules(self):
        """
        Get the number of molecules.
        """
        return self._num_molecules

    @num_molecules.setter
    def num_molecules(self, num_molecules: int):
        """
        Set the number of molecules.
        """
        if not isinstance(num_molecules, int):
            raise TypeError("Number of molecules should be an integer.")
        if num_molecules < 1:
            raise ValueError("Number of molecules should be greater than 0.")
        self._num_molecules = num_molecules

    @property
    def postprocess(self):
        """
        Get the postprocess flag.
        """
        return self._postprocess

    @postprocess.setter
    def postprocess(self, postprocess: bool):
        """
        Set the postprocess flag.
        """
        if not isinstance(postprocess, bool):
            raise TypeError("Postprocess should be a boolean.")
        self._postprocess = postprocess

    @property
    def write_xyz(self):
        """
        Get the write xyz flag.
        """
        return self._write_xyz

    @write_xyz.setter
    def write_xyz(self, write_xyz: bool):
        """
        Set the write xyz flag.
        """
        if not isinstance(write_xyz, bool):
            raise TypeError("Write xyz should be a boolean.")
        self._write_xyz = write_xyz

    @property
    def symmetrization(self):
        """
        Get the symmetrization flag.
        """
        return self._symmetrization

    @symmetrization.setter
    def symmetrization(self, symmetrization: bool):
        """
        Set the symmetrization flag.
        """
        if not isinstance(symmetrization, bool):
            raise TypeError("Symmetrization should be a boolean.")
        self._symmetrization = symmetrization

    @property
    def tmp_dir(self):
        """
        Get the temporary directory.
        """
        return self._tmp_dir

    @tmp_dir.setter
    def tmp_dir(self, tmp_dir: str | Path):
        """
        Set the temporary directory.
        """
        if not isinstance(tmp_dir, (str, Path)):
            raise TypeError("Temporary directory should be a string or a Path object.")
        if isinstance(tmp_dir, str):
            if tmp_dir == "" or tmp_dir.lower() == "none":
                self._tmp_dir = None
                return
            tmp_dir = Path(tmp_dir).resolve()
        elif isinstance(tmp_dir, Path):
            tmp_dir = tmp_dir.resolve()
        # raise error if value is a file
        if tmp_dir.is_file():
            raise ValueError("Temporary directory should not be a file.")
        self._tmp_dir = tmp_dir

    def check_config(self, verbosity: int = 1) -> None:
        ### GeneralConfig checks ###
        # lower number of the available cores and the configured parallelism
        num_cores = min(mp.cpu_count(), self.parallel)
        if self.parallel > mp.cpu_count() and verbosity > -1:
            warnings.warn(
                f"Number of cores requested ({self.parallel}) is greater "
                + f"than the number of available cores ({mp.cpu_count()})."
                + f"Using {num_cores} cores instead."
            )

        if num_cores > 1 and verbosity > 0:
            # raise warning that parallelization will disable verbosity
            warnings.warn(
                "Parallelization will disable verbosity during iterative search. "
                + "Set '--verbosity 0' or '-P 1' to avoid this warning, or simply ignore it."
            )

        # Symmetrization without postprocessing
        if self.symmetrization and not self.postprocess:
            if verbosity > 0:
                warnings.warn(
                    "Postprocessing is turned off. The structure will not be relaxed."
                )


class GenerateConfig(BaseConfig):
    """
    Configuration for the "generate" section responsible for setting up an initial Molecule type.
    """

    def __init__(self: GenerateConfig) -> None:
        self._min_num_atoms: int = 5
        self._max_num_atoms: int = 10
        self._init_coord_scaling: float = 3.0
        self._increase_scaling_factor: float = 1.1
        self._element_composition: dict[int, tuple[int | None, int | None]] = {}
        self._forbidden_elements: list[int] | None = None
        self._scale_fragment_detection: float = 1.25
        self._scale_minimal_distance: float = 0.8
        self._contract_coords: bool = True
        self._molecular_charge: int | None = None
        self._fixed_composition: bool = False

    def get_identifier(self) -> str:
        return "generate"

    @property
    def min_num_atoms(self):
        """
        Get the minimum number of atoms.
        """
        return self._min_num_atoms

    @min_num_atoms.setter
    def min_num_atoms(self, min_num_atoms: int):
        """
        Set the minimum number of atoms.
        """
        if not isinstance(min_num_atoms, int):
            raise TypeError("Min num atoms should be an integer.")
        if min_num_atoms < 1:
            raise ValueError("Min num atoms should be greater than 0.")
        self._min_num_atoms = min_num_atoms

    @property
    def max_num_atoms(self):
        """
        Get the maximum number of atoms.
        """
        return self._max_num_atoms

    @max_num_atoms.setter
    def max_num_atoms(self, max_num_atoms: int):
        """
        Set the maximum number of atoms.
        """
        if not isinstance(max_num_atoms, int):
            raise TypeError("Max num atoms should be an integer.")
        if max_num_atoms < 1:
            raise ValueError("Max num atoms should be greater than 0.")
        self._max_num_atoms = max_num_atoms

    @property
    def init_coord_scaling(self):
        """
        Get the initial coordinate scaling.
        """
        return self._init_coord_scaling

    @init_coord_scaling.setter
    def init_coord_scaling(self, init_coord_scaling: float):
        """
        Set the initial coordinate scaling.
        """
        if not isinstance(init_coord_scaling, float):
            raise TypeError("Initial coordinate scaling should be a float.")
        if init_coord_scaling <= 0:
            raise ValueError("Initial coordinate scaling should be greater than 0.")
        self._init_coord_scaling = init_coord_scaling

    @property
    def increase_scaling_factor(self):
        """
        Get the increase scaling factor.
        """
        return self._increase_scaling_factor

    @increase_scaling_factor.setter
    def increase_scaling_factor(self, increase_scaling_factor: float):
        """
        Set the increase scaling factor.
        """
        if not isinstance(increase_scaling_factor, float):
            raise TypeError("Increase scaling factor should be a float.")
        if increase_scaling_factor <= 1:
            raise ValueError("Increase scaling factor should be greater than 1.")
        self._increase_scaling_factor = increase_scaling_factor

    @property
    def element_composition(self):
        """
        Return the element composition.
        """
        return self._element_composition

    @element_composition.setter
    def element_composition(
        self, composition: None | str | dict[int | str, tuple[int | None, int | None]]
    ) -> None:
        """
        If composition: str:
            Parses the element_composition string and stores the parsed data
            in the _element_composition dictionary.
            Format: "C:2-10, H:10-20, O:1-5, N:1-*"
        If composition: dict:
            Should be a dictionary with integer/string keys and tuple values. Will be stored as is.

        Arguments:
            composition_str (str): String with the element composition
            composition_str (dict): Dictionary with integer/str keys and tuple values
        Raises:
            TypeError: If composition_str is not a string or a dictionary
            AttributeError: If the element is not found in the periodic table
            ValueError: If the minimum count is larger than the maximum count
        Returns:
            None
        """

        if not composition:
            return

        # Will return if composition dict does not contain either int or str keys and tuple[int | None, int | None] values
        # Will also return if dict is valid after setting property
        if isinstance(composition, dict):
            tmp = {}

            # Check validity and also convert str keys into atomic numbers
            for key, value in composition.items():
                if (
                    not (isinstance(key, int) or isinstance(key, str))
                    or not isinstance(value, tuple)
                    or len(value) != 2
                    or not all(isinstance(val, int) or val is None for val in value)
                ):
                    raise TypeError(
                        "Element composition dictionary should be a dictionary with either integer or string keys and tuple values (int, int)."
                    )

                # Convert str keys
                if isinstance(key, str):
                    element_number = PSE_NUMBERS.get(key.lower(), None)
                    if element_number is None:
                        raise KeyError(
                            f"Element {key} not found in the periodic table."
                        )
                    tmp[element_number - 1] = composition[key]
                # Check int keys
                else:
                    if key + 1 in PSE_SYMBOLS:
                        tmp[key] = composition[key]
                    else:
                        raise KeyError(
                            f"Element with atomic number {key + 1} (provided key: {key}) not found in the periodic table."
                        )
            self._element_composition = tmp
            return

        if not isinstance(composition, str):
            raise TypeError(
                "Element composition should be a string (will be parsed) or "
                + "a dictionary with integer/string keys and tuple values."
            )

        # Parsing composition string
        element_dict: dict[int, tuple[int | None, int | None]] = {}
        elements = composition.split(",")
        # remove leading and trailing whitespaces
        elements = [element.strip() for element in elements]

        min_count: int | str | None
        max_count: int | str | None
        for element in elements:
            element_type, range_str = element.split(":")
            min_count, max_count = range_str.split("-")
            element_number = PSE_NUMBERS.get(element_type.lower(), None)
            if element_number is None:
                raise AttributeError(
                    f"Element {element_type} not found in the periodic table."
                )
            # correct for 1- vs. 0-based indexing
            element_number = element_number - 1

            # Convert counts, handle wildcard '*'
            min_count = None if min_count == "*" else int(min_count)
            max_count = None if max_count == "*" else int(max_count)
            if (
                min_count is not None
                and max_count is not None
                and min_count > max_count
            ):
                raise ValueError(
                    f"Minimum count ({min_count}) is larger than maximum count ({max_count})."
                )

            element_dict[element_number] = (min_count, max_count)

        self._element_composition = element_dict

    @property
    def forbidden_elements(self):
        """
        Return sorted list of forbidden elements.
        """
        return self._forbidden_elements

    @forbidden_elements.setter
    def forbidden_elements(
        self: GenerateConfig, forbidden: None | str | list[int]
    ) -> None:
        """
        If forbidden: str:
            Parses the forbidden_elements string and stores the parsed data
            in the _forbidden_elements set.
            Format: "57-71, 8, 1" or "19-*"
        If forbidden: list:
            Stores the forbidden elements as is.

        Arguments:
            forbidden (str): String with the forbidden elements
            forbidden (list): List with integer values
        Raises:
            TypeError: If forbidden is not a string or a list of integers
            ValueError: If both start and end are wildcard '*'
        Returns:
            None
        """
        # if string is empty or None, set to None
        if not forbidden:
            self._forbidden_elements = None
            return
        if isinstance(forbidden, list):
            if all(isinstance(elem, int) for elem in forbidden):
                self._forbidden_elements = sorted(forbidden)
                return
            raise TypeError("Forbidden elements should be a list of integers.")
        if not isinstance(forbidden, str):
            raise TypeError(
                "Forbidden elements should be a string or a list of integers."
            )
        forbidden_set: set[int] = set()
        elements = forbidden.split(",")
        elements = [element.strip() for element in elements]

        for item in elements:
            if "-" in item:
                start: str | int
                end: str | int
                start, end = item.split("-")
                if end == "*" and start == "*":
                    raise ValueError("Both start and end cannot be wildcard '*'.")
                if end == "*":
                    end = 103  # Set to a the maximum atomic number in mindlessgen
                if start == "*":
                    start = 0
                forbidden_set.update(
                    range(int(start) - 1, int(end))
                )  # Subtract 1 to convert to 0-based indexing
            else:
                forbidden_set.add(
                    int(item) - 1
                )  # Subtract 1 to convert to 0-based indexing

        self._forbidden_elements = sorted(list(forbidden_set))

    @property
    def scale_fragment_detection(self):
        """
        Get the scaling factor for the fracment detection based on the van der Waals radii.
        """
        return self._scale_fragment_detection

    @scale_fragment_detection.setter
    def scale_fragment_detection(self, scale_fragment_detection: float):
        """
        Set the scaling factor for van der Waals radii.
        """
        if not isinstance(scale_fragment_detection, float):
            raise TypeError("Scale van der Waals radii should be a float.")
        if scale_fragment_detection <= 0:
            raise ValueError("Scale van der Waals radii should be greater than 0.")
        self._scale_fragment_detection = scale_fragment_detection

    @property
    def scale_minimal_distance(self):
        """
        Get the scaling factor for minimal distance between two atoms.
        """
        return self._scale_minimal_distance

    @scale_minimal_distance.setter
    def scale_minimal_distance(self, scale_minimal_distance: float):
        """
        Set the scaling factor for minimal distance between two atoms.
        """
        if not isinstance(scale_minimal_distance, float):
            raise TypeError("Scale minimal distance should be a float.")
        if scale_minimal_distance <= 0:
            raise ValueError("Scale minimal distance should be greater than 0.")
        self._scale_minimal_distance = scale_minimal_distance

    @property
    def contract_coords(self):
        """
        Get the contract_coords flag.
        """
        return self._contract_coords

    @contract_coords.setter
    def contract_coords(self, contract_coords: bool):
        """
        Set the contract_coords flag.
        """
        if not isinstance(contract_coords, bool):
            raise TypeError("Contract coords should be a boolean.")
        self._contract_coords = contract_coords

    @property
    def molecular_charge(self):
        """
        Get the molecular_charge.
        """
        return self._molecular_charge

    @molecular_charge.setter
    def molecular_charge(self, molecular_charge: str | int):
        """
        Set the molecular_charge.
        """
        if isinstance(molecular_charge, str):
            if molecular_charge.lower() == "none" or molecular_charge == "":
                self._molecular_charge = None
            else:
                self._molecular_charge = int(molecular_charge)
        elif isinstance(molecular_charge, int):
            self._molecular_charge = molecular_charge
        else:
            raise TypeError("Molecular charge should be a string or an integer.")

    @property
    def fixed_composition(self):
        """
        Get the fixed_composition flag.
        """
        return self._fixed_composition

    @fixed_composition.setter
    def fixed_composition(self, fixed_composition: bool):
        """
        Set the fixed_composition flag.
        """
        if not isinstance(fixed_composition, bool):
            raise TypeError("Fixed composition should be a boolean.")
        self._fixed_composition = fixed_composition

    def check_config(self, verbosity: int = 1) -> None:
        """
        GenerateConfig checks for any incompatibilities that are imaginable
        """
        # - Check if the minimum number of atoms is smaller than the maximum number of atoms
        if self.min_num_atoms is not None and self.max_num_atoms is not None:
            if self.min_num_atoms > self.max_num_atoms:
                raise ValueError(
                    "The minimum number of atoms is larger than the maximum number of atoms."
                )
        # - Check if the summed number of minimally required atoms from cfg.element_composition
        #   is larger than the maximum number of atoms
        if self.max_num_atoms is not None:
            if (
                np.sum(
                    [
                        (
                            self.element_composition.get(i, (0, 0))[0]
                            if self.element_composition.get(i, (0, 0))[0] is not None
                            else 0
                        )
                        for i in self.element_composition
                    ]
                )
                > self.max_num_atoms
            ):
                raise ValueError(
                    "The summed number of minimally required atoms "
                    + "from the fixed composition is larger than the maximum number of atoms."
                )
        if self.fixed_composition:
            # - Check if all defintions in cfg.element_composition
            #   are completely fixed (min and max are equal)
            for elem, count_range in self.element_composition.items():
                # check if for all entries: min and max are both not None, and if min and max are equal.
                if (
                    (count_range[0] is None)
                    or (count_range[1] is None)
                    or (count_range[0] != count_range[1])
                ):
                    raise ValueError(
                        f"Element {elem} is not completely fixed in the element composition. "
                        + "Usage together with fixed_composition is not possible."
                    )
            # Check if the summed number of fixed atoms
            # is within the defined overall limits
            sum_fixed_atoms = np.sum(
                [
                    self.element_composition.get(i, (0, 0))[0]
                    for i in self.element_composition
                ]
            )
            if self.min_num_atoms is not None and not (
                self.min_num_atoms <= sum_fixed_atoms
            ):
                raise ValueError(
                    "The summed number of fixed atoms from the fixed composition "
                    + "is not within the range of min_num_atoms and max_num_atoms."
                )


class RefineConfig(BaseConfig):
    """
    Configuration class for refinement settings.
    """

    def __init__(self: RefineConfig) -> None:
        self._max_frag_cycles: int = 10
        self._engine: str = "xtb"
        self._hlgap: float = 0.5
        self._debug: bool = False
        self._ncores: int = 2

    def get_identifier(self) -> str:
        return "refine"

    @property
    def max_frag_cycles(self):
        """
        Get the maximum number of fragment cycles.
        """
        return self._max_frag_cycles

    @max_frag_cycles.setter
    def max_frag_cycles(self, max_frag_cycles: int):
        """
        Set the maximum number of fragment cycles.
        """
        if not isinstance(max_frag_cycles, int):
            raise TypeError("Max fragment cycles should be an integer.")
        if max_frag_cycles < 1:
            raise ValueError("Max fragment cycles should be greater than 0.")
        self._max_frag_cycles = max_frag_cycles

    @property
    def engine(self):
        """
        Get the engine.
        """
        return self._engine

    @engine.setter
    def engine(self, engine: str):
        """
        Set the engine.
        """
        if not isinstance(engine, str):
            raise TypeError("Refinement engine should be a string.")
        if engine not in ["xtb", "orca", "turbomole"]:
            raise ValueError("Refinement engine can only be xtb, orca or turbomole.")
        self._engine = engine

    @property
    def hlgap(self):
        """
        Get the minimum HOM
        """
        return self._hlgap

    @hlgap.setter
    def hlgap(self, hlgap: float):
        """
        Set the minimum HOM
        """
        if not isinstance(hlgap, float):
            raise TypeError("Minimum HL gap should be a float.")
        if hlgap < 0:
            raise ValueError("Minimum HL gap should be greater than 0.")
        self._hlgap = hlgap

    @property
    def debug(self):
        """
        Get the debug flag for refinement.
        """
        return self._debug

    @debug.setter
    def debug(self, debug: bool):
        """
        Set the debug flag for refinement.
        """
        if not isinstance(debug, bool):
            raise TypeError("Debug should be a boolean.")
        self._debug = debug

    @property
    def ncores(self):
        """
        Get the number of cores to be used for geometry optimizations in refinement.
        """
        return self._ncores

    @ncores.setter
    def ncores(self, ncores: int):
        """
        Set the number of cores to be used for geometry optimizations in refinement.
        """
        if not isinstance(ncores, int):
            raise TypeError("Number of cores should be an integer.")
        self._ncores = ncores


class PostProcessConfig(BaseConfig):
    """
    Configuration class for post-processing settings.
    """

    def __init__(self: PostProcessConfig) -> None:
        self._engine: str = "orca"
        self._opt_cycles: int | None = None
        self._optimize: bool = True
        self._debug: bool = False
        self._ncores: int = 4

    def get_identifier(self) -> str:
        return "postprocess"

    @property
    def engine(self):
        """
        Get the postprocess engine.
        """
        return self._engine

    @engine.setter
    def engine(self, engine: str):
        """
        Set the postprocess engine.
        """
        if not isinstance(engine, str):
            raise TypeError("Postprocess engine should be a string.")
        if engine not in ["xtb", "orca", "gxtb", "turbomole"]:
            raise ValueError("Postprocess engine can only be xtb, orca or turbomole.")
        self._engine = engine

    @property
    def optimize(self):
        """
        Get the optimization flag for post-processing.
        """
        return self._optimize

    @optimize.setter
    def optimize(self, optimize: bool):
        """
        Set the optimization flag for post-processing.
        """
        if not isinstance(optimize, bool):
            raise TypeError("Optimize should be a boolean.")
        self._optimize = optimize

    @property
    def opt_cycles(self):
        """
        Get the optimization cycles for post-processing.
        """
        return self._opt_cycles

    @opt_cycles.setter
    def opt_cycles(self, opt_cycles: int):
        """
        Set the optimization cycles for post-processing.
        """
        if not isinstance(opt_cycles, (int, str)):
            raise TypeError("Optimization cycles can only be an integer or a string.")
        if isinstance(opt_cycles, str):
            if opt_cycles.lower() != "none":
                raise ValueError(
                    "Optimization cycles can only be an integer or 'none'."
                )
            self._opt_cycles = None
            return
        if opt_cycles == 0:
            self._opt_cycles = None
            return
        if opt_cycles < 0:
            raise ValueError("Optimization cycles can only be 0 or greater.")
        self._opt_cycles = opt_cycles

    @property
    def debug(self):
        """
        Get the debug flag for post-processing.
        """
        return self._debug

    @debug.setter
    def debug(self, debug: bool):
        """
        Set the debug flag for post-processing.
        """
        if not isinstance(debug, bool):
            raise TypeError("Debug should be a boolean.")
        self._debug = debug

    @property
    def ncores(self):
        """
        Get the number of cores to be used in post-processing.
        """
        return self._ncores

    @ncores.setter
    def ncores(self, ncores: int):
        """
        Set the number of cores to be used in post-processing.
        """
        if not isinstance(ncores, int):
            raise TypeError("Number of cores should be an integer.")
        self._ncores = ncores


class XTBConfig(BaseConfig):
    """
    Configuration class for XTB.
    """

    def __init__(self: XTBConfig) -> None:
        self._xtb_path: str | Path = "xtb"
        self._level: int = 2
        self._distance_constraints: list[DistanceConstraint] = []
        self._distance_constraint_force_constant: float | None = None

    def get_identifier(self) -> str:
        return "xtb"

    @property
    def xtb_path(self):
        """
        Get the xtb path.
        """
        return self._xtb_path

    @xtb_path.setter
    def xtb_path(self, xtb_path: str | Path):
        """
        Set the xtb path.
        """
        if not isinstance(xtb_path, str | Path):
            raise TypeError("xtb_path should be a string.")
        self._xtb_path = xtb_path

    @property
    def level(self):
        """
        Get the GFN<n>-xTB level.
        """
        return self._level

    @level.setter
    def level(self, level: int):
        """
        Set the GFN<n>-xTB level.
        """
        if not isinstance(level, int):
            raise TypeError("xtb level should be an integer.")
        if level < 0 or level > 2:
            raise ValueError("xtb level should be 0, 1, or 2.")
        self._level = level

    @property
    def distance_constraints(self) -> list[DistanceConstraint]:
        """
        Distance constraints applied during xTB optimizations.
        """
        return self._distance_constraints

    @distance_constraints.setter
    def distance_constraints(
        self,
        constraints: DistanceConstraint
        | Sequence[DistanceConstraint | Mapping[str, Any] | str]
        | None,
    ):
        """
        Set the distance constraints from CLI or TOML input.
        """
        if constraints is None:
            self._distance_constraints = []
            return
        if isinstance(constraints, DistanceConstraint):
            self._distance_constraints = [constraints]
            return
        if not isinstance(constraints, Sequence):
            raise TypeError("Distance constraints must be provided as a sequence.")

        parsed: list[DistanceConstraint] = []
        for entry in constraints:
            parsed.append(self._parse_distance_constraint(entry))
        self._distance_constraints = parsed

    @staticmethod
    def _parse_distance_constraint(
        entry: DistanceConstraint | Mapping[str, Any] | str,
    ) -> DistanceConstraint:
        if isinstance(entry, DistanceConstraint):
            return entry
        if isinstance(entry, Mapping):
            return DistanceConstraint.from_mapping(entry)
        if isinstance(entry, str):
            return DistanceConstraint.from_cli_string(entry)
        raise TypeError("Unsupported distance constraint entry.")

    @property
    def distance_constraint_force_constant(self) -> float | None:
        """
        Shared force constant applied to all distance constraints.
        """
        return self._distance_constraint_force_constant

    @distance_constraint_force_constant.setter
    def distance_constraint_force_constant(self, value: float | int | None):
        """
        Set the shared force constant for distance constraints.
        """
        if value is None:
            self._distance_constraint_force_constant = None
            return
        if not isinstance(value, (float, int)):
            raise TypeError("Force constant must be a float.")
        float_value = float(value)
        if float_value <= 0:
            raise ValueError("Force constant must be positive.")
        self._distance_constraint_force_constant = float_value


class ORCAConfig(BaseConfig):
    """
    Configuration class for ORCA.
    """

    def __init__(self: ORCAConfig) -> None:
        self._orca_path: str | Path = "orca"
        self._functional: str = "PBE"
        self._basis: str = "def2-SVP"
        self._gridsize: int = 1
        self._scf_cycles: int = 100
        self._use_xtb_driver: bool = False

    def get_identifier(self) -> str:
        return "orca"

    @property
    def orca_path(self):
        """
        Get the orca path.
        """
        return self._orca_path

    @orca_path.setter
    def orca_path(self, orca_path: str | Path):
        """
        Set the orca path.
        """
        if not isinstance(orca_path, str | Path):
            raise TypeError("orca_path should be a string or Path.")
        self._orca_path = orca_path

    @property
    def functional(self):
        """
        Get the ORCA functional/method.
        """
        return self._functional

    @functional.setter
    def functional(self, functional: str):
        """
        Set the ORCA functional/method.
        """
        if not isinstance(functional, str):
            raise TypeError("Functional should be a string.")
        self._functional = functional

    @property
    def basis(self):
        """
        Get the ORCA basis set.
        """
        return self._basis

    @basis.setter
    def basis(self, basis: str):
        """
        Set the ORCA basis set.
        """
        if not isinstance(basis, str):
            raise TypeError("Basis should be a string.")
        self._basis = basis

    @property
    def gridsize(self):
        """
        Get the ORCA gridsize for numerical integration.
        """
        return self._gridsize

    @gridsize.setter
    def gridsize(self, gridsize: int):
        """
        Set the ORCA gridsize for numerical integration.
        """
        if not isinstance(gridsize, int):
            raise TypeError("Gridsize should be an integer.")
        if gridsize not in [1, 2, 3]:
            raise ValueError(
                "Gridsize should be 1, 2, or 3. (-> DEFGRID1, DEFGRID2, DEFGRID3)"
            )
        self._gridsize = gridsize

    @property
    def scf_cycles(self):
        """
        Get the maximum number of SCF cycles.
        """
        return self._scf_cycles

    @scf_cycles.setter
    def scf_cycles(self, max_scf_cycles: int):
        """
        Set the maximum number of SCF cycles.
        """
        if not isinstance(max_scf_cycles, int):
            raise TypeError("Max SCF cycles should be an integer.")
        if max_scf_cycles < 1:
            raise ValueError("Max SCF cycles should be greater than 0.")
        self._scf_cycles = max_scf_cycles

    @property
    def use_xtb_driver(self) -> bool:
        """
        Determine whether the xTB external driver can be used during post-processing.
        """
        return self._use_xtb_driver

    @use_xtb_driver.setter
    def use_xtb_driver(self, enabled: bool):
        """
        Enable or disable the usage of the xTB external driver during post-processing.
        """
        if not isinstance(enabled, bool):
            raise TypeError("use_xtb_driver should be a boolean.")
        self._use_xtb_driver = enabled


class TURBOMOLEConfig(BaseConfig):
    """
    Configuration class for TURBOMOLE.
    """

    def __init__(self: TURBOMOLEConfig) -> None:
        self._ridft_path: str | Path = "ridft"
        self._jobex_path: str | Path = "jobex"
        self._functional: str = "pbe"
        self._basis: str = "def2-SVP"
        self._scf_cycles: int = 100

    def get_identifier(self) -> str:
        return "turbomole"

    @property
    def ridft_path(self):
        """
        Get the ridft path.
        """
        return self._ridft_path

    @ridft_path.setter
    def ridft_path(self, ridft_path: str | Path):
        """
        Set the ridft path.
        """
        if not isinstance(ridft_path, str | Path):
            raise TypeError("ridft_path should be a string or Path.")
        self._ridft_path = ridft_path

    @property
    def jobex_path(self):
        """
        Get the jobex path.
        """
        return self._jobex_path

    @jobex_path.setter
    def jobex_path(self, jobex_path: str | Path):
        """
        Set the jobex path.
        """
        if not isinstance(jobex_path, str | Path):
            raise TypeError("jobex_path should be a string or Path.")
        self._jobex_path = jobex_path

    @property
    def functional(self):
        """
        Get the TURBOMOLE functional/method.
        """
        return self._functional

    @functional.setter
    def functional(self, functional: str):
        """
        Set the TURBOMOLE functional/method.
        """
        if not isinstance(functional, str):
            raise TypeError("Functional should be a string.")
        self._functional = functional

    @property
    def basis(self):
        """
        Get the TURBOMOLE basis set.
        """
        return self._basis

    @basis.setter
    def basis(self, basis: str):
        """
        Set the TURBOMOLE basis set.
        """
        if not isinstance(basis, str):
            raise TypeError("Basis should be a string.")
        self._basis = basis

    @property
    def scf_cycles(self):
        """
        Get the maximum number of SCF cycles.
        """
        return self._scf_cycles

    @scf_cycles.setter
    def scf_cycles(self, max_scf_cycles: int):
        """
        Set the maximum number of SCF cycles.
        """
        if not isinstance(max_scf_cycles, int):
            raise TypeError("Max SCF cycles should be an integer.")
        if max_scf_cycles < 1:
            raise ValueError("Max SCF cycles should be greater than 0.")
        self._scf_cycles = max_scf_cycles


class GXTBConfig(BaseConfig):
    """
    Configuration class for g-xTB.
    """

    def __init__(self: GXTBConfig) -> None:
        self._gxtb_path: str | Path = "gxtb"
        self._scf_cycles: int = 100

    def get_identifier(self) -> str:
        return "gxtb"

    @property
    def gxtb_path(self):
        """
        Get the g-xTB path.
        """
        return self._gxtb_path

    @gxtb_path.setter
    def gxtb_path(self, gxtb_path: str | Path):
        """
        Set the g-xTB path.
        """
        if not isinstance(gxtb_path, str | Path):
            raise TypeError("gxtb_path should be a string or Path.")
        self._gxtb_path = gxtb_path

    @property
    def scf_cycles(self):
        """
        Get the maximum number of SCF cycles.
        """
        return self._scf_cycles

    @scf_cycles.setter
    def scf_cycles(self, max_scf_cycles: int):
        """
        Set the maximum number of SCF cycles.
        """
        if not isinstance(max_scf_cycles, int):
            raise TypeError("Max SCF cycles should be an integer.")
        if max_scf_cycles < 1:
            raise ValueError("Max SCF cycles should be greater than 0.")
        self._scf_cycles = max_scf_cycles


class SymmetrizationConfig(BaseConfig):
    """
    Configuration class for the generation of symmetric complexes.
    """

    def __init__(self: SymmetrizationConfig) -> None:
        self._distance: float = 3.0
        self._operation: str = "mirror"
        self._rotation: int | None = None

    def get_identifier(self) -> str:
        return "symmetrization"

    @property
    def distance(self):
        """
        Get the distance.
        """
        return self._distance

    @distance.setter
    def distance(self, distance: float):
        """
        Set the distance.
        """
        if not isinstance(distance, float):
            raise TypeError("Distance should be a float.")
        if distance <= 0:
            raise ValueError("Distance should be greater than 0.")
        self._distance = distance

    @property
    def operation(self):
        """
        Get the operation.
        """
        return self._operation

    @operation.setter
    def operation(self, operation: str):
        """
        Set the operation.
        """
        if not isinstance(operation, str):
            raise TypeError("Operation should be a string.")

        if operation.split("_")[0] not in ["c", "mirror", "inversion"]:
            raise ValueError(
                "Operation can only be translation, c_<n>_rotation, mirror or inversion."
            )

        if operation.endswith("rotation"):
            parts = operation.split("_")
            if len(parts) != 3 or not parts[1].isdigit():
                raise ValueError(
                    "Rotation operations should be in the format 'c_<n>_rotation' where <n> is an integer."
                )
            self.rotation = int(parts[1])
        else:
            if operation not in ["mirror", "inversion"]:
                raise ValueError(
                    "Operation can only be translation, c_<n>_rotation, mirror or inversion."
                )
            self.rotation = None

        self._operation = operation

    @property
    def rotation(self):
        """
        Get the rotation value.
        """
        return self._rotation

    @rotation.setter
    def rotation(self, rotation: int | None):
        """
        Set the rotation value.
        """
        if rotation is None:
            self._rotation = None
            return
        if not isinstance(rotation, int):
            raise TypeError("Rotation should be an integer or None.")
        if rotation < 2:
            raise ValueError("Rotation should be greater than or equal to 2.")
        self._rotation = rotation


class ConfigManager:
    """
    Overall configuration manager for the program.
    """

    def __init__(self, config_file: str | Path | None = None):
        """
        Initialize configuration sections with default values
        """
        self.general = GeneralConfig()
        self.xtb = XTBConfig()
        self.orca = ORCAConfig()
        self.turbomole = TURBOMOLEConfig()
        self.refine = RefineConfig()
        self.postprocess = PostProcessConfig()
        self.generate = GenerateConfig()
        self.gxtb = GXTBConfig()
        self.symmetrization = SymmetrizationConfig()

        if config_file:
            self.load_from_toml(config_file)

    def check_config(self, verbosity: int = 1) -> None:
        """
        Checks ConfigClass for any incompatibilities that are imaginable
        """

        ### Config-specific checks ###
        ##############################
        for attr_name in dir(self):
            attr_value = getattr(self, attr_name)
            if isinstance(attr_value, BaseConfig):
                if hasattr(attr_value, "check_config"):
                    attr_value.check_config(verbosity)

        ### Overlapping checks ###
        ##############################
        num_cores = min(mp.cpu_count(), self.general.parallel)
        if num_cores > 1 and self.postprocess.debug and verbosity > -1:
            # raise warning that debugging of postprocessing will disable parallelization
            warnings.warn(
                "Debug output might seem to be redundant due to the parallel processes "
                + "with possibly similar errors in parallel mode. "
                + "Don't be confused!"
            )

        # Assert that parallel configuration is valid
        if num_cores < self.refine.ncores:
            raise RuntimeError(
                f"Number of cores ({num_cores}) is too low to run refinement using {self.refine.ncores}."
            )
        if self.general.postprocess and num_cores < self.postprocess.ncores:
            raise RuntimeError(
                f"Number of cores ({num_cores}) is too low to run post-processing using {self.postprocess.ncores}."
            )

        # xtb-related checks
        if self.refine.engine == "xtb":
            # Check for f-block elements in forbidden elements
            if self.generate.forbidden_elements:
                if verbosity > 0:
                    f_block_elements = set(range(56, 71)) | set(range(88, 103))
                    lanthanides = set(range(56, 71))
                    if not all(
                        elem in self.generate.forbidden_elements for elem in lanthanides
                    ) or any(
                        elem in f_block_elements
                        for elem in self.generate.element_composition
                    ):
                        warnings.warn(
                            "f-block elements could be within the molecule. "
                            + "xTB does not treat f electrons explicitly. "
                            + "UHF is temporarily set to 0."
                        )

            # Check for super heavy elements in forbidden elements
            super_heavy_elements = set(range(86, 102))
            if self.generate.element_composition and any(
                elem in super_heavy_elements
                for elem in self.generate.element_composition
            ):
                if verbosity > 0:
                    warnings.warn(
                        "xTB does not treat super heavy elements. Atomic numbers are temporarily reduced by 32 to their lighter homologues and then replaced with the correct atom number."
                    )
                # Check if postprocessing is turned off
                if not self.general.postprocess:
                    if verbosity > 0:
                        warnings.warn(
                            "Postprocessing is turned off. The structure will not be relaxed."
                        )

            self._check_distance_constraint_requirements()

    def get_all_identifiers(self):
        """
        Returns the identifiers of all subconfiguration classes, e.g. "orca", "refinement", ...
        """
        identifiers = []
        for attr_name in dir(self):
            attr_value = getattr(self, attr_name)
            # Check if the attribute is an instance of BaseConfig
            if isinstance(attr_value, BaseConfig):
                identifiers.append(attr_value.get_identifier())
        return identifiers

    def load_from_toml(self, config_file: str | Path) -> None:
        """
        Load configuration from TOML file that is structured as follows:
        [general]
        verbosity = 1
        max_cycles = 100
        engine = "xtb"

        [xtb]
        xtb_option = "opt"

        [orca]
        orca_option = "opt"

        Arguments:
            config_file (str): Path to the configuration file

        """
        # Load the configuration file
        config_data = toml.load(config_file)
        self.load_from_dict(config_data)

    def load_from_dict(self, config_dict: dict) -> None:
        """
        Load configuration from a dictionary structured as follows:
        {
            "general": {
                "verbosity": 1,
                "max_cycles": 100,
                "input": "input.xyz"
            },
            "xtb": {
                "xtb_option": "opt"
            },
            "orca": {
                "orca_option": "opt"
            },
        }

        Arguments:
            config_dict (dict): Dictionary containing the configuration
        """
        # Check for unknown keys
        all_identifiers = self.get_all_identifiers()
        for key in config_dict:
            if key not in all_identifiers:
                raise KeyError(f"Unknown key in configuration file: {key}")

        for sub_config in all_identifiers:
            if sub_config not in config_dict:
                continue
            for config_key, config_value in config_dict[sub_config].items():
                # check if config_value is not None and if the attribute exists
                if config_value is not None and hasattr(
                    getattr(self, sub_config), config_key
                ):
                    setattr(getattr(self, sub_config), config_key, config_value)

    def __str__(self) -> str:
        """
        Automated method to display the current configuration.
        """
        configstr = ""
        for attr_name in dir(self):
            attr_value = getattr(self, attr_name)
            if isinstance(attr_value, BaseConfig):
                configstr += (
                    f"{attr_value.get_identifier().capitalize()} configuration:\n"
                )
                for key, value in attr_value.__dict__.items():
                    configstr += (
                        f"{key[1:]:>30}:   {value}\n"  # Skip the leading underscore
                    )
                configstr += "\n"
        return configstr

    def _check_distance_constraint_requirements(self) -> None:
        """
        Ensure distance constraints can be fulfilled with the configured composition.
        """
        constraints = self.xtb.distance_constraints
        if not constraints:
            return

        element_composition = self.generate.element_composition
        # When composition is not fixed, defer feasibility checks until runtime.
        if not self.generate.fixed_composition or not element_composition:
            return

        for constraint in constraints:
            counts = constraint.required_counts()
            for atomic_number, expected in counts.items():
                idx = atomic_number - 1
                element_symbol = constraint.symbol_for(atomic_number)

                if idx not in element_composition:
                    raise ValueError(
                        f"Distance constraint {constraint} requires element "
                        + f"{element_symbol}, but it is not part of the fixed composition."
                    )

                min_count, max_count = element_composition[idx]
                if min_count is None or max_count is None or min_count != max_count:
                    raise ValueError(
                        f"Distance constraint {constraint} requires a fixed number of "
                        + f"{element_symbol} atoms. Please set matching minimum and maximum values."
                    )

                actual = min_count
                if actual < expected:
                    raise ValueError(
                        f"Distance constraint {constraint} requires at least "
                        + f"{expected} {element_symbol} atom(s) but the composition fixes {actual}."
                    )