from Constants import *
from numpy import sqrt, abs, cos, sin, arcsin, searchsorted, loadtxt
from Read import ReadProfile

# Value for analytic expression
Ee_default=1e3 #GeV
Egamma_default=1 #GeV
B_default=1e-17 #Gauss
lambda_B_default=0.3 #Mpc

def Analytic_flux(E_ic):
   return  me*1e3/4*sqrt(3e9/Ecmb)*1e-9*E_ic**(1/2.) #GeV

def Analytic_theta(E_ic,fileId):
   Esource,Dsource,B=ReadProfile(fileId,[0,2,4]) #GeV #Mpc #Gauss
   delta_to_theta=lambda_gg(Esource)/Dsource
   delta_to_theta=1.9440974051313842/Dsource
   RL0=RL(Esource/2,B)
   Dic0=Dic(Esource/2)
   delta=Dic0/(2*RL0)*((Esource/(2*Ee(E_ic)))**2 -1)
   return abs(arcsin(delta_to_theta*sin(delta)))*degre

def Analytic_delay_vs_theta(theta,fileId):
   Esource,distSource=ReadProfile(fileId,[0,2]) #GeV #Mpc 
   E0_source = Esource*1e-3 #TeV
   lgg = 1.94 #lambda_gg(E0_source)
   delta_ic = arcsin(distSource/lgg*sin(theta))
   Dic0=Dic(E0_source/2)
   c_delta_t = lgg*(1-cos(delta_ic)) - distSource*(1-cos(theta))
   return c_delta_t *Mpc/c # sec.

def dt_approx(th):
   lgg = 1.94*Mpc
   D = 131.75*Mpc
   thetamax = arcsin(lgg/D)
   return D**2/lgg /c * th**2 / 2.

def Analytic_delay_vs_Egamma(Egamma, fileId):
   Esource,Dsource,B=ReadProfile(fileId,[0,2,4]) #GeV #Mpc #Gauss
   RL0=RL(Esource/2,B)
   Dic0=Dic(Esource/2)
   E_e = Ee(Egamma)
   delta_ic = Dic0/(2*RL0)*((Esource/2/E_e)**2 -1)
   theta = arcsin(lambda_gg(Esource)/Dsource*sin(delta_ic))
   c_delta_t = lambda_gg(Esource)*(1-cos(delta_ic)) - Dsource*(1-cos(theta))
   return c_delta_t *Mpc/c # sec.

# Compton accumulation
def ECompton_threshold(Compton_threshold = 0.005,z=0):
   return Compton_threshold/(4/3*Ecmb*(1+z)/me*1e-3) *me*1e-6 #GeV

# Compton scattering
def Dic(Ee=Ee_default,z=0): # Ee (GeV)
   return 3*(me*1e3)**2/(4*sigmaT*Ee*1e9*rhoCMB*(1+z)**4) /Mpc #Mpc

def lambdaIC(z=0):
   return 1/(sigmaT*Mpc*nCMB*(1+z)**3) #Mpc

def Eic(Ee=Ee_default,z=0):
   return 4*Ecmb*(1+z) *Ee**2/(3*me**2)*1e3 #GeV 

def Ee(Egamma=Egamma_default,z=0):
   return me*sqrt((3*Egamma*1e-3 )/(4*Ecmb*(1+z))) #GeV 

def tIC():
   return lambdaIC()/(c*yr/Mpc) #yr

# Larmor radius
def RL(Ee=Ee_default,B=B_default):
   return (Ee/erg_to_GeV)/(e*B) /Mpc #Mpc

# Magnetic deflection
def delta(Ee=Ee_default,B=B_default):
   return lambdaIC()/RL(Ee,B)*degre

def Delta(Ee=Ee_default,B=B_default,lambda_B=lambda_B_default):
   Delta1=Dic(Ee)/RL(Ee,B)*degre                #degre if Dic >> lambda_B 
   Delta2=sqrt(Dic(Ee)*lambda_B)/RL(Ee,B)*degre #degre if Dic << lambda_B
   return Delta1, Delta2

# Threshold energy when Dic = RL
def Ethreshold_ic(Ee=Ee_default,B=B_default):
   return Eic(Ee)*Dic(Ee)/RL(Ee,B) # GeV
# Pair production

def Ethreshold_gg():
   return (me)**2/Eebl *1e-3 #GeV

def lambda_gg(Egamma=1,z=0): # Egamma (GeV)
   # Bilinear interpolation 
   z_tab = [0,0.5,1,2,3]
   E_tab = loadtxt("lambda_e.dat",unpack=True,usecols=[0])*me*1e-6
   i2 = searchsorted(z_tab,z)
   if z<=0:
      fy=1
      i1=i2
   elif z>3:
      fy=1
      i2-=1
      i1=i2
   else:
      i1 = i2-1
      fy=(z-z_tab[i1])/(z_tab[i2]-z_tab[i1])
   j2 = searchsorted(E_tab,Egamma)
   j1 = j2-1
   fx=(Egamma-E_tab[j1])/(E_tab[j2]-E_tab[j1])
   lambda_e = loadtxt("lambda_e.dat",unpack=True,usecols=[i1+1,i2+1,i1+6,i2+6])
   lambda11 = lambda_e[0,j1]*((1-fx)*(1-fy))
   lambda12 = lambda_e[1,j1]*((1-fx)*fy)
   lambda21 = lambda_e[0,j2]*(fx*(1-fy))
   lambda22 = lambda_e[1,j2]*(fx*fy)
   lambda_proper = lambda11 + lambda12 + lambda21 + lambda22
   lambda11 = lambda_e[2,j1]*((1-fx)*(1-fy))
   lambda12 = lambda_e[3,j1]*((1-fx)*fy)
   lambda21 = lambda_e[2,j2]*(fx*(1-fy))
   lambda22 = lambda_e[3,j2]*(fx*fy)
   lambda_comobile = lambda11 + lambda12 + lambda21 + lambda22
   return lambda_proper, lambda_comobile, 800.e3/Egamma #Mpc (from Durrer and Neronov 2013)