# ----------------------------------------------------------------------------------------
# 					SIMULATION PARAMETERS FOR THE PIC-CODE SMILEI
# ----------------------------------------------------------------------------------------

import numpy as np
import cmath
import math
from scipy.special import genlaguerre,factorial

lambda0 = 0.8e-6               # m
c = 299792458                  # m/s
omega0 = 2*math.pi*c/lambda0   # rad/s
c_ov_omega0 = c/omega0         # m/rad
eps0=8.854187e-12              # Vacuum permittivity, F/m
e=1.60217646e-19               # Elementary charge, C
me=9.10938215e-31              # Electron mass, kg
ncrit = eps0*omega0**2*me/e**2 # Plasma critical number density [m-3]
normalized_electron_charge = -1 # For electrons
Lr=lambda0   #reference length
Tr=1./omega0                    #reference time
Er=me*c*omega0/e                #reference electric field
Br=me*omega0/e       #reference magnetic field
Nr=eps0*me*omega0**2/e**2  #reference particle density 
N0=6.02*10**23/(9*10**(-6))      ##SI  density 1/8(density of water)
n0=N0/Nr

dx = lambda0 / (2*math.pi*40) / c_ov_omega0
dtrans = lambda0/(2*math.pi*10)/ c_ov_omega0
dt = 0.95/math.sqrt(1/dx**2+2/dtrans**2)
ncells_x =256#1792
Lx = ncells_x * dx
ntrans =256
Ltrans = ntrans * dtrans
npatch_x =16
npatch_trans = 16
Nit =500
Main(
    geometry = "3Dcartesian",
    interpolation_order = 2,
    timestep = dt,
    simulation_time = dt*Nit,
    cell_length  = [dx, dtrans, dtrans],
    grid_length = [ Lx,  Ltrans, Ltrans],
    number_of_patches = [npatch_x, npatch_trans, npatch_trans],
    cluster_width =16,
    EM_boundary_conditions = [ ["PML"]],
    number_of_pml_cells=[[10,10],[10,10],[10,10]],	
    solve_poisson = False,
    solve_relativistic_poisson = True,
    relativistic_poisson_max_iteration = 50000,		
    reference_angular_frequency_SI=omega0,	
    print_every = 10,	
)




# We build a Laguerre-Gauss laser from scratch instead of using LaserGaussian3D
# The goal is to test the space_time_profile attribute
##definition of factorial function 
def factorial(n):
     if n==0 or n==1:
          return 1
     else:
          return (n*factorial(n-1))

p=0  #  radial index
m=1   # azimuthal index 
C_pm=math.sqrt(2*factorial(p)/math.pi/factorial(p+abs(m)))  #normalized constant
#beam waist 
w0=3e-6/Lr     #beam waist radius at focus x=0, in units of Lr   w0=3e-6m
k0=2*math.pi   #wavenumber 1/k0=Lr/2/math.pi[normalized Lr="1"->k0=2*math.pi]
Xr=w0**2*k0/2.  # Reyleigh length
def w(x):
    return w0*np.sqrt(1.+((x-focus[0])/Xr)**2)



omega=1
x0=0        
focus               = [x0, Main.grid_length[1]/2., Main.grid_length[2]/2.]
#Gouy phase
def Phi_G(x):
    return -math.atan((x-focus[0])/Xr)
a0=60./C_pm # normalized dimensionless laser vector potential amplitude  

def LG_beam_with_temporal_profile(x,y,z,t):
    r=np.sqrt( (y-focus[1])**2 + (z-focus[2])**2 )
    spatial_amplitude=a0*(w0/w(x))*(math.sqrt(2)*r/w(x))**abs(m)*genlaguerre(p,m)(2*r**2/w(x)**2)*math.exp(-r**2/w(x)**2)
    exponential_with_phase=cmath.exp(1j*((k0*x-omega*t+m*np.arctan2((z-Ltrans/2.), (y-Ltrans/2.)) +(2*p+abs(m)+1)*Phi_G(x)-x*r**2/Xr/w(x)**2)))
    space_time_profile=spatial_amplitude
    return space_time_profile*exponential_with_phase

LaserEnvelope(
    omega          = 1.,
    envelope_profile=LG_beam_with_temporal_profile,
    envelope_solver = 'explicit',
    Envelope_boundary_conditions = [ ["PML"] ],
    polarization_phi = 0.,
    ellipticity      = 1.
)




# Plasma density profile
def my_profile(x,y,z):
	center = [0.,Ltrans/2.,Ltrans/2.]
	Radius = 8. 
	Length = Lx
	if ((abs(x-center[0])<Length) and ((y-center[1])**2+(z-center[2])**2<Radius*Radius)):
		return n0
	else:
		return 0



# Add some test electrons
Species( 
    name = "electron",
    position_initialization = "regular",
    momentum_initialization = "cold",
    particles_per_cell = 1,
    mass = 1,
    charge = -1,
    number_density = my_profile,  
    mean_velocity = [0., 0., 0.],
    time_frozen = 0.0,
    pusher = "ponderomotive_boris", # pusher to interact with envelope
    boundary_conditions = [
       ["remove", "remove"],
       ["remove", "remove"],
       ["remove", "remove"],
    ],

)

# Diagnostics 

list_fields = ['Ex','Ey','Ez','Bx','By','Bz','Rho_electron','Env_A_abs','Env_E_abs','Env_Ex_abs']

DiagFields(
    every = 5,
    fields = list_fields
)


DiagTrackParticles(
    species = "electron",
    every =5,
    attributes = ["x","y", "z", "px", "py","pz","weight"]
)