"""
Example 10

- SASE FEL of 10-keV radiation using particle distribution data saved in an ASCII file
- adjust the injection condition to eliminate positional and angular errors at the lasing slice
"""

import sys
isjupyter = "ipykernel" in sys.modules

import simplex 

if isjupyter:
    simplex.Start(mode="c") 
    import plotly.io as pio
    pio.renderers.default = "notebook" # to enable exporting as an HTML file
    from IPython.display import display, HTML
else:
    simplex.Start() 

# open "sample.json" in the current directory
simplex.Open("sample_sase.json") 

# change "Bunch Profile" configuration
simplex.Set("ebeam", "bmprofile", "Particle Distribution")
simplex.Set("ebeam", "partfile", "./particles.dat")

# configure the format of the data file
simplex.PreProcess.LoadParticle()
simplex.PreProcess.ParticleDataFormat(unitE="gamma", pcharge=1e-15, bins=25) 
    # energy unit = gamma, charge/particle = 2fC, #bins/RMS bunch length = 25
if isjupyter: 
    display(HTML("<h2>E-t Phase Space"))
simplex.PreProcess.PlotParticles(x="s (m)", y="Energy (GeV)", max=50000)
if not isjupyter:
    simplex.PreProcess.DuplicatePlot("E-t Phase Space")
if isjupyter: 
    display(HTML("<h2>Expected Saturation Power"))
simplex.PreProcess.PlotSliceParameter("Saturation Power")
if not isjupyter:
    simplex.PreProcess.DuplicatePlot("Expected Saturation Power")

# export spatial profile of radiation
simplex.Set("datadump", "spatial", True)
simplex.Set("condition", "simrange", [-2e-6,3e-6])
simplex.Set("condition", "beamlets", 2e6)

# start simulation without change
simplex.StartSimulation(folder="./output", prefix="sample10", serial=1)
if isjupyter:
    display(HTML("<h2>Temporal Profile"))
    simplex.PostProcessCLI.Plot(item="Temporal Profile")
else:
    simplex.PostProcess.Plot("Temporal Profile")
    simplex.PostProcess.SetSlide(-1)
    simplex.PostProcess.DuplicatePlot()

# adjust the injection condition; roughly evaluated by the particles distribution
simplex.Set("dispersion", "einjec", True)
simplex.Set("dispersion", "exy", [-3e-2, 0]) # horizozntal position offset 30um
simplex.Set("dispersion", "exyp", [3e-3, 0]) # horizozntal angular offset -3urad
simplex.StartSimulation(folder="./output", prefix="sample10", serial=2)

# compare the two results: gain curve
if isjupyter:
    display(HTML("<h2>Gain Curve"))
    simplex.PostProcessCLI.Plot(data=["sample10-1", "sample10-2"], item="Pulse Energy")
else:
    simplex.PostProcess.SelectData("sample10-1")
    simplex.PostProcess.Plot("Pulse Energy")
    simplex.PostProcess.PlotScale(y="log")
    simplex.PostProcess.ComparativePlot("sample10-2")
    simplex.PostProcess.DuplicatePlot()

# compare the two results: growth of the spatial profile
if isjupyter:
    display(HTML("<h2>Spatial Profile"))
    simplex.PostProcessCLI.Plot(data=["sample10-1", "sample10-2"], item="Spatial Profile")
else:
    simplex.PostProcess.Plot("Spatial Profile")
    simplex.PostProcess.PlotScale(y="log")
    simplex.PostProcess.ComparativePlot("sample10-2")
    simplex.PostProcess.ComparativePlotCols(2)
    simplex.PostProcess.StartAnimation()

if not isjupyter:
    input("Completed. Press enter to exit. ")
    simplex.Exit()