from numpy import histogram, log10, shape, arange, pi
from matplotlib.pyplot import figure, show
from Read import ReadTime, ReadTiming, ReadExtraFile, ReadProfile, ReadEnergy
from Constants import yr, degre
from Analytic import Analytic_delay_vs_theta

def timing(time,weight,nbBins=100):
   '''
      Compute flux versus time delay
      Input: directory name
      Input (optional): generation (default = all)
      Output: time delay (sec), flux
   '''
   time=time[time>0]
   weight=weight[time>0]

   time = log10(time*yr)

   dN,dt=histogram(time,nbBins,weights=weight)
   timecenter=(dt[1:nbBins+1]+dt[0:nbBins])/2
   binSize=dt[1:nbBins+1]-dt[0:nbBins]
   dNdt=dN/(max(dN)*binSize)

   return 10**timecenter, dNdt

def drawTiming(files):
   '''
      Plot flux versus time delay
      Input: list of directories
      Output: graph of flux versus time delay 
   '''
   fig = figure()
   ax = fig.add_subplot(111)
   for fileId in files:
      i=1
      for powerlaw_index in [1,1.5,2,2.5]: 
         delta_t,dNdt = ReadTiming(fileId,[0,i])
         ax.plot(delta_t,dNdt,linestyle="steps-mid",label="p=%.1f"%powerlaw_index)

   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.grid(b=True,which='major')
   ax.legend(loc="best")
   ax.set_xlabel("Time delay [s]")
   ax.set_ylabel("$t\ dN/dt$ [arbitrary units]")

   show()

def drawDelay_vs_angle(fileId):
   '''
      Plot angle versus time delay, generation by generation
      Input: directory name
      Output: graph of angle versus time delay, generation by generation
   '''
   fig = figure()
   ax = fig.add_subplot(111)
   nbBins = 100

   theta  = ReadExtraFile(fileId,[4])
   energy = ReadEnergy(fileId)
   delay  = ReadTime(fileId)

   for n in arange(0,8,2):
      Emin = 10**(3-n) #GeV
      Emax = 10**(5-n) #GeV
      label = "%1.0e Gev"%Emin+"< $E_\\gamma$ < %1.0e Gev"%Emax
      cond= (energy>Emin) & (energy<Emax)
      ax.plot(theta[cond],delay[cond],".",label=label)

   cond= (theta!=0)
   theta=log10(theta[cond])
   dt,dtheta=histogram(theta,nbBins,weights=delay[cond])
   dN,dtheta=histogram(theta,nbBins)
   dtheta=10**dtheta
   thetacenter=(dtheta[1:nbBins+1]+dtheta[0:nbBins])/2
   dt=dt/dN
   ax.plot(thetacenter,dt,color='m',linestyle="steps-mid",linewidth=2)

   thmin = 1.e-8
   thmax = pi/2.
   nth = 1000
   th = thmin*(thmax/thmin)**(arange(nth)/(nth-1.))
   delay_fit = Analytic_delay_vs_theta(th,fileId)
   ax.plot(th,delay_fit,'--k',linewidth=2)

   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.grid(b=True,which='major')
   ax.legend(loc="best")
   ax.set_xlabel("$\\theta $[rad]")
   ax.set_ylabel("Time delay [s]")

   show()