def drawEnergy(ax, relative = False, divided = 1.0, popFirstStep = True, label = ''):
  """
  Draw the KS Energy and Total Energy in the ax
  """
  X, temp, E_KS, Etot, volume, pressure  =  tdp.getEnergyTemperaturePressure()
  
  refEnergy=E_KS[0]
  if relative:
    E_KS = E_KS - refEnergy
    Etot = Etot - refEnergy
  
  if popFirstStep:
    ax.plot(X[1:],E_KS[1:]/divided,'-',linewidth=3.0,label = 'KS'+ label)#,color=ma.colors[0]
    ax.plot(X[1:],Etot[1:]/divided,'--',linewidth=3.0,label = 'Total'+ label)#,color=ma.colors[2]
  else:
    ax.plot(X,E_KS/divided,'-',linewidth=3.0,label = 'KS' + label)#,color=ma.colors[0]
    ax.plot(X,Etot/divided,'--',linewidth=3.0,label = 'Total' + label)#,color=ma.colors[2]
  
  kargs = ma.getPropertyFromPosition(index=None, 
                                         xlabel=r'Time(fs)',ylabel=r'Energy(eV)',title='', 
                                         xticks=None, yticks=None, 
                                         xticklabels=None, yticklabels=None,
                                         xlimits=None, ylimits=None)
                                 
  ma.setProperty(ax, **kargs)

def case(indexCase, folderCase):
  data = scanFolder(getTime)
  selectedProcesses = ['evolve','DHSCF','PostSCF','siesta'] #,'TDbuildD'
  explainations = ['Propagation', 'Build H', 'Postprocess','Total'] #,'Build DM'
  image = data[-1][-1]['Prg.tot']
  pieData = image[selectedProcesses[:-1]]
  other = image[selectedProcesses[-1]] - np.sum(pieData.values)
  
  ax = axsall[1][indexCase]
  explode = np.zeros(len(pieData)) + 0.04
  explode[0]  = 0.07
  patches, texts, autotexts = ax.pie(pieData.values, explode=explode, labels=explainations[:-1], 
                                     labeldistance = 0.2, autopct='%.0f%%',pctdistance = 0.7,
                                     shadow=True, textprops={'fontsize':'xx-large','color':'gold'})
  for text in texts:
    text.set_fontsize(0) 
    
  kargs=ma.getPropertyFromPosition(indexCase + 3, title='')
  ma.setProperty(ax,**kargs) 
  
  for key,label in zip(selectedProcesses,explainations):
    totalTimes = np.array([time['Prg.tot'][key] for index, folder, time in data])/60.0
    situations = np.array([int(folder) for index, folder, time in data])
    ax = axsall[0][indexCase]
    ax.semilogy(situations, totalTimes,'-o', mew = 3, alpha=0.8 ,ms=12, label = label)
#    from scipy.optimize import curve_fit  
#    popt, pcov = curve_fit(errorFunc, situations, totalTimes)    
#    xfit = np.linspace(2,25,1000)
#    ax.semilogy(xfit, errorFunc(xfit,*popt),'--',lw=3, )
  ax.grid(which='major',axis=u'both')
  kargs=ma.getPropertyFromPosition(indexCase, ylabel=r'Clock Time (Mimute)',title=folderCase,
                                   #xlimits = [0,10]
                                   )
  ma.setProperty(ax,**kargs)

def plotExcitation(ax, label=''):
  time, exe = dp.getExcitedElectrons()  
  ax.plot(time, exe - exe[0],'-', lw=2)
  kargs=ma.getPropertyFromPosition(ylabel=r'n(e)', xlabel='Time (fs)',
                                   title='Excited Electrons')

  print exe[-1] - exe[0]
  ma.setProperty(ax,**kargs)

def plotRMSD(ax, label=''):
  dp.getTrajactory()
  import pyramids.process.struct as pps
  time, distance = pps.calculateRMSD()
  ax.plot(time, distance, lw=2, label=label)
  kargs=ma.getPropertyFromPosition(xlabel='Time (fs)', ylabel=r'$\langle u \rangle^\frac{1}{2}$ ($\AA$)', 
                                   title='RMSD')
  ma.setProperty(ax,**kargs)  

def plotEField(ax, label=''):
  time, Efield = dp.getEField()
  directions = ['x', 'y', 'z']
  for direct in range(3):
    if max(Efield[:,direct]) > 1E-10:
      ax.plot(time,Efield[:,direct],
              label=directions[direct], lw=2, alpha=1.0) 
  kargs=ma.getPropertyFromPosition(ylabel=r'$\varepsilon$ (a.u.)',xlabel='Time(fs)',
                                   title='Electric Field')
  ma.setProperty(ax,**kargs)

def plotKpoints(kpts):
  #Plot the K points
  from pyramids.plot.setting import setProperty, getPropertyFromPosition
  fig, ax = plt.subplots(1,1)
  ax.plot(kpts[:,0],kpts[:,1],'o')
  setProperty(ax,**getPropertyFromPosition(title='K points',xlabel=r'$k_x(1/\AA)$',ylabel=r'$k_y(1/\AA)$'))
  plt.axis('equal')
  plt.tight_layout()
  plt.savefig('Kpoints.pdf',dpi=600)
  plt.show()
  #plt.close()
  
#unit = reciprocal_vectors/grid
#KIndex = np.dot(kpts,np.linalg.inv(unit)) + [nkx/2, nky/2, nkz/2]
#kIndex = np.array(KIndex, dtype=int)
#skpts = np.array([kpts + np.dot([i,j,k],reciprocal_vectors)
#         for i in [0,-1,1] for j in [0,-1,1] for k in [0,-1,1]])
#skIndex = np.array(np.dot(skpts,np.linalg.inv(unit)) + [0.5,0.5,0], dtype=int)            
#print kIndex.shape, #skIndex.shape
#print kIndex

#Ef = 6.0
#T = 300
#qOrder = 2     
#susp = np.zeros([qOrder,qOrder,qOrder],dtype=float) 
#
#
#for ik, ikpt in enumerate(kpts):
#  kxy = kIndex[ik] 
#  for i in range(-qOrder,qOrder):
#    for j in range(-qOrder,qOrder):
#      for k in range(-qOrder,qOrder):
#        qxy = (kxy + np.array([i,j,k])) % np.array([nkx,nky,nkz])
#        iq = qxy[0]*nky*nkz + qxy[1]*nkz + qxy[2]
#        #susp[i,j,k] = susFunc(eigValueAllK,ik,iq,Ef,T)
#        pass
#print susp

#if __name__ == "__main__":
#  nkx = 12
#  nky = 12
#  nkz = 1
#  # Read the structure information from POSCAR
#  from ase.io import read
#  atoms = read('POSCAR')  
#  # Generate K points, 
#  # Selection: Monkhorst-Pack or Line-Mode
#  grid = np.array([nkx,nky,nkz],dtype=int)
#  kpts = getMPKpts(atoms,grid)
#  #kpts = getBandKpoints(atoms,npoints=50)
#  x, v, u = calculateEigenPairs(kpts)
#  plotBands(x, v)
#  plotKpoints(kpts)

def drawPartition(ax, label = ''):
  """
  Draw the KS Energy and Total Energy in the ax
  """
  X,qo  =  tdp.getPartition()
  ax.plot(X,qo,'-',linewidth=3.0,label = 'Partitions '+ label)#,color=ma.colors[0]
  kargs = ma.getPropertyFromPosition(index=None, 
                                         xlabel=r'Time(fs)',ylabel=r'Partition',title='', 
                                         xticks=None, yticks=None, 
                                         xticklabels=None, yticklabels=None,
                                         xlimits=None, ylimits=None)               
  ma.setProperty(ax, **kargs) 

def plotBands(xall,eigValueAllK):
  from pyramids.plot.setting import setProperty, getPropertyFromPosition
  fig, ax = plt.subplots(1,1)
  #Plot the eigenvalues 
  for index in range(eigValueAllK.shape[0]):
    ax.plot(xall[index,:],eigValueAllK[index,:],'ob',lw=1)
  setProperty(ax,**getPropertyFromPosition(title='Energy Bands',xlabel=r'',ylabel=r'Energy(eV)',xticklabels=[]))
  #Set tight layout and output 
  plt.axis('tight')
  plt.tight_layout()
  plt.savefig('Bands.pdf',dpi=600)
  plt.show()

def drawPressure(ax, label = ''):
  """
  Draw the KS Energy and Total Energy in the ax
  """
  X, temp, E_KS, Etot, volume, pressure  =  tdp.getEnergyTemperaturePressure()
  ax.plot(X,pressure,'-',linewidth=3.0,label = 'Pressure '+ label)#,color=ma.colors[0]

  kargs = ma.getPropertyFromPosition(index=None, 
                                         xlabel=r'Time(fs)',ylabel=r'Pressure(KBar)',title='', 
                                         xticks=None, yticks=None, 
                                         xticklabels=None, yticklabels=None,
                                         xlimits=None, ylimits=None)               
  ma.setProperty(ax, **kargs)

def plotDOS(ax, step, ylimits=[0,None], bins=100, **kargs):
  xlabel = 'Energy (eV)'
  ylabel = 'DOS'
  
  eDosInterp, yDosInterp, yParInterp = dp.calculateDOS(step,bins=bins)
  ax.fill_between(eDosInterp, yParInterp, color='b')
  ax.fill_between(eDosInterp,-yParInterp, color='r')
  ax.fill_between(eDosInterp, yDosInterp, lw=3, color='g',alpha=0.2)
  
  kargs = ma.getPropertyFromPosition(xlabel=xlabel, ylabel=ylabel, 
                                     ylimits=ylimits, yticklabels=[], 
                                     **kargs)
  ma.setProperty(ax,**kargs)

def animate(i): 
  i = i % len(dataAndAtoms)
  z = (dataAndAtoms[i][0]-pho0)[slide,:,:] 
  atoms = dataAndAtoms[i][1]
  print z.min(), z.max()
  axs[0].contourf(x, y, z, 100,vmin = vmin,vmax =vmax, cmap = cm.jet)
  for pos, Z in zip(atoms.positions, atoms.numbers):
    axs[0].scatter(pos[2],pos[1],c=tuple(cpk_colors[Z]),s=400*covalent_radii[Z])
  setProperty(axs[0],**getPropertyFromPosition(0,xlimits=[2,6],ylimits=[2,6],
              xlabel = r'distance($\AA$)',ylabel = r'distance($\AA$)')) 
  scatter.set_data(time[selectedTimeSteps[i]],light[selectedTimeSteps[i],2])
  setProperty(axs[1],**getPropertyFromPosition(1,
              xlabel = r'time(fs)',ylabel = r'E-Field(a.u.)'))
  plt.tight_layout()

def drawRDF(ax,step=None, label = ''):
  """
  Draw the KS Energy and Total Energy in the ax
  """
  if step == None:
    step = tdp.getNumStep()
  
  systemLabel,timestep = tdp.getSystemLabelTimpstep()
  tdp.splitMDCAR()
  hist,bin_edges = tdp.calculateRDF(step-1)
  ax.plot(bin_edges[:-1],hist,linewidth=3.0,label='RDF '+str((step)*timestep)+' fs'+ label) #,color=ma.colors[0]
  kargs = ma.getPropertyFromPosition(index=None, 
                                         xlabel=r'Radius($\AA$)',ylabel=r'RDF(a.u.)',title='', 
                                         xticks=None, yticks=None, 
                                         xticklabels=None, yticklabels=None,
                                         xlimits=None, ylimits=None)           
  ma.setProperty(ax, **kargs) 

def drawRMSD(ax,selected=None, label = '', init = 0):
  """
  Draw the KS Energy and Total Energy in the ax
  """
  systemLabel,timestep = tdp.getSystemLabelTimpstep()
  if selected == None:
    selected = range(0,tdp.getNumStep(),int(1.0/timestep))
    
  tdp.splitMDCAR()
  time,distance,velocity = tdp.calculateRMSD(selected, init=init)
  
  #os.remove('POSCAR')
  ax.plot(time,distance,linewidth=3.0,label='RMSD'+ label)#,color=ma.colors[0])
  kargs = ma.getPropertyFromPosition(index=None, 
                                         xlabel=r'Time(fs)',ylabel=r'RMSD($\AA$)',title='', 
                                         xticks=None, yticks=None, 
                                         xticklabels=None, yticklabels=None,
                                         xlimits=None, ylimits=None)           
  ma.setProperty(ax, **kargs) 

def plotDistribution(ax, step, bins=100, intepNum=2000, ylimits=[0,1], **kargs):
  xlabel = 'Energy (eV)'
  ylabel = 'FD'
  yticklabels=[]
  
  eDos, dos, par = dp.calculateDOS(step, ref=False, interp=False,bins=bins)
  distribution = ((par))/((dos)+1E-10)
  
  def interp(xin,yin,xout):
    from scipy.interpolate import interp1d
    spline = interp1d(xin, yin, kind='cubic')
    return spline(xout) 
  
  x = np.linspace(eDos[0], eDos[-1], intepNum)
  y = interp(eDos, distribution, x)
  
  #print (np.abs(yDosInterp)+0.001).min()
  ax.fill_between(x, y, color='b')
  #print yDosInterp.min()
  kargs = ma.getPropertyFromPosition(xlabel=xlabel, ylabel=ylabel, ylimits=ylimits, **kargs)
  ma.setProperty(ax,**kargs)

def action(index,folder):
  ls = ['-','-','-','-']
  ax = axs[0]
  time, Efield = dP.getEField()
  #directions = ['x', 'y', 'z']
  for direct in range(3):
    if max(Efield[:,direct]) > 1E-10:
      ax.plot(time,Efield[:,direct],
              label=folder,lw=2,alpha=0.5) 
  kargs=ma.getPropertyFromPosition(ylabel=r'$\varepsilon$(a.u.)',xlabel='Time(fs)',
                                   title='Electric Field')
  ma.setProperty(ax,**kargs)
  ax.ticklabel_format(style='sci',axis='y',scilimits=[0,0])
  #------------------------------------------------------------------------------
  ax = axs[1]
  time, T, E_ks, E_tot, Vol, P  = dP.getEnergyTemperaturePressure(ave=True)
#  for i in range(2,E_ks.shape[0]-1):
#    if E_ks[i+1] - (E_ks[i] + E_ks[i-1])*0.5 > 2.0:
#      E_ks[i+1] = (E_ks[i] + E_ks[i-1])*0.5
  ax.plot(time[2:], E_ks[2:] - E_ks[2],'-', lw=1, alpha=1, label=folder)
  kargs=ma.getPropertyFromPosition(ylabel=r'E(eV)',title='Excitation Energy')
  ma.setProperty(ax,**kargs)

def drawEigenvalue(ax, label = ''):
  """
  Draw the KS Energy and Total Energy in the ax
  """
  X,eo  =  tdp.getEigenvalues()
  X,qo  =  tdp.getPartition()

  countEleState = 0 
  countHolState = 0 
  countNormalState = 0
  for i in range(eo.shape[1]):
    if eo[0,i] > 0.03 and qo[0,i] > 0.2:
      if countEleState == 0:    
        ax.plot(X,eo[:,i],'-',linewidth=3.0,color=ma.colors[0],label='Electron'+ label)
      else:
        ax.plot(X,eo[:,i],'-',linewidth=3.0,color=ma.colors[0])
      countEleState += 1
    elif eo[0,i] < -0.03 and qo[0,i] < 1.8:
      if countHolState == 0:    
        ax.plot(X,eo[:,i],'-',linewidth=3.0,color=ma.colors[1],label='Hole'+ label)
      else:
        ax.plot(X,eo[:,i],'-',linewidth=3.0,color=ma.colors[1])
      countHolState +=1
    else:
      if countNormalState == 0:    
        ax.plot(X,eo[:,i],'--',linewidth=3.0,color='grey',alpha=0.7,label='Normal'+ label)
      else:
        ax.plot(X,eo[:,i],'--',linewidth=3.0,color='grey',alpha=0.7)
      countNormalState +=1

  #ax.plot(X,eo,'-',linewidth=qo,label = 'Eigenvalues')#,color=ma.colors[0]
  kargs = ma.getPropertyFromPosition(index=None, 
                                         xlabel=r'Time(fs)',ylabel=r'Eigenvalues(eV)',title='', 
                                         xticks=None, yticks=None, 
                                         xticklabels=None, yticklabels=None,
                                         xlimits=None, ylimits=None)               
  ma.setProperty(ax, **kargs)

def drawBands(ax,bandType='e', label = ''):
  """
  Draw the KS Energy and Total Energy in the ax
  """
  X, Ek, xticks, xticklabels = tdp.getBands()
  if bandType == 'e':
   eFermi = tdp.getFermiEnergy()
   Ek = Ek - eFermi
   ylabel = 'Energy(eV)'
  else:
   ylabel = r'$\omega$(cm$^{-1}$)'
 
  alpha = 1.0
  color = 'b'
  for ispin in range(Ek.shape[2]):
    for iband in range(1,Ek.shape[1]):
      ax[0].plot(X,Ek[:,iband,ispin],lw=3.0,ls='-',color=color,alpha=alpha,label=label)
  
  kargs = ma.getPropertyFromPosition(index=None, 
                                         xlabel=r'',ylabel=ylabel,title='', 
                                         xticks=xticks, yticks=None, 
                                         xticklabels=xticklabels, yticklabels=None,
                                         xlimits=None, ylimits=None)           
  ma.setProperty(ax, **kargs) 

fig=plt.figure(figsize=A4_LANDSCAPE)#_LANDSCAPE
plt.subplots_adjust(left=0.1, bottom=0.10, right=0.95, top=0.95, wspace=0.3, hspace=0.05)

axs = [fig.add_subplot(2,1,1),
      fig.add_subplot(2,1,2)]

colors = getColors(4)
chargeData=np.loadtxt('AECHARGE')
kargs=getPropertyFromPosition(1,'Charge(e)',r'r($\AA$)')
axs[0].plot(chargeData[:,0],chargeData[:,1],':',linewidth=3,label='Down',color=colors[1])
axs[0].plot(chargeData[:,0],chargeData[:,2],'--',linewidth=3,label='Up',color=colors[2])
axs[0].plot(chargeData[:,0],chargeData[:,3],'-',linewidth=3,label='Core',color=colors[0])
axs[0].legend(fontsize=16,loc=1)
kargs['xlimits'] = [0,6]
kargs['xticklabels'] = []
setProperty(axs[0],**kargs)

numWFfiles=int(os.popen('ls AEWFNR* |wc -l').readline())

print numWFfiles
waveData = []
lineTypes=['-','--',':']
orbitalType=['s','p','d','f']
for indexFile in range(min(numWFfiles,3)):
  waveData=np.loadtxt('AEWFNR'+str(indexFile))
  axs[1].plot(waveData[:,0],waveData[:,1],lineTypes[indexFile],linewidth=3,label=orbitalType[indexFile],color=colors[indexFile])


kargs=getPropertyFromPosition(2,'Wavefunction',r'r($\AA$)')
axs[1].legend(fontsize=16,loc=1)
kargs['xlimits'] = [0,15]

time, msd = dp.readMSD()
nocc = int(float(os.popen('grep "starting charge" result').readline().split()[-1])/2.0)

fig, axs = plt.subplots(2,2,sharex=True,figsize=(8,5))
axs = axs.flatten()
norm, kweight = readData('pwscf.norm.dat')
excite = (norm - norm[0,:,:])
for ib in range(excite.shape[2]):
    excite[:,:,ib] *= kweight
    #pass
step = min(time.shape[0]-1, excite.shape[0])
axs[0].plot(time[1:step+1], (excite[:step,:,:nocc]).sum(axis=(1,2)))
axs[0].plot(time[1:step+1], (excite[:step,:,nocc:]).sum(axis=(1,2)))
args = ma.getPropertyFromPosition(0,title='Excited electrons',ylabel='n (e)')
ma.setProperty(axs[0],**args)
Efield = [[float(i) for i in line.split()] for line in open('TDEFIELD')]
Efield = np.array(Efield)/1E5

ppu.plotTotalEnergy(axs[2])
ppu.plotEField(axs[1])


f = os.popen('grep "first current is " result')
current = np.array([[float(i) for i in line.split()[-3:]] for line in f.readlines()])
for idir in range(3):  
    axs[3].plot(time[1:], current[:,idir],'-',label=['x','y','z'][idir])

args = ma.getPropertyFromPosition(3,title='Current',ylabel='j (a.u.)')
ma.setProperty(axs[3],**args)

maxEnergies = []
minEnergies = []

for [(index, folder), (timeEn,deltaE, T)] in data:
  eField = float(folder) * 13.6/0.529
  ax = axs[exTemp]        
  ax.plot(timeEn, T,'-', c=c[index], 
          lw=2, alpha=1)
          
  ax = axs[exEnergy]
  ax.plot(timeEn, deltaE,'-', c=c[index],
          label=r'%5.2f $V/\AA$' % eField,
          lw=2, alpha=1)
          

  
kargs=ma.getPropertyFromPosition(exEnergy, ylabel=r'E(eV)',
                                 title='Excitation Energy', 
                                 xlabel='Time (fs)',
                                 #ylimits=[0,np.max(maxElectrons)],
                                 xlimits=None,)
ma.setProperty(axs[exEnergy],**kargs)
kargs=ma.getPropertyFromPosition(exTemp, ylabel=r'T(K)',                     
                                 #ylimits=[np.min(minEnergies),np.max(maxEnergies)],
                                 title='Temperature')
ma.setProperty(axs[exTemp],**kargs)

plt.tight_layout()
for save_type in ['.pdf','.png']:
  filename = SaveName + save_type
  plt.savefig(filename,dpi=800)