#!/bin/python
from numpy import zeros, size, nditer, average, savetxt, set_printoptions
from src.read import select_events
from src.analytic import degre, yr
set_printoptions(precision=2)

Redshifts=["0.04","0.14","0.4","1","2"] 

powerlaw_index=1.2

###############################################################################

#print "Compute mean obs vs EGMF strength ..."
#EGMFs=["12","13","14","15","16","17","18","19","20","21","22"]
#
#theta_mean  = zeros((size(Redshifts)+1,size(EGMFs)+1))
#theta_mean1 = theta_mean.copy()
#theta_mean2 = theta_mean.copy()
#dt_mean     = theta_mean.copy()
#dt_mean1    = theta_mean.copy()
#dt_mean2    = theta_mean.copy()
#ratio_gen   = theta_mean.copy()
#
#def compute_mean_delay_and_theta(i,j):
#   if i==0:
#      theta_mean[i,0] = 0
#      theta_mean1[i,0]= 0
#      theta_mean2[i,0]= 0
#      theta_mean[i,j] = 10**(-float(EGMFs[j-1]))
#      theta_mean1[i,j]= 10**(-float(EGMFs[j-1]))
#      theta_mean2[i,j]= 10**(-float(EGMFs[j-1]))
#      
#      dt_mean[i,0] = 0
#      dt_mean1[i,0]= 0
#      dt_mean2[i,0]= 0
#      dt_mean[i,j] = 10**(-float(EGMFs[j-1]))
#      dt_mean1[i,j]= 10**(-float(EGMFs[j-1]))
#      dt_mean2[i,j]= 10**(-float(EGMFs[j-1]))
#      
#      ratio_gen[i,0]= 0
#      ratio_gen[i,j]= 10**(-float(EGMFs[j-1]))
#      
#   else:
#      theta_mean[i,0]  = Redshifts[i-1]
#      theta_mean1[i,0] = Redshifts[i-1]
#      theta_mean2[i,0] = Redshifts[i-1]
#      dt_mean[i,0]     = Redshifts[i-1]
#      dt_mean1[i,0]    = Redshifts[i-1]
#      dt_mean2[i,0]    = Redshifts[i-1]
#      ratio_gen[i,0]   = Redshifts[i-1]
#      fileId="Simulations/z="+Redshifts[i-1]+"/EGMF=1e-"+EGMFs[j-1]+"G"

#      weight, energy, time, theta, phi, Esource, gen = select_events(fileId,Erange=[1,1e6],powerlaw_index=1.2)

#      dt_mean[i,j]=average(time,weights=weight)/yr
#      theta_mean[i,j]=average(theta,weights=weight)/degre
#      cond = (gen==2)
#      nb_gen1 = sum(weight[cond])
#      dt_mean1[i,j]=average(time[cond],weights=weight[cond])/yr
#      theta_mean1[i,j]=average(theta[cond],weights=weight[cond])/degre
#      cond = (gen==4)
#      nb_gen2 = sum(weight[cond])
#      dt_mean2[i,j]=average(time[cond],weights=weight[cond])/yr
#      theta_mean2[i,j]=average(theta[cond],weights=weight[cond])/degre
#      
#      ratio_gen[i,j-1] = nb_gen2/nb_gen1  
#
#it=nditer(theta_mean1, flags=['multi_index'], op_flags=['readwrite'])
#while not it.finished:
#   compute_mean_delay_and_theta(it.multi_index[0],it.multi_index[1])    
#   it.iternext()
#
#savetxt("Results/theta_mean_vs_EGMF.dat"      ,theta_mean ,fmt='%1.4e')
#savetxt("Results/theta_mean_vs_EGMF-gen=1.dat",theta_mean1,fmt='%1.4e')
#savetxt("Results/theta_mean_vs_EGMF-gen=2.dat",theta_mean2,fmt='%1.4e')
#savetxt("Results/delay_mean_vs_EGMF.dat"      ,dt_mean    ,fmt='%1.4e')
#savetxt("Results/delay_mean_vs_EGMF-gen=1.dat",dt_mean1   ,fmt='%1.4e')
#savetxt("Results/delay_mean_vs_EGMF-gen=2.dat",dt_mean2   ,fmt='%1.4e')
#savetxt("Results/ratio_gen_EGMF.dat"          ,ratio_gen  ,fmt='%1.4e')
#   
#print(ratio_gen)

###############################################################################
print "Compute mean obs vs EGMF coherence length ..."

L_Bs=["0.00001","0.0001","0.001","0.010","0.100","1","10","100","1000"]

theta_mean=zeros((size(Redshifts)+1,size(L_Bs)+1))
dt_mean=theta_mean.copy()

def compute_mean_delay_and_theta(i,j):
   if i==0:
      theta_mean[i,0]=0
      dt_mean[i,0]=0
      theta_mean[i,j]=L_Bs[j-1]
      dt_mean[i,j]=L_Bs[j-1]

   else:
      theta_mean[i,0]  = Redshifts[i-1]
      dt_mean[i,0]     = Redshifts[i-1]
      if (L_Bs[j-1]=="1"):
         fileId="Simulations/z="+Redshifts[i-1]+"/EGMF=1e-15G"
      else:
         fileId="Simulations/z="+Redshifts[i-1]+"/lambda_B="+L_Bs[j-1]+"Mpc"

         #fileId="Simulations/simple case/lambda_B="+L_Bs[j-1]+"Mpc"

      weight, energy, time, theta, phi, Esource, generation = select_events(fileId,Erange=[1,1e6],powerlaw_index=1.2)
      #weight, energy, time, theta, phi, Esource = select_events(fileId,Erange=[1,1e6],powerlaw_index=1.2, generation=2)

      dt_mean[i,j]=average(time,weights=weight)/yr
      theta_mean[i,j]=average(theta,weights=weight)/degre
   
it=nditer(theta_mean, flags=['multi_index'], op_flags=['readwrite'])
while not it.finished:
   compute_mean_delay_and_theta(it.multi_index[0],it.multi_index[1])    
   it.iternext()
   
set_printoptions(precision=2)
print(theta_mean)
print "========================================"
print(dt_mean)

savetxt("Results/theta_mean_vs_LB.dat",theta_mean,fmt='%1.4e')
savetxt("Results/delay_mean_vs_LB.dat",dt_mean,fmt='%1.4e')

###############################################################################

#print "Compute mean obs vs jet opening angle ..."
#opening_angle = [90,75,60,45,30,15,10,5,2,1,0.7,0.5,0.1,1e-2] # degres 
#
#def compute_theta_mean(i,j):
#   if i==0:
#      if j==0:
#         theta_mean[i,j]=0
#         dt_mean[i,0]=0
#      else:
#         theta_mean[i,j]=opening_angle[j-1]
#         dt_mean[i,j]=opening_angle[j-1]
#   else:
#      if j==0:
#         theta_mean[i,j]=float(Redshifts[i-1])
#      else:
#         fileId="Simulations/z="+Redshifts[i-1]+"/EGMF=1e-15G"

#         weight, energy, time, theta, phi, Esource, gen = select_events(fileId,Erange=[1,1e6],powerlaw_index=1.2,tjet=opening_angle[j-1])

#         theta_mean[i,j]=average(theta,weights=weight)/degre
#         dt_mean[i,j]=average(time,weights=weight)/yr
#   
#theta_mean=zeros((size(Redshifts)+1,size(opening_angle)+1))
#dt_mean=theta_mean.copy()
#it=nditer(theta_mean, flags=['multi_index'], op_flags=['readwrite'])
#while not it.finished:
#   compute_theta_mean(it.multi_index[0],it.multi_index[1])
#   it.iternext()
#    
#set_printoptions(precision=3)
#print(theta_mean)
#
#savetxt("Results/theta_mean_vs_jet_opening.dat",theta_mean,fmt='%1.4e')
#savetxt("Results/delay_mean_vs_jet_opening.dat",dt_mean,fmt='%1.4e')