from matplotlib.pylab import *
import numpy
import sys
import os 

fig_size = [800/72.27 ,700/72.27]
params = {'backend': 'ps', 'axes.labelsize': 25, 'font.size': 25, 'legend.fontsize': 16,
          'xtick.labelsize': 25, 'ytick.labelsize': 25, 
          'xtick.major.size': 15,'ytick.major.size': 15,
          'xtick.minor.size': 8,'ytick.minor.size': 8,
          'text.usetex': False, 'figure.figsize': fig_size, 'xtick.major.pad': 9,
          'figure.subplot.top': 0.94,'figure.subplot.left': 0.105,'figure.subplot.bottom': 0.095,'figure.subplot.right': 0.978}
#'figure.subplot.bottom': 0.097

rcParams.update(params)
a=subplot(111)

title_ = "First-Principles Equation of State Database";
if (len(sys.argv)>1):
   title_ = str(sys.argv[1])
title(title_)

scriptName = sys.argv[0][:-3]; # [:3] removes ".py" for file name 
try:
   os.remove(scriptName+'.pdf')
   os.remove(scriptName+'.png')
except:
   pass

##############################################################################################################################

hug = numpy.loadtxt('FPEOS_Hugoniot.txt',usecols=(1,3,5,7,9,13)); # T,rho/rho0,rho0,rho,P,E
rho0 = hug[0,2]
# print(rho0)

##############################################################################################################################

cx = 1
cy = 4
l = 9;
plot(hug[:,cx],hug[:,cy],'-', linewidth=l, markersize=0,color='r',mec='r',mew=2, mfc='white', dashes=(16/l,3/l,3/l,3/l), label='Hugoniot curve',zorder=20);

l = 8
hugRel = numpy.loadtxt('FPEOS_Hugoniot_rel.txt',usecols=(1,3,5,7,9,13)); # T,rho/rho0,rho0,rho,P,E
plot(hugRel[:,cx],hugRel[:,cy],'-', linewidth=l, markersize=0,color='k',mec='k',mew=2, mfc='white', dashes=(16/l,3/l,3/l,3/l),label='Hugoniot with\nrelativistic effects');

if os.path.isfile('FPEOS_Hugoniot_rad.txt'): # for hydrogen this file does not get written 
   l = 5
   hugRad = numpy.loadtxt('FPEOS_Hugoniot_rad.txt',usecols=(1,3,5,7,9,13)); # T,rho/rho0,rho0,rho,P,E
   plot(hugRad[:,cx],hugRad[:,cy],'-', linewidth=l, markersize=0,color='brown',mec='brown',mew=2, mfc='white', dashes=(16/l,3/l,3/l,3/l),label='Hugoniot with\nradiation effects');
   #fill_betweenx(hugRad[:,cy], hug[:,cx],hugRad[:,cx],facecolor='r',color='r',alpha=0.1)

   #careful: Since the Y axis is pressure, not temperature, we need to interpolate hug curve to give rho/rho0 for given P 
   from scipy import interpolate
   hugFindRatioGivenP = interpolate.interp1d(hug[:,cy],hug[:,cx])

   n1 = hugRad.shape[0]//2
   while(n1>0 and hugRad[n1,cy]>min(hug[:,cy])):
      n1 -= 1;

   n2 = hugRad.shape[0]//2
   while(n2<hugRad.shape[0] and hugRad[n2,cy]<max(hug[:,cy])):
      n2 += 1;

   fill_betweenx(hugRad[n1:n2,cy], hugFindRatioGivenP(hugRad[n1:n2,cy]),hugRad[n1:n2,cx],facecolor='r',color='r',alpha=0.1)
   a.set_xlim(2.0,ceil(5.01*max(hugRad[:,cx]))/5 )
else:
   a.set_xlim(2.0,6.0)

l = 2
hugLower = numpy.loadtxt('FPEOS_Hugoniot_lower_initial_density.txt',usecols=(1,3,5,7,9,13)); # T,rho/rho0,rho0,rho,P,E
plot(hugLower[:,cx],hugLower[:,cy],'-', linewidth=l, markersize=0,color='darkgreen',mec='brown',mew=2, mfc='white',label='Hugoniot with\n'+r'20% lower $\rho_0$');

l = 2
hugHigher = numpy.loadtxt('FPEOS_Hugoniot_higher_initial_density.txt',usecols=(1,3,5,7,9,13)); # T,rho/rho0,rho0,rho,P,E
plot(hugHigher[:,cx],hugHigher[:,cy],'-', linewidth=l, markersize=0,color='navy',mec='brown',mew=2, mfc='white',dashes=(2/l,1/l),label='Hugoniot with\n'+r'20% higher $\rho_0$');

##############################################################################################################################
d   = numpy.loadtxt('FPEOS_isotherms.txt'      ,usecols=(1,3,7,9,11)); # index,rho,T,P,E
dp  = numpy.loadtxt('FPEOS_isotherm_points.txt',usecols=(1,3,7,9,11)); # index,rho,T,P,E

iMin = int(min(d[:,0]))
iMax = int(max(d[:,0]))
#print(iMin,iMax)
iMin = int(min(dp[:,0]))
iMax = int(max(dp[:,0]))
#print(iMin,iMax)

for i in range(iMin,iMax+1):
    ii  = (d[:,0]==i)
    iip = (dp[:,0]==i)

    label = ''
    if (i==iMin): label='Isotherm'
    plot( dp[iip,1]/rho0,dp[iip,3],'sb-', linewidth=0, markersize=5,color='b',mec='b',mfc='lightblue',mew=1,zorder=-10);
    plot( d[ii,1] /rho0,d[ii,3],   'sb-', linewidth=1, markersize=0,color='b',mec='b',mfc='lightblue',mew=1,zorder=-11);
    plot(-d[ii,1]/rho0,d[ii,3],'sb-', linewidth=1, markersize=5,color='b',mec='b',mfc='lightblue',mew=1, label=label);

##############################################################################################################################

# plot(iso[:,2],iso[:,1],'k-', linewidth=lw, markersize=0,color='b',mec='k',mew=2, mfc='white',dashes=(10,4),label='Isobar');

#########################################################################################################################################

#xlabel(r'Pressure (GPa)')
#xlabel(r'Density (g$\,$cm$^{-3}$)')
xlabel(r'Compression ratio $\rho / \rho_0$')
#ylabel(r'Temperature (K)')
ylabel(r'Pressure (GPa)')

#print(min(hug[:,cy]))

#a.set_xscale('log')
a.set_ylim( 10.0**floor(np.log10(min(hug[:,cy]))) , 10.0**ceil(np.log10(max(hug[:,cy]))) )
#a.set_xlim( 2.0                                   , ceil(5.01*max(hugRad[:,cx]))/5 )
a.xaxis.set_minor_locator(MultipleLocator(0.2));
a.set_yscale('log')
a.xaxis.set_ticks_position('both')
a.get_xaxis().set_tick_params(which='both', direction='in')
a.yaxis.set_ticks_position('both')
a.get_yaxis().set_tick_params(which='both', direction='in')

legend(loc='upper left',frameon=True,framealpha=1.0,handlelength=3.5,ncol=1,handletextpad=0.4,numpoints=1)

savefig(scriptName+'.pdf', bbox_inches='tight')
savefig(scriptName+'.png', bbox_inches='tight', dpi=200)