from scipy import *
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib as mpl

mpl.rc('font',**{'family':'times new roman','sans-serif':['Helvetica']})

band = loadtxt('cell0/BGG_phband.bands')
dos = loadtxt('cell0/BGG_phdos.dos')

gs = gridspec.GridSpec(1,2, width_ratios = [2, 1])
gs.update(wspace = 0.1)
#ax_band = plt.subplot(111)

ax_band = plt.subplot(gs[0])
ax_dos = plt.subplot(gs[1])

ax_band.plot(band[:,0], band[:,1:], color='black', linestyle = '-', lw = 1)
ax_band.grid('on')
ax_band.axis([0, 0.781037, 0, 280])

kpname = ['$\Gamma$', '$\mathrm{X}$', '$\mathrm{M}$', '$\Gamma$', '$\mathrm{R}$']
kpval  = [ 0.000000,  0.151768,  0.303536,  0.518168,  0.781037]

ax_band.axes.set_xticks(kpval)
ax_band.axes.set_xticklabels(kpname, fontsize=18)
ax_band.axes.set_ylabel(r"Frequency (cm${}^{-1}$)", fontsize=18)
plt.setp(ax_band.get_yticklabels(), fontsize=18)

#plt.plot(dos[:,0], dos[:,1], label = "Total", color='black', zorder = 50)
#ax.plot(dos[:,0], dos_Ba1[:], label = "Ba(1)", color='red')
ax_dos.axes.get_yaxis().set_ticks([])
ax_dos.axes.get_xaxis().set_ticks([])

dos_Ba1 = np.zeros(len(dos[:,0]))
dos_Ba2 = np.zeros(len(dos[:,0]))
dos_cage= np.zeros(len(dos[:,0]))

for i in range(len(dos[:,0])):
    for j in range(0, 2):
        dos_Ba1[i] += dos[i, j + 2]

    for j in range(2, 8):
        dos_Ba2[i] += dos[i, j + 2]

    for j in range(8, 24):
        dos_cage[i] += dos[i, j + 2]

    for j in range(24, 54):
        dos_cage[i] += dos[i, j + 2]

ax_dos.plot(dos_Ba1[:], dos[:,0], label = "Ba(1)", color='green', zorder = 4, lw = 1.5)
ax_dos.plot(dos_Ba2[:], dos[:,0], label = "Ba(2)", color='magenta', zorder = 2, lw = 1.5)
ax_dos.plot(dos_cage[:], dos[:,0], label = "Cage", color='blue', zorder = 0, lw = 1.5)

xmin, xmax = ax_dos.get_xlim()
ymin, ymax = ax_dos.get_ylim()

print xmin, xmax
print ymin, ymax


ax_dos.axis([0, 3.0, 0, 280])
ax_dos.set_xlim(0.0, 3.0)
ax_dos.set_ylim(0, 280)

ax_dos.legend(loc="center right",prop={'size':12})
ax_dos.axes.set_xlabel(r"DOS", fontsize=18)

plt.savefig("BGG_band_dos.eps", format='eps', transparent=True)

plt.show()

