using Dates: Dates, now
using ProgressMeter: @showprogress
using Distributions: Normal
using CellListMap: Box, CellList, map_pairwise!, UpdateCellList!
using StaticArrays: SVector

import Base: wait

rand_normal01() = rand(Normal(0, 1))

function push_to_verlet_list(verlet_list, i, j)
    if i < j
        push!(verlet_list[i], Int64(j))
    else
        push!(verlet_list[j], Int64(i))
    end

    return nothing
end

function update_verlet_list!(args, cl)
    @simd for pv in args.verlet_list
        reset!(pv)
    end

    @simd for i in 1:(args.N)
        args.particles_c[i] = args.particles[i].c
    end

    cl = UpdateCellList!(args.particles_c, args.box, cl; parallel=false)

    map_pairwise!(
        (x, y, i, j, d2, output) -> push_to_verlet_list(args.verlet_list, i, j),
        nothing,
        args.box,
        cl;
        parallel=false,
    )

    return cl
end

function euler!(args)
    for i in 1:(args.N - 1)
        p1 = args.particles[i]
        verlet_list = args.verlet_list[p1.id]

        for j in 1:(verlet_list.last_ind)
            p2 = args.particles[verlet_list.v[j]]

            overlapping, r⃗₁₂, distance² = are_overlapping(
                p1, p2, args.interaction_r², args.l
            )

            if overlapping
                c = args.c₁ / (distance²^4) * (args.c₂ / (distance²^3) - 1)

                dc = c * r⃗₁₂
                p1.tmp_c -= dc
                p2.tmp_c += dc
            end
        end
    end

    @simd for p in args.particles
        p.tmp_c +=
            args.vδt * SVector(cos(p.φ), sin(p.φ)) +
            args.c₃ * SVector(rand_normal01(), rand_normal01())

        p.φ += args.c₄ * rand_normal01()

        restrict_coordinates!(p, args.l)

        p.c = p.tmp_c
    end

    return nothing
end

wait(n::Nothing) = n

function simulate(
    args,
    δt::Float64,
    T0::Float64,
    T::Float64,
    n_steps_before_verlet_list_update::Int64,
    n_steps_before_snapshot::Int64,
    n_bundles::Int64,
    dir::String,
    save_data::Bool,
)
    bundle_snapshot_counter = 0

    task::Union{Task,Nothing} = nothing

    cl = CellList(args.particles_c, args.box; parallel=false)

    start_time = now()
    println("Started simulation at $start_time.")

    @showprogress 0.6 for (integration_step, t) in enumerate(T0:δt:T)
        if (integration_step % n_steps_before_snapshot == 0) && save_data
            wait(task)
            bundle_snapshot_counter += 1
            save_snapshot!(args.bundle, bundle_snapshot_counter, t, args.particles)

            if bundle_snapshot_counter == args.n_bundle_snapshots
                task = @async begin
                    bundle_snapshot_counter = 0
                    n_bundles += 1

                    save_bundle(dir, args.bundle, n_bundles, t)
                end
            end
        end

        if integration_step % n_steps_before_verlet_list_update == 0
            cl = update_verlet_list!(args, cl)
        end

        euler!(args)
    end

    if bundle_snapshot_counter > 0
        bundle = first_n_snapshots(args.bundle, bundle_snapshot_counter)

        n_bundles += 1

        save_bundle(dir, bundle, n_bundles, T)
    end

    end_time = now()
    elapsed_time = Dates.canonicalize(Dates.CompoundPeriod(end_time - start_time))
    println("Simulation done at $end_time and took $elapsed_time.")

    return nothing
end