#!/usr/bin/python

from sys import argv
import os.path

import matplotlib.pyplot as plt

from Constantes import *
from scipy.integrate import quad

if argv[1] == "EGMF":
   fileName="Results_EGMF"
elif argv[1] == "EBL":
   fileName="Results_EBL"
   B=10**(-15)
   Model=["Kneiske and Doll - 'best fit'","Kneiske and Doll - 'lower limit'",
         "Franceschini","Finke and Al","Gilmore and Al","Dominguez and Al"]

elif argv[1] == "prec":
   fileName="Results_prec"
   B=10**(-17)

else:
   print "Used: ./Energy_distribution.py arg1 ind1 ind2 ..."
   print "        arg1 = EGMF or EBL "
   print "        ind1 = file number \n"
   print "Examples: "
   print "   to study EGMF14: ./Energy_distribution.py EGMF 14 "
   print "   to study EBL1 and EBL2: ./Energy_distribution.py EBL 1 2 \n"
   exit()

# figure: energy distribution
#=============================

zSource = 0.0308
nbBins = 100
convert= 180/pi

color=['b','r','g','c','m','y']

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

def properTime(z):
   return -1/(H0*(1+z)*sqrt(omegaM*(1+z)**3+omegaK*(1+z)**2+omegaL))

tlim = quad(properTime,zSource,0)[0]/s_to_yr
print tlim

time= loadtxt(fileName+argv[2]+"/Results_position",unpack=True,usecols=[0])
time=time[time>0]
time=time/tlim

tmax = max(time)
tmin = min(time) 
print tmin, tmax

ind = 0
for fileId in argv[2:]:

   time= loadtxt(fileName+fileId+"/Results_position",unpack=True,usecols=[0])
   charge,weight,gen=loadtxt(fileName+fileId+"/Results_extra",unpack=True,usecols=[0,2,3])
   time=time/tlim

   prop = float(shape(time[time<0])[0])/shape(time)[0]
   print "data shape",fileName+fileId,":", shape(time), "negative time:", shape(time[time<0]), "~", prop

   cond= (gen<10) & (time>0) & (charge==0)

   time=time[cond]
   weight=weight[cond]

   time=log10(time)

   dN,dt=histogram(time,nbBins,range=(log10(tmin),log10(tmax)),weights=weight)
   #time=10**time

   timecenter=(dt[1:nbBins+1]+dt[0:nbBins])/2
   binSize=dt[1:nbBins+1]-dt[0:nbBins]
   dNdt=dN/binSize

   if argv[1] == "EGMF":

      ax.hist(time,nbBins,weights=weight,range=(log10(tmin),log10(tmax)),log=1,facecolor=color[ind],
            alpha=.5, label="$10^{-%.0f"%float(fileId)+"}$Gauss")

      ax.plot(timecenter[dNdt!=0],dNdt[dNdt!=0],"-"+color[ind],
         label="$10^{-%.0f"%float(fileId)+"}$Gauss")

   elif argv[1] == "prec":
      ax.hist(time,nbBins,weights=weight,range=(log10(tmin),log10(tmax)),log=1,facecolor=color[ind],
            alpha=.5)#, label="prec = $10^{-%.0f"%float(fileId)+"}$")
      ax.plot(timecenter[dNdt!=0],dNdt[dNdt!=0],"-"+color[ind],
            label="prec = $10^{-%.0f"%float(fileId)+"}$")

   elif argv[1] == "EBL":
      #plt.hist(time,nbBins,weights=weight,range=(tmin,tmax),log=1,facecolor=color[ind],alpha=.5,
      ax.plot(timecenter,dNdt,"."+color[ind],
         label=Model[int(fileId)-1])

   ind=ind+1

#ax.set_xscale('log')
#ax.set_yscale('log')
ax.grid(b=True,which='major')
ax.legend(loc="best")

if argv[1]=="prec":
   ax.legend(loc="best",title="$10^{-17}$Gauss")
ax.set_xlabel("Time [$\Delta t / t$ log scale]")
ax.set_ylabel("$t\ dN/dt$")

plt.show()