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 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 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 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 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 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 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 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 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 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 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 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) 

#from ase.calculators.siesta.import_functions import xv_to_atoms
#atoms = xv_to_atoms('siesta.XV')

import os

atoms = dp.getStructrue()

  
# --------------------------------------a----------------------------------   
ax = axs[iStruct]

generateStructPNG(atoms,cell=True,repeat = [5,3,1])
insertImag(ax)
#confStructrueFigure(ax)

kargs = getPropertyFromPosition(iStruct,r'',"",title='Structure',)
setProperty(ax,**kargs)     
# --------------------------------------b----------------------------------   
ax = axs[iBZone]
reciprocal_vectors = 2*np.pi*atoms.get_reciprocal_cell()
#print reciprocal_vectors
points=np.array([(reciprocal_vectors[0,0:2]*i+
                  reciprocal_vectors[1,0:2]*j+
                  reciprocal_vectors[2,0:2]*k) 
                  for i in range(-1,2) 
                  for j in range(-1,2)
                  for k in range(-0,1)])

from scipy.spatial import Voronoi
from pyramids.plot.PlotUtility import voronoi_plot_2d
vor = Voronoi(points)

    
maxPeak = energyArray[np.argmax(absorbanceSum)], np.max(absorbanceSum)
maxMinusPeak = energyArray[np.argmin(absorbanceSum)], np.min(absorbanceSum)

xlimits = [0, 8]#maxPeak[0]*1.5] 
ylimits = None#[0, 50]#[0, maxPeak[1]]

from scipy.signal import argrelextrema
extrema = argrelextrema(absorbanceSum, np.greater)
maxPeakValue = max(absorbanceSum)
peaks = []
for extreme in extrema[0]:
  peakPosition, peakValue = energyArray[extreme]/4.3, absorbanceSum[extreme]
  if peakValue > 0.01 * maxPeakValue:
    peaks.append((peakPosition,peakValue))
#    axs.text(peakPosition, maxPeakValue*0.7,'%3.2f' % peakPosition, fontsize='large',rotation=90)

peaks = np.array(peaks)    
axs[0].semilogy(energyArray/guassFieldEnergy,absorbanceSumHHG,'-',linewidth=2.,color='black',label='Total',alpha=0.5)
axs[1].plot(energyArray/guassFieldEnergy,absorbanceSumHHG,'-',linewidth=2.,color='red',label='Total')
args1 = getPropertyFromPosition(xlimits=xlimits,ylimits=ylimits,
                               ylabel='absorbance',xlabel='Order')

args2 = getPropertyFromPosition(xlimits=xlimits,ylimits=None,
                               ylabel='absorbance',xlabel='Order')
setProperty(axs[0],**args1)
setProperty(axs[1],**args2)

#plt.tight_layout()
plt.savefig('HHG.pdf',dpi=600)

#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Tue Dec 20 06:28:23 2016

@author: cl-iop
"""
import numpy as np
import matplotlib.pyplot as plt
import pyramids.io.result as dp
import pyramids.plot.setting as ma
numSample = 20
thetas = np.linspace(0,np.pi/2.0,numSample)
x = 0.2*np.sin(thetas)
y1 = np.load('intense.npy')
y2 = np.load('angle.npy')
fig, ax = plt.subplots(1,1,sharex=False,sharey=False,figsize=(8,6))
ax.plot(x,y1,'-',label = 'intense')
ax.plot(x,y2,'o',label = 'angle')
kargs = ma.getPropertyFromPosition(ylabel=r'', xlabel=r'$A \sin\theta$',title='')
ma.setProperty(ax,**kargs)

plt.tight_layout()
SaveName = __file__.split('/')[-1].split('.')[0] 
for save_type in ['.pdf']:
  plt.savefig(SaveName+save_type,transparent=True,dpi=600)