#!/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

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

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

energy= np.loadtxt(fileName+argv[2]+"/Results_momentum",unpack=True,usecols=[0])
Edata =np.log10(energy)
Emax = max(Edata)
Emin = min(Edata) 

ind = 0
for fileId in argv[2:]:
   energy= np.loadtxt(fileName+fileId+"/Results_momentum",unpack=True,usecols=[0])
   weight= np.loadtxt(fileName+fileId+"/Results_extra",unpack=True,usecols=[2])
   print "data shape:", np.shape(energy) 

   Edata =np.log10(energy)
   flux,ener=np.histogram(Edata,nbBins,range=(Emin,Emax),weights=weight)
   ener = 10**ener
   enercenter=(ener[1:nbBins+1]+ener[0:nbBins])/2
   binSize=ener[1:nbBins+1]-ener[0:nbBins]
   flux = enercenter**2 * flux/binSize 

   if argv[1] == "EGMF":
      ax.plot(enercenter,flux[:]/max(flux),"."+color[ind],
            label="$J_\gamma$ $10^{-%.0f"%float(fileId)+"}$Gauss")
   elif argv[1] == "EBL":
      ax.plot(enercenter,flux[:]/max(flux),"."+color[ind],label=Model[int(fileId)-1])

   # Results from Elmag
   Elmag=fileName+fileId+"/spec_diff_test"
   if os.path.isfile(Elmag):
      energy,fluxGamma,fluxLepton = np.loadtxt(Elmag,unpack=True,usecols=[0,1,2])
      energy=energy*10**(-9)
      flux=fluxGamma+fluxLepton
      ax.plot(energy,flux/max(flux),"-"+color[ind],label="Elmag - "+Model[int(fileId)-1])
    
   ind=ind+1

def fit(x,a,b,c):
   return a* x**b +c
a=1
b=0.5
c=0

xfit = np.logspace(np.log10(10**(-3)),np.log10(2*10**4),200)
yfit = fit(xfit,a,b,c)
ax.plot(xfit,yfit,'--m',linewidth=2,label="$E^{%.2f"%b+"}$")

ax.set_xscale('log')
ax.set_yscale('log')
ax.grid(b=True,which='major')
ax.legend(loc="best")
ax.set_xlabel("energy [GeV]")
ax.set_ylabel("$E^2 \\frac{dN}{dE}$")

plt.show()