"""
Example 5

- optimization of undulator taper to maximize the pulse energy
- 2-stage simulation: 1st without taper, 2nd scan taper
"""
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")

# stop at the 6th segment and export the raw data
simplex.Set("datadump", "particle", True)
simplex.Set("datadump", "radiation", True)
simplex.Set("datadump", "expstep", "All Segments")

# start simulation with an output file of "./output/sample5-1.json"
simplex.StartSimulation(folder="./output", prefix="sample5", serial=1)

# import the above simulation result
simplex.Set("ebeam", "bmprofile", "SIMPLEX Output")
simplex.Set("seed", "seedprofile", "SIMPLEX Output")
simplex.Set("spxout", "spxfile", "./output/sample5-1.json")

# matching distance = 6.15 (segment interval) - 270 (#periods) x 0.018 (period) = 1.29 m
simplex.Set("spxout", "matching", 1.29)

# kill the raw data export
simplex.Set("datadump", "particle", False)
simplex.Set("datadump", "radiation", False)

# assume stair-like taper from the 1st segment
simplex.Set("undulator", "taper", "Stair-Like")
simplex.Set("undulator", "opttype", "N.A.")
simplex.Set("undulator", "initial", 1)

taperinit = [5, 6] # optimize taper starting after the 5-th or 6-th segment
Mtotal = 12 # total segments
Nt = 4 # taper scan points

datanames = []
for k in taperinit:
    # use the result at the exit of the k-th segment
    simplex.Set("spxout", "spxstep", k-Mtotal)

    # number of segments in the taper
    simplex.Set("undulator", "segments", Mtotal-k)

    # data names in this run
    for n in range(Nt):
        datanames.append("sample5-"+str(k)+"_"+str(n))

    if k%2 == 0: # starting from an odd-number segment
        simplex.Set("lattice", "ltype", "FUDU (QF-U-QD-U)")
    else: # starting from an even-number segment
        simplex.Set("lattice", "ltype", "DUFU (QD-U-QF-U)")
    simplex.PreProcess.OptimizeLattice([10,10])

    # scan the taper value
    simplex.Scan("undulator", "base", 0, -0.0012, Nt, folder="./output", serial=k, iniSN=0)

# check the result; compare with the former result
if isjupyter:
    display(HTML("<h2>Taper optimization starting from 5/6-th segment"))
    simplex.PostProcessCLI.Plot(data=["sample5-1", *datanames], item="Pulse Energy")
else:
    simplex.PostProcess.SelectData("sample5-1")
    simplex.PostProcess.Plot("Pulse Energy")
    simplex.PostProcess.ComparativePlot(*datanames)
    simplex.PostProcess.PlotScale(y="linear")
    simplex.PostProcess.DuplicatePlot("Taper optimization starting from 5/6-th segment")

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