# Explicit Runge-Kutta solvers
##############################
# (Hairer & Wanner 1992 p.134, p.165-169)

###########################################
# Tableaus for explicit Runge-Kutta methods
###########################################

struct TableauRKExplicit{Name, S, T} <: Tableau{Name, S, T}
    order::(Tuple{Vararg{Int}}) # the order of the methods
    a::Matrix{T}
    # one or several row vectors.  First row is used for the step,
    # second for error calc.
    b::Matrix{T}
    c::Vector{T}
    function TableauRKExplicit{Name, S, T}(order,a,b,c) where {Name, S, T}
        @assert isa(S,Integer)
        @assert isa(Name,Symbol)
        @assert c[1]==0
        @assert istril(a)
        @assert S==length(c)==size(a,1)==size(a,2)==size(b,2)
        @assert size(b,1)==length(order)
        # consistency:
        if T<:AbstractFloat
            @assert norm(sum(a,dims=2)-c'',Inf) < 100*eps(T)
        else
            @assert norm(sum(a,dims=2)-c'',Inf) < 100*eps(Float64)
        end
        new{Name, S, T}(order,a,b,c)
    end
end
function TableauRKExplicit(name::Symbol, order::(Tuple{Vararg{Int}}),
                   a::Matrix{T}, b::Matrix{T}, c::Vector{T}) where T
    TableauRKExplicit{name,length(c),T}(order, a, b, c)
end
function TableauRKExplicit(name::Symbol, order::(Tuple{Vararg{Int}}), T::Type,
                   a::Matrix, b::Matrix, c::Vector)
    TableauRKExplicit{name,length(c),T}(order, convert(Matrix{T},a),
                                        convert(Matrix{T},b), convert(Vector{T},c) )
end
conv_field(D,a::Array{T,N}) where {T,N} = convert(Array{D,N}, a)
function Base.convert(::Type{Tnew}, tab::TableauRKExplicit{Name,S,T}) where {Tnew<:Real,Name,S,T}
    # Converts the tableau coefficients to the new type Tnew
    newflds = ()
    for n in fieldnames(typeof(tab))
        fld = getfield(tab,n)
        if eltype(fld)==T
            newflds = tuple(newflds..., conv_field(Tnew, fld))
        else
            newflds = tuple(newflds..., fld)
        end
    end
    TableauRKExplicit{Name,S,Tnew}(newflds...) # TODO: could this be done more generically in a type-stable way?
end

isexplicit(b::TableauRKExplicit) = istril(b.a) # Test whether it's an explicit method
isadaptive(b::TableauRKExplicit) = size(b.b, 1)==2

# First same as last.  Means ks[:,end]=ks_nextstep[:,1], c.f. H&W p.167
isFSAL(btab::TableauRKExplicit) = btab.a[end,:]==btab.b[1,:] && btab.c[end]==1 # the latter is not needed really

## Tableaus for explicit RK methods
# Fixed step:
const bt_feuler = TableauRKExplicit(:feuler,(1,), Rational{Int64},
                                    zeros(Int,1,1),
                                    [1][:,:],
                                    [0]
                                    )
const bt_midpoint = TableauRKExplicit(:midpoint,(2,), Rational{Int64},
                                      [0  0
                                       1//2  0],
                                      [0  1],
                                      [0, 1//2]
                                      )
const bt_heun = TableauRKExplicit(:heun,(2,), Rational{Int64},
                                  [0  0
                                   1  0],
                                  [1//2  1//2],
                                  [0, 1])

const bt_rk4 = TableauRKExplicit(:rk4,(4,),Rational{Int64},
                                 [0    0    0 0
                                  1//2 0    0 0
                                  0    1//2 0 0
                                  0    0    1 0],
                                 [1//6  1//3  1//3  1//6],
                                 [0, 1//2, 1//2, 1])

# Adaptive step:
# Heun Euler https://en.wikipedia.org/wiki/Runge–Kutta_methods
const bt_rk21 = TableauRKExplicit(:heun_euler,(2,1), Rational{Int64},
                                  [0     0
                                   1     0],
                                  [1//2  1//2
                                   1     0],
                                  [0,    1])

# Bogacki–Shampine coefficients
const bt_rk23 = TableauRKExplicit(:bogacki_shampine,(2,3), Rational{Int64},
                                  [0           0      0      0
                                   1//2         0      0      0
                                   0         3//4      0      0
                                   2//9       1//3     4//9     0],
                                  [7//24 1//4 1//3 1//8
                                   2//9 1//3 4//9 0],
                                  [0, 1//2, 3//4, 1]
                         )

# Fehlberg https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta%E2%80%93Fehlberg_method
const bt_rk45 = TableauRKExplicit(:fehlberg,(4,5),Rational{Int64},
                          [  0           0           0            0         0     0
                             1//4        0           0            0         0     0
                             3//32       9//32       0            0         0     0
                          1932//2197 -7200//2197  7296//2197      0         0     0
                           439//216     -8        3680//513    -845//4104   0     0
                            -8//27       2       -3544//2565   1859//4104 -11//40 0 ],
                           [25//216      0        1408//2565   2197//4104  -1//5  0
                            16//135      0        6656//12825 28561//56430 -9//50 2//55],
                            [0,          1//4,       3//8,       12//13,    1,    1//2])

# Dormand-Prince https://en.wikipedia.org/wiki/Dormand%E2%80%93Prince_method
const bt_dopri5 = TableauRKExplicit(:dopri, (5,4), Rational{Int64},
                     [0           0            0                   0            0      0 0
                      1//5        0            0                   0            0      0 0
                      3//40       9//40        0                   0            0      0 0
                      44//45      -56//15      32//9               0            0      0 0
                      19372//6561 -25360//2187 64448//6561 -212//729            0      0 0
                      9017//3168  -355//33     46732//5247   49//176 -5103//18656      0 0
                      35//384     0            500//1113    125//192  -2187//6784 11//84 0],
                     [35//384         0     500//1113       125//192      -2187//6784         11//84      0
                      5179//57600     0     7571//16695     393//640     -92097//339200     187//2100     1//40],
                     [0, 1//5, 3//10, 4//5, 8//9, 1, 1]
                     )

# Fehlberg 7(8) coefficients
# Values from pag. 65, Fehlberg, Erwin. "Classical fifth-, sixth-, seventh-, and eighth-order Runge-Kutta formulas with stepsize control".
# National Aeronautics and Space Administration.
const bt_feh78 = TableauRKExplicit(:feh78, (7,8), Rational{Int64},
                            [     0       0      0       0         0          0        0        0      0       0     0 0 0
                                  2//27   0      0       0         0          0        0        0      0       0     0 0 0
                                  1//36   1//12  0       0         0          0        0        0      0       0     0 0 0
                                  1//24   0      1//8    0         0          0        0        0      0       0     0 0 0
                                  5//12   0    -25//16  25//16     0          0        0        0      0       0     0 0 0
                                  1//20   0      0       1//4      1//5       0        0        0      0       0     0 0 0
                                -25//108  0      0     125//108  -65//27    125//54    0        0      0       0     0 0 0
                                 31//300  0      0       0       61//225    -2//9     13//900   0      0       0     0 0 0
                                  2       0      0     -53//6    704//45   -107//9    67//90    3      0       0     0 0 0
                                -91//108  0      0      23//108 -976//135   311//54  -19//60   17//6  -1//12   0     0 0 0
                               2383//4100 0      0    -341//164 4496//1025 -301//82 2133//4100 45//82 45//164 18//41 0 0 0
                                  3//205  0      0       0        0        -6//41     -3//205  -3//41  3//41   6//41 0 0 0
                              -1777//4100 0      0    -341//164 4496//1025 -289//82 2193//4100 51//82 33//164 12//41 0 1 0],
                              [41//840 0 0 0 0 34//105 9//35 9//35 9//280 9//280 41//840 0 0
                               0 0 0 0 0 34//105 9//35 9//35 9//280 9//280 0     41//840 41//840],
                               [0,    2//27, 1//9, 1//6 , 5//12, 1//2 , 5//6 , 1//6 , 2//3 , 1//3 , 1 , 0, 1]
                            )

################################
# Fixed step Runge-Kutta methods
################################

# TODO: iterator method
ode1(fn, y0, tspan; kwargs...) = oderk_fixed(fn, y0, tspan, bt_feuler; kwargs...)
ode2_midpoint(fn, y0, tspan; kwargs...) = oderk_fixed(fn, y0, tspan, bt_midpoint; kwargs...)
ode2_heun(fn, y0, tspan; kwargs...) = oderk_fixed(fn, y0, tspan, bt_heun; kwargs...)
ode4(fn, y0, tspan; kwargs...) = oderk_fixed(fn, y0, tspan, bt_rk4; kwargs...)

function oderk_fixed(fn, y0, tspan, btab::TableauRKExplicit; kwargs...)
    # Non-arrays y0 treat as scalar
    fn_(t, y) = [fn(t, y[1])]
    t,y = oderk_fixed(fn_, [y0], tspan, btab)
    return t, vcat_nosplat(y)
end
function oderk_fixed(fn, y0::AbstractVector, tspan,
                     btab_::TableauRKExplicit{N,S}, alias_u0=false) where {N,S}
    # TODO: instead of AbstractVector use a Holy-trait

    # Needed interface:
    # On components:
    # On y0 container: length, deepcopy, similar, setindex!
    # On time container: getindex, convert. length

    Et, Eyf, Ty, btab = make_consistent_types(fn, y0, tspan, btab_)
    dof = length(y0)

    ys = Vector{Ty}(undef,length(tspan))
    allocate!(ys, y0, dof)
    if alias_u0
        ys[1] = y0
    else
        ys[1] = deepcopy(y0)
    end

    tspan = convert(Vector{Et}, tspan)
    # work arrays:
    ks = Vector{Ty}(undef,S)
    # allocate!(ks, y0, dof) # no need to allocate as fn is not in-place
    ytmp = similar(y0, Eyf, dof)
    for i=1:length(tspan)-1
        dt = tspan[i+1]-tspan[i]
        ys[i+1][:] = ys[i]
        for s=1:S
            calc_next_k!(ks, ytmp, ys[i], s, fn, tspan[i], dt, dof, btab)
            for d=1:dof
                ys[i+1][d] += dt * btab.b[s]*ks[s][d]
            end
        end
    end
    return tspan, ys
end

##############################
# Adaptive Runge-Kutta methods
##############################

ode21(fn, y0, tspan; kwargs...) = oderk_adapt(fn, y0, tspan, bt_rk21; kwargs...)
ode23(fn, y0, tspan; kwargs...) = oderk_adapt(fn, y0, tspan, bt_rk23; kwargs...)
ode45_fe(fn, y0, tspan; kwargs...) = oderk_adapt(fn, y0, tspan, bt_rk45; kwargs...)
ode45_dp(fn, y0, tspan; kwargs...) = oderk_adapt(fn, y0, tspan, bt_dopri5; kwargs...)
# Use Dormand-Prince version of ode45 by default
ode45(fn, y0, tspan; kwargs...) = ode45_dp(fn, y0, tspan; kwargs...)
ode78(fn, y0, tspan; kwargs...) = oderk_adapt(fn, y0, tspan, bt_feh78; kwargs...)

function oderk_adapt(fn, y0, tspan, btab::TableauRKExplicit; kwords...)
    # For y0 which don't support indexing.
    fn_ = (t, y) -> [fn(t, y[1])]
    t,y = oderk_adapt(fn_, [y0], tspan, btab; kwords...)
    return t, vcat_nosplat(y)
end
function oderk_adapt(fn, y0::AbstractVector, tspan, btab_::TableauRKExplicit{N,S};
                     reltol = 1.0e-5, abstol = 1.0e-8,
                     norm=LinearAlgebra.norm,
                     minstep=abs(tspan[end] - tspan[1])/1e18,
                     maxstep=abs(tspan[end] - tspan[1])/2.5,
                     initstep=0,
                     points=:all,
                     alias_u0=false
                     ) where {N,S}
    # Needed interface:
    # On components:
    #  - note that the type of the components might change!
    # On y0 container: length, similar, setindex!
    # On time container: getindex, convert, length

    # For y0 which support indexing.  Currently y0<:AbstractVector but
    # that could be relaxed with a Holy-trait.
    !isadaptive(btab_) && error("Can only use this solver with an adaptive RK Butcher table")
    Et, Eyf, Ty, btab = make_consistent_types(fn, y0, tspan, btab_)
    # parameters
    order = minimum(btab.order)
    timeout_const = 5 # after step reduction do not increase step for
                      # timeout_const steps

    ## Initialization
    dof = length(y0)
    tspan = convert(Vector{Et}, tspan)
    t = tspan[1]
    tstart = tspan[1]
    tend = tspan[end]

    # work arrays:
    if alias_u0
        y = y0                     # y at time t
    else
        y  = similar(y0, Eyf, dof) # y at time t
    end
    y[:]   = y0
    ytrial = similar(y0, Eyf, dof) # trial solution at time t+dt
    yerr   = similar(y0, Eyf, dof) # error of trial solution
    ks = Vector{Ty}(undef,S)
    # allocate!(ks, y0, dof) # no need to allocate as fn is not in-place
    ytmp   = similar(y0, Eyf, dof)

    # output ys
    nsteps_fixed = length(tspan) # these are always output
    ys = Vector{Ty}(undef,nsteps_fixed)
    allocate!(ys, y0, dof)
    ys[1] = y0

    # Option points determines where solution is returned:
    if points == :all
        tspan_fixed = tspan
        tspan = Et[tstart]
        iter_fixed = 2 # index into tspan_fixed
        sizehint!(tspan, nsteps_fixed)
    elseif points != :specified
        error("Unrecognized option points==$points")
    end
    # Time
    dt, tdir, ks[1] = hinit(fn, y, tstart, tend, order, reltol, abstol) # sets ks[1]=f0
    if initstep!=0
        dt = sign(initstep)==tdir ? initstep : error("initstep has wrong sign.")
    end
    # Diagnostics
    dts = Et[]
    errs = Float64[]
    steps = [0,0]  # [accepted, rejected]

    ## Integration loop
    islaststep = abs(t+dt-tend)<=eps(tend) ? true : false
    timeout = 0 # for step-control
    iter = 2 # the index into tspan and ys
    while true
        # do one step (assumes ks[1]==f0)
        rk_embedded_step!(ytrial, yerr, ks, ytmp, y, fn, t, dt, dof, btab)
        # Check error and find a new step size:
        err, newdt, timeout = stepsize_hw92!(dt, tdir, y, ytrial, yerr, order, timeout,
                                            dof, abstol, reltol, maxstep, norm)

        if err<=1 # accept step
            # diagnostics
            steps[1] +=1
            push!(dts, dt)
            push!(errs, err)

            # Output:
            f0 = ks[1]
            f1 = isFSAL(btab) ? ks[S] : fn(t+dt, ytrial)
            if points==:specified
                # interpolate onto given output points
                while iter-1<nsteps_fixed && (tdir*tspan[iter]<tdir*(t+dt) || islaststep) # output at all new times which are < t+dt
                    hermite_interp!(ys[iter], tspan[iter], t, dt, y, ytrial, f0, f1) # TODO: 3rd order only!
                    iter += 1
                end
            else
                # first interpolate onto given output points
                while iter_fixed-1<nsteps_fixed && tdir*t<tdir*tspan_fixed[iter_fixed]<tdir*(t+dt) # output at all new times which are < t+dt
                    yout = hermite_interp(tspan_fixed[iter_fixed], t, dt, y, ytrial, f0, f1)
                    index_or_push!(ys, iter, yout) # TODO: 3rd order only!
                    push!(tspan, tspan_fixed[iter_fixed])
                    iter_fixed += 1
                    iter += 1
                end
                # but also output every step taken
                index_or_push!(ys, iter, deepcopy(ytrial))
                push!(tspan, t+dt)
                iter += 1
            end
            ks[1] = f1 # load ks[1]==f0 for next step

            # Break if this was the last step:
            islaststep && break

            # Swap bindings of y and ytrial, avoids one copy
            y, ytrial = ytrial, y

            # Update t to the time at the end of current step:
            t += dt
            dt = newdt

            # Hit end point exactly if next step within 1% of end:
            if tdir*(t+dt+dt/100) >= tdir*tend
                dt = tend-t
                islaststep = true # next step is the last, if it succeeds
            end
        elseif abs(newdt)<minstep  # minimum step size reached, break
            println("Warning: dt < minstep.  Stopping.")
            break
        else # redo step with smaller dt
            islaststep = false
            steps[2] +=1
            dt = newdt
            timeout = timeout_const
        end
    end
    return tspan, ys
end

function rk_embedded_step!(ytrial, yerr, ks, ytmp, y, fn, t, dt, dof, btab::TableauRKExplicit{N,S}) where {N,S}
    # Does one embedded R-K step updating ytrial, yerr and ks.
    #
    # Assumes that ks[:,1] is already calculated!
    #
    # Modifies ytrial, yerr, ks, and ytmp

    # Needed interface:
    # On components: arithmetic, zero
    # On y0 container: fill!, setindex!, getindex

    fill!(ytrial, zero(eltype(ytrial)) )
    fill!(yerr, zero(eltype(ytrial)) )
    for d=1:dof
        ytrial[d] += btab.b[1,1]*ks[1][d]
        yerr[d]   += btab.b[2,1]*ks[1][d]
    end
    for s=2:S
        calc_next_k!(ks, ytmp, y, s, fn, t, dt, dof, btab)
        for d=1:dof
            ytrial[d] += btab.b[1,s]*ks[s][d]
            yerr[d]   += btab.b[2,s]*ks[s][d]
        end
    end
    for d=1:dof
        yerr[d]   = dt * (ytrial[d]-yerr[d])
        ytrial[d] = y[d] + dt * ytrial[d]
    end
end

function stepsize_hw92!(dt::T, tdir, x0, xtrial, xerr, order,
                    timeout, dof, abstol, reltol, maxstep, norm) where T
    # Estimates the error and a new step size following Hairer &
    # Wanner 1992, p167 (with some modifications)
    #
    # If timeout>0 no step size increase is allowed, timeout is
    # decremented in here.
    #
    # Returns the error, newdt and the number of timeout-steps
    #
    # The work-array xerr is modified in-place
    #
    # TODO:
    # - allow component-wise reltol and abstol?
    # - allow other norms

    # Needed interface:
    # On components: isoutofdomain
    # On y0 container: norm, get/setindex

    timout_after_nan = 5
    fac = T[0.8, 0.9, 0.25^(1/(order+1)), 0.38^(1/(order+1))][1]
    facmax = T(5) # maximal step size increase. 1.5-5
    facmin = T(1)/T(facmax)  # maximal step size decrease. ?
    maxstep = T(maxstep)
    abstol = T(abstol)
    reltol = T(reltol)

    # in-place calculate xerr./tol
    for d=1:dof
        # if outside of domain (usually NaN) then make step size smaller by maximum
        isoutofdomain(xtrial[d]) && return T(10), dt*facmin, timout_after_nan
        xerr[d] = xerr[d]/(abstol + max(norm(x0[d]), norm(xtrial[d]))*reltol) # Eq 4.10
    end
    err = norm(xerr, 2) # Eq. 4.11
    newdt = min(maxstep, tdir*dt*max(facmin, fac*(1/err)^(T(1/(order+1)))) ) # Eq 4.13 modified
    if timeout>0
        newdt = min(newdt, dt)
        timeout -= 1
    end
    return err, tdir*newdt, timeout
end

function calc_next_k!(ks::Vector, ytmp::Ty, y, s, fn, t, dt, dof, btab) where Ty
    # Calculates the next ks and puts it into ks[s]
    # - ks and ytmp are modified inside this function.

    # Needed interface:
    # On components: +, *
    # On y0 container: setindex!, getindex, fn

    ytmp[:] = y
    for ss=1:s-1, d=1:dof
        ytmp[d] += dt * ks[ss][d] * btab.a[s,ss]
    end
    ks[s] = fn(t + btab.c[s]*dt, ytmp)::Ty
    nothing
end

# Helper functions:
function allocate!(vec::Vector{T}, y0, dof) where T
    # Allocates all vectors inside a Vector{Vector} using the same
    # kind of container as y0 has and element type eltype(eltype(vec)).
    for s=1:length(vec)
        vec[s] = similar(y0, eltype(T), dof)
    end
end
function index_or_push!(vec, i, val)
    # Fills in the vector until there is no space, then uses push!
    # instead.
    if length(vec)>=i
        vec[i] = val
    else
        push!(vec, val)
    end
end
vcat_nosplat(y) =  eltype(y[1])[el[1] for el in y] # Does vcat(y...) without the splatting

function hermite_interp!(y, tquery,t,dt,y0,y1,f0,f1)
    # For dense output see Hairer & Wanner p.190 using Hermite
    # interpolation. Updates y in-place.
    #
    # f_0 = f(x_0 , y_0) , f_1 = f(x_0 + h, y_1 )
    # this is O(3). TODO for higher order.

    theta = (tquery-t)/dt
    for i=1:length(y0)
        y[i] = ((1-theta)*y0[i] + theta*y1[i] + theta*(theta-1) *
                ((1-2*theta)*(y1[i]-y0[i]) + (theta-1)*dt*f0[i] + theta*dt*f1[i]) )
    end
    nothing
end
function hermite_interp(tquery,t,dt,y0,y1,f0,f1)
    # Returns the y instead of in-place
    y = similar(y0)
    hermite_interp!(y,tquery,t,dt,y0,y1,f0,f1)
    return y
end


# Solves sparse regression problems with the SLOPE method of Bogdan et al [Bogdan2014].
# Since the prox step in solving SLOPE is a quadratic program we here use Convex.jl
# to solve it. This should in general be an inferior solution to the FastProxSL1 method
# described in [Bogdan2014]. But it is fun and simple and still performs well. :)

# Beware! This is not yet bug-free...

# The key step in both the proximal gradient and the accelerated proximal gradient algorithms
# to solve SLOPE is to solve the prox problem. According to 2.2 and 2.3 in [Bogdan2014] the prox
# can be solved by a quadratic program which we implement in Convex.jl:
using Convex

# Set up x once and for all so we need not do it on each prox optimization round.
setup_x_and_constraints(p) = begin
  x = Variable(p)
  constraints = Constraint[]
  for i in 1:(p-1)
    constraints += (x[i] >= x[i+1])
  end
  constraints += (x[p] >= 0)
  return x, constraints
end

# The prox problem we solve assumes y are sorted and positive. We keep the permutation vector
# around so we can revert later.
sort_and_normalize(y) = begin
  absy = abs(y)
  return sortperm(absy, rev = true), sign(y), absy
end

# Go back from a sortperm to orig order.
function revert_sortperm(y, yperm)
  p = length(y)
  yr = zeros(p)
  for i in 1:p
    yi = yperm[i]
    yr[yi] = y[i]
  end
  yr
end

solve_prox_problem_with_convexjl(yproxvals, lambdas, x, constraints) = begin
  yperm, sngy, absy = sort_and_normalize(yproxvals)
  ynormalized = absy[yperm]

  proxproblem = minimize(0.5 * norm(ynormalized .- x, 2)^2 + sum(lambdas .* x), constraints)
  solve!(proxproblem)
  if proxproblem.status != :Optimal
    throw(string(proxproblem))
  end

  # Reorder and put back the signs...
  sngy .* revert_sortperm(x.value, yperm)
end

# The default choice for lambdas is Benjamini-Hochberg sequence where q is the FDR:
using Distributions
NormalDistr = Normal()
bh_lambdafunc(i, p, q = 0.10) = quantile(NormalDistr, 1 - i*q/(2*p))

# Page 7 of the latest Candes2015 paper on SLOPE states that a very simple lambda sequence can be used:
# But we found they need to be smaller for Convex.jl and FastProxSL1 prox solvers!?
candes2015_lambdafunc(k, p, q = 0.10) = sqrt(2*log(p/k))

# Use a given lambdas sequence or set it up based on the simple sequence from
# Candes2015 paper if not given.
function check_and_setup_lambdasequence(p, lambdas = nothing, q = 0.05)
  if lambdas != nothing
    @assert p == length(lambdas)
  else
    lambdas = map(k -> bh_lambdafunc(k, p, q), 1:p)
  end
  return lambdas
end

# The simplest SLOPE solver is just a proximal gradient. It assumes we have a way to solve the prox
# problem.
function solve_SLOPE_w_proximal_gradient(X::Matrix{Float64}, y::Vector{Float64}, proxsolver::Function;
  stepsizefunc = nothing,       # Default is to use 1/norm(X)^2, so a fixed scheme...
  lambdas = nothing,            # Default is to use the simple sequence below from Candes2015 paper...
  maxiterations = int(5e4), stopval = 1e-3, verbose = true)

  # Assert valid inputs, setup and precalc
  n, p = size(X)
  @assert n == length(y)
  lambdas = check_and_setup_lambdasequence(p, lambdas)
  Xprim = X'        # Precalc for speed
  bnew = bprev = randn(p)  # Initial guess is random
  if stepsizefunc == nothing
    limitval = (2 / norm(X)^2) / 2 # Divide by 2 to ensure it is less than the limitval
    stepsizefunc = (i,b) -> limitval
  end

  # Now iterate for maxiterations steps or until stopval reached.
  for k in 1:maxiterations
    tk = stepsizefunc(k, bprev)
    yproxvals = bprev .- tk * Xprim * (X * bprev .- y)
    bnew = proxsolver(yproxvals, lambdas)
    delta = norm(bnew .- bprev)
    println("k = ", k, ", delta = ", delta)
    if delta < stopval
      return bnew
    end
    bprev = bnew
  end
  return bnew
end

function solve_SLOPE_w_convexjl(X::Matrix{Float64}, y::Vector{Float64}; options...)
  p = size(X, 2)
  x, constraints = setup_x_and_constraints(p)
  proxsolver(yvals, lambdas) = solve_prox_problem_with_convexjl(yvals, lambdas, x, constraints)
  solve_SLOPE_w_proximal_gradient(X, y, proxsolver; options...)
end

# Now generate a random, sparse regression problem
function rand_sparse_problem(p; a = iceil(log(p)), n = int(p/2), amplitude = 10.0, sigma = amplitude/100, shuffle = true)
  beta = vcat(amplitude*ones(a)+randn(a), zeros(p-a))
  indices = collect(1:p)
  if shuffle
    shuffle!(indices)
    beta = beta[indices]
  end
  X = randn(n, p)
  errors = sigma*randn(n)
  y = X * beta .+ errors
  return beta, X, y, indices, n, p, a, sigma, amplitude, errors
end

beta, X, y, indices, n, p, a, sigma, amplitude, errors = rand_sparse_problem(100; shuffle = false);

# And let SLOPE loose on it...
#@time betahat = solve_SLOPE_w_convexjl(X, y)
#println(beta[1:a])
#println(betahat[1:a])
#println(norm(beta .- betahat))

# There are bugs though...

# Lets implement the stack-based FastProx SL1 alg of Bogdan et al to see if that helps...
# This is Alg 4 in the Bogdan2014 paper.
function fast_prox_sl1(yproxvals, lambdas)
  p = length(yproxvals)
  tuples = Any[]
  t = 0

  # Find optimal group levels
  for k in 1:p
    t += 1
    i, j, s = k, k, (yproxvals[k] - lambdas[k])
    w = abs(s)
    push!(tuples, (i, j, s, w))
    while (t > 1) && (tuples[t-1][4] .<= w)
      itm1, jtm1, stm1, wtm1 = tuples[t-1]
      it, jt, st, wt = tuples[t]
      wnew = (jtm1 - itm1 + 1)/(jt - itm1 + 1) * stm1 .+ (jt - it + 1)/(jt - itm1 + 1) * st
      tuples[t-1] = (itm1, jt, stm1 .+ st, abs(wnew))
      pop!(tuples) # Delete last entry
      t -= 1
    end
  end

  # Set entries in x for each block
  x = zeros(p)
  for l in 1:t
    for k in tuples[l][1]:tuples[l][2]
      x[k] = tuples[l][4]
    end
  end
  x
end

solve_prox_problem_with_fastproxsl1(yproxvals, lambdas) = begin
  yperm, sngy, absy = sort_and_normalize(yproxvals)
  ynormalized = absy[yperm]

  xvals = fast_prox_sl1(ynormalized, lambdas)

  # Reorder and put back the signs...
  sngy .* revert_sortperm(xvals, yperm)
end

function solve_SLOPE_w_fastproxsl1(X::Matrix{Float64}, y::Vector{Float64}; options...)
  solve_SLOPE_w_proximal_gradient(X, y, solve_prox_problem_with_fastproxsl1; options...)
end

# And let SLOPE loose on it...
@time betahat = solve_SLOPE_w_fastproxsl1(X, y)
println(betahat)
println(beta[1:a])
println(betahat[1:a])
println(length(find(abs(betahat) .> 0)))
println(norm(beta .- betahat))


"""
    smoptimize(f::Function, search_range::Array{Tuple{Float64,Float64},1}; options=Options())

Optimize the function `f` in the range `search_range` using a Radial Basis Function based surrogate model.
"""
function smoptimize(f::Function, search_range::Array{Tuple{Float64,Float64},1}; options::Options=Options())

    #Load some option values
    @unpack num_start_samples, sampling_plan_opt_gens,
            iterations, trace, create_final_surrogate = options

    reset_timer!()

    #Create sampling plan
    @timeit_debug "lhc" lhc_plan = scaled_LHC_sampling_plan(search_range,num_start_samples,sampling_plan_opt_gens;trace=trace)

    #Evaluate sampling plan
    @timeit_debug "f_eval" lhc_samples = f_opt_eval(f,lhc_plan;trace=trace)

    #Initialize variables to be returned
    sm_interpolant = nothing
    infill_type = Array{Symbol,1}(undef,0)
    infill_prediction = Array{Float64,1}(undef,0)
    optres = nothing
    infill_plan = Array{Float64,2}(undef,size(lhc_plan,1),0)
    infill_sample = Array{Float64,2}(undef,1,0)

    #Run the entire optimization iterations number of times
    for i = 1:iterations

        #Create the optimized Radial Basis Function interpolant
        samples_all = [lhc_samples infill_sample]
        plan_all = [lhc_plan infill_plan]
        @timeit_debug "sm_opt" begin
            sm_interpolant, optres = surrogate_model(plan_all, samples_all; options=options)
        end

        #Points to add to the sampling plan to improve the interpolant
        @timeit_debug "model_infill" begin
            infill_plan_new, infill_type_new, infill_prediction_new, options  = model_infill(search_range,plan_all,
                                                                        samples_all,sm_interpolant;options=options)
        end

        #Evaluate the new infill points
        print_iteration(trace,i,iterations)
        @timeit_debug "f_eval" infill_sample_new = f_opt_eval(f,infill_plan_new,samples_all;trace=trace)

        #Add infill points
        infill_plan = [infill_plan infill_plan_new]
        infill_sample = [infill_sample infill_sample_new]
        infill_type = [infill_type; infill_type_new]
        infill_prediction = [infill_prediction; infill_prediction_new]

    end

    if create_final_surrogate
        #Create the optimized Radial Basis Function interpolant
        samples_all = [lhc_samples infill_sample]
        plan_all = [lhc_plan infill_plan]
        sm_interpolant, optres = surrogate_model(plan_all, samples_all; options=options)
    end

    timeit_debug_enabled() && print_timer()

    return SurrogateResult( lhc_samples, lhc_plan, sm_interpolant,
                            optres, infill_sample, infill_type,
                            infill_plan, infill_prediction,options)
end


Iterators.map(x -> x^2, 1:3) |> display
println()

collect(Iterators.map(x -> x^2, 1:3)) |> display

mapreduce(isodd, +, 1:5) |> display


#=

  Set covering in Julia.

  Example from Steven Skiena, The Stony Brook Algorithm Repository
  http://www.cs.sunysb.edu/~algorith/files/set-cover.shtml
  """
  Input Description: A set of subsets S_1, ..., S_m of the
  universal set U = {1,...,n}.

  Problem: What is the smallest subset of subsets T subset S such
  that \cup_{t_i in T} t_i = U?
  """
  Data is from the pictures INPUT/OUTPUT.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function set_covering_skiena(problem,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   # "all_solutions"=> all_solutions,
                                                            "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n, m = size(problem)

    @variable(model, x[1:n], Bin)
    @variable(model, 1 <= z <= n, Int)

    for b in eachrow(transpose(problem))
        @constraint(model, sum(x.*b) >= 1)
    end
    @constraint(model,z == sum(x))

    @objective(model,Min,z)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                z_val = convert.(Integer,JuMP.value.(z; result=sol))
                println("x:$x_val")
                println("z:$z_val")
                println()

            end
        end
    else
        println("status:$status")
    end

    return status
end

set_covering_skiena_problems = Dict(
:1 => resize_matrix([[1,1,0,0,0,0,0,0,0,0,0,0],
                    [0,1,0,0,0,0,0,1,0,0,0,0],
                    [0,0,0,0,1,1,0,0,0,0,0,0],
                    [0,0,0,0,0,1,1,0,0,1,1,0],
                    [0,0,0,0,0,0,0,0,1,1,0,0],
                    [1,1,1,0,1,0,0,0,1,1,1,0],
                    [0,0,1,1,0,0,1,1,0,0,1,1]])
)

@time set_covering_skiena(set_covering_skiena_problems[:1])


using PathwiseSampling
using AbstractGPs
using SparseGPs
using Distributions
using LinearAlgebra

k = 3 * (SqExponentialKernel() ∘ ScaleTransform(3.0))
gp = GP(k)
x = rand(9)
fx = gp(x, 0.01)
y = rand(fx)
C = rand(length(x), length(x))
ap = approx_posterior(SVGP(), fx, MvNormal(C'C))

n_feats = Int(1e6)
n_samples = 100

ϕ = prior_basis(k, n_feats)
k_approx(ϕ, x, y) = ϕ(x)'ϕ(y)
k_matrix_approx(ϕ, X, Y) = hcat(ϕ.(RowVecs(X))...)'hcat(ϕ.(RowVecs(Y))...)
# k_matrix_approx(ϕ, X, Y) = reshape(hcat(k_approx.(ϕ, X, Y')...), (length(X), length(Y)))
k_matrix_approx(ϕ, X) = k_matrix_approx(ϕ, X, X)

km_true = kernelmatrix(k, x)
km_approx = k_matrix_approx(ϕ, x) # TODO: too slow?

using Test
@test km_true ≈ km_approx rtol=1e-3 atol=1e-3

n_feats = 10
ϕ = prior_basis(k, n_feats)

f, w = sample_prior_functions(ϕ, n_feats, n_samples)

f_post = sample_posterior_functions(ap, n_feats, n_samples)

using Plots
x_plot = sort(rand(100))
plot(x_plot, hcat(f.(x_plot)...)', color=:blue, linealpha=0.1, label="")
plot!(x_plot, gp, label="True prior")

plot(x_plot, hcat(f_post.(x_plot)...)', label="", color=:red, linealpha=0.2)
plot!(x_plot, ap, color=:green, label="True posterior")

# # 2 dimensional input
x2 = RowVecs(rand(9, 2))
ψ = prior_basis(k, 10, 2)

p = hcat(ψ.(x2)...)

p*p'


using SparseArrays, Test, JuMP, Polyhedra

function _hypercube(n, base_name)
    model = Model()
    @variable(model, [1:n], lower_bound = 0, upper_bound = 1, base_name = base_name)
    h = hrep(model)
    @test dimension_names(h) == ["$base_name[$i]" for i in 1:n]
    return h
end
function hypercube_name()
    h = _hypercube(1, "x") * _hypercube(2, "y")
    @test dimension_names(h) == ["x[1]", "y[1]", "y[2]"]
    @test dimension_names(spzeros(3, 4) \ h) === nothing
end
@testset "Hypercube name" begin
    hypercube_name()
end


function loadoctopodes(filename::AbstractString)
    lines = readlines(filename)
    transpose(reshape(
        [parse(Int, c) for line in lines for c in line],
        (length(lines), length(first(lines)))
    ))
end

if basename(pwd()) == "aoc"
    cd("2021/11")
end

example = loadoctopodes("example.txt")
input = loadoctopodes("input.txt")

function neighbourcoords(A, c)
    (
        c + off
        for off in (CartesianIndices((-1:1, -1:1)))
        if !iszero(off) && (c + off) in CartesianIndices(A)
    )
end

function part1(octos, n::Integer = 1)
    octos = copy(octos)
    flashes = 0
    while n > 0
        # First, the energy level of each octopus increases by 1.
        octos .+= 1

        # Then, any octopus with an energy level greater than 9 flashes.
        # This increases the energy level of all adjacent octopuses by 1,
        # including octopuses that are diagonally adjacent. If this causes
        # an octopus to have an energy level greater than 9, it also flashes.
        # This process continues as long as new octopuses keep having their
        # energy level increased beyond 9. (An octopus can only flash at most
        # once per step.)
        canflash = ones(Bool, size(octos))
        while any(>(9), octos[canflash])
            coord = findfirst(i -> canflash[i] && octos[i] > 9, CartesianIndices(octos))
            octos[collect(neighbourcoords(octos, coord))] .+= 1
            canflash[coord] = false
            flashes += 1
        end

        # Finally, any octopus that flashed during this step has its energy
        # level set to 0, as it used all of its energy to flash.
        octos .*= canflash

        n -= 1
    end

    octos, flashes
end

function part2(octos)
    octos = copy(octos)
    n = 0
    while true
        n += 1

        # First, the energy level of each octopus increases by 1.
        octos .+= 1

        # Then, any octopus with an energy level greater than 9 flashes.
        # This increases the energy level of all adjacent octopuses by 1,
        # including octopuses that are diagonally adjacent. If this causes
        # an octopus to have an energy level greater than 9, it also flashes.
        # This process continues as long as new octopuses keep having their
        # energy level increased beyond 9. (An octopus can only flash at most
        # once per step.)
        canflash = ones(Bool, size(octos))
        while any(>(9), octos[canflash])
            coord = findfirst(i -> canflash[i] && octos[i] > 9, CartesianIndices(octos))
            octos[collect(neighbourcoords(octos, coord))] .+= 1
            canflash[coord] = false
        end

        # Finally, any octopus that flashed during this step has its energy
        # level set to 0, as it used all of its energy to flash.
        octos .*= canflash

        if !any(canflash)
            return n
        end
    end
end

println("Part 1: ", part1(input, 100))
println("Part 2: ", part2(input))


using FourierFlows, Printf, Random, FFTW

using Random: seed!

import GeophysicalFlows.TwoDTurb
import GeophysicalFlows.TwoDTurb: energy, enstrophy
import GeophysicalFlows: peakedisotropicspectrum

dev = CPU()     # Device (CPU/GPU)

# Parameters
  n = 256
  L = 2π
 nν = 2
  ν = 0.0
 dt = 1e-3
nstepsjit = 10 # call stepforward for nstepsjit to force compilation
nsteps = 5000
 nothingfunction(args...) = nothing

gr = TwoDGrid(dev, n, L, n, L; nthreads=1)
pr = TwoDTurb.Params(ν, nν)
vs = TwoDTurb.Vars(dev, gr)
eq = TwoDTurb.Equation(pr, gr)

prob = FourierFlows.Problem(eq, "FilteredAB3", dt, gr, vs, pr, dev)
filter = FourierFlows.makefilter(prob.eqn)

# some aliases
sol, cl, vs, gr = prob.sol, prob.clock, prob.vars, prob.grid

x, y = gridpoints(gr)

# Initial condition closely following pyqg barotropic example
# that reproduces the results of the paper by McWilliams (1984)
seed!(1234)
k0, E0 = 6, 0.5
zetai  = peakedisotropicspectrum(gr, k0, E0, mask=filter)
TwoDTurb.set_zeta!(prob, zetai)

for j=1:nstepsjit  #just-in-time compilation
  stepforward!(prob)
end

TwoDTurb.set_zeta!(prob, zetai)

startwalltime = time()
while cl.step < nsteps
  stepforward!(prob)
end

println(round((time()-startwalltime)/(nsteps-1)*1000, digits=3), " ms per time-step")

# using PyPlot
# TwoDTurb.updatevars!(prob)
# figure(figsize=(10, 4))
# subplot(121)
# pcolormesh(x, y, zetai)
# xlabel("x")
# ylabel("y")
# title("vorticity @ t=0")
# axis("square")
# subplot(122)
# pcolormesh(x, y, vs.zeta)
# axis("square")
# xlabel("x")
# ylabel("y")
# title("vorticity @ t="*string(round(cl.t, digits=2)))
# savefig("GeophysicalFlows_n"*string(n)*".png", dpi=400)


"""
$(TYPEDEF)
Abstract parent type for all distributions.
"""
abstract type Distribution end

"""
$(TYPEDEF)
One dimensional normal distribution.
"""
Base.@kwdef mutable struct Normal <: Distribution
  mu::Float64
  sigma::Float64
  _type::String = "IncrementalInference.PackedNormal"
end
Normal(mu::Number, sigma::Number) = Normal(;mu, sigma)

"""
$(TYPEDEF)
One dimensional Rayleigh distribution.
"""
Base.@kwdef mutable struct Rayleigh <: Distribution
  sigma::Float64
  _type::String = "IncrementalInference.PackedRayleigh"
end
Rayleigh(sigma::Number) = Rayleigh(sigma=sigma)

"""
$(TYPEDEF)
Multidimensional normal distribution specified by means and a covariance matrix.
"""
Base.@kwdef mutable struct FullNormal <: Distribution
  mu::Vector{Float64}
  cov::Vector{Float64}
  _type::String = "IncrementalInference.PackedFullNormal"
end
# TODO: Generalize this to any number type.
FullNormal(mu::Vector{Float64}, cov::Matrix{Float64}) = FullNormal(;mu, cov=vec(cov))

"""
$(TYPEDEF)
One dimensional uniform distribution.
"""
Base.@kwdef mutable struct Uniform <: Distribution
  a::Float64
  b::Float64
  _type::String = "IncrementalInference.PackedUniform"
end
Uniform(a::Number, b::Number) = Uniform(;a, b)

"""
$(TYPEDEF)
Categorical distribution specified by a set of probabilities summing up to 1.
"""
Base.@kwdef mutable struct Categorical <: Distribution
  p::Vector{Float64}
  _type::String = "IncrementalInference.PackedCategorical"
end
# TODO: Generalize this to any number type.
Categorical(p::Vector{Float64}) = Categorical(;p)

Base.@kwdef mutable struct ManifoldKernelDensity <: Distribution
  varType::String
  pts::Vector{Vector{Float64}}
  bw::Vector{Float64}   = Float64[]
  partial::Vector{Int}  = Int[]
  infoPerCoord::Vector{Float64} = zeros(length(pts[1]))
  _type::String         = "IncrementalInference.PackedManifoldKernelDensity"
end
ManifoldKernelDensity(varType::String, pts::AbstractVector{<:AbstractVector{<:Real}}; kw...) = ManifoldKernelDensity(;varType,pts,kw...)


function sample_time(
    rng::AbstractRNG, proj::Project, overall_start=:start, overall_finish=:finish
)
    realizations = Dict{Symbol, Float64}()
    realization(task_name::Symbol) = get!(realizations, task_name) do
        task_dist = proj.task_durations[task_name]
        if task_dist isa Real  # it is some constant
            return task_dist
        else  # it is a collection or a Distribution
            rand(rng, task_dist)
        end
    end

    function time_taken(start, finish)
        start == finish && return 0.0
        dependencies = (pre for (pre, this) in proj.links if this==finish)
        time_to_begin = maximum(time_taken.(start, dependencies))
        time_to_finish = realization(finish)
        return time_to_begin + time_to_finish
    end

    return time_taken(overall_start, overall_finish)
end

Base.eltype(::Type{<:Project}) = Float64
Base.eltype(::Type{<:Random.SamplerTrivial{<:Project}}) = Float64
function Random.Sampler(rng::AbstractRNG, proj::Project, ::Union{Val{1}, Val{Inf}})
    return Random.SamplerTrivial(proj)
end

Random.rand(rng::AbstractRNG, d::Random.SamplerTrivial{<:Project}) = sample_time(rng, d[])

# For Plots.jl users: should really be called with StatsPlot's `density` function
@recipe function f(proj::Project, nsamples=100_000, rng=Random.default_rng())
    title --> "Completion Time"
    rand(rng, proj, nsamples)
end

"""
    density(proj::Project, nsamples=100_000, rng=Random.default_rng())

Displayed a probability density function for how long it will take to complete the project.
`nsamples` controls how many samples are used to estimate the distribution, increasing this makes it take longer, but be more accurate.
`rng` controls the random number generator to use, you shouldn't normally need to touch this.

[Gadfly.jl](https://github.com/GiovineItalia/Gadfly.jl) is used for the displaying.
A [RecipesBase.jl](https://github.com/JuliaPlots/RecipesBase.jl) recipe is also provided and so this can be used with [StatsPlots.jl](https://github.com/JuliaPlots/StatsPlots.jl) similarly: as `StatsPlots.density(proj)`.
"""
function density(proj::Project, nsamples=100_000, rng=Random.default_rng(); kwargs...)
    # Our own Gadfly based density function
    if !isempty(kwargs)
        @warn "ProjectMangement.density does not support these keyword arguments. You may have intended to use StatsPlot.density." kwargs
    end
    samples = rand(rng, proj, nsamples)
    Gadfly.plot(
        x=samples,
        Gadfly.Geom.density,
        Gadfly.Guide.title("Completion Time PDF"),
        Gadfly.Guide.xlabel("Completed By"),
        Gadfly.Guide.ylabel("Probability Density"),
    )
end


const ∅ = BK_EMPTY

function PATH(s, t, i, source, sink, PARENT, INDEX)
    path = Int[]
    idxs = Int[]
    sᵢ = tᵢ = i
    while s != source
        pushfirst!(path, s)
        pushfirst!(idxs, sᵢ)
        s, sᵢ = PARENT[s], INDEX[s]
    end
    pushfirst!(path, s)
    pushfirst!(idxs, sᵢ)
    pop!(idxs)  # pop i out since we gonna push it again
    while t != sink
        push!(path, t)
        push!(idxs, tᵢ)
        t, tᵢ = PARENT[t], INDEX[t]
    end
    push!(path, t)
    push!(idxs, tᵢ)
    return path, idxs
end

function boykov_kolmogorov(source::Int, sink::Int, neighbors::Vector{Vector{Tuple{Int,Int}}}, weights::AbstractMatrix)
    # initialize: S = {s}, T = {t},  A = {s,t}, O = ∅
    vertexNum = length(neighbors)
    PARENT = zeros(Int, vertexNum)
    INDEX = zeros(Int, vertexNum)
    STATUS = fill(BK_FREE, vertexNum)
    STATUS[source] = BK_S_ACTIVE
    STATUS[sink] = BK_T_ACTIVE
    ORPHAN = fill(false, vertexNum)
    A = Int[source, sink]  # this queue could also contains inactive nodes which will be automatically skipped in the growth stage
    O = Int[]
    residual = copy(weights)
    flow = zero(Float64)
    # Boykov, Yuri, and Vladimir Kolmogorov. "An experimental comparison of min-cut/max-flow
    # algorithms for energy minimization in vision." Pattern Analysis and Machine Intelligence,
    # IEEE Transactions on 26.9 (2004): 1124-1137.
    while true
        # “growth” stage: search trees S and T grow until they touch giving an s → t path
        # grow S or T to find an augmenting path P from s to t
        P, IDX = growth_stage!(source, sink, neighbors, residual, A, STATUS, PARENT, INDEX)
        isempty(P) && break
        # “augmentation” stage: the found path is augmented, search tree(s) break into forest(s)
        flow += augmentation_stage!(neighbors, residual, P, IDX, O, STATUS, PARENT, INDEX)
        # “adoption” stage: trees S and T are restored
        adoption_stage!(source, sink, neighbors, residual, O, A, ORPHAN, STATUS, PARENT, INDEX)
    end
    return flow, STATUS
end

boykov_kolmogorov(source, sink, neighbors, weights::AbstractVector) = boykov_kolmogorov(source, sink, neighbors, reshape(weights, 2, :))

function growth_stage!(source, sink, neighbors, residual, A, STATUS, PARENT, INDEX)
    TREE(x) = STATUS[x] & (BK_S | BK_T)
    while !isempty(A)
        p = last(A)  # pick an active node p ∈ A ("First-In-First-Out"): enqueue -> queue -> dequeue
        STATUS[p] & BK_ACTIVE == BK_ACTIVE || (pop!(A); continue)  # automatically skip inactive node
        @inbounds for (q,qᵢ) in neighbors[p]
            tree_cap = TREE(p)==BK_S ? (p < q ? residual[1,qᵢ] : residual[2,qᵢ]) : (p < q ? residual[2,qᵢ] : residual[1,qᵢ])
            tree_cap > 0 || continue
            if TREE(q) == ∅
                # then add q to search tree as an active node
                STATUS[q] = STATUS[p]
                PARENT[q] = p
                pushfirst!(A, q)  # "First-In-First-Out": enqueue -> queue -> dequeue
                INDEX[q] = qᵢ   # cache index for referencing residual in augmentation stage
            end
            TREE(q)≠∅ && TREE(q)≠TREE(p) && return TREE(p)==BK_S ? PATH(p,q,qᵢ,source,sink,PARENT,INDEX) : PATH(q,p,qᵢ,source,sink,PARENT,INDEX)
        end
        # remove p from A ("First-In-First-Out"): enqueue -> queue -> dequeue
        STATUS[p] &= ~BK_ACTIVE
        pop!(A)
    end
    return Int[], Int[]
end

function augmentation_stage!(neighbors, residual, P, IDX, O, STATUS, PARENT, INDEX)
    TREE(x) = STATUS[x] & (BK_S | BK_T)
    # find the bottleneck capacity Δ on P
    Δ = Inf
    @inbounds for i = 1:length(P)-1
        p, q = P[i], P[i+1]
        if p < q
            Δ > residual[1,IDX[i]] && (Δ = residual[1,IDX[i]];)
        else
            Δ > residual[2,IDX[i]] && (Δ = residual[2,IDX[i]];)
        end
    end
    # update the residual graph by pushing flow Δ through P
    @inbounds for i = 1:length(P)-1
        p, q = P[i], P[i+1]
        if p < q
            residual[1,IDX[i]] -= Δ
            residual[2,IDX[i]] += Δ
            residual[1,IDX[i]] == 0 || continue
        else
            residual[2,IDX[i]] -= Δ
            residual[1,IDX[i]] += Δ
            residual[2,IDX[i]] == 0 || continue
        end
        if TREE(p) == TREE(q) == BK_S
            PARENT[q] = 0
            pushfirst!(O, q)
            INDEX[q] = 0  # clean up index cache
        end
        if TREE(p) == TREE(q) == BK_T
            PARENT[p] = 0
            pushfirst!(O, p)
            INDEX[p] = 0  # clean up index cache
        end
    end
    return Δ
end

function adoption_stage!(source, sink, neighbors, residual, O, A, ORPHAN, STATUS, PARENT, INDEX)
    TREE(x) = STATUS[x] & (BK_S | BK_T)
    fill!(ORPHAN, false)
    while !isempty(O)
        # pick an orphan node p ∈ O and remove it from O
        p = pop!(O)
        # find a new valid parent for p among its neighbors
        has_valid_parent = false
        @inbounds for (q,qᵢ) in neighbors[p]
            ORPHAN[q] && continue
            TREE(q) == TREE(p) || continue
            tree_cap = TREE(p)==BK_S ? (q < p ? residual[1,qᵢ] : residual[2,qᵢ]) : (q < p ? residual[2,qᵢ] : residual[1,qᵢ])
            tree_cap > 0 || continue
            # the “origin” of q should be either source or sink, it should not originates from orphan
            x = q
            while PARENT[x] ≠ 0
                x = PARENT[x]
                ORPHAN[x] && break
                ORPHAN[x] = true
            end
            if x == source || x == sink
                PARENT[p] = q
                INDEX[p] = qᵢ
                has_valid_parent = true
                fill!(ORPHAN, false)
                break
            end
        end
        has_valid_parent && continue
        @inbounds for (q,qᵢ) in neighbors[p]
            TREE(q) == TREE(p) || continue
            tree_cap = TREE(p)==BK_S ? (q < p ? residual[1,qᵢ] : residual[2,qᵢ]) : (q < p ? residual[2,qᵢ] : residual[1,qᵢ])
            tree_cap > 0   && (STATUS[q] |= BK_ACTIVE; pushfirst!(A, q);)
            PARENT[q] == p && (PARENT[q] = 0; INDEX[q] = 0; pushfirst!(O, q);)
        end
        # TREE(p) := ∅, A := A - {p}
        STATUS[p] = BK_FREE   # note that this also marks p as inactive node
    end
end


@testset "Rectifier square" begin

r = Rectifier("squared")
@test r.T == "squared"

    x = 2.5
@test abs(r(x) - x^2)<1e-10

# Test gradient
@test abs(ForwardDiff.derivative(y->r(y), x) - grad_x(r, x) ) < 1e-10

# Test hessian
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),x) - hess_x(r, x) ) < 1e-10

# Test gradient of log evaluation
@test abs(ForwardDiff.derivative(y->log(r(y)), x) - grad_x_logeval(r, x) ) < 1e-10

# Test hessian of log evaluation
@test abs(ForwardDiff.hessian(y->log(r(y[1])), [x])[1,1] - hess_x_logeval(r, x) ) < 1e-10

x = -2.5
@test abs(r(x) - x^2)<1e-10

# Test gradient
@test abs(ForwardDiff.derivative(y->r(y), x) - grad_x(r, x) ) < 1e-10

# Test hessian
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),x) - hess_x(r, x) ) < 1e-10

# Test gradient of log evaluation
@test abs(ForwardDiff.derivative(y->log(r(y)), x) - grad_x_logeval(r, x) ) < 1e-10

# Test hessian of log evaluation
@test abs(ForwardDiff.hessian(y->log(r(y[1])), [x])[1,1] - hess_x_logeval(r, x) ) < 1e-10

end

@testset "Rectifier exponential" begin

r = Rectifier("exponential")
@test r.T == "exponential"

    x = 2.46
@test abs(r(x) - exp(x))<1e-10

# Test gradient
@test abs(ForwardDiff.derivative(y->r(y), x) - grad_x(r, x) ) < 1e-10

# Test hessian
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),x) - hess_x(r, x) ) < 1e-10

# Test gradient of log evaluation
@test abs(ForwardDiff.derivative(y->log(r(y)), x) - grad_x_logeval(r, x) ) < 1e-10

# Test hessian of log evaluation
@test abs(ForwardDiff.hessian(y->log(r(y[1])), [x])[1,1] - hess_x_logeval(r, x) ) < 1e-10

x = -2.46
@test abs(r(x) - exp(x))<1e-10

# Test gradient
@test abs(ForwardDiff.derivative(y->r(y), x) - grad_x(r, x) ) < 1e-10

# Test hessian
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),x) - hess_x(r, x) ) < 1e-10

# Test gradient of log evaluation
@test abs(ForwardDiff.derivative(y->log(r(y)), x) - grad_x_logeval(r, x) ) < 1e-10

# Test hessian of log evaluation
@test abs(ForwardDiff.hessian(y->log(r(y[1])), [x])[1,1] - hess_x_logeval(r, x) ) < 1e-10

end

@testset "Rectifier softplus" begin

r = Rectifier("softplus")
@test r.T == "softplus"

    x = 2.5
    a = log(2)
@test abs(r(x) - (log(1 + exp(-abs(a*x))) + max(a*x, 0))/a)<1e-10

#Test inversion
@test abs(inverse(r,r(x))-x)<1e-10
@test abs(r(inverse(r,x))-x)<1e-10

#Test inversion with x large
xlarge = 10.0^6
@test abs(inverse(r,xlarge) -xlarge)<1e-10

# Test gradient
@test abs(ForwardDiff.derivative(y->r(y), x) - grad_x(r, x) ) < 1e-10

# Test hessian
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),x) - hess_x(r, x) ) < 1e-10

# Test gradient of log evaluation
@test abs(ForwardDiff.derivative(y->log(r(y)), x) - grad_x_logeval(r, x) ) < 1e-10

# Test hessian of log evaluation
@test abs(ForwardDiff.hessian(y->log(r(y[1])), [x])[1,1] - hess_x_logeval(r, x) ) < 1e-10

x = -2.5
a = log(2)
@test abs(r(x) - (log(1 + exp(-abs(a*x))) + max(a*x, 0))/a)<1e-10

#Test inversion with x large
xlarge = 10.0^6
@test abs(inverse(r,xlarge) -xlarge)<1e-10

# Test gradient
@test abs(ForwardDiff.derivative(y->r(y), x) - grad_x(r, x) ) < 1e-10

# Test hessian
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),x) - hess_x(r, x) ) < 1e-10

# Test gradient of log evaluation
@test abs(ForwardDiff.derivative(y->log(r(y)), x) - grad_x_logeval(r, x) ) < 1e-10

# Test hessian of log evaluation
@test abs(ForwardDiff.hessian(y->log(r(y[1])), [x])[1,1] - hess_x_logeval(r, x) ) < 1e-10

end

@testset "Rectifier explinearunit" begin

r = Rectifier("explinearunit")
@test r.T == "explinearunit"

## Test x<0
xneg = -0.5
@test abs(r(xneg) - exp(xneg))<1e-10

# Test gradient
@test abs(ForwardDiff.derivative(y->r(y), xneg) - grad_x(r, xneg) ) < 1e-10
@test abs(ForwardDiff.derivative(y->r(y), xneg) - exp(xneg) ) < 1e-10

# Test hessian
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),xneg) - hess_x(r, xneg) ) < 1e-10
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),xneg) - exp(xneg) ) < 1e-10

# Test gradient of log evaluation
@test abs(ForwardDiff.derivative(y->log(r(y)), xneg) - grad_x_logeval(r, xneg) ) < 1e-10

# Test hessian of log evaluation
@test abs(ForwardDiff.hessian(y->log(r(y[1])), [xneg])[1,1] - hess_x_logeval(r, xneg) ) < 1e-10

## Test x>=0
x = 2.5
@test abs(r(x) - (x + 1.0))<1e-10

#Test inversion
@test abs(inverse(r,r(x))-x)<1e-10
@test abs(r(inverse(r,x))-x)<1e-10

# Test gradient
@test abs(ForwardDiff.derivative(y->r(y), x) - grad_x(r, x) ) < 1e-10
@test abs(ForwardDiff.derivative(y->r(y), x) - 1.0 ) < 1e-10

# Test hessian
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),x) - hess_x(r, x) ) < 1e-10
@test abs(ForwardDiff.derivative(z->ForwardDiff.derivative(y->r(y), z),x) - 0.0) < 1e-10

# Test gradient of log evaluation
@test abs(ForwardDiff.derivative(y->log(r(y)), x) - grad_x_logeval(r, x) ) < 1e-10

# Test hessian of log evaluation
@test abs(ForwardDiff.hessian(y->log(r(y[1])), [x])[1,1] - hess_x_logeval(r, x) ) < 1e-10

end


"""
    Base.:+(a::UncertainDataset, b::UncertainDataset; n = 10000)

Addition operator for two uncertain value datasets.

To obtain the sum for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th added uncertain value is a kernel density estimate to those sums, i.e.
`added_val = UncertainValue(UnivariateKDE, sample_xᵢ .+ sample_yᵢ)`.
"""
function Base.:+(a::UncertainDataset, b::UncertainDataset; n = 10000)
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Dataset lengths do not match ($N, $n_vals_b)"))
    end

    UncertainDataset([a[i] + b[i] for i = 1:N])
end

"""
    Base.:+(a::Real, b::UncertainDataset;
        n = 10000)

Addition operator for scalars and uncertain value datasets.

To obtain the sum for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th added uncertain value is a kernel density estimate to those sums, i.e.
`added_val = UncertainValue(UnivariateKDE, sample_xᵢ .+ sample_yᵢ)`.
"""
function Base.:+(a::T, b::UncertainDataset; n = 10000) where {T <: Number}
    N = length(a)
    UncertainDataset([a + b[i] for i = 1:N])
end

"""
    Base.:+(a::Vector{Real}, b::UncertainDataset; n = 10000)

Addition operator for scalars and uncertain value datasets.

To obtain the sum for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th added uncertain value is a kernel density estimate to those sums, i.e.
`added_val = UncertainValue(UnivariateKDE, sample_xᵢ .+ sample_yᵢ)`.
"""
function Base.:+(a::Vector{T}, b::UncertainDataset; n = 10000) where {T <: Number}
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Dataset lengths do not match ($N, $n_vals_b)"))
    end

    UncertainDataset([a[i] + b[i] for i = 1:N])
end

"""
    Base.:+(a::UncertainDataset, b::Real; n = 10000)

Addition scalars to uncertain value datasets.

To obtain the sum for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th added uncertain value is a kernel density estimate to those sums, i.e.
`added_val = UncertainValue(UnivariateKDE, sample_xᵢ .+ sample_yᵢ)`.
"""
function Base.:+(a::UncertainDataset, b::T; n = 10000) where {T <: Number}
    N = length(a)

    UncertainDataset([a[i] + b for i = 1:N])
end

"""
    Base.:+(a::UncertainDataset, b::Real; n = 10000)

Addition scalars to uncertain value datasets.

To obtain the sum for a pair of uncertain values `(xᵢ ∈ a, yᵢ ∈ b)`, both `u₁` and `u₂` are
resampled `n` times each, so `sample_xᵢ = resample(x, n)` and `sample_yᵢ = resample(y, n)`.

The i-th added uncertain value is a kernel density estimate to those sums, i.e.
`added_val = UncertainValue(UnivariateKDE, sample_xᵢ .+ sample_yᵢ)`.
"""
function Base.:+(a::UncertainDataset, b::Vector{T}; n = 10000) where {T <: Number}
    N = length(a)
    n_vals_b = length(b)

    if N != n_vals_b
        throw(ArgumentError("Dataset lengths do not match ($N, $n_vals_b)"))
    end

    UncertainDataset([a[i] + b[i] for i = 1:N])
end


@testset "permutation" begin

N = [3,4]
M = prod(N)
Z = zeros(M,M)

@testset "create_πcmp" begin
    permute = [3,1,2]
    scale = [-1,1,-1]
    Kf = length(permute)
    KfM = Kf * M

    # order_cmpfirst = false
    P = create_πcmp(N, permute, scale=scale, order_cmpfirst=false)
    P2 = zeros(KfM,KfM)
    Id = Matrix{eltype(scale)}(I, M, M)

    for nw = 1:Kf
        iblk = permute[nw]
        jblk = nw
        I = (iblk-1)*M+1:iblk*M
        J = (jblk-1)*M+1:jblk*M
        P2[I,J] = scale[nw] * Id
    end
    @test P2 == P

    # order_cmpfirst = true
    P_cmpfirst = create_πcmp(N, permute, scale=scale)
    p = spzeros(eltype(scale), Kf, Kf)
    for nw = 1:Kf
        p[permute[nw], nw] = scale[nw]
    end
    P2_cmpfirst = blockdiag(ntuple(k->p, M)...)
    @test P2_cmpfirst == P_cmpfirst
end  # @testset "create_πcmp"

end  # @testset "permutation"


# The primes 3, 7, 109, and 673, are quite remarkable. By taking any two primes
# and concatenating them in any order the result will always be prime. For
# example, taking 7 and 109, both 7109 and 1097 are prime. The sum of these four
# primes, 792, represents the lowest sum for a set of four primes with this
# property.
#
# Find the lowest sum for a set of five primes for which any two primes
# concatenate to produce another prime.

using ProjectEulerSolutions

# Fairly brute-force approach, start with Seive of Eratosthenes to get an
# initial set of primes, then start building up a combination of 5 that each
# satisfy the concatenated condition given all of the previous numbers.
# Short-cutting can be aided by upper bounding some prime number ranges based
# on previous best results.  Finally, the order you pass through the pairs
# matters, looping with the largest value on the outside loop finds the minimum
# set faster.  A recursive solution would look cleaner, but it wouldn't
# necessarily be faster.
#
# This seems like the kind of thing you can do by building a graph with edges
# representing legal concatenated pairs and then solving for the fully connected
# cliques of size 5.
function p060solution(n::Integer=1_000)::Integer
    primes = sieve_eratosthenes(n, 3)

    min_sum = 10_000_000
    n = length(primes)

    # Use a bitarray, but might need to like a Dict of Sets if the number of
    # primes is very large
    pairs = falses(n, n)
    ashift = typeof(primes[1])(10)

    # Find concatenated prime pairs
    for a in 1:n
        # Break early if the sum is too large
        if 5 * primes[a] > min_sum
            break
        end

        # Fill in the pairs bitarray
        prime_a = primes[a]
        jshift = typeof(prime_a)(10)
        if prime_a > ashift
            ashift *= 10
        end
        for j in a+1:n
            prime_j = primes[j]
            if prime_j > jshift
                jshift *= 10
            end
            if isprime(ashift * prime_j + prime_a) && isprime(jshift * prime_a + prime_j)
                pairs[a,j] = true
            end
        end

        # Start looking for legal combinations of 2, 3, 4, and 5 primes
        for b in 4:a-1
            if !pairs[b,a]
                continue
            end

            for c in 3:b-1
                if !pairs[c,a] || !pairs[c,b]
                    continue
                end

                for d in 2:c-1
                    if !pairs[d,a] || !pairs[d,b] || !pairs[d,c]
                        continue
                    end

                    for e in 1:d-1
                        if pairs[e,a] && pairs[e,b] && pairs[e,c] && pairs[e,d]
                            min_sum = primes[a] + primes[b] + primes[c] + primes[d] + primes[e]
                        end
                    end
                end
            end
        end
    end

    return min_sum
end

p060 = Problems.Problem(p060solution)

Problems.benchmark(p060, 10_000)


#=
Continuation of data, to e.g. fill missing values
=#
using SignalDecomposition
export sinusoidal_continuation

_numbertype(T) = _numbertype(eltype(T))
_numbertype(::Type{Union{T, Missing}}) where T = T
_numbertype(::Type{T}) where {T <: Real} = T

# TODO: This function needs to become more general to allow for time sampling
# other than monthly (currently explicitly assumed)

# TODO: perhaps it is worth considering a version that makes true time in seconds
# instead of months, which in the future would allow support of diurnal cycles
# TODO: YES! Make true time be seconds!

"""
    sinusoidal_continuation(T::ClimArray, fs = [1, 2]; Tmin = -Inf, Tmax = Inf)
Fill-in the missing values of spatiotemporal field `T`, by fitting sinusoidals
to the non-missing values, and then using the fitted sinusoidals for the missing values.

Frequencies are given per year (frequency 2 means 1/2 of a year).

`Tmin, Tmax` limits are used to clamp the result into this range (no clamping in the
default case).
"""
function sinusoidal_continuation(T, frequencies = [1.0, 2.0]; Tmin = -Inf, Tmax = Inf)
    E = _numbertype(T)
    lpv = Sinusoidal(E.(frequencies ./ DAYS_IN_YEAR))
    fullT = copy(T)
    # TODO: this must be extended to a general "true time" function
    truetime = Float32.(cumsum(daysinmonth.(dims(T, Time))))
    for i in otheridxs(T, Time)
        x = T[i...]
        any(ismissing, x) || continue # this space needs no correction
        mi = findall(!ismissing, x)
        x0 = Array{Float32}(x[mi])
        t0 = truetime[mi]
        sea, res = decompose(t0, x0, lpv)
        c = cosines(truetime, lpv)
        imi = findall(ismissing, x)
        x[imi] .= c[imi]
        fullT[i...] .= x
    end
    if !isinf(Tmin) && !isinf(Tmin)
        return clamp.(fullT, Tmin, Tmax)
    else
        return fullT
    end
end

function cosines(t, lpv::Sinusoidal)
    s = fill(lpv.A[1], length(t))
    for i in 2:length(lpv.A)
        @. s += lpv.A[i] * cos(2π * lpv.fs[i-1] * t + lpv.φ[i])
    end
    return s
end


export solveConvexMAP

"""
function solveConvexMAP

Find global MAP estimate using Gurobi

Input:

	1. param - StrainReconParam

Keyword arguments:

	2. pwr - exponent in the likelihood (must be 1.0 or 2.0)
	3. out - output flag for Gurobi solver
	4. tol - MIPGap tolerance parameter for Gurobi solver

Output:

	1. M      - estimated binary matrix
	2. w      - estimated frequency vector
    3. status - return code of the Gurobi solver
"""

function solveConvexMAP(param::StrainReconParam;pwr::Float64=param.pwr,out::Int=0,tol::Float64=1e-6)

    pwr = Int(pwr)

    m = param.m
    n = param.n
    p = param.p
    q = m*(p-1)
    d = param.d
    dRp = diag(param.R).^pwr

    model = Model(solver=GurobiSolver(OutputFlag=out,MIPGap=tol))
    @variable(model, 0 <= w[1:n] <=1)
    @variable(model, 0 <= W[1:q,1:n] <=1)
    @variable(model, 0 <= M[1:q,1:n] <=1,Int)
    if pwr == 1
        @variable(model, zplus[1:q] >= 0)
        @variable(model, zminus[1:q] >= 0)
    end

    @constraint(model,W.<=M)
    for k=1:n
        @constraint(model,W[:,k].<= w[k])
        @constraint(model,W[:,k].>= M[:,k] + w[k]-1)
    end
    for k=1:n-1
        @constraint(model,w[k]>=w[k+1])
    end
    @constraint(model, sum(w)==1)
    if pwr == 1
        @constraint(model, M*w-d .== zminus+zplus)
    end

    if p > 2
        for k=1:n
            for j=1:m
                @constraint(model, sum(M[(j-1)*(p-1)+1 : j*(p-1), k]) <= 1)
            end
        end
    end

    e = ones(n)

    if pwr == 2
        JuMP.setobjective(model,:Min, dRp'*(W*e-d).^2 );
    elseif pwr == 1
        JuMP.setobjective(model,:Min, dRp'*(zplus+zminus) );
    end

    status = solve(model)

    return BitArray(getvalue(M)), getvalue(w), status

end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 0b19c578-762b-11eb-34b6-01e80cef1406
using Pkg, DrWatson

# ╔═╡ 3f814e9e-762b-11eb-1340-91617ca7b58a
begin
    using Distributions
    using StatsPlots
    using StatsBase
    using LaTeXStrings
    using CSV
    using DataFrames
    using LinearAlgebra
    using Random
    using StanSample
    using StatisticalRethinking
end

# ╔═╡ af1efcf4-779a-11eb-0d62-e3c63876a09b
md" ## Clip-09-25-27s.jl"

# ╔═╡ cb88768a-779b-11eb-1a4e-99f92ded29ca
begin
	y = rand(Normal(), 100)
	data = (N = length(y), y=y)
end;

# ╔═╡ e37331ea-779b-11eb-1b02-1fdb61aae124
stan9_4 = "
data {
	int N;
	vector[N] y;
}
parameters{
    real a1;
    real a2;
    real<lower=0> sigma;
}
model{
    real mu;
    sigma ~ exponential( 1 );
    a2 ~ normal( 0 , 1000 );
    a1 ~ normal( 0 , 1000 );
    mu = a1 + a2;
    y ~ normal( mu , sigma );
}
";

# ╔═╡ 3fafc7fe-762b-11eb-37da-254b6d75f5ca
begin
	m9_4s = SampleModel("m9.4s", stan9_4)
	rc9_4s = stan_sample(m9_4s; data)

	if success(rc9_4s)
		m9_4s_df = read_samples(m9_4s, :dataframe)
		PRECIS(m9_4s_df)
	end
end

# ╔═╡ 140d4962-7ebd-11eb-20f6-efc652cf9335
read_summary(m9_4s)

# ╔═╡ b4d020b8-762d-11eb-167a-9fd79d948647
trankplot(m9_4s, :a1)[1]

# ╔═╡ 7b4bc93e-7ebb-11eb-1912-f152a5e76df6
trankplot(m9_4s, :a2)[1]

# ╔═╡ 728ad3aa-7eb7-11eb-2fea-634a82832501
trankplot(m9_4s, :sigma)[1]

# ╔═╡ 059b3ada-7eb8-11eb-10aa-459e3ec2bc55
begin
	chns9_4s = read_samples(m9_4s, :mcmcchains)
	Text(sprint(show, "text/plain", summarize(chns9_4s)))
end

# ╔═╡ 453b6b34-8356-11eb-0049-51a7363b24d4
CHNS(chns9_4s)

# ╔═╡ 6fbfa790-7eb9-11eb-0360-ed6fc7fa11a3
plot(chns9_4s)

# ╔═╡ a93a71e4-7eb9-11eb-2a36-d33e7b88e2d6
stan9_5 = "
data {
	int N;
	vector[N] y;
}
parameters{
    real a1;
    real a2;
    real<lower=0> sigma;
}
model{
    real mu;
    sigma ~ exponential( 1 );
    a2 ~ normal( 0 , 10 );
    a1 ~ normal( 0 , 10 );
    mu = a1 + a2;
    y ~ normal( mu , sigma );
}
";

# ╔═╡ b4ea8d58-7eb9-11eb-172f-3941a1b8fe98
begin
	m9_5s = SampleModel("m9.5s", stan9_5)
	rc9_5s = stan_sample(m9_5s; data)

	if success(rc9_5s)
		m9_5s_df = read_samples(m9_5s, :dataframe)
		PRECIS(m9_5s_df)
	end
end

# ╔═╡ ef7953e8-7eb9-11eb-3344-8d762d7e3584
trankplot(m9_5s, :a1)[1]

# ╔═╡ f042277e-7ebb-11eb-33d4-b7edd9539943
trankplot(m9_5s, :a2)[1]

# ╔═╡ 29b1d698-7eba-11eb-00aa-e95796579cc6
trankplot(m9_5s, :sigma)[1]

# ╔═╡ 39bb0f76-7eba-11eb-1bc2-293259d12dfc
begin
	chns9_5s = read_samples(m9_5s, :mcmcchains)
	Text(sprint(show, "text/plain", summarize(chns9_5s)))
end

# ╔═╡ 6748f438-8380-11eb-312d-cbfe0500b586
CHNS(chns9_5s)

# ╔═╡ 43b362c6-7eba-11eb-389e-fb53b3783e8d
plot(chns9_5s)

# ╔═╡ ca61c622-779a-11eb-1292-3f272953867f
md" ## End of clip-09-25-27s.jl"

# ╔═╡ Cell order:
# ╟─af1efcf4-779a-11eb-0d62-e3c63876a09b
# ╠═0b19c578-762b-11eb-34b6-01e80cef1406
# ╠═3f814e9e-762b-11eb-1340-91617ca7b58a
# ╠═cb88768a-779b-11eb-1a4e-99f92ded29ca
# ╠═e37331ea-779b-11eb-1b02-1fdb61aae124
# ╠═3fafc7fe-762b-11eb-37da-254b6d75f5ca
# ╠═140d4962-7ebd-11eb-20f6-efc652cf9335
# ╠═b4d020b8-762d-11eb-167a-9fd79d948647
# ╠═7b4bc93e-7ebb-11eb-1912-f152a5e76df6
# ╠═728ad3aa-7eb7-11eb-2fea-634a82832501
# ╠═059b3ada-7eb8-11eb-10aa-459e3ec2bc55
# ╠═453b6b34-8356-11eb-0049-51a7363b24d4
# ╠═6fbfa790-7eb9-11eb-0360-ed6fc7fa11a3
# ╠═a93a71e4-7eb9-11eb-2a36-d33e7b88e2d6
# ╠═b4ea8d58-7eb9-11eb-172f-3941a1b8fe98
# ╠═ef7953e8-7eb9-11eb-3344-8d762d7e3584
# ╠═f042277e-7ebb-11eb-33d4-b7edd9539943
# ╠═29b1d698-7eba-11eb-00aa-e95796579cc6
# ╠═39bb0f76-7eba-11eb-1bc2-293259d12dfc
# ╠═6748f438-8380-11eb-312d-cbfe0500b586
# ╠═43b362c6-7eba-11eb-389e-fb53b3783e8d
# ╟─ca61c622-779a-11eb-1292-3f272953867f


using MetabolicSim
using Test

#function to create 1x1 matricies
function mat(sc)
    reshape([sc],1,1)
end

#create parameters
M,N = 1,1
u,Rm = mat(1.0), [1.0]
l,ρ = mat(0.2), [0.2]
p = make_parameters(N,M,u,Rm,l,ρ)

#test vs analytical prediction
C_eq = ρ[1] / Rm[1]
R_eq = -Rm[1] / (u[1] * (l[1] - 1))

@testset begin
    sol = @test_nowarn simulate(ones(N+M),p,t_end = 100)
    @test all(isapprox.(sol[end],[C_eq,R_eq], rtol = 1e-2))
end


using DataFrames
using StatPlots
using DataArrays
#Importing in dateset
cpi = readtable("cpi.csv")
#setting as date format
cpidate = Date(cpi[1])
#seperating dataset
cpidata = cpi[2]
#normalizing
cpinormal = cpi[2]/sum(cpi[2])
#reading in higher ed dataset
highered = readtable("highered.csv")
#seperating out date
highereddate = DataArray(highered[1])
#making as Date
highereddate = Date(highereddate)
#normalizing
higherednormal = highered[2]/sum(highered[2])
#plot unnormalized cpi vs highered
plot(cpidate, cpidata, label = "Urban CPI")
plot!(highereddate, highered[2], label = "Higher Education CPI")
title!("Urban CPI vs. Higher Education")
xaxis!("Date")
yaxis!("Price")

#plot normalized Plots
plot(cpidate, cpinormal)
plot!(highereddate, higherednormal)


module TestLines

include("data.jl")

using Test
import Pkg: TOML
import TreeLas
import GraphIdx

function instance(line; to_tree::Bool = false)
    y = Float64.(line["y"])
    n = length(y)
    mu = if haskey(line, "mu")
        GraphIdx.Weights(Float64.(line["mu"]))
    else
        GraphIdx.Weights(Float64)
    end
    lam = Float64.(line["lam"])
    if to_tree && lam isa Array
        prepend!(lam, NaN)
    end
    lam = GraphIdx.Weights(lam)
    return y, lam, mu
end

const lines_toml = data_dir("test", "lines.toml")
lines = TOML.parsefile(lines_toml)["test"]

@testset "TreeLas LineLas                " begin
    @test length(lines) >= 5
    for (i, line) in enumerate(lines)
        y, lam, mu = instance(line, to_tree=false)
        n = length(y)
        @testset "line_las[$i]: n = $n" begin
            # @info "\nLine[$i]"
            @test TreeLas.LineDP.line_las(y, lam, mu) ≈ line["x"]
        end
    end
end

@testset "TreeLas line by trees          " begin
    for (i, line) in enumerate(lines)
        y, lam, mu = instance(line, to_tree=true)
        n = length(y)
        @testset "Line[$i]: n = $n" begin
            parent = collect(0:n-1)
            parent[1] = 1
            t = GraphIdx.Tree.RootedTree(1, parent)
            @test mu[1] != 0.0
            @test length(parent) == n
            @test TreeLas.TreeDP.tree_dp(y, t, lam, mu) ≈ line["x"]
        end
    end
end

end


using Test, F1Method

using LinearAlgebra, DiffEqBase, ForwardDiff

# Set up:
# - overload `SteadyStateProblem` constructor
# - overload `solve` function
# - define solver algorithm (basic Newton here)
# - define type for that algorithm (here `MyAlg`)
include("simple_setup.jl")

@testset "quasi-Rosenbrock derivative" begin
    include("rosenbrock.jl")
end


using IterativeSolvers
using Base.Test

function randSPD(T, n)
    A = rand(T, n, n) + n * I
    A' * A
end

function approx_eigenvalue_bounds(A)
    λs = eigvals(A)
    mnv, mxv = minimum(real(λs)), maximum(real(λs))
    Δ = (mxv - mnv) / 100
    mnv - Δ, mxv + Δ
end

#Chebyshev
@testset "Chebyshev" begin

n = 10
srand(1234321)

@testset "Matrix{$T}" for T in (Float32, Float64, Complex64, Complex128)
    A = randSPD(T, n)
    b = rand(T, n)
    tol = √(eps(real(T)))

    # Without a preconditioner
    begin
        λ_min, λ_max = approx_eigenvalue_bounds(A)
        x0 = rand(n)
        x, history = chebyshev(A, b, λ_min, λ_max, tol=tol, maxiter=10n, log=true)
        @test isa(history, ConvergenceHistory)
        @test history.isconverged
        @test norm(A * x - b) / norm(b) ≤ tol
    end

    # With an initial guess
    begin
        λ_min, λ_max = approx_eigenvalue_bounds(A)
        x0 = rand(T, n)
        x, history = chebyshev!(x0, A, b, λ_min, λ_max, tol=tol, maxiter=10n, log=true)
        @test isa(history, ConvergenceHistory)
        @test history.isconverged
        @test x == x0
        @test norm(A * x - b) / norm(b) ≤ tol
    end

    # With a preconditioner
    begin
        B = randSPD(T, n)
        B_fact = cholfact!(B)
        λ_min, λ_max = approx_eigenvalue_bounds(B_fact \ A)
        x, history = chebyshev(A, b, λ_min, λ_max, Pl = B_fact, tol=tol, maxiter=10n, log=true)
        @test history.isconverged
        @test norm(A * x - b) / norm(b) ≤ tol
    end
end
end


# # Incompressible Elasticity
#
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`incompressible_elasticity.ipynb`](@__NBVIEWER_ROOT_URL__/examples/incompressible_elasticity.ipynb)

# ## Introduction
#
# Mixed elements can be used to overcome locking when the material becomes
# incompressible. However, for an element to be stable, it needs to fulfill
# the LBB condition.
# In this example we will consider two different element formulations
# - linear displacement with linear pressure approximation (does *not* fulfill LBB)
# - quadratic displacement with linear pressure approximation (does fulfill LBB)
# The quadratic/linear element is also known as the Taylor-Hood element.
# We will consider Cook's Membrane with an applied traction on the right hand side.
#-
# ## Commented Program
#
#md # The full program, without comments, can be found in the next
#md # [section](@ref incompressible_elasticity-plain-program).
using Ferrite
using BlockArrays, SparseArrays, LinearAlgebra

# First we generate a simple grid, specifying the 4 corners of Cooks membrane.
function create_cook_grid(nx, ny)
    corners = [Vec{2}((0.0,   0.0)),
               Vec{2}((48.0, 44.0)),
               Vec{2}((48.0, 60.0)),
               Vec{2}((0.0,  44.0))]
    grid = generate_grid(Triangle, (nx, ny), corners);
    ## facesets for boundary conditions
    addfaceset!(grid, "clamped", x -> norm(x[1]) ≈ 0.0);
    addfaceset!(grid, "traction", x -> norm(x[1]) ≈ 48.0);
    return grid
end;

# Next we define a function to set up our cell- and facevalues.
function create_values(interpolation_u, interpolation_p)
    ## quadrature rules
    qr      = QuadratureRule{2,RefTetrahedron}(3)
    face_qr = QuadratureRule{1,RefTetrahedron}(3)

    ## geometric interpolation
    interpolation_geom = Lagrange{2,RefTetrahedron,1}()

    ## cell and facevalues for u
    cellvalues_u = CellVectorValues(qr, interpolation_u, interpolation_geom)
    facevalues_u = FaceVectorValues(face_qr, interpolation_u, interpolation_geom)

    ## cellvalues for p
    cellvalues_p = CellScalarValues(qr, interpolation_p, interpolation_geom)

    return cellvalues_u, cellvalues_p, facevalues_u
end;

# We create a DofHandler, with two fields, `:u` and `:p`,
# with possibly different interpolations
function create_dofhandler(grid, ipu, ipp)
    dh = DofHandler(grid)
    push!(dh, :u, 2, ipu) # displacement
    push!(dh, :p, 1, ipp) # pressure
    close!(dh)
    return dh
end;

# We also need to add Dirichlet boundary conditions on the `"clamped"` faceset.
# We specify a homogeneous Dirichlet bc on the displacement field, `:u`.
function create_bc(dh)
    dbc = ConstraintHandler(dh)
    add!(dbc, Dirichlet(:u, getfaceset(dh.grid, "clamped"), (x,t) -> zero(Vec{2}), [1,2]))
    close!(dbc)
    t = 0.0
    update!(dbc, t)
    return dbc
end;

# The material is linear elastic, which is here specified by the shear and bulk moduli
struct LinearElasticity{T}
    G::T
    K::T
end

# Now to the assembling of the stiffness matrix. This mixed formulation leads to a blocked
# element matrix. Since Ferrite does not force us to use any particular matrix type we will
# use a `PseudoBlockArray` from `BlockArrays.jl`.
function doassemble(cellvalues_u::CellVectorValues{dim}, cellvalues_p::CellScalarValues{dim},
                    facevalues_u::FaceVectorValues{dim}, K::SparseMatrixCSC, grid::Grid,
                    dh::DofHandler, mp::LinearElasticity) where {dim}

    f = zeros(ndofs(dh))
    assembler = start_assemble(K, f)
    nu = getnbasefunctions(cellvalues_u)
    np = getnbasefunctions(cellvalues_p)

    fe = PseudoBlockArray(zeros(nu + np), [nu, np]) # local force vector
    ke = PseudoBlockArray(zeros(nu + np, nu + np), [nu, np], [nu, np]) # local stiffness matrix

    ## traction vector
    t = Vec{2}((0.0, 1/16))
    ## cache ɛdev outside the element routine to avoid some unnecessary allocations
    ɛdev = [zero(SymmetricTensor{2, dim}) for i in 1:getnbasefunctions(cellvalues_u)]

    for cell in CellIterator(dh)
        fill!(ke, 0)
        fill!(fe, 0)
        assemble_up!(ke, fe, cell, cellvalues_u, cellvalues_p, facevalues_u, grid, mp, ɛdev, t)
        assemble!(assembler, celldofs(cell), fe, ke)
    end

    return K, f
end;

# The element routine integrates the local stiffness and force vector for all elements.
# Since the problem results in a symmetric matrix we choose to only assemble the lower part,
# and then symmetrize it after the loop over the quadrature points.
function assemble_up!(Ke, fe, cell, cellvalues_u, cellvalues_p, facevalues_u, grid, mp, ɛdev, t)

    n_basefuncs_u = getnbasefunctions(cellvalues_u)
    n_basefuncs_p = getnbasefunctions(cellvalues_p)
    u▄, p▄ = 1, 2
    reinit!(cellvalues_u, cell)
    reinit!(cellvalues_p, cell)

    ## We only assemble lower half triangle of the stiffness matrix and then symmetrize it.
    @inbounds for q_point in 1:getnquadpoints(cellvalues_u)
        for i in 1:n_basefuncs_u
            ɛdev[i] = dev(symmetric(shape_gradient(cellvalues_u, q_point, i)))
        end
        dΩ = getdetJdV(cellvalues_u, q_point)
        for i in 1:n_basefuncs_u
            divδu = shape_divergence(cellvalues_u, q_point, i)
            δu = shape_value(cellvalues_u, q_point, i)
            for j in 1:i
                Ke[BlockIndex((u▄, u▄), (i, j))] += 2 * mp.G * ɛdev[i] ⊡ ɛdev[j] * dΩ
            end
        end

        for i in 1:n_basefuncs_p
            δp = shape_value(cellvalues_p, q_point, i)
            for j in 1:n_basefuncs_u
                divδu = shape_divergence(cellvalues_u, q_point, j)
                Ke[BlockIndex((p▄, u▄), (i, j))] += -δp * divδu * dΩ
            end
            for j in 1:i
                p = shape_value(cellvalues_p, q_point, j)
                Ke[BlockIndex((p▄, p▄), (i, j))] += - 1/mp.K * δp * p * dΩ
            end

        end
    end

    symmetrize_lower!(Ke)

    ## We integrate the Neumann boundary using the facevalues.
    ## We loop over all the faces in the cell, then check if the face
    ## is in our `"traction"` faceset.
    @inbounds for face in 1:nfaces(cell)
        if onboundary(cell, face) && (cellid(cell), face) ∈ getfaceset(grid, "traction")
            reinit!(facevalues_u, cell, face)
            for q_point in 1:getnquadpoints(facevalues_u)
                dΓ = getdetJdV(facevalues_u, q_point)
                for i in 1:n_basefuncs_u
                    δu = shape_value(facevalues_u, q_point, i)
                    fe[i] += (δu ⋅ t) * dΓ
                end
            end
        end
    end
end

function symmetrize_lower!(K)
    for i in 1:size(K,1)
        for j in i+1:size(K,1)
            K[i,j] = K[j,i]
        end
    end
end;

# Now we have constructed all the necessary components, we just need a function
# to put it all together.

function solve(ν, interpolation_u, interpolation_p)
    ## material
    Emod = 1.
    Gmod = Emod / 2(1 + ν)
    Kmod = Emod * ν / ((1+ν) * (1-2ν))
    mp = LinearElasticity(Gmod, Kmod)

    ## grid, dofhandler, boundary condition
    n = 50
    grid = create_cook_grid(n, n)
    dh = create_dofhandler(grid, interpolation_u, interpolation_p)
    dbc = create_bc(dh)

    ## cellvalues
    cellvalues_u, cellvalues_p, facevalues_u = create_values(interpolation_u, interpolation_p)

    ## assembly and solve
    K = create_sparsity_pattern(dh);
    K, f = doassemble(cellvalues_u, cellvalues_p, facevalues_u, K, grid, dh, mp);
    apply!(K, f, dbc)
    u = Symmetric(K) \ f;

    ## export
    filename = "cook_" * (isa(interpolation_u, Lagrange{2,RefTetrahedron,1}) ? "linear" : "quadratic") *
                         "_linear"
    vtk_grid(filename, dh) do vtkfile
        vtk_point_data(vtkfile, dh, u)
    end
    return u
end

# All that is left is to solve the problem. We choose a value of Poissons
# ratio that is near incompressibility -- $ν = 0.5$ -- and thus expect the
# linear/linear approximation to return garbage, and the quadratic/linear
# approximation to be stable.
linear    = Lagrange{2,RefTetrahedron,1}()
quadratic = Lagrange{2,RefTetrahedron,2}()

u1 = solve(0.4999999, linear, linear)
u2 = solve(0.4999999, quadratic, linear);

## test the result                 #src
using Test                         #src
@test norm(u2) ≈ 919.2122668839389 #src

#md # ## [Plain Program](@id incompressible_elasticity-plain-program)
#md #
#md # Below follows a version of the program without any comments.
#md # The file is also available here: [incompressible_elasticity.jl](incompressible_elasticity.jl)
#md #
#md # ```julia
#md # @__CODE__
#md # ```


using Stheno, Random, Plots
using Distributions

# Where to start and stop input observations
START = 0
STOP_OBS = 1
STOP_TRUE = 1.0  # might be different to see how they generalize

DATA_SAMPLES = 30
TRUE_SAMPLES = 1000 # discretised points used to model true function
NOISE_MU = 0  # mean of noise
NOISE_SIGMA = 0.05
SAMPLES = 1  # samples we take for the GP representation

function generate_data()
    f1 = (x -> -sin.(10 * pi .* (x .+ 1)) ./ (2 .* x .+ 1) - x .^ 4)
    f2 = ((x, y1) -> cos.(y1) .^ 2 + sin.(3 .* x))
    f3 = ((x, y1, y2) -> y2 .* (y1 .^ 2) + 3 * x)
    # Generate True function
    x_true = range(START, STOP_TRUE, length = TRUE_SAMPLES)
    y1_true = f1(x_true)
    y2_true = f2(x_true, y1_true)
    y3_true = f3(x_true, y1_true, y2_true)
    y_true = hcat(y1_true, y2_true, y3_true)

    # Create the noisy normal distribution
    normal_noise = Normal(NOISE_MU, NOISE_SIGMA^2)
    # Add the noise to the observations
    x = range(START, STOP_OBS, length = DATA_SAMPLES)
    y1 = f1(x) + rand(normal_noise, DATA_SAMPLES)
    y2 = f2(x, y1) + rand(normal_noise, DATA_SAMPLES)
    y3 = f3(x, y1, y2) + rand(normal_noise, DATA_SAMPLES)
    y_obs = hcat(y1, y2, y3)

    return x, y_obs, x_true, y_true
end

x, y_obs, x_true, y_true = generate_data();

# Generate the Independent GPs
f1_gp = GP(stretch(EQ(), 10), GPC())
f2_gp = GP(EQ(), GPC())
f3_gp = GP(EQ(), GPC())

f1_gp_posterior = f1_gp | (f1_gp(x, NOISE_SIGMA^2) ← y_obs[:, 1])
f2_gp_posterior = f2_gp | (f2_gp(x, NOISE_SIGMA^2) ← y_obs[:, 2])
f3_gp_posterior = f3_gp | (f3_gp(x, NOISE_SIGMA^2) ← y_obs[:, 3])

# # Attempt to create GP modelling all outputs
# full_gp = GP(EQ() + PerEQ(0.5) , GPC())
# full_gp_post =
#     full_gp |
#     (full_gp(repeat(x, 3), NOISE_SIGMA^2) ← vcat(y_obs[:, 1], y_obs[:, 2], y_obs[:, 3]))

# Create GPAR model
f1_gpar = GP(stretch(EQ(), 10), GPC())
f2_gpar = GP(EQ(), GPC())
f3_gpar = GP(EQ(), GPC())
# Compute the posteriors
f1_gpar_post = f1_gpar | (f1_gpar(x, NOISE_SIGMA^2) ← y_obs[:, 1])
# We represent the input locations as a 2 x 30 matrix
f2_data_locations = ColVecs(transpose(hcat(x, y_obs[:, 1])))

f2_gpar_post =
    f2_gpar | (
        f2_gpar(f2_data_locations, NOISE_SIGMA^2) ← y_obs[:, 2]
    )
# We represent the input locations as a 3 x 30 matrix
f3_data_locations = ColVecs(transpose(hcat(x, y_obs[:, 1], y_obs[:, 2])))
f3_gpar_post =
    f3_gpar | (
        f3_gpar(f3_data_locations, NOISE_SIGMA^2) ← y_obs[:, 3]
    )

# Plotting
gr();
overall_plot = plot(layout = (3, 1), legend = false);
# Plot data
plot!(overall_plot, x_true, y_true, color = :orange, label = "True")
scatter!(
    overall_plot,
    x,
    y_obs,
    color = :black,
    label = "Observations",
    markersize = 3,
    markeralpha = 0.8,
)

# Plot IGP results
plot!(
    overall_plot[1],
    f1_gp_posterior(x_true),
    samples = SAMPLES,
    color = :green,
    fillalpha = 0.5,
    linealpha = 1,
)
plot!(
    overall_plot[2],
    f2_gp_posterior(x_true),
    samples = SAMPLES,
    color = :green,
    fillalpha = 0.3,
    linealpha = 1,
)
plot!(
    overall_plot[3],
    f3_gp_posterior(x_true),
    samples = SAMPLES,
    color = :green,
    fillalpha = 0.3,
    linealpha = 1,
)

# Plot GPAR
# First plot is easy
y1_mean_vals = mean(f1_gpar_post(x_true))
plot!(
    overall_plot[1],
    x_true,
    y1_mean_vals,
    color = :blue,
    linealpha = 1,
)
# Use x and mean of y1 to feed into f2_gpar
x_y1_true = ColVecs(transpose(hcat(x_true, y1_mean_vals)))
y2_mean_vals = mean(f2_gpar_post(x_y1_true))
plot!(overall_plot[2], x_true, y2_mean_vals, color = :blue, linealpha = 1)

# Plot f3
# Again, collect x, y1, and y2 to feed as input location to f3
x_y1_y2_true = ColVecs(transpose(hcat(x_true, y1_mean_vals, y2_mean_vals)))
y3_mean_vals = mean(f3_gpar_post(x_y1_y2_true))
plot!(overall_plot[3], x_true, y3_mean_vals, color = :blue, linealpha = 1)

display(overall_plot)


using GeometricEquations: _idae_default_v̄, _ldae_default_v̄

function ode_v(t, x, ẋ, params)
    ẋ[1] = x[2]
    ẋ[2] = 2x[1]
end

function sode_v1(t, x, v, params)
    v[1] = x[2]
end

function sode_v2(t, x, v, params)
    v[2] = 2x[1]
end

function sode_q1(t, x̄, x, params)
    x[1] = x̄[1]
    x[2] = x̄[2]
end

function sode_q2(t, x̄, x, params)
    x[1] = x̄[1]
    x[2] = x̄[2]
end

sode_v = (sode_v1, sode_v2)
sode_q = (sode_q1, sode_q2)

function pode_v(t, q, p, v, params)
    v[1] = p[1]
end

function pode_f(t, q, p, f, params)
    f[1] = 2q[1]
end

function pode_h(t, q, p, params)
    p[1]^2/2 + cos(q[1])
end

pode_eqs = (pode_v, pode_f)

function hode_h(t, q, p, params)
    p[1]^2/2 + cos(q[1])
end

function hode_ω(t, q, p, ω, params)
    ω[1,1] = sin(q[1])
    ω[1,2] = 0
    ω[2,1] = 0
    ω[2,2] = p[1]
end

hode_eqs = (pode_v, pode_f, hode_ω, hode_h)

function iode_ϑ(t, q, v, p, params)
    p[1] = v[1]
end

function iode_f(t, q, v, f, params)
    f[1] = sin(q[1])
end

function iode_u(t, q, λ, u, params)
    u[1] = λ[1]
end

function iode_g(t, q, λ, g, params)
    g[1] = λ[1]
end

function iode_v(t, q, p, v, params)
    v[1] = p[1]
end

function iode_h(t, q, v, params)
    v[1]^2/2 + cos(q[1])
end

iode_eqs = (iode_ϑ, iode_f, iode_g)

function lode_l(t, q, v, params)
    v[1]^2/2 - cos(q[1])
end

function lode_ω(t, q, v, ω, params)
    ω[1,1] = sin(q[1])
    ω[1,2] = 0
    ω[2,1] = 0
    ω[2,2] = v[1]
end

lode_eqs = (iode_ϑ, iode_f, iode_g, lode_ω, lode_l)

function dae_v(t, x, v, params)
    v[1] = x[1]
    v[2] = x[2]
end

function dae_u(t, x, λ, u, params)
    u[1] = +λ[1]
    u[2] = -λ[1]
end

function dae_ū(t, x, λ, u, params)
    u[1] = +λ[1]
    u[2] = -λ[1]
end

function dae_ϕ(t, x, ϕ, params)
    ϕ[1] = x[2] - x[1]
end

function dae_ψ(t, x, v, ψ, params)
    ψ[1] = v[2] - v[1]
end

function pdae_v(t, q, p, v, params)
    v[1] = p[1]
end

function pdae_f(t, q, p, f, params)
    f[1] = q[1]
end

function pdae_p(t, q, v, p, params)
    p[1] = v[1]
end

function pdae_u(t, q, p, λ, u, params)
    u[1] = λ[1]
end

function pdae_g(t, q, p, λ, g, params)
    g[1] = λ[1]
end

function pdae_ϕ(t, q, p, λ, ϕ, params)
    ϕ[1] = p[1] - q[1]
end

function pdae_ψ(t, q, p, λ, μ, ψ, params)
    ψ[1] = μ[1] - λ[1]
end

function pdae_h(t, q, p)
    p[1]^2/2 + q[1]^2/2
end

hdae_ω = hode_ω

idae_ϑ = iode_ϑ
idae_f = iode_f
idae_v = _idae_default_v̄

ldae_l = lode_l
ldae_ω = lode_ω
ldae_v = _ldae_default_v̄

function sde_v(λ, t, q, v)
    v[1] = λ*q[1]
    v[2] = λ*q[2]
end

function sde_B(μ, t, q, u)
    u[1] = μ*q[1]
    u[2] = μ*q[2]
end

function psde_v(t, q, p, v)
    v[1] = p[1]
end

function psde_f(t, q, p, f)
    f[1] = - q[1]
end

function psde_B(t, q, p, B)
    B[1,1] = noise_intensity * p[1]
end

function psde_G(t, q, p, G)
    G[1,1] = - noise_intensity * q[1]
end

function spsde_v(t, q, p, v)
    v[1] = p[1]
end

function spsde_f1(t, q, p, f)
    f[1] = - q[1]
end

function spsde_f2(t, q, p, f)
    f[1] = 0
end

function spsde_B(t, q, p, B)
    B[1,1] = noise_intensity * p[1]
end

function spsde_G1(t, q, p, G)
    G[1,1] = - noise_intensity * q[1]
end

function spsde_G2(t, q, p, G)
    G[1,1] = 0
end