"""
Example 2

- SASE FEL of 10-keV radiation without/with the wakefield and compensation by undulator tapering
- compare the three results in the post-processor and export an ASCII file
"""

if __name__ == "__main__":
    import simplex 

# start with default settings
# Mode = interactive, Browser = Chrome, Source Files = remote
    simplex.Start(src="l") 

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

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

# plot the gain curve (pulse enerby vs. undulator length)
simplex.PostProcess.PlotGainCurve("Pulse Energy")

# turn on the wakefield
simplex.Set("wake", "wakeon", True)
simplex.PreProcess.Plot("Wakefield Temporal Profile")
simplex.PreProcess.DuplicatePlot("Wakefield along the electron bunch")
simplex.StartSimulation()

# compensating the energy loss by undulator tapering (with rough optimization)
simplex.Set("undulator", "taper", "Stair-Like")
simplex.Set("undulator", "opttype", "Compensate Wake")
simplex.StartSimulation()

# check the undulator tapering
simplex.PostProcess.PlotKTrend()
simplex.PostProcess.DuplicatePlot("K value variation to compensate for the energy loss")

# compare with the former results
simplex.PostProcess.SelectData("sample2-1")
simplex.PostProcess.PlotGainCurve("Pulse Energy")
simplex.PostProcess.ComparativePlot("sample2-2", "sample2-3")
simplex.PostProcess.DuplicatePlot("1/2: without/with wake, 3 compensation by tapering")

# export the comparative plot as an ASCII file
simplex.PostProcess.Export("./output/comparative.txt")

if __name__ == "__main__":
    input("Completed. Press enter to exit. ")
    simplex.Exit()
