# 通过读取Wannier90给出的能带数据wannier90_band.dat，与VASP结合vaspkit得到的数据BAND.dat
# 来绘制 wannier插值后的能带与VASP计算的能带之间的差距
#-----------------------------------------------------------------
from cProfile import label
import os
import numpy as np
import matplotlib.pyplot as plt
#-----------------------------------------------------------------
def dataread():
    da1 = "wannier90_band.dat"
    x0 = []
    y0 = []
    with open(da1) as file:
            da = file.readlines()
            for f1 in da:
                temp = f1.strip().split() # 移除字符串头尾指定的字符（默认为空格）,并以空格将这些字符串分开
                if len(temp) != 0:
                    x0.append(float(temp[0]))
                    y0.append(float(temp[1]))
    # plt.scatter(x0,y0,s = 1, color = 'red', alpha = 0.7,marker = '.')
    plt.plot(x0,y0,'ro',markersize = 1,label="Wannier90")

    da2 = "BAND.dat"
    fermi =  4.8774
    x1 = []
    y1 = []
    with open(da2) as file:
            da = file.readlines()
            for f1 in da:
                temp = f1.strip().split() # 移除字符串头尾指定的字符（默认为空格）,并以空格将这些字符串分开
                if len(temp) == 2:
                    x1.append(float(temp[0]))
                    y1.append(float(temp[1]) + fermi)
    # plt.scatter(x1,y1 ,s = 1, color = 'blue', alpha = 0.7, marker = '.')
    plt.plot(x1,y1,'bo',markersize = 1,label = "VASP")
    plt.yticks(fontproperties='Times New Roman', size = 15)
    plt.xticks(fontproperties='Times New Roman', size = 15)
    plt.xlim(np.min(x0),np.max(x0))
    plt.legend(loc = 0,ncol = 3)
    plt.ylim(-5,5)
    picname = "band-compare.png"
    plt.savefig(picname, dpi = 600, bbox_inches = 'tight')
#--------------------------------------------------------------
def main():
    os.chdir(os.getcwd())# 确定用户执行路径
    dataread()
#----------------------------------------------------------------
if __name__=='__main__':
    main()