from numpy import histogram, log10,logspace, shape, arange, pi, r_, diff, sqrt, correlate
from numpy import append, zeros#, conj
#from numpy.fft import rfft, irfft
from matplotlib.pyplot import figure, show
from read import ReadResults, ReadProfile, ReadTiming, ReadDelayVsAngle, ReadGeneration
from constants import yr, degre
from analytic import Analytic_delay_vs_theta

def timing(time,weight,nbBins=100000,max_dt=315576000): # max_dt=10yr
   '''
      Compute flux versus time delay
      Input: directory name
      Input (optional): generation (default = all)
      Output: time delay (sec), flux
   '''
   cond = (time>1e3) #& (time<max_dt)  
   time = time[cond]

   time=log10(time)

   weight=weight[cond]

   #dN,dt=histogram(time,nbBins,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

   return timecenter, dNdt#/max(dNdt)

def drawTiming(files,plot=""):

   '''
      Plot flux versus time delay
      Input: list of directories
      Output: graph of flux versus time delay 
   '''

   fig = figure(figsize=(12,9))

   ax = fig.add_subplot(111)
   for fileId in files:
      delta_t,dNdt = ReadTiming(fileId,[0,1]) # TeV band

      p=ax.plot(delta_t,dNdt,linestyle="steps-mid",label=fileId)

      if plot=="energy band":

         delta_t,dNdt1,dNdt2,dNdt3 = ReadTiming(fileId,[0,2,3,4])

         ax.plot(delta_t,dNdt1,drawstyle='steps-mid',linestyle='--',label="MeV band")
         ax.plot(delta_t,dNdt2,drawstyle='steps-mid',linestyle='--',label="GeV band")
         ax.plot(delta_t,dNdt3,drawstyle='steps-mid',linestyle='--',label="TeV band")

      elif plot=="generation":

         delta_t,dNdt1,dNdt2,dNdt3 = ReadTiming(fileId,[0,5,6,7])

         ax.plot(delta_t,dNdt1,drawstyle='steps-mid',linestyle='--',label="gen=1")
         ax.plot(delta_t,dNdt2,drawstyle='steps-mid',linestyle='--',label="gen=2")
         ax.plot(delta_t,dNdt3,drawstyle='steps-mid',linestyle='--',label="gen=3")

         #i=5
         #for gen in ReadGeneration(fileId,[0]):
         #   delta_t,dNdt = ReadTiming(fileId,[0,i])
         #   ax.plot(delta_t,dNdt,drawstyle='steps-mid',linestyle='--',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("$dN/dt$ [arbitrary units]")

   show()

def drawConvolveTiming(files):

   fig = figure(figsize=(12,9))

   ax = fig.add_subplot(111)
   for fileId in files:
      #print fileId, " ====================================="
      dt,N_MeV,N_GeV,N_TeV = ReadTiming(fileId,[0,2,3,4]) # MeV GeV TeV band
      delta_t = -dt[::-1]
      delta_t = append(delta_t,dt)
      dNdt_MeV = zeros(shape(dt)[0])
      dNdt_MeV = append(dNdt_MeV,N_MeV)
      dNdt_GeV = zeros(shape(dt)[0]) 
      dNdt_GeV = append(dNdt_GeV,N_GeV)
      dNdt_TeV = zeros(shape(dt)[0]) 
      dNdt_TeV = append(dNdt_TeV,N_TeV)
      #fft_MeV = rfft(dNdt_MeV)
      #fft_GeV = rfft(dNdt_GeV)
      #fft_TeV = rfft(dNdt_TeV)
      #CC_MeV_GeV = irfft(conj(fft_MeV) * fft_GeV)
      #CC_GeV_TeV = irfft(conj(fft_GeV) * fft_TeV)
      #CC_MeV_TeV = irfft(conj(fft_MeV) * fft_TeV)
      ac_MeV = correlate(dNdt_MeV,dNdt_MeV)[0]
      ac_GeV = correlate(dNdt_GeV,dNdt_GeV)[0]
      ac_TeV = correlate(dNdt_TeV,dNdt_TeV)[0]
      autocorr_MeV = correlate(dNdt_MeV,dNdt_MeV,'same')[0]
      autocorr_GeV = correlate(dNdt_GeV,dNdt_GeV,'same')[0] 
      autocorr_TeV = correlate(dNdt_TeV,dNdt_TeV,'same')[0]
      CC_MeV_GeV = correlate(dNdt_MeV,dNdt_GeV,'same')/sqrt(ac_GeV*ac_MeV)
      CC_GeV_TeV = correlate(dNdt_GeV,dNdt_TeV,'same')/sqrt(ac_TeV*ac_GeV)
      CC_MeV_TeV = correlate(dNdt_MeV,dNdt_TeV,'same')/sqrt(ac_TeV*ac_MeV)
      if shape(files)[0]==1:
         #ax.plot(delta_t,autocorr_MeV/ac_MeV,linestyle="steps-mid",label="MeV band")
         #ax.plot(delta_t,autocorr_GeV/ac_GeV,linestyle="steps-mid",label="GeV band")
         #ax.plot(delta_t,autocorr_TeV/ac_TeV,linestyle="steps-mid",label="TeV band")
         ax.plot(delta_t,CC_MeV_GeV,linestyle="steps-mid",label="MeV VS GeV band")
         #print "   MeV - GeV band: ",delta_t[r_[True, conv1[1:] > conv1[:-1]] & r_[conv1[:-1] > conv1[1:], True]] 
         ax.plot(delta_t,CC_GeV_TeV,linestyle="steps-mid",label="GeV VS TeV band")
         #print("   GeV - TeV band: ",delta_t[r_[True, conv2[1:] > conv2[:-1]] & r_[conv2[:-1] >   conv2[1:], True]])
         ax.plot(delta_t,CC_MeV_TeV,linestyle="steps-mid",label="MeV VS TeV band")
         #print "   MeV - TeV band: ",delta_t[r_[True, conv3[1:] > conv3[:-1]] & r_[conv3[:-1] > conv3[1:], True]]
      else:
         ax.plot(delta_t,CC_GeV_TeV,linestyle="steps-mid",label=fileId)

   #ax.set_xscale('log')
   #ax.set_yscale('log')
   ax.set_xlim([-1e16,1e16])
   ax.legend(loc="best")
   ax.grid(b=True,which='major')
   ax.set_xlabel("Time delay [s]")
   ax.set_ylabel("cross correlation [normalized]")

   show()