using FFTW
using Plots
using LinearAlgebra: mul!
using Test
using BenchmarkTools
using Random
using BenchmarkPlots
using StatsPlots

@info "Threads: $(FFTW.nthreads())"

Ns = [2^6, 3^4, 2 * 3 * 5 * 7, 2^8, 2^2 * 3^2 * 5^2, 2^10, 2^11, 2^2 * 3 * 5^2 * 7]

flags = ["ESTIMATE", "MEASURE", "PATIENT", "EXHAUSTIVE"]
ext_flags = ["NO PLAN", flags...]
nothing

L = 2π
κ₀ = 2π/L
nothing

rng = Xoshiro(123)
num_modes = 24
max_mode = 16

kxs = rand(rng, 1:max_mode, num_modes)
kys = rand(rng, 1:max_mode, num_modes)
ars = 10*randn(rng, num_modes)
ais = 10*randn(rng, num_modes)
nothing

suite = BenchmarkGroup()
for flag in ext_flags
    suite[flag] = BenchmarkGroup()
end

plan_stats = Dict{String, Dict{Int, Float64}}(flag => Dict() for flag in flags)

for N in Ns
    x = y = (L/N):(L/N):L
    vort = sum(
        [
            2κ₀^2 * (kx^2 + ky^2) * (
                ar * cos.(κ₀ * (kx * one.(y) * x' + ky * y * one.(x)'))
                - ai * sin.(κ₀ * (kx * one.(y) * x' + ky * y * one.(x)'))
            )
            for (kx, ky, ar, ai) in zip(kxs, kys, ars, ais)
        ]
    )

    vort_hat = rfft(vort)

    flag = "NO PLAN"
    @info "N = $N; flag: $flag"
    suite[flag][N] = @benchmarkable rfft(w) setup = (w = copy($vort));

    for flag in flags
        @info "N = $N; flag: $flag"
        planed, pstats... = @timed plan_rfft(vort, flags = eval(Meta.parse("FFTW.$flag")))
        plan_stats[flag][N] = pstats.time

        suite[flag][N] = @benchmarkable mul!(w_hat, p, w) setup = (
            w_hat = copy($vort_hat);
            p = $planed;
            w = copy($vort)
        );
    end
end

suite

for N in Ns
    @info "N = $N"
    for flag in flags
        @info "$flag: \t$(BenchmarkTools.prettytime(plan_stats[flag][N] * 1.0e+9))"
    end
end

results, stats... = @timed run(suite)
nothing

stats

results

plt = plot(
    title="Minimum times for different plans with vector fields of different sizes",
    xlabel = "N",
    ylabel = "time (ns)",
    xticks = Ns[1:4],
    rotation = 90,
    titlefont=10,
    legend=:topleft
)

for flag in ext_flags
    plot!(plt, Ns[1:4], N -> minimum(values(results[flag][N]).times),
    linestyle = :dash,
    markershape = :rect,
    label="$flag"
    )
end

plt

plt = plot(
    title="Minimum times for different plans with vector fields of different sizes",
    xlabel = "N",
    ylabel = "time (ns)",
    xticks = Ns,
    rotation = 90,
    titlefont=10,
    legend=:topleft
)

for flag in ext_flags
    plot!(plt, Ns, N -> minimum(values(results[flag][N]).times),
    linestyle = :dash,
    markershape = :rect,
    label="$flag"
    )
end

plt

plt = plot(
    title="Median times for different plans with vector fields of different sizes",
    xlabel = "N",
    ylabel = "time (ms)",
    xticks = Ns[1:4],
    rotation = 90,
    titlefont=10,
    legend=:topleft
)

for flag in ext_flags
    plot!(plt, Ns[1:4], N -> median(values(results[flag][N]).times),
    linestyle = :solid,
    markershape = :circle,
    label="$flag"
    )
end
plt

plt = plot(
    title="Median times for different plans with vector fields of different sizes",
    xlabel = "N",
    ylabel = "time (ms)",
    xticks = Ns,
    rotation = 90,
    titlefont=10,
    legend=:topleft
)

for flag in ext_flags
    plot!(plt, Ns, N -> median(values(results[flag][N]).times)./ 1.0e+6,
    linestyle = :solid,
    markershape = :circle,
    label="$flag"
    )
end

plt

plts = []

for N in Ns
    push!(
        plts,
        violin(
            [results[flag][N].times for flag in flags],
            title = "Trials with N = $N",
            titlefont = 12,
            xticks = (1:length(flags), string.(flags)),
            yaxis = "time (ns)",
            legend = nothing
        )
    )
end

if isodd(length(Ns))
    push!(plts, plot(border = :none))
end

plt = plot(plts..., layout = (div(length(plts), 2), 2), size = (800, 1200))

for N in Ns
    @info "N = $N"
    for flag in flags
        @info "median time for flag $flag: $(BenchmarkTools.prettytime(median(values(results[flag][N]).times)))"
    end
end

flag = "PATIENT"

twos = filter(N -> isinteger(log2(N)), Ns)

mat = [ones(length(twos)) [N^2 * log(N^2) for N in twos]]

a, b = mat \ [median(values(results[flag][N].times)) for N in twos]

plt = plot(
    title="Median times for plan $flag with vector fields of different sizes",
    xlabel = "N (stretched out as N²)",
    ylabel = "time (ms)",
    xticks = (Ns.^2, string.(Ns)),
    rotation = 90,
    titlefont=10,
    legend=:topleft
)

plot!(plt, Ns.^2, map(N -> median(values(results[flag][isqrt(N)]).times)./ 1.0e+6, Ns.^2),
linestyle = :solid,
markershape = :circle,
label="$flag"
)

plot!(plt, Ns.^2, K -> (a + b * K * log(K) )./ 1.0e+6, linestyle = :dash, markershape = :square, label="N²log(N²) fit")


plt

lowtwos = filter(N -> isinteger(log2(N)), Ns[1:4])

lowmat = [ones(length(lowtwos)) [N^2 * log(N^2) for N in lowtwos]]

c, d = lowmat \ [median(values(results[flag][N].times)) for N in lowtwos]

plt = plot(
    title="Median times for plan $flag with vector fields of different sizes",
    xlabel = "N (stretched out as N²)",
    ylabel = "time (ms)",
    xticks = (Ns[1:4].^2, string.(Ns[1:4])),
    rotation = 90,
    titlefont=10,
    legend=:topleft
)

plot!(plt, Ns[1:4].^2, map(N -> median(values(results[flag][isqrt(N)]).times)./ 1.0e+6, Ns[1:4].^2),
linestyle = :solid,
markershape = :circle,
label="$flag"
)

plot!(plt, Ns[1:4].^2, K -> (c + d * K * log(K) )./ 1.0e+6, linestyle = :dash, markershape = :square, label="N²log(N²) fit")


plt