# the output of plotting commands is displayed inline within frontends, 
# directly below the code cell that produced it
%matplotlib inline
from __future__ import print_function
# this python library provides generic shallow (copy) and deep copy (deepcopy) operations 
from copy import deepcopy

# import from Ocelot graphical modules
import sys, os
from ocelot import *
from ocelot.utils.xfel_utils import *
from ocelot.gui.accelerator import *
from ocelot.gui.genesis_plot import *
#from ocelot.optics.elements import Filter_freq

import numpy as np
from copy import copy
#import matplotlib.pyplot as plt
# load beam distribution

"""
Explanation: *This notebook was created by Svitozar Serkez. Source and license info is on GitHub. August 2016. *
Tutorial N5:  Genesis preprocessor. 
Ocelot is not only a particle tracking code, but a simulation toolkit.
Python is known to be a very good "glue" between different software
Interface to other codes may be (and was) developed.
GENESIS interface is under active development and is currently used for studies
-list of papers

This example will cover the following topics:

Initialization  of the library
Preparing the GENESIS simulation
running many-stage statistical simulation

Requirements

OCELOT - library
numpy, scipy, matplotlib
beam.txt - input beam file

Import of modules
End of explanation
"""


E_beam=8.5 #[GeV]
E_photon=250 #[eV]

"""
Explanation: Setting input parameters
electron beam energy and expected radiation photon energy
End of explanation
"""


# defining the undulator
lperiod=0.068
nperiods=73
und = Undulator(lperiod=lperiod, nperiods=nperiods, Kx=1.0);
und.Kx = Ephoton2K(E_photon, und.lperiod, E_beam)

# defining of the drifts
d2 = Drift (l=4*und.lperiod)
d3 = Drift (l=7*und.lperiod)

# defining of the quads
qf = Quadrupole (l=6*und.lperiod, k1=-7.3)
qd = Quadrupole (l=6*und.lperiod, k1=7.3)
qdh=deepcopy(qd)
qdh.l/=2

# creating of the cell
extra_fodo = (und, d2, qdh)
cell_ps = (und, d2, qf, d3, und, d2, qd, d3)
l_fodo= MagneticLattice(cell_ps).totalLen/2
sase3 = MagneticLattice((und, d2, qd, d3) + 11*cell_ps)

up = UndulatorParameters(und,E_beam)
up.printParameters()

"""
Explanation: Creating SASE3 lattice
with native ocelot objects
End of explanation
"""


beamf = read_beam_file('beam.dat')

"""
Explanation: Load beam file
End of explanation
"""


fig=plt.figure()
fig.set_size_inches((20,15))
plot_beam(fig, beamf)

"""
Explanation: Plot beamfile
End of explanation
"""


beta_av = 20.0
beam=get_beam_peak(beamf)
beam.E=E_beam
rematch(beta_av, l_fodo, qdh, sase3, extra_fodo, beam, qf, qd)
beamf = transform_beam_file(beamf ,transform = [ [beam.beta_x,beam.alpha_x], [beam.beta_y,beam.alpha_y] ], energy_new = beam.E, emit_scale = 1.0)
beamf = cut_beam(beamf,[-2e-6, 2e-6])

fig=plt.figure()
fig.set_size_inches((20,15))
plot_beam(fig, beamf)

"""
Explanation: Match beam file
End of explanation
"""


def f1(n, n0, a0, a1, a2):
    '''
    piecewise-quadratic tapering function
    '''
    for i in xrange(1,len(n0)):
        if n < n0[i]:
            return a0 + (n-n0[i-1])*a1[i-1] + (n-n0[i-1])**2 * a2[i-1]
        a0 += (n0[i]-n0[i-1])*a1[i-1] + (n0[i]-n0[i-1])**2 * a2[i-1]
    
    return 1.0

tap_start=3 #number of undulators
lin_tap=0.01 #taper step
quad_tap=0.0

n = 60
n0 = [0,tap_start,60]
a0 = und.Kx
a1 = [0,lin_tap*a0]
a2 = [0,quad_tap]

taper_func1 = lambda n : f1(n, n0, a0, a1, a2)

sase3= taper(sase3, taper_func1)

"""
Explanation: Tapering the undulator
linear, quadratic,power-law...
End of explanation
"""


run_dir = 'gen_stst' #directory to dump data
run_id=0 # run number (subdirectory 'run_#') for statistical studies
try:
    os.makedirs(run_dir)
except:
    pass
launcher = get_genesis_launcher('genesis') # launcher object to start genesis

"""
Explanation: specify the run_dir - directory into which the experimental results will be saved
End of explanation
"""


inp = generate_input(up, beam, itdp=False)
inp.lattice_str = generate_lattice(sase3, unit = up.lw, energy = beam.E) #generate Genesis lattice based on Ocelot lattice object
inp.beam_file_str = beam_file_str(beamf)
#inp.beamfile = 'tmp.beam'

inp.runid = run_id
inp.run_dir = run_dir
inp.ipseed = 17111+7*run_id # defines shot-noise, changes automatically 
# below other Genesis parameters may be specified, like prad0, dgrid, etc.....

"""
Explanation: generate Genesis input object
End of explanation
"""


print(inp.input())

print(inp.lattice_str)

print(inp.beam_file_str)

"""
Explanation: now all the genesis input files are created, such as lattice file, beam file, input file.
End of explanation
"""


#g = run(inp,launcher)

"""
Explanation: Genesis may be executed with the following command:
End of explanation
"""


exp_dir='/some_directory'

run_number=10
run_ids = xrange(0,run_number)

start_stage = 1
stop_stage = 4

# set simulation parameters
# prepare electron beam file

stage=1
if start_stage <= stage and stop_stage >= stage:
    for run_id in run_ids:   
        run_dir = exp_dir + 'run_' + str(run_id)    
        #prepare input, specify parameters
        #inp.ipseed = 17111*(run_id + 1)
        #
        #
        #
        #g = run(inp,launcher)
        print('run #',run_id, ' of stage ',stage)
        
stage=2
if start_stage <= stage and stop_stage >= stage:
    for run_id in run_ids:   
        run_dir = exp_dir + 'run_' + str(run_id)    
        #prepare input based on stage 1 output
        #inp.ipseed = 27222*(run_id + 1)
        #
        #inp.distfile = 'run.'+ str(inp.runid)+'.s1.gout.dist'
        #
        #g = run(inp,launcher)
        print('run #',run_id, ' of stage ',stage)
        
stage=3
if start_stage <= stage and stop_stage >= stage:
    for run_id in run_ids:   
        run_dir = exp_dir + 'run_' + str(run_id)    
        #prepare input based on stage 1 output
        #inp.ipseed = 37333*(run_id + 1)
        #
        #inp.distfile = 'run.'+ str(inp.runid)+'.s1.gout.dist'
        #
        #g = run(inp,launcher)
        print('run #',run_id, ' of stage ',stage)
    
    #stage=n .........

"""
Explanation: if "run" function is placed in a sctipt, the following post-processing code will be executed after the GENESIS simulation is finished
the following two python scripts would start ocelot/genesis several-stage simulation for many independent runs
Possible post-processing between stages:
* electron beam:  propagation through chicane via second-order tracking +CSR (in development)
* radiation: hard X-ray Self-seeding
* radiation: soft X-ray Self-seeding (in development)
End of explanation
"""


