"""
Example 1

- Gain curve of SASE FEL of 10-keV radiation
- Visualize the growth of the radiation profile
"""

import sys
isjupyter = "ipykernel" in sys.modules

import simplex 

if isjupyter:
    # start with CLI mode
    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:
    # start with default settings
    # Mode = interactive, Browser = Chrome, Source Files = remote
    simplex.Start() 

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

# export radiation profile data (longer time needed)
simplex.Set("datadump", "temporal", True)
simplex.Set("datadump", "spectral", True)
simplex.Set("datadump", "spatial", True)
simplex.Set("datadump", "angular", True)

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

if isjupyter:
    display(HTML("<h2>Gain Curve of SASE FEL"))
    simplex.PostProcessCLI.Plot(item="Pulse Energy", config={"yscale":"Logarithmic"})

    display(HTML("<h2>Growth of Temporal and Spectral Profiles"))
    simplex.PostProcessCLI.Plot(item=["Temporal Profile", "Spectral Profile"])

    display(HTML("<h2>Growth of Spatial and Angular Profiles"))
    simplex.PostProcessCLI.Plot(item=["Spatial Profile", "Angular Profile"])
else:
    # plot the gain curve (pulse energy vs. undulator length)
    simplex.PostProcess.Plot("Pulse Energy")

    # switch to semi-log scale
    simplex.PostProcess.PlotScale(y="log")

    # plot in a new window
    simplex.PostProcess.DuplicatePlot("Gain Curve: Pulse Energy vs Undulator Length")

    # visualize the growth of temporal and spectral profiles
    simplex.PostProcess.Plot("Temporal Profile")
    simplex.PostProcess.MultiPlot(item="Spectral Profile")
    simplex.PostProcess.MultiPlotCols(2)
    simplex.PostProcess.SetSlide(-1)
    simplex.PostProcess.DuplicatePlot("Growth of Temporal and Spectral Profiles")

    # visualize the growth of spatial and angular profiles
    simplex.PostProcess.Plot("Spatial Profile")
    simplex.PostProcess.MultiPlot(item="Angular Profile")
    simplex.PostProcess.SetSlide(-1)
    simplex.PostProcess.ShowColorScale(False)
    simplex.PostProcess.DuplicatePlot("Growth of Spatial and Angular Profiles")

if not isjupyter:
    # wait for the user's response
    input("Completed. Press enter to exit. ")
    # exit
    simplex.Exit()