
function average_run(T; dims=2)
    return cumsum(T; dims=dims) ./ [1:size(T, dims);]
end

function err_l2(f, z, p, u, zs, ps, us)
    Z = f.(z) .- f.(zs)
    P = f.(p) .- f.(ps)
    U = f.(u) .- f.(us)
    Zr = average_run(Z; dims=1)
    Pr = average_run(P; dims=1)
    Ur = average_run(U; dims=1)
    return Float64.(sqrt.(vec(sum(abs2, Zr; dims=2) .+ sum(abs2, Pr; dims=2) .+ Ur .^ 2)))
end

function average_run2(T)
    return cumsum(T; dims=2) ./ [1:size(T, 2);]'
end
function MCf(f, z, p, u, zs, ps, us)
    Z = sum(f.(z); dims=2)
    P = sum(f.(p); dims=2)
    U = sum(f.(u); dims=2)
    Zs = sum(f.(zs); dims=2)
    Ps = sum(f.(ps); dims=2)
    Us = sum(f.(us); dims=2)

    return Z .+ P .+ U, Zs .+ Ps .+ Us
end

function single_err_traj(o, a, n_ref)
    E_fwd = zeros(n_ref - 1)
    E_bwd = zeros(n_ref - 1)
    d = o.d
    μ, D = a.μ, a.D
    z0 = o.q_sampler(d) .* D .+ μ
    ρ0 = randn(d)
    u0 = rand()
    zs, ρs, us = MixFlow.flow_trace_fwd(
        o, ft.(a.leapfrog_stepsize), MixFlow.ref_coord, ft.(z0), ft.(ρ0), ft.(u0), n_ref
    )
    zs = Float64.(zs)
    ρs = Float64.(ρs)
    us = Float64.(us)

    @views for i in 2:n_ref
        zb, ρb, ub, _ = MixFlow.one_bwd(
            o, a.leapfrog_stepsize, MixFlow.inv_ref_coord, zs[i, :], ρs[i, :], us[i]
        )
        E_fwd[i - 1] = sqrt(
            sum(abs2, zs[i - 1, :] .- zb) +
            sum(abs2, ρs[i - 1, :] .- ρb) +
            (us[i - 1] - ub)^2,
        )
    end
    @views for i in 1:(n_ref - 1)
        zf, ρf, uf, _ = MixFlow.one_fwd(
            o, a.leapfrog_stepsize, MixFlow.ref_coord, zs[i, :], ρs[i, :], us[i]
        )
        E_bwd[i] = sqrt(
            sum(abs2, zs[i + 1, :] .- zf) +
            sum(abs2, ρs[i + 1, :] .- ρf) +
            (us[i + 1] - uf)^2,
        )
    end
    return E_fwd, E_bwd
end
