using DynamicPPL, Distributions, LogDensityProblems, ForwardDiff, ADTypes, Test, Chairmarks, Test

@model trivial(y=nothing) = x ~ Normal()
@model function threaded_asm_and_obs(y)
    x = Vector{Float64}(undef, length(y))
    Threads.@threads for i in eachindex(y)
        x[i] ~ Normal()
        y[i] ~ Normal(x[i])
    end
end
@model function threaded_obs_only(y)
    x ~ Normal()
    Threads.@threads for i in eachindex(y)
        y[i] ~ Normal(x)
    end
end

function expected_logpdf(x, y)
    p = sum(logpdf.(Normal(), x))
    y === nothing && return p # this is the trivial model
    # handle both models
    lk = if length(x) == length(y)
        sum(logpdf.(Normal.(x), y))
    else
        sum(logpdf.(Normal(x[1]), y))
    end
    return p + lk
end
function expected_grad(x, y)
    return ForwardDiff.gradient(Base.Fix2(expected_logpdf, y), x)
end

### Need to manually construct the ranges, because the default
### behaviour is to generate a VarInfo and extract ranges from that
### -- but generating the VarInfo would fail if the model has
### threaded assume.
struct NotAVarInfo <: DynamicPPL.AbstractVarInfo
    xs::Vector{Float64}
end
function DynamicPPL.Experimental.get_ranges_and_linked(v::NotAVarInfo)
    n_xs = length(v.xs)
    return if n_xs == 50
        all_iden_ranges = NamedTuple()
        all_dict_ranges = Dict{VarName,DynamicPPL.RangeAndLinked}()
        for i in 1:n_xs
            all_dict_ranges[@varname(x[i])] = DynamicPPL.RangeAndLinked(i:i, false)
        end
        all_iden_ranges, all_dict_ranges
    else
        (x=DynamicPPL.RangeAndLinked(1:1, false),), Dict{VarName,DynamicPPL.RangeAndLinked}()
    end
end
Base.getindex(v::NotAVarInfo, ::Colon) = v.xs

s = Threads.nthreads() > 1 ? "s" : ""
println("$(Threads.nthreads()) thread$s")
function test_atomic(model_constructor, xs, ys; adtype=AutoForwardDiff())
    model = model_constructor(ys)
    println("    $(model.f)")
    # Evaluation: correctness
    print("        eval ")
    ldf = DynamicPPL.Experimental.FastLDF(model, getlogjoint_internal, NotAVarInfo(xs))
    lp1 = LogDensityProblems.logdensity(ldf, xs)
    @test lp1 ≈ expected_logpdf(xs, ys)
    # Evaluation: perf
    display(median(@be LogDensityProblems.logdensity(ldf, xs)))
    # Gradient: correctness
    ldf = DynamicPPL.Experimental.FastLDF(model, getlogjoint_internal, NotAVarInfo(xs); adtype=adtype)
    lp2, grad = LogDensityProblems.logdensity_and_gradient(ldf, xs)
    @test lp2 ≈ expected_logpdf(xs, ys)
    @test grad ≈ expected_grad(xs, ys)
    # Gradient: perf
    print("        grad ")
    display(median(@be LogDensityProblems.logdensity_and_gradient(ldf, xs)))
end

test_atomic(trivial, ones(1), nothing)
test_atomic(threaded_asm_and_obs, ones(50), zeros(50))
test_atomic(threaded_obs_only, ones(1), zeros(50))

# Both PRs (no threadsafe stuff)
#
# 1 thread
#     trivial
#         eval 22.459 ns (1 allocs: 16 bytes)
#         grad 45.738 ns (4 allocs: 128 bytes)
#     threaded_asm_and_obs
#         eval 30.125 μs (770 allocs: 20.156 KiB)
#         grad 207.167 μs (4369 allocs: 449.125 KiB)
#     threaded_obs_only
#         eval 23.792 μs (422 allocs: 12.125 KiB)
#         grad 26.708 μs (428 allocs: 14.766 KiB)
#
# This PR (atomic accumulators)
#
# 4 threads
#     trivial
#         eval 148.840 ns (10 allocs: 272 bytes)
#         grad 116.071 ns (7 allocs: 272 bytes)
#     threaded_asm_and_obs
#         eval 30.125 μs (989 allocs: 28.938 KiB)
#         grad 159.916 μs (3985 allocs: 264.266 KiB)
#     threaded_obs_only
#         eval 17.458 μs (445 allocs: 14.062 KiB)
#         grad 24.375 μs (445 allocs: 14.953 KiB)
#
# PR #1132 (ThreadSafeVarInfo)
#
# 4 threads
#     trivial
#         eval 573.864 ns (19 allocs: 976 bytes)
#         grad 210.262 ns (12 allocs: 1.156 KiB)
#     threaded_asm_and_obs
#         eval 28.958 μs (1194 allocs: 35.047 KiB)
#         grad 117.833 μs (4015 allocs: 533.875 KiB)
#     threaded_obs_only
#         eval 13.354 μs (454 allocs: 15.578 KiB)
#         grad 13.802 μs (450 allocs: 19.156 KiB)