"""
Example 13

- Oscillator FEL at 13.3 nm
- perform simulation for a single turn, and use the radiation data for the next turn
"""

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

import numpy as np
import json
import math

# constants
Wavelen1eV = 1.23984247e-6
CC = 2.9979246e+8

# keys for radiation profiles
profkeys = ["Gain Curve", "Characteristics", "Temporal Profile", "Spectral Profile", "Spatial Profile", "Angular Profile"]

# compute the phase factor & mirror aperture
def SetApt(Apt, Ld, apt, wavel):
    bdr = math.sqrt(dDx*dDy)
    aptin = apt-bdr
    for m in range(M):
        mf = m if m <= M/2 else m-M
        qx = mf*dDx
        for n in range(N):
            nf = n if n <= N/2 else n-N
            qy = nf*dDy
            q = (math.hypot(qx, qy)-aptin)/bdr
            q = min(q, 1)
            if q >= 0:
                phase = -np.pi/wavel*Ld*(qx**2+qy**2) # phase advance in the distance of Ld
                Apt[m][2*n] = q*math.cos(phase)
                Apt[m][2*n+1] = q*math.sin(phase)

# transfer from the undulator exit to the entrance with outcoupling
def Transfer(Apt, E):
    for m in range(M):
        for n in range(N):
            dummy = E[m][2*n]
            E[m][2*n  ] = E[m][2*n]*Apt[m][2*n  ]-E[m][2*n+1]*Apt[m][2*n+1]
            E[m][2*n+1] =     dummy*Apt[m][2*n+1]+E[m][2*n+1]*Apt[m][2*n  ]

# export spatial profile (for debugging)
def Export(E, isfar, suf):
    Mh = M>>1
    Nh = N>>1
    dx = dDx if isfar else Dx
    dy = dDy if isfar else Dy
    with open("./temp/exy"+suf+".dat", "w") as f:
        for n in range(-Nh, Nh+1):
            nf = n if n >= 0 else n+N
            for m in range(-Mh, Mh+1):
                mf = m if m >= 0 else m+M
                P = E[mf][2*nf]**2+E[mf][2*nf+1]**2
                f.write(f'{m*dx*1000}\t{n*dy*1000}\t{E[mf][2*nf]}\t{E[mf][2*nf+1]}\t{P}\n')

# retrieve and store the result after a single turn
def RetrieveData(turn, fpath, profiles):
    with open(fpath, "r") as f:
        obj = json.load(f)
        for prof in profkeys:
            if prof in obj:
                dim = obj[prof]["dimension"]
                profiles[prof]["data"][dim-1].append(turn)
                ndata = len(profiles[prof]["data"])
                for n in range(dim, ndata):
                    data = obj[prof]["data"][n]
                    if dim == 1:
                        profiles[prof]["data"][n].append(data[-1])
                    else:
                        ndata = len(data)
                        nz = len(obj[prof]["data"][dim-1])
                        nslice = ndata//nz
                        sdata = np.array(data).reshape([nz, nslice])
                        profiles[prof]["data"][n] += sdata[-1].tolist()

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

# undulator specs.
Nu = simplex.Get("undulator", "periods")
lu = simplex.Get("undulator", "lu")
Lu = Nu*lu*1e-3

# start simulation with an output file of "./output/sample1.json"
simplex.StartSimulation(folder="./output", prefix="sample13", serial=0)

# get parameters used for the simulation
with open("./output/sample13-0.json", "r") as f:
    obj = json.load(f)
    config = obj["Raw Data Export"]
    K = len(config["Steps (m)"]) # steps exported
    L = len(config["Slices (m)"]) # total slices
    MN = config["Grid Points"]
    M = MN[0] # x grids
    N = MN[1] # y grids
    Nb = config["Beamlets"] # beamlets
    Np = config["Particles/Beamlet"] # particles/beamlet
    Dxy = config["Grid Intervals (m,rad)"]
    Dx = Dxy[0]
    Dy = Dxy[1]
    dDx = Dxy[2]
    dDy = Dxy[3]
    wavel = Wavelen1eV/config["Central Photon Energy (eV)"]

# timing between e-bunch and FEL pulse at the undulator entrance w/o cavity detuning
sdif = -wavel*Nu

# adjust cavity detuning (shoten cavity legth by half the slippage)
sdif += wavel*Nu/2
tdif = sdif/CC*1e15 # should be give in "fs"

# modify relevant parameters for the 1st and later turns
simplex.Set("seed", "seedprofile", "SIMPLEX Output")
simplex.Set("seed", "timing", tdif)
simplex.Set("spxout", "matching", 0)

# outcoupling: aperture with a radius of 25urad
apt = 25e-6
Apt = np.full((M, 2*N), 0.0)
SetApt(Apt, -Lu, apt, wavel)

# generate an object to store data at each turn
profiles = json.loads(json.dumps(obj))
for datakey in profkeys:
    if datakey in profiles:
        ndata = len(profiles[datakey]["titles"])
        dim = profiles[datakey]["dimension"]
        profiles[datakey]["titles"][dim-1] = "Turns" 
        profiles[datakey]["units"][dim-1] = "-"
        for n in range(dim-1, ndata):
            profiles[datakey]["data"][n] = []

# save the 0th turn
RetrieveData(0, f"./output/sample13-0.json", profiles)

# data counts per step
nc = 4*2*M*N*L # bytes = 4(byte)*2(re,im)*M(x)*N(y)*L(slice)

# number of turns
turns = 30

for nt in range(turns):
    print(f"-- {nt} turn started --")
    fdata = f"./output/sample13-{nt%2}"

    # retrive the radiation data from the binary file "*-1.fld"
    x = np.fromfile(fdata+"-1.fld", dtype=np.float32, count=nc)
    E = x.reshape([L, M, 2*N]) # E[l][m][2*n] real part, [2*n+1] imaginary part

    # transfer the radiation field to the entrance of the undulator
    for l in range(L):
        print(f"Transferring wavefront: {l+1}/{L} slice", end="\r")
        Transfer(Apt, E[l])
    print("")

    # save the data in the binary file
    x = E.reshape(2*L*M*N)
    x.tofile(fdata+"-1.fld")

    # modify the input SIMPLEX file
    simplex.Set("spxout", "spxfile", fdata+".json")

    # change the shot noise condition
    simplex.Set("condition", "randseed", nt+1)

    # start simulation for the (nt+1)-th turn
    simplex.StartSimulation(folder="./output", prefix="sample13", serial=(nt+1)%2)

    # save the (nt+1)-th turn result
    RetrieveData(nt+1, f"./output/sample13-{(nt+1)%2}.json", profiles)

# export the bundled data (growth of radiation vs. number of turns) and export in the post-processor
with open("./output/sample13.json", "w") as f:
    json.dump(profiles, f, indent=4, sort_keys=True, separators=(',', ': '))
simplex.PostProcess.Import("./output/sample13.json")

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

