from numpy import shape, arange, log10, histogram
from matplotlib.pyplot import figure, show, hist
from Read import ReadGeneration, ReadWeight, ReadEnergy


def drawWeightHisto(fileId):
   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):
   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):
   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()