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

if __name__ == "__main__":
    import spectra 
# start with local source files
    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
spectra.PreProcess.Plot("betatron Functions")

# 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
spectra.PostProcess.Plot("Flux")

# 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
spectra.PostProcess.Plot("Flux")
spectra.PostProcess.ComparativePlot("sample2-1")

if __name__ == "__main__":
    input("Completed. Press enter to exit. ")
    spectra.Exit()
