#!/usr/bin/python

from sys import argv
import matplotlib.pyplot as plt
from matplotlib import gridspec
from Constantes import *

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'",
         "Fraceschini","Finke and Al","Gilmore and Al","Dominguez and Al"]
else:
   print "Used: ./Angle_distribution.py arg1 ind1 ind2 ..."
   print "        arg1 = EGMF or EBL "
   print "        ind1 = file number \n"
   print "Examples: "
   print "   to study EGMF14: ./Angle_distribution.py EGMF 14 "
   print "   to study EBL1 and EBL2: ./Angle_distribution.py EBL 1 2 \n"
   exit()

nbBins = 100

Esource=100 #TeV
Dsource=135 #Mpc
delta_to_theta=(lambda_gg/Esource*1e3)/Dsource

def fit_theta(E,B):
   RL0=RL*(Esource/2)*(1e-17/B)
   Dic0=Dic/(Esource/2)
   delta=Dic0/(2*RL0)*((Esource/2)**2 *Eic/E-1)
   return abs(arcsin(delta_to_theta*sin(delta))*degre)

fig2 = plt.figure()
ax2 = fig2.add_subplot(111)

fig1 = plt.figure()
gs = gridspec.GridSpec(2, 1, height_ratios=[4,1]) 
fig1.subplots_adjust(hspace=0.001)
ax11 = fig1.add_subplot(gs[0])
ax12 = fig1.add_subplot(gs[1],sharex=ax11)

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

ind = 0
for fileId in argv[2:]:
   energy,dirx,diry,dirz = loadtxt(fileName+fileId+"/Results_momentum",unpack=True,usecols=[0,1,2,3])
   posx,posy,posz = loadtxt(fileName+fileId+"/Results_position",unpack=True,usecols=[1,2,3])

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

   hyp=sqrt(posx**2+posy**2+posz**2)
   posx=posx/hyp
   posy=posy/hyp
   posz=posz/hyp
   hyp=sqrt(dirx**2+diry**2+dirz**2)
   dirx=dirx/hyp
   diry=diry/hyp
   dirz=dirz/hyp
   costheta=dirx*posx+diry*posy+dirz*posz

   cond=(costheta<=1)&(costheta>=-1)

   theta = arccos(costheta[cond])*degre
   weight= weight[cond]
   energy= energy[cond]
   print fileId, ":", shape(energy)

   # Angle versus energy
   #===================== 
   nbBins = 40
   Edata =log10(energy)
   angle,ener =histogram(Edata,nbBins,weights=theta)
   nb,ener =histogram(Edata,nbBins)
   ener = 10**ener
   enercenter=(ener[1:nbBins+1]+ener[0:nbBins])/2
   angle = angle/nb 
   if argv[1] == "EGMF":
      ax11.plot(enercenter,angle,'.'+color[ind],
               label="$10^{-%.0f"%float(fileId)+"}$Gauss - MC")
      #ax11.plot(energy,theta,'.'+color[ind],label="$10^{-%.0f"%float(fileId)+"}$Gauss - MC")
      B=10**(-float(fileId))
      yfit = fit_theta(enercenter,B)

      ax11.plot(enercenter,yfit,'--'+color[ind],linewidth=2,label="$10^{-%.0f"%float(fileId)+"}$Gauss - analytic")

   elif argv[1] == "EBL":
      ax11.plot(enercenter,angle,'.'+color[ind],label=Model[int(fileId)-1])
      if ind==0: 
         yfit = fit_theta(enercenter,B)

         ax11.plot(enercenter,yfit,'--'+color[ind],linewidth=2,label="$10^{-15}$Gauss - analytic")

   error=angle/yfit-1
   ax12.plot(enercenter,error,'+'+color[ind])

   # figure: radial distribution
   #=============================

   nbBins = 20
   theta = theta[theta!=0]
   weight=weight[theta!=0]
   #theta = log10(theta**2)
   nb,angle1=histogram(theta**2,nbBins,  range=[0,1],weights=weight)

   #angle1=10**angle1

   anglecenter=(angle1[1:nbBins+1]+angle1[0:nbBins])/2
   binSize=angle1[1:nbBins+1]-angle1[0:nbBins]
   ax2.plot(anglecenter,nb,'-'+color[ind],
         label="$10^{-%.0f"%float(fileId)+"}$Gauss - MC")
   #ax2.hist(theta, nbBins, range=[0,0.5], weights=weight,log=True,

   #      facecolor=color[ind],alpha=.5,label="$10^{-%.0f"%float(fileId)+"}$Gauss")

   ind=ind+1

ax11.legend(loc="best")
ax11.set_xscale('log')
ax11.set_yscale('log')
ax11.set_ylim([0,180])
ax11.grid(b=True,which='major')
ax11.set_ylabel("$\\theta$ [degre]")
ax12.set_xscale('log')
ax12.grid(b=True,which='major')
ax12.set_xlabel("energy [GeV]")
ax12.set_ylabel("relative error")
step=0.2*(max(error)-min(error))
plt.ylim(min(error)-step, max(error)+step)
xticklabels = ax11.get_xticklabels()
plt.setp(xticklabels, visible=False)

ax2.legend(loc="best")

#ax2.set_xscale('log')

ax2.set_yscale('log')
ax2.grid(b=True,which='major')
ax2.set_xlabel("$\\theta^2$ [degre]")
ax2.set_ylabel("Nb events * weight")

plt.show()