# Installation
# using Pkg
# pkg"add MonteCarloMeasurements#gpu"
# pkg"add OrdinaryDiffEq ChangePrecision"

using CuArrays
using MonteCarloMeasurements, OrdinaryDiffEq, ChangePrecision, CuArrays, LinearAlgebra

function sim((±)::F, tspan) where F
    @changeprecision Float32 begin
        g = 9.79 ± 0.02; # Gravitational constant
        L = 1.00 ± 0.01; # Length of the pendulum
        u₀ = [0.0 ± 0.0, π / 3.0 ± 0.02] # Initial speed and initial angle
        gL = g/L
        function simplependulum(du,u,p,t)
            θ  = u[1]
            dθ = u[2]
            du[1] = dθ
            du[2] = -gL * sin(θ)
        end
        prob = ODEProblem(simplependulum, u₀, tspan)
        sol = solve(prob, Tsit5(), reltol = 1e-6, save_everystep=false, dense=false) # save_everystep=false, dense=false is required to not run out of GPU memory
    end
end

function naive_mc(tspan, n)
    for i = 1:n
        sim((x,y)->x+y*randn(), tspan)
    end
end

tspan = (0.0f0, 100f0) # Duration of simulation

CuArrays.reclaim()
CuArrays.allowscalar(false) # This will lead to an error in OrdinaryDiffEq/initdt.jl line 73, uncomment that line and hotpatch the function

# With 1e5 samples it goes about 2.5 times faster than the naive sim
@time sim((x,y)->MonteCarloMeasurements.CuParticles(fill(x,100_000) .+ y .* randn.()), tspan) # This takes around 19 seconds on my GPU, GTX 770
# 19.738116 seconds (30.17 M allocations: 33.190 GiB, 9.03% gc time)

# With 1e6 samples more than 50% of the time is spent in GC
@time sim((x,y)->MonteCarloMeasurements.CuParticles(fill(x,1_000_000) .+ y .* randn.()), tspan)
# 387.462061 seconds (30.24 M allocations: 321.369 GiB, 68.33% gc time)

@time naive_mc(tspan, 10000)
# 47.677040 seconds (46.87 M allocations: 3.223 GiB, 1.04% gc time)