#!/usr/bin/python

from sys import argv
import os.path
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
from Constantes import s_to_yr

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"]
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
#=============================
nbBins = 200
convert= 180/np.pi
color=['b','r','g','c','m','y']

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

time= np.loadtxt(fileName+argv[2]+"/Results_position",unpack=True,usecols=[0])
tmax = max(time)
tmin = min(time) 
print tmin, tmax

ind = 0
for fileId in argv[2:]:
   time= np.loadtxt(fileName+fileId+"/Results_position",unpack=True,usecols=[0])
   weight,gen=np.loadtxt(fileName+fileId+"/Results_extra",unpack=True,usecols=[2,3])

   cond=(gen<8)
   time=time[cond]
   weight=weight[cond]

   print "data shape",fileName+fileId,":", np.shape(time) 

   dN,dt=np.histogram(time,nbBins,range=(tmin,tmax),weights=weight)
   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=(tmin,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] == "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")
ax.set_xlabel("Time [yr]")
ax.set_ylabel("$dN/dt$ [yr$^{-1}$]")

plt.show()