import os
from numpy import loadtxt, select, log, sin, cos, tan, arctan, arccos, exp, sqrt, pi, size
from numpy import append, zeros, linspace, shape
from src.constants import degre, yr

resultsFile="/results.dat"
profileFile="/profile.dat"

def PSF_Taylor2011(E): # E in GeV -> PSF in degre
   return 1.7 * E**(-0.74) * (1+(E/15)**2)**0.37

# READ SIMULATIONS FILES ===================================================================
def select_events(simulation,Erange=[1e-3,1e6],delayrange=[-1e8,1e30],powerlaw_index=1,tjet=180.,tobs=0.,generation='all'):    
   '''
      Allowed to read results file of a simulation and to apply a preselection on events. 
  
      Input:
         1. simulation directory location
         OPTIONAL:
         2. energy range [min,max] in GeV
         3. delay range in second
         5. powerlaw index Gamma such as dN/dE = E^-Gamma
         6. jet opening angle
         7. observation angle 
         8. generation desired else 'all' 

      Output:
            0. weight
            1. energy [GeV]
            2. time delay [s]
            3. arrival angle [degre]
            4,5. theta, phi (arrival polar and azimuthal coordinates) [degre]
            (6. generation if generation = "all")
   '''

   # 0 - generation
   # 1 - weight
   # 2 - energy [GeV]
   # 3 - time delay [s]
   # 4,5,6 - theta_pos, theta_dir, phi_dir [rad]
   # 7 - source photon energy [GeV]
   results = loadtxt(simulation+resultsFile,unpack=True)
   Dsource,redshift,n_phot  = ReadProfile(simulation,[2,3,4]) 
   try:
      gen_select = int(generation)
      cond = (results[0] == gen_select)
   except:
      # select photons only
      cond = (results[0]%2==0) 
   # energy selection and time selection
   cond = cond & (results[2] >= Erange[0])     & (results[2] <= Erange[1])  
   cond = cond & (results[3] >= delayrange[0]) & (results[3] <= delayrange[1])
   results = results[:,cond]

   # Apply source spectrum
   results[1] /= n_phot 
   if "simple case" not in simulation: 
      # compute intrinsic luminosity 
      powerlaw_index = float(powerlaw_index) 
      results[1]*= results[7]**(1-powerlaw_index) * log(Erange[1]/Erange[0])
      if powerlaw_index == 2:
         results[1] /= log(Erange[1]/Erange[0])
      else:
         results[1] *= (2-powerlaw_index)/(Erange[1]**(2-powerlaw_index)-Erange[0]**(2-powerlaw_index))

   # misaligned jet & jet opening angle selection 
   results = append(results,zeros((2,size(results,axis=1))),axis=0)
   phi_dj = 0
   results[8] = phi_dj
   results[9] = cos(tobs*degre)*cos(results[4])+sin(tobs*degre)*sin(results[4])*cos(phi_dj-results[6])
   tabi = results
   Nphi=100
   for phi_dj in linspace(-pi,pi,Nphi-1):
   #for phi_dj in 2*pi*random.rand(Nphi):
      results[8] = phi_dj
      results[9] = cos(tobs*degre)*cos(results[4])+sin(tobs*degre)*sin(results[4])*cos(phi_dj-results[6])
      tabi = append(tabi,results,axis=1)
   results = tabi     
   cond = (arccos(results[9])/degre < tjet)
   results = results[:,cond]         
   #if tjet < 90:
   #    sigma = sin(tjet*degre)
   #    results[1] *= exp((sin(arccos(results[9]))/sigma)**2) 
   results[1] /= Nphi #*4*pi * Dsource**2 *(1-cos(tjet*degre))
   
   if generation == "all":
      return results[1],results[2],results[3],results[5],results[8],results[7],results[0]
   else:
      return results[1],results[2],results[3],results[5],results[8],results[7]

# READ RESULTS FILES =======================================================================
resultdir="Results/"

def ReadProfile(fileName,cols=[0,1,2,3,4,5]):
   '''
      Allowed to read simulation profile
         0 - magnetic field amplitude [Gauss]
         1 - maximum source energy [TeV]
         2 - source distance [Mpc]
         3 - source redshift
         4 - number of photons emitted
         5 - number of leptons produced during the cascade
   '''
   if not os.path.exists(resultdir+fileName):
      return loadtxt(fileName+profileFile,unpack=True,usecols=cols)
   else:
      return loadtxt(resultdir+fileName+profileFile,unpack=True,usecols=cols)

def ReadSpectrum(fileName,cols=[0,1,2,3,4,5,6,7,8]): 
   return loadtxt(resultdir+fileName+"/Spectrum.txt",unpack=True,usecols=cols)

def ReadSourceSpectrum(fileName): 
   return loadtxt(resultdir+fileName+"/Source_spectrum.txt",unpack=True,usecols=[0,1])

def ReadRadial(fileName,cols=[0,1,2,3,4]): 
   return loadtxt(resultdir+fileName+"/arrival_Angle_distribution.txt",unpack=True,usecols=cols)

def ReadTiming(fileName,cols=[0,1,2,3,4]): 
   return loadtxt(resultdir+fileName+"/Timing.txt",unpack=True,usecols=cols)

def ReadGeneration(fileName,cols=[0,1]):
   return loadtxt(resultdir+fileName+"/Generation.txt",unpack=True,usecols=cols)

def ReadAngleVsEnergy(fileName,cols=[0,1,2,3,4]): 
   return loadtxt(resultdir+fileName+"/Angle_versus_Energy.txt",unpack=True,usecols=cols)

def ReadDelayVsEnergy(fileName,cols=[0,1,2,3,4]): 
   return loadtxt(resultdir+fileName+"/Delay_versus_Energy.txt",unpack=True,usecols=cols)

def ReadDelayVsAngle(fileName,cols=[0,1,2,3,4]): 
   return loadtxt(resultdir+fileName+"/Delay_versus_angle.txt",unpack=True,usecols=cols)

def ReadElmag(fileName):
   E,E2dNdE = loadtxt(resultdir+fileName+"/Elmag.dat",unpack=True,usecols=[0,1])
   E_PSF,E2dNdE_PSF = loadtxt(resultdir+fileName+"/Elmag_PSF.dat",unpack=True,usecols=[0,1])
   return  E*1e-9, E2dNdE*1e-9, E_PSF*1e-9, E2dNdE_PSF*1e-9

def ReadTaylor(fileName):
   E,E2dNdE = loadtxt(resultdir+fileName+"/Taylor2011-fig1.csv",delimiter=",",unpack=True,usecols=[0,1])
   E *= 1e-9
   E2dNdE *= 1e-9 
   return E, E2dNdE