#!/usr/bin/env python3
## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab
## ---------------------------------------------------------------------
##
## Copyright (C) 2019 by the adcc authors
##
## This file is part of adcc.
##
## adcc is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published
## by the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## adcc is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with adcc. If not, see <http://www.gnu.org/licenses/>.
##
## ---------------------------------------------------------------------
import numpy as np

from libadcc import HartreeFockProvider

import psi4

from .EriBuilder import EriBuilder
from ..exceptions import InvalidReference
from ..ElectronicStates import EnergyCorrection
from ..OneParticleOperator import OneParticleOperator


class Psi4OperatorIntegralProvider:
    available: tuple[str, ...] = (
        "overlap", "electric_dipole", "electric_dipole_velocity", "magnetic_dipole",
        "electric_quadrupole", "electric_quadrupole_traceless",
        "pe_induction_elec", "pcm_potential_elec"
    )

    def __init__(self, wfn):
        self.wfn = wfn
        self.backend = "psi4"
        self.mints = psi4.core.MintsHelper(self.wfn)

    @property
    def overlap(self) -> np.ndarray:
        return np.asarray(self.mints.ao_overlap())

    @property
    def electric_dipole(self) -> tuple[np.ndarray, ...]:
        """-sum_i r_i"""
        return tuple(np.asarray(comp) for comp in self.mints.ao_dipole())

    def magnetic_dipole(self, gauge_origin="origin") -> tuple[np.ndarray, ...]:
        """
        The imaginary part of the integral is returned.
        -0.5 * sum_i r_i x p_i
        """
        # TODO: Gauge origin?
        if gauge_origin != (0.0, 0.0, 0.0) and gauge_origin != "origin":
            raise NotImplementedError(
                "For Psi4 only the origin (0.0, 0.0, 0.0) can be selected as"
                " gauge origin for the magnetic dipole operator. "
                f"{gauge_origin} is not valid."
            )
        return tuple(
            0.5 * np.asarray(comp)
            for comp in self.mints.ao_angular_momentum()
        )

    @property
    def electric_dipole_velocity(self) -> tuple[np.ndarray, ...]:
        """
        The imaginary part of the integral is returned.
        -sum_i p_i
        """
        return tuple(-1.0 * np.asarray(comp) for comp in self.mints.ao_nabla())

    def electric_quadrupole(self, gauge_origin="origin") -> tuple[np.ndarray, ...]:
        """-sum_i r_{i, alpha} r_{i, beta}"""
        # TODO: Gauge origin?
        if gauge_origin != (0.0, 0.0, 0.0) and gauge_origin != "origin":
            raise NotImplementedError(
                "For Psi4 only the origin (0.0, 0.0, 0.0) can be selected as"
                " gauge origin for the electric quadrupole operator. "
                f"{gauge_origin} is not valid."
            )
        # Expand 6 upper-triangular components to 9 symmetric matrix entries.
        u = [np.asarray(comp) for comp in self.mints.ao_quadrupole()]
        assert len(u) == 6
        return (u[0], u[1], u[2], u[1], u[3], u[4], u[2], u[4], u[5])

    def electric_quadrupole_traceless(self, gauge_origin="origin"
                                      ) -> tuple[np.ndarray, ...]:
        """
        -0.5 * sum_i (3 * r_{i, alpha} r_{i, beta}
        - delta_{alpha, beta} r_{i}^2)
        """
        # TODO: Gauge origin?
        if gauge_origin != (0.0, 0.0, 0.0) and gauge_origin != "origin":
            raise NotImplementedError(
                "For Psi4 only the origin (0.0, 0.0, 0.0) can be selected as"
                " gauge origin for the traceless electric quadrupole operator. "
                f"{gauge_origin} is not valid."
            )
        # Expand 6 upper-triangular components to 9 symmetric matrix entries.
        u = [np.asarray(comp) for comp in self.mints.ao_traceless_quadrupole()]
        assert len(u) == 6
        return (u[0], u[1], u[2], u[1], u[3], u[4], u[2], u[4], u[5])

    def pe_induction_elec(self, dm: OneParticleOperator) -> psi4.core.Matrix:
        if not hasattr(self.wfn, "pe_state"):
            raise RuntimeError("Can not compute the PE electronic induction "
                               "operator using the given psi4 object.")

        return self.wfn.pe_state.get_pe_contribution(
            psi4.core.Matrix.from_array(dm.to_ndarray()),
            elec_only=True
        )[1]

    def pcm_potential_elec(self, dm: OneParticleOperator) -> psi4.core.Matrix:
        if hasattr(self.wfn, "ddx"):
            return self.wfn.ddx.get_solvation_contributions(
                psi4.core.Matrix.from_array(dm.to_ndarray()),
                elec_only=True, nonequilibrium=True
            )[1]
        elif self.wfn.PCM_enabled():
            return psi4.core.PCM.compute_V(
                self.wfn.get_PCM(),
                psi4.core.Matrix.from_array(dm.to_ndarray())
            )
        raise RuntimeError("Can not compute the electronic PCM potential "
                           "operator using the given psi4 object.")


class Psi4EriBuilder(EriBuilder):
    def __init__(self, wfn, n_orbs, n_orbs_alpha, n_alpha, n_beta, restricted):
        self.wfn = wfn
        self.mints = psi4.core.MintsHelper(self.wfn)
        super().__init__(n_orbs, n_orbs_alpha, n_alpha, n_beta, restricted)

    @property
    def coefficients(self):
        return {
            "Oa": self.wfn.Ca_subset("AO", "OCC"),
            "Ob": self.wfn.Cb_subset("AO", "OCC"),
            "Va": self.wfn.Ca_subset("AO", "VIR"),
            "Vb": self.wfn.Cb_subset("AO", "VIR"),
        }

    def compute_mo_eri(self, blocks, spins):
        coeffs = tuple(self.coefficients[blocks[i] + spins[i]] for i in range(4))
        return np.asarray(self.mints.mo_eri(*coeffs))


class Psi4HFProvider(HartreeFockProvider):
    """
        This implementation is only valid
        if no orbital reordering is required.
    """
    def __init__(self, wfn):
        # Do not forget the next line,
        # otherwise weird errors result
        super().__init__()
        self.wfn = wfn
        self.eri_builder = Psi4EriBuilder(self.wfn, self.n_orbs, self.wfn.nmo(),
                                          wfn.nalpha(), wfn.nbeta(),
                                          self.restricted)
        self.operator_integral_provider = Psi4OperatorIntegralProvider(self.wfn)

        self.environment = None
        self.environment_implementation = None
        if hasattr(self.wfn, "pe_state"):
            self.environment = "pe"
            self.environment_implementation = "cppe"
        elif hasattr(self.wfn, "ddx"):
            self.environment = "pcm"
            self.environment_implementation = "ddx"
        elif self.wfn.PCM_enabled():
            self.environment = "pcm"
            self.environment_implementation = "pcmsolver"

    def pe_energy(self, dm, elec_only=True):
        density_psi = psi4.core.Matrix.from_array(dm.to_ndarray())
        e_pe, _ = self.wfn.pe_state.get_pe_contribution(density_psi,
                                                        elec_only=elec_only)
        return e_pe

    def pcm_energy(self, dm, elec_only=True):
        psi_dm = psi4.core.Matrix.from_array(dm.to_ndarray())
        # computes the Fock matrix contribution.
        # By contraction with the tdm, the electronic energy contribution is
        # calculated.
        if self.environment_implementation == "ddx":
            V_pcm = self.wfn.ddx.get_solvation_contributions(
                psi_dm, elec_only=elec_only, nonequilibrium=True
            )[1]
        elif self.environment_implementation == "pcmsolver":
            V_pcm = psi4.core.PCM.compute_V(self.wfn.get_PCM(), psi_dm)
        return np.einsum("uv,uv->", dm.to_ndarray(), V_pcm.to_array())

    @property
    def excitation_energy_corrections(self):
        ret = []
        if self.environment == "pe":
            ptlr = EnergyCorrection(
                "pe_ptlr_correction",
                lambda view: 2.0 * self.pe_energy(view.transition_dm_ao,
                                                  elec_only=True)
            )
            ptss = EnergyCorrection(
                "pe_ptss_correction",
                lambda view: self.pe_energy(view.state_diffdm_ao,
                                            elec_only=True)
            )
            ret.extend([ptlr, ptss])
        if self.environment == "pcm":
            ptlr = EnergyCorrection(
                "pcm_ptlr_correction",
                lambda view: self.pcm_energy(view.transition_dm_ao,
                                             elec_only=True)
            )
            ret.extend([ptlr])
        return {ec.name: ec for ec in ret}

    def get_backend(self):
        return "psi4"

    def get_conv_tol(self):
        conv_tol = psi4.core.get_option("SCF", "E_CONVERGENCE")
        # RMS value of the orbital gradient
        conv_tol_grad = psi4.core.get_option("SCF", "D_CONVERGENCE")
        threshold = max(conv_tol, conv_tol_grad)
        return threshold

    def get_restricted(self):
        return isinstance(self.wfn, (psi4.core.RHF, psi4.core.ROHF))

    def get_energy_scf(self):
        return self.wfn.energy()

    def get_spin_multiplicity(self):
        return self.wfn.molecule().multiplicity()

    def get_n_orbs_alpha(self):
        return self.wfn.nmo()

    def get_n_bas(self):
        return self.wfn.basisset().nbf()

    def get_nuclear_multipole(self, order, gauge_origin=(0, 0, 0)):
        molecule = self.wfn.molecule()
        if order == 0:
            # The function interface needs to be a np.array on return
            return np.array([sum(molecule.charge(i)
                                 for i in range(molecule.natom()))])
        elif order == 1:
            dip_nuclear = molecule.nuclear_dipole()
            return np.array([dip_nuclear[0], dip_nuclear[1], dip_nuclear[2]])
        else:
            raise NotImplementedError("get_nuclear_multipole with order > 1")

    def transform_gauge_origin_to_xyz(self, gauge_origin):
        raise NotImplementedError("transform_gauge_origin_to_xyz not implemented.")

    def fill_orbcoeff_fb(self, out):
        mo_coeff_a = np.asarray(self.wfn.Ca())
        mo_coeff_b = np.asarray(self.wfn.Cb())
        mo_coeff = (mo_coeff_a, mo_coeff_b)
        out[:] = np.transpose(
            np.hstack((mo_coeff[0], mo_coeff[1]))
        )

    def fill_occupation_f(self, out):
        out[:] = np.hstack((
            np.asarray(self.wfn.occupation_a()),
            np.asarray(self.wfn.occupation_b())
        ))

    def fill_orben_f(self, out):
        orben_a = np.asarray(self.wfn.epsilon_a())
        orben_b = np.asarray(self.wfn.epsilon_b())
        out[:] = np.hstack((orben_a, orben_b))

    def fill_fock_ff(self, slices, out):
        diagonal = np.empty(self.n_orbs)
        self.fill_orben_f(diagonal)
        out[:] = np.diag(diagonal)[slices]

    def fill_eri_ffff(self, slices, out):
        self.eri_builder.fill_slice_symm(slices, out)

    def fill_eri_phys_asym_ffff(self, slices, out):
        raise NotImplementedError("fill_eri_phys_asym_ffff not implemented.")

    def has_eri_phys_asym_ffff(self):
        return False

    def flush_cache(self):
        self.eri_builder.flush_cache()


def import_scf(wfn):
    if not isinstance(wfn, psi4.core.HF):
        raise InvalidReference(
            "Only psi4.core.HF and its subtypes are supported references in "
            "backends.psi4.import_scf. This indicates that you passed an "
            "unsupported SCF reference. Make sure you did a restricted or "
            "unrestricted HF calculation."
        )

    if not isinstance(wfn, (psi4.core.RHF, psi4.core.UHF)):
        raise InvalidReference("Right now only RHF and UHF references are "
                               "supported for Psi4.")

    # TODO This is not fully correct, because the core.Wavefunction object
    #      has an internal, but py-invisible Options structure, which contains
    #      the actual set of options ... theoretically they could differ
    scf_type = psi4.core.get_global_option('SCF_TYPE')
    # CD = Choleski, DF = density-fitting
    unsupported_scf_types = ["CD", "DISK_DF", "MEM_DF"]
    if scf_type in unsupported_scf_types:
        raise InvalidReference("Unsupported Psi4 SCF_TYPE, should not be one "
                               f"of {unsupported_scf_types}")

    if wfn.nirrep() > 1:
        raise InvalidReference("The passed Psi4 wave function object needs to "
                               "have exactly one irrep, i.e. be of C1 symmetry.")

    # Psi4 throws an exception if SCF is not converged, so there is no need
    # to assert that here.
    provider = Psi4HFProvider(wfn)
    return provider


def run_hf(xyz, basis, charge=0, multiplicity=1, conv_tol=1e-11,
           conv_tol_grad=1e-9, max_iter=150, pe_options=None):
    basissets = {
        "sto3g": "sto-3g",
        "def2tzvp": "def2-tzvp",
        "ccpvdz": "cc-pvdz",
    }

    # needed for PE and PCM tests
    psi4.core.clean_options()
    mol = psi4.geometry(f"""
        {charge} {multiplicity}
        {xyz}
        symmetry c1
        units au
        no_reorient
        no_com
    """)

    psi4.core.be_quiet()
    psi4.set_options({
        'basis': basissets.get(basis, basis),
        'scf_type': 'pk',
        'e_convergence': conv_tol,
        'd_convergence': conv_tol_grad,
        'maxiter': max_iter,
        'reference': "RHF",
    })
    if pe_options:
        psi4.set_options({"pe": "true"})
        psi4.set_module_options("pe", {"potfile": pe_options["potfile"]})

    if multiplicity != 1:
        psi4.set_options({
            'reference': "UHF",
            'maxiter': max_iter + 500,
            'soscf': 'true'
        })

    _, wfn = psi4.energy('SCF', return_wfn=True, molecule=mol)
    psi4.core.clean()
    return wfn
