#!/usr/bin/python

# phonopy-FHI-aims.py - phonopy wrapper for FHI-aims
# Copyright (C) 2009-2011 Joerg Meyer (jm)
# Contributions and testing by Christian Carbogno (cc)
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

# 2016/08/29  togo: Known bug. Support for modulation output doesn't work.
#
# 2016/08/29  togo: Minor fix to follow the change of Phonopy instantiation
#                   Now phonopy_obj.generate_displacements(distance=some_value)
#                   must be called after the Phonopy instantiation.
#
# 2014/09/23  christian.carbogno && shanghui: adding write_eigenfrequencies_and_eigenvectors_at_gamma 
#                                            (both non-nac and nac)                      
#
# 2014/09/22 shanghui: adding support for modulation output 
#                      (modulation setting in control.in)                      
#
# 2014/02/16 togo: Minor fix to follow the change of parse_BORN
# 2014/01/31 togo: Follow the changes of Phonopy and Mesh classes (phonopy-1.8.2)
#
# 2014/01/09 togo: Follow the changes of Mesh, BandStructure, and Animation
#                 classes (phonopy-1.8.0)
#
# 2011/11/13 jm : included support for initial magnetic moments -
#                 (initial_moment tag in geometry.in)
#
# 2011/06/06 jm : polished some formating of numbers printed in output
#                 added support for atom_frac tag in phonopy.interface.FHIaims
#                 verified that 
#                    phonon supercell  -2 1 4  2 -1 4  2 1 -4
#                 (now) yields the correct 2x4x1 mutliple (containing 64 atoms) 
#                 of the smallest simple cubic cell (containing 8 atoms) of a diamond structure
#                 (i.e. two fcc lattices shifted by 0.25*(1,1,1) with respect to each, primitive cell with 8 atoms)
#                 after 
#                    phonopy.structure.cells.get_surrounding_frame()
#                 has been updated
#
# 2011/02/04 jm : adapted to >=phonopy-0.9.2:
#                 IO functionality required for FHI-aims now rewritten in phonopy.FHIaims
#                 phonopy.ASE (and phonopy.ASE_FHIaims) removed since no longer needed
#                 adjusted to argument changes in write_arc and write_xyz
#
# 2010/07/23 jm : made qdensity mesh valued (optional - keeping backward compatibility)
#                 added frequency unit meV
#                 fixed wrong unit output for frequencies at Gamma
#
# 2010/06/10 jm : ported to phonopy-0.9.1.2:
#          & cc   adapted to extensions of dynamical_matrix with respect to non-analytical corrections
#                 added support for animation infrastructure
#                 moved several options to control.in
#
# 2010/05/12 jm : ported to phonopy-0.9.1:
#                 adapted to split of dos array into the two seperate omega, dos arrays in TotalDOS class
#
# 2010/04/10 jm : ported to phonopy-0.9.0
#                 added command line option -y to write yaml files which can be used with
#                 'native' post-processing tools supplied by phonopy
#                 added command line option -e to also calculate eigenvectors 
#                 (which are output to the yaml files only at present if combined with -y option)
#
# 2010/03/23 jm : added command line option -n to include non-analytical terms based on polarizabilities
#                 read from external file (VASP BORN format at present)
#                 removed 'hack' to use q-point distances calculated in BandStructure - now entirely done
#                 in post_process_band
#
# 2010/03/20 jm : ported to phonopy-0.7.4
#                 adapted for ASE-FHI-aims module (containing all FHI-aims specific ASE routines)
#                 exploit symmetry in q-point meshes by default (i.e. set default of symmetry option to True)
#                 Greek letters now substituted for corresponding English letter names as labels in dispersion plots 
#                 (can be disabled via new command line option -g)
#
# 2010/03/02 jm : one minor syntax change to be compatible with Python 2.5.2
#
# 2010/02/26 jm : switched to read_aims_output (instead of detour via aims calculator)
#                 integrated into phonopy-0.7.3
#
# 2010/02/22 jm : support for supercell matrices
#                 optional phonon subtags (according to FHI-aims) manual are now also treated as optional here
#
# 2010/02/19 jm : added new format of force constants matrix (aka Hessian) output 
#                 as required by thermodynamic integration code of Christian Carbogno 
#
# 2010/02/18 jm : included correct index mapping from primitive to supercell in write_Hessian()
#
# 2010/02/08 jm : added output of Hessian aka force constant matrix via command line option -H
#
# 2009/12/27 jm : initial version


class Control:
    def __init__(self, file=None):
        self.phonon = {}
        self.phonon["supercell"] = []
        self.phonon["displacement"] = 0.001
        self.phonon["symmetry_thresh"] = 1E-6
        self.phonon["frequency_unit"] = "cm^-1"
        self.phonon["nac"] = {}
        self.phonon["hessian"] = []
        self.phonon["band"] = []
        self.phonon["dos"] = {}
        self.phonon["free_energy"] = {}
        self.phonon["animation"] = []
        self.phonon["modulations"] = {}
        if file is None:
            self.file = "control.in"
        else:
            self.file = file
        self.read_phonon()
        
    def read_phonon(self):
        f = open(self.file, 'r')
        try:
            for line in f:
                if not line:
                    break
                fields = line.split()
                nfields = len(fields)
                if (nfields > 2) and (fields[0] == "phonon"):
                    if (fields[1] == "supercell"):
                        if (nfields >= 11):
                            s = map(int, fields[2:11])
                            Smat = numpy.array(s).reshape(3,3)
                        elif (nfields >= 5):
                            s = map(int, fields[2:5])
                            Smat = numpy.diag(s)
                        else:
                            raise Exception("invalid supercell")
                        det_Smat = numpy.linalg.det(Smat)
                        if (det_Smat == 0):
                            raise Exception("supercell matrix is not invertible")
                        # this consistency check is not strictly necessary, since int function above used to transform the input 
                        # already throws an exception when decimal numbers are encountered
                        # keep for consistency (and future changes to that spot?) nevertheless...
                        elif (abs(det_Smat-round(det_Smat)) > 1.0e-6):
                            raise Exception("determinant of supercell differs from integer by more than numerical tolerance of 1.0e-6")
                        self.phonon["supercell"] = s
                    if (fields[1] == "displacement"):
                        self.phonon["displacement"] = float(fields[2])
                    if (fields[1] == "symmetry_thresh"):
                        self.phonon["symmetry_thresh"] = float(fields[2])
                    if (fields[1] == "frequency_unit"):
                        self.phonon["frequency_unit"] = fields[2]
                    if (fields[1] == "nac") and (len(fields) >= 4):
                        if (len(fields) >= 7):
                            delta = tuple(map(float, fields[4:7]))
                        else: 
                            delta = (0.1, 0.1, 0.1)
                        if (delta[0] == 0.0) and (delta[1] == 0.0) and (delta[2] == 0.0):
                            raise Exception("evaluation of frequencies with non-analytic corrections must be shifted by delta away from Gamma")
                        parameters = { "file" : fields[2],
                                       "method" : fields[3].lower(),
                                       "delta" : delta }
                        self.phonon["nac"].update(parameters)
                    if (fields[1] == "hessian"):
                        self.phonon["hessian"] = fields[2:]
                    if (fields[1] == "band") and (len(fields) >= 11):
                        parameters = { "kstart" : map(float, fields[2:5]),
                                       "kend" : map(float, fields[5:8]),
                                       "npoints" : int(fields[8]),
                                       "startname" : fields[9],
                                       "endname" : fields[10] }
                        self.phonon["band"].append(parameters)
                    if (fields[1] == "dos") and (len(fields) >= 7):
                        parameters = { "fstart" : float(fields[2]),
                                       "fend" : float(fields[3]),
                                       "fpoints" : int(fields[4]),
                                       "broad" : float(fields[5]),
                                       "qdensity" : map(int, fields[6:]) }
                        self.phonon["dos"].update(parameters)
                    if (fields[1] == "free_energy") and (len(fields) >= 6):
                        parameters = { "Tstart" : float(fields[2]),
                                       "Tend" : float(fields[3]),
                                       "Tpoints" : int(fields[4]),
                                       "qdensity" : map(int, fields[5:]) }
                        self.phonon["free_energy"].update(parameters)
                    if (fields[1] == "animation") and (len(fields) >= 12):
                        parameters = { "q" : map(float, fields[2:5]),
                                       "band" : int(fields[5]),
                                       "amp" : float(fields[6]),
                                       "div" : int(fields[7]),
                                       "shift" : map(float, fields[8:11]),
                                       "files" : fields[11:] }
                        self.phonon["animation"].append(parameters)
                    if (fields[1] == "modulations") :
                        parameters = { "q" : map(float, fields[2:5]),
                                       "supercell" : map(int,fields[5:8]), 
                                       "delta" : float(fields[8]) }
                        self.phonon["modulations"].update(parameters)

        except Exception:
            print line,
            print "|-> line triggered exception: " + str(Exception)
            raise
        # supercell is mandatory for all what follows
        if not self.phonon["supercell"]:
            raise Exception("no supercell specified in %s" % self.file)
        f.close()
        
    def print_phonon(self, prefix=""):
        s = self.phonon["supercell"]
        if (len(s) == 3):
            print prefix + ( "phonon supercell  %d %d %d" % tuple(s) )
        elif (len(s) == 9):
            print prefix + ( "phonon supercell   %d %d %d  %d %d %d  %d %d %d" % tuple(s) )
        print prefix + ( "phonon displacement %f" % self.phonon["displacement"] )
        print prefix + ( "phonon symmetry_thresh %g" % self.phonon["symmetry_thresh"] )
        print prefix + ( "phonon frequency unit %s" % self.phonon["frequency_unit"] )
        nac = self.phonon["nac"]
        if (nac):
            print prefix + ( "phonon nac %s %s" % (nac["file"], nac["method"]) )
        hess = self.phonon["hessian"]
        if (hess):
            print prefix + ( "phonon %s" % " ".join(hess) )
        for nb in range(len(self.phonon["band"])):
            b = self.phonon["band"][nb]
            print prefix + ( "phonon band  %f %f %f  %f %f %f  %d  %s  %s" % \
                             tuple(b["kstart"] + b["kend"] + [b["npoints"]] + [b["startname"]] + [b["endname"]]) )
        d = self.phonon["dos"]
        if (d):
            print prefix + ( "phonon dos %f %f %d %f %s" % \
                             ( d["fstart"], d["fend"], d["fpoints"], d["broad"], \
                               ' '.join(map(str, d["qdensity"])) ) )
        fe = self.phonon["free_energy"]
        if (fe):
            print prefix + ( "phonon free_energy %f %f %d %s" % \
                             ( fe["Tstart"], fe["Tend"], fe["Tpoints"], \
                               ' '.join(map(str, fe["qdensity"])) ) )
        for na in range(len(self.phonon["animation"])):
            anim = self.phonon["animation"][na]
            print prefix + ( "phonon animation  %f %f %f  %d %d %d  %f %f %f  %s" % \
                             tuple(anim["q"] + [anim["band"], anim["amp"], anim["div"]]  
                                             + anim["shift"] + [' '.join(anim["files"])]) )


from phonopy.structure.cells import Primitive
from phonopy.harmonic.forces import Forces
from phonopy.harmonic.force_constants import get_force_constants
def write_Hessian(phonopy_obj, supercell_matrix_inv, set_of_forces):
    cells = (("initial cell", phonopy_obj.unitcell), ("supercell", phonopy_obj.supercell))
    primitive = Primitive(phonopy_obj.supercell, supercell_matrix_inv)
    forces = []
    for (i, disp) in enumerate(phonopy_obj.displacements):
        atom_number = disp[0]
        displacement = disp[1:]
        forces.append(Forces(atom_number, displacement, set_of_forces[i]))
    Hessian = get_force_constants(forces, phonopy_obj.symmetry, phonopy_obj.supercell)
    f = open("phonopy-FHI-aims-Hessian.dat", 'w')
    f.write("### Hessian from phonopy-FHI-aims \n")
    f.write("# \n")
    for (cell_name,cell) in cells:
        f.write("# %s: \n" % cell_name)
        for i in range(3):
            v = cell.get_cell()[i]
            f.write("#   lattice_vector %d   %11.6f % 11.6f %11.6f \n" % (i+1,v[0],v[1],v[2]))
        for n in range(cell.get_number_of_atoms()):
            r = cell.get_positions()[n]
            sym = cell.get_chemical_symbols()[n]
            f.write("#   atom %4d   %11.6f %11.6f %11.6f   %s \n" % (n+1,r[0],r[1],r[2],sym)) 
        f.write("# \n")
    f.write("# force constant matrix elements in eV/(Angstrom)^2 \n")	
    for (j_atom_count,j_atom) in enumerate(primitive.get_primitive_to_supercell_map()):
        for j_coord in range(3):
            for i_atom in range(phonopy_obj.supercell.get_number_of_atoms()):
                for i_coord in range(3):
                    f.write("%18.10e " % Hessian[i_atom,j_atom,i_coord,j_coord])
                f.write("  ")
            f.write("\t # displaced atom %d, coord %d \n" % (j_atom+1,j_coord+1))
    f.close()


from phonopy.harmonic.forces import Forces
from phonopy.harmonic.force_constants import get_force_constants
def write_force_constants(phonopy_obj, set_of_forces):
    cells = (("initial cell", phonopy_obj.unitcell,"#"), ("supercell", phonopy_obj.supercell,""))
    Nsuper = phonopy_obj.supercell.get_number_of_atoms()
    forces = []
    for (i, disp) in enumerate(phonopy_obj.displacements):
        atom_number = disp[0]
        displacement = disp[1:]
        forces.append(Forces(atom_number, displacement, set_of_forces[i]))
    force_constants = get_force_constants(forces, phonopy_obj.symmetry, phonopy_obj.supercell)
    f = open("phonopy-FHI-aims-force_constants.dat", 'w')
    f.write("### Hessian from phonopy-FHI-aims \n")
    f.write("# \n")
    for (cell_name,cell,prefix) in cells:
        f.write("# %s: \n" % cell_name)
        for i in range(3):
            v = cell.get_cell()[i]
            f.write(prefix + ("lattice_vector %11.6f %11.6f %11.6f   %4d \n" % (v[0],v[1],v[2],i+1)))
        for n in range(cell.get_number_of_atoms()):
            r = cell.get_positions()[n]
            sym = cell.get_chemical_symbols()[n]
            f.write(prefix + ("atom           %11.6f %11.6f %11.6f   %3s   %4d \n" % (r[0],r[1],r[2],sym,n+1))) 
        f.write("\n")
    f.write("# \n")
    f.write("# force constant matrix elements in eV/(Angstrom)^2 \n")
    f.write("#  1 2 3 means displaced atom 1, in coord 2, force on atom 3 in this line \n")
    for j_atom in range(Nsuper):
        for j_coord in range(3):
            for i_atom in range(Nsuper):
                fcs = [ force_constants[i_atom,j_atom,i_coord,j_coord] for i_coord in range(3) ]
                indeces = [ j_atom+1, j_coord+1, i_atom+1 ]
                f.write("force_constants %18.10e %18.10e %18.10e      %d %d %d \n" % 
                        (tuple(fcs)+tuple(indeces)))
        f.write("\n")
    f.close()


from phonopy.units import VaspToTHz as AimsToTHz, VaspToCm as AimsToCm, VaspToEv as AimsToEv, THzToCm, THzToEv
AimsFrequencyUnitFactors = { 'cm^-1' : AimsToCm, 'THz' : AimsToTHz, 'meV' : 1E3*AimsToEv } 
# use "\minus" instead of "-" here to avoid down arrows sometimes printed by matplotlib instead
AimsFrequencyUnitLabelsMatplotlib = { 'cm^-1' : "cm$^{\minus 1}$", 'THz' : "THz", 'meV' : "meV" }
AimsFrequencyUnitFactorsToTHz = { 'cm^-1' : 1.0 / THzToCm, 'THz' : 1.0, 'meV' : 1.0 / THzToEv / 1000 } 

BandStructureLabels = { 'gamma'    : "\Gamma", 
                        'delta'    : "\Delta", 
                        'lambda'   : "\Lambda",
                        'sigma'    : "\Sigma" }
    
from phonopy.phonon.band_structure import BandStructure
def post_process_band(phonopy_obj, parameters, frequency_unit_factor, is_eigenvectors=False, write_yaml=False, do_matplotlib=False, lookup_labels=False):
    bands = []
    # Distances calculated in phonopy.band_structure.BandStructure object
    # are based on absolute positions of q-points in reciprocal space
    # as calculated by using the cell which is handed over during instantiation.
    # Fooling that object by handing over a "unit cell" diag(1,1,1) instead clashes
    # with calculation of non-analytical terms.
    # Hence generate appropriate distances and special k-points list based on fractional
    # coordinates in reciprocal space (to keep backwards compatibility with previous
    # FHI-aims phonon implementation).
    bands_distances = []
    distance = 0.0
    bands_special_points = [distance]
    bands_labels = []
    label = parameters[0]["startname"]
    if lookup_labels:
        bands_labels.append(BandStructureLabels.get(label.lower(),label))
    else:
        bands_labels.append(label)
    for b in parameters:
        kstart = numpy.array(b["kstart"])
        kend = numpy.array(b["kend"])
        npoints = b["npoints"]
        dk = (kend-kstart)/(npoints-1)
        bands.append([(kstart + dk*n) for n in range(npoints)])
        dk_length = numpy.linalg.norm(dk)
        # one long list to simplify output
        for n in range(npoints):
            bands_distances.append(distance + dk_length*n)
        distance += dk_length * (npoints-1)
        bands_special_points.append(distance)
        # assuming that startname is the same as previous endname
        # (i.e. non-interrupted paths!)
        label = b["endname"]
        if lookup_labels:
            bands_labels.append(BandStructureLabels.get(label.lower(),label))
        else:
            bands_labels.append(label)
    bs_obj = BandStructure(bands, 
                           phonopy_obj.dynamical_matrix, 
                           is_eigenvectors=is_eigenvectors,
                           factor=frequency_unit_factor)
    # make band index first index (simpler for bands plotting !)
    eigenvalues = numpy.vstack(bs_obj.get_eigenvalues()).T
    frequencies = numpy.zeros_like(eigenvalues)
    for i, eigenvalues_band in enumerate(eigenvalues):
        frequencies_band = []
        for eigenvalue in eigenvalues_band:
            if eigenvalue < 0:
                frequencies_band.append(-numpy.sqrt(-eigenvalue))
            else: 
                frequencies_band.append(numpy.sqrt(eigenvalue))
        frequencies[i] = numpy.array(frequencies_band)
    fmin = numpy.min(frequencies)
    fmax = numpy.max(frequencies)
    for unit in AimsFrequencyUnitFactors:
        factor = AimsFrequencyUnitFactors[unit]
        if (factor == frequency_unit_factor):
            frequency_unit = unit
            break
    f = open("phonopy-FHI-aims-band_structure.dat", 'w')
    f.write("#\n")
    f.write("# Phonon bands from phonopy-FHI-aims!\n")
    f.write("#\n")
    lattice_real = phonopy_obj.unitcell.get_cell()
    lattice_reciprocal = 2.0 * numpy.pi * numpy.linalg.inv(lattice_real.transpose())
    for i in range(3):
        f.write("# Reciprocal lattice vector %11.6f %11.6f %11.6f \n" % tuple(lattice_reciprocal[i]))
    digits = int( math.ceil( math.log(len(parameters)+1,10) ) ) + 1
    digits_string = "%0" + str(digits) + "d"
    total_qpoints = 0
    total_qpoints_bands = []
    for i, b in enumerate(parameters):
        f.write("#\n")
        distance_start = bands_distances[total_qpoints]        
        f.write( ("# %5s point for band " + digits_string + ", %5s = (%9.5f,%9.5f,%9.5f) will be at real distance = %11.5f \n") % \
                  tuple(["Start"] + [i] + [b["startname"]] + b["kstart"] + [distance_start]) )
        total_qpoints += b["npoints"]
        distance_end = bands_distances[total_qpoints-1]
        f.write( ("# %5s point for band " + digits_string + ", %5s = (%9.5f,%9.5f,%9.5f) will be at real distance = %11.5f \n") % \
                  tuple(["End"] + [i] + [b["endname"]] + b["kend"] + [distance_end]) )                 
        total_qpoints_bands.append(total_qpoints)
    f.write("#\n")
    f.write("# number of phonon branches : %d \n" % len(frequencies))
    f.write("#\n")
    f.write("# q-distance(frac.)  frequencies(%s) \n" % frequency_unit)
    for i, dq in enumerate(bands_distances):    
        f.write("%f " % dq)
        frequencies_at_q = frequencies[:,i]
        for freq in frequencies_at_q:
            f.write(" %11.6f" % (freq*bs_obj.get_unit_conversion_factor()))
        f.write(" \n")
        if (i+1) in total_qpoints_bands:
            f.write("\n")
    f.close()
    if write_yaml:
        bs_obj.write_yaml()
    if do_matplotlib:
        params = { 'axes.labelsize': 20,
                   'xtick.labelsize' : 16,
                   'ytick.labelsize' : 16,
                   'text.fontsize': 24 }
        plt.rcParams.update(params)
        plt.cla()
        for frequencies_band in frequencies:
            plt.plot(bands_distances, frequencies_band*bs_obj.get_unit_conversion_factor(), 'r-')
        plt.xticks( bands_special_points,
                    tuple("$\mathrm{\mathsf{%s}}$" % bands_labels[i] for i in range(len(bands_labels))) )
        plt.xlabel('Wave vector')
        plt.ylabel("Frequency (%s)" % AimsFrequencyUnitLabelsMatplotlib[frequency_unit])
        plt.vlines(bands_special_points[1:-1], *plt.ylim())
        plt.xlim(xmin=0, xmax=bands_distances[-1])
        plt.ylim(ymin=round(fmin*bs_obj.get_unit_conversion_factor()))
        plt.savefig("phonopy-FHI-aims-band_structure.pdf")    

from phonopy.phonon.dos import TotalDos
def post_process_dos(phonopy_obj, mesh_obj, parameters, do_matplotlib=False):
    broad = parameters["broad"]
    dos_obj = TotalDos(mesh_obj, sigma=broad)
    fmin = parameters["fstart"]
    fmax = parameters["fend"] 
    fstep = ( fmax - fmin ) / parameters["fpoints"]
    dos_obj.set_draw_area(fmin, fmax, fstep)
    dos_obj.run()
    for unit in AimsFrequencyUnitFactors:
        factor = AimsFrequencyUnitFactors[unit]
        if (factor == frequency_unit_factor):
            frequency_unit = unit
            break
    f = open("phonopy-FHI-aims-dos.dat", 'w')
    f.write("# Phonon density of states from phonopy-FHI-aims\n")
    f.write("# %20s %20s \n" % ("Frequency (%s)" % frequency_unit, "DOS"))

    omegas, total_dos = dos_obj.get_dos()
    for freq, dos in zip(omegas, total_dos):
        f.write("  %20.10f %20.10f \n" % (freq, dos))
    f.close()
    if do_matplotlib:
        params = { 'axes.labelsize': 20,
                   'xtick.labelsize' : 16,
                   'ytick.labelsize' : 16,
                   'text.fontsize': 24 }
        plt.rcParams.update(params)
        plt.cla()
        plt.plot(omegas, total_dos, 'r-')
        plt.grid(True)
        plt.xlabel("Frequency (%s)" % AimsFrequencyUnitLabelsMatplotlib[frequency_unit])
        plt.ylabel("Density of states (a.u.)")
        plt.yticks([])
        plt.savefig("phonopy-FHI-aims-dos.pdf")

from phonopy.phonon.thermal_properties import ThermalProperties
from phonopy.units import EvTokJmol, Kb as kBoltzmann
def post_process_free_energy(phonopy_obj, mesh_obj, parameters, frequency_unit_to_THz, write_yaml=False, do_matplotlib=False):
    tp_obj = ThermalProperties(mesh_obj.get_frequencies() * frequency_unit_to_THz,
                               weights=mesh_obj.get_weights())
    Tmin = parameters["Tstart"]
    Tmax = parameters["Tend"]
    Tstep = ( Tmax - Tmin) / parameters["Tpoints"]
    tp_obj.run(Tstep, Tmax, Tmin)
    T, F, S, cv = tp_obj.get_thermal_properties()
    kJmolToEv = 1.0 / EvTokJmol
    JmolToEv = kJmolToEv / 1000
    T_aims = T
    F_aims = kJmolToEv * F
    JmolToEv = kJmolToEv / 1000
    TS_aims = T * (JmolToEv * S)
    U_aims = F_aims + TS_aims
    cv_aims = JmolToEv / kBoltzmann * cv    
    f = open("phonopy-FHI-aims-free_energy.dat", 'w')
    f.write("# Phonon free energy and specific heat from phonopy-FHI-aims\n")
    f.write("# Phonon zero point energy = %22.9f eV \n" % (F_aims[0]))
    f.write( ("#%20s " + "%30s"*4 + "\n") %
             ("Temperature (K)", "Free energy (eV/cell)", "Internal energy (eV/cell)",
              "c_v (kB/cell)", "-TS_vib(eV/cell)") )
    for i in range(T.shape[0]):
        f.write( ("%20.3f " + "%30.9f"*4 + "\n") %
                 (T_aims[i], F_aims[i], U_aims[i], cv_aims[i], -TS_aims[i]) )
    f.close()
    if write_yaml:
        tp_obj.write_yaml()
    if do_matplotlib:
        params = { 'axes.labelsize': 20,
                   'xtick.labelsize' : 16,
                   'ytick.labelsize' : 16,
                   'text.fontsize': 24 }
        plt.rcParams.update(params)
        plt.cla()
        plt.plot(T, F_aims, 'r-', label="Free energy (eV/cell)")
        plt.plot(T, U_aims, 'b-', label="Internal energy (eV/cell)")
        plt.plot(T, cv_aims, 'g-', label="c$_\mathrm{v}$ (kB/cell)")
        plt.plot(T, -TS_aims, "k--", label="-TS_vib(eV/cell)")
        plt.legend(loc='best', shadow=True)
        plt.grid(True)
        plt.xlabel("Temperature (K)")
        plt.savefig("phonopy-FHI-aims-free_energy.pdf")



#------------------begin for aims_modulation----------------------------
from phonopy.phonon.modulation import Modulation
from phonopy.phonon.degeneracy import get_eigenvectors
from phonopy.structure.cells import get_supercell
def run_aims_modulation(phonopy_obj):
        f_freq = open("freq.dat", 'w')
        f_norm = open("norm.dat", 'w')
        frequencies = []
        for ph_mode in phonopy_obj._phonon_modes:
            q, band_index, amplitude, argument = ph_mode
            eigvals, eigvecs = get_eigenvectors(
                q,
                phonopy_obj._dm,
                phonopy_obj._ddm,
                perturbation=phonopy_obj._delta_q,
                derivative_order=phonopy_obj._derivative_order,
                nac_q_direction=phonopy_obj._nac_q_direction)
            #print '---------------(1)d(q)----------------------'
            #print eigvecs[:, band_index]
            #print '---------------(2)d(q) expt(iqR)/sqrt(M)----------------------'
            u = phonopy_obj._get_delta(eigvecs[:, band_index], q)
            #print u
            phonopy_obj._eigvecs.append(eigvecs[:, band_index])
            phonopy_obj._eigvals.append(eigvals[band_index])
            #---------shanghui begin freq-------------------------------------
            if eigvals[band_index]  < 0:
                # Safety check for imaginary frequencies:
                neg_value = -numpy.sqrt(-eigvals[band_index])*AimsToCm
                if neg_value < -0.1:
                   print "* WARNING:"
                   print "* Imaginary Eigenfrequencies unphysical in modulation_aims.py"
                frequencies.append(neg_value)
            else:
                frequencies.append(numpy.sqrt(eigvals[band_index])*AimsToCm)  
            #---------shanghui end freq-------------------------------------

            #---------shanghui begin norm-------------------------------------
            norm =  numpy.linalg.norm(numpy.array(u).real)
            u = u/norm 
            f_norm.write("%d%18.12f \n" %(band_index+1, norm*norm))
            #----------shanghui end norm--------------------------------------
            phonopy_obj._delta_modulations.append(u * amplitude)

       
        for freq in frequencies:
            f_freq.write(" %11.6f" %freq)
        f_freq.close()
        f_norm.close()

def write_aims_modulation(self, filename="output_geometry"):
        deltas = []
        for i, delta_positions in enumerate(self._delta_modulations):
            cell = self._get_cell_with_modulation(delta_positions)
            write_aims((filename+"-%03d") % (i+1), cell)
            deltas.append(delta_positions)
    
        sum_of_deltas = numpy.sum(deltas, axis=0)
        no_modulations = numpy.zeros(sum_of_deltas.shape, dtype=complex)
        cell = self._get_cell_with_modulation(no_modulations)
        write_aims(filename+"-orig", cell)

def post_process_write_modulations(phonopy_obj,parameters):
        setting =  {}
        #-----set by reading contrin.in--------
        q = parameters['q']
        setting['dimension'] = parameters['supercell'] 
        amplitude = parameters['delta']

        #-----set to default values------------
        setting['order'] = None # This is derivtive of density matrix for group velocity
        argument=float(0.0) # This is phase factor, we set 0 here

        # add to phonon_modes
        vals=[]
        for band_index in range(3*cell.get_number_of_atoms()): 
            vals.append([q, band_index, amplitude,argument])

        setting['phonon_modes'] = vals 

        phonopy_obj._modulation = Modulation(phonopy_obj._dynamical_matrix,
                                      dimension=setting['dimension'],
                                      phonon_modes=setting['phonon_modes'],
                                      delta_q=None,
                                      derivative_order=None,
                                      nac_q_direction=None,
                                      factor=phonopy_obj._factor)

        run_aims_modulation(phonopy_obj._modulation)

        """Create output_geometry"""
        write_aims_modulation(phonopy_obj._modulation,"output_geometry")
#------------------end for aims_modulation----------------------------


from phonopy.phonon.animation import Animation
def post_process_animation(phonopy_obj, parameters, frequency_unit_factor):
    for anim in parameters:
        anim_obj = Animation(anim["q"],
                             phonopy_obj.dynamical_matrix)
        for file in anim["files"]:
            suffix = os.path.splitext(file)[-1]
            if (suffix == ".ascii"):
                anim_obj.write_v_sim(factor=frequency_unit_factor, filename=file)
            if (suffix == ".arc"):
                anim_obj.write_arc(band_index=anim["band"], amplitude=anim["amp"], 
                                   num_div=anim["div"], filename=file)
            if (suffix == ".xyz_jmol"):
                anim_obj.write_xyz_jmol(factor=frequency_unit_factor, filename=file)
            if (suffix == ".xyz"):
                anim_obj.write_xyz(band_index=anim["band"], amplitude=anim["amp"], 
                                   num_div=anim["div"], 
                                   factor=frequency_unit_factor, filename=file)

from phonopy.units import VaspToEv as AimsToEv
from phonopy.units import VaspToCm as AimsToCm
def write_eigenfrequencies_and_eigenvectors_at_gamma(phonopy_obj):
    # Only Gamma point is of interest
    q = numpy.zeros(3)


    # Check for unit cell == supercell
    # FIXME: CC an unfolding of the gamma frequencies
    #        and eigenvectors might be added later on 
    n_atoms =  phonopy_obj.unitcell.get_number_of_atoms() 
    if  n_atoms !=  phonopy_obj.supercell.get_number_of_atoms():
	print " * WARNING:"
	print " * The unitcell and the supercell are not identical."

    # Open file and write reference equilibrium geometry
    f = open("phonopy-FHI-aims-eigen-frequencies-and-vectors-at-Gamma.dat", 'w')
    f.write("### Eigenfrequencies and Eigenvectors at Gamma from phonopy-FHI-aims \n")
    f.write("# \n")
    f.write("# Equilibrium geometry\n")
    cell = phonopy_obj.unitcell
    for i in range(3):
        v = cell.get_cell()[i]
        f.write("lattice_vector %20.10f %20.10f %20.10f   \n" % (v[0],v[1],v[2]))
    for n in range(n_atoms):
        r = cell.get_positions()[n]
        sym = cell.get_chemical_symbols()[n]
        f.write("atom           %20.10f %20.10f %20.10f   %3s   %4d \n" % (r[0],r[1],r[2],sym,n+1)) 

    # Calculate Eigenvalues and Eigenfrequencies at Gamma
    f.write("\n# \n# Eigenfrequencies at Gamma in cm^-1 (ascending order) \n")


#------------------ shanghui-add for nac_gamma-------------------------------------------
#------------using referenc of phonopy/phonon/band_structure.py (set_band)------------------        
    q_direction= numpy.array([0.5,0.5,0.5]) 
 
    if requested_nac: # at gamma point, you need give a q_dirction, then the nac can be added.
       phonopy_obj.dynamical_matrix.set_dynamical_matrix(q,q_direction)
       dm = phonopy_obj.dynamical_matrix.get_dynamical_matrix()
    else :
       phonopy_obj.dynamical_matrix.set_dynamical_matrix(q)
       dm = phonopy_obj.dynamical_matrix.get_dynamical_matrix()

    frequencies = []
    eigfreq, eigvec =  numpy.linalg.eigh(dm)
#--------------------shanghui-end-add for nac_gamma----------------------------------------------


    # Write eigenvalues to file
    for eig in eigfreq:
            if eig < 0:
    		# Safety check for imaginary frequencies:
		neg_value = -numpy.sqrt(-eig)*AimsToEv*1000
		if neg_value < -0.1:
	           print "* WARNING:"  
	           print "* Imaginary Eigenfrequencies at the Gamma point are highly unphysical."
                frequencies.append(neg_value)
            else:
                frequencies.append(numpy.sqrt(eig)*AimsToCm)  #AimsToEv*1000)
    for i, freq in enumerate(frequencies):
        f.write("frequency %20.10f \t %4d \n" %  ( freq,i+1))

    # Write eigenvectors to file
    f.write("\n#\n# Eigenvectors at Gamma (orthonormalized) \n")
    for i, freq in enumerate(frequencies):
        eigvec_band_i = eigvec[:,i]
    	# Safety check for imaginary eigenvector components:
	if (abs(eigvec_band_i.imag) > 1E-7).any():
	   print "* WARNING:"
	   print "* Imaginary Eigenvectors at the Gamma point are highly unphysical."
			
    	f.write("#\n# Mode %4d  - Atoms 1 to %4d  - Eigenfrequency  %20.10f\n" % (i+1,n_atoms,frequencies[i]))
    	for n in range(n_atoms):
    		f.write("eigenvector %4d %4d  %20.15f %20.15f %20.15f \n" \
		% (i+1,n+1,eigvec_band_i[3*n],eigvec_band_i[3*n+1],eigvec_band_i[3*n+2]))
    f.close()

   

if __name__ == "__main__":

    # suppressing warnings, particularly from imported packages, to keep the output clean
    # -> comment out for debugging purposes!
    import warnings
    warnings.simplefilter('ignore')

    import sys
    import os
    import shutil
    import math
    import numpy

    from optparse import OptionParser

    from phonopy import Phonopy
#    from ase.io.aims import read_aims, write_aims, read_aims_output
    from phonopy.interface.FHIaims import read_aims, write_aims, read_aims_output
    from phonopy.phonon.mesh import Mesh

    print "### phonopy wrapper for FHI-aims ###"
    version = "20111113"
    print "# version " + version
    print "#"

    has_matplotlib = False
    try:
        import matplotlib
        matplotlib.use("PDF")
        import matplotlib.pyplot as plt
        has_matplotlib = True
    except ImportError:
        print "# matplotlib is not available"
        print "#"
    
    usage = "usage: %prog [options] [arguments are ignored] \n" + \
            "       run in directory with control.in and geometry.in files"
    parser = OptionParser(version=version, usage=usage)
    
    parser.add_option("-d", "--data-only",
                      action="store_true", dest="data", default=False,
                      help="only write .dat files (no plots even if matplotlib is available)")
    parser.add_option("-e", "--eigenvectors",
                      action="store_true", dest="eigenvectors", default=False,
                      help="also calculate Eigenvectors for every q-point occuring in current calculation")
    parser.add_option("-g", "--no-greek-labels",
                      action="store_false", dest="greek_labels", default=True,
                      help="turn off replacement of Greek letters for labels in band structure plots")
    parser.add_option("-s", "--no-symmetry",
                      action="store_false", dest="symmetry", default=True,
                      help="turn off symmetry wherever possible (at the moment only for q-point grid)")
    parser.add_option("-G", "--EigVecGamma",
                      action="store_true", dest="EigVecGamma", default=False,
                      help="output Eigenvalues and Eigenvectors at the Gamma point")
    parser.add_option("-y", "--yaml",
                      action="store_true", dest="yaml", default=False,
                      help="write .yaml data files for 'native' phonopy post-processing tools")
    (options, args) = parser.parse_args()
    do_matplotlib = has_matplotlib and not options.data
    
    print "# reading control.in"
    control = Control()
    control.print_phonon("# |   ")
    if (len(control.phonon["supercell"]) == 3):
        supercell_matrix = numpy.diag(control.phonon["supercell"])
    elif (len(control.phonon["supercell"]) == 9):
        supercell_matrix = numpy.array(control.phonon["supercell"]).reshape(3,3)
    supercell_matrix_inv = numpy.linalg.inv(supercell_matrix)
    print "# +-> specified supercell has %.1f times the volume of the initial cell" % numpy.linalg.det(supercell_matrix)
    print "#"

    frequency_unit = control.phonon["frequency_unit"]
    frequency_unit_factor = AimsFrequencyUnitFactors[frequency_unit]
    frequency_unit_to_THz = AimsFrequencyUnitFactorsToTHz[frequency_unit]

    print "# reading geometry.in"
    cell = read_aims("geometry.in")
    if cell.get_magnetic_moments() is not None:
        print "# | initial_moments found and about to be replicated in supercells"
        print "# | WARNING: No consideration in symmetry analysis, i.e."
        print "# |          if initial spin breaks the spatial symmetry,"
        print "# |          this is currently NOT detected."
        print "# |          You thus might want to consider the"
        print "# |          command line option to turn off symmetry."
    print "# generating supercells with displacements"
    phonopy_obj = Phonopy(cell, supercell_matrix, 
                          symprec=control.phonon["symmetry_thresh"],
                          factor=frequency_unit_factor)
    phonopy_obj.generate_displacements(distance=control.phonon["displacement"])
    print "# | Spacegroup: %s" % phonopy_obj.symmetry.get_international_table()
    supercells = phonopy_obj.get_supercells_with_displacements()
    digits = int( math.ceil( math.log(len(supercells)+1,10) ) ) + 1
    directories = []
    for i in range(len(supercells)):
        directories.append( ("phonopy-FHI-aims-displacement-%0" + str(digits) + "d") % (i+1))
    print "#"

    if not any(map(os.path.isdir, directories)):
        for (supercell, directory) in zip(supercells, directories):
            os.mkdir(directory)
            write_aims(os.path.join(directory, "geometry.in"), supercell)
            shutil.copy("control.in", directory)
        print "!   Please calculate forces with FHI-aims for the (supercell) structures"
        print "!   which have been generated in the subdirectories"
        for directory in directories:
            print "! \t %s" % directory
        print "!   redirecting the outputs into"
        print "!	<directory name>.out"
        print "!   in each directory." 
        quit()

    print "# reading forces from"
    set_of_forces = []
    for (directory, supercell) in zip(directories, supercells):
        aims_out = os.path.join(directory, directory + ".out")
        print "# | %s" % aims_out
        if not os.path.isfile(aims_out):
            print "!!! file not found: %s" % aims_out
            sys.exit(1)
        # FHI-aims calculation is supposed to be a single point - 
        # hence only use first structure contained in aims output
        supercell_calculated = read_aims_output(aims_out)
        # sanity check to verify that calculated structures from which forces are read in 
        # are those which are expected
        # cannot use builtin __eq__ of atoms object here because 
        #    - numerical tolerances are required when comparing read in positions and cell vectors
        #    - FHI-aims does not support other than 0 or 3 periodic directions
        tol = 1.0e-6
        if ( (supercell_calculated.get_number_of_atoms() == supercell.get_number_of_atoms()) and
             (supercell_calculated.get_atomic_numbers() == supercell.get_atomic_numbers()).all() and
             (abs(supercell_calculated.get_positions()-supercell.get_positions()) < tol).all() and
             (abs(supercell_calculated.get_cell()-supercell.get_cell()) < tol).all() ):
            # read_aims_output reads in forces from FHI-aims output as list structure, 
            # but further processing below requires numpy array
            forces = numpy.array(supercell_calculated.get_forces())
            drift_force = forces.sum(axis=0)
            print "#   | correcting drift : %11.5f %11.5f %11.5f" % tuple(drift_force)
            for force in forces:
                force -= drift_force / forces.shape[0]
            set_of_forces.append(forces)
        else:
            print "!!! calculated varies from expected supercell in FHI-aims output %s" % aims_out
            sys.exit(2)
    print "#"

    if "phono-perl" in control.phonon["hessian"]:
        print "# writing Hessian matrix (same format as Perl script)"
        write_Hessian(phonopy_obj, supercell_matrix_inv, set_of_forces)
        print "#"
    if "TDI" in control.phonon["hessian"]:
        print "# writing force constants"
        write_force_constants(phonopy_obj, set_of_forces)
        print "#"

    requested_nac = ( control.phonon["nac"] != {} )
    # Since phonopy-0.9.0, first argument only denotes axes of primitive cell relative to input unit cell 
    # (always unit matrix here), i.e. supercell_matrix inversion (as required for Primitive()) is done 
    # in Phonopy.set_post_process() directly.
    print "# preparing calculation of dynamical matrix"
    # phonopy_obj.set_post_process(numpy.eye(3), set_of_forces)
    phonopy_obj.set_forces(set_of_forces)
    phonopy_obj.produce_force_constants()
    if requested_nac:
        from phonopy.hphonopy.file_IO import parse_BORN
        nac_file = control.phonon["nac"]["file"]
        nac_method = control.phonon["nac"]["method"]
        print "# | reading polarizabilities from %s" % nac_file
        nac_params = parse_BORN(phonopy_obj.primitive, filename=nac_file)
        # print "# | setting method for non-analytical terms in dynamical matrix to \'%s\'" % nac_method
        phonopy_obj.set_nac_params(nac_params)
        print "#"

    if requested_nac:
        print "# phonon frequencies at Gamma + %s:" % str(control.phonon["nac"]["delta"])
        q = numpy.zeros(3) + numpy.array(control.phonon["nac"]["delta"])
    else:
        print "# phonon frequencies at Gamma:"
        q = numpy.zeros(3)
    # Since phonopy-0.9.0, Phonopy.get_frequencies already returns frequencies in units 
    # according to (unit) factor relayed at Phonopy object instantiation.
    for i, freq in enumerate(phonopy_obj.get_frequencies(q)):
        print "# |  %3d: %10.5f %s" %  (i+1, freq, frequency_unit)
    print "#"

    if (control.phonon["band"] != []):
        print "# post-processing band structure"
        post_process_band(phonopy_obj, control.phonon["band"], frequency_unit_factor, 
                          is_eigenvectors=options.eigenvectors, write_yaml=options.yaml, 
                          do_matplotlib=do_matplotlib, lookup_labels=options.greek_labels)
        print "#"
    
    if (control.phonon["dos"] != {}) or (control.phonon["free_energy"] != {}):
        meshes = []
        if "qdensity" in control.phonon["dos"]:
            qdensity_dos = control.phonon["dos"]["qdensity"]
            if len(qdensity_dos) >= 3:
                mesh_dos = qdensity_dos[0:3] 
            else: 
                mesh_dos = [ qdensity_dos[0] ] * 3
            meshes.append(mesh_dos)
        if "qdensity" in control.phonon["free_energy"]:
            qdensity_free_energy = control.phonon["free_energy"]["qdensity"]
            if len(qdensity_free_energy) >= 3:
                mesh_free_energy = qdensity_free_energy[0:3] 
            else: 
                mesh_free_energy = [ qdensity_free_energy[0] ] * 3
            meshes.append(mesh_free_energy)
        mesh = sorted(meshes, reverse=True)[0]
        print "# creating q-points grid " + str(mesh)
        print "#"
        mesh_shift = [0,0,0]
        mesh_time_reversal = options.symmetry
        mesh_symmetry = options.symmetry
        mesh_eigenvectors = options.eigenvectors
        mesh_obj = Mesh(phonopy_obj.dynamical_matrix,
                        mesh, shift=mesh_shift,
                        is_time_reversal=mesh_time_reversal,
                        is_mesh_symmetry=mesh_symmetry,
                        is_eigenvectors=mesh_eigenvectors,
                        factor=frequency_unit_factor)
        if options.yaml:
            mesh_obj.write_yaml()
                        
        if (len(control.phonon["dos"]) > 0):
            print "# post-processing density of states"
            post_process_dos(phonopy_obj, mesh_obj, control.phonon["dos"], do_matplotlib=do_matplotlib)
            print "#"
        if (len(control.phonon["free_energy"]) > 0):
            print "# post-processing free energy"
            post_process_free_energy(phonopy_obj, mesh_obj, control.phonon["free_energy"], 
                                     frequency_unit_to_THz, write_yaml=options.yaml, do_matplotlib=do_matplotlib)
            print "#"

    if (control.phonon["animation"] != []):
        print "# post-processing animation output"
        post_process_animation(phonopy_obj, control.phonon["animation"], frequency_unit_factor)
        print "#"

    
    if (control.phonon["modulations"] != {}):
        print "# post-processing modulations output"
        post_process_write_modulations(phonopy_obj,control.phonon["modulations"])
        print "#"

    if options.EigVecGamma:
        write_eigenfrequencies_and_eigenvectors_at_gamma(phonopy_obj)
    print "# finished successfully"
    
