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

if __name__ == "__main__":
    import spectra
    spectra.Start()

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

# 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", "aptx", 0.1)
spectra.Set("config", "aptdistx", 0.3)
spectra.Set("config", "apty", 0.5)
spectra.Set("config", "aptdisty", 0)
spectra.StartCalculation(folder="./output", prefix="sample9", serial=2)

# vertically separated double slit
spectra.Set("config", "aptx", 0.5)
spectra.Set("config", "aptdistx", 00)
spectra.Set("config", "apty", 0.02)
spectra.Set("config", "aptdisty", 0.06)
spectra.StartCalculation(folder="./output", prefix="sample9", serial=3)

# compare the two results
spectra.PostProcess.Import("./output/sample9-2.json")
spectra.PostProcess.Import("./output/sample9-3.json")
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 __name__ == "__main__":
    input("Completed. Press enter to exit. ")
    spectra.Exit()
