"""
Example 5

- optimization of undulator taper to maximize the pulse energy
- 2-stage simulation: 1st without taper, 2nd scan taper
"""

if __name__ == "__main__":
    import simplex 

# start with default settings
# Mode = interactive, Browser = Chrome, Source Files = remote
    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
simplex.PostProcess.SelectData("sample5-1")
simplex.PostProcess.PlotGainCurve("Pulse Energy")
simplex.PostProcess.ComparativePlot(*datanames)
simplex.PostProcess.PlotScale(y="linear")
simplex.PostProcess.DuplicatePlot("Taper optimization starting from 5/6-th segment")

if __name__ == "__main__":
# wait for the user's response
    input("Completed. Press enter to exit.")

# exit
    simplex.Exit()

