# Import the OSCARS SR module
import oscars.sr

# Import basic plot utilities (matplotlib).  You don't need these to run OSCARS, but it's used here for basic plots
from oscars.plots_mpl import *


# Create a new OSCARS object.  Default to 8 threads and always use the GPU if available
osr = oscars.sr.sr(nthreads=8, gpu=1)


# Clear any existing fields (just good habit in notebook style) and add an undulator field
osr.clear_bfields()
osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.049], nperiods=31)

# Just to check the field that we added seems visually correct
plot_bfield(osr)


# Setup beam similar to NSLSII
osr.set_particle_beam(
    energy_GeV=3,
    x0=[0, 0, -1],
    current=0.500,
    sigma_energy_GeV=0.001*3,
    beta=[1.5, 0.8],
    emittance=[0.9e-9, 0.008e-9]
)

# Set the start and stop times for the calculation
osr.set_ctstartstop(0, 2)


# Run the particle trajectory calculation on the ideal case
osr.set_new_particle(particle='ideal')
trajectory = osr.calculate_trajectory()

# Plot the trajectory position and velocity
plot_trajectory_position(trajectory)


# Single particle spectrum for the ideal case
osr.set_new_particle(particle='ideal')
spectrum_se = osr.calculate_spectrum(obs=[0, 0, 30], energy_range_eV=[130, 160], npoints=500)
spectrum_me = osr.calculate_spectrum(obs=[0, 0, 30], energy_range_eV=[130, 160], npoints=500, nparticles=100)
plot_spectra([spectrum_se, spectrum_me], ['single-electron', 'multi-electron'])


# Calcuate flux at 30 [m] for 152 eV
flux_multi = osr.calculate_flux_rectangle(
    plane='XY',
    energy_eV=152,
    width=[0.01, 0.01],
    npoints=[101, 101],
    translation=[0, 0, 30],
    nparticles=5
)

plot_flux(flux_multi, ofile='UndulatorFlux_MultiParticle.png')