"""
Example 6

- double-color SASE FEL by applying different K values in the first and second sections
- visualize the variation of spectral profile
"""

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

simplex.Set("undulator", "taper", "Custom")
simplex.Set("datadump", "spectral", True)

taper = []
for n in range(12):
    taper.append(["0.02" if n >= 5 else "-0.02", "0"])
    # first 5 segments with lower K, last 7 segments with higher K
simplex.Set("undulator", "tapercustom", taper)

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

if isjupyter:
    display(HTML("<h2>Example of double-color SASE FEL"))
    simplex.PostProcessCLI.Plot(item="Spectral Profile")
else:
    simplex.PostProcess.Plot("Spectral Profile")
    simplex.PostProcess.StartAnimation()

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