import logging
from functools import partial

import numpy as np
import torch
from torch import nn

from deepqmc import Molecule
from deepqmc.physics import pairwise_diffs, pairwise_distance
from deepqmc.torchext import sloglindet, triu_flat
from deepqmc.utils import NULL_DEBUG
from deepqmc.wf import WaveFunction

from .cusp import CuspCorrection, ElectronicAsymptotic
from .distbasis import DistanceBasis
from .gto import GTOBasis
from .molorb import MolecularOrbital
from .omni import OmniSchNet

__version__ = '0.1.0'
__all__ = ['PauliNet']

log = logging.getLogger(__name__)


def eval_slater(xs):
    if xs.shape[-1] == 0:
        return xs.new_ones(xs.shape[:-2])
    return torch.det(xs.contiguous())


def eval_log_slater(xs):
    if xs.shape[-1] == 0:
        return xs.new_ones(xs.shape[:-2]), xs.new_zeros(xs.shape[:-2])
    return xs.contiguous().slogdet()

def eval_vandermonde(xs):
    print('non-log vandermonde is evaluated')
    product_of_wavefunctions = xs.prod(dim=-1)
    total = 1.0
    if product_of_wavefunctions.shape[-1] == 1:
        # if only 1 basis func, just return it
        squeezed = product_of_wavefunctions.squeeze(dim=-1)
        return squeezed
    for i in range(product_of_wavefunctions.shape[-1]):
        for j in range(i+1, product_of_wavefunctions.shape[-1]):
            total = total * (product_of_wavefunctions[..., j] - product_of_wavefunctions[..., i])
    return total

def eval_log_vandermonde(xs):
    product_of_wavefunctions = xs.prod(dim=-1)
    # + torch.sum(torch.log(torch.abs(a[...,
    #                                   torch.triu(torch.ones(self.dim, self.dim), diagonal=1).nonzero(as_tuple=True)[0],
    #                                   torch.triu(torch.ones(self.dim, self.dim), diagonal=1).nonzero(as_tuple=True)[
    #                                       1]])), -1) )
    total = 0.0
    sign = 1.0
    if product_of_wavefunctions.shape[-1] == 1:
        # if only 1 basis func, just return it
        squeezed = product_of_wavefunctions.squeeze(dim=-1)
        return torch.sign(squeezed), torch.log(torch.abs(squeezed))
    # prod_all_sign = torch.sign(product_all_wavefunctions)
    for i in range(product_of_wavefunctions.shape[-1]):
        for j in range(i+1, product_of_wavefunctions.shape[-1]):
            diff = product_of_wavefunctions[..., j] - product_of_wavefunctions[..., i]
            total = total + torch.log(torch.abs(diff))
            sign = sign * torch.sign(diff)

    return sign, total


class PauliNet(WaveFunction):
    r"""Implements the PauliNet ansatz from [Hermann19]_.

    Derived from :class:`WaveFunction`. This constructor provides a fully
    flexible low-level interface. See the alternative constructors for the
    high-level interfaces.

    The PauliNet ansatz combines a multireference Slater determinant expansion,
    Gaussian-type cusp-corrected molecular orbitals (:class:`MolecularOrbital`),
    electronic cusp Jastrow factor (:class:`ElectronicAsymptotic`) and many-body
    Jastrow factor and backflow transformation represented by neural networks
    that use featurized particle distances,
    :math:`\mathbf e(|\mathbf r-\mathbf r'|)` (:class:`DistanceBasis`), as input,

    .. math::
        \psi_{\boldsymbol\theta}(\mathbf r)
          =\mathrm e^{\gamma(\mathbf r)+J_{\boldsymbol\theta}(\mathbf r)}
          \sum_{pq} c_p
          \det[\tilde\varphi_{\boldsymbol\theta,q{\mu_p}i}^\uparrow(\mathbf r)]
          \det[\tilde\varphi_{\boldsymbol\theta,q{\mu_p}i}^\downarrow(\mathbf r)] \\
        \tilde\varphi_{\boldsymbol\theta,q\mu i}(\mathbf r)
          :=\big(1+2\tanh(\kappa_{\boldsymbol\theta,q\mu i}(\mathbf r))\big)
          \varphi_\mu(\mathbf r_i)

    Here, :math:`c_p,\mu_p` define the multideterminant expansion,
    :math:`\varphi_\mu(\mathbf r)` are the baseline
    single-electron molecular orbitals, :math:`J_{\boldsymbol\theta}(\mathbf r)`
    is the permutation-invariant deep Jastrow factor,
    :math:`\kappa_{\boldsymbol\theta,q\mu i}(\mathbf r)` is the :math:`q`-th
    channel of the permutation-equivariant deep backflow, and :math:`\gamma`
    enforces correct electronic cusp conditions.

    Args:
        mol (:class:`~deepqmc.Molecule`): molecule whose wave function is represented
        basis (:class:`~deepqmc.wf.paulinet.GTOBasis`): basis for the molecular orbitals
        jastrow_factory (callable): constructor for a Jastrow factor,
            :math:`(M,\dim(\mathbf e),N^\uparrow,N^\downarrow)`
            :math:`\rightarrow(\mathbf e_{ij},\mathbf e_{iI})\rightarrow J`
        backflow_factory (callable): constructor for a backflow,
            :math:`(M,\dim(\mathbf e),N^\uparrow,N^\downarrow,N_\text{orb},C)`
            :math:`\rightarrow(\mathbf e_{ij},\mathbf e_{iI})`
            :math:`\rightarrow\kappa_{q\mu i}`
        omni_factory (callable): constructor for a combined Jastrow factor and backflow,
            with interface identical to :class:`~deepqmc.wf.paulinet.OmniSchNet`
        n_configurations (int): number of electron configurations
        n_orbitals (int): number of distinct molecular orbitals used across all
            configurations if given, otherwise the larger of the number of spin-up
            and spin-down electrons
        mo_factory (callable): passed to :class:`~deepqmc.wf.paulinet.MolecularOrbital`
            as ``net_factory``
        cusp_correction (bool): whether nuclear cusp correction is used
        cusp_electrons (bool): whether electronic cusp function is used
        dist_feat_dim (int): :math:`\dim(\mathbf e)`, number of distance features
        dist_feat_cutoff (float, a.u.): distance at which distance features
            go to zero
        backflow_channels (int): :math:`C`, number of backflow channels

    Attributes:
        jastrow: :class:`torch.nn.Module` representing the Jastrow factor
        backflow: :class:`torch.nn.Module` representing the backflow transformation
        conf_coeff: :class:`torch.nn.Linear` with no bias that represents
            the multireference coefficients :math:`c_p` via its :attr:`weight`
            variable of shape :math:`(1,N_\text{det})`
    """

    def __init__(
        self,
        mol,
        basis,
        jastrow_factory=None,
        backflow_factory=None,
        omni_factory=None,
        n_configurations=1,
        n_orbitals=None,
        mo_factory=None,
        return_log=True,
        use_sloglindet='training',
        use_vandermonde=False,
        *,
        cusp_correction=False,
        cusp_electrons=False,
        dist_feat_dim=32,
        dist_feat_cutoff=10.0,
        backflow_type='orbital',
        backflow_channels=1,
        backflow_transform='mult',
        rc_scaling=1.0,
        cusp_alpha=10.0,
        freeze_embed=False,
    ):
        assert use_sloglindet in {'never', 'training', 'always'}
        assert return_log or use_sloglindet == 'never'
        super().__init__(mol)
        n_up, n_down = self.n_up, self.n_down
        self.dist_basis = (
            DistanceBasis(dist_feat_dim, cutoff=dist_feat_cutoff, envelope='nocusp')
            if mo_factory or jastrow_factory or backflow_factory or omni_factory
            else None
        )
        n_orbitals = n_orbitals or max(n_up, n_down)
        confs = [list(range(n_up)) + list(range(n_down))] + [
            sum((torch.randperm(n_orbitals)[:n].tolist() for n in (n_up, n_down)), [])
            for _ in range(n_configurations - 1)
        ]
        self.register_buffer('confs', torch.tensor(confs))
        self.conf_coeff = (
            nn.Linear(n_configurations, 1, bias=False)
            if n_configurations > 1
            else nn.Identity()
        )
        self.mo = MolecularOrbital(
            mol,
            basis,
            n_orbitals,
            net_factory=mo_factory,
            dist_feat_dim=dist_feat_dim,
            cusp_correction=cusp_correction,
            rc_scaling=rc_scaling,
        )
        self.cusp_same, self.cusp_anti = (
            (ElectronicAsymptotic(cusp=cusp, alpha=cusp_alpha) for cusp in (0.25, 0.5))
            if cusp_electrons
            else (None, None)
        )
        self.jastrow = (
            jastrow_factory(len(mol), dist_feat_dim, n_up, n_down)
            if jastrow_factory
            else None
        )
        backflow_spec = {
            'orbital': [n_orbitals, backflow_channels],
            'det': [max(n_up, n_down), len(self.confs) * backflow_channels],
        }[backflow_type]
        if backflow_transform == 'both':
            backflow_spec[1] *= 2
        self.backflow_type = backflow_type
        self.backflow_transform = backflow_transform
        self.backflow = (
            backflow_factory(len(mol), dist_feat_dim, n_up, n_down, *backflow_spec)
            if backflow_factory
            else None
        )
        self.r_backflow = None
        if omni_factory:
            assert not backflow_factory and not jastrow_factory
            self.omni = omni_factory(mol, dist_feat_dim, n_up, n_down, *backflow_spec)
            self.backflow = self.omni.forward_backflow
            self.r_backflow = self.omni.forward_r_backflow
            self.jastrow = self.omni.forward_jastrow
        else:
            self.omni = None
        self.return_log = return_log
        if freeze_embed:
            self.requires_grad_embeddings_(False)
        self.n_determinants = len(self.confs) * backflow_channels
        self.n_backflows = 0 if not self.backflow else backflow_spec[1]
        self.use_vandermonde = use_vandermonde
        if n_up <= 1 or n_down <= 1:
            self.use_sloglindet = 'never'
            log.warning(
                'Setting use_sloglindet to "never" as not implemented for n=0 and n=1.'
            )
        if self.use_vandermonde:
            self.use_sloglindet = 'never'
            log.warning(
                'Setting use_sloglindet to "never" as incompatible with Vandermonde determinant.'
            )
        # TODO implement sloglindet for special cases n=0 and n=1
        else:
            self.use_sloglindet = use_sloglindet

    def requires_grad_classes_(self, classes, requires_grad):
        for m in self.modules():
            if isinstance(m, classes):
                for p in m.parameters(recurse=False):
                    p.requires_grad_(requires_grad)
        return self

    def requires_grad_cusps_(self, requires_grad):
        return self.requires_grad_classes_(CuspCorrection, requires_grad)

    def requires_grad_embeddings_(self, requires_grad):
        return self.requires_grad_classes_(nn.Embedding, requires_grad)

    def requires_grad_nets_(self, requires_grad):
        return self.requires_grad_classes_(nn.Linear, requires_grad)

    @classmethod
    def from_pyscf(
        cls,
        mf,
        *,
        init_weights=True,
        freeze_mos=True,
        freeze_confs=False,
        conf_cutoff=1e-2,
        conf_limit=None,
        **kwargs,
    ):
        r"""Construct a :class:`PauliNet` instance from a finished PySCF_ calculation.

        Args:
            mf (:class:`pyscf.scf.hf.RHF` | :class:`pyscf.mcscf.mc1step.CASSCF`):
                restricted (multireference) HF calculation
            init_weights (bool): whether molecular orbital coefficients and
                configuration coefficients are initialized from the HF calculation
            freeze_mos (bool): whether the MO coefficients are frozen for
                gradient optimization
            freeze_confs (bool): whether the configuration coefficients are
                frozen for gradient optimization
            conf_cutoff (float): determinants with a linear coefficient above
                this threshold are included in the determinant expansion
            conf_limit (int): if given, at maximum the given number of configurations
                with the largest linear coefficients are used in the ansatz
            kwargs: all other arguments are passed to the :class:`PauliNet`
                constructor

        .. _PySCF: http://pyscf.org
        """
        assert not (set(kwargs) & {'n_configurations', 'n_orbitals'})
        n_up, n_down = mf.mol.nelec
        if hasattr(mf, 'fcisolver'):
            if conf_limit:
                conf_cutoff = max(
                    np.sort(abs(mf.ci.flatten()))[-conf_limit:][0] - 1e-10, conf_cutoff,
                )
            for tol in [conf_cutoff, conf_cutoff + 2e-10]:
                conf_coeff, *confs = zip(
                    *mf.fcisolver.large_ci(
                        mf.ci, mf.ncas, mf.nelecas, tol=tol, return_strs=False
                    )
                )
                if not conf_limit or len(conf_coeff) <= conf_limit:
                    break
            else:
                raise AssertionError()
            # discard the last ci wave function if degenerate
            ns_dbl = n_up - mf.nelecas[0], n_down - mf.nelecas[1]
            conf_coeff = torch.tensor(conf_coeff)
            confs = [
                [
                    torch.arange(n_dbl, dtype=torch.long).expand(len(conf_coeff), -1),
                    torch.tensor(cfs, dtype=torch.long) + n_dbl,
                ]
                for n_dbl, cfs in zip(ns_dbl, confs)
            ]
            confs = [torch.cat(cfs, dim=-1) for cfs in confs]
            confs = torch.cat(confs, dim=-1)
            kwargs['n_configurations'] = len(confs)
            kwargs['n_orbitals'] = confs.max().item() + 1
        else:
            confs = None
        mol = Molecule(
            mf.mol.atom_coords().astype('float32'),
            mf.mol.atom_charges(),
            mf.mol.charge,
            mf.mol.spin,
        )
        basis = GTOBasis.from_pyscf(mf.mol)
        wf = cls(mol, basis, **kwargs)
        if init_weights:
            wf.mo.init_from_pyscf(mf, freeze_mos=freeze_mos)
            if confs is not None:
                wf.confs.detach().copy_(confs)
                if len(confs) > 1:
                    wf.conf_coeff.weight.detach().copy_(conf_coeff)
                if freeze_confs:
                    wf.conf_coeff.weight.requires_grad_(False)
        return wf

    @classmethod
    def from_hf(
        cls, mol, *, basis='6-311g', cas=None, pauli_kwargs=None, omni_kwargs=None
    ):
        r"""Construct a :class:`PauliNet` instance by running a HF calculation.

        This is the top-level interface.

        Args:
            mol (:class:`~deepqmc.Molecule`): molecule whose wave function
                is represented
            basis (str): basis of the internal HF calculation
            cas ((int, int)): tuple of the number of active orbitals and number of
                active electrons for a complete active space multireference
                HF calculation
            pauli_kwargs: arguments passed to :func:`PauliNet.from_pyscf`
            omni_kwargs: arguments passed to :class:`~deepqmc.wf.paulinet.OmniSchNet`
        """
        from pyscf import gto, lib, mcscf, scf

        mol = gto.M(
            atom=mol.as_pyscf(),
            unit='bohr',
            basis=basis,
            charge=mol.charge,
            spin=mol.spin,
            cart=True,
        )
        log.info('Running HF...')
        mf = scf.RHF(mol)
        mf.kernel()
        if cas:
            log.info('Running MCSCF...')
            mc = mcscf.CASSCF(mf, *cas)
            mc.kernel()
            lib.chkfile.dump(mc.chkfile, 'ci', mc.ci)
            lib.chkfile.dump(mc.chkfile, 'nelecas', mc.nelecas)
        wf = PauliNet.from_pyscf(
            mc if cas else mf,
            **{
                'omni_factory': partial(OmniSchNet, **(omni_kwargs or {})),
                'cusp_correction': True,
                'cusp_electrons': True,
                **(pauli_kwargs or {}),
            },
        )
        wf.mf = mf
        return wf

    def _backflow_op(self, xs, fs):
        if self.backflow_transform == 'mult':
            fs_mult, fs_add = fs, None
        elif self.backflow_transform == 'add':
            fs_mult, fs_add = None, fs
        elif self.backflow_transform == 'both':
            fs_mult, fs_add = fs[:, : fs.shape[1] // 2], fs[:, fs.shape[1] // 2 :]
        if fs_add is not None:
            envel = (xs ** 2).mean(dim=-1, keepdim=True).sqrt()
        if fs_mult is not None:
            xs = xs * (1 + 2 * torch.tanh(fs_mult / 4))
        if fs_add is not None:
            xs = xs + 0.1 * envel * torch.tanh(fs_add / 4)
        return xs

    def forward(self, rs, debug=NULL_DEBUG):  # noqa: C901
        batch_dim, n_elec = rs.shape[:2]
        assert n_elec == self.confs.shape[1]
        n_atoms = len(self.mol)
        coords = self.mol.coords
        diffs_nuc = pairwise_diffs(torch.cat([coords, rs.flatten(end_dim=1)]), coords)
        edges_nuc = (
            self.dist_basis(diffs_nuc[:, :, 3].sqrt())
            if self.jastrow or self.mo.net
            else None
        )
        if self.r_backflow or self.backflow or self.cusp_same or self.jastrow:
            dists_elec = pairwise_distance(rs, rs)
        if self.r_backflow or self.backflow or self.jastrow:
            edges_nuc = edges_nuc[n_atoms:].view(batch_dim, n_elec, n_atoms, -1)
            edges = self.dist_basis(dists_elec), edges_nuc
        if self.r_backflow:
            rs_flowed = self.r_backflow(rs, *edges, debug=debug)
            diffs_nuc = pairwise_diffs(
                torch.cat([coords, rs_flowed.flatten(end_dim=1)]), coords
            )
        with debug.cd('mos'):
            xs = self.mo(diffs_nuc, edges_nuc, debug=debug)
        # get orbitals as [bs, 1, i, mu]
        xs = debug['slaters'] = xs.view(batch_dim, 1, n_elec, -1)
        # this should be a product of one orbital evaluated on each electron.
        single_orbital = xs[..., 0].prod(dim=-1).squeeze(dim=-1)
        if self.backflow:
            with debug.cd('backflow'):
                fs = self.backflow(*edges, debug=debug)  # [bs, q, i, mu/nu]
            if self.backflow_type == 'orbital':
                xs = self._backflow_op(xs, fs)
        # form dets as [bs, q, p, i, nu]
        conf_up, conf_down = self.confs[:, : self.n_up], self.confs[:, self.n_up :]
        det_up = xs[:, :, : self.n_up, conf_up].transpose(-3, -2)
        det_down = xs[:, :, self.n_up :, conf_down].transpose(-3, -2)
        if self.backflow and self.backflow_type == 'det':
            n_conf = len(self.confs)
            fs = fs.unflatten(1, ((None, fs.shape[1] // n_conf), (None, n_conf)))
            det_up = self._backflow_op(det_up, fs[..., : self.n_up, : self.n_up])
            det_down = self._backflow_op(det_down, fs[..., self.n_up :, : self.n_down])
            # with open-shell systems, part of the backflow output is not used
        if self.use_sloglindet == 'always' or (
            self.use_sloglindet == 'training' and not self.sampling
        ):
            bf_dim = det_up.shape[-4]
            if isinstance(self.conf_coeff, nn.Linear):
                conf_coeff = self.conf_coeff.weight[0]
                conf_coeff = conf_coeff.expand(bf_dim, -1).flatten() / np.sqrt(bf_dim)
            else:
                conf_coeff = det_up.new_ones(1)
            det_up = det_up.flatten(start_dim=-4, end_dim=-3).contiguous()
            det_down = det_down.flatten(start_dim=-4, end_dim=-3).contiguous()
            sign, psi = sloglindet(conf_coeff, det_up, det_down)
            sign = sign.detach()
        else:
            if self.return_log:
                if self.use_vandermonde:
                    sign_up, det_up = eval_log_vandermonde(det_up)
                    sign_down, det_down = eval_log_vandermonde(det_down)
                    xs = det_up + det_down
                else:
                    sign_up, det_up = eval_log_slater(det_up)
                    sign_down, det_down = eval_log_slater(det_down)
                    xs = det_up + det_down

                # premultiply by diagonal to ensure cusps? make sure antisymmetry is maintained!
                xs_shift = xs.flatten(start_dim=1).max(dim=-1).values

                # the exp-normalize trick, to avoid over/underflow of the exponential
                xs = sign_up * sign_down * torch.exp(xs - xs_shift[:, None, None])
            else:
                if self.use_vandermonde:
                    det_up = debug['det_up'] = eval_vandermonde(det_up)
                    det_down = debug['det_down'] = eval_vandermonde(det_down)
                    xs = det_up * det_down
                else:
                    det_up = debug['det_up'] = eval_slater(det_up)
                    det_down = debug['det_down'] = eval_slater(det_down)
                    xs = det_up * det_down
            psi = self.conf_coeff(xs).squeeze(dim=-1).mean(dim=-1)
            if self.return_log:
                psi, sign = psi.abs().log() + xs_shift, psi.sign().detach()
        if self.cusp_same:
            cusp_same = self.cusp_same(
                torch.cat(
                    [triu_flat(dists_elec[:, idxs, idxs]) for idxs in self.spin_slices],
                    dim=1,
                )
            )
            cusp_anti = self.cusp_anti(
                dists_elec[:, : self.n_up, self.n_up :].flatten(start_dim=1)
            )
            if self.use_vandermonde:
                # here we multiply the final result by the single orbital
                psi = (
                    psi + cusp_same + cusp_anti + torch.log(torch.abs(single_orbital))
                    if self.return_log
                    else psi * torch.exp(cusp_same + cusp_anti) * single_orbital
                )
                if not self.return_log:
                    print('it gets called without log')
                else:
                    sign = sign * torch.sign(single_orbital)
            else:
                psi = (
                    psi + cusp_same + cusp_anti
                    if self.return_log
                    else psi * torch.exp(cusp_same + cusp_anti)
                )
        if self.jastrow:
            with debug.cd('jastrow'):
                J = self.jastrow(*edges, debug=debug)
                psi = psi + J if self.return_log else psi * torch.exp(J)
        if self.omni:
            self.omni.forward_close()
        return (psi, sign) if self.return_log else psi

import logging
from functools import partial

import numpy as np
import torch
from torch import nn

from deepqmc import Molecule
from deepqmc.physics import pairwise_diffs, pairwise_distance
from deepqmc.torchext import sloglindet, triu_flat
from deepqmc.utils import NULL_DEBUG
from deepqmc.wf import WaveFunction

from .cusp import CuspCorrection, ElectronicAsymptotic
from .distbasis import DistanceBasis
from .gto import GTOBasis
from .molorb import MolecularOrbital
from .omni import OmniSchNet

__version__ = '0.1.0'
__all__ = ['PauliNet']

log = logging.getLogger(__name__)


def eval_slater(xs):
    if xs.shape[-1] == 0:
        return xs.new_ones(xs.shape[:-2])
    return torch.det(xs.contiguous())


def eval_log_slater(xs):
    if xs.shape[-1] == 0:
        return xs.new_ones(xs.shape[:-2]), xs.new_zeros(xs.shape[:-2])
    return xs.contiguous().slogdet()

def eval_vandermonde(xs):
    print('non-log vandermonde is evaluated')
    product_of_wavefunctions = xs.prod(dim=-1)
    total = 1.0
    if product_of_wavefunctions.shape[-1] == 1:
        # if only 1 basis func, just return it
        squeezed = product_of_wavefunctions.squeeze(dim=-1)
        return squeezed
    for i in range(product_of_wavefunctions.shape[-1]):
        for j in range(i+1, product_of_wavefunctions.shape[-1]):
            total = total * (product_of_wavefunctions[..., j] - product_of_wavefunctions[..., i])
    return total

def eval_log_vandermonde(xs):
    product_of_wavefunctions = xs.prod(dim=-1)
    # + torch.sum(torch.log(torch.abs(a[...,
    #                                   torch.triu(torch.ones(self.dim, self.dim), diagonal=1).nonzero(as_tuple=True)[0],
    #                                   torch.triu(torch.ones(self.dim, self.dim), diagonal=1).nonzero(as_tuple=True)[
    #                                       1]])), -1) )
    total = 0.0
    sign = 1.0
    if product_of_wavefunctions.shape[-1] == 1:
        # if only 1 basis func, just return it
        squeezed = product_of_wavefunctions.squeeze(dim=-1)
        return torch.sign(squeezed), torch.log(torch.abs(squeezed))
    # prod_all_sign = torch.sign(product_all_wavefunctions)
    for i in range(product_of_wavefunctions.shape[-1]):
        for j in range(i+1, product_of_wavefunctions.shape[-1]):
            diff = product_of_wavefunctions[..., j] - product_of_wavefunctions[..., i]
            total = total + torch.log(torch.abs(diff))
            sign = sign * torch.sign(diff)

    return sign, total


class PauliNet(WaveFunction):
    r"""Implements the PauliNet ansatz from [Hermann19]_.

    Derived from :class:`WaveFunction`. This constructor provides a fully
    flexible low-level interface. See the alternative constructors for the
    high-level interfaces.

    The PauliNet ansatz combines a multireference Slater determinant expansion,
    Gaussian-type cusp-corrected molecular orbitals (:class:`MolecularOrbital`),
    electronic cusp Jastrow factor (:class:`ElectronicAsymptotic`) and many-body
    Jastrow factor and backflow transformation represented by neural networks
    that use featurized particle distances,
    :math:`\mathbf e(|\mathbf r-\mathbf r'|)` (:class:`DistanceBasis`), as input,

    .. math::
        \psi_{\boldsymbol\theta}(\mathbf r)
          =\mathrm e^{\gamma(\mathbf r)+J_{\boldsymbol\theta}(\mathbf r)}
          \sum_{pq} c_p
          \det[\tilde\varphi_{\boldsymbol\theta,q{\mu_p}i}^\uparrow(\mathbf r)]
          \det[\tilde\varphi_{\boldsymbol\theta,q{\mu_p}i}^\downarrow(\mathbf r)] \\
        \tilde\varphi_{\boldsymbol\theta,q\mu i}(\mathbf r)
          :=\big(1+2\tanh(\kappa_{\boldsymbol\theta,q\mu i}(\mathbf r))\big)
          \varphi_\mu(\mathbf r_i)

    Here, :math:`c_p,\mu_p` define the multideterminant expansion,
    :math:`\varphi_\mu(\mathbf r)` are the baseline
    single-electron molecular orbitals, :math:`J_{\boldsymbol\theta}(\mathbf r)`
    is the permutation-invariant deep Jastrow factor,
    :math:`\kappa_{\boldsymbol\theta,q\mu i}(\mathbf r)` is the :math:`q`-th
    channel of the permutation-equivariant deep backflow, and :math:`\gamma`
    enforces correct electronic cusp conditions.

    Args:
        mol (:class:`~deepqmc.Molecule`): molecule whose wave function is represented
        basis (:class:`~deepqmc.wf.paulinet.GTOBasis`): basis for the molecular orbitals
        jastrow_factory (callable): constructor for a Jastrow factor,
            :math:`(M,\dim(\mathbf e),N^\uparrow,N^\downarrow)`
            :math:`\rightarrow(\mathbf e_{ij},\mathbf e_{iI})\rightarrow J`
        backflow_factory (callable): constructor for a backflow,
            :math:`(M,\dim(\mathbf e),N^\uparrow,N^\downarrow,N_\text{orb},C)`
            :math:`\rightarrow(\mathbf e_{ij},\mathbf e_{iI})`
            :math:`\rightarrow\kappa_{q\mu i}`
        omni_factory (callable): constructor for a combined Jastrow factor and backflow,
            with interface identical to :class:`~deepqmc.wf.paulinet.OmniSchNet`
        n_configurations (int): number of electron configurations
        n_orbitals (int): number of distinct molecular orbitals used across all
            configurations if given, otherwise the larger of the number of spin-up
            and spin-down electrons
        mo_factory (callable): passed to :class:`~deepqmc.wf.paulinet.MolecularOrbital`
            as ``net_factory``
        cusp_correction (bool): whether nuclear cusp correction is used
        cusp_electrons (bool): whether electronic cusp function is used
        dist_feat_dim (int): :math:`\dim(\mathbf e)`, number of distance features
        dist_feat_cutoff (float, a.u.): distance at which distance features
            go to zero
        backflow_channels (int): :math:`C`, number of backflow channels

    Attributes:
        jastrow: :class:`torch.nn.Module` representing the Jastrow factor
        backflow: :class:`torch.nn.Module` representing the backflow transformation
        conf_coeff: :class:`torch.nn.Linear` with no bias that represents
            the multireference coefficients :math:`c_p` via its :attr:`weight`
            variable of shape :math:`(1,N_\text{det})`
    """

    def __init__(
        self,
        mol,
        basis,
        jastrow_factory=None,
        backflow_factory=None,
        omni_factory=None,
        n_configurations=1,
        n_orbitals=None,
        mo_factory=None,
        return_log=True,
        use_sloglindet='training',
        use_vandermonde=False,
        *,
        cusp_correction=False,
        cusp_electrons=False,
        dist_feat_dim=32,
        dist_feat_cutoff=10.0,
        backflow_type='orbital',
        backflow_channels=1,
        backflow_transform='mult',
        rc_scaling=1.0,
        cusp_alpha=10.0,
        freeze_embed=False,
    ):
        assert use_sloglindet in {'never', 'training', 'always'}
        assert return_log or use_sloglindet == 'never'
        super().__init__(mol)
        n_up, n_down = self.n_up, self.n_down
        self.dist_basis = (
            DistanceBasis(dist_feat_dim, cutoff=dist_feat_cutoff, envelope='nocusp')
            if mo_factory or jastrow_factory or backflow_factory or omni_factory
            else None
        )
        n_orbitals = n_orbitals or max(n_up, n_down)
        confs = [list(range(n_up)) + list(range(n_down))] + [
            sum((torch.randperm(n_orbitals)[:n].tolist() for n in (n_up, n_down)), [])
            for _ in range(n_configurations - 1)
        ]
        self.register_buffer('confs', torch.tensor(confs))
        self.conf_coeff = (
            nn.Linear(n_configurations, 1, bias=False)
            if n_configurations > 1
            else nn.Identity()
        )
        self.mo = MolecularOrbital(
            mol,
            basis,
            n_orbitals,
            net_factory=mo_factory,
            dist_feat_dim=dist_feat_dim,
            cusp_correction=cusp_correction,
            rc_scaling=rc_scaling,
        )
        self.cusp_same, self.cusp_anti = (
            (ElectronicAsymptotic(cusp=cusp, alpha=cusp_alpha) for cusp in (0.25, 0.5))
            if cusp_electrons
            else (None, None)
        )
        self.jastrow = (
            jastrow_factory(len(mol), dist_feat_dim, n_up, n_down)
            if jastrow_factory
            else None
        )
        backflow_spec = {
            'orbital': [n_orbitals, backflow_channels],
            'det': [max(n_up, n_down), len(self.confs) * backflow_channels],
        }[backflow_type]
        if backflow_transform == 'both':
            backflow_spec[1] *= 2
        self.backflow_type = backflow_type
        self.backflow_transform = backflow_transform
        self.backflow = (
            backflow_factory(len(mol), dist_feat_dim, n_up, n_down, *backflow_spec)
            if backflow_factory
            else None
        )
        self.r_backflow = None
        if omni_factory:
            assert not backflow_factory and not jastrow_factory
            self.omni = omni_factory(mol, dist_feat_dim, n_up, n_down, *backflow_spec)
            self.backflow = self.omni.forward_backflow
            self.r_backflow = self.omni.forward_r_backflow
            self.jastrow = self.omni.forward_jastrow
        else:
            self.omni = None
        self.return_log = return_log
        if freeze_embed:
            self.requires_grad_embeddings_(False)
        self.n_determinants = len(self.confs) * backflow_channels
        self.n_backflows = 0 if not self.backflow else backflow_spec[1]
        self.use_vandermonde = use_vandermonde
        if n_up <= 1 or n_down <= 1:
            self.use_sloglindet = 'never'
            log.warning(
                'Setting use_sloglindet to "never" as not implemented for n=0 and n=1.'
            )
        if self.use_vandermonde:
            self.use_sloglindet = 'never'
            log.warning(
                'Setting use_sloglindet to "never" as incompatible with Vandermonde determinant.'
            )
        # TODO implement sloglindet for special cases n=0 and n=1
        else:
            self.use_sloglindet = use_sloglindet

    def requires_grad_classes_(self, classes, requires_grad):
        for m in self.modules():
            if isinstance(m, classes):
                for p in m.parameters(recurse=False):
                    p.requires_grad_(requires_grad)
        return self

    def requires_grad_cusps_(self, requires_grad):
        return self.requires_grad_classes_(CuspCorrection, requires_grad)

    def requires_grad_embeddings_(self, requires_grad):
        return self.requires_grad_classes_(nn.Embedding, requires_grad)

    def requires_grad_nets_(self, requires_grad):
        return self.requires_grad_classes_(nn.Linear, requires_grad)

    @classmethod
    def from_pyscf(
        cls,
        mf,
        *,
        init_weights=True,
        freeze_mos=True,
        freeze_confs=False,
        conf_cutoff=1e-2,
        conf_limit=None,
        **kwargs,
    ):
        r"""Construct a :class:`PauliNet` instance from a finished PySCF_ calculation.

        Args:
            mf (:class:`pyscf.scf.hf.RHF` | :class:`pyscf.mcscf.mc1step.CASSCF`):
                restricted (multireference) HF calculation
            init_weights (bool): whether molecular orbital coefficients and
                configuration coefficients are initialized from the HF calculation
            freeze_mos (bool): whether the MO coefficients are frozen for
                gradient optimization
            freeze_confs (bool): whether the configuration coefficients are
                frozen for gradient optimization
            conf_cutoff (float): determinants with a linear coefficient above
                this threshold are included in the determinant expansion
            conf_limit (int): if given, at maximum the given number of configurations
                with the largest linear coefficients are used in the ansatz
            kwargs: all other arguments are passed to the :class:`PauliNet`
                constructor

        .. _PySCF: http://pyscf.org
        """
        assert not (set(kwargs) & {'n_configurations', 'n_orbitals'})
        n_up, n_down = mf.mol.nelec
        if hasattr(mf, 'fcisolver'):
            if conf_limit:
                conf_cutoff = max(
                    np.sort(abs(mf.ci.flatten()))[-conf_limit:][0] - 1e-10, conf_cutoff,
                )
            for tol in [conf_cutoff, conf_cutoff + 2e-10]:
                conf_coeff, *confs = zip(
                    *mf.fcisolver.large_ci(
                        mf.ci, mf.ncas, mf.nelecas, tol=tol, return_strs=False
                    )
                )
                if not conf_limit or len(conf_coeff) <= conf_limit:
                    break
            else:
                raise AssertionError()
            # discard the last ci wave function if degenerate
            ns_dbl = n_up - mf.nelecas[0], n_down - mf.nelecas[1]
            conf_coeff = torch.tensor(conf_coeff)
            confs = [
                [
                    torch.arange(n_dbl, dtype=torch.long).expand(len(conf_coeff), -1),
                    torch.tensor(cfs, dtype=torch.long) + n_dbl,
                ]
                for n_dbl, cfs in zip(ns_dbl, confs)
            ]
            confs = [torch.cat(cfs, dim=-1) for cfs in confs]
            confs = torch.cat(confs, dim=-1)
            kwargs['n_configurations'] = len(confs)
            kwargs['n_orbitals'] = confs.max().item() + 1
        else:
            confs = None
        mol = Molecule(
            mf.mol.atom_coords().astype('float32'),
            mf.mol.atom_charges(),
            mf.mol.charge,
            mf.mol.spin,
        )
        basis = GTOBasis.from_pyscf(mf.mol)
        wf = cls(mol, basis, **kwargs)
        if init_weights:
            wf.mo.init_from_pyscf(mf, freeze_mos=freeze_mos)
            if confs is not None:
                wf.confs.detach().copy_(confs)
                if len(confs) > 1:
                    wf.conf_coeff.weight.detach().copy_(conf_coeff)
                if freeze_confs:
                    wf.conf_coeff.weight.requires_grad_(False)
        return wf

    @classmethod
    def from_hf(
        cls, mol, *, basis='6-311g', cas=None, pauli_kwargs=None, omni_kwargs=None
    ):
        r"""Construct a :class:`PauliNet` instance by running a HF calculation.

        This is the top-level interface.

        Args:
            mol (:class:`~deepqmc.Molecule`): molecule whose wave function
                is represented
            basis (str): basis of the internal HF calculation
            cas ((int, int)): tuple of the number of active orbitals and number of
                active electrons for a complete active space multireference
                HF calculation
            pauli_kwargs: arguments passed to :func:`PauliNet.from_pyscf`
            omni_kwargs: arguments passed to :class:`~deepqmc.wf.paulinet.OmniSchNet`
        """
        from pyscf import gto, lib, mcscf, scf

        mol = gto.M(
            atom=mol.as_pyscf(),
            unit='bohr',
            basis=basis,
            charge=mol.charge,
            spin=mol.spin,
            cart=True,
        )
        log.info('Running HF...')
        mf = scf.RHF(mol)
        mf.kernel()
        if cas:
            log.info('Running MCSCF...')
            mc = mcscf.CASSCF(mf, *cas)
            mc.kernel()
            lib.chkfile.dump(mc.chkfile, 'ci', mc.ci)
            lib.chkfile.dump(mc.chkfile, 'nelecas', mc.nelecas)
        wf = PauliNet.from_pyscf(
            mc if cas else mf,
            **{
                'omni_factory': partial(OmniSchNet, **(omni_kwargs or {})),
                'cusp_correction': True,
                'cusp_electrons': True,
                **(pauli_kwargs or {}),
            },
        )
        wf.mf = mf
        return wf

    def _backflow_op(self, xs, fs):
        if self.backflow_transform == 'mult':
            fs_mult, fs_add = fs, None
        elif self.backflow_transform == 'add':
            fs_mult, fs_add = None, fs
        elif self.backflow_transform == 'both':
            fs_mult, fs_add = fs[:, : fs.shape[1] // 2], fs[:, fs.shape[1] // 2 :]
        if fs_add is not None:
            envel = (xs ** 2).mean(dim=-1, keepdim=True).sqrt()
        if fs_mult is not None:
            xs = xs * (1 + 2 * torch.tanh(fs_mult / 4))
        if fs_add is not None:
            xs = xs + 0.1 * envel * torch.tanh(fs_add / 4)
        return xs

    def forward(self, rs, debug=NULL_DEBUG):  # noqa: C901
        batch_dim, n_elec = rs.shape[:2]
        assert n_elec == self.confs.shape[1]
        n_atoms = len(self.mol)
        coords = self.mol.coords
        diffs_nuc = pairwise_diffs(torch.cat([coords, rs.flatten(end_dim=1)]), coords)
        edges_nuc = (
            self.dist_basis(diffs_nuc[:, :, 3].sqrt())
            if self.jastrow or self.mo.net
            else None
        )
        if self.r_backflow or self.backflow or self.cusp_same or self.jastrow:
            dists_elec = pairwise_distance(rs, rs)
        if self.r_backflow or self.backflow or self.jastrow:
            edges_nuc = edges_nuc[n_atoms:].view(batch_dim, n_elec, n_atoms, -1)
            edges = self.dist_basis(dists_elec), edges_nuc
        if self.r_backflow:
            rs_flowed = self.r_backflow(rs, *edges, debug=debug)
            diffs_nuc = pairwise_diffs(
                torch.cat([coords, rs_flowed.flatten(end_dim=1)]), coords
            )
        with debug.cd('mos'):
            xs = self.mo(diffs_nuc, edges_nuc, debug=debug)
        # get orbitals as [bs, 1, i, mu]
        xs = debug['slaters'] = xs.view(batch_dim, 1, n_elec, -1)
        # this should be a product of one orbital evaluated on each electron.
        single_orbital = xs[..., 0].prod(dim=-1).squeeze(dim=-1)
        if self.backflow:
            with debug.cd('backflow'):
                fs = self.backflow(*edges, debug=debug)  # [bs, q, i, mu/nu]
            if self.backflow_type == 'orbital':
                xs = self._backflow_op(xs, fs)
        # form dets as [bs, q, p, i, nu]
        conf_up, conf_down = self.confs[:, : self.n_up], self.confs[:, self.n_up :]
        det_up = xs[:, :, : self.n_up, conf_up].transpose(-3, -2)
        det_down = xs[:, :, self.n_up :, conf_down].transpose(-3, -2)
        if self.backflow and self.backflow_type == 'det':
            n_conf = len(self.confs)
            fs = fs.unflatten(1, ((None, fs.shape[1] // n_conf), (None, n_conf)))
            det_up = self._backflow_op(det_up, fs[..., : self.n_up, : self.n_up])
            det_down = self._backflow_op(det_down, fs[..., self.n_up :, : self.n_down])
            # with open-shell systems, part of the backflow output is not used
        if self.use_sloglindet == 'always' or (
            self.use_sloglindet == 'training' and not self.sampling
        ):
            bf_dim = det_up.shape[-4]
            if isinstance(self.conf_coeff, nn.Linear):
                conf_coeff = self.conf_coeff.weight[0]
                conf_coeff = conf_coeff.expand(bf_dim, -1).flatten() / np.sqrt(bf_dim)
            else:
                conf_coeff = det_up.new_ones(1)
            det_up = det_up.flatten(start_dim=-4, end_dim=-3).contiguous()
            det_down = det_down.flatten(start_dim=-4, end_dim=-3).contiguous()
            sign, psi = sloglindet(conf_coeff, det_up, det_down)
            sign = sign.detach()
        else:
            if self.return_log:
                if self.use_vandermonde:
                    sign_up, det_up = eval_log_vandermonde(det_up)
                    sign_down, det_down = eval_log_vandermonde(det_down)
                    xs = det_up + det_down
                else:
                    sign_up, det_up = eval_log_slater(det_up)
                    sign_down, det_down = eval_log_slater(det_down)
                    xs = det_up + det_down

                # premultiply by diagonal to ensure cusps? make sure antisymmetry is maintained!
                xs_shift = xs.flatten(start_dim=1).max(dim=-1).values

                # the exp-normalize trick, to avoid over/underflow of the exponential
                xs = sign_up * sign_down * torch.exp(xs - xs_shift[:, None, None])
            else:
                if self.use_vandermonde:
                    det_up = debug['det_up'] = eval_vandermonde(det_up)
                    det_down = debug['det_down'] = eval_vandermonde(det_down)
                    xs = det_up * det_down
                else:
                    det_up = debug['det_up'] = eval_slater(det_up)
                    det_down = debug['det_down'] = eval_slater(det_down)
                    xs = det_up * det_down
            psi = self.conf_coeff(xs).squeeze(dim=-1).mean(dim=-1)
            if self.return_log:
                psi, sign = psi.abs().log() + xs_shift, psi.sign().detach()
        if self.cusp_same:
            cusp_same = self.cusp_same(
                torch.cat(
                    [triu_flat(dists_elec[:, idxs, idxs]) for idxs in self.spin_slices],
                    dim=1,
                )
            )
            cusp_anti = self.cusp_anti(
                dists_elec[:, : self.n_up, self.n_up :].flatten(start_dim=1)
            )
            if self.use_vandermonde:
                # here we multiply the final result by the single orbital
                psi = (
                    psi + cusp_same + cusp_anti + torch.log(torch.abs(single_orbital))
                    if self.return_log
                    else psi * torch.exp(cusp_same + cusp_anti) * single_orbital
                )
                if not self.return_log:
                    print('it gets called without log')
                else:
                    sign = sign * torch.sign(single_orbital)
            else:
                psi = (
                    psi + cusp_same + cusp_anti
                    if self.return_log
                    else psi * torch.exp(cusp_same + cusp_anti)
                )
        if self.jastrow:
            with debug.cd('jastrow'):
                J = self.jastrow(*edges, debug=debug)
                psi = psi + J if self.return_log else psi * torch.exp(J)
        if self.omni:
            self.omni.forward_close()
        return (psi, sign) if self.return_log else psi

# if only 1 basis func, just return it # + torch.sum(torch.log(torch.abs(a[..., #                                   torch.triu(torch.ones(self.dim, self.dim), diagonal=1).nonzero(as_tuple=True)[0], #                                   torch.triu(torch.ones(self.dim, self.dim), diagonal=1).nonzero(as_tuple=True)[ #                                       1]])), -1) ) # if only 1 basis func, just return it # prod_all_sign = torch.sign(product_all_wavefunctions) Implements the PauliNet ansatz from [Hermann19]_.      Derived from :class:`WaveFunction`. This constructor provides a fully     flexible low-level interface. See the alternative constructors for the     high-level interfaces.      The PauliNet ansatz combines a multireference Slater determinant expansion,     Gaussian-type cusp-corrected molecular orbitals (:class:`MolecularOrbital`),     electronic cusp Jastrow factor (:class:`ElectronicAsymptotic`) and many-body     Jastrow factor and backflow transformation represented by neural networks     that use featurized particle distances,     :math:`\mathbf e(|\mathbf r-\mathbf r'|)` (:class:`DistanceBasis`), as input,      .. math::         \psi_{\boldsymbol\theta}(\mathbf r)           =\mathrm e^{\gamma(\mathbf r)+J_{\boldsymbol\theta}(\mathbf r)}           \sum_{pq} c_p           \det[\tilde\varphi_{\boldsymbol\theta,q{\mu_p}i}^\uparrow(\mathbf r)]           \det[\tilde\varphi_{\boldsymbol\theta,q{\mu_p}i}^\downarrow(\mathbf r)] \\         \tilde\varphi_{\boldsymbol\theta,q\mu i}(\mathbf r)           :=\big(1+2\tanh(\kappa_{\boldsymbol\theta,q\mu i}(\mathbf r))\big)           \varphi_\mu(\mathbf r_i)      Here, :math:`c_p,\mu_p` define the multideterminant expansion,     :math:`\varphi_\mu(\mathbf r)` are the baseline     single-electron molecular orbitals, :math:`J_{\boldsymbol\theta}(\mathbf r)`     is the permutation-invariant deep Jastrow factor,     :math:`\kappa_{\boldsymbol\theta,q\mu i}(\mathbf r)` is the :math:`q`-th     channel of the permutation-equivariant deep backflow, and :math:`\gamma`     enforces correct electronic cusp conditions.      Args:         mol (:class:`~deepqmc.Molecule`): molecule whose wave function is represented         basis (:class:`~deepqmc.wf.paulinet.GTOBasis`): basis for the molecular orbitals         jastrow_factory (callable): constructor for a Jastrow factor,             :math:`(M,\dim(\mathbf e),N^\uparrow,N^\downarrow)`             :math:`\rightarrow(\mathbf e_{ij},\mathbf e_{iI})\rightarrow J`         backflow_factory (callable): constructor for a backflow,             :math:`(M,\dim(\mathbf e),N^\uparrow,N^\downarrow,N_\text{orb},C)`             :math:`\rightarrow(\mathbf e_{ij},\mathbf e_{iI})`             :math:`\rightarrow\kappa_{q\mu i}`         omni_factory (callable): constructor for a combined Jastrow factor and backflow,             with interface identical to :class:`~deepqmc.wf.paulinet.OmniSchNet`         n_configurations (int): number of electron configurations         n_orbitals (int): number of distinct molecular orbitals used across all             configurations if given, otherwise the larger of the number of spin-up             and spin-down electrons         mo_factory (callable): passed to :class:`~deepqmc.wf.paulinet.MolecularOrbital`             as ``net_factory``         cusp_correction (bool): whether nuclear cusp correction is used         cusp_electrons (bool): whether electronic cusp function is used         dist_feat_dim (int): :math:`\dim(\mathbf e)`, number of distance features         dist_feat_cutoff (float, a.u.): distance at which distance features             go to zero         backflow_channels (int): :math:`C`, number of backflow channels      Attributes:         jastrow: :class:`torch.nn.Module` representing the Jastrow factor         backflow: :class:`torch.nn.Module` representing the backflow transformation         conf_coeff: :class:`torch.nn.Linear` with no bias that represents             the multireference coefficients :math:`c_p` via its :attr:`weight`             variable of shape :math:`(1,N_\text{det})` # TODO implement sloglindet for special cases n=0 and n=1 Construct a :class:`PauliNet` instance from a finished PySCF_ calculation.          Args:             mf (:class:`pyscf.scf.hf.RHF` | :class:`pyscf.mcscf.mc1step.CASSCF`):                 restricted (multireference) HF calculation             init_weights (bool): whether molecular orbital coefficients and                 configuration coefficients are initialized from the HF calculation             freeze_mos (bool): whether the MO coefficients are frozen for                 gradient optimization             freeze_confs (bool): whether the configuration coefficients are                 frozen for gradient optimization             conf_cutoff (float): determinants with a linear coefficient above                 this threshold are included in the determinant expansion             conf_limit (int): if given, at maximum the given number of configurations                 with the largest linear coefficients are used in the ansatz             kwargs: all other arguments are passed to the :class:`PauliNet`                 constructor          .. _PySCF: http://pyscf.org # discard the last ci wave function if degenerate Construct a :class:`PauliNet` instance by running a HF calculation.          This is the top-level interface.          Args:             mol (:class:`~deepqmc.Molecule`): molecule whose wave function                 is represented             basis (str): basis of the internal HF calculation             cas ((int, int)): tuple of the number of active orbitals and number of                 active electrons for a complete active space multireference                 HF calculation             pauli_kwargs: arguments passed to :func:`PauliNet.from_pyscf`             omni_kwargs: arguments passed to :class:`~deepqmc.wf.paulinet.OmniSchNet` # noqa: C901 # get orbitals as [bs, 1, i, mu] # this should be a product of one orbital evaluated on each electron. # [bs, q, i, mu/nu] # form dets as [bs, q, p, i, nu] # with open-shell systems, part of the backflow output is not used # premultiply by diagonal to ensure cusps? make sure antisymmetry is maintained! # the exp-normalize trick, to avoid over/underflow of the exponential # here we multiply the final result by the single orbital