import numpy as np
twopi = 2*np.pi
twopi2 = twopi**2.

mi = 100.0  # ion mass, unit me

# let the reference frequency wr = wci (ion cyclotron frequency)
# length normalized to c/wci = wpiwci * di
# time normalized to   1/wci
# B field normalized to me*wci/e
# density normalized to epsilon_0 * me * wci^2/e^2
# temperature normalized to me*c^2

# wpi/wci 
wpiwci = 100

Ti_Te = 1./4  # ratio Ti/Te
beta = 5e-4  # total beta
beta_e = beta / (1+Ti_Te)
beta_i = beta / (1+1/Ti_Te)
Te1 = 0.25*beta_e*mi/wpiwci**2.  # Te (one component)
Ti1 = 0.25*beta_i*mi/wpiwci**2.  # Ti (one component)
Te = [Te1, Te1, Te1]
Ti = [Ti1, Ti1, Ti1]

n0 = mi*wpiwci**2.
ppc = 100 # particles per cell, same for e, i


# ---------- basic setup -------------------
box = [30./wpiwci,0.5/wpiwci]  # c/wci
res = [3200,1600]   # resolution: cell number per unit
number_of_patches = [128,1]

tmax = 400.  # max run time, wci^-1

#********** DO NOT CHANGE **********************
dr =[ 1./i for i in res ]
dt_CFL = 1./np.sqrt(sum(i**2 for i in res))  # Courant condition
rest = (1./dt_CFL)/0.98
dt =1./rest
# timesteps to finish one time unit
Nt = int(rest)
# sim. dimension
dim = len(box)
# adjust box size (cells) to be divided by number_of_pathces
for i in range(dim):
	fac = float(box[i])*res[i]/number_of_patches[i]
	if fac!=fac//1: fac = int(fac)+1
	box[i] = float(fac)*number_of_patches[i]/res[i]

# total number of cells
cells = list( np.array(box)*np.array(res) )
#************************************************

dump_freq  = Nt  # dump every wci^-1  
flush_freq = 20*dump_freq

# ------------ global settings --------------------
Main(
  geometry = "2Dcartesian",
  interpolation_order = 2,
  grid_length  = box,
  cell_length = dr,
  # each must be power of 2; total must be greater than MPI process number
  number_of_patches = number_of_patches,
  #patch_arrangement = "linearized_ZYX",
  timestep = dt,
  simulation_time = tmax,
  EM_boundary_conditions = [
      ['PML'],
       ['periodic'],
#       ['silver-muller'],
#       ['silver-muller'],
  ],
  number_of_pml_cells=[[64]],
  solve_poisson = False,  # "False" to have a neutralizing background
  # reference_angular_frequency_SI = 3.e8/0.4e-6, #speed of light / laser wavelegth
  print_every = 1e4,
  # random_seed = smilei_mpi_rank
)


## ------------ load balancing --------------
#LoadBalancing(
#    initial_balance = True,
#    every = 5*Nt,
#    cell_load = 1.,
#    frozen_particle_load = 1.,
#)

#------------  vectorization ---------------
# Vectorization(
# 	mode = "adaptive",
# 	reconfigure_every = 20,
# 	initial_mode = "on"
# )

#---------- moving window --------------
# MovingWindow(
#     time_start = 0.,
#     velocity_x = 1.,
# )

# -------------- filtering ---------------
# CurrentFilter(
#     model = "binomial",
#     passes = [4, 2],
# )

# FieldFilter(
#     model = "Friedman",
#     theta = 0.,
# )


# ----------------- laser pulse  ----------------------
# def time_envelope(t):
#     if 0<t<tau: return( (np.sin(np.pi*t/tau))**2 )
#     else: return (0.)

# a0=1.0; tau=100.0; focus=[20.0, box[1]/2.0, box[2]/2.0]; waist=3.0

# LaserGaussian3D(box_side="xmin", a0=a0, omega=twopi, polarization_phi=0.0, 
#                 ellipticity=0, focus=focus, waist=waist, incidence_angle=[0.0, 0.0], 
#                 time_envelope=time_envelope,)

w0wci = 0.1*twopi   # w0/wci = 0.6283
t_upramp = 50.  # wci^-1
theta = 0

# B0 
B0 = mi
dB_B0 = 0.001*B0

def By_profile(y,t):
    if t<t_upramp: return dB_B0*np.sin(0.5*np.pi*t/t_upramp)**2*np.sin(w0wci*t+theta)
    else: return dB_B0*np.sin(w0wci*t+theta)
    
def Bz_profile(y,t):
    if t<t_upramp: return dB_B0*np.sin(0.5*np.pi*t/t_upramp)**2*np.cos(w0wci*t+theta)
    else: return dB_B0*np.cos(w0wci*t+theta)
    
Laser(
    box_side = "xmin",
    space_time_profile = [By_profile, Bz_profile]
)

def B0_profile(x,y,t):
    return B0

PrescribedField(
    field = "Bx_m",
    profile = B0_profile,
)

# ExternalField(
#     field = "Bx",
#     profile = B0,
# )

#--------------------- plasma -------------
# nmax = 0.02*twopi2
# # def number_density(x,y,z):
# #   if 10.0<x<190.0: return ( nmax*np.cos(np.pi*(x-100.0)/180) )
# #   else: return (0.)
# def number_density(x,y,z):
#     if x<=10.: return (0.0)
#     elif x<=20.: return ( nmax*np.sin(np.pi*(x-10.)/20.) )
#     elif x<=120.: return (nmax)
#     elif x<=130.: return ( nmax*np.cos(np.pi*(x-120.)/20.) )
#     else: return (0.0)
    

# ion
Species(name = 'ion',position_initialization = 'random',momentum_initialization = 'maxwell-juettner', temperature = Ti, particles_per_cell = ppc, mass = mi, charge = 1.0, number_density = n0, boundary_conditions = [['thermalize']], thermal_boundary_temperature = Ti, time_frozen=0.0,)

# electron
Species(name = 'elec',position_initialization = 'random', momentum_initialization = 'maxwell-juettner', temperature = Te, particles_per_cell = ppc, mass = 1.0, charge = -1.0, number_density = n0, boundary_conditions = [['thermalize']], thermal_boundary_temperature = Te, time_frozen=0.0,)


#------------ diagnostics ----------------------

# def angle_y(particles):
#     return (np.arctan2(particles.py, particles.px) * 180/np.pi)

# def angle_z(particles):
#     return (np.arctan2(particles.pz, particles.px) * 180/np.pi)
  
#---------- Scalar ---------------
DiagScalar(
    every = int(Nt),
    precision = 8,
)

#---------- Probe ---------------
# DiagProbe(
#     name = "probe_at_z5",
#     every = [0, 2],
#     flush_every = flush_freq,
#     origin   = [5.0],
# #     corners  = [
# #       [plasma_start-5.0, box[1], 0.], # along y
# #       [plasma_start-5.0, 0., box[2]], # along z
# #     ],
# #     number   = [cells[1], cells[2]],
#     fields   = ["Ex", "Ey", "Ez", "By", "Bz", "Rho_elec", "Rho_ion"],
# #     time_integral = True,
# )

# DiagProbe(
#     name = "reflection",
#     every = [0, int(Nt/20)],
#     flush_every = flush_freq,
#     origin   = [plasma_start-5.0, 0., 0.],
#     corners  = [
#       [plasma_start-5.0, box[1], 0.], # along y
#       [plasma_start-5.0, 0., box[2]], # along z
#     ],
#     number   = [cells[1], cells[2]],
#     fields   = ["Ey", "Ez", "By", "Bz"],
# #     time_integral = True,
# )

# DiagProbe(
#     name = "transmission",
#     every = [0, int(Nt/20)],
#     flush_every = flush_freq,
#     origin   = [plasma_start+5.0, 0., 0.],
#     corners  = [
#       [plasma_start+5.0, box[1], 0.],
#       [plasma_start+5.0, 0., box[2]],
#     ],
#     number   = [cells[1], cells[2]],
#     fields   = ["Ey", "Ez", "By", "Bz"],
# #     time_integral = True,
# )

#---------Field --------------

from numpy import s_
# central x-y plane
DiagFields(
  every = [0, dump_freq],
  flush_every = flush_freq,
#   subgrid = s_[:, :, int(cells[2]/2.)],
  fields = ["Ex", 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'Rho_elec', 'Rho_ion'],
)

# DiagFields(
#   every = [0, dump_freq],
#   flush_every = flush_freq,
#   subgrid = s_[:, :, int(cells[2]/2.)],
#   fields = ['Rho_ion', "Rho_elec"],
# )

# central x-z plane
# DiagFields(
#   every = [0, dump_freq],
#   flush_every = flush_freq,
#   subgrid = s_[:, int(cells[1]/2.), :],
#   fields = ["Ex", 'Ey', 'Ez', 'Bx', 'By', 'Bz', 'Rho_elec'],
# )

# DiagFields(
#   every = [0, dump_freq],
#   flush_every = flush_freq,
#   subgrid = s_[:, int(cells[1]/2.), :],
#   fields = ['Rho_ion', "Rho_elec"],
# )

#------------------------ ParticleBinning diagnostics  ---------------

# particle radial position relative to box center axis
# def particles_r(particles):
#   return ( np.sqrt((particles.y-box[1]/2.)**2.0 + (particles.z-box[2]/2.)**2.0) )

#0-3, ion momentum-density distributions
DiagParticleBinning(deposited_quantity="weight_px", every=[0, dump_freq], flush_every = flush_freq, species=['ion'], axes=[ ['x', 0, box[0], cells[0]],], )

DiagParticleBinning(deposited_quantity="weight_py", every=[0, dump_freq], flush_every = flush_freq, species=['ion'], axes=[ ['x', 0, box[0], cells[0]],], )

DiagParticleBinning(deposited_quantity="weight_pz", every=[0, dump_freq], flush_every = flush_freq, species=['ion'], axes=[ ['x', 0, box[0], cells[0]],], )

DiagParticleBinning(deposited_quantity="weight", every=[0, dump_freq], flush_every = flush_freq, species=['ion'], axes=[ ['x', 0, box[0], cells[0]],], )

    
#4-7, electron momentum-density distributions
DiagParticleBinning(deposited_quantity="weight_px", every=[0, dump_freq], flush_every = flush_freq, species=['elec'], axes=[ ['x', 0, box[0], cells[0]],], )

DiagParticleBinning(deposited_quantity="weight_py", every=[0, dump_freq], flush_every = flush_freq, species=['elec'], axes=[ ['x', 0, box[0], cells[0]],], )

DiagParticleBinning(deposited_quantity="weight_pz", every=[0, dump_freq], flush_every = flush_freq, species=['elec'], axes=[ ['x', 0, box[0], cells[0]],], )

DiagParticleBinning(deposited_quantity="weight", every=[0, dump_freq], flush_every = flush_freq, species=['elec'], axes=[ ['x', 0, box[0], cells[0]],], )

#4 local density / energy density averaged over a radius of half lambda around central laser axis
# DiagParticleBinning(deposited_quantity="weight", every=[0, dump_freq], flush_every = flush_freq, species=['elec_central', 'elec_outer'], axes=[ ["x", plasma_start-10.0, plasma_start+10.0, int(20*res[0])], [particles_r, 0., 0.5, 1] ],)

# #5 
# DiagParticleBinning(deposited_quantity="weight_ekin", every=[0, dump_freq], flush_every = flush_freq, species=['elec_central', 'elec_outer'], axes=[ ["x", plasma_start-10.0, plasma_start+10.0, int(20*res[0])], [particles_r, 0., 0.5, 1] ],)


#     particle energy spectrum 

# elec 
# DiagParticleBinning(deposited_quantity="weight", every=[0, dump_freq], flush_every = flush_freq, species=['elec'], axes=[ ['ekin', 0, 200., 800] ],)

# ion
# DiagParticleBinning(deposited_quantity="weight", every=[0, dump_freq], flush_every = flush_freq, species=['ion'], axes=[ ['ekin', 0, 200., 800] ],)


#    divergence of local (x=[13,15]) plasma 

# #4: elec
# DiagParticleBinning(deposited_quantity="weight", every=[0, dump_freq], flush_every = flush_freq, species=['elec_central', 'elec_outer'], axes=[ [angle_y, -50, 50, 1000], [angle_z, -50, 50, 800], ["x", 13., 15., 1], ["y", box[1]/2.-8., box[1]/2.+8., 1], ["z",box[2]/2.-8., box[2]/2.+8., 1] ], )

# #5: proton
# DiagParticleBinning(deposited_quantity="weight", every=[0, dump_freq], flush_every = flush_freq, species=['proton'], axes=[ [angle_y, -50, 50, 1000], [angle_z, -50, 50, 800], ["x", 13., 15., 1], ["y", box[1]/2.-8., box[1]/2.+8., 1], ["z",box[2]/2.-8., box[2]/2.+8., 1] ], )


# ------------- Screen diagnostics  -------------

# def angle_y(particles):
#     return (np.arctan2(particles.py, particles.px) * 180/np.pi)

# def angle_z(particles):
#     return (np.arctan2(particles.pz, particles.px) * 180/np.pi)

# #0: electron spatial dist. 
# DiagScreen(every = dump_freq, flush_every = flush_freq, shape = "plane", point=[120.0, box[1]/2., box[2]/2.], vector=[1.,0,0], direction = "canceling", deposited_quantity = "weight", species = ['elec'], axes=[ ['a', -box[1]/2, box[1]/2, int(box[1]*10)], ['b', -box[2]/2, box[2]/2, int(box[2]*10)] ] )

# #1: electron angular dist. 
# DiagScreen(every = dump_freq, flush_every = flush_freq, shape = "plane", point=[120.0, box[1]/2., box[2]/2.], vector=[1,0,0], direction = "canceling", deposited_quantity = "weight", species = ['elec'], axes=[ [angle_y, -50, 50, 1000], [angle_z, -50, 50, 800] ] )

# #2: electron spectrum 
# DiagScreen(every = dump_freq, flush_every = flush_freq, shape = "plane", point=[120.0, box[1]/2., box[2]/2.], vector=[1,0,0], direction = "canceling", deposited_quantity = "weight", species = ['elec'], axes=[ ['a', -box[1]/2, box[1]/2, 1], ['b', -box[2]/2, box[2]/2, 1], ['ekin', 0, 200., 800] ] )


# ----------- particle tracking ------------------------

# def filter_func(particles):
#     ekin = (1+particles.px**2.0+particles.py**2.0+particles.pz**2.0)**0.5
#     return (ekin>4.)

# DiagTrackParticles(species ="ion", every = [0, dump_freq], flush_every=flush_freq, attributes = ["x", "px", "py", "pz",] )
# DiagTrackParticles(species ="elec", every = [0, dump_freq], flush_every=flush_freq, attributes = ["x", "px", "py", "pz",] )

# ----------- checkpoints for restarts ------------------------
# Checkpoints(
#     # restart_dir = "dump1",
#     # restart_number = 1, 
#     dump_minutes = 15.5*60,
#     dump_step = 0,
#     exit_after_dump = True,
#     keep_n_dumps = 1,
# )