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

def drawGeneration(files,plot=""):
   '''
      Plot histogram of particles'generations
      Input: list of directories
      Output: histogram of generations
   '''
   fig = figure(figsize=(12,9))
   ax = fig.add_subplot(111)
   for fileId in files:
      if plot=="generation":
         i=1
         for powerlaw_index in [1,1.5,2,2.5]: 
            generation, ratio = ReadGeneration(fileId,[0,i])
            ax.plot(generation,ratio,drawstyle='steps-mid',label="p=%.1f"%powerlaw_index)
            i+=1
      else:
         generation, ratio = ReadGeneration(fileId,[0,3])
         ax.plot(generation,ratio,drawstyle='steps-mid',label=fileId)

   xtick=arange(9)
   ax.legend(loc='best')
   ax.set_yscale('log')
   ax.set_xlabel("Generation")
   ax.set_ylabel("Ratio of events [%]")

   show()