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

function initial_particle_grid_pos(
    i::Int64, j::Int64, grid_box_width::Float64, half_box_len::Float64
)
    term = -0.5 * grid_box_width - half_box_len
    return SVector(i * grid_box_width + term, j * grid_box_width + term)
end

function generate_particles(grid_n::Int64, grid_box_width::Float64, half_box_len::Float64)
    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,
                initial_particle_grid_pos(i, j, grid_box_width, half_box_len),
                rand(Uniform(-π, π)),
            )

            id += 1
        end
    end

    return particles
end

function generate_particles(bundle::Bundle, n_particles::Int64)
    particles = Vector{Particle}(undef, n_particles)

    @simd for id in 1:n_particles
        particles[id] = Particle(id, bundle.c[id, end], bundle.φ[id, end])
    end

    return particles
end

function init_sim(;
    n_particles::Int64,
    v::Float64,
    δt::Float64=1e-5,
    packing_ratio::Float64=0.5,
    comment::String="",
    parent_dir::String="",
)
    @assert n_particles > 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_particles)))

    n_particles = grid_n^2

    half_box_len = sqrt(n_particles * π / packing_ratio) * σ / 4
    grid_box_width = 2 * half_box_len / grid_n

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

    particles = generate_particles(grid_n, grid_box_width, half_box_len)
    bundle = Bundle(n_particles, 1)
    save_snapshot!(bundle, 1, 0.0, particles)

    particles = nothing

    dir = "exports"
    if length(parent_dir) > 0
        dir *= "/$parent_dir"
    end
    dir *= "/$(now())_N=$(n_particles)_v=$(v)_#$(rand(1000:9999))"
    if length(comment) > 0
        dir *= "_$comment"
    end

    mkpath(dir)

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

    save_bundle(dir, bundle, 1, 0.0)

    return dir
end