#  Copyright (C) 2023 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell 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 Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

import ctypes
from ..qm import _QM

try:
    from pyscf import __config__

    # This prevents a warning being thrown pyscf which causes chemshell not to run
    __config__.B3LYP_WITH_VWN5 = True
except (ModuleNotFoundError, ValueError):
    pyscf = None


__author__ = "Joseph Thacker <joseph.thacker@stfc.ac.uk>"


class PySCF(_QM):
    def _update_molecule(self, init: bool = False):
        import pyscf
        import pyscf.qmmm

        # define the molecule string
        coord_string = ""
        for i, znum in enumerate(self.frag.znums):
            coords = self.frag.coords[i, :]
            coord_string += f"{znum} {coords[0]} {coords[1]} {coords[2]}; "

        self.molecule_pyscf = pyscf.gto.M(
            atom=coord_string,
            unit="au",
            charge=self.charge,
            spin=self.mult - 1,
            symmetry=False,
            basis=self.basis,
        )

        if init:
            self.pyscf_method = self.pyscf_method(self.molecule_pyscf)
        else:
            self.pyscf_method.mol = self.molecule_pyscf

        pyscf_job = self.pyscf_method
        pyscf_job.conv_tol = self.threshold
        # Max number of iterations
        pyscf_job.max_cycle = self.maxiter
        # Name of SCF restart file
        pyscf_job.chkfile = "_pySCF.restart"
        # Control initial SCF (restart/guess)
        if self.restart:
            pyscf_job.init_guess = "chk"
        else:
            pyscf_job.init_guess = self.guess

        if hasattr(pyscf_job, "xc"):
            pyscf_job.xc = self.functional
            pyscf_job.grids.level = self.gridsize
            pyscf_job.grids.prune = None

        # Add the QM/MM point charges
        # if self.frag.nbqs:
        #    raise RuntimeError("The PySCF interface does not yet support electrostatic embedding")
        #    point_charge_coords = []
        #    point_charges = []
        #    for i, q in enumerate(self.frag.bqs.charges):
        #        point_charge_coords = point_charge_coords + [
        #            (
        #                self.frag.bqs.coords[i, 0],
        #                self.frag.bqs.coords[i, 1],
        #                self.frag.bqs.coords[i, 2],
        #            )
        #        ]
        #        point_charges = point_charges + [q]

        #    self.method_pyscf = pyscf.qmmm.mm_charge(
        #        pyscf_job, point_charge_coords, point_charges
        #    )
        ## No QM/MM
        # else:
        self.method_pyscf = pyscf_job

    def __init__(
        self,
        method,
        functional: str = "blyp",
        gridsize: int = 5,
        guess: str = "minao",
        maxiter: int = 100,
        restart: bool = False,
        fcidump: bool = False,
        threshold=1e-8,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        from ...base import errors
        from ...utils import testutils

        # Get chemshell to throw an unsupported error if pyscf not installed
        # or if it is installed but not importable
        try:
            import pyscf
        except ModuleNotFoundError:
            print(" >>> pyscf module not found.")
            testutils.setUnsupported()
            exit(errors.UNSUPPORTED)
        except ValueError:
            print(" >>> Error: pyscf module found but could not be imported.")
            testutils.setUnsupported()
            exit(errors.UNSUPPORTED)

        import pyscf.scf
        import pyscf.dft

        self.fcidump = fcidump
        self.guess = guess
        self.functional = functional
        self.gridsize = gridsize
        self.threshold = threshold
        self.restart = restart
        self.maxiter = maxiter
        if type(method) == str:
            if method.lower() == "hf" and self.scftype.lower() == "rhf":
                self.pyscf_method = pyscf.scf.RHF
            elif method.lower() == "hf" and self.scftype.lower() == "uhf":
                self.pyscf_method = pyscf.scf.UHF
            elif method.lower() == "hf" and self.scftype.lower() == "rohf":
                self.pyscf_method = pyscf.scf.ROHF
            elif method.lower() == "dft" and self.scftype.lower() == "rhf":
                self.pyscf_method = pyscf.dft.RKS
            elif method.lower() == "dft" and self.scftype.lower() == "uhf":
                self.pyscf_method = pyscf.dft.UKS
            elif method.lower() == "dft" and self.scftype.lower() == "rohf":
                self.pyscf_method = pyscf.dft.ROKS
            else:
                raise ValueError(f"Method: {method} unsupported")
        else:
            self.pyscf_method = method

        self._requires_init = True

    def run(self, **kwargs) -> ctypes.c_long:
        self._update_molecule(self._requires_init)
        self._requires_init = False

        try:  # Gradients can be set in 2 places?
            self._gradients = kwargs["_gradients"]
        except KeyError:
            pass

        if self.frag.nbqs:
            from .ComputeElectrostaticEmbedding import ComputeElectrostaticEmbedding

            if self._gradients:
                import numpy as np

                (
                    self._result.energy,
                    self._result.gradients,
                    self.frag.bqs.gradients,
                ) = ComputeElectrostaticEmbedding(
                    self.molecule_pyscf,
                    self.method_pyscf,
                    np.array(self.frag.bqs.coords),
                    np.array(self.frag.bqs.charges).flatten(),
                    self._gradients,
                )
            else:
                self._result.energy = ComputeElectrostaticEmbedding(
                    self.molecule_pyscf,
                    self.method_pyscf,
                    np.array(self.frag.bqs.coords),
                    np.array(self.frag.bqs.charges).flatten(),
                    self._gradients,
                )
        else:
            if self._gradients:
                grad_method = self.method_pyscf.nuc_grad_method().as_scanner()
                self._result.energy, self._result.gradients = grad_method(
                    self.molecule_pyscf
                )
            else:
                self._result.energy = self.method_pyscf.kernel()

        if self.fcidump:
            from pyscf.tools import fcidump

            fcidump.from_scf(self.method_pyscf, "_pyscf.fcidump")

        return ctypes.c_long(1)
