using StaticArrays: SVector
using JLD2: JLD2

struct Bundle
    t::Vector{Float64}
    c::Matrix{SVector{2,Float64}}
    φ::Matrix{Float64}
end

function Bundle(n_particles::Int64, n_snapshots::Int64)
    return Bundle(
        Vector{Float64}(undef, n_snapshots),
        Matrix{SVector{2,Float64}}(undef, (n_particles, n_snapshots)),
        Matrix{Float64}(undef, (n_particles, n_snapshots)),
    )
end

function first_n_snapshots(bundle::Bundle, n::Int64)
    return Bundle(bundle.t[1:n], bundle.c[:, 1:n], bundle.φ[:, 1:n])
end

function save_snapshot!(
    bundle::Bundle, n_snapshot::Int64, t::Float64, particles::Vector{Particle}
)
    bundle.t[n_snapshot] = t

    @simd for p in particles
        bundle.c[p.id, n_snapshot] = p.c

        bundle.φ[p.id, n_snapshot] = p.φ
    end

    return nothing
end

function set_sim_state(dir::String, n_bundles::Int64, T::Float64)
    open("$dir/sim_state.json", "w") do f
        JSON3.write(f, (n_bundles=n_bundles, T=round(T; digits=3)))
    end

    return nothing
end

function save_bundle(dir::String, bundle::Bundle, n::Int64, T::Float64)
    bundles_dir = "$dir/bundles"
    mkpath(bundles_dir)

    JLD2.save_object("$bundles_dir/bundle_$n.jld2", bundle)

    set_sim_state(dir, n, T)

    return nothing
end