"""Script for making thermodynamic comparison plots for the Lyman-alpha forest
between Arepo and Gadget"""

import matplotlib
matplotlib.use('PDF')

import numpy as np
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import plot_flux_power as pflux
plt.figure(1)
#Mean flux
ddir="/home/spb/scratch/ComparisonProject/"
tau_ar512= np.loadtxt(ddir+"tau_eff_arepo_512.txt")
tau_gad512=np.loadtxt(ddir+"tau_eff_gadget_512.txt")
tau_ar256= np.loadtxt(ddir+"tau_eff_arepo_256.txt")
tau_gad256=np.loadtxt(ddir+"tau_eff_gadget_256.txt")
#Tau figure; bottom plot shows normalised to tau_gadget.
gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1])
ax1=plt.subplot(gs[0])
plt.plot(tau_ar512[:,0], tau_ar512[:,1] ,lw=4,color="red",ls="-")
plt.plot(tau_gad512[:,0],tau_gad512[:,1],lw=4,color="blue",ls="--")
plt.plot(tau_ar256[:,0], tau_ar256[:,1] ,lw=4,color="red",ls=":")
plt.plot(tau_gad256[:,0],tau_gad256[:,1],lw=4,color="blue",ls="-.")
plt.ylabel(r"$\tau_\mathrm{eff}$",size=25)
plt.yticks((0,0.4,0.8,1.2,1.6),('0','0.4','0.8','1.2','1.6'))
plt.ylim(0,1.6)
ax2=plt.subplot(gs[1],sharex=ax1)
# plt.ylabel("Flux PDF")
plt.plot(tau_ar512[:,0], 100*(tau_ar512[:,1]/tau_gad512[:,1]-1) ,lw=4,color="red",ls="-")
plt.plot(tau_ar256[:,0], 100*(tau_ar256[:,1]/tau_gad512[:,1]-1) ,lw=4,color="red",ls=":")
plt.plot(tau_gad256[:,0],100*(tau_gad256[:,1]/tau_gad512[:,1]-1),lw=4,color="blue",ls="-.")
plt.subplots_adjust(hspace=0)
plt.setp( ax1.get_xticklabels(), visible=False)
plt.yticks((0,-5,-10),('0','-5','-10'))
plt.ylim(-11,0)
plt.xlabel("Redshift",size=22)
plt.xlim(2,4)
# plt.legend(loc=3,ncol=2)
plt.savefig(ddir+"tau_eff.pdf")
plt.clf()
gad=np.loadtxt(ddir+"Gadget/thermo.txt")
ar=np.loadtxt(ddir+"Arepo/thermo.txt")
#ar_m=np.loadtxt(ddir+"Arepo/thermo_m.txt")
#are=np.loadtxt(ddir+"Arepo_ENTROPY/thermo.txt")
ar2=np.loadtxt(ddir+"Arepo_256/thermo_m.txt")
# ar2=np.loadtxt(ddir+"Arepo_256/thermo.txt")
gad2=np.loadtxt(ddir+"Gadget_256/thermo.txt")
plt.plot(ar[:,0],ar[:,1]/1e3,label="Arepo",color="blue")
#plt.plot(ar_m[:,0],ar_m[:,1],label="Arepo Mass")
#plt.plot(are[:,0],are[:,1],label="Arepo_ENTROPY")
#plt.plot(ar2[:,0],ar2[:,1],label="Arepo 256")
plt.plot(ar2[:,0],ar2[:,1]/1e3,label="Arepo 256",ls="-.",color="black")
plt.plot(gad[:,0],gad[:,1]/1e3, label="Gadget",ls = '--',color="red")
ind = np.where((gad[:,0] < 4.5)*(gad[:,0] > 2))
max=np.max(np.abs(ar[ind,1]-gad[ind,1]))
ind2 = np.where(np.abs(ar[ind,1][0]-gad[ind,1][0]) == max)
print max,gad[ind,0][0][ind2],ar[ind,1][0][ind2]/gad[ind,1][0][ind2],np.abs(gad[ind,2][0][ind2]/ar[ind,2][0][ind2])
##QUICK_LYA plots
# arQ=np.loadtxt(ddir+"Arepo_256_Qlya/thermo.txt")
# ar2=np.loadtxt(ddir+"Arepo_256/thermo.txt")
# gadQ=np.loadtxt(ddir+"Gadget_256_Qlya/thermo.txt")
# plt.plot(arQ[:,0],arQ[:,1]/1e3,label="Arepo 256 Qlya",ls="-.",color="cyan")
# plt.plot(gadQ[:,0],gadQ[:,1]/1e3, label="Gadget 256 Qlya",ls = '--',color="magenta")
#plt.plot(gad2[:,0],gad2[:,1]/1e3,label="Gadget 256",ls = '-.',color="magenta")
plt.ylabel("$T_0$ ($10^3$ K)")
plt.xlabel("Redshift")
plt.xlim(0,8)
plt.ylim(6,12)
# plt.legend(loc=3,ncol=2)
plt.tight_layout()
plt.savefig(ddir+"temp.pdf")
plt.figure(4)
plt.plot(gad[:,0],gad[:,2], label="Gadget",color="red")
plt.plot(ar[:,0],ar[:,2],label="Arepo",color="blue")
plt.plot(ar2[:,0],ar2[:,2],label="Arepo 256",color="black")
# plt.plot(gad2[:,0],gad2[:,2],label="Gadget 256",ls = '--')
# plt.plot(arQ[:,0],arQ[:,2],label="Arepo 256 Q")
# plt.plot(gadQ[:,0],gadQ[:,2],label="Gadget 256 Q",ls = '--')
plt.ylabel("Temp-density relation")
plt.xlabel("Redshift")
plt.xlim(0,8)
plt.ylim(1.54,1.6)
# plt.legend(loc=0)
plt.tight_layout()
plt.savefig(ddir+"gamma.pdf")
f=plt.figure(2)
# pflux.pdfplots(119,color="black")
# Three subplots sharing both x/y axes
# f, (ax1, ax2) = plt.subplots(2, sharex=True)
# Fine-tune figure; make subplots close to each other and hide x ticks for
# all but bottom plot.
gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1])
ax1=plt.subplot(gs[0])
pflux.pdfplots(139)
pflux.pdfplots(159,ls="--")
pflux.pdfplots(189,ls=":")
plt.ylim(0.2,10)
plt.ylabel("Flux PDF",size=22)
plt.yticks((1,10),('1.0','10'))
# plt.xlabel("")
# plt.legend(loc=0)
#plt.tight_layout()
#plt.clf()
ax2=plt.subplot(gs[1],sharex=ax1)
# pflux.pdfrelplots(119,color="black")
pflux.pdfrelplots(139)
pflux.pdfrelplots(159,ls="--")
pflux.pdfrelplots(189,ls=":")
# plt.ylabel("Flux PDF")
plt.ylim(-8,8)
plt.yticks((-5,0,5),('-5','0','5'))
plt.subplots_adjust(hspace=0)
plt.setp( ax1.get_xticklabels(), visible=False)
plt.xlim(0,1)
plt.xlabel("Flux",size=22)
# plt.legend(loc=0)
# plt.tight_layout()
plt.savefig(ddir+"flux_pdf.pdf")


plt.figure(3)
gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1])
ax1=plt.subplot(gs[0])
pflux.pfplots(119,ls="-.")
pflux.pfplots(139)
pflux.pfplots(159,ls="--")
pflux.pfplots(189,ls=":")
plt.ylim(2e-2,0.5)
plt.yticks((0.1,0.5),('0.1','0.5'))
ax2=plt.subplot(gs[1],sharex=ax1)
pflux.pfrelplots(119,ls="-.")
pflux.pfrelplots(139)
pflux.pfrelplots(159,ls="--")
pflux.pfrelplots(189,ls=":")
plt.ylim(-6,6)
plt.yticks((-5,0,5),('-5','0','5'))
plt.subplots_adjust(hspace=0)
plt.setp( ax1.get_xticklabels(), visible=False)
plt.xlim(0.4,5)
plt.xticks((0.5,1,2,5),('0.5','1','2','5'))
plt.savefig(ddir+"flux_pow.pdf")


#plt.figure(5)
#gad=np.loadtxt(ddir+"Gadget/high-den-thermo.txt")
#ar=np.loadtxt(ddir+"Arepo/high-den-thermo.txt")
#plt.plot(gad[:,0],gad[:,1], label="Gadget")
#plt.plot(ar[:,0],ar[:,1],label="Arepo")
#plt.ylabel("Temp at mean density")
#plt.xlabel("Redshift")
#plt.legend(loc=0)
#plt.savefig(ddir+"hdtemp.pdf")
#plt.figure(6)
#plt.plot(gad[:,0],gad[:,2], label="Gadget")
#plt.plot(ar[:,0],ar[:,2],label="Arepo")
#plt.ylabel("Temp-density relation")
#plt.xlabel("Redshift")
#plt.legend(loc=0)
#plt.tight_layout()
#plt.savefig(ddir+"hdgamma.pdf")

#Histogram
#import re
#def load_multi_txt(snap_file):
#    """Load data from a text file split into multiple pieces"""
#    data=np.loadtxt(snap_file)
#    for i in np.arange(1,500):
#        file2=re.sub("_0_","_"+str(i)+"_",snap_file)
#        if snap_file == file2:
#            break
#        try:
#            data=np.append(data,np.loadtxt(file2))
#        except IOError:
#            break
#    return np.array(data)
#plt.clf()
#temp_ar=load_multi_txt(ddir+"/Arepo/124_tempdata.txt")
#temp_gad=load_multi_txt(ddir+"/Gadget/124_tempdata.txt")
#temp_ar_256=load_multi_txt(ddir+"/Arepo_256/124_0_tempdata.txt")
#temp_gad_256=load_multi_txt(ddir+"/Gadget_256/124_0_tempdata.txt")
## temp_ar_128=load_multi_txt(ddir+"/Arepo_128/124_0_tempdata.txt")
## temp_gad_128=load_multi_txt(ddir+"/Gadget_128/124_0_tempdata.txt")
#plt.hist(temp_ar_256,bins=40,log=True,histtype='stepfilled',label="Arepo 256")
#plt.hist(temp_gad_256,bins=40,log=True,histtype='stepfilled',label="Gadget 256")
## plt.hist(temp_ar_128,bins=40,log=True,histtype='step',label="Arepo 128")
## plt.hist(temp_gad_128,bins=40,log=True,histtype='step',label="Gadget 128")
#plt.hist(temp_ar,bins=40,log=True,histtype='step',label="Arepo 512")
#plt.hist(temp_gad,bins=40,log=True,histtype='step',label="Gadget 512")
#plt.legend(loc=0)
#
#plt.tight_layout()
#plt.savefig(ddir+"temp_hist.pdf")
#Positional plot
#plt.clf()
#pos_ar_256=load_multi_txt(ddir+"/Arepo_256/124_0_posdata.txt")
#pos_ar_256=np.reshape(pos_ar_256,[-1,3])/1000
#pos_gad_256=load_multi_txt(ddir+"/Gadget_256/124_0_posdata.txt")
#pos_gad_256=np.reshape(pos_gad_256,[-1,3])/1000
#ind=np.where(pos_ar_256[:,2]> 19.9)
#plt.scatter(pos_ar_256[ind,0],pos_ar_256[ind,1],c=temp_ar_256[ind],vmin=3.5,vmax=5)
#plt.xlim(0,20)
#plt.ylim(0,20)
#plt.xlabel("x (Mpc)")
#plt.ylabel("y (Mpc)")
#plt.title("Arepo Positions at 19.9 < z < 20 Mpc.")
#cb=plt.colorbar()
#cb.set_label("log T_0")
#plt.tight_layout()
#plt.savefig(ddir+"arepo_pos.pdf")
#plt.figure(7)
#ind2=np.where(pos_gad_256[:,2]> 19.9)
#plt.scatter(pos_gad_256[ind2,0],pos_gad_256[ind2,1],c=temp_gad_256[ind2],vmin=3.5,vmax=5)
#plt.xlim(0,20)
#plt.ylim(0,20)
#plt.xlabel("x (Mpc)")
#plt.ylabel("y (Mpc)")
#plt.title("Gadget Positions at 19.9 < z < 20 Mpc.")
#cb=plt.colorbar()
#cb.set_label("log T_0")
#plt.tight_layout()
#plt.savefig(ddir+"gadget_pos.pdf")
