SimulationDistributed.jl

using Distributed
using TOML

# Add worker processes if not already added
if nworkers() == 1
    addprocs(10)  # Add 10 workers
end

# Load packages on all workers
@everywhere begin
    using StaticArrays
    using LinearAlgebra
    using OrdinaryDiffEq
    using FastGaussQuadrature
    using Random
    using Printf

    include("InitialConditionGenerator.jl")
end

@everywhere function get_electric_field(t, gen, T_pulse, T_0)
    # Use envelope function from generator
    f = gen.envelope_func(t)

    # Analytical derivative for cos^2 envelope
    if t < T_0 - T_pulse || t > T_pulse
        df = 0.0
    elseif t < 0.0
        # Rising edge: f(t) = cos^2(π*t / (2*(T_pulse - T_0)))
        # df/dt = -π/(T_pulse - T_0) * sin(π*t / (T_pulse - T_0))
        df = -π / (T_pulse - T_0) * sin(π * t / (T_pulse - T_0))
    elseif t < T_0
        df = 0.0
    else  # T_0 <= t < T_pulse
        # Falling edge: f(t) = cos^2(π*(t - T_0) / (2*(T_pulse - T_0)))
        # df/dt = -π/(T_pulse - T_0) * sin(π*(t - T_0) / (T_pulse - T_0))
        df = -π / (T_pulse - T_0) * sin(π * (t - T_0) / (T_pulse - T_0))
    end

    α = gen.ω * t + gen.ϕ_cep
    sa, ca = sincos(α)
    Ex = gen.A0 * (df * sa + f * gen.ω * ca)
    Ey = -gen.χ * gen.A0 * (df * ca - f * gen.ω * sa)
    return SVector(Ex, Ey, 0.0)
end

@everywhere function electron_eom(u, p, t)
    gen, T_pulse, T_0, Z_eff, soft_core = p

    r = SVector(u[1], u[2], u[3])
    v = SVector(u[4], u[5], u[6])

    r2 = dot(r, r)
    dist_soft = sqrt(r2 + soft_core)
    Potential = -Z_eff / dist_soft
    Fc = -Z_eff * r / (dist_soft^3)
    E_field = get_electric_field(t, gen, T_pulse, T_0)

    dr = v
    dv = Fc - E_field

    kinetic = 0.5 * dot(v, v)
    integrand = kinetic + Potential + dot(r, Fc)
    dS = -integrand

    return SVector(dr[1], dr[2], dr[3], dv[1], dv[2], dv[3], dS)
end

# Callback to terminate integration when electron gets too close to core
@everywhere function condition_core_collision(u, t, integrator)
    r = sqrt(u[1]^2 + u[2]^2 + u[3]^2)
    return r - 1e-4  # Event occurs when this crosses zero
end

@everywhere function affect_core_collision!(integrator)
    terminate!(integrator)
end

# Smooth cos^2 envelope function
@everywhere function make_trapezoidal_envelope(T_pulse, T_0)
    return function(t)
        if t < T_0 - T_pulse
            return 0.0
        elseif t < 0.0
            # Rising edge: cos^2(π*t / (2*(T_pulse - T_0)))
            return cos(π * t / (2 * (T_pulse - T_0)))^2
        elseif t < T_0
            return 1.0
        elseif t < T_pulse
            # Falling edge: cos^2(π*(t - T_0) / (2*(T_pulse - T_0)))
            return cos(π * (t - T_0) / (2 * (T_pulse - T_0)))^2
        else
            return 0.0
        end
    end
end

@everywhere function run_worker_batch(worker_id, N_traj, ω, ϕ_cep, χ, E_peak, Ip, T_pulse, T_0, Z_eff, soft_core, pp_range)
    my_envelope = make_trapezoidal_envelope(T_pulse, T_0)
    gen = InitialConditionGenerator(ω, ϕ_cep, χ, E_peak, Ip, my_envelope; N_quad=20)

    params = (gen, T_pulse, T_0, Z_eff, soft_core)

    filename = "sfa_phase_worker_$(worker_id).txt"
    open(filename, "w") do io
        @printf(io, "tr,p_perp,p_z,r0_x,r0_y,r0_z,v0_x,v0_y,v0_z,rf_x,rf_y,rf_z,vf_x,vf_y,vf_z,S_sub_real,S_sub_imag,S_dyn\n")
    end

    n_valid = 0
    n_attempts = 0

    open(filename, "a") do io
        while n_valid < N_traj
            # Generate random inputs
            tr = rand() * T_0
            pp = (rand() * 2 - 1) * pp_range
            pz = (rand() * 2 - 1) * pp_range
            n_attempts += 1

            # Get initial conditions
            r0, v0, t_i, p_lab, S_sub = gen(pp, pz, tr)
            if norm(r0) < 1e-6; continue; end

            # Setup and solve ODE
            u0 = SVector(r0[1], r0[2], r0[3], v0[1], v0[2], v0[3], 0.0)
            prob = ODEProblem{false}(electron_eom, u0, (tr, T_pulse), params)
            
            # Create callback to terminate when electron gets too close to core
            cb = ContinuousCallback(condition_core_collision, affect_core_collision!)
            sol = solve(prob, Tsit5(), save_everystep=false, callback=cb)

            if sol.retcode != :Success; continue; end

            # Extract results
            uf = sol[end]
            rf = SVector(uf[1], uf[2], uf[3])
            vf = SVector(uf[4], uf[5], uf[6])
            S_dyn = uf[7]

            # Check if electron is ionized (positive energy)
            p = sqrt(vf[1]^2 + vf[2]^2 + vf[3]^2)
            r = sqrt(rf[1]^2 + rf[2]^2 + rf[3]^2)
            Ek = p^2 / 2.0 - Z_eff / r
            if Ek <= 0.0; continue; end

            S_sub_real = real(S_sub)
            S_sub_imag = imag(S_sub)

            # Write immediately
            @printf(io, "%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e,%.5e\n",
                    tr, pp, pz, r0[1], r0[2], r0[3], v0[1], v0[2], v0[3],
                    rf[1], rf[2], rf[3], vf[1], vf[2], vf[3],
                    S_sub_real, S_sub_imag, S_dyn)
            n_valid += 1

            if n_valid % 10000 == 0
                println("Worker $worker_id: $n_valid/$N_traj valid ($n_attempts attempts)")
                flush(io)
            end
        end
    end

    println("Worker $worker_id completed: $n_valid valid trajectories from $n_attempts attempts")
    return n_valid
end

function run_distributed_simulation()
    # Load configuration from TOML file
    config = TOML.parsefile("config.toml")

    # Extract parameters from config
    ω = config["laser"]["omega"]
    Ip = config["atom"]["Ip"]
    E_peak = config["laser"]["E_peak"]
    ϕ_cep = config["laser"]["phi_cep"]
    χ = config["laser"]["chi"]
    T_cycles = config["pulse"]["T_cycles"]
    T_0_cycles = config["pulse"]["T_0_cycles"]
    T_pulse = T_cycles * 2π / ω
    T_0 = T_0_cycles * 2π / ω
    Z_eff = config["atom"]["Z_eff"]
    soft_core = config["atom"]["soft_core"]
    pp_range = config["initial_conditions"]["pp_range"]

    N_total = 1000_000_000  # 1 billion
    N_workers = nworkers()
    N_per_worker = div(N_total, N_workers)

    println("Starting distributed simulation with $N_workers workers")
    println("Total trajectories: $N_total")
    println("Per worker: $N_per_worker")

    # Distribute work
    results = pmap(1:N_workers) do worker_id
        run_worker_batch(worker_id, N_per_worker, ω, ϕ_cep, χ, E_peak, Ip, T_pulse, T_0, Z_eff, soft_core, pp_range)
    end

    total_valid = sum(results)
    println("\nSimulation complete!")
    println("Total valid trajectories: $total_valid / $N_total")
    println("Output files: sfa_phase_worker_1.txt to sfa_phase_worker_$N_workers.txt")
end

run_distributed_simulation()