# Import oscars modules and plotting tools
import oscars.sr
from oscars.plots_mpl import *

# Other python imports
import numpy as np


# Create an oscars.sr object for calculations (use the GPU if you can!)
osr = oscars.sr.sr(nthreads=10, gpu=1)

# Define a beam before we import any trajectories
# For imported trajectories ctstartstop is not needed, but is if you use internal propogation
osr.set_particle_beam(type='electron', current=0.5)


# Importation of a python list trajectory.  This is just a straight trajectory.
beta = 0.9999999854933347
my_trajectory = [[t, [0, 0, t*beta*osr.c()], [0, 0, beta], [0, 0, 0]] for t in np.linspace(0, 1e-7, 100)]
osr.set_trajectory(trajectory=my_trajectory)

# Plot the trajectory
plot_trajectory_position(osr.get_trajectory())


# It is possible to let OSCARS calculate the beta and beta' terms, just leave them out
# of the trajectory
beta = 0.9999999854933347
my_trajectory = [[t, [1e-6 * np.sin(2.*osr.pi()*t*beta*osr.c()/0.050), 0, t*beta*osr.c()]] for t in np.linspace(0, 1e-8, 10000)]
osr.set_trajectory(trajectory=my_trajectory)

# Plot the trajectory
plot_trajectory_position(osr.get_trajectory())
plot_trajectory_velocity(osr.get_trajectory())


# As an example we calculate the power density
pd = osr.calculate_power_density_rectangle(plane='XY', width=[0.01, 0.01], npoints=[51, 51], translation=[0, 0, 30])
plot_power_density(pd)
print('Total Power:', osr.calculate_total_power(), '[W]')


# Create an EPU trajectory.  Here we need to define a beam with ctstartstop
osr.set_particle_beam(beam='NSLSII', x0=[0, 0, -1], ctstartstop=[0, 2])
osr.clear_bfields()
osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.042], nperiods=31, phase=-0.25*osr.pi())
osr.add_bfield_undulator(bfield=[1, 0, 0], period=[0, 0, 0.042], nperiods=31, phase=+0.25*osr.pi())

# Calculate trajectory and export it to a file
t = osr.calculate_trajectory(ofile='Example_007_Trajectory1.txt')
plot_trajectory_position(t)


# Read trajectory in default format
osr.set_trajectory(ifile='Example_007_Trajectory1.txt')

# Get and plot the trajectory
t = osr.get_trajectory()
plot_trajectory_position(t)
plot_trajectory_velocity(t)


# As an example we can now calculate the power density
pd = osr.calculate_power_density_rectangle(plane='XY', width=[0.1, 0.1], npoints=[51, 51], translation=[0, 0, 30])
plot_power_density(pd)
print('Total Power:', osr.calculate_total_power(), '[W]')


# Read trajectory in default format
osr.set_trajectory(ifile='Example_007_Trajectory1.txt', iformat='T X Y Z * * * * * *')

# Get and plot the trajectory
t = osr.get_trajectory()
plot_trajectory_position(t)
plot_trajectory_velocity(t)


# As an example we can now calculate the power density
pd = osr.calculate_power_density_rectangle(plane='XY', width=[0.1, 0.1], npoints=[51, 51], translation=[0, 0, 30])
plot_power_density(pd)
print('Total Power:', osr.calculate_total_power(), '[W]')


# Read trajectory in default format
osr.set_trajectory(ifile='Example_007_Trajectory1.txt', iformat='T X * Z')

# Get and plot the trajectory
t = osr.get_trajectory()
plot_trajectory_position(t)
plot_trajectory_velocity(t)


# Create an EPU trajectory and save it to a file
osr.set_particle_beam(beam='NSLSII', x0=[0, 0, -1], ctstartstop=[0, 2])
osr.clear_bfields()
osr.add_bfield_undulator(bfield=[0, 1, 0], period=[0, 0, 0.042], nperiods=31, phase=-0.25*osr.pi())
osr.add_bfield_undulator(bfield=[1, 0, 0], period=[0, 0, 0.042], nperiods=31, phase=+0.25*osr.pi())

t = osr.calculate_trajectory(bofile='Example_007_Trajectory1.dat')


# Read trajectory in default format
osr.set_trajectory(bifile='Example_007_Trajectory1.dat')

# Get and plot the trajectory
t = osr.get_trajectory()
plot_trajectory_position(t)
plot_trajectory_velocity(t)


# As an example we can now calculate the power density
pd = osr.calculate_power_density_rectangle(plane='XY', width=[0.1, 0.1], npoints=[51, 51], translation=[0, 0, 30])
plot_power_density(pd)
print('Total Power:', osr.calculate_total_power(), '[W]')