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,generation=[],nbBins=100):
   '''
      Compute flux versus energy
      Input: events energy and weight (optional: number of bins)
      Output: energy, flux (E^2 dN/dE)
   '''
   Emin=min(energy)
   Emax=max(energy)
   energy =log10(energy)
   flux,ener=histogram(energy,nbBins,range=[log10(Emin),log10(Emax)],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 

   if generation == []:
      return enercenter,flux
   else:
      cond=(generation==0)
      flux_0,ener_0=histogram(energy[cond],nbBins,range=[log10(Emin),log10(Emax)],weights=weight[cond])
      ener_0 = 10**ener_0
      enercenter_0=(ener_0[1:nbBins+1]+ener_0[0:nbBins])/2
      binSize=ener_0[1:nbBins+1]-ener_0[0:nbBins]
      flux_0 = enercenter_0**2 * flux_0/binSize 
      return enercenter,flux,flux_0


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

   for fileId in files:
      if all_source_spectrum:
         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,linestyle=':')
            # draw full spectrum
            ener,flux,flux_0 = ReadSpectrum(fileId,[0,j,j+1])
            ax1.plot(ener,flux,color=p[0].get_color(),drawstyle='steps-mid',label="p=%.1f"%powerlaw_index)
            # primary gamma-rays contribution
            #ax1.plot(ener,flux_0,color=p[0].get_color(),linestyle='--',drawstyle='steps-mid')
            #contrib=flux_0/flux
            #ax2.plot(ener,contrib,drawstyle='steps-mid',color=p[0].get_color())
            i+=1
            j+=2
      elif all_jets:
         i=5
         j=9
         for jet_opening in ["isotrop","60 deg","30 deg"]: 
            # read files
            Es,Fs = ReadSourceSpectrum(fileId,[0,i])
            p = ax1.plot(Es,Fs,linestyle=':')
            # draw full spectrum
            ener,flux,flux_0 = ReadSpectrum(fileId,[0,j,j+1])
            ax1.plot(ener,flux,color=p[0].get_color(),drawstyle='steps-mid',label=jet_opening)
            # primary gamma-rays contribution
            #ax1.plot(ener,flux_0,color=p[0].get_color(),linestyle='--',drawstyle='steps-mid')
            #contrib=flux_0/flux
            #ax2.plot(ener,contrib,drawstyle='steps-mid',color=p[0].get_color())
            i+=1
            j+=2

      else:
         # read files
         #Es,Fs = ReadSourceSpectrum(fileId,[0,3])
         #p = ax1.plot(Es,Fs,linestyle=':')
         # draw full spectrum
         ener,flux,flux_0 = ReadSpectrum(fileId,[0,5,6])
         ax1.plot(ener,flux,drawstyle='steps-mid')
         #ax1.plot(ener,flux,color=p[0].get_color(),drawstyle='steps-mid',label=fileId)
         # primary gamma-rays contribution
         #ax1.plot(ener,flux_0,color=p[0].get_color(),linestyle='--',drawstyle='steps-mid')
         #contrib=flux_0/flux
         #ax2.plot(ener,contrib,drawstyle='steps-mid',color=p[0].get_color())
         
   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.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()