"""
Example 12

- EEHG (echo enhanced harmonic generation) FEL
- 3-stage simulation: 1st/2nd modulator, 3rd 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 1&2
simplex.Open("sample_mod.json")

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

# reduce the bunch charge and length
simplex.Set("ebeam", "bunchleng", 5e-6)
simplex.Set("ebeam", "bunchcharge", 0.05)

# reduce the seed power and lengthen its pulse length
simplex.Set("seed", "pulselen", 100)
simplex.Set("seed", "pulseenergy", 5e-7)

# disable the quiet loading scheme, because it does not 
# work properly after a large longitudinal dispersion (R56)
simplex.Set("condition", "simoption", "Disable Quiet Loading")

# increase the number of beamlets (=macroparticles) to suppress artifacts
simplex.Set("condition", "beamlets", 1e7)

# shrink the simulation temporal window
simplex.Set("condition", "simrange", [-15e-6,15e-6])

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

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

# verify the energy modulation
simplex.PostProcess.SetDataProcessing("item", "Particle Motion")
simplex.PostProcess.SetDataProcessing("timerange", "Set Window")
simplex.PostProcess.SetDataProcessing("timewindow", [0, 0])
simplex.PostProcess.SetDataProcessing("r56pp", 5e-3)
simplex.PostProcess.RunDataProcessing()
simplex.PostProcess.SymbolPlot(1)
simplex.PostProcess.DuplicatePlot("Particle distribution at the 1st modulator exit with R<sub>56</sub>=5mm")

# 2nd modulator
# import the particle data at the 1st modulator exit with R56=5mm
simplex.Set("condition", "simoption", "N.A.") # disable any option to enable importing SIMPLEX output (next command)
simplex.Set("ebeam", "bmprofile", "SIMPLEX Output")
simplex.Set("ebeam", "r56", 5e-3)
simplex.Set("spxout", "spxfile", "./output/sample12-1.json")

# increase the seed power; other configurations are kept
simplex.Set("seed", "pulseenergy", 1e-6)

# start simulation in the 2nd modulator
simplex.StartSimulation(folder="./output", prefix="sample12", serial=2)

# verify the energy modulation
simplex.PostProcess.SetDataProcessing("r56pp", 2.4e-4)
simplex.PostProcess.RunDataProcessing()
simplex.PostProcess.SymbolPlot(1)
simplex.PostProcess.DuplicatePlot("Particle distribution at the 2nd modulator exit with R<sub>56</sub>=0.24mm")

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

# import the simulation result at the 2nd modulator exit
simplex.Set("ebeam", "bmprofile", "SIMPLEX Output")
simplex.Set("spxout", "spxfile", "./output/sample12-2.json")

# shrink the simulation temporal window
simplex.Set("condition", "simrange", [-15e-6,15e-6])

harmonics = [18, 23] # evaluate the 18th and 23rd harmonics
nr = 5 # number of R56 scan for each harmonic
r56 = [1.6e-4, 3.2e-4] # scan R56 from 160um to 320um
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="sample12", serial=harmonic)
    datanames = []
    for sn in range(1, nr):
        datanames.append("sample12-"+str(harmonic)+"_"+str(sn))
    simplex.PostProcess.SelectData("sample12-"+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()

