module ReCoAnimation

export animate

using GLMakie
using ColorSchemes: cyclic_mrybm_35_75_c68_n256_s25
using LaTeXStrings: @L_str
using JSON3: JSON3
using JLD2: JLD2
using ProgressMeter: @showprogress

using ReCo: Bundle

function animate_bundle!(args)
    bundle_t = args.bundle.t
    bundle_c = args.bundle.c
    bundle_φ = args.bundle.φ

    π2 = 2 * π

    for frame in 1:length(bundle_t)
        @simd for i in 1:(args.n_particles)
            c = bundle_c[i, frame]
            args.circles[][i] = Circle(Point2(c[1], c[2]), args.particle_radius)

            color = get(args.color_scheme, rem2pi(bundle_φ[i, frame], RoundDown) / π2)
            args.colors[][i] = RGBAf(color)

            if args.debug
                args.interaction_circles[][i] = Circle(
                    Point2(bundle_c[i, frame]), args.interaction_r
                )
                args.skin_circles[][i] = Circle(Point2(bundle_c[i, frame]), args.skin_r)

                args.interaction_colors[][i] = RGBAf(color, 0.08)
                args.skin_colors[][i] = RGBAf(color, 0.04)
            end
        end

        if args.n_bundle == 1
            poly!(args.ax, args.circles; color=args.colors)

            if args.debug
                poly!(args.ax, args.interaction_circles; color=args.interaction_colors)
                poly!(args.ax, args.skin_circles; color=args.skin_colors)
            end

            println("Recording started!")
        else
            if args.debug && frame > 1
                @simd for i in 1:(args.n_particles)
                    first_ind = 2 * i - 1
                    second_ind = 2 * i
                    frame_min_1 = frame - 1

                    args.segments_x[][first_ind] = bundle_c[i, frame_min_1][1]
                    args.segments_x[][second_ind] = bundle_c[i, frame][1]

                    args.segments_y[][first_ind] = bundle_c[i, frame_min_1][2]
                    args.segments_y[][second_ind] = bundle_c[i, frame][2]
                end

                if frame == 2
                    linesegments!(
                        args.ax, args.segments_x, args.segments_y; color=args.colors
                    )
                end
            end

            notify(args.circles)
            notify(args.colors)

            if args.debug && frame > 1
                notify(args.interaction_circles)
                notify(args.interaction_colors)

                notify(args.skin_circles)
                notify(args.skin_colors)

                notify(args.segments_x)
                notify(args.segments_y)
            end
        end

        args.ax.title = "t = $(round(bundle_t[frame], digits=3))"

        recordframe!(args.io)
    end

    return nothing
end

function animate_after_loading(dir, animation_path, sim_consts, framerate, debug)
    set_theme!(theme_black())

    fig = Figure(; resolution=(1080, 1080))
    ax = Axis(
        fig[1, 1];
        limits=(
            -sim_consts.half_box_len,
            sim_consts.half_box_len,
            -sim_consts.half_box_len,
            sim_consts.half_box_len,
        ),
        aspect=AxisAspect(1),
        xlabel=L"x",
        ylabel=L"y",
    )

    color_scheme = cyclic_mrybm_35_75_c68_n256_s25

    Colorbar(fig[1, 2]; limits=(0, 2), colormap=color_scheme, label=L"\frac{\varphi}{\pi}")

    n_particles = sim_consts.n_particles

    record(fig, animation_path; framerate=framerate) do io
        circles = Observable(Vector{Circle}(undef, n_particles))
        colors = Observable(Vector{RGBAf}(undef, n_particles))

        segments_x =
            segments_y =
                interaction_circles =
                    skin_circles = interaction_colors = skin_colors = nothing

        if debug
            segments_x = Observable(zeros(2 * n_particles))
            segments_y = Observable(zeros(2 * n_particles))

            interaction_circles = Observable(Vector{Circle}(undef, n_particles))
            skin_circles = Observable(Vector{Circle}(undef, n_particles))

            interaction_colors = Observable(Vector{RGBAf}(undef, n_particles))
            skin_colors = Observable(Vector{RGBAf}(undef, n_particles))
        end

        particle_radius = sim_consts.particle_radius
        interaction_r = sim_consts.interaction_r

        bundle_paths = readdir("$dir/bundles"; join=true, sort=false)

        sort_perm = sortperm([
            parse(Int64, s[(findfirst("bundle_", s).stop + 1):(end - length(".jld2"))]) for
            s in bundle_paths
        ])

        bundle_paths = bundle_paths[sort_perm]
        sort_perm = nothing

        skin_r = 0.0

        @showprogress 1 for (n_bundle, bundle_path) in enumerate(bundle_paths)
            bundle::Bundle = JLD2.load_object(bundle_path)

            run_params_file = "$dir/runs/run_params_$n_bundle.json"
            if debug && isfile(run_params_file)
                skin_r::Float64 = JSON3.read(read(run_params_file, String)).skin_r
            end

            args = (;
                io,
                ax,
                bundle,
                sim_consts,
                debug,
                circles,
                colors,
                segments_x,
                segments_y,
                interaction_circles,
                skin_circles,
                interaction_colors,
                skin_colors,
                particle_radius,
                skin_r,
                n_particles,
                interaction_r,
                color_scheme,
                n_bundle,
            )

            animate_bundle!(args)
        end
    end

    return nothing
end

function animate(dir::String; framerate::Int64=1, debug::Bool=false)
    println("Generating animation...")

    sim_consts = JSON3.read(read("$dir/sim_consts.json", String))

    animation_path = "$dir/animation.mkv"

    animate_after_loading(dir, animation_path, sim_consts, framerate, debug)

    println("Animation done and saved to $animation_path.")

    return nothing
end

end # module