import builtins
import os
from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
from typing import IO, Any, overload, Literal as L, SupportsIndex, TypedDict

import numpy as np
from scm.plams import Atom, Bond, PDBHandler, Settings

class _AtomDict(TypedDict):
    atnum: int
    bonds: list[int]
    coords: tuple[float, float, float]
    properties: Settings[Any, Any]

class _BondDict(TypedDict):
    atom1: int
    atom2: int
    order: float
    properties: Settings[Any, Any]

class _MolDictBase(TypedDict, total=False):
    charge: int

class _MolDict(_MolDictBase):
    atoms: list[_AtomDict]
    bonds: list[_BondDict]
    lattice: list[Sequence[float]]
    properties: Settings[Any, Any]

class _LabelDict(TypedDict, total=False):
    BO: bool
    RS: bool
    EZ: bool
    DH: bool
    CO: bool
    twist_tol: float
    bend_tol: float

class Molecule:
    atoms: list[Atom]
    bonds: list[Bond]
    lattice: list[Sequence[float]]
    properties: Settings[Any, Any]
    def __init__(
        self,
        filename: None | builtins.str | os.PathLike[builtins.str] = ...,
        inputformat: L[None, "xyz", "pdb", "mol", "mol2", "rkf"] = ...,
        positions: None | Sequence[Sequence[float]] | np.ndarray = ...,
        numbers: None | Collection[int] = ...,
        lattice: None | Sequence[Sequence[float]] | np.ndarray = ...,
        **other: Any,
    ) -> None: ...
    def copy(self, atoms: None | Collection[Atom] = ...) -> Molecule: ...
    def add_molecule(self, other: Molecule, copy: bool = ...) -> None: ...
    def add_atom(self, atom: Atom, adjacent: None | Iterable[Atom] = ...) -> None: ...
    def delete_atom(self, atom: Atom) -> None: ...
    @overload
    def add_bond(self, arg1: Atom, arg2: Atom, order: float = ...) -> None: ...
    @overload
    def add_bond(self, arg1: Bond, arg2: None = ..., order: float = ...) -> None: ...
    @overload
    def delete_bond(self, arg1: Atom, arg2: Atom) -> None: ...
    @overload
    def delete_bond(self, arg1: Bond, arg2: None = ...) -> None: ...
    def delete_all_bonds(self) -> None: ...
    def find_bond(self, atom1: Atom, atom2: Atom) -> None | Bond: ...
    def set_atoms_id(self, start: int = ...) -> None: ...
    def unset_atoms_id(self) -> None: ...
    def neighbors(self, atom: Atom) -> list[Atom]: ...
    def bond_matrix(self) -> np.ndarray: ...
    def separate(self) -> list[Molecule]: ...
    def guess_bonds(self, atom_subset: None | Sequence[Atom] = ..., dmax: float = ...) -> None: ...
    def in_ring(self, arg: Atom | Bond) -> bool: ...
    def supercell(self, *args: float | Sequence[float]) -> Molecule: ...
    def cell_lengths(self, unit: builtins.str) -> list[float]: ...
    def cell_angles(self, unit: builtins.str) -> list[float]: ...
    def unit_cell_volume(self, unit: builtins.str = ...) -> float: ...
    def set_integer_bonds(self, action: builtins.str = ..., tolerance: float = ...) -> None: ...
    @overload
    def index(self, value: Atom, start: int = ..., stop: None | int = ...) -> int: ...
    @overload
    def index(self, value: Bond, start: int = ..., stop: None | int = ...) -> tuple[int, int]: ...
    @overload
    def round_coords(self, decimals: int = ..., *, inplace: L[False] = ...) -> Molecule: ...
    @overload
    def round_coords(self, decimals: int = ..., *, inplace: L[True]) -> None: ...
    def translate(self, vector: tuple[float, float, float], unit: builtins.str = ...) -> None: ...
    def rotate_lattice(self, matrix: Sequence[Sequence[float]] | np.ndarray) -> None: ...
    def rotate(self, matrix: Sequence[Sequence[float]] | np.ndarray, lattice: bool = ...) -> None: ...
    def align_lattice(self, convention: L["AMS", "reax"] = ..., zero: float = ...) -> bool: ...
    def rotate_bond(self, bond: Bond, moving_atom: Atom, angle: float, unit: builtins.str = ...) -> None: ...
    def resize_bond(self, bond: Bond, moving_atom: Atom, length: float, unit: builtins.str = ...) -> None: ...
    def closest_atom(self, point: Atom | tuple[float, float, float], unit: builtins.str = ...) -> Atom: ...
    def distance_to_point(self, point: Sequence[float], unit: builtins.str = ..., result_unit: builtins.str = ...) -> float: ...
    @overload
    def distance_to_mol(self, other: Molecule, result_unit: builtins.str = ..., *, return_atoms: L[False] = ...) -> float: ...
    @overload
    def distance_to_mol(
        self, other: Molecule, result_unit: builtins.str = ..., *, return_atoms: L[True]
    ) -> tuple[float, Atom, Atom]: ...
    def wrap(
        self, length: float, angle: float = ..., length_unit: builtins.str = ..., angle_unit: builtins.str = ...
    ) -> None: ...
    def get_center_of_mass(self, unit: builtins.str = ...) -> tuple[float, float, float]: ...
    def get_mass(self) -> float: ...
    def get_density(self) -> float: ...
    @overload
    def get_formula(self, as_dict: L[False] = ...) -> builtins.str: ...
    @overload
    def get_formula(self, as_dict: L[True]) -> dict[builtins.str, int]: ...
    def apply_strain(self, strain: np.ndarray | Sequence[float], voigt_form: bool = ...) -> None: ...
    def map_to_central_cell(self, around_origin: bool = ...) -> Molecule: ...
    def perturb_atoms(
        self, max_displacement: float = ..., unit: builtins.str = ..., atoms: None | Iterable[Atom] = ...
    ) -> None: ...
    def perturb_lattice(self, max_displacement: float = ..., unit: builtins.str = ..., ams_convention: bool = ...) -> None: ...
    def substitute(
        self,
        connector: Bond | tuple[Atom, Atom],
        ligand: Molecule,
        ligand_connector: Bond | tuple[Atom, Atom],
        bond_length: None | float = ...,
        steps: int = ...,
        cost_func_mol: None | Callable[[Molecule, Molecule], float] = ...,
        cost_func_array: None | Callable[[np.ndarray, np.ndarray], float] = ...,
    ) -> None: ...
    def __len__(self) -> int: ...
    def __str__(self) -> builtins.str: ...
    def str(self, decimal: int = ...) -> builtins.str: ...
    def __iter__(self) -> Iterator[Atom]: ...
    @overload
    def __getitem__(self, key: SupportsIndex) -> Atom: ...
    @overload
    def __getitem__(self, key: tuple[SupportsIndex, SupportsIndex]) -> Bond: ...
    def __add__(self, other: Molecule) -> Molecule: ...
    def __iadd__(self, other: Molecule) -> Molecule: ...
    def __copy__(self) -> Molecule: ...
    def __deepcopy__(self, memo: object) -> Molecule: ...
    def __round__(self, ndigits: None | int = ...) -> Molecule: ...
    def __getstate__(self) -> _MolDict: ...
    def __setstate__(self, state: _MolDict) -> None: ...
    def as_dict(self) -> _MolDict: ...
    @classmethod
    def from_dict(cls, dictionary: _MolDict) -> Molecule: ...
    @classmethod
    def from_elements(cls, elements: Iterable[builtins.str]) -> Molecule: ...
    def as_array(self, atom_subset: None | Iterable[Atom] = ...) -> np.ndarray: ...
    def from_array(
        self, xyz_array: Molecule | Iterable[tuple[float, float, float]] | np.ndarray, atom_subset: None | Iterable[Atom] = ...
    ) -> None: ...
    def __array__(self, dtype: Any = ...) -> np.ndarray: ...
    def readxyz(self, f: IO[builtins.str], geometry: int = ..., **other: Any) -> None: ...
    def writexyz(self, f: IO[builtins.str], **other: Any) -> None: ...
    def readmol(self, f: IO[builtins.str], **other: Any) -> None: ...
    def writemol(self, f: IO[builtins.str], **other: Any) -> None: ...
    def readmol2(self, f: IO[builtins.str], **other: Any) -> None: ...
    def writemol2(self, f: IO[builtins.str], **other: Any) -> None: ...
    def readpdb(self, f: IO[builtins.str], geometry: int = ..., **other: Any) -> PDBHandler: ...
    def writepdb(self, f: IO[builtins.str], **other: Any) -> None: ...
    def read(
        self, filename: builtins.str | os.PathLike[builtins.str], inputformat: None | builtins.str = ..., **other: Any
    ) -> None | PDBHandler: ...
    def write(
        self, filename: builtins.str | os.PathLike[builtins.str], outputformat: None | builtins.str = ..., **other: Any
    ) -> None: ...
    def add_hatoms(self) -> Molecule: ...
    @overload
    @staticmethod
    def rmsd(mol1: Molecule, mol2: Molecule, *, return_rotmat: L[False] = ..., check: bool = ...) -> np.ndarray: ...
    @overload
    @staticmethod
    def rmsd(mol1: Molecule, mol2: Molecule, *, return_rotmat: L[True], check: bool = ...) -> tuple[np.ndarray, np.ndarray]: ...
    # `label` and `reorder` are defined in `scm.plams.mol.identity`
    @overload
    def label(self, level: int = ..., keep_labels: bool = ..., flags: None | _LabelDict = ...) -> builtins.str: ...
    @overload
    def label(
        self, level: Sequence[int], keep_labels: bool = ..., flags: None | _LabelDict = ...
    ) -> tuple[builtins.str, ...]: ...
    def reorder(self, other: Molecule) -> None: ...
