# Copyright 2023 The NetKet Authors - All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import numpy as np

import jax
import jax.numpy as jnp

from netket.operator import DiscreteOperator
from netket.experimental.hilbert import SpinOrbitalFermions
from netket.utils.optional_deps import import_optional_dependency
from ._fermion_operator_2nd_numba import FermionOperator2nd


def compute_pyscf_integrals(mol, mo_coeff):
    t_ij_ao = jnp.asarray(mol.intor("int1e_kin") + mol.intor("int1e_nuc"))
    v_ijkl_ao = jnp.asarray(mol.intor("int2e").transpose(0, 2, 3, 1))
    c = jnp.asarray(mo_coeff.T)
    t_ij_mo = jax.jit(jnp.einsum, static_argnums=0)("ij,ai,bj->ab", t_ij_ao, c, c)
    v_ijkl_mo = jax.jit(jnp.einsum, static_argnums=0)(
        "ijkl,ai,bj,ck,dl->abcd", v_ijkl_ao, c, c, c, c
    )
    const = jnp.asarray(float(mol.energy_nuc()))
    return const, t_ij_mo, v_ijkl_mo


def spinorb_from_spatial_sparse_coo2(tij_sparse, interleave=False, spin_values=[0, 1]):
    sparse = import_optional_dependency("sparse", descr="TV_from_pyscf_molecule")

    # Σ_ijσ t_ij c†_iσ c_jσ
    # for σ ∈ spin_values
    # interleave=True -> 2i+spin
    # interleave=False -> n*spin+i

    if isinstance(spin_values, int):
        spin_values = list(spin_values)
    # assert spin_values in [[0], [1], [0,1]]

    assert isinstance(tij_sparse, sparse.COO)
    n = tij_sparse.shape[0]

    if interleave:
        t = lambda x: 2 * x
        tp1 = lambda x: 2 * x + 1
    else:
        t = lambda x: x
        tp1 = lambda x: x + n

    i, j = tij_sparse.coords
    a = tij_sparse.data
    a2 = []
    i2 = []
    j2 = []
    if 0 in spin_values:
        a2.append(a)
        i2.append(t(i))
        j2.append(t(j))
    if 1 in spin_values:
        a2.append(a)
        i2.append(tp1(i))
        j2.append(tp1(j))
    i2 = np.concatenate(i2)
    j2 = np.concatenate(j2)
    a2 = np.concatenate(a2)
    return sparse.COO(np.array([i2, j2]), a2, shape=(2 * n, 2 * n))


def spinorb_from_spatial_sparse_coo4(
    vijkl_sparse, interleave=False, _order_preserving=False
):
    sparse = import_optional_dependency("sparse", descr="TV_from_pyscf_molecule")

    # Σ_ijklμσ v_ijkl c†_iμ c†_jσ c_kμ c_lσ
    # interleave=True -> 2i+spin
    # interleave=False -> n*spin+i
    assert isinstance(vijkl_sparse, sparse.COO)
    n = vijkl_sparse.shape[0]

    if interleave:
        t = lambda x: 2 * x
        tp1 = lambda x: 2 * x + 1
    else:
        t = lambda x: x
        tp1 = lambda x: x + n

    p, q, r, s = vijkl_sparse.coords
    b = vijkl_sparse.data

    if _order_preserving:
        # swap i/j and k/l
        a2 = np.concatenate([-b, -b, b, b])
        p2 = np.concatenate([tp1(q), tp1(p), t(p), tp1(p)])
        q2 = np.concatenate([t(p), t(q), t(q), tp1(q)])
        r2 = np.concatenate([tp1(r), tp1(s), t(r), tp1(r)])
        s2 = np.concatenate([t(s), t(r), t(s), tp1(s)])
    else:
        # same as openfermion
        a2 = np.concatenate([b, b, b, b])
        p2 = np.concatenate([t(p), tp1(p), t(p), tp1(p)])
        q2 = np.concatenate([tp1(q), t(q), t(q), tp1(q)])
        r2 = np.concatenate([tp1(r), t(r), t(r), tp1(r)])
        s2 = np.concatenate([t(s), tp1(s), t(s), tp1(s)])

    return sparse.COO(
        np.array([p2, q2, r2, s2]), a2, shape=(2 * n, 2 * n, 2 * n, 2 * n)
    )


def to_desc_order_sparse(vijkl_sparse, cutoff, set_zero_same=True):
    sparse = import_optional_dependency("sparse", descr="TV_from_pyscf_molecule")

    # !! use this only after adding spin, spinorb_from_spatial_sparse will be wrong
    # because the (implicitly assumed) symmetries of the tensor are not conserved

    # assume (1,1,0,0) are already index order (daggers are left/ij)
    # swap i/j k/l so that i>j and k>l

    # now swap the larger one to the left, will cause lots of them to cancel
    i, j, k, l = vijkl_sparse.coords
    a = vijkl_sparse.data.copy()
    a[np.where(i < j)] *= -1
    a[np.where(k < l)] *= -1
    if set_zero_same:
        # set to zero all those where we try to create / destroy two on the same orbital
        a[np.where(i == j)] *= 0
        a[np.where(k == l)] *= 0
    new_coords = np.array(
        [np.maximum(i, j), np.minimum(i, j), np.maximum(k, l), np.minimum(k, l)]
    )
    # use coo to merge same indices
    vijkl_sparse2 = sparse.COO(new_coords, a, shape=vijkl_sparse.shape)
    # we might have some new almost zeros from the cancellations, make sure they are 0
    new_coords2 = vijkl_sparse2.coords
    a2 = vijkl_sparse2.data
    mask = np.abs(a2) > cutoff
    return sparse.COO(new_coords2[:, mask], a2[mask], shape=vijkl_sparse.shape)


def spinorb_from_spatial_sparse(tij_sparse, vijkl_sparse, interleave=False):
    return spinorb_from_spatial_sparse_coo2(
        tij_sparse, interleave
    ), spinorb_from_spatial_sparse_coo4(vijkl_sparse, interleave)


def arrays_to_terms(
    const, tij_sparse, vijkl_sparse, term_conj2=(1, 0), term_conj4=(1, 0, 1, 0)
):
    ij = np.array(
        [
            tij_sparse.coords[0],
            np.ones_like(tij_sparse.coords[0], dtype=int) * term_conj2[0],
            tij_sparse.coords[1],
            np.ones_like(tij_sparse.coords[1], dtype=int) * term_conj2[1],
        ]
    ).T.reshape(-1, 2, 2)

    ijkl = np.array(
        [
            vijkl_sparse.coords[0],
            np.ones_like(vijkl_sparse.coords[0], dtype=int) * term_conj4[0],
            vijkl_sparse.coords[1],
            np.ones_like(vijkl_sparse.coords[1], dtype=int) * term_conj4[1],
            vijkl_sparse.coords[2],
            np.ones_like(vijkl_sparse.coords[2], dtype=int) * term_conj4[2],
            vijkl_sparse.coords[3],
            np.ones_like(vijkl_sparse.coords[3], dtype=int) * term_conj4[3],
        ]
    ).T.reshape(-1, 4, 2)
    terms = ij.tolist() + ijkl.tolist()
    weights = tij_sparse.data.tolist() + vijkl_sparse.data.tolist()
    return terms, weights, const


def operator_from_arrays(
    const,
    tij_sparse,
    vijkl_sparse,
    n_electrons,
    hi=None,
    cls=FermionOperator2nd,
    term_conj2=(1, 0),
    term_conj4=(1, 0, 1, 0),
):
    """
    H = const + Σ tᵢⱼ cᵢ†cⱼ + Σ vᵢⱼₖₗ cᵢ†cⱼcₖ†cₗ
    use term_conj4=(1,1,0,0) if you want cᵢ†cⱼ†cₖcₗ instead.
    use a zero matrix/tensor of the correct shape for tij and/or vijkl if you only want to use one of them
    """
    if hi is None:
        if isinstance(n_electrons, tuple):
            assert len(n_electrons) == 2
            hi = SpinOrbitalFermions(
                n_orbitals=tij_sparse.shape[0] // 2,
                s=1 / 2,
                n_fermions_per_spin=n_electrons,
            )
        else:
            hi = SpinOrbitalFermions(
                n_orbitals=tij_sparse.shape[0], n_fermions_per_spin=n_electrons
            )
    terms, weights, constant = arrays_to_terms(
        const, tij_sparse, vijkl_sparse, term_conj2=term_conj2, term_conj4=term_conj4
    )
    return cls(hi, terms, weights, constant)


def TV_from_pyscf_molecule(
    molecule,  # type: pyscf.gto.mole.Mole  # noqa: F821
    mo_coeff: np.ndarray,
    *,
    cutoff: float = 1e-11,
) -> tuple[...]:
    r"""
    Computes the nuclear repulsion energy :math:`E_{nuc}`, and the T and
    V tensors encoding the 1-body and 2-body terms in the electronic
    hamiltonian of a pyscf molecule using the specified molecular orbitals.

    The tensors returned correspond to the following expressions:

    .. math::

        \hat{H} = E_{nuc} + \sum_{ij} T_{ij} \hat{c}^\dagger_i\hat{c}_j +
            \sum_{ijkl} V_{ijkl} \hat{c}^\dagger_i\hat{c}_\dagger_j\hat{c}_k\hat{c}_l

    The electronic spin degree of freedom is encoded following the *NetKet convention*
    where the first :math:`N_{\downarrow}` values of the indices :math:`i,j,k,l` represent
    the spin down electrons, and the following :math:`N_{\uparrow}` values represent the
    spin up.

    .. note::

        In the `netket.experimental.operator.pyscf` module you can find some utility
        functions to convert from normal ordering to other orderings, but
        those are all internals so if you need them do copy-paste them somewhere else.

    Example:
        Constructs the hamiltonian for a Li-H molecule, using the `sto-3g` basis
        and the Boys orbitals

        >>> from pyscf import gto, scf, lo
        >>> import netket as nk; import netket.experimental as nkx
        >>>
        >>> geometry = [('Li', (0., 0., -1.5109/2)), ('H', (0., 0., 1.5109/2))]
        >>> mol = gto.M(atom=geometry, basis='STO-3G')
        >>>
        >>> # compute the boys orbitals
        >>> mf = scf.RHF(mol).run()  # doctest:+ELLIPSIS
            converged SCF energy = -7.86338...
        >>> mo_coeff = lo.Boys(mol).kernel(mf.mo_coeff)
        >>> ha = nkx.operator.pyscf.TV_from_pyscf_molecule(mol, mo_coeff)

    Args:
        molecule: The pyscf :class:`~pyscf.gto.mole.Mole` object describing the
            Hamiltonian
        mo_coeff: The molecular orbital coefficients determining the
            linear combination of atomic orbitals to produce the
            molecular orbitals. If unspecified this defaults to
            the hartree fock orbitals computed using :class:`~pyscf.scf.HF`.
        cutoff: Ignores all matrix elements in the `V` and `T` matrix that have
            magnitude less than this value. Defaults to :math:`10^{-11}`

    Returns:
        :code:`E,T,V`: a scalar and two numpy arrays, the first with 2 dimensions
        and the latter with 4 dimensions.

    """
    sparse = import_optional_dependency("sparse", descr="TV_from_pyscf_molecule")

    # Compute the T and V matrices
    E_nuc, tij, vijkl = compute_pyscf_integrals(molecule, mo_coeff)
    tij_sparse = sparse.COO.from_numpy(tij * (np.abs(tij) >= cutoff))
    vijkl_sparse = sparse.COO.from_numpy(vijkl * (np.abs(vijkl) >= cutoff))
    tij_sparse_spin, vijkl_sparse_spin = spinorb_from_spatial_sparse(
        tij_sparse, vijkl_sparse
    )
    vijkl_sparse_spin = to_desc_order_sparse(vijkl_sparse_spin, cutoff)

    return E_nuc, tij_sparse_spin, vijkl_sparse_spin


def from_pyscf_molecule(
    molecule,  # type: pyscf.gto.mole.Mole  # noqa: F821
    mo_coeff: np.ndarray | None = None,
    *,
    cutoff: float = 1e-11,
    implementation: DiscreteOperator = FermionOperator2nd,
) -> DiscreteOperator:
    r"""
    Construct a netket operator encoding the electronic hamiltonian of a pyscf
    molecule in a chosen orbital basis.

    Example:
        Constructs the hamiltonian for a Li-H molecule, using the `sto-3g` basis
        and the default Hartree-Fock molecular orbitals.

        >>> from pyscf import gto, scf, fci
        >>> import netket as nk; import netket.experimental as nkx
        >>>
        >>> bond_length = 1.5109
        >>> geometry = [('Li', (0., 0., -bond_length/2)), ('H', (0., 0., bond_length/2))]
        >>> mol = gto.M(atom=geometry, basis='STO-3G')
        >>>
        >>> mf = scf.RHF(mol).run()  # doctest:+ELLIPSIS
            converged SCF energy = -7.86338...
        >>> E_hf = sum(mf.scf_summary.values())
        >>>
        >>> E_fci = fci.FCI(mf).kernel()[0]
        >>>
        >>> ha = nkx.operator.from_pyscf_molecule(mol)  # doctest:+ELLIPSIS
            converged SCF energy = -7.86338...
        >>> E0 = float(nk.exact.lanczos_ed(ha))
        >>> print(f"{E0 = :.5f}, {E_fci = :.5f}")
        E0 = -7.88253, E_fci = -7.88253

    Example:
        Constructs the hamiltonian for a Li-H molecule, using the `sto-3g` basis
        and the Boys orbitals using :class:`~pyscf.lo.Boys`.

        >>> from pyscf import gto, scf, lo
        >>> import netket as nk; import netket.experimental as nkx
        >>>
        >>> bond_length = 1.5109
        >>> geometry = [('Li', (0., 0., -bond_length/2)), ('H', (0., 0., bond_length/2))]
        >>> mol = gto.M(atom=geometry, basis='STO-3G')
        >>>
        >>> # compute the boys orbitals
        >>> mf = scf.RHF(mol).run()  # doctest:+ELLIPSIS
            converged SCF energy = -7.86338...
        >>> mo_coeff = lo.Boys(mol).kernel(mf.mo_coeff)
        >>> # use the boys orbitals to construct the netket hamiltonian
        >>> ha = nkx.operator.from_pyscf_molecule(mol, mo_coeff=mo_coeff)

    Args:
        molecule: The pyscf :class:`~pyscf.gto.mole.Mole` object describing the
            Hamiltonian
        mo_coeff: The molecular orbital coefficients determining the
            linear combination of atomic orbitals to produce the
            molecular orbitals. If unspecified this defaults to
            the hartree fock orbitals computed using :class:`~pyscf.scf.HF`.
        cutoff: Ignores all matrix elements in the `V` and `T` matrix that have
            magnitude less than this value. Defaults to :math:`10^{-11}`
        implementation: The particular implementation to use for the operator.
            Different fermionic operator implementation might have different
            performances. Defaults to
            :class:`netket.experimental.operator.FermionOperator2nd` (this might
            change in the future).

    Returns:
        A netket second quantised operator that encodes the electronic hamiltonian.
    """
    # If unspecified, use HF molecular orbitals
    if mo_coeff is None:
        pyscf = import_optional_dependency("pyscf", descr="pyscf_molecule_to_arrays")

        mf = pyscf.scf.HF(molecule).run()
        mo_coeff = mf.mo_coeff

    E_nuc, Tij, Vijkl = TV_from_pyscf_molecule(molecule, mo_coeff, cutoff=cutoff)

    ha = operator_from_arrays(
        E_nuc,
        Tij,
        0.5 * Vijkl,
        molecule.nelec,
        term_conj4=(1, 1, 0, 0),
        cls=implementation,
    )
    # TODO maybe run setup and set _max_conn_size here estimating it analytially
    return ha
