local json = require 'json'
local os = require 'os'
local math = require 'math'

local printf in MAD.utility
local beam, element, sequence, track in MAD
local quadrupole in element
local fringe in MAD.dynmap

-- Load the particles generated by xtrack

local file = assert(io.open('initial_particles.json', 'r'))
local config_file = assert(io.open('config.json', 'r'))
local json_str = file:read("*all")
local config_str = config_file:read("*all")
file:close()
local xt_particles = json.decode(json_str)
local num_particles = #(xt_particles.x)
local element_name = json.decode(config_str).element
local kn_param_name = json.decode(config_str).param

-- Make the beam and the particles

local q0 = xt_particles.q0
local energy0 = (xt_particles.p0c[1] ^ 2 + xt_particles.mass0 ^ 2) ^ 0.5 / 1e9 -- GeV
local mass0 = xt_particles.mass0 / 1e9 -- GeV
local beta0 = xt_particles.beta0[1]
local p0c = xt_particles.p0c[1] / 1e9 -- GeV

printf("loaded particle: '%s', charge = %e, energy = %e, mass = %e, beta = %e, pc = %e\n",
    'xtrack', q0, energy0, mass0, beta0, p0c)

local beam = MAD.beam {
    particle = 'xtrack',
    energy = energy0,
    charge = q0,
    mass = mass0,
    -- beta = beta0,  -- this errors out with 'innacurate beta', even though
                      -- madng calculates the same value
    pc = p0c,
}

printf("madng beam: '%s', charge = %e, energy = %e, mass = %e, beta = %e, pc = %e\n",
    beam.particle, beam.charge, beam.energy, beam.mass, beam.beta, beam.pc)

local particles = {}
local idx = 1
for i = 1, num_particles do
    particles[idx] = {
        x = xt_particles.x[i],
        px = xt_particles.px[i],
        y = xt_particles.y[i],
        py = xt_particles.py[i],
        t = xt_particles.zeta[i] / beta0,
        pt = xt_particles.ptau[i],
    }
    idx = idx + 1
end

local function do_test(fringe_type, name)
    local knl = 1
    local length = 0.01

    local elem = MAD.element[element_name] { at = 0, l = length, fringe = fringe_type, frngmax = 4}
    elem[kn_param_name] = knl / length

    printf("element '%s' with k1 = %f, k2 = %f, k3 = %f, frngmax = %d\n", element_name, elem.k1 or 0, elem.k2 or 0, elem.k3 or 0, elem.frngmax or 0)

    local seq = sequence 'seq' {
        l = length,
        refer = 'entry',
        beam = beam,
        elem,
    }

    printf("will track %d particles\n", #particles)

    local trk, mfl = track {
        sequence = seq,
        beam = beam,
        X0 = particles,
        nturn = 1,
        debug = 4,
        method = 2,
    }

    printf("Tracked %d particles\n", #trk)

    local out = { x = {}, px = {}, y = {}, py = {}, t = {}, pt = {} }
    for i_row = 1, #trk do
        local row = trk:getrow(i_row)
        out.x[i_row] = row.x
        out.px[i_row] = row.px
        out.y[i_row] = row.y
        out.py[i_row] = row.py
        out.t[i_row] = row.t
        out.pt[i_row] = row.pt
    end

    local json_str = json.encode(out)
    local file, err = io.open(name, 'w+')
    if file then
        file:write(json_str)
        file:close()
    else
        printf("Error: %s\n", err)
    end

    return trk
end

do_test(fringe.mult, 'track_fringe.json')
do_test(fringe.none, 'track_no_fringe.json')
