from pylab import *
import numpy as np
from matplotlib.ticker import ScalarFormatter
import matplotlib.gridspec as gridspec
from mpl_toolkits.axes_grid1.inset_locator import mark_inset
from matplotlib.ticker import AutoMinorLocator
from matplotlib.colors import LogNorm
from matplotlib.patches import Ellipse
rcParams["mathtext.fontset"]='cm'

############################### figure ###########################
#fig=figure(figsize=(15,10))     #give dimensions to the figure
##################################################################

################################ INPUT #######################################
#axes range
##############################################################################

############################ subplots ############################
#gs = gridspec.GridSpec(2,1,height_ratios=[5,2])
#ax1=plt.subplot(gs[0])
#ax2=plt.subplot(gs[1])

#make a subplot at a given position and with some given dimensions
#ax2=axes([0.4,0.55,0.25,0.1])

#gs.update(hspace=0.0,wspace=0.4,bottom=0.6,top=1.05)
#subplots_adjust(left=None, bottom=None, right=None, top=None,
#                wspace=0.5, hspace=0.5)

#set minor ticks
#ax1.xaxis.set_minor_locator(AutoMinorLocator(4))
#ax1.yaxis.set_minor_locator(AutoMinorLocator(4))


#ax1.xaxis.set_major_formatter( NullFormatter() )   #unset x label 
#ax1.yaxis.set_major_formatter( NullFormatter() )   #unset y label

# custom xticks 
#ax1.set_xticks([0.25, 0.5, 1.0])
#ax1.set_yticks([0.25, 0.5, 1.0])
#ax1.get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter()) #for log 


#ax1.get_yaxis().set_label_coords(-0.2,0.5)  #align y-axis for multiple plots
##################################################################

##################### special behaviour stuff ####################
#to show error missing error bars in log scale
#ax1.set_yscale('log',nonposy='clip')  #set log scale for the y-axis

#set the x-axis in %f format instead of %e
#ax1.xaxis.set_major_formatter(ScalarFormatter()) 

#set size of ticks
#ax1.tick_params(axis='both', which='major', labelsize=10)
#ax1.tick_params(axis='both', which='minor', labelsize=8)

#set the position of the ylabel 
#ax1.yaxis.set_label_coords(-0.2, 0.4)

#set yticks in scientific notation
#ax1.ticklabel_format(axis='y',style='sci',scilimits=(1,4))

#set the x-axis in %f format instead of %e
#formatter = matplotlib.ticker.FormatStrFormatter('$%.2e$') 
#ax1.yaxis.set_major_formatter(formatter) 

#add two legends in the same plot
#ax5 = ax1.twinx()
#ax5.yaxis.set_major_formatter( NullFormatter() )   #unset y label 
#ax5.legend([p1,p2],['0.0 eV','0.3 eV'],loc=3,prop={'size':14},ncol=1)

#set points to show in the yaxis
#ax1.set_yticks([0,1,2])

#highlight a zoomed region
#mark_inset(ax1, ax2, loc1=2, loc2=4, fc="none",edgecolor='purple')
##################################################################

############################ plot type ###########################
#standard plot
#p1,=ax1.plot(x,y,linestyle='-',marker='None')

#error bar plot with the minimum and maximum values of the error bar interval
#p1=ax1.errorbar(r,xi,yerr=[delta_xi_min,delta_xi_max],lw=1,fmt='o',ms=2,
#               elinewidth=1,capsize=5,linestyle='-') 

#filled area
#p1=ax1.fill_between([x_min,x_max],[1.02,1.02],[0.98,0.98],color='k',alpha=0.2)

#hatch area
#ax1.fill([x_min,x_min,x_max,x_max],[y_min,3.0,3.0,y_min],#color='k',
#         hatch='X',fill=False,alpha=0.5)

#scatter plot
#p1=ax1.scatter(k1,Pk1,c='b',edgecolor='none',s=8,marker='*')

#plot with markers
#pl4,=ax1.plot(ke3,Pk3/Pke3,marker='.',markevery=2,c='r',linestyle='None')

#set size of dashed lines
#ax.plot([0, 1], [0, 1], linestyle='--', dashes=(5, 1)) #length of 5, space of 1

#image plot
#cax = ax1.imshow(densities,cmap=get_cmap('jet'),origin='lower',
#           extent=[x_min, x_max, y_min, y_max],
#           #vmin=min_density,vmax=max_density)
#           norm = LogNorm(vmin=min_density,vmax=max_density))
#cbar = fig.colorbar(cax, ax2, ax=ax1, ticks=[-1, 0, 1]) #in ax2 colorbar of ax1
#cbar.set_label(r"$M_{\rm CSF}\/[h^{-1}M_\odot]$",fontsize=14,labelpad=-50)
#cbar.ax.tick_params(labelsize=10)  #to change size of ticks

#make a polygon
#polygon = Rectangle((0.4,50.0), 20.0, 20.0, edgecolor='purple',lw=0.5,
#                    fill=False)
#ax1.add_artist(polygon)
####################################################################

x_min, x_max = 0.3, 50.0
y_min, y_max = 0.0, 1.6


#fig=figure()
fig=figure(figsize=(10,6))     #give dimensions to the figure
ax1=fig.add_subplot(231) 
ax2=fig.add_subplot(232) 
ax3=fig.add_subplot(233) 
ax4=fig.add_subplot(234) 
ax5=fig.add_subplot(235) 
ax6=fig.add_subplot(236) 

subplots_adjust(left=None, bottom=None, right=None, top=None,
                wspace=0.05, hspace=0.05)

for ax in [ax1,ax2,ax3,ax4,ax5,ax6]:
    ax.set_xscale('log')
    #ax1.set_yscale('log')

    ax.set_xlim([x_min,x_max])
    ax.set_ylim([y_min,y_max])

for ax in [ax4,ax5,ax6]:
    ax.set_xlabel(r'$k\/[h\/{\rm Mpc}^{-1}]$',fontsize=18)

for ax in [ax1,ax4]:
    ax.set_ylabel(r'$P_{\rm hydro}(k)/P_{\rm N-body}(k)$',fontsize=18)


for ax in [ax1,ax2,ax3]:
    ax.xaxis.set_major_formatter( NullFormatter() )   #unset x label 

for ax in [ax2,ax3,ax5,ax6]:
    ax.yaxis.set_major_formatter( NullFormatter() )   #unset y label


f_out='Pk_ratio.pdf'

root = '/mnt/ceph/users/camels/Results/Pk_ratio'

# IllustrisTNG z=0
count = 0
for i in range(1000):
    fin = '%s/IllustrisTNG/%d/Pk_ratio_m_z=0.00.txt'%(root,i)
    k, Pk_ratio = np.loadtxt(fin, unpack=True)
    p1,=ax1.plot(k,Pk_ratio,linestyle='-',marker='None',c='k',alpha=0.1)
    count += 1
print('Found %d realizations'%count)

# IllustrisTNG z=2
count = 0
for i in range(1000):
    fin = '%s/IllustrisTNG/%d/Pk_ratio_m_z=2.00.txt'%(root,i)
    k, Pk_ratio = np.loadtxt(fin, unpack=True)
    p1,=ax2.plot(k,Pk_ratio,linestyle='-',marker='None',c='k',alpha=0.1)
    count += 1
print('Found %d realizations'%count)

# IllustrisTNG z=5
count = 0
for i in range(1000):
    fin = '%s/IllustrisTNG/%d/Pk_ratio_m_z=5.00.txt'%(root,i)
    k, Pk_ratio = np.loadtxt(fin, unpack=True)
    p1,=ax3.plot(k,Pk_ratio,linestyle='-',marker='None',c='k',alpha=0.1)
    count += 1
print('Found %d realizations'%count)

# SIMBA z=0
count = 0
for i in range(1000):
    fin = '%s/SIMBA/%d/Pk_ratio_m_z=0.00.txt'%(root,i)
    k, Pk_ratio = np.loadtxt(fin, unpack=True)
    p1,=ax4.plot(k,Pk_ratio,linestyle='-',marker='None',c='k',alpha=0.1)
    count += 1
print('Found %d realizations'%count)

# SIMBA z=2
count = 0
for i in range(1000):
    fin = '%s/SIMBA/%d/Pk_ratio_m_z=2.00.txt'%(root,i)
    k, Pk_ratio = np.loadtxt(fin, unpack=True)
    p1,=ax5.plot(k,Pk_ratio,linestyle='-',marker='None',c='k',alpha=0.1)
    count += 1
print('Found %d realizations'%count)

# SIMBA z=5
count = 0
for i in range(1000):
    fin = '%s/SIMBA/%d/Pk_ratio_m_z=5.00.txt'%(root,i)
    k, Pk_ratio = np.loadtxt(fin, unpack=True)
    p1,=ax6.plot(k,Pk_ratio,linestyle='-',marker='None',c='k',alpha=0.1)
    count += 1
print('Found %d realizations'%count)

BoxSize = 25.0
grid = 256
kN = 2.0*np.pi/BoxSize*(grid/2.0)
for ax in [ax1,ax2,ax3,ax4,ax5,ax6]:
    ax.plot([kN,kN],[y_min,y_max],linestyle='--',marker='None',c='b',alpha=1.0)

#f1 = 'prueba.txt'

#x,y=np.loadtxt(f1,unpack=True) 

#p1,=ax1.plot(x,y,linestyle='-',marker='None')


#place a label in the plot
ax1.text(0.4,1.4, r"${\rm IllustrisTNG\,:\,z=0}$", fontsize=12, color='k')
ax2.text(0.4,1.4, r"${\rm IllustrisTNG\,:\,z=2}$", fontsize=12, color='k')
ax3.text(0.4,1.4, r"${\rm IllustrisTNG\,:\,z=5}$", fontsize=12, color='k')
ax4.text(0.4,1.4, r"${\rm SIMBA\,:\,z=0}$", fontsize=12, color='k')
ax5.text(0.4,1.4, r"${\rm SIMBA\,:\,z=2}$", fontsize=12, color='k')
ax6.text(0.4,1.4, r"${\rm SIMBA\,:\,z=5}$", fontsize=12, color='k')

#legend
#ax1.legend([p1,p2],
#           [r"$z=3$",
#            r"$z=4$"],
#           loc=0,prop={'size':18},ncol=1,frameon=True)
            
            #columnspacing=2,labelspacing=2)




#ax1.set_title(r'$\sum m_\nu=0.0\/{\rm eV}$',position=(0.5,1.02),size=18)
#title('About as simple as it gets, folks')
#suptitle('About as simple as it gets, folks')  #for title with several panels
#grid(True)
#show()
savefig(f_out, bbox_inches='tight')
close(fig)










###############################################################################
#some useful colors:

#'darkseagreen'
#'yellow'
#"hotpink"
#"gold
#"fuchsia"
#"lime"
#"brown"
#"silver"
#"cyan"
#"dodgerblue"
#"darkviolet"
#"magenta"
#"deepskyblue"
#"orchid"
#"aqua"
#"darkorange"
#"coral"
#"lightgreen"
#"salmon"
#"bisque"
