"""
Example 9

- calculate a Wigner function in the 4D phase space
- verify how the spatial profile varies after a double slit using the Wigner function
"""

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

spectra.Open("sample.json")
spectra.SelectBL("X-ray-BL")

# enable parallel computing with multithread
spectra.SetParallel("Multithread", 4)

# flux spatial profile
spectra.SelectCalculation("srcpoint", "wigner", "phasespace", "XXpYYp")

# transverse range auto configuration
spectra.Set("config", "autot", True) 

# compute the Wigner function
spectra.StartCalculation(folder="./output", prefix="sample9", serial=1)

# load the Wigner function for CMD
spectra.LoadWigner4Propagation("./output/sample9-1.json")

# horizontally separated double slit
spectra.Set("config", "zrange", [18, 50])
spectra.Set("config", "zmesh", 17)
spectra.Set("config", "optics", "Double Slit")
spectra.Set("config", "optpos", 20)
spectra.Set("config", "aptxy", [0.1, 0.5])
spectra.Set("config", "aptdistxy", [0.3, 0])
spectra.StartCalculation(folder="./output", prefix="sample9", serial=2)

# vertically separated double slit
spectra.Set("config", "aptxy", [0.5, 0.02])
spectra.Set("config", "aptdistxy", [0, 0.06])
spectra.StartCalculation(folder="./output", prefix="sample9", serial=3)

# compare the two results
if isjupyter:
    from IPython.display import display, HTML
    display(HTML("<h2>Wavefront propagation after the horizontal and vertical double slit"))
    spectra.ConfigurePlot(type2d="Contour")
    spectra.PostProcessCLI.Plot(data=["sample9-2", "sample9-3"], xrange=[-0.8,0.8], yrange=[-0.4,0.4])
else:
    spectra.PostProcess.SelectData("sample9-2")
    spectra.PostProcess.ContourPlot()
    spectra.PostProcess.ComparativePlot("sample9-3")
    spectra.PostProcess.ComparativePlotCols(2)
    spectra.PostProcess.PlotRange(x=[-0.8,0.8], y=[-0.4,0.4])
    spectra.PostProcess.NormalizePlot(True)
    spectra.PostProcess.StartAnimation()

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