from numpy import log10, histogram
from matplotlib.pyplot import figure, show
from matplotlib import gridspec
from Read import ReadRadial, ReadAngleVsEnergy
from Analytic import Analytic_theta
from Constants import degre

Eliminf = 1e0 #GeV
Elimsup = 1e5 #GeV

def angle_vs_energy(theta,energy,weight,nbPhotonsEmitted=1,nbBins=40):
   '''
      Compute arrival angle versus energy
      Input: directory name
      Output: energy, arrival angle
   '''
   energy =log10(energy)
   angle,ener =histogram(energy,nbBins,weights=weight*theta)
   nb,ener =histogram(energy,nbBins,weights=weight)
   ener = 10**ener
   enercenter=(ener[1:nbBins+1]+ener[0:nbBins])/2
   angle /= nb * nbPhotonsEmitted

   return enercenter, angle

def radial(theta,weight,nbPhotonsEmitted=1,nbBins=50,thetamax=90):
   '''
      Compute flux versus arrival angle
      Input: directory name
      Input (optionnal): maximal angle desired (by default full sky)
      Output: arrival angle, flux (dn/dtheta)
   '''

   theta=log10(theta)
   dn,angle = histogram(theta,nbBins,range=[log10(1e-6),log10(thetamax)],weights=weight)
   angle=10**angle
   thetacenter = (angle[1:nbBins+1]+angle[0:nbBins])/2
   dtheta = angle[1:nbBins+1]-angle[0:nbBins]
   dndtheta = dn/dtheta*thetacenter
   dndtheta /= nbPhotonsEmitted

   return thetacenter, dndtheta

def drawAngle_vs_energy(files):
   '''
      Plot arrival angle versus energy + analytic expression
      Input: list of directories
      Output: graph of arrival angle versus energy
   '''
   fig = figure()
   ax1 = fig.add_subplot(111)

   for fileId in files:
      i=1
      for powerlaw_index in [1,1.5,2,2.5]: 
         # apply source spectrum
         ener,angle = ReadAngleVsEnergy(fileId,[0,i])
         p=ax1.plot(ener,angle,'-',label="p=%.1f"%powerlaw_index)

      yfit = Analytic_theta(ener,fileId)
      ax1.plot(ener,yfit,'--',color="m",linewidth=2,label="analytic")

   ax1.legend(loc="best")
   ax1.set_xscale('log')
   ax1.set_yscale('log')
   ax1.set_ylim([0,180])
   ax1.grid(b=True,which='major')
   ax1.set_ylabel("$\\theta$ [rad]")
   ax1.set_xlabel("energy [GeV]")

   show()

def drawRadial(files):
   '''
      Plot flux versus arrival angle
      Input: list of directories
      Output: graph of flux versus angle
   '''
   fig = figure()
   ax = fig.add_subplot(111)

   for fileId in files:
      i=1
      for powerlaw_index in [1,1.5,2,2.5]: 
         theta,dndtheta2 = ReadRadial(fileId,[0,i])
         ax.plot(theta,dndtheta2,drawstyle='steps-mid',label="p=%.1f"%powerlaw_index)

   ax.legend(loc="best")
   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.grid(b=True,which='major')
   ax.set_xlabel("$\\theta$ [rad]")
   ax.set_ylabel("$\\theta*dN/d\\theta$")
   
   show()