"""
    update_ista!(B, resid, grad, X, Y, Z, norms, lambda, reg, stepsize)

Updates coefficient estimates in place for each ISTA iteration when `X` and
`Z` are both standardized.

# Arguments

- B = 2d array of floats consisting of coefficient estimates
- resid = 2d array of floats consisting of the residuals
- grad = 2d array of floats consisting of the gradient
- X = 2d array of floats consisting of the row covariates, with all
  categorical variables coded in appropriate contrasts
- Y = 2d array of floats consisting of the multivariate response observations
- Z = 2d array of floats consisting of the column covariates, with all
  categorical variables coded in appropriate contrasts
- norms = `nothing`
- lambda = lambda penalty, a floating scalar
- reg = 2d array of bits, indicating whether or not to regularize each of the
  coefficients
- stepsize = 1d array consisting of a float; step size of updates

# Value

None; updates coefficients in place

"""
function update_ista!(B::AbstractArray{Float64,2},
                      resid::AbstractArray{Float64,2},
                      grad::AbstractArray{Float64,2},
                      X::AbstractArray{Float64,2},
                      Y::AbstractArray{Float64,2},
                      Z::AbstractArray{Float64,2},
                      norms::Nothing, lambda::Float64, reg::BitArray{2},
                      stepsize::AbstractArray{Float64,1})

    # Update gradient
    calc_grad!(grad, X, Z, resid)

    # Cycle through all the coefficients to perform assignments
    for j = 1:size(B,2), i = 1:size(B,1)
        b2update = B[i,j] - stepsize[1] * grad[i,j] # L2 update
        b2sign = sign(b2update) # L2 update sign

        # Update the current coefficient
        b = reg[i,j] ? prox(B[i,j], grad[i,j], b2sign, lambda, norms,
                            stepsize[1]) : b2update

        # Update coefficient estimate in the current copy
        if b != B[i,j]
            B[i,j] = b
        end
    end

    # Update residuals
    calc_resid!(resid, X, Y, Z, B)
end

"""
    update_ista!(B, resid, grad, X, Y, Z, norms, lambda, reg, stepsize)

Updates coefficient estimates in place for each ISTA iteration when `X` and
`Z` are not standardized.

# Arguments

- B = 2d array of floats consisting of coefficient estimates
- resid = 2d array of floats consisting of the residuals
- grad = 2d array of floats consisting of the gradient
- X = 2d array of floats consisting of the row covariates, with all
  categorical variables coded in appropriate contrasts
- Y = 2d array of floats consisting of the multivariate response observations
- Z = 2d array of floats consisting of the column covariates, with all
  categorical variables coded in appropriate contrasts
- norms = 2d array of floats consisting of the norms corresponding to each
  coefficient
- lambda = lambda penalty, a floating scalar
- reg = 2d array of bits, indicating whether or not to regularize each of the
  coefficients
- stepsize = 1d array consisting of a float; step size of updates

# Value

None; updates coefficients in place

"""
function update_ista!(B::AbstractArray{Float64,2},
                      resid::AbstractArray{Float64,2},
                      grad::AbstractArray{Float64,2},
                      X::AbstractArray{Float64,2},
                      Y::AbstractArray{Float64,2},
                      Z::AbstractArray{Float64,2},
                      norms::AbstractArray{Float64,2},
                      lambda::Float64, reg::BitArray{2},
                      stepsize::AbstractArray{Float64,1})

    # Update gradient
    calc_grad!(grad, X, Z, resid)

    # Cycle through all the coefficients to perform assignments
    for j = 1:size(B,2), i = 1:size(B,1)
        b2update = B[i,j] - stepsize[1] * grad[i,j]./norms[i,j] # L2 updates
        b2sign = sign(b2update) # L2 update signs

        # Update the current coefficient
        b = reg[i,j] ? prox(B[i,j], grad[i,j], b2sign, lambda, norms[i,j],
                            stepsize[1]) : b2update

        # Update coefficient estimate in the current copy
        if b != B[i,j]
            B[i,j] = b
        end
    end

    # Update residuals
    calc_resid!(resid, X, Y, Z, B)
end

"""
    ista!(X, Y, Z, lambda, B, regXidx, regZidx, reg, norms;
          isVerbose, stepsize, thresh, maxiter)

Performs ISTA with fixed step size.

# Arguments

- X = 2d array of floats consisting of the row covariates, with all
  categorical variables coded in appropriate contrasts
- Y = 2d array of floats consisting of the multivariate response
- Z = 2d array of floats consisting of the column covariates, with all
  categorical variables coded in appropriate contrasts
- lambda = lambda penalty, a floating scalar
- B = 2d array of floats consisting of starting coefficient estimates
- regXidx = 1d array of indices corresponding to regularized X covariates
- regZidx = 1d array of indices corresponding to regularized Z covariates
- reg = 2d array of bits, indicating whether or not to regularize each of the
  coefficients
- norms = 2d array of floats consisting of the norms corresponding to each
  coefficient or `nothing`

# Keyword arguments

- isVerbose = boolean flag indicating whether or not to print messages.
  Defaults to `true`.
- stepsize = float; step size for updates. Defaults to `0.01`.
- thresh = threshold at which the coefficients are considered to have
  converged, a floating scalar. Defaults to `10^(-7)`.
- maxiter = maximum number of update iterations. Defaults to `10^10`.

# Value

None; updates coefficients in place

# Some notes

Convergence is determined as when the log ratio of the current and previous
criteria is less than the threshold `thresh`.

The default method for choosing the fixed step size for `ista!` is to use the
reciprocal of the product of the maximum eigenvalues of `X*transpose(X)`
and `Z*transpose(Z)`. This is computed by the `mlmnet` function when `ista!`
is passed into the `fun` argument and `setStepSize` is set to `true`. If
`setStepSize` is set to `false`, the value of the `stepsize` argument will be
used as the fixed step size. Note that obtaining the eigenvalues when `X`
and/or `Z` are very large may exceed computational limitations.

"""
function ista!(X::AbstractArray{Float64,2}, Y::AbstractArray{Float64,2},
               Z::AbstractArray{Float64,2}, lambda::Float64,
               B::AbstractArray{Float64,2},
               regXidx::AbstractArray{Int64,1},
               regZidx::AbstractArray{Int64,1}, reg::BitArray{2}, norms;
               isVerbose::Bool=true, stepsize::Float64=0.01,
               thresh::Float64=10.0^(-7), maxiter::Int=10^10)

    # Pre-allocating arrays to store residuals using and gradient.
    resid = Array{Float64}(undef, size(Y))
    grad = Array{Float64}(undef, size(B))

    # Calculate residuals
    calc_resid!(resid, X, Y, Z, B)

    # Denominators of criterion
    crit_denom = [size(X,1)*size(Z,1), size(X,2)*size(Z,2)]
    # Placeholder to store the old criterion
    oldcrit = 1.0
    # Calculate the current criterion
    crit = criterion(B[regXidx, regZidx], resid, lambda, crit_denom)

    stepsize = [stepsize]
    iter = 0
    # Iterate until coefficients converge or maximum iterations have been
    # reached.
    while (abs(log(crit/oldcrit)) > thresh) && (iter < maxiter)
        # Store the current criterion
        oldcrit = crit
        # Update the coefficient estimates
        update_ista!(B, resid, grad, X, Y, Z, norms, lambda, reg, stepsize)
        # Calculate the criterion after updating
        crit = criterion(B[regXidx, regZidx], resid, lambda, crit_denom)

        iter += 1 # Increment the number of iterations
        # Warning message if coefficient estimates do not converge.
        if iter==(maxiter)
            println("Estimates did not converge.")
        end
    end

    println_verbose(string("Criterion: ", crit), isVerbose)
    println_verbose(string("Number of iterations: ", iter), isVerbose)
end


@polly function kernel_doitgen(A,C4,sum)
    nr,nq,np = size(A)

    for r = 1:nr, q = 1:nq
        for p = 1:np
            sum[p] = zero(eltype(sum))
            for s = 1:np
                sum[p] += A[r,q,s] * C4[s,p]
            end
        end
        for p = 1:np
            A[r,q,p] = sum[p]
        end
    end
end

let
    nq = 140
    nr = 150
    np = 160

    A = zeros(Float32,nr,nq,np)
    sum = zeros(Float32,np)
    C4 = zeros(Float32,np,np)

    for i = 1:nr, j = 1:nq, k = 1:np
        A[i,j,k] = ((i*j + k)%np) / np
    end
    for i = 1:np, j = 1:np
        C4[i,j] = (i*j % np) / np
    end

    SUITE["doitgen"] = @benchmarkable kernel_doitgen(A, C4, sum) setup = (A = copy($A); C4 = copy($C4); sum = copy($sum))
end


struct Twist{T} <: StaticVector{6,T}
    angular::AngularVelocity{T}
    linear::SVector{3,T}
end
function (::Type{T})(xs::NTuple{6}) where {T<:Twist}
    (x1, x2, x3, x4, x5, x6) = Base.promote(xs...)
    ω = AngularVelocity{typeof(x1)}(x1, x2, x3)
    v = SVector{3,typeof(x4)}(x4, x5, x6)
    return T(ω, v)
end

Base.@propagate_inbounds function Base.getindex(x::Twist, i::Int)
    return i < 4 ? getindex(x.angular, i) : getindex(x.linear, i - 3)
end

@inline Base.Tuple(t::Twist) = (Tuple(t.angular)..., Tuple(t.linear)...)

function groupexp(t::Twist)
    ω, b = t.angular, t.linear
    θ = rotation_angle(ω)
    rotation = groupexp(t.angular)
    translation = if iszero(θ)
        b
    else
        sinθ, cosθ = sincos(θ)
        θ² = θ^2
        p₁, p₂ = (1 - cosθ) / θ², (1 - sinθ / θ) / θ²
        ωxb = cross(ω, b)
        b + p₁ * ωxb + p₂ * cross(ω, ωxb)
    end
    return AffineMap(rotation, translation)
end

function grouplog(x::RigidMotion{3})
    r, t = x.linear, x.translation
    angular = grouplog(r)
    θ = rotation_angle(angular)
    linear = if iszero(θ)
        t
    else
        sinθ, cosθ = sincos(θ)
        θ² = θ^2
        p₂ = 1 / θ² - (1 + cosθ) / 2 / θ / sinθ
        ωxt = cross(angular, t)
        t - ωxt / 2 + p₂ * cross(angular, ωxt)
    end
    return Twist(angular, linear)
end


"""
    base_ip_model(nteams::Int, isphased::Bool)

Builds the base IP model, adding all base constraints for a RR2-timetable with
`nteams` teams and enforces a phased structure if `isphased == true`.

# Arguments

- `nteams::Int`: the number of teams participating
- `isphased::Bool`: sets the format of the tournament to be phased, enforcing
    each pair of teams to meet once in each half of the tournament

# Returns

- `Model` - Gurobi Optimizer model

# Model Description

The model uses the following binary variables:
- `x[i,j,s]` - Team `i` plays team `j` at home in slot `s`.
- `bh[i,s]` - Team `i` has a home break in slot `s`.
- `ba[i,s]` - Team `i` has an away break in slot `s`.
"""
function base_ip_model(nteams::Int, isphased::Bool)
    nslots = 2 * nteams - 2
    model = direct_model(Gurobi.Optimizer(GRB_ENV))

    @variable(model, x[1:nteams, 1:nteams, 1:nslots], Bin)
    @variable(model, bh[1:nteams, 1:nslots], Bin)
    @variable(model, ba[1:nteams, 1:nslots], Bin)

    @constraint(model, [i = 1:nteams], x[i, i, :] .== 0)
    @constraint(model, hgames[s = 1:nteams], sum(x[:, :, s]) == div(nteams, 2))
    @constraint(
        model,
        onefactor[i = 1:nteams, s = 1:nslots],
        sum(x[i, :, s]) + sum(x[:, i, s]) == 1
    )
    @constraint(
        model,
        alldiffh[i = 1:nteams, j = 1:nteams; i != j],
        sum(x[i, j, :]) == 1
    )
    @constraint(
        model,
        [i = 1:nteams, j = 1:nteams, s = 1:nslots],
        x[i, j, s] + x[j, i, s] <= 1
    )
    if isphased
        @constraint(
            model,
            phased[i = 1:nteams, j = 1:nteams; i != j],
            sum(x[i, j, 1:div(nslots, 2)]) + sum(x[j, i, 1:div(nslots, 2)]) ==
            1
        )
    end

    @constraint(
        model,
        [i = 1:nteams, s = 2:nslots],
        ba[i, s] - bh[i, s] + sum(x[i, :, s]) + sum(x[i, :, s-1]) == 1
    )

    @constraint(model, [i = 1:nteams, s = 2:nslots], bh[i, s] + ba[i, s] <= 1)

    return model
end

"""
    ip_model(inst::Instance)

Builds the base IP-model and populates it with the constraints given by `inst`.
"""
function ip_model(inst::Instance)

    @assert(length(inst.constrSE) <= 1)
    @assert(length(inst.constrFA) <= 1)

    rr2 = inst.rr2
    model = base_ip_model(rr2.nteams, rr2.isphased)

    # expression used to build objective function iterativly
    @expression(model, obj, AffExpr(0))

    # adding capacity constraints
    length(inst.constrCA) == 0 || @constraint(
        model,
        ca[i = 1:length(inst.constrCA)],
        expr!(model, inst.constrCA[i]) <= inst.constrCA[i].ub
    )

    # adding break constraints
    length(inst.constrBR) == 0 || @constraint(
        model,
        br[i = 1:length(inst.constrBR)],
        expr!(model, inst.constrBR[i]) <= inst.constrBR[i].ub
    )

    # adding game constraints
    if length(inst.constrGA) != 0
        @constraint(
            model,
            gaub[i = 1:length(inst.constrGA)],
            expr!(model, inst.constrGA[i], true) <= inst.constrGA[i].ub
        )
        @constraint(
            model,
            galb[i = 1:length(inst.constrGA)],
            inst.constrGA[i].lb <= expr!(model, inst.constrGA[i], false)
        )
    end

    # adding fairness constraints
    if length(inst.constrFA) > 0
        cfa = only(inst.constrFA)
        for i in cfa.teams
            for j in cfa.teams
                j <= i && continue
                offset = add_offset_var(model, cfa)
                for s in cfa.slots
                    @constraint(
                        model,
                        expr!(model, cfa, i, j, s) <= cfa.ub + offset
                    )
                    @constraint(
                        model,
                        expr!(model, cfa, j, i, s) <= cfa.ub + offset
                    )
                end
            end
        end
    end

    # adding separation constraints
    if length(inst.constrSE) > 0
        cse = only(inst.constrSE)
        for i in cse.teams
            for j in cse.teams
                j <= i && continue
                homefirst = @variable(model, binary = true)
                @constraint(
                    model,
                    expr!(model, cse, i, j, true) >=
                    cse.lb + 2 * rr2.nslots * homefirst - 2 * rr2.nslots
                )
                @constraint(
                    model,
                    expr!(model, cse, i, j, false) >=
                    cse.lb - 2 * rr2.nslots * homefirst
                )
            end
        end
    end

    @objective(model, Min, model[:obj])
    model
end

"""
    game_alloc_model(
        rr2::RR2,
        patternset::Vector{Int64},
        cons::Vector{Constraint};
        feas_relax = false,
    )

Creates the game allocation model for a given `patternset`. If
`feas_relax == true` a feasibility relaxation is computed, minimizing the offset
of the generic constraints `cons`.
"""
function game_alloc_model(
    rr2::RR2,
    patternset::Vector{Int64},
    cons::Vector{Constraint};
    feas_relax = false,
)
    ps(j, s) = patternset[j] >>> (s - 1) & 1
    gam = direct_model(Gurobi.Optimizer())

    @variable(gam, x[1:rr2.nteams, 1:rr2.nteams, 1:rr2.nslots], Bin)

    # setting pattern set
    for s = 1:rr2.nslots
        for i = 1:rr2.nteams
            for j = 1:rr2.nteams
                if ps(i, s) == 1
                    fix(x[j, i, s], 0; force = true)
                else
                    fix(x[i, j, s], 0; force = true)
                end
            end
        end
    end
    # one-factor
    @constraint(
        gam,
        [i = 1:rr2.nteams, s = 1:rr2.nslots],
        sum(x[i, :, s]) + sum(x[:, i, s]) == 1
    )
    # alldiff
    @constraint(
        gam,
        [i = 1:rr2.nteams, j = 1:rr2.nteams; i != j],
        sum(x[i, j, :]) == 1
    )

    # phased cons
    if rr2.isphased
        @constraint(
            gam,
            [i = 1:rr2.nteams, j = i+1:rr2.nteams],
            sum(x[i, j, 1:rr2.nslots÷2]) + sum(x[j, i, 1:rr2.nslots÷2]) == 1
        )
        @constraint(
            gam,
            [i = 1:rr2.nteams, j = i+1:rr2.nteams],
            sum(x[i, j, rr2.nslots÷2+1:end]) +
            sum(x[j, i, rr2.nslots÷2+1:end]) == 1
        )
    end

    offsets = Vector{VariableRef}()
    for con in cons
        @assert isa(con, CAcon) || isa(con, GAcon)
        if isa(con, CAcon)
            os = feas_relax ? @variable(gam, lower_bound = 0) : 0
            feas_relax && push!(offsets, os)
            @constraint(gam, expr!(gam, con) <= con.ub + os)
        elseif isa(con, GAcon)
            os1 = feas_relax ? @variable(gam, lower_bound = 0) : 0
            os2 = feas_relax ? @variable(gam, lower_bound = 0) : 0
            feas_relax && push!(offsets, os1, os2)
            @constraint(gam, expr!(gam, con, true) <= con.ub + os1)
            @constraint(gam, con.lb <= expr!(gam, con, false) + os2)
        end
    end
    if feas_relax
        @objective(gam, Min, sum(offsets))
    end
    return gam
end

"""
    add_offset_var(model::Model, con::Constraint)

Adds an offset variable for a given soft constraint and updates the objective
expression.
"""
function add_offset_var(model::Model, con::Constraint)
    offset = @variable(model, lower_bound = 0)
    add_to_expression!(model[:obj], con.pen, offset)
    offset
end

"""
    expr!(model::Model, con::Constraint, [args...])

Adds an expression to `model` representing the left-hand-side of the
IP-constraint modelling `con`, or if additional `args` are given, one of the
expresions used to model `con`. If the constraint is `soft` offset variables are
added accordingly.

# Arguments

- `model::Model` - The direct model of the Gurobi Optimizer
- `con::Constraint` - The constraint to be modeled
"""
function expr! end

"""
    expr!(model::Model, ca::CAcon)

Adds a capacity constraint expression for `ca`.
"""
function expr!(model::Model, ca::CAcon)
    x = model[:x]
    offset = ca.pen == 0 ? 0 : add_offset_var(model, ca)
    if ca.mode == MODE_H
        return @expression(
            model,
            sum(
                x[i, j, s] for i in ca.teams1 for j in ca.teams2 for
                s in ca.slots
            ) - offset
        )
    elseif ca.mode == MODE_A
        return @expression(
            model,
            sum(
                x[j, i, s] for i in ca.teams1 for j in ca.teams2 for
                s in ca.slots
            ) - offset
        )
    elseif ca.mode == MODE_HA
        return @expression(
            model,
            sum(
                x[j, i, s] + x[i, j, s] for i in ca.teams1 for j in ca.teams2
                for s in ca.slots
            ) - offset
        )
    end
end

"""
    expr!(model::Model, br::BRcon)

Adds a break constraint expression for `br`.
"""
function expr!(model::Model, br::BRcon)
    bh, ba = model[:bh], model[:ba]
    offset = br.pen == 0 ? 0 : add_offset_var(model, br)
    if br.mode == MODE_H
        return @expression(
            model,
            sum(bh[i, s] for i in br.teams for s in br.slots) - offset
        )
    elseif br.mode == MODE_A
        return @expression(
            model,
            sum(ba[i, s] for i in br.teams for s in br.slots) - offset
        )
    elseif br.mode == MODE_HA
        return @expression(
            model,
            sum(bh[i, s] + ba[i, s] for i in br.teams for s in br.slots) -
            offset
        )
    end
end

"""
    expr!(model::Model, ga::GAcon, upperbound)

Adds upper (`upperbound == true`) or lower (`upperbound == false`) bound
expession for `ga`.
"""
function expr!(model::Model, ga::GAcon, upperbound)
    x = model[:x]
    offset = ga.pen == 0 ? 0 : add_offset_var(model, ga)
    if upperbound
        return @expression(
            model,
            sum(
                x[ga.meetings[i][1], ga.meetings[i][2], s] for
                i = 1:length(ga.meetings) for s in ga.slots
            ) - offset
        )
    else
        return @expression(
            model,
            sum(
                x[ga.meetings[i][1], ga.meetings[i][2], s] for
                i = 1:length(ga.meetings) for s in ga.slots
            ) + offset
        )
    end
end

"""
    expr!(model::Model, se::SEcon, i::Int, j::Int, homefirst::Bool)

Adds expression for the separation of games between teams `i` and `j`.
`homefirst == true` is used to model the case that `i` plays the first game at
home, whereas `homefirst == false` is used to model the case that `i` plays the
first game away.
"""
function expr!(model::Model, se::SEcon, i::Int, j::Int, homefirst::Bool)
    x = model[:x]
    offset = se.pen == 0 ? 0 : add_offset_var(model, se)
    if homefirst
        return @expression(
            model,
            -sum(s * x[i, j, s] - s * x[j, i, s] for s = 1:size(x, 3)) +
            offset - 1
        )
    else
        return @expression(
            model,
            sum(s * x[i, j, s] - s * x[j, i, s] for s = 1:size(x, 3)) + offset -
            1
        )
    end
end

"""
    expr!(model::Model, fa::FAcon, i::Int, j::Int, s::Int)

Adds expression of the difference in home games played after slot `s` between
teams `i` and `j`.
"""
function expr!(model::Model, fa::FAcon, i::Int, j::Int, s::Int)
    x = model[:x]
    return @expression(model, sum(x[i, :, 1:s]) - sum(x[j, :, 1:s]))
end

"""
    parse_solution!(sched::Matrix{Int}, sol, model)

parses solution data `sol` into a schedule `sched`.
"""
function parse_solution!(sched::Matrix{Int}, sol, model)
    for s = 1:size(sched, 2)
        for i = 1:size(sched, 1)
            for j = 1:size(sched, 1)
                if sol[i, j, s] == 1.0
                    sched[i, s] = j
                    sched[j, s] = -i
                end
            end
        end
    end
    return sched
end

"""
    parse_solution(model)

Retrieves a solution schedule from `model`.
"""
function parse_solution(model)
    sol = value.(model[:x])
    sched = Array{Int64}(undef, size(sol, 1), size(sol, 3))
    return parse_solution!(sched, sol, model)
end


import CLBLAS
import OpenCL
    const clblas = CLBLAS
    clblas.setup()
    using OpenCL: cl

    device, ctx, queue = clblas.get_next_compute_context()

    N = Unsigned(5)
    alpha1 = cl.cl_float(5)
    alpha2 = cl.cl_float(2)
    data = Vector{Float32}(5)
    data[1] = 0.1
    data[2] = 0.22
    data[3] = 0.333
    data[4] = 0.4444
    data[5] = 0.55555

    X = cl.Buffer(Float32, ctx, (:rw, :copy), hostbuf=data)
    offx = Unsigned(0)
    incx = cl.cl_int(1)
    ncq = cl.cl_uint(1)
    clq = queue.id
    newl = cl.cl_uint(0)

    ptrEvent = clblas.clblasSscal(N, alpha1, X.id, offx, incx, [queue])

    ptrEvent2 = clblas.clblasSscal(N, alpha2, X.id, offx, incx,[queue], ptrEvent)

    cl.api.clWaitForEvents(cl.cl_uint(1), ptrEvent2)

    result = Vector{Float32}(length(X))
    cl.enqueue_read_buffer(queue, X, result, Unsigned(0), nothing, true)

    expected = data * alpha1 * alpha2
    println(result)
    println("------------")
    println(expected)

    try
        @assert(expected == result)
        info("success!")
    finally
        clblas.teardown()
    end


using Distributions, StatsBase

export imcmc_multinomial_edit_accept_reject!, unif_multinomial_sample_tester
export imcmc_trans_dim_accept_reject!, draw_sample!, draw_sample
export rand_delete!, rand_insert!, accept_reject!

function rand_delete!(x::Path{Int}, d::Int)

    n = length(x)
    k = d
    i = 0
    live_index = 0
    while k > 0
        u = rand()
        q = (n - k) / n
        while q > u  # skip
            i += 1
            n -= 1
            q *= (n - k) / n
        end
        i+=1
        # i is now index to delete
        deleteat!(x, i-live_index)
        live_index += 1
        n -= 1
        k -= 1
    end

end

function rand_insert!(x::Path{Int}, a::Int, V::UnitRange)

    n = length(x)+a
    k = a
    i = 0
    while k > 0
        u = rand()
        q = (n - k) / n
        while q > u  # skip
            i += 1
            n -= 1
            q *= (n - k) / n
        end
        i+=1
        # i is now index to insert
        insert!(x, i, rand(V))
        n -= 1
        k -= 1
    end
end

# Edit Allocation Move
# --------------------

function imcmc_multinomial_edit_accept_reject!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mode::InteractionSequence{Int},
    γ::Float64,
    dist::Metric,
    V::UnitRange,
    K_inner::DimensionRange,
    K_outer::DimensionRange,
    mcmc::Union{SisMcmcInsertDelete,SisMcmcSplitMerge}
    )

    dist_curr = mcmc.dist_curr

    N = length(S_curr)
    K_in_lb = K_inner.l
    K_in_ub = K_inner.u
    δ = rand(1:mcmc.ν_ed)  # Number of edits to enact
    rem_edits = δ # Remaining edits to allocate
    len_diffs = 0
    j = 0 # Keeps track how many interaction have been edited
    log_prod_term = 0.0

    # println("Making $δ edits in total...")

    for i in eachindex(S_curr)

        # If at end we just assign all remaining edits to final interaction
        if i == N
            δ_tmp = rem_edits
        # Otherwise we sample the number of edits via rescaled Binomial
        else
            p = 1/(N-i+1)
            δ_tmp = rand(Binomial(rem_edits, p)) # Number of edits to perform on ith interaction
        end

        # println("   Index $i getting $δ_tmp edits...")
        # If we sampled zero edits we skip to next iteration
        if δ_tmp == 0
            continue
        else
            j += 1 # Increment j
            # Make edits ....
            @inbounds n = length(S_curr[i])
            # a,b = (lb(n, δ_tmp, model), ub(n, δ_tmp))
            d = rand(0:min(n,δ_tmp))
            m = n + δ_tmp - 2*d

            # Catch invalid proposals
            if (m < K_in_lb) | (m > K_in_ub)
                # Here we just reject the proposal
                for i in 1:N
                    copy!(S_prop[i], S_curr[i])
                end
                return 0
            end

            I_tmp = S_prop[i]
            rand_delete!(I_tmp, d)
            rand_insert!(I_tmp, δ_tmp-d, V)

            mcmc.ind_update[j] = i # Store which interaction was updated

            log_prod_term += log(min(n, δ_tmp)+1) - log(min(m, δ_tmp)+1)
            len_diffs += m-n  # How much bigger the new interaction is
        end

        # Update rem_edits
        rem_edits -= δ_tmp

        # If no more left terminate
        if rem_edits == 0
            break
        end

    end

    # # Add final part of log_ratio term
    log_ratio = log(length(V)) * len_diffs + log_prod_term
    # log_ratio = log_dim_diff + log_prod_term
    dist_prop = dist(mode, S_prop)
    log_lik_ratio = -γ * (
        dist_prop - dist_curr[1]
    )

    # Log acceptance probability
    log_α = log_lik_ratio + log_ratio

    # Accept-reject step. Use info in mcmc.ind_update to know which interaction are to be copied over
    if log(rand()) < log_α
        for i in view(mcmc.ind_update, 1:j)
            copy!(S_curr[i], S_prop[i])
        end
        dist_curr[1] = dist_prop # Update stored distance
        return 1
    else
        for i in view(mcmc.ind_update, 1:j)
            copy!(S_prop[i], S_curr[i])
        end
        return 0
    end
end

function imcmc_multinomial_edit_accept_reject!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    model::SIS,
    mcmc::Union{SisMcmcInsertDelete,SisMcmcSplitMerge}
    )

    return imcmc_multinomial_edit_accept_reject!(
        S_curr, S_prop,
        model.mode, model.γ,
        model.dist, model.V,
        model.K_inner, model.K_outer,
        mcmc
    )
end

# To check the above code is doing as required. Output should be samples from a uniform
# with n trials and k bins.
function unif_multinomial_sample_tester(k, n)
    n_rem = n
    x = Int[]
    for i in 1:(k-1)
        p = 1/(k-i+1)
        z = rand(Binomial(n_rem, p))
        push!(x, z)
        n_rem -= z
    end
    push!(x, n_rem)
    return x
end

# Trans-dimensional Move
# ----------------------

function migrate!(
    y::Vector{Vector{Int}}, x::Vector{Vector{Int}},
    j::Int, i::Int)
    @inbounds insert!(y, j, x[i])
    @inbounds deleteat!(x, i)
end

function imcmc_multi_insert_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mcmc::T,
    ind::AbstractVector{Int},
    V::UnitRange,
    K_in_ub::Int
    ) where {T<:Union{SisMcmcSampler,SimMcmcSampler}}

    prop_pointers = mcmc.prop_pointers
    ν_td = mcmc.ν_td
    N = length(S_curr)
    len_dist = mcmc.len_dist

    log_ratio = 0.0
    for i in ind
        tmp = popfirst!(prop_pointers)
        m = rand(len_dist)
        resize!(tmp, m)
        sample!(V, tmp)
        @inbounds insert!(S_prop, i, tmp)
        log_ratio += - logpdf(len_dist, m) + m*log(length(V)) - Inf * (m > K_in_ub)
    end
    return log_ratio

end

function imcmc_multi_delete_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mcmc::T,
    ind::AbstractVector{Int},
    V::UnitRange
    ) where {T<:Union{SisMcmcSampler,SimMcmcSampler}}

    prop_pointers = mcmc.prop_pointers
    ν_td = mcmc.ν_td
    N = length(S_curr)
    len_dist = mcmc.len_dist

    log_ratio = 0.0

    for i in Iterators.reverse(ind)
        @inbounds tmp = popat!(S_prop, i)
        pushfirst!(prop_pointers, tmp)
        m = length(tmp)
        log_ratio += logpdf(len_dist, m) - m * log(length(V))
    end
    return log_ratio

end

function imcmc_multi_insert_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mcmc::T,
    ε::Int,
    V::UnitRange,
    K_in_ub::Int
    ) where {T<:Union{SisMcmcSampler,SimMcmcSampler}}

    prop_pointers = mcmc.prop_pointers
    ν_td = mcmc.ν_td
    N = length(S_curr)
    len_dist = mcmc.len_dist
    log_ratio = 0.0
    ind = mcmc.ind_td_add

    n = length(S_prop)+ε
    k = ε
    i = 0
    j = 0
    while k > 0
        u = rand()
        q = (n - k) / n
        while q > u  # skip
            i += 1
            n -= 1
            q *= (n - k) / n
        end
        i += 1
        j += 1
        # Insert path at index i
        # i is now index to insert
        ind[j] = i
        tmp = popfirst!(prop_pointers) # Get storage
        m = rand(len_dist)  # Sample length
        resize!(tmp, m) # Resize
        sample!(V, tmp) # Sample new entries uniformly
        @inbounds insert!(S_prop, i, tmp) # Insert path into S_prop
        log_ratio += - logpdf(len_dist, m) + m*log(length(V)) - Inf * (m > K_in_ub)  # Add to log_ratio term
        n -= 1
        k -= 1
    end
    return log_ratio

end

function imcmc_multi_delete_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mcmc::T,
    ε::Int,
    V::UnitRange
    ) where {T<:Union{SisMcmcSampler,SimMcmcSampler}}

    prop_pointers = mcmc.prop_pointers
    N = length(S_curr)
    len_dist = mcmc.len_dist
    log_ratio = 0.0
    ind = mcmc.ind_td_del # Store which entries were deleted

    n = length(S_prop)
    k = ε
    i = 0
    j = 0
    live_index = 0
    while k > 0
        u = rand()
        q = (n - k) / n
        while q > u  # skip
            i += 1
            n -= 1
            q *= (n - k) / n
        end
        i+=1
        j+=1
        # Delete path
        # i is now index to delete
        @inbounds ind[j] = i
        @inbounds tmp = popat!(S_prop, i - live_index)
        pushfirst!(prop_pointers, tmp)
        m = length(tmp)
        log_ratio += logpdf(len_dist, m) - m * log(length(V))
        live_index += 1
        n -= 1
        k -= 1
    end
    return log_ratio
end

function imcmc_trans_dim_accept_reject!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mode::InteractionSequence{Int},
    γ::Float64,
    dist::Metric,
    V::UnitRange,
    K_inner::DimensionRange,
    K_outer::DimensionRange,
    mcmc::Union{SisMcmcInsertDeleteGibbs,SisMcmcInsertDelete}
    )

    K_out_lb = K_outer.l
    K_out_ub = K_outer.u
    K_in_ub = K_inner.u
    ν_td = mcmc.ν_td
    curr_pointers = mcmc.curr_pointers
    prop_pointers = mcmc.prop_pointers

    dist_curr = mcmc.dist_curr

    log_ratio = 0.0

    # Enact insertion / deletion
    N = length(S_curr)
    ε = rand(1:ν_td)
    d = rand(0:min(ε, N))
    a = ε - d

    # Catch invalid proposal (outside dimension bounds)
    M = N - d + a  # Number of paths in proposal
    if (M < K_out_lb) | (M > K_out_ub)
        return false
    end
    if d > 0
        log_ratio +=  imcmc_multi_delete_prop_sample!(
            S_curr, S_prop,
            mcmc,
            d,
            V
        ) # Enact deletion move and catch log ratio term
    end
    if a > 0  # Number of insertions
        log_ratio += imcmc_multi_insert_prop_sample!(
            S_curr, S_prop,
            mcmc,
            a,
            V, K_in_ub
        )
    end

    log_ratio += log(min(ε, N) + 1) - log(min(ε, M) + 1)

    # Now do accept-reject step
    dist_prop = dist(mode, S_prop)
    log_α = - γ * (
        dist_prop - dist_curr[1]
    ) + log_ratio

    # Note that we copy interactions between S_prop (resp. S_curr) and prop_pointers (resp .curr_pointers) by hand.
    if log(rand()) < log_α
        # Do deletions to S_curr
        ind = view(mcmc.ind_td_del, 1:d)
        for i in Iterators.reverse(ind) # If we didnt do reverse would have to update indices
            @inbounds tmp = popat!(S_curr, i)
            pushfirst!(curr_pointers, tmp)
        end
        # Do insertions
        ind = view(mcmc.ind_td_add, 1:a)
        for i in ind
            tmp = popfirst!(curr_pointers)
            copy!(tmp, S_prop[i])
            insert!(S_curr, i, tmp)
        end
        dist_curr[1] = dist_prop  # Update stored distance
        return true
    else
        # Remove insertions and return to prop_pointers
        ind = view(mcmc.ind_td_add, 1:a)
        for i in Iterators.reverse(ind)
            @inbounds tmp = popat!(S_prop, i)
            pushfirst!(prop_pointers, tmp)
        end
        # Re-insert deletions
        ind = view(mcmc.ind_td_del, 1:d)
        for i in ind
            @inbounds tmp = popfirst!(prop_pointers)
            copy!(tmp, S_curr[i])
            insert!(S_prop, i, tmp)
        end
        return false
    end
end

function imcmc_trans_dim_accept_reject!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    model::SIS,
    mcmc::Union{SisMcmcInsertDeleteGibbs,SisMcmcInsertDelete}
    )
    return imcmc_trans_dim_accept_reject!(
        S_curr, S_prop,
        model.mode, model.γ,
        model.dist, model.V,
        model.K_inner, model.K_outer,
        mcmc
    )
end

function accept_reject!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mode::InteractionSequence{Int},
    γ::Float64,
    dist::Metric,
    V::UnitRange,
    K_inner::DimensionRange,
    K_outer::DimensionRange,
    mcmc::Union{SisMcmcInsertDeleteGibbs,SisMcmcInsertDelete}
    )

    β = mcmc.β
    if rand() < β
        return imcmc_multinomial_edit_accept_reject!(
            S_curr, S_prop,
            mode, γ, dist, V, K_inner, K_outer,
            mcmc
        )
    else
        return imcmc_trans_dim_accept_reject!(
            S_curr, S_prop,
            mode, γ, dist, V, K_inner, K_outer,
            mcmc
        )
    end

end

function accept_reject!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    model::SIS,
    mcmc::SisMcmcInsertDelete
    )
    β = mcmc.β
    if rand() < β
        return imcmc_multinomial_edit_accept_reject!(
            S_curr, S_prop,
            model, mcmc
        )
    else
        return imcmc_trans_dim_accept_reject!(
            S_curr, S_prop,
            model, mcmc
        )
    end

end

function accept_reject!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    model::SIS,
    mcmc::SisMcmcInsertDelete,
    count::Vector{Int},
    acc_count::Vector{Int}
    )
    β = mcmc.β
    if rand() < β

        was_acc = imcmc_multinomial_edit_accept_reject!(
            S_curr, S_prop,
            model, mcmc
        )
        count[1] += 1
        acc_count[1] += was_acc
    else
        was_acc = imcmc_trans_dim_accept_reject!(
            S_curr, S_prop,
            model, mcmc
        )
        count[2] += 1
        acc_count[2] += was_acc
    end

end

# Sampler Functions
# -----------------
"""
    draw_sample!(
        sample_out::InteractionSequenceSample,
        mcmc::SisMcmcInsertDelete,
        model::SIS;
        burn_in::Int=mcmc.burn_in,
        lag::Int=mcmc.lag,
        init::InteractionSequence=get_init(model, mcmc.init)
    )

Draw sample in-place from given SIS model `model::SIS` via MCMC algorithm with edit allocation and interaction insertion/deletion, storing output in `sample_out::InteractionSequenceSample`.

Accepts keyword arguments to change MCMC output, including burn-in, lag and initial values. If not given, these are set to the default values of the passed MCMC sampler `mcmc::SisMcmcInsertDelete`.
"""
function draw_sample!(
    sample_out::Union{InteractionSequenceSample{Int}, SubArray},
    mcmc::SisMcmcInsertDelete,
    model::SIS;
    burn_in::Int=mcmc.burn_in,
    lag::Int=mcmc.lag,
    init::InteractionSequence{Int}=get_init(mcmc.init, model)
    )

    # Define aliases for pointers to the storage of current vals and proposals
    curr_pointers = mcmc.curr_pointers
    prop_pointers = mcmc.prop_pointers

    S_curr = InteractionSequence{Int}()
    S_prop = InteractionSequence{Int}()
    for i in 1:length(init)
        migrate!(S_curr, curr_pointers, i, 1)
        migrate!(S_prop, prop_pointers, i, 1)
        copy!(S_curr[i], init[i])
        copy!(S_prop[i], init[i])
    end

    sample_count = 1 # Keeps which sample to be stored we are working to get
    i = 0 # Keeps track all samples (included lags and burn_ins)
    count = [0,0]
    acc_count = [0,0]

    mcmc.dist_curr[1] = model.dist(S_curr, model.mode) # Initialise stored distance

    while sample_count ≤ length(sample_out)
        i += 1
        # Store value
        if (i > burn_in) & (((i-1) % lag)==0)
            sample_out[sample_count] = deepcopy(S_curr)
            sample_count += 1
        end
        # Accept reject
        accept_reject!(S_curr, S_prop, model, mcmc, count, acc_count)
    end
    for i in 1:length(S_curr)
        migrate!(curr_pointers, S_curr, 1, 1)
        migrate!(prop_pointers, S_prop, 1, 1)
    end
    return count, acc_count
end

"""
    draw_sample(
        mcmc::SisMcmcInsertDelete,
        model::SIS;
        desired_samples::Int=mcmc.desired_samples,
        burn_in::Int=mcmc.burn_in,
        lag::Int=mcmc.lag,
        init::Vector{Path{T}}=get_init(model, mcmc.init)
        )
"""
function draw_sample(
    mcmc::SisMcmcInsertDelete,
    model::SIS;
    desired_samples::Int=mcmc.desired_samples,
    burn_in::Int=mcmc.burn_in,
    lag::Int=mcmc.lag,
    init::InteractionSequence{Int}=get_init(mcmc.init, model)
    )

    sample_out = InteractionSequenceSample{Int}(undef, desired_samples)
    draw_sample!(sample_out, mcmc, model, burn_in=burn_in, lag=lag, init=init)
    return sample_out

end

function (mcmc::SisMcmcInsertDelete)(
    model::SIS;
    desired_samples::Int=mcmc.desired_samples,
    burn_in::Int=mcmc.burn_in,
    lag::Int=mcmc.lag,
    init::InteractionSequence{Int}=get_init(mcmc.init, model)
    )

    sample_out = InteractionSequenceSample{Int}(undef, desired_samples)
    (count, acc_count) = draw_sample!(sample_out, mcmc, model, burn_in=burn_in, lag=lag, init=init)

    p_measures = Dict(
            "Proportion Update Moves" => count[1]/sum(count),
            "Update Move Acceptance Probability" => acc_count[1] / count[1],
            "Trans-Dimensional Move Acceptance Probability" => acc_count[2] / count[2]
        )
    output = SisMcmcOutput(
            model,
            sample_out,
            p_measures
            )

    return output

end


#=
Objective
---------
Use customized root-finding method.

Finding
-------
Name tuple is allowed
=#

include("testutils.jl")

de = DEsteady(func=bistable_ode!, p=Vector(p_), u0= Vector(u_1), method=SSRootfind(nlsolve = (f,u0,abstol) -> (res=NLsolve.nlsolve(f,u0,autodiff=:forward,method=:newton,iterations=Int(1e6),ftol=abstol);res.zero) ))

sol = solve(de)
sol.retcode == :Success

#=
With boundary
=#
de_b = DEsteady(de, method=SSRootfind(nlsolve = (f,u0,abstol) -> (res=NLsolve.nlsolve(f,u0, zeros(length(u0)), fill(Inf, length(u0)) ,autodiff=:forward,method=:newton,iterations=Int(1e6),ftol=abstol);res.zero) )
)


struct KrylovTrustRegion{T <: Real} <: SecondOrderOptimizer
    initial_radius::T
    max_radius::T
    eta::T
    rho_lower::T
    rho_upper::T
    cg_tol::T
end

KrylovTrustRegion(; initial_radius::Real = 1.0,
                    max_radius::Real = 100.0,
                    eta::Real = 0.1,
                    rho_lower::Real = 0.25,
                    rho_upper::Real = 0.75,
                    cg_tol::Real = 0.01) =
                    KrylovTrustRegion(initial_radius, max_radius, eta,
                                  rho_lower, rho_upper, cg_tol)

update_h!(d, state, method::KrylovTrustRegion) = nothing

# TODO: support x::Array{T,N} et al.?
mutable struct KrylovTrustRegionState{T} <: AbstractOptimizerState
    x::Vector{T}
    x_previous::Vector{T}
    f_x_previous::T
    s::Vector{T}
    interior::Bool
    accept_step::Bool
    radius::T
    m_diff::T
    f_diff::T
    rho::T
    r::Vector{T}  # residual vector
    d::Vector{T}  # direction to consider
    cg_iters::Int
end

function initial_state(method::KrylovTrustRegion, options, d, initial_x::Array{T}) where T
    n = length(initial_x)
    # Maintain current gradient in gr
    @assert(method.max_radius > 0)
    @assert(0 < method.initial_radius < method.max_radius)
    @assert(0 <= method.eta < method.rho_lower)
    @assert(method.rho_lower < method.rho_upper)
    @assert(method.rho_lower >= 0)

    value_gradient!!(d, initial_x)

    KrylovTrustRegionState(copy(initial_x),    # Maintain current state in state.x
                           copy(initial_x), # x_previous
                           zero(T),            # f_x_previous
                           similar(initial_x), # Maintain current search direction in state.s
                           true,               # interior
                           true,               # accept step
                           method.initial_radius,
                           zero(T),            # model change
                           zero(T),            # observed f change
                           zero(T),            # state.rho
                           Vector{T}(undef, n),       # residual vector
                           Vector{T}(undef, n),       # direction to consider
                           0)                  # cg_iters
end

function trace!(tr, d, state, iteration, method::KrylovTrustRegion, options, curr_time=time())
    dt = Dict()
    dt["time"] = curr_time
    if options.extended_trace
        dt["radius"] = copy(state.radius)
        dt["interior"] = state.interior
        dt["accept_step"] = state.accept_step
        dt["norm(s)"] = norm(state.s)
        dt["rho"] = state.rho
        dt["m_diff"] = state.m_diff
        dt["f_diff"] = state.f_diff
        dt["cg_iters"] = state.cg_iters
    end
    g_norm = norm(gradient(d), Inf)
    update!(tr,
            iteration,
            value(d),
            g_norm,
            dt,
            options.store_trace,
            options.show_trace,
            options.show_every,
            options.callback)
end

function cg_steihaug!(objective::TwiceDifferentiableHV,
                      state::KrylovTrustRegionState{T},
                      method::KrylovTrustRegion) where T
    n = length(state.x)
    x, g, d, r, z, Hd = state.x, gradient(objective), state.d, state.r, state.s, hv_product(objective)

    fill!(z, 0.0)  # the search direction is initialized to the 0 vector,
    r .= g  # so at first the whole gradient is the residual.
    d .= -r # the first direction is the direction of steepest descent.
    rho0 = 1e100  # just a big number

    state.cg_iters = 0
    for i in 1:n
        state.cg_iters += 1
        hv_product!(objective, x, d)
        dHd = dot(d, Hd)
        if -1e-15 < dHd < 1e-15
            break
        end

        alpha = dot(r, r) / dHd

        if dHd < 0. || norm(z .+ alpha .* d) >= state.radius
            a_ = dot(d, d)
            b_ = 2 * dot(z, d)
            c_ = dot(z, z) - state.radius^2
            tau = (-b_ + sqrt(b_ * b_ - 4 * a_ * c_)) / (2 * a_)
            z .+= tau .* d
            break
        end

        z .+= alpha .* d
        rho_prev = dot(r, r)
        if i == 1
            rho0 = rho_prev
        end
        r .+= alpha * Hd
        rho_next = dot(r, r)
        r_sqnorm_ratio = rho_next / rho_prev
        d[:] = -r + r_sqnorm_ratio * d

        if (rho_next / rho0) < method.cg_tol^2
            break
        end
    end

    hv_product!(objective, x, z)
    return dot(g, z) + 0.5 * dot(z, Hd)
end

function update_state!(objective::TwiceDifferentiableHV,
                          state::KrylovTrustRegionState,
                          method::KrylovTrustRegion)
    state.m_diff = cg_steihaug!(objective, state, method)
    @assert state.m_diff <= 0

    state.f_diff = value(objective, state.x .+ state.s) - value(objective)
    state.rho = state.f_diff / state.m_diff
    state.interior = norm(state.s) < 0.9 * state.radius

    if state.rho < method.rho_lower
        state.radius *= 0.25
    elseif (state.rho > method.rho_upper) && (!state.interior)
        state.radius = min(2 * state.radius, method.max_radius)
    end

    state.accept_step = state.rho > method.eta
    if state.accept_step
        state.x .+= state.s
    end

    return false
end

function update_g!(objective, state::KrylovTrustRegionState, method::KrylovTrustRegion)
    if state.accept_step
        # Update the function value and gradient
        state.f_x_previous = value(objective)
        value_gradient!(objective, state.x)
    end
end

function assess_convergence(state::KrylovTrustRegionState, d, options::Options)
    if !state.accept_step
        return state.radius < options.x_abstol, false, false, false, false
    end

    x_converged, f_converged, f_increased, g_converged = false, false, false, false

    if norm(state.s, Inf) < options.x_abstol
        x_converged = true
    end

    # Absolute Tolerance
    # if abs(f_x - f_x_previous) < f_tol
    # Relative Tolerance
    if abs(state.f_diff) < max(options.f_reltol * (abs(value(d)) + options.f_reltol), eps(abs(value(d))+abs(state.f_x_previous)))
        f_converged = true
    end

    if norm(gradient(d), Inf) < options.g_abstol
        g_converged = true
    end

    converged = x_converged || f_converged || g_converged

    return x_converged, f_converged, g_converged, converged, false
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

function boykov_kolmogorov_cut(A::AbstractArray{T,N}, B::AbstractArray{T,N}, dim::Integer) where {N,T}
  @assert size(A) == size(B) "arrays must have the same size for cut"

  # size and number of voxels
  sz   = size(A)
  nvox = prod(sz)

  # source and sink terminals
  s = nvox + 1; t = nvox + 2

  # lattice graph and capacity matrix
  G = DiGraph(nvox+2)
  C = spzeros(nvox+2, nvox+2)

  # fill lattice graph with original vertices
  for d=1:N
    # loop over all indices except the borders using an increment direction
    inds = CartesianIndices(ntuple(i -> i == d ? (1:sz[d]-1) : (1:sz[i]), N))
    incr = CartesianIndex(ntuple(i -> i == d ? 1 : 0, N))

    for ind in inds
      # adjacent vertices along direction
      u = cart2lin(sz, ind)
      v = cart2lin(sz, ind + incr)

      # simple difference
      Du = abs(A[u] - B[u])
      Dv = abs(A[v] - B[v])

      # gradient along direction
      ∇Au = abs(A[v] - A[u])
      ∇Bu = abs(B[v] - B[u])

      # next vertex along direction
      w = ind + 2incr

      # repeat gradient on the border when
      # outside valid indices for arrays
      ok = all(w.I .≤ sz)
      ∇Av = ok ? abs(A[w] - A[v]) : ∇Au
      ∇Bv = ok ? abs(B[w] - B[v]) : ∇Bu

      add_edge!(G, u, v)
      add_edge!(G, v, u)

      C[u,v] = C[v,u] = (Du + Dv) / (∇Au + ∇Av + ∇Bu + ∇Bv + eps())
    end
  end

  # fill edges from source terminal to left slice
  linds = CartesianIndices(ntuple(i -> i == dim ? (1:1) : (1:sz[i]), N))
  for ind in linds
    u = cart2lin(sz, ind)
    add_edge!(G, s, u)
    C[s,u] = Inf
  end

  # fill edges from right slice to sink terminal
  rinds = CartesianIndices(ntuple(i -> i == dim ? (sz[dim]:sz[dim]) : (1:sz[i]), N))
  for ind in rinds
    v = cart2lin(sz, ind)
    add_edge!(G, v, t)
    C[v,t] = Inf
  end

  # Boykov-Kolmogorov max-flow/min-cut
  _, __, labels = maximum_flow(G, s, t, C, algorithm=BoykovKolmogorovAlgorithm())

  # remove source and sink terminals
  labels = labels[1:end-2]

  # cut mask
  M = falses(size(A))
  M[labels .== 1] .= true
  M[labels .== 0] .= true

  M
end


using Test

using LinearAlgebra, Statistics
using LowRankVortex
using TransportBasedInference
using ForwardDiff
using PotentialFlow
import PotentialFlow.Plates: Plate, Points, Blobs
import PotentialFlow.Motions: reset_velocity!
import PotentialFlow.Points
import PotentialFlow.Elements
import PotentialFlow.Properties: @property
import PotentialFlow.Elements: jacobian_position, jacobian_strength, jacobian_param

include("vortex.jl")
include("forecast.jl")
include("convective_complexpotential.jl")
include("pressure.jl")
include("symmetric_pressure.jl")
include("symmetric_analytical_jacobian_pressure_freestream.jl")
include("symmetric_analytical_jacobian_pressure.jl")
include("AD_pressure.jl")
include("analytical_pressure.jl")


# This file is a part of Julia. License is MIT: http://julialang.org/license

module TridiagTest
using Base.Test
debug = false

# basic tridiagonal operations
n = 5

srand(123)

d = 1 .+ rand(n)
dl = -rand(n-1)
du = -rand(n-1)
v = randn(n)
B = randn(n,2)

for elty in (Float32, Float64, Complex64, Complex128, Int)
    if elty == Int
        srand(61516384)
        d = rand(1:100, n)
        dl = -rand(0:10, n-1)
        du = -rand(0:10, n-1)
        v = rand(1:100, n)
        B = rand(1:100, n, 2)
    else
        d = convert(Vector{elty}, d)
        dl = convert(Vector{elty}, dl)
        du = convert(Vector{elty}, du)
        v = convert(Vector{elty}, v)
        B = convert(Matrix{elty}, B)
    end
    ε = eps(abs2(float(one(elty))))
    T = Tridiagonal(dl, d, du)
    Ts = SymTridiagonal(d, dl)
    @test_throws ArgumentError size(Ts,0)
    @test size(Ts,3) == 1
    @test size(T, 1) == n
    @test size(T) == (n, n)
    F = diagm(d)
    for i = 1:n-1
        F[i,i+1] = du[i]
        F[i+1,i] = dl[i]
    end
    @test Array(T) == F

    # elementary operations on tridiagonals
    @test conj(T) == Tridiagonal(conj(dl), conj(d), conj(du))
    @test transpose(T) == Tridiagonal(du, d, dl)
    @test ctranspose(T) == Tridiagonal(conj(du), conj(d), conj(dl))

    @test abs.(T) == Tridiagonal(abs.(dl),abs.(d),abs.(du))
    if elty <: Real
        @test round.(T) == Tridiagonal(round.(dl),round.(d),round.(du))
        @test isa(round.(T), Tridiagonal)
        @test trunc.(T) == Tridiagonal(trunc.(dl),trunc.(d),trunc.(du))
        @test isa(trunc.(T), Tridiagonal)
        @test floor.(T) == Tridiagonal(floor.(dl),floor.(d),floor.(du))
        @test isa(floor.(T), Tridiagonal)
        @test ceil.(T) == Tridiagonal(ceil.(dl),ceil.(d),ceil.(du))
        @test isa(ceil.(T), Tridiagonal)
    end
    @test real(T) == Tridiagonal(real(dl),real(d),real(du))
    @test imag(T) == Tridiagonal(imag(dl),imag(d),imag(du))
    @test abs.(Ts) == SymTridiagonal(abs.(d),abs.(dl))
    if elty <: Real
        @test round.(Ts) == SymTridiagonal(round.(d),round.(dl))
        @test isa(round.(Ts), SymTridiagonal)
        @test trunc.(Ts) == SymTridiagonal(trunc.(d),trunc.(dl))
        @test isa(trunc.(Ts), SymTridiagonal)
        @test floor.(Ts) == SymTridiagonal(floor.(d),floor.(dl))
        @test isa(floor.(Ts), SymTridiagonal)
        @test ceil.(Ts) == SymTridiagonal(ceil.(d),ceil.(dl))
        @test isa(ceil.(Ts), SymTridiagonal)
    end
    @test real(Ts) == SymTridiagonal(real(d),real(dl))
    @test imag(Ts) == SymTridiagonal(imag(d),imag(dl))

    # test interconversion of Tridiagonal and SymTridiagonal
    @test Tridiagonal(dl, d, dl) == SymTridiagonal(d, dl)
    @test SymTridiagonal(d, dl) == Tridiagonal(dl, d, dl)
    @test Tridiagonal(dl, d, du) + Tridiagonal(du, d, dl) == SymTridiagonal(2d, dl+du)
    @test SymTridiagonal(d, dl) + Tridiagonal(dl, d, du) == Tridiagonal(dl + dl, d+d, dl+du)
    @test convert(SymTridiagonal,Tridiagonal(Ts)) == Ts
    @test Array(convert(SymTridiagonal{Complex64},Tridiagonal(Ts))) == convert(Matrix{Complex64}, Ts)
    if elty == Int
        vv = rand(1:100, n)
        BB = rand(1:100, n, 2)
    else
        vv = convert(Vector{elty}, v)
        BB = convert(Matrix{elty}, B)
    end
    let Bs = BB, vs = vv
        for atype in ("Array", "SubArray")
            if atype == "Array"
                BB = Bs
                vv = vs
            else
                BB = view(Bs, 1:n, 1)
                vv = view(vs, 1:n)
            end
        end

        # tridiagonal linear algebra
        @test T*vv ≈ F*vv
        invFv = F\vv
        @test T\vv ≈ invFv
        # @test Base.solve(T,v) ≈ invFv
        # @test Base.solve(T, B) ≈ F\B
        Tlu = factorize(T)
        x = Tlu\vv
        @test x ≈ invFv
    end
    @test det(T) ≈ det(F)

    @test T*Base.LinAlg.UnitUpperTriangular(eye(n)) ≈ F*eye(n)
    @test T*Base.LinAlg.UnitLowerTriangular(eye(n)) ≈ F*eye(n)
    @test T*UpperTriangular(eye(n)) ≈ F*eye(n)
    @test T*LowerTriangular(eye(n)) ≈ F*eye(n)

    # symmetric tridiagonal
    if elty <: Real
        Ts = SymTridiagonal(d, dl)
        Fs = Array(Ts)
        Tldlt = factorize(Ts)
        @test_throws DimensionMismatch Tldlt\rand(elty,n+1)
        @test size(Tldlt) == size(Ts)
        if elty <: AbstractFloat
            @test typeof(convert(Base.LinAlg.LDLt{Float32},Tldlt)) == Base.LinAlg.LDLt{Float32,SymTridiagonal{elty}}
        end
        let vs = vv
            for atype in ("Array", "SubArray")
                if atype == "Array"
                    vv = vs
                else
                    vv = view(vs, 1:n)
                end
            end
            invFsv = Fs\vv
            x = Ts\vv
            @test x ≈ invFsv
            @test Array(AbstractArray(Tldlt)) ≈ Fs
        end

        # similar
        @test isa(similar(Ts), SymTridiagonal{elty})
        @test isa(similar(Ts, Int), SymTridiagonal{Int})
        @test isa(similar(Ts, Int, (3,2)), Matrix{Int})
    end

    # eigenvalues/eigenvectors of symmetric tridiagonal
    if elty === Float32 || elty === Float64
        DT, VT = @inferred eig(Ts)
        @inferred eig(Ts, 2:4)
        @inferred eig(Ts, 1.0, 2.0)
        D, Vecs = eig(Fs)
        @test DT ≈ D
        @test abs.(VT'Vecs) ≈ eye(elty, n)
        @test eigvecs(Ts) == eigvecs(Fs)
        #call to LAPACK.stein here
        Test.test_approx_eq_modphase(eigvecs(Ts,eigvals(Ts)),eigvecs(Fs))
    elseif elty != Int
        # check that undef is determined accurately even if type inference
        # bails out due to the number of try/catch blocks in this code.
        @test_throws UndefVarError Fs
    end

    # Test det(A::Matrix)
    # In the long run, these tests should step through Strang's
    #  axiomatic definition of determinants.
    # If all axioms are satisfied and all the composition rules work,
    #  all determinants will be correct except for floating point errors.

    # The determinant of the identity matrix should always be 1.
    for i = 1:10
        A = eye(elty, i)
        @test det(A) ≈ one(elty)
    end

    # The determinant of a Householder reflection matrix should always be -1.
    for i = 1:10
        A = eye(elty, 10)
        A[i, i] = -one(elty)
        @test det(A) ≈ -one(elty)
    end

    # The determinant of a rotation matrix should always be 1.
    if elty != Int
        for theta = convert(Vector{elty}, pi ./ [1:4;])
            R = [cos(theta) -sin(theta);
                 sin(theta) cos(theta)]
            @test convert(elty, det(R)) ≈ one(elty)
        end

    # issue #1490
    @test det(ones(elty,3,3)) ≈ zero(elty) atol=3*eps(real(one(elty)))

    @test det(SymTridiagonal(elty[],elty[])) == one(elty)

    #tril/triu
    @test_throws ArgumentError tril!(SymTridiagonal(d,dl),n+1)
    @test_throws ArgumentError tril!(Tridiagonal(dl,d,du),n+1)
    @test tril(SymTridiagonal(d,dl))    == Tridiagonal(dl,d,zeros(dl))
    @test tril(SymTridiagonal(d,dl),1)  == Tridiagonal(dl,d,dl)
    @test tril(SymTridiagonal(d,dl),-1) == Tridiagonal(dl,zeros(d),zeros(dl))
    @test tril(SymTridiagonal(d,dl),-2) == Tridiagonal(zeros(dl),zeros(d),zeros(dl))
    @test tril(Tridiagonal(dl,d,du))    == Tridiagonal(dl,d,zeros(du))
    @test tril(Tridiagonal(dl,d,du),1)  == Tridiagonal(dl,d,du)
    @test tril(Tridiagonal(dl,d,du),-1) == Tridiagonal(dl,zeros(d),zeros(du))
    @test tril(Tridiagonal(dl,d,du),-2) == Tridiagonal(zeros(dl),zeros(d),zeros(du))

    @test_throws ArgumentError triu!(SymTridiagonal(d,dl),n+1)
    @test_throws ArgumentError triu!(Tridiagonal(dl,d,du),n+1)
    @test triu(SymTridiagonal(d,dl))    == Tridiagonal(zeros(dl),d,dl)
    @test triu(SymTridiagonal(d,dl),-1) == Tridiagonal(dl,d,dl)
    @test triu(SymTridiagonal(d,dl),1)  == Tridiagonal(zeros(dl),zeros(d),dl)
    @test triu(SymTridiagonal(d,dl),2)  == Tridiagonal(zeros(dl),zeros(d),zeros(dl))
    @test triu(Tridiagonal(dl,d,du))    == Tridiagonal(zeros(dl),d,du)
    @test triu(Tridiagonal(dl,d,du),-1) == Tridiagonal(dl,d,du)
    @test triu(Tridiagonal(dl,d,du),1)  == Tridiagonal(zeros(dl),zeros(d),du)
    @test triu(Tridiagonal(dl,d,du),2)  == Tridiagonal(zeros(dl),zeros(d),zeros(du))

    @test !istril(SymTridiagonal(d,dl))
    @test !istriu(SymTridiagonal(d,dl))
    @test istriu(Tridiagonal(zeros(dl),d,du))
    @test istril(Tridiagonal(dl,d,zeros(du)))
    end
end

#Test equivalence of eigenvectors/singular vectors taking into account possible phase (sign) differences
function test_approx_eq_vecs{S<:Real,T<:Real}(a::StridedVecOrMat{S}, b::StridedVecOrMat{T}, error=nothing)
    n = size(a, 1)
    @test n==size(b,1) && size(a,2)==size(b,2)
    error===nothing && (error=n^3*(eps(S)+eps(T)))
    for i=1:n
        ev1, ev2 = a[:,i], b[:,i]
        deviation = min(abs(norm(ev1-ev2)),abs(norm(ev1+ev2)))
        if !isnan(deviation)
            @test deviation ≈ 0.0 atol=error
        end
    end
end

let n = 12 #Size of matrix problem to test
    srand(123)
    debug && println("SymTridiagonal (symmetric tridiagonal) matrices")
    for relty in (Float32, Float64), elty in (relty, Complex{relty})
        debug && println("elty is $(elty), relty is $(relty)")
        a = convert(Vector{elty}, randn(n))
        b = convert(Vector{elty}, randn(n-1))
        if elty <: Complex
            a += im*convert(Vector{elty}, randn(n))
            b += im*convert(Vector{elty}, randn(n-1))
        end

        @test_throws DimensionMismatch SymTridiagonal(a, ones(n+1))
        @test_throws ArgumentError SymTridiagonal(rand(n,n))

        A = SymTridiagonal(a, b)
        fA = map(elty <: Complex ? Complex128 : Float64, Array(A))

        debug && println("getindex")
        @test_throws BoundsError A[n+1,1]
        @test_throws BoundsError A[1,n+1]
        @test A[1,n] == convert(elty,0.0)
        @test A[1,1] == a[1]

        debug && println("setindex!")
        @test_throws ArgumentError A[n,1] = 1
        @test_throws ArgumentError A[1,n] = 1
        A[3,3] = A[3,3]
        A[2,3] = A[2,3]
        A[3,2] = A[3,2]
        @test A == fA

        debug && println("Diagonal extraction")
        @test diag(A,1) == b
        @test diag(A,-1) == b
        @test diag(A,0) == a
        @test diag(A,n-1) == zeros(elty,1)
        @test_throws ArgumentError diag(A,n+1)

        debug && println("Idempotent tests")
        for func in (conj, transpose, ctranspose)
            @test func(func(A)) == A
        end

        debug && println("Simple unary functions")
        for func in (det, inv)
            @test func(A) ≈ func(fA) atol=n^2*sqrt(eps(relty))
        end

        debug && println("Rounding to Ints")
        if elty <: Real
            @test round.(Int,A) == round.(Int,fA)
            @test isa(round.(Int,A), SymTridiagonal)
            @test trunc.(Int,A) == trunc.(Int,fA)
            @test isa(trunc.(Int,A), SymTridiagonal)
            @test ceil.(Int,A) == ceil.(Int,fA)
            @test isa(ceil.(Int,A), SymTridiagonal)
            @test floor.(Int,A) == floor.(Int,fA)
            @test isa(floor.(Int,A), SymTridiagonal)
        end

        debug && println("Tridiagonal/SymTridiagonal mixing ops")
        B = convert(Tridiagonal{elty},A)
        @test B == A
        @test B + A == A + B
        @test B - A == A - B

        debug && println("Multiplication with strided vector")
        @test A*ones(n) ≈ Array(A)*ones(n)

        debug && println("Multiplication with strided matrix")
        @test A*ones(n, 2) ≈ Array(A)*ones(n, 2)

        debug && println("Eigensystems")
        if elty <: Real
            zero, infinity = convert(elty, 0), convert(elty, Inf)
            debug && println("This tests eigenvalue and eigenvector computations using stebz! and stein!")
            w, iblock, isplit = LAPACK.stebz!('V', 'B', -infinity, infinity, 0, 0, zero, a, b)
            evecs = LAPACK.stein!(a, b, w)

            (e, v) = eig(SymTridiagonal(a, b))
            @test e ≈ w
            test_approx_eq_vecs(v, evecs)

            debug && println("stein! call using iblock and isplit")
            w, iblock, isplit = LAPACK.stebz!('V', 'B', -infinity, infinity, 0, 0, zero, a, b)
            evecs = LAPACK.stein!(a, b, w, iblock, isplit)
            test_approx_eq_vecs(v, evecs)

            debug && println("stegr! call with index range")
            F = eigfact(SymTridiagonal(a, b),1:2)
            fF = eigfact(Symmetric(Array(SymTridiagonal(a, b))),1:2)
            Test.test_approx_eq_modphase(F[:vectors], fF[:vectors])
            @test F[:values] ≈ fF[:values]

            debug && println("stegr! call with value range")
            F = eigfact(SymTridiagonal(a, b),0.0,1.0)
            fF = eigfact(Symmetric(Array(SymTridiagonal(a, b))),0.0,1.0)
            Test.test_approx_eq_modphase(F[:vectors], fF[:vectors])
            @test F[:values] ≈ fF[:values]
        end

        debug && println("Binary operations")
        a = convert(Vector{elty}, randn(n))
        b = convert(Vector{elty}, randn(n - 1))
        if elty <: Complex
            a += im*convert(Vector{elty}, randn(n))
            b += im*convert(Vector{elty}, randn(n - 1))
        end

        B = SymTridiagonal(a, b)
        fB = map(elty <: Complex ? Complex128 : Float64, Array(B))

        for op in (+, -, *)
            @test Array(op(A, B)) ≈ op(fA, fB)
        end
        α = rand(elty)
        @test Array(α*A) ≈ α*Array(A)
        @test Array(A*α) ≈ Array(A)*α
        @test Array(A/α) ≈ Array(A)/α

        debug && println("A_mul_B!")
        @test_throws DimensionMismatch A_mul_B!(zeros(elty,n,n),B,ones(elty,n+1,n))
        @test_throws DimensionMismatch A_mul_B!(zeros(elty,n+1,n),B,ones(elty,n,n))
        @test_throws DimensionMismatch A_mul_B!(zeros(elty,n,n+1),B,ones(elty,n,n))

    end

    debug && println("Tridiagonal matrices")
    for relty in (Float32, Float64), elty in (relty, Complex{relty})
        debug && println("relty is $(relty), elty is $(elty)")
        a = convert(Vector{elty}, randn(n - 1))
        b = convert(Vector{elty}, randn(n))
        c = convert(Vector{elty}, randn(n - 1))
        if elty <: Complex
            a += im*convert(Vector{elty}, randn(n - 1))
            b += im*convert(Vector{elty}, randn(n))
            c += im*convert(Vector{elty}, randn(n - 1))
        end

        @test_throws ArgumentError Tridiagonal(a,a,a)
        A = Tridiagonal(a, b, c)
        fA = map(elty <: Complex ? Complex128 : Float64, Array(A))

        debug && println("Similar, size, and copy!")
        B = similar(A)
        @test size(B) == size(A)
        copy!(B,A)
        @test B == A
        @test isa(similar(A), Tridiagonal{elty})
        @test isa(similar(A, Int), Tridiagonal{Int})
        @test isa(similar(A, Int, (3,2)), Matrix{Int})
        @test size(A,3) == 1
        @test_throws ArgumentError size(A,0)

        debug && println("Diagonal extraction")
        @test diag(A,-1) == a
        @test diag(A,0) == b
        @test diag(A,1) == c
        @test diag(A,n-1) == zeros(elty,1)
        @test_throws ArgumentError diag(A,n+1)

        debug && println("Simple unary functions")
        for func in (det, inv)
            @test func(A) ≈ func(fA) atol=n^2*sqrt(eps(relty))
        end

        debug && println("Rounding to Ints")
        if elty <: Real
            @test round.(Int,A) == round.(Int,fA)
            @test isa(round.(Int,A), Tridiagonal)
            @test trunc.(Int,A) == trunc.(Int,fA)
            @test isa(trunc.(Int,A), Tridiagonal)
            @test ceil.(Int,A) == ceil.(Int,fA)
            @test isa(ceil.(Int,A), Tridiagonal)
            @test floor.(Int,A) == floor.(Int,fA)
            @test isa(floor.(Int,A), Tridiagonal)
        end

        debug && println("Binary operations")
        a = convert(Vector{elty}, randn(n - 1))
        b = convert(Vector{elty}, randn(n))
        c = convert(Vector{elty}, randn(n - 1))
        if elty <: Complex
            a += im*convert(Vector{elty}, randn(n - 1))
            b += im*convert(Vector{elty}, randn(n))
            c += im*convert(Vector{elty}, randn(n - 1))
        end

        debug && println("Multiplication with strided vector")
        @test A*ones(n) ≈ Array(A)*ones(n)

        debug && println("Multiplication with strided matrix")
        @test A*ones(n, 2) ≈ Array(A)*ones(n, 2)

        B = Tridiagonal(a, b, c)
        fB = map(elty <: Complex ? Complex128 : Float64, Array(B))

        for op in (+, -, *)
            @test Array(op(A, B)) ≈ op(fA, fB)
        end
        α = rand(elty)
        @test Array(α*A) ≈ α*Array(A)
        @test Array(A*α) ≈ Array(A)*α
        @test Array(A/α) ≈ Array(A)/α

        @test_throws ArgumentError convert(SymTridiagonal{elty},A)

        debug && println("A_mul_B!")
        @test_throws DimensionMismatch Base.LinAlg.A_mul_B!(zeros(fA),A,ones(elty,n,n+1))
        @test_throws DimensionMismatch Base.LinAlg.A_mul_B!(zeros(fA),A,ones(elty,n+1,n))

        debug && println("getindex")
        @test_throws BoundsError A[n+1,1]
        @test_throws BoundsError A[1,n+1]

        debug && println("setindex!")
        @test_throws ArgumentError A[n,1] = 1
        @test_throws ArgumentError A[1,n] = 1
        A[3,3] = A[3,3]
        A[2,3] = A[2,3]
        A[3,2] = A[3,2]
        @test A == fA
    end
end

# Issue 12068
SymTridiagonal([1, 2], [0])^3 == [1 0; 0 8]

#test convert for SymTridiagonal
@test convert(SymTridiagonal{Float64},SymTridiagonal(ones(Float32,5),ones(Float32,4))) == SymTridiagonal(ones(Float64,5),ones(Float64,4))
@test convert(AbstractMatrix{Float64},SymTridiagonal(ones(Float32,5),ones(Float32,4))) == SymTridiagonal(ones(Float64,5),ones(Float64,4))

# Test constructors from matrix
@test SymTridiagonal([1 2 3; 2 5 6; 0 6 9]) == [1 2 0; 2 5 6; 0 6 9]
@test Tridiagonal([1 2 3; 4 5 6; 7 8 9]) == [1 2 0; 4 5 6; 0 8 9]

# Test constructors with range and other abstract vectors
@test SymTridiagonal(1:3, 1:2) == [1 1 0; 1 2 2; 0 2 3]
@test Tridiagonal(4:5, 1:3, 1:2) == [1 1 0; 4 2 2; 0 5 3]
end


# include("../../src/SFFM.jl")
using LinearAlgebra, Plots, SFFM

## define the model(s)
include("exampleModelDef.jl")

## section 4.3: the marginal stationary distribution of X
## mesh
Δtemp = 0.4
Nodes = collect(approxBounds[1, 1]:Δtemp:approxBounds[1, 2])
nBases = 2
Basis = "lagrange"
mesh = SFFM.DGMesh(
    approxModel,
    Nodes,
    nBases,
    Basis = Basis,
)
let q = SFFM.plot(approxModel)
    ## analytic version for comparison
    # construction
    Ψₓ = SFFM.PsiFunX(approxModel)
    ξₓ = SFFM.MakeXiX(approxModel, Ψₓ)
    pₓ, πₓ, Πₓ, Kₓ = SFFM.StationaryDistributionX( approxModel, Ψₓ, ξₓ)

    # evaluate the distribution
    analyticX = SFFM.SFFMDensity(
        [pₓ[:];0;0],
        πₓ(SFFM.CellNodes(mesh)),
        SFFM.CellNodes(mesh),
    )

    # plot it
    q = SFFM.plot!(q,
        approxModel,
        mesh,
        analyticX,
        color = :red,
        label = "Analytic",
        marker = :x,
        seriestype = :line,
        jitter = 0.5,
    )

    # now DG it
    c = 0
    colours = [:green;:blue]
    for nBases in 1:2
        c = c+1
        mesh = SFFM.DGMesh(
            approxModel,
            Nodes,
            nBases,
            Basis=Basis,
        )

        # compute the marginal via DG
        All = SFFM.MakeAll( approxModel, mesh, approxType = "projection")
        Ψ = SFFM.PsiFun(All.D)

        # the distribution of X when Y first returns to 0
        ξ = SFFM.MakeXi(All.B, Ψ)

        marginalX, p, K = SFFM.MakeLimitDistMatrices(
            All.B,
            All.D,
            All.R.RDict,
            Ψ,
            ξ,
            mesh,
            approxModel,
        )
        # convert marginalX to a distribution for plotting
        Dist = SFFM.Coeffs2Dist(
            approxModel,
            mesh,
            marginalX,
            SFFM.SFFMDensity,
        )
        # plot it
        q = SFFM.plot!(
            q,
            approxModel,
            mesh,
            Dist,
            color = colours[c],
            label = "DG: N_k = "*string(nBases),
            seriestype = :line,
            jitter = 0.5,
        )

    end
    q = plot!(windowsize=(600,600))
    titles = ["Phase 11" "Phase 10" "Phase 01" "Phase 00"]
    for sp in 1:4
        q = plot!(
            subplot = sp,
            xlims = (-0.5,5),
            title = titles[sp],
            xlabel = "x",
            ylabel = "Density / Probability",
        )
    end
    display(q)
    # savefig(pwd()*"/examples/paperNumerics/dump/marginalStationaryDistX.png")
end


include("../DegeneracyHunter.jl")
using JuMP
#using Ipopt
using AmplNLWriter

#s = IpoptSolver(bound_relax_factor=0.0)
s = AmplNLSolver("ipopt")

m = Model(solver=s)

@variable(m, -10 <= a[1:5] <= 10, start=1.0)
@variable(m, -100 <= z <= 100, start=5.0)

# Fix this variable by setting upper bound = lower bound
@variable(m, 0 <= y <= 0)

@constraint(m, a[1] + a[2] - a[3] >= 10)
@NLconstraint(m, a[1]*a[4] - a[2]*y + a[2] >= 0)
@NLconstraint(m, a[3]*a[2] + a[1]*a[4] - a[5] == 0)
@NLconstraint(m, z == sum{a[i], i=1:5} - 0.001*sum{a[i]^2, i=1:5})
@constraint(m, a[3] + a[4] == 10*y)

@objective(m, Min, z)

solve(m)

DegeneracyHunter.printInfeasibleEquations(m, 1E-10)

println("Manually check constrain violations")

r = zeros(5)
a_ = zeros(5)

for i = 1:5
	a_[i] = getvalue(a[i])
end

z_ = getvalue(z)
y_ = getvalue(y)

r[1] = max(10 - a_[1] - a_[2] + a_[3], 0)
r[2] = max(-a_[2] + a_[2]*y_ - a_[1]*a_[4], 0)
r[3] = a_[3]*a_[2] + a_[1]*a_[4] - a_[5]
r[4] = z_ - sum(a_) + 0.001*sum(a_.^2)
r[5] = a_[3] + a_[4] - 10*y_

for i = 1:5
	println("r[",i,"] = ",r[i])
end


using Base.Test, Whittaker

@testset "whittaker2" begin
  @testset "out-of-place" begin
    tol = 0.0001

    y = [1.0, 1.2, 0, 1.6, 1.8]
    w = [1.0, 1.0, 0, 1.0, 1.0]
    lambda = 0.0001
    z = whittaker2(y, w, lambda)
    @test_approx_eq_eps z [0.999997, 1.20001, 1.4305, 1.64575, 1.80004] tol

    y = [1., 2, 0, 4, 5]
    w = [1, 1, 0, 1, 1]
    lambda = 0.0001
    z = whittaker2(y, w, lambda)
    @test_approx_eq_eps z [0.999992, 2.00002, 3.07626, 4.11438, 5.00005] tol
  end

  @testset "in-place" begin
    tol = 0.0001

    y = [1.0, 1.2, 0, 1.6, 1.8]
    w = [1.0, 1.0, 0, 1.0, 1.0]
    lambda = 0.0001
    whittaker2!(y, w, lambda)
    @test_approx_eq_eps y [0.999997, 1.20001, 1.4305, 1.64575, 1.80004] tol

    y = [1., 2, 0, 4, 5]
    w = [1, 1, 0, 1, 1]
    lambda = 0.0001
    whittaker2!(y, w, lambda)
    @test_approx_eq_eps y [0.999992, 2.00002, 3.07626, 4.11438, 5.00005] tol
  end
end


infstep(x0, d) = convert(eltype(d), Inf)

#=
A generic way to create a mutable N×N matrix from a length-N vector.
Behaves like `similar(vec)` but returns `MMatrix` if the
argument is a `StaticVector` (`similar(staticvec, (n,n))` creates a
`Base.Matrix` which is less performant, and I don't want an explicit
dependency on `StaticArrays.jl` - V.P.)
=#
"""
    sqmatr(vec::AbstractVector, [element_type = eltype(vec)])

Create an uninitialized mutable N×N matrix with the given element type, given a
vector of length N.
"""
function sqmatr(vec::AbstractVector{T}, F::DataType = T) where {T}
    sz = axes(vec)
    return similar(vec, F, (sz..., sz...))
end

"""
Return the maximum absolute values of diagonal and off-diagonal elements of `A`.
"""
function γξ(A::AbstractMatrix)
    n = LinearAlgebra.checksquare(A)
    γ = ξ = zero(eltype(A))
    for j in 1:n
        for i in 1:j-1
            ξ = max(ξ, abs(A[i,j]))
        end
        γ = max(γ, abs(A[j,j]))
    end
    return γ, ξ
end

"""
    mcholesky!(A::AbstractMatrix)

Perform an in-place modified Cholesky decomposition on matrix `A`
(Gill, Murray, Wright, Practical optimization (1981), p.111)
"""
function mcholesky!(A::AbstractMatrix{T}; δ = convert(T, 1e-3)) where T
    n = LinearAlgebra.checksquare(A)
    γ, ξ = γξ(A)
    ν = max(1, sqrt(n^2 - 1))
    β² = max(γ, ξ / ν, eps(T))
    θ = zero(T)
    u = UpperTriangular(A)
    c = u
    @inbounds for j in 1:n
        c_jj = A[j,j]
        θ = zero(c_jj)
        for k in 1:j-1
            u[k,j] /= u[k,k]
        end
        for i in j+1:n
            c_ij = c[j,i]
            for k in 1:j-1
                c_ij -= u[k,j] * c[k,i] / u[k,k]
            end
            θ = max(abs(c_ij), θ)
            c[j,i] = c_ij
        end
        d_j = max(abs(c_jj), θ^2 / β², δ)
        u[j,j] = sqrt(d_j)
        for i in j+1:n
            c[i,i] -= (c[j,i] / u[j,j])^2
        end
    end
    return Cholesky(A, 'U', 0)
end


module SingleFieldFESpacesTests

using Test
using Gridap.Arrays
using Gridap.TensorValues
using Gridap.ReferenceFEs
using Gridap.Geometry
using Gridap.FESpaces

domain =(0,1,0,1,0,1)
partition = (3,3,3)
model = CartesianDiscreteModel(domain,partition)

order = 2
grid_topology = get_grid_topology(model)
polytopes = get_polytopes(grid_topology)
reffes = [LagrangianRefFE(Float64,p,order) for p in polytopes]

dirichlet_tags = ["tag_24","tag_25"]
V0 = GradConformingFESpace(reffes,model,dirichlet_tags)
matvecdata = ([],[],[])
matdata = ([],[],[])
vecdata = ([],[])
test_single_field_fe_space(V0,matvecdata,matdata,vecdata)

cell_map = get_cell_map(V0)

f(x) = sin(4*pi*(x[1]-x[2]^2))+1

fh = interpolate_everywhere(f, V0)

fh = interpolate_dirichlet(f, V0)

dirichlet_values = compute_dirichlet_values_for_tags(V0,[1,2])

r=[2,2,1,1,2,2,1,1,2,2,2,2,1,2,1,1,1,1,2,2,2,2,1,2,1,1,1,1,2,2,1,1,2,2,1,2,1,1,2,2,1,2,1,1,2,2,1,2,1,1]

@test dirichlet_values == r

dirichlet_values = compute_dirichlet_values_for_tags(V0,2)
@test dirichlet_values == fill(2,num_dirichlet_dofs(V0))

free_values = zero_free_values(V0)
uh = FEFunction(V0,free_values,dirichlet_values)

#using Gridap.Visualization
#
#trian = get_triangulation(model)
#
#writevtk(model,"model")
#
#writevtk(trian,"trian",cellfields=["fh" => fh, "uh" => uh])

#
#
#strian = SkeletonTriangulation(model)
#
#fh_Γ = restrict(get_array(fh),strian)
#
#writevtk(strian,"strian",nsubcells=10,cellfields=["jump_fh" => jump(fh_Γ)])

end # module


abstract type AbstractGrid{dim, T} end

"""
```
struct RectilinearGrid{dim, T, N, M} <: AbstractGrid{dim, T}
    grid::JuAFEM.Grid{dim, N, T, M}
    nels::NTuple{dim, Int}
    sizes::NTuple{dim, T}
    corners::NTuple{2, JuAFEM.Vec{dim, T}}
end
```

`dim`: dimension of the problem

`T`: number type for computations and coordinates

`N`: number of nodes in a cell of the grid

`M`: number of faces in a cell of the grid

`grid`: a JuAFEM.Grid struct

`nels`: number of elements in every dimension

`sizes`: dimensions of each rectilinear cell

`corners`: 2 corner points of the rectilinear grid

API:

`function RectilinearGrid(nels::NTuple{dim,Int}, sizes::NTuple{dim,T}) where {dim,T}`

Example:

`rectgrid = RectilinearGrid((60,20), (1.0,1.0))`
"""
struct RectilinearGrid{dim, T, N, M} <: AbstractGrid{dim, T}
    grid::JuAFEM.Grid{dim, N, T, M}
    nels::NTuple{dim, Int}
    sizes::NTuple{dim, T}
    corners::NTuple{2, JuAFEM.Vec{dim, T}}
end
nnodespercell(::RectilinearGrid{dim,T,N,M}) where {dim, T, N, M} = N
nfacespercell(::RectilinearGrid{dim,T,N,M}) where {dim, T, N, M} = M

left(rectgrid::RectilinearGrid, x) = x[1] ≈ rectgrid.corners[1][1]
right(rectgrid::RectilinearGrid, x) = x[1] ≈ rectgrid.corners[2][1]
bottom(rectgrid::RectilinearGrid, x) = x[2] ≈ rectgrid.corners[1][2]
top(rectgrid::RectilinearGrid, x) = x[2] ≈ rectgrid.corners[2][2]
back(rectgrid::RectilinearGrid, x) = x[3] ≈ rectgrid.corners[1][3]
front(rectgrid::RectilinearGrid, x) = x[3] ≈ rectgrid.corners[2][3]
middlex(rectgrid::RectilinearGrid, x) = x[1] ≈ (rectgrid.corners[1][1] + rectgrid.corners[2][1]) / 2
middley(rectgrid::RectilinearGrid, x) = x[2] ≈ (rectgrid.corners[1][2] + rectgrid.corners[2][2]) / 2
middlez(rectgrid::RectilinearGrid, x) = x[3] ≈ (rectgrid.corners[1][3] + rectgrid.corners[2][3]) / 2

nnodes(cell::Type{JuAFEM.Cell{dim,N,M}}) where {dim, N, M} = N
nnodes(cell::JuAFEM.Cell) = nnodes(typeof(cell))

function RectilinearGrid(nels::NTuple{dim,Int}, sizes::NTuple{dim,T}) where {dim,T}
    geoshape = dim === 2 ? Quadrilateral : Hexahedron
    corner1 = JuAFEM.Vec{dim}(fill(T(0), dim))
    corner2 = JuAFEM.Vec{dim}((nels .* sizes))
    grid = generate_grid(geoshape, nels, corner1, corner2);

    N = nnodes(geoshape)
    M = JuAFEM.nfaces(geoshape)
    ncells = prod(nels)
    return RectilinearGrid{dim, T, N, M}(grid, nels, sizes, (corner1, corner2))
end


mutable struct BackwardFlatInterpolation <: Interpolation
  x_vals::Vector{Float64}
  y_vals::Vector{Float64}
  primitive::Vector{Float64}
end

function BackwardFlatInterpolation()
  x_vals = Vector{Float64}()
  y_vals = Vector{Float64}()
  primitive = Vector{Float64}()

  return BackwardFlatInterpolation(x_vals, y_vals, primitive)
end

# Backward Flat initialize #
function initialize!(interp::BackwardFlatInterpolation, x_vals::Vector{Float64}, y_vals::Vector{Float64})
  interp.x_vals = x_vals
  interp.y_vals = y_vals
  interp.primitive = zeros(length(y_vals))

  return interp
end

# BackwardFlatInterpolation update
function update!(interp::BackwardFlatInterpolation, idx::Int)
  interp.primitive[1] = 0.0
  @simd for i = 2:idx
    @inbounds dx = interp.x_vals[i] - interp.x_vals[i - 1]
    @inbounds interp.primitive[i] = interp.primitive[i - 1] + dx * interp.y_vals[i]
  end

  return interp
end

update!(interp::BackwardFlatInterpolation) = update!(interp, length(interp.x_vals))

function value(interp::BackwardFlatInterpolation, x::Float64)
  if x <= interp.x_vals[1]
    return interp.y_vals[1]
  end

  i = locate(interp, x)
  if x == interp.x_vals[i]
    return interp.y_vals[i]
  else
    return i >= length(interp.y_vals) ? 0.0 : interp.y_vals[i + 1]
  end
end

function get_primitive(interp::BackwardFlatInterpolation, x::Float64, ::Bool)
  i = locate(interp, x)
  # if i > 4
  #   println(x)
  #   println(interp.x_vals)
  #   println(interp.y_vals)
  # end
  dx = x - interp.x_vals[i]
  return interp.primitive[i] + dx * (i >= length(interp.y_vals) ? 0.0 : interp.y_vals[i + 1])
end


#=
Code for simulating, visualizing and manipulating additive and
multiplicative functionals.

@authors: Shunsuke Hori
=#
using QuantEcon
using PyPlot
using Distributions

"""
This type transforms an additive (multipilcative)
functional into a QuantEcon linear state space system.
"""
struct AMF_LSS_VAR{TF<:AbstractFloat, TI<:Integer}
    A::Array{TF, 2}
    B::Array{TF, 2}
    D::Array{TF, 2}
    F::Array{TF, 2}
    nu::Array{TF, 2}
    nx::TI
    nk::TI
    nm::TI
    lss::LSS
end

function AMF_LSS_VAR(A::Array, B::Array,
                D::Union{RowVector, Array},
                F::Union{Void, Array}=nothing;
                nu::Union{Void, Array}=nothing)

    if typeof(B) <: Vector
        B = reshape(B, length(B), 1)
    end
    # Unpack required elements
    nx, nk = size(B)

    # checking the dimension of D (extended from the scalar case)
    if ndims(D) > 1
        nm = size(D, 1)
        if typeof(D) <: RowVector
            D = convert(Matrix, D)
        end
    else
        nm = 1
        D = reshape(D, 1, length(D))
    end

    # Set F
    if F==nothing
        F = zeros(nk, 1)
    elseif ndims(F) == 1
        F = reshape(F, length(F), 1)
    end

    # Set nu
    if nu==nothing
        nu = zeros(nm, 1)
    elseif ndims(nu) == 1
        nu = reshape(nu, length(nu), 1)
    else
        throw(ArgumentError("nu must be column vector!"))
    end

    if size(nu, 1) != size(D, 1)
        error("The size of nu is inconsistent with D!")
    end

    # Construct BIG state space representation
    lss = construct_ss(A, B, D, F, nu, nx, nk, nm)

    return AMF_LSS_VAR(A, B, D, F, nu, nx, nk, nm, lss)
end

AMF_LSS_VAR(A::Array, B::Array, D::Union{RowVector, Array}) =
    AMF_LSS_VAR(A, B, D, nothing, nu=nothing)
AMF_LSS_VAR(A::Array, B::Array, D::Union{RowVector, Array}, F::Real, nu::Real) =
    AMF_LSS_VAR(A, B, D, [F], nu=[nu])

"""
This creates the state space representation that can be passed
into the quantecon LSS class.
"""
function construct_ss(A::Array, B::Array, D::Union{RowVector, Array}, F::Array,
                      nu, nx::TI, nk::TI, nm::TI) where TI <: Integer

    H, g = additive_decomp(A, B, D, F, nx)

    # Auxiliary blocks with 0's and 1's to fill out the lss matrices
    nx0c = zeros(nx, 1)
    nx0r = zeros(1, nx)
    nx1 = ones(1, nx)
    nk0 = zeros(1, nk)
    ny0c = zeros(nm, 1)
    ny0r = zeros(1, nm)
    ny1m = eye(nm)
    ny0m = zeros(nm, nm)
    nyx0m = zeros(D)

    # Build A matrix for LSS
    # Order of states is: [1, t, xt, yt, mt]
    A1 = hcat(1, 0, nx0r, ny0r, ny0r)          # Transition for 1
    A2 = hcat(1, 1, nx0r, ny0r, ny0r)          # Transition for t
    A3 = hcat(nx0c, nx0c, A, nyx0m', nyx0m')   # Transition for x_{t+1}
    A4 = hcat(nu, ny0c, D, ny1m, ny0m)         # Transition for y_{t+1}
    A5 = hcat(ny0c, ny0c, nyx0m, ny0m, ny1m)   # Transition for m_{t+1}
    Abar = vcat(A1, A2, A3, A4, A5)

    # Build B matrix for LSS
    Bbar = vcat(nk0, nk0, B, F, H)

    # Build G matrix for LSS
    # Order of observation is: [xt, yt, mt, st, tt]
    G1 = hcat(nx0c, nx0c, eye(nx), nyx0m', nyx0m')    # Selector for x_{t}
    G2 = hcat(ny0c, ny0c, nyx0m, ny1m, ny0m)          # Selector for y_{t}
    G3 = hcat(ny0c, ny0c, nyx0m, ny0m, ny1m)          # Selector for martingale
    G4 = hcat(ny0c, ny0c, -g, ny0m, ny0m)             # Selector for stationary
    G5 = hcat(ny0c, nu, nyx0m, ny0m, ny0m)            # Selector for trend
    Gbar = vcat(G1, G2, G3, G4, G5)

    # Build LSS type
    x0 = hcat(1, 0, nx0r, ny0r, ny0r)
    S0 = zeros(length(x0), length(x0))
    lss = LSS(Abar, Bbar, Gbar, x0, S0)

    return lss
end

"""
Return values for the martingale decomposition
    - nu        : unconditional mean difference in Y
    - H         : coefficient for the (linear) martingale component (kappa_a)
    - g         : coefficient for the stationary component g(x)
    - Y_0       : it should be the function of X_0 (for now set it to 0.0)
"""
function additive_decomp(A::Array, B::Array, D::Array, F::Union{Array, Real},
                          nx::Integer)
    I = eye(nx)
    A_res = \(I-A, I)
    g = D * A_res
    H = F + D * A_res * B

    return H, g
end

"""
Return values for the multiplicative decomposition (Example 5.4.4.)
    - nu_tilde  : eigenvalue
    - H         : vector for the Jensen term
"""
function multiplicative_decomp(A::Array, B::Array, D::Array, F::Union{Array, Real},
                                nu::Union{Array, Real}, nx::Integer)
    H, g = additive_decomp(A, B, D, F, nx)
    nu_tilde = nu + 0.5*diag(H*H')

    return H, g, nu_tilde
end

function loglikelihood_path(amf::AMF_LSS_VAR, x::Array, y::Array)
    A, B, D, F = amf.A, amf.B, amf.D, amf.F
    k, T = size(y)
    FF = F*F'
    FFinv = inv(FF)
    temp = y[:, 2:end]-y[:, 1:end-1] - D*x[:, 1:end-1]
    obs =  temp .* FFinv .* temp
    obssum = cumsum(obs)
    scalar = (log(det(FF)) + k*log(2*pi))*collect(1:T)

    return -0.5*(obssum + scalar)
end

function loglikelihood(amf::AMF_LSS_VAR, x::Array, y::Array)
    llh = loglikelihood_path(amf, x, y)

    return llh[end]
end

"""
Plots for the additive decomposition
"""
function plot_additive(amf::AMF_LSS_VAR, T::Integer;
                       npaths::Integer=25, show_trend::Bool=true)

    # Pull out right sizes so we know how to increment
    nx, nk, nm = amf.nx, amf.nk, amf.nm

    # Allocate space (nm is the number of additive functionals - we want npaths for each)
    mpath = Array{Real}(nm*npaths, T)
    mbounds = Array{Real}(nm*2, T)
    spath = Array{Real}(nm*npaths, T)
    sbounds = Array{Real}(nm*2, T)
    tpath = Array{Real}(nm*npaths, T)
    ypath = Array{Real}(nm*npaths, T)

    # Simulate for as long as we wanted
    moment_generator = moment_sequence(amf.lss)
    state = start(moment_generator)
    # Pull out population moments
    for t in 1:T
        tmoms, state = next(moment_generator, state)
        ymeans = tmoms[2]
        yvar = tmoms[4]

        # Lower and upper bounds - for each additive functional
        for ii in 1:nm
            li, ui = (ii-1)*2+1, ii*2
            if sqrt(yvar[nx+nm+ii, nx+nm+ii]) != 0.0
                madd_dist = Normal(ymeans[nx+nm+ii], sqrt(yvar[nx+nm+ii, nx+nm+ii]))
                mbounds[li, t] = quantile(madd_dist, 0.01)
                mbounds[ui, t] = quantile(madd_dist, 0.99)
            elseif sqrt(yvar[nx+nm+ii, nx+nm+ii]) == 0.0
                mbounds[li, t] = ymeans[nx+nm+ii]
                mbounds[ui, t] = ymeans[nx+nm+ii]
            else
                error("standard error is negative")
            end

            if sqrt(yvar[nx+2*nm+ii, nx+2*nm+ii]) != 0.0
                sadd_dist = Normal(ymeans[nx+2*nm+ii], sqrt(yvar[nx+2*nm+ii, nx+2*nm+ii]))
                sbounds[li, t] = quantile(sadd_dist, 0.01)
                sbounds[ui, t] = quantile(sadd_dist, 0.99)
            elseif sqrt(yvar[nx+2*nm+ii, nx+2*nm+ii]) == 0.0
                sbounds[li, t] = ymeans[nx+2*nm+ii]
                sbounds[ui, t] = ymeans[nx+2*nm+ii]
            else
                error("standard error is negative")
            end

        end
    end

    # Pull out paths
    for n in 1:npaths
        x, y = simulate(amf.lss,T)
        for ii in 0:nm-1
            ypath[npaths*ii+n, :] = y[nx+ii+1, :]
            mpath[npaths*ii+n, :] = y[nx+nm + ii+1, :]
            spath[npaths*ii+n, :] = y[nx+2*nm + ii+1, :]
            tpath[npaths*ii+n, :] = y[nx+3*nm + ii+1, :]
        end
    end

    add_figs = Array{Any}(nm)

    for ii in 0:nm-1
        li, ui = npaths*(ii), npaths*(ii+1)
        LI, UI = 2*(ii), 2*(ii+1)
        add_figs[ii+1] =
            plot_given_paths(T, ypath[li+1:ui, :], mpath[li+1:ui, :], spath[li+1:ui, :],
                             tpath[li+1:ui, :], mbounds[LI+1:UI, :], sbounds[LI+1:UI, :],
                             show_trend=show_trend)

        add_figs[ii+1][:suptitle]( L"Additive decomposition of $y_{$(ii+1)}$", fontsize=14 )
    end

    return add_figs
end

"""
Plots for the multiplicative decomposition
"""
function plot_multiplicative(amf::AMF_LSS_VAR, T::Integer,
                            npaths::Integer=25, show_trend::Bool=true)
    # Pull out right sizes so we know how to increment
    nx, nk, nm = amf.nx, amf.nk, amf.nm
    # Matrices for the multiplicative decomposition
    H, g, nu_tilde = multiplicative_decomp(A, B, D, F, nu, nx)

    # Allocate space (nm is the number of functionals - we want npaths for each)
    mpath_mult = Array{Real}(nm*npaths, T)
    mbounds_mult = Array{Real}(nm*2, T)
    spath_mult = Array{Real}(nm*npaths, T)
    sbounds_mult = Array{Real}(nm*2, T)
    tpath_mult = Array{Real}(nm*npaths, T)
    ypath_mult = Array{Real}(nm*npaths, T)

    # Simulate for as long as we wanted
    moment_generator = moment_sequence(amf.lss)
    state = start(moment_generator)
    # Pull out population moments
    for t in 1:T
        tmoms, state = next(moment_generator, state)
        ymeans = tmoms[2]
        yvar = tmoms[4]

        # Lower and upper bounds - for each multiplicative functional
        for ii in 1:nm
            li, ui = (ii-1)*2+1, ii*2
            if yvar[nx+nm+ii, nx+nm+ii] != 0.0
                Mdist = LogNormal(ymeans[nx+nm+ii]- t*0.5*diag(H * H')[ii],
                                sqrt(yvar[nx+nm+ii, nx+nm+ii]))
                mbounds_mult[li, t] = quantile(Mdist, 0.01)
                mbounds_mult[ui, t] = quantile(Mdist, 0.99)
            elseif yvar[nx+nm+ii, nx+nm+ii] == 0.0
                mbounds_mult[li, t] = exp.(ymeans[nx+nm+ii]- t*0.5*diag(H * H')[ii])
                mbounds_mult[ui, t] = exp.(ymeans[nx+nm+ii]- t*0.5*diag(H * H')[ii])
            else
                error("standard error is negative")
            end
            if yvar[nx+2*nm+ii, nx+2*nm+ii] != 0.0
                Sdist = LogNormal(-ymeans[nx+2*nm+ii],
                                sqrt(yvar[nx+2*nm+ii, nx+2*nm+ii]))
                sbounds_mult[li, t] = quantile(Sdist, 0.01)
                sbounds_mult[ui, t] = quantile(Sdist, 0.99)
            elseif yvar[nx+2*nm+ii, nx+2*nm+ii] == 0.0
                sbounds_mult[li, t] = exp.(-ymeans[nx+2*nm+ii])
                sbounds_mult[ui, t] = exp.(-ymeans[nx+2*nm+ii])
            else
                error("standard error is negative")
            end
        end
    end

    # Pull out paths
    for n in 1:npaths
        x, y = simulate(amf.lss,T)
        for ii in 0:nm-1
            ypath_mult[npaths*ii+n, :] = exp.(y[nx+ii+1, :])
            mpath_mult[npaths*ii+n, :] =
                exp.(y[nx+nm + ii+1, :] - collect(1:T)*0.5*diag(H * H')[ii+1])
            spath_mult[npaths*ii+n, :] = 1./exp.(-y[nx+2*nm + ii+1, :])
            tpath_mult[npaths*ii+n, :] =
                exp.(y[nx+3*nm + ii+1, :] + collect(1:T)*0.5*diag(H * H')[ii+1])
        end
    end

    mult_figs = Array{Any}(nm)

    for ii in 0:nm-1
        li, ui = npaths*(ii), npaths*(ii+1)
        LI, UI = 2*(ii), 2*(ii+1)
        mult_figs[ii+1] =
            plot_given_paths(T, ypath_mult[li+1:ui, :], mpath_mult[li+1:ui, :],
                             spath_mult[li+1:ui, :], tpath_mult[li+1:ui, :],
                             mbounds_mult[LI+1:UI, :], sbounds_mult[LI+1:UI, :],
                             horline = 1.0, show_trend=show_trend)
        mult_figs[ii+1][:suptitle]( L"Multiplicative decomposition of $y_{$(ii+1)}$",
                                    fontsize=14)
    end

    return mult_figs
end

function plot_martingales(amf::AMF_LSS_VAR, T::Integer, npaths::Integer=25)

    # Pull out right sizes so we know how to increment
    nx, nk, nm = amf.nx, amf.nk, amf.nm
    # Matrices for the multiplicative decomposition
    H, g, nu_tilde = multiplicative_decomp(amf.A, amf.B, amf.D, amf.F, amf.nu, amf.nx)

    # Allocate space (nm is the number of functionals - we want npaths for each)
    mpath_mult = Array{Real}(nm*npaths, T)
    mbounds_mult = Array{Real}(nm*2, T)

    # Simulate for as long as we wanted
    moment_generator = moment_sequence(amf.lss)
    state = start(moment_generator)
    # Pull out population moments
    for t in 1:T
        tmoms, state = next(moment_generator, state)
        ymeans = tmoms[2]
        yvar = tmoms[4]

        # Lower and upper bounds - for each functional
        for ii in 1:nm
            li, ui = (ii-1)*2+1, ii*2
            if yvar[nx+nm+ii, nx+nm+ii] != 0.0
                Mdist = LogNormal(ymeans[nx+nm+ii]-t*(.5)*diag(H*H')[ii],
                            sqrt(yvar[nx+nm+ii, nx+nm+ii]))
                mbounds_mult[li, t] = quantile(Mdist, 0.01)
                mbounds_mult[ui, t] = quantile(Mdist, 0.99)
            elseif yvar[nx+nm+ii, nx+nm+ii] == 0.0
                mbounds_mult[li, t] = ymeans[nx+nm+ii]-t*(.5)*diag(H*H')[ii]
                mbounds_mult[ui, t] = ymeans[nx+nm+ii]-t*(.5)*diag(H*H')[ii]
            else
                error("standard error is negative")
            end
        end
    end

    # Pull out paths
    for n in 1:npaths
        x, y = simulate(amf.lss, T)
        for ii in 0:nm-1
            mpath_mult[npaths*ii+n, :] =
                exp.(y[nx+nm + ii+1, :] - (1:T)*0.5*diag(H*H')[ii+1])
        end
    end

    mart_figs = Array{Any}(nm)

    for ii in 0:nm-1
        li, ui = npaths*(ii), npaths*(ii+1)
        LI, UI = 2*(ii), 2*(ii+1)
        mart_figs[ii+1] = plot_martingale_paths(T, mpath_mult[li+1:ui, :],
                                                    mbounds_mult[LI+1:UI, :], horline=1)
        mart_figs[ii+1][:suptitle](L"Martingale components for many paths of $y_{ii+1}$",
                                   fontsize=14)
    end

    return mart_figs
end

function plot_given_paths(T::Integer,
                        ypath::Array, mpath::Array, spath::Array, tpath::Array,
                        mbounds::Array, sbounds::Array; horline::Real=0.0,
                        show_trend::Bool = true)

    # Allocate space
    trange = 1:T

    # Create figure
    fig, ax = subplots(2, 2, sharey=true, figsize=(15, 8))

    # Plot all paths together
    ax[1, 1][:plot](trange, ypath[1, :], label=L"$y_t$", color="k")
    ax[1, 1][:plot](trange, mpath[1, :], label=L"$m_t$", color="m")
    ax[1, 1][:plot](trange, spath[1, :], label=L"$s_t$", color="g")
    if show_trend == true
        ax[1, 1][:plot](trange, tpath[1, :], label=L"t_t", color="r")
    end
    ax[1, 1][:axhline](horline, color="k", linestyle = "-.")
    ax[1, 1][:set_title]("One Path of All Variables")
    ax[1, 1][:legend](loc="top left")

    # Plot Martingale Component
    ax[1, 2][:plot](trange, mpath[1, :], "m")
    ax[1, 2][:plot](trange, mpath', alpha=0.45, color="m")
    ub = mbounds[2, :]
    lb = mbounds[1, :]
    ax[1, 2][:fill_between](trange, lb, ub, alpha=0.25, color="m")
    ax[1, 2][:set_title]("Martingale Components for Many Paths")
    ax[1, 2][:axhline](horline, color="k", linestyle = "-.")

    # Plot Stationary Component
    ax[2, 1][:plot](spath[1, :], color="g")
    ax[2, 1][:plot](spath', alpha=0.25, color="g")
    ub = sbounds[2, :]
    lb = sbounds[1, :]
    ax[2, 1][:fill_between](trange, lb, ub, alpha=0.25, color="g")
    ax[2, 1][:axhline](horline, color="k", linestyle = "-.")
    ax[2, 1][:set_title]("Stationary Components for Many Paths")

    # Plot Trend Component
    if show_trend == true
        ax[2, 2][:plot](tpath', color="r")
    end
    ax[2, 2][:set_title]("Trend Components for Many Paths")
    ax[2, 2][:axhline](horline, color="k", linestyle = "-.")

    return fig
end

function plot_martingale_paths(T::Integer,
                        mpath::Array, mbounds::Array;
                        horline::Real=1,
                        show_trend::Bool = false)
    # Allocate space
    trange = 1:T

    # Create figure
    fig, ax = subplots(1, 1, figsize=(10, 6))

    # Plot Martingale Component
    ub = mbounds[2, :]
    lb = mbounds[1, :]
    ax[:fill_between](trange, lb, ub, color="#ffccff")
    ax[:axhline](horline, color="k", linestyle = "-.")
    ax[:plot](trange, mpath', linewidth=0.25, color="#4c4c4c")

    return fig
end


"""
This yields:
0: 3
1: 0
2: 0
3: 0
4: 4
5: 4
6: 3
7: 2
8: 3
9: 3
10: 4
11: 5
12: 0
13: 0
14: 0
15: 0
16: 6
17: 7
18: 7

At this point we need to understand how to go from these values to
bit strings. This is done with Huffman coding.
https://en.wikipedia.org/wiki/Huffman_coding

I won't go into too much detail but the important idea is that we can git a
bunch of bit strings whereby each string can be uniquely identified.

For example the following, when read left to right are unique:
010
001
000

The following however would not be ok:
01
01011

Since if we saw 01 we wouldn't know if it was its own code or if it was a
subset of 01011

The algorithm is as follows:

- Start with 0 at the minimum length.
- Increment by 1 for every code at that length
- Once done go back to the next length and start at 2*last_value + 1
"""


# ===== VSB Verifications: Velocity Profiles =====

println("=== BEGINNING Verifications4.jl ===")

# === Imports ===
println("   IMPORTING PACKAGES...")
using VSB
using Plots
pyplot()

cd("/Users/Damyn/Documents/BYU/FLOW Lab/VSB/verifications")

# === Geomerty of problem ===
println("   GENERATING GEOMETRY...")
NPTS = 250                              # Number of Points on Boundary
del_theta = (2*pi)/(NPTS - 1)           # Step in theta around body
R = 1.0                                 # Cylinder Radius
X_coor = zeros(NPTS)                    # X-coordinates
Y_coor = zeros(NPTS)                    # Y-coordinates
X_tHat = zeros(NPTS)                    # X Tangent Vector
Y_tHat = zeros(NPTS)                    # X Tangent Vector
X_nHat = zeros(NPTS)                    # Y Normal Vector
Y_nHat = zeros(NPTS)                    # Y Normal Vector
theta = zeros(NPTS)                     # Angle from positive X-axis
for i=1:NPTS
    X_coor[i] = R*cos((i-1)*del_theta)
    Y_coor[i] = R*sin((i-1)*del_theta)

    X_tHat[i] = sin((i-1)*del_theta)
    Y_tHat[i] = -cos((i-1)*del_theta)

    X_nHat[i] = cos((i-1)*del_theta)
    Y_nHat[i] = sin((i-1)*del_theta)

    theta[i] = (i-1)*del_theta
end
body_pts = [[X_coor[i], Y_coor[i], 0.0] for i in 1:NPTS]
t_hats = [[X_tHat[i], Y_tHat[i], 0.0] for i in 1:NPTS]
n_hats = [[X_nHat[i], Y_nHat[i], 0.0] for i in 1:NPTS]
body = VSB.Boundary(body_pts,t_hats,n_hats)

# === Vortex sheet calculations ===
println("   CALCULATING PROBLEM PARAMETERS...")
magU_inf = 1.0
U_inf(X) = magU_inf .* [1.0, 0.0, 0.0]

params = VSB.Parameters(body, U_inf)
alphas = params.alphas
U(X) = VSB.CalcVSVelocityCirlce(body, alphas, X)

# === Generate velocity field ===
println("   CALCULATING VELOCITY FIELD...")

println("=== END OF Verifications1.jl ===")


#=
Plots consumption, income and debt for the simple infinite horizon LQ
permanent income model with Gaussian iid income.

@author : Spencer Lyon

@date: 07/09/2014
=#

using PyPlot

const r = 0.05
const beta = 1.0 / (1.0 + r)
const T = 60
const sigma = 0.15
const mu = 1.0

function time_path()
    w = randn(T+1)
    w[1] =  0.0
    b = Array(Float64, T+1)
    for t=2:T+1
        b[t] = sum(w[1:t])
    end
    b .*= -sigma
    c = mu + (1.0 - beta) .* (sigma .* w .- b)
    return w, b, c
end

# == Figure showing a typical realization == #
function single_realization()
    fig, ax = subplots()

    ax[:grid]()
    ax[:set_xlabel]("Time")
    bbox = [0.0, 1.02, 1.0, 0.102]
    p_args = {:lw=> 2, :alpha => 0.7}

    w, b, c = time_path()
    ax[:plot](0:T, mu + sigma .* w, "g-", label="non-financial income"; p_args...)
    ax[:plot](0:T, c, "k-", label="consumption"; p_args...)
    ax[:plot](0:T, b, "b-", label="debt"; p_args...)
    ax[:legend](ncol=3, bbox_to_anchor=bbox, loc="upper left", mode="expand")

    return Void
end

# == Figure showing multiple consumption paths == #
function consumption_paths(n=250)  # n is number of paths
    fix, ax = subplots()
    ax[:grid]()
    ax[:set_xlabel]("Time")
    ax[:set_ylabel]("Consumption")
    colors = ["c-", "g-", "k-", "b-"]
    for i=1:n
        ax[:plot](0:T, time_path()[3], colors[rand(1:4)], lw=0.8, alpha=0.7)
    end
    return Void
end