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


def drawGeneration(files):
   '''
      Plot histogram of particles'generations
      Input: list of directories
      Output: histogram of generations
   '''
   fig = figure()
   ax = fig.add_subplot(111)
   nbBins = 7
   for fileId in files:
      weight = ReadWeight(fileId)
      generation = ReadGeneration(fileId)
      print "file", fileId, "->", shape(generation)[0], "events" 
      hist(generation,nbBins,range=[2,9],log=1,alpha=.5, weights=weight,label=fileId)

   xtick=arange(9)
   ax.legend(loc='best')
   ax.set_xticks(xtick+0.5)
   ax.set_xticklabels(xtick)
   ax.set_xlabel("Generation")
   ax.set_ylabel("Number of events")

   show()