import numpy as np
import os
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt

# Parameters
FWHM = 0.4   # Full width half maximum
Emin=-16.    # Starting energy
Emax= 4.     # -----
DE = 0.01    # Energy intervals
dE=np.arange(Emin,Emax,DE)
E_shift = 0
filename = 'liquid_water-k1-1.pdos'

def plot_pdos(txtfile):

    if not os.path.isfile(txtfile):
        print("Warning: File does not exist!")
        return np.zeros(np.shape(dE)[0])
        
    PDOS = np.asarray(open(txtfile).read().split())

    if 'f' in PDOS:
        start = 26
        Norb = 4+3
    elif 'd' in PDOS:
        start = 25
        Norb = 3+3
    elif 'p' in PDOS:
        start = 24
        Norb = 2+3
    else:
        start = 23
        Norb = 1+3
        
    L = int(PDOS[-Norb])
    N = np.where(PDOS=='Projected')[0]-1
    
    for i in range(N.shape[0]-1):
        pdos = PDOS[N[i]:N[i+1]]
        pdos = pdos[start:]
        pdos = np.asarray(np.split(pdos,L))
        E = pdos[:,1].astype(np.float)*27.2114 #a.u. -> eV
        pdos = pdos[:,3:].astype(float)
        
        if i == 0:
            pdos_tot = np.zeros((pdos.shape))
        pdos_tot += pdos / N.shape[0]
        
        s, p = convolve(E,pdos)
        if i == 0:
            yshift = np.max(np.append(s,p))
        plt.plot(dE+E_shift,s-i*yshift, c='blue',linestyle='-',linewidth=2.,alpha=0.7)
        plt.plot(dE+E_shift,p-i*yshift, c='red',linestyle='-',linewidth=2.,alpha=0.7)
        
    s, p = convolve(E,pdos_tot)
    plt.text(-14,-(i+1.5)*yshift,'Sum')
    plt.plot(dE+E_shift,s-(i+2)*yshift, c='blue',linestyle='-',linewidth=2.,alpha=0.7)
    plt.plot(dE+E_shift,p-(i+2)*yshift, c='red',linestyle='-',linewidth=2.,alpha=0.7)
 
    plt.yticks([])                                                                                              
    plt.ylabel('PDOS (arbitrary units)')
    plt.xlabel('Binding energy (eV)')
    plt.xticks(np.arange(-15,5,5),-np.arange(-15,5,5))
    plt.xlim([-16,3])
    plt.title('FWHM = ' + str(FWHM) + ' eV')

    plt.plot([100,101],[0,0], c='blue',linestyle='-',linewidth=2., label='O s')
    plt.plot([100,101],[0,0], c='red',linestyle='-',linewidth=2., label = 'O p')
    plt.legend(loc=1)

    plt.savefig('pdos.png')
    plt.show()    
            
def convolve(E,PDOS):
    
    sigma = 1 / (2 * np.sqrt(2 * np.log(2) ) ) * FWHM
                
    conv = np.zeros((len(dE),2))
    for j in range(conv.shape[1]):
        W = PDOS[:,j]
        for i,e in enumerate(E):
            conv[:,j] += np.exp(-(dE-e)**2 / (2 * sigma**2))*W[i]
    return conv[:,0], conv[:,1]

plot_pdos(filename)