using LAP_julia: load_H, make_flow_from_H, gen_rand_points, get_valid_landmarks
using DataFrames, CSV

## create a birl like structure for with the oxford affine dataset
base_path = "/Users/MrTrololord/Downloads/mri_head_homography/"
target_images_path = joinpath(base_path, "images", "target")
target_images_fnames = filter(x -> !isfile(x), readdir(target_images_path))
target_images_paths = map(fname -> joinpath(target_images_path, fname), target_images_fnames)

COUNT = 500
max_disp = 40
# flow_gen_func(img_size) = gen_quad_flow(img_size, max_disp)
flow_gen_func(img_size) = gen_homo_flow(img_size, max_disp)

for target_img_path in target_images_paths
    k = parse(Int, target_img_path[findfirst(r"[0-9]+", target_img_path)])
    target_img = load(target_img_path)
    # create save paths
    target_land_path = joinpath(base_path, "landmarks", "target", "landmarks$(k).csv")
    source_land_path = joinpath(base_path, "landmarks", "source", "landmarks$(k).csv")
    source_img_path = joinpath(base_path, "images", "source", "img$(k).png")
    # generate rand inds in target
    target_inds = gen_rand_points(target_img, COUNT, "Gridded")
    cur_flow = flow_gen_func(size(target_img))
    # edit target inds and save
    target_inds = LAP_julia.get_valid_landmarks(cur_flow, target_inds)
    LAP_julia.save_landmarks(target_inds, target_land_path)
    # warp target inds and save
    source_points = LAP_julia.move_landmarks(transpose(LAP_julia.inds_to_points(target_inds)), cur_flow)
    # rng = extrema.(indices_spatial(cur_flow))
    # @info source_points[:,1] size(cur_flow) rng
    # @assert all(LAP_julia.is_in_bounds.(source_points[:,1], rng[1]...) .& LAP_julia.is_in_bounds.(source_points[:,2], rng[2]...))
    LAP_julia.save_landmarks(source_points, source_land_path)
    # LAP_julia.save_shift_landmarks(target_land_path, cur_flow.*(-1), source_land_path)
    # create source img
    source = warp_img(target_img, -real.(cur_flow), -imag.(cur_flow))
    save(source_img_path, source)
    @info k target_img_path size(target_inds)
end

# create loc table
tmp_df = CSV.read("/Volumes/davidkunz/Documents/bakalarka/birl_small/dataset_small_repaired.csv") |> DataFrame

df = deepcopy(tmp_df)
delete!(df, 1:size(tmp_df,1))

# rename!(df, Symbol.(replace.(string.(names(df)), Ref(" "=> "_"))))
# rename!(df, Symbol.(names(df)))

names(df)
eltypes(df)

abc = [:a , :b, :c, :d, :e, :f, :g, :h, :i, :j, :k]
new_df = DataFrame(a = eltypes(df)[1][],
                   b = eltypes(df)[2][],
                   c = eltypes(df)[3][],
                   d = eltypes(df)[4][],
                   e = eltypes(df)[5][],
                   f = eltypes(df)[6][],
                   g = eltypes(df)[7][],
                   h = eltypes(df)[8][],
                   i = eltypes(df)[9][],
                   j = eltypes(df)[10][],
                   k = eltypes(df)[11][])

function rename_dict_keys(from_name, to_name, dict)
    new_dict = Dict()
    rename_dict = Dict(from_name .=> to_name)
    for (key, value) in rename_dict
        new_dict[value] = dict[key]
    end
    return new_dict
end

experiment_count = count(x -> !isfile(x), readdir(target_images_path))

for k in 1:experiment_count
    row = Dict{String}{Any}()
    row["Column1"] = k
    # images
    row["Target image"] = joinpath("images", "target", "img$k.png")
    row["Source image"] = joinpath("images", "source", "img$k.png")
    # lnds
    row["Target landmarks"] = joinpath("landmarks", "target", "landmarks$k.csv")
    row["Source landmarks"] = joinpath("landmarks", "source", "landmarks$k.csv")
    # sizes
    source_img = load(joinpath(base_path, "images", "source", "img$k.png"))
    row["Image size [pixels]"] = string(size(source_img))
    row["Image diagonal [pixels]"] = round(sqrt(sum(size(source_img).^2)), digits=1)
    row["status"] = "training"
    # empty values
    row["Warped target landmarks"] = missing
    row["Warped source landmarks"] = missing
    row["Execution time [minutes]"] = missing

    @info rename_dict_keys(names(tmp_df), abc, row)
    push!(new_df, rename_dict_keys(names(tmp_df), abc, row))
end

new_df

# rename back
renamed = rename!(new_df, names(tmp_df))

CSV.write(joinpath(base_path, "loc_table.csv"), renamed)


module VolumeRHS

using BenchmarkTools
using CUDA
using StableRNGs
using StaticArrays

function loopinfo(name, expr, nodes...)
    if expr.head != :for
        error("Syntax error: pragma $name needs a for loop")
    end
    push!(expr.args[2].args, Expr(:loopinfo, nodes...))
    return expr
end

macro unroll(expr)
    expr = loopinfo("@unroll", expr, (Symbol("llvm.loop.unroll.full"),))
    return esc(expr)
end

# HACK: module-local versions of core arithmetic; needed to get FMA
for (jlf, f) in zip((:+, :*, :-), (:add, :mul, :sub))
    for (T, llvmT) in ((Float32, "float"), (Float64, "double"))
        ir = """
            %x = f$f contract nsz $llvmT %0, %1
            ret $llvmT %x
        """
        @eval begin
            # the @pure is necessary so that we can constant propagate.
            @inline Base.@pure function $jlf(a::$T, b::$T)
                Base.llvmcall($ir, $T, Tuple{$T, $T}, a, b)
            end
        end
    end
    @eval function $jlf(args...)
        Base.$jlf(args...)
    end
end

let (jlf, f) = (:div_arcp, :div)
    for (T, llvmT) in ((Float32, "float"), (Float64, "double"))
        ir = """
            %x = f$f fast $llvmT %0, %1
            ret $llvmT %x
        """
        @eval begin
            # the @pure is necessary so that we can constant propagate.
            @inline Base.@pure function $jlf(a::$T, b::$T)
                Base.llvmcall($ir, $T, Tuple{$T, $T}, a, b)
            end
        end
    end
    @eval function $jlf(args...)
        Base.$jlf(args...)
    end
end
rcp(x) = div_arcp(one(x), x) # still leads to rcp.rn which is also a function call

# div_fast(x::Float32, y::Float32) = ccall("extern __nv_fast_fdividef", llvmcall, Cfloat, (Cfloat, Cfloat), x, y)
# rcp(x) = div_fast(one(x), x)

# note the order of the fields below is also assumed in the code.
const _nstate = 5
const _ρ, _U, _V, _W, _E = 1:_nstate
const stateid = (ρ = _ρ, U = _U, V = _V, W = _W, E = _E)

const _nvgeo = 14
const _ξx, _ηx, _ζx, _ξy, _ηy, _ζy, _ξz, _ηz, _ζz, _MJ, _MJI,
_x, _y, _z = 1:_nvgeo
const vgeoid = (ξx = _ξx, ηx = _ηx, ζx = _ζx,
                ξy = _ξy, ηy = _ηy, ζy = _ζy,
                ξz = _ξz, ηz = _ηz, ζz = _ζz,
                MJ = _MJ, MJI = _MJI,
                x = _x,   y = _y,   z = _z)

const N = 4
const nmoist = 0
const ntrace = 0

Base.@irrational grav 9.81 BigFloat(9.81)
Base.@irrational gdm1 0.4 BigFloat(0.4)

function volumerhs!(rhs, Q, vgeo, gravity, D, nelem)
    Q = Base.Experimental.Const(Q)
    vgeo = Base.Experimental.Const(vgeo)
    D = Base.Experimental.Const(D)

    nvar = _nstate + nmoist + ntrace

    Nq = N + 1

    s_D = CuStaticSharedArray(eltype(D), (Nq, Nq))
    s_F = CuStaticSharedArray(eltype(Q), (Nq, Nq, _nstate))
    s_G = CuStaticSharedArray(eltype(Q), (Nq, Nq, _nstate))

    r_rhsρ = MArray{Tuple{Nq}, eltype(rhs)}(undef)
    r_rhsU = MArray{Tuple{Nq}, eltype(rhs)}(undef)
    r_rhsV = MArray{Tuple{Nq}, eltype(rhs)}(undef)
    r_rhsW = MArray{Tuple{Nq}, eltype(rhs)}(undef)
    r_rhsE = MArray{Tuple{Nq}, eltype(rhs)}(undef)

    e = blockIdx().x
    j = threadIdx().y
    i = threadIdx().x

    @inbounds begin
        for k in 1:Nq
            r_rhsρ[k] = zero(eltype(rhs))
            r_rhsU[k] = zero(eltype(rhs))
            r_rhsV[k] = zero(eltype(rhs))
            r_rhsW[k] = zero(eltype(rhs))
            r_rhsE[k] = zero(eltype(rhs))
        end

        # fetch D into shared
        s_D[i, j] = D[i, j]
        @unroll for k in 1:Nq
            sync_threads()

            # Load values will need into registers
            MJ = vgeo[i, j, k, _MJ, e]
            ξx, ξy, ξz = vgeo[i,j,k,_ξx,e], vgeo[i,j,k,_ξy,e], vgeo[i,j,k,_ξz,e]
            ηx, ηy, ηz = vgeo[i,j,k,_ηx,e], vgeo[i,j,k,_ηy,e], vgeo[i,j,k,_ηz,e]
            ζx, ζy, ζz = vgeo[i,j,k,_ζx,e], vgeo[i,j,k,_ζy,e], vgeo[i,j,k,_ζz,e]
            z = vgeo[i,j,k,_z,e]

            U, V, W = Q[i, j, k, _U, e], Q[i, j, k, _V, e], Q[i, j, k, _W, e]
            ρ, E = Q[i, j, k, _ρ, e], Q[i, j, k, _E, e]

            # GPU performance trick
            # Allow optimizations to use the reciprocal of an argument rather than perform division.
            # IEEE floating-point division is implemented as a function call
            ρinv = rcp(ρ)
            ρ2inv = rcp(2ρ)
            # ρ2inv = 0.5f0 * pinv

            P = gdm1*(E - (U^2 + V^2 + W^2)*ρ2inv - ρ*gravity*z)

            fluxρ_x = U
            fluxU_x = ρinv * U * U + P
            fluxV_x = ρinv * U * V
            fluxW_x = ρinv * U * W
            fluxE_x = ρinv * U * (E + P)

            fluxρ_y = V
            fluxU_y = ρinv * V * U
            fluxV_y = ρinv * V * V + P
            fluxW_y = ρinv * V * W
            fluxE_y = ρinv * V * (E + P)

            fluxρ_z = W
            fluxU_z = ρinv * W * U
            fluxV_z = ρinv * W * V
            fluxW_z = ρinv * W * W + P
            fluxE_z = ρinv * W * (E + P)

            s_F[i, j,  _ρ] = MJ * (ξx * fluxρ_x + ξy * fluxρ_y + ξz * fluxρ_z)
            s_F[i, j,  _U] = MJ * (ξx * fluxU_x + ξy * fluxU_y + ξz * fluxU_z)
            s_F[i, j,  _V] = MJ * (ξx * fluxV_x + ξy * fluxV_y + ξz * fluxV_z)
            s_F[i, j,  _W] = MJ * (ξx * fluxW_x + ξy * fluxW_y + ξz * fluxW_z)
            s_F[i, j,  _E] = MJ * (ξx * fluxE_x + ξy * fluxE_y + ξz * fluxE_z)

            s_G[i, j,  _ρ] = MJ * (ηx * fluxρ_x + ηy * fluxρ_y + ηz * fluxρ_z)
            s_G[i, j,  _U] = MJ * (ηx * fluxU_x + ηy * fluxU_y + ηz * fluxU_z)
            s_G[i, j,  _V] = MJ * (ηx * fluxV_x + ηy * fluxV_y + ηz * fluxV_z)
            s_G[i, j,  _W] = MJ * (ηx * fluxW_x + ηy * fluxW_y + ηz * fluxW_z)
            s_G[i, j,  _E] = MJ * (ηx * fluxE_x + ηy * fluxE_y + ηz * fluxE_z)

            r_Hρ = MJ * (ζx * fluxρ_x + ζy * fluxρ_y + ζz * fluxρ_z)
            r_HU = MJ * (ζx * fluxU_x + ζy * fluxU_y + ζz * fluxU_z)
            r_HV = MJ * (ζx * fluxV_x + ζy * fluxV_y + ζz * fluxV_z)
            r_HW = MJ * (ζx * fluxW_x + ζy * fluxW_y + ζz * fluxW_z)
            r_HE = MJ * (ζx * fluxE_x + ζy * fluxE_y + ζz * fluxE_z)

            # one shared access per 10 flops
            for n = 1:Nq
                Dkn = s_D[k, n]

                r_rhsρ[n] += Dkn * r_Hρ
                r_rhsU[n] += Dkn * r_HU
                r_rhsV[n] += Dkn * r_HV
                r_rhsW[n] += Dkn * r_HW
                r_rhsE[n] += Dkn * r_HE
            end

            r_rhsW[k] -= MJ * ρ * gravity

            sync_threads()

            # loop of ξ-grid lines
            @unroll for n = 1:Nq
                Dni = s_D[n, i]
                Dnj = s_D[n, j]

                r_rhsρ[k] += Dni * s_F[n, j, _ρ]
                r_rhsρ[k] += Dnj * s_G[i, n, _ρ]

                r_rhsU[k] += Dni * s_F[n, j, _U]
                r_rhsU[k] += Dnj * s_G[i, n, _U]

                r_rhsV[k] += Dni * s_F[n, j, _V]
                r_rhsV[k] += Dnj * s_G[i, n, _V]

                r_rhsW[k] += Dni * s_F[n, j, _W]
                r_rhsW[k] += Dnj * s_G[i, n, _W]

                r_rhsE[k] += Dni * s_F[n, j, _E]
                r_rhsE[k] += Dnj * s_G[i, n, _E]
            end
        end # k

        @unroll for k in 1:Nq
            MJI = vgeo[i, j, k, _MJI, e]

            # Updates are a performance bottleneck
            # primary source of stall_long_sb
            rhs[i, j, k, _U, e] += MJI*r_rhsU[k]
            rhs[i, j, k, _V, e] += MJI*r_rhsV[k]
            rhs[i, j, k, _W, e] += MJI*r_rhsW[k]
            rhs[i, j, k, _ρ, e] += MJI*r_rhsρ[k]
            rhs[i, j, k, _E, e] += MJI*r_rhsE[k]
        end
    end
    return
end

function main()
    DFloat = Float32
    nelem = 240_000

    rng = StableRNG(123)

    Nq = N + 1
    nvar = _nstate + nmoist + ntrace

    Q = 1 .+ CuArray(rand(rng, DFloat, Nq, Nq, Nq, nvar, nelem))
    Q[:, :, :, _E, :] .+= 20
    vgeo = CuArray(rand(rng, DFloat, Nq, Nq, Nq, _nvgeo, nelem))

    # make sure the entries of the mass matrix satisfy the inverse relation
    vgeo[:, :, :, _MJ, :] .+= 3
    vgeo[:, :, :, _MJI, :] .= 1 ./ vgeo[:, :, :, _MJ, :]

    D = CuArray(rand(rng, DFloat, Nq, Nq))

    rhs = CuArray(zeros(DFloat, Nq, Nq, Nq, nvar, nelem))

    threads=(N+1, N+1)

    kernel = @cuda launch=false volumerhs!(rhs, Q, vgeo, DFloat(grav), D, nelem)
    # XXX: should we print these for all kernels? maybe upload them to Codespeed?
    @info """volumerhs! details:
              - $(CUDA.registers(kernel)) registers, max $(CUDA.maxthreads(kernel)) threads
              - $(Base.format_bytes(CUDA.memory(kernel).local)) local memory,
                $(Base.format_bytes(CUDA.memory(kernel).shared)) shared memory,
                $(Base.format_bytes(CUDA.memory(kernel).constant)) constant memory"""
    results = @benchmark begin
        CUDA.@sync $kernel($rhs, $Q, $vgeo, $(DFloat(grav)), $D, $nelem;
                           threads=$threads, blocks=$nelem)
    end

    # BenchmarkTools captures inputs, JuliaCI/BenchmarkTools.jl#127, so forcibly free them
    CUDA.unsafe_free!(rhs)
    CUDA.unsafe_free!(Q)
    CUDA.unsafe_free!(vgeo)
    CUDA.unsafe_free!(D)

    results
end

end

VolumeRHS.main()


#cd("/home/david/Documents/gitclones/AlphaZero.jl")
cd("/home/davidyang/Documents/gitclones/AlphaZero.jl")
using Pkg
Pkg.activate(".")
using AlphaZero
#include("/home/david/Documents/gitclones/AlphaZero.jl/games/amazon_full/main.jl")
include("/home/davidyang/Documents/gitclones/AlphaZero.jl/games/amazon_full/main.jl")

gs = Amazon_Full.GameSpec()
g = GI.init(gs)
GI.vectorize_state(gs,GI.current_state(g))
GI.play!(g,g|>GI.actions_mask|>findall|>rand)
GI.vectorize_state(gs,GI.current_state(g))

g = GI.init(gs)
GI.render(g)
while !g.finished
  action = g|>GI.actions_mask|>findall|>rand
  println(GI.action_string(gs,action))
  GI.play!(g,action)
  GI.render(g)
end

cd("/home/davidyang/Documents/gitclones/AlphaZero.jl")
#cd("/home/david/Documents/gitclones/AlphaZero.jl")
#@everywhere using Pkg
#@everywhere Pkg.activate(".")
#@everywhere using AlphaZero
using Pkg
Pkg.activate(".")
using AlphaZero
experiment = Examples.experiments["amazon_full"]
session = Session(experiment, dir="sessions/amazon_full4")
resume!(session)

duel1 = session.benchmark[1]
#duel1 equivalent to eval
progress = UI.Log.Progress(session.logger, duel1.sim.num_games)
env = session.env
params = env.params.self_play
net() = Network.copy(env.bestnn, on_gpu=duel1.sim.use_gpu, test_mode=true)
simulator = Simulator(net, record_trace) do net
  player = Benchmark.instantiate(duel1.player, env.gspec, net)
  baseline = Benchmark.instantiate(duel1.baseline, env.gspec, net)
  return TwoPlayers(player, baseline)
end
#=
samples, elapsed = @timed simulate(
    simulator, env.gspec, duel1.sim,
    game_simulated=(() -> next!(progress)))
=#
oracles = simulator.make_oracles()
spawn_oracles, done =
    AlphaZero.batchify_oracles(oracles; duel1.sim.num_workers, duel1.sim.batch_size, duel1.sim.fill_batches)
oracles = spawn_oracles()
player = simulator.make_player(oracles)

println("debugging play game")
trace = @time play_game(env.gspec, simulator.make_player(oracles), flip_probability=0.0)

colors_flipped = true
report = simulator.measure(trace, colors_flipped, player)

results, elapsed = @timed simulate_distributed(
    simulator, env.gspec, params.sim,
    game_simulated=()->Handlers.game_played(handler))

#seems stuck at the above function, so expand to see what's going on

simulate(simulator,
        gspec,
        SimParams(p; num_games=(w == workers[1] ? num_each + rem : num_each)),
        game_simulated=remote_game_simulated)

using Distributed
gspec = env.gspec
p = params.sim
chan = Distributed.RemoteChannel(()->Channel{Nothing}(1))
Util.@tspawn_main begin
  for i in 1:p.num_games
    take!(chan)
    game_simulated()
  end
end
remote_game_simulated() = put!(chan, nothing)
# Distributing the simulations across workers
num_each, rem = divrem(p.num_games, Distributed.nworkers())
@assert num_each >= 1
workers = Distributed.workers()
tasks = map(workers) do w
  Distributed.@spawnat w begin
    Util.@printing_errors begin
      simulate(
        simulator,
        gspec,
        SimParams(p; num_games=(w == workers[1] ? num_each + rem : num_each)),
        game_simulated=remote_game_simulated)
      end
  end
end
results = fetch.(tasks)
# If one of the worker raised an exception, we print it
for r in results
  if isa(r, Distributed.RemoteException)
    showerror(stderr, r, catch_backtrace())
  end
end

flip_probability = 0.0
player = simulator.make_player(oracles)
game = GI.init(env.gspec)
trace = Trace(GI.current_state(game))
#actions, π_target = think(player, game)
env = player.white.mcts
#MCTS.explore!(env,game,player.white.niters)
#η = MCTS.dirichlet_noise(game, env.noise_α)
#MCTS.run_simulation!(env,GI.clone(game),η=η)
#MCTS.state_info(env,GI.current_state(game))
#env.oracle(state)
#Network.evaluate(env.oracle,state)
nn = env.oracle
gspec = Network.game_spec(nn)
actions_mask = GI.actions_mask(GI.init(gspec, state))
x = GI.vectorize_state(gspec, state)
a = Float32.(actions_mask)
xnet, anet = Network.to_singletons.(Network.convert_input_tuple(nn, (x, a)))
#net_output = Network.forward_normalized(nn, xnet, anet)
p, v = Network.forward(nn, xnet)

while true
  if GI.game_terminated(game)
    return trace
  end
  if !iszero(flip_probability) && rand() < flip_probability
    GI.apply_random_symmetry!(game)
  end
  actions, π_target = think(player, game)
  τ = player_temperature(player, game, length(trace))
  π_sample = apply_temperature(π_target, τ)
  a = actions[Util.rand_categorical(π_sample)]
  GI.play!(game, a)
  push!(trace, π_target, GI.white_reward(game), GI.current_state(game))
end


"""
# module DataStructure

# Examples

```jldoctest
julia>
```
"""
module DataStructure

export Cell

struct Cell
    lattice::AbstractMatrix
    positions::AbstractMatrix
    numbers::AbstractVector
end

end


mutable struct TwoBodyDB <: Database
    m::Int
    n::Int

    rawDB
    chart::Vector{Tuple}
    degeneracy::Dict

    fx::Vector{Number}
    fy::Vector{Number}
    fz::Vector{Number}

    idx::Dict
    idxset::Dict

    xi::Dict
    yi::Dict
    zi::Dict

    ri::Dict

    xi1::Dict
    yi1::Dict
    zi1::Dict

    ri1::Dict

    xi_i1::Dict
    yi_i1::Dict
    zi_i1::Dict

    rii1::Dict
end

function TwoBodyDB(numbers::Array{Int, 1})
    chart, degeneracy = getchart(Tuple(numbers), 2)
    m = length(chart)
    n = 0

    fx = Vector{Number}([])
    fy = Vector{Number}([])
    fz = Vector{Number}([])

    # Indexing to force
    idx = Dict{Tuple{Int, Int}, Vector{Int}}([(key, Vector{Int}([])) for key in chart])
    idxset = Dict{Tuple{Int, Int}, Set}([(key, Set()) for key in chart])

    xi = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])

    ri = Dict{Tuple{Int, Int}, Array{Array{Float64, 1}, 1}}([(key, Array{Array{Float64, 1}, 1}([])) for key in chart])

    xi1 = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi1 = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi1 = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])

    # ri1 = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    ri1 = Dict{Tuple{Int, Int}, Array{Array{Float64, 1}, 1}}([(key, Array{Array{Float64, 1}, 1}([])) for key in chart])

    xi_i1 = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi_i1 = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi_i1 = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])

    rii1 = Dict{Tuple{Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])

    TwoBodyDB(m, n, [], chart, degeneracy, fx, fy, fz, idx, idxset,
              xi, yi, zi, ri, xi1, yi1, zi1, ri1, xi_i1, yi_i1, zi_i1, rii1)
end

update!(db::TwoBodyDB,
        coords::Cartesian{T1, 2},
        forces::Array{T2, 2}; kwargs...) where {T1<:Number, T2<:Number} = #=
        =# update!(db, coords[nax, :, :], forces[nax, :, :]; kwargs...)
update!(db::TwoBodyDB, coords::Cartesian{T1, 3},
        forces::Array{T2, 3}; kwargs...) where {T1<:Number, T2<:Number}#=
    =# = update!(db, convert(TwoBody, coords; kwargs...), forces)

function update!(db::TwoBodyDB, coords::SparseAtomic, forces::Array{T, 3}) where {T<:Number}
    M, A, _ = size(forces)
    # size between coords and forces is inconsistent
    @assert (M, A) == Tuple(size(coords)[1:2])
    # "No data found"
    @assert M > 0

    chart = keys(db.idx)
    numbers = coords.cache["numbers"]

    count = db.n + 1

    for m=1:M
        for a=1:A
            push!(db.rawDB, coords[m, a])
            push!(db.fx, forces[m, a, 1]);
            push!(db.fy, forces[m, a, 2]);
            push!(db.fz, forces[m, a, 3])
            datum = coords[m, a].descriptor
            for cluster in chart
                append!(db.xi[cluster],    datum[cluster]["xi"])
                append!(db.yi[cluster],    datum[cluster]["yi"])
                append!(db.zi[cluster],    datum[cluster]["zi"])
                append!(db.ri[cluster],    datum[cluster]["ri"])
                append!(db.xi1[cluster],   datum[cluster]["xi1"])
                append!(db.yi1[cluster],   datum[cluster]["yi1"])
                append!(db.zi1[cluster],   datum[cluster]["zi1"])
                append!(db.ri1[cluster],   datum[cluster]["ri1"])
                append!(db.xi_i1[cluster], datum[cluster]["xi_i1"])
                append!(db.yi_i1[cluster], datum[cluster]["yi_i1"])
                append!(db.zi_i1[cluster], datum[cluster]["zi_i1"])
                append!(db.rii1[cluster],  datum[cluster]["rii1"])

                idx = similar(Vector{Int}, length(datum[cluster]["xi"]))
                fill!(idx, count)
                append!(db.idx[cluster], idx)
                union!(db.idxset[cluster], count)
            end

            count += 1
            db.n += 1
        end
    end
end

"""
 idx refers to the number where its enviroment located in the array.
"""
mutable struct ThreeBodyDB <: Database
    m::Int                  # Number of rows
    n::Int                  # Number of columns

    rawDB
    chart::Vector{Tuple}
    degeneracy::Dict

    fx::Vector{Number}
    fy::Vector{Number}
    fz::Vector{Number}

    idx::Dict
    idxset::Dict

    xi::Dict
    yi::Dict
    zi::Dict
    ri::Dict

    xi1::Dict
    yi1::Dict
    zi1::Dict
    ri1::Dict

    xi2::Dict
    yi2::Dict
    zi2::Dict
    ri2::Dict

    xi_i1::Dict
    yi_i1::Dict
    zi_i1::Dict

    xi_i2::Dict
    yi_i2::Dict
    zi_i2::Dict

    xi1_i2::Dict
    yi1_i2::Dict
    zi1_i2::Dict

    rii1::Dict
    rii2::Dict
    ri1i2::Dict

end

function ThreeBodyDB(numbers::Array{Int, 1})
    chart, degeneracy = getchart(Tuple(numbers), 3)
    m = length(chart)
    n = 0

    fx = Vector{Number}([])
    fy = Vector{Number}([])
    fz = Vector{Number}([])

    # Indexing to force
    idx = Dict{Tuple{Int, Int, Int}, Vector{Int}}([(key, Vector{Int}([])) for key in chart])
    idxset = Dict{Tuple{Int, Int, Int}, Set}([(key, Set()) for key in chart])

    xi = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    ri = Dict{Tuple{Int, Int, Int}, Array{Array{Float64, 1}, 1}}([(key, Array{Array{Float64, 1}, 1}([])) for key in chart])

    xi1 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi1 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi1 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    ri1 = Dict{Tuple{Int, Int, Int}, Array{Array{Float64, 1}, 1}}([(key, Array{Array{Float64, 1}, 1}([])) for key in chart])

    xi2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    ri2 = Dict{Tuple{Int, Int, Int}, Array{Array{Float64, 1}, 1}}([(key, Array{Array{Float64, 1}, 1}([])) for key in chart])

    xi_i1 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi_i1 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi_i1 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])

    xi_i2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi_i2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi_i2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])

    xi1_i2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    yi1_i2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    zi1_i2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])

    rii1 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    rii2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])
    ri1i2 = Dict{Tuple{Int, Int, Int}, Vector{Number}}([(key, Vector{Number}([])) for key in chart])

    ThreeBodyDB(m, n, [], chart, degeneracy, fx, fy, fz, idx, idxset, xi, yi, zi, ri, xi1, yi1, zi1, ri1, xi2, yi2, zi2, ri2,
                xi_i1, yi_i1, zi_i1, xi_i2, yi_i2, zi_i2, xi1_i2, yi1_i2, zi1_i2,
                rii1, rii2, ri1i2)
end

update!(db::ThreeBodyDB, coords::Cartesian{T, 2}, forces::Array{Number, 2}; kwargs...) where {T<:Number} = update!(db, coords[nax, :, :], forces[nax, :, :]; kwargs...)
update!(db::ThreeBodyDB, coords::Cartesian{T, 3}, forces::Array{Number, 3}; kwargs...) where {T<:Number} = update!(db, convert(TwoBody, coords; kwargs...), forces)

function update!(db::ThreeBodyDB, coords::SparseAtomic, forces::Array{T, 3}) where {T<:Number}
    M, A, _ = size(forces)
    # "size between coords and forces is inconsistent"
    @assert (M, A) == Tuple(size(coords)[1:2])
    # No data found"
    @assert M > 0

    chart = keys(db.idx)
    numbers = coords.cache["numbers"]

    count = db.n + 1
    for m=1:M
        for a=1:A
            push!(db.rawDB, coords[m, a])
            push!(db.fx, forces[m, a, 1]); push!(db.fy, forces[m, a, 2]);
            push!(db.fz, forces[m, a, 3])

            datum = coords[m, a].descriptor
            for cluster in chart
                append!(db.xi[cluster],    datum[cluster]["xi"])
                append!(db.yi[cluster],    datum[cluster]["yi"])
                append!(db.zi[cluster],    datum[cluster]["zi"])
                append!(db.ri[cluster],    datum[cluster]["ri"])
                append!(db.xi1[cluster],   datum[cluster]["xi1"])
                append!(db.yi1[cluster],   datum[cluster]["yi1"])
                append!(db.zi1[cluster],   datum[cluster]["zi1"])
                append!(db.ri1[cluster],   datum[cluster]["ri1"])
                append!(db.xi2[cluster],   datum[cluster]["xi2"])
                append!(db.yi2[cluster],   datum[cluster]["yi2"])
                append!(db.zi2[cluster],   datum[cluster]["zi2"])
                append!(db.ri2[cluster],   datum[cluster]["ri2"])
                append!(db.xi_i1[cluster], datum[cluster]["xi_i1"])
                append!(db.yi_i1[cluster], datum[cluster]["yi_i1"])
                append!(db.zi_i1[cluster], datum[cluster]["zi_i1"])
                append!(db.xi_i2[cluster], datum[cluster]["xi_i2"])
                append!(db.yi_i2[cluster], datum[cluster]["yi_i2"])
                append!(db.zi_i2[cluster], datum[cluster]["zi_i2"])
                append!(db.xi1_i2[cluster],datum[cluster]["xi1_i2"])
                append!(db.yi1_i2[cluster],datum[cluster]["yi1_i2"])
                append!(db.zi1_i2[cluster],datum[cluster]["zi1_i2"])
                append!(db.rii1[cluster],  datum[cluster]["rii1"])
                append!(db.rii2[cluster],  datum[cluster]["rii2"])
                append!(db.ri1i2[cluster], datum[cluster]["ri1i2"])

                idx = similar(Vector{Int}, length(datum[cluster]["xi"]))
                fill!(idx, count)
                append!(db.idx[cluster], idx)
                union!(db.idxset[cluster], count)
            end

            count += 1
            db.n += 1
        end
    end
end

mutable struct TwoThreeBodyDB <: Database
    twobody::TwoBodyDB
    threebody::ThreeBodyDB
end

function TwoThreeBodyDB(numbers)
    twobody = TwoBodyDB(numbers)
    threebody = ThreeBodyDB(numbers)
    TwoThreeBodyDB(twobody, threebody)
end

update!(db::TwoThreeBodyDB, coords::Cartesian{T, 2},
        forces::Array{Number, 2}; kwargs...) where {T<:Number}#=
        =# = update!(db, coords[nax, :, :], forces[nax, :, :]; kwargs...)
update!(db::TwoThreeBodyDB, coords::Cartesian{T, 3},
        forces::Array{Number, 3}; kwargs...) where {T<:Number}#=
        =# = update!(db, convert(TwoBody, coords; kwargs...), forces)

function update!(db::TwoThreeBodyDB, coords, forces)
    update!(db.twobody, coords, forces)
    update!(db.threebody, coords, forces)
end


using MetaGraphs
using LightGraphs
using GraphMLReader
using GraphIO
using TimerOutputs

abstract type MyAbstractPopulation end

export Population
```
Struct Population

```
struct Population <: MyAbstractPopulation

    nb_pop::Int32
    id_ind::Array{Int32,1}
    id_house::Array{Int32,1}
    statut::Array{Int32,1}
    indiv_in_out::Array{Int32,1}
    id_ind_work::Array{Int32,1}
    tissu_work::Array{Int32,1}
    cat_age::Array{Int32,1}

    function Population()

        name_file = ParamFile.GRAPH_FILE ;
        graph = loadgraphml("Data/$(lpad(name_file,1)).graphml") ;
        n_node = nv(graph) ;
        id_ind, id_house, cat_age = individual(graph,n_node) ;
        id_work, area_work = working_place(graph,n_node) ;
        nb_pop_graph = length(id_ind) ;
        statut = def_statut(cat_age) ;
        indiv_in_out,statut,id_house = indiv_inout(statut,id_house) ;
        nb_pop = nb_pop_graph + ParamPop.NB_IN ;
        id_ind_work,tissu_work = choix_travail(statut,indiv_in_out,id_work,area_work) ;

        return new(
        nb_pop,
        id_ind,
        id_house,
        statut,
        indiv_in_out,
        id_ind_work,
        tissu_work,
        cat_age
        )

    end

end

"""
Individual(graph, n_node)

    # Function used to find the id node of all individuals
    # Input:
    # - graph: undirected Int64 metagraph with Float64 weights
    # - n_node ({Int32}): total number of nodes
    # Output:
    # - id_indiv (Array{Int32,.}): Vector containing the id of all individuals
    # - id_house (Array{Int32,.}): Vector containing the id of the corresponding house
    # - cat_age (Array{Int32,.}): Vector containing the age category

"""
function individual(graph::MetaGraph,n_node::Int)

    # Initialisation of the vectors
    id_indiv = Array{Int32}([]) ;
    id_house = Array{Int32}([]) ;
    cat_age = Array{Int32}([]) ;
    x = ["a0_2","a3_5","a6_10","a11_17","a18_24","a25_39","a40_54","a55_64","a65_79","a80p"] ;

    # Find the nodes corresponding to the individuals
    @inbounds for i=1:n_node
        if(get_prop(graph,i,:type)=="Individual")
            append!(id_indiv,i) ;
        end
    end

        # Find the households of each individuals and the age category
    @inbounds for i=1:length(id_indiv)
        ind_age = findall(x.==get_prop(graph,id_indiv[i],:AgeClass)) ;
        push!(cat_age,ind_age[1]) ;
        neigh = neighbors(graph,id_indiv[i])[1] ;
        # Pour les individus ayant un ménage :
        if(get_prop(graph,neigh,:type)=="Household")
            push!(id_house,neighbors(graph,id_indiv[i])[1]) ;
        else # Les individus sont en hébergement
            push!(id_house,-id_indiv[i]-1) ;
        end
    end

    return(id_indiv,id_house,cat_age) ;

end

"""
Function DefStatut
#Fonction permettant de définir les statuts sociales des individus : actif, étudiant, écolier, retraité
#Entrées :
#- cat_age::Array{Int32,1}
#Sorties :
#- statut::Array{Int32,1}
"""
function def_statut(cat_age::Array{Int32,1})

    # Actif == 1
    statut = ones(Int32,length(cat_age)) ;
    # Etudiant == 2
    statut[findall(cat_age .== 5)] .= 2 ;
    # Ecolier == 3
    statut[findall(cat_age .<=4)] .= 3 ;
    # Retraite == 4
    statut[findall(cat_age .>= 9)] .= 4 ;

    return(statut)

end

"""
Function Indiv_InOut
#Fonction permettant de mettre en place les individus entrant et sortant du quariter
#Entrées :
#- statut::Array{Int32,1}
#- id_house::Array{Int32,1}
#Sorties :
#- indiv_in_out::Array{Int32,1}
#- statut::Array{Int32,1}
#- id_house::Array{Int32,1}
"""
function indiv_inout(statut::Array{Int32,1},id_house::Array{Int32,1})

    # Définition des sorties (travail + resto = -1)
    if((ParamPop.NB_OUT == 0) & (ParamPop.NB_IN == 0))
        indiv = zeros(Int32,length(statut)) ;
        return(indiv,statut,id_house) ;
    elseif((ParamPop.NB_OUT > 0) & (ParamPop.NB_IN == 0))
        indiv = def_indiv_out(statut) ;
        return(indiv,statut,id_house) ;
    elseif((ParamPop.NB_OUT == 0) & (ParamPop.NB_IN > 0))
        indiv_out = zeros(Int32,length(statut)) ;
        indiv_in_out,statut,id_house = def_indiv_in(indiv_out,statut,id_house) ;
        return(indiv_in_out,statut,id_house) ;
    elseif((ParamPop.NB_OUT > 0) & (ParamPop.NB_IN > 0))
        indiv = def_indiv_out(statut) ;
        indiv_in_out,statut,id_house = def_indiv_in(indiv,statut,id_house) ;
        return(indiv_in_out,statut,id_house) ;
    end

end

"""
Function Indiv_In
#Fonction permettant de mettre en place les individus qui entrent pour travailler au sein du quartier
#Entrées :
#- indiv_out::Array{Int32,1}
#- statut::Array{Int32,1}
#- id_house::Array{Int32,1}
#Sorties :
#- indiv_in_out::Array{Int32,1}
#- statut::Array{Int32,1}
#- id_house::Array{Int32,1}

"""
function def_indiv_in(indiv_out::Array{Int32,1},statut::Array{Int32,1},id_house::Array{Int32,1})

    indiv_in_out = zeros(Int32,length(statut)+ParamPop.NB_IN) ;
    indiv_in_out[1:length(statut)] .= indiv_out ;
    # Il faudra mettre à -1 les gens qui entrent pour travailler.
    indiv_in_out[length(statut)+1:length(statut)+ParamPop.NB_IN] .= -1 ;
    # Mettre à jour le statut
    nb_etud = Int(floor(ParamPop.NB_IN*ParamPop.TAUX_ETUD)) ;
    statut_out = ones(Int32,ParamPop.NB_IN) ;
    statut_out[1:nb_etud] .= 2 ;

    id_house_out = zeros(Int32,ParamPop.NB_IN) ;
    id_house_out .= -1 ;

    @inbounds for i=1:length(statut_out)
        push!(statut,statut_out[i]) ;
        push!(id_house,id_house_out[i]) ;
    end

    return(indiv_in_out, statut, id_house) ;

end

"""
Function Indiv_Out
#Fonction permettant de mettre en place les individus qui sortent pour travailler en dehors du quartier
#Entrée :
#- statut::Array{Int32,1}
#Sorties :
#- indiv_out::Array{Int32,1}
"""
function def_indiv_out(statut::Array{Int32,1})

    etud_actif = findall(statut .<= 2) ;
    indiv_out = zeros(Int32,length(statut)) ;
    out = zeros(Int32,length(etud_actif)) ;
    out[1:ParamPop.NB_OUT] .= 1 ;
    shuffle!(out) ;
    indiv_out[etud_actif] = out ;

    return(indiv_out) ;

end

"""
Function ChoixTravail
#Fonction permettant d'associer une entreprise à un individu ainsi que le tissu social (~collègues proches)
#Entrées :
#- statut::Array{Int32,1}
#- indiv_in_out::Array{Int32,1}
#- id_work::Array{Int32,1}
#- area_work::Array{Float64,1}
#Sorties :
#- id_ind_work::Array{Int32,1}
#- tissu_work::Array{Int32,1}
"""
function choix_travail(statut::Array{Int32,1},indiv_in_out::Array{Int32,1},id_work::Array{Int32,1},area_work::Array{Float64,1})

    actif = findall(statut .== 1) ;
    ecolier = findall(statut .== 3) ;

    etudiant_interieur = findall((indiv_in_out .< 1) .& (statut .== 2)) ;

    id_ind_work = zeros(Int32,length(statut)) ;

    q_etud = div(length(etudiant_interieur),ParamPop.NBPERAMPHI) ;
    q_ecol = div(length(ecolier), ParamPop.NBPERCLASS) ;

    Repartition = distribution_workingplace(actif,id_work,area_work) ;

    i_class = 300000 ;
    i_amphi = 400000 ;
    cpt = 1 ;

    @inbounds for i=1:length(statut)

        if((statut[i] == 1) & (indiv_in_out[i] <= 0)) #actif au sein du 5ème
            id_ind_work[i] = Repartition[cpt] ;
            cpt = cpt + 1 ;
        elseif((statut[i] == 1) & (indiv_in_out[i] == 1)) #actif extérieur
            id_ind_work[i] = -1 ;
        elseif((statut[i] == 2) & (indiv_in_out[i] < 1)) # etudiant au sein du 5ème
            id_ind_work[i] = floor(Int,rand(Uniform(1,q_etud+1))) + i_amphi ;
        elseif((statut[i] == 2) & (indiv_in_out[i] == 1)) # etudiant extérieur
            id_ind_work[i] = -1 ;
        elseif(statut[i] == 4) # retraité
            id_ind_work[i] = -1 ;
        elseif(statut[i] == 3)
            id_ind_work[i] = floor(Int,rand(Uniform(1,q_ecol+1))) + i_class ;
        end

    end

    # Définition du tissu social
    tissu_work = tissu(id_ind_work) ;

    return(id_ind_work, tissu_work) ;

end

"""
Function RepartitionEntreprise
#Fonction permettant de repartir les individus au sein des entreprises
#Entrées :
#- actif::Array{Int64,1}
#- id_work::Array{Int32,1}
#- area_work::Array{Float64,1}
#Sorties :
#- Repartition::Array{Int64,1}
"""
function distribution_workingplace(actif::Array{Int64,1},id_work::Array{Int32,1},area_work::Array{Float64,1})

    nb_entr = length(id_work) ;
    nb_mean_per_entr = floor.(Int32,div.(area_work,ParamPop.M2_EMPLOY)) ;
    nb_total_mean = sum(nb_mean_per_entr) ;
    repartition = [] ;

    @inbounds for i=1:nb_entr
        append!(repartition,ones(Int32,nb_mean_per_entr[i])*id_work[i]) ;
    end

    reste = length(actif) - nb_total_mean ;

    if(reste > 0)
        for j=1:reste
            append!(repartition, id_work[floor(Int,rand(Uniform(1,nb_entr)))]) ;
        end
    end

    shuffle!(repartition) ;

    return(repartition) ;

end

"""
Function Tissu_social
#Fonction permettant de mettre en place un tissu social au sein de l'entreprise. Ce sont des personnes qui seront plus proches entre eux.
#Entrées :
#- Structure::Array{Int32,1}
#Sorties :
#- Structure::Array{Int32,1}
"""
function tissu_social(Structure::Array{Int32,1})

    nb_actif = length(Structure) ;
    cpt = nb_actif ;
    cpt2 = 0 ;
    fin = 1 ;

    while(cpt >= 0)

        s = rand(Poisson(ParamPop.LAMBDA_TISSU)) ;
        cpt = cpt - s ;
        deb = fin ;
        fin = deb+s ;
        if(fin>nb_actif)
            fin = nb_actif ;
        end

        Structure[deb:fin] .= cpt2 ;
        cpt2 = cpt2 + 1 ;

    end

    return(Structure) ;

end

"""
Function Tissu

"""
function tissu(id_ind_work::Array{Int32,1})

    tissu = zeros(Int32,length(id_ind_work)) ;
    d = Dict([(i,count(x->x.==i,id_ind_work)) for i in unique(id_ind_work)]) ;
    dr = collect(keys(d)) ;

    @inbounds for i=1:length(dr)

        if(dr[i] > 0)
            Ind_work = findall(id_ind_work .== dr[i]) ;
            Structure = zeros(Int32,length(Ind_work)) ;
            tissu_s = tissu_social(Structure) ;
            tissu[Ind_work] = tissu_s ;
        end
    end

    return(tissu) ;

end


############################################################
# judiModeling ###############################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: January 2017

export judiModeling, judiModelingException, judiSetupModeling, judiModelingAdjoint, judiModelingAdjointException, judiSetupModelingAdjoint, subsample

############################################################

# Type for linear operator representing  Pr*A(m)^-1*Ps,
# i.e. it includes source and receiver projections
struct judiModeling{DDT<:Number,RDT<:Number} <: joAbstractLinearOperator{DDT,RDT}
    name::String
    m::Integer
    n::Integer
    info::Info
    model::Model
    options::Options
	fop::Function              # forward
	fop_T::Function  # transpose
end

mutable struct judiModelingException <: Exception
    msg :: String
end

struct judiModelingAdjoint{DDT,RDT} <: joAbstractLinearOperator{DDT,RDT}
    name::String
    m::Integer
    n::Integer
    info::Info
    model::Model
    options::Options
	fop::Function              # forward
	fop_T::Function  # transpose
end

mutable struct judiModelingAdjointException <: Exception
    msg :: String
end

############################################################
## outer constructors
"""
    judiModeling(info, model; options=Options())
    judiModeling(info, model, src_geometry, rec_geometry; options=Options())

Create seismic modeling operator for a velocity model given as a `Model` structure. `info` is an `Info` structure\\
containing necessary dimensions to set up the operator. The function also takes the source and receiver geometries\\
as additional input arguments, which creates a combined operator `judiProjection*judiModeling*judiProjection'`.

Example
=======

`Pr` and `Ps` are projection operatos of type `judiProjection` and\\
`q` is a data vector of type `judiVector`:

    F = judiModeling(info, model)
    dobs = Pr*F*Ps'*q

    F = judiModeling(info, model, q.geometry, rec_geometry)
    dobs = F*q

"""
function judiModeling(info::Info, model::Model; options=Options(), DDT::DataType=Float32, RDT::DataType=DDT)
# JOLI wrapper for nonlinear forward modeling
	(DDT == Float32 && RDT == Float32) || throw(judiModelingException("Domain and range types not supported"))
	m = info.n * sum(info.nt)
	n = m
	srcnum = 1:info.nsrc

	return F = judiModeling{Float32,Float32}("forward wave equation", m, n, info, model, options,
							  args -> time_modeling(model, args[1], args[2], args[3], args[4], args[5], srcnum, 'F', 1, options),
							  args_T -> time_modeling(model, args_T[1], args_T[2], args_T[3], args_T[4], args_T[5], srcnum, 'F', -1, options)
							  )
end

function judiModelingAdjoint(info::Info, model::Model; options=Options(), DDT::DataType=Float32, RDT::DataType=DDT)
# JOLI wrapper for nonlinear forward modeling
	(DDT == Float32 && RDT == Float32) || throw(judiModelingAdjointException("Domain and range types not supported"))
	m = info.n * sum(info.nt)
	n = m
	srcnum = 1:info.nsrc

	return F = judiModelingAdjoint{Float32,Float32}("adjoint wave equation", m, n, info, model, options,
							  args_T -> time_modeling(model, args_T[1], args_T[2], args_T[3], args_T[4], args_T[5], srcnum, 'F', -1, options),
							  args -> time_modeling(model, args[1], args[2], args[3], args[4], args[5], srcnum, 'F', 1, options)
							  )
end

############################################################
## overloaded Base functions

# conj(jo)
conj(A::judiModeling{DDT,RDT}) where {DDT,RDT} =
    judiModeling{DDT,RDT}("conj("*A.name*")",A.m,A.n,A.info,A.model,A.options,A.fop, A.fop_T)

# transpose(jo)
transpose(A::judiModeling{DDT,RDT}) where {DDT,RDT} =
    judiModelingAdjoint{DDT,RDT}("adjoint wave equation",A.n,A.m,A.info,A.model,A.options,A.fop_T, A.fop)

adjoint(A::judiModeling{DDT,RDT}) where {DDT,RDT} =
	judiModelingAdjoint{DDT,RDT}("adjoint wave equation",A.n,A.m,A.info,A.model,A.options,A.fop_T, A.fop)

# conj(jo)
conj(A::judiModelingAdjoint{DDT,RDT}) where {DDT,RDT} =
    judiModelingAdjoint{DDT,RDT}("conj("*A.name*")",A.m,A.n,A.info,A.model,A.options,A.fop, A.fop_T)

# transpose(jo)
transpose(A::judiModelingAdjoint{DDT,RDT}) where {DDT,RDT} =
    judiModeling{DDT,RDT}("forward wave equation",A.n,A.m,A.info,A.model,A.options,A.fop_T, A.fop)

adjoint(A::judiModelingAdjoint{DDT,RDT}) where {DDT,RDT} =
	judiModeling{DDT,RDT}("forward wave equation",A.n,A.m,A.info,A.model,A.options,A.fop_T, A.fop)

############################################################
## Additional overloaded functions

# *(judiModeling,judiWavefield)
function *(A::judiModeling{ADDT,ARDT}, v::judiWavefield{vDT}) where {ADDT,ARDT,vDT}
	A.n == size(v,1) || throw(judiModelingException("Shape mismatch: A:$(size(A)), v: $(size(v))"))
	jo_check_type_match(ADDT,vDT,join(["DDT for *(judiModeling,judiWavefield):",A.name,typeof(A),vDT]," / "))
	args = (nothing,v.data,nothing,nothing,nothing)
	V = A.fop(args)
	jo_check_type_match(ARDT,eltype(V),join(["RDT from *(judiModeling,judiWavefield):",A.name,typeof(A),eltype(V)]," / "))
	return V
end

# *(judiModelingAdjoint,judiWavefield)
function *(A::judiModelingAdjoint{ADDT,ARDT}, v::judiWavefield{vDT}) where {ADDT,ARDT,vDT}
	A.n == size(v,1) || throw(judiModelingAdjointException("Shape mismatch: A:$(size(A)), v: $(size(v))"))
	jo_check_type_match(ADDT,vDT,join(["DDT for *(judiModeling,judiWavefield):",A.name,typeof(A),vDT]," / "))
	args = (nothing,nothing,nothing,v.data,nothing)
	V = A.fop(args)
	jo_check_type_match(ARDT,eltype(V),join(["RDT from *(judiModeling,judiWavefield):",A.name,typeof(A),eltype(V)]," / "))
	return V
end

# *(judiModeling,judiRHS)
function *(A::judiModeling{ADDT,ARDT}, v::judiRHS{vDT}) where {ADDT,ARDT,vDT}
	A.n == size(v,1) || throw(judiModelingException("Shape mismatch: A:$(size(A)), v: $(size(v))"))
	jo_check_type_match(ADDT,vDT,join(["DDT for *(judiModeling,judiRHS):",A.name,typeof(A),vDT]," / "))
	args = (v.geometry,v.data,nothing,nothing,nothing)
	V = A.fop(args)
	jo_check_type_match(ARDT,eltype(V),join(["RDT from *(judiModeling,judiRHS):",A.name,typeof(A),eltype(V)]," / "))
	return V
end

# *(judiModelingAdjoint,judiRHS)
function *(A::judiModelingAdjoint{ADDT,ARDT}, v::judiRHS{vDT}) where {ADDT,ARDT,vDT}
	A.n == size(v,1) || throw(judiModelingException("Shape mismatch: A:$(size(A)), v: $(size(v))"))
	jo_check_type_match(ADDT,vDT,join(["DDT for *(judiModeling,judiRHS):",A.name,typeof(A),vDT]," / "))
	args = (nothing, nothing, v.geometry,v.data,nothing)
	V = A.fop(args)
	jo_check_type_match(ARDT,eltype(V),join(["RDT from *(judiModeling,judiRHS):",A.name,typeof(A),eltype(V)]," / "))
	return V
end

############################################################
## Additional overloaded functions

# Subsample Modeling function
function subsample(F::judiModeling, srcnum)
    info = Info(F.info.n, length(srcnum), F.info.nt[srcnum])
    return judiModeling(info, F.model, options=F.options)
end

function subsample(F::judiModelingAdjoint, srcnum)
    info = Info(F.info.n, length(srcnum), F.info.nt[srcnum])
    return judiModelingAdjoint(info, F.model, options=F.options)
end

getindex(F::judiModeling,a) = subsample(F,a)
getindex(F::judiModelingAdjoint,a) = subsample(F,a)


# Define symbol inputs for different variable types
const Parameter = :Parameter

# Extend Base.copy for infinite parameters
function Base.copy(p::ParameterRef, new_model::InfiniteModel)::ParameterRef
    return ParameterRef(new_model, p.index)
end

# Internal structure for building InfOptParameters
mutable struct _ParameterInfoExpr
    has_lb::Bool
    lower_bound::Any
    has_ub::Bool
    upper_bound::Any
    has_dist::Bool
    distribution::Any
    has_set::Bool
    set::Any
end

# Default constructor
function _ParameterInfoExpr(; lower_bound = NaN, upper_bound = NaN,
                            distribution = NaN, set = NaN)
    # isnan(::Expr) is not defined so we need to do !== NaN
    return _ParameterInfoExpr(lower_bound !== NaN, lower_bound,
                              upper_bound !== NaN, upper_bound,
                              distribution !== NaN, distribution,
                              set !== NaN, set)
end

# Internal function for use in processing valid key word arguments
function _is_set_keyword(kw::Expr)
    return kw.args[1] in [:set, :lower_bound, :upper_bound, :distribution]
end

# Extend to assist in building InfOptParameters
function JuMP._set_lower_bound_or_error(_error::Function,
                                        info::_ParameterInfoExpr, lower)
    info.has_lb && _error("Cannot specify parameter lower_bound twice")
    info.has_dist && _error("Cannot specify parameter lower_bound and " *
                            "distribution")
    info.has_set && _error("Cannot specify parameter lower_bound and set")
    info.has_lb = true
    info.lower_bound = lower
    return
end

# Extend to assist in building InfOptParameters
function JuMP._set_upper_bound_or_error(_error::Function,
                                        info::_ParameterInfoExpr, upper)
    info.has_ub && _error("Cannot specify parameter upper_bound twice")
    info.has_dist && _error("Cannot specify parameter upper_bound and " *
                            "distribution")
    info.has_set && _error("Cannot specify parameter upper_bound and set")
    info.has_ub = true
    info.upper_bound = upper
    return
end

# Extend to assist in building InfOptParameters
function _dist_or_error(_error::Function, info::_ParameterInfoExpr, dist)
    info.has_dist && _error("Cannot specify parameter distribution twice")
    (info.has_lb || info.has_ub) && _error("Cannot specify parameter " *
                                           "distribution and upper/lower bounds")
    info.has_set && _error("Cannot specify parameter distribution and set")
    info.has_dist = true
    info.distribution = dist
    return
end

# Extend to assist in building InfOptParameters
function _set_or_error(_error::Function, info::_ParameterInfoExpr, set)
    info.has_set && _error("Cannot specify variable fixed value twice")
    (info.has_lb || info.has_ub) && _error("Cannot specify parameter set and " *
                                           "upper/lower bounds")
    info.has_dist && _error("Cannot specify parameter set and distribution")
    info.has_set = true
    info.set = set
    return
end

# Construct an expression to build an infinite set (use with @infinite_macro)
function _constructor_set(_error::Function, info::_ParameterInfoExpr)
    if (info.has_lb || info.has_ub) && !(info.has_lb && info.has_ub)
        _error("Must specify both an upper bound and a lower bound")
    elseif info.has_lb
        check = :(isa($(info.lower_bound), Number))
        return :($(check) ? IntervalSet($(info.lower_bound), $(info.upper_bound)) : error("Bounds must be a number."))
    elseif info.has_dist
        check = :(isa($(info.distribution), Distributions.NonMatrixDistribution))
        return :($(check) ? DistributionSet($(info.distribution)) : error("Distribution must be a subtype of Distributions.NonMatrixDistribution."))
    elseif info.has_set
        check1 = :(isa($(info.set), AbstractInfiniteSet))
        check2 = :(isa($(info.set), Distributions.NonMatrixDistribution))
        return :($(check1) ? $(info.set) : ($(check2) ? DistributionSet($(info.set)) : error("Set must be a subtype of AbstractInfiniteSet.")))
    else
        _error("Must specify upper/lower bounds, a distribution, or a set")
    end
end

# Check that supports don't violate the set bounds
function _check_supports_in_bounds(_error::Function,
                                   supports::Union{Number, Vector{<:Number}},
                                   set::AbstractInfiniteSet)
    if !supports_in_set(supports, set)
        _error("Supports violate the set domain bounds.")
    end
    return
end

"""
    build_parameter(_error::Function, set::AbstractInfiniteSet,
                    [num_params::Int = 1; num_supports::Int = 0,
                    supports::Union{Number, Vector{<:Number}} = Number[],
                    independent::Bool = false,
                    sig_fig::Int = 5])::InfOptParameter

Returns a [`InfOptParameter`](@ref) given the appropriate information. This is
analagous to `JuMP.build_variable`. Errors if supports violate the bounds
associated `set`. Also errors if `set` contains a multivariate distribution with
a different dimension than `num_params`. This is meant to primarily serve as a
helper method for [`@infinite_parameter`](@ref).

**Example**
```jldoctest; setup = :(using InfiniteOpt)
julia> build_parameter(error, IntervalSet(0, 3), supports = Vector(0:3))
InfOptParameter{IntervalSet}([0, 3], [0, 1, 2, 3], false)
```
"""
function build_parameter(_error::Function, set::AbstractInfiniteSet,
                         num_params::Int = 1;
                         num_supports::Int = 0,
                         supports::Union{Number, Vector{<:Number}} = Number[],
                         independent::Bool = false,
                         sig_fig::Int = 5,
                         extra_kw_args...)::InfOptParameter
    for (kwarg, _) in extra_kw_args
        _error("Unrecognized keyword argument $kwarg")
    end
    length_supports = length(supports)
    if num_supports == 0 && length_supports != 0
        _check_supports_in_bounds(_error, supports, set)
    elseif num_supports != 0 && length_supports != 0
        @warn("Ignoring num_supports since supports is not empty.")
        _check_supports_in_bounds(_error, supports, set)
    elseif num_supports != 0 && length_supports == 0
        if isa(set, DistributionSet{<:Distributions.MultivariateDistribution})
            _error("Support generation is not available for multivariate " *
                   "distributions.")
        end
        supports = generate_support_values(set, num_supports = num_supports,
                                           sig_fig = sig_fig)
    end
    if isa(set, DistributionSet{<:Distributions.MultivariateDistribution})
        if num_params != length(set.distribution)
            _error("Multivariate distribution dimension must match dimension " *
                   "of parameter.")
        end
    end
    if num_params == 1 || independent # double check num_params for other types of containers
        unique_supports = unique(supports)
        if length(unique_supports) != length(supports)
            @warn("Support points are not unique, eliminating redundant points.")
        end
        return InfOptParameter(set, unique_supports, independent)
    end
    if isa(supports, Number)
        supports = [supports]
    end
    return InfOptParameter(set, supports, independent)
end

# Check the number of supports of one dimension matches the other dimension
# within the same multi-dimensional parameter
function _check_supports_dimensions(model::InfiniteModel, p::InfOptParameter,
    index::Int)
    if haskey(model.params, index - 1) &&
       model.param_to_group_id[index - 1] == model.next_param_id
        if length(p.supports) != length(model.params[index - 1].supports)
            error("Support dimension mismatch. Make sure support dimension of " *
                  "each parameter element match.")
        end
    end
    return
end

"""
    add_parameter(model::InfiniteModel, p::InfOptParameter,
                  [name::String = ""])::ParameterRef

Returns a [`ParameterRef`](@ref) associated with the parameter `p` that is added
to `model`. This adds a parameter to the model in a manner similar to
`JuMP.add_variable`. This can be used to add parameters with the use of
[`@infinite_parameter`](@ref). [`build_parameter`](@ref) should be used to
construct `p`.

**Example**
```jldoctest; setup = :(using InfiniteOpt; model = InfiniteModel())
julia> p = build_parameter(error, IntervalSet(0, 3), supports = Vector(0:3))
InfOptParameter{IntervalSet}([0, 3], [0, 1, 2, 3], false)

julia> param_ref = add_parameter(model, p, "name")
name
```
"""
function add_parameter(model::InfiniteModel, p::InfOptParameter,
                       name::String=""; multi_dim = false,
                       macro_call = false)::ParameterRef
    index = model.next_param_index += 1
    pref = ParameterRef(model, index)
    if !macro_call
        model.next_param_id += 1
    elseif !(p.independent) && multi_dim
        _check_supports_dimensions(model, p, index)
    end
    model.params[JuMP.index(pref)] = p
    model.param_to_group_id[JuMP.index(pref)] = model.next_param_id
    JuMP.set_name(pref, name)
    return pref
end

"""
    used_by_constraint(pref::ParameterRef)::Bool

Return true if `pref` is used by a constraint or false otherwise.

**Example**
```julia-repl
julia> used_by_constraint(t)
true
```
"""
function used_by_constraint(pref::ParameterRef)::Bool
    return haskey(JuMP.owner_model(pref).param_to_constrs, JuMP.index(pref))
end

"""
    used_by_measure(pref::ParameterRef)::Bool

Return true if `pref` is used by a measure or false otherwise.

**Example**
```julia-repl
julia> used_by_measure(t)
false
```
"""
function used_by_measure(pref::ParameterRef)::Bool
    return haskey(JuMP.owner_model(pref).param_to_meas, JuMP.index(pref))
end

"""
    used_by_variable(pref::ParameterRef)::Bool

Return true if `pref` is used by an infinite variable or false otherwise.

**Example**
```julia-repl
julia> used_by_variable(t)
true
```
"""
function used_by_variable(pref::ParameterRef)::Bool
    return haskey(JuMP.owner_model(pref).param_to_vars, JuMP.index(pref))
end

"""
    is_used(pref::ParameterRef)::Bool

Return true if `pref` is used in the model or false otherwise.

**Example**
```julia-repl
julia> is_used(t)
true
```
"""
function is_used(pref::ParameterRef)::Bool
    return used_by_measure(pref) || used_by_constraint(pref) || used_by_variable(pref)
end

# Check if parameter is used by measure data and error if it is to prevent bad
# deleting behavior
function _check_param_in_data(pref::ParameterRef, data::AbstractMeasureData)
    prefs = parameter_refs(data)
    if (pref == prefs || pref in prefs)
        error("Unable to delete `$pref` since it is used to evaluate measures.")
    end
    return
end

## Determine if a tuple element contains a particular parameter
# ParameterRef
function _contains_pref(search_pref::ParameterRef, pref::ParameterRef)::Bool
    return search_pref == pref
end

# SparseAxisArray
function _contains_pref(arr::JuMPC.SparseAxisArray{<:ParameterRef},
                        pref::ParameterRef)::Bool
    return pref in collect(values(arr.data))
end

# Return parameter tuple without a particular parameter and return location of
# where it was
function _remove_parameter(prefs::Tuple, delete_pref::ParameterRef)::Tuple
    for i in eachindex(prefs)
        if _contains_pref(prefs[i], delete_pref)
            if isa(prefs[i], ParameterRef)
                return Tuple(prefs[j] for j = 1:length(prefs) if j != i), (i, )
            else
                for (k, v) in prefs[i].data
                    if v == delete_pref
                        new_dict = filter(x -> x.second != delete_pref,
                                         prefs[i].data)
                        if length(new_dict) != 0
                            pref_list = [prefs...]
                            pref_list[i] = JuMPC.SparseAxisArray(new_dict)
                            return Tuple(pref_list[j] for j = 1:length(prefs)), (i, k)
                        else
                            return Tuple(prefs[j] for j = 1:length(prefs) if j != i), (i, )
                        end
                    end
                end
            end
        end
    end
end

# Used to update infinite variable when one of its parameters is deleted
function _update_infinite_variable(vref::InfiniteVariableRef, new_prefs::Tuple)
    _update_variable_param_refs(vref, new_prefs)
    JuMP.set_name(vref, _root_name(vref))
    if used_by_measure(vref)
        for mindex in JuMP.owner_model(vref).var_to_meas[JuMP.index(vref)]
            JuMP.set_name(MeasureRef(JuMP.owner_model(vref), mindex),
                       _make_meas_name(JuMP.owner_model(vref).measures[mindex]))
        end
    end
    return
end

# Return a parameter value tuple without element at a particular location
function _remove_parameter_values(pref_vals::Tuple, location::Tuple)::Tuple
    # removed parameter was a scalar parameter
    if length(location) == 1
        return Tuple(pref_vals[i] for i in eachindex(pref_vals) if i != location[1])
    # removed parameter was part of an array
    else
        new_dict = filter(x -> x.first != location[2], pref_vals[location[1]].data)
        val_list = [pref_vals...]
        val_list[location[1]] = JuMPC.SparseAxisArray(new_dict)
        return Tuple(val_list[i] for i in eachindex(pref_vals))
    end
end

# Update point variable for which a parameter is deleted
function _update_point_variable(pvref::PointVariableRef, pref_vals::Tuple)
    _update_variable_param_values(pvref, pref_vals)
    # update name if no alias was provided
    if !isa(findfirst(isequal('('), JuMP.name(pvref)), Nothing)
        JuMP.set_name(pvref, "")
    end
    if used_by_measure(pvref)
        for mindex in JuMP.owner_model(pvref).var_to_meas[JuMP.index(pvref)]
            JuMP.set_name(MeasureRef(JuMP.owner_model(pvref), mindex),
                      _make_meas_name(JuMP.owner_model(pvref).measures[mindex]))
        end
    end
    return
end

# Update a reduced variable associated with an infinite variable whose parameter
# was removed
function _update_reduced_variable(vref::ReducedInfiniteVariableRef,
                                  location::Tuple)
    eval_supps = eval_supports(vref)
    # removed parameter was a scalar
    if length(location) == 1
        new_supports = Dict{Int, Union{Number, JuMPC.SparseAxisArray}}()
        for (index, support) in eval_supps
            if index < location[1]
                new_supports[index] = support
            elseif index > location[1]
                new_supports[index - 1] = support
            end
        end
        JuMP.owner_model(vref).reduced_info[JuMP.index(vref)] = ReducedInfiniteInfo(infinite_variable_ref(vref), new_supports)
    # removed parameter was part of an array and was reduced previously
    elseif haskey(eval_supps, location[1])
        new_dict = filter(x -> x.first != location[2], eval_supps[location[1]].data)
        new_supports = copy(eval_supps)
        new_supports[location[1]] = JuMPC.SparseAxisArray(new_dict)
        JuMP.owner_model(vref).reduced_info[JuMP.index(vref)] = ReducedInfiniteInfo(infinite_variable_ref(vref), new_supports)
    end
    if used_by_measure(vref)
        for mindex in JuMP.owner_model(vref).reduced_to_meas[JuMP.index(vref)]
            JuMP.set_name(MeasureRef(JuMP.owner_model(vref), mindex),
                       _make_meas_name(JuMP.owner_model(vref).measures[mindex]))
        end
    end
    return
end

"""
    JuMP.delete(model::InfiniteModel, pref::ParameterRef)

Extend [`JuMP.delete`](@ref JuMP.delete(::JuMP.Model, ::JuMP.VariableRef)) to delete
infinite parameters and their dependencies. All variables, constraints, and
measure functions that depend on `pref` are updated to exclude it. Errors if the
parameter is contained in an `AbstractMeasureData` datatype that is employed by
a measure since the measure becomes invalid otherwise. Thus, measures that
contain this dependency must be deleted first. Note that
```_check_param_in_data(pref, measure_data)``` needs to be extended to allow
deletion of parameters when custom `AbstractMeasureData` datatypes are used.

**Example**
```julia-repl
julia> print(model)
Min measure(g(t, x)*t + x) + z
Subject to
 z >= 0.0
 g(t, x) + z >= 42.0
 g(0.5, x) == 0
 t in [0, 6]
 x in [0, 1]

julia> delete(model, x)

julia> print(model)
Min measure(g(t)*t) + z
Subject to
 g(t) + z >= 42.0
 g(0.5) == 0
 z >= 0.0
 t in [0, 6]
```
"""
function JuMP.delete(model::InfiniteModel, pref::ParameterRef)
    @assert JuMP.is_valid(model, pref) "Parameter reference is invalid."
    # update optimizer model status
    if is_used(pref)
        set_optimizer_model_ready(model, false)
    end
    # update measures
    if used_by_measure(pref)
        # ensure deletion is okay (pref isn't used by measure data)
        for mindex in model.param_to_meas[JuMP.index(pref)]
            _check_param_in_data(pref, model.measures[mindex].data)
        end
        # delete dependence of measures on pref
        for mindex in model.param_to_meas[JuMP.index(pref)]
            if isa(model.measures[mindex].func, ParameterRef)
                model.measures[mindex] = Measure(zero(JuMP.AffExpr),
                                                 model.measures[mindex].data)
            else
                _remove_variable(model.measures[mindex].func, pref)
            end
            JuMP.set_name(MeasureRef(model, mindex),
                          _make_meas_name(model.measures[mindex]))
        end
        # delete mapping
        delete!(model.param_to_meas, JuMP.index(pref))
    end
    # update variables
    if used_by_variable(pref)
        # update infinite variables that depend on pref
        for vindex in model.param_to_vars[JuMP.index(pref)]
            # find location of parameter in storage tuple
            prefs, location = _remove_parameter(model.vars[vindex].parameter_refs,
                                                pref)
            vref = InfiniteVariableRef(model, vindex)
            # remove the parameter dependence
            _update_infinite_variable(vref, prefs)
            # update any point variables that depend on vref accordingly
            if used_by_point_variable(vref)
                for pindex in model.infinite_to_points[vindex]
                    pvref = PointVariableRef(model, pindex)
                    pref_vals = _remove_parameter_values(parameter_values(pvref),
                                                         location)
                    _update_point_variable(pvref, pref_vals)
                end
            end
            # update any reduced variables that depend on vref accordingly
            if used_by_reduced_variable(vref)
                for rindex in model.infinite_to_reduced[vindex]
                    _update_reduced_variable(ReducedInfiniteVariableRef(model,
                                                              rindex), location)
                end
            end
        end
        # delete mapping
        delete!(model.param_to_vars, JuMP.index(pref))
    end
    # update constraints
    if used_by_constraint(pref)
        # update constraints in mapping to remove the parameter
        for cindex in model.param_to_constrs[JuMP.index(pref)]
            if isa(model.constrs[cindex].func, ParameterRef)
                model.constrs[cindex] = JuMP.ScalarConstraint(zero(JuMP.AffExpr),
                                                      model.constrs[cindex].set)
            else
                _remove_variable(model.constrs[cindex].func, pref)
            end
        end
        # delete mapping
        delete!(model.param_to_constrs, JuMP.index(pref))
    end
    # delete parameter information stored in model
    delete!(model.params, JuMP.index(pref))
    delete!(model.param_to_name, JuMP.index(pref))
    delete!(model.param_to_group_id, JuMP.index(pref))
    return
end

"""
    JuMP.is_valid(model::InfiniteModel, pref::ParameterRef)::Bool

Extend the [`JuMP.is_valid`](@ref JuMP.is_valid(::JuMP.Model, ::JuMP.VariableRef))
function to accomodate infinite parameters.
Returns true if the `InfiniteModel` stored in `pref` matches `model` and if
the parameter index is used by `model`. It returns false otherwise.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> is_valid(model, t)
true
```
"""
function JuMP.is_valid(model::InfiniteModel, pref::ParameterRef)::Bool
    check1 = model === JuMP.owner_model(pref)
    check2 = JuMP.index(pref) in keys(model.params)
    return check1 && check2
end

"""
    JuMP.name(pref::ParameterRef)::String

Extend the [`JuMP.name`](@ref JuMP.name(::JuMP.VariableRef)) function to
accomodate infinite parameters. Returns the name string associated with `pref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> name(t)
"t"
```
"""
function JuMP.name(pref::ParameterRef)::String
    return JuMP.owner_model(pref).param_to_name[JuMP.index(pref)]
end

"""
    JuMP.set_name(pref::ParameterRef, name::String)

Extend the [`JuMP.set_name`](@ref JuMP.set_name(::JuMP.VariableRef, ::String))
function to accomodate infinite parameters. Set a new base name to be associated
with `pref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> set_name(t, "time")

julia> name(t)
"time"
```
"""
function JuMP.set_name(pref::ParameterRef, name::String)
    JuMP.owner_model(pref).param_to_name[JuMP.index(pref)] = name
    JuMP.owner_model(pref).name_to_param = nothing
    return
end

"""
    num_parameters(model::InfiniteModel)::Int

Return the number of infinite parameters currently present in `model`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> num_parameters(model)
1
```
"""
function num_parameters(model::InfiniteModel)::Int
    return length(model.params)
end

# Internal functions
_parameter_set(pref::ParameterRef) = JuMP.owner_model(pref).params[JuMP.index(pref)].set
_parameter_supports(pref::ParameterRef) = JuMP.owner_model(pref).params[JuMP.index(pref)].supports
function _update_parameter_set(pref::ParameterRef, set::AbstractInfiniteSet)
    supports = JuMP.owner_model(pref).params[JuMP.index(pref)].supports
    independent = JuMP.owner_model(pref).params[JuMP.index(pref)].independent
    JuMP.owner_model(pref).params[JuMP.index(pref)] = InfOptParameter(set, supports, independent)
    if is_used(pref)
        set_optimizer_model_ready(JuMP.owner_model(pref), false)
    end
    return
end
function _update_parameter_supports(pref::ParameterRef, supports::Vector{<:Number})
    set = JuMP.owner_model(pref).params[JuMP.index(pref)].set
    independent = JuMP.owner_model(pref).params[JuMP.index(pref)].independent
    JuMP.owner_model(pref).params[JuMP.index(pref)] = InfOptParameter(set, supports, independent)
    if is_used(pref)
        set_optimizer_model_ready(JuMP.owner_model(pref), false)
    end
    return
end

"""
    infinite_set(pref::ParameterRef)::AbstractInfiniteSet

Return the infinite set associated with `pref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> infinite_set(t)
[0, 1]
```
"""
function infinite_set(pref::ParameterRef)::AbstractInfiniteSet
    return _parameter_set(pref)
end

"""
    set_infinite_set(pref::ParameterRef, set::AbstractInfiniteSet)

Specify the infinite set of `pref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> set_infinite_set(t, IntervalSet(0, 2))

julia> infinite_set(t)
[0, 2]
```
"""
function set_infinite_set(pref::ParameterRef, set::AbstractInfiniteSet)
    _update_parameter_set(pref, set)
    return
end

"""
    JuMP.has_lower_bound(pref::ParameterRef)::Bool

Extend the `JuMP.has_lower_bound` function to accomodate infinite parameters.
Return true if the set associated with `pref` has a defined lower bound or if a
lower bound can be found. Extensions with user-defined infinite set types
should extend `JuMP.has_lower_bound(set::NewType)`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> has_lower_bound(t)
true
```
"""
function JuMP.has_lower_bound(pref::ParameterRef)::Bool
    set = _parameter_set(pref)
    return JuMP.has_lower_bound(set)
end

"""
    JuMP.lower_bound(pref::ParameterRef)::Number

Extend the `JuMP.lower_bound` function to accomodate infinite parameters.
Returns the lower bound associated with the infinite set. Errors if such a bound
is not well-defined.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> lower_bound(t)
0.0
```
"""
function JuMP.lower_bound(pref::ParameterRef)::Number
    set = _parameter_set(pref)
    if !JuMP.has_lower_bound(pref)
        error("Parameter $(pref) does not have a lower bound.")
    end
    return JuMP.lower_bound(set)
end

"""
    JuMP.set_lower_bound(pref::ParameterRef, lower::Number)

Extend the `JuMP.set_lower_bound` function to accomodate infinite parameters.
Updates the infinite set lower bound if such an operation is supported. Set
extensions that seek to employ this should extend
`JuMP.set_lower_bound(set::NewType, lower::Number)`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> set_lower_bound(t, -1)

julia> lower_bound(t)
-1.0
```
"""
function JuMP.set_lower_bound(pref::ParameterRef, lower::Number)
    set = _parameter_set(pref)
    new_set = JuMP.set_lower_bound(set, lower)
    _update_parameter_set(pref, new_set)
    return
end

"""
    JuMP.has_upper_bound(pref::ParameterRef)::Bool

Extend the `JuMP.has_upper_bound` function to accomodate infinite parameters.
Return true if the set associated with `pref` has a defined upper bound or if a
upper bound can be found. Extensions with user-defined sets should extend
`JuMP.has_upper_bound(set::NewType)`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> has_upper_bound(t)
true
```
"""
function JuMP.has_upper_bound(pref::ParameterRef)::Bool
    set = _parameter_set(pref)
    return JuMP.has_upper_bound(set)
end

"""
    JuMP.upper_bound(pref::ParameterRef)::Number

Extend the `JuMP.upper_bound` function to accomodate infinite parameters.
Returns the upper bound associated with the infinite set. Errors if such a bound
is not well-defined. Extensions with user-defined set types should extend
`JuMP.has_upper_bound(set::NewType)` and `JuMP.upper_bound(set::NewType)` if
appropriate.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> upper_bound(t)
1.0
```
"""
function JuMP.upper_bound(pref::ParameterRef)::Number
    set = _parameter_set(pref)
    if !JuMP.has_upper_bound(pref)
        error("Parameter $(pref) does not have a upper bound.")
    end
    return JuMP.upper_bound(set)
end

"""
    JuMP.set_upper_bound(pref::ParameterRef, lower::Number)

Extend the `JuMP.set_upper_bound` function to accomodate infinite parameters.
Updates the infinite set upper bound if and only if it is an IntervalSet. Errors
otherwise. Extensions with user-defined infinite sets should extend
`JuMP.set_upper_bound(set::NewType, upper::Number)` if appropriate.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> set_upper_bound(t, 2)

julia> upper_bound(t)
2.0
```
"""
function JuMP.set_upper_bound(pref::ParameterRef, upper::Number)
    set = _parameter_set(pref)
    new_set = JuMP.set_upper_bound(set, upper)
    _update_parameter_set(pref, new_set)
    return
end

"""
    num_supports(pref::ParameterRef)::Int

Return the number of support points associated with `pref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> num_supports(t)
2
```
"""
function num_supports(pref::ParameterRef)::Int
    return length(_parameter_supports(pref))
end

"""
    has_supports(pref::ParameterRef)::Bool

Return true if `pref` has supports or false otherwise.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> has_supports(t)
true
```
"""
has_supports(pref::ParameterRef)::Bool = num_supports(pref) > 0

"""
    supports(pref::ParameterRef)::Vector

Return the support points associated with `pref`. Errors if there are no
supports.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> supports(t)
2-element Array{Int64,1}:
 0
 1
```
"""
function supports(pref::ParameterRef)::Vector
    !has_supports(pref) && error("Parameter $pref does not have supports.")
    return _parameter_supports(pref)
end

"""
    supports(prefs::AbstractArray{<:ParameterRef})::Vector

Return the support points associated with an array of `prefs` formatted as a
vector of SparseAxisArrays following the format of the input array. If the
parameters are not independent then the supports of each parameter are simply
spliced together. Alternatively can call `supports.` to more efficiently obtain
an array of the same input format whose parameter references have been replaced
with their supports. Errors if all the parameter references do not have the same
group ID number (were intialized together as an array) or if the nonindependent
parameters have support vectors of different lengths. If the parameters are
independent then all the unique combinations are identified and returned as
supports. Warning this operation is computationally expensive if there exist a
large number of combinations.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel())
julia> x = @infinite_parameter(model, [i = 1:2], set = IntervalSet(-1, 1),
                               base_name = "x", independent = true)
2-element Array{ParameterRef,1}:
 x[1]
 x[2]

julia> for i = 1:length(x)
           set_supports(x[i], [-1, 1])
       end

julia> supports(x)
4-element Array{JuMP.Containers.SparseAxisArray,1}:
   [2]  =  -1
  [1]  =  -1
   [2]  =  1
  [1]  =  -1
   [2]  =  -1
  [1]  =  1
   [2]  =  1
  [1]  =  1
```
"""
function supports(prefs::AbstractArray{<:ParameterRef})::Vector
    prefs = convert(JuMPC.SparseAxisArray, prefs)
    !_only_one_group(prefs) && error("Array contains parameters from multiple" *
                                     " groups.")
    for (k, pref) in prefs.data
        !has_supports(pref) && error("Parameter $pref does not have supports.")
    end
    lengths = [num_supports(pref) for (k, pref) in prefs.data]
    if !is_independent(collect(values(prefs.data))[1])
        length(unique(lengths)) != 1 && error("Each nonindependent parameter " *
                                              "must have the same number of " *
                                              "support points.")
        support_list = Vector{JuMPC.SparseAxisArray}(undef, lengths[1])
        for i in eachindex(support_list)
            support_list[i] = JuMPC.SparseAxisArray(Dict(k => supports(pref)[i] for (k, pref) in prefs.data))
        end
        # unique-nize support_list using the dictionary unique method
        # TODO: make it work with JuMPC.SparseAxisArray
        dict_list = unique([arr.data for arr in support_list])
        support_list = [JuMPC.SparseAxisArray(dict) for dict in dict_list]
    else
        all_keys = collect(keys(prefs))
        all_supports = [supports(pref) for (k, pref) in prefs.data]
        support_list = Vector{JuMPC.SparseAxisArray}(undef, prod(lengths))
        counter = 1
        for combo in Iterators.product(all_supports...)
            support_list[counter] = JuMPC.SparseAxisArray(Dict(all_keys[i] => combo[i] for i in eachindex(combo)))
            counter += 1
        end
    end
    return support_list
end

"""
    set_supports(pref::ParameterRef, supports::Vector{<:Number}; [force = false])

Specify the support points for `pref`. Errors if the supports violate the bounds
associated with the infinite set. Warns if the points are not unique. If `force`
this will overwrite exisiting supports otherwise it will error if there are
existing supports.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1]))
julia> set_supports(t, [0, 1])

julia> supports(t)
2-element Array{Int64,1}:
 0
 1
```
"""
function set_supports(pref::ParameterRef, supports::Vector{<:Number};
                      force = false)
    set = _parameter_set(pref)
    _check_supports_in_bounds(error, supports, set)
    if has_supports(pref) && !force
        error("Unable set supports for $pref since it already has supports." *
              " Consider using `add_supports` or use set `force = true` to " *
              "overwrite the existing supports.")
    end
    if !(is_independent(pref)) &&
       sum(values(pref.model.param_to_group_id) .== group_id(pref)) > 1
        _update_parameter_supports(pref, supports)
    else
        unique_supports = unique(supports)
        if length(unique_supports) != length(supports)
            @warn("Support points are not unique, eliminating redundant points.")
        end
        _update_parameter_supports(pref, unique_supports)
    end
    return
end

"""
    add_supports(pref::ParameterRef, supports::Union{Number, Vector{<:Number}})

Add additional support points for `pref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> add_supports(t, 0.5)

julia> supports(t)
3-element Array{Float64,1}:
 0.0
 1.0
 0.5

julia> add_supports(t, [0.25, 1])

julia> supports(t)
4-element Array{Float64,1}:
 0.0
 1.0
 0.5
 0.25
```
"""
function add_supports(pref::ParameterRef, supports::Union{Number,
                                                          Vector{<:Number}})
    set = _parameter_set(pref)
    _check_supports_in_bounds(error, supports, set)
    current_supports = _parameter_supports(pref)
    if !(is_independent(pref)) &&
       sum(values(pref.model.param_to_group_id) .== group_id(pref)) > 1
        new_supports = [current_supports; supports]
    else
        new_supports = unique([current_supports; supports])
    end
    _update_parameter_supports(pref, new_supports)
    return
end

"""
    delete_supports(pref::ParameterRef)

Delete the support points for `pref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> delete_supports(t)

julia> supports(t)
ERROR: Parameter t does not have supports.
```
"""
function delete_supports(pref::ParameterRef)
    _update_parameter_supports(pref, Int[])
    return
end

"""
    is_finite_parameter(pref::ParameterRef)::Bool

Return a `Bool` indicating if `pref` is a finite parameter.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @finite_parameter(model, cost, 42))
julia> is_finite_parameter(cost)
true
```
"""
function is_finite_parameter(pref::ParameterRef)::Bool
    set = infinite_set(pref)
    if isa(set, IntervalSet) && set.lower_bound == set.upper_bound
        return true
    end
    return false
end

"""
    JuMP.value(pref::ParameterRef)::Number

Return the value of `pref` so long as it is a finite parameter. Errors if it is
an infinite parameter.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @finite_parameter(model, cost, 42))
julia> value(cost)
42
```
"""
function JuMP.value(pref::ParameterRef)::Number
    !is_finite_parameter(pref) && error("$pref is an infinite parameter.")
    return supports(pref)[1]
end

"""
    JuMP.set_value(pref::ParameterRef, value::Number)

Set the value of `pref` so long as it is a finite parameter. Errors if it is
an infinite parameter.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @finite_parameter(model, cost, 42))
julia> set_value(cost, 27)

julia> value(cost)
27
```
"""
function JuMP.set_value(pref::ParameterRef, value::Number)
    !is_finite_parameter(pref) && error("$pref is an infinite parameter.")
    set_infinite_set(pref, IntervalSet(value, value))
    set_supports(pref, [value], force = true)
    return
end

"""
    fill_in_supports!(model::InfiniteModel; [num_supports::Int = 10,
                      sig_fig::Int = 5])

Automatically generate support points for all infinite parameters in model. User
can specify the number of significant figures kept after decimal point for the
auto-generated supports wtih `sig_fig`. This calls
[`fill_in_supports!`](@ref fill_in_supports!(::ParameterRef)) for each parameter
in the model. See [`fill_in_supports!`](@ref fill_in_supports!(::ParameterRef))
for more information. Errors if one of the infinite set types is unrecognized.

**Example**
```jldoctest; setup = :(using InfiniteOpt; model = InfiniteModel(); @infinite_parameter(model, 0 <= x <= 1);)
julia> fill_in_supports!(model, num_supports = 4, sig_fig = 3)

julia> supports(x)
4-element Array{Number,1}:
 0.0
 0.333
 0.667
 1.0
```
"""
function fill_in_supports!(model::InfiniteModel; num_supports::Int = 10,
                           sig_fig::Int = 5)
    for key in keys(model.params)
        pref = ParameterRef(model, key)
        fill_in_supports!(pref, num_supports = num_supports, sig_fig = sig_fig)
    end
    return
end

"""
    fill_in_supports!(pref::ParameterRef; [num_supports::Int = 10,
                                           sig_fig::Int = 5])

Automatically generate support points for a particular infinite parameter `pref`.
Generating `num_supports` for the parameter. The supports are generated uniformly
if the underlying infinite set is an `IntervalSet` or they are generating randomly
accordingly to the distribution if the set is a `DistributionSet`.
User can specify the number of digits kept after decimal point for the
auto-generated supports wtih `sig_fig`. Extensions that use user defined
set types should extend [`generate_and_add_supports!`](@ref) and/or
[`generate_support_values`](@ref) as needed. Errors if the infinite set type is
not recognized.

**Example**
```jldoctest; setup = :(using InfiniteOpt; model = InfiniteModel(); @infinite_parameter(model, 0 <= x <= 1);)
julia> fill_in_supports!(x, num_supports = 4, sig_fig = 3)

julia> supports(x)
4-element Array{Number,1}:
 0.0
 0.333
 0.667
 1.0

```
"""
function fill_in_supports!(pref::ParameterRef; num_supports::Int = 10,
                           sig_fig::Int = 5)
    p = JuMP.owner_model(pref).params[JuMP.index(pref)]
    if length(p.supports) == 0
        generate_and_add_supports!(pref, p.set, num_supports = num_supports,
                                   sig_fig = sig_fig)
    end
    return
end

"""
    generate_and_add_supports!(pref::ParameterRef, set::AbstractInfiniteSet;
                               [num_supports::Int = 10, sig_fig::Int = 5])

Generate supports for `pref` via [`generate_support_values`](@ref) and add them
to `pref`. This is intended as an extendable internal method for
[`fill_in_supports!`](@ref fill_in_supports!(::ParameterRef)). Note that if
`pref` is part of a `DistributionSet` that features a multivariate distribution,
all the associated parameters with `pref` will also have supports added to them.
Most extensions that empoy user-defined infinite sets can typically enable this
by extending [`generate_support_values`](@ref). However, in some cases it may be
necessary to extend this when more complex operations need to take place then just
adding supports to a single infinite parameter (e.g., how we enable multivariate
distribution sets). Errors if the infinite set type is not recognized.
"""
function generate_and_add_supports!(pref::ParameterRef, set::AbstractInfiniteSet;
                                    num_supports::Int = 10, sig_fig::Int = 5)
    add_supports(pref, generate_support_values(set, num_supports = num_supports,
                                               sig_fig = sig_fig))
    return
end

# Multivariate distribution sets
function generate_and_add_supports!(pref::ParameterRef,
                                    set::DistributionSet{<:Distributions.MultivariateDistribution};
                                    num_supports::Int = 10, sig_fig::Int = 5)
    pref_group_id = group_id(pref)
    model = JuMP.owner_model(pref)
    associated_p_index = sort([i for i in 1:length(model.params)
                               if model.param_to_group_id[i] == pref_group_id])
    new_supports = generate_support_values(set, num_supports = num_supports, sig_fig = sig_fig)

    for i in 1:length(associated_p_index)
        pref_i = ParameterRef(model, associated_p_index[i])
        add_supports(pref_i, new_supports[i, :])
    end
    return
end

"""
    group_id(pref::ParameterRef)::Int

Return the group ID number for `pref`.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> group_id(t)
1
```
"""
function group_id(pref::ParameterRef)::Int
    return JuMP.owner_model(pref).param_to_group_id[JuMP.index(pref)]
end

"""
    group_id(prefs::AbstractArray{<:ParameterRef})::Int

Return the group ID number for a group of `prefs`. Error if contains multiple
groups.

**Example**
```julia-repl
julia> group_id([x[1], x[2]])
2
```
"""
function group_id(prefs::AbstractArray{<:ParameterRef})::Int
    groups = group_id.(prefs)
    length(unique(groups)) != 1 && error("Array contains parameters from " *
                                         "multiple groups.")
    return first(groups)
end

"""
    is_independent(pref::ParameterRef)::Bool

Returns true for `pref` if it is independent or false otherwise.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> is_independent(t)
false
```
"""
function is_independent(pref::ParameterRef)::Bool
    return JuMP.owner_model(pref).params[JuMP.index(pref)].independent
end

"""
    set_independent(pref::ParameterRef)

Specify that `pref` be independent.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> set_independent(t)

julia> is_independent(t)
true
```
"""
function set_independent(pref::ParameterRef)
    old_param = JuMP.owner_model(pref).params[JuMP.index(pref)]
    new_param = InfOptParameter(old_param.set, old_param.supports, true)
    JuMP.owner_model(pref).params[JuMP.index(pref)] = new_param
    return
end

"""
    unset_independent(pref::ParameterRef)

Specify that `pref` be not independent.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> unset_independent(t)

julia> is_independent(t)
false
```
"""
function unset_independent(pref::ParameterRef)
    old_param = JuMP.owner_model(pref).params[JuMP.index(pref)]
    new_param = InfOptParameter(old_param.set, old_param.supports, false)
    JuMP.owner_model(pref).params[JuMP.index(pref)] = new_param
    return
end

"""
    parameter_by_name(model::InfiniteModel, name::String)::Union{ParameterRef,
                                                                 Nothing}

Return the parameter reference assoociated with a parameter name. Errors if
multiple parameters have the same name. Returns nothing if no such name exists.

**Example**
```jldoctest; setup = :(using InfiniteOpt, JuMP; model = InfiniteModel(); @infinite_parameter(model, t in [0, 1], supports = [0, 1]))
julia> parameter_by_name(model, "t")
t
```
"""
function parameter_by_name(model::InfiniteModel,
                           name::String)::Union{ParameterRef, Nothing}
    if model.name_to_param === nothing
        # Inspired from MOI/src/Utilities/model.jl
        model.name_to_param = Dict{String, Int}()
        for (param, param_name) in model.param_to_name
            if haskey(model.name_to_param, param_name)
                # -1 is a special value that means this string does not map to
                # a unique variable name.
                model.name_to_param[param_name] = -1
            else
                model.name_to_param[param_name] = param
            end
        end
    end
    index = get(model.name_to_param, name, nothing)
    if index isa Nothing
        return
    elseif index == -1
        error("Multiple parameters have the name $name.")
    else
        return ParameterRef(model, index)
    end
end

"""
    all_parameters(model::InfiniteModel)::Vector{ParameterRef}

Return all of the infinite parameter references currently in `model`.

**Example**
```julia-repl
julia> all_parameters(model)
3-element Array{ParameterRef,1}:
 t
 x[1]
 x[2]
```
"""
function all_parameters(model::InfiniteModel)::Vector{ParameterRef}
    pref_list = Vector{ParameterRef}(undef, num_parameters(model))
    indexes = sort([index for index in keys(model.params)])
    counter = 1
    for index in indexes
        pref_list[counter] = ParameterRef(model, index)
        counter += 1
    end
    return pref_list
end

## Define functions to extract the names of parameters
# Extract the root name of a parameter reference
function _root_name(pref::ParameterRef)::String
    name = JuMP.name(pref)
    first_bracket = findfirst(isequal('['), name)
    if first_bracket == nothing
        return name
    else
        # Hacky fix to handle invalid Unicode
        try
            return name[1:first_bracket-1]
        catch
            return name[1:first_bracket-2]
        end
    end
end

# Return the root names of a tuple parameter of references
function _root_names(prefs::Tuple)::Tuple
    return _root_name.(first.(prefs))
end

## Internal functions for group checking
# Return group id of ParameterRef
function _group(pref::ParameterRef)::Int
    return group_id(pref)
end

# Return the group if of the first element in an array (assuming all same)
function _group(arr::AbstractArray{<:ParameterRef})::Int
    return group_id(first(arr))
end

# Return true if SparseAxisArray only has one group
function _only_one_group(arr::JuMPC.SparseAxisArray{<:ParameterRef})::Bool
    return length(unique(group_id.(arr))) == 1
end

# Return true to have one group ID since is singular
_only_one_group(pref::ParameterRef)::Bool = true

## Internal function for extracting parameter references
# Return a vector of parameter references from a tuple of references
function _list_parameter_refs(prefs::Tuple)
    list = ParameterRef[]
    for pref in prefs
        if isa(pref, ParameterRef)
            push!(list, pref)
        else
            for k in keys(pref.data)
                push!(list, pref[k])
            end
        end
    end
    return list
end