# flake8: noqa

"""Infrared and Raman intensities using siesta and MBPT_LCAO"""

import pickle
from math import sqrt
from sys import stdout
import numpy as np
import ase.units as units
from ase.parallel import parprint, paropen
from ase.vibrations import Vibrations
import warnings

# XXX This class contains much repeated code.  FIXME


class SiestaRaman(Vibrations):
    """
    Class for calculating vibrational modes, infrared and
    non-resonant Raman intensities using finite difference.

    The vibrational modes are calculated from a finite difference
    approximation of the Dynamical matrix and the IR and Ram intensities from
    a finite difference approximation of the gradient of the dipole
    moment. The method is described in:

      D. Porezag, M. R. Pederson:
      "Infrared intensities and Raman-scattering activities within
      density-functional theory",
      Phys. Rev. B 54, 7830 (1996)

    The calculator object (calc) must be Siesta, and the
    pyscf program (nao branch: https://github.com/cfm-mpc/pyscf/tree/nao)
    must be installed.

    >>> calc.get_dipole_moment(atoms)

    In addition to the methods included in the ``Vibrations`` class
    the ``Raman`` as the ``Infrared`` class introduces two methods;
    *get_spectrum()* and *write_spectra()*. The *summary()*, *get_energies()*,
    *get_frequencies()*, *get_spectrum()* and *write_spectra()*
    methods all take an optional *method* keyword.  Use
    method='Frederiksen' to use the method described in:

      T. Frederiksen, M. Paulsson, M. Brandbyge, A. P. Jauho:
      "Inelastic transport theory from first-principles: methodology
      and applications for nanoscale devices",
      Phys. Rev. B 75, 205413 (2007)

    atoms: Atoms object
        The atoms to work on.
    siesta: Siesta calculator
    mbpt_inp: dict
        dictionary containing the input for the mbpt_lcao program
    indices: list of int
        List of indices of atoms to vibrate.  Default behavior is
        to vibrate all atoms.
    name: str
        Name to use for files.
    delta: float
        Magnitude of displacements.
    nfree: int
        Number of displacements per degree of freedom, 2 or 4 are
        supported. Default is 2 which will displace each atom +delta
        and -delta in each cartesian direction.
    directions: list of int
        Cartesian coordinates to calculate the gradient
        of the dipole moment in.
        For example directions = 2 only dipole moment in the z-direction will
        be considered, whereas for directions = [0, 1] only the dipole
        moment in the xy-plane will be considered. Default behavior is to
        use the dipole moment in all directions.

    freq_pol: float or array of float
        frequency at which the Raman intensity is computed, can be float or array

    Example:

    See the example test/siesta/mbpt_lcao/script_raman.py
    This example calculate the Raman signal for a CO2 molecule.
    You should get something like,

    ---------------------------------------------------------------------------------------------------------------------------
    Mode    Frequency        Intensity IR        Intensity Raman (real)          Intensity Raman (imag)          Raman Ehanced
    #    meV     cm^-1   (D/Å)^2 amu^-1          A^4 amu^-1                     A^4 amu^-1                      A^4 amu^-1
    ---------------------------------------------------------------------------------------------------------------------------
    0   23.2i    187.1i     0.0005                1.0810                 0.0000                 0.0000
    1   22.7i    183.2i     0.0007                1.1471                 0.0000                 0.0000
    2    4.0i     32.6i     0.0001                0.0054                 0.0000                 0.0000
    3   19.6     158.0      0.0001                1.1790                 0.0000                 0.0000
    4   21.0     169.2      0.0004                1.0894                 0.0000                 0.0000
    5   77.9     628.3      0.4257                0.0060                 0.0000                 0.0000
    6   79.7     642.6      0.4354                0.0022                 0.0000                 0.0000
    7  163.5    1319.0      0.0000               21.7631                 0.0000                 0.0000
    8  294.0    2371.0     12.1479                0.0002                 0.0000                 0.0000
    ---------------------------------------------------------------------------------------------------------------------------

    It can be compared to calculations done with Quantum Espresso (see  test/siesta/mbpt_lcao/raman_espresso)
    that give something like,

    # mode   [cm-1]    [THz]      IR          Raman   depol.fact
    1     -0.01   -0.0002    0.0000         0.4930    0.7500
    2     -0.00   -0.0000    0.0000         0.0018    0.7500
    3      0.00    0.0001    0.0000         0.8202    0.7500
    4      0.00    0.0001    0.0000         0.9076    0.7500
    5      0.01    0.0002    0.0000         1.8576    0.7499
    6      0.07    0.0021    0.0000         0.0001    0.7500
    7    717.64   21.5144    0.5303         0.0000    0.0862
    8   1244.37   37.3052    0.0000        23.8219    0.1038
    9   2206.78   66.1575   12.6139         0.0000    0.6417
    """

    def __init__(self, atoms, siesta, indices=None, name='ram',
                 delta=0.01, nfree=2, directions=None, freq_pol=0.0, **kw):

        Vibrations.__init__(self, atoms, indices=indices, name=name,
                            delta = delta, nfree=nfree)
        if atoms.constraints:
            warnings.warn('WARNING! \n Your Atoms object is constrained. ' +
                  'Some forces may be unintended set to zero. \n')
        self.name = name + '-d%.3f' % delta
        self.calc = atoms.calc

        if directions is None:
            self.directions = np.asarray([0, 1, 2])
        else:
            self.directions = np.asarray(directions)
        self.ir = True
        self.ram = True
        self.siesta = siesta

        if isinstance(freq_pol, list):
            self.freq_pol = np.array(freq_pol)
        elif isinstance(freq_pol, float):
            self.freq_pol = np.array([freq_pol])
        elif isinstance(freq_pol, float) or isinstance(freq_pol, np.ndarray):
            self.freq_pol = freq_pol
        else:
            raise ValueError("wrong type for freq_pol, only float, list or array")

        self.pyscf_arg = kw

    def get_polarizability(self):
        if "tddft_iter_tol" in list(self.pyscf_arg.keys()):
            if self.pyscf_arg["tddft_iter_tol"] > 1e-4:
                warnings.warn("tddft_iter_tol > 1e-4, polarizability may not have " +
                              "enough precision. The Raman intensity will not be precise.")
        else:
            self.pyscf_arg["tddft_iter_tol"] = 1e-4

        self.siesta.lrtddft(Edir=np.array([1.0, 1.0, 1.0]), **self.pyscf_arg)
        return self.siesta.results["freq range"], \
               self.siesta.results["polarizability nonin"], \
               self.siesta.results["polarizability inter"]

    def read(self, method='standard', direction='central', inter = True):
        self.method = method.lower()
        self.direction = direction.lower()
        assert self.method in ['standard', 'frederiksen']
        if direction != 'central':
            raise NotImplementedError(
                'Only central difference is implemented at the moment.')

        # Get "static" dipole moment polarizability and forces
        name = '%s.eq.pckl' % self.name
        [forces_zero, dipole_zero, freq_zero,
            noninPol_zero, pol_zero] = pickle.load(open(name, "rb"))
        self.dipole_zero = (sum(dipole_zero**2)**0.5) / units.Debye
        self.force_zero = max([sum((forces_zero[j])**2)**0.5
                               for j in self.indices])
        self.noninPol_zero = noninPol_zero * (units.Bohr)**3  # Ang**3
        self.pol_zero = pol_zero * (units.Bohr)**3  # Ang**3

        ndof = 3 * len(self.indices)
        H = np.empty((ndof, ndof))
        dpdx = np.empty((ndof, 3))
        dadx = np.empty((ndof, self.pol_zero.shape[0], 3, 3), dtype=complex)

        r = 0
        for a in self.indices:
            for i in 'xyz':
                name = '%s.%d%s' % (self.name, a, i)
                [fminus, dminus, frminus, noninpminus, pminus] = pickle.load(
                    open(name + '-.pckl', "rb"))
                [fplus, dplus, frplus, noninpplus, pplus] = pickle.load(
                    open(name + '+.pckl', "rb"))
                if self.nfree == 4:
                    [fminusminus, dminusminus, frminusminus,
                            noninpminusminus, pminusminus] =\
                                    pickle.load(open(name + '--.pckl', "rb"))
                    [fplusplus, dplusplus, frplusplus,
                            noninpplusplus, pplusplus] =\
                                    pickle.load(open(name + '++.pckl', "rb"))
                if self.method == 'frederiksen':
                    fminus[a] += -fminus.sum(0)
                    fplus[a] += -fplus.sum(0)
                    if self.nfree == 4:
                        fminusminus[a] += -fminus.sum(0)
                        fplusplus[a] += -fplus.sum(0)
                if self.nfree == 2:
                    H[r] = (fminus - fplus)[self.indices].ravel() / 2.0
                    dpdx[r] = (dminus - dplus)
                    if inter:
                        dadx[r] = (pminus - pplus)
                    else:
                        dadx[r] = (noninpminus - noninpplus)
                if self.nfree == 4:
                    H[r] = (-fminusminus + 8 * fminus - 8 * fplus +
                            fplusplus)[self.indices].ravel() / 12.0
                    dpdx[r] = (-dplusplus + 8 * dplus - 8 * dminus +
                               dminusminus) / 6.0
                    if inter:
                        dadx[r] = (-pplusplus + 8 * pplus - 8 * pminus +
                                   pminusminus) / 6.0
                    else:
                        dadx[r] = (-noninpplusplus + 8 * noninpplus - 8 * noninpminus +
                                   noninpminusminus) / 6.0
                H[r] /= 2 * self.delta
                dpdx[r] /= 2 * self.delta
                dadx[r] /= 2 * self.delta  # polarizability in Ang
                for n in range(3):
                    if n not in self.directions:
                        dpdx[r][n] = 0
                        dpdx[r][n] = 0
                r += 1
        # Calculate eigenfrequencies and eigenvectors
        m = self.atoms.get_masses()
        H += H.copy().T
        self.H = H
        m = self.atoms.get_masses()
        self.im = np.repeat(m[self.indices]**-0.5, 3)
        omega2, modes = np.linalg.eigh(self.im[:, None] * H * self.im)
        self.modes = modes.T.copy()

        # infrared
        # Calculate intensities
        dpdq = np.array([dpdx[j] / sqrt(m[self.indices[j // 3]] *
                                        units._amu / units._me)
                         for j in range(ndof)])
        dpdQ = np.dot(dpdq.T, modes)
        dpdQ = dpdQ.T
        intensities = np.array([sum(dpdQ[j]**2) for j in range(ndof)])
        # Conversion factor:
        s = units._hbar * 1e10 / sqrt(units._e * units._amu)
        self.hnu = s * omega2.astype(complex)**0.5
        # Conversion factor from atomic units to (D/Angstrom)^2/amu.
        conv = (1.0 / units.Debye)**2 * units._amu / units._me
        self.intensities_ir = intensities * conv

        # Raman
        dadq = np.array([(dadx[j, :, :, :] / (units.Bohr**2)) /
                         sqrt(m[self.indices[j // 3]] * units._amu / units._me)
                         for j in range(ndof)])
        dadQ = np.zeros((ndof, 3, 3, dadq.shape[1]), dtype=complex)
        for w in range(dadq.shape[1]):
            dadQ[:, :, :, w] = np.dot(dadq[:, w, :, :].T, modes).T

        ak = (dadQ[:, 0, 0, :] + dadQ[:, 1, 1, :] + dadQ[:, 2, 2, :]) / 3.0
        gk2 = ((dadQ[:, 0, 0, :] - dadQ[:, 1, 1, :])**2 + (dadQ[:, 1, 1, :] -
                dadQ[:, 2, 2, :])**2 + (dadQ[:, 2, 2, :] -
                dadQ[:, 0, 0, :])**2 + 6 * (dadQ[:, 0, 1, :]**2 +
                dadQ[:, 1, 2, :]**2 + dadQ[:, 2, 0, :]**2))

        intensities = np.zeros((ndof, self.freq_pol.size), dtype=np.complex128)
        intensities_ram_enh = np.zeros((ndof, self.freq_pol.size), dtype=np.complex128)

        # calculate the coefficients for calculating Raman Signal
        for j in range(ndof):
            aj = self.get_nearrest_value(self.freq_pol, freq_zero, ak[j, :])
            gj2 = self.get_nearrest_value(self.freq_pol, freq_zero, gk2[j, :])

            intensities[j, :] = (45 * aj**2 + 7 * gj2) / 45.0

        self.intensities_ram = intensities  # Bohr**4 .me**-1
        self.intensities_ram_enh = intensities_ram_enh  # Bohr**4 .me**-1

    def get_nearrest_value(self, val, x_range, y, kind="cubic", prec = 1e-3):
        """
            return the closest value of y at val (interpolating the function)
            val may be a number or an array
        """
        import scipy.interpolate as interp


        func = interp.interp1d(x_range, y, kind = kind)
        new_range = np.arange(x_range[0], x_range[x_range.size-1], prec)
        interpol = func(new_range)

        if isinstance(val, np.ndarray):
            mult = np.zeros(val.shape, dtype=interpol.dtype)

            for i, va in enumerate(val):
                idx = (np.abs(new_range-va)).argmin()
                mult[i] = interpol[idx]
            return mult
        else:
            idx = (np.abs(new_range-val)).argmin()
            return np.array([interpol[idx]])

    def intensity_prefactor(self, intensity_unit):
        if intensity_unit == '(D/A)2/amu':
            return 1.0, '(D/Å)^2 amu^-1'
        elif intensity_unit == 'km/mol':
            # conversion factor from Porezag PRB 54 (1996) 7830
            return 42.255, 'km/mol'
        elif intensity_unit == 'au':
            return 1.0, '   a.u.'
        elif intensity_unit == 'A^4 amu^-1':
            # Quantum espresso units
            return (units.Bohr**4) / (units._me / units._amu), '   A^4 amu^-1'
        else:
            raise RuntimeError('Intensity unit >' + intensity_unit +
                               '< unknown.')

    def summary(self, method='standard', direction='central', freq_pol = 0.0,
                intensity_unit_ir='(D/A)2/amu', intensity_unit_ram='au', log=stdout,
                inter = True):
        hnu = self.get_energies(method, direction, inter=inter)
        s = 0.01 * units._e / units._c / units._hplanck
        iu_ir, iu_string_ir = self.intensity_prefactor(intensity_unit_ir)
        iu_ram, iu_string_ram = self.intensity_prefactor(intensity_unit_ram)
        arr = []

        freq_idx = (np.abs(self.freq_pol-freq_pol)).argmin()
        print("index: ", freq_idx)

        if intensity_unit_ir == '(D/A)2/amu':
            iu_format_ir = '%9.4f             '
        elif intensity_unit_ir == 'km/mol':
            iu_string_ir = '   ' + iu_string_ir
            iu_format_ir = ' %7.1f            '
        elif intensity_unit_ir == 'au':
            iu_format_ir = '%.6e              '
        elif intensity_unit_ir == 'A^4 amu^-1':
            iu_format_ir = '%9.4f             '

        if intensity_unit_ram == '(D/A)2/amu':
            iu_format_ram = '%9.4f'
        elif intensity_unit_ram == 'km/mol':
            iu_string_ram = '   ' + iu_string_ram
            iu_format_ram = ' %7.1f'
        elif intensity_unit_ram == 'au':
            iu_format_ram = '%.6e              '
        elif intensity_unit_ram == 'A^4 amu^-1':
            iu_format_ram = '%9.4f              '

        if isinstance(log, str):
            log = paropen(log, 'a')

        parprint('---------------------------------------------------------------------------------------------------------------------------', file=log)
        parprint(' Mode    Frequency        Intensity IR        Intensity Raman (real)          Intensity Raman (imag)          Raman Ehanced', file=log)
        parprint('  #    meV     cm^-1   ' + iu_string_ir + '       ' + iu_string_ram +
                 '                  ' + iu_string_ram + '                   ' + iu_string_ram, file=log)
        parprint('---------------------------------------------------------------------------------------------------------------------------', file=log)
        for n, e in enumerate(hnu):
            if e.imag != 0:
                c = 'i'
                e = e.imag
            else:
                c = ' '
                e = e.real
                arr.append([n, 1000 * e, s * e, iu_ir * self.intensities_ir[n],
                            iu_ram * self.intensities_ram[n, freq_idx].real, iu_ram *
                            self.intensities_ram[n, freq_idx].imag])
            parprint(('%3d %6.1f%s  %7.1f%s  ' + iu_format_ir + iu_format_ram +
                      iu_format_ram + iu_format_ram) %
                     (n, 1000 * e, c, s * e, c, iu_ir * self.intensities_ir[n],
                      iu_ram * self.intensities_ram[n, freq_idx].real, iu_ram *
                      self.intensities_ram[n, freq_idx].imag, iu_ram *
                      self.intensities_ram_enh[n, freq_idx].real), file=log)
        parprint(
            '-----------------------------------------------------------------------------------------',
            file=log)
        parprint('Zero-point energy: %.3f eV' % self.get_zero_point_energy(),
                 file=log)
        parprint('Static dipole moment: %.3f D' % self.dipole_zero, file=log)
        parprint('Maximum force on atom in `equilibrium`: %.4f eV/Å' %
                 self.force_zero, file=log)
        parprint(file=log)
        np.savetxt('ram-summary.txt', np.array(arr))

    def write_latex_array(self, fname='vib_latex_array.tex.table', caption='', nb_column=5,
                          hline=False, method='standard', direction='central',
                          intensity_unit_ir='(D/A)2/amu', intensity_unit_ram='au',
                          label='tab_vib', log=stdout, freq_pol=0.0):
        """
        Write the summary into a latex table that can be easily incorporate into a latex file.
        """

        hnu = self.get_energies(method, direction)
        s = 0.01 * units._e / units._c / units._hplanck
        iu_ir, iu_string_ir = self.intensity_prefactor(intensity_unit_ir)
        iu_ram, iu_string_ram = self.intensity_prefactor(intensity_unit_ram)

        freq_idx = (np.abs(self.freq_pol-freq_pol)).argmin()


        if intensity_unit_ir == '(D/A)2/amu':
            iu_format_ir = '%9.4f             '
        elif intensity_unit_ir == 'km/mol':
            iu_string_ir = '   ' + iu_string_ir
            iu_format_ir = ' %7.1f            '
        elif intensity_unit_ir == 'au':
            iu_format_ir = '%.6e              '
        elif intensity_unit_ir == 'A^4 amu^-1':
            iu_format_ir = '%9.4f             '

        if intensity_unit_ram == '(D/A)2/amu':
            iu_format_ram = '%9.4f'
        elif intensity_unit_ram == 'km/mol':
            iu_string_ram = '   ' + iu_string_ram
            iu_format_ram = ' %7.1f'
        elif intensity_unit_ram == 'au':
            iu_format_ram = '%.6e              '
        elif intensity_unit_ram == 'A^4 amu^-1':
            iu_format_ram = '%9.4f              '

        if isinstance(log, str):
            log = paropen(log, 'a')

        if hline:
            column = "|"
        else:
            column = ""

        for i in range(nb_column + 1):
            if hline:
                column = column + "c|"
            else:
                column = column + "c"

        f = open(fname, 'w')
        f.write("\\begin{table}[h] \n")
        f.write("  \\caption{" + caption + "} \n")
        f.write("  \\begin{center}\n")
        f.write("    \\begin{tabular}{" + column + "} \n")

        if hline:
            f.write("     \\hline \n")

        f.write('     Mode & Frequency (meV) & Frequency ($cm^{-1}$) & Intensity IR  (' +
                iu_string_ir + ')  & Intensity Raman (' + iu_string_ram + ') \n')
        if hline:
            f.write("   \\hline \n")
        for n, e in enumerate(hnu):
            if e.imag != 0:
                c = ' + i'
                e = e.imag
            else:
                c = ' '
                e = e.real
            f.write(('     %3d & %6.1f %s  & %7.1f %s  & ' + iu_format_ir +
                     '  & ' + iu_format_ram + ' \n') % (n, 1000 * e, c, s * e, c,
                                                        iu_ir * self.intensities_ir[n], iu_ram *
                                                        self.intensities_ram[n, freq_idx].real))
            if hline:
                f.write(r"      \hline \n")
        f.write("    \\end{tabular} \n")
        f.write("  \\end{center} \n")
        f.write(" \\label{" + label + "} \n")
        f.write("\\end{table}\n")

        f.close()

    def get_spectrum(self, start=800, end=4000, npts=None, width=4,
                     type='Gaussian', method='standard', direction='central',
                     intensity_unit='(D/A)2/amu', normalize=False, freq_pol=0.0):
        """Get raman spectrum.

        The method returns wavenumbers in cm^-1 with corresponding
        absolute infrared intensity.
        Start and end point, and width of the Gaussian/Lorentzian should
        be given in cm^-1.
        normalize=True ensures the integral over the peaks to give the
        intensity.
        """
        name = '%s.eq.pckl' % self.name
        [forces_zero, dipole_zero, freq_zero, noninPol_zero,
            pol_zero] = pickle.load(open(name, "rb"))

        freq_idx = (np.abs(self.freq_pol-freq_pol)).argmin()

        frequencies = self.get_frequencies(method, direction).real
        intensities_ir = self.intensities_ir
        intensities_ram = self.intensities_ram[:, freq_idx]
        intensities_ram_enh = self.intensities_ram_enh[:, freq_idx]
        energies, spectrum_ir = self.fold(frequencies, intensities_ir,
                                          start, end, npts, width, type,
                                          normalize)
        energies, spectrum_ram = self.fold(frequencies, intensities_ram.real,
                                           start, end, npts, width, type,
                                           normalize)
        energies, spectrum_ram_enh = self.fold(frequencies, intensities_ram_enh.real,
                                               start, end, npts, width, type,
                                               normalize)
        return energies, spectrum_ir, spectrum_ram, spectrum_ram_enh

    def write_spectra(self, out='ram-spectra.dat', start=800, end=4000,
                      npts=None, width=10, type='Gaussian',
                      method='standard', direction='central',
                      intensity_unit_ir='(D/A)2/amu',
                      intensity_unit_ram='au', normalize=False):
        """Write out raman spectrum to file.
        First column is the wavenumber in cm^-1, the second column the
        absolute infrared intensities, and
        the third column the absorbance scaled so that data runs
        from 1 to 0.  idem for the Rahman spectrum for columns 4 and 5. Start and end
        point, and width of the Gaussian/Lorentzian should be given
        in cm^-1."""
        energies, spectrum_ir, spectrum_ram, spectrum_ram_enh =\
            self.get_spectrum(start, end, npts, width, type, method,
                              direction, normalize)

        # Write out spectrum in file. First column is absolute intensities.
        # Second column is absorbance scaled so that data runs from 1 to 0
        spectrum2_ir = 1. - spectrum_ir / spectrum_ir.max()
        spectrum2_ram = 1. - spectrum_ram / spectrum_ram.max()
        if abs(spectrum_ram_enh.max()) > 0.0:
            spectrum2_ram_enh = 1. - spectrum_ram_enh / spectrum_ram_enh.max()
        else:
            spectrum2_ram_enh = np.zeros(spectrum_ram.shape, dtype = np.float64)

        outdata = np.empty([len(energies), 7])
        outdata.T[0] = energies
        outdata.T[1] = spectrum_ir
        outdata.T[2] = spectrum2_ir
        outdata.T[3] = spectrum_ram
        outdata.T[4] = spectrum2_ram
        outdata.T[5] = spectrum_ram_enh
        outdata.T[6] = spectrum2_ram_enh

        fd = open(out, 'w')
        fd.write('# %s folded, width=%g cm^-1\n' % (type.title(), width))
        iu_ir, iu_string_ir = self.intensity_prefactor(intensity_unit_ir)
        iu_ram, iu_string_ram = self.intensity_prefactor(intensity_unit_ram)

        # if normalize:
        #    iu_string = 'cm ' + iu_string_ir + iu_string_ram
        fd.write('# [cm^-1] %14s\n' %
                 ('[' + iu_string_ir + iu_string_ram + iu_string_ram + ']'))
        for row in outdata:
            fd.write('%.3f  %15.5e  %15.5e %15.5e  %15.5e   %15.5e    %15.5e\n' %
                     (row[0], iu_ir * row[1], row[2], iu_ram * row[3], row[4],
                      iu_ram * row[5], row[6]))
        fd.close()