from numpy import histogram, log10,logspace, shape, arange, pi, r_, diff, sqrt, correlate
from numpy import append, zeros, sin, size
from matplotlib.pyplot import figure, show, setp, legend, gca
from matplotlib import gridspec, rcParams
from read import ReadProfile, ReadTiming, ReadSpectrum, ReadSourceSpectrum, ReadElmag, ReadTaylor
from read import ReadRadial, ReadAngleVsEnergy
from analytic import Ethreshold_gg, Analytic_flux, Eic, Analytic_dNdt, Analytic_dNdtheta
from constants import yr , degre

label_size = 20
rcParams['xtick.labelsize'] = label_size
rcParams['ytick.labelsize'] = label_size

#====================================================================================
#       Compute distributions (dN/dE, dN/dtheta, dN/dt)
#====================================================================================
def distribution(x,weight,nbBins=100,x_range=[]):
   '''
      Compute distribution dN/dx vs x (x could be E, theta or time)
      Input: events x and weight
             (optional: number of bins and x range)
      Output: mean x [x_unit], dN/dx [x_unit^-1]
   '''
   if x_range==[]:
      x_range[0]=min(x)
      x_range[1]=max(x)
   dN,xi = histogram(log10(x[x>0]),nbBins,range=log10(x_range),weights=weight[x>0])
   xi = 10**xi
   xc = (xi[1:nbBins+1]+xi[0:nbBins])/2
   dx = xi[1:nbBins+1]-xi[0:nbBins]
   return xc, dN/dx

#====================================================================================
#       Draw distributions (dN/dE, dN/dtheta, dN/dt)
#====================================================================================
def which_label(fileId):
      if "Gamma" in fileId:
         return "$\\"+fileId+"$"
      elif "lambda" in fileId:
         return "$\\"+fileId+"$"
      elif "theta" in fileId:
         return "$\\"+fileId+"^\circ$"
      elif "tobs" in fileId:
         return "$\\theta_{obs}="+fileId.split("tobs=")[1]+"^\circ$"
      else:
         return fileId

def drawSpectrum(files,plot_analytic=False,plot_other_codes=False,plot="",Erange=[1e0,1e5],yrange=[1e-4,1e0],colors=['k','b','g','r','c','m','y']):
   '''
      Plot flux versus energy
      Input: list of directories
      Input (optional, default=False): plot analytic expression
      Output: graph spectrum normalize to 1 source photon
   '''
   fig = figure(figsize=(12,9))
   #gs = gridspec.GridSpec(2, 1, height_ratios=[3,1]) 
   #fig.subplots_adjust(hspace=0.001)
   #ax1 = fig.add_subplot(gs[0])
   #ax2 = fig.add_subplot(gs[1],sharex=ax1)
   ax1 = fig.add_subplot(111)
   i=0


   for fileId in files:
      lab = which_label(fileId)
      z = ReadProfile(fileId,[3])

      if plot=="1month":
         ener,dNdE = ReadSpectrum(fileId,[0,8])
         ax1.plot(ener,ener**2*dNdE *(1+z),drawstyle='steps-mid',color=colors[i],linewidth=2,label=lab)
      else:
         # draw full spectrum
         ener,dNdE = ReadSpectrum(fileId,[0,1])
         if plot=="cascade contribution":
            p, = ax1.plot(ener,ener**2*dNdE *(1+z),drawstyle='steps-mid',linewidth=2,color=colors[i],label=lab)
            ener,dNdE0,dNdE1,dNdE2,dNdE3,dNdE4 = ReadSpectrum(fileId,[0,2,3,4,5,6])
            p1,=ax1.plot(ener,ener**2 *(dNdE1+dNdE2+dNdE3+dNdE4) *(1+z),color=p.get_color(),linewidth=1,linestyle='--')
            p2,=ax1.plot(ener,ener**2 *dNdE0 *(1+z),color=p.get_color(),linewidth=1,linestyle=':')
            leg1 = legend([p,p1,p2], ["all photons","cascade photons","source photons"],fontsize=label_size, loc=3)
         elif plot=="generation":
            p, = ax1.plot(ener,ener**2*dNdE *(1+z),drawstyle='steps-mid',linewidth=2,color=colors[i],label=lab)#"all")
            ener,dNdE1,dNdE2,dNdE3 = ReadSpectrum(fileId,[0,3,4,5])  
            ax1.plot(ener,ener**2*dNdE1 *(1+z),color=p.get_color(),drawstyle='steps-mid',linestyle='--',linewidth=2)#,label="gen=1")
            ax1.plot(ener,ener**2*dNdE2 *(1+z),color=p.get_color(),drawstyle='steps-mid',linestyle='-.',linewidth=2)#,label="gen=2")
            ax1.plot(ener,ener**2*dNdE3 *(1+z),color=p.get_color(),drawstyle='steps-mid',linestyle=':',linewidth=2)#,label="gen=3")
         elif plot=="obs time":
            p, = ax1.plot(ener,ener**2 *dNdE *(1+z),drawstyle='steps-mid',linewidth=2,color='k',label="$t_{obs}=\infty$")
            ener,dNdE1,dNdE2,dNdE3,dNdE4,dNdE5 = ReadSpectrum(fileId,[0,7,8,9,10,11])
            ax1.plot(ener,ener**2 *dNdE5 *(1+z),drawstyle='steps-mid',linestyle='-',linewidth=2,label="$t_{obs}$=$10^6$ yr")
            ax1.plot(ener,ener**2 *dNdE4 *(1+z),drawstyle='steps-mid',linestyle='-',linewidth=2,label="$t_{obs}$=$10^3$ yr")
            ax1.plot(ener,ener**2 *dNdE3 *(1+z),drawstyle='steps-mid',linestyle='-',linewidth=2,label="$t_{obs}$=1 yr")
            ax1.plot(ener,ener**2 *dNdE2 *(1+z),drawstyle='steps-mid',linestyle='-',linewidth=2,label="$t_{obs}$=1 month")
            ax1.plot(ener,ener**2 *dNdE1 *(1+z),drawstyle='steps-mid',linestyle='-',linewidth=2,label="$t_{obs}$=1 day")
         else:
            p, = ax1.plot(ener,ener**2 *dNdE *(1+z),drawstyle='steps-mid',linewidth=2,color=colors[i],label=lab)

      i+=1
         
   #minEner=min(ener)
   #maxEner=100e3
   #Ec = 309e3
   #alpha = (1-(4*ener*Ec/maxEner**2)**0.1) *3e-2 #dNdE[ener==minEner]*minEner**3*(1+z)/(minEner**(2-1.6))
   #ax1.plot(ener,alpha*ener**(2-1.6),'--m',linewidth=2)
   #alpha=2e-2 #dNdE[ener==minEner]*minEner**2*(1+z)/(minEner**(2-1.8))
   #ax1.plot(ener,alpha*ener**(2-1.8),'-.m',linewidth=2)

   if plot_analytic:
      #minEner=min(ener)
      #alpha=flux[ener==minEner]/sqrt(minEner)
      #ax1.plot(ener,alpha*sqrt(ener),'--m')
      #=== Gen=2 ===
      #ax1.plot(ener,2*Analytic_flux(ener),'--g')
      #=== Gen=4 ===
      Elim= Eic(Ethreshold_gg()/2)
      alpha = 0.75*9.262e3/Elim**(1/4.)
      alpha = dNdE[0]*ener[0]**2/sqrt(ener[0])
      ax1.plot(ener,2*ener**2*Analytic_flux(ener,z),'--b',linewidth=2)
      #alpha = 2*29.29e3
      #ax1.plot(ener,alpha*(ener/100)**(1/4.),'--m',linewidth=1)
      #ax1.plot(ener,ener**(0)*max(flux)*0.95,'--m',linewidth=2)
      #ax1.axvline(x=Ethreshold_gg(), ymin=0., ymax = 1., color='r', linewidth=2)
      #ax1.axvline(x=1e4, ymin=0., ymax = 1., color='g', linewidth=2)
        
   if plot_other_codes:
      coef = max(ener**2*dNdE)
      ener,E2dNdE,ener_PSF,E2dNdE_PSF = ReadElmag(fileId)
      E2dNdE *= coef/max(E2dNdE)
      E2dNdE_PSF *= coef/max(E2dNdE_PSF)
      ax1.plot(ener,E2dNdE*(1+z),"k-.",linewidth=2,label="Elmag")
      #ax1.plot(ener_PSF,E2dNdE_PSF,"b.",label="Elmag - Fermi/LAT PSF")
      ener,E2dNdE = ReadTaylor(fileId)
      E2dNdE *= coef/max(E2dNdE)
      ax1.plot(ener,E2dNdE*(1+z),"k:",linewidth=2,label="Taylor 2011")

   ax1.legend(loc="best",fontsize=label_size)
   #ax1.legend(loc=3,fontsize=label_size)
   ax1.set_xscale('log')
   ax1.set_yscale('log')
   ax1.grid(b=True,which='major')
   if "simple case" in fileId:
       ax1.set_ylabel("$E^2 dN/dE$ [GeV]",fontsize=label_size)
   else:
       ax1.set_ylabel("$E^2 dN/dE / L_{0}$",fontsize=label_size)
   ax1.set_xlabel("$E$ [GeV]",fontsize=label_size)
   ax1.set_xlim(Erange)
   ax1.set_ylim(yrange)
   #ax2.set_xscale('log')
   #ax2.grid(b=True,which='major')
   #ax2.set_xlabel("energy [GeV]",fontsize=label_size)
   #xticklabels = ax1.get_xticklabels()
   #setp(xticklabels, visible=False)
   if plot=="cascade contribution":
      ax1 = gca().add_artist(leg1)

   show()

def drawArrivalAngle(files,plot_analytic=False,plot="",colors=['k','b','g','r','c','m','y'],yrange=[1e0,1e7]):
   '''
      Plot flux versus arrival angle
      Input: list of directories
      Output: graph of flux versus angle
   '''
   fig = figure(figsize=(12,9))
   ax = fig.add_subplot(111)
   i=0 
   for fileId in files:
      lab = which_label(fileId)

      if plot=="generation":
         theta,dNdtheta,dNdtheta1,dNdtheta2,dNdtheta3 = ReadRadial(fileId,[0,1,3,4,5])
         theta /= degre
         p,=ax.plot(theta,theta*dNdtheta*degre,drawstyle='steps-mid',color=colors[i],linewidth=2,label=lab)#,label="all gen")
         ax.plot(theta,theta*dNdtheta1*degre,color=p.get_color(),drawstyle='steps-mid',linestyle='--',linewidth=2)#,label="gen=1")
         ax.plot(theta,theta*dNdtheta2*degre,color=p.get_color(),drawstyle='steps-mid',linestyle='-.',linewidth=2)#,label="gen=2")
         ax.plot(theta,theta*dNdtheta3*degre,color=p.get_color(),drawstyle='steps-mid',linestyle=':',linewidth=2)#,label="gen=3")

      elif plot=="obs time":
         theta,dNdtheta,dNdtheta1,dNdtheta2,dNdtheta3,dNdtheta4 = ReadRadial(fileId,[0,1,7,8,9,10])
         theta /= degre
         ax.plot(theta,theta*dNdtheta1*degre,drawstyle='steps-mid',linestyle='--',label="$t_{obs}$=1 day")
         ax.plot(theta,theta*dNdtheta2*degre,drawstyle='steps-mid',linestyle='--',label="$t_{obs}$=10 days")
         ax.plot(theta,theta*dNdtheta3*degre,drawstyle='steps-mid',linestyle='--',label="$t_{obs}$=1 month")
         ax.plot(theta,theta*dNdtheta3*degre,drawstyle='steps-mid',linestyle='--',label="$t_{obs}$=1 yr")
         ax.plot(theta,theta*dNdtheta*degre,color="k",drawstyle='steps-mid',linewidth=2,label="$t_{obs}$=\\infty")
        
      else:
         theta,dndtheta = ReadRadial(fileId,[0,1])
         theta /= degre
         ax.plot(theta,theta*dndtheta*degre,drawstyle='steps-mid',linewidth=2,color=colors[i],label=lab)

      i+=1   
        
   if plot_analytic:         
      B,Ds = ReadProfile(fileId,[0,2]) #Gauss #Mpc 
      lgg = 1.32 #lambda_gg(Eg0*1e3,zs)[1] #Mpc
      tau_gg = Ds/lgg
      ax.plot(theta,2*theta*Analytic_dNdtheta(theta,B,tau_gg),'--b',linewidth=2) 

   ax.legend(loc="best",fontsize=label_size)
   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.set_xlim([1e-5,180])
   ax.set_ylim(yrange)
   ax.grid(b=True,which='major')
   ax.set_xlabel("$\\theta$ [deg]",fontsize=label_size)
   ax.set_ylabel("$\\theta dN/d\\theta$ [/ primary phot.]",fontsize=label_size)
   
   show()

def drawTiming(files,plot_analytic=False,plot="",colors=['k','b','g','r','c','m','y'],yrange=[1e0,5e3]):
   '''
      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)
   i=0
   for fileId in files:
      lab = which_label(fileId)

      if plot=="generation":
         delta_t,dNdt,dNdt1,dNdt2,dNdt3 = ReadTiming(fileId,[0,1,3,4,5])
         p,=ax.plot(delta_t/yr,delta_t*dNdt,linestyle="steps-mid",color=colors[i],linewidth=2,label=lab)#"all gen")
         ax.plot(delta_t/yr,delta_t*dNdt1,color=p.get_color(),drawstyle='steps-mid',linestyle='--',linewidth=2)#,label="gen=1")
         ax.plot(delta_t/yr,delta_t*dNdt2,color=p.get_color(),drawstyle='steps-mid',linestyle='-.',linewidth=2)#,label="gen=2")
         ax.plot(delta_t/yr,delta_t*dNdt3,color=p.get_color(),drawstyle='steps-mid',linestyle=':',linewidth=2)#,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
      else:
         delta_t,dNdt = ReadTiming(fileId,[0,1])
         p,=ax.plot(delta_t/yr,delta_t*dNdt,linestyle="steps-mid",linewidth=2,color=colors[i],label=lab)

      i+=1   
        
   if plot_analytic:         
      B = ReadProfile(fileId,[0]) #Gauss  
      lgg = 1.32 #lambda_gg(Eg0*1e3,zs)[1] #Mpc
      ax.plot(delta_t,2*Analytic_dNdt(delta_t,B,lgg),'--b',linewidth=2) 
            
   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.set_xlim([1e-4,1e10])
   ax.set_ylim(yrange)
   ax.grid(b=True,which='major')
   ax.legend(loc="best",fontsize=label_size)
   ax.set_xlabel("$\\Delta t$ [yr]",fontsize=label_size)
   ax.set_ylabel("$\\Delta t\  dN/dt$ [/ primary phot.]",fontsize=label_size)

   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",fontsize=label_size)
   ax.grid(b=True,which='major')
   ax.set_xlabel("Time delay [s]",fontsize=label_size)
   ax.set_ylabel("cross correlation [normalized]",fontsize=label_size)

   show()