from numpy import histogram, log10, shape, arange, pi, convolve, r_
from matplotlib.pyplot import figure, show
from Read import ReadTime, ReadTiming, ReadExtraFile, ReadProfile, ReadEnergy, ReadGeneration
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
   '''
   cond = time>0
   time = time[cond]
   time=log10(time)
   weight=weight[cond]

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

   return timecenter, dNdt#/max(dNdt)

def drawTiming(files,plot_gen=False,plot_energy_range=False):
   '''
      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:
      delta_t,dNdt = ReadTiming(fileId,[0,1])
      p=ax.plot(delta_t,dNdt,linestyle="steps-mid")

      if plot_energy_range:
         labels = ["MeV band","GeV band","TeV band"]
         i = 2
         for n in arange(0,3,1):
            delta_t,dNdt = ReadTiming(fileId,[0,i])
            ax.plot(delta_t,dNdt,linestyle="steps-mid",label=labels[n])
            i+=1

      if plot_gen:
         generation = ReadGeneration(fileId,[0])
         i=5
         for gen in generation:
            delta_t,dNdt = ReadTiming(fileId,[0,i])
            ax.plot(delta_t,dNdt,linestyle="steps-mid",label="gen=%.0f"%gen)
            i+=1

   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 drawConvolveTiming(files):
   fig = figure()
   ax = fig.add_subplot(111)
   for fileId in files:
      delta_t,dNdt1 = ReadTiming(fileId,[0,2]) # MeV band
      delta_t,dNdt2 = ReadTiming(fileId,[0,3]) # GeV band
      delta_t,dNdt3 = ReadTiming(fileId,[0,4]) # TeV band
      conv1 = convolve(dNdt1,dNdt2,'same')
      conv2 = convolve(dNdt2,dNdt3,'same')
      conv3 = convolve(dNdt1,dNdt3,'same')
      ax.plot(delta_t,conv1/max(conv1),linestyle="steps-mid",label="MeV - GeV band")
      ax.plot(delta_t,conv2/max(conv2),linestyle="steps-mid",label="GeV - TeV band")
      ax.plot(delta_t,conv3/max(conv3),linestyle="steps-mid",label="MeV - TeV band")
      print fileId, " ====================================="
      print "   MeV - GeV band: ",delta_t[r_[True, conv1[1:] > conv1[:-1]] & r_[conv1[:-1] > conv1[1:], True]] 
      print "   GeV - TeV band: ",delta_t[r_[True, conv2[1:] > conv2[:-1]] & r_[conv2[:-1] > conv2[1:], True]]
      print "   MeV - TeV band: ",delta_t[r_[True, conv3[1:] > conv3[:-1]] & r_[conv3[:-1] > conv3[1:], True]]

   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("cross correlation [arbitrary units]")

   show()

def delay_vs_theta(theta,delay,nbBins=100):
   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
   return dt, thetacenter

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

   generation, theta = ReadExtraFile(fileId,cols=[3,4])
   theta *= degre
   energy = ReadEnergy(fileId)
   delay  = ReadTime(fileId)

   Emin = [1e-3,1e0,1e3] #GeV
   Emax = [1e0,1e3,1e5] #GeV

   for gen in list(set(generation)):
      cond = generation==gen
      ax.plot(theta[cond],delay[cond],",",label="gen = %.0f"%gen)
   #for n in arange(0,3,1):
   #   cond= (energy>Emin[n]) & (energy<Emax[n])
   #   if Emin[n]==1e-3:
   #      label="MeV band"
   #   if Emin[n]==1e0:
   #      label="GeV band"
   #   if Emin[n]==1e3:
   #      label="TeV band"
   #   ax.plot(theta[cond],delay[cond],".",label=label)

   dt,angle=delay_vs_theta(theta,delay,nbBins)
   ax.plot(angle,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*degre,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$ [deg]")
   ax.set_ylabel("Time delay [s]")

   show()