# [1] https://aip.scitation.org/doi/pdf/10.1063/5.0004046
#     Efficient implementation of the superposition of atomic potentials initial
#     guess for electronic structure calculations in Gaussian basis sets
#     Lehtola, Visscher, Engel,  2020

import itertools as it
from math import sqrt, log, pi
from pathlib import Path
import textwrap
from typing import List, Literal


import h5py
from jinja2 import Template
from joblib import Memory
import numpy as np
from numpy.typing import NDArray
import scipy as sp


from pysisyphus.config import L_MAX, L_AUX_MAX
from pysisyphus.elem_data import (
    ATOMIC_NUMBERS,
    INV_ATOMIC_NUMBERS,
    nuc_charges_for_atoms,
)
from pysisyphus.helpers_pure import file_or_str
from pysisyphus.linalg import multi_component_sym_mat
from pysisyphus.wavefunction.helpers import (
    canonical_order,
    get_l,
    get_shell_shape,
    L_MAP_INV,
    permut_for_order,
)

from pysisyphus.wavefunction.ints import (
    cart_gto3d,
    coulomb3d,
    diag_quadrupole3d,
    dipole3d,
    kinetic3d,
    ovlp3d,
    quadrupole3d,
    int2c2e3d,
    int3c2e3d_sph,
)

from pysisyphus.wavefunction.cart2sph import cart2sph_coeffs
from pysisyphus.wavefunction.normalization import norm_cgto_lmn


class Shell:
    def __init__(
        self,
        L: int,
        center: NDArray[float],
        coeffs: NDArray[float],
        exps: NDArray[float],
        center_ind: int,
        atomic_num=None,
        shell_index=None,
        index=None,
        # min_coeff: float = 1e-8,
    ):
        self.L = get_l(L)
        self.center = np.array(center, dtype=float)  # (x, y, z), 1d array
        coeffs = np.array(coeffs, dtype=float)
        # cur_min_coeff = np.abs(coeffs).min()
        # assert (
        # cur_min_coeff >= min_coeff
        # ), f"Got invalid contraction coeff of '{cur_min_coeff:.6f}'"
        # Store original contraction coefficients
        self.coeffs_org = coeffs.copy()
        # Orbital exponents, 1d array
        self.exps = np.array(exps, dtype=float)
        assert self.coeffs_org.size == self.exps.size
        coeffs = np.array(coeffs)
        assert coeffs.size == self.exps.size
        self.center_ind = int(center_ind)
        if atomic_num is None:
            atomic_num = -1
        self.atomic_num = int(atomic_num)
        if shell_index is not None:
            self.shell_index = shell_index
        if index is not None:
            self.index = index

        self.coeffs, _ = norm_cgto_lmn(coeffs, self.exps, self.L)

        try:
            self.atom = INV_ATOMIC_NUMBERS[self.atomic_num]
        except KeyError:
            self.atom = "X"
        self.atom = self.atom.lower()

    def as_tuple(self):
        return self.L, self.center, self.coeffs, self.exps, self.index, self.size

    def exps_coeffs_iter(self):
        return zip(self.exps, self.coeffs)

    @property
    def contr_depth(self):
        return self.coeffs.size

    @property
    def cart_powers(self):
        return np.array(canonical_order(self.L), dtype=int)

    @property
    def size(self):
        """Number of cartesian basis functions in the shell."""
        return self.cart_size

    @property
    def cart_size(self):
        """Number of cartesian basis functions in the shell."""
        L = self.L
        return (L + 1) * (L + 2) // 2

    @property
    def sph_size(self):
        """Number of cartesian basis functions in the shell."""
        return 2 * self.L + 1

    @property
    def index(self) -> int:
        return self._index

    @index.setter
    def index(self, index: int):
        self._index = int(index)

    @property
    def shell_index(self) -> int:
        return self._shell_index

    @shell_index.setter
    def shell_index(self, shell_index: int):
        self._shell_index = shell_index

    def to_pyscf_shell(self):
        key = L_MAP_INV[self.L]
        lines = [
            f"{self.atom.title()}    {key.upper()}",
        ]
        lines += [
            f"{exp_:> 18.10f}    {coeff:>18.10f}"
            for exp_, coeff in zip(self.exps, self.coeffs)
        ]
        return "\n".join(lines)

    def to_sap_shell(self):
        """Fix contraction coefficients, for use in SAP-initial guess calculation.

        See [1]."""
        # SAP shells must always be s-shells
        assert self.L == 0
        self.coeffs = self.coeffs_org
        coeffs_sum = self.coeffs.sum()
        # Check sum rule; sum should equal -Z.
        assert (coeffs_sum + ATOMIC_NUMBERS[self.atom.lower()]) <= 1e-5
        # The potential fits were carried out for functions of the form:
        #   g_p(r) = (α_p / π)**1.5 * exp(-α_p / r) .
        # So we multiply the prefactor onto the contraction coefficients.
        # See [1], p. 152, just above Eq. (2).
        self.coeffs *= (self.exps / np.pi) ** 1.5

    def __str__(self):
        try:
            center_str = f", at atom {self.center_ind}"
        except AttributeError:
            center_str = ""
        return f"Shell(L={self.L}, {self.contr_depth} pGTO{center_str})"

    def __repr__(self):
        return self.__str__()


Ordering = Literal["native", "pysis"]


class Shells:
    sph_Ps = {l: np.eye(2 * l + 1) for l in range(max(L_MAX, L_AUX_MAX) + 1)}

    def __init__(
        self,
        shells: List[Shell],
        screen: bool = False,
        cache: bool = False,
        cache_path: str = "./cache",
        ordering: Ordering = "native",
    ):
        self.shells = shells
        self.ordering = ordering
        self.screen = screen
        self.cache = cache
        self.cache_path = Path(cache_path)
        """
        'native' refers to the original ordering, as used in the QC program. The
        ordering will be reflected in the MO coefficient ordering. With 'native'
        the integrals calculated by pysisyphus must be reorderd, to match the native
        ordering of the MO coefficients.
        """
        assert ordering in ("pysis", "native")

        # Now that we have all Shell objects, we can set their starting indices
        index = 0
        shell_index = 0
        for shell in self.shells:
            shell.shell_index = shell_index
            shell.index = index
            index += shell.size
            shell_index += 1

        # Try to construct Cartesian permutation matrix from cart_order, if defnied.
        try:
            self.cart_Ps = permut_for_order(self.cart_order)
        except AttributeError:
            pass

        self.atoms, self.coords3d = self.atoms_coords3d
        # Precontract & store coefficients for reordering spherical basis functions
        # and converting them from Cartesian basis functions.
        self.reorder_c2s_coeffs = self.P_sph @ self.cart2sph_coeffs

        # Enable disk cache for 1el-integrals
        if self.cache:
            self.memory = Memory(self.cache_path, verbose=0)
            self.get_1el_ints_cart = self.memory.cache(self.get_1el_ints_cart)
            self.get_1el_ints_sph = self.memory.cache(self.get_1el_ints_sph)

    def __len__(self):
        return len(self.shells)

    def __getitem__(self, key):
        return self.shells[key]

    def print_shells(self):
        for shell in self.shells:
            print(shell)

    def as_tuple(self):
        # Ls, centers, contr_coeffs, exponents, indices, sizes
        return zip(*[shell.as_tuple() for shell in self.shells])

    @property
    def cart_size(self):
        """Number of cartesian basis functions."""
        return sum([shell.cart_size for shell in self.shells])

    @property
    def sph_size(self):
        """Number of spherical basis."""
        return sum([shell.sph_size for shell in self.shells])

    def as_arrays(self, fortran=False):
        """Similar to the approach in libcint."""
        # bas_centers
        # One entry per shell, integer array.
        #   center_ind, atomic number, pointer to center coordinates in bas_data (3 integers)
        bas_centers = list()
        # bas_spec
        # One entry per shell, integer array.
        #   shell_ind, total angmom, N_pgto, N_cgto, \
        #   pointer to contraction coefficients and exponents in bas_data \
        #   (2*N_pgto floats)
        bas_spec = list()
        # bas_data
        # Float array. Starts with 3 * N_centers floats, containing the center
        # coordinates. Continues with alternating contraction coefficient and
        # orbital exponent data.
        bas_data = list()
        pointer = 1 if fortran else 0

        shells = self.shells
        # Store center data, i.e., where the basis functions are located.
        for shell in shells:
            center_ind = shell.center_ind
            bas_data.extend(shell.center)  # Store coordinates
            atomic_num = shell.atomic_num
            if atomic_num is None:
                atomic_num = -1
            bas_centers.append((center_ind, atomic_num, pointer))
            pointer += 3

        # Store contraction coefficients and primitive exponents.
        for shell_ind, shell in enumerate(self.shells):
            # L, center, coeffs, exponents, index, size = shell.as_tuple()
            L, _, _, exponents, _, _ = shell.as_tuple()
            contr_coeffs = shell.coeffs_org
            nprim = len(contr_coeffs)
            # ncontr is hardcoded to 1 for now, as there are no special functions
            # to handle general contractions.
            ncontr = 1
            bas_spec.append((shell_ind, L, nprim, ncontr, pointer))
            bas_data.extend(contr_coeffs)
            pointer += nprim
            bas_data.extend(exponents)
            pointer += nprim

        bas_centers = np.array(bas_centers, dtype=int)
        bas_spec = np.array(bas_spec, dtype=int)
        bas_data = np.array(bas_data, dtype=float)
        return bas_centers, bas_spec, bas_data

    def to_sympleints_arrays(self):
        shells = self.shells
        # center_ind, atomic_num, L, nprims
        centers = np.zeros((len(shells), 3))
        shell_data = list()
        coefficients = list()
        exponents = list()
        for i, shell in enumerate(shells):
            coeffs = shell.coeffs_org
            coefficients.extend(coeffs.tolist())
            exponents.extend(shell.exps.tolist())
            nprims = len(coeffs)

            centers[i] = shell.center
            atomic_num = shell.atomic_num
            if atomic_num is None:
                atomic_num = -1
            shell_data.append((shell.center_ind, atomic_num, shell.L, nprims))
        shell_data = np.array(shell_data, dtype=np.int32)
        coefficients = np.array(coefficients)
        exponents = np.array(exponents)
        return shell_data, centers, coefficients, exponents

    def dump_to_h5_group(self, h5_handle, group_name):
        group = h5_handle.create_group(group_name)
        shell_data, centers, coefficients, exponents = self.to_sympleints_arrays()
        group.create_dataset("shell_data", data=shell_data)
        group.create_dataset("centers", data=centers)
        group.create_dataset("coefficients", data=coefficients)
        group.create_dataset("exponents", data=exponents)

    @property
    def l_max(self):
        return max([shell.L for shell in self.shells])

    @property
    def atoms_coords3d(self):
        atoms = list()
        coords3d = list()
        center_inds = list()
        for shell in self.shells:
            center_ind = shell.center_ind
            # Skip cycle if we already registered this center
            if center_ind in center_inds:
                continue
            else:
                center_inds.append(center_ind)
            try:
                atom = INV_ATOMIC_NUMBERS[shell.atomic_num]
            # Use dummy atom when atomic_num is not set / None
            except KeyError:
                atom = "X"
            atoms.append(atom)
            center = shell.center
            coords3d.append(center)
        coords3d = np.array(coords3d)
        return atoms, coords3d

    @property
    def cart_bf_num(self):
        return sum([shell.size for shell in self.shells])

    def from_basis(self, name, shells_cls=None, **kwargs):
        from pysisyphus.wavefunction.Basis import shells_with_basis

        atoms, coords3d = self.atoms_coords3d
        if shells_cls is None:
            shells_cls = type(self)
        return shells_with_basis(
            atoms, coords3d, name=name, shells_cls=shells_cls, **kwargs
        )

    @staticmethod
    @file_or_str(".in")
    def from_aomix(text):
        # import here, to avoid cyclic imports
        from pysisyphus.io.aomix import shells_from_aomix

        shells = shells_from_aomix(text)
        return shells

    @staticmethod
    @file_or_str(".json")
    def from_orca_json(text):
        # import here, to avoid cyclic imports
        from pysisyphus.io.orca import shells_from_json

        shells = shells_from_json(text)
        return shells

    @staticmethod
    @file_or_str(".fchk")
    def from_fchk(text):
        # import here, to avoid cyclic imports
        from pysisyphus.io.fchk import shells_from_fchk

        shells = shells_from_fchk(text)
        return shells

    @staticmethod
    def from_sympleints_arrays(shell_data, centers, coefficients, exponents, **kwargs):
        shells = list()
        pointer = 0
        for i, (center_ind, atomic_num, L, nprims) in enumerate(shell_data):
            coeffs = coefficients[pointer : pointer + nprims]
            exps = exponents[pointer : pointer + nprims]
            shell = Shell(L, centers[i], coeffs, exps, center_ind, atomic_num)
            pointer += nprims
            shells.append(shell)
        return Shells(shells, **kwargs)

    @staticmethod
    def from_sympleints_h5(h5_fn, group_name="shells"):
        with h5py.File(h5_fn, "r") as handle:
            group = handle[group_name]
            shell_data = group["shell_data"][:]
            centers = group["centers"][:]
            coefficients = group["coefficients"][:]
            exponents = group["exponents"][:]
        return Shells.from_sympleints_arrays(
            shell_data, centers, coefficients, exponents
        )

    @staticmethod
    def from_pyscf_mol(mol):
        shells = list()
        for bas_id in range(mol.nbas):
            L = mol.bas_angular(bas_id)
            center = mol.bas_coord(bas_id)
            coeffs = mol.bas_ctr_coeff(bas_id).flatten()
            exps = mol.bas_exp(bas_id)
            assert (
                coeffs.size == exps.size
            ), "General contractions are not yet supported."
            center_ind = mol.bas_atom(bas_id)
            atom_symbol = mol.atom_symbol(center_ind)
            atomic_num = ATOMIC_NUMBERS[atom_symbol.lower()]
            shell = Shell(L, center, coeffs, exps, center_ind, atomic_num)
            shells.append(shell)
        return PySCFShells(shells)

    def to_pyscf_mol(self):
        # TODO: this would be better suited as a method of Wavefunction,
        # as pyscf Moles must have sensible spin & charge etc.
        # TODO: currently, this does not support different basis sets
        # for the same element.
        try:
            from pyscf import gto
        except ModuleNotFoundError:
            return None

        unique_atoms = set()
        basis = dict()
        atoms_coords = list()
        for _, center_shells in self.center_shell_iter():
            center_shells = list(center_shells)
            shell0 = center_shells[0]
            atom = shell0.atom
            x, y, z = shell0.center
            atoms_coords.append(f"{atom} {x} {y} {z}")
            if atom not in unique_atoms:
                unique_atoms.add(atom)
            else:
                continue
            atom_shells = "\n".join([shell.to_pyscf_shell() for shell in center_shells])
            basis[atom] = gto.basis.parse(atom_shells)

        mol = gto.Mole()
        mol.atom = "; ".join(atoms_coords)
        mol.unit = "Bohr"
        mol.basis = basis
        mol.build()
        return mol

    def center_shell_iter(self):
        sorted_shells = sorted(self.shells, key=lambda shell: shell.center_ind)
        return it.groupby(sorted_shells, key=lambda shell: shell.center_ind)

    def fragment_ao_map(self, fragments):
        frag_map = dict()
        for i, frag in enumerate(fragments):
            for center_ind in frag:
                frag_map[center_ind] = i

        frag_ao_map = dict()
        for i, ao_center in enumerate(self.sph_ao_centers):
            frag_ind = frag_map[ao_center]
            frag_ao_map.setdefault(frag_ind, list()).append(i)
        return frag_ao_map

    def as_gau_gbs(self) -> str:
        def dfmt(num):
            return f"{num: 12.10e}".replace("e", "D")

        gbs_tpl = Template(
            """
        {% for center_ind, shells in center_shell_iter %}
            {{ center_ind+1 }} 0
        {%- for shell in shells %}
        {{ ("S", "P", "D", "F", "G")[shell.L] }}   {{ shell.exps.size}}    1.00 0.000000000000
            {%- for exp_, coeff in shell.exps_coeffs_iter() %}
            {{ dfmt(exp_) }} {{ dfmt(coeff) }}
            {%- endfor -%}
        {% endfor %}
         ****
        {%- endfor %}
        """
        )
        rendered = gbs_tpl.render(center_shell_iter=self.center_shell_iter(), dfmt=dfmt)
        rendered = textwrap.dedent(rendered)
        return rendered

    def _ao_center_iter(self, func):
        i = 0
        for shell in self.shells:
            for _ in func(shell.L):
                yield shell.center_ind
                i += 0

    @property
    def cart_ao_centers(self):
        return self._ao_center_iter(canonical_order)

    @property
    def sph_ao_centers(self):
        return self._ao_center_iter(lambda l: range(2 * l + 1))

    @property
    def cart2sph_coeffs(self):
        cart2sph = cart2sph_coeffs(self.l_max)
        C = sp.linalg.block_diag(*[cart2sph[shell.L] for shell in self.shells])
        return C

    @property
    def P_sph(self):
        """Permutation matrix for spherical basis functions."""
        return sp.linalg.block_diag(*[self.sph_Ps[shell.L] for shell in self.shells])

    @property
    def P_cart(self):
        """Permutation matrix for Cartesian basis functions."""
        try:
            P_cart = sp.linalg.block_diag(
                *[self.cart_Ps[shell.L] for shell in self.shells]
            )
        except AttributeError:
            P_cart = np.eye(self.cart_size)
        return P_cart

    def eval(self, xyz, spherical=True):
        """Evaluate all basis functions at points xyz using generated code.

        A possibly more efficient approach is discussed in III C of
        https://doi.org/10.1063/1.469408.
        """
        if spherical:
            precontr = self.cart2sph_coeffs.T @ self.P_sph.T
        else:
            precontr = self.P_cart.T
        ncbfs = precontr.shape[0]
        npoints = len(xyz)

        coords3d = self.coords3d
        grid_vals = np.zeros((npoints, ncbfs))
        for center_ind, shells in self.center_shell_iter():
            # Create list from the shells iterator, otherwise it will be empty
            # after the first pass over all points.
            shells = list(shells)
            for i, R in enumerate(xyz):
                for shell in shells:
                    La, A, da, ax, a_ind, a_size = shell.as_tuple()
                    grid_vals[i, a_ind : a_ind + a_size] = cart_gto3d.cart_gto3d[(La,)](
                        ax, da, A, R
                    )
        return grid_vals @ precontr

    def get_1el_ints_cart(
        self,
        func_dict,
        other=None,
        can_reorder=True,
        components=0,
        screen_func=None,
        **kwargs,
    ):
        shells_a = self
        symmetric = other is None
        shells_b = shells_a if symmetric else other
        cart_bf_num_a = shells_a.cart_bf_num
        cart_bf_num_b = shells_b.cart_bf_num

        # components = 0 indicates, that a plain 2d matrix is desired.
        if is_2d := (components == 0):
            components = 1

        # Preallocate empty matrices and directly assign the calculated values
        integrals = np.zeros((components, cart_bf_num_a, cart_bf_num_b))

        # Dummy screen function that always returns True
        if (not self.screen) or (screen_func is None):
            screen_func = lambda *args: True

        for i, shell_a in enumerate(shells_a):
            La, A, da, ax, a_ind, a_size = shell_a.as_tuple()
            a_slice = slice(a_ind, a_ind + a_size)
            # When we don't deal with symmetric matrices we have to iterate over
            # all other basis functions in shells_b, so we reset i to 0.
            if not symmetric:
                i = 0
            a_min_exp = ax.min()  # Minimum exponent used for screening
            for shell_b in shells_b[i:]:
                Lb, B, db, bx, b_ind, b_size = shell_b.as_tuple()
                b_slice = slice(b_ind, b_ind + b_size)
                # Screen shell pair
                b_min_exp = bx.min()  # Minimum exponent used for screening
                R_ab = np.linalg.norm(A - B)  # Shell center distance
                b_size = get_shell_shape(Lb)[0]
                if not screen_func(a_min_exp, b_min_exp, R_ab):
                    continue
                integrals[:, a_slice, b_slice] = func_dict[(La, Lb)](
                    ax[:, None],
                    da[:, None],
                    A,
                    bx[None, :],
                    db[None, :],
                    B,
                    **kwargs,
                )
                # Fill up the lower tringular part of the matrix in the symmetric
                # case.
                # TODO: this may be (?) better done outside of this loop, after
                # everything is calculated...
                if symmetric:
                    integrals[:, b_slice, a_slice] = np.transpose(
                        integrals[:, a_slice, b_slice], axes=(0, 2, 1)
                    )

        # Return plain 2d array if components is set to 0, i.e., remove first axis.
        if is_2d:
            integrals = np.squeeze(integrals, axis=0)

        # Reordering will be disabled, when spherical integrals are desired. They
        # are reordered outside of this function. Reordering them already here
        # would mess up the results after the 2nd reordering.
        if can_reorder and self.ordering == "native":
            integrals = np.einsum(
                "ij,...jk,kl->...il",
                self.P_cart,
                integrals,
                shells_b.P_cart.T,
                optimize="greedy",
            )
        return integrals

    def get_1el_ints_sph(
        self, int_func=None, int_matrix=None, other=None, **kwargs
    ) -> NDArray:
        if int_matrix is None:
            # Disallow reordering of the Cartesian integrals
            int_matrix = self.get_1el_ints_cart(
                int_func, other=other, can_reorder=False, **kwargs
            )

        shells_b = self if other is None else other
        # Reordering (P) and Cartesian to spherical conversion (C).
        PC_a = self.reorder_c2s_coeffs
        PC_b = shells_b.reorder_c2s_coeffs
        int_matrix_sph = np.einsum(
            "ij,...jk,kl->...il", PC_a, int_matrix, PC_b.T, optimize="greedy"
        )
        return int_matrix_sph

    def get_2c2el_ints_cart(self):
        return self.get_1el_ints_cart(int2c2e3d.int2c2e3d)

    def get_2c2el_ints_sph(self):
        return self.get_1el_ints_sph(int2c2e3d.int2c2e3d)

    def get_3c2el_ints_cart(self, shells_aux):
        """Cartesian 3-center-2-electron integrals.

        DO NOT USE THESE INTEGRALS AS THEY ARE RETURNED FROM THIS METHOD.
        These integrals utilize recurrence relations that are only valid,
        when the resulting Cartesian integrals are transformed into spherical
        integrals.

        Contrary to the general function 'get_1el_ints_cart', that supports
        different 'func_dict' arguments and cross-integrals between two
        different shells this function is less general. This function is
        restricted to '_3center2el_sph' and always uses 'self.shells' as well as
        'self.shells_aux'.
        """
        shells_a = self
        cart_bf_num_a = shells_a.cart_bf_num
        cart_bf_num_aux = shells_aux.cart_bf_num

        integrals = np.zeros((cart_bf_num_a, cart_bf_num_a, cart_bf_num_aux))
        # func_dict = _3center2el3d_sph._3center2el3d_sph
        func_dict = int3c2e3d_sph.int3c2e3d_sph

        for i, shell_a in enumerate(shells_a):
            La, A, da, ax, a_ind, a_size = shell_a.as_tuple()
            a_slice = slice(a_ind, a_ind + a_size)
            # As noted in the docstring, we iterate over pairs of self.shells (shells_a)
            for shell_b in shells_a[i:]:
                Lb, B, db, bx, b_ind, b_size = shell_b.as_tuple()
                b_slice = slice(b_ind, b_ind + b_size)
                for shell_c in shells_aux:
                    Lc, C, dc, cx, c_ind, c_size = shell_c.as_tuple()
                    c_slice = slice(c_ind, c_ind + c_size)
                    integrals[a_slice, b_slice, c_slice] = func_dict[(La, Lb, Lc)](
                        ax[:, None, None],
                        da[:, None, None],
                        A,
                        bx[None, :, None],
                        db[None, :, None],
                        B,
                        cx[None, None, :],
                        dc[None, None, :],
                        C,
                    )
                    integrals[b_slice, a_slice, c_slice] = np.transpose(
                        integrals[a_slice, b_slice, c_slice], axes=(1, 0, 2)
                    )
        return integrals

    def get_3c2el_ints_sph(self, shells_aux):
        int_matrix = self.get_3c2el_ints_cart(shells_aux)

        PC_a = self.reorder_c2s_coeffs
        PC_c = shells_aux.reorder_c2s_coeffs
        int_matrix_sph = np.einsum(
            "ij,jlm,nl,om->ino", PC_a, int_matrix, PC_a, PC_c, optimize="greedy"
        )
        return int_matrix_sph

    #####################
    # Overlap integrals #
    #####################

    @staticmethod
    def screen_S(a, b, R, k=10):
        """
        Returns False when distance R is below calculated threshold.
        Derived from the sympy code below.

        from sympy import *

        a, b, S, R = symbols("a b S R", real=True, positive=True)
        k = symbols("k", integer=True, positive=True)
        apb = a + b
        # 0 = S - 10**(-k)
        ovlp = (pi / apb)**Rational(3,2) * exp(-a*b / apb * R**2) - 10**(-k)
        print("ss-overlap:", ovlp)
        # Distance R, when ss-overlaps drops below 10**(-k)
        R = solve(ovlp, R)[0]
        print("R:", R)

        Parameters
        ----------
        a
            Minimum exponent in shell a.
        b
            Minimum exponent in shell b.
        R
            Real space distance of shells a and b.
        """
        return R < (
            sqrt(a + b)
            * sqrt(log(10**k * pi ** (3 / 2) / (a + b) ** (3 / 2)))
            / (sqrt(a) * sqrt(b))
        )

    def get_S_cart(self, other=None) -> NDArray:
        # return self.get_1el_ints_cart(ovlp, other=other, screen_func=Shells.screen_S)
        return self.get_1el_ints_cart(
            ovlp3d.ovlp3d, other=other, screen_func=Shells.screen_S
        )

    def get_S_sph(self, other=None) -> NDArray:
        return self.get_1el_ints_sph(
            ovlp3d.ovlp3d, other=other, screen_func=Shells.screen_S
        )

    @property
    def S_cart(self):
        return self.get_S_cart()

    @property
    def S_sph(self):
        return self.get_S_sph()

    ############################
    # Kinetic energy integrals #
    ############################

    def get_T_cart(self, other=None) -> NDArray:
        # return self.get_1el_ints_cart(kinetic, other=other)
        return self.get_1el_ints_cart(kinetic3d.kinetic3d, other=other)

    def get_T_sph(self, other=None) -> NDArray:
        return self.get_1el_ints_sph(kinetic3d.kinetic3d, other=other)

    @property
    def T_cart(self):
        return self.get_T_cart()

    @property
    def T_sph(self):
        return self.get_T_sph()

    ################################
    # Nuclear attraction integrals #
    ################################

    def get_V_cart(self, coords3d=None, charges=None, **kwargs):
        """Nuclear attraction integrals.

        Coordinates and charges must be either be omitted or given together.
        Alternatively, this function could also take one array that combines
        coords3c and charges (not yet implemented).
        """
        assert ((coords3d is None) and (charges is None)) or (
            (coords3d is not None) and (charges is not None)
        )
        if coords3d is None:
            atoms, coords3d = self.atoms_coords3d
            charges = nuc_charges_for_atoms(atoms)

        # def boys_1c1el(n_max, ax, A, bx, B, C):
        # ax = ax[:, None]
        # bx = bx[None, :]
        # p = ax + bx
        # P = -(ax * A + bx * B) / p
        # R_PC2 = ((P - C)**2).sum()
        # pR_PC2 = p * R_PC2
        # return np.array([boys.boys(n, pR_PC2) for n in range(n_max+1)])

        cart_bf_num = self.cart_bf_num
        V_nuc = np.zeros((cart_bf_num, cart_bf_num))
        # Loop over all centers and add their contributions
        for R, Z in zip(coords3d, charges):
            # -Z = -1 * Z, because electrons have negative charge.
            V_nuc += -Z * self.get_1el_ints_cart(
                coulomb3d.coulomb3d,
                R=R,
                **kwargs,
            )
        return V_nuc

    def get_V_sph(self, coords3d=None, charges=None) -> NDArray:
        V_cart = self.get_V_cart(coords3d, charges, can_reorder=False)
        return self.get_1el_ints_sph(int_func=None, int_matrix=V_cart)

    @property
    def V_cart(self):
        return self.get_V_cart()

    @property
    def V_sph(self):
        return self.get_V_sph()

    ###########################
    # Dipole moment integrals #
    ###########################

    def get_dipole_ints_cart(self, origin):
        return self.get_1el_ints_cart(
            dipole3d.dipole3d, components=3, R=origin, screen_func=Shells.screen_S
        )

    def get_dipole_ints_sph(self, origin) -> NDArray:
        return self.get_1el_ints_sph(
            dipole3d.dipole3d, components=3, R=origin, screen_func=Shells.screen_S
        )

    ##################################################
    # Quadrupole moment integrals, diagonal elements #
    ##################################################

    def get_diag_quadrupole_ints_cart(self, origin):
        return self.get_1el_ints_cart(
            diag_quadrupole3d.diag_quadrupole3d,
            components=3,
            R=origin,
            screen_func=Shells.screen_S,
        )

    def get_diag_quadrupole_ints_sph(self, origin):
        return self.get_1el_ints_sph(
            diag_quadrupole3d.diag_quadrupole3d,
            components=3,
            R=origin,
            screen_func=Shells.screen_S,
        )

    ###############################
    # Quadrupole moment integrals #
    ###############################

    def get_quadrupole_ints_cart(self, origin):
        ints_flat = self.get_1el_ints_cart(
            quadrupole3d.quadrupole3d,
            components=6,
            R=origin,
            screen_func=Shells.screen_S,
        )
        return multi_component_sym_mat(ints_flat, 3)

    def get_quadrupole_ints_sph(self, origin) -> NDArray:
        ints_flat = self.get_1el_ints_sph(
            quadrupole3d.quadrupole3d,
            components=6,
            R=origin,
            screen_func=Shells.screen_S,
        )
        return multi_component_sym_mat(ints_flat, 3)

    def to_sap_shells(self):
        for shell in self.shells:
            shell.to_sap_shell()

    def __str__(self):
        return f"{self.__class__.__name__}({len(self.shells)} shells, ordering={self.ordering})"


class AOMixShells(Shells):
    cart_order = (
        ("",),
        ("x", "y", "z"),
        ("xx", "yy", "zz", "xy", "xz", "yz"),
        ("xxx", "yyy", "zzz", "xyy", "xxy", "xxz", "xzz", "yzz", "yyz", "xyz"),
        (
            "zzzz",
            "yzzz",
            "yyzz",
            "yyyz",
            "yyyy",
            "xzzz",
            "xyzz",
            "xyyz",
            "xyyy",
            "xxzz",
            "xxyz",
            "xxyy",
            "xxxz",
            "xxxy",
            "xxxx",
        ),
    )


class ORCAShells(Shells):
    sph_Ps = {
        0: [[1]],  # s
        1: [[0, 1, 0], [1, 0, 0], [0, 0, 1]],  # pz px py
        2: [
            [0, 0, 1, 0, 0],  # dz²
            [0, 1, 0, 0, 0],  # dxz
            [0, 0, 0, 1, 0],  # dyz
            [1, 0, 0, 0, 0],  # dx² - y²
            [0, 0, 0, 0, 1],  # dxy
        ],
        3: [
            [0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0],
            [-1, 0, 0, 0, 0, 0, 0],  # ORCA, why you do this to me?
            [0, 0, 0, 0, 0, 0, -1],  # sign flip, as in the line above
        ],
        4: [
            [0, 0, 0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, -1, 0, 0, 0, 0, 0, 0, 0],  # sign flip
            [0, 0, 0, 0, 0, 0, 0, -1, 0],  # sign flip
            [-1, 0, 0, 0, 0, 0, 0, 0, 0],  # sign flip
            [0, 0, 0, 0, 0, 0, 0, 0, -1],  # sign flip
        ],
    }


class FCHKShells(Shells):
    cart_order = (
        ("",),
        ("x", "y", "z"),
        ("xx", "yy", "zz", "xy", "xz", "yz"),
        ("xxx", "yyy", "zzz", "xyy", "xxy", "xxz", "xzz", "yzz", "yyz", "xyz"),
        (
            "zzzz",
            "yzzz",
            "yyzz",
            "yyyz",
            "yyyy",
            "xzzz",
            "xyzz",
            "xyyz",
            "xyyy",
            "xxzz",
            "xxyz",
            "xxyy",
            "xxxz",
            "xxxy",
            "xxxx",
        ),
    )

    sph_Ps = {
        0: [[1]],  # s
        1: [[1, 0, 0], [0, 0, 1], [0, 1, 0]],  # px pz py
        2: [
            [0, 0, 1, 0, 0],  # dz²
            [0, 1, 0, 0, 0],  # dxz
            [0, 0, 0, 1, 0],  # dyz
            [1, 0, 0, 0, 0],  # dx² - y²
            [0, 0, 0, 0, 1],  # dxy
        ],
        # Similar to ORCA, but w/o the sign flips ;)
        3: [
            [0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0],
            [1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1],
        ],
        4: [
            [0, 0, 0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 1, 0],
            [1, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 1],
        ],
    }


class MoldenShells(Shells):
    sph_Ps = {
        0: [[1]],  # s
        1: [[1, 0, 0], [0, 0, 1], [0, 1, 0]],  # px, py, pz
        2: [
            [0, 0, 1, 0, 0],  # dz²
            [0, 1, 0, 0, 0],  # dxz
            [0, 0, 0, 1, 0],  # dyz
            [1, 0, 0, 0, 0],  # dx² - y²
            [0, 0, 0, 0, 1],  # dxy
        ],
        3: [
            [0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0],
            [1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1],
        ],
        4: [
            [0, 0, 0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 1, 0],
            [1, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, 1],
        ],
    }


class ORCAMoldenShells(Shells):
    sph_Ps = {
        0: [[1]],  # s
        1: [[1, 0, 0], [0, 0, 1], [0, 1, 0]],  # px, py, pz
        2: [
            [0, 0, 1, 0, 0],  # dz²
            [0, 1, 0, 0, 0],  # dxz
            [0, 0, 0, 1, 0],  # dyz
            [1, 0, 0, 0, 0],  # dx² - y²
            [0, 0, 0, 0, 1],  # dxy
        ],
        3: [
            [0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0],
            [0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0],
            [-1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, -1],
        ],
        4: [
            [0, 0, 0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, -1, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, -1, 0],
            [-1, 0, 0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0, 0, -1],
        ],
    }


def pyscf_cart_order(l):
    order = list()
    for lx in reversed(range(l + 1)):
        for ly in reversed(range(l + 1 - lx)):
            lz = l - lx - ly
            order.append("x" * lx + "y" * ly + "z" * lz)
    return tuple(order)


class PySCFShells(Shells):

    """
    Cartesian bfs >= d angular momentum are not normalized!
    S_ref = mol.intor("int1e_ovlp_cart")
    N = 1 / np.diag(S_ref)**0.5
    ao *= N
    """

    cart_order = [pyscf_cart_order(l) for l in range(5)]

    sph_Ps = {
        0: [[1]],  # s
        1: [[1, 0, 0], [0, 0, 1], [0, 1, 0]],  # px py pz
        2: [
            [0, 0, 0, 0, 1],  # dxy
            [0, 0, 0, 1, 0],  # dyz
            [0, 0, 1, 0, 0],  # dz²
            [0, 1, 0, 0, 0],  # dxz
            [1, 0, 0, 0, 0],  # dx² - y²
        ],
        3: [
            [0, 0, 0, 0, 0, 0, 1],
            [0, 0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 1, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0, 0],
            [1, 0, 0, 0, 0, 0, 0],
        ],
        4: [
            [0, 0, 0, 0, 0, 0, 0, 0, 1],
            [0, 0, 0, 0, 0, 0, 0, 1, 0],
            [0, 0, 0, 0, 0, 0, 1, 0, 0],
            [0, 0, 0, 0, 0, 1, 0, 0, 0],
            [0, 0, 0, 0, 1, 0, 0, 0, 0],
            [0, 0, 0, 1, 0, 0, 0, 0, 0],
            [0, 0, 1, 0, 0, 0, 0, 0, 0],
            [0, 1, 0, 0, 0, 0, 0, 0, 0],
            [1, 0, 0, 0, 0, 0, 0, 0, 0],
        ],
    }
