from ast import keyword
from asyncore import read
import os
import shutil
import subprocess
import sys
import textwrap
from typing import Dict, List, TextIO

import numpy as np

from diceplayer.DPpack.Environment.Atom import Atom
from diceplayer.DPpack.Environment.Molecule import Molecule
from diceplayer.DPpack.Utils.Misc import *
from diceplayer.DPpack.Utils.PTable import *
from diceplayer.DPpack.Utils.StepDTO import StepDTO
from diceplayer.DPpack.Utils.Validations import NotNull


class Gaussian:

    mem = None
    chgmult = [0, 1]
    gmiddle = None  # In each case, if a filename is given, its content will be
    gbottom = None  # inserted in the gaussian input
    pop = "chelpg"

    keywords = ""

    def __init__(self) -> None:
        pass

    @NotNull(requiredArgs=["qmprog","level"])
    def updateKeywords(self, **data):
        self.__dict__.update(**data)
        self.checkKeywords()

    def checkKeywords(self):

        if self.pop not in ["chelpg", "mk", "nbo"]:
            self.pop = "chelpg"

    def run_formchk(self, cycle: int, fh: TextIO):

        simdir = "simfiles"
        stepdir = "step{:02d}".format(cycle)
        path = simdir + os.sep + stepdir + os.sep + "qm"

        work_dir = os.getcwd()
        os.chdir(path)

        fh.write("Formatting the checkpoint file... \n")

        exit_status = subprocess.call(["formchk", "asec.chk"], stdout=fh)

        fh.write("Done\n")

        os.chdir(work_dir)

    def readChargesFromFchk(self, file: str, fh: TextIO) -> List[float]:

        try:
            with open(file) as fchk:
                fchkfile = fchk.readlines()
        except:
            sys.exit("Error: cannot open file {}".format(file))

        if self.pop in ["chelpg", "mk"]:
            CHARGE_FLAG = "ESP Charges"
        else:
            CHARGE_FLAG = "ESP Charges"

        start = fchkfile.pop(0).strip()
        while start.find(CHARGE_FLAG) != 0:  # expression in begining of line
            start = fchkfile.pop(0).strip()

        charges: List[float] = []
        while len(charges) < len(self.step.molecule[0].atom):
            charges.extend([float(x) for x in fchkfile.pop(0).split()])

        return charges

    def read_forces_fchk(self, file: str, fh: TextIO) -> np.ndarray:

        forces = []
        try:
            with open(file) as tmpfh:
                fchkfile = tmpfh.readlines()
        except:
            sys.exit("Error: cannot open file {}".format(file))

        start = fchkfile.pop(0).strip()
        while start.find("Cartesian Gradient") != 0:  # expression in begining of line
            start = fchkfile.pop(0).strip()

        degrees = 3 * len(self.step.molecule[0].atom)
        count = 0
        while len(forces) < degrees:
            values = fchkfile.pop(0).split()
            forces.extend([float(x) for x in values])
            count += len(values)
            if count >= degrees:
                forces = forces[:degrees]
                break

        gradient = np.array(forces)

        fh.write("\nGradient read from file {}:\n".format(file))
        fh.write(
            "-----------------------------------------------------------------------\n"
            "Center     Atomic                     Forces (Hartree/Bohr)\n"
            "Number     Number              X                Y                Z\n"
            "-----------------------------------------------------------------------\n"
        )
        for i in range(len(self.step.molecule[0].atom)):
            fh.write(
                "  {:>5d}     {:>3d}        {:>14.9f}   {:>14.9f}   {:>14.9f}\n".format(
                    i + 1,
                    self.step.molecule[0].atom[i].na,
                    forces.pop(0),
                    forces.pop(0),
                    forces.pop(0),
                )
            )

        fh.write(
            "-----------------------------------------------------------------------\n"
        )

        force_max = np.amax(np.absolute(gradient))
        force_rms = np.sqrt(np.mean(np.square(gradient)))

        fh.write(
            "  Max Force = {:>14.9f}      RMS Force = {:>14.9f}\n\n".format(
                force_max, force_rms
            )
        )

        return gradient

    def read_hessian_fchk(self, file: str) -> np.ndarray:

        force_const = []
        try:
            with open(file) as tmpfh:
                fchkfile = tmpfh.readlines()
        except:
            sys.exit("Error: cannot open file {}".format(file))

        start = fchkfile.pop(0).strip()
        while start.find("Cartesian Force Constants") != 0:
            start = fchkfile.pop(0).strip()

        degrees = 3 * len(self.step.molecule[0].atom)
        last = round(degrees * (degrees + 1) / 2)
        count = 0

        while len(force_const) < last:

            value = fchkfile.pop(0).split()
            force_const.extend([float(x) for x in value])

        # while len(force_const) < last:

        # 	values = fchkfile.pop(0).split()
        # 	force_const.extend([ float(x) for x in values ])
        # 	count += len(values)
        # 	if count >= last:
        # 		force_const = force_const[:last]
        # 		break

        hessian = np.zeros((degrees, degrees))
        for i in range(degrees):
            for j in range(i + 1):
                hessian[i, j] = force_const.pop(0)
                hessian[j, i] = hessian[i, j]

        return hessian

    def read_hessian_log(self, file: str) -> np.ndarray:

        try:
            with open(file) as tmpfh:
                logfile = tmpfh.readlines()
        except:
            sys.exit("Error: cannot open file {}".format(file))

        start = logfile.pop(0).strip()
        while start.find("The second derivative matrix:") != 0:
            start = logfile.pop(0).strip()

        degrees = 3 * len(self.step.molecule[0].atom)
        hessian = np.zeros((degrees, degrees))

        k = 0
        while k < degrees:
            logfile.pop(0)
            for i in range(k, degrees):
                values = logfile.pop(0).split()[1:]
                for j in range(k, min(i + 1, k + 5)):
                    hessian[i, j] = float(values.pop(0))
                    hessian[j, i] = hessian[i, j]
            k += 5

        return hessian

    def print_grad_hessian(
        self, cycle: int, cur_gradient: np.ndarray, hessian: np.ndarray
    ) -> None:

        try:
            fh = open("grad_hessian.dat", "w")
        except:
            sys.exit("Error: cannot open file grad_hessian.dat")

        fh.write("Optimization cycle: {}\n".format(cycle))
        fh.write("Cartesian Gradient\n")
        degrees = 3 * len(self.step.molecule[0].atom)
        for i in range(degrees):
            fh.write(" {:>11.8g}".format(cur_gradient[i]))
            if (i + 1) % 5 == 0 or i == degrees - 1:
                fh.write("\n")

        fh.write("Cartesian Force Constants\n")
        n = int(np.sqrt(2 * degrees))
        last = degrees * (degrees + 1) / 2
        count = 0
        for i in range(n):
            for j in range(i + 1):
                count += 1
                fh.write(" {:>11.8g}".format(hessian[i, j]))
                if count % 5 == 0 or count == last:
                    fh.write("\n")

        fh.close()

    # Change the name to make_gaussian_input
    def make_gaussian_input(self, cycle: int, asec_charges: List[Dict]) -> None:

        simdir = "simfiles"
        stepdir = "step{:02d}".format(cycle)
        path = simdir + os.sep + stepdir + os.sep + "qm"

        file = path + os.sep + "asec.gjf"

        try:
            fh = open(file, "w")
        except:
            sys.exit("Error: cannot open file {}".format(file))

        fh.write("%Chk=asec.chk\n")
        if self.mem != None:
            fh.write("%Mem={}MB\n".format(self.mem))
        fh.write("%Nprocs={}\n".format(self.step.nprocs * self.step.ncores))

        kword_line = "#P " + str(self.level)

        if self.keywords != "":
            kword_line += " " + self.keywords

        if self.step.opt == "yes":
            kword_line += " Force"

        # kword_line += " Charge"
        kword_line += " NoSymm"
        kword_line += " Pop={} Density=Current".format(self.pop)

        if cycle > 1:
            kword_line += " Guess=Read"

        fh.write(textwrap.fill(kword_line, 90))
        fh.write("\n")

        fh.write("\nForce calculation - Cycle number {}\n".format(cycle))
        fh.write("\n")
        fh.write("{},{}\n".format(self.chgmult[0], self.chgmult[1]))

        for atom in self.step.molecule[0].atom:
            symbol = atomsymb[atom.na]
            fh.write(
                "{:<2s}    {:>10.5f}   {:>10.5f}   {:>10.5f}\n".format(
                    symbol, atom.rx, atom.ry, atom.rz
                )
            )

        fh.write("\n")

        for charge in asec_charges:
            fh.write(
                "{:>10.5f}   {:>10.5f}   {:>10.5f}     {:>11.8f}\n".format(
                    charge['rx'], charge['ry'], charge['rz'], charge['chg']
                )
            )

        fh.write("\n")
        
        fh.close()

    def read_charges(self, file: str, fh: TextIO) -> None:

        try:
            with open(file) as tmpfh:
                glogfile = tmpfh.readlines()
        except:
            sys.exit("Error: cannot open file {}".format(file))

        start = glogfile.pop(0).strip()
        while start != "Fitting point charges to electrostatic potential":
            start = glogfile.pop(0).strip()

        glogfile = glogfile[3:]  # Consume 3 more lines

        fh.write("\nAtomic charges:\n")
        fh.write("------------------------------------\n")
        for atom in self.step.molecule[0].atom:
            line = glogfile.pop(0).split()
            atom_str = line[1]
            charge = float(line[2])
            atom.chg = charge
            fh.write(" {:<2s}      {:>10.6f}\n".format(atom_str, charge))

        # if self.pop == "chelpg":
        # 	for ghost in ghost_atoms:
        # 		line = glogfile.pop(0).split()
        # 		atom_str = line[1]
        # 		charge = float(line[2])
        # 		ghost['chg'] = charge
        # 		fh.write(" {:<2s}      {:>10.6f}\n".format(atom_str, charge))

        # 	for lp in lp_atoms:
        # 		line = glogfile.pop(0).split()
        # 		atom_str = line[1]
        # 		charge = float(line[2])
        # 		lp['chg'] = charge
        # 		fh.write(" {:<2s}      {:>10.6f}\n".format(atom_str, charge))

        fh.write("------------------------------------\n")

    def executeOptimizationRoutine(self, cycle: int, outfile: TextIO, readhessian: str):

        try:
            gradient
            old_gradient = gradient
        except:
            pass

        gradient = self.read_forces_fchk(file, outfile)

        # If 1st step, read the hessian
        if cycle == 1:

            if readhessian == "yes":

                file = "grad_hessian.dat"
                outfile.write(
                    "\nReading the hessian matrix from file {}\n".format(file)
                )
                hessian = self.read_hessian_log(file)

            else:

                file = (
                    "simfiles"
                    + os.sep
                    + "step01"
                    + os.sep
                    + "qm"
                    + os.sep
                    + "asec.fchk"
                )
                outfile.write(
                    "\nReading the hessian matrix from file {}\n".format(file)
                )
                hessian = self.read_hessian_fchk(file)

        # From 2nd step on, update the hessian
        else:
            outfile.write("\nUpdating the hessian matrix using the BFGS method... ")
            hessian = self.step.molecule[0].update_hessian(
                step, gradient, old_gradient, hessian
            )
            outfile.write("Done\n")

        # Save gradient and hessian
        self.print_grad_hessian(cycle, gradient, hessian)

        # Calculate the step and update the position
        step = self.calculate_step(cycle, gradient, hessian)

        position += step

        ## If needed, calculate the charges
        if cycle < self.step.switchcyc:

            # internal.gaussian.make_charge_input(cycle, asec_charges)
            self.run_gaussian(cycle, "charge", outfile)

            file = (
                "simfiles"
                + os.sep
                + "step{:02d}".format(cycle)
                + os.sep
                + "qm"
                + os.sep
                + "asec2.log"
            )
            self.read_charges(file, outfile)
        else:
            file = (
                "simfiles"
                + os.sep
                + "step{:02d}".format(cycle)
                + os.sep
                + "qm"
                + os.sep
                + "asec.log"
            )
            self.read_charges(file, outfile)

        self.outfile.write("\nNew values for molecule type 1:\n\n")
        self.step.molecule[0].print_mol_info(outfile)

    def run_gaussian(self, cycle: int, type: str, fh: TextIO) -> None:

        simdir = "simfiles"
        stepdir = "step{:02d}".format(cycle)
        path = simdir + os.sep + stepdir + os.sep + "qm"
        work_dir = os.getcwd()
        os.chdir(path)

        # if type == "force":
        # 	infile = "asec.gjf"
        # elif type == "charge":
        # 	infile = "asec2.gjf"

        infile = "asec.gjf"

        fh.write(
            "\nCalculation of {}s initiated with Gaussian on {}\n".format(
                type, date_time()
            )
        )

        if shutil.which("bash") != None:
            exit_status = subprocess.call(
                [
                    "bash",
                    "-c",
                    "exec -a {}-step{} {} {}".format(
                        self.qmprog, cycle, self.qmprog, infile
                    ),
                ]
            )
        else:
            exit_status = subprocess.call([self.qmprog, infile])

        if exit_status != 0:
            sys.exit("Gaussian process did not exit properly")

        fh.write("Calculation of {}s finished on {}\n".format(type, date_time()))

        os.chdir(work_dir)

    # def calculate_step(
    #     self, cycle: int, gradient: np.ndarray, hessian: np.ndarray
    # ) -> np.ndarray:

    #     invhessian = np.linalg.inv(hessian)
    #     pre_step = -1 * np.matmul(invhessian, gradient.T).T
    #     maxstep = np.amax(np.absolute(pre_step))
    #     factor = min(1, self.player.maxstep / maxstep)
    #     step = factor * pre_step

    #     self.outfile.write("\nCalculated step-{}:\n".format(cycle))
    #     pre_step_list = pre_step.tolist()

    #     self.outfile.write(
    #         "-----------------------------------------------------------------------\n"
    #         "Center     Atomic                          Step (Bohr)\n"
    #         "Number     Number              X                Y                Z\n"
    #         "-----------------------------------------------------------------------\n"
    #     )
    #     for i in range(len(self.system.molecule[0].atom)):
    #         self.outfile.write(
    #             "  {:>5d}     {:>3d}        {:>14.9f}   {:>14.9f}   {:>14.9f}\n".format(
    #                 i + 1,
    #                 self.system.molecule[0].atom[i].na,
    #                 pre_step_list.pop(0),
    #                 pre_step_list.pop(0),
    #                 pre_step_list.pop(0),
    #             )
    #         )

    #     self.outfile.write(
    #         "-----------------------------------------------------------------------\n"
    #     )

    #     self.outfile.write("Maximum step is {:>11.6}\n".format(maxstep))
    #     self.outfile.write("Scaling factor = {:>6.4f}\n".format(factor))
    #     self.outfile.write("\nFinal step (Bohr):\n")
    #     step_list = step.tolist()

    #     self.outfile.write(
    #         "-----------------------------------------------------------------------\n"
    #         "Center     Atomic                          Step (Bohr)\n"
    #         "Number     Number              X                Y                Z\n"
    #         "-----------------------------------------------------------------------\n"
    #     )
    #     for i in range(len(self.system.molecule[0].atom)):
    #         self.outfile.write(
    #             "  {:>5d}     {:>3d}        {:>14.9f}   {:>14.9f}   {:>14.9f}\n".format(
    #                 i + 1,
    #                 self.system.molecule[0].atom[i].na,
    #                 step_list.pop(0),
    #                 step_list.pop(0),
    #                 step_list.pop(0),
    #             )
    #         )

    #     self.outfile.write(
    #         "-----------------------------------------------------------------------\n"
    #     )

    #     step_max = np.amax(np.absolute(step))
    #     step_rms = np.sqrt(np.mean(np.square(step)))

    #     self.outfile.write(
    #         "  Max Step = {:>14.9f}      RMS Step = {:>14.9f}\n\n".format(
    #             step_max, step_rms
    #         )
    #     )

    #     return step

    def configure(self, step: StepDTO):

        self.step = step

    def start(self, cycle: int, outfile: TextIO, asec_charges: List[Dict], readhessian: str) -> StepDTO:

        make_qm_dir(cycle)

        if cycle > 1:

            src = (
                "simfiles"
                + os.sep
                + "step{:02d}".format(cycle - 1)
                + os.sep
                + "qm"
                + os.sep
                + "asec.chk"
            )
            dst = (
                "simfiles"
                + os.sep
                + "step{:02d}".format(cycle)
                + os.sep
                + "qm"
                + os.sep
                + "asec.chk"
            )
            shutil.copyfile(src, dst)

        self.make_gaussian_input(cycle, asec_charges)
        self.run_gaussian(cycle, "force", outfile)
        self.run_formchk(cycle, outfile)

        ## Read the gradient
        file = (
            "simfiles"
                + os.sep
                + "step{:02d}".format(cycle)
                + os.sep
                + "qm"
                + os.sep
                + "asec.fchk"
            )

        if self.step.opt:

            pass
            # position = self.executeOptimizationRoutine(cycle, outfile, readhessian)
            # self.step.position = position

        else:

            charges = self.readChargesFromFchk(file, outfile)
            self.step.charges = charges

        return self.step

    def reset(self):

        del self.step