text:
using ChainMap

"""
    code_lines(file_in)

A function which extracts just the lines containing code from markdown.

```julia
file_in = joinpath("C:\\", "Users", "jsnot",
                   ".julia", "v0.4", "ChainMap",
                   "src", "chain.jl")
```
"""
function code_lines(file_in)
  text = readlines(file_in)

  starts = @chain begin
      text
      @map @chain begin
          chomp(_)
          ismatch(r"^```.+", _)
      end
  end

  ends = @chain begin
      text
      @map @chain begin
          chomp(_)
          _ == "```"
      end
   end

  Test.@test sum(starts) == sum(ends)

  @chain begin
      cumsum(starts) - cumsum(ends) .== 1
      text[_]
      @lambda filter !startswith(_, "```")
      @map replace(_, r"\\", "")
  end

end

"""
    make_tests(path, head = "")

Populates your runtest.jl file with all of the julia code in your docstrings.

An optional header string, or vector of header strings, `head`, can be added.
`path` is the path to your package.

# Examples
```julia
path_in = "C:\\Users\\jsnot\\.julia\\v0.4\\ChainMap"
head = ["using ChainMap", "import DataStructures", "import NullableArrays",
        "import DataFrames", "using Base.Test"]
make_tests(path, head)
```
"""
function make_tests(path_in, head = "")
    head_cat = @chain begin
        head
        @map _ * "\n"
        *(_...)
    end

    @chain begin
        path_in
        joinpath(_, "src")
        readdir(_)
        @map joinpath(path_in, "src", _)
        [_; joinpath(path_in, "README.md") ]
        @map code_lines(_)
        vcat(head_cat, _...)
        write(joinpath(path_in, "test", "runtests.jl"), _)
    end
end

path = "C:\\Users\\jsnot\\.julia\\v0.4\\ChainMap"
head = ["using ChainMap", "import DataStructures", "import NullableArrays",
        "import DataFrames", "using Base.Test"]
make_tests(path, head)


text:
import Test

include("slq.jl")

ext: jl


text:
using Printf

function bench_time(fctn, times, name)
    real_time = @elapsed for i = 1 : times
	fctn()
    end
    return name, real_time, real_time, 0, 0
end

function run_(name, times, fctn)
    name, rt, ut, sys, consed = bench_time(fctn, times, name)
    @printf ";;; running #<benchmark %s for %s runs>\n" name times
    @printf "(\"%s\" %s %s %d %d)\n" name rt ut sys consed
end

const fib_iter = 25

function fact(n)
    if 0==n
        1
    else
        n * fact(n - 1)
    end
end

function run_factorial()
    fact(big(500))
end

function fib(n)
    if n < 2
        1
    else
        fib(n-1)+fib(n-2)
    end
end

function run_fib()
    fib(fib_iter)
end

function ackermann(m, n)
    if 0==m
        n+1
    elseif 0==n
        ackermann(m-1, 1)
    else
        ackermann(m-1, ackermann(m, n-1))
    end
end

function run_ackermann()
    ackermann(3, 11)
end

function tak(x, y, z)
    if !(y < x)
        z
    else
        tak(tak(x-1, y, z),
            tak(y-1, z, x),
            tak(z-1, x, y))
    end
end

function run_tak()
    tak(18, 12, 6)
end

function bench_1d_arrays(size, runs)
    ones = fill(UInt16(1), size)
    twos = fill(UInt16(2), size)
    threes = Vector{UInt16}(undef, size)
    for run = 1 : runs
	for pos = 1 : size
	    threes[pos] = ones[pos] + twos[pos]
	end
	for elt = threes
	    @assert elt == 3 @sprintf "The assertion %s == 3 failed." elt
        end
    end
end

function run_bench_1d_arrays()
    bench_1d_arrays(100000, 10)
end

function bench_2d_arrays(size, runs)
    ones = fill(UInt16(1), (size, size))
    twos = fill(UInt16(2), (size, size))
    threes = Array{UInt16, 2}(undef, (size, size))
    for run = 1 : runs
        for i = 1 : size
            for j = 1 : size
                threes[i, j] = ones[i, j] + twos[i, j]
            end
        end
        for i = 1 : size
            for j = 1 : size
                elt = threes[i, j]
                @assert elt == 3 @sprintf "The assertion %s == 3 failed." elt
            end
        end
    end
end

function run_bench_2d_arrays()
    bench_2d_arrays(2000, 10)
end

function fib_single_float(n)
    if n < Float32(2)
        Float32(1)
    else
	fib_single_float(n-Float32(1)) + fib_single_float(n-Float32(2))
    end
end

function run_fib_single_float()
    fib_single_float(Float32(fib_iter))
end

function fib_double_float(n)
    if n < Float64(2)
	Float64(1)
    else
        fib_double_float(n-Float64(1)) + fib_double_float(n-Float64(2))
    end
end

function run_fib_double_float()
    fib_double_float(Float64(fib_iter))
endn

function fib_ratio(n)
    function fr(n)
        if n == 1
            1
        else
            1 + Rational{BigInt}(1//fr(n-1))
        end
    end
    numerator(fr(n))
end

function run_fib_ratio()
    fib_ratio(150)
end

function listn(n)
    fill(0, n)
end

function shorterp(x, y)
    0<length(y) && (0==length(x) || shorterp(x[2:end], y[2:end]))
end

function mas(x, y, z)
    if !shorterp(y, x)
        z
    else
        mas(mas(x[2:end], y, z),
            mas(y[2:end], z, x),
            mas(z[2:end], x, y))
    end
end

function run_takl()
    l18 = listn(18)
    l12 = listn(12)
    l6 = listn(6)
    mas(l18, l12, l6)
end

function bench_string_concat(size, runs)
    for _ = 1 : runs
        buf = IOBuffer()
        for _ = 1 : size
            write(buf, "hi there!")
        end
        buflen = length(String(take!(buf)))
        # buflen = length(repeat("hi there!", size)) # much faster
        @assert buflen == length("hi there!")*size
                @sprintf "buflen:%s == size*runs:%s" buflen length("hi there!")*size
    end
end

function crc_division_step(bit, rmdr, poly, msb_mask)
    # Shift in the bit into the LSB of the register (rmdr)
    new_rmdr = Int64(bit | (rmdr * 2))
    # Divide by the polynomial, and return the new remainder
    if (0 == (msb_mask & new_rmdr))
        new_rmdr
    else
        xor(new_rmdr, poly)
    end
end

function integer_length(n)
    if n < 0
        n = -n
    else
        n += 1
    end
    ans = log2(n)
    Int64(ceil(ans))
end

function compute_adjustment(poly,n)
    # Precompute X^(n-1) mod poly
    poly_len_mask = 1 << (integer_length(poly) - 1)
    rmdr = crc_division_step(1, 0, poly, poly_len_mask)
    for k = 1 : n-1
        rmdr = crc_division_step(0, rmdr, poly, poly_len_mask)
    end
    rmdr
end

function calculate_crc40(iterations)
    crc_poly = 1099587256329
    len = 3014633
    answer = 0
    for k = 1 : iterations
        answer = compute_adjustment(crc_poly, len)
    end
  return answer;
end

function run_crc40()
    calculate_crc40(10) #549793628164
end

function main()
    @printf "\n(\"Julia %s\"\n" VERSION
    run_("CRC40", 2, run_crc40)
    run_("1D-ARRAYS", 1, run_bench_1d_arrays)
    run_("2D-ARRAYS", 1, run_bench_2d_arrays)
    run_("FIB", 50, run_fib)
    run_("FIB-RATIO", 500, run_fib_ratio)
    run_("FIB-SINGLE-FLOAT", 50, run_fib_single_float)
    run_("FIB-DOUBLE-FLOAT", 50, run_fib_double_float)
    run_("ACKERMANN", 1, run_ackermann)
    run_("TAK", 1000, run_tak)
    run_("TAKL", 150, run_takl)
    run_("FACTORIAL", 1000, run_factorial)
    run_("STRING-CONCAT", 1, (() -> bench_string_concat(1000000, 100)))
    @printf ")\n"
end

#main()


text:
"""
    ADVector{T}

A cache of vectors for making it easier to ForwardDiff. A vector of the desired type
    is extracted by indexing:

    vec = adv[eltype(x)]

Since ForwardDiff Duals are parameterized on the chunk size, this type will store any
number of vectors for different chunk sizes. A new vector of duals will be automatically
created (once) if it doesn't exist yet.

For best performance (by avoiding type-instability), annotate the indexing operation:

    vec = adv[T]::Vector{T}

While the annotation could have been included in the indexing operation itself, some
benchmark results found it significantly faster to apply the annotation at the call site.

# Constructor

    ADVector{T}(n)

where `n` is the size of the vectors, and `T` is the numeric type (typically `Float64`).
"""
struct ADVector{T}
    v::Vector{T}
    d::Dict{DataType,Vector{D} where {D<:ForwardDiff.Dual{Nothing,T}}}
    function ADVector{T}(n::Integer) where {T}
        new{T}(zeros(T, n), Dict{DataType,Vector{D} where D}())
    end

    # Copy constructor
    # Does a deep copy of the cache dictionary
    function ADVector(adv::ADVector{T}) where T
        new{T}(copy(adv.v), deepcopy(adv.d))
    end
end
Base.getindex(adv::ADVector, ::Type{<:Number}) = adv.v
function Base.getindex(adv::ADVector, T::Type{<:ForwardDiff.Dual})
    if haskey(adv.d, T)
        return adv.d[T]
    else
        d = zeros(T, length(adv.v))
        adv.d[T] = d
        return d
    end
end
Base.copy(adv::ADVector{T}) where T = ADVector(adv)  # calls copy constructor

########################################
# Explicit Methods
########################################

"""
    Euler

Explicit Euler integration:

```math
x_{k+1} = x_k + h f(x_k, u_k)
```
where ``h`` is the time step.

!!! warning
    In general, explicit Euler integration **SHOULD NOT BE USED!** It is the worst possible
    integration method since it is very inaccurate and can easily go unstable.
"""
struct Euler <: Explicit
    Euler(n::Integer, m::Integer) = new()
end

function integrate(::Euler, model, x, u, t, h)
    xdot = dynamics(model, x, u, t)
    return x + h * xdot
end

function integrate!(int::Euler, model, xn, x, u, t, h)
    dynamics!(model, xn, x, u, t)
    xn .*= h
    xn .+= x
    return nothing
end

function jacobian!(int::Euler, sig::FunctionSignature, model::ContinuousDynamics, J, xn, x, u, t, h)
    # Call the user-defined Continuous-time Jacobian
    jacobian!(model, J, xn, x, u, t)
    J .*= h
    for i = 1:state_dim(model)
        J[i, i] += 1.0
    end
    return nothing
end

@doc raw"""
    RK3

A third-order explicit Runge-Kutta method:

```math
\begin{aligned}
&k_1 = f(x_k, u_k, t) h \\
&k_2 = f(x_k + \frac{1}{2} k_1, u_k, t, + \frac{1}{2} h) h \\
&k_3 = f(x_k - k_1 + 2 k_2, u_k, t, + h) h \\
&x_{k+1} = x_k + \frac{1}{6} (k_1 + 4 k_2 + k_3)
\end{aligned}
```
"""
struct RK3 <: Explicit
    k1::ADVector{Float64}
    k2::ADVector{Float64}
    k3::ADVector{Float64}
    A::Vector{Matrix{Float64}}
    B::Vector{Matrix{Float64}}
    dA::Vector{Matrix{Float64}}
    dB::Vector{Matrix{Float64}}
end
function RK3(n::Integer, m::Integer)
    k1, k2 = ADVector{Float64}(n), ADVector{Float64}(n)
    k3 = ADVector{Float64}(n)
    A = [zeros(n, n) for i = 1:3]
    B = [zeros(n, m) for i = 1:3]
    dA = [zeros(n, n) for i = 1:3]
    dB = [zeros(n, m) for i = 1:3]
    RK3(k1, k2, k3, A, B, dA, dB)
end
getks(int::RK3, ::Type{T}) where {T} =
    int.k1[T]::Vector{T}, int.k2[T]::Vector{T}, int.k3[T]::Vector{T}

function integrate(::RK3, model, x, u, t, h)
    k1 = dynamics(model, x,            u, t      ) * h
    k2 = dynamics(model, x + k1 / 2,   u, t + h/2) * h
    k3 = dynamics(model, x - k1 + 2k2, u, t + h  ) * h
    return x + (k1 + 4k2 + k3) / 6
end

function integrate!(int::RK3, model, xn, x, u, t, h)
    T = eltype(xn)
    k1, k2, k3 = getks(int, T)
    dynamics!(model, k1, x, u, t)
    @. xn = x + k1 * h / 2
    dynamics!(model, k2, xn, u, t + h / 2)
    @. xn = x - k1 * h + 2 * k2 * h
    dynamics!(model, k3, xn, u, t + h)
    @. xn = x + h * (k1 + 4k2 + k3) / 6
    return nothing
end

function jacobian!(int::RK3, sig::StaticReturn, model, J, xn, x, u, t, h)
    n, m = dims(model)
    ix = SVector{n}(1:n)
    iu = SVector{m}(n+1:n+m)
    k1 = dynamics(model, x,            u, t      ) * h
    k2 = dynamics(model, x + k1 / 2,   u, t + h/2) * h

    jacobian!(model, J, xn, x, u, t)
    A1, B1 = J[ix, ix], J[ix, iu]

    jacobian!(model, J, xn, x + k1 / 2, u, t + h / 2)
    A2, B2 = J[ix, ix], J[ix, iu]

    jacobian!(model, J, xn, x - k1 + 2k2, u, t + h)
    A3, B3 = J[ix, ix], J[ix, iu]

    dA1 = A1 * h
    dA2 = A2 * (I + 0.5 * dA1) * h
    dA3 = A3 * (I - dA1 + 2 * dA2) * h

    dB1 = B1 * h
    dB2 = B2 * h + 0.5 * A2 * dB1 * h
    dB3 = B3 * h + A3 * (2dB2 - dB1) * h

    J[ix, ix] .= I + (dA1 + 4dA2 + dA3) / 6
    J[ix, iu] .= (dB1 + 4dB2 + dB3) / 6

    return nothing
end

function jacobian!(int::RK3, sig::InPlace, model, J, xn, x, u, t, h)
    # x,u,t,h = state(z), control(z), time(z), timestep(z)
    k1, k2, k3 = getks(int, Float64)
    A1, A2, A3 = int.A[1], int.A[2], int.A[3]
    B1, B2, B3 = int.B[1], int.B[2], int.B[3]
    dA1, dA2, dA3 = int.dA[1], int.dA[2], int.dA[3]
    dB1, dB2, dB3 = int.dB[1], int.dB[2], int.dB[3]
    n, m = dims(model)
    ix, iu = 1:n, n+1:n+m

    jacobian!(model, J, k1, x, u, t)
    dynamics!(model,    k1, x, u, t)
    A1 .= @view J[ix, ix]
    B1 .= @view J[ix, iu]

    @. xn = x + k1 * h / 2
    jacobian!(model, J, k2, xn, u, t + h / 2)
    dynamics!(model,    k2, xn, u, t + h / 2)
    A2 .= @view J[ix, ix]
    B2 .= @view J[ix, iu]

    @. xn = x - k1 * h + 2 * k2 * h
    jacobian!(model, J, k3, xn, u, t + h)
    dynamics!(model,    k3, xn, u, t + h)
    A3 .= @view J[ix, ix]
    B3 .= @view J[ix, iu]

    # dA = A1 * h
    dA1 .= A1 .* h

    # dA2 = A2 * (I + 0.5 * dA1) * h
    mul!(dA2, A2, dA1, 0.5, 0.0)
    dA2 .+= A2
    dA2 .*= h

    # dA3 = A3 * (I - dA1 + 2 * dA2) * h
    mul!(dA3, A3, dA2, 2.0, 0.0)
    mul!(dA3, A3, dA1,-1.0, 1.0)
    dA3 .+= A3
    dA3 .*= h

    # dB1 = B1 * h
    dB1 .= B1 .* h

    # dB2 = B2 * h + 0.5 * A2 * dB1 * h
    dB2 .= B2
    mul!(dB2, A2, dB1, 0.5, 1.0)
    dB2 .*= h

    # dB3 = B3 * h + A3 * (2dB2 - dB1) * h
    dB3 .= B3
    mul!(dB3, A3, dB2, 2.0, 1.0)
    mul!(dB3, A3, dB1,-1.0, 1.0)
    dB3 .*= h

    @. J[ix, ix] = (dA1 + 4dA2 + dA3) / 6
    for i = 1:n
        J[i, i] += 1.0
    end
    @. J[ix, iu] = (dB1 + 4dB2 + dB3) / 6

    return nothing
end

@doc raw"""
    RK4

The classic fourth-order explicit Runge-Kutta method.

```math
\begin{aligned}
& k_1 = f(x_k, u_k, t) h \\
& k_2 = f(x_k + \frac{1}{2} k_1, u_k, t + \frac{1}{2} h) h \\
& k_3 = f(x_k + \frac{1}{2} k_2, u_k, t + \frac{1}{2} h) h \\
& k_4 = f(x_k +  k_3, u_k, t + h) h \\
& x_{k+1} = x_k + \frac{1}{6}(k_1 + 2k_2 + 2k_3 + k_4)
\end{aligned}
```
"""
struct RK4 <: Explicit
    k1::ADVector{Float64}
    k2::ADVector{Float64}
    k3::ADVector{Float64}
    k4::ADVector{Float64}
    A::Vector{Matrix{Float64}}
    B::Vector{Matrix{Float64}}
    dA::Vector{Matrix{Float64}}
    dB::Vector{Matrix{Float64}}
end
function RK4(n::Integer, m::Integer)
    k1, k2 = ADVector{Float64}(n), ADVector{Float64}(n)
    k3, k4 = ADVector{Float64}(n), ADVector{Float64}(n)
    A = [zeros(n, n) for i = 1:4]
    B = [zeros(n, m) for i = 1:4]
    dA = [zeros(n, n) for i = 1:4]
    dB = [zeros(n, m) for i = 1:4]
    RK4(k1, k2, k3, k4, A, B, dA, dB)
end
getks(int::RK4, ::Type{T}) where {T} =
    int.k1[T]::Vector{T}, int.k2[T]::Vector{T}, int.k3[T]::Vector{T}, int.k4[T]::Vector{T}

function integrate(::RK4, model, x, u, t, h)
    k1 = dynamics(model, x, u, t) * h
    k2 = dynamics(model, x + k1 / 2, u, t + h / 2) * h
    k3 = dynamics(model, x + k2 / 2, u, t + h / 2) * h
    k4 = dynamics(model, x + k3, u, t + h) * h
    x + (k1 + 2k2 + 2k3 + k4) / 6
end

function integrate!(int::RK4, model, xn, x, u, t, h)
    T = eltype(xn)
    k1, k2, k3, k4 = getks(int, T)
    dynamics!(model, k1, x, u, t)
    @. xn = x + k1 * h / 2
    dynamics!(model, k2, xn, u, t + h / 2)
    @. xn = x + k2 * h / 2
    dynamics!(model, k3, xn, u, t + h / 2)
    @. xn = x + k3 * h
    dynamics!(model, k4, xn, u, t + h)
    @. xn = x + h * (k1 + 2k2 + 2k3 + k4) / 6
    return nothing
end

function jacobian!(int::RK4, sig::StaticReturn, model, J, xn, x, u, t, h)
    n, m = dims(model)
    ix = SVector{n}(1:n)
    iu = SVector{m}(n+1:n+m)
    k1 = dynamics(model, x, u, t) * h
    k2 = dynamics(model, x + k1 / 2, u, t + h / 2) * h
    k3 = dynamics(model, x + k2 / 2, u, t + h / 2) * h

    jacobian!(model, J, xn, x, u, t)
    A1, B1 = J[ix, ix], J[ix, iu]

    jacobian!(model, J, xn, x + k1 / 2, u, t + h / 2)
    A2, B2 = J[ix, ix], J[ix, iu]

    jacobian!(model, J, xn, x + k2 / 2, u, t + h / 2)
    A3, B3 = J[ix, ix], J[ix, iu]

    jacobian!(model, J, xn, x + k3, u, t + h)
    A4, B4 = J[ix, ix], J[ix, iu]

    dA1 = A1 * h
    dA2 = A2 * (I + 0.5 * dA1) * h
    dA3 = A3 * (I + 0.5 * dA2) * h
    dA4 = A4 * (I + dA3) * h

    dB1 = B1 * h
    dB2 = B2 * h + 0.5 * A2 * dB1 * h
    dB3 = B3 * h + 0.5 * A3 * dB2 * h
    dB4 = B4 * h + A4 * dB3 * h

    J[ix, ix] .= I + (dA1 + 2dA2 + 2dA3 + dA4) / 6
    J[ix, iu] .= (dB1 + 2dB2 + 2dB3 + dB4) / 6

    return nothing
end

function jacobian!(int::RK4, sig::InPlace, model, J, xn, x, u, t, h)
    # x,u,t,h = state(z), control(z), time(z), timestep(z)
    k1, k2, k3, k4 = getks(int, Float64)
    A1, A2, A3, A4 = int.A[1], int.A[2], int.A[3], int.A[4]
    B1, B2, B3, B4 = int.B[1], int.B[2], int.B[3], int.B[4]
    dA1, dA2, dA3, dA4 = int.dA[1], int.dA[2], int.dA[3], int.dA[4]
    dB1, dB2, dB3, dB4 = int.dB[1], int.dB[2], int.dB[3], int.dB[4]
    n, m = dims(model)
    ix, iu = 1:n, n+1:n+m

    jacobian!(model, J, k1, x, u, t)
    dynamics!(model,    k1, x, u, t)
    A1 .= @view J[ix, ix]
    B1 .= @view J[ix, iu]

    @. xn = x + k1 * h / 2
    jacobian!(model, J, k2, xn, u, t + h / 2)
    dynamics!(model,    k2, xn, u, t + h / 2)
    A2 .= @view J[ix, ix]
    B2 .= @view J[ix, iu]

    @. xn = x + k2 * h / 2
    jacobian!(model, J, k3, xn, u, t + h / 2)
    dynamics!(model,    k3, xn, u, t + h / 2)
    A3 .= @view J[ix, ix]
    B3 .= @view J[ix, iu]

    @. xn = x + k3 * h
    jacobian!(model, J, k4, xn, u, t + h)
    A4 .= @view J[ix, ix]
    B4 .= @view J[ix, iu]

    # dA = A1 * h
    dA1 .= A1 .* h

    # dA2 = A2 * (I + 0.5 * dA1) * h
    mul!(dA2, A2, dA1, 0.5, 0.0)
    dA2 .+= A2
    dA2 .*= h

    # dA3 = A3 * (I + 0.5 * dA2) * h
    mul!(dA3, A3, dA2, 0.5, 0.0)
    dA3 .+= A3
    dA3 .*= h

    # dA4 = A4 * (I + dA3) * h
    mul!(dA4, A4, dA3, 1.0, 0.0)
    dA4 .+= A4
    dA4 .*= h

    # dB1 = B1 * h
    dB1 .= B1 .* h

    # dB2 = B2 * h + 0.5 * A2 * dB1 * h
    dB2 .= B2
    mul!(dB2, A2, dB1, 0.5, 1.0)
    dB2 .*= h

    # dB3 = B3 * h + 0.5 * A3 * dB2 * h
    dB3 .= B3
    mul!(dB3, A3, dB2, 0.5, 1.0)
    dB3 .*= h

    # dB4 = B4 * h + A4 * dB3 * h
    mul!(dB4, A4, dB3)
    dB4 .+= B4
    dB4 .*= h

    @. J[ix, ix] = (dA1 + 2dA2 + 2dA3 + dA4) / 6
    for i = 1:n
        J[i, i] += 1.0
    end
    @. J[ix, iu] = (dB1 + 2dB2 + 2dB3 + dB4) / 6

    return nothing
end

########################################
# Implicit Methods
########################################

# Use Newton's method to solve for the next state for implicit dynamics
function integrate(integrator::Implicit, model::ImplicitDynamicsModel,
                   z::AbstractKnotPoint{Nx,Nu}) where {Nx,Nu}
    cache = getnewtoncache(integrator)
    newton_iters = cache.newton_iters
    tol = cache.newton_tol

    J2 = cache.J2
    J1 = cache.J1
    y2 = cache.y2
    y1 = cache.y1
    z1 = z
    z2 = StaticKnotPoint(z)
    ix = SVector{Nx}(1:Nx)

    # Use the current state as the current guess
    xn = state(z1)

    diff = default_diffmethod(model.continuous_dynamics)
    for iter = 1:newton_iters
        # Set the guess for the next state
        z2 = setstate(z2, xn)

        # Calculate the residual
        r = dynamics_error(integrator, model.continuous_dynamics, z2, z1)

        # Calculate the Jacobian wrt x2
        dynamics_error_jacobian!(StaticReturn(), diff, model,
                                 J2, J1, y2, y1, z2, z1)
        A = J2[ix, ix]

        if norm(r) < tol
            break
        end

        # Get the step
        F = lu(A)
        dx = F \ r
        xn -= dx
    end
    setdata!(cache.z2, getdata(z))  # copy input into cache for Jacobian check
    return xn
end

function integrate!(integrator::Implicit, model::ImplicitDynamicsModel, xn,
                    z::AbstractKnotPoint)
    integrate!(integrator, model, xn, state(z), control(z), z.t, z.dt)
end
function integrate!(integrator::Implicit, model::ImplicitDynamicsModel, xn,
                    x, u, t, h)
    cache = getnewtoncache(integrator)
    newton_iters = cache.newton_iters
    tol = cache.newton_tol

    z1 = cache.z1
    z2 = cache.z2
    # copyto!(z1, z)
    setstate!(z1, x)
    setcontrol!(z1, u)
    z1.t = t
    z1.dt = h
    copyto!(z2, z1)

    n,m = dims(model)
    J2 = cache.J2
    J1 = cache.J1
    r = cache.y2
    dx = cache.y1
    ipiv = cache.ipiv
    A = cache.A
    Aview = @view J2[:, 1:n]

    # Use the current state as the guess
    copyto!(xn, state(z1))

    diff = default_diffmethod(model.continuous_dynamics)
    for iter = 1:newton_iters
        # Set the guess for the next state
        setstate!(z2, xn)

        # Calculate the residual
        dynamics_error!(integrator, model.continuous_dynamics, r, dx, z2, z1)

        if norm(r) < tol
            break
        end

        # Calculate the Jacobian wrt x2
        dynamics_error_jacobian!(InPlace(), diff, model, J2, J1, r, dx, z2, z1)
        A .= Aview

        # Get the step
        dx .= r
        F = lu!(A, ipiv)
        ldiv!(F, dx)

        # Apply the step
        xn .-= dx
    end
    setdata!(cache.z2, getdata(z1))  # copy input to cache for Jacobian check
end

# Use Implicit Function Theorem to calculate the dynamics Jacobians
function jacobian!(integrator::Implicit, ::StaticReturn, diff::DiffMethod,
                   model::ImplicitDynamicsModel, J, y, z::AbstractKnotPoint{Nx,Nu}
                   ) where {Nx,Nu}
    cache = getnewtoncache(integrator)
    J2 = cache.J2
    J1 = cache.J1
    ix = SVector{Nx}(1:Nx)

    # Update Jacobian
    aresame = maxdiff(cache.z2, z) < √eps()
    if !aresame
        @debug "Solving for next state to get Jacobian using IFT."
        evaluate(model, z)
    else
        @debug "Using cached Static Factorization"
    end
    A2 = J2[ix,ix]
    Jstatic = SMatrix{Nx,Nx+Nu}(J1)
    J .= A2 \ Jstatic
    J .*= -1
    return
end

function jacobian!(integrator::Implicit, ::InPlace, diff::DiffMethod,
                   model::ImplicitDynamicsModel, J, y, z::AbstractKnotPoint)
    n,m = dims(z)
    cache = getnewtoncache(integrator)
    J1 = cache.J1

    aresame = maxdiff(cache.z2, z) < √eps()
    local F
    if !aresame
        @debug "Solving for next state to get Jacobian using IFT."
        y .= state(z)
        evaluate!(model, y, z)
    else
        @debug "Using cached Factorization"
    end
    F = cache.F
    J .= J1
    J .*= -1
    ldiv!(F, J)
    return
end

"""
    ImplicitNewtonCache

A cache for the temporary variables needed while solving for the next state
    using Newton's method for implicit dynamics. Also used to calculate the
    Jacobians using the implicit function theorem.

Also provides a couple options for controlling the behavior of the Newton solve.

Every implicit integrator should store this cache internally and define it's
getter method `getnewtoncache(::Implicit)`.
"""
mutable struct ImplicitNewtonCache
    J2::Matrix{Float64}
    J1::Matrix{Float64}
    y2::Vector{Float64}
    y1::Vector{Float64}
    z2::StaticKnotPoint{Any,Any,Vector{Float64},Float64}
    z1::KnotPoint{Any,Any,Vector{Float64},Float64}
    ipiv::Vector{BlasInt}
    A::Matrix{Float64}
    F::LinearAlgebra.LU{Float64, Matrix{Float64}}
    newton_iters::Int    # number of newton iterations
    newton_tol::Float64  # Newton tolerance
end
function ImplicitNewtonCache(n::Integer, m::Integer)
    J2 = zeros(n,n+m)
    J1 = zeros(n,n+m)
    y2 = zeros(n)
    y1 = zeros(n)
    v = zeros(n+m)
    z2 = StaticKnotPoint{Any,Any}(n, m, v, 0.0, NaN)
    z1 = KnotPoint{Any,Any}(n, m, copy(v), 0.0, NaN)
    ipiv = zeros(BlasInt, n)
    A = zeros(n,n)
    F = lu!(A, check=false)
    iters = 10    # Default number of Newton iterations
    tol = 1e-12   # Default Newton tolerance
    ImplicitNewtonCache(J2, J1, y2, y1, z2, z1, ipiv, A, F, iters, tol)
end

"""
    ImplicitMidpoint

A symplectic method with second-order accuracy. A great option for those wanting
good performance with few calls to the dynamics.

```math
x_1 + h f(\\frac{1}{2}(x_1 + x_2), u_1, t + \\frac{1}{2} h) - x_2 = 0
````
"""
struct ImplicitMidpoint <: Implicit
    xmid::ADVector{Float64}
    cache::ImplicitNewtonCache
end
function ImplicitMidpoint(n::Integer, m::Integer)
    cache = ImplicitNewtonCache(n, m)
    ImplicitMidpoint(ADVector{Float64}(n), cache)
end

getnewtoncache(integrator::ImplicitMidpoint) = integrator.cache

function dynamics_error(
    ::ImplicitMidpoint,
    model::ContinuousDynamics,
    z2::AbstractKnotPoint,
    z1::AbstractKnotPoint,
)
    t, h = time(z1), timestep(z1)
    x1, u1 = state(z1), control(z1)
    x2, u2 = state(z2), control(z2)
    xmid = (x1 + x2) / 2
    fmid = dynamics(model, xmid, u1, t + h / 2)
    x1 + h * fmid - x2
end

function dynamics_error!(
    int::ImplicitMidpoint,
    model::ContinuousDynamics,
    y2,
    y1,
    z2::AbstractKnotPoint,
    z1::AbstractKnotPoint,
)
    T = eltype(y2)
    t, h = time(z1), timestep(z1)
    x1, u1 = state(z1), control(z1)
    x2, u2 = state(z2), control(z2)
    xmid = int.xmid[T]::Vector{T}

    @. xmid = (x1 + x2) / 2
    dynamics!(model, y2, xmid, u1, t + h / 2)
    @. y2 = x1 + h * y2 - x2
    return nothing
end

function dynamics_error_jacobian!(
    int::ImplicitMidpoint,
    ::StaticReturn,
    model::ContinuousDynamics,
    J2,
    J1,
    y2,
    y1,
    z2::AbstractKnotPoint,
    z1::AbstractKnotPoint,
)
    n, m = dims(model)
    ix = SVector{n}(1:n)
    iu = SVector{m}(n+1:n+m)
    t, h = time(z1), timestep(z1)
    x1, u1 = state(z1), control(z1)
    x2, u2 = state(z2), control(z2)

    xmid = (x1 + x2) / 2
    jacobian!(model, J1, y1, xmid, u1, t)

    J1 .*= h
    J1[ix,ix] ./= 2
    J2[ix,ix] .= J1[ix,ix]
    J2[ix,iu] .= 0
    for i = 1:n
        J1[i,i] += 1.0
        J2[i,i] -= 1.0
    end
    return nothing
end

function dynamics_error_jacobian!(
    int::ImplicitMidpoint,
    ::InPlace,
    model::ContinuousDynamics,
    J2,
    J1,
    y2,
    y1,
    z2::AbstractKnotPoint,
    z1::AbstractKnotPoint,
)
    n, m = dims(model)
    ix = 1:n
    iu = n+1:n+m
    t, h = time(z1), timestep(z1)
    x1, u1 = state(z1), control(z1)
    x2, u2 = state(z2), control(z2)
    xmid = int.xmid[Float64]

    @. xmid = (x1 + x2) / 2
    jacobian!(model, J1, y1, xmid, u1, t)
    J1 .*= h
    A1 = view(J1, ix, ix)
    A2 = view(J2, ix, ix)
    A1 ./= 2
    A2 .= A1
    view(J2, ix, iu) .= 0
    for i = 1:n
        J1[i,i] += 1.0
        J2[i,i] -= 1.0
    end
    return nothing
end


text:
"""
    Wishart(ν, S)
```julia
ν::Real           degrees of freedom (whole number or a real number greater than p - 1)
S::AbstractPDMat  p x p scale matrix
```
The [Wishart distribution](http://en.wikipedia.org/wiki/Wishart_distribution)
generalizes the gamma distribution to ``p\\times p`` real, positive semidefinite
matrices ``\\mathbf{H}``.

If ``\\nu>p-1``, then ``\\mathbf{H}\\sim \\textrm{W}_p(\\nu, \\mathbf{S})``
has rank ``p`` and its probability density function is

```math
f(\\mathbf{H};\\nu,\\mathbf{S}) = \\frac{1}{2^{\\nu p/2} \\left|\\mathbf{S}\\right|^{\\nu/2} \\Gamma_p\\left(\\frac {\\nu}{2}\\right ) }{\\left|\\mathbf{H}\\right|}^{(\\nu-p-1)/2} e^{-(1/2)\\operatorname{tr}(\\mathbf{S}^{-1}\\mathbf{H})}.
```

If ``\\nu\\leq p-1``, then ``\\mathbf{H}`` is rank ``\\nu`` and it has
a density with respect to a suitably chosen volume element on the space of
positive semidefinite matrices. See [here](https://doi.org/10.1214/aos/1176325375).

For integer ``\\nu``, a random matrix given by

```math
\\mathbf{H} = \\mathbf{X}\\mathbf{X}^{\\rm{T}},
\\quad\\mathbf{X} \\sim \\textrm{MN}_{p,\\nu}(\\mathbf{0}, \\mathbf{S}, \\mathbf{I}_{\\nu})
```

has ``\\mathbf{H}\\sim \\textrm{W}_p(\\nu, \\mathbf{S})``.
For non-integer ``\\nu``, Wishart matrices can be generated via the
[Bartlett decomposition](https://en.wikipedia.org/wiki/Wishart_distribution#Bartlett_decomposition).
"""
struct Wishart{T<:Real, ST<:AbstractPDMat, R<:Integer} <: ContinuousMatrixDistribution
    df::T          # degree of freedom
    S::ST          # the scale matrix
    logc0::T       # the logarithm of normalizing constant in pdf
    rank::R        # rank of a sample
    singular::Bool # singular of nonsingular wishart?
end

#  -----------------------------------------------------------------------------
#  Constructors
#  -----------------------------------------------------------------------------

function Wishart(df::T, S::AbstractPDMat{T}) where T<:Real
    df > 0 || throw(ArgumentError("df must be positive. got $(df)."))
    p = dim(S)
    singular = df <= p - 1
    if singular
        isinteger(df) || throw(
            ArgumentError("df of a singular Wishart distribution must be an integer (got $df)")
        )
    end
    rnk::Integer = ifelse(singular, df, p)
    logc0 = wishart_logc0(df, S, rnk)
    _df, _logc0 = promote(df, logc0)
    Wishart{typeof(_df), typeof(S), typeof(rnk)}(_df, S, _logc0, rnk, singular)
end

function Wishart(df::Real, S::AbstractPDMat)
    T = Base.promote_eltype(df, S)
    Wishart(T(df), convert(AbstractArray{T}, S))
end

Wishart(df::Real, S::Matrix) = Wishart(df, PDMat(S))
Wishart(df::Real, S::Cholesky) = Wishart(df, PDMat(S))

#  -----------------------------------------------------------------------------
#  REPL display
#  -----------------------------------------------------------------------------

show(io::IO, d::Wishart) = show_multline(io, d, [(:df, d.df), (:S, Matrix(d.S))])

#  -----------------------------------------------------------------------------
#  Conversion
#  -----------------------------------------------------------------------------

function convert(::Type{Wishart{T}}, d::Wishart) where T<:Real
    P = convert(AbstractArray{T}, d.S)
    Wishart{T, typeof(P), typeof(d.rank)}(T(d.df), P, T(d.logc0), d.rank, d.singular)
end
function convert(::Type{Wishart{T}}, df, S::AbstractPDMat, logc0, rnk, singular) where T<:Real
    P = convert(AbstractArray{T}, S)
    Wishart{T, typeof(P), typeof(rnk)}(T(df), P, T(logc0), rnk, singular)
end

#  -----------------------------------------------------------------------------
#  Properties
#  -----------------------------------------------------------------------------

insupport(::Type{Wishart}, X::AbstractMatrix) = ispossemdef(X)
function insupport(d::Wishart, X::AbstractMatrix)
    size(X) == size(d) || return false
    if d.singular
        return ispossemdef(X, rank(d))
    else
        return isposdef(X)
    end
end

dim(d::Wishart) = dim(d.S)
size(d::Wishart) = (p = dim(d); (p, p))
rank(d::Wishart) = d.rank
params(d::Wishart) = (d.df, d.S)
@inline partype(d::Wishart{T}) where {T<:Real} = T

mean(d::Wishart) = d.df * Matrix(d.S)

function mode(d::Wishart)
    r = d.df - dim(d) - 1.0
    r > 0.0 || throw(ArgumentError("mode is only defined when df > p + 1"))
    return Matrix(d.S) * r
end

function meanlogdet(d::Wishart)
    d.singular && return -Inf
    p = dim(d)
    df = d.df
    v = logdet(d.S) + p * logtwo
    for i = 1:p
        v += digamma(0.5 * (df - (i - 1)))
    end
    return v
end

function entropy(d::Wishart)
    d.singular && throw(ArgumentError("entropy not defined for singular Wishart."))
    p = dim(d)
    df = d.df
    -d.logc0 - 0.5 * (df - p - 1) * meanlogdet(d) + 0.5 * df * p
end

#  Gupta/Nagar (1999) Theorem 3.3.15.i
function cov(d::Wishart, i::Integer, j::Integer, k::Integer, l::Integer)
    S = Matrix(d.S)
    d.df * (S[i, k] * S[j, l] + S[i, l] * S[j, k])
end

function var(d::Wishart, i::Integer, j::Integer)
    S = Matrix(d.S)
    d.df * (S[i, i] * S[j, j] + S[i, j] ^ 2)
end

#  -----------------------------------------------------------------------------
#  Evaluation
#  -----------------------------------------------------------------------------

function wishart_logc0(df::Real, S::AbstractPDMat, rnk::Integer)
    p = dim(S)
    if df <= p - 1
        return singular_wishart_logc0(p, df, S, rnk)
    else
        return nonsingular_wishart_logc0(p, df, S)
    end
end

function logkernel(d::Wishart, X::AbstractMatrix)
    if d.singular
        return singular_wishart_logkernel(d, X)
    else
        return nonsingular_wishart_logkernel(d, X)
    end
end

#  Singular Wishart pdf: Theorem 6 in Uhlig (1994 AoS)
function singular_wishart_logc0(p::Integer, df::Real, S::AbstractPDMat, rnk::Integer)
    h_df = df / 2
    -h_df * (logdet(S) + p * typeof(df)(logtwo)) - logmvgamma(rnk, h_df) + (rnk*(rnk - p) / 2)*typeof(df)(logπ)
end

function singular_wishart_logkernel(d::Wishart, X::AbstractMatrix)
    df = d.df
    p = dim(d)
    r = rank(d)
    L = eigvals(Hermitian(X), (p - r + 1):p)
    0.5 * ((df - (p + 1)) * sum(log.(L)) - tr(d.S \ X))
end

#  Nonsingular Wishart pdf
function nonsingular_wishart_logc0(p::Integer, df::Real, S::AbstractPDMat)
    h_df = df / 2
    -h_df * (logdet(S) + p * typeof(df)(logtwo)) - logmvgamma(p, h_df)
end

function nonsingular_wishart_logkernel(d::Wishart, X::AbstractMatrix)
    df = d.df
    p = dim(d)
    Xcf = cholesky(X)
    0.5 * ((df - (p + 1)) * logdet(Xcf) - tr(d.S \ X))
end

#  -----------------------------------------------------------------------------
#  Sampling
#  -----------------------------------------------------------------------------

function _rand!(rng::AbstractRNG, d::Wishart, A::AbstractMatrix)
    if d.singular
        A .= zero(eltype(A))
        A[:, 1:rank(d)] = randn(rng, dim(d), rank(d))
    else
        _wishart_genA!(rng, dim(d), d.df, A)
    end
    unwhiten!(d.S, A)
    A .= A * A'
end

function _wishart_genA!(rng::AbstractRNG, p::Int, df::Real, A::AbstractMatrix)
    # Generate the matrix A in the Bartlett decomposition
    #
    #   A is a lower triangular matrix, with
    #
    #       A(i, j) ~ sqrt of Chisq(df - i + 1) when i == j
    #               ~ Normal()                  when i > j
    #
    A .= zero(eltype(A))
    for i = 1:p
        @inbounds A[i,i] = rand(rng, Chi(df - i + 1.0))
    end
    for j in 1:p-1, i in j+1:p
        @inbounds A[i,j] = randn(rng)
    end
end

#  -----------------------------------------------------------------------------
#  Test utils
#  -----------------------------------------------------------------------------

function _univariate(d::Wishart)
    check_univariate(d)
    df, S = params(d)
    α = df / 2
    β = 2Matrix(S)[1]
    return Gamma(α, β)
end

function _rand_params(::Type{Wishart}, elty, n::Int, p::Int)
    n == p || throw(ArgumentError("dims must be equal for Wishart"))
    ν = elty( n - 1 + abs(10randn()) )
    S = (X = 2rand(elty, n, n) .- 1; X * X')
    return ν, S
end


text:
"""
    DiscriminativeLRs(paramgroups, factors)

Use different learning rates based on `paramgroups`. `factors` maps
each group to a factor that the learning rate is multiplied by, so
for a parameter `x` the factor is
`get(factors, getgroup(paramgroups, x), 1)`.

See [`ParamGroups`](#).

## Examples

Combining with regular gradient descent, but only training a part of
the model.

```julia
using Flux.Optimise: Descent, Optimiser

model = Chain(Dense(3, 5), Dense(5, 3))
paramgroups = ParamGroups(IndexGrouper([1, 2]), model)

dlro = DiscriminativeLRs(paramgroups, Dict(1 => 0., 2 => 1.))
o = Optimiser(dlro, Descent(0.1))
```
"""
struct DiscriminativeLRs
    pg::ParamGroups
    factorfn
end

function DiscriminativeLRs(pg::ParamGroups, factors::Dict)
    return DiscriminativeLRs(pg, group -> get(factors, group, 1))
end

function apply!(o::DiscriminativeLRs, x, Δ::AbstractArray{T}) where T
    factor = convert(T, o.factorfn(getgroup(o.pg, x)))

    if factor == one(T)
        return Δ
    else
        @. Δ *= factor
    end
end

function FluxTraining.setlearningrate!(optimizer::Optimiser, value)
    FluxTraining.setlearningrate!(optimizer.os[end], value)
end


text:
#=

  Ski assignment in Julia ConstraintSolver.jl

  From
  Jeffrey Lee Hellrung, Jr.: PIC 60, Fall 2008, Final Review, December 12, 2008
  http://www.math.ucla.edu/~jhellrun/course_files/Fall%25202008/PIC%252060%2520-%2520Data%2520Structures%2520and%2520Algorithms/final_review.pdf
  """
  5. Ski Optimization! Your job at Snapple is pleasant but in the winter you've
  decided to become a ski bum. You've hooked up with the Mount Baldy Ski Resort.
  They'll let you ski all winter for free in exchange for helping their ski rental
  shop with an algorithm to assign skis to skiers. Ideally, each skier should
  obtain a pair of skis whose height matches his or her own height exactly.
  Unfortunately, this is generally not possible. We define the disparity between
  a skier and his or her skis to be the absolute value of the difference between
  the height of the skier and the pair of skis. Our objective is to find an
  assignment of skis to skiers that minimizes the sum of the disparities.
  ...
  Illustrate your algorithm by explicitly filling out the A[i, j] table for the
  following sample data:
    * Ski heights: 1, 2, 5, 7, 13, 21.
    * Skier heights: 3, 4, 7, 11, 18.
  """

  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 ski_assignment(print_solutions=true,all_solutions=false, all_optimal_solutions=false)

    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_optimal_solutions,
                                                            "logging"=>[],

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

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

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

                                                            "time_limit"=>6,

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

    ski_heights   =  [1, 2, 5, 7, 13, 21]
    skier_heights =  [3, 4, 7, 11, 18]

    num_skis = length(ski_heights)
    num_skiers = length(skier_heights)

    @variable(model, 1 <= x[1:num_skiers] <= num_skis, Int)
    @constraint(model, x in CS.AllDifferent())

    # minimize the differences of ski height and skier's height (z)
    t = 100
    @variable(model, 0 <= z <= t, Int)
    @variable(model, sxi[1:num_skiers], CS.Integers(ski_heights))
    @variable(model, 0 <= abs_diff[1:num_skiers] <= t, Int)

    # We must use element here...
    for i in 1:num_skiers
        # find the ski height for skier i
        my_element(model,x[i],ski_heights,sxi[i])
        # the difference beteen ski height and skier's height
        # my_abs(model,sxi[i],skier_heights[i],abs_diff[i])
        my_abs(model,sxi[i]-skier_heights[i],abs_diff[i])
    end
    @constraint(model,z==sum(abs_diff))

    @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")
                xx = convert.(Integer,JuMP.value.(x; result=sol))
                zz = convert.(Integer,JuMP.value.(z; result=sol))
                sxix = convert.(Integer,JuMP.value.(sxi; result=sol))
                abs_diffx = convert.(Integer,JuMP.value.(abs_diff; result=sol))
                println("x:$xx  z:$zz")
                println("sxi:$sxix")
                println("abs_diff:$abs_diffx")
                println()
            end
        end
    else
        println("status:$status")
    end

    return status
end

@time ski_assignment()

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
using Documenter, MakieLayout, Makie, Animations, ColorSchemes, Colors

# don't open windows while generating animations
Makie.AbstractPlotting.inline!(true)

makedocs(
    sitename="MakieLayout.jl",
    pages = [
        "index.md",
        "GridLayout" => "grids.md",
        "LAxis" => "laxis.md",
        "How layouting works" => "layouting.md",
        "Frequently Asked Questions" => "faq.md",
    ],
    format = Documenter.HTML(
            prettyurls = get(ENV, "CI", nothing) == "true"
        )
    )

struct Local <: Documenter.DeployConfig end

function Documenter.deploy_folder(cfg::Local;
        repo, devbranch, push_preview, devurl, kwargs...)

    folder = if ENV["PUSH_LOCAL_BUILD"] == "true"
        @warn("Setting ENV[\"PUSH_LOCAL_BUILD\"] = \"false\", remember to set it to true for the next push.")
        ENV["PUSH_LOCAL_BUILD"] = "false"
        devurl
    else
        @warn("Set ENV[\"PUSH_LOCAL_BUILD\"] = \"true\" if you want your local build to be pushed to Github Pages.")
        nothing
    end
end

Documenter.authentication_method(::Local) = Documenter.SSH

function Documenter.documenter_key(::Local)
    open(readline, expanduser("~/.ssh/documenter_makielayout"))
end

deploydocs(
    repo = "github.com/jkrumbiegel/MakieLayout.jl.git",
    deploy_config = Local(),
)


text:
# Test circle constraint plotting
xc = [1,2,3.]
yc = [-1,-3,-2.]
rc = [0.5, 0.25, 0.75]
con = CircleConstraint(3, xc, yc, rc)
length(con)

con2 = GoalConstraint(SA[1,-2,3])

plot(con, circle_color=:red)
plot!(con2, xind=1, yind=2, color=:green)
plot!(con2, xind=3, yind=2, seriescolor=:yellow, markershape=:diamond, label="start")

prob = DubinsCar()
plot(prob.constraints)
add_constraint!(prob.constraints, con, 1:prob.N-1)
plot(prob.constraints)

# Make sure you can add it an existing plot
t = range(0,3,length=21)
p = plot(t, (t./2).^2)
plot!(prob.constraints)
@test length(p.series_list) == 12

# Dubins Car
prob = DubinsCar()
U0 = [SA[1.0, k/100] for k = 1:prob.N-1]
initial_controls!(prob, U0)
rollout!(prob)
plot(prob.model, prob.Z)
plot!(prob.constraints)
plot(prob)

# Double Integrator
prob = DoubleIntegrator()
initial_controls!(prob, SA[0.1])
rollout!(prob)
plot(prob.model, prob.Z)
plot(prob)


text:
function starting_point!(model::Model, data; μ::Float64 = -5.0, δλ = 10.0)
    randn!(model.x)
    model.x ./ model.L
    model.x[end] = μ
    for t in 1:model.T-1
        model.λ[t] = λmin(data.counts, t) + δλ
    end
end

function set_prior(variance::Float64, model::Model)
    var = fill(variance, fields_length(model.epistasis))
    numJ = (model.A)^2 * model.L*(model.L-1)÷2
    numh = model.A*model.L
    for l in 1:numJ
        var[l] /= numJ
    end
    for l in 1:numh
        var[numJ+l] /= numh
    end

    return GaussianPrior(var)
end

function set_optimizer(algorithm::Symbol, model::Model, data; ftol_rel = 0.0, ftol_abs = 0.0, xtol_rel = 0.0, xtol_abs = 0.0, Δλ = 1e-10)
    opt = Opt(algorithm, Model_length(model.A, model.L, model.T))
    opt.ftol_rel = ftol_rel
    opt.ftol_abs = ftol_abs
    opt.xtol_rel = xtol_rel
    opt.xtol_abs = xtol_abs
    lb = opt.lower_bounds
    lb[end-model.T+2:end] .= [λmin(data.counts, t) + Δλ for t in 1:model.T-1]
    opt.lower_bounds = lb
    return opt
end

export starting_point!, set_prior, set_optimizer


text:
#=
Given an array of integers where every integer occurs three times except for one integer, which only occurs once, find and return the non-duplicated integer.

For example, given [6, 1, 3, 3, 3, 6, 6], return 1. Given [13, 19, 13, 13], return 19.

Do this in O(N) time and O(1) space.=
=#

function nonduplicated(arr::Array{Int, 1})
    searching_arr = copy(arr)
    non_duplicated_int = nothing
    while isnothing(non_duplicated_int)
        current_int = searching_arr[1]
        if sum(searching_arr .== current_int) == 1
            non_duplicated_int = current_int
        else
            deleteat!(searching_arr, findall(isequal(current_int), searching_arr))
        end
    end
    return non_duplicated_int
end

max_stars_repo_name: DominiqueCaron/daily-coding-problem

max_issues_repo_name: DominiqueCaron/daily-coding-problem

max_forks_repo_name: DominiqueCaron/daily-coding-problem


text:
module QuasiDefinite

# package code goes here

# imports
import Base.blasfunc
import Base.LinAlg: BlasReal, BlasComplex, BlasFloat, BlasChar, BlasInt,
    blas_int, DimensionMismatch, chkstride1, chksquare

# load the libQuasiDefinite
include("../deps/deps.jl")

#Generic LAPACK error handlers, copied from https://github.com/JuliaLang/julia/blob/master/base/linalg/lapack.jl
macro assertargsok() #Handle only negative info codes - use only if positive info code is useful!
    :(info[1]<0 && throw(ArgumentError("invalid argument #$(-info[1]) to LAPACK call")))
end

#Check that upper/lower (for special matrices) is correctly specified
macro chkuplo()
    :((uplo=='U' || uplo=='L') ||
      throw(ArgumentError(string("uplo argument must be 'U' (upper) or 'L' (lower), got ", repr(uplo)))))
end

# from SCS.jl
function __init__()
    # default binaries need access to Julia's lapack symbols
    @unix_only dlopen(Base.liblapack_name, RTLD_LAZY|RTLD_DEEPBIND|RTLD_GLOBAL)
end

# see: https://github.com/JuliaLang/julia/blob/master/base/linalg/lapack.jl#L1735
@eval begin
    function qdtrf!(uplo::BlasChar, A::StridedMatrix{Float64})
        chkstride1(A)
        chksquare(A)
        @chkuplo
        lda = max(1,stride(A,2))
        lda==0 && return A, 0
        info = Array(BlasInt, 1)
        ccall(($(blasfunc(:dqdtrf_)), libQuasiDefinite), Void,
              (Ptr{BlasChar}, Ptr{BlasInt}, Ptr{Float64}, Ptr{BlasInt}, Ptr{BlasInt}),
              &uplo, &size(A,1), A, &lda, info)
        @assertargsok
        return A, info[1]
    end
end

@eval begin
    function qdtf2!(uplo::BlasChar, A::StridedMatrix{Float64})
        chkstride1(A)
        chksquare(A)
        @chkuplo
        lda = max(1,stride(A,2))
        lda==0 && return A, 0
        info = Array(BlasInt, 1)
        ccall(($(blasfunc(:dqdtf2_)), libQuasiDefinite), Void,
              (Ptr{BlasChar}, Ptr{BlasInt}, Ptr{Float64}, Ptr{BlasInt}, Ptr{BlasInt}),
              &uplo, &size(A,1), A, &lda, info)
        @assertargsok
        return A, info[1]
    end
end

@eval begin
    function nptf2!(A::StridedMatrix{Float64})
        chkstride1(A)
        lda = max(1,stride(A,2))
        lda==0 && return A, 0
        info = Array(BlasInt, 1)
        ccall(($(blasfunc(:dnptf2_)), libQuasiDefinite), Void,
              (Ptr{BlasInt}, Ptr{BlasInt}, Ptr{Float64}, Ptr{BlasInt}, Ptr{BlasInt}),
              &size(A,1), &size(A,2), A, &lda, info)
        @assertargsok
        return A, info[1]
    end
end

function run()
    A = Float64[2 -1; -1 2]
    B, info = qdtrf!('U',A)
    println(B)
    println(info)
end

end # module


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    SPredicatePartition(pred)

A method for partitioning spatial objects with a given spatial
predicate function `pred(x, y)`.
"""
struct SPredicatePartition <: SPredicatePartitionMethod
  pred::Function
end

(p::SPredicatePartition)(x, y) = p.pred(x, y)


text:
using Compose, DataFrames, Gadfly

set_default_plot_size(9inch, 3.3inch)

theme1 = Theme(point_size=3mm)
coord1 = Coord.cartesian(xmin=0.0, xmax=6.0)
D = DataFrame(x=1:5,
              y=[0.768448, 0.940515, 0.673959, 0.395453, 0.313244],
              V1=["A","A","B","B","D"],
              V2 = string.([1,2,2,3,3])  )

pa = plot(x=1:5, y=[0.77, 0.94, 0.67, 0.39, 0.31], shape=["A","A","B","B","D"], theme1, coord1,
    Guide.shapekey(title="Key",labels=["α","β","δ"]),
    Guide.title("Guide.shapekey") )

pb = plot(D, x=:x, y=:y, shape=:V1, color=:V1,
        Scale.shape_discrete(levels=["D","A","B"]),
        theme1, coord1, Guide.title("Shape==Color") )

pc = plot(D, x=:x, y=:y, shape=:V1, color=:V2,  coord1,
        Guide.colorkey(title="Color"),
        Guide.shapekey(title="Shape ", pos=[0.74w,-0.27h]),
        Theme(point_size=3mm, key_swatch_color="slategrey"),
        Guide.title("Shape!=Color") )

hstack(pa,pb,pc)


text: boolean2string(b::Bool)::String = string(b)

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
VERSION >= v"0.4.0-dev+6521" && __precompile__()
module BayesianDataFusion

include("ROC.jl")
include("normal_wishart.jl")
include("RelationData.jl")
include("parallel_matrix.jl")
include("sparsebin_csr.jl")
include("parallel_csr.jl")
include("parallel_cg.jl")
include("sampling.jl")
include("macau.jl")
include("macau_blocked.jl")
include("macau_vb.jl")
include("macau_hmc.jl")
include("macau_sgd.jl")

include("data_reading.jl")

end # module


text:
const stableGIDs = (UInt128, UInt64, UInt32, UInt16, UInt8, Int128, Int64, Int32, Int16, Int8)
const stablePIDs = (UInt128, UInt64, UInt32, UInt16, UInt8, Int128, Int64, Int32, Int16, Int8, Bool)
const stableLIDs = (UInt128, UInt64, UInt32, UInt16, UInt8, Int128, Int64, Int32, Int16, Int8)

const stableReals= (Float64, Float32, Float16, UInt128, UInt64, UInt32, UInt16, UInt8, Int128, Int64, Int32, Int16, Int8, Bool)
const stableDatas= union(stableReals, [Complex{r} for r in stableReals])


text:
@testset "Copeland" begin
    """
    This example is interpreted from the results of

    Arslan, Rahim, and Hüdaverdi Bircan. "Çok Kriterli Karar Verme Teknikleriyle Elde Edilen
    Sonuçların Copeland Yöntemiyle Birleştirilmesi ve Karşılaştırılması." Yönetim ve Ekonomi:
    Celal Bayar Üniversitesi İktisadi ve İdari Bilimler Fakültesi Dergisi 27.1 (2020): 109-127.

    for testing the method.
    """
     topsis_ord =
        [7, 17, 3, 14, 6, 18, 10, 19, 8, 15, 22, 12, 2, 9, 23, 11, 16, 20, 5, 13, 1, 21, 4]
     gia_ord =
        [9, 17, 3, 14, 6, 18, 7, 19, 8, 16, 22, 13, 1, 10, 23, 11, 15, 20, 5, 12, 2, 21, 4]
     vikor_ord =
        [7, 17, 3, 14, 5, 18, 10, 19, 8, 16, 22, 13, 1, 9, 23, 11, 15, 20, 6, 12, 2, 21, 4]
     mooraref_ord =
        [3, 14, 2, 15, 16, 21, 10, 23, 8, 9, 19, 6, 7, 5, 22, 17, 13, 18, 12, 11, 1, 20, 4]
     copras_ord =
        [7, 17, 3, 14, 6, 18, 11, 19, 8, 15, 22, 12, 1, 9, 23, 10, 16, 20, 5, 13, 2, 21, 4]
     moora_ord =
        [7, 17, 3, 14, 5, 18, 11, 19, 8, 15, 22, 12, 1, 9, 23, 10, 16, 20, 6, 13, 2, 21, 4]
     aras_ord =
        [7, 17, 3, 14, 6, 18, 11, 19, 8, 15, 22, 12, 1, 9, 23, 10, 16, 20, 5, 13, 2, 21, 4]

     mat = hcat(topsis_ord, gia_ord, vikor_ord, mooraref_ord)

    c = copeland(mat)

    d = c |> sortperm

    #show(d)
    #@info d

    @test d == [15, 11, 22, 18, 8, 6, 2, 10, 17, 4, 12, 20, 16, 7, 14, 9, 1, 5, 19, 23, 3, 13, 21]

end

max_stars_repo_name: jbytecode/jMCDM

max_issues_repo_name: jbytecode/jMCDM

max_forks_repo_name: jbytecode/jMCDM


text:
struct TestAcvModel <: WhittleLikelihoodInference.TimeSeriesModel{2,Float64} end
struct TestAcvModelUni <: WhittleLikelihoodInference.TimeSeriesModel{1,Float64} end
struct TestUnknownAcv <: WhittleLikelihoodInference.UnknownAcvTimeSeriesModel{2,Float64} end
struct TestUnknownAcvUni <: WhittleLikelihoodInference.UnknownAcvTimeSeriesModel{1,Float64} end
WhittleLikelihoodInference.npars(::Type{TestAcvModel}) = 2
WhittleLikelihoodInference.npars(::Type{TestAcvModelUni}) = 2
WhittleLikelihoodInference.npars(::Type{TestUnknownAcv}) = 2
WhittleLikelihoodInference.npars(::Type{TestUnknownAcvUni}) = 2
WhittleLikelihoodInference.add_sdf!(out,::TestUnknownAcv,ω) = (out .+= 0.2; nothing)
WhittleLikelihoodInference.grad_add_sdf!(out,::TestUnknownAcv,ω) = (out .+= 0.2; nothing)
WhittleLikelihoodInference.hess_add_sdf!(out,::TestUnknownAcv,ω) = (out .+= 0.2; nothing)
WhittleLikelihoodInference.sdf(::TestUnknownAcvUni,ω) = 0.2
WhittleLikelihoodInference.grad_add_sdf!(out,::TestUnknownAcvUni,ω) = (out .+= 0.2; nothing)
WhittleLikelihoodInference.hess_add_sdf!(out,::TestUnknownAcvUni,ω) = (out .+= 0.2; nothing)
struct TestAcv2 <: WhittleLikelihoodInference.TimeSeriesModel{2,Float64} end
struct TestAcvUni2 <: WhittleLikelihoodInference.TimeSeriesModel{1,Float64} end
struct TestAcv3 <: WhittleLikelihoodInference.TimeSeriesModel{2,Float64} end
struct TestAcvUni3 <: WhittleLikelihoodInference.TimeSeriesModel{1,Float64} end
TestAcv2(x) = TestAcv2() # just to allow additive constructor to work
TestAcvUni2(x) = TestAcvUni2()
TestAcv3(x) = TestAcv3() # just to allow additive constructor to work
TestAcvUni3(x) = TestAcvUni3()
WhittleLikelihoodInference.npars(::Type{TestAcv2}) = 4
WhittleLikelihoodInference.npars(::Type{TestAcvUni2}) = 4
WhittleLikelihoodInference.npars(::Type{TestAcv3}) = 4
WhittleLikelihoodInference.npars(::Type{TestAcvUni3}) = 4
WhittleLikelihoodInference.acv!(out,::TestAcv2,τ::Number) = (out.=1.0;nothing)
WhittleLikelihoodInference.grad_acv!(out,::TestAcv2,τ::Number) = (out.=1.0;nothing)
WhittleLikelihoodInference.hess_acv!(out,::TestAcv2,τ::Number) = (out.=1.0;nothing)
WhittleLikelihoodInference.acv(::TestAcvUni2,τ::Number) = 1.0
WhittleLikelihoodInference.grad_acv!(out,::TestAcvUni2,τ::Number) = (out.=1.0;nothing)
WhittleLikelihoodInference.hess_acv!(out,::TestAcvUni2,τ::Number) = (out.=1.0;nothing)
WhittleLikelihoodInference.acv!(out,::TestAcv3,τ::Number) = (out.=2.0;nothing)
WhittleLikelihoodInference.grad_acv!(out,::TestAcv3,τ::Number) = (out.=2.0;nothing)
WhittleLikelihoodInference.hess_acv!(out,::TestAcv3,τ::Number) = (out.=2.0;nothing)
WhittleLikelihoodInference.acv(::TestAcvUni3,τ::Number) = 2.0
WhittleLikelihoodInference.grad_acv!(out,::TestAcvUni3,τ::Number) = (out.=2.0;nothing)
WhittleLikelihoodInference.hess_acv!(out,::TestAcvUni3,τ::Number) = (out.=2.0;nothing)
const AcvDouble = TestAcv2 + TestAcv3
const AcvDoubleUni = TestAcvUni2 + TestAcvUni3
@testset "autocovariance" begin
    @testset "Function" begin
        @testset "Error handling" begin
            @testset "Acv unspecified error" begin
                @test_throws ErrorException WhittleLikelihoodInference.acv!(ones(ComplexF64,3), TestAcvModel(), 1.0)
                @test_throws ErrorException acv(TestAcvModelUni(), 1.0)
            end
            @testset "Lags different size to preallocated memory error" begin
                store1 = allocate_memory_EI_F(TestAcvModel, 1000, 1)
                store2 = allocate_memory_EI_F(TestAcvModel, 10, 1)
                store1uni = allocate_memory_EI_F(TestAcvModelUni, 1000, 1)
                store2uni = allocate_memory_EI_F(TestAcvModelUni, 10, 1)
                @test_throws ArgumentError WhittleLikelihoodInference.acv!(store1.funcmemory, TestAcvModel(), store2.encodedtime)
                @test_throws ArgumentError WhittleLikelihoodInference.acv!(store1uni.funcmemory, TestAcvModelUni(), store2uni.encodedtime)
            end
            @testset "Custom lag when acv unknown error" begin
                @test_throws ArgumentError acv(TestUnknownAcv(), 1.0)
                @test_throws ArgumentError acv(TestUnknownAcvUni(), 1.0)
            end
            @testset "UnknownAcv" begin
                store = allocate_memory_EI_FGH(TestUnknownAcv, 100, 1)
                @test WhittleLikelihoodInference.acv!(store, TestUnknownAcv()) == nothing
                @test WhittleLikelihoodInference.grad_acv!(store, TestUnknownAcv()) == nothing
                @test WhittleLikelihoodInference.hess_acv!(store, TestUnknownAcv()) == nothing
                storeuni = allocate_memory_EI_FGH(TestUnknownAcvUni, 100, 1)
                @test WhittleLikelihoodInference.acv!(storeuni, TestUnknownAcvUni()) == nothing
                @test WhittleLikelihoodInference.grad_acv!(storeuni, TestUnknownAcvUni()) == nothing
                @test WhittleLikelihoodInference.hess_acv!(storeuni, TestUnknownAcvUni()) == nothing
            end
        end
        @testset "Value" begin
            @testset "Simple model" begin
                store1 = allocate_memory_EI_F(TestAcv2, 10, 1)
                store1uni = allocate_memory_EI_F(TestAcvUni2, 10, 1)
                WhittleLikelihoodInference.acv!(store1,TestAcv2())
                WhittleLikelihoodInference.acv!(store1uni,TestAcvUni2())
                @test all(store1.funcmemory.allocatedarray.==1.0)
                @test all(store1uni.funcmemory.allocatedarray.==1.0)
            end
            @testset "Additive model" begin
                store2 = allocate_memory_EI_F(AcvDouble, 10, 1)
                store2uni = allocate_memory_EI_F(AcvDoubleUni, 10, 1)
                WhittleLikelihoodInference.acv!(store2,AcvDouble(ones(8)))
                WhittleLikelihoodInference.acv!(store2uni,AcvDoubleUni(ones(8)))
                @test all(store2.store1.funcmemory.allocatedarray.==3.0)
                @test all(store2uni.store1.funcmemory.allocatedarray.==3.0)
            end
        end
    end
    @testset "Gradient" begin
        @testset "Error handling" begin
            @testset "Acv unspecified error" begin
                @test_throws ErrorException WhittleLikelihoodInference.grad_acv!(ones(ComplexF64,3,4), TestAcvModel(), 1.0)
                @test_throws ErrorException WhittleLikelihoodInference.grad_acv!(ones(ComplexF64,4),TestAcvModelUni(), 1.0)
            end
            @testset "Lags different size to preallocated memory error" begin
                store1 = allocate_memory_EI_FG(TestAcvModel, 1000, 1)
                store2 = allocate_memory_EI_FG(TestAcvModel, 10, 1)
                store1uni = allocate_memory_EI_FG(TestAcvModelUni, 1000, 1)
                store2uni = allocate_memory_EI_FG(TestAcvModelUni, 10, 1)
                @test_throws ArgumentError WhittleLikelihoodInference.grad_acv!(store1.gradmemory, TestAcvModel(), store2.encodedtime)
                @test_throws ArgumentError WhittleLikelihoodInference.grad_acv!(store1uni.gradmemory, TestAcvModelUni(), store2uni.encodedtime)
            end
            @testset "Custom lag when acv unknown error" begin
                @test_throws ArgumentError grad_acv(TestUnknownAcv(), 1.0)
                @test_throws ArgumentError grad_acv(TestUnknownAcvUni(), 1.0)
            end
        end
        @testset "Value" begin
            @testset "Simple model" begin
                store1 = allocate_memory_EI_FG(TestAcv2, 10, 1)
                store1uni = allocate_memory_EI_FG(TestAcvUni2, 10, 1)
                WhittleLikelihoodInference.grad_acv!(store1,TestAcv2())
                WhittleLikelihoodInference.grad_acv!(store1uni,TestAcvUni2())
                @test all(store1.gradmemory.allocatedarray.==1.0)
                @test all(store1uni.gradmemory.allocatedarray.==1.0)
            end
            @testset "Additive model" begin
                store2 = allocate_memory_EI_FG(AcvDouble, 10, 1)
                store2uni = allocate_memory_EI_FG(AcvDoubleUni, 10, 1)
                WhittleLikelihoodInference.grad_acv!(store2,AcvDouble(ones(8)))
                WhittleLikelihoodInference.grad_acv!(store2uni,AcvDoubleUni(ones(8)))
                @test all(store2.store1.gradmemory.allocatedarray.==1.0)
                @test all(store2.store2.gradmemory.allocatedarray.==2.0)
                @test all(store2uni.store1.gradmemory.allocatedarray.==1.0)
                @test all(store2uni.store2.gradmemory.allocatedarray.==2.0)
            end
        end
    end
    @testset "Hessian" begin
        @testset "Error handling" begin
            @testset "Acv unspecified error" begin
                @test_throws ErrorException WhittleLikelihoodInference.hess_acv!(ones(ComplexF64,3,4), TestAcvModel(), 1.0)
                @test_throws ErrorException WhittleLikelihoodInference.hess_acv!(ones(ComplexF64,4),TestAcvModelUni(), 1.0)
            end
            @testset "Lags different size to preallocated memory error" begin
                store1 = allocate_memory_EI_FGH(TestAcvModel, 1000, 1)
                store2 = allocate_memory_EI_FGH(TestAcvModel, 10, 1)
                store1uni = allocate_memory_EI_FGH(TestAcvModelUni, 1000, 1)
                store2uni = allocate_memory_EI_FGH(TestAcvModelUni, 10, 1)
                @test_throws ArgumentError WhittleLikelihoodInference.hess_acv!(store1.hessmemory, TestAcvModel(), store2.encodedtime)
                @test_throws ArgumentError WhittleLikelihoodInference.hess_acv!(store1uni.hessmemory, TestAcvModelUni(), store2uni.encodedtime)
            end
            @testset "Custom lag when acv unknown error" begin
                @test_throws ArgumentError grad_acv(TestUnknownAcv(), 1.0)
                @test_throws ArgumentError grad_acv(TestUnknownAcvUni(), 1.0)
            end
        end
        @testset "Value" begin
            @testset "Simple model" begin
                store1 = allocate_memory_EI_FGH(TestAcv2, 10, 1)
                store1uni = allocate_memory_EI_FGH(TestAcvUni2, 10, 1)
                WhittleLikelihoodInference.hess_acv!(store1,TestAcv2())
                WhittleLikelihoodInference.hess_acv!(store1uni,TestAcvUni2())
                @test all(store1.hessmemory.allocatedarray.==1.0)
                @test all(store1uni.hessmemory.allocatedarray.==1.0)
            end
            @testset "Additive model" begin
                store2 = allocate_memory_EI_FGH(AcvDouble, 10, 1)
                store2uni = allocate_memory_EI_FGH(AcvDoubleUni, 10, 1)
                WhittleLikelihoodInference.hess_acv!(store2,AcvDouble(ones(8)))
                WhittleLikelihoodInference.hess_acv!(store2uni,AcvDoubleUni(ones(8)))
                @test all(store2.store1.hessmemory.allocatedarray.==1.0)
                @test all(store2.store2.hessmemory.allocatedarray.==2.0)
                @test all(store2uni.store1.hessmemory.allocatedarray.==1.0)
                @test all(store2uni.store2.hessmemory.allocatedarray.==2.0)
            end
        end
    end
end


text:
function Message(conn::Connection)
    msg_props_handle_ref = Ref{Ptr{Cvoid}}()
    result = dpiConn_newMsgProps(conn.handle, msg_props_handle_ref)
    error_check(context(conn), result)
    return Message(conn, msg_props_handle_ref[])
end

function set_payload_bytes!(msg::Message, bytes::Vector{UInt8})
    result = dpiMsgProps_setPayloadBytes(msg.handle, pointer(bytes), length(bytes))
    error_check(context(msg), result)
    nothing
end

function get_payload_bytes(msg::Message) :: Vector{UInt8}
    dpi_object_handle_ref = Ref{Ptr{Cvoid}}(C_NULL)
    bytes_array_ptr_ref = Ref{Ptr{UInt8}}(C_NULL)
    bytes_array_len_ref = Ref{UInt32}(0)
    result = dpiMsgProps_getPayload(msg.handle, dpi_object_handle_ref, bytes_array_ptr_ref, bytes_array_len_ref)
    error_check(context(msg), result)

    @assert dpi_object_handle_ref[] == C_NULL "Message payload is object"
    @assert bytes_array_ptr_ref[] != C_NULL

    return deepcopy(unsafe_wrap(Vector{UInt8}, bytes_array_ptr_ref[], bytes_array_len_ref[]))
end

function set_correlation!(msg::Message, correlation::AbstractString)
    err = dpiMsgProps_setCorrelation(msg.handle, pointer(correlation), sizeof(correlation))
    error_check(context(msg), err)
    nothing
end

#function clear_correlation!(msg::Message)
#    err = dpiMsgProps_setCorrelation(msg.handle, Ptr{UInt8}(C_NULL), 0)
#    error_check(context(msg), err)
#    nothing
#end

function get_correlation(msg::Message) :: Union{Nothing, String}
    val_ref = Ref{Ptr{UInt8}}()
    val_length_ref = Ref{UInt32}()
    err = dpiMsgProps_getCorrelation(msg.handle, val_ref, val_length_ref)
    error_check(context(msg), err)
    if val_ref[] == C_NULL
        return nothing
    else
        return unsafe_string(val_ref[], val_length_ref[])
    end
end


text:
pcount = 4
psize = 2

srand(1)

# TODO: Links gets set up properly on instantiation.

context("metapopulation.jl") do
  landscapes = [
    NK.NKLandscape(4, 1),
    NK.NKqLandscape(4, 1, 2),
    NK.NKpLandscape(4, 1, 0.90),
  ]

  for ls = landscapes
    mpr = rand(NK.MetaPopulation, ls, psize, pcount)
    mpz = zeros(NK.MetaPopulation, ls, psize, pcount)

    context("$(ls)") do
      context("popsize(...)") do
        @fact NK.popsize(mpr) --> psize
        @fact NK.popsize(mpz) --> psize
      end

      context("popct(...)") do
        @fact NK.popct(mpr) --> pcount
        @fact NK.popct(mpz) --> pcount
      end

      context("popfits(...)") do
        @fact size(NK.popfits(mpr))[1] --> psize
        @fact size(NK.popfits(mpz))[1] --> psize
        @fact size(NK.popfits(mpr))[2] --> pcount
        @fact size(NK.popfits(mpz))[2] --> pcount
      end
    end
  end
end


text:
using CauseMap

function calcFullseries(vec1, vec2, E, tau_s, tau_p, npred, pred_start, lib_start; nboots=0)

    shadowmat_dict_vec1, distmat_dict_vec1 = precalc_manif_dists(3:3, 1:1, vec1)
    shadowmat_dict_vec2, distmat_dict_vec2 = precalc_manif_dists(3:3, 1:1, vec2)

    librange12, yval_12 = calcCCM(vec1, vec2,
                                                        shadowmat_dict_vec1, distmat_dict_vec1,
                                                        E, tau_s, tau_p; nboots=nboots,
                                                        libsizemin=libsizemin, libsizemax=libsizemax,
                                                        npred=npred, pred_start=pred_start)

    librange21, yval_21 = calcCCM(vec2, vec1,
                                                        shadowmat_dict_vec2, distmat_dict_vec2,
                                                        E, tau_s, tau_p; nboots=nboots,
                                                        libsizemin=libsizemin, libsizemax=libsizemax,
                                                        npred=npred, pred_start=pred_start)

    return librange12, librange21, yval_12, yval_21
end

function calcshortseries(vec1, vec2, E, tau_s, tau_p, npred, pred_start, lib_start, libsizemin, libsizemax, factor::Int64; nboots=0)
    npred = int(npred/factor)
    libsizemax = int(libsizemax/factor)
    xval = libsizemin:libsizemax
    vec1 = vec1[1:factor:end]
    vec2 = vec2[1:factor:end]

    shadowmat_dict_vec1, distmat_dict_vec1 = precalc_manif_dists(E_vals, tau_s_vals, vec1)
    shadowmat_dict_vec2, distmat_dict_vec2 = precalc_manif_dists(E_vals, tau_s_vals, vec2)

    lib12, res12 = calcCCM(vec1, vec2,
                                            shadowmat_dict_vec1, distmat_dict_vec1,
                                            E, 1, 0;
                                            nboots=nboots, lib_start=lib_start,
                                            libsizemin=libsizemin, libsizemax=libsizemax,
                                            npred=npred, pred_start=pred_start)

    lib21, res21 = calcCCM(vec2, vec1,
                                            shadowmat_dict_vec2, distmat_dict_vec2,
                                            E, 1, 0;
                                            nboots=nboots, lib_start=lib_start,
                                            libsizemin=libsizemin, libsizemax=libsizemax,
                                            npred=npred, pred_start=pred_start)
    return xval, res12, res21

end

E = 3
tau_s = 1
tau_p = 0
lib_start = 1
libsizemin = 5
libsizemax = 58
pred_start = 13
npred = 60
E_vals = E:E
tau_s_vals = tau_s:tau_s
nboots = 0
nboots2 = 50

ds   = readdlm("vr.raw_fixed.txt", '\t', Float64, header=false)  # read in data
para = ds[:, 2]
didi = ds[:, 3]

librange12, librange21, yval_12, yval_21 = calcFullseries(para, didi, E, tau_s, tau_p, npred, pred_start, lib_start; nboots=nboots)

xval_f2, yval_12_f2, yval_21_f2 = calcshortseries(para, didi, E, tau_s,
                                                                            tau_p, npred, pred_start,
                                                                            lib_start, libsizemin, libsizemax, 2; nboots=nboots)

xval_f3, yval_12_f3, yval_21_f3 = calcshortseries(para, didi, E, tau_s,
                                                                            tau_p, npred, pred_start,
                                                                            lib_start, libsizemin, libsizemax, 3; nboots=nboots)
fac = 3
left_E = 2
right_E = 6
right_tau_p = int(12/fac)
E_vals = left_E:right_E
tau_p_vals = 0:right_tau_p

# makeoptimizationplots(para[1:fac:end], didi[1:fac:end], libsizemin, int(libsizemax/fac), E_vals,
#             tau_s_vals, tau_p_vals, int(npred/fac), pred_start,
#             "Para.", "Didi."; nreps=1, left_E=3, left_tau_p=0,
#                                     right_E=right_E, right_tau_p=right_tau_p, lagunit=.5 * fac, unit="days", nboots=nboots, legend=true)

using PyCall
@pyimport matplotlib.pyplot as plt
@pyimport seaborn as sns
sns.set_style("white")
sns.set_context("poster")

plt.plot(librange12, yval_12, color="k", alpha=.7, label="Full (60 points)")
plt.plot(librange21, yval_21, color="k", alpha=.5)
plt.plot(xval_f2, yval_12_f2, color="b", label="1/2 thinned (30 points)", alpha=.9)
plt.plot(xval_f2, yval_21_f2, color="b",  alpha=.5)
plt.plot(xval_f3, yval_12_f3, color="r", label="1/3 thinned (20 points)",  alpha=.9)
plt.plot(xval_f3, yval_21_f3, color="r",  alpha=.5)
plt.xlabel("Library size")
plt.ylabel("\$\\rho_{ccm}\$", fontsize=40)
plt.locator_params(axis="y", nbins=4)
plt.legend(loc=0)
plt.show()


text:
"""
Main code for running the branching process. This invokes other Julia scripts in order
to run:
    - read_in_TP.jl:
        Code for reading the TP csv files saved from the fitting and
        forecasting parts of the model.
    - helper_functions.jl:
        Code which is commonly reused in the model. This involves non-specialised functions
        for sampling onset and infection times, and for sampling from a NegativeBinomial
        and Binomial distribution using asymptotic limiting properties.
    - forecast_types.jl
        Code for some user defined types. This lets us group variables of importance. There
        is a struct for the dynamical constants, properties of the data and a
        simple mapping for moving between types and indices.
    - simulation_consistency.jl
        Code for checking consistency of simulations against the data. The code in this file
        is also responsible for injecting cases into the simulation if we are below some
        threshold.
    - assumptions.jl
        One of the most important files. Keeps track of all assumptions in a function that
        is used to initialise the system constants.

The key logic of this simulation is featured in the function simulate_branching_process().
This is written in such a way that it is relatively straightforward to follow the logic of
the processes and functions are called in a systematic way. simulate_branching_process()
will return reshaped arrays that are the good simulated realisations based off
the consistency checks.
"""

using Distributed
using Distributions
using Random
using Statistics
using LinearAlgebra
using ProgressBars
using StaticArrays
using SharedArrays
using ProgressMeter

include("helper_functions.jl")
include("forecast_types.jl")
include("simulation_consistency.jl")

##### Methods
function initialise_population!(
    forecast::Forecast,
    D0,
)
    """
    Fill the infection array for each simulation, inititialising the states of
    all the simulations.
    """
    Z = forecast.sim_realisation.Z
    Z_historical = forecast.sim_realisation.Z_historical
    D = forecast.sim_realisation.D
    U = forecast.sim_realisation.U

    # extract prob of being symptomatic and prob of symptoms given detection
    # NOTE: this currently assumes we are starting in the Delta wave (which is true as
    # of 31/1/2022) but we will need to update this if we shift the starting date into
    # an only Omicron period.
    T_end = forecast.sim_features.T_end
    individual_type_map = forecast.individual_type_map

    # assume detection probabilities for initialisation is whatever it was at day 0
    p_detect_given_symp = forecast.sim_constants.p_detect_given_symp[map_day_to_index_p(0)]
    p_symp = forecast.sim_constants.p_symp[map_day_to_index_p(0)]

    # infer some initial undetected symptomatic
    num_S_undetected = 0

    if D0.S == 0
        num_S_undetected = rand(NegativeBinomial(1, p_detect_given_symp))
    else
        num_S_undetected = rand(NegativeBinomial(D0.S, p_detect_given_symp))
    end

    # infer some initial undetected asymptomatic
    total_S = D0.S + num_S_undetected
    num_A_undetected = 0

    if total_S == 0
        num_A_undetected = rand(NegativeBinomial(1, p_symp))
    else
        num_A_undetected = rand(NegativeBinomial(total_S, p_symp))
    end

    # total initial cases include the undetected
    total_initial_cases = sum(D0)
    total_initial_cases += num_S_undetected + num_A_undetected

    start_day = 0

    # reset arrays
    D .= 0
    U .= 0
    Z .= 0
    Z_historical .= 0

    if total_initial_cases > 0
        inf_times = MVector{total_initial_cases}(
            zeros(Int, total_initial_cases)
        )

        # sample negative infection times as these individuals are assumed to occur
        # before the simulation
        for i in 1:total_initial_cases
            inf_times[i] = -sample_onset_time(false)
        end

        # the starting day is the earliest infection time deemed to contribute to the sim.
        start_day = minimum(inf_times)
        idx = 1

        for _ in 1:D0.S
            Z[map_day_to_index_Z(inf_times[idx]), individual_type_map.S] += 1
            idx += 1
            D[1, individual_type_map.S] += 1
        end

        for _ in 1:D0.A
            Z[map_day_to_index_Z(inf_times[idx]), individual_type_map.A] += 1
            idx += 1
            D[1, individual_type_map.A] += 1
        end

        for _ in 1:num_S_undetected
            Z[map_day_to_index_Z(inf_times[idx]), individual_type_map.S] += 1
            idx += 1
            U[1, individual_type_map.S] += 1
        end

        for _ in 1:num_A_undetected
            Z[map_day_to_index_Z(inf_times[idx]), individual_type_map.A] += 1
            idx += 1
            U[1, individual_type_map.A] += 1
        end
    end

    return start_day

end

function import_cases_model!(
	forecast::Forecast,
    import_cases,
	forecast_start_date,
)
	"""
	A model to handle international imports in the forecasting of cases. This is run
    prior to the standard forecasting model and will pre-fill D, U and Z for each
    simulation.
	We assume that cases arise from
	D[t] ~ NegBin(a[t], 1/(b+1))
	where
	a[t] = alpha + f(I[t])
	b = beta + 1
	are the parameters. The function f() is an exponential weighted moving average such
    that
	f(I[t]) = phi I[t] + (1-phi) f(I[t-1])
	and alpha, beta are hyperparameters fixed at 0.5 and 0.2, respectively. We
    take b = β + 1 as we assume a fixed period of 1 day for estimating the posterior.
	"""

    Z = forecast.sim_realisation.Z
    D = forecast.sim_realisation.D
    U = forecast.sim_realisation.U

    T_observed = forecast.sim_features.T_observed
    T_end = forecast.sim_features.T_end

	# number of days after t = 0 (2020/03/01)
	days_into_forecast = Day(Date(forecast_start_date) - Date("2020-03-01")).value

    # grab the correct parameters for the particular dominant strain
    p_detect_import = 0.0

	# the number of days the simulation is run for
	forecast_days_plus_padding = -10:T_end
	# posterior a, b for sampling number of imports on day t
	a_post = MVector{length(forecast_days_plus_padding)}(
        zeros(Float64, length(forecast_days_plus_padding))
    )

	# current ema is taken 11 days before the forecast start date
	current_ema = zero(Float64)

    # imports show up in the data before
    # import_infections = import_cases[5:end]
    inc_period_shift = 5
    current_ema += import_cases[inc_period_shift + days_into_forecast - 11]

	for (i, t) in enumerate(forecast_days_plus_padding)
		# if within observation period, update the posterior. Otherwise use the last
		# available measurement
		if t <= T_observed - inc_period_shift
			count_on_day = import_cases[inc_period_shift + days_into_forecast + t]
			current_ema = forecast.sim_constants.ϕ * count_on_day +
                (1 - forecast.sim_constants.ϕ) * current_ema
			a_post[i] = forecast.sim_constants.prior_alpha + current_ema
		else
			a_post[i] = forecast.sim_constants.prior_alpha + current_ema
		end
	end

    # assuming a fixed period of 1 day (so this is the same for the whole period)
    b_post = forecast.sim_constants.prior_beta + 1

    for (i, t) in enumerate(forecast_days_plus_padding)
        # sample detected imports
        D_I = rand(NegativeBinomial(a_post[i], b_post / (1 + b_post)))
        # see whether they're observed
        unobserved_I = D_I == 0 ? 1 : D_I
        # based on the switch to omicron (in terms of infections so -5 days)
        if map_day_to_index_p(t) <= length(forecast.sim_constants.p_detect_import)
            p_detect_import = forecast.sim_constants.p_detect_import[map_day_to_index_p(t)]
        else
            p_detect_import = forecast.sim_constants.p_detect_import[end]
        end
        U_I = rand(NegativeBinomial(unobserved_I, p_detect_import))

        # day of infection is time t which is why the infection time in
        if D_I + U_I > 0
            assign_to_arrays_and_times!(
                forecast,
                t,
                is_omicron = false,
                num_I_detected = D_I,
                num_I_undetected = U_I,
            )
        end
    end

	return nothing

end

# function calculate_sus_dep_factor(
#     susceptible_depletion_sim,
#     Z_historical,
#     day,
#     N,
# )

#     ψ = 0.0

#     for (i, τ) in enumerate((30, 60, 90, 120))
#         # guard clauses for when we do not have enough observations to consider the longer term
#         # effects (in other words, all the infections are relatively recent)
#         map_day_to_index_Z(day) < 30 && continue

#         # Get starting and ending index. Start must be at least the beginning of the array and end
#         # must be at least the forecast length.
#         τ_start = max(map_day_to_index_Z(day) - (τ - 1), 1)
#         τ_end = min(map_day_to_index_Z(day) - (τ - 30), size(Z_historical, 1))

#         # sum infections over the time period of interest
#         infections = sum(@view Z_historical[τ_start:τ_end])

#         # we index from 4 - 1 as the largest reduction occurs as a result of recent infections
#         prop_inf = min(1.0, infections / N)
#         ψ += susceptible_depletion_sim[5 - i] * prop_inf
#     end

#     scale_factor = max(0.0, 1.0 - ψ)

#     return scale_factor

# end

function calculate_sus_dep_factor(
    susceptible_depletion_sim,
    Z_historical,
    day,
    N,
)

    ψ = 0.0

    τ_horizons_start = (0, 14, 28, 42, 72, 102)
    τ_horizons_end = (0, 13, 27, 41, 71, 101, 131)

    for i in 1:length(τ_horizons_start)
        # guard clause to ensure we have a reasonable number of observations to calculate the
        # contribution to ψ from period i
        map_day_to_index_Z(day) <= τ_horizons_start[i] && continue

        # Get starting and ending index. Start must be at least the beginning of the array and end
        # must be at least the forecast length.
        τ_start = max(map_day_to_index_Z(day) - τ_horizons_end[i + 1], 1)
        τ_end = min(map_day_to_index_Z(day) - τ_horizons_start[i], size(Z_historical, 1))

        # sum infections over the time period of interest
        infections = sum(@view Z_historical[τ_start:τ_end])

        # we index from 4 - 1 as the largest reduction occurs as a result of recent infections
        prop_inf = min(1.0, infections / N)
        ψ += susceptible_depletion_sim[7 - i] * prop_inf
    end

    scale_factor = max(0.0, 1.0 - ψ)

    return scale_factor

end

# function calculate_sus_dep_factor(
#     susceptible_depletion_sim,
#     Z_historical,
#     day,
#     N,
# )

#     ψ = 0.0

#     τ_end = map_day_to_index_Z(day)

#     infections = sum(@view Z_historical[begin:τ_end])
#     ψ = susceptible_depletion_sim * infections / N

#     scale_factor = max(0.0, 1.0 - ψ)

#     return scale_factor

# end

function sigmoid(t, prop_pars)
    (m₀, m₁, r, τ) = prop_pars
    res = m₀ + (m₁ - m₀) / (1 + exp(-r * (t - τ)))

    return res

end

function sample_offspring!(
    forecast::Forecast,
    day,
    TP,
    prop_pars,
    susceptible_depletion_sim,
    local_cases;
    adjust_TP = false,
)
    """
    Sample offspring for all adults at time t. This uses the fact that the sum of
    N NB(s, p) is NB(N*s, p) and hence we can sample ALL offspring for a particular type
    of adults.
    """

    Z = forecast.sim_realisation.Z
    Z_historical = forecast.sim_realisation.Z_historical
    D = forecast.sim_realisation.D
    U = forecast.sim_realisation.U
    omicron_start_day = forecast.sim_features.omicron_start_day
    omicron_only_day = forecast.sim_features.omicron_only_day
    T_observed = forecast.sim_features.T_observed
    state = forecast.sim_features.state

    # extract the CAR related quantities
    p_detect = forecast.sim_constants.p_detect
    p_symp = forecast.sim_constants.p_symp[map_day_to_index_Z(day)]
    p_detect_given_symp = forecast.sim_constants.p_detect_given_symp[map_day_to_index_Z(day)]
    p_detect_given_asymp = forecast.sim_constants.p_detect_given_asymp[map_day_to_index_Z(day)]

    # extract fixed simulation values
    N = forecast.sim_features.N

    # pointer to the current day's infections
    Z_tmp = @view Z[map_day_to_index_Z(day), :]

    total_parents = sum(Z_tmp)

    if total_parents > 0
        (S_parents, A_parents, I_parents) = Z_tmp

        # sum up the number local infections on the current day
        Z_historical[map_day_to_index_Z(day)] += S_parents + A_parents
        # Z_tmp gets emptied at the end of sampling so we won't be double counting
        # zero out infections on day t so we can restart sim and keep the progress
        Z_tmp .= 0

        S_parents_omicron = 0
        A_parents_omicron = 0
        I_parents_omicron = 0

        if day >= omicron_start_day && day < omicron_only_day
            # prior to 2022-02-01, there is some probability of non-Omicron (Delta) cases
            # number of days into the omicron wave
            t_omicron = day - omicron_start_day
            prop_omicron = sigmoid(t_omicron, prop_pars)
            # sample a random number of Omicron parents
            S_parents_omicron = rand(Binomial(S_parents, prop_omicron))
            A_parents_omicron = rand(Binomial(A_parents, prop_omicron))
            I_parents_omicron = rand(Binomial(I_parents, prop_omicron))
        elseif day >= omicron_only_day
            # after 2022-02-01 everything is just Omicron
            S_parents_omicron = S_parents
            A_parents_omicron = A_parents
            I_parents_omicron = I_parents
        end

        # put the number of parents in a tuple to iterate over
        num_parents = (
            S_parents - S_parents_omicron,
            A_parents - A_parents_omicron,
            I_parents - I_parents_omicron,
            S_parents_omicron,
            A_parents_omicron,
            I_parents_omicron,
        )

        # adjust the TP/Reff by susceptible depletion based off the current simulation state
        scale_factor = calculate_sus_dep_factor(
            susceptible_depletion_sim,
            Z_historical,
            day,
            N,
        )

        for (i, n_parents) in enumerate(num_parents)
            # a guard clause for skipping the current iteration
            n_parents == 0 && continue

            # reset counters
            TP_parent = 0.0
            k = 0.0
            # relative infectiousness of asymptomatic
            γ = 0.5

            # essentially a (gross) switch statement
            if i == 1
                α_s = 1 / (p_symp + γ * (1 - p_symp))
                TP_parent = TP.TP_local_delta_parent * α_s
                k = 0.15
            elseif i == 2
                α_s = 1 / (p_symp + γ * (1 - p_symp))
                α_a = γ * α_s
                TP_parent = TP.TP_local_delta_parent * α_a
                k = 0.15
            elseif i == 3
                TP_parent = TP.TP_import_delta_parent
                k = 0.15
            elseif i == 4
                α_s = 1 / (p_symp + γ * (1 - p_symp))
                α_a = γ * α_s
                TP_parent = TP.TP_local_omicron_parent * α_s
                k = 0.5
            elseif i == 5
                α_s = 1 / (p_symp + γ * (1 - p_symp))
                α_a = γ * α_s
                TP_parent = TP.TP_local_omicron_parent * α_a
                k = 0.6
            elseif i == 6
                TP_parent = TP.TP_import_omicron_parent
                k = 0.6
            end

            # scale the TP by depletion of susceptibles
            # println("Day: ", day)
            # println("Scale factor: ", scale_factor)
            TP_parent = max(0, TP_parent * scale_factor)

            TP_parent == 0 && continue

            num_offspring = sample_negative_binomial_limit(TP_parent * n_parents, k * n_parents)

            if num_offspring > 0
                is_omicron = false

                if i in 1:3
                    is_omicron = false
                elseif i in 4:6
                    is_omicron = true
                end

                num_S_offspring = sample_binomial_limit(num_offspring, p_symp)
                num_A_offspring = num_offspring - num_S_offspring

                num_S_offspring_detected = sample_binomial_limit(
                    num_S_offspring, p_detect_given_symp
                )
                num_A_offspring_detected = sample_binomial_limit(
                    num_A_offspring, p_detect_given_asymp
                )

                num_S_offspring_undetected = num_S_offspring - num_S_offspring_detected
                num_A_offspring_undetected = num_A_offspring - num_A_offspring_detected

                assign_to_arrays_and_times!(
                    forecast,
                    day,
                    is_omicron = is_omicron,
                    num_S_detected = num_S_offspring_detected,
                    num_S_undetected = num_S_offspring_undetected,
                    num_A_detected = num_A_offspring_detected,
                    num_A_undetected = num_A_offspring_undetected,
                )
            end
        end
    end

    return nothing

end

function assign_to_arrays_and_times!(
    forecast::Forecast,
    day;
    is_omicron = false,
    num_S_detected = 0,
    num_S_undetected = 0,
    num_A_detected = 0,
    num_A_undetected = 0,
    num_I_detected = 0,
    num_I_undetected = 0,
)
    """
    Increment the arrays D and U based on sampled offsrping counts. The number of
    each offspring are keyword parameters which are default 0. This allows us to call
    this function from multiple spots.
    """

    Z = forecast.sim_realisation.Z
    D = forecast.sim_realisation.D
    U = forecast.sim_realisation.U
    individual_type_map = forecast.individual_type_map
    T_end = forecast.sim_features.T_end

    # sampling detected cases
    for _ in 1:num_S_detected
        (inf_time, onset_time) = sample_times(day, is_omicron)

        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.S] += 1
        end

        if onset_time < T_end
            D[map_day_to_index_UD(onset_time), individual_type_map.S] += 1
        end
    end

    for _ in 1:num_A_detected
        (inf_time, onset_time) = sample_times(day, is_omicron)

        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.A] += 1
        end

        if onset_time < T_end
            D[map_day_to_index_UD(onset_time), individual_type_map.A] += 1
        end
    end

    # sampling undetected cases
    for _ in 1:num_S_undetected
        (inf_time, onset_time) = sample_times(day, is_omicron)

        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.S] += 1
        end

        if onset_time < T_end
            U[map_day_to_index_UD(onset_time), individual_type_map.S] += 1
        end
    end

    for _ in 1:num_A_undetected
        (inf_time, onset_time) = sample_times(day, is_omicron)

        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.A] += 1
        end

        if onset_time < T_end
            U[map_day_to_index_UD(onset_time), individual_type_map.A] += 1
        end
    end

    # Sampling imports. These are treated slightly differently as the infection
    # time is inferred through the import model during initialisation
    for _ in 1:num_I_detected
        inf_time = day
        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.I] += 1
        end

        onset_time = inf_time + sample_onset_time(is_omicron)
        if onset_time < T_end
            # time+1 as first index is day 0
            D[map_day_to_index_UD(onset_time), individual_type_map.I] += 1
        end
    end

    for _ in 1:num_I_undetected
        inf_time = day
        if inf_time < T_end
            Z[map_day_to_index_Z(inf_time), individual_type_map.I] += 1
        end

        onset_time = inf_time + sample_onset_time(is_omicron)
        if onset_time < T_end
            # time+1 as first index is day 0
            U[map_day_to_index_UD(onset_time), individual_type_map.I] += 1
        end
    end

    return nothing

end

function clean_up_sim_results(
    Z_historical_results,
    D_results,
    U_results,
    good_sims,
)
    """
    A simple function that takes in the simulation output, cleans it up and spits out the
    processed format. We also account for the effect of susceptible depletion in the TP
    here so that the output includes the adjusted TP.
    """
    # return (
    #     Z_historical_results,
    #     D_results,
    #     U_results,
    # )
    return (
        Z_historical_results[:, 1:good_sims - 1],
        D_results[:, :, 1:good_sims - 1],
        U_results[:, :, 1:good_sims - 1],
    )
end

function simulate_branching_process(
    D0,
    N,
    nsims,
    local_cases,
    import_cases,
    forecast_start_date,
    date,
    state;
    adjust_TP = false,
)
    """
    Simulate branching process nsims times conditional on the cumulative observed cases,
    with initial state X0 = (S, A, I) (a named tuple). T is the duration of the simulation
    and N is the population size which is currently unused.
    """
    # read in the TPs for the state and strains
    (TP_indices, TP_local_delta, TP_import_delta) = get_single_state_TP(
        forecast_start_date,
        date,
        state,
        strain = "Delta",
        adjust_TP = adjust_TP,
    )

    (_, TP_local_omicron, TP_import_omicron) = get_single_state_TP(
        forecast_start_date,
        date,
        state,
        strain = "Omicron",
        adjust_TP = adjust_TP,
    )

    susceptible_depletion = read_in_susceptible_depletion(date)

    omicron_prop_pars = read_in_prop_omicron(date, state)

    max_restarts = 10
    good_sims = ones(Bool, nsims)

    # calculate the upper and lower limits for checking sims against data
    sim_features = get_simulation_limits(
        local_cases,
        forecast_start_date,
        TP_indices,
        N,
        state,
    )

    # the range of days for the simulation -1 as the "days", day ↦ day + 1
    day_range = -30:sim_features.T_end - 1

    earliest_start_date = Dates.Date(forecast_start_date) - Dates.Day(30)
    forecast_end_date = Dates.Date(forecast_start_date) + Dates.Day(sim_features.T_end - 1)

    # get simulation constants
    (sim_constants, individual_type_map) = set_simulation_constants(
        earliest_start_date,
        forecast_end_date,
        state,
    )

    # initialise state arrays (T_end + 1 as we include 0 day)
    sim_realisation = Realisation(sim_features.T_end)
    # always targeting 2000 sims
    sim_results = Results(sim_features.T_end, 2000)

    # put everything into a forecast object
    forecast = Forecast(
        sim_features,
        sim_realisation,
        sim_results,
        sim_constants,
        individual_type_map,
    )

    good_TPs_inds = zeros(Int, 2000)

    # Applying a reduction in NSW TP based off observations of the model fit. This issue is
    # at the beginning of the third wave so doesn't really influence the results for
    # Omicron. I DON'T LIKE THIS BUT FOR NOW IT'S CORRECTING SOME ISSUE IN THE TP/REFF
    # affecting NSW only.
    if state == "NSW" || state == "VIC"
        TP_ind = findfirst(40 == ind for ind in TP_indices)
        TP_local_delta[1:TP_ind, :] *= 0.5
        TP_local_omicron[1:TP_ind, :] *= 0.5
    end

    # number of TP samples saved
    num_TP = size(TP_local_delta, 2)

    # counts in each window are the same size as the limit arrays
    max_cases = forecast.sim_features.max_cases
    case_counts = zero(similar(max_cases))

    # a counter for the good sims, will terminate if good_sims > 2000
    good_sims = 1

    Z_results = forecast.sim_results.Z
    D_results = forecast.sim_results.D
    U_results = forecast.sim_results.U
    Z_historical_results = forecast.sim_results.Z_historical

    @showprogress for sim in 1:nsims
        # sample the TP/susceptible_depletion for this sim
        TP_ind = sim % num_TP == 0 ? num_TP : sim % num_TP
        susceptible_depletion_sim = susceptible_depletion[TP_ind, :]
        # susceptible_depletion_sim = susceptible_depletion[TP_ind]

        # reset boolean flags
        bad_sim = false
        injected_cases = false

        start_day = initialise_population!(
            forecast,
            D0,
        )

        import_cases_model!(
            forecast,
            import_cases,
            forecast_start_date,
        )

        # initialise counter to first day, when earliest infection occured based on
        # initialisation
        day = start_day

        # find index of the first TP to use based on the indices representing time
        # from forecast origin
        TP_start_day_ind = findfirst(start_day == d for d in TP_indices)
        TP_day_ind = TP_start_day_ind

        last_injection = 0
        n_restarts = 0

        while day <= day_range[end]
            # get TP's for the current day under the two strains and put in a named tuple
            TP_day = (
                TP_local_delta_parent = TP_local_delta[TP_day_ind, TP_ind],
                TP_import_delta_parent = TP_import_delta[TP_day_ind, TP_ind],
                TP_local_omicron_parent = TP_local_omicron[TP_day_ind, TP_ind],
                TP_import_omicron_parent = TP_import_omicron[TP_day_ind, TP_ind],
            )
            # get the sampled parameters on a given day
            prop_pars_day = (
                m0 = omicron_prop_pars["m0"][TP_ind],
                m1 = omicron_prop_pars["m1"][TP_ind],
                # m1 = 1.0,
                r = omicron_prop_pars["r"][TP_ind],
                tau = omicron_prop_pars["tau"][TP_ind],
            )

            sample_offspring!(
                forecast,
                day,
                TP_day,
                prop_pars_day,
                susceptible_depletion_sim,
                local_cases;
                adjust_TP = adjust_TP
            )

            reinitialise_allowed = day >= 3 &&
                day < forecast.sim_features.T_observed - 3 &&
                n_restarts < max_restarts &&
                day - last_injection >= 3
            # reinitialise_allowed = day >= 7 &&
            #     day < forecast.sim_features.T_observed - 7 &&
            #     n_restarts < max_restarts

            (bad_sim, injected_cases) = check_sim!(
                forecast,
                case_counts,
                local_cases,
                reinitialise_allowed;
                day = day,
                sim = sim,
            )

            if !injected_cases
                day += 1
                TP_day_ind += 1
            else
                n_restarts += 1
                if n_restarts > max_restarts
                    bad_sim = true
                end
                last_injection = day
                day = start_day
                TP_day_ind = TP_start_day_ind
                injected_cases = false
            end

            bad_sim && break

        end

        if !bad_sim
            Z_historical_results[:, good_sims] .= forecast.sim_realisation.Z_historical
            D_results[:, :, good_sims] .= forecast.sim_realisation.D
            U_results[:, :, good_sims] .= forecast.sim_realisation.U
            good_sims += 1
        end

        if good_sims > 2000
            break
        end

    end

    println("##########")
    println(state, " finished with ", good_sims - 1, " good simulations.")
    println("##########")

    results = clean_up_sim_results(
        Z_historical_results,
        D_results,
        U_results,
        good_sims,
    )

    return results

end