using SciMLExpectations, OrdinaryDiffEq, Distributions, IntegralsCubature

function eom!(du, u, p, t, A)
    du .= A * u
end

u0 = [1.0, 1.0]
tspan = (0.0, 3.0)
p = [1.0; 2.0]
A = [0.0 1.0; -p[1] -p[2]]
prob = ODEProblem((du, u, p, t) -> eom!(du, u, p, t, A), u0, tspan, p)
u0s_dist = (Uniform(1, 10), truncated(Normal(3.0, 1), 0.0, 6.0))
gd = GenericDistribution(u0s_dist...)
cov(x, u, p) = x, p

sm = SystemMap(prob, Tsit5(), save_everystep=false)

analytical = (exp(A * tspan[end]) * [mean(d) for d in u0s_dist])
analytical
#=
julia> analytical
2-element Vector{Float64}:
  1.5433991194037804
 -1.120209038276938
 =#



g(sol, p) = sol[:, end]
exprob = ExpectationProblem(sm, g, cov, gd; nout=length(u0))
sol = solve(exprob, Koopman(); quadalg=CubatureJLh(),
    ireltol=1e-3, iabstol=1e-3)

# Expectation of the states 1 and 2 at the final time point
sol.u

#=
2-element Vector{Float64}:
  1.5433860531082695
 -1.1201922503747408
=#

# Approximate error on the expectation
sol.resid
#=
2-element Vector{Float64}:
 7.193424502016654e-5
 5.2074632876847327e-5
=#