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):
   '''
      Compute flux versus energy
      Input: directory name, generation (optional, by default: all generation)
      Output: energy, flux
   '''
   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,gen=-1,PlotElmag=False,PlotAnalytic=False):
   '''
      Plot flux versus energy
      Input: list of directories
      Input (optional, default=all): selection on the generation
      Input (optional, default=False): Plot Elmag spectrum if exists, plot analytic expression
      Output: graph spectrum normalize to 1 source photon
   '''
   fig = figure()
   ax = fig.add_subplot(111)

   for fileId in files:
      energy,flux = spectrum(fileId,gen)
      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:
      minEner=min(energy)
      alpha=flux[energy==minEner]/sqrt(minEner)
      ax.plot(energy,alpha*sqrt(energy),'--m',linewidth=2)
      ax.plot(energy,energy**(0)*max(flux)*0.95,'--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):
   '''
      Plot flux versus energy compare with analytic expression
      Input: list of directories
      Output: graph spectrum normalize to 1 lepton generated
   '''
   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):
   '''
      Plot flux versus energy generation by generation
      Input: directory name
      Output: graph spectrum normalize to 1 photon source
   '''
   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")

   minEner=min(energy)
   alpha=flux[energy==minEner]/sqrt(minEner)
   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()