# fmt: off

"""Filters"""
from functools import cached_property
from itertools import product
from warnings import warn

import numpy as np

from ase.calculators.calculator import PropertyNotImplementedError
from ase.stress import full_3x3_to_voigt_6_stress, voigt_6_to_full_3x3_stress
from ase.utils import deprecated
from ase.utils.abc import Optimizable

__all__ = [
    'Filter', 'StrainFilter', 'UnitCellFilter', 'FrechetCellFilter',
    'ExpCellFilter'
]


class OptimizableFilter(Optimizable):
    def __init__(self, filterobj):
        self.filterobj = filterobj

    def get_x(self):
        return self.filterobj.get_positions().ravel()

    def set_x(self, x):
        self.filterobj.set_positions(x.reshape(-1, 3))

    def get_gradient(self):
        return self.filterobj.get_forces().ravel()

    @cached_property
    def _use_force_consistent_energy(self):
        # This boolean is in principle invalidated if the
        # calculator changes.  This can lead to weird things
        # in multi-step optimizations.
        try:
            self.filterobj.get_potential_energy(force_consistent=True)
        except PropertyNotImplementedError:
            return False
        else:
            return True

    def get_value(self):
        force_consistent = self._use_force_consistent_energy
        return self.filterobj.get_potential_energy(
            force_consistent=force_consistent)

    def ndofs(self):
        return 3 * len(self.filterobj)

    def iterimages(self):
        return self.filterobj.iterimages()


class Filter:
    """Subset filter class."""

    def __init__(self, atoms, indices=None, mask=None):
        """Filter atoms.

        This filter can be used to hide degrees of freedom in an Atoms
        object.

        Parameters
        ----------
        indices : list of int
           Indices for those atoms that should remain visible.
        mask : list of bool
           One boolean per atom indicating if the atom should remain
           visible or not.

        If a Trajectory tries to save this object, it will instead
        save the underlying Atoms object.  To prevent this, override
        the iterimages method.
        """

        self.atoms = atoms
        self.constraints = []
        # Make self.info a reference to the underlying atoms' info dictionary.
        self.info = self.atoms.info

        if indices is None and mask is None:
            raise ValueError('Use "indices" or "mask".')
        if indices is not None and mask is not None:
            raise ValueError('Use only one of "indices" and "mask".')

        if mask is not None:
            self.index = np.asarray(mask, bool)
            self.n = self.index.sum()
        else:
            self.index = np.asarray(indices, int)
            self.n = len(self.index)

    def iterimages(self):
        # Present the real atoms object to Trajectory and friends
        return self.atoms.iterimages()

    def get_cell(self):
        """Returns the computational cell.

        The computational cell is the same as for the original system.
        """
        return self.atoms.get_cell()

    def get_pbc(self):
        """Returns the periodic boundary conditions.

        The boundary conditions are the same as for the original system.
        """
        return self.atoms.get_pbc()

    def get_positions(self):
        'Return the positions of the visible atoms.'
        return self.atoms.get_positions()[self.index]

    def set_positions(self, positions, **kwargs):
        'Set the positions of the visible atoms.'
        pos = self.atoms.get_positions()
        pos[self.index] = positions
        self.atoms.set_positions(pos, **kwargs)

    positions = property(get_positions, set_positions,
                         doc='Positions of the atoms')

    def get_momenta(self):
        'Return the momenta of the visible atoms.'
        return self.atoms.get_momenta()[self.index]

    def set_momenta(self, momenta, **kwargs):
        'Set the momenta of the visible atoms.'
        mom = self.atoms.get_momenta()
        mom[self.index] = momenta
        self.atoms.set_momenta(mom, **kwargs)

    def get_atomic_numbers(self):
        'Return the atomic numbers of the visible atoms.'
        return self.atoms.get_atomic_numbers()[self.index]

    def set_atomic_numbers(self, atomic_numbers):
        'Set the atomic numbers of the visible atoms.'
        z = self.atoms.get_atomic_numbers()
        z[self.index] = atomic_numbers
        self.atoms.set_atomic_numbers(z)

    def get_tags(self):
        'Return the tags of the visible atoms.'
        return self.atoms.get_tags()[self.index]

    def set_tags(self, tags):
        'Set the tags of the visible atoms.'
        tg = self.atoms.get_tags()
        tg[self.index] = tags
        self.atoms.set_tags(tg)

    def get_forces(self, *args, **kwargs):
        return self.atoms.get_forces(*args, **kwargs)[self.index]

    def get_stress(self, *args, **kwargs):
        return self.atoms.get_stress(*args, **kwargs)

    def get_stresses(self, *args, **kwargs):
        return self.atoms.get_stresses(*args, **kwargs)[self.index]

    def get_masses(self):
        return self.atoms.get_masses()[self.index]

    def get_potential_energy(self, **kwargs):
        """Calculate potential energy.

        Returns the potential energy of the full system.
        """
        return self.atoms.get_potential_energy(**kwargs)

    def get_chemical_symbols(self):
        return self.atoms.get_chemical_symbols()

    def get_initial_magnetic_moments(self):
        return self.atoms.get_initial_magnetic_moments()

    def get_calculator(self):
        """Returns the calculator.

        WARNING: The calculator is unaware of this filter, and sees a
        different number of atoms.
        """
        return self.atoms.calc

    @property
    def calc(self):
        return self.atoms.calc

    def get_celldisp(self):
        return self.atoms.get_celldisp()

    def has(self, name):
        'Check for existence of array.'
        return self.atoms.has(name)

    def __len__(self):
        'Return the number of movable atoms.'
        return self.n

    def __getitem__(self, i):
        'Return an atom.'
        return self.atoms[self.index[i]]

    def __ase_optimizable__(self):
        return OptimizableFilter(self)


class StrainFilter(Filter):
    """Modify the supercell while keeping the scaled positions fixed.

    Presents the strain of the supercell as the generalized positions,
    and the global stress tensor (times the volume) as the generalized
    force.

    This filter can be used to relax the unit cell until the stress is
    zero.  If MDMin is used for this, the timestep (dt) to be used
    depends on the system size. 0.01/x where x is a typical dimension
    seems like a good choice.

    The stress and strain are presented as 6-vectors, the order of the
    components follow the standard engingeering practice: xx, yy, zz,
    yz, xz, xy.

    """

    def __init__(self, atoms, mask=None, include_ideal_gas=False):
        """Create a filter applying a homogeneous strain to a list of atoms.

        The first argument, atoms, is the atoms object.

        The optional second argument, mask, is a list of six booleans,
        indicating which of the six independent components of the
        strain that are allowed to become non-zero.  It defaults to
        [1,1,1,1,1,1].

        """

        self.strain = np.zeros(6)
        self.include_ideal_gas = include_ideal_gas

        if mask is None:
            mask = np.ones(6)
        else:
            mask = np.array(mask)

        Filter.__init__(self, atoms=atoms, mask=mask)
        self.mask = mask
        self.origcell = atoms.get_cell()

    def get_positions(self):
        return self.strain.reshape((2, 3)).copy()

    def set_positions(self, new):
        new = new.ravel() * self.mask
        eps = np.array([[1.0 + new[0], 0.5 * new[5], 0.5 * new[4]],
                        [0.5 * new[5], 1.0 + new[1], 0.5 * new[3]],
                        [0.5 * new[4], 0.5 * new[3], 1.0 + new[2]]])

        self.atoms.set_cell(np.dot(self.origcell, eps), scale_atoms=True)
        self.strain[:] = new

    def get_forces(self, **kwargs):
        stress = self.atoms.get_stress(include_ideal_gas=self.include_ideal_gas)
        return -self.atoms.get_volume() * (stress * self.mask).reshape((2, 3))

    def has(self, x):
        return self.atoms.has(x)

    def __len__(self):
        return 2


class UnitCellFilter(Filter):
    """Modify the supercell and the atom positions. """

    def __init__(self, atoms, mask=None,
                 cell_factor=None,
                 hydrostatic_strain=False,
                 constant_volume=False,
                 orig_cell=None,
                 scalar_pressure=0.0):
        """Create a filter that returns the atomic forces and unit cell
        stresses together, so they can simultaneously be minimized.

        The first argument, atoms, is the atoms object. The optional second
        argument, mask, is a list of booleans, indicating which of the six
        independent components of the strain are relaxed.

        - True = relax to zero
        - False = fixed, ignore this component

        Degrees of freedom are the positions in the original undeformed cell,
        plus the deformation tensor (extra 3 "atoms"). This gives forces
        consistent with numerical derivatives of the potential energy
        with respect to the cell degreees of freedom.

        For full details see:
            E. B. Tadmor, G. S. Smith, N. Bernstein, and E. Kaxiras,
            Phys. Rev. B 59, 235 (1999)

        You can still use constraints on the atoms, e.g. FixAtoms, to control
        the relaxation of the atoms.

        >>> # this should be equivalent to the StrainFilter
        >>> atoms = Atoms(...)
        >>> atoms.set_constraint(FixAtoms(mask=[True for atom in atoms]))
        >>> ucf = UnitCellFilter(atoms)

        You should not attach this UnitCellFilter object to a
        trajectory. Instead, create a trajectory for the atoms, and
        attach it to an optimizer like this:

        >>> atoms = Atoms(...)
        >>> ucf = UnitCellFilter(atoms)
        >>> qn = QuasiNewton(ucf)
        >>> traj = Trajectory('TiO2.traj', 'w', atoms)
        >>> qn.attach(traj)
        >>> qn.run(fmax=0.05)

        Helpful conversion table:

        - 0.05 eV/A^3   = 8 GPA
        - 0.003 eV/A^3  = 0.48 GPa
        - 0.0006 eV/A^3 = 0.096 GPa
        - 0.0003 eV/A^3 = 0.048 GPa
        - 0.0001 eV/A^3 = 0.02 GPa

        Additional optional arguments:

        cell_factor: float (default float(len(atoms)))
            Factor by which deformation gradient is multiplied to put
            it on the same scale as the positions when assembling
            the combined position/cell vector. The stress contribution to
            the forces is scaled down by the same factor. This can be thought
            of as a very simple preconditioners. Default is number of atoms
            which gives approximately the correct scaling.

        hydrostatic_strain: bool (default False)
            Constrain the cell by only allowing hydrostatic deformation.
            The virial tensor is replaced by np.diag([np.trace(virial)]*3).

        constant_volume: bool (default False)
            Project out the diagonal elements of the virial tensor to allow
            relaxations at constant volume, e.g. for mapping out an
            energy-volume curve. Note: this only approximately conserves
            the volume and breaks energy/force consistency so can only be
            used with optimizers that do require do a line minimisation
            (e.g. FIRE).

        scalar_pressure: float (default 0.0)
            Applied pressure to use for enthalpy pV term. As above, this
            breaks energy/force consistency.
        """

        Filter.__init__(self, atoms=atoms, indices=range(len(atoms)))
        self.atoms = atoms
        if orig_cell is None:
            self.orig_cell = atoms.get_cell()
        else:
            self.orig_cell = orig_cell
        self.stress = None

        if mask is None:
            mask = np.ones(6)
        mask = np.asarray(mask)
        if mask.shape == (6,):
            self.mask = voigt_6_to_full_3x3_stress(mask)
        elif mask.shape == (3, 3):
            self.mask = mask
        else:
            raise ValueError('shape of mask should be (3,3) or (6,)')

        if cell_factor is None:
            cell_factor = float(len(atoms))
        self.hydrostatic_strain = hydrostatic_strain
        self.constant_volume = constant_volume
        self.scalar_pressure = scalar_pressure
        self.cell_factor = cell_factor
        self.copy = self.atoms.copy
        self.arrays = self.atoms.arrays

    def deform_grad(self):
        return np.linalg.solve(self.orig_cell, self.atoms.cell).T

    def get_positions(self):
        """
        this returns an array with shape (natoms + 3,3).

        the first natoms rows are the positions of the atoms, the last
        three rows are the deformation tensor associated with the unit cell,
        scaled by self.cell_factor.
        """

        cur_deform_grad = self.deform_grad()
        natoms = len(self.atoms)
        pos = np.zeros((natoms + 3, 3))
        # UnitCellFilter's positions are the self.atoms.positions but without
        # the applied deformation gradient
        pos[:natoms] = np.linalg.solve(cur_deform_grad,
                                       self.atoms.positions.T).T
        # UnitCellFilter's cell DOFs are the deformation gradient times a
        # scaling factor
        pos[natoms:] = self.cell_factor * cur_deform_grad
        return pos

    def set_positions(self, new, **kwargs):
        """
        new is an array with shape (natoms+3,3).

        the first natoms rows are the positions of the atoms, the last
        three rows are the deformation tensor used to change the cell shape.

        the new cell is first set from original cell transformed by the new
        deformation gradient, then the positions are set with respect to the
        current cell by transforming them with the same deformation gradient
        """

        natoms = len(self.atoms)
        new_atom_positions = new[:natoms]
        new_deform_grad = new[natoms:] / self.cell_factor
        deform = (new_deform_grad - np.eye(3)).T * self.mask
        # Set the new cell from the original cell and the new
        # deformation gradient.  Both current and final structures should
        # preserve symmetry, so if set_cell() calls FixSymmetry.adjust_cell(),
        # it should be OK
        newcell = self.orig_cell @ (np.eye(3) + deform)

        self.atoms.set_cell(newcell,
                            scale_atoms=True)
        # Set the positions from the ones passed in (which are without the
        # deformation gradient applied) and the new deformation gradient.
        # This should also preserve symmetry, so if set_positions() calls
        # FixSymmetyr.adjust_positions(), it should be OK
        self.atoms.set_positions(new_atom_positions @ (np.eye(3) + deform),
                                 **kwargs)

    def get_potential_energy(self, force_consistent=True):
        """
        returns potential energy including enthalpy PV term.
        """
        atoms_energy = self.atoms.get_potential_energy(
            force_consistent=force_consistent)
        return atoms_energy + self.scalar_pressure * self.atoms.get_volume()

    def get_forces(self, **kwargs):
        """
        returns an array with shape (natoms+3,3) of the atomic forces
        and unit cell stresses.

        the first natoms rows are the forces on the atoms, the last
        three rows are the forces on the unit cell, which are
        computed from the stress tensor.
        """

        stress = self.atoms.get_stress(**kwargs)
        atoms_forces = self.atoms.get_forces(**kwargs)

        volume = self.atoms.get_volume()
        virial = -volume * (voigt_6_to_full_3x3_stress(stress) +
                            np.diag([self.scalar_pressure] * 3))
        cur_deform_grad = self.deform_grad()
        atoms_forces = atoms_forces @ cur_deform_grad
        virial = np.linalg.solve(cur_deform_grad, virial.T).T

        if self.hydrostatic_strain:
            vtr = virial.trace()
            virial = np.diag([vtr / 3.0, vtr / 3.0, vtr / 3.0])

        # Zero out components corresponding to fixed lattice elements
        if (self.mask != 1.0).any():
            virial *= self.mask

        if self.constant_volume:
            vtr = virial.trace()
            np.fill_diagonal(virial, np.diag(virial) - vtr / 3.0)

        natoms = len(self.atoms)
        forces = np.zeros((natoms + 3, 3))
        forces[:natoms] = atoms_forces
        forces[natoms:] = virial / self.cell_factor

        self.stress = -full_3x3_to_voigt_6_stress(virial) / volume
        return forces

    def get_stress(self):
        raise PropertyNotImplementedError

    def has(self, x):
        return self.atoms.has(x)

    def __len__(self):
        return (len(self.atoms) + 3)


class FrechetCellFilter(UnitCellFilter):
    """Modify the supercell and the atom positions."""

    def __init__(self, atoms, mask=None,
                 exp_cell_factor=None,
                 hydrostatic_strain=False,
                 constant_volume=False,
                 scalar_pressure=0.0):
        r"""Create a filter that returns the atomic forces and unit cell
        stresses together, so they can simultaneously be minimized.

        The first argument, atoms, is the atoms object. The optional second
        argument, mask, is a list of booleans, indicating which of the six
        independent components of the strain are relaxed.

        - True = relax to zero
        - False = fixed, ignore this component

        Degrees of freedom are the positions in the original undeformed cell,
        plus the log of the deformation tensor (extra 3 "atoms"). This gives
        forces consistent with numerical derivatives of the potential energy
        with respect to the cell degrees of freedom.

        You can still use constraints on the atoms, e.g. FixAtoms, to control
        the relaxation of the atoms.

        >>> # this should be equivalent to the StrainFilter
        >>> atoms = Atoms(...)
        >>> atoms.set_constraint(FixAtoms(mask=[True for atom in atoms]))
        >>> ecf = FrechetCellFilter(atoms)

        You should not attach this FrechetCellFilter object to a
        trajectory. Instead, create a trajectory for the atoms, and
        attach it to an optimizer like this:

        >>> atoms = Atoms(...)
        >>> ecf = FrechetCellFilter(atoms)
        >>> qn = QuasiNewton(ecf)
        >>> traj = Trajectory('TiO2.traj', 'w', atoms)
        >>> qn.attach(traj)
        >>> qn.run(fmax=0.05)

        Helpful conversion table:

        - 0.05 eV/A^3   = 8 GPA
        - 0.003 eV/A^3  = 0.48 GPa
        - 0.0006 eV/A^3 = 0.096 GPa
        - 0.0003 eV/A^3 = 0.048 GPa
        - 0.0001 eV/A^3 = 0.02 GPa

        Additional optional arguments:

        exp_cell_factor: float (default float(len(atoms)))
            Scaling factor for cell variables. The cell gradients in
            FrechetCellFilter.get_forces() is divided by exp_cell_factor.
            By default, set the number of atoms. We recommend to set
            an extensive value for this parameter.

        hydrostatic_strain: bool (default False)
            Constrain the cell by only allowing hydrostatic deformation.
            The virial tensor is replaced by np.diag([np.trace(virial)]*3).

        constant_volume: bool (default False)
            Project out the diagonal elements of the virial tensor to allow
            relaxations at constant volume, e.g. for mapping out an
            energy-volume curve.

        scalar_pressure: float (default 0.0)
            Applied pressure to use for enthalpy pV term. As above, this
            breaks energy/force consistency.

        Implementation note:

        The implementation is based on that of Christoph Ortner in JuLIP.jl:
        https://github.com/JuliaMolSim/JuLIP.jl/blob/master/src/expcell.jl

        The initial implementation of ExpCellFilter gave inconsistent gradients
        for cell variables (matrix log of the deformation tensor). If you would
        like to keep the previous behavior, please use ExpCellFilter.

        The derivation of gradients of energy w.r.t positions and the log of the
        deformation tensor is given in
        https://github.com/lan496/lan496.github.io/blob/main/notes/cell_grad.pdf
        """

        Filter.__init__(self, atoms=atoms, indices=range(len(atoms)))
        UnitCellFilter.__init__(self, atoms=atoms, mask=mask,
                                hydrostatic_strain=hydrostatic_strain,
                                constant_volume=constant_volume,
                                scalar_pressure=scalar_pressure)

        # We defer the scipy import to avoid high immediate import overhead
        from scipy.linalg import expm, expm_frechet, logm
        self.expm = expm
        self.logm = logm
        self.expm_frechet = expm_frechet

        # Scaling factor for cell gradients
        if exp_cell_factor is None:
            exp_cell_factor = float(len(atoms))
        self.exp_cell_factor = exp_cell_factor

    def get_positions(self):
        pos = UnitCellFilter.get_positions(self)
        natoms = len(self.atoms)
        pos[natoms:] = self.logm(pos[natoms:]) * self.exp_cell_factor
        return pos

    def set_positions(self, new, **kwargs):
        natoms = len(self.atoms)
        new2 = new.copy()
        new2[natoms:] = self.expm(new[natoms:] / self.exp_cell_factor)
        UnitCellFilter.set_positions(self, new2, **kwargs)

    def get_forces(self, **kwargs):
        # forces on atoms are same as UnitCellFilter, we just
        # need to modify the stress contribution
        stress = self.atoms.get_stress(**kwargs)
        volume = self.atoms.get_volume()
        virial = -volume * (voigt_6_to_full_3x3_stress(stress) +
                            np.diag([self.scalar_pressure] * 3))

        cur_deform_grad = self.deform_grad()
        cur_deform_grad_log = self.logm(cur_deform_grad)

        if self.hydrostatic_strain:
            vtr = virial.trace()
            virial = np.diag([vtr / 3.0, vtr / 3.0, vtr / 3.0])

        # Zero out components corresponding to fixed lattice elements
        if (self.mask != 1.0).any():
            virial *= self.mask

        # Cell gradient for UnitCellFilter
        ucf_cell_grad = virial @ np.linalg.inv(cur_deform_grad.T)

        # Cell gradient for FrechetCellFilter
        deform_grad_log_force = np.zeros((3, 3))
        for mu, nu in product(range(3), repeat=2):
            dir = np.zeros((3, 3))
            dir[mu, nu] = 1.0
            # Directional derivative of deformation to (mu, nu) strain direction
            expm_der = self.expm_frechet(
                cur_deform_grad_log,
                dir,
                compute_expm=False
            )
            deform_grad_log_force[mu, nu] = np.sum(expm_der * ucf_cell_grad)

        # Cauchy stress used for convergence testing
        convergence_crit_stress = -(virial / volume)
        if self.constant_volume:
            # apply constraint to force
            dglf_trace = deform_grad_log_force.trace()
            np.fill_diagonal(deform_grad_log_force,
                             np.diag(deform_grad_log_force) - dglf_trace / 3.0)
            # apply constraint to Cauchy stress used for convergence testing
            ccs_trace = convergence_crit_stress.trace()
            np.fill_diagonal(convergence_crit_stress,
                             np.diag(convergence_crit_stress) - ccs_trace / 3.0)

        atoms_forces = self.atoms.get_forces(**kwargs)
        atoms_forces = atoms_forces @ cur_deform_grad

        # pack gradients into vector
        natoms = len(self.atoms)
        forces = np.zeros((natoms + 3, 3))
        forces[:natoms] = atoms_forces
        forces[natoms:] = deform_grad_log_force / self.exp_cell_factor
        self.stress = full_3x3_to_voigt_6_stress(convergence_crit_stress)
        return forces


class ExpCellFilter(UnitCellFilter):

    @deprecated(DeprecationWarning(
        'Use FrechetCellFilter for better convergence w.r.t. cell variables.'
    ))
    def __init__(self, atoms, mask=None,
                 cell_factor=None,
                 hydrostatic_strain=False,
                 constant_volume=False,
                 scalar_pressure=0.0):
        r"""Create a filter that returns the atomic forces and unit cell
            stresses together, so they can simultaneously be minimized.

            The first argument, atoms, is the atoms object. The optional second
            argument, mask, is a list of booleans, indicating which of the six
            independent components of the strain are relaxed.

            - True = relax to zero
            - False = fixed, ignore this component

            Degrees of freedom are the positions in the original undeformed
            cell, plus the log of the deformation tensor (extra 3 "atoms"). This
            gives forces consistent with numerical derivatives of the potential
            energy with respect to the cell degrees of freedom.

            For full details see:
                E. B. Tadmor, G. S. Smith, N. Bernstein, and E. Kaxiras,
                Phys. Rev. B 59, 235 (1999)

            You can still use constraints on the atoms, e.g. FixAtoms, to
            control the relaxation of the atoms.

            >>> # this should be equivalent to the StrainFilter
            >>> atoms = Atoms(...)
            >>> atoms.set_constraint(FixAtoms(mask=[True for atom in atoms]))
            >>> ecf = ExpCellFilter(atoms)

            You should not attach this ExpCellFilter object to a
            trajectory. Instead, create a trajectory for the atoms, and
            attach it to an optimizer like this:

            >>> atoms = Atoms(...)
            >>> ecf = ExpCellFilter(atoms)
            >>> qn = QuasiNewton(ecf)
            >>> traj = Trajectory('TiO2.traj', 'w', atoms)
            >>> qn.attach(traj)
            >>> qn.run(fmax=0.05)

            Helpful conversion table:

            - 0.05 eV/A^3   = 8 GPA
            - 0.003 eV/A^3  = 0.48 GPa
            - 0.0006 eV/A^3 = 0.096 GPa
            - 0.0003 eV/A^3 = 0.048 GPa
            - 0.0001 eV/A^3 = 0.02 GPa

            Additional optional arguments:

            cell_factor: (DEPRECATED)
                Retained for backwards compatibility, but no longer used.

            hydrostatic_strain: bool (default False)
                Constrain the cell by only allowing hydrostatic deformation.
                The virial tensor is replaced by np.diag([np.trace(virial)]*3).

            constant_volume: bool (default False)
                Project out the diagonal elements of the virial tensor to allow
                relaxations at constant volume, e.g. for mapping out an
                energy-volume curve.

            scalar_pressure: float (default 0.0)
                Applied pressure to use for enthalpy pV term. As above, this
                breaks energy/force consistency.

        Implementation details:

        The implementation is based on that of Christoph Ortner in JuLIP.jl:
        https://github.com/libAtoms/JuLIP.jl/blob/expcell/src/Constraints.jl#L244

        We decompose the deformation gradient as

            F = exp(U) F0
            x =  F * F0^{-1} z  = exp(U) z

        If we write the energy as a function of U we can transform the
        stress associated with a perturbation V into a derivative using a
        linear map V -> L(U, V).

        \phi( exp(U+tV) (z+tv) ) ~ \phi'(x) . (exp(U) v) + \phi'(x) .
                                                    ( L(U, V) exp(-U) exp(U) z )

        where

                \nabla E(U) : V  =  [S exp(-U)'] : L(U,V)
                                =  L'(U, S exp(-U)') : V
                                =  L(U', S exp(-U)') : V
                                =  L(U, S exp(-U)) : V     (provided U = U')

        where the : operator represents double contraction,
        i.e. A:B = trace(A'B), and

            F = deformation tensor - 3x3 matrix
            F0 = reference deformation tensor - 3x3 matrix, np.eye(3) here
            U = cell degrees of freedom used here - 3x3 matrix
            V = perturbation to cell DoFs - 3x3 matrix
            v = perturbation to position DoFs
            x = atomic positions in deformed cell
            z = atomic positions in original cell
            \phi = potential energy
            S = stress tensor [3x3 matrix]
            L(U, V) = directional derivative of exp at U in direction V, i.e
            d/dt exp(U + t V)|_{t=0} = L(U, V)

        This means we can write

            d/dt E(U + t V)|_{t=0} = L(U, S exp (-U)) : V

        and therefore the contribution to the gradient of the energy is

            \nabla E(U) / \nabla U_ij =  [L(U, S exp(-U))]_ij

        .. deprecated:: 3.23.0
            Use :class:`~ase.filters.FrechetCellFilter` for better convergence
            w.r.t. cell variables.
        """
        Filter.__init__(self, atoms=atoms, indices=range(len(atoms)))
        UnitCellFilter.__init__(self, atoms=atoms, mask=mask,
                                cell_factor=cell_factor,
                                hydrostatic_strain=hydrostatic_strain,
                                constant_volume=constant_volume,
                                scalar_pressure=scalar_pressure)
        if cell_factor is not None:
            # cell_factor used in UnitCellFilter does not affect on gradients of
            # ExpCellFilter.
            warn("cell_factor is deprecated")
        self.cell_factor = 1.0

        # We defer the scipy import to avoid high immediate import overhead
        from scipy.linalg import expm, logm
        self.expm = expm
        self.logm = logm

    def get_forces(self, **kwargs):
        forces = UnitCellFilter.get_forces(self, **kwargs)

        # forces on atoms are same as UnitCellFilter, we just
        # need to modify the stress contribution
        stress = self.atoms.get_stress(**kwargs)
        volume = self.atoms.get_volume()
        virial = -volume * (voigt_6_to_full_3x3_stress(stress) +
                            np.diag([self.scalar_pressure] * 3))

        cur_deform_grad = self.deform_grad()
        cur_deform_grad_log = self.logm(cur_deform_grad)

        if self.hydrostatic_strain:
            vtr = virial.trace()
            virial = np.diag([vtr / 3.0, vtr / 3.0, vtr / 3.0])

        # Zero out components corresponding to fixed lattice elements
        if (self.mask != 1.0).any():
            virial *= self.mask

        deform_grad_log_force_naive = virial.copy()
        Y = np.zeros((6, 6))
        Y[0:3, 0:3] = cur_deform_grad_log
        Y[3:6, 3:6] = cur_deform_grad_log
        Y[0:3, 3:6] = - virial @ self.expm(-cur_deform_grad_log)
        deform_grad_log_force = -self.expm(Y)[0:3, 3:6]
        for (i1, i2) in [(0, 1), (0, 2), (1, 2)]:
            ff = 0.5 * (deform_grad_log_force[i1, i2] +
                        deform_grad_log_force[i2, i1])
            deform_grad_log_force[i1, i2] = ff
            deform_grad_log_force[i2, i1] = ff

        # check for reasonable alignment between naive and
        # exact search directions
        all_are_equal = np.all(np.isclose(deform_grad_log_force,
                                          deform_grad_log_force_naive))
        if all_are_equal or \
            (np.sum(deform_grad_log_force * deform_grad_log_force_naive) /
             np.sqrt(np.sum(deform_grad_log_force**2) *
                     np.sum(deform_grad_log_force_naive**2)) > 0.8):
            deform_grad_log_force = deform_grad_log_force_naive

        # Cauchy stress used for convergence testing
        convergence_crit_stress = -(virial / volume)
        if self.constant_volume:
            # apply constraint to force
            dglf_trace = deform_grad_log_force.trace()
            np.fill_diagonal(deform_grad_log_force,
                             np.diag(deform_grad_log_force) - dglf_trace / 3.0)
            # apply constraint to Cauchy stress used for convergence testing
            ccs_trace = convergence_crit_stress.trace()
            np.fill_diagonal(convergence_crit_stress,
                             np.diag(convergence_crit_stress) - ccs_trace / 3.0)

        # pack gradients into vector
        natoms = len(self.atoms)
        forces[natoms:] = deform_grad_log_force
        self.stress = full_3x3_to_voigt_6_stress(convergence_crit_stress)
        return forces