from numpy import shape, arange, log10, histogram
from matplotlib.pyplot import figure, show, hist

from read import ReadGeneration, ReadWeight, ReadEnergy

def drawWeightHisto(fileId):

   '''
      Plot particles'weights histogram, generation by generation
      Input: directory name
      Output: particles'weights histogram
   '''

   fig = figure()
   ax = fig.add_subplot(111)
   nbBins =100 

   weight = ReadWeight(fileId)
   generation = ReadGeneration(fileId)
   print "file", fileId, "->", shape(generation)[0], "events" 

   for gen in list(set(generation)):
      hist(log10(weight[generation==gen]),nbBins,log=1,range=[0,7],alpha=.5,
            label="gen="+str(int(gen)))

   ax.legend(loc='best')
   ax.set_xlabel("Weight [log]")
   ax.set_ylabel("Number of events")

   show()

#######################################################################################
def weightVsEnergy(fileId,gen=-1):

   '''
      Compute particles'weight versus energy
      Input: directory name
      Input (optional): generation (default = all)
   '''

   energy = ReadEnergy(fileId)
   weight = ReadWeight(fileId)
   nbBins = 100 

   if gen != -1:
      generation = ReadGeneration(fileId)
      energy=energy[generation==gen]
      weight=weight[generation==gen]
      nbBins =40 

   weight,ener=histogram(log10(energy),nbBins,weights=log10(weight))
   nb,ener=histogram(log10(energy),nbBins)
   ener=10**ener
   enercenter=(ener[1:nbBins+1]+ener[0:nbBins])/2
   weight /= nb

   return enercenter, weight

def drawWeightVsEnergy(fileId):

   '''
      Plot particles'weight versus energy, generation by generation
      Input: directory name
      Output: graph of particles'weight versus energy
   '''

   fig = figure()
   ax = fig.add_subplot(111)

   generation = ReadGeneration(fileId)
   print "file", fileId, "->", shape(generation)[0], "events" 

   for gen in list(set(generation)):
      energy,weight = weightVsEnergy(fileId,gen)
      ax.plot(energy,weight,".",label="gen="+str(int(gen)))

   energy,weight = weightVsEnergy(fileId)
   ax.plot(energy,weight,"k",label="gen=all")

   ax.set_xscale('log')
   ax.grid(b=True,which='major')
   ax.legend(loc="best")
   ax.set_title(fileId)
   ax.set_xlabel("energy [GeV]")
   ax.set_ylabel("mean weight [log]")

   show()