"""
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
"""

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") 

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

if isjupyter:
    display(HTML("<h2>Gain Curve of SASE FEL"))
    simplex.PostProcessCLI.Plot(item="Pulse Energy", config={"yscale":"Logarithmic"})
else:
    # plot the gain curve (pulse energy vs. undulator length)
    simplex.PostProcess.Plot("Pulse Energy")
    simplex.PostProcess.PlotScale(y="log")

# turn on the wakefield
simplex.Set("wake", "wakeon", True)
simplex.PreProcess.Plot("Wakefield Temporal Profile")
if not isjupyter:
    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
if isjupyter:
    display(HTML("<h2>K Value Trend"))
    simplex.PostProcessCLI.Plot(item="K Value Trend")
else:
    simplex.PostProcess.Plot("K Value Trend")
    simplex.PostProcess.DuplicatePlot("K value variation to compensate for the energy loss")

# compare with the former results
if isjupyter:
    display(HTML("<h2>Comparison: 1/2: without/with wake, 3 compensation by tapering"))
    simplex.PostProcessCLI.Plot(data=["sample2-1", "sample2-2", "sample2-3"], item="Pulse Energy")
else:
    simplex.PostProcess.SelectData("sample2-1")
    simplex.PostProcess.Plot("Pulse Energy")
    simplex.PostProcess.PlotScale(y="linear")
    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 not isjupyter:
    input("Completed. Press enter to exit. ")    
    simplex.Exit()