from numpy import histogram, log10, sqrt, array, arange
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 = flux/binSize #*enercenter**2
   return enercenter,flux

def drawSpectrum(files,plot_gen=False,plot_source=False,plot_analytic=False,plot_time_range=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,dNdE = ReadSpectrum(fileId,[0,1])
      p = ax1.plot(ener,ener**2 *dNdE,drawstyle='steps-mid',label=fileId)#"$t_{obs}=\infty$")

      #minEner=min(ener)
      #alpha=flux[ener==minEner]/sqrt(minEner)
      #ax1.plot(ener,alpha*sqrt(ener),'--m')
      if plot_time_range:
         ener,dNdE1,dNdE2,dNdE3 = ReadSpectrum(fileId,[0,2,3,4])
         ax1.plot(ener,ener**2 *dNdE1,drawstyle='steps-mid',linestyle='--',label="$t_{obs}$=30days")
         ax1.plot(ener,ener**2 *dNdE2,drawstyle='steps-mid',linestyle='--',label="$t_{obs}$=1yr")
         ax1.plot(ener,ener**2 *dNdE3,drawstyle='steps-mid',linestyle='--',label="$t_{obs}$=5yrs")

      if plot_gen:
         ener,dNdE1,dNdE2,dNdE3 = ReadSpectrum(fileId,[0,5,6,7])
         ax1.plot(ener,ener**2 *dNdE1,drawstyle='steps-mid',linestyle='--',label="gen=1")
         ax1.plot(ener,ener**2 *dNdE2,drawstyle='steps-mid',linestyle='--',label="gen=2")
         ax1.plot(ener,ener**2 *dNdE3,drawstyle='steps-mid',linestyle='--',label="gen=3")
         #j=5
         #for gen in ReadGeneration(fileId,[0]):
         #   ener,flux = ReadSpectrum(fileId,[0,j])
         #   ax1.plot(ener,ener**2 *flux,drawstyle='steps-mid',linestyle='--',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:
      #=== Gen=2 ===
      ax1.plot(ener,2*Analytic_flux(ener),'--g')
      #=== Gen=4 ===
      #Elim= Eic(Ethreshold_gg()/2)
      #alpha = 2*9.262e3/Elim**(1/4.)
      #ax1.plot(ener,alpha*sqrt(ener),'--r',linewidth=1)
      #alpha = 2*29.29e3
      #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()