module MDL

using LinearAlgebra
import Statistics: mean
import StatsBase: fit, Histogram
import ..LMCLUS: Manifold, indim, outdim, projection, threshold, points, hasfullbasis, project, LMCLUSResult

# Various types for MDL calculation
abstract type MethodType end

struct Raw <: MethodType end
struct ZeroDim <: MethodType end

struct Uniform <: MethodType end
struct Gaussian <: MethodType end
struct Empirical <: MethodType end
struct OptimalQuant <: MethodType end    # ICPR-2016 Eq. 6
struct SizeIndependent <: MethodType end # ICPR-2016 Eq. 8

"""Optimal number of bins given constant C over intervals"""
function optbins(intervals::Vector{T}, C::T) where T<:AbstractFloat
    ns = ceil.(intervals * exp( (C - sum(intervals))/length(intervals) ), digits=0)
    nsi = zeros(UInt, size(ns))
    for i in 1:length(nsi)
        nsi[i] = if ns[i] >= typemax(UInt)
            typemax(UInt)
        else
            round(UInt, ns[i])
        end
    end
    return nsi
end

"""Variance of a uniform distribution over interval"""
univar(interval::Vector{T}) where T<:AbstractFloat = (1//12)*(interval).^2

"""Quantizaton error"""
quanterror(intervals, N_k_opt) = sum(univar(intervals./N_k_opt))

"""Optimal quantization of the interval"""
function optquant(intervals::Vector{T}, ɛ::Real; tot::Int=1000, tol::Real=1e-6) where {T<:AbstractFloat}
    intervals[isnan.(intervals)]  .= eps(T) # remove nans
    intervals[intervals .< eps()] .= eps(T) # remove 0s
    intervals[isinf.(intervals)]  .= 1.0   # remove inf

    # Setup C bounds
    K = length(intervals)
    C = zero(T)
    Cmin = zero(T)
    Cmax = broadcast(+, log.(typemax(UInt)./intervals).*K, sum(log.(intervals))) |> minimum

    i = 1
    N_k_opt = ones(UInt, K)
    err_opt = Inf
    while i < tot
        C = Cmax - (Cmax - Cmin)/2
        N_k_opt = optbins(intervals, C)
        err_opt = quanterror(intervals, N_k_opt)
        err_diff = err_opt - ɛ^2
        if abs(err_diff) < tol
            break # error is small enough to stop search
        elseif err_diff > 0.
            # error is too big
            Cmin = C
        elseif err_diff < 0.
            # error is too small
            Cmax = C
        end
        i += 1
    end
    return N_k_opt, err_opt, C, i
end

"""Calculate LM cluster bounding box using PCA

    Ref: Jonathon Shlens, A Tutorial on Principal Component Analysis, 2005
"""
function boundingbox(X, m)
    n, l = size(X)

    # PCA data
    F = svdfact(X./sqrt(n))

    # Take an orthogonal compliment subspace (OCS) basis
    r = if m < min(n,l)
        (m+1):min(n,l)
    else
        1:min(n,l)
    end
    BC = F[:U][:,r]

    # Project data to OCS
    Y = BC'*X

    # Calculate data spread in OCS
    Ymin = vec(minimum(Y, dims=2))
    Ymax = vec(maximum(Y, dims=2))
    range = abs.(Ymax - Ymin)

    return range, Y, Ymin, Ymax
end

function boundingbox(Y::AbstractMatrix{T}) where T<:AbstractFloat
    Ymin = vec(minimum(Y, dims=2))
    Ymax = vec(maximum(Y, dims=2))
    Yrange = abs.(Ymax - Ymin)
    return Yrange, Ymin, Ymax
end

#=
    Minimum Description Length Encoding
=#

"""Raw representation of the data"""
modeldl(::Type{Raw}, M::Manifold, X::AbstractMatrix, P::Int) = P*length(X)
datadl(::Type{Raw}, aggs...) = 0

function raw(M::Manifold, Pm::Int)
    return Pm*size(M)*indim(M)
end

"""Zero dimensional model description length: L(ZD)

    From: O. Georgieva, K. Tschumitschew, and F. Klawonn,
    “Cluster validity measures based on the minimum description length principle”
"""
function modeldl(::Type{ZeroDim}, M::Manifold, X::AbstractMatrix, P::Int)
    return P*size(X,1)
end

"""Zero dimensional data description length: L(X|ZD)

    Ref: O. Georgieva, K. Tschumitschew, and F. Klawonn, “Cluster validity
    measures based on the minimum description length principle”
"""
function datadl(::Type{ZeroDim}, C::Manifold{T}, X::AbstractMatrix{T},
                P::Int, ɛ::Real, tot::Int, tol::Real) where T<:AbstractFloat
    N = size(X,1)  # space dimension
    μ = mean(C)    # manifold translation
    n = size(C)  # size of cluster

    E = 0.0
    for i in 1:n
        for j in 1:N
            E += log2(abs(X[j,i] - μ[j]) + 1)
        end
    end
    return round(Int, E)
end

#=
    Various models of MDL calculation (model MDL)
=#

#=
""" General model description length (v1)

    Only encodes linear manifold basis vectors (no orthogonal complement space)
"""
function modeldl{MT<:MethodType}(::Type{MT}, C::Manifold, X::Matrix, P::Int)
    N = size(X,1)  # space dimension
    M = outdim(C)   # manifold dimension
    return P*(N + M*(N - (M+1)>>1))
end
=#

""" General model description length (v2)

    Encodes linear manifold & orthogonal complement space basis vectors
"""
function modeldl(::Type{MT}, C::Manifold, X::AbstractMatrix, P::Int) where MT <: MethodType
    N = size(X,1)  # space dimension
    return if outdim(C) != 0
        (P*N*(N+1))>>1
    else # for 0D manifold no basis encoding required
        P*N
    end
end

"""Size independent data description length: L(X|SI)
"""
function datadl(::Type{SizeIndependent}, C::Manifold{T}, X::AbstractMatrix{T},
                P::Int, ɛ::Real, tot::Int, tol::Real) where T<:AbstractFloat
    M = outdim(C)  # manifold dimension
    N = size(X,1)  # space dimension

    # Project points to the orthogonal compliment space (OCS) of the manifold
    Y = project(C, X)

    # Get OCS bounding box
    rngs, _ = boundingbox(Y)

    # Estimate number of bins required for each component of OSC
    bins, _ = optquant(rngs, ɛ, tot=tot, tol=tol)

    return convert(Int, P*(N + 2*sum(bins)))
end

#=
    Various models of MDL calculation (data MDL)
=#

""" Data MDL: Uniform encoding

    Points in the orthogonal complement space of the linear manifold cluster
    are uniformly distributed around the cluster manifold.
"""
function datadl(::Type{Uniform}, C::Manifold{T}, X::AbstractMatrix{T},
                P::Int, ɛ::Real, tot::Int, tol::Real) where T<:AbstractFloat
    M = outdim(C)  # manifold dimension
    n = size(X,2)  # size of cluster

    # Encoding size of point projected on manifold
    PR = P*M

    # Entropy of points in the orthogonal complement space
    E = -float(n)*log(threshold(C)[1])

    # Number of bits of two parts for every point
    return round(Int, (PR + E)*n)
end

""" Data MDL: Gaussian encoding

    Points in the orthogonal complement space of the linear manifold cluster
    are normally distributed around the cluster manifold.
"""
function datadl(::Type{Gaussian}, C::Manifold{T}, X::AbstractMatrix{T},
                P::Int, ɛ::Real, tot::Int, tol::Real) where T<:AbstractFloat
    N, n = size(X)  # space dimension & size
    M = outdim(C)   # manifold dimension
    μ = mean(C)     # manifold translation
    B = projection(C) # manifold basis

    # Encoding size of point projected on manifold
    PR = P*M

    # Change points basis to of orthogonal complement subspace
    OP = (Matrix{T}(I, N, N) - B*B')*(X.-μ)

    # multivariate normal distribution entropy
    Σ = OP*OP'
    Σd = det(Σ)
    E = (M/2)*(1+log(2π)) + (Σd > eps() ? log(Σd) : 0.0)/2

    # Number of bits of two parts for every point
    return round(Int, (PR + E)*n)
end

""" Data MDL: Empirical entropy encoding

    Distribution of points in the orthogonal complement space of the linear manifold cluster,
    constructed by optimal quantizing with a fixed length bin, is used to
    calculate entropy, thus number of bits.
"""
function datadl(::Type{Empirical}, C::Manifold{T}, X::AbstractMatrix{T},
                P::Int, ɛ::Real, tot::Int, tol::Real) where T<:AbstractFloat
    N, n = size(X) # space dimension & size
    M = outdim(C)  # manifold dimension

    # Encoding size of point projected on manifold
    PR = P*M

    # Project points to the orthogonal compliment space (OCS) of the manifold
    Y = project(C, X)

    # Get OCS bounding box
    rngs, Ymin, Ymax = boundingbox(Y)

    # Estimate number of bins required for each competent of OSC
    if ɛ < 1. # using optimal quantization
        bins, _ = optquant(rngs, ɛ, tot=tot, tol=tol)
    else      # constant number of bins per component
        bins = fill(round(Int,ɛ), N-M)
    end

    # Cumulative entropy of data spread in OSC (calculated from quantization)
    E = 0.0
    for i in 1:N-M
        Yb = range(Ymin[i], stop=Ymax[i], length=bins[i]+1)
        H = fit(Histogram, vec(Y[i,:]), Yb, closed=:left)
        E += -sum(map(x-> x > 0. ? (x/(n-1))*log2(x/(n-1)) : 0., H.weights))
    end

    # Number of bits of two parts for every point
    return round(Int, (PR + E)*n)
end

""" Data MDL: Optimal quantization encoding

    Distribution of points in the orthogonal complement space of the linear manifold cluster,
    constructed by optimal quantizing with a fixed length bin, is used to
    calculate entropy, thus number of bits.
"""
function datadl(::Type{OptimalQuant}, C::Manifold{T}, X::AbstractMatrix{T},
                P::Int, ɛ::Real, tot::Int, tol::Real) where T<:AbstractFloat
    M = outdim(C)  # manifold dimension
    n = size(X,2)  # size of cluster

    # Encoding size of a point projected on manifold
    PR = P*M

    # Project points to the orthogonal compliment space (OCS) of the manifold
    Y = project(C, X)

    # Get OCS bounding box
    rngs, _ = boundingbox(Y)

    # Estimate number of bins required for each component of OSC
    bins, sqerr, C = optquant(rngs, ɛ, tot=tot, tol=tol)

    # Number of bins required to encode data spread in OSC given optimal quantization
    # C = Σ_k log(N_k)
    # where N_k number of bins in quantized interval of component K of OSC
    E = C/log(2)

    # Number of bits of two parts for every point
    return round(Int, (PR + E)*n)
end

#=
    Main MDL call
=#

"Calculate MDL for the manifold"
function calculate(::Type{MT}, M::Manifold{T}, X::AbstractMatrix{T}, Pm::Int, Pd::Int;
                   ɛ::Real = 1e-3, tot::Int=1000, tol::Real=1e-18) where {MT<:MethodType, T<:AbstractFloat}
    return modeldl(MT, M, X, Pm) + datadl(MT, M, X, Pd, ɛ, tot, tol)
end

"Calculate MDL for the clustering"
function calculate(::Type{MT}, Ms::Vector{Manifold}, X::AbstractMatrix{T}, Pm::Int, Pd::Int;
                   ɛ::Real=1e-3, tot::Int=1000, tol::Real=1e-18) where {MT<:MethodType, T<:AbstractFloat}
    return sum(calculate(MT, m, X[:,points(m)], Pm, Pd, ɛ=ɛ, tot=tot, tol=tol) for m in Ms)
end

function calculate!(::Type{MT}, C::LMCLUSResult, X::AbstractMatrix{T}, Pm::Int, Pd::Int;
                   ɛ::Real=1e-3, d=-1) where {MT<:MethodType, T<:AbstractFloat}
    mdl = 0
    for (i, m) in enumerate(C.manifolds)
        pts = points(m)
        MX = view(X, :, pts)
        μ = mean(MX, dims=2)
        Svd = svd(MX .- μ, full=true)
        U = Svd.U
        mtmp = Manifold((d >= 0 ? d : m.d), vec(μ), U, pts)
        mdl += calculate(MT, mtmp, MX, Pm, Pd, ɛ=ɛ)
    end
    mdl
end

function compression(::Type{MT}, C::LMCLUSResult, X::AbstractMatrix{T}, Pm::Int,
                     Pd::Int; kwargs...) where {MT<:MethodType, T<:AbstractFloat}
    raw = Pm * length(X)
    mdl = calculate!(MT, C, X, Pm, Pd; kwargs...)
    raw/mdl
end

end


#=
backtrack.jl

    Backtracking line search algorithms based on the Armijo or sufficient
    decrease condition to determine the optimal step length, using either no,
    quadratic, or cubic safeguarded interpolation.

@author: Quint Wiersma <q.wiersma@vu.nl>

@date: 2022/02/02
=#

# Backtracking struct
Base.@kwdef mutable struct BackTrack{Tf, Gf, Ti}
    α::Tf	            # stepsize
	α_init::Tf 	        # initial stepsize
    ρ::Gf = .5		    # contraction factor
    c_1::Tf = 1e-4      # sufficient decrease constant
    order::Ti = 0       # interpolation order
    ρ_hi::Gf = .5       # safeguard upper bound
    ρ_lo::Gf = .1       # safeguard lower bound
    max_iter::Ti = 1000 # max iterations
end

"""
    fixed!(ls, x, x_prev, p, f_prev, ∇f_prev, f)

Backtracking line search wtih a fixed contraction parameter ``ρ`` to find ``α``,
storing the result in `x`.

#### Arguments
  - `ls::BackTrack`             : line search parameters
  - `x_prev::AbstractVector`    : previous state
  - `p::AbstractVector`         : search direction
  - `f_prev::Real`              : previous objective function value
  - `∇f_prev::AbstractVector`   : previous gradient
  - `f::Function`               : ``f(x)``

#### Returns
  - `x::AbstractVector` : updated state
  - `f_new::Real`       : objective function value
"""
function fixed!(ls::BackTrack, x::AbstractVector, x_prev::AbstractVector,
                p::AbstractVector, f_prev::Real, ∇f_prev::AbstractVector,
                f::Function)
    # Unpack
    α= ls.α_init
    c_1= ls.c_1
    ρ= ls.ρ
    max_iter= ls.max_iter

    # Objective function
    x.= x_prev .+ α .* p
    f_new= f(x)

    # Directional derivative
    d= dot(∇f_prev, p)

    # Iteration counter
    iter= 1
    # Backtracking
    while f_new > f_prev + c_1 * α * d && iter < max_iter
        # Contract
        α*= ρ

        # Update objective function
        x.= x_prev .+ α .* p
        f_new= f(x)

        # Update iteration counter
        iter+= 1
    end

    # Store stepsize
    ls.α= α

    return f_new
end

"""
    quad_interp!(ls, x, x_prev, p, f_prev, ∇f_prev, f)

Backtracking line search wtih  quadratic interpolation to find ``α``, storing
the result in `x`.

#### Arguments
  - `ls::BackTrack`             : line search parameters
  - `x_prev::AbstractVector`    : previous state
  - `p::AbstractVector`         : search direction
  - `f_prev::Real`              : previous objective function value
  - `∇f_prev::AbstractVector`   : previous gradient
  - `f::Function`               : ``f(x)``

#### Returns
  - `x::AbstractVector` : updated state
  - `f_new::Real`       : objective function value
"""
function quad_interp!(ls::BackTrack, x::AbstractVector, x_prev::AbstractVector,
                        p::AbstractVector, f_prev::Real, ∇f_prev::AbstractVector,
                        f::Function)
    # Unpack
    α_1= ls.α_init
    c_1= ls.c_1
    ρ_hi, ρ_lo= ls.ρ_hi, ls.ρ_lo
    max_iter= ls.max_iter

    # Objective function
    x.= x_prev .+ α_1 .* p
    f_new= f(x)

    # Directional derivative
    d= dot(∇f_prev, p)

    # Iteration counter
    iter= 1
    # Backtracking
    while f_new > f_prev + c_1 * α_1 * d && iter < max_iter
        # Quadratic interpolation using f(x), ∇f(x)'p, and f(x + α_1 * p)
        α_tmp= -d * α_1^2 * inv(2 * (f_new - f_prev - d * α_1))

        # Safeguard
        α_tmp= min(α_tmp, ρ_hi * α_1)   # avoid too small reductions
        α_1= max(α_tmp, ρ_lo * α_1) # avoid too big reductions

        # Update objective function
        x.= x_prev .+ α_1 .* p
        f_new= f(x)

        # Update iteration counter
        iter+= 1
    end

    # Store stepsize
    ls.α= α_1

    return f_new
end

"""
    cubic_interp!(ls, x, x_prev, p, f_prev, ∇f_prev, f)

Backtracking line search wtih cubic interpolation to find ``α``, storing the
result in `x`.

#### Arguments
  - `ls::BackTrack`             : line search parameters
  - `x_prev::AbstractVector`    : previous state
  - `p::AbstractVector`         : search direction
  - `f_prev::Real`              : previous objective function value
  - `∇f_prev::AbstractVector`   : previous gradient
  - `f::Function`               : ``f(x)``

#### Returns
  - `x::AbstractVector` : updated state
  - `f_new::Real`       : objective function value
"""
function cubic_interp!(ls::BackTrack, x::AbstractVector, x_prev::AbstractVector,
                        p::AbstractVector, f_prev::Real, ∇f_prev::AbstractVector,
                        f::Function)
    # Infer type
    Tα= typeof(ls.α_init)

    # Unpack
    α_2, α_1= ls.α_init, ls.α_init
    c_1= ls.c_1
    ρ_hi, ρ_lo= ls.ρ_hi, ls.ρ_lo
    max_iter= ls.max_iter

    # Objective function values
    x.= x_prev .+ α_2 .* p
    f_new= f(x)
    f_tmp= f_new

    # Directional derivative
    d= dot(∇f_prev, p)

    # Iteration counter
    iter= 1
    # Backtracking
    while f_new > f_prev + c_1 * α_2 * d && iter < max_iter
        if iter == 1
            # Quadratic interpolation using f(x), ∇f(x)'p, and f(x + α_2 * p)
            α_tmp= -d * α_2^2 * inv(2 * (f_new - f_prev - d * α_2))
        else
            # Cubic interpolation using f(x), ∇f(x)'p, f(x + α_1 * p), and f(x + α_2 * p)
            div= inv(α_1^2 * α_2^2 * (α_2 - α_1))
            tmp_2= f_new - f_prev - d * α_2
            tmp_1= f_tmp - f_prev - d * α_1
            a= (α_1^2 * tmp_2  - α_2^2 * tmp_1) * div
            b= (-α_1^3 * tmp_2  + α_2^3 * tmp_1) * div

            if isapprox(a, zero(a), atol=eps(Tα))
                # approximate quadratic interpolation
                α_tmp= -d * inv(2 * b)
            else
                # discriminant
                D= max(b^2 - 3 * a * d, Tα(0))
                # quadratic equation root
                α_tmp= (-b + sqrt(D)) * inv(3 * a)
            end
        end

        # Store stepsize
        α_1= α_2

        # Safeguard
        α_tmp= min(α_tmp, ρ_hi * α_2)   # avoid too small reductions
        α_2= max(α_tmp, ρ_lo * α_2) # avoid too big reductions

        # Update objective function values
        x.= x_prev .+ α_2 .* p
        f_tmp, f_new= f_new, f(x)

        # Update iteration counter
        iter+= 1
    end

    # Store stepsize
    ls.α= α_2

    return f_new
end

"""
    backtrack!(ls, x, x_prev, p, f_prev, ∇f_prev, f)

Backtracking line search to find optimal step size ``α``, storing the result in
`x`.

#### Arguments
  - `ls::BackTrack`             : line search parameters
  - `x_prev::AbstractVector`    : previous state
  - `p::AbstractVector`         : search direction
  - `f_prev::Real`              : previous objective function value
  - `∇f_prev::AbstractVector`   : previous gradient
  - `f::Function`               : ``f(x)``

#### Returns
  - `x::AbstractVector` : updated state
  - `f_new::Real`       : objective function value
"""
function backtrack!(ls::BackTrack, x::AbstractVector, x_prev::AbstractVector,
                    p::AbstractVector, f_prev::Real, ∇f_prev::AbstractVector,
                    f::Function)
    # Check interpolation
    if ls.order == 0
        # Fixed backtracking
        f_new= fixed!(ls, x, x_prev, p, f_prev, ∇f_prev, f)
    elseif ls.order == 2
        # Quadratic interpolation
        f_new= quad_interp!(ls, x, x_prev, p, f_prev, ∇f_prev, f)
    elseif ls.order == 3
        # Cubic interpolation
        f_new= cubic_interp!(ls, x, x_prev, p, f_prev, ∇f_prev, f)
    end

    return f_new
end


"""
    SeqDPP()

Sequential sampling via Determinantal Point Processes
"""
struct SeqDPP <: OnIPSA end

Base.show(io::IO, ::SeqDPP) = print(io, "Sequential DPP")

function initZ(
    rng::AbstractRNG,
    ::SeqDPP,
    X::AbstractVector;
    kernel::Kernel,
    arraytype=Vector{Float64},
    kwargs...,
)
    length(X) > 2 || throw(ArgumentError("First batch should contain at least 3 elements"))
    K = kernelmatrix(kernel, X) + jitt * I
    dpp = DPP(K)
    samp = []
    while length(samp) < 3 # Sample from a normal DPP until at least 3 elements are sampled
        samp = rand(rng, dpp)
    end
    Z = to_vec_of_vecs(X[samp], arraytype)
    return Z
end

function add_point!(
    rng::AbstractRNG,
    Z::AbstractVector{T},
    ::SeqDPP,
    X::AbstractVector;
    kernel::Kernel,
    kwargs...,
) where {T}
    L = kernelmatrix(kernel, vcat(Z, X)) + jitt * I
    Iₐ = diagm(vcat(zeros(length(Z)), ones(length(X))))
    Lₐ = Symmetric(
        inv(inv(L + Iₐ)[(length(Z) + 1):size(L, 1), (length(Z) + 1):size(L, 1)]) - I
    )
    new_dpp = DPP(Lₐ)
    new_samp = rand(rng, new_dpp)
    return push!(Z, X[new_samp]...)
end

function add_point(
    rng::AbstractRNG,
    Z::AbstractVector{T},
    ::SeqDPP,
    X::AbstractVector;
    kernel::Kernel,
    kwargs...,
) where {T}
    L = kernelmatrix(kernel, vcat(Z, X)) + jitt * I
    Iₐ = diagm(vcat(zeros(length(Z)), ones(length(X))))
    Lₐ = Symmetric(
        inv(inv(L + Iₐ)[(length(Z) + 1):size(L, 1), (length(Z) + 1):size(L, 1)]) - I
    )
    new_dpp = DPP(Lₐ)
    new_samp = rand(rng, new_dpp)
    return vcat(Z, T.(X[new_samp]))
end


using ForwardDiff

"""
Evaluates the approximate contouring error ϵc = sin(Φ)(X - X_ref(s)) - cos(Φ)(Y - Y_ref(s)). Argument
P should represent the positional states of the systems (X, Y, s) corresponding to the position of the vehicle
(X, Y) and the arc length of the (approximate) projection point of the vehicle onto the reference path.
"""
function approx_contouring_error(P, path::SplinePath, spline_idx::Integer)
    X = P[1]; Y = P[2]; s = P[3]

    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    x_ref = spline_x(s, path, spline_idx)
    y_ref = spline_y(s, path, spline_idx)
    Φ = heading(s, path, spline_idx)
    sin(Φ) * (X - x_ref) - cos(Φ) * (Y - y_ref)
end

"""
Evaluates the approximate lag error ϵl = -cos(Φ)(X - X_ref(s)) - sin(Φ)(Y - Y_ref(s)). Argument
P should represent the positional states of the systems (X, Y, s) corresponding to the position of the vehicle
(X, Y) and the arc length of the (approximate) projection point of the vehicle onto the reference path.
"""
function approx_lag_error(P, path::SplinePath, spline_idx::Integer)
    X = P[1]; Y = P[2]; s = P[3]

    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    x_ref = spline_x(s, path, spline_idx)
    y_ref = spline_y(s, path, spline_idx)
    Φ = heading(s, path, spline_idx)
    -cos(Φ) * (X - x_ref) - sin(Φ) * (Y - y_ref)
end

"""
Returns [dϵc/dX, dϵc/dY, dϵc/ds].
"""
function contouring_error_gradient(P0::Array{Float64,1}, path::SplinePath, spline_idx::Integer)
    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    ForwardDiff.gradient(P -> approx_contouring_error(P, path, spline_idx), P0)
end

"""
Returns [dϵl/dX, dϵl/dY, dϵl/ds].
"""
function lag_error_gradient(P0::Array{Float64,1}, path::SplinePath, spline_idx::Integer)
    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    ForwardDiff.gradient(P -> approx_lag_error(P, path, spline_idx), P0)
end

#=
NOTE regarding the following two functions:
Define the tracking error cost at some point (X, Y, s) as (qc * ϵc(X, Y, s)^2 + ql * ϵl(X, Y, s)^2). As this is
nonlinear, approximate this cost by linearizing ϵc and ϵl individually about some point P0 = (X0, Y0, s0) and
plugging into the tracking error cost expression. The approximate tracking cost may then be arranged as the
sum of a linear term, cᵀP, and a quadratic term, PᵀQP.
=#

"""
Returns the coefficient vector, c, of the linear term described above.
"""
function tracking_linear_term(P0::Array{Float64,1}, qc::Float64, ql::Float64, path::SplinePath, spline_idx::Integer)
    X0 = P0[1]; Y0 = P0[2]; s0 = P0[3]

    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    ϵc0 = approx_contouring_error(P0, path, spline_idx)
    ϵl0 = approx_lag_error(P0, path, spline_idx)
    dϵc_dx, dϵc_dy, dϵc_ds = contouring_error_gradient(P0, path, spline_idx)
    dϵl_dx, dϵl_dy, dϵl_ds = lag_error_gradient(P0, path, spline_idx)

    cc = zeros(3)
    cc[1] = 2*qc*ϵc0*dϵc_dx - 2*qc*X0*dϵc_dx^2 - 2*qc*Y0*dϵc_dx*dϵc_dy - 2*qc*s0*dϵc_dx*dϵc_ds
    cc[2] = 2*qc*ϵc0*dϵc_dy - 2*qc*Y0*dϵc_dy^2 - 2*qc*X0*dϵc_dx*dϵc_dy - 2*qc*s0*dϵc_dy*dϵc_ds
    cc[3] = 2*qc*ϵc0*dϵc_ds - 2*qc*s0*dϵc_ds^2 - 2*qc*X0*dϵc_dx*dϵc_ds - 2*qc*Y0*dϵc_dy*dϵc_ds

    cl = zeros(3)
    cl[1] = 2*ql*ϵl0*dϵl_dx - 2*ql*X0*dϵl_dx^2 - 2*ql*Y0*dϵl_dx*dϵl_dy - 2*ql*s0*dϵl_dx*dϵl_ds
    cl[2] = 2*ql*ϵl0*dϵl_dy - 2*ql*Y0*dϵl_dy^2 - 2*ql*X0*dϵl_dx*dϵl_dy - 2*ql*s0*dϵl_dy*dϵl_ds
    cl[3] = 2*ql*ϵl0*dϵl_ds - 2*ql*s0*dϵl_ds^2 - 2*ql*X0*dϵl_dx*dϵl_ds - 2*ql*Y0*dϵl_dy*dϵl_ds

    c = cc + cl
end

"""
Returns the coefficient matrix, Q, of the quadratic term described above.
"""
function tracking_quadratic_term(P0::Array{Float64,1}, qc::Float64, ql::Float64, path::SplinePath, spline_idx::Integer)
    X0 = P0[1]; Y0 = P0[2]; s0 = P0[3]
    ϵc0 = approx_contouring_error(P0, path, spline_idx)
    ϵl0 = approx_lag_error(P0, path, spline_idx)
    dϵc_dx, dϵc_dy, dϵc_ds = contouring_error_gradient(P0, path, spline_idx)
    dϵl_dx, dϵl_dy, dϵl_ds = lag_error_gradient(P0, path, spline_idx)

    Qc = zeros(3, 3)
    a = qc*dϵc_dx^2
    b = d = (2*qc*dϵc_dx*dϵc_dy)/2
    c = g = (2*qc*dϵc_dx*dϵc_ds)/2
    e = qc*dϵc_dy^2
    f = h = (2*qc*dϵc_dy*dϵc_ds)/2
    i = qc*dϵc_ds^2
    Qc = [a b c; d e f; g h i]

    Ql = zeros(3, 3)
    a = ql*dϵl_dx^2
    b = d = (2*ql*dϵl_dx*dϵl_dy)/2
    c = g = (2*ql*dϵl_dx*dϵl_ds)/2
    e = ql*dϵl_dy^2
    f = h = (2*ql*dϵl_dy*dϵl_ds)/2
    i = ql*dϵl_ds^2
    Ql = [a b c; d e f; g h i]

    Q = Qc + Ql
end

function tracking_cost_matrices!(qs, vals)
    S_q = vals.k_c + vals.n_modes * (vals.N - vals.k_c)

    for k in 1:S_q
        P0 = [qs[1, k], qs[2, k], qs[5, k]] #X, Y, s of the initial guess
        spline_idx = find_spline_interval(P0[3], vals.path)
        if k in end_horizon_idces(vals)
            ck = tracking_linear_term(P0, vals.qcterm, vals.ql, vals.path, spline_idx)
            Γk = tracking_quadratic_term(P0, vals.qcterm, vals.ql, vals.path, spline_idx) # TODO: fix function signature
        else
            ck = tracking_linear_term(P0, vals.qc, vals.ql, vals.path, spline_idx)
            Γk = tracking_quadratic_term(P0, vals.qc, vals.ql, vals.path, spline_idx) # TODO: fix function signature
        end
        vals.cs[k] = ck
        vals.Γs[k] = Γk
    end

    vals.c = vcat(vals.cs...)
    vals.Γ = Array(blockdiag(sparse.(vals.Γs)...))
end


#! /usr/bin/env julia
# Parsing
str = read("./3.in", String)
input = split(str, "\n")[1:end-1]
squares = [
    begin
        # Thanks to u/mserrano!
        id, x, y, w, h = Tuple(map(x -> parse(Int64, x.match), eachmatch(r"\d+", i)));
        (
            id = id,
            # Plus one 'cuz of Julia's indexing starting at one not zero
            idx = (x + 1 : x + w, y + 1 : y + h)
        )
    end
    for i in input
]

# Calculating the fabric with claims per square inch
fabric = zeros(Int64, 1000, 1000)
for s in squares
    v = view(fabric, s.idx...)
    v.+= 1
end

println("Solution for part 1:")
println(length(fabric[fabric .> 1]))

println("Solution for part 2:")
for s in squares
    if sum(fabric[s.idx[1], s.idx[2]] .- 1) == 0
        println(s.id)
    end
end

# Solution 1: 107663
# Solution 2: 1166


using Printf
using CUDA
using Oceananigans
using Oceananigans.TurbulenceClosures: Explicit, VerticallyImplicit
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary
using LinearAlgebra

function boundary_clustered(N, L, ini)
    Δz(k)   = k < N / 2 + 1 ? 2 / (N - 1) * (k - 1) + 1 : - 2 / (N - 1) * (k - N) + 1
    z_faces = zeros(N+1)
    for k = 2:N+1
        z_faces[k] = z_faces[k-1] + Δz(k-1)
    end
    z_faces = z_faces ./ z_faces[end] .* L .+ ini
    return z_faces
end

function center_clustered(N, L, ini)
    Δz(k)   = k < N / 2 + 1 ? 2 / (N - 1) * (k - 1) + 1 : - 2 / (N - 1) * (k - N) + 1
    z_faces = zeros(N+1)
    for k = 2:N+1
        z_faces[k] = z_faces[k-1] + 3 - Δz(k-1)
    end
    z_faces = z_faces ./ z_faces[end] .* L .+ ini
    return z_faces
end

grid = RectilinearGrid(GPU(), size=(512, 256),
                       x = (-10, 10),
                       z = (0, 5),
                topology = (Periodic, Flat, Bounded))

# Gaussian bump of width "1"
bump(x, y, z) = z < exp(-x^2)

@inline show_name(t) = t isa Explicit ? "explicit" : "implicit"

grid_with_bump = ImmersedBoundaryGrid(grid, GridFittedBoundary(bump))

# Tidal forcing
tidal_forcing(x, y, z, t) = 1e-4 * cos(t)

for time_stepper in (ExplicitTimeDiscretization, VerticallyImplicitTimeDiscretization)

    model = HydrostaticFreeSurfaceModel(grid = grid_with_bump,
                                        momentum_advection = CenteredSecondOrder(),
                                        free_surface = ExplicitFreeSurface(gravitational_acceleration=10),
                                        closure = ScalarDiffusivity(time_stepper, ν=1e-2, κ=1e-2),
                                        tracers = :b,
                                        buoyancy = BuoyancyTracer(),
                                        coriolis = FPlane(f=sqrt(0.5)),
                                        forcing = (u = tidal_forcing,))

    # Linear stratification
    set!(model, b = (x, y, z) -> 4 * z)

    progress_message(s) = @info @sprintf("[%.2f%%], iteration: %d, time: %.3f, max|w|: %.2e",
                                100 * s.model.clock.time / s.stop_time, s.model.clock.iteration,
                                s.model.clock.time, maximum(abs, model.velocities.w))

    gravity_wave_speed = sqrt(model.free_surface.gravitational_acceleration * grid.Lz)

    Δt = CUDA.@allowscalar 0.1 * minimum(grid.Δxᶜᵃᵃ) / gravity_wave_speed

    simulation = Simulation(model, Δt = Δt, stop_time = 50000Δt)

    serialize_grid(file, model) = file["serialized/grid"] = model.grid.grid

    simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.velocities, model.tracers),
                                                        schedule = TimeInterval(0.1),
                                                        prefix = "internal_tide_$(show_name(time_stepper))",
                                                        init = serialize_grid,
                                                        force = true)

    simulation.callbacks[:progress] = Callback(progress_message, IterationInterval(10))

    run!(simulation)

    @info """
        Simulation complete.
        Output: $(abspath(simulation.output_writers[:fields].filepath))
    """
end

using JLD2
using Plots
ENV["GKSwstype"] = "100"

function nice_divergent_levels(c, clim; nlevels=20)
    levels = range(-clim, stop=clim, length=nlevels)
    cmax = maximum(abs, c)
    clim < cmax && (levels = vcat([-cmax], levels, [cmax]))
    return (-clim, clim), levels
end

function nan_solid(x, z, u, bump)
    Nx, Nz = size(u)
    x2 = reshape(x, Nx, 1)
    z2 = reshape(z, 1, Nz)
    u[bump.(x2, 0, z2)] .= NaN
    return nothing
end

function visualize_internal_tide_simulation(prefix)

    filename = prefix * ".jld2"
    file = jldopen(filename)

    grid = adapt(CPU(), file["serialized/grid"])

    bump(x, y, z) = z < exp(-x^2)

    xu, yu, zu = nodes((Face, Center, Center), grid)
    xw, yw, zw = nodes((Center, Center, Face), grid)
    xb, yb, zb = nodes((Center, Center, Center), grid)

    b₀ = file["timeseries/b/0"][:, 1, :]

    iterations = parse.(Int, keys(file["timeseries/t"]))

    anim = @animate for (i, iter) in enumerate(iterations)

        @info "Plotting iteration $iter of $(iterations[end])..."

        u = file["timeseries/u/$iter"][:, 1, :]
        w = file["timeseries/w/$iter"][:, 1, :]
        b = file["timeseries/b/$iter"][:, 1, :]
        t = file["timeseries/t/$iter"]

        b′ = b .- b₀

        wlims, wlevels = nice_divergent_levels(w, 1e-4)
        ulims, ulevels = nice_divergent_levels(u, 1e-3)
        blims, blevels = nice_divergent_levels(b′, 1e-4)

        nan_solid(xu, zu, u, bump)
        nan_solid(xw, zw, w, bump)
        nan_solid(xb, zb, b, bump)

        u_title = @sprintf("x velocity, t = %.2f", t)

        u_plot = contourf(xu, zu, u'; title = u_title,                  color = :balance, aspectratio = :equal, linewidth = 0, levels = ulevels, clims = ulims)
        w_plot = contourf(xw, zw, w'; title = "z velocity",             color = :balance, aspectratio = :equal, linewidth = 0, levels = wlevels, clims = wlims)
        b_plot = contourf(xb, zb, b′'; title = "buoyancy perturbation", color = :balance, aspectratio = :equal, linewidth = 0, levels = blevels, clims = blims)

        plot(u_plot, w_plot, b_plot, layout = (3, 1), size = (1200, 1200))
    end

    mp4(anim, prefix * ".mp4", fps = 16)

    close(file)
end

function plot_implicit_explicit_difference(filename)

    file_explicit = jldopen(filename * "_explicit.jld2")
    file_implicit = jldopen(filename * "_implicit.jld2")

    iterations = parse.(Int, keys(file_explicit["timeseries/t"]))

    comparison_u = zeros(length(iterations))
    comparison_w = zeros(length(iterations))
    comparison_b = zeros(length(iterations))

    for (i, iter) in enumerate(iterations)

        u_explicit = file_explicit["timeseries/u/$iter"][:, 1, :]
        w_explicit = file_explicit["timeseries/w/$iter"][:, 1, :]
        b_explicit = file_explicit["timeseries/b/$iter"][:, 1, :]

        u_implicit = file_implicit["timeseries/u/$iter"][:, 1, :]
        w_implicit = file_implicit["timeseries/w/$iter"][:, 1, :]
        b_implicit = file_implicit["timeseries/b/$iter"][:, 1, :]

        comparison_u[i] = norm(u_explicit .- u_implicit)
        comparison_w[i] = norm(w_explicit .- w_implicit)
        comparison_b[i] = norm(b_explicit .- b_implicit)
    end

    kwargs = (linewidth = 2, foreground_color_legend = nothing, legendfontsize = 12, legend = :right, grid = false,
              xtickfontsize = 12, ytickfontsize=12, xlabel = "time", ylabel = "norm of difference")

     plot(iterations, comparison_u, label = "u"; kwargs...)
    plot!(iterations, comparison_w, label = "w"; kwargs...)
    plot!(iterations, comparison_b, label = "b"; kwargs...)

    savefig(filename * "_comparison_implicit_explicit.png")

    close(file_explicit)
    close(file_implicit)
end

visualize_internal_tide_simulation("internal_tide_explicit")
visualize_internal_tide_simulation("internal_tide_implicit")
plot_implicit_explicit_difference("internal_tide")


A = rand(1:4,3,3)

x = fill(1.0, (3))

b = A*x

Asym = A + A'

Apd = A'A

A\b

Atall = rand(3, 2)
display(Atall)
Atall\b

v = randn(3)
rankdef = [v v]

rankdef\b

Ashort = rand(2, 3)
display(Ashort)
Ashort\b[1:2]

es-es/intro-to-julia-ES/plutos/10. Algebra lineal basica.jl


@testset "triangles" begin
    A = load_matrix_network("clique-10")
    mytriangles = triangles(A)
    z = collect(mytriangles)
    @test length(z) == 120
    ei,ej,ek = unzip_triangles(z)
    @test length(ei) == 120

    A = erdos_renyi_undirected(100,0.2)
    z = collect(triangles(A;symmetries = true))
    @test sum(Diagonal(sparse(A)^3)) == length(z)
end


@testset "lattice" begin
    @test_throws AssertionError lattice(:chain, 0)

    @test lattice(:chain, 1) == [0]

    @test lattice(:chain, 2) == [0 2
                                 2 0]
#
    @test lattice(:chain, 2, ispbc=false) == [
                                            0 1
                                            1 0]

    @test lattice(:chain, 3) == [0 1 1
                                 1 0 1
                                 1 1 0]

    @test lattice(:chain, 3, ispbc=false) == [
                                            0 1 0
                                            1 0 1
                                            0 1 0
                                            ]
#
    @test_throws DomainError issquare(-1)
    @test issquare(0)
    @test issquare(1)
    @test !issquare(2)
    @test !issquare(3)
    @test issquare(4)
    @test issquare(9)

    @test_throws AssertionError lattice(:square, 0)
    @test_throws AssertionError lattice(:square, 3)

    @test lattice(:square, 1) == [0]

    @test lattice(:square, 4) == [
                                0 2 2 0
                                2 0 0 2
                                2 0 0 2
                                0 2 2 0
                                ]
#
    @test lattice(:square, 4, ispbc=false) == [
                                0 1 1 0
                                1 0 0 1
                                1 0 0 1
                                0 1 1 0
                                ]

    @test lattice(:square, 9) == [
                                0 1 1 1 0 0 1 0 0
                                1 0 1 0 1 0 0 1 0
                                1 1 0 0 0 1 0 0 1
                                1 0 0 0 1 1 1 0 0
                                0 1 0 1 0 1 0 1 0
                                0 0 1 1 1 0 0 0 1
                                1 0 0 1 0 0 0 1 1
                                0 1 0 0 1 0 1 0 1
                                0 0 1 0 0 1 1 1 0
                                ]
#
    @test lattice(:square, 9, ispbc=false) == [
                                0 1 0 1 0 0 0 0 0
                                1 0 1 0 1 0 0 0 0
                                0 1 0 0 0 1 0 0 0
                                1 0 0 0 1 0 1 0 0
                                0 1 0 1 0 1 0 1 0
                                0 0 1 0 1 0 0 0 1
                                0 0 0 1 0 0 0 1 0
                                0 0 0 0 1 0 1 0 1
                                0 0 0 0 0 1 0 1 0
                                ]

end


@with_kw mutable struct AdadeltaParameters{T <: AbstractFloat} <: AbstractPolicyParameters
    ρ::T = 0.95
    ϵ::T = 1e-6
end

struct Adadelta <: AbstractSmoothing
    params::AdadeltaParameters{Float64}
    grms::Vector{Float64}
    xrms::Vector{Float64}
    xprev::Vector{Float64}

    function (::Type{Adadelta})(; kwargs...)
        return new(AdadeltaParameters(; kwargs...),Vector{Float64}(),Vector{Float64}(),Vector{Float64}())
    end
end

function initialize!(adadelta::Adadelta,x₀::Vector{Float64})
    resize!(adadelta.grms,length(x₀))
    resize!(adadelta.xrms,length(x₀))
    resize!(adadelta.xprev,length(x₀))
    adadelta.grms .= zero(length(x₀))
    adadelta.xrms .= zero(length(x₀))
    adadelta.xprev .= x₀
end

function smooth!(adadelta::Adadelta,klocal::Integer,kglobal::Integer,x::AbstractVector,gprev::AbstractVector,gcurr::AbstractVector)
    @unpack ρ,ϵ = adadelta.params
    Δx = x - adadelta.xprev
    adadelta.grms .= ρ*adadelta.grms + (1-ρ)*(gprev .* gprev)
    adadelta.xrms .= ρ*adadelta.xrms + (1-ρ)*(Δx .* Δx)
    gcurr .= gprev .* ((.√adadelta.xrms .+ ϵ) ./ (.√adadelta.grms .+ ϵ))
    adadelta.xprev .= x
end


function ls_params(ϕ::AbstractVector{T}, c::Bool, p::AbstractVector{<:Integer}, q::AbstractVector{<:Integer}) where T
    P = isempty(p) ? 0 : maximum(p)
    Np = length(p)
    Q = isempty(q) ? 0 : maximum(q)
    Nq = length(q)
    ϕ0 = zeros(T, P)
    θ0 = zeros(T, Q)
    μ0 = c ? ϕ[1] : zero(T)
    for (i, l) in enumerate(p)
        ϕ0[l] = ϕ[i + c]
    end
    for (i, l) in enumerate(q)
        θ0[l] = ϕ[i + Np + c]
    end
    β0 = ϕ[Nq + Np + 1 + c:end]
    return μ0, ϕ0, θ0, β0
end

function ls_matrix!(Z::AbstractMatrix{T}, z::AbstractVector{T}, a::AbstractVector{T}, x::AbstractMatrix{T}, c::Bool, p::AbstractVector{<:Integer}, q::AbstractVector{<:Integer}) where T
    all(q .!= 0) || throw(ArgumentError("order of ma can't be zero"))
    all(p .!= 0) || throw(ArgumentError("order of ar can't be zero"))
    N = length(z)
    P = isempty(p) ? 0 : maximum(p)
    Np = length(p)
    Q = isempty(q) ? 0 : maximum(q)
    Nq = length(q)
    X = size(x, 2)
    if c
        Z[:,1] .= 1
    end
    for (i, l) in enumerate(p)
        @views Z[:,i + c] = z[P + Q - l + 1:N - l]
    end
    for (i, l) in enumerate(q)
        @views Z[:,i + Np + c] = -a[P + Q - l + 1:N - l]
    end
    if X > 0 @views Z[:, Np + Nq + 1 + c:end] = x[P:N - 1, :] end
    return Z, z[P + Q + 1:N]
end

function ls_matrix(z::AbstractVector{T}, a::AbstractVector{T}, x::AbstractMatrix{T}, c::Bool, p::AbstractVector{<:Integer}, q::AbstractVector{<:Integer}) where T
    N = length(z)
    P = isempty(p) ? 0 : maximum(p)
    Np = length(p)
    Q = isempty(q) ? 0 : maximum(q)
    Nq = length(q)
    X = size(x, 2)
    Z = Matrix{T}(undef, N - P - Q, Np + Nq + X + c)
    ls_matrix!(Z, z, a, x, c, p, q)
end

function least_squares(z::AbstractVector{T}, a::AbstractVector{T}, x::AbstractMatrix{T}, c::Bool, p::AbstractVector{<:Integer}, q::AbstractVector{<:Integer}) where T
    Z, z = ls_matrix(z, a, x, c, p, q)
    ϕ = Z \ z
    ε = z - Z * ϕ
    σ2 = dot(ε, ε) / length(ε)
    μ, ϕ, θ, β = ls_params(ϕ, c, p, q)
    return μ, ϕ, θ, β, σ2
end

least_squares(z::AbstractVector{T}, a::AbstractVector{T}, x::AbstractMatrix{T}, c::Bool, p::Integer, q::Integer) where T = least_squares(z, a, x, c, collect(0:p), collect(1:q))


using JuMP, Ipopt
# using Gurobi
using SparseArrays # ybus
# using DataFrames # for comparison

include(string(pwd(), "/src/struct_network.jl"))
include(string(pwd(), "/src/readnetwork.jl"))
include(string(pwd(), "/src/ybus.jl"))
#----------------------------------------------------------------------------------#
solver = "Ipopt"
# solver = "Gurobi"

testsystem = "case_ieee30"
buses, lines, generators, datamat = readnetwork(testsystem)

lineset = 1:length(lines)
busset = 1:length(buses)
genset = 1:length(generators)

B_g = []
for g in genset
   push!(B_g, generators[g].location)
end

B_gn = Array{Array{Int64,1}}(undef, length(buses))
for ii in 1:length(buses)
   B_gn[ii] = Int64[]
end
for g in genset
   push!(B_gn[generators[g].location], g)
end

############
Ybus, yff, yft, ytf, ytt = ybus( buses, lines )

yff_r = real.(yff)
yff_i = imag.(yff)
ytt_r = real.(ytt)
ytt_i = imag.(ytt)
yft_r = real.(yft)
yft_i = imag.(yft)
ytf_r = real.(ytf)
ytf_i = imag.(ytf)

#---------------------------------------------------------------------------------#
if solver == "Gurobi"
   m = Model(Gurobi.Optimizer)
   # set_optimizer_attribute(m, "MIPGap", 1e-4)
   # set_optimizer_attribute(m, "Method", 5)
   # set_optimizer_attribute(m, "Presolve", 0) # max 2
   # set_optimizer_attribute(m, "NonConvex", 2)
#    set_optimizer_attribute(m, "NumericFocus", 3) # max 3
   # set_optimizer_attribute(m, "Threads", 6)
elseif solver == "Mosek"
   m = Model(Mosek.Optimizer)
elseif solver == "KNITRO"
   m = Model(KNITRO.Optimizer)
elseif solver == "Ipopt"
   m = Model(Ipopt.Optimizer)
end

@variables(m, begin
   - π <= θ[b in busset] <= π
   generators[g].Pmin <= pg[g in genset] <= generators[g].Pmax  # real power output of generator
end)

for b in busset
    if buses[b].btype == 3
        fix(θ[b], 0; force=true) # θ[1] = 0 at root node
    end
end

@expression(m, p_ft[l in lineset], + yft_i[l]*(θ[lines[l].fbus]-θ[lines[l].tbus]))
@expression(m, p_tf[l in lineset], + ytf_i[l]*(θ[lines[l].tbus]-θ[lines[l].fbus]))

@constraint(m, PBalance[b in busset],
    + sum(p_ft[l] for l in buses[b].outline)
    + sum(p_tf[l] for l in buses[b].inline)
    - sum(pg[g] for g in B_gn[b])
    + buses[b].Pd
    + buses[b].Gs*1.0^2
        == 0)

@constraint(m, LineCap_ub[l in lineset; lines[l].u != 0], p_ft[l] <= + lines[l].u)
@constraint(m, LineCap_lb[l in lineset; lines[l].u != 0], p_ft[l] >= - lines[l].u)
@constraint(m, LineCapBW_ub[l in lineset; lines[l].u != 0], p_tf[l] <= + lines[l].u)
@constraint(m, LineCapBW_lb[l in lineset; lines[l].u != 0], p_tf[l] >= - lines[l].u)

@constraint(m, AngDiffmax[l in lineset], θ[lines[l].fbus] - θ[lines[l].tbus] <= min(lines[l].angmax, deg2rad(60)))
@constraint(m, AngDiffmin[l in lineset], θ[lines[l].fbus] - θ[lines[l].tbus] >= max(lines[l].angmin, -deg2rad(60)))

@expression(m, OCgen[g in genset],
    generators[g].cost[1]*(pg[g]*datamat["baseMVA"])^2
    + generators[g].cost[2]*(pg[g]*datamat["baseMVA"])
    + generators[g].cost[3])
@expression(m, OC, sum(OCgen[g] for g in genset))

@objective(m, Min, OC)

optimize!(m)
@show termination_status(m)
@show objective_value(m)


@testset "Neighborhoods" begin
  @testset "NormBall" begin
    # Euclidean metric
    ball = NormBall(T(1/2))
    @test evaluate(metric(ball), T[0], T[0]) ≤ radius(ball)
    @test evaluate(metric(ball), T[0], T[1]) > radius(ball)

    ball = NormBall(T(1))
    @test evaluate(metric(ball), T[0,0], T[0,0]) ≤ radius(ball)
    @test evaluate(metric(ball), T[0,0], T[1,0]) ≤ radius(ball)
    @test evaluate(metric(ball), T[0,0], T[0,1]) ≤ radius(ball)

    # Chebyshev metric
    ball = NormBall(T(1/2), Chebyshev())
    @test evaluate(metric(ball), T[0], T[0]) ≤ radius(ball)
    @test evaluate(metric(ball), T[0], T[1]) > radius(ball)

    for r in [1.,2.,3.,4.,5.]
      ball = NormBall(r, Chebyshev())
      for i in 0.0:1.0:r, j in 0.0:1.0:r
        @test evaluate(metric(ball), T[0,0], T[i,j]) ≤ radius(ball)
      end
    end

    ball = NormBall(T(1))
    if T == Float32
      @test sprint(show, ball) == "NormBall{Float32}(1.0, Euclidean(0.0))"
    else
      @test sprint(show, ball) == "NormBall{Float64}(1.0, Euclidean(0.0))"
    end
  end

  @testset "Ellipsoid" begin
    # 2D ellipsoid rotated 45 degrees in GSLIB convention
    ellipsoid = Ellipsoid(T[2,1], T[45], convention=GSLIB)

    # tests along main semiaxes, slightly below threshold
    @test evaluate(metric(ellipsoid), T[0,0], T(1.9) * T[√2/2, √2/2]) ≤ T(1)
    @test evaluate(metric(ellipsoid), T[0,0], T(0.9) * T[√2/2,-√2/2]) ≤ T(1)

    # tests along main semiaxes, slightly above threshold
    @test evaluate(metric(ellipsoid), T[0,0], T(2.1) * T[√2/2, √2/2]) > T(1)
    @test evaluate(metric(ellipsoid), T[0,0], T(1.1) * T[√2/2,-√2/2]) > T(1)

    # 3D ellipsoid rotated (45, -45, 0) in GSLIB convention
    ellipsoid = Ellipsoid(T[3,2,1], T[45,-45,0], convention=GSLIB)

    # tests along main semiaxes, slightly below threshold
    @test evaluate(metric(ellipsoid), T[0,0,0], T(2.9) * T[0.5,0.5,-√2/2]) ≤ T(1)
    @test evaluate(metric(ellipsoid), T[0,0,0], T(1.9) * T[√2/2,-√2/2,0.0]) ≤ T(1)
    @test evaluate(metric(ellipsoid), T[0,0,0], T(0.9) * T[0.5,0.5,√2/2]) ≤ T(1)

    # Tests along main semiaxes, slightly above threshold
    @test evaluate(metric(ellipsoid), T[0,0,0], T(3.1) * T[0.5,0.5,-√2/2]) > T(1)
    @test evaluate(metric(ellipsoid), T[0,0,0], T(2.1) * T[√2/2,-√2/2,0.0]) > T(1)
    @test evaluate(metric(ellipsoid), T[0,0,0], T(1.1) * T[0.5,0.5,√2/2]) > T(1)

    # 2D simple test of default convention
    d₁ = metric(Ellipsoid([1.,1.], [0.]))
    d₂ = metric(Ellipsoid([1.,2.], [0.]))
    @test evaluate(d₁, [1.,0.], [0.,0.]) == evaluate(d₁, [0.,1.], [0.,0.])
    @test evaluate(d₂, [1.,0.], [0.,0.]) != evaluate(d₂, [0.,1.], [0.,0.])

    # 3D simple test of default convention
    d₃ = metric(Ellipsoid([1.,.5,.5], [π/4,0.,0.]))
    @test evaluate(d₃, [1.,1.,0.], [0.,0.,0.]) ≈ √2
    @test evaluate(d₃, [-1.,1.,0.], [0.,0.,0.]) ≈ √8

    # test of intrinsic conventions
    gslib = metric(Ellipsoid([50.,25.,5.], [30.,-30.,30.], convention=GSLIB))
    tait  = metric(Ellipsoid([25.,50.,5.], [-π/6,-π/6,π/6], convention=TaitBryanIntr))
    euler = metric(Ellipsoid([50.,25.,5.], deg2rad.([-78,-41,-50]), convention=EulerIntr))
    lpf   = metric(Ellipsoid([50.,25.,5.], [78.,41.,50.], convention=Leapfrog))
    dm    = metric(Ellipsoid([50.,25.,5.], [78.,41.,50.], convention=Datamine))

    @test evaluate(gslib, [1.,0.,0.], [0.,0.,0.]) ≈ 0.1325707358356285
    @test evaluate(gslib, [0.,1.,0.], [0.,0.,0.]) ≈ 0.039051248379533283
    @test evaluate(gslib, [0.,0.,1.], [0.,0.,0.]) ≈ 0.15132745950421558

    @test evaluate(gslib, [0.,0.,1.], [0.,0.,0.]) ≈ evaluate(tait, [0.,0.,1.], [0.,0.,0.])
    @test evaluate(euler, [0.,0.,1.], [0.,0.,0.]) ≈ evaluate(dm, [0.,0.,1.], [0.,0.,0.])
    @test evaluate(euler, [0.,0.,1.], [0.,0.,0.]) ≈ evaluate(lpf, [0.,0.,1.], [0.,0.,0.])
    @test evaluate(euler, [0.,0.,1.], [0.,0.,0.]) - evaluate(gslib, [0.,0.,1.], [0.,0.,0.]) < 10^-3

    # test for https://github.com/JuliaEarth/GeoStats.jl/issues/197
    gslib = metric(Ellipsoid([1.0, 0.5, 0.3], [100., -10., -20.], convention = GSLIB))

    @test evaluate(gslib, [1.,0.,0.], [0.,0.,0.]) ≈ 1.233956165693094
    @test evaluate(gslib, [0.,1.,0.], [0.,0.,0.]) ≈ 2.14219475359467
    @test evaluate(gslib, [0.,0.,1.], [0.,0.,0.]) ≈ 3.1621930570302914

    # test of extrinsic conventions
    xtait  = metric(Ellipsoid([50.,25.,5.], [π,0,π/2], convention=TaitBryanExtr))
    xeuler = metric(Ellipsoid([50.,25.,5.], [-π/2,-π/2,-π/2], convention=EulerExtr))

    @test evaluate(xtait, [1.,0.,0.], [0.,0.,0.]) ≈ 0.20
    @test evaluate(xtait, [0.,1.,0.], [0.,0.,0.]) ≈ 0.04
    @test evaluate(xtait, [0.,0.,1.], [0.,0.,0.]) ≈ 0.02

    @test evaluate(xtait, [1.,0.,0.], [0.,0.,0.]) ≈ evaluate(xeuler, [1.,0.,0.], [0.,0.,0.])
    @test evaluate(xtait, [0.,1.,0.], [0.,0.,0.]) ≈ evaluate(xeuler, [0.,1.,0.], [0.,0.,0.])
    @test evaluate(xtait, [0.,0.,1.], [0.,0.,0.]) ≈ evaluate(xeuler, [0.,0.,1.], [0.,0.,0.])

    ellipsoid = Ellipsoid(T.((2,1)), T.((45,)), convention=GSLIB)
    if T == Float32
      @test sprint(show, ellipsoid) == "Ellipsoid{Float32}((2.0f0, 1.0f0), (45.0f0,), GSLIB)"
    else
      @test sprint(show, ellipsoid) == "Ellipsoid{Float64}((2.0, 1.0), (45.0,), GSLIB)"
    end
  end
end


using Test
using CLIMA
using StaticArrays
using LinearAlgebra
using KernelAbstractions
using CLIMA.MPIStateArrays: device

include("ode_tests_common.jl")

CLIMA.init()
const ArrayType = CLIMA.array_type()

@testset "ODE Solvers" begin
    @testset "Convergence/extensive" begin
        @testset "1-rate ODE" begin
            function rhs!(dQ, Q, ::Nothing, time; increment)
                if increment
                    dQ .+= Q * cos(time)
                else
                    dQ .= Q * cos(time)
                end
            end
            exactsolution(q0, time) = q0 * exp(sin(time))

            @testset "Explicit methods" begin
                finaltime = 20.0
                dts = [2.0^(-k) for k in 0:7]
                errors = similar(dts)
                q0 =
                    ArrayType === Array ? [1.0] : range(-1.0, 1.0, length = 303)
                for (method, expected_order) in explicit_methods
                    for (n, dt) in enumerate(dts)
                        Q = ArrayType(q0)
                        solver = method(rhs!, Q; dt = dt, t0 = 0.0)
                        solve!(Q, solver; timeend = finaltime)
                        Q = Array(Q)
                        errors[n] =
                            maximum(@. abs(Q - exactsolution(q0, finaltime)))
                    end
                    rates = log2.(errors[1:(end - 1)] ./ errors[2:end])
                    @test isapprox(rates[end], expected_order; atol = 0.17)
                end
            end
        end

        @testset "Two-rate ODE with a linear stiff part" begin
            c = 100.0
            function rhs_full!(dQ, Q, ::Nothing, time; increment)
                if increment
                    dQ .+= im * c * Q .+ exp(im * time)
                else
                    dQ .= im * c * Q .+ exp(im * time)
                end
            end

            function rhs_nonlinear!(dQ, Q, ::Nothing, time; increment)
                if increment
                    dQ .+= exp(im * time)
                else
                    dQ .= exp(im * time)
                end
            end
            rhs_slow! = rhs_nonlinear!

            function rhs_linear!(dQ, Q, ::Nothing, time; increment)
                if increment
                    dQ .+= im * c * Q
                else
                    dQ .= im * c * Q
                end
            end
            rhs_fast! = rhs_linear!

            function exactsolution(q0, time)
                q0 * exp(im * c * time) +
                (exp(im * time) - exp(im * c * time)) / (im * (1 - c))
            end

            @testset "IMEX methods" begin
                finaltime = pi / 2
                dts = [2.0^(-k) for k in 2:13]
                errors = similar(dts)

                q0 = ArrayType <: Array ? [1.0] : range(-1.0, 1.0, length = 303)
                for (method, expected_order) in imex_methods
                    for split_explicit_implicit in (false, true)
                        for variant in (LowStorageVariant(), NaiveVariant())
                            for (n, dt) in enumerate(dts)
                                Q = ArrayType{ComplexF64}(q0)
                                rhs! =
                                    split_explicit_implicit ? rhs_nonlinear! :
                                    rhs_full!
                                solver = method(
                                    rhs!,
                                    rhs_linear!,
                                    LinearBackwardEulerSolver(
                                        DivideLinearSolver(),
                                        isadjustable = true,
                                    ),
                                    Q;
                                    dt = dt,
                                    t0 = 0.0,
                                    split_explicit_implicit = split_explicit_implicit,
                                    variant = variant,
                                )
                                solve!(Q, solver; timeend = finaltime)
                                Q = Array(Q)
                                errors[n] = maximum(@. abs(
                                    Q - exactsolution(q0, finaltime),
                                ))
                            end

                            rates = log2.(errors[1:(end - 1)] ./ errors[2:end])
                            @test errors[1] < 2.0
                            @test isapprox(
                                rates[end],
                                expected_order;
                                atol = 0.1,
                            )
                        end
                    end
                end
            end

            @testset "MRRK methods (no substeps)" begin
                finaltime = pi / 2
                dts = [2.0^(-k) for k in 2:11]
                errors = similar(dts)
                for (slow_method, slow_expected_order) in slow_mrrk_methods
                    for (fast_method, fast_expected_order) in fast_mrrk_methods
                        q0 = ArrayType === Array ? [1.0] :
                            range(-1.0, 1.0, length = 303)
                        for (n, dt) in enumerate(dts)
                            Q = ArrayType{ComplexF64}(q0)
                            solver = MultirateRungeKutta(
                                (
                                    slow_method(rhs_slow!, Q),
                                    fast_method(rhs_fast!, Q),
                                );
                                dt = dt,
                                t0 = 0.0,
                            )
                            solve!(Q, solver; timeend = finaltime)
                            Q = Array(Q)
                            errors[n] = maximum(@. abs(
                                Q - exactsolution(q0, finaltime),
                            ))
                        end

                        rates = log2.(errors[1:(end - 1)] ./ errors[2:end])
                        min_order =
                            min(slow_expected_order, fast_expected_order)
                        max_order =
                            max(slow_expected_order, fast_expected_order)
                        @test (
                            isapprox(rates[end], min_order; atol = 0.1) ||
                            isapprox(rates[end], max_order; atol = 0.1) ||
                            min_order <= rates[end] <= max_order
                        )
                    end
                end
            end

            @testset "MRRK methods (with substeps)" begin
                finaltime = pi / 2
                dts = [2.0^(-k) for k in 2:15]
                errors = similar(dts)
                for (slow_method, slow_expected_order) in slow_mrrk_methods
                    for (fast_method, fast_expected_order) in fast_mrrk_methods
                        q0 = ArrayType === Array ? [1.0] :
                            range(-1.0, 1.0, length = 303)
                        for (n, fast_dt) in enumerate(dts)
                            slow_dt = c * fast_dt
                            Q = ArrayType{ComplexF64}(q0)
                            solver = MultirateRungeKutta((
                                slow_method(rhs_slow!, Q; dt = slow_dt),
                                fast_method(rhs_fast!, Q; dt = fast_dt),
                            ))
                            solve!(Q, solver; timeend = finaltime)
                            Q = Array(Q)
                            errors[n] = maximum(@. abs(
                                Q - exactsolution(q0, finaltime),
                            ))
                        end

                        rates = log2.(errors[1:(end - 1)] ./ errors[2:end])
                        min_order =
                            min(slow_expected_order, fast_expected_order)
                        max_order =
                            max(slow_expected_order, fast_expected_order)
                        @test (
                            isapprox(rates[end], min_order; atol = 0.1) ||
                            isapprox(rates[end], max_order; atol = 0.1) ||
                            min_order <= rates[end] <= max_order
                        )
                    end
                end
            end

            @testset "MIS methods (with substeps)" begin
                finaltime = pi / 2
                dts = [2.0^(-k) for k in 2:11]
                errors = similar(dts)
                for (mis_method, mis_expected_order) in mis_methods
                    for fast_method in (LSRK54CarpenterKennedy,)
                        q0 = ArrayType === Array ? [1.0] :
                            range(-1.0, 1.0, length = 303)
                        for (n, dt) in enumerate(dts)
                            Q = ArrayType{ComplexF64}(q0)
                            solver = mis_method(
                                rhs_slow!,
                                rhs_fast!,
                                fast_method,
                                4,
                                Q;
                                dt = dt,
                                t0 = 0.0,
                            )
                            solve!(Q, solver; timeend = finaltime)
                            Q = Array(Q)
                            errors[n] = maximum(@. abs(
                                Q - exactsolution(q0, finaltime),
                            ))
                        end
                        rates = log2.(errors[1:(end - 1)] ./ errors[2:end])
                        @test isapprox(
                            rates[end],
                            mis_expected_order;
                            atol = 0.1,
                        )
                    end
                end
            end
        end

        #=
        Test problem (4.2) from RobertsSarsharSandu2018arxiv
        @article{RobertsSarsharSandu2018arxiv,
          title={Coupled Multirate Infinitesimal GARK Schemes for Stiff Systems with
                 Multiple Time Scales},
          author={Roberts, Steven and Sarshar, Arash and Sandu, Adrian},
          journal={arXiv preprint arXiv:1812.00808},
          year={2019}
        }

        Note: The actual rates are all over the place with this test and passing largely
              depends on final dt size
        =#
        @testset "2-rate ODE from RobertsSarsharSandu2018arxiv" begin
            ω = 100
            λf = -10
            λs = -1
            ξ = 1 // 10
            α = 1
            ηfs = ((1 - ξ) / α) * (λf - λs)
            ηsf = -ξ * α * (λf - λs)
            Ω = @SMatrix [
                λf ηfs
                ηsf λs
            ]

            function rhs_fast!(dQ, Q, param, t; increment)
                @inbounds begin
                    increment || (dQ .= 0)
                    yf = Q[1]
                    ys = Q[2]
                    gf = (-3 + yf^2 - cos(ω * t)) / 2yf
                    gs = (-2 + ys^2 - cos(t)) / 2ys
                    dQ[1] += Ω[1, 1] * gf + Ω[1, 2] * gs - ω * sin(ω * t) / 2yf
                end
            end

            function rhs_slow!(dQ, Q, param, t; increment)
                @inbounds begin
                    increment || (dQ .= 0)
                    yf = Q[1]
                    ys = Q[2]
                    gf = (-3 + yf^2 - cos(ω * t)) / 2yf
                    gs = (-2 + ys^2 - cos(t)) / 2ys
                    dQ[2] += Ω[2, 1] * gf + Ω[2, 2] * gs - sin(t) / 2ys
                end
            end

            exactsolution(t) = [sqrt(3 + cos(ω * t)); sqrt(2 + cos(t))]

            @testset "MRRK methods (no substeps)" begin
                finaltime = 5π / 2
                dts = [2.0^(-k) for k in 2:8]
                error = similar(dts)
                for (slow_method, slow_expected_order) in slow_mrrk_methods
                    for (fast_method, fast_expected_order) in fast_mrrk_methods
                        for (n, dt) in enumerate(dts)
                            Q = exactsolution(0)
                            solver = MultirateRungeKutta(
                                (
                                    slow_method(rhs_slow!, Q),
                                    fast_method(rhs_fast!, Q),
                                );
                                dt = dt,
                                t0 = 0.0,
                            )
                            solve!(Q, solver; timeend = finaltime)
                            error[n] = norm(Q - exactsolution(finaltime))
                        end

                        rate = log2.(error[1:(end - 1)] ./ error[2:end])
                        min_order =
                            min(slow_expected_order, fast_expected_order)
                        max_order =
                            max(slow_expected_order, fast_expected_order)
                        @test (
                            isapprox(rate[end], min_order; atol = 0.3) ||
                            isapprox(rate[end], max_order; atol = 0.3) ||
                            min_order <= rate[end] <= max_order
                        )
                    end
                end
            end

            @testset "MRRK methods (with substeps)" begin
                finaltime = 5π / 2
                dts = [2.0^(-k) for k in 2:9]
                error = similar(dts)
                for (slow_method, slow_expected_order) in slow_mrrk_methods
                    for (fast_method, fast_expected_order) in fast_mrrk_methods
                        for (n, fast_dt) in enumerate(dts)
                            Q = exactsolution(0)
                            slow_dt = ω * fast_dt
                            solver = MultirateRungeKutta((
                                slow_method(rhs_slow!, Q; dt = slow_dt),
                                fast_method(rhs_fast!, Q; dt = fast_dt),
                            ))
                            solve!(Q, solver; timeend = finaltime)
                            error[n] = norm(Q - exactsolution(finaltime))
                        end

                        rate = log2.(error[1:(end - 1)] ./ error[2:end])
                        min_order =
                            min(slow_expected_order, fast_expected_order)
                        max_order =
                            max(slow_expected_order, fast_expected_order)
                        @test (
                            isapprox(rate[end], min_order; atol = 0.3) ||
                            isapprox(rate[end], max_order; atol = 0.3) ||
                            min_order <= rate[end] <= max_order
                        )
                    end
                end
            end

            @testset "MRRK methods with IMEX fast solver" begin
                function rhs_zero!(dQ, Q, param, t; increment)
                    if !increment
                        dQ .= 0
                    end
                end

                finaltime = 5π / 2
                dts = [2.0^(-k) for k in 2:9]
                error = similar(dts)
                for (slow_method, slow_expected_order) in slow_mrrk_methods
                    for (fast_method, fast_expected_order) in imex_methods
                        for (n, fast_dt) in enumerate(dts)
                            Q = exactsolution(0)
                            slow_dt = ω * fast_dt
                            solver = MultirateRungeKutta((
                                slow_method(rhs_slow!, Q; dt = slow_dt),
                                fast_method(
                                    rhs_fast!,
                                    rhs_zero!,
                                    LinearBackwardEulerSolver(
                                        DivideLinearSolver(),
                                        isadjustable = true,
                                    ),
                                    Q;
                                    dt = fast_dt,
                                    split_explicit_implicit = false,
                                ),
                            ))
                            solve!(Q, solver; timeend = finaltime)
                            error[n] = norm(Q - exactsolution(finaltime))
                        end

                        rate = log2.(error[1:(end - 1)] ./ error[2:end])
                        min_order =
                            min(slow_expected_order, fast_expected_order)
                        max_order =
                            max(slow_expected_order, fast_expected_order)
                        atol = fast_method == ARK2GiraldoKellyConstantinescu ?
                            0.5 : 0.37
                        @test (
                            isapprox(rate[end], min_order; atol = atol) ||
                            isapprox(rate[end], max_order; atol = atol) ||
                            min_order <= rate[end] <= max_order
                        )
                    end
                end
            end

            @testset "MIS methods (with substeps)" begin
                finaltime = 5π / 2
                dts = [2.0^(-k) for k in 2:10]
                error = similar(dts)
                for (mis_method, mis_expected_order) in mis_methods
                    for fast_method in (LSRK54CarpenterKennedy,)
                        for (n, dt) in enumerate(dts)
                            Q = exactsolution(0)
                            solver = mis_method(
                                rhs_slow!,
                                rhs_fast!,
                                fast_method,
                                4,
                                Q;
                                dt = dt,
                                t0 = 0.0,
                            )
                            solve!(Q, solver; timeend = finaltime)
                            error[n] = norm(Q - exactsolution(finaltime))
                        end

                        rate = log2.(error[1:(end - 1)] ./ error[2:end])
                        @test isapprox(
                            rate[end],
                            mis_expected_order;
                            atol = 0.1,
                        )
                    end
                end
            end
        end

        # Simple 3-rate problem based on test of RobertsSarsharSandu2018arxiv
        #
        # NOTE: Since we have no theory to say this ODE solver is accurate, the rates
        #      suggest that things are really only 2nd order.
        #
        # TODO: This is not great, but no theory to say we should be accurate!
        @testset "3-rate ODE" begin
            ω1, ω2, ω3 = 10000, 100, 1
            λ1, λ2, λ3 = -100, -10, -1
            β1, β2, β3 = 2, 3, 4

            ξ12 = λ2 / λ1
            ξ13 = λ3 / λ1
            ξ23 = λ3 / λ2

            α12, α13, α23 = 1, 1, 1

            η12 = ((1 - ξ12) / α12) * (λ1 - λ2)
            η13 = ((1 - ξ13) / α13) * (λ1 - λ3)
            η23 = ((1 - ξ23) / α23) * (λ2 - λ3)

            η21 = ξ12 * α12 * (λ2 - λ1)
            η31 = ξ13 * α13 * (λ3 - λ1)
            η32 = ξ23 * α23 * (λ3 - λ2)

            Ω = @SMatrix [
                λ1 η12 η13
                η21 λ2 η23
                η31 η32 λ3
            ]

            function rhs1!(dQ, Q, param, t; increment)
                @inbounds begin
                    increment || (dQ .= 0)
                    y1, y2, y3 = Q[1], Q[2], Q[3]
                    g = @SVector [
                        (-β1 + y1^2 - cos(ω1 * t)) / 2y1,
                        (-β2 + y2^2 - cos(ω2 * t)) / 2y2,
                        (-β3 + y3^2 - cos(ω3 * t)) / 2y3,
                    ]
                    dQ[1] += Ω[1, :]' * g - ω1 * sin(ω1 * t) / 2y1
                end
            end
            function rhs2!(dQ, Q, param, t; increment)
                @inbounds begin
                    increment || (dQ .= 0)
                    y1, y2, y3 = Q[1], Q[2], Q[3]
                    g = @SVector [
                        (-β1 + y1^2 - cos(ω1 * t)) / 2y1,
                        (-β2 + y2^2 - cos(ω2 * t)) / 2y2,
                        (-β3 + y3^2 - cos(ω3 * t)) / 2y3,
                    ]
                    dQ[2] += Ω[2, :]' * g - ω2 * sin(ω2 * t) / 2y2
                end
            end
            function rhs3!(dQ, Q, param, t; increment)
                @inbounds begin
                    increment || (dQ .= 0)
                    y1, y2, y3 = Q[1], Q[2], Q[3]
                    g = @SVector [
                        (-β1 + y1^2 - cos(ω1 * t)) / 2y1,
                        (-β2 + y2^2 - cos(ω2 * t)) / 2y2,
                        (-β3 + y3^2 - cos(ω3 * t)) / 2y3,
                    ]
                    dQ[3] += Ω[3, :]' * g - ω3 * sin(ω3 * t) / 2y3
                end
            end
            function rhs12!(dQ, Q, param, t; increment)
                rhs1!(dQ, Q, param, t; increment = increment)
                rhs2!(dQ, Q, param, t; increment = true)
            end

            exactsolution(t) = [
                sqrt(β1 + cos(ω1 * t)),
                sqrt(β2 + cos(ω2 * t)),
                sqrt(β3 + cos(ω3 * t)),
            ]

            @testset "MRRK methods (no substeps)" begin
                finaltime = π / 2
                dts = [2.0^(-k) for k in 2:10]
                error = similar(dts)
                for (rate3_method, rate3_order) in slow_mrrk_methods
                    for (rate2_method, rate2_order) in slow_mrrk_methods
                        for (rate1_method, rate1_order) in fast_mrrk_methods
                            for (n, dt) in enumerate(dts)
                                Q = exactsolution(0)
                                solver = MultirateRungeKutta(
                                    (
                                        rate3_method(rhs3!, Q),
                                        rate2_method(rhs2!, Q),
                                        rate1_method(rhs1!, Q),
                                    );
                                    dt = dt,
                                    t0 = 0.0,
                                )
                                solve!(Q, solver; timeend = finaltime)
                                error[n] = norm(Q - exactsolution(finaltime))
                            end

                            rate = log2.(error[1:(end - 1)] ./ error[2:end])
                            min_order =
                                min(rate3_order, rate2_order, rate1_order)
                            max_order =
                                max(rate3_order, rate2_order, rate1_order)
                            @test 2 <= rate[end]
                        end
                    end
                end
            end

            @testset "MRRK methods (with substeps)" begin
                finaltime = π / 2
                dts = [2.0^(-k) for k in 10:17]
                error = similar(dts)
                for (rate3_method, rate3_order) in slow_mrrk_methods
                    for (rate2_method, rate2_order) in slow_mrrk_methods
                        for (rate1_method, rate1_order) in fast_mrrk_methods
                            for (n, dt1) in enumerate(dts)
                                Q = exactsolution(0)
                                dt2 = (ω1 / ω2) * dt1
                                dt3 = (ω2 / ω3) * dt2
                                solver = MultirateRungeKutta((
                                    rate3_method(rhs3!, Q; dt = dt3),
                                    rate2_method(rhs2!, Q; dt = dt2),
                                    rate1_method(rhs1!, Q; dt = dt1),
                                ))
                                solve!(Q, solver; timeend = finaltime)
                                error[n] = norm(Q - exactsolution(finaltime))
                            end

                            rate = log2.(error[1:(end - 1)] ./ error[2:end])
                            min_order =
                                min(rate3_order, rate2_order, rate1_order)
                            max_order =
                                max(rate3_order, rate2_order, rate1_order)

                            @test 2 <= rate[end]
                        end
                    end
                end
            end
        end

        #=
        Test problem (8.2) from Sandu (2019) for MRI-GARK Schemes
            @article{Sandu2019,
                title={A class of multirate infinitesimal gark methods},
                author={Sandu, Adrian},
                journal={SIAM Journal on Numerical Analysis},
                volume={57},
                number={5},
                pages={2300--2327},
                year={2019},
                publisher={SIAM},
                doi={10.1137/18M1205492}
            }
        =#
        @testset "2-rate problem" begin
            ω = 20
            λf = -10
            λs = -1
            ξ = 1 // 10
            α = 1
            ηfs = ((1 - ξ) / α) * (λf - λs)
            ηsf = -ξ * α * (λf - λs)
            Ω = @SMatrix [
                λf ηfs
                ηsf λs
            ]

            function rhs_fast!(dQ, Q, param, t; increment)
                @kernel function knl!(dQ, Q, t, increment)
                    @inbounds begin
                        increment || (dQ .= 0)
                        yf = Q[1]
                        ys = Q[2]
                        gf = (-3 + yf^2 - cos(ω * t)) / 2yf
                        gs = (-2 + ys^2 - cos(t)) / 2ys
                        dQ[1] +=
                            Ω[1, 1] * gf + Ω[1, 2] * gs - ω * sin(ω * t) / 2yf
                    end
                end
                event = Event(device(Q))
                event = knl!(device(Q), 1)(
                    dQ,
                    Q,
                    t,
                    increment;
                    ndrange = 1,
                    dependencies = (event,),
                )
                wait(device(Q), event)
            end

            function rhs_slow!(dQ, Q, param, t; increment)
                @kernel function knl!(dQ, Q, t, increment)
                    @inbounds begin
                        increment || (dQ .= 0)
                        yf = Q[1]
                        ys = Q[2]
                        gf = (-3 + yf^2 - cos(ω * t)) / 2yf
                        gs = (-2 + ys^2 - cos(t)) / 2ys
                        dQ[2] += Ω[2, 1] * gf + Ω[2, 2] * gs - sin(t) / 2ys
                    end
                end
                event = Event(device(Q))
                event = knl!(device(Q), 1)(
                    dQ,
                    Q,
                    t,
                    increment;
                    ndrange = 1,
                    dependencies = (event,),
                )
                wait(device(Q), event)
            end

            struct ODETestConvNonLinBE <: AbstractBackwardEulerSolver end
            ODESolvers.Δt_is_adjustable(::ODETestConvNonLinBE) = true
            function (::ODETestConvNonLinBE)(Q, Qhat, α, p, t)
                @kernel function knl!(Q, Qhat, α, p, t)
                    @inbounds begin
                        # Slow RHS has zero tendency of yf so just copy Qhat
                        Q[1] = yf = Qhat[1]

                        gf = (-3 + yf^2 - cos(ω * t)) / 2yf

                        # solves: Q = Qhat[2] + α * rhs_slow(Q, t)
                        # (simplifies to a quadratic equation)
                        a = 2 - α * Ω[2, 2]
                        b = -2 * (Qhat[2] + α * Ω[2, 1] * gf)
                        c = α * (Ω[2, 2] * (2 + cos(t)) + sin(t))
                        Q[2] = (-b + sqrt(b^2 - 4 * a * c)) / (2a)
                    end
                end
                event = Event(device(Q))
                event = knl!(device(Q), 1)(
                    Q,
                    Qhat,
                    α,
                    p,
                    t;
                    ndrange = 1,
                    dependencies = (event,),
                )
                wait(device(Q), event)
            end

            exactsolution(t) =
                ArrayType([sqrt(3 + cos(ω * t)); sqrt(2 + cos(t))])

            finaltime = 1
            dts = [2.0^(-k) for k in 2:7]
            error = similar(dts)
            @testset "Explicit" begin
                for (mri_method, mri_expected_order) in mrigark_erk_methods
                    for (fast_method, fast_expected_order) in
                        fast_mrigark_methods
                        for (n, slow_dt) in enumerate(dts)
                            Q = exactsolution(0)
                            fast_dt = slow_dt / ω
                            fastsolver = fast_method(rhs_fast!, Q; dt = fast_dt)
                            solver = mri_method(
                                rhs_slow!,
                                fastsolver,
                                Q,
                                dt = slow_dt,
                            )
                            solve!(Q, solver; timeend = finaltime)
                            error[n] = norm(Q - exactsolution(finaltime))
                        end

                        rate = log2.(error[1:(end - 1)] ./ error[2:end])
                        order = mri_expected_order
                        @test isapprox(
                            rate[end],
                            mri_expected_order;
                            atol = 0.3,
                        )
                    end
                end
            end

            @testset "Implicit" begin
                for (mri_method, mri_expected_order) in mrigark_irk_methods
                    for (fast_method, fast_expected_order) in
                        fast_mrigark_methods
                        for (n, slow_dt) in enumerate(dts)
                            Q = exactsolution(0)
                            fast_dt = slow_dt / ω
                            fastsolver = fast_method(rhs_fast!, Q; dt = fast_dt)
                            solver = mri_method(
                                rhs_slow!,
                                ODETestConvNonLinBE(),
                                fastsolver,
                                Q,
                                dt = slow_dt,
                            )
                            solve!(Q, solver; timeend = finaltime)
                            error[n] = norm(Q - exactsolution(finaltime))
                        end

                        rate = log2.(error[1:(end - 1)] ./ error[2:end])
                        order = mri_expected_order
                        @test isapprox(
                            rate[end],
                            mri_expected_order;
                            atol = 0.3,
                        )
                    end
                end
            end

            @testset "IMEX" begin
                for (method, expected_order) in imex_methods
                    for (n, dt) in enumerate(dts)
                        Q = exactsolution(0)
                        solver = method(
                            rhs_fast!,
                            rhs_slow!,
                            ODETestConvNonLinBE(),
                            Q;
                            dt = dt,
                            t0 = 0.0,
                            split_explicit_implicit = true,
                            variant = NaiveVariant(),
                        )
                        solve!(Q, solver; timeend = finaltime)
                        error[n] = norm(Q - exactsolution(finaltime))
                    end
                    rate = log2.(error[1:(end - 1)] ./ error[2:end])
                    order = expected_order
                    @test isapprox(rate[end], expected_order; atol = 0.3)
                end
            end
        end

        # 3-rate modification of the above test problem
        @testset "3-rate problem" begin
            ω1, ω2, ω3 = 20, 5, 1
            λ1, λ2, λ3 = -20, -5, -1
            β1, β2, β3 = 2, 2, 2

            ξ12 = λ2 / (λ1 + λ2)
            ξ13 = λ3 / (λ1 + λ3)
            ξ23 = λ3 / (λ2 + λ3)

            α12, α13, α23 = 1, 1, 1

            η12 = ((1 - ξ12) / α12) * (λ1 - λ2)
            η13 = 0#((1 - ξ13) / α13) * (λ1 - λ3)
            η23 = ((1 - ξ23) / α23) * (λ2 - λ3)

            η21 = ξ12 * α12 * (λ2 - λ1)
            η31 = 0#ξ13 * α13 * (λ3 - λ1)
            η32 = ξ23 * α23 * (λ3 - λ2)

            Ω = @SMatrix [
                λ1 η12 η13
                η21 λ2 η23
                η31 η32 λ3
            ]

            function rhs1!(dQ, Q, param, t; increment)
                @kernel function knl!(dQ, Q, t, increment)
                    @inbounds begin
                        increment || (dQ .= 0)
                        y1, y2, y3 = Q[1], Q[2], Q[3]
                        g = @SVector [
                            (-β1 + y1^2 - cos(ω1 * t)) / 2y1,
                            (-β2 + y2^2 - cos(ω2 * t)) / 2y2,
                            (-β3 + y3^2 - cos(ω3 * t)) / 2y3,
                        ]
                        dQ[1] += Ω[1, :]' * g - ω1 * sin(ω1 * t) / 2y1
                    end
                end
                event = Event(device(Q))
                event = knl!(device(Q), 1)(
                    dQ,
                    Q,
                    t,
                    increment;
                    ndrange = 1,
                    dependencies = (event,),
                )
                wait(device(Q), event)
            end
            function rhs2!(dQ, Q, param, t; increment)
                @kernel function knl!(dQ, Q, t, increment)
                    @inbounds begin
                        increment || (dQ .= 0)
                        y1, y2, y3 = Q[1], Q[2], Q[3]
                        g = @SVector [
                            (-β1 + y1^2 - cos(ω1 * t)) / 2y1,
                            (-β2 + y2^2 - cos(ω2 * t)) / 2y2,
                            (-β3 + y3^2 - cos(ω3 * t)) / 2y3,
                        ]
                        dQ[2] += Ω[2, :]' * g - ω2 * sin(ω2 * t) / 2y2
                    end
                end
                event = Event(device(Q))
                event = knl!(device(Q), 1)(
                    dQ,
                    Q,
                    t,
                    increment;
                    ndrange = 1,
                    dependencies = (event,),
                )
                wait(device(Q), event)
            end
            function rhs3!(dQ, Q, param, t; increment)
                @kernel function knl!(dQ, Q, t, increment)
                    @inbounds begin
                        increment || (dQ .= 0)
                        y1, y2, y3 = Q[1], Q[2], Q[3]
                        g = @SVector [
                            (-β1 + y1^2 - cos(ω1 * t)) / 2y1,
                            (-β2 + y2^2 - cos(ω2 * t)) / 2y2,
                            (-β3 + y3^2 - cos(ω3 * t)) / 2y3,
                        ]
                        dQ[3] += Ω[3, :]' * g - ω3 * sin(ω3 * t) / 2y3
                    end
                end
                event = Event(device(Q))
                event = knl!(device(Q), 1)(
                    dQ,
                    Q,
                    t,
                    increment;
                    ndrange = 1,
                    dependencies = (event,),
                )
                wait(device(Q), event)
            end
            struct ODETestConvNonLinBE3Rate <: AbstractBackwardEulerSolver end
            ODESolvers.Δt_is_adjustable(::ODETestConvNonLinBE3Rate) = true
            function (::ODETestConvNonLinBE3Rate)(Q, Qhat, α, p, t)
                @kernel function knl!(Q, Qhat, α, p, t)
                    @inbounds begin
                        # Slower RHS has zero tendency of yf so just copy Qhat
                        Q[1] = y1 = Qhat[1]
                        Q[2] = y2 = Qhat[2]

                        g = @SVector [
                            (-β1 + y1^2 - cos(ω1 * t)) / 2y1,
                            (-β2 + y2^2 - cos(ω2 * t)) / 2y2,
                        ]

                        # solves: Q = Qhat + α * rhs_slow(Q, t)
                        # (simplifies to a quadratic equation)
                        a = 2 - α * Ω[3, 3]
                        b =
                            -2 *
                            (Qhat[3] + α * Ω[3, 1] * g[1] + α * Ω[3, 2] * g[2])
                        c = α * (Ω[3, 3] * (β3 + cos(t)) + sin(t))
                        Q[3] = (-b + sqrt(b^2 - 4 * a * c)) / (2a)
                    end
                end
                event = Event(device(Q))
                event = knl!(device(Q), 1)(
                    Q,
                    Qhat,
                    α,
                    p,
                    t;
                    ndrange = 1,
                    dependencies = (event,),
                )
                wait(device(Q), event)
            end

            exactsolution(t) = ArrayType([
                sqrt(β1 + cos(ω1 * t)),
                sqrt(β2 + cos(ω2 * t)),
                sqrt(β3 + cos(ω3 * t)),
            ])

            finaltime = 1
            dts = [2.0^(-k) for k in 1:7]
            error = similar(dts)
            @testset "Explicit" begin
                for (slow_method, slow_order) in mrigark_erk_methods
                    for (mid_method, mid_order) in mrigark_erk_methods
                        for (fast_method, fast_order) in fast_mrigark_methods
                            for (n, slow_dt) in enumerate(dts)
                                Q = exactsolution(0)
                                mid_dt = slow_dt / ω2
                                fast_dt = slow_dt / ω1
                                fastsolver = fast_method(rhs1!, Q; dt = fast_dt)
                                midsolver = mid_method(
                                    rhs2!,
                                    fastsolver,
                                    Q,
                                    dt = mid_dt,
                                )
                                slowsolver = slow_method(
                                    rhs3!,
                                    midsolver,
                                    Q,
                                    dt = slow_dt,
                                )
                                solve!(Q, slowsolver; timeend = finaltime)
                                error[n] = norm(Q - exactsolution(finaltime))
                            end

                            rate = log2.(error[1:(end - 1)] ./ error[2:end])
                            expected_order =
                                min(slow_order, mid_order, fast_order)
                            @test isapprox(
                                rate[end],
                                expected_order;
                                atol = 0.3,
                            )
                        end
                    end
                end
            end
            @testset "Implicit-Explicit" begin
                for (slow_method, slow_order) in mrigark_irk_methods
                    for (mid_method, mid_order) in mrigark_erk_methods
                        for (fast_method, fast_order) in fast_mrigark_methods
                            for (n, slow_dt) in enumerate(dts)
                                Q = exactsolution(0)
                                mid_dt = slow_dt / ω2
                                fast_dt = slow_dt / ω1
                                fastsolver = fast_method(rhs1!, Q; dt = fast_dt)
                                midsolver = mid_method(
                                    rhs2!,
                                    fastsolver,
                                    Q,
                                    dt = mid_dt,
                                )
                                slowsolver = slow_method(
                                    rhs3!,
                                    ODETestConvNonLinBE3Rate(),
                                    midsolver,
                                    Q,
                                    dt = slow_dt,
                                )
                                solve!(Q, slowsolver; timeend = finaltime)
                                error[n] = norm(Q - exactsolution(finaltime))
                            end

                            rate = log2.(error[1:(end - 1)] ./ error[2:end])
                            expected_order =
                                min(slow_order, mid_order, fast_order)
                            @test isapprox(
                                rate[end],
                                expected_order;
                                atol = 0.4,
                            )
                        end
                    end
                end
            end
        end
    end
end


function sample_1_trajectory_from_joint_smoothing_CIR_logweights(δ, γ, σ, Λ_of_t, logwms_of_t, θ_of_t, Λ_pred_of_t, logwms_pred_of_t, θ_pred_of_t, data)

    ntimes = length(Λ_of_t)
    times = Λ_of_t |> keys |> collect |> sort
    sample_trajectory = zeros(ntimes)

    xT = sample_from_Gamma_mixture(δ, θ_of_t[times[end]], Λ_of_t[times[end]], exp.(logwms_of_t[times[end]]))
    sample_trajectory[end] = xT
    for i in (ntimes-1):-1:1
        # println("i=$i, xp1=$(sample_trajectory[i+1])")
        ti = times[i]
        tip1 = times[i+1]
        Δt = tip1-ti
        # predictive_dens_ip1 = DualOptimalFiltering.create_Gamma_mixture_density(δ, θ_pred_of_t[tip1], Λ_pred_of_t[tip1], exp.(logwms_pred_of_t[tip1]))(sample_trajectory[i+1])
        predictive_dens_ip1 = compute_normalisation_constant_joint_smoothing_CIR(sample_trajectory[i+1], θ_of_t[ti], θ_primeΔ(Δt, γ, σ), exp.(logwms_of_t[ti]), Λ_of_t[ti], Δt, δ, γ, σ)
        sample_trajectory[i] = backward_sampling_CIR_logw(sample_trajectory[i+1], logwms_of_t[ti], Λ_of_t[ti], predictive_dens_ip1, Δt, δ, γ, σ, θ_of_t[ti])
    end
    return sample_trajectory
end

function sample_1_trajectory_from_joint_smoothing_CIR_precompute(δ, γ, σ, Λ_of_t, wms_of_t, θ_of_t, Λ_pred_of_t, wms_pred_of_t, θ_pred_of_t, data)

    ntimes = length(Λ_of_t)
    times = Λ_of_t |> keys |> collect |> sort
    sample_trajectory = zeros(ntimes)
    mmax = maximum(maximum.(values(Λ_of_t)))

    log_pochammer_precomputed_terms = Array{Float64, 2}(undef, mmax+1, 10^5)#Might need to be updated for larger datasets

    last_k = 0

    xT = sample_from_Gamma_mixture(δ, θ_of_t[times[end]], Λ_of_t[times[end]], wms_of_t[times[end]])
    sample_trajectory[end] = xT
    for i in (ntimes-1):-1:1
        # println("i=$i, xp1=$(sample_trajectory[i+1])")
        ti = times[i]
        tip1 = times[i+1]
        Δt = tip1-ti
        predictive_dens_ip1, last_k = compute_normalisation_constant_adaptive_precomputation(sample_trajectory[i+1], θ_of_t[ti], θ_primeΔ(Δt, γ, σ), wms_of_t[ti], Λ_of_t[ti], Δt, δ, γ, σ, log_pochammer_precomputed_terms, last_k, mmax)
        sample_trajectory[i] = backward_sampling_CIR_precomputed(sample_trajectory[i+1], wms_of_t[ti], Λ_of_t[ti], predictive_dens_ip1, Δt, δ, γ, σ, θ_of_t[ti], log_pochammer_precomputed_terms)
    end
    return sample_trajectory
end

function compute_normalisation_constant_joint_smoothing_CIR(xip1, θ, θ_primeΔt, filtering_weights, filtering_indices, Δt, δ, γ, σ)
    # Careful that this can fail if the increments are not monotonous, i.e. Δs=0 for some values, then starts increasing. Could add a lower limit on k
    s = 0
    Δs = 1
    k=0
    while k <= 10^2 || (Δs > 0 && k <= 10^5)
        xx = pdf(Gamma(δ/2+k, 1/θ_primeΔt*exp(-2*γ*Δt)), xip1)
        Δs = 0
        for idx in eachindex(filtering_indices)
            m = filtering_indices[idx]
            Δs += xx * filtering_weights[idx] * μmθk(k, m, θ, δ, θ_primeΔt)
        end
        s += Δs
        k += 1
        # println("k=$k, Δs=$Δs")
    end
    return s
end

function compute_normalisation_constant_adaptive_precomputation(xip1, θ, θ_primeΔt, filtering_weights, filtering_indices, Δt, δ, γ, σ, precomputed_terms, last_k, mmax)
    # Careful that this can fail if the increments are not monotonous, i.e. Δs=0 for some values, then starts increasing. Added a lower limit on k
    s = 0
    Δs_rel = 1
    k=0
    while s <= 0 || k <= 10^2 || (Δs_rel > 0 && k <= 10^5)

        if k > last_k
            last_k = k
            precompute_log_pochammer_for_logμmθk!(precomputed_terms, δ, mmax, k; mmin = 0, kmin = last_k)
        end

        xx = pdf(Gamma(δ/2+k, 1/θ_primeΔt*exp(-2*γ*Δt)), xip1)
        Δs = 0
        for idx in eachindex(filtering_indices)
            m = filtering_indices[idx]
            Δs += xx * filtering_weights[idx] * exp(logμmθk5(k, m, θ, δ, θ_primeΔt, precomputed_terms))
        end
        s += Δs
        Δs_rel = Δs/s
        k += 1
    end
    # @info "k=$k, Δs_rel=$Δs_rel"
    return s, last_k
end

function precompute_log_pochammer_for_logμmθk!(precomputed_terms, δ, mmax, kmax; mmin = 0, kmin = 0)
    for m in mmin:mmax
        for k in kmin:kmax
            precomputed_terms[m+1, k+1] = log_pochammer_rec(δ/2+m, k)
        end
    end
end

log_pochammer_precomputed(m, k, precomputed_terms::Array{Float64,2}) = precomputed_terms[m+1,k+1]

function logμmθk5(k, m, θ, δ, θ_primeΔt, precomputed_terms)
    # precomputed
    α = δ/2+m
    β = θ
    λ = θ_primeΔt
    if k==0
        return α*(log(β)-log(λ+β))
    else
        # return sum(log(α + i) for i in 0:(k-1)) - (α+k)*log(λ+β) + k*log(λ) + α*log(β) - lgamma_local(k+1)
        return log_pochammer_precomputed(m, k, precomputed_terms) - (α+k)*log(λ+β) + k*log(λ) + α*log(β) - lgamma_local(k+1)
    end
end

function backward_sampling_CIR_logw(xip1, log_filtering_weights, filtering_indices, νip1_xp1, Δt, δ, γ, σ, θ)
    θ_primeΔt = θ_primeΔ(Δt, γ, σ)
    U = rand()

    # pred_dens_val::Float64 = predictive_dens_ip1(xip1)

    κ, M = select_κM_logw(xip1, θ, θ_primeΔt, U, log_filtering_weights, filtering_indices, Δt, δ, γ, σ, νip1_xp1)

    return rand(Gamma(δ/2+κ+M, 1/(θ+θ_primeΔt)))

end

function backward_sampling_CIR_precomputed(xip1, filtering_weights, filtering_indices, νip1_xp1, Δt, δ, γ, σ, θ, precomputed_terms)
    θ_primeΔt = θ_primeΔ(Δt, γ, σ)
    U = rand()

    # pred_dens_val::Float64 = predictive_dens_ip1(xip1)

    κ, M = select_κM_precomputed(xip1, θ, θ_primeΔt, U, filtering_weights, filtering_indices, Δt, δ, γ, σ, νip1_xp1, precomputed_terms)

    return rand(Gamma(δ/2+κ+M, 1/(θ + θ_primeΔt)))

end

function θ_primeΔ(Δt, γ, σ)
    γ/σ^2 / (exp(2*γ*Δt)-1)
end

function select_κM_logw(xip1, θ, θ_primeΔt, U, log_filtering_weights, filtering_indices, Δt, δ, γ, σ, pred_dens_val)
    k = 0
    logs = -Inf
    while k < 10^5
        logw_tilde_k = logpdf(Gamma(δ/2+k, 1/θ_primeΔt*exp(-2*γ*Δt)), xip1) - log(pred_dens_val)# Careful about the parametrisation of Gamma
        for idx in eachindex(filtering_indices)
            m = filtering_indices[idx]
            logs = logaddexp(logs, logw_tilde_k + log_filtering_weights[idx] + logμmθk(k, m, θ, δ, θ_primeΔt))
            # println("k=$k, m=$m, s=$(exp(logs)), logs=$(logs), U=$U, xip1=$xip1")
            # println("w_tilde_k=$(exp(logw_tilde_k)), wm=$(exp(log_filtering_weights[idx])), mumk=$(μmθk(k, m, θ, δ, θ_primeΔt)), logmumk=$(logμmθk(k, m, θ, δ, θ_primeΔt))")
            # println("mumk=$(μmθk(k, m, θ, δ, θ_primeΔt)), with parameters $((k, m, θ, δ, θ_primeΔt)), max at $(θ/(θ+θ_primeΔt)*(δ/2+m)/(1-θ/(θ+θ_primeΔt)))")
            if exp(logs) >= U
                return k, m
            end
        end
        k += 1
    end
    error("did not manage to select k, weights do not sum to 1")
end

function select_κM_precomputed(xip1, θ, θ_primeΔt, U, filtering_weights, filtering_indices, Δt, δ, γ, σ, pred_dens_val, precomputed_terms)
    k = 0
    s = 0
    while k < 10^5
        w_tilde_k = pdf(Gamma(δ/2+k, 1/θ_primeΔt*exp(-2*γ*Δt)), xip1)/pred_dens_val# Careful about the parametrisation of Gamma
        for idx in eachindex(filtering_indices)
            m = filtering_indices[idx]
            s += w_tilde_k * filtering_weights[idx] * exp(logμmθk5(k, m, θ, δ, θ_primeΔt, precomputed_terms))
            if s >= U
                return k, m
            end
        end
        k += 1
    end
    error("did not manage to select k, weights do not sum to 1")
end

function μmθk(k, m, θ, δ, θ_primeΔt)
    pdf(NegativeBinomial(δ/2 + m, θ/(θ_primeΔt+θ)), k)
end

log_pochammer_mem(a, b) = log_pochammer_rec(a,b)

# function logμmθk4(k, m, θ, δ, θ_primeΔt)
#     # Note that there is some room for optimisation by pre-computing the lgamma(α+k) and using lgamma(α+1) = lgamma(α) + ln(z)
#     α = δ/2+m
#     β = θ
#     λ = θ_primeΔt
#     if k==0
#         return α*(log(β)-log(λ+β))
#     else
#         # return sum(log(α + i) for i in 0:(k-1)) - (α+k)*log(λ+β) + k*log(λ) + α*log(β) - lgamma_local(k+1)
#         return log_pochammer_mem(α, k) - (α+k)*log(λ+β) + k*log(λ) + α*log(β) - lgamma_local(k+1)
#     end
# end

function logμmθk(k, m, θ, δ, θ_primeΔt)
    logpdf(NegativeBinomial(δ/2 + m, θ/(θ_primeΔt+θ)), k)
end


"""
    struct GadgetCylinder{T} <: AbstractGadgetGeometry
        pos_start::Vector{T}
        pos_end::Vector{T}
        radius::T
    end

Defines a cylinder by two end points `pos_start` and `pos_start` as well as a `radius`.
To be used for [`read_particles_in_geometry`](@ref)
"""
struct GadgetCylinder{T} <: AbstractGadgetGeometry
    pos_start::Vector{T}
    pos_end::Vector{T}
    radius::T
end

"""
    function get_geometry_center(cylinder::GadgetCylinder)

Returns the center of `cylinder`.
"""
get_geometry_center(cylinder::GadgetCylinder) = 1 // 2 .* cylinder.pos_start .+ cylinder.pos_end

"""
    get_geometry_box_corners(cylinder::GadgetCylinder)

Returns a tuple with the lower left and upper right corner of a box which contains the `cylinder`.
"""
function get_geometry_box_corners(cylinder::GadgetCylinder)
    min.(cylinder.pos_start, cylinder.pos_end) .- cylinder.radius,
    max.(cylinder.pos_start, cylinder.pos_end) .+ cylinder.radius
end

"""
    get_geometry_mask(cylinder::GadgetCylinder, pos::Matrix{T}) where T

Returns the indices of all particles contained in the `cylinder`.
"""
function get_geometry_mask(cylinder::GadgetCylinder, pos::Matrix{T}) where T

    # allocate mask array
    mask = Vector{Bool}(undef, size(pos,2))

    # check every position array if it's contained in the cylinder
    # and update the mask array
    @inbounds for i = 1:size(pos,2)
        mask[i] = check_in_cylinder(pos[:,i],
                                    cylinder.pos_start,
                                    cylinder.pos_end,
                                    cylinder.radius)
    end

    return findall(mask)
end