from QuantumATK import *
import numpy as np
import pylab
from pylab import *

#-----------
# User input
#-----------

# List of c parameters used in the TB09 meta-GGA bandstructure calculations
tb09=[1.05,1.055,1.06,1.065,1.07,1.075,1.08,1.085,1.09,1.095,1.1]

# File containing the different bandstructures
fname = "Si_bulk_fitc.nc"

# Experimetal band gap at 0 K
SK=1.17 # (from Kittel)

#---------------------------------
# Calculate band gaps and make fit
#---------------------------------

# Read the bandstructures objects calculated at different tb09 parameters and calculate the
# indirect band gap.
ind_gap=[]
for n in range(len(tb09)):
    bandstructure=nlread(fname, Bandstructure)[n]
    ind_gap.append(bandstructure._indirectBandGap().inUnitsOf(eV))

for n in range(len(tb09)):
      print(tb09[n], ind_gap[n])

# Perform a linear fit of the calculated TB09 meta-GGA band gap vs c parameter and
# calculate the c parameters which will give the SK band gap.
(m,b)=polyfit(tb09,ind_gap,1)
fitted_c = (SK - b) / m

print("Fitted c parameter: ",  fitted_c)

#------------
# Plot figure
#------------

pylab.figure()
pylab.figsize=(2,2)

ax = pylab.subplot(111)

ax.spines['top'].set_linewidth(3)
ax.spines['bottom'].set_linewidth(3)
ax.spines['right'].set_linewidth(3)
ax.spines['left'].set_linewidth(3)
ax.xaxis.set_tick_params(width=1.5)
ax.yaxis.set_tick_params(width=1.5)
ax.tick_params(direction='in', pad=15)

yp=polyval([m,b],tb09)
xmin = min(tb09)
xmax = max(tb09)
ymin = min(yp)-0.02
ymax = max(yp)+0.05

ax.set_xlim(xmin,xmax)
ax.set_ylim(ymin,ymax)

pylab.plot([xmin,xmax],[SK,SK] , linewidth=2.5, linestyle=':', color='black', label="$\mathrm{E^{EXP}_{gap}}$" )
pylab.plot(tb09,ind_gap,marker='s',color='b',linestyle='none',markersize=10,label="$\mathrm{E^{TB09}_{gap}}$")
pylab.plot(tb09,yp, linestyle='--', color='b',linewidth=1.5,label="fit")
pylab.plot([fitted_c,fitted_c],[ymin,ymax],linestyle='-.',color='red',linewidth=2.5,label="opt-$c$")

pylab.xticks(fontsize=20)
pylab.yticks(fontsize=20)
pylab.xlabel("$c$",fontsize=30)
pylab.ylabel("E$_\mathrm{gap}$ (eV)",fontsize=20)
pylab.legend(loc="upper left",fontsize=20,numpoints=1,ncol=2)

pylab.tight_layout()
pylab.savefig('c_fit.png',dpi=300)