"""
Example 3

- SASE FEL of 10-keV radiation with three different combinations of electron energies and undulator K values
- compare the three results in the post-processor
"""

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

EGeV = [6, 8, 10]

for E in EGeV:
    simplex.Set("ebeam", "eenergy", E) # change the electron energy
    simplex.Set("felprm", "e1st", 10000) # set the photon energy; K value is automatically set
    print(f"K@{E}GeV = {simplex.Get("undulator", "K"):.3f}")
    simplex.PreProcess.OptimizeLattice([10,10]) # adjust the beta-matching condition
    simplex.StartSimulation(folder="./output", prefix="sample3", serial=E)

# plot the gain curve (pulse energy vs. undulator length)
if isjupyter:
    display(HTML("<h2>10-keV SASE FEL with different energies and K values"))
    simplex.PostProcessCLI.Plot(data=["sample3-6", "sample3-8", "sample3-10"], 
        item="Pulse Energy", config={"yscale":"Logarithmic"})
else:
    simplex.PostProcess.SelectData("sample3-6")
    simplex.PostProcess.Plot("Pulse Energy")
    simplex.PostProcess.ComparativePlot("sample3-8", "sample3-10")
    simplex.PostProcess.PlotScale(y="log")

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

K@6GeV = 1.341

K@8GeV = 2.180

K@10GeV = 2.924