# Source code for psi4.driver.diatomic

    #
    # @BEGIN LICENSE
    #
    # Psi4: an open-source quantum chemistry software package
    #
    # Copyright (c) 2007-2019 The Psi4 Developers.
    #
    # The copyrights for code used from other parties are included in
    # the corresponding files.
    #
    # This file is part of Psi4.
    #
    # Psi4 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, version 3.
    #
    # Psi4 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 Psi4; if not, write to the Free Software Foundation, Inc.,
    # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
    #
    # @END LICENSE
    #

    import numpy as np

    from psi4 import core
    from psi4.driver import constants
    from psi4.driver.p4util.exceptions import *

    def least_squares_fit_polynomial(xvals, fvals, localization_point, no_factorials=True, weighted=True, polynomial_order=4):
        """Performs and unweighted least squares fit of a polynomial, with specified order
           to an array of input function values (fvals) evaluated at given locations (xvals).
           See http://dx.doi.org/10.1063/1.4862157, particularly eqn (7) for details. """
        xpts = np.array(xvals) - localization_point
        if weighted:
            R = 1.0
            p_nu = 1
            epsilon = 1e-3
            zvals = np.square(xpts/R)
            weights = np.exp(-zvals) / (zvals**p_nu + epsilon**p_nu)
        else:
            weights = None
        fit = np.polynomial.polynomial.polyfit(xpts, fvals, polynomial_order, w=weights)
        # Remove the 1/n! coefficients
        if no_factorials:
            scalefac = 1.0
            for n in range(2,polynomial_order+1):
                scalefac *= n
                fit[n] *= scalefac
        return fit

    def anharmonicity(rvals, energies, plot_fit='', mol = None):
        """Generates spectroscopic constants for a diatomic molecules.
           Fits a diatomic potential energy curve using a weighted least squares approach
           (c.f. http://dx.doi.org/10.1063/1.4862157, particularly eqn. 7), locates the minimum
           energy point, and then applies second order vibrational perturbation theory to obtain spectroscopic
           constants.  Any number of points greater than 4 may be provided, and they should bracket the minimum.
           The data need not be evenly spaced, and can be provided in any order.  The data are weighted such that
           those closest to the minimum have highest impact.

           A dictionary with the following keys, which correspond to spectroscopic constants, is returned:

           :type rvals: list
           :param rvals: The bond lengths (in Angstrom) for which energies are
               provided, of length at least 5 and equal to the length of the energies array

           :type energies: list
           :param energies: The energies (Eh) computed at the bond lengths in the rvals list

           :type plot_fit: string
           :param plot_fit: A string describing where to save a plot of the harmonic and anharmonic fits, the
               inputted data points, re, r0 and the first few energy levels, if matplotlib
               is available.  Set to 'screen' to generate an interactive plot on the screen instead. If a filename is
               provided, the image type is determined by the extension; see matplotlib for supported file types.

           :returns: (*dict*) Keys: "re", "r0", "we", "wexe", "nu", "ZPVE(harmonic)", "ZPVE(anharmonic)", "Be", "B0", "ae", "De"
                     corresponding to the spectroscopic constants in cm-1
        """

        angstrom_to_bohr = 1.0 / constants.bohr2angstroms
        angstrom_to_meter = 10e-10;

        # Make sure the input is valid
        if len(rvals) != len(energies):
            raise ValidationError("The number of energies must match the number of distances")
        npoints = len(rvals)
        if npoints < 5:
            raise ValidationError("At least 5 data points must be provided to compute anharmonicity")
        core.print_out("\n\nPerforming a fit to %d data points\n" % npoints)

        # Make sure the molecule the user provided is the active one
        molecule = mol if mol is not None else core.get_active_molecule()
        molecule.update_geometry()
        natoms = molecule.natom()
        if natoms != 2:
            raise Exception("The current molecule must be a diatomic for this code to work!")
        m1 = molecule.mass(0)
        m2 = molecule.mass(1)

        # Optimize the geometry, refitting the surface around each new geometry
        core.print_out("\nOptimizing geometry based on current surface:\n\n");
        re = np.mean(rvals)
        maxit = 30
        thres = 1.0e-9
        for i in range(maxit):
            derivs = least_squares_fit_polynomial(rvals,energies,localization_point=re)
            e,g,H = derivs[0:3]
            core.print_out("       E = %20.14f, x = %14.7f, grad = %20.14f\n" % (e, re, g))
            if abs(g) < thres:
                break
            re -= g/H;
            if i == maxit-1:
                raise ConvergenceError("diatomic geometry optimization", maxit)
        core.print_out(" Final E = %20.14f, x = %14.7f, grad = %20.14f\n" % (e, re, g));
        if re < min(rvals):
            raise Exception("Minimum energy point is outside range of points provided.  Use a lower range of r values.")
        if re > max(rvals):
            raise Exception("Minimum energy point is outside range of points provided.  Use a higher range of r values.")

        # Convert to convenient units, and compute spectroscopic constants
        d0,d1,d2,d3,d4 = derivs*constants.hartree2aJ
        core.print_out("\nEquilibrium Energy %20.14f Hartrees\n" % e)
        core.print_out("Gradient           %20.14f\n" % g)
        core.print_out("Quadratic Force Constant %14.7f MDYNE/A\n" % d2)
        core.print_out("Cubic Force Constant     %14.7f MDYNE/A**2\n" % d3)
        core.print_out("Quartic Force Constant   %14.7f MDYNE/A**3\n" % d4)

        hbar = constants.h / (2.0 * np.pi)
        mu = ((m1*m2)/(m1+m2))*constants.amu2kg
        we = 5.3088375e-11 * np.sqrt(d2/mu)
        wexe = (1.2415491e-6)*(we/d2)**2 * ((5.0*d3*d3)/(3.0*d2)-d4)

        # Rotational constant: Be
        I = ((m1*m2)/(m1+m2)) * constants.amu2kg * (re * angstrom_to_meter)**2
        B = constants.h / (8.0 * np.pi**2 * constants.c * I)

        # alpha_e and quartic centrifugal distortion constant
        ae = -(6.0 * B**2 / we) * ((1.05052209e-3*we*d3)/(np.sqrt(B * d2**3))+1.0)
        de = 4.0*B**3 / we**2

        # B0 and r0 (plus re check using Be)
        B0 = B - ae / 2.0
        r0 = np.sqrt(constants.h / (8.0 * np.pi**2 * mu * constants.c * B0))
        recheck = np.sqrt(constants.h / (8.0 * np.pi**2 * mu * constants.c * B))
        r0 /= angstrom_to_meter;
        recheck /= angstrom_to_meter;

        # Fundamental frequency nu
        nu = we - 2.0 * wexe;
        zpve_nu = 0.5 * we - 0.25 * wexe;

        # Generate pretty pictures, if requested
        if(plot_fit):
            try:
                import matplotlib.pyplot as plt
            except ImportError:
                msg = "\n\tPlot not generated; matplotlib is not installed on this machine.\n\n"
                print(msg)
                core.print_out(msg)

            # Correct the derivatives for the missing factorial prefactors
            dvals = np.zeros(5)
            dvals[0:5] = derivs[0:5]
            dvals[2] /= 2
            dvals[3] /= 6
            dvals[4] /= 24

            # Default plot range, before considering energy levels
            minE = np.min(energies)
            maxE = np.max(energies)
            minR = np.min(rvals)
            maxR = np.max(rvals)

            # Plot vibrational energy levels
            we_au = we / constants.hartree2wavenumbers
            wexe_au = wexe / constants.hartree2wavenumbers
            coefs2 = [ dvals[2], dvals[1], dvals[0] ]
            coefs4 = [ dvals[4], dvals[3], dvals[2], dvals[1], dvals[0] ]
            for n in range(3):
                Eharm = we_au*(n+0.5)
                Evpt2 = Eharm - wexe_au*(n+0.5)**2
                coefs2[-1] = -Eharm
                coefs4[-1] = -Evpt2
                roots2 = np.roots(coefs2)
                roots4 = np.roots(coefs4)
                xvals2 = roots2 + re
                xvals4 = np.choose(np.where(np.isreal(roots4)), roots4)[0].real + re
                Eharm += dvals[0]
                Evpt2 += dvals[0]
                plt.plot(xvals2, [Eharm, Eharm], 'b', linewidth=1)
                plt.plot(xvals4, [Evpt2, Evpt2], 'g', linewidth=1)
                maxE = Eharm
                maxR = np.max([xvals2,xvals4])
                minR = np.min([xvals2,xvals4])

            # Find ranges for the plot
            dE = maxE - minE
            minE -= 0.2*dE
            maxE += 0.4*dE
            dR = maxR - minR
            minR -= 0.2*dR
            maxR += 0.2*dR

            # Generate the fitted PES
            xpts = np.linspace(minR, maxR, 1000)
            xrel = xpts - re
            xpows = xrel[:, None] ** range(5)
            fit2 = np.einsum('xd,d', xpows[:,0:3], dvals[0:3])
            fit4 = np.einsum('xd,d', xpows, dvals)

            # Make / display the plot
            plt.plot(xpts, fit2, 'b', linewidth=2.5, label='Harmonic (quadratic) fit')
            plt.plot(xpts, fit4, 'g', linewidth=2.5, label='Anharmonic (quartic) fit')
            plt.plot([re, re], [minE, maxE], 'b--', linewidth=0.5)
            plt.plot([r0, r0], [minE, maxE], 'g--', linewidth=0.5)
            plt.scatter(rvals, energies, c='Black', linewidth=3, label='Input Data')
            plt.legend()

            plt.xlabel('Bond length (Angstroms)')
            plt.ylabel('Energy (Eh)')
            plt.xlim(minR, maxR)
            plt.ylim(minE, maxE)
            if plot_fit == 'screen':
                plt.show()
            else:
                plt.savefig(plot_fit)
                core.print_out("\n\tPES fit saved to %s.\n\n" % plot_fit)

        core.print_out("\nre     = %10.6f A  check: %10.6f\n" % (re, recheck))
        core.print_out("r0       = %10.6f A\n" % r0)
        core.print_out("we       = %10.4f cm-1\n" % we)
        core.print_out("wexe     = %10.4f cm-1\n" % wexe)
        core.print_out("nu       = %10.4f cm-1\n" % nu)
        core.print_out("ZPVE(nu) = %10.4f cm-1\n" % zpve_nu)
        core.print_out("Be       = %10.4f cm-1\n" % B)
        core.print_out("B0       = %10.4f cm-1\n" % B0)
        core.print_out("ae       = %10.4f cm-1\n" % ae)
        core.print_out("De       = %10.7f cm-1\n" % de)
        results = {
                   "re"               :  re,
                   "r0"               :  r0,
                   "we"               :  we,
                   "wexe"             :  wexe,
                   "nu"               :  nu,
                   "ZPVE(harmonic)"   :  zpve_nu,
                   "ZPVE(anharmonic)" :  zpve_nu,
                   "Be"               :  B,
                   "B0"               :  B0,
                   "ae"               :  ae,
                   "De"               :  de
                  }
        return results