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()