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

def spectrum(energy,weight,nbBins=100,E_range=[1e-3,1e5]):
   '''
      Compute flux versus energy
      Input: events energy and weight (optional: number of bins)
      Output: energy, flux (E^2 dN/dE)
   '''
   Emin=E_range[0]#min(energy)
   Emax=E_range[1]#max(energy)
   energy =log10(energy)
   flux,ener=histogram(energy,nbBins,range=log10(E_range),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 
   return enercenter,flux

def drawSpectrum(files,plot_gen=False,plot_source=False,plot_analytic=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)
   ax1 = fig.add_subplot(111)

   for fileId in files:
      # draw full spectrum
      ener,flux = ReadSpectrum(fileId,[0,1])
      p = ax1.plot(ener,flux,drawstyle='steps-mid')

      if plot_gen:
         generation = ReadGeneration(fileId,[0])
         j=2
         for gen in generation:
            # read files
            ener,flux = ReadSpectrum(fileId,[0,j])
            ax1.plot(ener,flux,drawstyle='steps-mid',label="gen=%.0f"%gen)
            j+=1

      if plot_source:
         Es,Fs = ReadSourceSpectrum(fileId,[0,1])
         ax1.plot(Es,Fs,color=p[0].get_color(),linestyle=':')
         
   if plot_analytic:
      minEner=min(ener)
      #alpha=flux[ener==minEner]/sqrt(minEner)
      #ax1.plot(ener,alpha*sqrt(ener/100),'.-g',linewidth=2)
      ax1.plot(ener,2*Analytic_flux(ener),'--g',linewidth=1)
      Elim= Eic(Ethreshold_gg()/2)
      alpha = 2*4.63e3/Elim**(1/4.)
      ax1.plot(ener,alpha*sqrt(ener),'--r',linewidth=1)
      alpha = 2*14.6e3
      ax1.plot(ener,alpha*(ener/100)**(1/4.),'--m',linewidth=1)
      #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.axvline(x=1e4, ymin=0., ymax = 1., color='g', 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]")
   ax1.set_xlabel("energy [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()