ase

@dataclass
class AimNet2Calculator(ASECalculator, MachineLearnedInteratomicPotentialCalculator, MSONable):
    """AimNet2 neural network potential calculator.

    AimNet2 is a neural network-based calculator for computing molecular energies
    and atomic partial charges. It provides fast predictions for molecules containing
    H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, and I atoms.

    The calculator requires the 'aimnet' package from:
    https://github.com/cfarm6/aimnetcentral.git

    Attributes:
        name: Name of the calculator (default: "AimNet2 Calculator").
        model: AimNet2 model to use (default: "aimnet2").
        charge: Molecular charge override. If None, uses charge from structure.
        multiplicity: Spin multiplicity override. If None, uses spin from structure.

    """

    name: str = "AimNet2 Calculator"
    model: str = field(default="aimnet2_2025", metadata={"description": "AimNet2 model to use"})
    _properties_model: type[AimNet2Properties] = AimNet2Properties

    def _set_calculator(self, atoms: Atoms, charge: float = 0, spin_multiplicity: int = 1) -> Atoms:
        """Set the AimNet2 calculator on the atoms object.

        Attaches the AimNet2 ASE calculator to the atoms object with the specified
        charge and spin multiplicity. Validates that all atoms are supported by AimNet2.

        Args:
            atoms: ASE Atoms object to attach calculator to.
            charge: Total molecular charge (default: 0). Overridden by self.charge if set.
            spin_multiplicity: Spin multiplicity 2S+1 (default: 1). Overridden by
                self.multiplicity if set.

        Returns:
            ASE Atoms object with AimNet2 calculator attached.

        Raises:
            ImportError: If the 'aimnet' package is not installed.
            ValueError: If molecule contains atoms not supported by AimNet2.

        Examples:
            >>> from ase import Atoms # doctest: +SKIP
            >>> calc = AimNet2Calculator() # doctest: +SKIP
            >>> atoms = Atoms('H2O', positions=[[0,0,0], [1,0,0], [0,1,0]]) # doctest: +SKIP
            >>> atoms = calc.set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
            >>> energy = atoms.get_potential_energy() # doctest: +SKIP
        """
        try:
            from aimnet.calculators import AIMNet2ASE
        except ImportError as e:
            raise ImportError(
                "The 'aimnet' package is required to use AimNet2Calculator but is not available. "
                "Please install it from: https://github.com/cfarm6/aimnetcentral.git"
            ) from e
        if self.charge is not None:
            charge = self.charge
        if self.spin_multiplicity is not None:
            spin_multiplicity = self.spin_multiplicity
        atoms.calc = AIMNet2ASE(self.model, charge, spin_multiplicity)

        aimnet2_atomtypes = [1, 6, 7, 8, 9, 17, 16, 5, 14, 15, 33, 34, 35, 53]
        atomic_nums = atoms.get_atomic_numbers()
        if not all(atom in aimnet2_atomtypes for atom in atomic_nums):
            raise ValueError(
                f"Unsupport atomtype by AimNet2. Supported atom types are {aimnet2_atomtypes}"
            )

        return atoms

    def _get_properties(self, atoms: Atoms) -> AimNet2Properties:
        """Extract computed properties from the AimNet2 calculation.

        Retrieves the total energy and atomic partial charges from the AimNet2
        calculation on the atoms object.

        Args:
            atoms: ASE Atoms object with AimNet2 calculator attached and calculation
                completed.

        Returns:
            Dictionary with structure:
                - "Global": {"Total Energy [eV]": float}
                - "Atomic": {"AimNet2 Partial Charges [e]": array}

        Examples:
            >>> calc = AimNet2Calculator() # doctest: +SKIP
            >>> atoms = calc.set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
            >>> atoms.get_potential_energy()  # Trigger calculation # doctest: +SKIP
            >>> props = calc.get_properties(atoms) # doctest: +SKIP
        """
        energy = atoms.get_total_energy()
        system_property = AimNet2SystemProperties(
            total_energy=SystemProperty(
                name="Total Energy",
                value=energy * ureg.eV,
                description=f"Total energy prediction from {self.model} model",
            ),
        )

        charge = atoms.get_charges()
        atomic_property = AimNet2AtomicProperties(
            aimnet2_partial_charges=AtomicProperty(
                name="AimNet2 Partial Charges",
                value=charge * ureg.e,
                description=f"Partial charges predicted by {self.model} model",
            ),
            forces=AtomicProperty(
                name="AimNet2 Forces",
                value=atoms.get_forces() * ureg.eV / ureg.angstrom,
                description=f"Forces predicted by {self.model} model",
            ),
        )

        return AimNet2Properties(
            atomic=atomic_property,
            system=system_property,
        )

@dataclass
class FairChemCalculator(ASECalculator, MachineLearnedInteratomicPotentialCalculator, MSONable):
    """FairChemital Materials FairChem machine learning force field calculator.

    FairChem models are graph neural network-based force fields developed by FairChem
    Materials for fast and accurate molecular property predictions. The calculator
    supports both conservative and direct versions of the FairChem-v3 model.

    Attributes:
        name: Name of the calculator (default: "FairChem Model Calculator").
        model: FairChem model variant to use. Options:
            - "FairChem-v3-conservative-omol": Conservative model (recommended)
            - "FairChem-v3-direct-omol": Direct model
        charge: Molecular charge override. If None, uses charge from structure.
        multiplicity: Spin multiplicity override. If None, uses spin from structure.
        device: Computation device ("cpu" or "cuda"). Default: "cpu".
        precision: Numerical precision for calculations. Options:
            - "float32-high": Standard precision (default)
            - "float32-highest": Higher precision float32
            - "float64": Double precision
        compile: Whether to compile the model for faster inference (default: False).

    Examples:
        >>> from jfchemistry.calculators import FairChemModelCalculator # doctest: +SKIP
        >>>
        >>> # Create calculator with GPU acceleration
        >>> calc = FairChemModelCalculator(
        ...     model="FairChem-v3-conservative-omol", # doctest: +SKIP
        ...     device="cuda", # doctest: +SKIP
        ...     precision="float32-highest" # doctest: +SKIP
        ... ) # doctest: +SKIP
        >>>
        >>> # Setup on structure
        >>> atoms = molecule.to_ase_atoms() # doctest: +SKIP
        >>> atoms = calc._set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
        >>>
        >>> # Get properties
        >>> props = calc.get_properties(atoms) # doctest: +SKIP
        >>> energy = props["Global"]["Total Energy [eV]"] # doctest: +SKIP
    """

    name: str = "FairChem Model Calculator"
    model: model_types = field(
        default="uma-s-1p1",
        metadata={"description": "The FairChem model to use"},
    )
    device: Literal["cpu", "cuda"] = field(
        default="cuda" if torch.cuda.is_available() else "cpu",
        metadata={"description": "The device to use"},
    )
    task: task_type = field(default="omol", metadata={"description": "The task to use"})
    workers: int = field(default=1, metadata={"description": "The number of workers to use"})
    turbo: bool = field(
        default=False,
        metadata={
            "description": "For long rollout trajectory\
         use-cases, such as molecular dynamics (MD) or relaxations, we provide a special mode\
          called turbo, which optimizes for speed but restricts the user to using a single system\
             where the atomic composition is held constant. Turbo mode is approximately 1.5-2x\
                 faster than default mode, depending on the situation. However, batching is not\
                         supported in this mode. It can be easily activated as shown below."
        },
    )
    tf32: bool = field(
        default=False,
        metadata={
            "description": "Flag to enable or disable the use of\
         tf32 data type for inference. TF32 will slightly reduce accuracy compared to FP32 but will\
             still keep energy conservation in most cases."
        },
    )
    activation_checkpointing: Optional[bool] = field(
        default=None,
        metadata={
            "description": "Flag to enable or disable activation checkpointing\
             during inference. This will dramatically decrease the memory footprint especially for\
                 large number of atoms (ie 10+) at a slight cost to inference speed. If set to None\
                    , the setting from the model checkpoint will be used."
        },
    )
    merge_mole: bool = field(
        default=False,
        metadata={
            "description": "Flag to enable or disable the merging of MOLE experts during inference.\
                 If this is used, the input composition, total charge and spin MUST remain constant\
                     throughout the simulation this will slightly increase speed and reduce memory\
                         footprint used by the parameters significantly"
        },
    )
    compile: bool = field(
        default=False,
        metadata={
            "description": "Flag to enable or disable the compilation of the inference model."
        },
    )

    _properties_model: type[FairChemProperties] = FairChemProperties

    def _set_calculator(self, atoms: Atoms, charge: float = 0, spin_multiplicity: int = 1) -> Atoms:
        """Set the FairChem model calculator on the atoms object.

        Loads the specified FairChem model and attaches it as an ASE calculator to the
        atoms object. Stores charge and spin information in atoms.info dictionary.

        Args:
            atoms: ASE Atoms object to attach calculator to.
            charge: Total molecular charge (default: 0). Overridden by self.charge if set.
            spin_multiplicity: Spin multiplicity 2S+1 (default: 1). Overridden by
                self.multiplicity if set.

        Returns:
            ASE Atoms object with FairChem calculator attached and charge/spin set.

        Raises:
            ImportError: If the 'FairChem-models' package is not installed.

        Examples:
            >>> calc = FairChemModelCalculator(device="cuda", compile=True) # doctest: +SKIP
            >>> atoms = molecule.to_ase_atoms() # doctest: +SKIP
            >>> atoms = calc._set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
            >>> # Charge and spin are stored in atoms.info
            >>> print(atoms.info["charge"]) # doctest: +SKIP
            0
        """
        atoms.info.update({"charge": charge})
        atoms.info.update({"spin": spin_multiplicity})

        predictor = pretrained_mlip.get_predict_unit(
            model_name=self.model,
            device=self.device,
            workers=self.workers,
            inference_settings=InferenceSettings(
                tf32=self.tf32,
                activation_checkpointing=self.activation_checkpointing,
                merge_mole=self.merge_mole,
                compile=self.compile,
            )
            if not self.turbo
            else "turbo",
        )

        atoms.calc = FAIRChemCalculator(predictor, task_name=self.task)
        return atoms

    def _get_properties(self, atoms: Atoms) -> FairChemProperties:
        """Extract computed properties from the FairChem calculation.

        Retrieves the total energy from the FairChem model calculation.

        Args:
            atoms: ASE Atoms object with FairChem calculator attached and calculation
                completed.

        Returns:
            Dictionary with structure:
                - "Global": {"Total Energy [eV]": float}

        Examples:
            >>> calc = FairChemModelCalculator() # doctest: +SKIP
            >>> atoms = calc._set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
            >>> atoms.get_potential_energy()  # Trigger calculation # doctest: +SKIP
            >>> props = calc.get_properties(atoms) # doctest: +SKIP
            >>> print(props["Global"]["Total Energy [eV]"]) # doctest: +SKIP
            -234.567
        """
        print(atoms.calc)
        energy = atoms.get_potential_energy()
        forces = atoms.get_forces()
        atomic_properties = FairChemAtomicProperties(
            forces=AtomicProperty(
                name="FairChem Forces",
                value=forces * ureg.eV / ureg.angstrom,
                description=f"Forces predicted by {self.model} model and {self.task} task",
            ),
        )
        system_properties = FairChemSystemProperties(
            total_energy=SystemProperty(
                name="Potential Energy",
                value=energy * ureg.eV,
                description=f"Potential energy prediction from {self.model} model and {self.task}\
                     task",
            ),
        )
        return FairChemProperties(
            atomic=atomic_properties,
            system=system_properties,
        )

@dataclass
class ORBCalculator(ASECalculator, MachineLearnedInteratomicPotentialCalculator, MSONable):
    """Orbital Materials ORB machine learning force field calculator.

    ORB models are graph neural network-based force fields developed by Orbital
    Materials for fast and accurate molecular property predictions. The calculator
    supports both conservative and direct versions of the ORB-v3 model.

    The calculator requires the 'orb-models' package from:
    https://github.com/orbital-materials/orb-models

    Attributes:
        name: Name of the calculator (default: "ORB Model Calculator").
        model: ORB model variant to use. Options:
            - "orb-v3-conservative-omol": Conservative model (recommended)
            - "orb-v3-direct-omol": Direct model
        charge: Molecular charge override. If None, uses charge from structure.
        multiplicity: Spin multiplicity override. If None, uses spin from structure.
        device: Computation device ("cpu" or "cuda"). Default: "cpu".
        precision: Numerical precision for calculations. Options:
            - "float32-high": Standard precision (default)
            - "float32-highest": Higher precision float32
            - "float64": Double precision
        compile: Whether to compile the model for faster inference (default: False).

    Examples:
        >>> from jfchemistry.calculators import ORBModelCalculator # doctest: +SKIP
        >>>
        >>> # Create calculator with GPU acceleration
        >>> calc = ORBModelCalculator(
        ...     model="orb-v3-conservative-omol", # doctest: +SKIP
        ...     device="cuda", # doctest: +SKIP
        ...     precision="float32-highest" # doctest: +SKIP
        ... ) # doctest: +SKIP
        >>>
        >>> # Setup on structure
        >>> atoms = molecule.to_ase_atoms() # doctest: +SKIP
        >>> atoms = calc._set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
        >>>
        >>> # Get properties
        >>> props = calc.get_properties(atoms) # doctest: +SKIP
        >>> energy = props["Global"]["Total Energy [eV]"] # doctest: +SKIP
    """

    name: str = "ORB Model Calculator"
    model: Literal[
        "orb-v3-conservative-omol",
        "orb-v3-direct-omol",
        "orb-v3-direct-20-omat",
        "orb-v3-direct-20-mpa",
        "orb-v3-direct-inf-omat",
        "orb-v3-direct-inf-mpa",
        "orb-v3-conservative-20-omat",
        "orb-v3-conservative-20-mpa",
        "orb-v3-conservative-inf-omat",
        "orb-v3-conservative-inf-mpa",
    ] = field(default="orb-v3-conservative-omol", metadata={"description": "The ORB model to use"})
    device: Literal["cpu", "cuda"] = field(
        default="cpu", metadata={"description": "The device to use"}
    )
    precision: Literal["float32-high", "float32-highest", "float64"] = field(
        default="float32-high", metadata={"description": "The precision to use"}
    )
    compile: bool = field(default=False, metadata={"description": "Whether to compile the model"})

    _properties_model: type[OrbProperties] = OrbProperties

    def _set_calculator(self, atoms: Atoms, charge: float = 0, spin_multiplicity: int = 1) -> Atoms:
        """Set the ORB model calculator on the atoms object.

        Loads the specified ORB model and attaches it as an ASE calculator to the
        atoms object. Stores charge and spin information in atoms.info dictionary.

        Args:
            atoms: ASE Atoms object to attach calculator to.
            charge: Total molecular charge (default: 0). Overridden by self.charge if set.
            spin_multiplicity: Spin multiplicity 2S+1 (default: 1). Overridden by
                self.multiplicity if set.

        Returns:
            ASE Atoms object with ORB calculator attached and charge/spin set.

        Raises:
            ImportError: If the 'orb-models' package is not installed.

        Examples:
            >>> calc = ORBModelCalculator(device="cuda", compile=True) # doctest: +SKIP
            >>> atoms = molecule.to_ase_atoms() # doctest: +SKIP
            >>> atoms = calc._set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
            >>> # Charge and spin are stored in atoms.info
            >>> print(atoms.info["charge"]) # doctest: +SKIP
            0
        """
        if (self.model in {"orb-v3-conservative-omol", "orb-v3-direct-omol"}) and all(atoms.pbc):
            raise ValueError(
                "ORB OMol models do not support periodic boundary conditions.\
                Please remove the cell from the atoms object."
            )
        try:
            from orb_models.forcefield import pretrained
            from orb_models.forcefield.calculator import ORBCalculator
        except ImportError as e:
            raise ImportError(
                "The 'orb-models' package is required to use ORBCalculator but is not available."
                "Please install it from: https://github.com/orbital-materials/orb-models"
            ) from e
        if self.charge is not None:
            charge = self.charge
        if self.spin_multiplicity is not None:
            spin_multiplicity = self.spin_multiplicity

        orbff = getattr(pretrained, self.model.replace("-", "_"))(
            device=self.device,
            precision=self.precision,
            compile=self.compile,
        )

        atoms.calc = ORBCalculator(orbff, device=self.device)
        atoms.info["charge"] = charge
        atoms.info["spin"] = spin_multiplicity
        return atoms

    def _get_properties(self, atoms: Atoms) -> OrbProperties:
        """Extract computed properties from the ORB calculation.

        Retrieves the total energy from the ORB model calculation.

        Args:
            atoms: ASE Atoms object with ORB calculator attached and calculation
                completed.

        Returns:
            Dictionary with structure:
                - "Global": {"Total Energy [eV]": float}

        Examples:
            >>> calc = ORBModelCalculator() # doctest: +SKIP
            >>> atoms = calc._set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
            >>> atoms.get_potential_energy()  # Trigger calculation # doctest: +SKIP
            >>> props = calc.get_properties(atoms) # doctest: +SKIP
            >>> print(props["Global"]["Total Energy [eV]"]) # doctest: +SKIP
            -234.567
        """
        energy = atoms.get_total_energy()
        forces = atoms.get_forces()
        atomic_properties = OrbAtomicProperties(
            forces=AtomicProperty(
                name="ORB Forces",
                value=forces * ureg.eV / ureg.angstrom,
                description=f"Forces predicted by {self.model} model",
            ),
        )
        system_properties = OrbSystemProperties(
            total_energy=SystemProperty(
                name="Total Energy",
                value=energy * ureg.eV,
                description=f"Total energy prediction from {self.model} model",
            ),
        )
        return OrbProperties(
            atomic=atomic_properties,
            system=system_properties,
        )

@dataclass
class TBLiteCalculator(ASECalculator, SemiempiricalCalculator, MSONable):
    """TBLite calculator for GFN-xTB semi-empirical methods.

    TBLite provides implementations of the GFN (Geometrical-dependent
    Forcefield for Noncovalent interactions) extended tight-binding methods
    developed by the Grimme group. These methods offer a good balance between
    accuracy and computational efficiency for large molecular systems.

    The calculator computes extensive molecular properties including energies,
    partial charges, bond orders, molecular orbitals, dipole/quadrupole moments,
    and HOMO-LUMO gaps.

    Attributes:
        name: Name of the calculator (default: "TBLite Calculator").
        method: Semi-empirical method to use. Options:
            - "GFN2-xTB": GFN2-xTB method (default, recommended for most cases)
            - "GFN1-xTB": GFN1-xTB method
            - "IPEA1-xTB": IPEA1-xTB method
        charge: Molecular charge override. If None, uses charge from structure.
        multiplicity: Spin multiplicity override. If None, uses spin from structure.
        accuracy: Numerical accuracy parameter (default: 1.0).
        electronic_temperature: Electronic temperature in Kelvin for Fermi smearing
            (default: 300.0).
        max_iterations: Maximum SCF iterations (default: 250).
        initial_guess: Initial guess for electronic structure. Options:
            - "sad": Superposition of atomic densities (default)
            - "eeq": Electronegativity equilibration
        mixer_damping: Damping parameter for SCF mixing (default: 0.4).
        verbosity: Output verbosity level (default: 0).

    Examples:
        >>> from jfchemistry.calculators import TBLiteCalculator # doctest: +SKIP
        >>> from ase.build import molecule # doctest: +SKIP
        >>> from pymatgen.core import Molecule # doctest: +SKIP
        >>> mol = Molecule.from_ase_atoms(molecule("C2H6")) # doctest: +SKIP
        >>> # Create calculator with custom settings
        >>> calc = TBLiteCalculator(  # doctest: +SKIP
        ...     method="GFN2-xTB", # doctest: +SKIP
        ...     accuracy=0.1,  # Tighter convergence # doctest: +SKIP
        ...     max_iterations=500 # doctest: +SKIP
        ... ) # doctest: +SKIP
        >>>
        >>> # Compute properties
        >>> atoms = mol.to_ase_atoms() # doctest: +SKIP
        >>> atoms = calc.set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
        >>> props = calc.get_properties(atoms) # doctest: +SKIP

    """

    name: str = "TBLite Calculator"
    method: Literal["GFN1-xTB", "GFN2-xTB", "IPEA1-xTB"] = field(
        default="GFN2-xTB", metadata={"description": "The method to use"}
    )
    accuracy: float = field(default=1.0, metadata={"description": "The accuracy to use"})
    electronic_temperature: float = field(
        default=300.0,
        metadata={
            "description": "The electronic temperature to use",
            "unit": "K",
        },
    )
    max_iterations: int = field(
        default=250, metadata={"description": "The maximum number of iterations to use"}
    )
    initial_guess: Literal["sad", "eeq"] = field(
        default="sad", metadata={"description": "The initial guess to use"}
    )
    mixer_damping: float = field(default=0.4, metadata={"description": "The mixer damping to use"})
    verbosity: int = field(default=0, metadata={"description": "The verbosity to use"})
    _properties_model: type[TBLiteProperties] = TBLiteProperties
    solvation: Optional[TBLiteSolvationType] = field(
        default=None, metadata={"description": "The solvation model to use"}
    )
    solvent: Optional[TBLiteSolventType] = field(
        default=None, metadata={"description": "The solvent to use"}
    )
    implicit_solvent: Optional[ImplicitSolventConfig] = field(
        default=None,
        metadata={"description": "Unified implicit-solvent configuration override."},
    )

    def __post_init__(self):
        """Apply unified implicit-solvent overrides when provided."""
        if self.implicit_solvent is None:
            return
        mapped_solvation, mapped_solvent = to_tblite(self.implicit_solvent)
        self.solvation = mapped_solvation  # type: ignore[assignment]
        self.solvent = mapped_solvent  # type: ignore[assignment]

    def _set_calculator(self, atoms: Atoms, charge: float = 0, spin_multiplicity: int = 1) -> Atoms:
        """Set the TBLite calculator on the atoms object.

        Configures and attaches a TBLite GFN-xTB calculator to the atoms object
        with the specified method, charge, and SCF parameters.

        Args:
            atoms: ASE Atoms object to attach calculator to.
            charge: Total molecular charge (default: 0). Overridden by self.charge if set.
            spin_multiplicity: Spin multiplicity 2S+1 (default: 1). Overridden by
                self.multiplicity if set.

        Returns:
            ASE Atoms object with TBLite calculator attached.

        Raises:
            ImportError: If the 'tblite' package is not installed.

        Examples:
            >>> from ase.build import molecule # doctest: +SKIP
            >>> from jfchemistry.calculators import TBLiteCalculator # doctest: +SKIP
            >>>
            >>> # Create calculator and set up a water molecule
            >>> calc = TBLiteCalculator(method="GFN2-xTB") # doctest: +SKIP
            >>> atoms = molecule("H2O") # doctest: +SKIP
            >>> atoms = calc.set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
            >>>
            >>> props = atoms.get_properties() # doctest: +SKIP
            >>> energy = props["Global"]["Total Energy [eV]"] # doctest: +SKIP
        """
        try:
            from tblite.ase import TBLite  # doctest: +SKIP
        except ImportError as e:
            raise ImportError(
                "The 'tblite' package is required to use TBLiteCalculator but is not available. "
                "Please install it with: pip install tblite or conda install tblite-python"
            ) from e
        if self.charge is not None:
            charge = self.charge
        if self.spin_multiplicity is not None:
            spin_multiplicity = self.spin_multiplicity
        atoms.calc = TBLite(
            method=self.method,
            charge=charge,
            spin_multiplicity=spin_multiplicity,
            accuracy=self.accuracy,
            electronic_temperature=self.electronic_temperature,
            max_iterations=self.max_iterations,
            initial_guess=self.initial_guess,
            mixer_damping=self.mixer_damping,
            verbosity=self.verbosity,
            solvation=(self.solvation, self.solvent) if self.solvation and self.solvent else None,
        )

        return atoms

    def _get_properties(self, atoms: Atoms) -> TBLiteProperties:
        """Extract comprehensive properties from the TBLite calculation.

        Computes and organizes a wide range of molecular properties from the
        GFN-xTB calculation including energies, multipole moments, partial charges,
        bond orders, and molecular orbital information.

        Args:
            atoms: ASE Atoms object with TBLite calculator attached and calculation
                completed.

        Returns:
            Dictionary with four main sections:
                - "Global": System-level properties
                    - "Total Energy [eV]": Total molecular energy
                    - "Dipole Moment [D]": Dipole moment vector (x, y, z)
                    - "Quadrupole Moment [D]": Quadrupole tensor (xx, yy, zz, xy, xz, yz)
                    - "Density Matrix": Electronic density matrix
                    - "HOMO-LUMO Gap [eV]": Energy gap between frontier orbitals
                    - "HOMO Energy [eV]": Highest occupied molecular orbital energy
                    - "LUMO Energy [eV]": Lowest unoccupied molecular orbital energy
                - "Atomic": Per-atom properties
                    - "Mulliken Partial Charges [e]": Mulliken population analysis charges
                - "Bond": Pairwise bond properties
                    - "Wiberg Bond Order": List of bond orders for all atom pairs
                        Each entry: {"i": atom_index, "j": atom_index, "value": bond_order}
                - "Orbital": Molecular orbital information
                    - "Orbital Energies [eV]": Energy of each molecular orbital
                    - "Orbital Occupations": Occupation number of each orbital
                    - "Orbital Coefficients": Coefficients of molecular orbitals

        Examples:
            >>> from ase.build import molecule # doctest: +SKIP
            >>> from pymatgen.core import Molecule # doctest: +SKIP
            >>> from jfchemistry.calculators import TBLiteCalculator # doctest: +SKIP
            >>>
            >>> # Create a simple molecule (ethane) and set up calculator
            >>> mol = Molecule.from_ase_atoms(molecule("C2H6")) # doctest: +SKIP
            >>> calc = TBLiteCalculator(method="GFN2-xTB") # doctest: +SKIP
            >>> atoms = mol.to_ase_atoms() # doctest: +SKIP
            >>> atoms = calc.set_calculator(atoms, charge=0, spin_multiplicity=1) # doctest: +SKIP
            >>> props = calc.get_properties(atoms) # doctest: +SKIP
        """
        atoms.get_potential_energy()
        forces = atoms.get_forces()
        props = atoms.calc._res.dict()
        energy = props["energy"] * Hartree
        dipole = props["dipole"] * Bohr
        quadrupole = props["quadrupole"] * Bohr**2
        charges = props["charges"]
        bond_orders = props["bond-orders"]
        orbital_energies = props["orbital-energies"] * Hartree
        orbital_occupations = props["orbital-occupations"]
        # Report HOMO-LUMO Gap
        homo_index = orbital_occupations.nonzero()[0][-1]
        lumo_index = homo_index + 1
        homo_energy = orbital_energies[homo_index]
        lumo_energy = orbital_energies[lumo_index]
        homo_lumo_gap = lumo_energy - homo_energy
        orbital_coefficients = props["orbital-coefficients"]
        density_matrix = props["density-matrix"]
        wbo = np.array([], dtype=float)
        atoms_i = np.array([], dtype=int)
        atoms_j = np.array([], dtype=int)
        for i in range(len(bond_orders)):
            for j in range(len(bond_orders[i])):
                if bond_orders[i][j] > BOND_ORDER_THRESHOLD:
                    wbo = np.append(wbo, bond_orders[i][j])
                    atoms_i = np.append(atoms_i, i)
                    atoms_j = np.append(atoms_j, j)
        system_properties = TBLiteSystemProperties(
            total_energy=SystemProperty(
                name="Total Energy",
                value=energy * ureg.eV,
                description="Total energy of the system",
            ),
            dipole_moment=SystemProperty(
                name="Dipole Moment",
                value=dipole * ureg.D,
                description="Dipole moment of the system",
            ),
            quadrupole_moment=SystemProperty(
                name="Quadrupole Moment",
                value=quadrupole * ureg.D * ureg.angstrom,
                description="Quadrupole moment of the system",
            ),
            density_matrix=SystemProperty(
                name="Density Matrix",
                value=density_matrix,
                description="Density matrix of the system",
            ),
            homo_lumo_gap=SystemProperty(
                name="HOMO-LUMO Gap",
                value=homo_lumo_gap * ureg.eV,
                description="HOMO-LUMO gap of the system",
            ),
            homo_energy=SystemProperty(
                name="HOMO Energy",
                value=homo_energy * ureg.eV,
                description="HOMO energy of the system",
            ),
            lumo_energy=SystemProperty(
                name="LUMO Energy",
                value=lumo_energy * ureg.eV,
                description="LUMO energy of the system",
            ),
        )
        atomic_properties = TBLiteAtomicProperties(
            tblite_partial_charges=AtomicProperty(
                name="Mulliken Partial Charges",
                value=(charges * ureg.e).tolist(),
                description="Mulliken partial charges of the system",
            ),
            tblite_forces=AtomicProperty(
                name="Forces",
                value=forces.tolist() * (ureg.eV / ureg.angstrom),
                description="Forces of the system",
            ),
        )
        bond_properties = TBLiteBondProperties(
            tblite_wiberg_bond_orders=BondProperty(
                name="Wiberg Bond Order",
                value=wbo.tolist(),
                atoms1=atoms_i.tolist(),
                atoms2=atoms_j.tolist(),
                description="Wiberg bond order of the system",
            ),
        )
        orbital_properties = TBLiteOrbitalProperties(
            tblite_orbital_energies=OrbitalProperty(
                name="Orbital Energies",
                value=orbital_energies.tolist() * ureg.eV,
                description="Orbital energies of the system",
            ),
            tblite_orbital_occupations=OrbitalProperty(
                name="Orbital Occupations",
                value=orbital_occupations.tolist(),
                description="Orbital occupations of the system",
            ),
            tblite_orbital_coefficients=OrbitalProperty(
                name="Orbital Coefficients",
                value=orbital_coefficients.tolist(),
                description="Orbital coefficients of the system",
            ),
        )
        properties = TBLiteProperties(
            system=system_properties,
            atomic=atomic_properties,
            bond=bond_properties,
            orbital=orbital_properties,
        )

        return properties