from numpy import histogram, log10, sqrt, array
from matplotlib.pyplot import figure, show, setp
from matplotlib import gridspec
from Read import ReadSpectrum, ReadSourceSpectrum
from Constants import degre
from Analytic import Ethreshold_gg

def spectrum(energy,weight,nbBins=100):
   '''
      Compute flux versus energy
      Input: events energy and weight (optional: number of bins)
      Output: energy, flux (E^2 dN/dE)
   '''
   energy =log10(energy)
   flux,ener=histogram(energy,nbBins,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 
   flux /= max(flux) #nbPhotonsEmitted
   return enercenter,flux

def drawSpectrum(files,PlotAnalytic=False):
   '''
      Plot flux versus energy
      Input: list of directories
      Input (optional, default=False): plot analytic expression
      Output: graph spectrum normalize to 1 source photon
   '''
   fig = figure()
   gs = gridspec.GridSpec(2, 1, height_ratios=[3,1]) 
   fig.subplots_adjust(hspace=0.001)
   ax1 = fig.add_subplot(gs[0])
   ax2 = fig.add_subplot(gs[1],sharex=ax1)

   for fileId in files:
      i=1
      j=1
      for powerlaw_index in [1,1.5,2,2.5]: 
         # read files
         Es,Fs = ReadSourceSpectrum(fileId,[0,i])
         p = ax1.plot(Es,Fs,drawstyle='--')
         # draw full spectrum
         ener,flux,flux_0 = ReadSpectrum(fileId,[0,j,j+1])
         ax1.plot(ener,flux,color=p[0].get_color(),drawstyle='-',label="p=%.1f"%powerlaw_index)
         # primary gamma-rays contribution
         ax1.plot(ener,flux_0,color=p[0].get_color(),linestyle='-.')
         contrib=flux_0/flux
         ax2.plot(ener,contrib,'-',color=p[0].get_color())
         
         i+=1
         j+=2

      # add spectrum for Fermi/LAT and CTA (FoV?)

   if PlotAnalytic:
      minEner=min(ener)
      alpha=flux[ener==minEner]/sqrt(minEner)
      ax1.plot(ener,alpha*sqrt(ener),'--m',linewidth=2)
      ax1.plot(ener,ener**(0)*max(flux)*0.95,'--m',linewidth=2)
      ax1.axvline(x=Ethreshold_gg(), ymin=0., ymax = 1., color='r', linewidth=2)

   ax1.legend(loc="best")
   ax1.set_xscale('log')
   ax1.set_yscale('log')
   ax1.grid(b=True,which='major')
   ax1.set_ylabel("$E^2 dN/dE$ [GeV]")
   ax2.set_xscale('log')
   ax2.grid(b=True,which='major')
   ax2.set_xlabel("energy [GeV]")
   xticklabels = ax1.get_xticklabels()
   setp(xticklabels, visible=False)

   show()