from numpy import histogram, log10, shape, logspace, sqrt
from matplotlib.pyplot import figure, show, ylim, setp
from matplotlib import gridspec
from os.path import isfile
from Read import ReadEnergy, ReadWeight, ReadElmag, ReadNbLeptonsProd, resultDirectory, ElmagFile
from Read import ReadNphot_source, ReadGeneration
from Analytic import Analytic_flux, Ethreshold_gg

def spectrum(fileId,gen=-1):
   energy = ReadEnergy(fileId)
   weight = ReadWeight(fileId)

   if gen != -1:
      generation = ReadGeneration(fileId)
      energy=energy[generation==gen]
      weight=weight[generation==gen]
      print "file", fileId, "gen=", gen,"->", shape(energy)[0], "events" 
   else:
      print "file", fileId, "->", shape(energy)[0], "events" 

   nbBins=100
   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 
   maxflux = ReadNphot_source(fileId)
   flux /= maxflux
   return enercenter,flux

###########################################################################################
def drawSpectrum(files,PlotElmag=False,PlotAnalytic=False):
   fig = figure()
   ax = fig.add_subplot(111)

   for fileId in files:
      energy,flux = spectrum(fileId)
      if PlotElmag:
         maxflux = max(flux)*(.8)
      p=ax.plot(energy,flux,drawstyle='steps-mid',label=fileId)

      if PlotElmag & isfile(resultDirectory+fileId+ElmagFile):
         energy, flux = ReadElmag(fileId)
         maxflux = max(flux)
         flux /= maxflux
         ax.plot(energy,flux,"--",color=p[0].get_color(),label="Elmag - "+fileId)
            
   if PlotAnalytic:
      alpha=5e3
      ax.plot(energy,alpha*sqrt(energy),'--m',linewidth=2)
      ax.plot(energy,energy**(0)*max(flux),'--m',linewidth=2)
      ax.axvline(x=Ethreshold_gg(), ymin=0., ymax = 1., color='r', linewidth=2)

   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.grid(b=True,which='major')
   ax.legend(loc="best")
   if PlotElmag:
      ax.set_ylabel("$E^2 \\frac{dN}{dE}$ [normalized]")
   else:
      ax.set_ylabel("$E^2 \\frac{dN}{dE}$ [GeV / photon]")
   ax.set_xlabel("energy [GeV]")

   show()

def drawSpectrum_normaLepton(files):
   fig = figure()
   gs = gridspec.GridSpec(2, 1, height_ratios=[4,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:
      energy,flux = spectrum(fileId)
      maxflux = ReadNbLeptonsProd(fileId)
      flux /= maxflux
      p=ax1.plot(energy,flux,drawstyle='steps-mid',label=fileId)

      yfit = Analytic_flux(energy)
      ax1.plot(energy,yfit,'--m',linewidth=2)

      error=flux/yfit-1
      ax2.plot(energy,error,'+',color=p[0].get_color())

   ax1.set_xscale('log')
   ax1.set_yscale('log')
   ax1.grid(b=True,which='major')
   ax1.legend(loc="best")
   ax1.set_ylabel("$E^2 \\frac{dN}{dE}$ [GeV] / lepton")
   ax2.set_xscale('log')
   ax2.grid(b=True,which='major')
   ax2.set_xlabel("energy [GeV]")
   ax2.set_ylabel("error")
   step=0.2*(max(error)-min(error))
   ylim(min(error)-step, max(error)+step)
   xticklabels = ax1.get_xticklabels()
   setp(xticklabels, visible=False)

   show()

########################################################################################
def drawSpectrumGen(fileId):
   fig = figure()
   ax = fig.add_subplot(111)

   for gen in list(set(ReadGeneration(fileId))):
      energy,flux = spectrum(fileId,gen)
      p=ax.plot(energy,flux,drawstyle='steps-mid',label="gen="+str(int(gen)))

   energy,flux = spectrum(fileId)
   p=ax.plot(energy,flux,drawstyle='steps-mid',color='k',label="gen = all")
   alpha=2e3

   ax.plot(energy,sqrt(energy)*alpha,'--m',linewidth=2)
   ax.plot(energy,energy**(0)*max(flux),'--m',linewidth=2)
   ax.axvline(x=Ethreshold_gg(), ymin=0., ymax = 1., color='r', linewidth=2)

   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.grid(b=True,which='major')
   ax.legend(loc="best")
   ax.set_title(fileId)
   ax.set_ylabel("$E^2 \\frac{dN}{dE}$ [GeV / photon]")
   ax.set_xlabel("energy [GeV]")

   show()