import numpy as np
from scipy import signal
# import yaml_phonons as ph
from src import phonon_lifetimes
import math
import scipy.constants as const
# import matplotlib.pyplot as plt
import h5py as h5
import spglib as spg
from phonopy import load, Phonopy
from phonopy.file_IO import parse_BORN
from phonopy.units import THz, AMU, Hartree, Bohr
from phonopy.interface.vasp import read_vasp
from phono3py.api_phono3py import Phono3py
from phono3py.file_IO import (parse_disp_fc3_yaml,
                              parse_FORCES_FC3)
from spglib import get_ir_reciprocal_mesh
from scipy.interpolate import interp1d
from src.utils import Gamma, Phono3pyInputs, BrillouinZone, ImaginarySelfEnergy, PhononEigenvalues
#from .utils import Gamma, Phono3pyInputs, BrillouinZone

AngstromsToMeters = 1e-10
# from numba import jit

"""
Goal of program is to compute S(q,w) with q-points (as vectors) as user-defined inputs.
Definition of S(q,w):
F_i(q,w) = \sum_j (q * u_{ij})^2 exp[-\sum_j (q * u_ij)^2] \delta(\hbar w - \hbar w_j)
S_i(q,w) = \sum_n conv^n/n! * F_i(q,w)
S(q,w) = \sum_i S_i(q,w)
which is perfectly defined for the "incoherent" approximation in neutron scattering.
Other words: this is the space and time Fourier transform of the self van-Hove correlation function
To include "coherent" effects (e.g. compute the full van-Hove correlation function):
f_i = \sum_j (q * u_{ij}) \delta(\hbar w - \hbar w_j)
DWF_i = exp[-\sum_j (q * u_ij)^2/2] 
F_ab(q,w) = conj(f_a) * f_b
S(q,w) = \sum_ab F_ab * DWF_a * DWF_b

What happens when you have two degenerate phonon modes (w_i == w_j)
    |(q * u_i) + (q * u_j)|^2 = |q*u_i|^2 + |q*u_j|^2 + conj(q*u_i) * (q*u_j) + c.c.
"""


class AnharmonicPhonons(object):
    """
    Class containing all information on anharmonicities/matrix elements for a specific material.
    """
    def __init__(self,
                 poscar,
                 fc3_file,
                 disp_file,
                 mesh,
                 supercell,
                 born=None,
                 temperature=0,
                 param_flag=False):
        """
        Initialized by Phono3py inputs.

        Parameters
        -----
        poscar
            POSCAR file of material in standard VASP format
        fc3_file
            File containing the third order force constants in one of the Phono3py standard formats
        disp_file
            File containing the information about the displacements generated by the `phono3py -d` command.
        mesh
            List containing the Monkhorst-Pack grid discretizations in each direction of the Brillouin zone. For example, a 5x5x5 MP grid is specified by [5, 5, 5]. Gamma centered is always assumed to be true
        supercell
            List containing the supercell dimensions with respect to the cell given in the POSCAR file.
        born
            BORN file built from `phonopy-vasp-born` auxiliary command. Optional because some materials do not have significant LO-TO splitting.
        temperature
            Temperature at which you want to calculate the anharmonic properties
        param_flag
            Boolean specifing whether or not the anharmonic broadening is parametrized by a Lorentzian function.
        """
        self.cell = read_vasp(poscar)
        self.mesh = mesh
        self.born = born
        self.phono3py = Phono3py(self.cell,
                                 supercell_matrix=np.diag(supercell),
                                 primitive_matrix='auto',
                                 mesh=mesh,
                                 log_level=1)
        self.disp_data = parse_disp_fc3_yaml(filename=disp_file)
        self.fc3_data = parse_FORCES_FC3(self.disp_data, filename=fc3_file)
        self.phono3py.produce_fc3(self.fc3_data,
                                  displacement_dataset=self.disp_data,
                                  symmetrize_fc3r=True)

        ## initialize phonon-phonon interaction instance
        self.phono3py.init_phph_interaction()
        self.mapping = None
        self.grid = None
        self.irr_BZ_gridpoints = None
        self.phonon_freqs = None
        self.temperature = temperature
        # self.phono3py.run_imag_self_energy(np.unique(self.mapping), temperatures=temperature)
        if self.born is not None:
            primitive = self.phono3py.get_phonon_primitive()
            nac_params = parse_BORN(primitive, filename=born)
            nac_params['factor'] = Hartree * Bohr
            self.phono3py.nac_params = nac_params
            nac_q_direction = [1, 0, 0]
            self.phono3py.init_phph_interaction(nac_q_direction=nac_q_direction)
        if self.born is None:
            nac = False
        else:
            nac = True
        self.inputs = Phono3pyInputs(poscar=poscar,
                                     fc3_file=fc3_file,
                                     disp_file=disp_file,
                                     mesh=mesh,
                                     supercell=supercell,
                                     nac=nac,
                                     born_file=born
                                     )
        if param_flag:
            self.gamma = Gamma(self.inputs)
        else:
            self.gamma = ImaginarySelfEnergy(self.inputs)
        self.param_flag = param_flag
        # set freqs, and prepare to remove the set_freqs() method
        self.phonon_freqs = PhononEigenvalues(self.inputs)

    def set_irr_BZ_gridpoints(self):
        """
        Setter of irreducible Brillouin zone q-points
        """
        self.mapping, grid = get_ir_reciprocal_mesh(mesh=self.mesh, cell=self.cell)
        self.grid = {tuple(k / self.mesh): v for (v, k) in enumerate(grid)}
        irr_BZ_gridpoints = np.unique(self.mapping)
        self.irr_BZ_gridpoints = {k: v for v, k in enumerate(irr_BZ_gridpoints)}

    def set_self_energies(self):
        """
        Setter of imaginary self energies. The imaginary self energy objects are stored in phono3py object in a weird way; the object is a list of a list of nparrays                                                            
        first index = grid-points, which are "addresses" of the irreducible q-points in the Brillouin zone. On the surface these are not obviously related to the q-points, but the grid points are given by:                         
              np.unique(mapping)                                                                              
              The actual q-points are stored in grid with a 1-to-1 correspondence to mapping                  
        second index = sigma values, there is only one sigma=None in default tetrahedron method               
        third index is the nparray, which is arranged as:                                                     
              (temperatures, frequency points, band index)                                                    
        """
        self.phono3py.init_phph_interaction()
        if self.mapping is None:
            self.set_irr_BZ_gridpoints()
        if type(self.temperature) is not list:
            self.temperature = [self.temperature]
        self.phono3py.run_imag_self_energy(np.unique(self.mapping), temperatures=self.temperature)

    def get_imag_self_energy(self, gridpoint, band_index):
        """
        Getter of imaginary self energies.
        Note that this returns the imaginary self energy function does not return the self energy evaluated at the phonon frequency. As a result, a tuple is returned with the frequency points, and self energy

        Parameters
        -----
        gridpoint
            Gridpoint index used in phonopy/phono3py API
        band_index
            Index of phonon band.
        """
        irr_gp = self.mapping[gridpoint]
        grid_index = self.irr_BZ_gridpoints[irr_gp]
        return self.phono3py._frequency_points, self.phono3py._gammas[grid_index][0][0, band_index, :]

    def get_imag_self_energy_interpolator(self, gridpoint, band_index):
        """
        Getter of interpolator object for the imaginary self energy function. The resulting object can be used as "y = ufunc(x)". The reason for creating an interpolator is because the frequency discretization is different in phono3py compared to my code below.

        Parameters
        -----
        gridpoint
            Gridpoint index used in phonopy/phono3py API
        band_index
            Index of phonon band.
        """
        freqs, ise = self.get_imag_self_energy(gridpoint, band_index)
        return interp1d(freqs, ise, kind='linear')

    def get_gridpoint(self, qpoint):
        """
        Convenience function used to translate a q-point to a grid-point used in phonopy API functions.

        Parameters
        -----
        qpoint
            List/array that specifies the q-point in reduced coordinates.
        """
        for i, q in enumerate(qpoint):
            if q > 0.5:
                shift = np.ceil(q).astype(np.int)
                qpoint[i] = q - shift
            elif q < -0.5:
                shift = np.floor(q).astype(np.int)
                qpoint[i] = q - shift
        # make sure qpoint is exactly a key
        key = np.round(np.array(qpoint) * np.array(self.mesh)).astype(int) / np.array(self.mesh)
        return self.grid[tuple(key)]

    def get_broadening_function(self, qpoint, band_index, max_freq=None):
        """
        Construct a broadening Lorentzian-type function with a frequency-dependent width. Since the frequency-dependent widths are still relatively small compared to the oscillator frequency, the broadening function will look like a "normal" Lorentzian near the resonance, but may appear non-Lorentzian far away from the resonance. One issue is that Lorentzians with widths smaller than the frequency spacing are not well defined which primarily means they are not normalized. Normalize the integral by adding in a delta function with the remaining weight 
        
        Parameters
        -----
        qpoint
            List/array specifying the q-point in reduced coordinates
        band_index
            Integer specifying the phonon band/mode index
        max_freq
            Optional value to specify the maximum frequency considered in the broadening function. Default is None, which means the code will automatically decide the maximum frequency based on the details in the Phono3py object.
        """
        
        if self.param_flag:
            gamma = self.gamma.get_property_value(qpoint, band_index)
        else:
            gamma = self.gamma.get_property_value(qpoint, band_index)
        freqs = self.gamma.freqs
        if freqs is None:
            gridpoint = self.get_gridpoint(qpoint)
            freqs, _ = self.get_imag_self_energy(gridpoint, band_index)
        # Old way:
        #phonon_freq = self.phonon_freqs[gridpoint, band_index]
        # New way
        phonon_freq = self.phonon_freqs.get_property_value(qpoint, band_index)

        broadening_func = gamma / (np.pi * ((phonon_freq - freqs) ** 2 + gamma ** 2))
        f_index_minus = np.floor(phonon_freq / (freqs[1] - freqs[0])).astype(int)
        if not self.param_flag:
            avg_gamma_at_freq = (gamma[f_index_minus] + gamma[f_index_minus + 1]) / 2
        else:
            avg_gamma_at_freq = gamma
        # print('original integral of broad func =', np.trapz(broadening_func, freqs))

        if max_freq is None:
            freqs = np.append(freqs, freqs[-1] * 1000)
        else:
            freqs = np.append(freqs, max_freq + 1.0)
        broadening_func = np.append(broadening_func, 0.0)

        if np.trapz(broadening_func, freqs) < 1.0:
            if avg_gamma_at_freq < freqs[1] - freqs[0]:
                # Case: the width is small, and the area is underestimated
                # Solution: add a delta function
                broadening_func += self.create_delta(phonon_freq,
                                                     len(freqs),
                                                     freqs[1] - freqs[0]) * \
                                   (1.0 - np.trapz(broadening_func, freqs))
            else:
                # Case: width is large enough, but the area is underestimated
                # Solution: scale the function by the trapz area
                broadening_func /= np.trapz(broadening_func, freqs)
        else:
            if avg_gamma_at_freq < freqs[1] - freqs[0]:
                # Case: width is too small to be resolved, area is overestimated by tails
                # Solution: subtract a delta function
                # Note: form is exactly the same as above because the (1-integral) term is now negative
                broadening_func += self.create_delta(phonon_freq,
                                                     len(freqs),
                                                     freqs[1] - freqs[0]) * \
                                   (1.0 - np.trapz(broadening_func, freqs))
            else:
                # Case: width large enough to be resolved, but area is overestimated
                # Solution: scale by integral
                broadening_func /= np.trapz(broadening_func, freqs)
        # Notes: - can actually remove the first if statement as the Solution is always the same
        #        - Also, the function is a linear interpolation because the later convolution operations are compatible
        #          are discrete, and thus only compatible with linear interpolations.
        #        - For example, if we used a 'cubic' interpolator, the norm of the interpolated function is not
        #          guaranteed to be 1.0, which creates large problems when expanding to higher orders of perturbation.

        return interp1d(freqs, broadening_func, kind='linear')

    def set_phonon_freqs(self):
        """
        Setter for phonon frequencies
        """
        if self.phonon_freqs is None:
            self.phonon_freqs, _, _ = self.phono3py.get_phonon_data()

    def get_phonon_freqs(self):
        """
        Getter for phonon frequencies
        """
        if self.phonon_freqs is None:
            self.set_phonon_freqs()
        return self.phonon_freqs

    def create_delta(self, energy, num_bins, delta_e):
        """
        Create a delta function at a specified energy with a discretized space defined by `num_bins` and `delta_e`.
    
        Parameters
        -----
        energy
            Energy at which the delta function exists. As the energy value can be somewhere between the discretized points, we represent the delta function as a trapezoid.
        num_bins
            Number of bins in the discrete frequency space
        delta_e
            Spacing between the bins in the discrete frequency space

        Returns
        -----
        delta_fcn
            A spectrum representing a delta function with a shape of (num_bins,).
        """
        # units in meV
        delta_fcn = np.zeros(num_bins)

        if energy < 0:
            return delta_fcn

        e_bin_minus = int(np.floor(energy / delta_e))
        e_bin_plus = int(np.ceil(energy / delta_e))

        alpha_minus = np.abs(e_bin_minus * delta_e - energy) / delta_e
        alpha_plus = np.abs(e_bin_plus * delta_e - energy) / delta_e

        delta_fcn[e_bin_minus] = (1 - alpha_minus) / delta_e
        delta_fcn[e_bin_plus] = (1 - alpha_plus) / delta_e
        return delta_fcn

    def set_gamma_interpolator(self):
        """
        Setter of a Gamma interpolator. This function sets the interpolator of the Gamma class using the Interpolator class.
        """
        from src.Interpolation import Interpolator

        gamma_data, qpts, freqs = self.get_gamma_data()
        interp = Interpolator(gamma=gamma_data, qpoints=qpts, freqs=freqs)

    def get_gamma_data(self, band_index):
        """
        Get Gamma data from phono3py API.
    
        Parameters
        -----
        band_index
            Integer specifying the band index.
        """
        gamma_data = []
        q_pts = []
        for q in self.phono3py.get_phonon_data()[2]:
            q /= self.phono3py.mesh_numbers
            freqs, gamma = self.get_imag_self_energy(self.get_gridpoint(q), band_index)
            # gamma_data += list(gamma)
            # q_data += [list(q) + [freq] for freq in freqs]
            q_pts.append(q)
            gamma_data.append(gamma)
        return gamma_data, q_pts, freqs


class DynamicStructureFactor(object):
    '''
    Class that calculates the dynamic structure factor from a set of Phono3py/Phonopy inputs with the option to print the output to HDF5 files. The inputs can also augment the dynamic structure factor to calculate the differential scattering cross section for a variety of dark matter models and neutron scattering.
    '''
    def __init__(self,
                 poscar_file,
                 mesh,
                 supercell,
                 fc_file=None,
                 q_point_list=[],
                 q_point_shift=[0.0, 0.0, 0.0],
                 meshG=[1, 1, 1],
                 weights=None,
                 fc2_disp=None,
                 fc3_file=None,
                 fc3_disp=None,
                 delta_e=0.01,
                 max_e=30,
                 num_overtones=10,
                 temperature=4,
                 freq_min=1e-3,
                 primitive_flag='auto',
                 is_nac=False,
                 born_file=None,
                 scalar_mediator_flag=True,
                 dark_photon_flag=False,
                 param_flag=False,
                 fold_BZ=True,
                 lowq_scaling=False,
                 scaling_qpoints=None,
                 no_anh_flag=False):
        '''
        Constructor for DynamicStructureFactor with several input parameter options.

        Parameters
        -------
        poscar_file : str
            POSCAR file in VASP standard format representing the unit cell of material.
        mesh : List
            List of integers representing the gamma-centered Monkhorst-Pack grid (e.g. 5x5x5 grid is mesh=[5, 5, 5])
        supercell : List
            List of integers representing the supercell matrix needed to make a supercell from the POSCAR.
        fc_file : str
            String representing the FORCE_SETS (or equivalent) file created from Phonopy (not Phono3py).
        q_point_list : List
            Array representing the list of q-points requested for calculation of dynamic structure factor.
        q_point_shift : List
            Optional array specifying the shift away from the Gamma center for the MP grid. Note that by default theMP grid in this code is gamma-centered, and specifying a non-centered grid must be set with this parameter, which is different from standard formulations of MP grids.
        meshG : List
            Optional list of integers describing the MP grid of reciprocal lattice vectors (called G vectors) to use in conjunction with `q_point_list`.
        weights : List
            Optional list allowing the user to specify the weights for each q-point in `q_point_list`. This feature should be used if the user wants to calculate the dynamic structure factor for the irreducible Brillouin zone, and knows the weights associated with each q-point in the irreducible zone.
        fc2_disp : str
            File specifying the displacements used in Phonopy to find the 2nd order force sets (e.g. disp.yaml)
        fc3_disp : str
            File specifying the displacements used in Phono3py to find the 3rd order force sets.
        delta_e : float
            Float specifying the energy spacing to discretize the dynamic structure factor in frequency/energy space in THz. Default is 0.01 THz
        max_e : float
            Float specifying the maximum energy to calculate the dynamic structure factor for in THz. Default is 30 THz.
        num_overtones : int
            Integer specifying the number of overtones, or multi-phonon processes to include in the contact interaction. For example, if one wants to include up to 5-phonon processes, then num_overtones=5.
        temperature : float
            Temperature to use for Phono3py internal calculations. In the future, this will be used for the dynamic structure factor, but at present is not implemented in that way.
        freq_min : float
            Minimum frequency, below which the contribution of that phonon is ignored. E.g. if a phonon frequency is 1e-5 THz, its contributions to the dynamic structure factor is ignored. Primarily used to eliminate artifacts from unstable modes with negative frequencies, and to ignore the gamma-point acoustic contributions.
        primitive_flag : str
            Flag used to find the primitive cell in the Phono3py calculations. String must match Phono3py conventions.
        is_nac : bool
            Boolean flag specifying whether non-analytic correction (NAC) is to be used.
        born_file : str
            the BORN file which provides the Born effective charges. File has to be generated from `phonopy-vasp-born`.
        scalar_mediator_flag : bool
            Boolean specifying whether to calculate the differential scattering rate (instead of the dynamic structure factor) for scalar mediated dark matter interaction. For this, the atomic contributions to the dynamic structure factor are weighted by the atomic mass.
        dark_photon_flag : bool
            Boolean specifying whether to calculate the diff scatt. rate (instead of DSF) for dark photon mediated dark matter interactions. This functionality only works with a BORN file given, as this is needed to calculate the diff scattering rate in this case.
        param_flag : bool
            Boolean specifying if the anharmonic processes are modeled by a parametrized Lorentzian function instead of the imaginary self-energy function. 
        fold_BZ : bool
            Flag specifying if the q-points given are to be folded into the irreducible Brillouin zone.
        lowq_scaling : bool
            Flag specifying if dynamic structure factor is to be heavily sampled near the Gamma point. This is done by creating a pseudo MP grid within the voxel centered at the gamma point of the original MP grid. The values of the dynamic structure factor at these q-points within the central voxel are determined via linear interpolation, and the non-periodicity of these q-points means that the contact interaction is ignored here.
        scaling_qpoints : List
            List allowing the user to specify the scaling q-points directly instead of letting the program determine the scaling q-points itself.
        no_anh_flag : bool
            Flag specifying whether or not to exclude anharmonic interactions.
        '''
        self.param_flag = param_flag
        self.mesh = mesh
        self.meshG = meshG
        self.qpoints = q_point_list
        self.supercell = supercell
        self.delta_e = delta_e
        self.max_e = max_e
        self._num_bins = int(np.ceil(self.max_e / self.delta_e))
        self.num_overtones = num_overtones
        self.temperature = temperature
        self.freq_min = freq_min
        self.is_nac = is_nac
        self.fold_BZ = fold_BZ
        self.lowq_scaling = lowq_scaling
        self._nac_q_direction = [1, 0, 0]

        self.disp_data = parse_disp_fc3_yaml(filename=fc3_disp)
        self.fc3_data = parse_FORCES_FC3(self.disp_data, filename=fc3_file)
        phono3py = Phono3py(read_vasp(poscar_file),
                            supercell_matrix=np.diag(supercell),
                            primitive_matrix='auto',
                            mesh=mesh,
                            log_level=1)
        phono3py.produce_fc3(self.fc3_data, self.disp_data)

        # if qpoint list not given, then load from mesh and shift
        self.qpoint_shift = q_point_shift
        self._brillouinzone = None
        self._mesh_brillouinzone = None
        if self.fold_BZ:
            self._mesh_brillouinzone = BrillouinZone(mesh=self.mesh, shift=self.qpoint_shift, poscar=poscar_file)
            if weights is None:
                combined_mesh = np.array(self.mesh) * np.array(self.meshG)
                combined_shift = np.array(self.qpoint_shift) / np.array(self.meshG)
                self._brillouinzone = BrillouinZone(mesh=combined_mesh, poscar=poscar_file, shift=combined_shift)
                # need to correct the keys of the weights dict
                new_weights_dict = {tuple(np.array(k) * np.array(meshG)): v for k, v in self._brillouinzone.weights.items()}
                self._brillouinzone.weights = new_weights_dict
                self.weights = new_weights_dict
            else:
                self.weights = weights
        #self.kernel_qpoints = self.get_qpoint_list(self.mesh)
        if fc_file is not None:
            if fc_file[-4:] == 'hdf5' or fc_file[-15:] == 'FORCE_CONSTANTS':
                phonon = load(supercell_matrix=supercell,
                              primitive_matrix=primitive_flag,
                              unitcell_filename=poscar_file,
                              force_constants_filename=fc_file,
                              is_nac=is_nac,
                              born_filename=born_file)
            elif fc_file[-10:] == 'FORCE_SETS' or fc_file[-3:] == 'FC2' or fc_file is None:
                phonon = load(supercell_matrix=supercell,
                              primitive_matrix=primitive_flag,
                              unitcell_filename=poscar_file,
                              force_sets_filename=fc_file,
                              is_nac=is_nac,
                              born_filename=born_file)
            else:
                print(fc_file, 'is not a recognized filetype!\nProgram exiting...')
                raise FileNotFoundError
        else:
            if self.is_nac:
                nac_params = parse_BORN(phono3py.get_phonon_primitive(), filename=born_file)
                nac_params['factor'] = Hartree * Bohr
            else:
                nac_params=None
            phonon = Phonopy(read_vasp(poscar_file),
                             supercell_matrix=supercell,
                             primitive_matrix=primitive_flag,
                             nac_params=nac_params)

            fc2 = phono3py.get_fc2()
            phonon.force_constants = fc2

        self.phonon = phonon
        self.eigenvectors = None
        self.frequencies = None
        self._weights = None
        self.kernel_qpoints = None
        self.run_mesh()
        self.rec_lat = np.linalg.inv(self.phonon.primitive.get_cell()).T

        self.sqw = []
        self.scaling_sqw = []
        self.exp_DWF = []
        self.dxdydz = 0.0
        self.dxdydzdw = 0.0
        self.skw_kernel = []
        self.anharmonicities = None
        self.born_charges = None
        self.dielectric = None

        # scaling properties
        self._scaling_qpoints = scaling_qpoints
        self._scaling_eigs = None
        self._scaling_freqs = None
        self._scaling_weights = None
        # NAC settings
        if self.is_nac is True or born_file is not None:
            self.set_born_charges()
            self.set_dielectric()
        # Dark matter scattering flags
        self._scalar_mediator_flag = scalar_mediator_flag
        self._dark_photon_flag = dark_photon_flag
        self._no_anh_flag = no_anh_flag
        #Anharmonic flags
        if fc3_file is not None and fc3_disp is not None:
            self.set_anharmonicities(poscar=poscar_file,
                                     fc3_file=fc3_file,
                                     disp_file=fc3_disp,
                                     born=born_file)

    #@property
    #def qpoints(self):
    #    if self._brillouinzone is None:
    #        return self.qpoints
    #    else:
    #        if self._qpoints is None:
    #            self.qpoints = self._brillouinzone.qpoints
    #        return self._qpoints

    @property
    def weights(self):
        '''
        Getter function of weights parameter.
        '''
        if self._weights is None:
            if self._brillouinzone is None:
                self._weights = np.ones(len(self.qpoints))
            else:
                self._weights = self._brillouinzone.weights
        return self._weights

    #@qpoints.setter
    #def qpoints(self, new_qpoints):
    #    self.qpoints = new_qpoints

    @weights.setter
    def weights(self, new_weights):
        '''
        Setter function of weights parameter.
        '''
        self._weights = new_weights

    @property
    def scaling_weights(self):
        '''
        Getter function of scaling_weights parameter.
        '''
        if self._scaling_weights is None:
            if self._mesh_brillouinzone is None:
                self._scaling_weights = np.ones(len(self._scaling_qpoints))
            else:
                self._scaling_weights = self._mesh_brillouinzone.weights
        return self._scaling_weights

    @scaling_weights.setter
    def scaling_weights(self, new_weights):
        '''
        Setter function of scaling_weights parameter
        '''
        self._scaling_weights = new_weights

    def sample_longwave_phonons(self):
        '''
        Sample dynamic structure factor near the Gamma point by dividing the current q-points in the mesh by the mesh list again to get small values of q-points that exist near the Gamma point. Then, the phonon properties are sampled at arbitrary q-points using Phonopy API.
        
        Returns
        -------
        Nothing, the sampled values are stored in `self._scaling_freqs` and `self._scaling_eigs`
        '''
        sq = self._scaling_qpoints
        if sq is None:
            scaling_qpoints = []
            for q, w in self._mesh_brillouinzone.weights.items():
                scaling_qpoints.append(q)
            # Convert from dict_keys to np.array for further math operations
            self._scaling_qpoints = np.array(list(scaling_qpoints))
            self._scaling_qpoints /= self.mesh
        self.sample_longwave_phonons_at_q()

    def sample_longwave_phonons_at_q(self, qpoints=None):
        '''
        Sample low-q phonon directly by running Phonopy API.

        Returns
        -------
        Nothing, the sampled values are stored in `self._scaling_freqs` and `self._scaling_eigs`
        '''
        if qpoints is None:
            self.phonon.run_qpoints(q_points=self._scaling_qpoints,
                                    with_eigenvectors=True,
                                    nac_q_direction=self._nac_q_direction)
        else:
            self.phonon.run_qpoints(q_points=qpoints,
                                    with_eigenvectors=True,
                                    nac_q_direction=self._nac_q_direction)
        qpoints_dict = self.phonon.get_qpoints_dict()
        self._scaling_freqs = qpoints_dict['frequencies']
        self._scaling_eigs = qpoints_dict['eigenvectors']

    def single_phonon_dsf(self):
        '''
        Calculate the dynamic structure factor for the single phonon and 3-phonon anharmonic contributions, ignoring the contact interaction terms.

        Returns
        -------
        Nothing, the dynamic structure factor is stored in self.scaling_sqw
        '''
        #DWF = self.test_exp_DWF_at_q(qpoint=qpoint)
        positions = self.phonon.primitive.get_scaled_positions()
        masses = self.phonon.masses
        natoms = len(masses)

        # Loop over scaling qpoints
        for sq_index, qpoint in enumerate(self._scaling_qpoints):
            # Test if qpoint is [0, 0, 0]
            if np.linalg.norm(qpoint) == 0.:
                scaling_sqw_at_q = np.zeros(self._num_bins, dtype=np.complex)
                self.scaling_sqw.append(scaling_sqw_at_q)
                continue
            # initialize scaling_sqw_at_q
            if self.anharmonicities is None:
                scaling_sqw_at_q = 0.
            else:
                scaling_sqw_at_q = np.zeros(self._num_bins, dtype=np.complex)
            DWF = self.test_exp_DWF_at_q(qpoint=qpoint)
            q_cart = np.dot(qpoint, self.rec_lat) * (2 * np.pi / AngstromsToMeters)
            for i, freq in enumerate(self._scaling_freqs[sq_index]):
                eigvec = self._scaling_eigs[sq_index, :, i]
                outer_eig = self.get_outer_eig(eigvec=eigvec, masses=masses, freq=freq, reduced_q=qpoint, pos=positions)

                for tau_1 in range(natoms):
                    for tau_2 in range(natoms):
                        if self.anharmonicities is None or self._no_anh_flag:
                            function = self.create_delta(freq)
                        else:
                            dist = self.create_anharmonic_distribution(q_point=self._scaling_qpoints[sq_index],
                                                                         band_index=i)
                            #interp = self.anharmonicities.get_broadening_function(qpoint=self._scaling_qpoints[sq_index],
                            #                                                        band_index=i,
                            #                                                        max_freq=self.max_e)
                            function = np.array(dist).astype(complex)
                            #function = np.array(interp(self.get_bin_energies())).astype(complex)
                        factor = np.dot(q_cart, np.dot(outer_eig[tau_1, tau_2, :, :], q_cart)) * function *\
                                 np.exp(2j * np.pi * np.vdot(qpoint, (positions[tau_1] - positions[tau_2]))) \
                                 * DWF[tau_1] * DWF[tau_2]
                        if self._scalar_mediator_flag is True:
                            factor *= masses[tau_1] * masses[tau_2]
                        elif self._dark_photon_flag is True:
                            factor *= np.abs(np.linalg.norm(q_cart) ** 2 /
                                             (np.dot(q_cart, np.dot(self.dielectric, q_cart)))
                                             ) ** 2
                        scaling_sqw_at_q += factor
            self.scaling_sqw.append(scaling_sqw_at_q)

    def run_mesh(self):
        '''
        Calculate the phonon properties using Phonopy API for the mesh of q-points.
        '''
        self.phonon.run_mesh(self.mesh,
                             shift=self.qpoint_shift,
                             is_mesh_symmetry=False,
                             with_eigenvectors=True,
                             is_gamma_center=True)
        mesh_dict = self.phonon.get_mesh_dict()
        self.eigenvectors = mesh_dict['eigenvectors']
        self.frequencies = mesh_dict['frequencies']
        #self.weights = mesh_dict['weights']
        self.kernel_qpoints = mesh_dict['qpoints']

    def set_born_charges(self):
        '''
        Setter of born charges to get array of Born charges from phonopy.
        '''
        self.born_charges = self.phonon.nac_params['born']

    def set_dielectric(self):
        '''
        Setter of dielectric to get 3x3 array of dielectric tensor from phonopy.
        '''
        self.dielectric = self.phonon.nac_params['dielectric']

    def set_anharmonicities(self, poscar, fc3_file, disp_file, born):
        '''
        Set anharmonicities using AnharmonicPhonons object.

        Parameters
        -------
        poscar : str
            VASP POSCAR file.
        fc3_file : str
            FORCES_SETS_FC3 file from phono3py.
        disp_file : str
            disp.yaml file created from the `phono3py -d` command option.
        born : str
            BORN file created from `phonopy-vasp-born` command.
        '''
        self.anharmonicities = AnharmonicPhonons(poscar=poscar,
                                                 fc3_file=fc3_file,
                                                 disp_file=disp_file,
                                                 mesh=self.mesh,
                                                 supercell=self.supercell,
                                                 temperature=self.temperature,
                                                 born=born,
                                                 param_flag=self.param_flag
                                                 )
        self.anharmonicities.set_self_energies()

    def get_bin_energies(self):
        '''
        Get a list of discretized energies used to calculate the dynamic structure factor.
        '''
        num_bins = int(np.ceil(self.max_e / self.delta_e))
        return np.arange(num_bins) * self.delta_e

    def get_qpoint_list(self, mesh):
        '''
        Get q-point list from a given Gamma-centered MP mesh. Ignore the meshG parameter and do not fold into the irreducible zone.

        Parameters
        -------
        mesh : List
            List of mesh indices. This could be removed to make this module use self.mesh instead.

        Returns
        -------
        qpoints
            List of qpoints requested from mesh indices. Shape = (prod(mesh), 3)
        '''
        qpoints = np.zeros([np.prod(mesh), 3])
        count = 0
        curr_qx = 0.0
        curr_qy = 0.0
        curr_qz = 0.0
        spacing = 1.0 / np.array(mesh)
        for z in range(mesh[2]):
            for y in range(mesh[1]):
                for x in range(mesh[0]):
                    qpoints[count, :] = [curr_qx, curr_qy, curr_qz]
                    count += 1
                    curr_qx += spacing[0]
                    if curr_qx > 0.5:
                        curr_qx -= 1.0
                curr_qy += spacing[1]
                if curr_qy > 0.5:
                    curr_qy -= 1.0
            curr_qz += spacing[2]
            if curr_qz > 0.5:
                curr_qz -= 1.0
        return qpoints

    def get_outer_eig(self, eigvec, masses, freq, reduced_q, pos):
        '''
        Calculate the outer products of an eigenvectors with itself.

        Parameters
        -------
        eigvec : numpy.array
            Phonon eigenvector that is a list of 3 vectors for each atom in the unit cell.
        masses : List
            List of masses for each atom in the unit cell.
        freq : float
            Frequency of phonon mode associated with the eigenvector provided.
        reduced_q : List
            The q-point in reduced coordinates that matches with the eigenvector provided.
        pos : List
            The positions of the atoms in the unit cell in reduced coordinates. This input could be removed in favor of saving the atomic positions into a field in the class itself.

        Returns
        -------
        outer_eig
            List of outer products of eigenvectors.
        '''
        outer_eig = np.zeros([len(masses), len(masses), 3, 3], dtype=np.complex)
        eigvec = eigvec.reshape(-1, 3)
        phase = np.exp(-2j * np.pi * np.dot(pos, reduced_q))
        for i in range(len(masses)):
            for j in range(i, len(masses)):
                if self._dark_photon_flag is True:
                    outer_eig[i, j, :, :] = np.outer(np.conj(np.dot(eigvec[i, :], self.born_charges[i].T)),
                                                     np.dot(self.born_charges[j], eigvec[j, :].T)) * const.hbar
                else:
                    outer_eig[i, j, :, :] = np.outer(np.conj(eigvec[i, :]), eigvec[j, :]) * const.hbar
                outer_eig[i, j, :, :] *= phase[i] * np.conj(phase[j]) / (2 * AMU * np.sqrt(masses[i] * masses[j])) / (
                        2 * np.pi * freq * THz)

                if i is not j:
                    outer_eig[j, i, :, :] = np.conj(outer_eig[i, j, :, :])
        return outer_eig

    def get_outer_eigs_at_q(self, q_index):
        '''
        Calculate the outer product of an eigenvector at a given q-point.

        Parameters
        -------
        q_index : int
            Integer specifying the index of the q-point in the self.qpoints array. This uses the get_outer_eig module above to get the list of outer_eigs. The code loops over all branch indices for a given q-point to create a list of outer product of eigenvectors

        Returns
        -------
        outer_eig_list
            List of outer product of eigenvectors for each branch index at a given q-point. Shape = (num_masses, num_masses, 3, 3, num phonon modes)
        '''
        eigvecs = self.eigenvectors[q_index]
        masses = self.phonon.masses
        frequencies = self.frequencies[q_index]
        qpoint = self.kernel_qpoints[q_index]
        positions = self.phonon.primitive.get_scaled_positions()

        outer_eig_list = np.zeros([len(masses), len(masses), 3, 3, len(frequencies)], dtype=np.complex)
        for i, f in enumerate(frequencies):
            if q_index == 0 and i < 3:
                continue
            if self.freq_min < f:
                outer_eig_list[:, :, :, :, i] = self.get_outer_eig(eigvecs[:, i], masses, f, qpoint, positions)
        return outer_eig_list

    def set_dxdydzdw(self):
        '''
        Set the jacobian of the dynamic structure factor. Defined as the dq_x * dq_y * dq_z * dw, where d_w is differential for the frequency.
        '''
        dxdydz_matrix = np.empty([3, 3])
        for i in range(3):
            dxdydz_matrix[i, :] = 2 * np.pi * self.rec_lat[i, :] / self.mesh[i]
        dxdydz = np.abs(np.linalg.det(dxdydz_matrix))
        dxdydzdw = dxdydz * self.delta_e
        self.dxdydz = dxdydz
        self.dxdydzdw = dxdydzdw

    def build_skw_kernel(self):
        '''
        Internal routine that builds the kernel used for calculating the dynamic structure factor. This is important to pre-calculate for the contact interaction terms.
        '''
        q_index = 0
        freqs = self.frequencies
        num_atoms = len(self.phonon.primitive.masses)
        num_bins = int(np.ceil(self.max_e / self.delta_e))
        skw_kernel = np.zeros(self.mesh + [num_atoms, num_atoms, 3, 3, num_bins], dtype=np.complex)
        norm_factor = 1 / np.prod(self.mesh)
        # norm_factor = 1
        if self.dxdydz == 0:
            self.set_dxdydzdw()

        for k in range(self.mesh[2]):
            for j in range(self.mesh[1]):
                for i in range(self.mesh[0]):
                    outer_eigs_at_q = self.get_outer_eigs_at_q(q_index)
                    freqs_at_q = freqs[q_index]

                    for outer_counter in range(outer_eigs_at_q.shape[-1]):
                        if self.anharmonicities is None:
                            skw_kernel[i, j, k, :, :, :, :, :] += norm_factor * np.tensordot(
                                outer_eigs_at_q[:, :, :, :, outer_counter].reshape((1,) + outer_eigs_at_q.shape[:-1]),
                                self.get_spectrum([self.dxdydz ** -1], [freqs_at_q[outer_counter]]).reshape(
                                    [1, num_bins]),
                                axes=[0, 0]
                            )
                        else:
                            skw_kernel[i, j, k, :, :, :, :, :] += norm_factor * np.tensordot(
                                outer_eigs_at_q[:, :, :, :, outer_counter].reshape((1,) + outer_eigs_at_q.shape[:-1]),
                                self.get_spectrum([self.dxdydz ** -1],
                                                  [freqs_at_q[outer_counter]],
                                                  q_point=self.kernel_qpoints[q_index],
                                                  band_index=outer_counter
                                                  ).reshape([1, num_bins]),
                                axes=[0, 0]
                            )
                    q_index += 1
        self.skw_kernel = skw_kernel

    def create_delta(self, energy):
        '''
        Module that creates a delta function at a specified energy/frequency. Uses self.delta_e and self.max_e to determine the discretized field over which to create the delta function. The delta function is a trapezoid function where the largest side of the trapezoid coincides with the discretized energy value closest to the provided energy value.

        Returns
        -------
        delta_fcn : List
            The delta function list. The list is mostly zeros, so it could be possible to build a delta function class that allows for more of a sparse representation easing the memory constraints of the code. Although the memory constraints are not too restrictive.
        '''
        # units in meV
        num_bins = int(np.ceil(self.max_e / self.delta_e))
        delta_fcn = np.zeros(num_bins)

        if energy < 0:
            return delta_fcn

        e_bin_minus = int(np.floor(energy / self.delta_e))
        e_bin_plus = int(np.ceil(energy / self.delta_e))

        alpha_minus = np.abs(e_bin_minus * self.delta_e - energy) / self.delta_e
        alpha_plus = np.abs(e_bin_plus * self.delta_e - energy) / self.delta_e

        delta_fcn[e_bin_minus] = (1 - alpha_minus) / self.delta_e
        delta_fcn[e_bin_plus] = (1 - alpha_plus) / self.delta_e
        return delta_fcn

    def create_anharmonic_distribution(self, q_point, band_index):
        '''
        Calculate and return the anharmonic distribution function. Calculate from the Phono3py API. Can return a parametrized Lorentzian from phonon lifetime.

        Parameters
        -------
        q_point : List
            Vector representing the q-point for which to calculate the anharmonic distribution.
        band_index : int
            Integer representing the band/mode index.

        Returns
        -------
        broadening_func : List
            List representing the broadening function calculated from Phono3py API.
        '''
        freqs = self.get_bin_energies()
        anh_dist_func = self.anharmonicities.get_broadening_function(q_point, band_index, max_freq=self.max_e)
        gamma = self.anharmonicities.gamma.get_property_value(q_point, band_index)
        omega = self.anharmonicities.phonon_freqs.get_property_value(q_point, band_index)

        f_index_minus = np.floor(omega / (freqs[1] - freqs[0])).astype(int)
        if not self.param_flag:
            avg_gamma_at_freq = (gamma[f_index_minus] + gamma[f_index_minus + 1]) / 2
        else:
            avg_gamma_at_freq = gamma

        broadening_func = anh_dist_func(freqs)
        # Normalize in case Phono3py sampling is denser than delta_e here; I'm looking at you CsI
        if avg_gamma_at_freq < freqs[1] - freqs[0]:
            # Case: the width is small, and the area is underestimated
            # Solution: add a delta function
            broadening_func += self.create_delta(omega) * (1.0 - np.trapz(broadening_func, freqs))
        else:
            # Case: width is large enough, but the area is underestimated
            # Solution: scale the function by the trapz area
            broadening_func /= np.trapz(broadening_func, freqs)

        return broadening_func
        #return anh_dist_func(freqs)

    def create_lorentzian(self, energy, gamma):
        '''
        Module that calculates a Lorentzian. Very likely deprecated at this point as the parametrized Lorentzian is now calculated from a different piece of code.

        Parameters
        ------
        energy : float
            Energy representing the center of the Lorentzian function.
        gamma : float
            Float representing the width of the Lorentzian.

        Returns
        -------
        lorentzian : List
            List representing the Lorentzian function
        '''
        # create a Lorentzian function instead of a delta function for peaks broadened by anharmonicities
        # units of energy are THz for compatibility with phono3py

        num_bins = int(np.ceil(self.max_e / self.delta_e))
        lorentzian = np.zeros(num_bins)

        # check if energy is negative, return zero vector if true
        if energy < 0:
            return lorentzian

        # check if gamma is less than resolution of grid, if so return delta fcn instead of lorentzian
        if gamma < self.delta_e:
            return self.create_delta(energy)

        # in an effort to speed up section
        n_sigma = 10
        # x_vals = np.arange(num_bins) * delta_e
        # lorentzian = np.array([(1 / np.pi * 0.5 * gamma / ((energy - x_val)**2 + (0.5 * gamma)**2))
        #                       if (x_val < energy + n_sigma * gamma) and (x_val > energy + n_sigma * gamma)
        #                       else 0.0 for x_val in x_vals])
        bin_spread = int(np.ceil(n_sigma * gamma / self.delta_e))
        e_bin = int(np.round(energy / self.delta_e))
        x_vals = np.arange(e_bin - bin_spread, e_bin + bin_spread + 1) * self.delta_e
        for i in range(len(x_vals)):
            lorentzian[i + e_bin - bin_spread] = 1 / np.pi * 0.5 * gamma / (
                    (energy - x_vals[i]) ** 2 + (0.5 * gamma) ** 2)
        return lorentzian

    def circ_conv(self, signal, kernel):
        '''
        purpose is to perform a circular convolution (e.g. signal is periodic) of a 3-dimensional object. In this case, the 3-d object is a function proportional to phonon eigenvector as a function of q in the BZ
        TODO Incorporate my own padding to the frequency dimension of the circular convolution; use appropriate tags

        Parameters
        ------
        signal : List
            Initial signal function to be convolved.
        kernel : List
            Kernel function with which to convolve the signal.

        Returns
        -------
        The convolution of the signal and kernel using the convolution theorem using the Fourier transforms.
        '''
        return np.fft.ifftn(np.fft.fftn(signal) * np.fft.fftn(kernel))

    def convolve_f_i(self, f_i, coh_flag=False):
        '''
        Convolve the current version of the scattering signal. The kernel that convolves with the signal is determined from the self.skw_kernel object.

        Parameters
        -------
        f_i : List
            Current version of signal to be convolved.
        coh_flag : bool
            Boolean specifying if coherent scattering version of dynamic structure factor is to be used.

        Returns
        -------
        total_f_i : List
            The convolved signal
        '''
        # norm_constant = 1 / sum(phonons.weights)
        norm_constant = 1
        curr_f = f_i  # / sum(phonons.weights)
        f_i = f_i * norm_constant
        total_f_i = curr_f
        # norm_constant = 1 / 8
        # norm_constant = 1 / (8 * sum(phonons.weights))

        for i in range(1, self.num_overtones):
            # print('overtone ='), print(i + 1)
            if i > 1:
                curr_f = curr_f * norm_constant
            # print('int of curr_f'), print(np.trapz(curr_f[1,0,0,:]) * self.delta_e)
            if coh_flag:
                if i == 1:
                    curr_f *= 1
                curr_f = self.circ_conv(curr_f, f_i) * self.delta_e * self.dxdydz
            else:
                curr_f = signal.fftconvolve(curr_f, f_i, mode='full')[:len(f_i)] * self.delta_e
            # print('int of curr_f'), print(np.trapz(curr_f[1,0,0,:]) * self.delta_e)
            # total_f_i += curr_f / (np.sqrt(math.factorial(i + 1)))

            total_f_i += curr_f / float(math.factorial(i + 1))
        return total_f_i

    def get_spectrum(self, f_ab, frequencies, q_point=None, band_index=None):
        '''
        Get the spectrum for a given frequency.

        Parameters
        -------
        f_ab : List
            List of weights that multiply the spectra by.
        frequencies : List
            List of phonon frequencies for which to calculate spectra. 
        q_point : List
            Optional specification of q_points for which to calculate the spectrum.
        band_index : int
            Optional specification of band index for use with anharmonic part of code.

        Returns
        -------
        spectrum : List
            Returns the spectrum
        '''
        num_bins = int(np.ceil(self.max_e / self.delta_e))
        spectrum = np.zeros(num_bins, dtype=np.complex)
        freqs = self.get_bin_energies()
        if q_point is not None:
            # print('using anharmonic gammas')
            for i in range(len(frequencies)):
                # print('integral of broad func in DSF code =', np.trapz(self.create_anharmonic_distribution(q_point, band_index))*0.1)
                if self._no_anh_flag:
                    anh_dist = self.create_delta(frequencies[i])
                else:
                    anh_dist = self.create_anharmonic_distribution(q_point, band_index)
                integral = np.trapz(anh_dist, freqs)
                # print('integral =', integral)
                if integral != 0:
                    spectrum += f_ab[i] * anh_dist / integral
                else:
                    if frequencies[i] > self.freq_min:
                        # Add delta in case there is a purely harmonic mode that has a positive frequency (very rare)
                        spectrum += f_ab[i] * self.create_delta(frequencies[i])
        else:
            for i in range(len(frequencies)):
                spectrum += f_ab[i] * self.create_delta(frequencies[i])
        # print('integral of spectrum =', np.trapz(spectrum, freqs))
        return spectrum

    def test_exp_DWF_at_q(self, q_index=None, qpoint=None):
        '''
        Calculate the Debye-Waller factor (DWF) for a given q_index and qpoint. This is a test module to attempt to speed up DWF calculation.

        Parameters
        -------
        q_index : int
            Integer specifying q-index.
        q_point : List
            List specifying the q-point.

        Returns
        -------
        exp_DWF : List
            List of DWF for each atomic index.
        '''
        eigvecs = self.eigenvectors
        if self.fold_BZ:
            weights = np.array(list(self.weights.values()))
        else:
            weights = self.weights

        if qpoint is None:
            q_cart = np.dot(self.qpoints[q_index], self.rec_lat) * (2 * np.pi / AngstromsToMeters)
        else:
            q_cart = np.dot(qpoint, self.rec_lat) * (2 * np.pi / AngstromsToMeters)
        norm_constant = 1 / np.sum(weights)
        frequencies = self.frequencies
        masses = self.phonon.primitive.masses
        # eigvecs = np.reshape(eigvecs, [num_qpts, num_bands, len(masses), 3])
        exp_DWF = np.zeros(masses.shape)
        for m in range(len(masses)):
            DWF = 0.0
            for i, eigs_at_k in enumerate(eigvecs):
                for s, eig in enumerate(eigs_at_k.T):
                    if i == 0 and s < 3:
                        continue
                    if frequencies[i, s] > self.freq_min:
                        eig = np.reshape(eig, [len(masses), 3])
                        DWF += np.abs(np.dot(q_cart, eig[m, :])) ** 2 * const.hbar / (
                                4 * np.pi * frequencies[i, s] * THz * masses[m] * AMU)
            exp_DWF[m] = np.exp(-norm_constant * DWF / 2)
        return exp_DWF

    def compute_exp_DWF_at_q(self, q_index):
        '''
        Calculate the Debye-Waller factor (DWF) for a given q_index.

        Parameters
        -------
        q_index : int
            Integer specifying the q-point index for which to calculate the Debye-Waller factor.

        Returns
        -------
        Returns the calculated Debye waller factor
        '''

        eigvecs = self.eigenvectors
        weights = self.weights
        num_bands = eigvecs.shape[1]
        num_qpts = eigvecs.shape[0]

        q_cart = np.dot(self.qpoints[q_index], self.rec_lat) * (2 * np.pi / AngstromsToMeters)
        norm_constant = 1 / np.sum(weights)
        frequencies = self.frequencies
        masses = self.phonon.primitive.masses
        eigvecs = np.reshape(eigvecs, [num_qpts, num_bands, len(masses), 3])
        # eigvecs now has shape (#qpts, #bands, #atoms, 3)
        # want to contract q with the cartesian indices, use np.dot
        q_dot_e = np.dot(eigvecs, q_cart)
        # find absolute value and square of all numbers in resulting matrix
        q_dot_e_sq = np.abs(q_dot_e) ** 2
        # multiply weights by norm_factor
        # multiply abs square by weights using transpose and np.multiply
        weighted_q_dot_e_sq = (weights * q_dot_e_sq.T * norm_constant * const.hbar).T
        weighted_q_dot_e_sq /= masses * AMU
        # remove negative frequencies
        inv_freqs = np.zeros(frequencies.shape)
        for i, freqs_at_q in enumerate(frequencies):
            for j, f in enumerate(freqs_at_q):
                if f > self.freq_min:
                    inv_freqs[i, j] = 1 / (4 * np.pi * THz * f)  # Represents 1/2w, where w is the angular frequency
        weighted_q_dot_e_sq = (np.reshape(inv_freqs, np.prod(frequencies.shape)) * \
                               np.reshape(weighted_q_dot_e_sq, [np.prod(frequencies.shape), len(masses)]).T).T

        return np.exp(-1 / 2 * np.sum(weighted_q_dot_e_sq, axis=0))

    def get_exp_DWF(self):
        '''
        Calculate Debye Waller factor for each q-point from the internal list.

        Returns
        -------
        Nothing, saves the list of Debye Waller factors into self.exp_DWF
        '''
        for i in range(len(self.qpoints)):
            # self.exp_DWF.append(self.compute_exp_DWF_at_q(i))
            self.exp_DWF.append(self.test_exp_DWF_at_q(q_index=i))

    def get_coherent_sqw_at_q(self, q_index):
        '''
        Calculate the coherent dynamic structure factor at a given q-point index.

        Parameters
        -------
        q_index : int
            Integer specifying the q-point to use from the self.qpoints list.

        Returns
        -------
        sqw : List
            Returns the dynamic structure factor function calculated a given q-point.
        '''
        if len(self.skw_kernel) == 0:
            self.build_skw_kernel()
        s_qw = np.zeros(self.skw_kernel.shape[:3] + (self.skw_kernel.shape[-1],), dtype=np.complex)
        # dot skw_kernel by q_point
        q_point = self.qpoints[q_index]
        q_cart = np.dot(q_point, self.rec_lat) * 2 * np.pi / AngstromsToMeters

        masses = self.phonon.primitive.masses
        natoms = len(masses)

        contracted_kernel = np.tensordot(self.skw_kernel, q_cart, axes=[[5], [0]])
        contracted_kernel = np.tensordot(contracted_kernel, q_cart, axes=[[5], [0]])
        # problem HERE with tensordot flipping the sign of the outerproducts
        # now compute full scattering function from convolutions of s_fcn

        if len(self.exp_DWF) == 0:
            self.get_exp_DWF()

        norm_factor = np.prod(self.mesh)
        positions = self.phonon.primitive.get_scaled_positions()
        for tau_1 in range(natoms):
            for tau_2 in range(natoms):
                factor = self.convolve_f_i(contracted_kernel[:, :, :, tau_1, tau_2, :], coh_flag=True) * \
                         np.exp(2j * np.pi * np.vdot(q_point, (positions[tau_1] - positions[tau_2]))) \
                         * self.exp_DWF[q_index][tau_1] * self.exp_DWF[q_index][tau_2] * norm_factor
                if self._scalar_mediator_flag is True:
                    factor *= masses[tau_1] * masses[tau_2]
                elif self._dark_photon_flag is True:
                    factor *= np.abs(np.linalg.norm(q_cart) ** 2 /
                                     (np.dot(q_cart, np.dot(self.dielectric, q_cart)))
                                     ) ** 2

                s_qw[:, :, :, :] += factor
        return s_qw

    def get_indices_from_qpoint(self, q_point):
        '''
        Helper function to get index from q-point vector.

        Parameters
        ------
        q_point : List
            List representing q-point for which to get index.
        '''
        indices = np.round(q_point * self.mesh - self.qpoint_shift).astype(np.int)
        indices = indices % self.mesh
        return indices

    def interpolate_sqw(self, sqw, q_point):
        '''
        Module for interpolating sqw, but will be deprecated soon and removed because interpolation was never implemented for the full dynamic structure factor here. Also, interpolation may be implemented at some point, but will not be implemented in this module; would most likely use the Interpolater class from the `src` directory.

        Parameters
        -------
        sqw : List
            Array representing the dynamic structure factor to interpolate.
        q_point : List
            Arbitrary q-point to calculate the dynamic structure factor for.

        Returns
        ------
        The interpolated function for the dynamic structure factor
        '''
        # print('q-point =', q_point)
        # TODO save for later, implement actual interpolation scheme, right now will set up for the exact grid points
        indices = np.round(q_point * self.mesh - self.qpoint_shift).astype(np.int)
        indices = indices % self.mesh
        # print('indices =', indices)
        return sqw[indices[0], indices[1], indices[2], :]

    def get_coherent_sqw(self, start_q_index=0, stop_q_index=None):
        '''
        Calculate the coherent dynamic structure factor optionally starting with a given q-point index and stopping at another specified one.

        Parameters
        -------
        start_q_index : int
            Optional integer specifying which q-point index to start at. Helpful if user doesn't want to calculate for all q-points in the Brillouin zone.
        stop_q_index : int
            Optional integer specifying which q-point index to stop at.
        '''
        if stop_q_index is None:
            stop_q_index = len(self.qpoints)
        if self.fold_BZ and self._brillouinzone is not None and self.qpoints is None:
            qpoints = list(self.weights.keys())[start_q_index:stop_q_index]
        elif self.fold_BZ and self._brillouinzone is None and self.qpoints is None:
            print('NOTHING IMPLEMENTED YET FOR THE CASE OF BZ FOLDING BUT NO GRID, ENJOY YOUR ERROR MESSAGE! >:D')
        else:
            qpoints = self.qpoints[start_q_index:stop_q_index]
        print('testing qpoints =', qpoints)
        q_indices = np.arange(start_q_index, stop_q_index)
        for i, qpoint in zip(q_indices, qpoints):
            if np.linalg.norm(qpoints[i]) == 0.:
                if len(self.skw_kernel) == 0:
                    self.build_skw_kernel()
                self.sqw.append(np.zeros(self.skw_kernel.shape[-1], dtype=np.complex))
            else:
                self.sqw.append(self.interpolate_sqw(self.get_coherent_sqw_at_q(i), qpoint))
            print('Progress =', i, '/', len(self.qpoints))
        # set qpoints to what was actually calculated
        self.qpoints = self.qpoints[start_q_index:stop_q_index]
        if self.lowq_scaling:
            if self._scaling_qpoints is None:
                self.sample_longwave_phonons()
            else:
                self.sample_longwave_phonons_at_q()
            self.single_phonon_dsf()

    def write_coherent_sqw(self, filename, ftype='hdf5'):
        '''
        Write the final results to a file. Only works for an HDF5 file right now, so it's not advised to change the ftype command.

        Parameters
        -------
        filename : str
            Name of output file (given without suffix)
        ftype : str
            Type of output file to write. Only 'hdf5' is supported at this time. Subject to change in the future.
        '''
        if ftype == 'txt':
            fw = open(filename, 'w')
            fw.write('# first column is energy in THz, second column is the value of S(q, w)\n')
            for i, sqw_at_q in enumerate(self.sqw):
                fw.write('q-point = {}\n'.format(self.qpoints[i]))
                energy = 0.0
                for val in sqw_at_q:
                    fw.write('{:.3f}\t{:.10e}\n'.format(energy, val))
                    energy += self.delta_e
                fw.write('\n')
            fw.close()
        elif ftype == 'hdf5':

            with h5.File(filename, 'w') as fw:
                fw['q-points'] = self.qpoints
                # Write down connected q-points and weights if the BZ was folded
                if self.fold_BZ:
                    import json
                    from yaml import dump
                    try:
                        from yaml import CDumper as Dumper
                    except ImportError:
                        from yaml import Dumper
                    fw['weights'] = [self.weights[tuple(q)] for q in self.qpoints]
                    # connected q-points
                    #new_symm_dict = {k: np.array(v) * np.array(self.meshG)
                    #                 for k, v in self._brillouinzone.symm_qpoints.items()}
                    irr_gps = np.unique(self._brillouinzone.mapping)
                    symm_points = []
                    for gps in irr_gps:
                        q = self._brillouinzone.qpoints[gps]
                        symm_points.append(list(np.array(self._brillouinzone.symm_qpoints[tuple(q)]) *
                                                np.array(self.meshG)))
                    #symm_points = [np.array(self._brillouinzone.symm_qpoints[tuple(q)]) * np.array(self.meshG)
                    #               for q in irr_qpoints]
                    #fw['equivalent q-points'] = dump(symm_points, Dumper=Dumper)
                    symm_points = {'{:.8f},{:.8f},{:.8f}'.format(*(np.array(k) * np.array(self.meshG))): [list(v * np.array(self.meshG)) for v in val] 
                                    for k, val in self._brillouinzone.symm_qpoints.items()}
                    dt = h5.string_dtype(encoding='utf-8')
                    fw.create_dataset('equivalent q-points', shape=(), dtype=dt, data=json.dumps(symm_points))
                else:
                    if self.weights is None:
                        self.weights = np.ones(len(self.qpoints))
                    fw['weights'] = self.weights
                fw['sqw'] = np.abs(self.sqw)
                fw['reclat'] = self.rec_lat * 2 * np.pi
                volume = np.linalg.det(self.phonon.primitive.get_cell())
                total_mass = np.sum(self.phonon.masses) * const.atomic_mass
                density = total_mass / volume * 1e24 * 1000         # Convert kg/Ang^3 to g/cm^3
                fw['density'] = density
                fw['max freq'] = np.max(np.array(self.frequencies).reshape(-1))
                fw['frequencies'] = self.get_bin_energies()
                fw['delta_w'] = self.delta_e
                fw['dxdydz'] = self.dxdydz
                if self.lowq_scaling or self._scaling_qpoints is not None:
                    fw['scaling_sqw'] = np.abs(self.scaling_sqw)
                    fw['scaling_q-points'] = self._scaling_qpoints
                    if self.fold_BZ:
                        #symm_points = [np.array(self._mesh_brillouinzone.symm_qpoints[tuple(q * self.mesh)]) / np.array(self.mesh)
                        #               for q in self._scaling_qpoints]
                        #irr_gps = np.unique(self._mesh_brillouinzone.mapping)
                        #symm_points = []
                        #for gps in irr_gps:
                        #    q = self._mesh_brillouinzone.qpoints[gps]
                        #    symm_points.append(list(np.array(self._mesh_brillouinzone.symm_qpoints[tuple(q)]) /
                        #                            np.array(self.mesh)))
                        symm_points = {'{:.8f},{:.8f},{:.8f}'.format(*(np.array(k) / np.array(self.mesh))): [list(v / np.array(self.mesh)) for v in val] 
                                        for k, val in self._mesh_brillouinzone.symm_qpoints.items()}

                        #fw['equivalent scaling_q-points'] = dump(symm_points, Dumper=Dumper)
                        #fw['equivalent scaling_q-points'] = json.dumps(symm_points)
                        dt = h5.string_dtype(encoding='utf-8')
                        fw.create_dataset('equivalent scaling_q-points', shape=(), dtype=dt, data=json.dumps(symm_points))
                        if self.scaling_weights is None:
                            # Case where scaling points are set but want is to include folded weights
                            # TODO won't work with G-vectors in the weights object, need separate scaling BZ
                            fw['scaling_weights'] = [self._mesh_brillouinzone.weights[tuple(q * self.mesh)]
                                                     for q in self._scaling_qpoints]
                        else:
                            fw['scaling_weights'] = list(self.scaling_weights.values())
                    else:
                        fw['scaling_weights'] = np.ones(len(self._scaling_qpoints))
                    fw['scaling_dxdydz'] = self.dxdydz / np.prod(self.mesh)
                    fw['anharmonicity_flag'] = self.anharmonicities is not None
        else:
            print('ERROR: Unrecognized filetype used: ftype =', ftype)


'''
def compute_incoherent_Sqw(phonons, q_point, delta_e, max_e, num_overtones, anh_flag=False, gammas=[0]):
    # assume q-point is in reduced coordinates
    if np.abs(q_point[0]) > 0.5 or np.abs(q_point[1]) > 0.5 or np.abs(q_point[2]) > 0.5:
        print("WARNING: q-point is outside the first Brillouin zone. Did you enter the q-point in reduced coordinates?")
    # convert reduced qpt to coordinates with real units (units = 2pi * m^-1)
    q_point = np.dot(q_point, phonons.rlattice) * 10**10

    phonons.unnormalize_eigvecs()
    num_bins = int(np.ceil(max_e / delta_e))
    #f_vec = np.zeros([phonons.natoms, num_bins], dtype=np.complex)
    f_vec = np.zeros([phonons.natoms, len(phonons.frequencies)], dtype=np.complex)

    # calculate vector of partial scattering functions called f_i * exp_DWF
    exp_DWF = np.zeros(phonons.natoms)
    for i in range(phonons.natoms):
        exp_DWF[i] = compute_f_DWF(q_point, phonons, i)
        f_i = compute_f_i(q_point, phonons, i)
        # don't convolve yet, just store into f_vec
        f_vec[i, :] = f_i
    # now f_vec contains all information necessary to compute S_qw
    # need to multiply two different f_i's and convolve, then multiply by DWF's, and normalization constant (not implemented yet)
    s_qw = np.zeros(num_bins)
    # now compute full scattering function from pairs of partial scattering functions
    print('gammas =', gammas)
    for i in range(phonons.natoms):
        f_ab = np.abs(np.conj(f_vec[i, :]) * f_vec[i, :])
        print('sum of f_i = '), print(sum(f_ab))
        f_ab_spectrum = get_spectrum(f_ab, phonons.frequencies, delta_e, max_e, anh_flag, gammas)
        f_ab_spectrum = np.abs(f_ab_spectrum)
        print('int of f_i = '), print(np.trapz(f_ab_spectrum) * delta_e)
        f_ab_conv = convolve_f_i(f_ab_spectrum, num_overtones, delta_e)
        print('DWF')
        print(exp_DWF ** 2)
        s_qw += np.real(f_ab_conv) * exp_DWF[i] * exp_DWF[i]
    return s_qw

def compute_decoherence_time(phonons, q_point, delta_e, max_e, num_overtones, gammas, is_normalized=True):
    # assume q-point is in reduced coordinates
    if np.abs(q_point[0]) > 0.5 or np.abs(q_point[1]) > 0.5 or np.abs(q_point[2]) > 0.5:
        print("WARNING: q-point is outside the first Brillouin zone. Did you enter the q-point in reduced coordinates?")
    # convert reduced qpt to coordinates with real units (units = 2pi * m^-1)
    q_point = np.dot(q_point, phonons.rlattice) * 10**10
    if is_normalized:
        phonons.unnormalize_eigvecs()
    num_bins = int(np.ceil(max_e / delta_e))
    #f_vec = np.zeros([phonons.natoms, num_bins], dtype=np.complex)
    f_vec = np.zeros([phonons.natoms, len(phonons.frequencies)], dtype=np.complex)

    # calculate vector of partial scattering functions called f_i * exp_DWF
    exp_DWF = np.zeros(phonons.natoms)
    for i in range(phonons.natoms):
        exp_DWF[i] = compute_f_DWF(q_point, phonons, i)
        f_i = compute_f_i(q_point, phonons, i)
        # don't convolve yet, just store into f_vec
        f_vec[i, :] = f_i
    # now f_vec contains all information necessary to compute S_qw
    # need to multiply two different f_i's and convolve, then multiply by DWF's, and normalization constant (not implemented yet)
    s_qw = np.zeros(num_bins)
    gamma_s_qw = np.zeros(num_bins)
    # now compute full scattering function from pairs of partial scattering functions
    print('gammas =', gammas)
    for i in range(phonons.natoms):
        f_ab = np.abs(np.conj(f_vec[i, :]) * f_vec[i, :])
        gamma_f_ab = np.multiply(f_ab, np.exp(gammas))
        print('sum of f_i = '), print(sum(f_ab))
        f_ab_spectrum = get_spectrum(f_ab, phonons.frequencies, delta_e, max_e, True, gammas)
        f_ab_spectrum = np.abs(f_ab_spectrum)

        gamma_f_ab_spectrum = get_spectrum(gamma_f_ab, phonons.frequencies, delta_e, max_e, True, gammas)
        gamma_f_ab_spectrum = np.abs(gamma_f_ab_spectrum)
        print('int of f_i = '), print(np.trapz(f_ab_spectrum) * delta_e)
        f_ab_conv = convolve_f_i(f_ab_spectrum, num_overtones, delta_e)
        gamma_f_ab_conv = convolve_f_i(gamma_f_ab_spectrum, num_overtones, delta_e)
        print('DWF')
        print(exp_DWF ** 2)
        s_qw += np.real(f_ab_conv) * exp_DWF[i] * exp_DWF[i]
        gamma_s_qw += np.real(gamma_f_ab_conv) * exp_DWF[i] * exp_DWF[i]
    # need to take ln of each element in spectrum; cant use np.log because some elements in spectrum may be zero
    log_sqw = np.zeros(len(s_qw))
    log_gamma_sqw = np.zeros(len(gamma_s_qw))
    for i in range(len(log_sqw)):
        if s_qw[i] <= 0 or gamma_s_qw[i] <= 0:
            continue
        else:
            log_sqw[i] = np.log(s_qw[i])
            log_gamma_sqw[i] = np.log(gamma_s_qw[i])
    decoherence_time = 1 / (2 * np.trapz((log_gamma_sqw - log_sqw) * s_qw) / np.trapz(s_qw))
    return s_qw, decoherence_time;

def smear_spectrum(spectrum, sigma, delta_e):
    n_sigma = 10
    num_bins_half = int(np.ceil(n_sigma * sigma / delta_e))
    e_vals = np.arange(-num_bins_half, num_bins_half + 1) * delta_e
    gaussian = [1 / np.sqrt(2 * np.pi * sigma**2) * np.exp(-e**2 / (2 * sigma**2)) for e in e_vals]
    return signal.fftconvolve(spectrum, gaussian, mode='full')[num_bins_half:(num_bins_half + len(spectrum))] * delta_e

class Runner:
    def __init__(self, yaml_file, hdf5_file, q_point, num_overtones, max_e, delta_e):
        self.yaml_file = yaml_file
        self.hdf5_file = hdf5_file
        self.num_overtones = num_overtones
        self.max_e = max_e
        self.delta_e = delta_e
        self.q_point = q_point
    def get_coherent_spectrum(self):
        anh_phonons = phonon_lifetimes.Anh_System(self.yaml_file, self.hdf5_file)
        s_qw = compute_Sqw(anh_phonons.dyn_system, self.q_point, self.delta_e, self.max_e, self.num_overtones)
        mesh = anh_phonons.dyn_system.mesh
        q_bin = tuple(np.round((self.q_point % 1) * mesh).astype(np.int))
        print('int of s_qw for q =', self.q_point, 'is =', np.trapz(s_qw[q_bin]) * self.delta_e)
        return s_qw
    def get_phonons(self):
        return phonon_lifetimes.Anh_System(self.yaml_file, self.hdf5_file)
'''


def set_qpoints(mesh, num_Gpoints=1, q_shift=[0, 0, 0], stride=1, stride_G=1, max_q=None):
    qpoints = np.zeros([0, 3])
    count = 0
    curr_qx = q_shift[0] / mesh[0]
    curr_qy = q_shift[1] / mesh[1]
    curr_qz = q_shift[2] / mesh[2]
    spacing = float(stride) / np.array(mesh)
    for gz in range(0, num_Gpoints, stride_G):
        for gy in range(0, num_Gpoints, stride_G):
            for gx in range(0, num_Gpoints, stride_G):
                for z in range(0, mesh[2], stride):
                    for y in range(0, mesh[1], stride):
                        for x in range(0, mesh[0], stride):
                            if max_q is not None:
                                if np.abs(curr_qx) <= max_q and np.abs(curr_qy) <= max_q and np.abs(curr_qz) <= max_q:
                                    qpoints = np.append(qpoints, [[curr_qx + gx, curr_qy + gy, curr_qz + gz]], axis=0)
                            else:
                                qpoints = np.append(qpoints, [[curr_qx + gx, curr_qy + gy, curr_qz + gz]], axis=0)
                            count += 1
                            curr_qx += spacing[0]
                            if curr_qx > 0.5:
                                curr_qx -= 1.0
                        curr_qy += spacing[1]
                        if curr_qy > 0.5:
                            curr_qy -= 1.0
                    curr_qz += spacing[2]
                    if curr_qz > 0.5:
                        curr_qz -= 1.0
        print('gz =', gz)
    return qpoints


if __name__ == '__main__':
    num_overtones = 1
    delta_e = 0.1
    # units are unfortunately in THz
    max_e = 100

    # yaml_file = sys.argv[1]
    # hdf5_file = sys.argv[2]
    # anh_phonons = phonon_lifetimes.Anh_System(yaml_file, hdf5_file)
    # phonons = anh_phonons.dyn_system
    # phonons = ph.Dyn_System(yaml_file)

    # conv_THz_to_meV = 4.13567

    # test with arbitrary single q-point
    reduced_qpt = np.outer(np.arange(1, 12) / 23.0, np.array([1, 0, 0])) + 0
    reduced_qpts = reduced_qpt
    reduced_qpts = np.append(reduced_qpts, np.outer(np.arange(3, 4) / 5.0, np.array([1, 1, 1])) + 10, axis=0)
    reduced_qpts = np.append(reduced_qpts, np.outer(np.arange(3, 4) / 5.0, np.array([1, 1, 1])) + 20, axis=0)
    reduced_qpts = np.append(reduced_qpts, np.outer(np.arange(3, 4) / 5.0, np.array([1, 1, 1])) + 50, axis=0)
    # test with entire BZ

    # print(qpoints)
    # mapping, testgrid = get_BZ_map(phonons.mesh, phonons.lattice, phonons.positions, phonons.masses, magmoms=[])
    mesh = [9, 9, 9]
    supercell = [2, 2, 2]

    qpoints = set_qpoints(mesh=mesh)  # , num_Gpoints=10)#, num_Gpoints=300, stride_G=50)
    qpoints = qpoints[:10]
    print(qpoints)
    # supercell = [1, 1, 1]
    poscar = "/Volumes/GoogleDrive/My Drive/Cori_backup/GaAs/phono3py/2x2x2/POSCAR"
    # poscar = "/Volumes/GoogleDrive/My Drive/multiphonon/rubrene_POSCAR"
    fc = "/Volumes/GoogleDrive/My Drive/Cori_backup/GaAs/phono3py/2x2x2/fc2.hdf5"
    # fc = "/Volumes/GoogleDrive/My Drive/multiphonon/rubrene_FORCE_SETS"
    fc3 = "/Volumes/GoogleDrive/My Drive/Cori_backup/GaAs/phono3py/2x2x2/FORCES_FC3"
    disp = "/Volumes/GoogleDrive/My Drive/Cori_backup/GaAs/phono3py/2x2x2/disp_fc3.yaml"
    #fc3 = None
    #disp = None
    born = "/Volumes/GoogleDrive/My Drive/Cori_backup/GaAs/phono3py/2x2x2/BORN"

    poscar = "/Users/tfharrelson/PycharmProjects/compute_Sqw/data/CsI/POSCAR"
    fc = "/Users/tfharrelson/PycharmProjects/compute_Sqw/data/CsI/FORCES_FC2"
    fc3 = "/Users/tfharrelson/PycharmProjects/compute_Sqw/data/CsI/FORCES_FC3"
    disp = "/Users/tfharrelson/PycharmProjects/compute_Sqw/data/CsI/disp_fc3.yaml"
    # fc3 = None
    # disp = None
    born = "/Users/tfharrelson/PycharmProjects/compute_Sqw/data/CsI/BORN"
    dynamic_structure_factor = DynamicStructureFactor(poscar,
                                                      fc,
                                                      mesh,
                                                      supercell,
                                                      # q_point_list=reduced_qpt[0, :].reshape(-1, 3),
                                                      q_point_list=qpoints,
                                                      # q_point_list=reduced_qpt,
                                                      max_e=max_e,
                                                      delta_e=delta_e,
                                                      num_overtones=num_overtones,
                                                      fc3_file=fc3,
                                                      fc3_disp=disp,
                                                      is_nac=True,
                                                      born_file=born,
                                                      scalar_mediator_flag=False,
                                                      dark_photon_flag=True)
    dynamic_structure_factor.get_coherent_sqw()
    #dynamic_structure_factor.write_coherent_sqw('csi_darkphoton_sqw_m999_anh.hdf5')
    # s_qw = compute_Sqw(phonons, reduced_qpt, delta_e, max_e, num_overtones)
    # s_qw = compute_incoherent_Sqw(phonons, reduced_qpt, delta_e, max_e, num_overtones)
    # s_qw, decoherence_time = compute_decoherence_time(phonons, reduced_qpt, delta_e, max_e, num_overtones, anh_phonons.gammas, True)
