text:
#=============================================

Definitions of various similarity functions

=============================================#

using QuadGK
using LinearAlgebra: dot, norm

#====================
Definitions of built-in similarity functions
====================#

#====================
Cosine similarity
====================#

@doc raw"""
    cossim(x,y)

Computes the cosine similarity between two inputs ``x`` and ``y``. Cosine similarity is defined as

``\text{cossim}(x,y) = \frac{\left\langle x,y\right\rangle}{\|x\|\cdot\|y\|}``

where ``\left\langle\cdot,\cdot\right\rangle`` is an inner product (e.g. dot product) and ``\|\cdot\|`` is its derived norm. This is roughly interpreted as being related to the angle between the inputs ``x`` and ``y``: when ``x`` and ``y`` have low angle between them, `cossim(x,y)` is high (close to ``1``). When ``x`` and ``y`` have large angle between them, `cossim(x,y)` is low (close to ``-1``).

# Arguments
- `x` and `y`: two inputs for which `dot(x,y)`, `norm(x)`, and `norm(y)` are defined.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> using LinearAlgebra: dot, norm;

julia> x, y = rand(4), rand(4);

julia> cossim(x,y) == dot(x,y) / (norm(x) * norm(y))
true

julia> z = rand(5);

julia> cossim(x,z)
ERROR: DimensionMismatch("dot product arguments have lengths 4 and 5")
```

See also: [`SimHash`](@ref)
"""
function cossim(x::AbstractVector, y::AbstractVector)
    norm_x = norm(x)
    norm_y = norm(y)

    if norm_x == 0 || norm_y == 0
        "x and y must be nonzero" |> ErrorException |> throw
    end

    dot(x,y) / (norm_x * norm_y)
end

function cossim(f, g, interval::LSHFunctions.RealInterval)
    norm_f = L2_norm(f, interval)
    norm_g = L2_norm(g, interval)

    if norm_f == 0 || norm_g == 0
        "f and g must be nonzero" |> ErrorException |> throw
    end

    inner_prod(f, g, interval) / (norm_f * norm_g)
end

#====================
L^p distance
====================#

@doc raw"""
    ℓp(x::AbstractVector, y::AbstractVector, p::Real=2)
    ℓ1(x::AbstractVector, y::AbstractVector)
    ℓ2(x::AbstractVector, y::AbstractVector)

Computes the ``\ell^p`` distance between a pair of vectors, given by

``\ell^p(x,y) \coloneqq \|x - y\|_p = \left(\sum_i \left|x_i - y_i\right|^p\right)^{1/p}``

`ℓ1(x,y)` is the same as `ℓp(x,y,1)`, and `ℓ2(x,y)` is the same as `ℓp(x,y,2)`.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = [1, 2, 3];

julia> y = [4, 5, 6];

julia> ℓp(x,y,2) == (abs(1-4)^2 + abs(2-5)^2 + abs(3-6)^2)^(1/2)
true

julia> ℓp(x,y,3) == (abs(1-4)^3 + abs(2-5)^3 + abs(3-6)^3)^(1/3)
true
```

See also: [`ℓp_norm`](@ref), [`L1Hash`](@ref), [`L2Hash`](@ref)
"""
ℓp(x::AbstractVector, y::AbstractVector, p::Real=2) = Lp(x, y, p)

@doc (@doc ℓp)
ℓ1(x::AbstractVector, y::AbstractVector) = L1(x, y)

@doc (@doc ℓp)
ℓ2(x::AbstractVector, y::AbstractVector) = L2(x, y)

@doc raw"""
    Lp(x::AbstractVector, y::AbstractVector, p::Real=2)
    L1(x::AbstractVector, y::AbstractVector)
    L2(x::AbstractVector, y::AbstractVector)

Computes the ``ℓ^p`` distance between a pair of vectors ``x`` and ``y``. Identical to `ℓp(x,y,p)`, `ℓ1(x,y)`, and `ℓ2(x,y)`, respectively.

See also: [`ℓp`](@ref)
"""
function Lp(x::AbstractVector{T}, y::AbstractVector, p::Real=2) where T
    if p ≤ 0
        "p must be positive" |> ErrorException |> throw
    elseif length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs(x[ii] - y[ii])^p
    end

    return result^(1/p)
end

@doc (@doc Lp)
function L1(x::AbstractVector{T}, y::AbstractVector) where T
    if length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs(x[ii] - y[ii])
    end

    return result
end

@doc (@doc Lp)
function L2(x::AbstractVector{T}, y::AbstractVector) where T
    if length(x) != length(y)
        "length(x) != length(y)" |> DimensionMismatch |> throw
    end

    result = T(0)
    @inbounds @simd for ii = 1:length(x)
        result += abs2(x[ii] - y[ii])
    end

    return √result
end

# Function space L^p distances

@doc raw"""
    Lp(f, g, interval::LSHFunctions.RealInterval, p)
    L1(f, g, interval::LSHFunctions.RealInterval)
    L2(f, g, interval::LSHFunctions.RealInterval)

Computes the ``L^p`` distance between two functions, given by

``L^p(f,g) \coloneqq \|f - g\|_p = \left(\int_a^b \left|f(x) - g(x)\right|^p \hspace{0.15cm} dx\right)^{1/p}``

# Examples
Below we compute the ``L^1``, ``L^2``, and ``L^3`` distances between ``f(x) = x^2 + 1`` and ``g(x) = 2x`` over the interval ``[0,1]``. The distances are computed by evaluating the integral

``\left(\int_0^1 \left|f(x) - g(x)\right|^p \hspace{0.15cm}dx\right)^{1/p} = \left(\int_0^1 \left|x^2 - 2x + 1\right|^p \hspace{0.15cm}dx\right)^{1/p} = \left(\int_0^1 (x - 1)^{2p} \hspace{0.15cm}dx\right)^{1/p}``

for ``p = 1``, ``p = 2``, and ``p = 3``.

```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = x^2 + 1; g(x) = 2x;

julia> interval = @interval(0 ≤ x ≤ 1);

julia> Lp(f, g, interval, 1) ≈ L1(f, g, interval) ≈ 3^(-1)
true

julia> Lp(f, g, interval, 2) ≈ L2(f, g, interval) ≈ 5^(-1/2)
true

julia> Lp(f, g, interval, 3) ≈ 7^(-1/3)
true
```

See also: [`Lp_norm`](@ref), [`ℓp`](@ref)
"""
Lp(f, g, interval::LSHFunctions.RealInterval, p::Real=2) =
    Lp_norm(x -> f(x) - g(x), interval, p)

@doc (@doc Lp)
L1(f, g, interval::LSHFunctions.RealInterval) = L1_norm(x -> f(x) - g(x), interval)

@doc (@doc Lp)
L2(f, g, interval::LSHFunctions.RealInterval) = L2_norm(x -> f(x) - g(x), interval)

#====================
Jaccard similarity
====================#

@doc raw"""
    jaccard(A::Set, B::Set) :: Float64

Computes the Jaccard similarity between sets ``A`` and ``B``, which is defined as

``\text{Jaccard}(A,B) = \frac{\left|A \cap B\right|}{\left|A \cup B\right|}``

# Arguments
- `A::Set`, `B::Set`: two sets whose Jaccard similarity we would like to compute.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> A, B = Set([1, 2, 3]), Set([2, 3, 4]);

julia> jaccard(A,B)
0.5

julia> jaccard(A,B) == length(A ∩ B) / length(A ∪ B)
true
```

See also: [`MinHash`](@ref)
"""
function jaccard(A::Set, B::Set) :: Float64
    if isempty(A)
        # Use the convention that if A = B = ∅, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        length(A ∩ B) / length(A ∪ B)
    end
end

@doc raw"""
    function jaccard(x::BitArray{1}, y::BitArray{1})

Computes the Jaccard similarity between a pair of binary vectors:

``J(x, y) = \frac{\sum_{i} \min{(x_i,y_i)}}{\sum_{i} \max{(x_i,y_i)}}``

# Arguments
- `x::BitArray{1}`, `y::BitArray{1}`: two binary vectors, in the form of `BitArray`s.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = BitArray([true, false, true, true, false]);

julia> y = BitArray([false, false, true, true, true]);

julia> jaccard(x,y)
0.5
```
"""
function jaccard(x::BitArray{1}, y::BitArray{1}) :: Float64
    union = sum(x .| y)
    if union == 0
        # Use the convention that if x and y are full of zeros, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        intersection = sum(x .& y)
        intersection / union
    end
end

@doc raw"""
    function jaccard(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})

Computes the Jaccard similarity between a pair of vectors of real numbers:

``J(x, y) = \frac{\sum_{i} \min{(x_i,y_i)}}{\sum_{i} \max{(x_i,y_i)}}``

# Arguments
- `x::AbstractVector{<:Real}`, `y::AbstractVector{<:Real}`: a pair of vectors containing real numbers (subtypes of `Real`).

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> x = [0.8, 0.1, 0.3, 0.4, 0.1];

julia> y = [1.0, 0.6, 0.0, 0.4, 0.5];

julia> jaccard(x,y)
0.5
```
"""
function jaccard(x::AbstractVector{T},
                 y::AbstractVector) :: Float64 where {T <: Real}
    if length(x) != length(y)
        DimensionMismatch("dimensions must match") |> throw
    end

    intersection = T(0)
    union = T(0)

    @inbounds @simd for ii = 1:length(x)
        if 0 ≤ x[ii] ≤ y[ii]
            intersection += x[ii]
            union += y[ii]
        elseif 0 ≤ y[ii] < x[ii]
            intersection += y[ii]
            union += x[ii]
        else
            ErrorException("vectors must have non-negative elements") |> throw
        end
    end

    if union == T(0)
        # Use the convention that if x and y are full of zeros, their Jaccard
        # similarity is zero.
        Float64(0)
    else
        Float64(intersection / union)
    end
end

jaccard(x::AbstractVector{<:Integer}, y::AbstractVector{<:AbstractFloat}) =
    jaccard(y, x)

@doc raw"""
    function jaccard(A::Set{<:K},
                     B::Set{<:K},
                     weights::Dict{K,V}) where {K,V<:Number}

Computes the weighted Jaccard similarity between two sets:

``J(x, y) = \frac{\sum_{x\in A\cap B} w_x}{\sum_{y\in A\cup B} w_y}``

# Arguments
- `A::Set`, `B::Set`: two sets whose Jaccard similarity we would like to compute.
- `weights::Dict`: a dictionary mapping symbols in the sets `A` and `B` to numerical weights. These weights must be positive.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> A = Set(["a", "b", "c"]);

julia> B = Set(["b", "c", "d"]);

julia> W = Dict("a" => 0.2, "b" => 2.4, "c" => 0.6, "d" => 1.8);

julia> jaccard(A,B,W)
0.6
```
"""
function jaccard(A::Set{<:K},
                 B::Set{<:K},
                 weights::Dict{K,V}) :: Float64 where {K,V<:Real}

    union_weight = V(0)

    for el in A ∪ B
        w = weights[el]
        if w < 0
            ErrorException("weights must be non-negative") |> throw
        end
        union_weight += w
    end

    intersection_weight = sum(weights[el] for el in A ∩ B)

    # By convention, if A = B = ∅, their Jaccard similarity is zero
    if union_weight == V(0)
        Float64(0)
    else
        Float64(intersection_weight / union_weight)
    end
end

#====================
Inner product and norms
====================#

### Inner products

@doc raw"""
    inner_prod(x::AbstractVector, y::AbstractVector)

Computes the ``\ell^2`` inner product (dot product)

``\left\langle x, y\right\rangle = \sum_i x_iy_i``

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> using LinearAlgebra: dot;

julia> x, y = randn(4), randn(4);

julia> inner_prod(x,y) ≈ dot(x,y)
true
```
"""
inner_prod(x::AbstractVector, y::AbstractVector) = dot(x,y)

# 1-dimensional inner product between L^2 functions
@doc raw"""
    inner_prod(f, g, interval::LSHFunctions.RealInterval)

Computes the ``L^2`` inner product

``\left\langle f, g\right\rangle = \int_a^b f(x)g(x) \hspace{0.15cm} dx``

where the interval we're integrating over is specified by the `interval` argument.

# Examples
```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = cos(x); g(x) = sin(x);

julia> inner_prod(f, g, @interval(0 ≤ x ≤ π/2)) ≈ 1/2
true
```
"""
inner_prod(f, g, interval::LSHFunctions.RealInterval) =
    quadgk(x -> f(x)g(x), interval.lower, interval.upper)[1]

### L^p norms
@doc raw"""
    Lp_norm(x::AbstractVector, p::Real = 2)
    L1_norm(x::AbstractVector)
    L2_norm(x::AbstractVector)

Compute the ``\ell^p`` norm of a vector ``x``. Identical to `ℓp_norm(x,p)`, `ℓ1_norm(x)`, and `ℓ2_norm(x)`, respectively.

See also: [`ℓp_norm`](@ref)
"""
Lp_norm(x::AbstractVector, p::Real = 2) = norm(x,p)

@doc (@doc Lp_norm)
L1_norm(x::AbstractVector) = norm(x,1)

@doc (@doc Lp_norm)
L2_norm(x::AbstractVector) = norm(x)

@doc raw"""
    ℓp_norm(x::AbstractVector, p::Real = 2)
    ℓ1_norm(x::AbstractVector)
    ℓ2_norm(x::AbstractVector)

Compute the ``\ell^p`` norm of a point ``x``, defined as

``\|x\|_p = \left(\sum_i \left|x_i\right|^p\right)^{1/p}``

# Examples

```jldoctest; setup = :(using LSHFunctions)
julia> x = randn(4);

julia> ℓp_norm(x, 1) ≈ ℓ1_norm(x) ≈ (map(u -> abs(u)^1, x) |> sum)^(1/1)
true

julia> ℓp_norm(x, 2) ≈ ℓ2_norm(x) ≈ (map(u -> abs(u)^2, x) |> sum)^(1/2)
true

julia> ℓp_norm(x, 3) ≈ (map(u -> abs(u)^3, x) |> sum)^(1/3)
true
```

See also: [`ℓp`](@ref), [`Lp_norm`](@ref)
"""
ℓp_norm(x::AbstractVector, p::Real = 2) = Lp_norm(x, p)

@doc (@doc ℓp_norm)
ℓ1_norm(x::AbstractVector) = L1_norm(x)

@doc (@doc ℓp_norm)
ℓ2_norm(x::AbstractVector) = L2_norm(x)

# 1-dimensional L^p norms

@doc raw"""
    Lp_norm(f, interval::LSHFunctions.RealInterval, p::Real=2)
    L1_norm(f, interval::LSHFunctions.RealInterval)
    L2_norm(f, interval::LSHFunctions.RealInterval)

Computes the ``L^p`` function-space norm of a function ``f``, which is given by the equation

``\|f\|_p = \left(\int_a^b \left|f(x)\right|^p \hspace{0.15cm} dx\right)^{1/p}``

`L1_norm(f, interval)` is the same as `Lp_norm(f, interval, 1)`, and `L2_norm(f, interval)` is the same as `Lp_norm(f, interval, 2)`.

# Examples

```jldoctest; setup = :(using LSHFunctions)
julia> f(x) = x;

julia> interval = @interval(0 ≤ x ≤ 1);

julia> Lp_norm(f, interval, 1) ≈ L1_norm(f, interval) ≈ 2^(-1/1)
true

julia> Lp_norm(f, interval, 2) ≈ L2_norm(f, interval) ≈ 3^(-1/2)
true

julia> Lp_norm(f, interval, 3) ≈ 4^(-1/3)
true
```
"""
Lp_norm(f, interval::LSHFunctions.RealInterval, p::Real=2) = (quadgk(x -> abs(f(x)).^p, interval.lower, interval.upper)[1])^(1/p)

@doc (@doc Lp_norm)
L1_norm(f, interval::LSHFunctions.RealInterval) = quadgk(x -> abs(f(x)), interval.lower, interval.upper)[1]

@doc (@doc Lp_norm)
L2_norm(f, interval::LSHFunctions.RealInterval) = √quadgk(x -> abs2(f(x)), interval.lower, interval.upper)[1]


text:
using ArrayLayouts, BlockBandedMatrices, BandedMatrices, BlockArrays, LinearAlgebra, Test
import BlockBandedMatrices: AbstractBandedBlockBandedMatrix, AbstractBlockBandedMatrix, block, blockindex, blockcolsupport, blockrowsupport
import BandedMatrices: bandwidths, AbstractBandedMatrix, BandedStyle, bandeddata, BandedColumns, _BandedMatrix

struct MyBandedBlockBandedMatrix <: AbstractBandedBlockBandedMatrix{Float64}
    A::BlockMatrix{Float64}
end

BlockBandedMatrices.blockbandwidths(::MyBandedBlockBandedMatrix) = (1,1)
BlockBandedMatrices.subblockbandwidths(::MyBandedBlockBandedMatrix) = (1,1)
Base.axes(A::MyBandedBlockBandedMatrix) = axes(A.A)
function Base.getindex(A::MyBandedBlockBandedMatrix, k::Int, j::Int)
    Kk, Jj = findblockindex(axes(A,1), k), findblockindex(axes(A,2), j)
    -1 ≤ Int(block(Kk)-block(Jj)) ≤ 1 || return 0.0
    -1 ≤ Int(blockindex(Kk)-blockindex(Jj)) ≤ 1 || return 0.0
    A.A[k,j]
end

struct MyBlockBandedMatrix <: AbstractBlockBandedMatrix{Float64}
    A::BlockMatrix{Float64}
end

BlockBandedMatrices.blockbandwidths(::MyBlockBandedMatrix) = (1,1)
Base.axes(A::MyBlockBandedMatrix) = axes(A.A)
function Base.getindex(A::MyBlockBandedMatrix, k::Int, j::Int)
    Kk, Jj = findblockindex(axes(A,1), k), findblockindex(axes(A,2), j)
    -1 ≤ Int(block(Kk)-block(Jj)) ≤ 1 || return 0.0
    A.A[k,j]
end

@testset "Interfaces" begin
    @testset "MyBandedBlockBandedMatrix" begin
        A = MyBandedBlockBandedMatrix(BlockMatrix(randn(6,6), 1:3, 1:3))
        @test MemoryLayout(A) isa BlockBandedMatrices.BandedBlockBandedLayout
        @test A[Block(3,3)] == BandedMatrix(A.A[Block(3,3)],(1,1))
        @test A[Block.(1:3),Block.(1:3)] == A

        @test A[Block(3,3)] isa BandedMatrix
        @test A[Block(3)[2:3],Block(3)[1:2]] isa BandedMatrix
        @test A[Block(3)[2:3],Block(3)] isa BandedMatrix
        @test A[Block(3),Block(3)[1:2]] isa BandedMatrix
        @test A[Block.(1:3),Block.(1:3)] isa BandedBlockBandedMatrix
        @test A[Block(1),Block.(2:3)] isa PseudoBlockArray
        @test A[Block.(2:3),Block(1)] isa PseudoBlockArray
        @test A[Block.(2:3),Block(2)[1:2]] isa PseudoBlockArray
        @test A[Block(2)[1:2],Block.(2:3)] isa PseudoBlockArray
    end

    @testset "MyBlockBandedMatrix" begin
        A = MyBlockBandedMatrix(BlockMatrix(randn(6,6), 1:3, 1:3))
        @test MemoryLayout(A) isa BlockBandedMatrices.BlockBandedLayout
        @test MemoryLayout(A') isa BlockBandedMatrices.BlockBandedLayout
        @test MemoryLayout(transpose(A)) isa BlockBandedMatrices.BlockBandedLayout
        @test A[Block(3,3)] == A.A[Block(3,3)]
        @test A[Block.(1:3),Block.(1:3)] == A

        @test A[Block(3,3)] isa Matrix
        @test A[Block(3)[2:3],Block(3)[1:2]] isa Matrix
        @test A[Block(3)[2:3],Block(3)] isa Matrix
        @test A[Block(3),Block(3)[1:2]] isa Matrix
        @test A[Block.(1:3),Block.(1:3)] isa BlockBandedMatrix
        @test A[Block(1),Block.(2:3)] isa PseudoBlockArray
        @test A[Block.(2:3),Block(1)] isa PseudoBlockArray
        @test A[Block.(2:3),Block(2)[1:2]] isa PseudoBlockArray
        @test A[Block(2)[1:2],Block.(2:3)] isa PseudoBlockArray
    end

    @testset "Zeros" begin
        Z = Zeros(blockedrange(1:3), blockedrange(1:3))
        B = BandedBlockBandedMatrix(Z)
        @test B == Z
        @test blockisequal(axes(Z), axes(B))
        @test blockbandwidths(B) == blockbandwidths(Z) == (-1,-1)
        @test subblockbandwidths(B) == subblockbandwidths(Z) == (-1,-1)
    end

    @testset "Diagonal interface" begin
        n = 10
        D = Diagonal(randn(n^2))
        @test blockbandwidths(D) == subblockbandwidths(D) == (0,0)

        PD = PseudoBlockArray(D, Fill(n,n), Fill(n,n))
        @test blockbandwidths(PD) == bandwidths(PD) == (0,0)
        @test MemoryLayout(typeof(PD)) isa DiagonalLayout{DenseColumnMajor}
        @test bandeddata(PD) == bandeddata(D)
        V = view(PD, Block(1,1))

        @test bandwidths(V) == (0,0)
        @test_broken Broadcast.BroadcastStyle(typeof(V)) == BandedStyle()
        @test MemoryLayout(typeof(V)) isa BandedColumns{DenseColumnMajor}
        @test bandeddata(V) == bandeddata(PD)[:,1:n]
        @test BandedMatrix(V) == V
    end

    @testset "Block Diagonal" begin
        A = BlockBandedMatrices.BlockDiagonal(fill([1 2],3))
        @test blockbandwidths(A) == (0,0)
        @test isblockbanded(A)
        @test A[Block(1,1)] == [1 2]
        @test @inferred(A[Block(1,2)]) == [0 0]
        @test_throws DimensionMismatch A+I
        A = BlockBandedMatrices.BlockDiagonal(fill([1 2; 1 2],3))
        @test A+I == I+A == mortar(Diagonal(fill([2 2; 1 3],3))) == Matrix(A) + I
    end

    @testset "Block Bidiagonal" begin
        Bu = BlockBidiagonal(fill([1 2],4), fill([3 4],3), :U)
        Bl = BlockBidiagonal(fill([1 2],4), fill([3 4],3), :L)
        @test blockbandwidths(Bu) == (0,1)
        @test blockbandwidths(Bl) == (1,0)
        @test isblockbanded(Bu)
        @test isblockbanded(Bl)
        @test Bu[Block(1,1)] == Bl[Block(1,1)] == [1 2]
        @test @inferred(Bu[Block(1,2)]) == @inferred(Bl[Block(2,1)]) == [3 4]
        @test @inferred(view(Bu,Block(1,3))) == @inferred(Bu[Block(1,3)]) == [0 0]
        @test_throws DimensionMismatch Bu+I
        Bu = BlockBidiagonal(fill([1 2; 1 2],4), fill([3 4; 3 4],3), :U)
        Bl = BlockBidiagonal(fill([1 2; 1 2],4), fill([3 4; 3 4],3), :L)
        @test Bu+I == I+Bu == mortar(Bidiagonal(fill([2 2; 1 3],4), fill([3 4; 3 4],3), :U)) == Matrix(Bu) + I
        @test Bl+I == I+Bl == mortar(Bidiagonal(fill([2 2; 1 3],4), fill([3 4; 3 4],3), :L)) == Matrix(Bl) + I
        @test Bu-I == mortar(Bidiagonal(fill([0 2; 1 1],4), fill([3 4; 3 4],3), :U)) == Matrix(Bu) - I
        @test I-Bu == mortar(Bidiagonal(fill([0 -2; -1 -1],4), fill(-[3 4; 3 4],3), :U)) == I - Matrix(Bu)
    end

    @testset "Block Tridiagonal" begin
        A = BlockTridiagonal(fill([1 2],3), fill([3 4],4), fill([4 5],3))
        @test blockbandwidths(A) == (1,1)
        @test isblockbanded(A)
        @test A[Block(1,1)] == [3 4]
        @test @inferred(A[Block(1,2)]) == [4 5]
        @test @inferred(view(A,Block(1,3))) == @inferred(A[Block(1,3)]) == [0 0]
        @test_throws DimensionMismatch A+I
        A = BlockTridiagonal(fill([1 2; 1 2],3), fill([3 4; 3 4],4), fill([4 5; 4 5],3))
        @test A+I == I+A == mortar(Tridiagonal(fill([1 2; 1 2],3), fill([4 4; 3 5],4), fill([4 5; 4 5],3))) == Matrix(A) + I
        @test A-I == mortar(Tridiagonal(fill([1 2; 1 2],3), fill([2 4; 3 3],4), fill([4 5; 4 5],3))) == Matrix(A) - I
        @test I-A == mortar(Tridiagonal(fill(-[1 2; 1 2],3), fill([-2 -4; -3 -3],4), fill(-[4 5; 4 5],3))) == I - Matrix(A)
    end

    @testset "DiagonalBlock" begin
        D = Diagonal(PseudoBlockVector(randn(6), 1:3))
        D̃ = BandedBlockBandedMatrix(D)
        @test D̃ == D
        @test blockbandwidths(D) == blockbandwidths(D̃) == subblockbandwidths(D) == subblockbandwidths(D̃) == (0,0)

        @test @inferred(D[Block(2,2)]) isa BandedMatrix
        @test D[Block(2,3)] isa BandedMatrix
        @test bandwidths(D[Block(2,2)]) == (0,0)
        @test bandwidths(D[Block(2,3)]) == (-720,-720)
        @test D[Block(2,2)] == D[2:3,2:3]
        @test D[Block(2,3)] == zeros(2,3)

        @test D[Block.(1:2),Block.(2:3)] isa BandedBlockBandedMatrix
        @test D[Block.(1:2),Block.(2:3)] == D[1:3,2:6]
        @test blockbandwidths(D[Block.(1:2),Block.(2:3)]) == (1,-1)
        @test subblockbandwidths(D[Block.(1:2),Block.(2:3)]) == (0,0)
    end

    @testset "Block-BandedMatrix" begin
        a = blockedrange(1:5)
        B = _BandedMatrix(PseudoBlockArray(randn(5,length(a)),(Base.OneTo(5),a)), a, 3, 1)
        @test blockcolsupport(B,1) == Block.(1:3)
        @test blockcolsupport(B,3) == Block.(2:4)
        @test blockrowsupport(B,1) == Block.(1:2)
        @test blockrowsupport(B,4) == Block.(3:5)

        Q = Eye((a,))[:,Block(2)]
        @test Q isa BandedMatrix
        @test blockcolsupport(Q,1) == Block.(2:2)

        Q = Eye((a,))[Block(2),:]
        @test Q isa BandedMatrix
        @test blockrowsupport(Q,1) == Block.(2:2)

        @testset "constant blocks" begin
            a = blockedrange(Fill(2,5))
            Q = Eye((a,))[:,Block(2)]
            @test blockbandwidths(Q) == (1,-1)

            B = _BandedMatrix(randn(5,length(a)), a, 3, 1)
            @test blockbandwidths(B) == (4,0)
        end
    end
end

ext: jl


text:
using ArgParse
using PyTest

argparse_settings = ArgParseSettings()
@add_arg_table argparse_settings begin
    "testpaths"
        help = "Paths to tests which should be run"
        nargs = '*'
        arg_type = AbstractString
end

# this line will make all @pytest sections obey the invocation's args
PyTest.set_parsed_args!(parse_args(argparse_settings))

println(joinpath(pwd(), "test/runtests.jl"))

include(joinpath(pwd(), "test/runtests.jl"))


text:
################################################################################
#                   CORE DISPATCHVARIABLEREF METHOD EXTENSIONS
################################################################################
# Extend dispatch_variable_ref
function dispatch_variable_ref(
    model::InfiniteModel,
    index::MeasureIndex
    )::MeasureRef
    return MeasureRef(model, index)
end

# Extend _add_data_object
function _add_data_object(
    model::InfiniteModel,
    object::MeasureData
    )::MeasureIndex
    return MOIUC.add_item(model.measures, object)
end

# Extend _data_dictionary (type based)
function _data_dictionary(
    model::InfiniteModel,
    ::Type{Measure}
    )::MOIUC.CleverDict{MeasureIndex, MeasureData{<:Measure}}
    return model.measures
end

# Extend _data_dictionary (reference based)
function _data_dictionary(
    mref::MeasureRef
    )::MOIUC.CleverDict{MeasureIndex, MeasureData{<:Measure}}
    return JuMP.owner_model(mref).measures
end

# Extend _data_object
function _data_object(mref::MeasureRef)
    object = get(_data_dictionary(mref), JuMP.index(mref), nothing)
    if isnothing(object)
        error("Invalid measure reference, cannot find ",
              "corresponding measure in the model. This is likely ",
              "caused by using the reference of a deleted measure.")
    end
    return object
end

# Extend _core_variable_object
function _core_variable_object(mref::MeasureRef)
    return _data_object(mref).measure
end

# Extend _object_numbers
function _object_numbers(mref::MeasureRef)::Vector{Int}
    return _core_variable_object(mref).object_nums
end

# Extend _parameter_numbers
function _parameter_numbers(mref::MeasureRef)::Vector{Int}
    return _core_variable_object(mref).parameter_nums
end

## Set helper methods for adapting data_objects with parametric changes
# No change needed
function _adaptive_data_update(
    mref::MeasureRef,
    m::M,
    data::MeasureData{M}
    )::Nothing where {M <: Measure}
    data.measure = m
    return
end

# Reconstruction is necessary
function _adaptive_data_update(
    mref::MeasureRef,
    m::M1,
    data::MeasureData{M2}
    )::Nothing  where {M1, M2}
    new_data = MeasureData(m, data.name, data.measure_indices,
                           data.constraint_indices, data.derivative_indices,
                           data.in_objective)
    _data_dictionary(mref)[JuMP.index(mref)] = new_data
    return
end

# Update the core measure object
function _set_core_variable_object(
    mref::MeasureRef,
    meas::Measure
    )::Nothing
    _adaptive_data_update(mref, meas, _data_object(mref))
    return
end

# Extend _measure_dependencies
function _measure_dependencies(mref::MeasureRef)::Vector{MeasureIndex}
    return _data_object(mref).measure_indices
end

# Extend _constraint_dependencies
function _constraint_dependencies(mref::MeasureRef)::Vector{InfOptConstraintIndex}
    return _data_object(mref).constraint_indices
end

# Extend _derivative_dependencies
function _derivative_dependencies(mref::MeasureRef)::Vector{DerivativeIndex}
    return _data_object(mref).derivative_indices
end

################################################################################
#                              MEASURE DATA METHODS
################################################################################
# Extend Base.copy for DiscreteMeasureData
function Base.copy(data::D)::D where {D <: DiscreteMeasureData}
    return DiscreteMeasureData(copy(data.parameter_refs), copy(data.coefficients),
                               copy(data.supports), data.label,
                               data.weight_function, copy(data.lower_bounds),
                               copy(data.upper_bounds), data.is_expect)
end

# Extend Base.copy for FunctionalDiscreteMeasureData
function Base.copy(data::D)::D where {D <: FunctionalDiscreteMeasureData}
    return FunctionalDiscreteMeasureData(copy(data.parameter_refs),
                                         data.coeff_function, data.min_num_supports,
                                         data.label, copy(data.generative_supp_info),
                                         data.weight_function,
                                         copy(data.lower_bounds),
                                         copy(data.upper_bounds), data.is_expect)
end

"""
    default_weight(t) = 1

Default weight function for [`DiscreteMeasureData`](@ref) and
[`FunctionalDiscreteMeasureData`](@ref). Returns 1 regardless of the input value.
"""
default_weight(t) = 1

# Ensure that a pref is valid for use in measure data
function _check_params(pref::GeneralVariableRef)::Nothing
    if !(_index_type(pref) <: InfiniteParameterIndex)
        error("$pref is not an infinite parameter.")
    end
    return
end

"""
    DiscreteMeasureData(pref::GeneralVariableRef,
        coefficients::Vector{<:Real},
        supports::Vector{<:Real};
        [label::Type{<:AbstractSupportLabel} = generate_unique_label(),
        weight_function::Function = [`default_weight`](@ref),
        lower_bound::Real = NaN,
        upper_bound::Real = NaN,
        is_expect::Bool = false]
        )::DiscreteMeasureData

Returns a 1-dimensional `DiscreteMeasureData` object that can be utilized
to define measures using [`measure`](@ref). This accepts input for a scalar (single)
infinite parameter. A description of the other arguments is provided in the
documentation for [`DiscreteMeasureData`](@ref). Errors if supports are out
bounds or an unequal number of supports and coefficients are given. Note that by
default a unique `label` is generated via `generate_unique_label` to ensure the supports can
be located in the infinite parameter support storage. Advanced implementations,
may choose a different behavior but should do so with caution.

**Example**
```julia-repl
julia> data = DiscreteMeasureData(pref, [0.5, 0.5], [1, 2]);
```
"""
function DiscreteMeasureData(pref::GeneralVariableRef,
    coefficients::Vector{<:Real},
    supports::Vector{<:Real};
    label::Type{<:AbstractSupportLabel} = generate_unique_label(),
    weight_function::F = default_weight,
    lower_bound::Real = NaN,
    upper_bound::Real = NaN,
    is_expect::Bool = false
    )::DiscreteMeasureData{GeneralVariableRef, 1, Float64, F} where {F <: Function}
    _check_params(pref)
    if length(coefficients) != length(supports)
        error("The amount of coefficients must match the amount of " *
              "support points.")
    end
    domain = infinite_domain(pref)
    if !supports_in_domain(supports, domain)
        error("Support points violate parameter domain.")
    end
    return DiscreteMeasureData(pref, coefficients, supports, label,
                               weight_function, lower_bound, upper_bound,
                               is_expect)
end

# Ensure that an array of prefs is valid for use in multi-dimensional data
function _check_params(prefs::AbstractArray{GeneralVariableRef})::Nothing
    only_dep = all(_index_type(pref) == DependentParameterIndex for pref in prefs)
    only_indep = !only_dep && all(_index_type(pref) == IndependentParameterIndex
                                  for pref in prefs)
    if !only_dep && !only_indep
        error("Cannot specify a mixture of infinite parameter types.")
    elseif only_dep && any(_raw_index(pref) != _raw_index(first(prefs)) for pref in prefs)
        error("Cannot specify multiple dependent parameter groups for one measure.")
    elseif only_dep && length(prefs) != _num_parameters(first(dispatch_variable_ref.(prefs)))
        error("Cannot specify a subset of dependent parameters, consider using " *
              "nested one-dimensional measures instead.")
    end
    return
end

## Define function to intelligently format supports properly
# Vectors
function _prepare_supports(
    supps::Vector{Vector{T}},
    inds::Collections.ContainerIndices{1, Nothing}
    )::Matrix{T} where {T <: Real}
    return reduce(hcat, supps)
end

# DenseAxisArrays and Arrays
function _prepare_supports(
    supps::Union{Array{<:Array{T}},Vector{<:JuMPC.DenseAxisArray{T}}},
    inds::Collections.ContainerIndices
    )::Matrix{T} where {T <: Real}
    return [supp[i] for i in eachindex(first(supps)), supp in supps]
end

# SparseAxisArray
function _prepare_supports(
    supps::Vector{<:JuMPC.SparseAxisArray},
    inds::Collections.ContainerIndices{1, <:Vector}
    )::Matrix{Float64}
    new_supps = Matrix{Float64}(undef, length(inds), length(supps))
    for i in eachindex(supps)
        new_supps[:, i] = Collections.vectorize(supps[i], inds)
    end
    return new_supps
end

## Check if a matrix of supports repsect infinite domain(s)
# IndependentParameterRefs
function _check_supports_in_bounds(
    prefs::Vector{IndependentParameterRef},
    supports::Matrix{<:Real}
    )::Nothing
    for i in eachindex(prefs)
        domain = infinite_domain(prefs[i])
        if !supports_in_domain(supports[i, :], domain)
            error("Support points violate parameter domain.")
        end
    end
    return
end

# DependentParameterRefs
function _check_supports_in_bounds(prefs::Vector{DependentParameterRef},
                                   supports::Matrix{<:Real})::Nothing
    domain = _parameter_domain(first(prefs))
    if !supports_in_domain(supports, domain)
        error("Support points violate parameter domain.")
    end
    return
end

"""
    DiscreteMeasureData(prefs::AbstractArray{GeneralVariableRef},
        coefficients::Vector{<:Real},
        supports::Vector{<:AbstractArray{<:Real}};
        label::Type{<:AbstractSupportLabel} = generate_unique_label(),
        weight_function::Function = [`default_weight`](@ref),
        lower_bounds::AbstractArray{<:Real} = [NaN...],
        upper_bounds::AbstractArray{<:Real} = [NaN...],
        is_expect::Bool = false
        )::DiscreteMeasureData

Returns a `DiscreteMeasureData` object that can be utilized to
define measures using [`measure`](@ref). This accepts input for an array (multi)
parameter. The inner arrays in the supports vector need to match the formatting
of the array used for `parameter_refs`. A description of the other arguments is
provided in the
documentation for [`DiscreteMeasureData`](@ref). Errors if supports are out
bounds, an unequal number of supports and coefficients are given, the array
formats do not match, or if mixed infinite parameter types are given. Note that by
default a unique `label` is generated via `generate_unique_label` to ensure the supports can
be located in the infinite parameter support storage. Advanced implementations,
may choose a different behavior but should do so with caution.

**Example**
```julia-repl
julia> data = DiscreteMeasureData(prefs, [0.5, 0.5], [[1, 1], [2, 2]]);
```
"""
function DiscreteMeasureData(
    prefs::AbstractArray{GeneralVariableRef},
    coefficients::Vector{<:Real},
    supports::Vector{<:AbstractArray{<:Real}};
    label::Type{<:AbstractSupportLabel} = generate_unique_label(),
    weight_function::F = default_weight,
    lower_bounds::AbstractArray{<:Real} = map(e -> NaN, prefs),
    upper_bounds::AbstractArray{<:Real} = map(e -> NaN, prefs),
    is_expect::Bool = false
    )::DiscreteMeasureData{Vector{GeneralVariableRef}, 2, Vector{Float64}, F} where {F <: Function}
    _check_params(prefs) # ensures that prefs are valid
    if _keys(prefs) != _keys(first(supports))
        error("Parameter references and supports must use same container type.")
    elseif _keys(prefs) != _keys(lower_bounds)
        error("Parameter references and bounds must use same container type.")
    end
    if length(coefficients) != length(supports)
        error("The amount of coefficients must match the amount of " *
              "support points.")
    end
    inds = Collections.indices(prefs)
    vect_prefs = Collections.vectorize(prefs, inds)
    supps = _prepare_supports(supports, inds)
    _check_supports_in_bounds(dispatch_variable_ref.(vect_prefs), supps)
    lbs = Collections.vectorize(lower_bounds, inds)
    ubs = Collections.vectorize(upper_bounds, inds)
    return DiscreteMeasureData(vect_prefs, coefficients, supps, label,
                               weight_function, lbs, ubs, is_expect)
end

"""
    FunctionalDiscreteMeasureData(pref::GeneralVariableRef,
        coeff_func::Function,
        min_num_supports::Int,
        label::Type{<:AbstractSupportLabel};
        [weight_function::Function = [`default_weight`](@ref),
        lower_bound::Real = NaN,
        upper_bound::Real = NaN,
        is_expect::Bool = false,
        generative_support_info::AbstractGenerativeInfo = NoGenerativeSupports()]
        )::FunctionalDiscreteMeasureData

Returns a 1-dimensional `FunctionalDiscreteMeasureData` object that can be utilized
to define measures using [`measure`](@ref). This accepts input for a scalar (single)
infinite parameter. A description of the other arguments is provided in the
documentation for [`FunctionalDiscreteMeasureData`](@ref). Errors if `pref` is
not an infinite parameter. Built-in choices for `label` include:
- `All`: Use all of the supports stored in `pref`
- `MCSample`: Use Monte Carlo samples associated with `pref`
- `WeightedSample`: Use weighted Monte Carlo samples associated with `pref`
- `UniformGrid`: Use uniform grid points associated with `pref`.

**Example**
```julia-repl
julia> data = FunctionalDiscreteMeasureData(pref, my_func, 20, UniformGrid);
```
"""
function FunctionalDiscreteMeasureData(
    pref::GeneralVariableRef,
    coeff_func::F1,
    min_num_supports::Int,
    label::Type{<:AbstractSupportLabel};
    weight_function::F2 = default_weight,
    lower_bound::Real = NaN,
    upper_bound::Real = NaN,
    is_expect::Bool = false,
    generative_support_info::I = NoGenerativeSupports()
    )::FunctionalDiscreteMeasureData{GeneralVariableRef, Float64, I, F1, F2} where {F1 <: Function, F2 <: Function, I <: AbstractGenerativeInfo}
    _check_params(pref)
    if _index_type(pref) == DependentParameterIndex && min_num_supports != 0
        error("`min_num_supports` must be 0 for individual dependent parameters.")
    end
    min_num_supports >= 0 || error("Number of supports must be nonnegative.")
    if !isnan(lower_bound) && !isnan(upper_bound)
        if !supports_in_domain([lower_bound, upper_bound], infinite_domain(pref))
            error("Bounds violate infinite domain bounds.")
        end
    end
    return FunctionalDiscreteMeasureData(pref, coeff_func, min_num_supports,
                                         label, generative_support_info,
                                         weight_function, lower_bound,
                                         upper_bound, is_expect)
end

## Check if the integral bounds satisfy the parameter domain(s)
# IndependentParameterRefs
function _check_bounds_in_domain(prefs::Vector{IndependentParameterRef}, lbs,
                              ubs)::Nothing
    for i in eachindex(prefs)
        if !supports_in_domain([lbs[i], ubs[i]], infinite_domain(prefs[i]))
            error("Bounds violate the infinite domain.")
        end
    end
    return
end

# DependentParameterRefs
function _check_bounds_in_domain(prefs::Vector{DependentParameterRef}, lbs,
                              ubs)::Nothing
    if !supports_in_domain(hcat(lbs, ubs), infinite_domain(prefs))
        error("Bounds violate the infinite domain.")
    end
    return
end

"""
    FunctionalDiscreteMeasureData(prefs::AbstractArray{GeneralVariableRef},
        coeff_func::Function,
        min_num_supports::Int,
        label::Type{<:AbstractSupportLabel};
        [weight_function::Function = [`default_weight`](@ref),
        lower_bounds::AbstractArray{<:Real} = [NaN...],
        upper_bounds::AbstractArray{<:Real} = [NaN...],
        is_expect::Bool = false]
        )::FunctionalDiscreteMeasureData

Returns a multi-dimensional `FunctionalDiscreteMeasureData` object that can be utilized
to define measures using [`measure`](@ref). This accepts input for an array of
infinite parameters. A description of the other arguments is provided in the
documentation for [`FunctionalDiscreteMeasureData`](@ref). Errors if `prefs` are
not infinite parameters or if the mixed parameter types are provided.
Built-in choices for `label` include:
- `All`: Use all of the supports stored in `prefs`
- `MCSample`: Use Monte Carlo samples associated with `prefs`
- `WeightedSample`: Use weighted Monte Carlo samples associated with `prefs`
- `UniformGrid`: Use uniform grid points associated with `prefs`.

**Example**
```julia-repl
julia> data = FunctionalDiscreteMeasureData(prefs, my_func, 20, MCSample);
```
"""
function FunctionalDiscreteMeasureData(
    prefs::AbstractArray{GeneralVariableRef},
    coeff_func::F1,
    min_num_supports::Int,
    label::Type{<:AbstractSupportLabel};
    weight_function::F2 = default_weight,
    lower_bounds::AbstractArray{<:Real} = map(e -> NaN, prefs),
    upper_bounds::AbstractArray{<:Real} = map(e -> NaN, prefs),
    is_expect::Bool = false
    )::FunctionalDiscreteMeasureData{Vector{GeneralVariableRef}, Vector{Float64}, NoGenerativeSupports, F1, F2} where {F1 <: Function, F2 <: Function}
    _check_params(prefs)
    if _keys(prefs) != _keys(lower_bounds)
        error("Parameter references and bounds must use same container type.")
    end
    min_num_supports >= 0 || error("Number of supports must be nonnegative.")
    inds = Collections.indices(prefs)
    vect_prefs = Collections.vectorize(prefs, inds)
    vect_lbs = Collections.vectorize(lower_bounds, inds)
    vect_ubs = Collections.vectorize(upper_bounds, inds)
    if !isnan(first(vect_lbs)) && !isnan(first(vect_ubs))
        dprefs = map(p -> dispatch_variable_ref(p), vect_prefs)
        _check_bounds_in_domain(dprefs, vect_lbs, vect_ubs)
    end
    return FunctionalDiscreteMeasureData(vect_prefs, coeff_func, min_num_supports,
                                         label, weight_function, vect_lbs,
                                         vect_ubs, is_expect)
end

"""
    parameter_refs(data::AbstractMeasureData)::Union{GeneralVariableRef,
                                                     AbstractArray{GeneralVariableRef}}

Return the infinite parameter reference(s) in `data`. This is intended as an
internal function to be used with measure addition. User-defined measure data types
will need to extend this function otherwise an error is thrown.
"""
function parameter_refs(data::AbstractMeasureData)
    error("Function `parameter_refs` not extended for measure data of type $(typeof(data)).")
end

# DiscreteMeasureData
function parameter_refs(data::DiscreteMeasureData{T})::T where {T}
    return data.parameter_refs
end

# FunctionalDiscreteMeasureData
function parameter_refs(data::FunctionalDiscreteMeasureData{T})::T where {T}
    return data.parameter_refs
end

"""
    support_label(data::AbstractMeasureData)::Type{<:AbstractSupportLabel}

Return the label stored in `data` associated with its supports.
This is intended as en internal method for measure creation and ensures any
new supports are added to parameters with such a label.
User-defined measure data types should extend this functionif supports are used,
otherwise an error is thrown.
"""
function support_label(data::AbstractMeasureData)
    error("Function `support_label` not defined for measure data of type " *
          "$(typeof(data)).")
end

# DiscreteMeasureData and FunctionalDiscreteMeasureData
function support_label(data::Union{FunctionalDiscreteMeasureData,
                       DiscreteMeasureData})::DataType
    return data.label
end

"""
    JuMP.lower_bound(data::AbstractMeasureData)::Union{Float64, Vector{Float64}}

Return the lower bound associated with `data` that defines its domain. This is
intended as an internal method, but may be useful for extensions. User-defined
measure data types should extend this function if desired, otherwise `NaN` is
returned
"""
function JuMP.lower_bound(data::AbstractMeasureData)::Float64
    return NaN
end

# DiscreteMeasureData and FunctionalDiscreteMeasureData
function JuMP.lower_bound(data::Union{FunctionalDiscreteMeasureData,
                                      DiscreteMeasureData})
    return data.lower_bounds
end

"""
    JuMP.upper_bound(data::AbstractMeasureData)::Union{Float64, Vector{Float64}}

Return the lower bound associated with `data` that defines its domain. This is
intended as an internal method, but may be useful for extensions. User-defined
measure data types should extend this function if desired, otherwise `NaN` is
returned.
"""
function JuMP.upper_bound(data::AbstractMeasureData)::Float64
    return NaN
end

# DiscreteMeasureData and FunctionalDiscreteMeasureData
function JuMP.upper_bound(data::Union{FunctionalDiscreteMeasureData,
                                      DiscreteMeasureData})
    return data.upper_bounds
end

## Indicate is the data is from an expectation call
# Fallback
function _is_expect(data::AbstractMeasureData)::Bool
    return false
end

# DiscreteMeasureData and FunctionalDiscreteMeasureData
function _is_expect(data::Union{FunctionalDiscreteMeasureData,
                                DiscreteMeasureData})::Bool
    return data.is_expect
end

"""
    generative_support_info(data::AbstractMeasureData)::AbstractGenerativeInfo

Return the generative support creation info that corresponds to `data`. This is
intended as an internal method and only needs to be extended for user-defined
measure data types that use generative supports.
"""
function generative_support_info(data::AbstractMeasureData)::NoGenerativeSupports
    return NoGenerativeSupports()
end

# FunctionalDiscreteMeasureData
function generative_support_info(data::FunctionalDiscreteMeasureData)
    return data.generative_supp_info
end

"""
    supports(data::AbstractMeasureData)::Array{Float64}

Return the supports associated with `data` and its infinite parameters.
This is intended as en internal method for measure creation and ensures any
new supports are added to parameters. User-defined measure data types should
extend this function if appropriate, otherwise an empty vector is returned.
"""
function supports(data::AbstractMeasureData)::Vector{Float64}
    return Float64[]
end

# DiscreteMeasureData
function supports(data::DiscreteMeasureData{T, N})::Array{Float64, N} where {T, N}
    return data.supports
end

## Define helper functions for generative functional data
# Include generative and uses All
function _get_supports(pref, gen::Val{true}, info, label::Type{All})::Vector{Float64}
    return supports(pref, label = label)
end

# Include generative and does not use All
function _get_supports(pref, gen::Val{true}, info, label)::Vector{Float64}
    return supports(pref, label = Union{support_label(info), label})
end

# Exclude generative and uses All
function _get_supports(pref, gen::Val{false}, info, label::Type{All})::Vector{Float64}
    gen_label = support_label(info)
    return findall(x -> any(v -> v != gen_label, x), _parameter_supports(pref))
end

# Exclude generative and does not use All
function _get_supports(pref, gen::Val{false}, info, label)::Vector{Float64}
    gen_label = support_label(info)
    return supps = findall(x -> any(v -> v <: label && v != gen_label, x),
                           _parameter_supports(pref))
end

# 1D FunctionalDiscreteMeasureData (not NoGenerativeSupports)
function supports(
    data::FunctionalDiscreteMeasureData{GeneralVariableRef};
    include_generative::Bool = true
    )::Vector{Float64}
    info = generative_support_info(data)
    label = support_label(data)
    pref = dispatch_variable_ref(parameter_refs(data))
    supps = _get_supports(pref, Val(include_generative), info, label)
    lb = JuMP.lower_bound(data)
    ub = JuMP.upper_bound(data)
    if isnan(lb) && isnan(ub)
        return supps
    else
        return filter!(i -> lb <= i <= ub, supps)
    end
end

# 1D FunctionalDiscreteMeasureData (NoGenerativeSupports)
function supports(
    data::FunctionalDiscreteMeasureData{GeneralVariableRef, Float64, NoGenerativeSupports};
    include_generative::Bool = true
    )::Vector{Float64}
    supps = supports(parameter_refs(data), label = support_label(data))
    lb = JuMP.lower_bound(data)
    ub = JuMP.upper_bound(data)
    if isnan(lb) && isnan(ub)
        return supps
    else
        return filter!(i -> lb <= i <= ub, supps)
    end
end

# Multi FunctionalDiscreteMeasureData
function supports(
    data::FunctionalDiscreteMeasureData{Vector{GeneralVariableRef}};
    include_generative::Bool = true
    )::Matrix{Float64}
    supps = supports(parameter_refs(data), label = support_label(data))
    lb = JuMP.lower_bound(data)
    ub = JuMP.upper_bound(data)
    if isnan(first(lb)) && isnan(first(ub))
        return supps
    else
        inds = [all(lb .<= s .<= ub) for s in eachcol(supps)]
        return supps[:, inds]
    end
end

"""
    num_supports(data::AbstractMeasureData)::Int

Return the number supports associated with `data` and its infinite parameters.
This is intended as an internal method for measure creation. User-defined
measure data types should extend this function if appropriate, otherwise
0 is returned.
"""
function num_supports(data::AbstractMeasureData)::Int
    return 0
end

# DiscreteMeasureData/FunctionalDiscreteMeasureData
function num_supports(data::Union{FunctionalDiscreteMeasureData,
                                  DiscreteMeasureData})::Int
    return size(supports(data))[end]
end

"""
    min_num_supports(data::AbstractMeasureData)::Int

Return the minimum number of supports associated with `data`. By fallback, this
will just return `num_supports(data)`. This is primarily intended for internal
queries of `FunctionalDiscreteMeasureData`, but can be extended for other
measure data types if needed.
"""
function min_num_supports(data::AbstractMeasureData)::Int
    return num_supports(data)
end

# FunctionalDiscreteMeasureData
function min_num_supports(data::FunctionalDiscreteMeasureData)::Int
    return data.min_num_supports
end

"""
    coefficient_function(data::AbstractMeasureData)::Function

Return the coefficient function stored in `data` associated with its
expansion abstraction is there is such a function. This is intended as an
internal method for measure creation. User-defined measure
data types should extend this function if appropriate, otherwise an error is
thrown for unsupported types.
"""
function coefficient_function(data::AbstractMeasureData)
    error("Function `coefficient_function` not defined for measure data of type " *
          "$(typeof(data)).")
end

# FunctionalDiscreteMeasureData
function coefficient_function(
    data::FunctionalDiscreteMeasureData{P, B, I, F}
    )::F where {P, B, I, F}
    return data.coeff_function
end

"""
    coefficients(data::AbstractMeasureData)::Vector{<:Real}

Return the coefficients associated with `data` associated with its expansion abstraction.
This is intended as en internal method for measure creation. User-defined measure
data types should extend this function if appropriate, otherwise an empty vector
is returned.
"""
function coefficients(data::AbstractMeasureData)::Vector{Float64}
    return Float64[]
end

# DiscreteMeasureData
function coefficients(data::DiscreteMeasureData)::Vector{Float64}
    return data.coefficients
end

# FunctionalDiscreteMeasureData
function coefficients(data::FunctionalDiscreteMeasureData)::Vector{Float64}
    return data.coeff_function(supports(data, include_generative = false))
end

"""
    weight_function(data::AbstractMeasureData)::Function

Return the weight function stored in `data` associated with its expansion abstraction.
This is intended as en internal method for measure creation. User-defined measure
data types should extend this function if appropriate, otherwise an error is thrown.
"""
function weight_function(data::AbstractMeasureData)::Function
    error("Function `weight_function` not defined for measure data of type " *
          "$(typeof(data)).")
end

# DiscreteMeasureData and FunctionalDiscreteMeasureData
function weight_function(
    data::Union{DiscreteMeasureData, FunctionalDiscreteMeasureData}
    )
    return data.weight_function
end

################################################################################
#                            MEASURE CONSTRUCTION METHODS
################################################################################
"""
    build_measure(expr::JuMP.AbstractJuMPScalar,
                  data::AbstractMeasureData)::Measure

Build and return a [`Measure`](@ref) given the expression to be measured `expr`
using measure data `data`. This principally serves as an internal method for
measure definition. Errors if the supports associated with `data` violate
an finite variable parameter bounds of finite variables that are included in the
measure.
"""
function build_measure(
    expr::T,
    data::D;
    )::Measure{T, D} where {T <: JuMP.AbstractJuMPScalar, D <: AbstractMeasureData}
    vrefs = _all_function_variables(expr)
    model = _model_from_expr(vrefs)
    expr_obj_nums = _object_numbers(expr)
    expr_param_nums = _parameter_numbers(expr)
    prefs = parameter_refs(data)
    data_obj_nums = _object_numbers(prefs)
    data_param_nums = [_parameter_number(pref) for pref in prefs]
    # NOTE setdiff! cannot be used here since it modifies object_nums of expr if expr is a single infinite variable
    obj_nums = sort(setdiff(expr_obj_nums, data_obj_nums))
    param_nums = sort(setdiff(expr_param_nums, data_param_nums))
    # check if analytic method should be applied
    lb_nan = isnan(first(JuMP.lower_bound(data)))
    ub_nan = isnan(first(JuMP.upper_bound(data)))
    # NOTE intersect! cannot be used here since it modifies parameter_nums of expr if expr is a single infinite variable
    constant_func = isempty(intersect(expr_param_nums, data_param_nums)) &&
                    ((!lb_nan && !ub_nan) || _is_expect(data))
    return Measure(expr, data, obj_nums, param_nums, constant_func)
end

################################################################################
#                               DEFINITION METHODS
################################################################################
## Define methods to check a measure's generative info and add it if needed
# DependentParameterRef and NoGenerativeSupports
function _check_and_set_generative_info(
    pref::DependentParameterRef,
    info::NoGenerativeSupports
    )::Nothing
    return
end

# DependentParameterRef and not NoGenerativeSupports
function _check_and_set_generative_info(
    pref::DependentParameterRef,
    info
    )::Nothing
    error("Cannot use a measure that requires generative supports with respect to " *
          "a dependent infinite parameter.")
end

# IndependentParameterRef and NoGenerativeSupports
function _check_and_set_generative_info(
    pref::IndependentParameterRef,
    info::NoGenerativeSupports
    )::Nothing
    return
end

# IndependentParameterRef w/ NoGenerativeSupports
function _check_and_set_generative_info(
    pref,
    info1::NoGenerativeSupports,
    info2
    )::Nothing
    _set_generative_support_info(pref, info2)
    return
end

# IndependentParameterRef w/ info already
function _check_and_set_generative_info(pref, info1, info2)::Nothing
    info1 != info2 && error("The generative support structure associated with " *
                            "the desired measure conflicts with the current scheme. " *
                            "This is likely due to the derivative method and/or other measures.")
    return
end

# IndependentParameterRef and not NoGenerativeSupports
function _check_and_set_generative_info(
    pref::IndependentParameterRef,
    info::AbstractGenerativeInfo
    )::Nothing
    current_info = generative_support_info(pref)
    _check_and_set_generative_info(pref, current_info, info)
    return
end

## Helper methods for adding supports
function _add_supports_to_multiple_parameters(
    prefs::Vector{DependentParameterRef},
    supps::Array{Float64, 2},
    label::Type{<:AbstractSupportLabel}
    )::Nothing
    add_supports(prefs, supps, label = label, check = false)
    return
end

function _add_supports_to_multiple_parameters(
    prefs::Vector{IndependentParameterRef},
    supps::Array{Float64, 2},
    label::Type{<:AbstractSupportLabel}
    )::Nothing
    for i in eachindex(prefs)
        add_supports(prefs[i], supps[i, :], label = label, check = false)
    end
    return
end

"""
    add_supports_to_parameters(data::AbstractMeasureData)::Nothing

Add supports as appropriate with `data` to the underlying infinite parameters.
This is an internal method with by [`add_measure`](@ref) and should be defined
for user-defined measure data types.
"""
function add_supports_to_parameters(data::AbstractMeasureData)
    error("`add_supports_to_parameters` not defined for measures with " *
          "measure data type $(typeof(data)).")
end

## Internal functions for adding measure data supports to the parameter supports
# scalar DiscreteMeasureData
function add_supports_to_parameters(
    data::DiscreteMeasureData{GeneralVariableRef, 1}
    )::Nothing
    pref = parameter_refs(data)
    supps = supports(data)
    label = support_label(data)
    add_supports(pref, supps, label = label, check = false)
    return
end

# multi-dimensional DiscreteMeasureData
function add_supports_to_parameters(
    data::DiscreteMeasureData{Vector{GeneralVariableRef}, 2}
    )::Nothing
    prefs = map(p -> dispatch_variable_ref(p), parameter_refs(data))
    supps = supports(data)
    label = support_label(data)
    _add_supports_to_multiple_parameters(prefs, supps, label)
    return
end

# scalar FunctionalDiscreteMeasureData
function add_supports_to_parameters(
    data::FunctionalDiscreteMeasureData{GeneralVariableRef}
    )::Nothing
    lb = JuMP.lower_bound(data)
    ub = JuMP.upper_bound(data)
    label = support_label(data)
    pref = dispatch_variable_ref(parameter_refs(data))
    if label == All && !isnan(lb) && !isnan(ub) && !isinf(lb) && !isinf(ub) && pref isa IndependentParameterRef
        add_supports(pref, [lb, ub], label = MeasureBound, check = false)
    end
    # determine if we need to add more supports
    num_supps = min_num_supports(data)
    curr_num_supps = num_supports(data)
    if curr_num_supps < num_supps
        # prepare the parameter reference
        if pref isa DependentParameterRef # This is just a last line of defense
            error("min_num_supports must be 0 for individual dependent parameters.")
        end
        # prepare the generation domain
        if isnan(lb) || isnan(ub)
            domain = infinite_domain(pref)
        else
            domain = IntervalDomain(lb, ub) # assumes lb and ub are in the domain
        end
        # generate the needed supports
        generate_and_add_supports!(pref, domain, label,
                                   num_supports = num_supps - curr_num_supps)
    end
    return
end

## generate more supports as needed in accordance with a label
# DependentParameterRefs
function _generate_multiple_functional_supports(
    prefs::Vector{DependentParameterRef},
    num_supps::Int, label::Type{<:AbstractSupportLabel},
    lower_bounds::Vector{<:Number},
    upper_bounds::Vector{<:Number}
    )::Nothing
    # prepare the generation domain
    if isnan(first(lower_bounds)) || isnan(first(upper_bounds))
        domain = infinite_domain(prefs)
    else
        # assumes we have valid bounds
        domain = CollectionDomain([IntervalDomain(lower_bounds[i], upper_bounds[i])
                             for i in eachindex(lower_bounds)])
    end
    # generate the supports
    generate_and_add_supports!(prefs, domain, label, num_supports = num_supps)
    return
end

# IndependentParameterRefs
function _generate_multiple_functional_supports(
    prefs::Vector{IndependentParameterRef},
    num_supps::Int, label::Type{<:AbstractSupportLabel},
    lower_bounds::Vector{<:Number},
    upper_bounds::Vector{<:Number}
    )::Nothing
    # we are gauranteed that each have the same number of supports
    for i in eachindex(prefs)
        # prepare the generation domain
        if isnan(lower_bounds[i]) || isnan(upper_bounds[i])
            domain = infinite_domain(prefs[i])
        else
            # assumes lb and ub are in the domain
            domain = IntervalDomain(lower_bounds[i], upper_bounds[i])
        end
        # generate the supports
        generate_and_add_supports!(prefs[i], domain, label, num_supports = num_supps)
    end
    return
end

# multi-dimensional FunctionalDiscreteMeasureData
function add_supports_to_parameters(
    data::FunctionalDiscreteMeasureData{Vector{GeneralVariableRef}}
    )::Nothing
    prefs = map(p -> dispatch_variable_ref(p), parameter_refs(data))
    label = support_label(data)
    lbs = JuMP.lower_bound(data)
    ubs = JuMP.upper_bound(data)
    if label == All && !any(isnan.(lbs)) && !any(isnan.(ubs)) && !any(isinf.(lbs)) && !any(isinf.(ubs))
        _add_supports_to_multiple_parameters(prefs, hcat(lbs, ubs), MeasureBound)
    end
    min_num_supps = min_num_supports(data)
    curr_num_supps = num_supports(data) # this will error check the support dims
    needed_supps = min_num_supps - curr_num_supps
    if needed_supps > 0
        _generate_multiple_functional_supports(prefs, needed_supps, label, lbs,
                                               ubs)
    end
    return
end

## Helper methods for updating the generative measure dependencies
# NoGenerativeSupports
function _update_generative_measures(
    prefs,
    info::NoGenerativeSupports,
    mindex
    )::Nothing
    return
end

# not NoGenerativeSupports and must be an IndependentParameterRef
function _update_generative_measures(pref, info, mindex)::Nothing
    push!(_generative_measures(pref), mindex)
    return
end

"""
    add_measure(model::InfiniteModel, meas::Measure,
                name::String = "measure")::GeneralVariableRef

Add a measure to `model` and return the corresponding measure reference. This
operates in a manner similar to `JuMP.add_variable`. Note this intended
as an internal method.
"""
function add_measure(
    model::InfiniteModel,
    meas::Measure,
    name::String = "measure"
    )::GeneralVariableRef
    # get the expression variables and check validity
    vrefs = _all_function_variables(meas.func)
    for vref in vrefs
        JuMP.check_belongs_to_model(vref, model)
    end
    # get the measure data info and check validity
    data = meas.data
    prefs = parameter_refs(data)
    for pref in prefs
        JuMP.check_belongs_to_model(pref, model)
    end
    # check generative info is compatible and update model if needed
    info = generative_support_info(data)
    _check_and_set_generative_info(dispatch_variable_ref(first(prefs)), info)
    # add supports to the model as needed
    if !meas.constant_func
        add_supports_to_parameters(data)
    end
    # add the measure to the model
    object = MeasureData(meas, name)
    mindex = _add_data_object(model, object)
    mref = _make_variable_ref(model, mindex)
    # update mappings
    for vref in union!(vrefs, prefs)
        push!(_measure_dependencies(vref), mindex)
    end
    _update_generative_measures(prefs, info, mindex)
    return mref
end

"""
    measure_function(mref::MeasureRef)::JuMP.AbstractJuMPScalar

Return the function associated with `mref`.

**Example**
```julia-repl
julia> measure_function(meas)
y(x, t) + 2
```
"""
function measure_function(mref::MeasureRef)::JuMP.AbstractJuMPScalar
    return _core_variable_object(mref).func
end

"""
    measure_data(mref::MeasureRef)::AbstractMeasureData

Return the measure data associated with `mref`.

**Example**
```julia-repl
julia> data = measure_data(meas);

julia> typeof(data)
FunctionalDiscreteMeasureData{Vector{GeneralVariableRef},Vector{Float64}}
```
"""
function measure_data(mref::MeasureRef)::AbstractMeasureData
    return _core_variable_object(mref).data
end

"""
    is_analytic(mref::MeasureRef)::Bool

Return if `mref` is evaluated analytically.

**Example**
```julia-repl
julia> is_analytic(meas)
false
```
"""
function is_analytic(mref::MeasureRef)::Bool
    return _core_variable_object(mref).constant_func
end

## Return an element of a parameter reference tuple given the model, index, and parameter numbers
# IndependentParameterIndex
function _make_param_tuple_element(
    model::InfiniteModel,
    idx::IndependentParameterIndex,
    param_nums::Vector{Int}
    )::GeneralVariableRef
    return _make_parameter_ref(model, idx)
end

# DependentParametersIndex
function _make_param_tuple_element(
    model::InfiniteModel,
    idx::DependentParametersIndex,
    param_nums::Vector{Int}
    )::Union{GeneralVariableRef, Vector{GeneralVariableRef}}
    dpref = DependentParameterRef(model, DependentParameterIndex(idx, 1))
    el_param_nums = _data_object(dpref).parameter_nums
    prefs = [GeneralVariableRef(model, idx.value, DependentParameterIndex, i)
             for i in eachindex(el_param_nums) if el_param_nums[i] in param_nums]
    return length(prefs) > 1 ? prefs : first(prefs)
end

"""
    parameter_refs(mref::MeasureRef)::Tuple

Return the tuple of infinite parameters that the measured expression associated
`mref` depends on once the measure has been evaluated. Note that this will
correspond to the parameter dependencies of the measure function excluding those
included in the measure data.

**Example**
```julia-repl
julia> parameter_refs(meas)
(t,)
```
"""
function parameter_refs(mref::MeasureRef)
    model = JuMP.owner_model(mref)
    obj_indices = _param_object_indices(model)[_object_numbers(mref)]
    param_nums = _parameter_numbers(mref)
    return Tuple(_make_param_tuple_element(model, idx, param_nums)
                 for idx in obj_indices)
end

# Extend raw_parameter_refs (this is helpful for defining derivatives)
function raw_parameter_refs(mref::MeasureRef)
    return Collections.VectorTuple(parameter_refs(mref))
end

# Extend parameter_list (this is helpful for defining derivatives)
function parameter_list(mref::MeasureRef)::Vector{GeneralVariableRef}
    return raw_parameter_refs(mref).values
end

"""
    measure(expr::JuMP.AbstractJuMPScalar,
            data::AbstractMeasureData;
            [name::String = "measure"])::GeneralVariableRef

Return a measure reference that evaluates `expr` using according to `data`.
The measure data `data` determines how the measure is to be evaluated.
Typically, the `DiscreteMeasureData` and the `FunctionalDiscreteMeasureData`
constructors can be used to for `data`. The variable expression `expr` can contain
`InfiniteOpt` variables, infinite parameters, other measure references (meaning
measures can be nested), and constants. Typically, this is called inside of
`JuMP.@expression`, `JuMP.@objective`, and `JuMP.@constraint` in a manner similar
to `sum`. Note measures are not explicitly evaluated until
[`build_optimizer_model!`](@ref) is called or unless they are expanded via
[`expand`](@ref) or [`expand_all_measures!`](@ref).

**Example**
```julia-repl
julia> tdata = DiscreteMeasureData(t, [0.5, 0.5], [1, 2]);

julia> xdata = DiscreteMeasureData(xs, [0.5, 0.5], [[-1, -1], [1, 1]]);

julia> constr_RHS = @expression(model, measure(g - s + 2, tdata) + s^2)
measure{t}[g(t) - s + 2] + s²

julia> @objective(model, Min, measure(g - 1  + measure(T, xdata), tdata))
measure{xs}[g(t) - 1 + measure{xs}[T(t, x)]]
```
"""
function measure(
    expr::JuMP.AbstractJuMPScalar,
    data::AbstractMeasureData;
    name::String = "measure"
    )::GeneralVariableRef
    model = _model_from_expr(expr)
    if isnothing(model)
        error("Expression contains no variables or parameters.")
    end
    meas = build_measure(expr, data)
    return add_measure(model, meas, name)
end

"""
    @measure(expr::JuMP.AbstractJuMPScalar,
             data::AbstractMeasureData;
             [name::String = "measure"])::GeneralVariableRef

An efficient wrapper for [`measure`](@ref), please see its doc string for more
information.
"""
macro measure(expr, data, args...)
    _error(str...) = _macro_error(:measure, (expr, args...), str...)
    extra, kwargs, _, _ = _extract_kwargs(args)
    if length(extra) != 0
        _error("Incorrect number of arguments. Must be of form " *
               "@measure(expr, data, name = ...).")
    end
    if !isempty(filter(kw -> kw.args[1] != :name, kwargs))
        _error("Invalid keyword arguments. Must be of form " *
               "@measure(expr, data, name = ...).")
    end
    expression = _MA.rewrite_and_return(expr)
    esc_kwargs = map(i -> esc(i), kwargs)
    mref = :( measure($expression, $(esc(data)); ($(esc_kwargs...))) )
    return mref
end

################################################################################
#                               NAMING METHODS
################################################################################
"""
    JuMP.name(mref::MeasureRef)::String

Extend `JuMP.name` to return the name associated with a measure
reference.
"""
function JuMP.name(mref::MeasureRef)::String
    object = get(_data_dictionary(mref), JuMP.index(mref), nothing)
    return isnothing(object) ? "" : object.name
end

"""
    JuMP.set_name(mref::MeasureRef, name::String)::Nothing

Extend `JuMP.set_name` to specify the name of a measure reference.
"""
function JuMP.set_name(mref::MeasureRef, name::String)::Nothing
    _data_object(mref).name = name
    return
end

################################################################################
#                              DEPENDENCY METHODS
################################################################################
"""
    used_by_measure(mref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used by a measure.

**Example**
```julia-repl
julia> used_by_measure(mref)
true
```
"""
function used_by_measure(mref::MeasureRef)::Bool
    return !isempty(_measure_dependencies(mref))
end

"""
    used_by_constraint(mref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used by a constraint.

**Example**
```julia-repl
julia> used_by_constraint(mref)
false
```
"""
function used_by_constraint(mref::MeasureRef)::Bool
    return !isempty(_constraint_dependencies(mref))
end

"""
    used_by_objective(mref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used by the objective.

**Example**
```julia-repl
julia> used_by_objective(mref)
true
```
"""
function used_by_objective(mref::MeasureRef)::Bool
    return _data_object(mref).in_objective
end

"""
    used_by_derivative(mref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used by a derivative.

**Example**
```julia-repl
julia> used_by_derivative(mref)
true
```
"""
function used_by_derivative(mref::MeasureRef)::Bool
    return !isempty(_derivative_dependencies(mref))
end

"""
    is_used(mref::MeasureRef)::Bool

Return a `Bool` indicating if `mref` is used in the model.

**Example**
```julia-repl
julia> is_used(mref)
true
```
"""
function is_used(mref::MeasureRef)::Bool
    return used_by_measure(mref) || used_by_constraint(mref) ||
           used_by_objective(mref) || used_by_derivative(mref)
end

################################################################################
#                                GENERAL QUERIES
################################################################################
"""
    num_measures(model::InfiniteModel)::Int

Return the number of measures defined in `model`.

**Example**
```julia-repl
julia> num_measures(model)
2
```
"""
function num_measures(model::InfiniteModel)::Int
    return length(_data_dictionary(model, Measure))
end

"""
    all_measures(model::InfiniteModel)::Vector{GeneralVariableRef}

Return the list of all measures added to `model`.

**Examples**
```julia-repl
julia> all_measures(model)
2-element Array{GeneralVariableRef,1}:
 ∫{t ∈ [0, 6]}[w(t, x)]
 𝔼{x}[w(t, x)]
```
"""
function all_measures(model::InfiniteModel)::Vector{GeneralVariableRef}
    vrefs_list = Vector{GeneralVariableRef}(undef, num_measures(model))
    for (i, (index, _)) in enumerate(_data_dictionary(model, Measure))
        vrefs_list[i] = _make_variable_ref(model, index)
    end
    return vrefs_list
end

################################################################################
#                                   DELETION
################################################################################
"""
    JuMP.delete(model::InfiniteModel, mref::MeasureRef)::Nothing

Extend `JuMP.delete` to delete measures. Errors if measure is invalid,
meaning it does not belong to the model or it has already been deleted.

**Example**
```julia-repl
julia> print(model)
Min ∫{t ∈ [0, 6]}[g(t)] + z
Subject to
 z ≥ 0.0
 ∫{t ∈ [0, 6]}[g(t)] = 0
 g(t) + z ≥ 42.0, ∀ t ∈ [0, 6]
 g(0.5) = 0

julia> delete(model, meas)

julia> print(model)
Min z
Subject to
 z ≥ 0.0
 0 = 0
 g(t) + z ≥ 42.0, ∀ t ∈ [0, 6]
 g(0.5) = 0
```
"""
function JuMP.delete(model::InfiniteModel, mref::MeasureRef)::Nothing
    @assert JuMP.is_valid(model, mref) "Invalid measure reference."
    # Reset the transcription status
    if is_used(mref)
        set_optimizer_model_ready(model, false)
    end
    gvref = _make_variable_ref(JuMP.owner_model(mref), JuMP.index(mref))
    # Remove from dependent measures if there are any
    for mindex in _measure_dependencies(mref)
        meas_ref = dispatch_variable_ref(model, mindex)
        func = measure_function(meas_ref)
        data = measure_data(meas_ref)
        if func isa GeneralVariableRef
            new_func = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
            new_meas = Measure(new_func, data, Int[], Int[], false)
        else
            _remove_variable(func, gvref)
            new_meas = build_measure(func, data)
        end
        _set_core_variable_object(meas_ref, new_meas)
    end
    # Remove from dependent constraints if there are any
    for cindex in copy(_constraint_dependencies(mref))
        cref = _make_constraint_ref(model, cindex)
        func = JuMP.jump_function(JuMP.constraint_object(cref))
        if func isa GeneralVariableRef
            set = JuMP.moi_set(JuMP.constraint_object(cref))
            new_func = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
            new_constr = JuMP.ScalarConstraint(new_func, set)
            _set_core_constraint_object(cref, new_constr)
            empty!(_object_numbers(cref))
            empty!(_measure_dependencies(cref))
        elseif func isa AbstractArray && any(isequal(gvref), func)
            JuMP.delete(model, cref)
        else
            _remove_variable(func, gvref)
            _data_object(cref).object_nums = sort(_object_numbers(func))
            filter!(e -> e != JuMP.index(mref), _measure_dependencies(cref))
        end
    end
    # Remove from objective if used there
    if used_by_objective(mref)
        if JuMP.objective_function(model) isa GeneralVariableRef
            new_func = zero(JuMP.GenericAffExpr{Float64, GeneralVariableRef})
            JuMP.set_objective_function(model, new_func)
            JuMP.set_objective_sense(model, MOI.FEASIBILITY_SENSE)
        else
            _remove_variable(JuMP.objective_function(model), gvref)
        end
    end
    # delete associated derivative variables and mapping
    for index in _derivative_dependencies(mref)
        JuMP.delete(model, dispatch_variable_ref(model, index))
    end
    # Update that the variables used by it are no longer used by it
    prefs = parameter_refs(measure_data(mref))
    vrefs = _all_function_variables(measure_function(mref))
    union!(vrefs, prefs)
    for vref in vrefs
        filter!(e -> e != JuMP.index(mref), _measure_dependencies(vref))
    end
    # Remove any unique supports associated with this measure
    if num_supports(measure_data(mref)) > 0
        label = support_label(measure_data(mref))
        if label <: UniqueMeasure
            delete_supports(prefs, label = label)
        end
    end
    # update the generative info if needed
    info = generative_support_info(measure_data(mref))
    if prefs isa GeneralVariableRef && !(info isa NoGenerativeSupports)
        filter!(e -> e != JuMP.index(mref), _generative_measures(prefs))
        if isempty(_generative_measures(prefs)) && derivative_method(prefs) isa NonGenerativeDerivativeMethod
            _set_generative_support_info(dispatch_variable_ref(prefs),
                                         NoGenerativeSupports())
        end
    end
    # delete remaining measure information
    _delete_data_object(mref)
    return
end


text:
### A Pluto.jl notebook ###
# v0.17.5

using Markdown
using InteractiveUtils

# ╔═╡ 84739f36-65fb-4a36-a342-ae0985a759a7
begin
    import Pkg
    # activate the shared project environment|
    Pkg.activate(Base.current_project())
    using Omega, Distributions, UnicodePlots
	using Plots
	Pkg.add("FreqTables")
	using FreqTables
end

# ╔═╡ 0042d812-7771-11eb-0aff-dd1d6767a96a
md""" # Open World Model
Omega supports so-called open-world models (also called trans/variable dimensional models) which have a variable number of parameters.

BLOG was perhaps the first probabilistic language to support open-world models.  The BLOG paper motivates them through a number of examples such as the following:

*An  unknown  number  of  aircraft  exist  in  some volume of airspace.  An aircraft’s state (position and veloc-ity) at each time step depends on its state at the previous timestep. We observe the area with radar: aircraft may appear as identical blips on a radar screen. Each blip gives the approximate position of the aircraft that generated it. However, some blips may be false detections, and some aircraft may not be detected. What aircraft exist, and what are their trajectories? Are there any aircraft that are not observed?*
"""

# ╔═╡ 05e99288-7787-11eb-2d06-f96f8b7f16e0
md"## The generative model"

# ╔═╡ 3ee1698e-9b28-4693-80ca-2ebcb51e7d7f
md"An aircraft is defined by a position and velocity"

# ╔═╡ 2d7e1730-7771-11eb-1d47-c9e05121d38b
struct Aircraft{P, V}
  position::P
  velocity::V
end

# ╔═╡ 15aeba40-7787-11eb-36de-25338cdfbca3
md"The number of aircraft is unknown, and for each aircraft the position and velocity is uncertain"

# ╔═╡ 66794654-7771-11eb-3044-ad6cf0b566f0
num_aircraft = @~ Poisson(5)

# ╔═╡ a4756686-7773-11eb-01bc-b1231fd5418a
xlb, xub, ylb, yub = 0, 1, 0, 1

# ╔═╡ aac835ae-7771-11eb-2792-838c80957a3d
function init_state(i, ω)
  x = @~ i Uniform(xlb, xub)
  y = @~ i Uniform(ylb, yub)
  vx = @~ i Normal(0, 1)
  vy = @~ i Normal(0, 1)
  Aircraft((x(ω), y(ω)), (vx(ω), vy(ω)))
end

# ╔═╡ eefa9a2c-778b-11eb-37d1-c39cc9cc53f7

# ╔═╡ 1ff3e3a2-7775-11eb-2c0c-7d255b3404ef
aircraft(ω) = manynth(init_state, num_aircraft(ω))(ω)

# ╔═╡ 2d60c906-b5cb-4e44-951c-e52f32568b73
randsample(aircraft)

# ╔═╡ af7af430-7774-11eb-3133-bfe35d03939d
state_transition(s; dt = 0.1) =
  Aircraft(s.position .+ dt .* s.velocity, s.velocity)

# ╔═╡ 6de7d402-7774-11eb-109d-5f8fb82851c7
"simulate the aircraft for `t` timesteps from initial state `s0`"
function simulate(s, t)
  ss = [s]
  for i = 1:t - 1
	s = state_transition(s)
	push!(ss, s)
  end
  ss
end

# ╔═╡ b415b1c8-7784-11eb-1a8b-0b437ab2432d
T = 20

# ╔═╡ ebed910c-7774-11eb-3bac-f7d7b7fbc2f1
simall(ω) = simulate(aircraft(ω), T)

# ╔═╡ 0b55508a-7783-11eb-2678-5f480c20c23f
traj_sample = randsample(simall)

# ╔═╡ 684c13f0-7788-11eb-1d92-ed4736e555a7
xtraj(ta) = [t.position[1] for t in ta]

# ╔═╡ 832d5de4-7788-11eb-3833-ad03ab2f9bf9
ytraj(ta) = [t.position[2] for t in ta]

# ╔═╡ 32bfdb5e-7783-11eb-35d7-53e6e75cb44e
begin
	plt = Plots.plot(xtraj(traj_sample), ytraj(traj_sample))
	if length(traj_sample) > 1
		for i = 2:length(traj_sample)
			@show i
			Plots.plot!(plt, xtraj(traj_sample), ytraj(traj_sample))
		end
	end
	plt
end

# ╔═╡ a6b85cf0-778a-11eb-0d3b-395d8678a185
md"""### Observation Model
Every $r$ seconds we'll send out a radar pulse and sometime later (at a time dependent on the distance from the source to the airraft) we'll recieve a pulse back.  There are a few caveats:

- false positive: some other object causes a pulse
- false negatives: sometimes a radar pulse willl scattar off an aircraft
"""

# ╔═╡ cdc5380a-7789-11eb-237f-57098bc805d4
function sim_radar(a_series)
  for i = 1:length(a_series)
  end
end

# ╔═╡ ec3c6662-7786-11eb-35e3-87845ca04b0a
md"## Inference"

# ╔═╡ 5ddb4910-778b-11eb-1954-93b4bf68e7df
md"""Given a model we can do inference, sample from the posterior of any of the variables given evidence"""

# ╔═╡ bc35478c-7785-11eb-1ec9-aff81d4aa339
length(traj_sample[1])

# ╔═╡ Cell order:
# ╠═84739f36-65fb-4a36-a342-ae0985a759a7
# ╟─0042d812-7771-11eb-0aff-dd1d6767a96a
# ╟─05e99288-7787-11eb-2d06-f96f8b7f16e0
# ╟─3ee1698e-9b28-4693-80ca-2ebcb51e7d7f
# ╠═2d7e1730-7771-11eb-1d47-c9e05121d38b
# ╟─15aeba40-7787-11eb-36de-25338cdfbca3
# ╠═66794654-7771-11eb-3044-ad6cf0b566f0
# ╠═a4756686-7773-11eb-01bc-b1231fd5418a
# ╠═aac835ae-7771-11eb-2792-838c80957a3d
# ╠═eefa9a2c-778b-11eb-37d1-c39cc9cc53f7
# ╠═1ff3e3a2-7775-11eb-2c0c-7d255b3404ef
# ╠═2d60c906-b5cb-4e44-951c-e52f32568b73
# ╠═af7af430-7774-11eb-3133-bfe35d03939d
# ╠═6de7d402-7774-11eb-109d-5f8fb82851c7
# ╠═b415b1c8-7784-11eb-1a8b-0b437ab2432d
# ╠═ebed910c-7774-11eb-3bac-f7d7b7fbc2f1
# ╠═0b55508a-7783-11eb-2678-5f480c20c23f
# ╠═684c13f0-7788-11eb-1d92-ed4736e555a7
# ╠═832d5de4-7788-11eb-3833-ad03ab2f9bf9
# ╠═32bfdb5e-7783-11eb-35d7-53e6e75cb44e
# ╟─a6b85cf0-778a-11eb-0d3b-395d8678a185
# ╠═cdc5380a-7789-11eb-237f-57098bc805d4
# ╟─ec3c6662-7786-11eb-35e3-87845ca04b0a
# ╟─5ddb4910-778b-11eb-1954-93b4bf68e7df
# ╠═bc35478c-7785-11eb-1ec9-aff81d4aa339

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
"""im a doc in Junk"""
const imwithdoc = nothing

module SubJunk

"""im a doc in SubJunk"""
const imwithdoc = nothing

end

baremodule BareJunk

"""im a doc in BareJunk"""
const imwithdoc = nothing

end


text:
@testset "GenCB_02" begin
    env = CPLEX.Env()
    CPLEX.set_param!(env, "CPX_PARAM_THREADS", 1)
    # Change some settings to force CPLEX to call the callback.
    CPLEX.set_param!(env, "CPX_PARAM_PREIND", 0)
    CPLEX.set_param!(env, "CPX_PARAM_HEURFREQ", -1)
    model = CPLEX.Model(env, "GenCB_02")
    #  max 0.5x + y
    # s.t. 0 <= x <= 1
    #      0 <= y <= 1
    CPLEX.add_var!(model, 0.5, 0, 1)
    CPLEX.add_var!(model, 1.0, 0, 1)
    CPLEX.set_sense!(model, :Max)
    CPLEX.set_vartype!(model, ['B', 'B'])
    CPLEX.add_constr!(model, [1.0, 1.0], '<', 1.5)
    was_heuristic_called = false
    function my_round_down_heur(
            cb_context::CPLEX.CallbackContext, context_id::Clong)
        if context_id == CPLEX.CPX_CALLBACKCONTEXT_RELAXATION
            relaxed_solution = Vector{Float64}(undef, 2)
            relaxed_objective = Ref{Float64}(0.0)
            CPLEX.cbgetrelaxationpoint(
                cb_context, relaxed_solution, 1, 2, relaxed_objective)
            objective_value = relaxed_objective[]
            for j in 1:2
                frac = relaxed_solution[j] - floor(relaxed_solution[j])
                if frac > 1.0e-6
                    objective_value -= frac
                    relaxed_solution[j] -= frac
                end
            end
            was_heuristic_called = true
            return CPLEX.cbpostheursoln(
                cb_context, 2, [1, 2], relaxed_solution, objective_value,
                CPLEX.CPXCALLBACKSOLUTION_PROPAGATE)
        else
            error("Heuristic should not be called from $(context_id).")
        end
    end
    CPLEX.cbsetfunc(
        model, CPLEX.CPX_CALLBACKCONTEXT_RELAXATION, my_round_down_heur)
    CPLEX.optimize!(model)
    @test was_heuristic_called
    @test CPLEX.get_status(model) == :CPXMIP_OPTIMAL
    @test CPLEX.get_objval(model) ≈ 1.0
end


text:
include("./imports.jl")

include("base.jl")
include("buffered.jl")
include("projective/base.jl")
include("projective/bounds.jl")
include("projective/compose.jl")
include("projective/crop.jl")
include("projective/affine.jl")
include("items/image.jl")
include("items/keypoints.jl")
include("items/mask.jl")
include("preprocessing.jl")


text:
################################################################################
### GENERAL-PURPOSE OBJECTIVE FUNCTIONS
################################################################################
"""
    likelihood(X, P, normalized=true)

Computes the likelihood of `X` given distribution `P`
"""
function likelihood(X, P::AbstractBinaryVectorDistribution, normalized=true)
    Px = [pdf(P, X[:,k]) for k = 1:size(X,2)]
    if normalized
        Px = Px .^ (1 / size(X,2))
    end
    return prod(Px)
end

"""
    loglikelihood(X, P, normalized=true)

Computes the log likelihood of `X` given distribution `P`. Returns `0.0` if `P`
assigns probability 0 to any column of `X`
"""
function loglikelihood(X, P::AbstractBinaryVectorDistribution, normalized=true)
    Px = [pdf(P, X[:,k]) for k = 1:size(X,2)]
    if any(Px .== 0.0)
        return 0.0
    end
    return sum_kbn(map(log,Px)) / (normalized ? size(X,2) : 1)
end

################################################################################
### ISING OBJECTIVE FUNCTIONS - NLopt.jl
################################################################################
"""
    loglikelihood(X, Jtilde, grad)

Computes the loglikelihood of the data given the Ising distribution generated by
`Jtilde`. If `grad` is a vector of length > 0, it is modified in-place with the
gradient of the function.

"""
function loglikelihood(X, Jtilde::Vector, grad::Vector=[]; mu_X=(X*X'/size(X,2)), kwargs...)
    (N_neurons, N_samples) = size(X)

    Jtilde = reshape(Jtilde, N_neurons,N_neurons)
    P = IsingDistribution(Jtilde)

    if length(grad) > 0
        # performing this computation first, if it's necessary, results in all
        # the pdf values being cached.
        # mu_X = X * X' / N_samples
        mu_P = expectation_matrix(P)
        grad[:] = -(mu_P - mu_X)[:]
        grad[1:(N_neurons+1):end] = diag(mu_P - mu_X)[:]

    end
    # return sum_kbn([log(pdf(P, X[:,k])) for k = 1:N_samples]) / N_samples
    return -log(_get_Z(P)) - (sum_kbn([_E_Ising(P, X[:,i]) for i = 1:N_samples])) / N_samples
end

"""
    MPF_objective(X, J, grad)

Uses function `K` from the MPF paper (or maybe from the MPF sample code for
Matlab). Using this to fit J will typically result in values close to optimal,
but the advantage is that this method does not involve computing the partition
function Z.

"""
function MPF_objective(X, Jtilde::Vector, grad::Vector=[])
    (N_neurons, N_samples) = size(X)
    Jtilde = reshape(Jtilde, N_neurons, N_neurons)
    theta = diag(Jtilde)
    J = Jtilde - Diagonal(Jtilde)
    DeltaX = 2 * X - 1 # this is Δx_l for each codeword x, for each index l
    # Kfull = exp.((-0.5 * DeltaX .* (J * X) + DeltaX .* theta)/2)
    Kfull = exp.((DeltaX .* theta - DeltaX .* (J * X)) / 2)
    K = sum_kbn(Kfull[:]) / N_samples
    if length(grad) > 0
        # M = zeros(N_neurons, N_neurons)
        # M[1:(N_neurons+1):end] = sum(0.5 * Kfull .* DeltaX, 2) / N_samples
        # for p = 1:(N_neurons - 1)
        #     for q = (p+1):N_neurons
        #         M[p,q] = M[q,p] = -0.5 * sum([(Kfull[p,w] * DeltaX[p,w] * X[q,w] + Kfull[q,w] * DeltaX[q,w] * X[p,w]) for w = 1:N_samples]) / N_samples
        #     end
        # end
        DK = Kfull .* DeltaX
        dJ = -0.5 * DK * X'
        dJ = (dJ + dJ') / 2
        dJ[1:(n+1):end] = 0.5 * sum(DK, 2)
        grad[:] = dJ[:] / N_samples
    end
    return K
end

################################################################################
### ISING OBJECTIVE FUNCTIONS - Optim.jl
################################################################################
"""
    negloglikelihood(X, Jtilde)

Computes negative LogLikelihood of the data given the Ising distribution generated by
`Jtilde`. See also `dnegloglikelihood!` for gradient. These functions are intended for use
with `Optim.jl` package

"""
function negloglikelihood(W, W_prev, buf)
    negloglikelihood_common!(W, W_prev, buf)
    N = size(buf[1], 2)
    return log(get_Z(buf[3])) - sum_kbn([F_Ising(buf[3], buf[1][:,α]) for α = 1:N]) / N
end
function dnegloglikelihood!(G, W, W_prev, buf)
    negloglikelihood_common!(W, W_prev, buf)
    #TODO compute gradient, modify G in place
    mu_P = expectation_matrix(buf[3])
    G[:] = lt2v(mu_P) - lt2v(buf[2])
end
function Hnegloglikelihood!(H, W, W_prev, buf)
    negloglikelihood_common!(W, W_prev, buf)
    # if we use lt2v on mu_P and take its outer product, we'll get the matrix of
    # E[x_s]E[x_t] that we subtract from E[x_s x_t] to get the Hessian matrix. So the only
    # thing we need to be careful about is making sure we line up terms correctly.
    N_neurons = n_bits(buf[3])
    p = get_pdf(buf[3])
    for k in sortperm(p)
        x = digits(Bool, k-1, 2, N_neurons)
        xx = lt2v(x * x')
        H[:] = H[:] + (p[k] * (xx * xx'))[:]
    end
    mm = lt2v(buf[4])
    H[:] = H[:] - (mm * mm')[:]
end
function negloglikelihood_common!(W, W_prev, buf)
    # buf[1] = X
    # buf[2] = mu_X
    # buf[3] = P
    # buf[4] = mu_P
    if W != W_prev
        copy!(W_prev, W)
        #TODO set the appropriate buffer entry to be P
        buf[3] = IsingDistribution(W)
        buf[4] = expectation_matrix(buf[3])
    end
end

"""
    K_MPF(J, J_prev, buf)

Computes the function ``K_X`` for data in `buf[1]` and parameters `J`. This version is meant
for use with the `Optim` package, which in particular uses a separate function for the
optimization function and gradient function. See also `dK_MPF!`

"""
function K_MPF(J, J_prev, buf)
    # N_neurons, N_samples = size(X)
    # J = reshape(Jtilde, N_neurons, N_neurons)
    # theta = diag(J)
    # J[1:(N_neurons+1):end] = 0.0
    # ΔX = 2X - 1 # flip bits in X
    # Kfull = exp.(0.5 * (ΔX .* theta - ΔX .* (J * X)))
    # return sum_kbn(Kfull) / N_samples
    K_common!(J, J_prev, buf)
    # N_samples = size(buf[1], 2)
    return sum_kbn(buf[3]) # / N_samples
end
"""
    dK_MPF!(G, J, J_prev, buf)

Computes the gradient of `K_MPF(J, J_prev, buf)` and modifies `G` in-place.
"""
function dK_MPF!(G, J, J_prev, buf)
    # N_neurons, N_samples = size(X)
    # J = reshape(Jtilde, N_neurons, N_neurons)
    # theta = diag(J)
    # J[1:(N_neurons+1):end] = 0.0
    # ΔX = 2X - 1 # flip bits in X
    # Kfull = exp.(0.5 * (ΔX .* theta - ΔX .* (J * X)))
    # # M = zeros(N_neurons, N_neurons)
    # # M[1:(N_neurons+1):N_neurons] = sum(0.5 * ΔX .* Kfull, 2)
    # DK = Kfull .* ΔX
    # dJ = -0.5 * (Kfull .* ΔX) * X'
    # dJ = (dJ + dJ') / 2
    # dJ[1:(N_neurons+1):end] = 0.5 * sum(DK, 2)
    # G[:] = dJ[:] / N_samples
    K_common!(J, J_prev, buf)
    N_neurons, N_samples = size(buf[1])
    DK = buf[3] .* buf[2]
    dJ = -0.5 * DK * buf[1]'
    dJ = (dJ + dJ') / 2
    dJ[1:(N_neurons+1):end] = 0.5 * sum(DK, 2)
    G[:] = dJ[:] # / N_samples
end
function K_common!(J, J_prev, buf)
    # buf[1] = X # fixed ahead of time
    # buf[2] = ΔX # fixed ahead of time
    # buf[3] = Kfull / N_samples
    # buf[4] = theta
    # buf[5] = J (no diagonal)
    if J != J_prev
        copy!(J_prev, J)

        N_neurons, N_samples = size(buf[1])
        J = reshape(J, N_neurons, N_neurons)
        theta = diag(J)
        J[1:(N_neurons+1):N_neurons] = 0.0
        buf[4] = theta
        buf[5] = J
        buf[3] = exp.(0.5 * (buf[2] .* theta - buf[2] .* (J * buf[1]))) / N_samples
    end
end

max_stars_repo_name: sekunder/BinaryVectorProbability

max_issues_repo_name: sekunder/BinaryVectorProbability

max_forks_repo_name: sekunder/BinaryVectorProbability


text:
## --------------  BoundedVector

"""
	$(SIGNATURES)

Initialize the `ParamVector`. Requires `dxV` to be set.
Note that bounds on `dxV` must be between 0 and 1.
This is called after the `BoundedVector` has been constructed with an empty `ParamVector` but with the values of the field `dxV` set to reasonable defaults.

# Example
```
dxV = [0.3, 0.2, 0.8];
b = BoundedVector(objId, ParamVector(objId), true, 1.0, 2.0, dxV);
set_pvector!(b; description = "Gradient", symbol = "g(x)");
```
"""
function set_pvector!(iv :: BoundedVector{T1};
    description = "Increments",
    symbol = "dxV",
    isCalibrated :: Bool = true) where T1

    defaultValueV = iv.dxV;
    n = length(defaultValueV);
    p = Param(:dxV, description, symbol, defaultValueV, defaultValueV,
        zeros(T1, n), ones(T1, n), isCalibrated);
    iv.pvec = ParamVector(iv.objId, [p]);
    return nothing
end

is_increasing(iv :: BoundedVector{T1}) where T1 = (iv.increasing == :increasing);
is_decreasing(iv :: BoundedVector{T1}) where T1 = (iv.increasing == :decreasing);
is_nonmonotone(iv :: BoundedVector{T1}) where T1 = (iv.increasing == :nonmonotone);
lb(iv :: BoundedVector{T1}) where T1 = iv.lb;
ub(iv :: BoundedVector{T1}) where T1 = iv.ub;
Base.length(iv :: BoundedVector{T1}) where T1 =  Base.length(iv.dxV);

# Scalar bounds across all elements
scalar_lb(iv :: BoundedVector{T1}) where T1 = iv.lb;
scalar_ub(iv :: BoundedVector{T1}) where T1 = iv.ub;

function calibrate_values!(iv :: BoundedVector{T1}) where T1
    calibrate!(iv, :dxV);
end

"""
	$(SIGNATURES)

Switches calibration toggle off. Sets values and default values everywhere. The end result is a `BoundedVector` with fixed (not calibrated) increments that result in values of `valueV`.
"""
function fix_values!(iv :: BoundedVector{T1},
    valueV :: AbstractVector{T1}) where T1

    p = retrieve(iv.pvec, :dxV);
    set_value!(p, valueV);
    set_default_value!(p, valueV);
    fix!(p);
    dxV = values_to_dx(iv, valueV);
    iv.dxV = dxV;
end

"""
	$(SIGNATURES)

Returns all values of a `BoundedVector`.
"""
function values(iv :: BoundedVector{T1}) where T1
    valueV = dx_to_values(iv, iv.dxV);
    return valueV
end

function dx_to_values(iv :: Union{BoundedVector{T1}, BVector{T1}}, dxV) where T1
    # n = length(iv);
    if is_increasing(iv)
        valueV = dx_to_values_increasing(dxV, scalar_lb(iv), scalar_ub(iv))
    elseif is_decreasing(iv)
        valueV = dx_to_values_decreasing(dxV, scalar_lb(iv), scalar_ub(iv))
    elseif is_nonmonotone(iv)
        valueV = copy(dxV);
    else
        error("Invalid");
    end
    return valueV
end

function dx_to_values_increasing(dxV, lb, ub)
    T1 = eltype(dxV);
    n = length(dxV);
    valueV = zeros(T1, n);
    valueV[1] = lb + dxV[1] * (ub - lb);
    if n > 1
        for j = 2 : n
            valueV[j] = valueV[j-1] + dxV[j] * (ub - valueV[j-1]);
        end
    end
    return valueV
end

function dx_to_values_decreasing(dxV, lb, ub)
    T1 = eltype(dxV);
    n = length(dxV);
    valueV = zeros(T1, n);
    valueV[1] = ub - dxV[1] * (ub - lb);
    if n > 1
        for j = 2 : n
            valueV[j] = valueV[j-1] - dxV[j] * (valueV[j-1] - lb);
        end
    end
    return valueV
end

"""
	$(SIGNATURES)

Returns a subset of the values of a `BoundedVector`.
"""
values(iv :: BoundedVector{T1}, idx) where T1 =
    values(iv)[idx];

"""
	$(SIGNATURES)

Set the default values of a `BoundedVector`.
"""
function set_default_value!(iv :: BoundedVector{T1},
    valueV :: AbstractVector{T1}) where T1

    dxV = values_to_dx(iv, valueV);
    p = retrieve(iv.pvec, :dxV);
    set_default_value!(p, valueV)
end

function values_to_dx(iv :: Union{BVector{T1}, BoundedVector{T1}}, valueV :: AbstractVector{T1}) where T1
    @assert check_values(iv, valueV);
    n = length(iv);

    if is_increasing(iv)
        dxV = values_to_dx_increasing(valueV, scalar_lb(iv), scalar_ub(iv));
    elseif is_decreasing(iv)
        dxV = values_to_dx_decreasing(valueV, scalar_lb(iv), scalar_ub(iv));
    elseif is_nonmonotone(iv)
        dxV = copy(valueV);
    else
        error("Invalid");
    end
    return dxV
end

function values_to_dx_increasing(valueV :: AbstractVector{T1}, lb, ub) where T1
    n = length(valueV);
    dxV = zeros(T1, n);
    dxV[1] = (valueV[1] - lb) / (ub - lb);
    if n > 1
        for j = 2 : n
            dxV[j] = (valueV[j] - valueV[j-1]) / (ub - valueV[j-1]);
        end
    end
    return dxV
end

function values_to_dx_decreasing(valueV :: AbstractVector{T1}, lb, ub) where T1
    n = length(valueV);
    dxV = zeros(T1, n);
    dxV[1] = (ub - valueV[1]) / (ub - lb);
    if n > 1
        for j = 2 : n
            dxV[j] = (valueV[j-1] - valueV[j]) / (valueV[j-1] - lb);
        end
    end
    return dxV
end

function check_values(iv :: Union{BVector{T1}, BoundedVector{T1}}, valueV :: AbstractVector{T1}) where T1
    isValid = true;
    isValid = isValid  &&  isequal(length(valueV), length(iv));
    if any(valueV .< lb(iv))
        isValid = false;
        @warn "Values too low: $valueV";
    end
    if any(valueV .> ub(iv))
        isValid = false;
        @warn "Values too high: $valueV";
    end
    isValid = isValid  &&  all(valueV .<= ub(iv));
    if is_increasing(iv)
        isValid = isValid  &&  all(diff(valueV) .> 0.0);
    elseif is_decreasing(iv)
        isValid = isValid  &&  all(diff(valueV) .< 0.0);
    end
    return isValid
end

function param_table(iv :: BoundedVector{T1}, isCalibrated :: Bool) where T1
    p = retrieve(iv.pvec, :dxV);
    if isCalibrated == is_calibrated(p)
        pt = ParamTable(1);
        set_row!(pt, 1, string(p.name), p.symbol, p.description,
            formatted_value(values(iv)));
    else
        pt = nothing;
    end
    return pt
end

# ----------------


text:
function get_job(@nospecialize(func), @nospecialize(A), @nospecialize(types);
                 run_enzyme::Bool=true, mode::API.CDerivativeMode=API.DEM_ReverseModeCombined, dupClosure::Bool=false, kwargs...)

    primal, adjoint = fspec(func, types)

    rt = Core.Compiler.return_type(primal.f, primal.tt)
    rt = A{rt}

    target = Compiler.EnzymeTarget()
    params = Compiler.EnzymeCompilerParams(adjoint, mode, rt, run_enzyme, dupClosure)
    return Compiler.CompilerJob(target, primal, params)
end

function reflect(@nospecialize(func), @nospecialize(A), @nospecialize(types);
                 optimize::Bool=true, second_stage::Bool=true, kwargs...)

    job = get_job(func, A, types; kwargs...)
    # Codegen the primal function and all its dependency in one module
    mod, meta = Compiler.codegen(:llvm, job, optimize=optimize, #= validate=false =#)

    if second_stage
        post_optimze!(mod, JIT.get_tm())
    end

    llvmf = meta.adjointf

    return llvmf, mod
end

function enzyme_code_llvm(io::IO, @nospecialize(func), @nospecialize(A), @nospecialize(types);
                          optimize::Bool=true, run_enzyme::Bool=true, second_stage::Bool=true,
                          raw::Bool=false, debuginfo::Symbol=:default, dump_module::Bool=false)
    llvmf, _ = reflect(func, A, types; optimize,run_enzyme, second_stage)

    str = ccall(:jl_dump_function_ir, Ref{String},
                (LLVM.API.LLVMValueRef, Bool, Bool, Ptr{UInt8}),
                llvmf, !raw, dump_module, debuginfo)
    print(io, str)
end
enzyme_code_llvm(@nospecialize(func), @nospecialize(A), @nospecialize(types); kwargs...) = enzyme_code_llvm(stdout, func, A, types; kwargs...)

function enzyme_code_native(io::IO, @nospecialize(func), @nospecialize(A), @nospecialize(types))
    _, mod = reflect(func, A, types)
    str = String(LLVM.emit(JIT.get_tm(), mod, LLVM.API.LLVMAssemblyFile))
    print(io, str)
end
enzyme_code_native(@nospecialize(func), @nospecialize(A), @nospecialize(types); kwargs...) = enzyme_code_native(stdout, func, A, types; kwargs...)

function enzyme_code_typed(@nospecialize(func), @nospecialize(A), @nospecialize(types); kwargs...)
    job = get_job(func, A, types; kwargs...)
    GPUCompiler.code_typed(job; kwargs...)
end


text:
using Images, ImageView

# load an image and create a grayscale copy
img = load("sample-images/busan-night-scene.jpg");
img_gray = RGB.(Gray.(img))

# get channels representation
img_channel_view = channelview(img);
img_gray_channel_view = channelview(img_gray);

# make channel dimension last and crop the required are
img_arr = permuteddimsview(img_channel_view, (2, 3, 1));
img_gray_arr = permuteddimsview(img_gray_channel_view, (2, 3, 1));

# create a mask with all values being true
img_mask = fill(true, size(img));

# spot are with colors to retain
img_spot_height = 430:460
img_spot_width = 430:460

# preview color area (optional)
imshow(img[img_spot_height, img_spot_width])

for channel_id = 1:3

    # select current channel and crop ares of interest
    current_channel = view(img_arr, :, :, channel_id)
    current_channel_area = current_channel[img_spot_height, img_spot_width, :]

    # identify min and max values in a cropped area
    channel_min = minimum(current_channel_area)
    channel_max = maximum(current_channel_area)

    # merge existing mask with a channel specific mask
    channel_mask = channel_min .< current_channel .< channel_max
    img_mask = img_mask .& channel_mask
end

# apply mask
img_masked = img_arr .* img_mask .+ img_gray_arr .* .~(img_mask);
final_image = colorview(RGB, permutedims(img_masked, (3, 1, 2)))
imshow(final_image)