# written by He Xinyi, modified by T. Kamiya

import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt
import matplotlib.ticker as ticker
import sys


from tklib.tkapplication import tkApplication
from tklib.tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg

#------------------ FONT_setup ----------------------
font = {'family' : 'arial', 
    'color'  : 'black',
    'weight' : 'normal',
    'size' : 18.0,
    }

file = 'TDOS.dat'
xmin = -2
xmax = 2
ymax = 10

save_figure = 1
plot_figure = 1

app    = tkApplication()
file        = getarg  (1, file)
save_figure = getintarg  (5, save_figure)
plot_figure = getintarg  (6, plot_figure)


if not plot_figure:
    mpl.use('Agg') #silent mode


#------------------- Data Read ----------------------
print("")
print(f"Read DOS from [{file}]")

with open(file,"r") as reader:
	legend = reader.readline()
legends=legend.split()[1:]
legends=[i.replace("_"," ") for i in legends]
legend_s=tuple(legends)
datas=np.loadtxt(file,dtype=np.float64,skiprows=1)

#--------------------- PLOTs ------------------------
axe = plt.subplot(111)
plt.subplots_adjust(left=0.18, right=0.95, top=0.95, bottom=0.18)
# Color methods! choose only one of the two methods
axe.plot(datas[:,0],datas[:,1:],linewidth=1.0) #auto colors

axe.set_xlabel(r'${E}$-$E_{F}$ (eV)',fontdict=font)
axe.set_ylabel(r'PDOS (states/eV)',fontdict=font)
#my_y_ticks=np.arange(0, ymax,ymax/2)
my_x_ticks=np.arange(-5, 5,1)
plt.xticks(my_x_ticks,fontsize=font['size']-2)
plt.yticks(fontsize=font['size']-2,fontname=font['family'])
plt.legend(legend_s,loc='upper right')
plt.xlim(( xmin,  xmax)) # set y limits manually
plt.ylim(( 0,  ymax))
leg = plt.gca().get_legend()
ltext = leg.get_texts()
plt.setp(ltext, fontsize=font['size']) 
fig = plt.gcf()
fig.set_size_inches( 5, 4)

if save_figure:
    outfile = 'dos.png'
    print("")
    print(f"save DOS fig: {outfile}")
    plt.savefig(outfile, dpi = 300, transparent = True)

if plot_figure:
    print("")
    print(f"plot")
    plt.pause(0.0001)


#plt.close()
app.terminate(pause = True)