from numpy import histogram, log10, shape, sqrt, arccos

from matplotlib.pyplot import figure, show
from scipy.integrate import quad

from Read import ReadTime, ReadWeight, ReadGeneration, ReadE_source, ReadD_source

from Read import ReadNphot_source,ReadPosition,ReadMomentum,ReadNbLeptonsProd,ReadEnergy

from Constants import yr, degre

from Integrand import comobileTime

from Analytic import Analytic_delay_vs_theta,Analytic_delay_vs_Egamma

def timing(fileId,gen=-1):
   '''
      Compute flux versus time delay
      Input: directory name
      Input (optional): generation (default = all)
      Output: time delay (sec), flux
   '''
   time=ReadTime(fileId)
   weight=ReadWeight(fileId)

   if gen != -1:
      generation = ReadGeneration(fileId)
      time=time[generation==gen]
      weight=weight[generation==gen]
      prop = float(shape(time[time<0])[0])/shape(time)[0]
      print "gen", int(gen), "->", shape(time)[0], "events", "negative time:",shape(time[time<0])[0], "~", prop
   else:
      prop = float(shape(time[time<0])[0])/shape(time)[0]
      print "file", fileId, "->", shape(time)[0], "events", "negative time:",shape(time[time<0])[0], "~", prop

   time=time[time>0]
   weight=weight[time>0]

   nbBins = 100

   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]

   Nmax=ReadNphot_source(fileId) # Nb photons emitted by the source

   dNdt=dN/(Nmax*binSize)

   maxflux = ReadNphot_source(fileId)

   dNdt /= maxflux

   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:
      time,dNdt = timing(fileId) 
      ax.plot(time[dNdt!=0],dNdt[dNdt!=0],linestyle="steps-mid",label=fileId)

   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$")

   show()

def drawTimingGen(fileId):
   fig = figure()
   ax = fig.add_subplot(111)
   for gen in list(set(ReadGeneration(fileId))):
      time,dNdt = timing(fileId,gen) 
      ax.plot(time[dNdt!=0],dNdt[dNdt!=0],linestyle="steps-mid",label="gen="+str(int(gen)))

   time,dNdt = timing(fileId) 
   ax.plot(time[dNdt!=0],dNdt[dNdt!=0],color='k',linestyle="steps-mid",label="gen=all")

   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$")

   show()

######## delay versus angle ###############
def delay_vs_angle(fileId,gen=-1):
   '''
      Compute time delay versus arrival angle
      Input: directory name
      Input (optional): generation (default = all)
      Output: arrival angle [degre], time delay [sec]
   '''
   position = ReadPosition(fileId)
   momentum = ReadMomentum(fileId)

   energy = ReadEnergy(fileId)

   weight = ReadWeight(fileId)
   time=ReadTime(fileId)

   generation = ReadGeneration(fileId)

   hyp=sqrt(position[0]**2+position[1]**2+position[2]**2)
   position /= hyp
   hyp=sqrt(momentum[0]**2+momentum[1]**2+momentum[2]**2)
   momentum /= hyp
   costheta=position[0]*momentum[0]+position[1]*momentum[1]+position[2]*momentum[2]

   if gen != -1 :

      cond=(abs(costheta)<=1) & (generation==gen) & (time>0) & (energy>1)

   else:

      cond=(abs(costheta)<=1) & (time>0) & (energy>1)

   prop = float(shape(time[time<0])[0])/shape(time)[0]

   print "gen", int(gen), "->", shape(time[time<0])[0], "negative time events among ",shape(time)[0], "~", prop

   theta = arccos(costheta[cond])*degre
   weight= weight[cond]
   time=time[cond]*yr

   cond=(theta!=0)

   return theta[cond], time[cond], weight[cond], generation[cond]

def delay_vs_angle_histo(fileId,gen=-1):

   theta, time, weight, generation = delay_vs_angle(fileId,gen)

   nbBins = 100

   theta=log10(theta)
   dt,dtheta=histogram(theta,nbBins,weights=time)
   dN,dtheta=histogram(theta,nbBins,weights=weight)
   dtheta=10**dtheta
   thetacenter=(dtheta[1:nbBins+1]+dtheta[0:nbBins])/2
   timing = dt/dN

   return thetacenter, timing

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)
   

   angle,delay,weight,generation = delay_vs_angle(fileId)
   #for gen in list(set(generation)):
   for gen in [2]:
      cond= (generation==gen)
      ax.plot(angle[cond],delay[cond]/weight[cond],".",label="gen="+str(int(gen)))

   angle,delay = delay_vs_angle_histo(fileId)

   #ax.plot(angle,delay,color='k',linestyle="steps-mid",label="gen=all")

   delay_fit = Analytic_delay_vs_theta(angle,fileId)
   ax.plot(angle,delay_fit,'--r',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()

######## delay versus angle ###############
def delay_vs_energy(fileId,gen=-1):
   '''
      Compute time delay versus arrival angle
      Input: directory name
      Input (optional): generation (default = all)
      Output: arrival angle [degre], time delay [sec]
   '''
   energy = ReadEnergy(fileId)
   weight = ReadWeight(fileId)
   time=ReadTime(fileId)
   generation = ReadGeneration(fileId)

   if gen != -1 :
      cond= (generation==gen) & (time>0) & (energy>1)
   else:
      cond= (time>0) & (energy>1)

   prop = float(shape(time[time<0])[0])/shape(time)[0]
   print "gen", int(gen), "->", shape(time[time<0])[0], "negative time events among ",shape(time)[0], "~", prop
   return energy[cond], time[cond]*yr, weight[cond], generation[cond]

def delay_vs_energy_histo(fileId,gen=-1):
   energy, time, weight, generation = delay_vs_energy(fileId,gen)
   nbBins = 100

   energy=log10(energy)
   dt,dE=histogram(energy,nbBins,weights=time)
   dN,dE=histogram(energy,nbBins,weights=weight)
   dE=10**dE
   Ecenter=(dE[1:nbBins+1]+dE[0:nbBins])/2
   timing = dt/dN

   return Ecenter, timing

def drawDelay_vs_energy(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)
   
   energy,delay,weight,generation = delay_vs_energy(fileId)
   #for gen in list(set(generation)):
   for gen in [2]:
      cond= (generation==gen)
      ax.plot(energy[cond],delay[cond]/weight[cond],".",label="gen="+str(int(gen)))

   energy,delay = delay_vs_energy_histo(fileId)
   #ax.plot(energy,delay,color='k',linestyle="steps-mid",label="gen=all")

   delay_fit = Analytic_delay_vs_Egamma(energy,fileId)
   ax.plot(energy,delay_fit,'--r',linewidth=2)

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

   show()