"""
Example 10

- SASE FEL of 10-keV radiation using particle distribution data saved in an ASCII file
- adjust the injection condition to eliminate positional and angular errors at the lasing slice
"""

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

# open "sample.json" in the current directory
simplex.Open("sample_sase.json") 

# change "Bunch Profile" configuration
simplex.Set("ebeam", "bmprofile", "Particle Distribution")
simplex.Set("ebeam", "partfile", "./particles.dat")

# configure the format of the data file
simplex.PreProcess.LoadParticle()
simplex.PreProcess.ParticleDataFormat(unitE="gamma", pcharge=1e-15, bins=25) 
    # energy unit = gamma, charge/particle = 2fC, #bins/RMS bunch length = 25
simplex.PreProcess.PlotParticles(x="s (m)", y="Energy (GeV)", max=50000)
simplex.PreProcess.DuplicatePlot("E-t Phase Space")
simplex.PreProcess.PlotSliceParameter("Saturation Power")
simplex.PreProcess.DuplicatePlot("Expected Saturated Power")

# export spatial profile of radiation
simplex.Set("datadump", "spatial", True)
simplex.Set("condition", "simrange", [-2e-6,3e-6])
simplex.Set("condition", "beamlets", 2e6)

# start simulation without change
simplex.StartSimulation(folder="./output", prefix="sample10", serial=1)
simplex.PostProcess.TemporalProfile()
simplex.PostProcess.SetSlide(-1)
simplex.PostProcess.DuplicatePlot()

# adjust the injection condition; roughly evaluated by the particles distribution
simplex.Set("dispersion", "einjec", True)
simplex.Set("dispersion", "exy", [-3e-2, 0]) # horizozntal position offset 30um
simplex.Set("dispersion", "exyp", [3e-3, 0]) # horizozntal angular offset -3urad
simplex.StartSimulation(folder="./output", prefix="sample10", serial=2)

# compare the two results: gain curve
simplex.PostProcess.SelectData("sample10-1")
simplex.PostProcess.PlotGainCurve("Pulse Energy")
simplex.PostProcess.ComparativePlot("sample10-2")
simplex.PostProcess.DuplicatePlot()

# compare the two results: growth of the spatial profile
simplex.PostProcess.SpatialProfile()
simplex.PostProcess.ComparativePlot("sample10-2")
simplex.PostProcess.ComparativePlotCols(2)
simplex.PostProcess.StartAnimation()

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