"""
Example 11

- HGHG (high gain harmonic generation) FEL
- 2-stage simulation: 1st modulator, 2nd radiator
"""

if __name__ == "__main__":
    import simplex 

# start with default settings
# Mode = interactive, Browser = Chrome, Source Files = remote
    simplex.Start(src="l") 

# open "sample_mod.json" to simulate the particle motion in the modulator
simplex.Open("sample_mod.json")

# seed wavelength: 266 nm
simplex.Set("felprm", "l1st", 266)

# 2(h+2) = 44 particles needed to simulate up to 20th harmonic (h=20)
simplex.Set("condition", "particles", 44)

# dump the particle data at the final step
simplex.Set("datadump", "particle", True)
simplex.Set("datadump", "expstep", "Final Step")

# start simulation in the modulator
simplex.StartSimulation(folder="./output", prefix="sample11", serial=1)

# verify the energy modulation
simplex.PostProcess.SetDataProcessing("item", "Particle Motion")
simplex.PostProcess.SetDataProcessing("timerange", "Set Window")
simplex.PostProcess.SetDataProcessing("timewindow", [-1, 1])
simplex.PostProcess.SetDataProcessing("r56pp", 1e-5)
simplex.PostProcess.RunDataProcessing()
simplex.PostProcess.SymbolPlot(1)
simplex.PostProcess.DuplicatePlot("Particle distribution in E-t space with 10um R<sub>56</sub>")

# open "sample_rad.json" to simulate the harmonic generation in the radiator
simplex.Open("sample_rad.json")

# import the simulation result in the modulator
simplex.Set("ebeam", "bmprofile", "SIMPLEX Output")
simplex.Set("spxout", "spxfile", "./output/sample11-1.json")

harmonics = [10, 15, 20] # 10th, 15th and 20th harmonics
nr = 5 # number of R56 scan for each harmonic
r56 = [0, 2e-5] # scan R56 from 0 to 20um
fr = 0.01  # spectral range (+-1%)  to show

for harmonic in harmonics:
    simplex.Set("felprm", "l1st", 266/harmonic)
    simplex.Scan("ebeam", "r56", r56[0], r56[1], nr, folder="./output", prefix="sample11", serial=harmonic)
    datanames = []
    for sn in range(1, nr):
        datanames.append("sample11-"+str(harmonic)+"_"+str(sn))
    simplex.PostProcess.SelectData("sample11-"+str(harmonic)+"_0")
    simplex.PostProcess.SpectralProfile()
    simplex.PostProcess.ComparativePlot(*datanames)
    enh = 1240/266*harmonic # central photon energy
    simplex.PostProcess.PlotRange(x=[enh*(1-fr),enh*(1+fr)])
    simplex.PostProcess.SetSlide(-1)
    simplex.PostProcess.DuplicatePlot("Spectrum of "+str(harmonic)+"th-harmonic radiation")

if __name__ == "__main__":
# wait for the user's response
    input("Completed. Press enter to exit.")

# exit
    simplex.Exit()

