"""
    ConvergenceSimulation(solutions, convergence_axis;
        auxdata = nothing, additional_errors = nothing, expected_value = nothing)

Collect solutions and error series from a convergence experiment. The constructor
combines the error measurements from `solutions` with `additional_errors` and estimates
the observed convergence order between adjacent points on `convergence_axis`.

Set `expected_value` when the error series are supplied through `additional_errors`, as
done by ensemble convergence tests.
"""
mutable struct ConvergenceSimulation{SolType}
    solutions::Array{SolType}
    errors::Any
    N::Any
    auxdata::Any
    𝒪est::Any
    convergence_axis::Any
end

function ConvergenceSimulation(
        solutions, convergence_axis;
        auxdata = nothing, additional_errors = nothing,
        expected_value = nothing
    )
    N = size(solutions, 1)
    errors = Dict() #Should add type information
    if expected_value == nothing
        if isnothing(solutions[1].errors) || isempty(solutions[1].errors)
            error("Errors dictionary is empty. No analytical solution set. If you used `test_convergence` you may consider `analyticless_test_convergence` instead.")
        end
        for k in keys(solutions[1].errors)
            errors[k] = [mean(sol.errors[k]) for sol in solutions]
        end
    end
    if additional_errors != nothing
        for k in keys(additional_errors)
            errors[k] = additional_errors[k]
        end
    end

    𝒪est = Dict((calc𝒪estimates(p) for p in pairs(errors)))
    #𝒪est = Dict(map(calc𝒪estimates,errors))
    𝒪esttmp = Dict() #Makes Dict of Any to be more compatible
    for (k, v) in 𝒪est
        if length(v) == 1
            push!(𝒪esttmp, Pair(k, v[1]))
        else
            push!(𝒪esttmp, Pair(k, v))
        end
    end
    𝒪est = 𝒪esttmp
    return (ConvergenceSimulation(solutions, errors, N, auxdata, 𝒪est, convergence_axis))
end

"""
    ConvergenceTrajectory(t, u, u_analytic, errors)

The part of a trajectory's solution that a convergence study actually consumes: its
error measurements and the endpoint values the weak errors are formed from.

[`test_convergence`](@ref) stores one of these per trajectory in place of the full
solution unless `retain_solutions = true`, which is what keeps a Monte-Carlo study's
memory proportional to the errors it reports rather than to the solver state it
happened to allocate along the way.
"""
struct ConvergenceTrajectory{tType, uType, EType <: NamedTuple}
    t::Tuple{tType}
    u::Tuple{uType}
    u_analytic::Tuple{uType}
    errors::EType
end

"""
    _convergence_output_func(sol, ctx) -> (ConvergenceTrajectory, false)

`output_func` that reduces a trajectory to its errors and endpoints as soon as it is
solved, so the ensemble never holds the full solutions at once.

Only valid when the weak timeseries and dense errors are switched off, since those
need the whole timeseries and the noise process; [`test_convergence`](@ref) checks
that before installing this.
"""
function _convergence_output_func(sol, ctx)
    # One-tuples rather than one-element vectors, and a NamedTuple rather than the
    # solver's error `Dict`: `[end]` and key lookup work the same, but a `Dict` alone
    # costs several hundred bytes per trajectory, which dominates everything else
    # retained here once the study runs to millions of trajectories.
    u_analytic = sol.u_analytic === nothing ? sol.u[end] : sol.u_analytic[end]
    reduced = ConvergenceTrajectory(
        (sol.t[end],), (sol.u[end],), (u_analytic,), NamedTuple(sol.errors)
    )

    return (reduced, false)
end

"""
    _drop_trajectories(sim::EnsembleTestSolution) -> EnsembleTestSolution

Discard every trajectory but the first, keeping the error statistics that were already
computed from the full ensemble.

Reducing at `output_func` bounds the peak, but the reduced trajectories still add up
across step sizes, and once `calculate_ensemble_errors` has run they are dead weight —
a [`ConvergenceSimulation`](@ref) reads only the error dictionaries. One is kept rather
than none so that anything reaching for a representative trajectory still finds one.
"""
function _drop_trajectories(sim::SciMLBase.EnsembleTestSolution{T, N}) where {T, N}
    u = sim.u[1:min(1, length(sim.u))]
    return SciMLBase.EnsembleTestSolution{T, N, typeof(u)}(
        u, sim.errors, sim.weak_errors, sim.error_means, sim.error_medians,
        sim.elapsedTime, sim.converged
    )
end

"""
    test_convergence(dts, prob, alg[, ensemblealg]; kwargs...)
    test_convergence(probs, convergence_axis, alg; kwargs...)
    test_convergence(setup::ConvergenceSetup, alg; kwargs...)

Solve a problem at each point on a convergence axis and return a
[`ConvergenceSimulation`](@ref). For time-step convergence tests, each value in `dts`
is passed to the solver as `dt`. Remaining keyword arguments are forwarded to `solve`.

The computed solutions must contain error measurements, usually obtained from an
analytic solution. Use [`analyticless_test_convergence`](@ref) when only a numerical
reference solution is available.

## Keyword Arguments

  - `retain_solutions`: whether `ConvergenceSimulation.solutions` keeps the trajectory
    solutions (default `false`). A Monte-Carlo study over many trajectories otherwise
    retains one complete solution object — solver cache, noise process, problem and
    interpolation — per trajectory per step size, which for large `trajectories` is
    orders of magnitude more memory than the error estimates it is computing, and is
    what put the weak-convergence test groups past a 16 GB runner.

    By default each trajectory is reduced to a [`ConvergenceTrajectory`](@ref) by the
    ensemble's `output_func` as it is solved, so the full solutions never coexist, and
    the reduced ensemble is stripped to one representative trajectory once the errors
    have been computed from it. `errors`, `weak_errors`, `error_means` and `𝒪est` are
    computed from the full ensemble either way and are unaffected; what changes is that
    `solutions[i].u` holds one entry rather than `trajectories` of them.

    Pass `true` when the trajectories themselves are the point, for example to compare
    two algorithms path by path.

    The reduction is skipped, and the solutions retained, where it cannot apply: when
    you supply your own `EnsembleProblem` (set its `output_func` instead), when
    `expected_value` is given (the weak error is formed from the trajectory values
    directly), and when `weak_timeseries_errors` or `weak_dense_errors` is set (both
    need the whole timeseries).
"""
function test_convergence(
        dts::AbstractArray,
        prob::Union{
            AbstractRODEProblem, AbstractSDEProblem,
            AbstractEnsembleProblem,
        },
        alg, ensemblealg = EnsembleThreads();
        trajectories, save_start = true, save_everystep = true,
        timeseries_errors = save_everystep, adaptive = false,
        weak_timeseries_errors = false, weak_dense_errors = false,
        expected_value = nothing, retain_solutions = false, kwargs...
    )
    N = length(dts)

    reduce_trajectories = !retain_solutions &&
        !(prob isa AbstractEnsembleProblem) &&
        expected_value === nothing &&
        !weak_timeseries_errors && !weak_dense_errors

    if prob isa AbstractEnsembleProblem
        ensemble_prob = prob
    elseif reduce_trajectories
        ensemble_prob = EnsembleProblem(prob; output_func = _convergence_output_func)
    else
        ensemble_prob = EnsembleProblem(prob)
    end

    _solutions = Array{Any}(undef, length(dts))
    for i in 1:length(dts)
        sol = solve(
            ensemble_prob, alg, ensemblealg; dt = dts[i], adaptive = adaptive,
            save_start = save_start, save_everystep = save_everystep,
            timeseries_errors = timeseries_errors,
            weak_timeseries_errors = weak_timeseries_errors,
            weak_dense_errors = weak_dense_errors, trajectories = Int(trajectories),
            kwargs...
        )
        @info "dt: $(dts[i]) ($i/$N)"
        # Summarise each ensemble before solving the next, so that only one step size's
        # trajectories are ever reachable rather than the whole study's.
        _solutions[i] = if expected_value === nothing
            summarised = SciMLBase.calculate_ensemble_errors(
                sol;
                weak_timeseries_errors = weak_timeseries_errors,
                weak_dense_errors = weak_dense_errors
            )
            reduce_trajectories ? _drop_trajectories(summarised) : summarised
        else
            sol
        end
    end

    auxdata = Dict("dts" => dts)

    if expected_value == nothing
        solutions = _solutions
        # Now Calculate Weak Errors
        additional_errors = Dict()
        for k in keys(solutions[1].weak_errors)
            additional_errors[k] = [sol.weak_errors[k] for sol in solutions]
        end

    else
        additional_errors = Dict()
        if length(expected_value) == 1 || expected_value isa Number
            additional_errors[:weak_final] = []
        else
            additional_errors[:weak_l2] = []
        end
        for sol in _solutions
            if length(expected_value) == 1 || expected_value isa Number
                weak_final = LinearAlgebra.norm(Statistics.mean(sol.u .- expected_value))
                push!(additional_errors[:weak_final], weak_final)
            else
                weak_l2 = LinearAlgebra.norm(Statistics.mean(sol .- expected_value))
                push!(additional_errors[:weak_l2], weak_l2)
            end
        end
        solutions = _solutions
    end

    return ConvergenceSimulation(
        solutions, dts, auxdata = auxdata,
        additional_errors = additional_errors,
        expected_value = expected_value
    )
end

"""
    analyticless_test_convergence(dts, prob, alg, reference; kwargs...)

Estimate convergence against a numerical reference solution and return a
[`ConvergenceSimulation`](@ref). For ODE problems, `reference` is a solver setup whose
`:alg` entry selects the reference algorithm. For stochastic problems, `reference` is
the finer reference step size, and each tested solution reuses the reference noise
realization.
"""
function analyticless_test_convergence(
        dts::AbstractArray,
        prob::Union{
            AbstractRODEProblem, AbstractSDEProblem,
            AbstractSDDEProblem,
        },
        alg, test_dt; trajectories = 100,
        save_everystep = true,
        timeseries_errors = save_everystep, adaptive = false,
        weak_timeseries_errors = false,
        weak_dense_errors = false, use_noise_grid = true,
        verbose = true,
        kwargs...
    )
    _solutions = []
    tmp_solutions = Array{Any}(undef, trajectories, length(dts))
    for j in 1:trajectories
        if verbose
            @info "Monte Carlo iteration: $j/$trajectories"
        end
        t = prob.tspan[1]:test_dt:prob.tspan[2]
        if use_noise_grid
            if prob.noise_rate_prototype === nothing
                brownian_values = cumsum(
                    [
                        [zeros(size(prob.u0))];
                        [
                            sqrt(test_dt) * randn(size(prob.u0))
                                for i in 1:(length(t) - 1)
                        ]
                    ]
                )
                brownian_values2 = cumsum(
                    [
                        [zeros(size(prob.u0))];
                        [
                            sqrt(test_dt) * randn(size(prob.u0))
                                for i in 1:(length(t) - 1)
                        ]
                    ]
                )
            else
                brownian_values = cumsum(
                    [
                        [zeros(size(prob.noise_rate_prototype, 2))];
                        [
                            sqrt(test_dt) *
                                randn(size(prob.noise_rate_prototype, 2))
                                for i in 1:(length(t) - 1)
                        ]
                    ]
                )
                brownian_values2 = cumsum(
                    [
                        [zeros(size(prob.noise_rate_prototype, 2))];
                        [
                            sqrt(test_dt) *
                                randn(size(prob.noise_rate_prototype, 2))
                                for i in 1:(length(t) - 1)
                        ]
                    ]
                )
            end
            np = NoiseGrid(t, brownian_values, brownian_values2)

            if prob isa AbstractSDDEProblem
                _prob = SDDEProblem(
                    prob.f, prob.g, prob.u0, prob.h, prob.tspan, prob.p,
                    noise = np,
                    noise_rate_prototype = prob.noise_rate_prototype,
                    constant_lags = prob.constant_lags,
                    dependent_lags = prob.dependent_lags,
                    neutral = prob.neutral,
                    order_discontinuity_t0 = prob.order_discontinuity_t0,
                    prob.kwargs...
                )
            else
                _prob = SDEProblem(
                    prob.f, prob.g, prob.u0, prob.tspan, prob.p,
                    noise = np,
                    noise_rate_prototype = prob.noise_rate_prototype
                )
            end

            true_sol = solve(_prob, alg; adaptive = adaptive, dt = test_dt)

            for i in 1:length(dts)
                sol = solve(_prob, alg; dt = dts[i], adaptive = adaptive)
                err_sol = appxtrue(sol, true_sol)
                tmp_solutions[j, i] = err_sol
            end
        else
            # using NoiseWrapper doesn't lead to constant true_sol
            true_sol = solve(
                prob, alg; adaptive = adaptive, dt = test_dt,
                save_noise = true
            )
            _sol = deepcopy(true_sol)
            W1 = NoiseWrapper(_sol.W)

            _prob = remake(
                prob, u0 = prob.u0, p = prob.p, tspan = prob.tspan, noise = W1,
                noise_rate_prototype = prob.noise_rate_prototype
            )

            for i in 1:length(dts)
                W1 = NoiseWrapper(_sol.W)
                _prob = remake(
                    prob, u0 = prob.u0, p = prob.p, tspan = prob.tspan,
                    noise = W1, noise_rate_prototype = prob.noise_rate_prototype
                )
                sol = solve(_prob, alg; dt = dts[i], adaptive = adaptive)
                err_sol = appxtrue(sol, true_sol)
                tmp_solutions[j, i] = err_sol
            end
        end
    end
    _solutions = [EnsembleSolution(tmp_solutions[:, i], 0.0, true) for i in 1:length(dts)]
    solutions = [
        SciMLBase.calculate_ensemble_errors(
                sim;
                weak_timeseries_errors = weak_timeseries_errors,
                weak_dense_errors = weak_dense_errors
            )
            for sim in _solutions
    ]
    auxdata = Dict("dts" => dts)
    # Now Calculate Weak Errors
    additional_errors = Dict()
    for k in keys(solutions[1].weak_errors)
        additional_errors[k] = [sol.weak_errors[k] for sol in solutions]
    end
    return ConvergenceSimulation(
        solutions, dts, auxdata = auxdata,
        additional_errors = additional_errors
    )
end

function test_convergence(
        dts::AbstractArray,
        prob::Union{AbstractODEProblem, AbstractDAEProblem}, alg;
        save_everystep = true, adaptive = false, kwargs...
    )
    N = length(dts)
    solutions = [
        solve(
                prob, alg; dt = dts[i], save_everystep = save_everystep,
                adaptive = adaptive, kwargs...
            ) for i in 1:N
    ]
    auxdata = Dict(:dts => dts)
    return ConvergenceSimulation(solutions, dts, auxdata = auxdata)
end

function analyticless_test_convergence(
        dts::AbstractArray, prob::AbstractODEProblem,
        alg, appxsol_setup;
        save_everystep = true, adaptive = false, kwargs...
    )
    true_sol = solve(prob, appxsol_setup[:alg]; appxsol_setup...)
    N = length(dts)
    _solutions = [
        solve(
                prob, alg; dt = dts[i], save_everystep = save_everystep,
                adaptive = adaptive, kwargs...
            ) for i in 1:N
    ]
    solutions = [appxtrue(sol, true_sol) for sol in _solutions]
    auxdata = Dict(:dts => dts)
    return ConvergenceSimulation(solutions, dts, auxdata = auxdata)
end

function test_convergence(probs, convergence_axis, alg; kwargs...)
    return ConvergenceSimulation([solve(prob, alg; kwargs...) for prob in probs], convergence_axis)
end

function test_convergence(c::ConvergenceSetup, alg::AbstractDEAlgorithm; kwargs...)
    return test_convergence(c.probs, c.convergence_axis, alg; kwargs...)
end

function calc𝒪estimates(error::Pair)
    key = error.first
    error = error.second

    if ndims(error) > 1
        error = mean(error, 1)
    end
    S = Vector{eltype(error)}(undef, length(error) - 1)
    for i in 1:(length(error) - 1)
        S[i] = log2(error[i + 1] / error[i])
    end
    return (Pair(key, abs.(mean(S, dims = 1))))
end

"""
length(simres::ConvergenceSimulation)

Returns the number of simultations in the Convergence Simulation
"""
Base.length(sim::ConvergenceSimulation) = sim.N
Base.getindex(sim::ConvergenceSimulation, i::Int) = sim.solutions[i]
Base.getindex(sim::ConvergenceSimulation, i::Int, I::Int...) = sim.solutions[i][I]
Base.lastindex(sim::ConvergenceSimulation) = lastindex(sim.solutions)
Base.firstindex(sim::ConvergenceSimulation) = firstindex(sim.solutions)
Base.length(sim::ConvergenceSetup) = sim.probs
