#  Copyright (C) 2022 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/>.

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

from ..qm import _QM
from ...dl_py2f import utils


class TURBOMOLE(_QM):
    """Interface to TURBOMOLE"""

    _kwargs = {
        "dscf_out": "output_turbomole",
        "grad_out": "_grad.out",
        "basis": "SVP",
        # input, output and _syscmd/path must all be same length tuples
        "output": "output_turbomole",
        "input": ("control", "coord"),
        "gradfile": "gradient",
        "point_charge_gradfile": "point_charge_gradients",
        "convergence": 8,
        "functional": "b-lyp",
        "maxiter": 30,
        "grid": "medium",
        "rohf_control": None,
        "energyfile": "energy",
        "use_ri": False,
        "self_gradients": False,
        "nprocs": 1,
        "user_control": None,
    }

    _internals = {
        "_syscmd": ("dscf"),
    }

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        if self.use_ri:
            self._syscmd = "ridft"

        if self.user_control != None:
            if type(self.user_control) == str:
                print(f"User supplied control file: {self.user_control}")
                if self.user_control.lower() == "control":
                    raise RuntimeError(
                        f"Cannot name user defined control file '{self.user_control}'. Please rename file."
                    )
            else:
                raise TypeError(
                    f"Cannot accept control file with type '{type(self.user_control)}'"
                )

    @property
    def _gradients(self):
        """grad getter"""
        return self.self_gradients

    @_gradients.setter
    def _gradients(self, do_grad: bool):
        "grad setter"

        # only set if changing the gradients flag
        if self.self_gradients != do_grad:
            # if we require gradients
            if do_grad:
                # Set the command path
                if self._use_ri:
                    self._syscmd = ("ridft", "rdgrad")
                else:
                    self._syscmd = ("dscf", "grad")
                # Set the outputs to include gradient output
                self.output = ("output_turbomole", "_grad.out")
                # Set the inputs to require multiple sets of files
                self.input = (
                    ("control", "coord"),
                    None,
                )

            # otherwise we have no need to use gradients
            else:
                if self._use_ri:
                    self._syscmd = "ridft"
                else:
                    self._syscmd = "dscf"

        self.self_gradients = do_grad

    @property
    def grid(self):
        """Grid accuracy"""

        """
        TODO: Turbomole has 2 types of grids:
            -- 1-5: which are single grids from coarse to dense
            -- m3-m5: multigrids which are also from coarse to dense
            It is generally recommended to use 'm3' in most cases,
            so we set this to medium
        """
        selectcases = {
            "low": "1",
            "medium": "m3",
            "high": "m4",
            "veryhigh": "m5",  # These are the same!
            "ultrahigh": "m5",  # These are the same!
        }

        # Use grids we define
        if self._grid.lower() in selectcases.keys():
            return selectcases.get(self._grid.lower(), "")
        # Allow the user to input their own grid!
        else:
            return self._grid.lower()

    @grid.setter
    def grid(self, val):
        """Setter of grid"""

        self._grid = val

    @property
    def use_ri(self):
        """Resolution of the Identity"""

        return self._use_ri

    @use_ri.setter
    def use_ri(self, val):
        """Setter of RI"""

        self._use_ri = val

    @property
    def functional(self):
        """DFT functional"""

        return self._functional.lower()

    @functional.setter
    def functional(self, val):
        """Setter of functional"""

        self._functional = val

    def getInpBuff(self, input="null"):
        """
        Get the contents of the two input files 'control' and 'config'
        """
        if input == "null":
            raise RuntimeError(
                "getInpBuff called without arguments -- not allowed for TURBOMOLE"
            )

        if input not in ["coord", "control", None]:
            raise IOError(
                "This is an internal error in ChemShell, the input file member variable (self.input) is not consistent for TURBOMOLE."
            )

        # archive old files
        if input != None:
            utils.fileutils.archiveFile(
                [input, self.output, self.gradfile, self.point_charge_gradfile]
            )

        # Write the coordinates file
        if input == "coord":
            configbuff = "$coord\n"
            # get the coordinates in standard form
            coordstr = self.frag.coords2str(charges=None, unit="a.u.")

            # reformat the coordinates so the element type is on the end
            reform_xyz = []
            for s in coordstr.split("\n"):
                splt = s.split()
                if len(splt) != 4:
                    continue
                reform_xyz = reform_xyz + [
                    splt[1],
                    splt[2],
                    splt[3],
                    splt[0].lower(),
                    "\n",
                ]
            # Write the coords to the config file
            configbuff += " ".join(reform_xyz)
            configbuff += "$end"
            strbuff = configbuff

        elif input == "control" and self.user_control != None:

            # Check the user control file
            with open(self.user_control, "r") as usrctrl:
                lines = usrctrl.readlines()

            for _, l in enumerate(lines):
                if "coord" in l:
                    config_name = l.split("=")[-1].strip()
                    if config_name != self.input[1]:
                        print(config_name, self.input[1])
                        raise RuntimeError(
                            f"In user defined control file please ensure the 'coord' file is named '{self.input[1]}', not {config_name}"
                        )
                if "symm" in l:
                    symm_name = l.split()[-1].strip()
                    if symm_name.lower() != "c1":
                        raise RuntimeError(
                            f"Symmetry must be 'c1' in user defined control file, not {symm_name}"
                        )
                if "$grad" in l:
                    raise RuntimeError(
                        f'In user defined control file, please do not set "$grad file=..."'
                    )
                if "$point_charge_gradients" in l:
                    raise RuntimeError(
                        'In user defined control file, please do not set "$point_charge_gradients file=..."'
                    )

            strbuff = "".join(lines)
            strbuff = strbuff.replace("$end", "")

        # Write the control file
        elif input == "control" and self.user_control == None:
            controlbuff = ""

            # Set the number of threads
            controlbuff += f"$smp_cpus {self.nprocs}\n"

            # Add the basis set and the coordinates of the atoms
            if "$atoms" in self.basis.lower():
                controlbuff += self.basis + "\n"
            else:
                controlbuff += "$atoms\n" + f"basis = {self.basis}\n"
            controlbuff += f"$coord file={self.input[1]}\n"
            controlbuff += f"$scfconv {self.convergence}\n"
            if self.use_ri:
                controlbuff += "$rij\n"

            scftypes = {"rhf": "rhf", "uhf": "uhf"}

            # uhf/rhf/rohf
            try:
                method = str(scftypes.get(self.scftype))
            except KeyError:
                print(
                    ' >>> ChemShell ERROR: unknown/unsupported method "{}"'.format(
                        self.method.lower()
                    )
                )
                raise

            # Add the charge and number of unpaired electrons to the input file
            nunpaired = self.mult - 1
            controlbuff += f"$eht charge={self.charge} unpaired={nunpaired}\n"

            # Sometimes just setting the charge and unpaired electrons is not enough!
            # So we have to explicity state the closed shells!
            nelec = self.frag.znums.sum() - self.charge
            npaired = nelec - nunpaired
            if npaired % 2 != 0:
                raise ValueError(
                    f"Number of paired electrons is odd: {npaired}"
                )

            selectcases = {
                "hf": f"${method}\n",
                # "tdhf": "hf",
                # "td-hf": "hf",
                "dft": "$dft\n"
                + f" functional {self._functional}\n"
                + f" gridsize {self.grid}\n"
                + f" ${method}\n",
                # "tddft": self.functional,
                # "td-dft": self.functional,
            }

            # Set the closed shells -- this bit is easy!
            controlbuff += "$symmetry c1\n"  # remove symmetry
            nclosedshells = int(npaired / 2)
            controlbuff += "$closed shells\n"
            controlbuff += f"a 1-{nclosedshells} ( 2 )\n"
            controlbuff += selectcases[self.method.lower()]
            controlbuff += f"$scfiterlimit {self.maxiter}\n"

            strbuff = controlbuff

        else:
            raise NotImplementedError(
                f"No write function implemented for file {input}"
            )
        if input == "control":
            # Add the QM/MM point charges to the control file
            if self.frag.nbqs:
                strbuff += "$point_charges nocheck list longprint=3\n"
                for i, q in enumerate(self.frag.bqs.charges):
                    strbuff += f"{self.frag.bqs.coords[i,0]} {self.frag.bqs.coords[i,1]} {self.frag.bqs.coords[i,2]} {q}\n"
                if self._gradients:
                    strbuff += "$drvopt point charges\n"
                    strbuff += f"$point_charge_gradients file={self.point_charge_gradfile}\n"

            if self._gradients:
                strbuff += f"$grad file={self.gradfile}\n"

            strbuff += "$end"

        return strbuff

    def parseMainOutput(self):
        import numpy as np

        """Parse main output file self.output"""

        # energy (only parse when the output is from dscf)
        if self.output == self.dscf_out:
            energy = None
            with open(self.output, "r") as f:
                for l_no, line in enumerate(f):
                    if "|  total energy" in line:
                        energy = line.split()[-2]
                        break
            if energy is None:
                raise ValueError(
                    "Total energy not found in TURBOMOLE output file."
                )
            self._result.energy = float(energy)

        # gradients (only parse when the output is from grad)
        elif self._gradients and self.output == self.grad_out:
            gradients = []
            with open(self.gradfile, "r") as f:
                linerange = [-1]
                for l_no, line in enumerate(f):
                    # cycle should ALWAYS be on the second line
                    if "cycle" in line:
                        start_l = l_no + self.natoms + 1
                        end_l = start_l + self.natoms
                        linerange = list(range(start_l, end_l))
                        gradients = []

                    if l_no in linerange:
                        gradients = gradients + [line]
                gradients = np.array(
                    [
                        g.replace("D", "E").replace("\n", "").split()
                        for g in gradients
                    ],
                    dtype=np.float64,
                )

            if len(gradients) == 0:
                raise ValueError("Gradients not found for TURBOMOLE")

            gradients = np.array(gradients)

            self._result.gradients = np.copy(gradients)

        # gradients + point-charge gradients#
        # (only parse when the output is from grad)
        if self._gradients and self.output == self.grad_out and self.frag.nbqs:
            with open(self.point_charge_gradfile, "r") as f:
                lines = f.readlines()[1:-1]
                pcgradients = np.array(
                    [
                        np.array(
                            l.replace("\n", " ").replace("D", "E").split(),
                            dtype=np.float64,
                        )
                        for l in lines
                    ]
                )
            self.frag.bqs.gradients = np.copy(pcgradients)
