import numpy as np
import math
from matplotlib import rc
rc('text', usetex=True)
rc('font', family='serif')

import pylab as pl
#from scipy import *
#import scipy.signal
from h5py import File
#import pyfits
import os
import string
from matplotlib import pyplot 
from matplotlib import pyplot as plt
from matplotlib import pylab
from matplotlib import colors
#from pyreadcol import *
import numpy as np
from matplotlib.ticker import NullFormatter

massmin = 1.0e6
bins =31  # 300
binwidth = 0.1 #math.log10(massmax/massmin)/bins
massmax = 10**(math.log10(massmin)+bins*binwidth)
numbins = 101
#numbins = 12
#ybinwidth = 2.0/12

f=open('PopIII_haloes_nooverlap_allPopIII.txt','r')

#fo = open('PopIII_histogram.txt','w')

halolist = []
PopIIInumber = []

line = f.readline()
while line != '':
   halomass = float(line.split()[0])
   halolist.append(halomass)
   PopIIInumber.append(int(line.split()[1]))
   line = f.readline()


totalPopIII = np.zeros([bins+1,numbins+1])

for halomass,PopIII in zip(halolist,PopIIInumber):
   index1 = int((math.log10(halomass/massmin)+0.0*binwidth)/binwidth)
   totalPopIII[index1][PopIII] = totalPopIII[index1][PopIII] + 1
   #index2  = int(math.log10((PopIII+1.0)/1.0)/ybinwidth)
   #if index2 < 0: index2=0
   #totalPopIII[index1][index2] = totalPopIII[index1][index2] + 1

f.close()

data = totalPopIII
xvals = 10**(math.log10(massmin)+np.array(range(bins+1))*binwidth)
yvals = np.array(range(numbins+1)) 
yvals = 10**(np.log10(np.array(range(numbins+1))-0.0))
yvals[0] = 0.8
#yvals = 10**(math.log10(1)+(np.array(range(numbins+1))-1.0)*ybinwidth)

ylabel = r'$PopIII Number$'
xlabel = '$Halo Mass [M_{\odot}]$'

#dlogy = np.log10(yvals[1])-np.log10(yvals[0])
dy = yvals[1]-yvals[0]
dlogx = np.log10(xvals[1])-np.log10(xvals[0])

#data /= dlogx*dlogy

#intmin = int(np.log10(data[1:,1:].min())) + 1
#intmax = int(np.log10(data[1:,1:].max())) + 1
intmin = data.min()
intmax = data.max()
intmin = 0
intmax =np.log10(data.max())

print data
print intmin, intmax

fig = pl.figure()
ax = fig.add_subplot(111)

mynorm = colors.LogNorm([intmin, intmax])
im = ax.contourf(xvals, yvals, data.T, np.logspace(intmin,intmax,3), cmap=pl.cm.jet, norm=mynorm)

#if xvals.size == data.T.shape[0]:
#   xvals = np.append(xvals, 2.*10.**np.log10(xvals[-1]) - 10.**np.log10(xvals[-2]))
#if yvals.size == data.T.shape[0]:
#   yvals = np.append(yvals, yvals[-1] - yvals[-2])

fontsize=24

#X, Y = np.meshgrid(xvals,yvals)
im = ax.pcolor(xvals,yvals,data.T,cmap=pl.cm.jet, norm=mynorm)
ax.set_xscale('log')
ax.set_yscale('log')
cbar = fig.colorbar(im, fraction=0.05, ticks=(1,10,100,1000), pad=0.01, shrink=0.9)
cbar.set_label('Halo Number')

pl.xlim(xvals.min(), xvals.max())
#pl.ylim(yvals.min(), yvals.max())
#pl.ylim(yvals.min(), 10)

pl.xlabel('%s' % xlabel)
pl.ylabel('%s'% ylabel)
pl.savefig('PopIII_halo_phase_nooverlapi_log.png',format='png')
fig.clf()