module Animation

export animate

using GLMakie
using ColorSchemes: ColorSchemes
using LaTeXStrings: @L_str
using JLD2: JLD2
using ProgressMeter: ProgressMeter

using ..ReCo: ReCo

include("common.jl")

const DEFAULT_FRAMERATE = 10
const DEFAULT_SHOW_CENTER_OF_MASS = false
const DEFAULT_SHOW_INTERACTION_CIRCLE = false
const DEFAULT_SHOW_SKIN_CIRCLE = false
const DEFAULT_SHOW_FRAME_DIFF = false
const DEFAULT_SHOW_PROGRESS = false

function animate_bundle!(args, sim_consts::ReCo.SimConsts)
    bundle_t = args.bundle.t
    bundle_c = args.bundle.c
    bundle_φ = args.bundle.φ

    for bundle_snapshot in 1:length(bundle_t)
        @simd for particle_ind in 1:(sim_consts.n_particles)
            c = bundle_c[particle_ind, bundle_snapshot]
            φ = bundle_φ[particle_ind, bundle_snapshot]

            args.particle_xs[][particle_ind] = c[1]
            args.particle_ys[][particle_ind] = c[2]

            color = angle_color(φ, args.color_scheme)
            args.particle_colors[][particle_ind] = color

            if args.show_interaction_circle
                args.interaction_colors[][particle_ind] = ColorSchemes.ColorTypes.RGBA(
                    color, 0.08
                )
            end

            if args.show_skin_circle
                args.skin_colors[][particle_ind] = ColorSchemes.ColorTypes.RGBA(color, 0.04)
            end
        end

        if args.show_center_of_mass
            center_of_mass = ReCo.center_of_mass(
                view(bundle_c, :, bundle_snapshot), sim_consts.half_box_len
            )
            args.center_of_mass_point[] = Point(center_of_mass)
        end

        if args.n_bundle == 1 && bundle_snapshot == 1
            scatter!(
                args.ax,
                args.particle_xs,
                args.particle_ys;
                markersize=2 * sim_consts.particle_radius,
                markerspace=SceneSpace,
                color=args.particle_colors,
            )

            if args.show_center_of_mass
                scatter!(
                    args.ax,
                    args.center_of_mass_point;
                    markersize=6 * sim_consts.particle_radius,
                    markerspace=SceneSpace,
                    color=ColorSchemes.ColorTypes.RGBA(1.0, 1.0, 1.0, 0.6),
                )
            end

            if args.show_interaction_circle
                scatter!(
                    args.ax,
                    args.particle_xs,
                    args.particle_ys;
                    markersize=2 * sim_consts.interaction_radius,
                    markerspace=SceneSpace,
                    color=args.interaction_colors,
                )
            end

            if args.show_skin_circle
                scatter!(
                    args.ax,
                    args.particle_xs,
                    args.particle_ys;
                    markersize=2 * sim_consts.skin_radius,
                    markerspace=SceneSpace,
                    color=args.skin_colors,
                )
            end

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

                    args.segment_xs[][first_ind] = bundle_c[i, frame_min_1][1]
                    args.segment_xs[][second_ind] = bundle_c[i, bundle_snapshot][1]

                    args.segment_ys[][first_ind] = bundle_c[i, frame_min_1][2]
                    args.segment_ys[][second_ind] = bundle_c[i, bundle_snapshot][2]
                end

                if bundle_snapshot == 2
                    linesegments!(
                        args.ax,
                        args.segment_xs,
                        args.segment_ys;
                        color=args.particle_colors,
                    )
                end

                notify(args.segment_xs)
                notify(args.segment_ys)
            end

            notify(args.particle_xs)
            notify(args.particle_ys)
            notify(args.particle_colors)

            if args.show_center_of_mass
                notify(args.center_of_mass_point)
            end

            if args.show_interaction_circle
                notify(args.interaction_colors)
            end

            if args.show_skin_circle
                notify(args.skin_colors)
            end
        end

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

        recordframe!(args.io)
    end

    return nothing
end

"""
    animate(sim_dir::String; <keyword arguments>)

Animate a simulation.

The output is `sim_dir/animation.mkv`.

Return `nothing`.

# Arguments
- `sim_dir::String`: Simulation directory.
- `framerate::Int64=$DEFAULT_FRAMERATE`: Framerate
- `show_center_of_mass::Bool=$DEFAULT_SHOW_CENTER_OF_MASS`: Show the center of mass as transparent white circle.
- `show_interaction_circle::Bool=$DEFAULT_SHOW_INTERACTION_CIRCLE`: Show the interaction radius with a circle around every particle.
- `show_skin_circle::Bool=$DEFAULT_SHOW_SKIN_CIRCLE`: Show the skin radius with a circle around every particle.
- `show_frame_diff::Bool=$DEFAULT_SHOW_FRAME_DIFF`: Show the translation of particles between two frames as lines connecting the old and new position. This is helpful to recognize unwanted jumps.
- `show_progress::Bool=$DEFAULT_SHOW_PROGRESS`: Show animation progress bar.
"""
function animate(
    sim_dir::String;
    framerate::Int64=DEFAULT_FRAMERATE,
    show_center_of_mass::Bool=DEFAULT_SHOW_CENTER_OF_MASS,
    show_interaction_circle::Bool=DEFAULT_SHOW_INTERACTION_CIRCLE,
    show_skin_circle::Bool=DEFAULT_SHOW_SKIN_CIRCLE,
    show_frame_diff::Bool=DEFAULT_SHOW_FRAME_DIFF,
    show_progress::Bool=DEFAULT_SHOW_PROGRESS,
)
    println("Initializing GLMakie...")

    GLMakie.activate!()
    set_theme!(theme_black())

    fig = Figure(; resolution=(1080, 1080))

    sim_consts = ReCo.load_sim_consts(sim_dir)

    ax, color_scheme = gen_axis_and_colorbar(fig, sim_consts)

    n_particles = sim_consts.n_particles

    animation_path = "$sim_dir/animation.mkv"

    record(fig, animation_path; framerate=framerate) do io
        particle_xs = Observable(Vector{Float64}(undef, n_particles))
        particle_ys = Observable(Vector{Float64}(undef, n_particles))

        particle_colors = Observable(
            Vector{ColorSchemes.ColorTypes.RGB{Float64}}(undef, n_particles)
        )

        center_of_mass_point =
            segment_xs =
                segment_ys =
                    interaction_circle_xs =
                        interaction_circle_ys =
                            skin_circle_xs =
                                skin_circle_ys = interaction_colors = skin_colors = nothing

        if show_center_of_mass
            center_of_mass_point = Observable(Point2(0.0, 0.0))
        end

        if show_interaction_circle
            interaction_circle_xs = Observable(Vector{Float64}(undef, n_particles))
            interaction_circle_ys = Observable(Vector{Float64}(undef, n_particles))

            interaction_colors = Observable(
                Vector{ColorSchemes.ColorTypes.RGBA{Float64}}(undef, n_particles)
            )
        end

        if show_skin_circle
            skin_circle_xs = Observable(Vector{Float64}(undef, n_particles))
            skin_circle_ys = Observable(Vector{Float64}(undef, n_particles))

            skin_colors = Observable(
                Vector{ColorSchemes.ColorTypes.RGBA{Float64}}(undef, n_particles)
            )
        end

        if show_frame_diff
            segment_xs = Observable(zeros(Float64, 2 * n_particles))
            segment_ys = Observable(zeros(Float64, 2 * n_particles))
        end

        bundle_paths = ReCo.sorted_bundle_paths(sim_dir)

        progress = ProgressMeter.Progress(
            length(bundle_paths); dt=2, enabled=show_progress, desc="Animation: "
        )

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

            args = (;
                # Input
                show_center_of_mass,
                show_interaction_circle,
                show_skin_circle,
                show_frame_diff,
                # Internal
                io,
                ax,
                bundle,
                particle_xs,
                particle_ys,
                particle_colors,
                center_of_mass_point,
                segment_xs,
                segment_ys,
                interaction_circle_xs,
                interaction_circle_ys,
                skin_circle_xs,
                skin_circle_ys,
                interaction_colors,
                skin_colors,
                color_scheme,
                n_bundle,
            )

            animate_bundle!(args, sim_consts)

            ProgressMeter.next!(progress)
        end
    end

    println("Animation done.")

    return nothing
end

end # module