using Dates: Dates
using Distributions: Uniform
using JLD2: JLD2
using StaticArrays: SVector

const DEFAULT_PACKING_FRACTION = 0.5
const DEFAULT_δt = 1e-5
const DEFAULT_SKIN_TO_INTERACTION_RADIUS_RATIO = 3.0
const DEFAULT_EXPORTS_DIR = "exports"
const DEFAULT_PARENT_DIR = ""
const DEFAULT_COMMENT = ""
const DEFAULT_PARTICLE_RADIUS = 0.5
const DEFAULT_Dₜ = 1.0
const DEFAULT_μₜ = 1.0
const DEFAULT_ϵ = 100.0
const DEFAULT_HALF_BOX_LEN = 0.0

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 gen_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 gen_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 gen_sim_consts(
    n_particles::Int64,
    v₀::Float64;
    δt::Float64=DEFAULT_δt,
    packing_fraction::Float64=DEFAULT_PACKING_FRACTION,
    skin_to_interaction_radius_ratio::Float64=DEFAULT_SKIN_TO_INTERACTION_RADIUS_RATIO,
    half_box_len::Float64=0.0,
)
    @assert n_particles > 0
    @assert v₀ >= 0
    @assert δt in 1e-7:1e-7:1e-4
    @assert packing_fraction > 0

    μₜ = DEFAULT_μₜ
    Dₜ = DEFAULT_Dₜ
    particle_radius = DEFAULT_PARTICLE_RADIUS
    Dᵣ = 3 * Dₜ / ((2 * particle_radius)^2)
    σ = 2 * particle_radius * 2^(-1 / 6)
    ϵ = DEFAULT_ϵ
    interaction_radius = 2^(1 / 6) * σ

    skin_radius = skin_to_interaction_radius_ratio * interaction_radius

    buffer = 3
    max_approach_after_one_integration_step = buffer * (2 * v₀ * δt)
    @assert skin_radius >= interaction_radius + max_approach_after_one_integration_step

    if v₀ != 0.0
        n_steps_before_verlet_list_update = floor(
            Int64,
            (skin_radius - interaction_radius) / max_approach_after_one_integration_step,
        )
    else
        n_steps_before_verlet_list_update = 1000
    end

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

    n_particles = grid_n^2

    if half_box_len == 0.0
        half_box_len = sqrt(n_particles * π / packing_fraction) * σ * 2^(-11 / 6)
    elseif packing_fraction != DEFAULT_PACKING_FRACTION
        error("You can not specify half_box_len and packing_fraction at the same time!")
    else
        @assert half_box_len > 0.0
    end

    grid_box_width = 2 * half_box_len / grid_n

    return SimConsts(
        # Input
        n_particles,
        v₀,
        δt,
        packing_fraction,
        # Internal
        μₜ,
        Dₜ,
        particle_radius,
        Dᵣ,
        σ,
        ϵ,
        interaction_radius,
        skin_radius,
        n_steps_before_verlet_list_update,
        grid_n,
        half_box_len,
        grid_box_width,
    )
end

function init_sim_with_sim_consts(
    sim_consts::SimConsts;
    exports_dir::String=DEFAULT_EXPORTS_DIR,
    parent_dir::String=DEFAULT_PARENT_DIR,
    comment::String=DEFAULT_COMMENT,
)
    particles = gen_particles(
        sim_consts.grid_n, sim_consts.grid_box_width, sim_consts.half_box_len
    )
    bundle = Bundle(sim_consts.n_particles, 1)
    save_snapshot!(bundle, 1, 0.0, particles)

    sim_dir = exports_dir

    if length(parent_dir) > 0
        sim_dir *= "/$parent_dir"
    end

    start_datetime = Dates.now()
    sim_dir *= "/$(start_datetime)_N=$(sim_consts.n_particles)_v=$(sim_consts.v₀)_R$(rand(1000:9999))"

    if length(comment) > 0
        sim_dir *= "_$comment"
    end

    mkpath(sim_dir)

    task = @async JLD2.save_object("$sim_dir/sim_consts.jld2", sim_consts)

    save_bundle(sim_dir, bundle, 1, 0.0)

    runs_dir = "$sim_dir/runs"
    mkpath(runs_dir)

    wait(task)

    return sim_dir
end

"""
    init_sim(n_particles::Int64, v₀::Float64; <keyword arguments>)

Initialize a simulation and return the relative path of the simulation directory.

Return `nothing`.

# Arguments
- `n_particles::Int64`: Number of particles.
- `v₀::Float64`: Self-propulsion velocity. Only values in the interval [0.0, 80.0] are tested.
- `δt::Float64=$DEFAULT_δt`: Integration time step.
- `packing_fraction::Float64=$DEFAULT_PACKING_FRACTION`: Packing fraction.
- `skin_to_interaction_radius_ratio::Float64=$DEFAULT_SKIN_TO_INTERACTION_RADIUS_RATIO`: Ratio of skin radius to interaction radius.
- `exports_dir::String="$DEFAULT_EXPORTS_DIR"`: Path to exports directory relative to the directory `ReCo.jl`.
- `parent_dir::String="$DEFAULT_PARENT_DIR"`: Directory relative to `exports_dir` where the simulation directory is placed.
- `comment::String="$DEFAULT_COMMENT"`: Comment to append to the simulation directory name.
- `half_box_len::Float64=$DEFAULT_HALF_BOX_LEN` Half box length. The default of 0.0 means that the half box length will be calculated from the packing fraction. Otherwise, the provided half box length will be used. It is not possible to provide a half box length and a packing fraction at the same time.
"""
function init_sim(;
    n_particles::Int64,
    v₀::Float64,
    δt::Float64=DEFAULT_δt,
    packing_fraction::Float64=DEFAULT_PACKING_FRACTION,
    skin_to_interaction_radius_ratio::Float64=DEFAULT_SKIN_TO_INTERACTION_RADIUS_RATIO,
    exports_dir::String=DEFAULT_EXPORTS_DIR,
    parent_dir::String=DEFAULT_PARENT_DIR,
    comment::String=DEFAULT_COMMENT,
    half_box_len::Float64=DEFAULT_HALF_BOX_LEN,
)
    sim_consts = gen_sim_consts(
        n_particles,
        v₀;
        δt=δt,
        packing_fraction=packing_fraction,
        skin_to_interaction_radius_ratio=skin_to_interaction_radius_ratio,
        half_box_len,
    )

    return init_sim_with_sim_consts(
        sim_consts; exports_dir=exports_dir, parent_dir=parent_dir, comment=comment
    )
end