# Copyright (C) 2015 Henrique Pereira Coutada Miranda, Alejandro Molina Sanchez
# All rights reserved.
#
# This file is part of yambopy
#
#
import os
import re
from math import sqrt
from numpy import array
from   qepy.auxiliary import *

meVtocm = 8.06573

class DynmatIn():
    """ A class to generate an manipulate quantum espresso input files for matdyn.x
    """
    def __init__(self):
        self.variable = dict()
        self.qpoints     = [] 

    def write(self,filename):
        f = open(filename,'w')
        f.write(str(self))
        f.close()

    def __str__(self):
        s = '&input\n'
        s += self.stringify_group('',self.variable) #print variable
        if len(self.qpoints) > 0:
          s+=("%d\n"%len(self.qpoints))
          for q in self.qpoints:
            s+=("%12.8lf %12.8lf %12.8lf")%tuple(q[:3])+"\n"
        return s

    def __setitem__(self,key,value):
        self.variable[key] = value

    def __getitem__(self,key):
        return self.variable[key]

    def stringify_group(self, keyword, group):
        if group != {}:
            string='\n'
            for keyword in group:
                string += "%20s = %s\n" % (keyword, group[keyword])
            string += "/\n"
            return string
        else:
            return ''


class Matdyn():
    """ Class to read and plot the data from matdyn.modes files 
    """
    _datafile = 'matdyn.modes'
    def __init__(self,natoms,path,folder='.'):
        self.folder   = folder
        self.path     = path
        data_phon     = open("%s/%s"%(folder, self._datafile),'r').readlines()
        self.nmodes   = natoms*3
        self.nqpoints = len(path.get_klist()) 
        self.eigen, self.modes = [], []

        for j in xrange(self.nqpoints):
          frec, v_frec = [], []
          for i in xrange(self.nmodes):
            k=4 + j*(self.nmodes*(natoms+1)+5) + i*(natoms+1)
            y = float_from_string(data_phon[k])
            v_mode = []
            for ii in xrange(1,natoms+1):
              z      = float_from_string(data_phon[k+ii])
              v_atom = array([complex(z[0],z[1]),complex(z[2],z[3]),complex(z[4],z[5])])
              v_mode.append(v_atom)
            v_frec.append(array(v_mode))
            frec.append(y[1])
          self.eigen.append(frec)
          self.modes.append(array(v_frec))

    def plot_eigen(self,path=[]):
        """ plot the phonon frequencies using matplotlib
        """
        import matplotlib.pyplot as plt

        if self.eigen is None:
            print('Error')
            #self.get_eigen()
      
        if path:
            if isinstance(path,Path):
                path = path.get_indexes()
            plt.xticks( *zip(*path) )
        plt.ylabel('\\omega (cm$^{-1}$)')

        #plot vertical line
        for point in path:
            x, label = point
            plt.axvline(x)

        #plot bands
        eigen = array(self.eigen)
        for ib in range(self.nmodes):
           plt.plot(xrange(self.nqpoints),eigen[:,ib], 'r-', lw=2)
        plt.show()
  

@@ -1708,16 +1708,42 @@
  %5B%5D, %5B%5D%0A
+        self.qpoints  = %5B%5D
 %0A       
@@ -1799,32 +1799,140 @@
 v_frec = %5B%5D, %5B%5D%0A
+          k=2 + j*(self.nmodes*(natoms+1)+5)%0A          self.qpoints.append(float_from_string(data_phon%5Bk%5D))%0A
           for i 
@@ -3173,8 +3173,526 @@
 ow()%0A  %0A
+    def __str__(self):%0A        s = ''%0A        for nq in xrange(self.nqpoints):%0A            s+=%22%5Cn%5Cn q = %22+(%22%2512.8lf %22*3)%25tuple(self.qpoints%5Bnq%5D)+%22%5Cn%22%0A            for n in xrange(self.nmodes):%0A                s+= %22freq (cm-1): %254.3lf%5Cn%22%25self.eigen%5Bnq%5D%5Bn%5D%0A                for na in xrange(self.nmodes/3):%0A                    xr = self.modes%5Bnq%5D%5Bn%5D%5Bna%5D.real%0A                    xi = self.modes%5Bnq%5D%5Bn%5D%5Bna%5D.imag%0A                    s+=(%22%2512.8lf %2512.8lfj    %22*3)%25(xr%5B0%5D,xi%5B0%5D,xr%5B1%5D,xi%5B1%5D,xr%5B2%5D,xi%5B2%5D)+%22%5Cn%22%0A        return s%0A%0A

print list of phonons dynmat