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

# Parameters
FWHM = 2.5   # Full width half maximum
Emin=-16.    # Starting energy
Emax= 4.     # -----
DE = 0.01    # Fineness of the energies to convolve to
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])
    
    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)
    
    s, p = convolve(E,pdos)
    plt.plot(dE+E_shift,s, c='blue',linestyle='-',linewidth=2., label='O s')
    plt.plot(dE+E_shift,p, c='red',linestyle='-',linewidth=2., label = 'O p')
        
    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.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)