using BenchmarkTools, Random
using LinearAlgebra, RecursiveFactorization, VectorizationBase
nc = min(Int(VectorizationBase.num_cores()), Threads.nthreads())
BLAS.set_num_threads(nc)
BenchmarkTools.DEFAULT_PARAMETERS.seconds = 0.5

function luflop(m, n=m; innerflop=2)
    sum(1:min(m, n)) do k
        invflop = 1
        scaleflop = isempty(k+1:m) ? 0 : sum(k+1:m)
        updateflop = isempty(k+1:n) ? 0 : sum(k+1:n) do j
            isempty(k+1:m) ? 0 : sum(k+1:m) do i
                innerflop
            end
        end
        invflop + scaleflop + updateflop
    end
end

bas_mflops = Float64[]
rec_mflops = Float64[]
rec4_mflops = Float64[]
rec800_mflops = Float64[]
ref_mflops = Float64[]
ns = 4:8:500
for n in ns
    @info "$n × $n"
    rng = MersenneTwister(123)
    global A = rand(rng, n, n)
    bt = @belapsed LinearAlgebra.lu!(B) setup=(B = copy(A))
    push!(bas_mflops, luflop(n)/bt/1e9)

    rt = @belapsed RecursiveFactorization.lu!(B) setup=(B = copy(A))
    push!(rec_mflops, luflop(n)/rt/1e9)

    rt4 = @belapsed RecursiveFactorization.lu!(B; threshold=4) setup=(B = copy(A))
    push!(rec4_mflops, luflop(n)/rt4/1e9)

    rt800 = @belapsed RecursiveFactorization.lu!(B; threshold=800) setup=(B = copy(A))
    push!(rec800_mflops, luflop(n)/rt800/1e9)

    ref = @belapsed LinearAlgebra.generic_lufact!(B) setup=(B = copy(A))
    push!(ref_mflops, luflop(n)/ref/1e9)
end

using DataFrames, VegaLite
blaslib = if VERSION ≥ v"1.7.0-beta2"
  config = BLAS.get_config().loaded_libs
  occursin("libmkl_rt", config[1].libname) ? :MKL : :OpenBLAS
else
  BLAS.vendor() === :mkl ? :MKL : :OpenBLAS
end
df = DataFrame(Size = ns,
               Reference = ref_mflops)
setproperty!(df, blaslib, bas_mflops)
setproperty!(df, Symbol("RF with default threshold"), rec_mflops)
setproperty!(df, Symbol("RF fully recursive"), rec4_mflops)
setproperty!(df, Symbol("RF fully iterative"), rec800_mflops)
df = stack(df, [Symbol("RF with default threshold"),
                Symbol("RF fully recursive"),
                Symbol("RF fully iterative"),
                blaslib,
                :Reference], variable_name = :Library, value_name = :GFLOPS)
plt = df |> @vlplot(
                    :line, color = {:Library, scale={scheme="category10"}},
                    x = {:Size}, y = {:GFLOPS},
                    width = 1000, height = 600
                   )
save(joinpath(homedir(), "Pictures", "lu_float64_$(VERSION)_$(Sys.CPU_NAME)_$(nc)cores_$blaslib.png"), plt)

#=
using Plot
plt = plot(ns, bas_mflops, legend=:bottomright, lab="OpenBLAS", title="LU Factorization Benchmark", marker=:auto, dpi=300)
plot!(plt, ns, rec_mflops, lab="RecursiveFactorization", marker=:auto)
plot!(plt, ns, ref_mflops, lab="Reference", marker=:auto)
xaxis!(plt, "size (N x N)")
yaxis!(plt, "GFLOPS")
savefig("lubench.png")
savefig("lubench.pdf")
=#
