text:
using Distributed
addprocs(length(Sys.cpu_info()))
using LinearAlgebra
@everywhere basepath = "/gscratch/astro/tagordon/celmcmc/"
@everywhere srcpath = basepath*"src/"
@everywhere datapath = basepath*"data/full/"
@everywhere outpath = basepath*"data/full/tmp"
@everywhere inputpath = basepath*"input/full/"
@everywhere include(srcpath*"GPdist/celerite.jl")
@everywhere include(srcpath*"TransitModel/transit.jl")

## Define a new multivariate Distribution type for Mamba.
## The definition must be placed within an unevaluated quote block.
@everywhere extensions = quote

    ## Load needed packages and import methods to be extended
    using Distributions
    import Distributions: length, insupport, _logpdf

    ## Type declaration
    mutable struct CeleriteDist <: ContinuousMultivariateDistribution
        t::Array{Real, 1}
        c::Array{Real, 1}
        log_s::Real
        log_q::Real
        log_w::Real
        log_wn::Array{Real, 1}
        mu::Array{Real, 1}
        gp::celerite.Celerite

        function CeleriteDist(t, log_s, log_q, log_w; log_wn=-12, c=[1], mu=zeros(Base.length(c), Base.length(t)))
            gp = celerite.Celerite(celerite.SHOTerm(log_s, log_q, log_w), broadcast(*, c, c'))
            new(t, c, log_s, log_q, log_w, log_wn, mu, gp)
        end
    end

    ## The following method functions must be implemented

    ## Dimension of the distribution
    length(d::CeleriteDist) = length(d.t)*length(d.c)

    ## Logical indicating whether x is in the support
    function insupport(d::CeleriteDist, x::AbstractVector{T}) where {T<:Real}
        length(d) == length(x) && all(isfinite.(x))
    end

    function _logpdf(d::CeleriteDist, x::AbstractVector{T}) where {T<:Real}
        m = Base.length(d.c)
        n = length(d.t)
        wn_vec = zeros(n*m)
	for i in 1:m
		wn_vec[i:m:end] .= exp.(log_wn[i])
	end
        try
            celerite.compute!(d.gp, d.t, wn_vec)
            return celerite.log_likelihood(d.gp, x.-d.mu)
        catch
            return -Inf
        end
    end

    function _rand!(d::CeleriteDist, x::AbstractVector{T}) where {T<:Real}
        m = Base.length(d.c)
        n = length(d.t)
        wn_vec = zeros(n*m)
        for i in 1:m
            wn_vec[i:m:end] .= exp(d.log_wn[i])
        end
        celerite.compute!(d.gp, d.t, wn_vec)
        u = randn(Base.length(d.t)*Base.length(d.c))
        x[:] = celerite.simulate_gp(d.gp, u)
        for i in 1:m
            x[i:m:end] .+= d.mu[i,:]
        end
    end

    mean(d::CeleriteDist) = convert(Array{Real, 2}, d.mu)

end

@everywhere include(srcpath*"Mamba.jl/src/Mamba.jl")
@everywhere using Main.Mamba
@everywhere eval(extensions)

@everywhere using DelimitedFiles
@everywhere codefile = inputpath*"codes.in"
@everywhere paramfile = inputpath*"params_poly.in"
@everywhere runcodes = readdlm(codefile, String)
@everywhere params = readdlm(paramfile)
@everywhere paramdict = Dict(runcodes[i] => params[i,:] for i in 1:length(runcodes))

# build model

@everywhere TransitModel = Model(

    # defines the probability of y given mu and sig as a multivariate normal likelihood function
    y = Stochastic(1,
        (t, c, log_s, log_w, log_q, log_wn, trans) -> CeleriteDist(t, log_s, log_q, log_w, log_wn=log_wn, mu=trans, c=[c[1], c[2]]),
        false
    ),

    trans = Logical(1,
        (t0, b0, rp, d, u, t) -> transit(t, t0, b0, rp, d, u)[1],
        false,
    ),

    # transit parameters
    t0 = Stochastic(
        () -> Uniform(-1, 1)
    ),
    b0 = Stochastic(
        () -> Uniform(0, 1)
    ),

    log_rp = Stochastic(
        () -> Uniform(-20, 0),
        false
    ),
    rp = Logical(
        (log_rp) -> exp(log_rp)
    ),
    d = Stochastic(
        () -> Uniform(0, 10)
    ),
    q = Stochastic(1,
        () -> UnivariateDistribution[Uniform(0, 1), Uniform(0, 1)],
        false
    ),
    u = Logical(1,
        (q) -> [2*sqrt(q[1])*q[2], sqrt(q[1])*(1-2*q[2])]
    ),

    #noise parameters
    log_s = Stochastic(
        () -> Uniform(-30, 0)
    ),
    log_q = Stochastic(
        () -> Uniform(-10, 10)
    ),
    log_w = Stochastic(
        () -> Uniform(-10, 10)
    ),
    log_wn = Stochastic(1,
        () -> UnivariateDistribution[Uniform(-40, 10), Uniform(-40, 10)]
    ),
    c = Stochastic(1,
        () -> UnivariateDistribution[Uniform(0, 10), Uniform(0, 10)]
    )
)

# define the sampling scheme: No U-Turn Sampling for the three parameters
@everywhere scheme = [AMWG([:t0, :b0, :log_rp, :d, :q, :log_s, :log_q, :log_w, :log_wn, :c], 0.1*ones(13))]

# define the initial values for the walkers
@everywhere epsilon = 1e-8
@everywhere nchains = 4
if nprocs() > length(runcodes) nchains = Int(floor(nprocs()/length(runcodes))) end

@everywhere function f(code)
    p = paramdict[code]
    log_wn1, log_wn2, c1, c2, b0, rp, d, q1, q2, log_s, log_q, log_w, t0 = p
    transfile = datapath*"transit_poly_"*code
    fluxdata = readdlm(transfile)

    data = Dict{Symbol, Any}(
        :t => fluxdata[:,1],
        :y => vec(permutedims(fluxdata[:, 2:end]))
    )
    inits = [
        Dict{Symbol, Any}(
            :y => data[:y],

            #transit parameters
            :t0 => t0 + epsilon*randn(),
            :b0 => b0 + epsilon*randn(),
            :log_rp => log(rp) + epsilon*randn(),
            :d => d + epsilon*randn(),
            :q => [q1, q2] + epsilon*randn(2),

            #noise parameters
            :log_s => log_s + epsilon*randn(),
            :log_q => log_q + epsilon*randn(),
            :log_w => log_w + epsilon*randn(),
            :log_wn => [log_wn1, log_wn2] + epsilon*randn(2),
            :c => [c1, c2] + epsilon*randn(2)
        )
        for i in 1:nchains
    ]

    setsamplers!(TransitModel, scheme)
    simfile = outpath*"simpoly"*code*".jls"
    sim = mcmc(TransitModel, data, inits, 10, burnin=0, chains=nchains)
    write(simfile, sim)
end
pmap(f, runcodes)


text:
module Libjsp

include("jdft.jl")
include("jwindow.jl")

export dft
export window

end  # module libjsp

ext: jl


text:
peel(t::Tuple) = map(peel, t)

# macro peel(args...)
#     ex = :nothing
#     for arg in args
#         exi = :($arg = peel($arg))
#         ex = :($exi; $ex)
#     end
#     ex
# end

macro peel(args...)
    if length(args) < 1; error("@peel: at least one arg!"); end
    body = args[end]
    args = args[1:(end-1)]
    peels = {:(($arg) = peel($arg)) | arg in args }
    expr(:let, body, peels...)
end

abstract Partial{T}

# -- Input --------------------------------------------------------------------

type Input{T} <: Partial{T}
    name::Symbol
end

#show{T}(arg::Input{T}) = print("Input{$T}(:$(arg.name))")

index(name::Symbol) = Input{Int}(name)
input_array(T, N, name::Symbol) = Input{Array{T,N}}(name)

peel(arg::Input) = arg.name

# -- Operation ----------------------------------------------------------------

type Operation{T} <: Partial{T}
#    eval::Function
    ex::Expr
end

peel(op::Partial) = op.ex

function ref{T,N}(A::Partial{Array{T,N}}, inds::Partial...)
#    Operation{T}(args->( peval(A, args)[peval(inds, args)...] ))
#    Operation{T}(ev->( ev(A)[ev(inds)] ))
#    Operation{T}(quote; $(peel(A)) [$(peel(inds))]; end)

#    @peel A inds Operation{T}(:( ($A)[$inds] ))
    @peel A inds Operation{T}( expr(:ref, A, inds...) )
end

function *{Tx,Ty}(x::Partial{Tx}, y::Partial{Ty})
    T = promote_type(Tx, Ty)
#    Operation{T}(ev->( ev(x)*ev(y) ))
#    Operation{T}(quote; $(peel(x)) * $(peel(y)); end)
    @peel x y Operation{T}(:( ($x)*($y) ))
end

# -- Output -------------------------------------------------------------------

type Output{T}
    writes::Vector
    name::Symbol
    Output(name::Symbol) = new({}, name)
end

#output_array(T, N, name) = Output{Array{T,N}}({}, name)
output_array{T}(::Type{T}, N, name) = Output{Array{T,N}}(name)

peel(arg::Output) = arg.name

function assign{T,N}(A::Output{Array{T,N}}, x::Partial{T}, inds::Partial...)
#    write_op = ev->( ev(A)[ev(inds)] = ev(x) )
#    write_op = quote $(peel(A)) [ $(peel(inds)) ] = $(peel(x)); end

#    write_op = @peel A x inds :( ($A)[($inds)...] = ($x) )
#    write_op = @peel A x inds :( ($A)[($inds...)] = ($x) )

    write_op = @peel A x inds begin
        ref = expr(:ref, A, inds...)
        :( ($ref) = $x )
    end
    append!(A.writes, {write_op})
end

# -- Test code ----------------------------------------------------------------

VType = Float32

i = index(:i)
j = index(:j)

A = input_array(VType, 2, :A)
B = input_array(VType, 2, :B)
dest = output_array(VType, 2, :dest)

a = A[j,i]
b = B[j,i]
c = a*b
dest[i,j] = c

@eval f(dest,A,B,i,j) = ($dest.writes[1])

X = [1 2; 3 4]
d=zeros(2,2)
f(d,X,X,1,1)

max_stars_repo_name: toivoh/julia-kernels

max_issues_repo_name: toivoh/julia-kernels

max_forks_repo_name: toivoh/julia-kernels


text:
struct ParseError <: Exception
    position::Int
    mode::PARSE_MODE
end

struct CloseTokenError <: Exception
    position::Int
end

max_stars_repo_name: haverzard/Husk

max_issues_repo_name: haverzard/Husk

max_forks_repo_name: haverzard/Husk


text:
using DiffEqDevDocs
using Base.Test


text:
@doc raw"""
    function make_kernel(f, node::N, d::D) where {N<:Node, D<:ContinuousUnivariateDistribution}

Creates a single site proposal kernel for `node` that has an associated continuous univariate distribution.
The resulting proposal kernel is a normal distribution, ``z' | z \sim \mathrm{Normal}(z, (0.5 \sigma)^2)`` where
``\sigma`` is the standard deviation of the node's distribution.
"""
function make_kernel(f, node::N, d::D) where {N<:Node, D<:ContinuousUnivariateDistribution}
    ex = extrema(node.dist)
    sd = std(node.dist)
    function k(t0::Trace, t1::Trace, params...)
        n = truncated(Normal(t0[node.address].value, sd * 0.5), ex[1], ex[2])
        propose(t1, node.address, n)
    end
    k
end

@doc raw"""
    function make_kernel(f, node::N, d::D) where {N<:Node, D<:DiscreteUnivariateDistribution}

Creates a single site proposal kernel for `node` that has an associated discrete univariate distribution.
This method assumes that the distribution is defined over some subset of the integers; a support that does not meet
this criteria *may* result in runtime inference errors. The resulting proposal kernel is a discrete nonparametric
distribution, ``z' | z \sim \mathrm{DiscreteNonParametric}([z - 1, z, z + 1])`` if ``z`` is not on the boundary of
the distribution's support; or ``z' | z \sim \mathrm{DiscreteNonParametric}([z \pm 1, z])`` if ``z`` is on the
boundary of the distribution's support, with the sign depending on the left or right of the support.
"""
function make_kernel(f, node::N, d::D) where {N<:Node, D<:DiscreteUnivariateDistribution}
    ex = extrema(node.dist)
    make_kernel(f, node, d, ex, Val(ex[2]))
end

function make_kernel(f, node::N, d::D, ex, v) where {N<:Node, D<:DiscreteUnivariateDistribution}
    function k(t0::Trace, t1::Trace, params...)
        last = t0[node.address].value
        left = last - 1
        right = last + 1
        if last == ex[1]
            propose(t1, node.address, DiscreteNonParametric([last, right], [0.5, 0.5]))
        elseif last == ex[2]
            propose(t1, node.address, DiscreteNonParametric([left, last], [0.5, 0.5]))
        else
            w = ones(3)
            w = w ./ sum(w)
            propose(t1, node.address, DiscreteNonParametric([left, last, right], w))
        end
    end
    k
end
#=
function make_kernel(f, node::N, d::D, ex, ::Val{Inf}) where {N<:Node, D<:DiscreteUnivariateDistribution}
    ex = extrema(node.dist)
    k1(t0::Trace, t1::Trace, params...) = propose(t1, node.address, truncated(Poisson(t0[node.address].value), t0[node.address].value - 10, t0[node.address].value + 10))
    k2 = make_kernel(f, node, d, ex, nothing)
    k(t0::Trace, t1::Trace, params...) = rand() < 0.1 ? k1(t0, t1, params...) : k2(t0, t1, params...)
    k
end
=#
@doc raw"""
    function make_kernel(f, node::N, d::D) where {N<:Node, D<:ContinuousMultivariateDistribution}

Creates a single site proposal kernel for `node` that has an associated continuous multivariate distribution.
Assumes that this site is *unconstrained* in ``\mathbb{R}^N``; using this method will result in inference runtime errors if the site
is constrained. The resulting proposal kernel is a multivariate normal,
``z' | z \sim \mathrm{MultivariateNormal}(z, 1 / \sqrt{|D|})`` where ``D`` is the dimensionality of the site.
"""
function make_kernel(f, node::N, d::D) where {N<:Node, D<:ContinuousMultivariateDistribution}
    dim = size(node.dist)[1]
    k(t0::Trace, t1::Trace, params...) = propose(t1, node.address, MvNormal(t0[node.address].value, 1.0 / sqrt(dim)))
    k
end

@doc raw"""
    function make_kernel(f, node::N, d) where {N<:Node}

Create a single site proposal kernel for an arbitrary site. This method draws from the prior site distribution
and is therefore inefficient (both computationally and statistically).
"""
function make_kernel(f, node::N, d) where {N<:Node}
    function k(t0::Trace, t1::Trace, params...)
        t = trace()
        f(t, params...)
        propose(t1, node.address, t[node.address].dist)
    end
    k
end

make_kernel(f, node::N) where N <: Node = make_kernel(f, node, node.dist)

@doc raw"""
    function make_kernels(f, t::Trace, addresses::Vector{T}; params = (), include::Bool=true) where T

Create a vector of single-site proposal functions for the stochastic function `f`.
This method assumes static model structure and generates a proposal either for each site in `addresses` (`include = true`)
or for each site *not* in `addresses` (`include = false`). The tuple of params are any necessary parameters needed to execute
`f`.
"""
function make_kernels(f, t::Trace, addresses::Vector{T}; params = (), include::Bool=true) where T
    t = trace()
    f(t, params...)
    kernels = []
    cond = (a, addresses) -> include ? a in addresses : !(a in addresses)
    for (a, n) in t.trace
        if cond(a, addresses) && !t.trace[a].observed
            push!(kernels, make_kernel(f, n))
        end
    end
    kernels
end

@doc raw"""
    function make_kernels(f; params = ())

Create a vector of single-site proposal functions for the stochastic function `f`.
This method assumes static model structure and generates a proposal kernel for each
model site. The tuple of params are any necessary parameters needed to execute
`f`.
"""
function make_kernels(f; params = ())
    t = trace()
    f(t, params...)
    a = collect(keys(t))
    make_kernels(f, t, a; params = params)
end

export make_kernel, make_kernels


text:
# Test the TranscriptionData datatype
@testset "TranscriptionData" begin
    @test TranscriptionData isa DataType
    @test TranscriptionData().finvar_mappings isa Dict
end

# Test basic definition and queries
@testset "Basic Definition and Queries" begin
    # initialize needed data
    mockoptimizer = () -> MOIU.MockOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()),
                                             eval_objective_value=false)
    # test TranscriptionModel (no optimizer)
    @testset "TranscriptionModel (Default)" begin
        @test haskey(TranscriptionModel().ext, :TransData)
    end
    # test TranscriptionModel (with optimizer)
    @testset "TranscriptionModel (Optimizer)" begin
        @test haskey(TranscriptionModel(mockoptimizer).ext, :TransData)
    end
    # test is_transcription_model
    @testset "is_transcription_model" begin
        @test is_transcription_model(TranscriptionModel())
        @test !is_transcription_model(Model())
    end
    # test transcription_data
    @testset "transcription_data" begin
        @test transcription_data(TranscriptionModel()).semi_infinite_vars isa Vector
        @test_throws ErrorException transcription_data(Model())
    end
    # test has_internal_supports
    @testset "has_internal_supports" begin
        m = TranscriptionModel()
        @test !IOTO.has_internal_supports(m)
        transcription_data(m).has_internal_supports = true
        @test IOTO.has_internal_supports(m)
    end
end

# Test query helper functions
@testset "Query Formatters" begin
    # initialize models
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1], supports = [0, 1])
    @infinite_parameter(m, pars[1:2] in [0, 1])
    @variable(m, x, Infinite(par, pars))
    @variable(m, q, Infinite(pars, par))
    @variable(m, w, Infinite(par))
    @variable(m, x0, Point(x, 0, [0, 0]))
    @variable(m, y)
    add_supports(par, 0.5, label = InternalLabel)
    @variable(m, xrv, SemiInfinite(x, par, [1, pars[2]]))
    tm = optimizer_model(m)
    data = transcription_data(tm)
    data.has_internal_supports = true
    data.supports = ([0., 0.5, 1., NaN], [[0., 0.], [1., 1.], [NaN, NaN]])
    s1 = Set([UserDefined])
    s2 = Set([InternalLabel])
    sd = Set{DataType}()
    data.support_labels = ([s1, s2, s1, sd], [s1, s1, sd])
    # test _get_array_type
    @testset "_get_array_type" begin
        @test IOTO._get_array_type(ones(Bool, 2)) == Bool
    end
    # test _get_object_numbers
    @testset "_get_object_numbers" begin
        @test IOTO._get_object_numbers(x) == [1, 2]
        @test IOTO._get_object_numbers(y) == []
        @test IOTO._get_object_numbers(y + x) == [1, 2]
    end
    # test make_ndarray
    @testset "make_ndarray" begin
        # test finite variable
        @test IOTO.make_ndarray(tm, y, [1], PublicLabel) == [1]
        # test ordered infinite variable
        @test IOTO.make_ndarray(tm, x, collect(1:6), All) == [1 4; 2 5; 3 6]
        @test IOTO.make_ndarray(tm, x, collect(1:6), PublicLabel) == [1 4; 3 6]
        # test unordered infinite variable
        @test IOTO.make_ndarray(tm, q, collect(1:6), All) == [1 2 3; 4 5 6]
        @test IOTO.make_ndarray(tm, q, collect(1:6), PublicLabel) == [1 3; 4 6]
        # test infinite variable with single parameter
        @test IOTO.make_ndarray(tm, w, collect(1:3), All) == [1, 2, 3]
        # test expression
        @test IOTO.make_ndarray(tm, w + x, collect(1:6), All) == [1 4; 2 5; 3 6]
        # test error
        @test_throws ErrorException IOTO.make_ndarray(tm, q, collect(1:3), All)
    end
end

# Test variable mapping queries
@testset "Variable Mapping Queries" begin
    # initialize models
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1], supports = [0, 1])
    @infinite_parameter(m, pars[1:2] in [0, 1])
    @variable(m, x, Infinite(par, pars))
    @variable(m, x0, Point(x, 0, [0, 0]))
    @variable(m, y)
    f1 = parameter_function(sin, par)
    f2 = parameter_function((a, b) -> 1, (par, pars))
    add_supports(par, 0.5, label = InternalLabel)
    supps = Dict{Int, Float64}(2 => 1)
    xrv = add_variable(m, build_variable(error, x, supps))
    tm = optimizer_model(m)
    data = transcription_data(tm)
    data.has_internal_supports = true
    data.supports = ([0., 0.5, 1., NaN], [[0., 0.], [1., 1.], [NaN, NaN]])
    s1 = Set([UserDefined])
    s2 = Set([InternalLabel])
    s3 = union(s1, s2)
    sd = Set{DataType}()
    data.support_labels = ([s1, s2, s1, sd], [s1, s1, sd])
    @variable(tm, a)
    @variable(tm, b)
    @variable(tm, c)
    @variable(tm, d)
    @variable(tm, e)
    @variable(tm, f)
    # test _ignore_label
    @testset "_ignore_label" begin
        @test IOTO._ignore_label(tm, All)
        @test !IOTO._ignore_label(tm, PublicLabel)
        @test !IOTO._ignore_label(tm, UserDefined)
    end
    # test transcription_variable (finite variables)
    @testset "transcription_variable (Finite)" begin
        # test error
        @test_throws ErrorException transcription_variable(tm, y)
        @test_throws ErrorException transcription_variable(tm, x0)
        # test normal
        data.finvar_mappings[y] = a
        @test transcription_variable(tm, y) == a
        data.finvar_mappings[x0] = b
        @test transcription_variable(tm, x0, ndarray = true, label = All) == b
    end
    # test transcription_variable (Infinite, semi-infinite, and derivative)
    @testset "transcription_variable (Infinite)" begin
        # test error
        @test_throws ErrorException transcription_variable(tm, x)
        @test_throws ErrorException transcription_variable(tm, xrv)
        # test normal
        data.infvar_mappings[x] = [a, b, c, d, e, f]
        data.infvar_support_labels[x] = [s1, s3, s1, s1, s3, s1]
        @test transcription_variable(tm, x) == [a, b, c, d, e, f]
        @test transcription_variable(tm, x, label = All) == [a, b, c, d, e, f]
        @test transcription_variable(tm, x, label = InternalLabel) == [b, e]
        data.infvar_mappings[xrv] = [d, e, f]
        data.infvar_support_labels[xrv] = [s1, s3, s1]
        @test transcription_variable(tm, xrv) == [d, e, f]
        @test transcription_variable(tm, xrv, label = All) == [d, e, f]
        @test transcription_variable(tm, xrv, label = InternalLabel) == [e]
        # test ndarray
        @test transcription_variable(tm, x, label = All, ndarray = true) == [a d; b e; c f]
        @test transcription_variable(tm, x, ndarray = true) == [a d; c f]
        @test_throws ErrorException transcription_variable(tm, xrv, ndarray = true)
    end
    # test transcription_variable (Parameter Function)
    @testset "transcription_variable (Parameter Function)" begin
        # test normal
        @test transcription_variable(tm, f1) == [sin(0), sin(1)]
        @test transcription_variable(tm, f1, label = All) == sin.([0, 0.5, 1])
        # test ndarray
        @test transcription_variable(tm, f1, label = All, ndarray = true) == sin.([0, 0.5, 1])
        @test transcription_variable(tm, f2, ndarray = true) == ones(2, 2)
    end
    # test transcription_variable (Fallback)
    @testset "transcription_variable (Fallback)" begin
        @test_throws ErrorException transcription_variable(tm, par)
    end
    # test transcription_variable (Single argument)
    @testset "transcription_variable (Single)" begin
        @test transcription_variable(y) == a
        @test transcription_variable(x, label = All) == [a, b, c, d, e, f]
        @test transcription_variable(x0) == b
        @test transcription_variable(f2, ndarray = true) == ones(2, 2)
    end
    # test optimizer_model_variable extension
    @testset "optimizer_model_variable" begin
        @test optimizer_model_variable(y, Val(:TransData), label = All) == a
        @test optimizer_model_variable(x, Val(:TransData), label = All) == [a, b, c, d, e, f]
        @test optimizer_model_variable(x, Val(:TransData)) == [a, b, c, d, e, f]
        @test optimizer_model_variable(x0, Val(:TransData)) == b
    end
    # test variable_supports for infinite variable with 2 inputs
    @testset "variable_supports (Model, Infinite)" begin
        # test not in mappings error
        dvref = dispatch_variable_ref(x)
        delete!(data.infvar_mappings, x)
        @test_throws ErrorException InfiniteOpt.variable_supports(tm, dvref)
        data.infvar_mappings[x] = [a, b, c, d, e, f]
        # test supports are empty
        lookups = Dict{Vector{Float64}, Int}([0, 0, 0] => 1, [0.5, 0, 0] => 2, [1, 0, 0] => 3,
                                             [0, 1, 1] => 4, [0.5, 1, 1] => 5, [1, 1, 1] => 6)
        data.infvar_lookup[x] = lookups
        expected = [(0., [0., 0.]), (0.5, [0., 0.]), (1., [0., 0.]),
                    (0., [1., 1.]), (0.5, [1., 1.]), (1., [1., 1.])]
        @test InfiniteOpt.variable_supports(tm, dvref) == expected
        # test normal
        @test InfiniteOpt.variable_supports(tm, dvref, label = All) == expected
        @test InfiniteOpt.variable_supports(tm, dvref, label = InternalLabel) == expected[[2, 5]]
        # test ndarray
        expected = permutedims([(0., [0., 0.]) (0.5, [0., 0.]) (1., [0., 0.]);
                                (0., [1., 1.]) (0.5, [1., 1.]) (1., [1., 1.])], (2, 1))
        @test InfiniteOpt.variable_supports(tm, dvref, ndarray = true) == expected[[1, 3], :]
        @test InfiniteOpt.variable_supports(tm, dvref, ndarray = true, label = All) == expected
        # test with semi-infinite variable
        lookups = Dict{Vector{Float64}, Int}([0, 1] => 1, [0.5, 1] => 2, [1, 1] => 3)
        data.infvar_lookup[xrv] = lookups
        dvref = dispatch_variable_ref(xrv)
        expected = [(0., 1.), (0.5, 1.), (1., 1.)]
        @test InfiniteOpt.variable_supports(tm, dvref) == expected
        @test InfiniteOpt.variable_supports(tm, dvref, label = InternalLabel) == [expected[2]]
    end
    # test variable_supports for infinite parameter functions with 2 inputs
    @testset "variable_supports (Model, Parameter Function)" begin
        # test normal
        df1 = dispatch_variable_ref(f1)
        @test InfiniteOpt.variable_supports(tm, df1) == [(0.,), (1.,)]
        @test InfiniteOpt.variable_supports(tm, df1, label = All) == [(0.,), (0.5,), (1.,)]
        # test ndarray
        df2 = dispatch_variable_ref(f2)
        @test InfiniteOpt.variable_supports(tm, df1, label = All, ndarray = true) == [(0.,), (0.5,), (1.,)]
        @test InfiniteOpt.variable_supports(tm, df2, ndarray = true) isa Array
    end
    # test supports for infinite variable
    @testset "supports (Infinite)" begin
        @test supports(x, label = InternalLabel) == [(0.5, [0., 0.]), (0.5, [1., 1.])]
        @test supports(xrv) == [(0., 1.), (0.5, 1.), (1., 1.)]
        @test supports(f1, label = All) == [(0.,), (0.5,), (1.,)]
    end
    # test lookup_by_support (infinite vars)
    @testset "lookup_by_support (Infinite)" begin
        # test errors
        @variable(m, x2, Infinite(par))
        @test_throws ErrorException IOTO.lookup_by_support(tm, x2, [0.])
        @test_throws ErrorException IOTO.lookup_by_support(tm, x, [0.9, 0., 0.])
        @test_throws ErrorException IOTO.lookup_by_support(tm, xrv, [0., 0., 0.])
        # test normal
        @test IOTO.lookup_by_support(tm, x, [0., 0., 0.]) == a
        @test IOTO.lookup_by_support(tm, x, [0., 1., 1.]) == d
        @test IOTO.lookup_by_support(tm, x, [1., 1., 1.]) == f
        @test IOTO.lookup_by_support(tm, xrv, [0., 1.]) == d
        @test IOTO.lookup_by_support(tm, xrv, [1., 1.]) == f
    end
    # test lookup_by_support (infinite parameter functions)
    @testset "lookup_by_support (Parameter Function)" begin
        @test IOTO.lookup_by_support(tm, f1, [0.]) == 0
        @test IOTO.lookup_by_support(tm, f2, [0., 0., 1.]) == 1
    end
    # test lookup_by_support (finite vars)
    @testset "lookup_by_support (Finite)" begin
        # test errors
        @variable(m, z2)
        @test_throws ErrorException IOTO.lookup_by_support(tm, z2, [0.])
        # test normal
        @test IOTO.lookup_by_support(tm, x0, [0., 0., 0.]) == b
        @test IOTO.lookup_by_support(tm, y, [0., 0., 1.]) == a
    end
    # test internal_semi_infinite_variable
    @testset "internal_semi_infinite_variable" begin
        rv = SemiInfiniteVariableRef(m, SemiInfiniteVariableIndex(-1))
        # test errors
        @test_throws ErrorException InfiniteOpt.internal_semi_infinite_variable(rv, Val(:TransData))
        # test normal
        var = build_variable(error, x, supps)
        push!(transcription_data(tm).semi_infinite_vars, var)
        @test InfiniteOpt.internal_semi_infinite_variable(rv, Val(:TransData)) == var
        eval_supports(rv) == supps
    end
end

# Test the measure mapping methods
@testset "Measure Queries" begin
    # initialize tbe needed info
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1], supports = [0, 1])
    @infinite_parameter(m, pars[1:2] in [0, 1])
    @variable(m, x, Infinite(par, pars))
    @variable(m, x0, Point(x, 0, [0, 0]))
    @variable(m, y)
    meas1 = support_sum(2par -2, par)
    meas2 = support_sum(x^2 - y, pars)
    tm = optimizer_model(m)
    data = transcription_data(tm)
    data.has_internal_supports = true
    s1 = Set([UserDefined])
    s2 = Set([InternalLabel])
    sd = Set{DataType}()
    data.supports = ([0., 1., NaN], [[0., 0.], [1., 1.], [NaN, NaN]])
    data.support_labels = ([s1, s2, sd], [s1, s1, sd])
    @variable(tm, a)
    @variable(tm, b)
    @variable(tm, c)
    @variable(tm, d)
    # test transcription_variable
    @testset "transcription_variable" begin
        # test errors
        @test_throws ErrorException IOTO.transcription_variable(tm, meas1)
        @test_throws ErrorException IOTO.transcription_variable(tm, meas2)
        # test normal
        data.measure_mappings[meas1] = [-2 * zero(AffExpr)]
        data.measure_support_labels[meas1] = [s1]
        data.measure_mappings[meas2] = [a^2 + c^2 - 2a, b^2 + d^2 - 2a]
        data.measure_support_labels[meas2] = [s1, s2]
        @test IOTO.transcription_variable(tm, meas1) == -2 * zero(AffExpr)
        expected = [a^2 + c^2 - 2a, b^2 + d^2 - 2a]
        @test IOTO.transcription_variable(tm, meas2) == [expected[1]]
        @test IOTO.transcription_variable(tm, meas2, label = All) == expected
        # test ndarray
        @test IOTO.transcription_variable(tm, meas2, ndarray = true) == [expected[1]]
        @test IOTO.transcription_variable(tm, meas2, label = All, ndarray = true) == expected
    end
    # test lookup_by_support
    @testset "lookup_by_support" begin
        # test errors
        @test_throws ErrorException IOTO.lookup_by_support(tm, meas1, Float64[])
        @test_throws ErrorException IOTO.lookup_by_support(tm, meas2, [0.])
        # test normal
        data.measure_lookup[meas1] = Dict(Float64[] => 1)
        data.measure_lookup[meas2] = Dict{Vector{Float64}, Int}([0] => 1, [1] => 2)
        @test IOTO.lookup_by_support(tm, meas1, Float64[]) == -2 * zero(AffExpr)
        @test IOTO.lookup_by_support(tm, meas2, [1.]) == b^2 + d^2 - 2a
    end
    # test InfiniteOpt.variable_supports
    @testset "InfiniteOpt.variable_supports" begin
        # test not in mappings error
        dvref = dispatch_variable_ref(meas1)
        delete!(data.measure_mappings, meas1)
        @test_throws ErrorException InfiniteOpt.variable_supports(tm, dvref)
        data.measure_mappings[meas1] = [-2 * zero(AffExpr)]
        # test supports are empty
        @test InfiniteOpt.variable_supports(tm, dvref) == ()
        dvref2 = dispatch_variable_ref(meas2)
        @test InfiniteOpt.variable_supports(tm, dvref2, label = All) == [(0.,), (1.,)]
        # test normal
        @test InfiniteOpt.variable_supports(tm, dvref) == ()
        @test InfiniteOpt.variable_supports(tm, dvref2) == [(0.,)]
        # test ndarray
        @test InfiniteOpt.variable_supports(tm, dvref, ndarray = true) == [()]
        @test InfiniteOpt.variable_supports(tm, dvref2, ndarray = true) == [(0.,)]
        @test InfiniteOpt.variable_supports(tm, dvref2, label = All, ndarray = true) == [(0.,), (1.,)]
    end
end

# Test the support iteration methods
@testset "Support Iteration Methods" begin
    # initialize the needed info
    m = InfiniteModel()
    @infinite_parameter(m, pars[1:2] in [0, 1], num_supports = 2)
    @infinite_parameter(m, par in [0, 1], supports = [0, 1],
                        derivative_method = OrthogonalCollocation(3))
    @variable(m, y, Infinite(par))
    d1 = @deriv(y, par)
    tm = optimizer_model(m)
    # test _temp_parameter_ref
    @testset "_temp_parameter_ref" begin
        @test IOTO._temp_parameter_ref(m, index(par)) == dispatch_variable_ref(par)
        @test IOTO._temp_parameter_ref(m, index(pars[1]).object_index) == dispatch_variable_ref(pars[1])
    end
    # test _collected_supports
    @testset "_collected_supports" begin
        # independent parameters
        expected = Float64[0., 1., NaN]
        @test isequal(IOTO._collected_supports(dispatch_variable_ref(par)), expected)
        # dependent parameters
        expected = sort!([Float64[0., 0.], Float64[1., 1.], Float64[NaN, NaN]])
        @test isequal(sort!(IOTO._collected_supports(dispatch_variable_ref(pars[1]))), expected)
    end
    # test _collected_support_labels
    @testset "_collected_support_labels" begin
        # independent parameters
        supps = Float64[0., 1., NaN]
        expected = [Set([UserDefined]), Set([UserDefined]), Set{DataType}()]
        @test isequal(IOTO._collected_support_labels(dispatch_variable_ref(par), supps), expected)
        # dependent parameters
        supps = [Float64[0., 0.], Float64[1., 1.], Float64[NaN, NaN]]
        expected = [Set([UniformGrid]), Set([UniformGrid]), Set{DataType}()]
        @test isequal(IOTO._collected_support_labels(dispatch_variable_ref(pars[1]), supps), expected)
    end
    # test set_parameter_supports
    @testset "set_parameter_supports" begin
        add_supports(par, 0.6, label = InternalLabel)
        @test IOTO.set_parameter_supports(tm, m) isa Nothing
        expected = ([[0., 0.], [1., 1.], [NaN, NaN]], [0., 0.3, 0.6, 0.8, 1., NaN])
        @test isequal(sort.(transcription_data(tm).supports), expected)
        @test IOTO.has_internal_supports(tm)
        expected = ([Set([UniformGrid]), Set([UniformGrid]), Set{DataType}()],
                    [Set([UserDefined]), Set([InternalGaussLobatto]),
                     Set([InternalLabel]), Set([InternalGaussLobatto]),
                     Set([UserDefined]), Set{DataType}()])
        @test isequal(transcription_data(tm).support_labels, expected)
        @test supports(par, label = All) == [0., 0.3, 0.6, 0.8, 1.]
        @test has_generative_supports(par)
    end
    # test parameter_supports
    @testset "parameter_supports" begin
    expected = ([[0., 0.], [1., 1.], [NaN, NaN]], [0., 0.3, 0.6, 0.8, 1., NaN])
        @test isequal(sort.(IOTO.parameter_supports(tm)), expected)
    end
    # test support_index_iterator with 1 argument
    @testset "support_index_iterator (1 Arg)" begin
        @test IOTO.support_index_iterator(tm) == CartesianIndices((1:2, 1:5))
    end
    # test support_index_iterator with 2 argument2
    @testset "support_index_iterator (2 Args)" begin
        @test IOTO.support_index_iterator(tm, [1, 2]) == CartesianIndices((1:2, 1:5))
        @test IOTO.support_index_iterator(tm, [1]) == CartesianIndices((1:2, 6:6))
    end
    # test index_to_support
    @testset "index_to_support" begin
        @test IOTO.index_to_support(tm, first(CartesianIndices((1:2, 1:5)))) isa Vector
        @test isnan(IOTO.index_to_support(tm, last(IOTO.support_index_iterator(tm, [1])))[3])
    end
    # test index_to_labels
    @testset "index_to_labels" begin
        idxs = CartesianIndices((1:2, 1:5))
        @test IOTO.index_to_labels(tm, first(idxs)) == Set([UserDefined, UniformGrid])
        @test IOTO.index_to_labels(tm, idxs[3]) == Set([UniformGrid, InternalGaussLobatto])
        @test IOTO.index_to_labels(tm, last(IOTO.support_index_iterator(tm, Int[]))) == Set{DataType}()
    end
end

# Test the expression mappings
@testset "Expression Queries" begin
    # initialize tbe needed info
    m = InfiniteModel()
    @infinite_parameter(m, pars[1:2] in [0, 1])
    @infinite_parameter(m, par in [0, 1], supports = [0])
    @variable(m, x, Infinite(par, pars))
    @finite_parameter(m, finpar == 42)
    @variable(m, x0, Point(x, 0, [0, 0]))
    @variable(m, y)
    f = parameter_function((a,b) -> 1, (par, pars))
    add_supports(par, 1, label = InternalLabel)
    meas1 = support_sum(2par -2, par)
    meas2 = support_sum(x^2 - y, pars)
    tm = optimizer_model(m)
    @variable(tm, a)
    @variable(tm, b)
    @variable(tm, c)
    @variable(tm, d)
    # transcribe the variables and measures
    data = transcription_data(tm)
    data.has_internal_supports = true
    s1 = Set([UserDefined])
    s2 = Set([InternalLabel])
    data.finvar_mappings[y] = a
    data.finvar_mappings[x0] = b
    data.infvar_mappings[x] = [b, c, d]
    data.infvar_support_labels[x] = [s1, s2, s1]
    data.measure_mappings[meas1] = [-2 * zero(AffExpr)]
    data.measure_support_labels[meas1] = [s1]
    data.measure_mappings[meas2] = [a^2 + c^2 - 2a, b^2 + d^2 - 2a]
    data.measure_support_labels[meas2] = [s1, s2]
    lookups = Dict{Vector{Float64}, Int}([0, 0, 0] => 1, [0, 0, 1] => 2, [1, 1, 1] => 3)
    data.infvar_lookup[x] = lookups
    data.measure_lookup[meas1] = Dict(Float64[] => 1)
    data.measure_lookup[meas2] = Dict{Vector{Float64}, Int}([0] => 1, [1] => 2)
    @test IOTO.set_parameter_supports(tm, m) isa Nothing
    # test transcription_expression in accordance with the methods defined in transcribe.jl
    @testset "transcription_expression (Fallback)" begin
        @test_throws ErrorException IOTO.transcription_expression(tm, a, Float64[])
    end
    # test transcription expression for infinite variables with 3 args
    @testset "transcription_expression (Infinite Variable)" begin
        @test IOTO.transcription_expression(tm, x, [0., 1., 0.]) == c
        @test IOTO.transcription_expression(tm, meas1, [0., 0., 1.]) == -2 * zero(AffExpr)
        @test IOTO.transcription_expression(tm, f, [0., 1., 0.]) == 1
    end
    # test transcription expression for semi_infinite variables with 3 args
    @testset "transcription_expression (Semi-Infinite Variable)" begin
        # semi_infinite of parameter function
        rv = add_variable(m, build_variable(error, f, Dict(1=>1.)),
                          add_support = false)
        @test IOTO.transcription_expression(tm, rv, [0., 1., 0.]) == 1
        # semi_infinite of infinite variable
        rv = add_variable(m, build_variable(error, x, Dict(1=>1.)),
                          add_support = false)
        data.infvar_mappings[rv] = [b, c]
        lookups = Dict{Vector{Float64}, Int}([0, 0] => 1, [1, 0] => 2)
        data.infvar_lookup[rv] = lookups
        @test IOTO.transcription_expression(tm, rv, [1., 0., 0.]) == c
    end
    # test transcription expression for finite variables with 3 args
    @testset "transcription_expression (Finite Variable)" begin
        @test IOTO.transcription_expression(tm, x0, [0., 1., 0.]) == b
        @test IOTO.transcription_expression(tm, y, [0., 0., 1.]) == a
    end
    # test transcription expression for infinite parameters with 3 args
    @testset "transcription_expression (Infinite Parameter)" begin
        @test IOTO.transcription_expression(tm, par, [0.5, 1., 0.]) == 0.
        @test IOTO.transcription_expression(tm, pars[1], [0.5, 0., 1.]) == 0.5
    end
    # test transcription expression for finite parameters with 3 args
    @testset "transcription_expression (Finite Parameter)" begin
        @test IOTO.transcription_expression(tm, finpar, [0.5, 1., 0.]) == 42
    end
    # test transcription expression for AffExprs with 3 args
    @testset "transcription_expression (AffExpr)" begin
        @test IOTO.transcription_expression(tm, x0 - y + 2x - 2.3, [1., 1., 1.]) == b - a + 2d - 2.3
    end
    # test transcription expression for QuadExprs with 3 args
    @testset "transcription_expression (QuadExpr)" begin
        # test normal
        expr = meas2 - 3y^2 - x0 - 2.3
        expected = b^2 + d^2 - 2a - 3a^2 - b - 2.3
        @test IOTO.transcription_expression(tm, expr, [1., 1., 1.]) == expected
        # test becomes a nonlinear expression
        expr = meas2 * x0
        @test IOTO.transcription_expression(tm, expr, [1., 1., 1.]) isa NonlinearExpression
        @test nl_expr_string(tm, REPLMode, tm.nlp_data.nlexpr[end]) == "+((+(-2.0 * a) + b * b + d * d) * b)"
    end
    # test transcription expression for NLPExprs with 3 args
    @testset "transcription_expression (NLPExpr)" begin
        expr = sin(y)
        @test IOTO.transcription_expression(tm, expr, [1., 1., 1.]) isa NonlinearExpression
        @test nl_expr_string(tm, REPLMode, tm.nlp_data.nlexpr[end]) == "sin(a)"
    end
    # test transcription expression for numbers with 3 args
    @testset "transcription_expression (Real)" begin
        expected = zero(AffExpr) + 42
        @test IOTO.transcription_expression(tm, 42, [1., 1., 1.]) == expected
    end
    # test transcription expression for Exprs with 2 args
    @testset "transcription_expression (Expr 2 Args)" begin
        # test infinite expr
        expr = meas2 - 3y^2 - x0 - 2.3
        expected = [-2a^2 + c^2 - 2a - b - 2.3, -3a^2 + b^2 + d^2 - 2a - b - 2.3]
        @test IOTO.transcription_expression(tm, expr, label = All) == expected
        @test IOTO.transcription_expression(tm, expr, label = All, ndarray = true) == expected
        expected = [-2a^2 + c^2 - 2a - b - 2.3]
        @test IOTO.transcription_expression(tm, expr) == expected
        @test IOTO.transcription_expression(tm, expr, ndarray = true) == expected
        # test finite expr
        expr = 2x0 -y
        expected = 2b- a
        @test IOTO.transcription_expression(tm, expr) == expected
        @test IOTO.transcription_expression(tm, expr, ndarray = true) == [expected]
        # test NLPExpr
        @test IOTO.transcription_expression(tm, sin(x0)) isa NonlinearExpression
        @test nl_expr_string(tm, REPLMode, tm.nlp_data.nlexpr[end]) == "sin(b)"
    end
    # test transcription expression for variables with 2 args
    @testset "transcription_expression (Variable 2 Args)" begin
        @test IOTO.transcription_expression(tm, x0) == b
        @test IOTO.transcription_expression(tm, x) == [b, d]
        @test IOTO.transcription_expression(tm, x, label = All) == [b, c, d]
    end
    # test transcription expression with 1 argument
    @testset "transcription_expression (1 Arg)" begin
        @test IOTO.transcription_expression(x0) == b
        @test IOTO.transcription_expression(x0 - y) == b - a
        @test IOTO.transcription_expression(zero(QuadExpr) + 2) == zero(AffExpr) + 2
    end
    # test optimizer_model_expression
    @testset "optimizer_model_expression" begin
        @test optimizer_model_expression(x0) == b
        @test optimizer_model_expression(x0 - y) == b - a
        @test optimizer_model_expression(zero(QuadExpr) + 2) == zero(AffExpr) + 2
    end
    # test expression_supports
    @testset "expression_supports" begin
        @test supports(x0) == ()
        @test supports(x0 - y) == ()
        @test supports(x0 - y, ndarray = true) == [()]
        @test supports(meas2 + y) == [(0.,)]
        @test supports(meas2 + y, ndarray = true) == [(0.,)]
        @test supports(meas2 + y, label = All) == [(0.,), (1.,)]
        @test supports(meas2 + y, label = All, ndarray = true) == [(0.,), (1.,)]
    end
end

# Test constraint mapping queries
@testset "Constraint Mapping Queries" begin
    # initialize models
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1], supports = [0])
    add_supports(par, 1, label = InternalLabel)
    @variable(m, x, Infinite(par))
    @variable(m, x0, Point(x, 0))
    @variable(m, y)
    @constraint(m, c1, x + y - 2 <= 0)
    @constraint(m, c2, support_sum(x, par) == 0)
    @constraint(m, c3, x0 + y == 5)
    tm = optimizer_model(m)
    data = transcription_data(tm)
    data.has_internal_supports = true
    s1 = Set([UserDefined])
    s2 = Set([InternalLabel])
    @variable(tm, a)
    @variable(tm, b)
    @variable(tm, c)
    @constraint(tm, tc1, b + a <= 2)
    @constraint(tm, tc2, c + a <= 2)
    @constraint(tm, tc3, 0.5b + 0.5c == 0)
    @constraint(tm, tc4, b + a == 5)
    @test IOTO.set_parameter_supports(tm, m) isa Nothing
    # test transcription_constraint (2 Args)
    @testset "transcription_constraint (Infinite)" begin
        # test error
        @test_throws ErrorException transcription_constraint(tm, c1)
        # test normal
        data.constr_mappings[c1] = [tc1, tc2]
        data.constr_support_labels[c1] = [s1, s2]
        @test transcription_constraint(tm, c1, label = All) == [tc1, tc2]
        @test transcription_constraint(tm, c1) == [tc1]
        @test transcription_constraint(tm, c1, label = All, ndarray = true) == [tc1, tc2]
        @test transcription_constraint(tm, c1, ndarray = true) == [tc1]
        # test error
        @test_throws ErrorException transcription_constraint(tm, c2)
        # test normal
        data.constr_mappings[c2] = [tc3]
        data.constr_support_labels[c2] = [s1]
        @test transcription_constraint(tm, c2) == tc3
        @test transcription_constraint(tm, c2, ndarray = true) == [tc3]
        # test error
        @test_throws ErrorException transcription_constraint(tm, c3)
        # test normal
        data.constr_mappings[c3] = [tc4]
        data.constr_support_labels[c3] = [s1]
        @test transcription_constraint(tm, c3) == tc4
    end
    # test transcription_constraint (Single argument)
    @testset "transcription_constraint (1 Arg)" begin
        @test transcription_constraint(c1) == [tc1]
        @test transcription_constraint(c1, label = All) == [tc1, tc2]
        @test transcription_constraint(c2) == tc3
        @test transcription_constraint(c3, label = All) == tc4
    end
    # test optimizer_model_constraint extension
    @testset "optimizer_model_constraint" begin
        @test optimizer_model_constraint(c1, Val(:TransData), label = All) == [tc1, tc2]
        @test optimizer_model_constraint(c2, Val(:TransData)) == tc3
        @test optimizer_model_constraint(c3, Val(:TransData)) == tc4
    end
    # test constraint_supports
    @testset "constraint_supports" begin
        # test error
        @test_throws ErrorException InfiniteOpt.constraint_supports(tm, c1)
        # test normal
        data.constr_supports[c1] = [(0.,), (1.,)]
        @test InfiniteOpt.constraint_supports(tm, c1) == [(0.,)]
        @test InfiniteOpt.constraint_supports(tm, c1, label = All) == [(0.,), (1.,)]
        @test InfiniteOpt.constraint_supports(tm, c1, ndarray = true) == [(0.,)]
        @test InfiniteOpt.constraint_supports(tm, c1, label = All, ndarray = true) == [(0.,), (1.,)]
        # test finite
        data.constr_supports[c3] = [()]
        @test InfiniteOpt.constraint_supports(tm, c3) == ()
        @test InfiniteOpt.constraint_supports(tm, c3, ndarray = true) == [()]
    end
    # test supports
    @testset "supports" begin
        @test supports(c1, label = All) == [(0.,), (1.,)]
        @test supports(c1, label = All, ndarray = true) == [(0.,), (1.,)]
        @test supports(c1, label = InternalLabel) == [(1.,)]
        @test supports(c3) == ()
    end
end


text:
#=
Given an array of integers, return a new array such that each element at index i of the new array is the product of all the numbers in the original array except the one at i.

For example, if our input was [1, 2, 3, 4, 5], the expected output would be [120, 60, 40, 30, 24]. If our input was [3, 2, 1], the expected output would be [2, 3, 6].

Follow-up: what if you can't use division?
=#

using Test
include("Solutions/problem2_product_except_i.jl")
@test product_excepti([1, 2, 3, 4, 5]) == [120, 60, 40, 30, 24]
@test product_excepti([3, 2, 1]) == [2, 3, 6]

max_stars_repo_name: DominiqueCaron/daily-coding-problem

max_issues_repo_name: DominiqueCaron/daily-coding-problem

max_forks_repo_name: DominiqueCaron/daily-coding-problem


text:
export
    DDproblem,
    value_iteration_step,
    transition_matrix,
    policy_iteration_step,
    DDsolution,
    solve_iteratively,
    simulate_transitions

# discrete state, discrete time dynamic programming

"""
Discrete state, discrete time dynamic programing problem with `N`
states and `M` actions.

`β` is the discount factor (infinite horizon).

`U` is an (N,M) matrix that describes the period payoffs for state `N`
and action `M`. For actions that are now allowed, use `-Inf`.

`P` is an `N`-element vector of `(M,N)` transition matrices, ie
`P[n][m,n′]` is the probability of state `n′` given state `n` and
action `m`.
"""
immutable DDproblem{Tβ <: Real, TU <: AbstractMatrix, TP <: AbstractMatrix}
    β::Tβ
    U::TU
    P::Vector{TP}
    function DDproblem(β, U, P)
        (N, M) = size(U)
        @assert length(P) == N "wrong number of elements in P"
        @assert all(size(p) == (M,N) for p in P) "non-conformable matrix in P"
        new(β, U, P)
    end
end

DDproblem{Tβ, TU, TP}(β::Tβ, U::TU, P::Vector{TP}) = DDproblem{Tβ, TU, TP}(β, U, P)

function show(io::IO, dd::DDproblem)
    (N,M) = size(dd.U)
    println(io, "Discrete state, discrete time dynamic programming problem")
    println(io, "  with $(N) states and $(M) actions, discount factor $(dd.β)")
end

"""
Value iteration step. Given value function `V`, return the new value
function and the indexes of optimal policies.
"""
function value_iteration_step(dd::DDproblem, V)
    @unpack β, U, P = dd
    optimal_Vms = [findmax(@view(U[n, :]) + β*P[n]*V) for n in indices(U, 1)]
    (first.(optimal_Vms), last.(optimal_Vms))
end

"""
Return a transition matrix between states given the optimal policies.

It will be the same type as the elements of `dd.P`, tus
"""
function transition_matrix(dd::DDproblem, optimal_ms)
    @unpack U, P = dd
    N = size(U, 1)
    M = zeros(N, N)
    for (n,m) in enumerate(optimal_ms)
        M[n, :] = P[n][m, :]
    end
    convert(eltype(P), M)
end

"""
Policy iteration step. Given a value function `V`, return the new
value function and the indexes of optimal policies.
"""
function policy_iteration_step(dd::DDproblem, V)
    optimal_ms = value_iteration_step(dd, V)[2]
    u = [dd.U[n,m] for (n,m) in enumerate(optimal_ms)]
    M = transition_matrix(dd, optimal_ms)
    V = (I - dd.β * M) \ u
    (V, optimal_ms)
end

"""
Structure for storing the result of iterative solution techniques.

`V` is the value function.
`ms` is the index of actions (only optimal when converged).
`iterations` is the number of iterations performed.
`converged` is a boolean which is `true` when convergence is achieved.
"""
immutable DDsolution
    V
    ms
    iterations
    converged
end

function show(io::IO, dd::DDsolution)
    print(io, "DDsolution ")
    if dd.converged
        print_with_color(:green, io, "converged")
    else
        print_with_color(:red, io, "did not converge")
    end
    println(io, " after $(dd.iterations) iterations")
end

"""
Solve a problem using value and policy iteration, using initial values
`V`. Successive value functions are compared with the L¹ norm,
convergence is achieved the difference between value functions is
below `ϵ`.

For the first `value_iter` steps, use value iteration, and after that,
policy iteration.
"""
function solve_iteratively(dd::DDproblem;
                           V = zeros(size(dd.U, 1)),
                           value_iter = 20,
                           max_iter = 100,
                           ϵ = √eps(Float64))
    for iteration in 1:max_iter
        (V′, ms) = iteration ≤ value_iter ? value_iteration_step(dd, V) :
            policy_iteration_step(dd, V)
        if norm(V′-V, 1) ≤ ϵ
            return DDsolution(V′, ms, iteration, true)
        elseif iteration == max_iter
            return DDsolution(V′, ms, iteration, false)
        else
            V .= V′
        end
    end
end

"""
Given a transition matrix `M` between states, an initial state `n0`,
return a sequence of `T` simulated transitions.
"""
function simulate_transitions(M, n0, T)
    ps = [Categorical(collect(M[n, :])) for n in indices(M, 1)]
    n = n0
    [(n_prev = n; n = rand(ps[n]); n_prev) for i in 1:T]
end


text:
using SparseArrays
using LinearAlgebra

CSCMAT_FORMAT_VERSION = v"0.1.0"  # track version of our custom CSCMAT file format.

"""
write the three arrays defining compressed sparse column (CSC) storage of a matrix into a file.

If `try_hex`, integers in arrays are stored as hexadecimals (without 0x prefix!)
If `allow_omit_entries_if_only_stored_ones`, the `stored values` array is omitted if all stored values compare equal to 1.
"""
function save_to_cscmat(
    mat::SparseMatrixCSC, destination_file_path::String
    ;
    additional_header_lines="",
    try_hex::Bool=false,
    allow_omit_entries_if_only_stored_ones=false,
    )

    try
        additional_header_lines = "#"*join(split(additional_header_lines, "\n"), "\n# ")
    catch
        @warn "Failed to process additional header lines. Discarding them."
        additional_header_lines = ""
    end

    if file_extension(destination_file_path) != ".cscmat"
        @warn "Writing to sparse column storage file with with extension
            '$(file_extension(destination_file_path))', expected '.cscmat'. (path '$(destination_file_path)')"
    end

    number_map(x::Real) = x
    number_map(n::Integer) = string(n, base=try_hex ? 16 : 10)

    open(destination_file_path, "w+") do file
        println(file, "# $CSCMAT_FORMAT_VERSION")
        println(file, "# Compressed sparse column storage of matrix (arrays `colptr`, `rowval`, `stored_values`"
            *" as space separated $(try_hex ? "hexadecimal" : "decimal") integers. Stored entries may be zero.).")
        println(file, additional_header_lines)
        println(file, "# n_rows n_columns n_stored_entries")
        println(file, "$(mat.m) $(mat.n) $(nnz(mat))\n")

        println(file, join(number_map.(mat.colptr .- 1), " "))  # convert to zero-based indices
        println(file, "")
        println(file, join(number_map.(mat.rowval .- 1), " "))  # convert to zero-based indices
        println(file, "")

        if !allow_omit_entries_if_only_stored_ones || any(x->x!=1, mat.nzval)
            println(file, join(number_map.(mat.nzval), " "))
        end
    end

    return nothing
end

function read_file_header(file_path; comment_marker = '#')
    header = ""
    open(file_path, "r") do file
        next_line = ""
        while true
            header *= (next_line*"\n")
            next_line = readline(file)

            (length(next_line) > 0 && next_line[1] == comment_marker) || break
        end
    end

    return header
end

"""
read the three arrays defining compressed sparse column (CSC) storage of a matrix into a file.

If `try_hex`, integers are stored as hexadecimals (without 0x prefix!)
"""
function load_cscmat(file_path::String;
    print_file_header=false)
    expected_file_extension = ".cscmat"
    if file_extension(file_path) != expected_file_extension
        @warn "load_cscmat called on file '$(file_path)'
            with extension '$(file_extension(file_path))', expected $expected_file_extension"
    end

    file = open(file_path, "r")
    header = ""
    next_line = ""
    while true
        header *= (next_line*"\n")
        next_line = readline(file)

        (length(next_line) > 0 && next_line[1] == '#') || break
    end

    try
        header[1] == "# $CSCMAT_FORMAT_VERSION" && @warn "File written in format $(header[1][2:end]) is being read in format $CSCMAT_FORMAT_VERSION"
    catch e
        @warn "Failed to verify CSC file format version: $e"
    end

    if contains(header, "hexadecimal")
        base = 16
    else
        base = 10
    end

    print_file_header && print(header)

    n_rows, n_cols, n_nnz = space_sep_ints(next_line)
    _ = readline(file)  # empty line

    colptr = space_sep_ints(readline(file); base)
    _ = readline(file)  # empty line

    rowval = space_sep_ints(readline(file); base)
    _ = readline(file)  # empty line

    stored_entries = space_sep_ints(readline(file); base)
    _ = readline(file)  # empty line
    remaining_lines = readlines(file)
    close(file)

    if length(remaining_lines) > 0
        @warn "Ignoring additional lines:\n`$remaining_lines`"
    end

    if length(stored_entries) == 0
        stored_entries = ones(Int8, n_nnz)
    end

    # convert from zero-based to one-based indices.
    return SparseMatrixCSC(n_rows, n_cols, colptr .+1, rowval .+1, stored_entries)
end

"""
convert matrix of exponents for QC LDPC matrix to the actual binary LDPC matrix.

The resulting QC-LDPC matrix is a block matrix where each block is either zero,
or a circ-shifted identity matrix of size `expansion_factor`x`expansion_factor`.
Each entry of the matrix Hqc denotes the amount of circular shift in the QC-LDPC matrix.
No entry (i.e., a zero but not a stored one) at a given position in Hqc means the associated block is zero.
"""
function Hqc_to_pcm(
    Hqc::SparseMatrixCSC{T,Int} where T <: Integer,
    expansion_factor::Integer,
    )
    scale_idx(idx::Integer, expansion_factor::Integer) = (idx - 1) * expansion_factor + 1
    shifted_identity(N::Integer, shift::Integer, scalar_one=Int8(1)) = circshift(Matrix(scalar_one*I, N, N), (0, shift))

    H = spzeros(Int8, size(Hqc, 1) * expansion_factor, size(Hqc, 2) * expansion_factor)

    Is, Js, Vs = findnz(Hqc)

    for (i, j, v) in zip(Is, Js, Vs)
        i_start = scale_idx(i, expansion_factor)
        i_end = scale_idx(i+1,expansion_factor) - 1
        j_start = scale_idx(j, expansion_factor)
        j_end = scale_idx(j+1, expansion_factor) - 1
        H[i_start:i_end, j_start:j_end] = shifted_identity(expansion_factor, v, Int8(1))
    end

    return H
end

"""
Load exponents for a QC-LDPC matrix from a `.CSCMAT` file and return the binary LDPC matrix.

Not every input `.cscmat` file will give a meaninful result.
The `.cscmat` format allows to store general sparse matrices in text format.
Meanwhile, this function expects that the file stores exponents for a quasi-cyclic LDPC matrix.
The exponent matrix is read and expanded using the expansion factor.

If the expansion factor is not provided, the CSCMAT file must contain a line specifying it.
"""
function load_matrix_from_qc_cscmat_file(file_path::AbstractString; expansion_factor=nothing)
    if isnothing(expansion_factor)
        header = ""
        next_line = ""

        open(file_path, "r") do f
            while true
                header *= (next_line * "\n")
                next_line = readline(f)

                (length(next_line) > 0 && next_line[1] == '#') || break
            end
        end

        m = match(r"Quasi cyclic exponents for a binary LDPC matrix with expansion factor ([0-9]*)\.", header)
        if isnothing(m)
            error("Failed to infer expansion factor! No header line found containing it.")
        else
            expansion_factor = parse(Int, m.captures[1])
            @info "Inferred expansion factor from file header: $expansion_factor"
        end
    end

    Hqc = load_cscmat(file_path)

    H = Hqc_to_pcm(Hqc, expansion_factor)

    return H
end


text:
# FunctionalForm.jl - explicit form of pair-potentials

# Goodwin--Skinner--Pettifor
macro GSP(A,n, nc,r0,rc)
    return :(r -> $A*($r0/r)^$n * exp( $n*(-(r/$rc)^$nc + ($r0/$rc)^$nc)) )
end

# Exponential x power law
#   Written as a crazy summation expansion for the joy of it.
macro EPL(r0,An,mn,pn)
    sum=:(0.0)
    for (A,m,p) in zip(An.args,mn.args,pn.args)
        sum=:($sum + $A * ($r0/r)^$m * exp(-$p *(r-$r0)) )
    end
    println("sum: ",sum)
    return :(r -> $sum)
end

# Chadi's quadratic
macro QUAD(r0,A1,A2)
    return :(r -> $A1*(r-$r0) + $A2*(r-$r0)*(r-$r0) )
end


text:
module UnROOT

using LazyArrays
export ROOTFile, LazyBranch, LazyTree

import Base: close, keys, get, getindex, getproperty, show, length, iterate, position, ntoh, lock, unlock, reinterpret
ntoh(b::Bool) = b

import AbstractTrees: children, printnode, print_tree

using CodecLz4, CodecXz, CodecZstd, StaticArrays, LorentzVectors, ArraysOfArrays
using Mixers, Parameters, Memoization, LRUCache
import IterTools: groupby

import LibDeflate: zlib_decompress!, Decompressor

import Tables, TypedTables, PrettyTables

@static if VERSION < v"1.6"
    Base.first(a::AbstractVector{S}, n::Integer) where S<: AbstractString = a[1:(length(a) > n ? n : end)]
    Base.first(a::S, n::Integer) where S<: AbstractString = a[1:(length(a) > n ? n : end)]
end

include("constants.jl")
include("io.jl")
include("types.jl")
include("utils.jl")
include("streamers.jl")
include("bootstrap.jl")
include("root.jl")
include("iteration.jl")
include("custom.jl")
include("displays.jl")

end # module


text:
"""
    MeasureSNR(signal, noisy; db=false)

Measure the signal-to-noise ratio between the clean input `signal` and the
contaminated input `noisy`.

# Arguments
- `signal::Array{Real, N}`: N-dimensional clean signal. `N` must be <= 5.
- `noisy::Array{Real, N}`: N-dimensional noisy signal of same size as `signal`.
- `db::Bool=false`: Flag is false if the signal-to-noise ratio is measured by amplitude. Flag is true if snr is in dB.

# Example
```julia
julia> d = SeisHypEvents(); dnoisy = SeisAddNoise(d, 2);
MeasureSNR(d, dnoisy)
```
"""
function MeasureSNR(signal::Array{Ts, N},
            noisy::Array{Tn, N}; db::Bool=false) where {Ts <: Real, Tn <: Real, N}

    if db==false
        snr = norm(signal)/norm(signal-noisy)
    elseif db==true
        snr = 20*log10(norm(signal)/norm(signal-noisy))
    end
    return snr

end


text:
using debspkg
using Documenter

makedocs(;
    modules=[debspkg],
    authors="Debraj bhal",
    repo="https://github.com/krishna-vasudev/debspkg.jl/blob/{commit}{path}#L{line}",
    sitename="debspkg.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://krishna-vasudev.github.io/debspkg.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/krishna-vasudev/debspkg.jl",
)


text:
push!(LOAD_PATH, "../src/")

using Documenter, EmpiricalOrthogonalFunctions

pages = ["Home" => "index.md", "API" => "api.md"]

makedocs(;
    modules = [EmpiricalOrthogonalFunctions],
    authors = "Kel Markert",
    repo = "https://github.com/KMarkert/EmpiricalOrthogonalFunctions.jl/blob/{commit}{path}#L{line}",
    sitename = "EmpiricalOrthogonalFunctions.jl",
    pages = pages,
)

deploydocs(;
    repo = "github.com/KMarkert/EmpiricalOrthogonalFunctions.jl.git",
    devbranch = "main",
)


text:
################################################################################
# Player
################################################################################

mutable struct Player{T,SMQ,SMR,SVx,SVu}
	x0::SVx        # Initial state
	id::Int # ID of the player
    lane_id::Int   # ID of the lane the vehicle drives in
    Q::SMQ         # LQR Cost Q
    R::SMR         # LQR Cost R
    xf::SVx        # LQR Cost xf
    uf::SVu        # LQR Cost uf
    u_min::SVu     # Bound constraint on the controls
    u_max::SVu     # Bound constraint on the controls
	v_min::T       # Bound constraint on the velocity
	v_max::T       # Bound constraint on the velocity
    r_col::T       # Collision radius of the car
    r_cost::T      # Radius of the collision avoidance cost
    μ::T           # Amplitude of the collision avoidance cost
end

function Player(model::AbstractGameModel, lane::Lane{T};
				   id::Int=1,
				   x0::VehicleState=lane.start.x_nom,
				   xf::VehicleState=lane.start.x_max,
				   uf::SVector{mi,T}=zeros(SVector{model.mi[1],T}),
				   Q::Diagonal{T,SVector{ni,T}}=Diagonal(1.0*ones(SVector{model.ni[1],T})),
				   R::Diagonal{T,SVector{mi,T}}=Diagonal(0.1*ones(SVector{model.mi[1],T})),
				   r_col::T=0.08,
				   r_cost::T=0.22,
				   μ::T=20.0,
				   u_min::SVector{mi,T}=-Inf*ones(SVector{model.mi[1],T}),
				   u_max::SVector{mi,T}=Inf*ones(SVector{model.mi[1],T}),
				   v_min::T=-Inf,
				   v_max::T=Inf,
		) where {ni,mi,T}

	x0_ = specialize(model, x0)
	xf_ = specialize(model, xf)

	player = Player(
		x0_,
		id,
		lane.id,
		Q,
		R,
		xf_,
		uf,
		u_min,
		u_max,
		v_min,
		v_max,
		r_col,
		r_cost,
		μ,
		)
	return player
end


text:
let sys = DiagonalSystem(test_params_diag...),
    beta = 12.34,
    analytical = Analytical(sys, beta)

    @test isapprox(analytical.Z, 6.92896484e-9)
    @test isapprox(analytical.E, 1.52249605e0)
    @test isapprox(analytical.Cv, 5.35769920e-4)
end


text:
abstract type AbstractMarkovChain end

function check(state_space, transition_matrix, type)
    if length(state_space) != size(transition_matrix)[1]
        error(
            "The state space, $(state_space), and "*
            "transition matrix should be the same size."
        )
    end

    if length(unique(state_space)) != length(state_space)
        error("The state space, $(state_space), must have unique elements.")
    end

    rrs = required_row_sum(type)
    try
        rrs = required_row_sum(type, eltype(transition_matrix))
    catch e
    end
    if !is_row_stochastic(transition_matrix, rrs)
        error(
            "The transition matrix, $(transition_matrix), should be row-stochastic "*
            "(each row must sum up to $(required_row_sum(type)))."
        )
    end
end

"""
    state_space(x)

# Definitions
The state space of a Markov chain is the (ordered)
set of values that the process is able to take on.
For example, in a Sunny-Cloudy-Rainy weather model,
the state space is `["Sunny", "Cloudy", "Rainy"]`.

# Arguments
- `x`: some kind of Markov chain.

# Returns
The state space of the Markov chain.
"""
state_space(x::AbstractMarkovChain) = x.state_space

"""
    transition_matrix(x)

# Definitions
The one-step transition matrix, ``T``, of a discrete Markov chain
or the generator matrix of a continuous Markov chain.

# Arguments
- `x`: some kind of Markov chain.

# Returns
The transition matrix of the Markov chain.
"""
transition_matrix(x::AbstractMarkovChain) = x.transition_matrix

"""
    probability_matrix(x)

# Definitions
The one-step transition matrix, ``T``, of a Markov chain, ``\\{X_t\\}``
is a matrix whose ``(i,j)``th entry is the probability of the process
being in state ``j`` at time 1 given that the process started
in state ``i`` at time 0. That is
```math
T = \\{\\{p^{(1)}_{i,j}\\}\\}_{i,j ∈ S} = \\{\\{\\mathbb{P}(X_1=j | X_0=i)\\}\\}_{i,j ∈ S}
```
where ``S`` is the state space of ``\\{X_t\\}``.

# Arguments
- `x`: some kind of Markov chain.

# Returns
The one-step probability matrix of the Markov chain.
"""
probability_matrix(x::AbstractMarkovChain) = x.transition_matrix

"""
    embedded(x)

# Arguments
- `x`: some kind of Markov chain.

# Returns
If the Markv chain is continuous, then it returns the
embedded Markov chain. If the Markov chain is discrete,
then it returns the given chain, `x`.

# Notes
If the equivalent chain is preffered rather than the
embedded chain, then use `DiscreteMarkovChain(x)`.

"""
embedded(x::AbstractMarkovChain) = x

"""
    state_index(x)

# Arguments
- `x`: some kind of Markov chain.

# Returns
A dictionary mapping each state in a Markov chain to its position in the state space.
It is essentially the inverse of `state_space(x)`.
"""
state_index(x::AbstractMarkovChain) = Dict(
    state => index for (index, state) in enumerate(state_space(x))
)

"""
    digraph(x)

Creates a digraph (directed graph) representation of a Markov chain.

# Arguments
- `x`: some kind of Markov chain.

# Returns
A 1D array of 2-tuples. An element ``(i, j)`` is in the array
iff the transition matrix at ``(i,j)`` is nonzero.
"""
function digraph(x::AbstractMarkovChain)
    T = transition_matrix(x)
    V = 1:(size(T)[1])
    return [(i, j) for i in V for j in V if T[i, j] != zero(eltype(T))]
end

"""
    is_row_stochastic(mat, row_sum=1)

Tests whether a matrix, x, is row stochasitc.
The desired sum of each row can be specified as well.

# Definitions
A matrix is said to be row stochasic if all its rows sum to 1.
This definition is extened so that all its rows sum to `row_sum`.

# Arguments
- `mat`: a matrix that we want to check.
- `row_sum`: the desired value that each row should total to.

# Returns
`true` if the given matrix, `mat`, is row-stochasitc.
"""
function is_row_stochastic(mat, row_sum=one(eltype(mat)))
    n, p = size(mat)
    if n == 0
        return true
    end

    # We only want to throw when we are sure it is incorrect.
    # Float16 seemed like a good compromise.
    lhs = repeat([row_sum], n)
    rhs = mat*repeat([one(eltype(mat))], p)

    for i in 1:p
        if !isapprox(lhs[i], rhs[i], atol=eps(Float16), rtol=eps(Float16))
            return false
        end
    end
    return true
end

Base.one(::Type{Array{T,2}}) where T = LinearAlgebra.I
Base.one(::Type{LinearAlgebra.Adjoint{T,Array{T,2}}}) where T = LinearAlgebra.I
Base.zero(::Type{Array{T,2}}) where T = LinearAlgebra.UniformScaling(0)
function Base.zero(::Type{LinearAlgebra.Adjoint{T,Array{T,2}}}) where T
    return LinearAlgebra.UniformScaling(0)
end

"""
    communication_classes(x)

# Definitions
A state ``j`` is accessible from state ``i`` if it is possible
to eventually reach state ``j`` given that the process started
in state ``i``. That is, ``∃ \\, t ∈ \\mathbb{N}`` such that
``p^{(t)}_{i,j} > 0``.

States ``i`` and ``j`` communicate if ``i`` is accessible
from ``j`` and ``j`` is accessible from ``i``.

A communication class is the set of states in a Markov chain
that are all accessible from one another. Communication classes
form a class in the mathematical sense. They also form a
partition of the state space.

A state ``i`` is recurrent if the process is guarenteed to
eventually return to state ``i`` given that the process
started in state ``i``. If a state ``i`` is not recurrent,
then it is said to be transient.

# Arguments
- `x`: some kind of Markov chain.

# Returns
A tuple containing 2 arrays.
- This first array contains C arrays which store the states that communicate.
- The second array is an array of Bool where the ith value is true if the
  ith communication class is recurrent.

# Examples
```jldoctest communication_classes
using DiscreteMarkovChains
T = [
    1 0;
    0 1;
]
X = DiscreteMarkovChain(T)

communication_classes(X)

# output

([[1], [2]], Any[true, true])
```

So the Markov chain has two communication classes and both are recurrent.

```jldoctest communication_classes
T = [
    .5 .5 .0;
    .2 .8 .0;
    .1 .2 .7;
]
X = DiscreteMarkovChain(["Sunny", "Cloudy", "Rainy"], T)

communication_classes(X)

# output

([["Sunny", "Cloudy"], ["Rainy"]], Any[true, false])
```

So the Sunny and Cloudy states communicate and are recurrent.
The Rainy state is transient.
Note that this is not a very good weather model since once it
stops raining, it will never rain again.
"""
function communication_classes(x::AbstractMarkovChain)
    S = state_space(x)
    T = transition_matrix(x)

    V = 1:length(S)
    E = digraph(x)

    classes = strongly_connected_components(V, E)
    recurrence = []

    for class in classes
        submatrix = T[class, class]
        push!(recurrence, is_row_stochastic(
            submatrix,
            required_row_sum(typeof(x), eltype(submatrix))
        ))
    end

    classes = [[S[i] for i in class] for class in classes]
    return classes, recurrence
end

"""
    decompose(x)

All transition matrices can be reordered so that we have
```math
T = \\begin{pmatrix}
A & 0\\\\
B & C
\\end{pmatrix}
```
Where ``A``, ``B`` and ``C`` are as described below.
``0`` is a matrix of zeros.

# Arguments
- `x`: some kind of Markov chain.

# Returns
A tuple of four values
- `states`: the first value is an array of the new states.
- `A`: the second value is a matrix of recurrent states to recurrent states.
- `B`: the third value is a matrix of transient states to recurrent states.
- `C`: the fourth value is a matrix of transient to transient states.

# Examples
```jldoctest decompose
using DiscreteMarkovChains
T = [
    0.0 1.0 0.0;
    1.0 0.0 0.0;
    0.1 0.2 0.7;
]
X = DiscreteMarkovChain(["Sunny", "Cloudy", "Rainy"], T)

decompose(X)

# output

(Any["Sunny", "Cloudy", "Rainy"], [0.0 1.0; 1.0 0.0], [0.1 0.2], [0.7])
```
"""
function decompose(x::AbstractMarkovChain)
    T = transition_matrix(x)
    classes, recurrence = communication_classes(x)
    r_states = []
    t_states = []
    for (i, recurrent) in enumerate(recurrence)
        if recurrent
            append!(r_states, classes[i])
        else
            append!(t_states, classes[i])
        end
    end

    states = vcat(r_states, t_states)
    indexes = [state_index(x)[state] for state in states]

    A = [
        T[indexes[i], indexes[j]]
        for i in 1:length(r_states), j in 1:length(r_states)
    ]
    B = [
        T[indexes[length(r_states)+i], indexes[j]]
        for i in 1:length(t_states), j in 1:length(r_states)
    ]
    C = [
        T[indexes[length(r_states)+i], indexes[length(r_states)+j]]
        for i in 1:length(t_states), j in 1:length(t_states)
    ]

    return states, A, B, C
end

"""
    canonical_form(x)

Reorders the states of the transition matrix of `x` so that
recurrent states appear first and transient states appear last.

# Arguments
- `x`: some kind of Markov chain.

# Returns
A tuple with the new states and the new transition matrix.

# Examples
We will use the same matrix as previous examples but change the order of it.

```jldoctest canonical_form
using DiscreteMarkovChains
T = [
    0.7 0.2 0.1;
    0.0 0.0 1.0;
    0.0 1.0 0.0;
]
X = DiscreteMarkovChain(["Rainy", "Cloudy", "Sunny"], T)

canonical_form(X)

# output

(Any["Cloudy", "Sunny", "Rainy"], [0.0 1.0 0.0; 1.0 0.0 0.0; 0.2 0.1 0.7])
```

Here, Cloudy and Sunny are recurrent states so they appear first.
Rainy is a transient state so it appears last.
"""
function canonical_form(x::AbstractMarkovChain)
    states, A, B, C = decompose(x)
    O = zeros(Int, size(A)[1], size(C)[2])
    result = vcat(
        hcat(A, O),
        hcat(B, C),
    )
    return states, result
end

"""
    is_ergodic(x)

# Definitions
A Markov chain is called an ergodic chain if it is possible to go
from every state to every state (not necessarily in one move).
That is, every state communicates and the whole tranistion
matrix forms one communication class.

In many books, ergodic Markov chains are called irreducible.

If a Markov chain is not irreducible then it is reducible.
This means that the Markov chain can be broken down into
smaller irreducible chains that do not communicate. One
chain might still be accessible from another though.

# Arguments
- `x`: some kind of Markov chain.

# Returns
`true` if the Markov chain, `x`, is ergodic.
This is, if every state can be accessed from every other state.
Another term for this is irreducible.

# Examples
We will set up a matrix with 2 communication classes and
show that it is not ergodic.

```jldoctest is_ergodic
using DiscreteMarkovChains
T = [
    0 1 0;
    1 0 0;
    0 0 1;
]
X = DiscreteMarkovChain(T)

is_ergodic(X)

# output

false
```

Repeat the above but now all states communicate.

```jldoctest is_ergodic
T = [
    0.0 0.5 0.5;
    0.0 0.0 1.0;
    1.0 0.0 0.0;
]
X = DiscreteMarkovChain(T)

is_ergodic(X)

# output

true
```

Notice how a periodic chain is regular no matter the periodicity.

```jldoctest is_ergodic
T = [
    0 1 0;
    0 0 1;
    1 0 0;
]
X = DiscreteMarkovChain(T)

is_ergodic(X)

# output

true
```
"""
function is_ergodic(x::AbstractMarkovChain)
    classes, _ = communication_classes(x)
    return length(classes) == 1
end

"""
    is_absorbing(x)

# Definitions
A Markov chain is absorbing if it has at least one absorbing
state, and if from every state it is possible to go to an absorbing
state (not necessarily in one step).

# Arguments
- `x`: some kind of Markov chain.

# Returns
`true` if the Markov chain, `x`, is an absorbing chain.
So the process is guarenteed to be absorbed eventually.

# Examples
The following is a typical example of an absorbing chain.

```jldoctest is_absorbing
using DiscreteMarkovChains
T = [
    1.0 0.0 0.0;
    0.1 0.8 0.1;
    0.0 0.3 0.7;
]
X = DiscreteMarkovChain(T)

is_absorbing(X)

# output

true
```

If a Markov chain does not have an absorbing state then the chain
is not absorbing. The converse is not true as we will see soon.

```jldoctest is_absorbing
T = [
    0.5 0.5 0.0;
    0.5 0.0 0.5;
    0.0 0.5 0.5;
]
X = DiscreteMarkovChain(T)

is_absorbing(X)

# output

false
```

In the following, the chain has multiple absorbing states but the
process is not guarenteed to be absorbed into those states.

```jldoctest is_absorbing
T = [
    1 0 0 0;
    0 1 0 0;
    0 0 0 1;
    0 0 1 0;
]
X = DiscreteMarkovChain(T)

is_absorbing(X)

# output

false
```

# References
1. [Dartmouth College](https://www.dartmouth.edu/~chance/teaching_aids/books_articles/probability_book/Chapter11.pdf)
"""
function is_absorbing(x::AbstractMarkovChain)
    states, A, B, C = decompose(x)
    r = size(A)[1]
    return (r > 0) &&
    (LinearAlgebra.diag(A) ≈ repeat([required_row_sum(typeof(x), eltype(A))], r))
end

"""
    is_reversible(x)

# Definitions
A Markov chain is said to be reversible if it satisfies the equation
``π_i p_{ij} = π_j p_{ji}`` where ``π`` is the stationary distribution
and ``p`` are the elements of the transition matrix of the Markov chain.

# Arguments
- `x`: some kind of Markov chain.

# Returns
`true` if the Markov chain is reversible.
"""
function is_reversible(x::AbstractMarkovChain)
    T = transition_matrix(x)
    n  = size(T)[1]
    l = stationary_distribution(x)
    L = repeat([l], n)'

    for i in 2:n
        for j in 1:(i-1)
            if !(l[i]*T[i, j] ≈ l[j]*T[j, i])
                return false
            end
        end
    end

    return true
end

"""
    stationary_distribution(x)

# Definitions
A stationary distribution of a Markov chain is a probability distribution
that remains unchanged in the Markov chain as time progresses.
It is a row vector, ``w`` such that its elements sum to 1 and it satisfies
``wT = w``. ``T`` is the one-step transiton matrix of the Markov chain.

In other words, ``w`` is invariant by the matrix ``T``.

For simplicity, this function returns a column vector instead of a row vector.

For continuous Markov chains, the stationary_distribution is given by the
solution to ``wQ = 0`` where ``Q`` is the transition intensity matrix.

# Arguments
- `x`: some kind of Markov chain.

# Returns
A column vector, ``w``, that satisfies the equation ``w'T = w'``.

# Examples
The stationary distribution will always exist. However, it might not be unique.

If it is unique there are no problems.
```jldoctest stationary_distribution
using DiscreteMarkovChains
T = [
    4 2 4;
    1 0 9;
    3 5 2;
]//10
X = DiscreteMarkovChain(T)

stationary_distribution(X)

# output

3-element Array{Rational{Int64},1}:
 35//129
 12//43
 58//129
```

If there are infinite solutions then the principle solution is taken
(every free variable is set to 0). A Moore-Penrose inverse is used.

```jldoctest stationary_distribution
T = [
    0.4 0.6 0.0;
    0.6 0.4 0.0;
    0.0 0.0 1.0;
]
X = DiscreteMarkovChain(T)

stationary_distribution(X)

# output

3-element Array{Float64,1}:
 0.33333333333333337
 0.33333333333333337
 0.33333333333333337
```

# References
1. [Brilliant.org](https://brilliant.org/wiki/stationary-distributions/#:~:text=A%20stationary%20distribution%20of%20a,transition%20matrix%20P%2C%20it%20satisfies)
"""
function stationary_distribution(x::AbstractMarkovChain)
    T = transition_matrix(x)
    n = length(state_space(x))

    if n == 0
        return Any[]
    end

    a = (T - characteristic_matrix(x))'

    a[1, :] = repeat([one(eltype(T))], n)
    b = repeat([zero(eltype(T))], n)
    b[1] = one(eltype(T))

    try
        return a\b
    catch e
        # This is a retarded way to handle this specific
        # exception but I don't know how else to deal with it.
        if e isa LinearAlgebra.SingularException
            # The pinv returns floats even if inputs are
            # rational. We will thus use it as a last resort.
            return LinearAlgebra.pinv(a)*b
        else
            rethrow(e)
        end
    end
end

"""
    fundamental_matrix(x)

# Definitions
The fundamental matrix of a markov chain is defined to be ``(I-C)^{-1}``
where ``C`` is the sub-transition matrix that takes transient states
to transient states.

The ``(i, j)``th entry of the fundamental matrix is the expected
number of times the chain is in state ``j`` over the whole process
given that the chain started in state ``i``.

# Arguments
- `x`: some kind of Markov chain.

# Returns
The fundamental matrix of the Markov chain.

# Notes
It is advised to have `x` in canonical form already.
This is to avoid confusion of what states make up each
element of the array ouput.

# Examples
Here is a typical example of the fundamental matrix.
```jldoctest fundamental_matrix
using DiscreteMarkovChains
T = [
    1.0 0.0 0.0;
    0.0 0.4 0.6;
    0.2 0.2 0.6;
]
X = DiscreteMarkovChain(T)

fundamental_matrix(X)

# output

2×2 Array{Float64,2}:
 3.33333  5.0
 1.66667  5.0
```

If the chain is ergodic, then the fundamental matrix is a 0x0
matrix (since there are no transient states).

```jldoctest fundamental_matrix
T = [
    0.0 1.0 0.0;
    0.5 0.0 0.5;
    0.0 1.0 0.0;
]
X = DiscreteMarkovChain(T)

fundamental_matrix(X)

# output

0×0 Array{Any,2}
```
"""
function fundamental_matrix(x::AbstractMarkovChain)
    _, _, _, C = decompose(x)

    if size(C)[1] == 0
        return Array{Any}(undef, 0, 0)
    end

    return LinearAlgebra.inv(characteristic_matrix(x) - C)
end

"""
    mean_time_to_absorption(x)

# Definitions
The expected time to absorption is the expected number of steps that
the process will take while in the transient super class given that
the process started in state ``i``.

# Arguments
- `x`: some kind of Markov chain.

# Returns
A 1D array where element ``i`` is the total number of revisits to
transient state ``i`` before leaving the transient super class.

# Notes
It is advised to have `x` in canonical form already.
This is to avoid confusion of what states make up each
element of the array ouput.

# Examples
See the following where we have a chain with 2 transient states
(states 3 and 4) that go between eachother. One of those states
will enter a pre-absorbing state (state 2). And then the
pre-absorbing state will enter the absorbing state (state 1)
on the next step.

So state 2 should spend no more steps in the transient states
and hence should have a time to absorption of 0. State 4 should
have 1 more step than state 3 since state 4
must enter state 3 to exit the transient states.

```jldoctest mean_time_to_absorption
using DiscreteMarkovChains
T = [
    1.0 0.0 0.0 0.0;
    1.0 0.0 0.0 0.0;
    0.0 0.2 0.0 0.8;
    0.0 0.0 1.0 0.0;
]
X = DiscreteMarkovChain(T)

mean_time_to_absorption(X)

# output

3-element Array{Float64,1}:
  0.0
  9.000000000000002
 10.000000000000002
```
"""
function mean_time_to_absorption(x::AbstractMarkovChain)
    M = fundamental_matrix(x)
    EV = (M-characteristic_matrix(x)) * ones(Int, size(M)[1])
    return EV
end

"""
    exit_probabilities(x)

# Arguments
- `x`: some kind of Markov chain.

# Returns
An array where element ``(i, j)`` is the probability that transient
state ``i`` will enter recurrent state ``j`` on its first step
out of the transient states. That is, ``e_{i,j}``.

# Examples
The following should be fairly obvious. States 1, 2
and 3 are the recurrent states and state 4 is the single
transient state that must enter one of these 3 on the next
time step. There is no randomness at play here.

```jldoctest exit_probabilities
using DiscreteMarkovChains
T = [
    0.2 0.2 0.6 0.0;
    0.5 0.4 0.1 0.0;
    0.6 0.2 0.2 0.0;
    0.2 0.3 0.5 0.0;
]
X = DiscreteMarkovChain(T)

exit_probabilities(X)

# output

1×3 Array{Float64,2}:
 0.2  0.3  0.5
```

So state 4 has probabilities 0.2, 0.3 and 0.5 of reaching
states 1, 2 and 3 respectively on the first step out of
the transient states (consisting only of state 4).

The following is less obvious.
```jldoctest exit_probabilities
T = [
    1.0 0.0 0.0 0.0;
    0.0 1.0 0.0 0.0;
    0.1 0.3 0.3 0.3;
    0.2 0.3 0.4 0.1;
]
X = DiscreteMarkovChain(T)

exit_probabilities(X)

# output

2×2 Array{Float64,2}:
 0.294118  0.705882
 0.352941  0.647059
```

So state 3 has a 29% chance of entering state 1 on the
first time step out (and the remaining 71% chance of
entering state 2). State 4 has a 35% chance of reaching
state 1 on the first time step out.
"""
function exit_probabilities(x::AbstractMarkovChain)
    M = fundamental_matrix(x)
    states, A, B, C = decompose(x)
    return M * B
end

"""
    mean_recurrence_time(x)

# Definitions
This is the expected number of steps for the process to
return to state i given that the process started in state ``i``.

# Arguments
- `x`: some kind of Markov chain.

# Returns
A 1D array where the ith entry is the mean recurrence time of state ``i``.

# Notes
`x` must be irreducible (i.e. ergodic).

# Examples
```jldoctest mean_recurrence_time
using DiscreteMarkovChains
T = [
    0.1 0.2 0.7;
    0.3 0.0 0.7;
    0.4 0.4 0.2;
]
X = DiscreteMarkovChain(T)

mean_recurrence_time(X)

# output

3-element Array{Float64,1}:
 3.4615384615384603
 4.090909090909091
 2.1428571428571432
```

If we have a reducible chain, then no error will be
raised but the output will be nonsense.

```jldoctest mean_recurrence_time
T = [
    1.0 0.0 0.0;
    0.1 0.5 0.4;
    0.0 0.3 0.7;
]
X = DiscreteMarkovChain(T)

mean_recurrence_time(X)

# output

3-element Array{Float64,1}:
   1.0
 -Inf
 -Inf
```
"""
function mean_recurrence_time(x::AbstractMarkovChain)
    return 1 ./ stationary_distribution(x)
end

"""
    mean_first_passage_time(x)

# Definitions
This is the expected number of steps for the process to
reach state ``j`` given that the process started in state ``i``.
We say that the mean first passage time between a state and itself is 0.

# Arguments
- `x`: some kind of Markov chain.

# Returns
A matrix where the ``(i,j)``th entry is the mean recurrence time of state ``i``.
Diagonal elements are 0.
The diagonals would have represented the mean recurrence time.

# Notes
`x` must be irreducible (i.e. ergodic).

# Examples
```jldoctest mean_first_passage_time
using DiscreteMarkovChains
T = [
    0.1 0.2 0.7;
    0.3 0.0 0.7;
    0.4 0.4 0.2;
]
X = DiscreteMarkovChain(T)

mean_first_passage_time(X)

# output

3×3 Array{Float64,2}:
 0.0      3.40909  1.42857
 2.88462  0.0      1.42857
 2.69231  2.95455  0.0
```

If we have a reducible chain, then no error will be
raised but the output will be nonsense.

```jldoctest mean_first_passage_time
T = [
    1.0 0.0 0.0;
    0.1 0.5 0.4;
    0.0 0.3 0.7;
]
X = DiscreteMarkovChain(T)

mean_first_passage_time(X)

# output

3×3 Array{Float64,2}:
  0.0     -Inf  -Inf
 23.3333  NaN   -Inf
 26.6667  -Inf  NaN
```
"""
function mean_first_passage_time(x::AbstractMarkovChain)
    n = length(state_space(x))
    if n == 0
        return transition_matrix(x)  # Keeps eltype the same
    end

    T = probability_matrix(x)
    W = repeat(stationary_distribution(x)', n)
    Z = LinearAlgebra.inv(LinearAlgebra.I - T + W)
    Zjj = repeat(LinearAlgebra.diag(Z)', n)

    M = (Zjj .- Z) ./ W
    return M
end