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