using Distributions: Uniform
using Dates: now
using JSON3: JSON3

function initial_particle_grid_pos(i, j; grid_box_width, l)
    term = -0.5 * grid_box_width - l
    return [k * grid_box_width + term for k in (i, j)]
end

function generate_particles(grid_n, grid_box_width, l)
    particles = Vector{Particle}(undef, grid_n^2)

    id = 1

    for i in 1:grid_n
        for j in 1:grid_n
            particles[id] = Particle(;
                id=id,
                c=initial_particle_grid_pos(i, j; grid_box_width=grid_box_width, l=l),
                φ=rand(Uniform(-π, π)),
            )

            id += 1
        end
    end

    return particles
end

function init_sim(;
    N::Int64, v::Float64, δt::Float64=1e-5, packing_ratio::Float64=0.5, comment::String=""
)
    @assert N > 0
    @assert v >= 0
    @assert δt in 1e-7:1e-7:1e-4
    @assert packing_ratio > 0

    μ = 1.0
    D₀ = 1.0
    particle_diameter = 1.0
    Dᵣ = 3 * D₀ / (particle_diameter^2)
    σ = 1.0
    ϵ = 100.0
    interaction_r = 2^(1 / 6) * σ

    grid_n = round(Int64, ceil(sqrt(N)))

    N = grid_n^2

    l = sqrt(N * π / packing_ratio) * σ / 4
    grid_box_width = 2 * l / grid_n

    sim_consts = (;
        # Input
        N,
        v,
        δt,
        packing_ratio,
        # Calculated
        μ,
        D₀,
        particle_diameter,
        Dᵣ,
        σ,
        ϵ,
        interaction_r,
        grid_n,
        l,
        grid_box_width,
    )

    particles = generate_particles(grid_n, grid_box_width, l)
    bundle = Bundle(N, 1)
    save_snapshot!(bundle, 1, 0.0, particles)

    dir = "exports/$(now())_N=$(N)_v=$(v)_#$(rand(1000:9999))$comment"
    mkpath(dir)

    open("$dir/sim_consts.json", "w") do f
        JSON3.pretty(f, sim_consts)
    end

    save_bundle(dir, bundle, 1, 0.0)

    return dir
end