"""
Example 2

- plot betatron functions in the pre-processor and export as an ASCII file
- calculate spectra of BM radiation with two different magnetic strengths
"""

import sys
isjupyter = "ipykernel" in sys.modules

import spectra 
if isjupyter:
    spectra.Start(mode="c") 
    import plotly.io as pio
    pio.renderers.default = "notebook"
else:
    spectra.Start(src="l") 

# open "std.json" in the current directory
spectra.Open("sample.json") 

# switch to BM beamline
spectra.SelectBL("BM-BL")

# try pre-processing: betatron functions
if  isjupyter:
    from IPython.display import display, HTML
    display(HTML("<h2>Pre-Processing example: &beta; function"))
spectra.PreProcess.Plot("betatron Functions")

if not isjupyter:
    # export the pre-processing result
    spectra.PreProcess.Export("./output/beta.txt")

# select calculation: "Far Field & Ideal Condition::Energy Dependence::Partial Flux::Rectangular Slit"
spectra.SelectCalculation("far", "energy", "pflux", "slitrect")

# start calculation with an output file of "./ouutput/sample2-1.json"
spectra.StartCalculation(folder="./output", prefix="sample2", serial=1)

# check the result; plot the spectrum of angular flux density
if  isjupyter:
    display(HTML("<h2>Spectrum of bending magnet radiation"))
    spectra.PostProcessCLI.Plot(item="Flux")
else:
    spectra.ShowPostProcessor()

# enhance the magnetic field (1 Tesla)
spectra.Set("src", "b", 1)

# start calculation with an output file of "sample2-2.json"
spectra.StartCalculation(prefix="sample2")

# check the result; compare with the former result
if  isjupyter:
    display(HTML("<h2>Comparison with a stronger-field condition"))
    spectra.PostProcessCLI.Plot(data=["sample2-2", "sample2-1"], item="Flux")
else:
    spectra.ShowPostProcessor()
    spectra.PostProcess.ComparativePlot("sample2-1")

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