import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.ticker import MultipleLocator, FormatStrFormatter

################################################################
## DOS ##
################################################################
# In order to use this code, create a folder DOS\\ 
# and put results espresso.dos and espresso.dos.in into subfolders 01, 02 etc.
# This is meant to make it easy to compare multiple calculations
# This code also saves all created figures in the folder of the inserted raw data as .png
# Search for '!!' and do all necessary adjustments

path = 'X:\\Your\\Path\\DOS\\' #!!
Nr = '01' #!!!! insert the subfolder name 01, 02 ...

filenameDOS = '\\espresso.dos'
filenameDISC = '\\espresso.dos.in'

data_file = np.loadtxt(path+Nr+filenameDOS)
data_file_txt = np.loadtxt(path+Nr+filenameDOS,max_rows =1, dtype=np.str, comments = 'c')

E_fermi_txt = data_file_txt[8]
#print(E_fermi)
E_fermi = E_fermi_txt.astype(np.float)
#print(E_fermi)

with open (path+Nr+filenameDISC, "r") as data_file_disc:
    data_file_disc = data_file_disc.read()
    data_file_disc = 'Calculated E_Fermi = '+ E_fermi_txt + '\n\n' + 'INPUT FILE:\n\n' + data_file_disc

Energy= data_file[:,0]
DOS = data_file[:,1]

x0 = Energy - E_fermi
y0 = DOS

fig = plt.figure(figsize=(15.00,8.80),dpi=150)


## FIRST SUBPLOT -> Graph ------------------------------
ax0 = plt.subplot2grid((1, 5), (0, 0), colspan=4)
ax0.plot(x0, y0, label='DOS',color='r',linewidth=1.0)

matplotlib.rcParams.update({'font.size': 12})

plt.xlabel('Energy [eV]    E$_{Fermi}$ is at 0 eV',fontsize=10,fontweight='bold')
plt.ylabel('DOS [states / eV]',fontsize=10,fontweight='bold')

plt.title("Density Of States")


#plt.legend()
#ax.legend(frameon=False)

#Ticks and Grid
plt.grid(b=True, which='major', axis='x',linewidth=1, color='k')
ax0.xaxis.set_major_locator(MultipleLocator(2))
ax0.xaxis.set_major_formatter(FormatStrFormatter('%d'))


plt.grid(b=True, which='minor', axis='x',linewidth=0.3, color='gray')
ax0.xaxis.set_minor_locator(MultipleLocator(0.5))
ax0.yaxis.major.formatter._useMathText = True
#plt.xticks(np.arange(min(x), max(x)+1, 1.0))

plt.grid(True)
plt.ylim((0,7))
plt.xlim((-10,10))


## SECOND SUBPLOT -> Text -----------------------------------
ax1 = plt.subplot2grid((1, 5), (0, 4),colspan=1)
ax1.text(0, -0.05, data_file_disc, fontsize = 7)
ax1.axis('off')
#ax1.table(cellText =data_file_disc,loc='bottom')

figname1 = '\\DOS_'
figname2 = '.png'

plt.savefig(path+Nr+figname1+Nr+figname2, format='png', dpi=150)

#plt.show()

################################################################
## BAND ##
################################################################
# In order to use this code, create a folder BAND\\ 
# and put results espresso.band1.gnu and espresso.band.in into subfolders 01, 02 etc.
# !! Insert the NRofBands Line 98
# !! Adjust the 'for' loop which adjusts coloring (Line 126). Insert the number of data pints per band
# !! Adjust k-vector sites in plt.xticks line 142 to get the wanted k-site describtion at the x-axis

path = 'X:\\Your\\Path\\BAND\\'

NRofBands = 16 #!!!!! insert the number of bands you have calculated
Nr = '01' #!!!! insert the subfolder name 01, 02 ...


filenameBAND = '\\espresso.band1.gnu'
filenameDISC2 = '\\espresso.band.in'

data_file = np.loadtxt(path+Nr+filenameBAND)


with open (path+Nr+filenameDISC2, "r") as data_file_disc2:
    data_file_disc2 = data_file_disc2.read()
    data_file_disc2 = 'INPUT FILE:\n\n' + data_file_disc2


k_vector= data_file[:,0]
Energy = data_file[:,1]

x1 = k_vector
y1 = Energy-E_fermi

fig = plt.figure(figsize=(15.00,8.80),dpi=150)

## FIRST SUBPLOT -> Graph ----------------------------------
ax1 = plt.subplot2grid((1, 5), (0, 0), colspan=4)

# !!!!! Adjust length of data for one band, here the data goes from 0 to 415 for one band, please adjust according to your calculation.
for i in range(1,NRofBands):
    ax1.plot(x1[0+415*i:415+415*i], y1[0+415*i:415+415*i], label = 'Electronic Bands'+ str(i),linewidth=1.0)

matplotlib.rcParams.update({'font.size': 12})

plt.xlabel('$\mathbf{Vector}$ $\mathbf{\overrightarrow{k}}$ [$\mathbf{2 \pi{}/a}$]',fontsize=10,fontweight='bold')
plt.ylabel('Energy [eV]    ($\mathbf{E_{Fermi}}$ is at 0 eV)',fontsize=10,fontweight='bold')

plt.title("Electronic Band Structure")

#plt.legend()
#ax.legend(frameon=False)

#Ticks and Grid X-Axis MAJOR
plt.grid(b=True, which='major', axis='x',linewidth=1, color='k')
ax1.xaxis.set_major_locator(MultipleLocator(10))
# !!!! Adjust k-sites, with k-path and name
plt.xticks(list([0, 1.0, 1.5, 2.21, 3.07, 4.13]), ['$\Gamma$ [0.0]','X [1.0]','W [1.5]','L [2.21]','$\Gamma$ [3.07]','K [4.13]'])
#ax1.xaxis.set_major_formatter(FormatStrFormatter('%d'))

#Ticks and Grid X-Axis MINOR
plt.grid(b=True, which='minor', axis='x',linewidth=0.3, color='gray')
ax1.xaxis.set_minor_locator(MultipleLocator(0.1))
#ax1.yaxis.major.formatter._useMathText = True

#Ticks and Grid y-Axis MAJOR
plt.grid(b=True, which='major', axis='y',linewidth=1.0, color='k')
ax1.yaxis.set_major_locator(MultipleLocator(2.0))

#Ticks and Grid y-Axis Minor
plt.grid(b=True, which='minor', axis='y',linewidth=0.3, color='gray')
ax1.yaxis.set_minor_locator(MultipleLocator(0.5))

plt.xlim((0,4.13))
plt.ylim((-10,10))

plt.grid(True)

## SECOND SUBPLOT -> Text -----------------------------------
ax2 = plt.subplot2grid((1, 5), (0, 4),colspan=1)
ax2.text(0, -0.05, data_file_disc2, fontsize = 7)
ax2.axis('off')
#ax1.table(cellText =data_file_disc,loc='bottom')

figname1 = '\\BAND_'
figname2 = '.png'

plt.savefig(path+Nr+figname1+Nr+figname2, format='png', dpi=150)

#plt.show()




################################################################
## BAND + DOS ##
################################################################
# !!!! Do all the adjustments from above also in this code



fig = plt.figure(figsize=(15.00,8.80),dpi=150)

## FIRST SUBPLOT -> Graph ----------------------------------
ax1 = plt.subplot2grid((1, 5), (0, 0), colspan=3)

for i in range(1,NRofBands): # !!! Adjust number of data points per band
    ax1.plot(x1[0+415*i:415+415*i], y1[0+415*i:415+415*i], label = 'Electronic Bands'+ str(i),linewidth=1.0)

matplotlib.rcParams.update({'font.size': 12})

plt.xlabel('$\mathbf{Vector}$ $\mathbf{\overrightarrow{k}}$ [$\mathbf{2 \pi{}/a}$]',fontsize=10,fontweight='bold')
plt.ylabel('Energy [eV]    ($\mathbf{E_{Fermi}}$ is at 0 eV)',fontsize=10,fontweight='bold')

plt.title("Electronic Band Structure")

#plt.legend()
#ax.legend(frameon=False)

#Ticks and Grid X-Axis MAJOR
plt.grid(b=True, which='major', axis='x',linewidth=1, color='k')
ax1.xaxis.set_major_locator(MultipleLocator(10))
plt.xticks(list([0, 1.0, 1.5, 2.21, 3.07, 4.13]), ['$\Gamma$ [0.0]','X [1.0]','W [1.5]','L [2.21]','$\Gamma$ [3.07]','K [4.13]      ']) #!!!
#ax1.xaxis.set_major_formatter(FormatStrFormatter('%d'))

#Ticks and Grid X-Axis MINOR
plt.grid(b=True, which='minor', axis='x',linewidth=0.3, color='gray')
ax1.xaxis.set_minor_locator(MultipleLocator(0.1))
#ax1.yaxis.major.formatter._useMathText = True

#Ticks and Grid y-Axis MAJOR
plt.grid(b=True, which='major', axis='y',linewidth=1.0, color='k')
ax1.yaxis.set_major_locator(MultipleLocator(2.0))

#Ticks and Grid y-Axis Minor
plt.grid(b=True, which='minor', axis='y',linewidth=0.3, color='gray')
ax1.yaxis.set_minor_locator(MultipleLocator(0.5))

plt.xlim((0,4.13))
plt.ylim((-10,10))

plt.grid(True)



## SECOND SUBPLOT -> DOS -----------------------------------

ax2 = plt.subplot2grid((1, 5), (0, 3), colspan=2)
ax2.plot(y0, x0, label='DOS',color='r',linewidth=1.0)

matplotlib.rcParams.update({'font.size': 12})

#plt.ylabel('Energy [eV]    E$_{Fermi}$ is at 0 eV',fontsize=10,fontweight='bold')
plt.xlabel('DOS [states / eV]',fontsize=10,fontweight='bold')

plt.title("Density Of States")


#plt.legend()
#ax.legend(frameon=False)

#Ticks and Grid
plt.grid(b=True, which='major', axis='y',linewidth=1, color='k')
ax2.yaxis.set_major_locator(MultipleLocator(2))
ax2.yaxis.set_major_formatter(FormatStrFormatter('%d'))


#Ticks and Grid y-Axis MAJOR
plt.grid(b=True, which='major', axis='y',linewidth=1.0, color='k')
ax2.yaxis.set_major_locator(MultipleLocator(2.0))

#Ticks and Grid y-Axis Minor
plt.grid(b=True, which='minor', axis='y',linewidth=0.3, color='gray')
ax2.yaxis.set_minor_locator(MultipleLocator(0.5))
ax2.set_yticklabels([])

plt.grid(True)
plt.xlim((0,7.5))
plt.ylim((-10,10))


#--------------------

figname1 = '\\BAND+DOS_'
figname2 = '.png'

plt.savefig(path+Nr+figname1+Nr+figname2, format='png', dpi=150)

plt.show()

#!!! If you have no python distribution, download microsoft visual studio (it is free (at least in 2019)) and istall
# a python package. It is very convenient and compiles inside visual studio and also allows interactive editing.
# Hope everything works fine for you,
# Elias Henögl