"
    _ftIdx(ft, grp)

Finds the index of grp (string) in ft.labels which corresponds to ft.mat.
For Index i for the grp  returned by this function ft[i, :, :] returns the
2D array [[TP, FP], [FN, TN]] for that group.
"
function _ftIdx(ft::FairTensor, grp)
    idx = findfirst(x->x==string.(grp), ft.labels)
    if idx==nothing throw(ArgumentError("$grp not found in the fairness tensor")) end
    return idx
end

# Helper function for calculating TruePositive, FalseNegative, etc.
# If grp is :, it calculates combined value for all groups, else the specified group
_calcmetric(ft::FairTensor, grp, inds...) = typeof(grp)==Colon ? sum(ft[:, inds...]) : sum(ft[_ftIdx(ft, grp), inds...])

(::TruePositive)(ft::FairTensor; grp=:) = _calcmetric(ft, grp, 1, 1)
(::FalsePositive)(ft::FairTensor; grp=:) = _calcmetric(ft, grp, 1, 2)
(::FalseNegative)(ft::FairTensor; grp=:) = _calcmetric(ft, grp, 2, 1)
(::TrueNegative)(ft::FairTensor; grp=:) = _calcmetric(ft, grp, 2, 2)

# The functions true_positive, false_negative, etc are instances of TruePositive, FalseNegative, etc.
# So on using true_positive(fair_tensor) will use the above defined functions
(::TPR)(ft::FairTensor; grp=:) = 1/(1+false_negative(ft; grp=grp)/(1e-15 + true_positive(ft; grp=grp)))
(::TNR)(ft::FairTensor; grp=:) = 1/(1+false_positive(ft; grp=grp)/(1e-15 + true_negative(ft; grp=grp)))
(::FPR)(ft::FairTensor; grp=:) = 1-true_negative_rate(ft; grp=grp)
(::FNR)(ft::FairTensor; grp=:) = 1-true_positive_rate(ft; grp=grp)

(::FDR)(ft::FairTensor; grp=:) = 1/(1+false_positive(ft; grp=grp)/(1e-15 + true_positive(ft; grp=grp)))
(::Precision)(ft::FairTensor; grp=:) = 1 - false_discovery_rate(ft; grp=grp)
(::NPV)(ft::FairTensor; grp=:) =  1/(1+false_negative(ft; grp=grp)/(1e-15 + true_negative(ft; grp=grp)))

(::MLJBase.Accuracy)(ft::FairTensor; grp=:) = (true_positive(ft; grp = grp) + true_negative(ft; grp = grp))/(
        true_positive(ft; grp = grp) + true_negative(ft; grp = grp) +
        false_positive(ft; grp = grp) + false_negative(ft; grp = grp))

struct PredictedPositiveRate <: Measure end
PP = PredictedPositiveRate
predicted_positive_rate = PP()
ppr = predicted_positive_rate
(::PP)(ft::FairTensor; grp=:) = _calcmetric(ft, grp, 1, :)/_calcmetric(ft, grp, :, :)

struct FalseOmissionRate <: Measure end
FOR = FalseOmissionRate
false_omission_rate = FOR()
foar = false_omission_rate
(::FOR)(ft::FairTensor; grp=:) = 1/(1+true_negative(ft; grp=grp)/(1e-15 + false_negative(ft; grp=grp)))

struct TruePositiveRateDifference <: Measure end
TPRD = TruePositiveRateDifference
true_positive_rate_difference = TPRD()
tprd = true_positive_rate_difference
(::TPRD)(ft::FairTensor; grp1=:,grp2=:)=abs(true_positive_rate(ft;grp=grp1)-true_positive_rate(ft;grp=grp2))

struct FalsePositiveRateDifference <: Measure end
FPRD = FalsePositiveRateDifference
false_positive_rate_difference = FPRD()
fprd = false_positive_rate_difference
(::FPRD)(ft::FairTensor; grp1=:,grp2=:)=abs(false_positive_rate(ft;grp=grp1)-false_positive_rate(ft;grp=grp2))

struct FalseNegativeRateDifference <: Measure end
FNRD = FalseNegativeRateDifference
false_negative_rate_difference = FNRD()
fnrd = false_negative_rate_difference
(::FNRD)(ft::FairTensor; grp1=:, grp2=:)=abs(false_negative_rate(ft;grp=grp1)-false_negative_rate(ft;grp=grp2))

struct FalseOmissionRateDifference <: Measure end
FORD = FalseOmissionRateDifference
false_omission_rate_difference = FORD()
ford = false_omission_rate_difference
(::FORD)(ft::FairTensor; grp1=:, grp2=:)=abs(false_omission_rate(ft;grp=grp1)-false_omission_rate(ft;grp=grp2))

struct FalseDiscoveryRateDifference <: Measure end
FDRD = FalseDiscoveryRateDifference
false_discovery_rate_difference = FDRD()
fdrd = false_discovery_rate_difference
(::FDRD)(ft::FairTensor; grp1=:, grp2=:)=abs(false_discovery_rate(ft;grp=grp1)-false_discovery_rate(ft;grp=grp2))

struct FalsePositiveRateRatio <: Measure end
FPRR = FalsePositiveRateRatio
false_positive_rate_ratio = FPRR()
fprr = false_positive_rate_ratio
(::FPRR)(ft::FairTensor; gpr1=:, grp2=:)=false_positve_rate(ft;grp=grp1)/false_positive_rate(ft;grp=grp2)

struct FalseNegativeRateRatio <: Measure end
FNRR = FalseNegativeRateRatio
false_negative_rate_ratio = FNRR()
fnrr = false_negative_rate_ratio
(::FNRR)(ft::FairTensor; gpr1=:, grp2=:)=false_negative_rate(ft;grp=grp1)/false_negative_rate(ft;grp=grp2)

struct FalseOmissionRateRatio <: Measure end
FORR = FalseOmissionRateRatio
false_omission_rate_ratio = FORR()
forr = false_omission_rate_ratio
(::FORR)(ft::FairTensor; grp1=:, grp2=:)=false_omission_rate(ft;grp=grp1)/false_omission_rate(ft;grp=grp2)

struct FalseDiscoveryRateRatio <: Measure end
FDRR = FalseDiscoveryRateRatio
false_discovery_rate_ratio = FDRR()
fdrr = false_discovery_rate_ratio
(::FDRR)(ft::FairTensor; grp1=:, grp2=:)=false_discovery_rate(ft;grp=grp1)/false_discovery_rate(ft;grp=grp2)

struct AverageOddsDifference <: Measure end
AOD = AverageOddsDifference
average_odds_difference = AOD()
aod = average_odds_difference
(::AOD)(ft::FairTensor; grp1=:, grp2=:)=0.5*(false_positive_rate_difference(ft;grp1=grp1,grp2=grp2)+true_positive_rate_difference(ft;grp1=grp1,grp2=grp2))


#=
    hilbert
    Copyright © 2015 jaidev <jaidev@newton>

    Distributed under terms of the MIT license.
=#

module hht

function hilbert(x)
    N = length(x)
    Xf = fft(x)
    h = zeros(N)
    h[1] = 1
    if N % 2 == 0
        h[floor(N / 2)] = 1
        h[1:floor(N / 2)] = 2
    else
        h[1:floor((N + 1) / 2)] = 2
    end
    x = ifft(Xf .* h)
    return x
end

end


using Revise
using Dice
using Dice: num_flips, num_nodes, to_dice_ir

code = @dice begin
    function discrete(p::Vector{Float64})
        mb = length(p)
        v = Vector(undef, mb)
        sum = 1
        for i=1:mb
            v[i] = p[i]/sum
            sum = sum - p[i]
        end

        # println(v)
        ans = DistInt(dicecontext(), mb-1)
        for i=mb-1:-1:1
            ans = if flip(v[i]) DistInt(dicecontext(), i-1) else ans end
        end
        return ans
    end

    function sendChar(key, observation)
        gen = discrete([0.5, 0.25, 0.125, 0.125])

        enc = key + gen
        d = prob_equals(observation, enc[1])
        return d
    end

    key = discrete([0.25, 0.25, 0.25, 0.25])
    d = true
    for i=1:4
        d&=sendChar(key, 2)
    end
    key, d
end

# BDD analysis
bdd = compile(code)
infer(code, :bdd)

# IR analysis
# to_dice_ir(code)
# has_dice_binary() && rundice(code)
# has_dice_binary() && infer(code, :ocaml)


using PIPG
import MathOptInterface
using LinearAlgebra
using StaticArrays
using SparseArrays
using Profile
const MOI = MathOptInterface
using Test

@testset "Linear algebra routines" begin
    for i=1:100
        res = MVector{10,Float64}(zeros(10))
        mat = sprand(10,10,0.5)
        inp = SVector{10, Float64}(rand(10))
        PIPG.spmul!(res, mat, inp, 1.0, 0.0)
        @test all(abs.(res .- mat*inp) .< 1e-10)
        PIPG.spmul!(res, mat, inp, 1.0, -1.0)
        @test all(abs.(res) .< 1e-10)
        PIPG.spmul!(res, mat, inp, -1.0, 0.0)
        @test all(abs.(res .+ mat*inp) .< 1e-10)
    end
end

@testset "Basic optimizer examples" begin
	γ = 0.9
	# QP example
    function make_qp()
    	P_ex = spdiagm([1/MathConstants.e, 2.0])
    	q = MVector(1/MathConstants.e,2.0)
    	H_ex = sparse([-1.0 -1.0])
    	g = MVector(0.0)
    	return PIPG.Problem(PIPG.NOCone{Float64, 1}(), PIPG.Reals{Float64, 2}(), H_ex, P_ex, q, g, 0.0)
    end

    # no scaling
    prob = make_qp()
    state = PIPG.State(prob; scaling=())
	PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(1.0, 1.0), SVector(0.0))
    result = state.primal
	@test norm(result .- [0.689, -0.689]) < 0.001

	# LP example
    make_prob() = PIPG.Problem(PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.POCone{Float64, 1}(), PIPG.POCone{Float64, 1}(),
        PIPG.POCone{Float64, 2}())), PIPG.Reals{Float64, 2}(),
        sparse([-100.0 -200.0; -10.0 -30.0; -1.0 -1.0; 1.0 0.0; 0.0 1.0]),
        spzeros(2,2),
        -1 .* MVector(50.0, 120.0),
        MVector(-10000.0, -1200.0, -110, -0.0, -0.0), 0.0)
	prob = make_prob()
	state = PIPG.State(prob; scaling=())
	PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
	@test norm(state.primal .- [60, 20]) < 0.01

    # geometric scaling
    prob = make_prob()
    state = PIPG.State(prob; scaling=(PIPG.GeoMean(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
    @test norm(state.primal .* state.col_scale .- [60, 20]) < 0.01

    prob = make_qp()
    state = PIPG.State(prob; scaling=(PIPG.GeoMean(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0))
    @test norm(state.primal .* state.col_scale .- [0.689, -0.689]) < 0.001

    # arithmetic scaling
    prob = make_prob()
    state = PIPG.State(prob; scaling=(PIPG.ArithMean(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
    @test norm(state.primal .* state.col_scale .- [60, 20]) < 0.01

    prob = make_qp()
    state = PIPG.State(prob; scaling=(PIPG.ArithMean(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.00001, SVector(0.0, 0.0), SVector(0.0))
    @test norm(state.primal .* state.col_scale .- [0.689, -0.689]) < 0.001

    # equilibration scaling
    prob = make_prob()
    state = PIPG.State(prob; scaling=(PIPG.Equilibration(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
    @test norm(state.primal .* state.col_scale .- [60, 20]) < 0.01

    prob = make_qp()
    state = PIPG.State(prob; scaling=(PIPG.Equilibration(prob), ))
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.00001, SVector(0.0, 0.0), SVector(0.0))
    @test norm(state.primal .* state.col_scale .- [0.689, -0.689]) < 0.001

    # default scaling (geomean + equlibration)
    prob = make_prob()
    state = PIPG.State(prob)
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.0001, SVector(0.0, 0.0), SVector(0.0, 0.0, 0.0, 0.0, 0.0))
    @test norm(state.primal .* state.col_scale .- [60, 20]) < 0.01

    prob = make_qp()
    state = PIPG.State(prob)
    PIPG.scale(prob, state)
    niters = PIPG.pipg(prob, state, 2000000, PIPG.compute_α(prob, γ), 0.00001, SVector(0.0, 0.0), SVector(0.0))
    @test norm(state.primal .* state.col_scale .- [0.689, -0.689]) < 0.001
end

#=
@testset "Profiling" begin
    make_prob() = PIPG.Problem(PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.POCone{Float64, 1}(), PIPG.POCone{Float64, 1}(),
        PIPG.POCone{Float64, 2}())), PIPG.Reals{Float64, 2}(),
        sparse([-100.0 -200.0; -10.0 -30.0; -1.0 -1.0; 1.0 0.0; 0.0 1.0]),
        spzeros(2,2),
        -1 .* MVector(50.0, 120.0),
        MVector(-10000.0, -1200.0, -110, -0.0, -0.0), 0.0)
    prob = make_prob()
    state = PIPG.State(prob)
    γ = 0.9
    PIPG.scale(prob, state)
    scale_a_bunch(prob,state,niters) = for i=1:niters PIPG.scale(prob, state) end
    scale_a_bunch(prob, state, 1)
    Profile.clear()
    @profile scale_a_bunch(prob, state, 1000000)
    @time scale_a_bunch(prob, state, 1000000)
    Profile.print(C=true)
end
=#

function basic_lp()
    opt = PIPG.Optimizer(ϵ=1e-7)
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
        MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([100.0, 200.0], x), 0.0),
        MOI.LessThan(10000.0))
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([10.0, 30.0], x), 0.0),
        MOI.LessThan(1200.0))
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], x), 0.0),
        MOI.LessThan(110.0))
    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))
    return opt, model, x
end

@testset "Basic LP" begin
    opt,model,x = basic_lp()
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [60, 20]) < 0.01
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 5400

    opt,model,x = basic_lp()
    opt.scaling = (p->(PIPG.GeoMean(p), ))
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [60, 20]) < 0.01
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 5400
end

@testset "Basic LP - VectorAffineFunction" begin
    opt = PIPG.Optimizer()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    o = MOI.set(model,
    	MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
    	MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    cv = MOI.add_constraint(model,
    	MOI.VectorAffineFunction(MOI.VectorAffineTerm.(
    	[1,1,2,2,3,3],
    	[MOI.ScalarAffineTerm.([100.0, 200.0], x);
    	 MOI.ScalarAffineTerm.([10.0, 30.0], x);
    	 MOI.ScalarAffineTerm.([1.0, 1.0], x)]), [-10000.0, -1200.0, -110.0]),
    	MOI.Nonpositives(3))
    pos = MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [60, 20]) < 0.01
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 5400
end

@testset "Basic SOCP" begin
    opt,model,x = basic_lp()
    opt.scaling = p->(PIPG.ArithMean(p),)
    tip = MOI.add_variable(model)
    MOI.add_constraint(model, tip, MOI.EqualTo(50.0))
    MOI.add_constraint(model, MOI.VectorOfVariables([tip, x...]),
        MOI.SecondOrderCone(3)) # |x|_2 < 50.0
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [42.881, 25.706]) < 0.01
end

@testset "Trivial QP" begin
    opt = PIPG.Optimizer()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )

	P_ex = spdiagm([1/MathConstants.e, 2.0])
	q = SVector(1/MathConstants.e,2.0)
	H_ex = sparse([-1.0 -1.0])
	g = SVector(0.0)

    x = MOI.add_variables(model, 2)
    MOI.set(model,
    	MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
    	MOI.ScalarQuadraticFunction(
    		MOI.ScalarQuadraticTerm{Float64}.([1/MathConstants.e, 2.0], x, x),
    		MOI.ScalarAffineTerm.([1/MathConstants.e,2.0], x), 0.0))
    MOI.add_constraint(model,
    	MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], x), 0.0),
    	MOI.GreaterThan(0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
	@test norm(MOI.get.(model, MOI.VariablePrimal(), x) .- [0.6893, -0.6893]) < 0.001
end

@testset "Basic QP" begin
    opt = PIPG.Optimizer()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )

    x = MOI.add_variables(model, 3)
    MOI.set(model,
    	MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
    	MOI.ScalarQuadraticFunction(
    		MOI.ScalarQuadraticTerm{Float64}.([5.0, -2, -1, 4, 3, 5],
    			[repeat([x[1]], 3); repeat([x[2]], 2); repeat([x[3]], 1)],
    			[x; x[2:3]; x[3]]),
    		MOI.ScalarAffineTerm.([2.0, -35, -47], x), 5.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
	@test norm(MOI.get.(model, MOI.VariablePrimal(), x) .- [3, 5, 7]) < 0.001
end

@testset "LP scalar 2" begin
    opt = PIPG.Optimizer(niters=100000000)
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
        MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([100.0, 200.0], x), 0.0),
        MOI.LessThan(10000.0))
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([10.0, 30.0], x), 0.0),
        MOI.LessThan(1200.0))
    sum_cstr = MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], x), 0.0),
        MOI.LessThan(60.0))
    lb_cstr = MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(0.3, x[1])], 0.0), MOI.GreaterThan(10.0))
    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))

    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [33.333, 26.666]) < 0.01
end
#=
prob = PIPG.Problem(
           PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.NOCone{Float64, 5}())),
           PIPG.Reals{Float64, 2}(),
sparse([   0.3    0
         100.0  200.0
          10.0   30.0
           1.0    1.0
          -1.0    0
           0     -1.0]), spzeros(2,2) , MVector(-50.0, -120.0), MVector(10.0, 10000.0, 1200.0, 60.0, -0.0, -0.0), 0.0)
state = PIPG.State(prob, PIPG.LogDiagnostics{Float64}())
α = PIPG.compute_α(prob, 0.9)
PIPG.pipg(prob, state, 100000000, α, 1e-5, SVector{2, Float64}(zeros(2)), SVector{6, Float64}(zeros(6)))
# baseline run (opt) 15270124

prob = PIPG.Problem(
           PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.NOCone{Float64, 3}())),
           PIPG.POCone{Float64, 2}(),
sparse([   0.3    0
         100.0  200.0
          10.0   30.0
           1.0    1.0]), spzeros(2,2) , MVector(-50.0, -120.0), MVector(10.0, 10000.0, 1200.0, 60.0), 0.0)
PIPG.State(prob, PIPG.LogDiagnostics{Float64}(), PIPG.GeoMean{Float64, 4, 2}())
α = PIPG.compute_α(prob, 0.9)
PIPG.pipg(prob, state, 100000000, α, 1e-5, SVector{2, Float64}(zeros(2)), SVector{4, Float64}(zeros(4)))

prob = PIPG.Problem(
           PIPG.PTCone{Float64}((PIPG.POCone{Float64, 1}(), PIPG.NOCone{Float64, 3}())),
           PIPG.POCone{Float64, 2}(),
sparse([ 0.3 0
         .1  .2
         .1   .3
         1.0    1.0]), spzeros(2,2) , MVector(-50.0, -120.0), MVector(10.0, 10.0, 12.0, 60.0), 0.0)
state = PIPG.State(prob, PIPG.LogDiagnostics{Float64}())
α = PIPG.compute_α(prob, 0.9)
PIPG.pipg(prob, state, 100000000, α, 1e-5, SVector{2, Float64}(zeros(2)), SVector{4, Float64}(zeros(4)))
=#

@testset "LP scalar modification" begin
    opt = PIPG.Optimizer()
    opt.niters *= 10
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    MOI.set(model,
        MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
        MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([100.0, 200.0], x), 0.0),
        MOI.LessThan(10000.0))
    MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([10.0, 30.0], x), 0.0),
        MOI.LessThan(1200.0))
    sum_cstr = MOI.add_constraint(model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, 1.0], x), 0.0),
        MOI.LessThan(110.0))
    lb_cstr = MOI.add_constraint(model,
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x[1])], 0.0), MOI.GreaterThan(0.0))
    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))
    MOI.optimize!(model)
    MOI.modify(model, sum_cstr, MOI.ScalarConstantChange(60.0))
    MOI.modify(model, lb_cstr, MOI.ScalarConstantChange(10.0))
    MOI.optimize!(model)
    change = MOI.ScalarCoefficientChange(x[1], 0.3)
    MOI.modify(model, lb_cstr, change)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [33.333, 26.666]) < 0.01
end

@testset "LP vector modification" begin
    opt = PIPG.Optimizer()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    c = [50.0, 120.0]
    x = MOI.add_variables(model, length(c))
    MOI.set(model,
    	MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}(),
    	MOI.ScalarQuadraticFunction(MOI.ScalarQuadraticTerm{Float64}[], MOI.ScalarAffineTerm.(c, x), 0.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    cv = MOI.add_constraint(model,
    	MOI.VectorAffineFunction(MOI.VectorAffineTerm.(
    	[1,1,2,2,3,3],
    	[MOI.ScalarAffineTerm.([100.0, 200.0], x);
    	 MOI.ScalarAffineTerm.([10.0, 30.0], x);
    	 MOI.ScalarAffineTerm.([1.0, 1.0], x)]), [-10000.0, -1200.0, -110.0]),
    	MOI.Nonpositives(3))
    MOI.add_constraint(model, MOI.VectorOfVariables(x), MOI.Nonnegatives(2))
    MOI.optimize!(model)
    cst = MOI.VectorConstantChange([-10000.0, -1200.0, -60.0])
    MOI.modify(model, cv, cst)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [30, 30]) < 0.01
    cst = MOI.VectorConstantChange([-10000.0, -600.0, -60.0])
    MOI.modify(model, cv, cst)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.DualStatus()) == MOI.FEASIBLE_POINT
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.OPTIMAL
    @test norm(MOI.get.(model, MOI.VariablePrimal(), x)  .- [60, 0]) < 0.01
end

@testset "Geometric mean cone clone" begin
    opt = PIPG.Optimizer()
    #opt.scaling = PIPG.NoScaling
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    n = 3
    t = MOI.add_variable(model)
    x = MOI.add_variables(model, n)
    vov = MOI.VectorOfVariables([t; x])
    gmc = MOI.add_constraint(model, vov, MOI.GeometricMeanCone(n + 1))
    c = MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(Float64(1), x), Float64(0)),
        MOI.LessThan(Float64(n)),
    )
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(Float64(1), t)], Float64(0)),
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 1
end

@testset "Geometric mean cone clone2" begin
    opt = PIPG.Optimizer()
    opt.scaling = x->(PIPG.Equilibration(x), )
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    T = Float64
    n = 3
    n = 9
    t = MOI.add_variable(model)
    x = MOI.add_variables(model, n)
    vov = MOI.VectorOfVariables([t; x])
    gmc = MOI.add_constraint(model, vov, MOI.GeometricMeanCone(n + 1))
    cx =
        Vector{MOI.ConstraintIndex{MOI.ScalarAffineFunction{T},MOI.EqualTo{T}}}(
            undef,
            n,
        )
    for i in 1:n
        cx[i] = MOI.add_constraint(
            model,
            MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(T(1), x[i])], T(0)),
            MOI.EqualTo(T(1)),
        )
    end
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
        MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(T(1), t)], T(0)),
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    MOI.optimize!(model)

    @test MOI.get(model, MOI.PrimalStatus()) == MOI.FEASIBLE_POINT
end

@testset "Infeasible SOC2" begin
    opt = PIPG.Optimizer()
    opt.scaling = x->(PIPG.Equilibration(x), )
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    T = Float64
    t = MOI.add_variable(model)
    x = MOI.add_variables(model, 2)
    MOI.add_constraints(model, x, MOI.GreaterThan.(T[3, 4]))
    MOI.set(model, MOI.ObjectiveFunction{MOI.VariableIndex}(), t)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.DUAL_INFEASIBLE
end
@testset "Infeasible SOC" begin
    opt = PIPG.Optimizer()
    opt.scaling = x->(PIPG.GeoMean(x), )
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    T = Float64
    t = MOI.add_variable(model)
    x = MOI.add_variables(model, 2)
    MOI.add_constraints(model, x, MOI.GreaterThan.(T[3, 4]))
    c_soc = MOI.add_constraint(
        model,
        MOI.VectorOfVariables([t; x]),
        MOI.SecondOrderCone(3),
    )
    MOI.set(model, MOI.ObjectiveFunction{MOI.VariableIndex}(), t)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(model)
    @test isapprox(MOI.get(model, MOI.VariablePrimal(), t), T(5))
    MOI.delete(model, c_soc)
    MOI.optimize!(model)
    println(opt.problem)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.DUAL_INFEASIBLE
end

@testset "Quadratic test copies" begin
    opt = PIPG.Optimizer()
    opt.scaling = x->()
    model = MOI.Bridges.full_bridge_optimizer(
        MOI.Utilities.CachingOptimizer(
            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
            opt,
        ),
        Float64,
    )
    T = Float64
    obj_attr = MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{T}}()
    x = MOI.add_variables(model, 2)
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    vc1 = MOI.add_constraint(model, x[1], MOI.GreaterThan(T(1)))
    @test vc1.value == x[1].value
    vc2 = MOI.add_constraint(model, x[2], MOI.GreaterThan(T(2)))
    @test vc2.value == x[2].value
    # Basic model
    # min x^2 + y^2 | x>=1, y>=2
    MOI.set(
        model,
        obj_attr,
        MOI.ScalarQuadraticFunction(
            MOI.ScalarQuadraticTerm.(T(2), x, x),  # quad
            MOI.ScalarAffineTerm{T}[],  # affine terms
            T(0),  # constant
        ),
    )
    MOI.optimize!(model)
    println(opt.problem.P)
    variable_primal = [(x[1], T(1)), (x[2], T(2))]
    for (index, solution_value) in variable_primal
        @test isapprox(
            MOI.get(model, MOI.VariablePrimal(), index),
            solution_value
        )
    end
    @test isapprox(
            MOI.get(model, MOI.ObjectiveValue()),
            T(5),
        )
    MOI.set(
        model,
        obj_attr,
        MOI.ScalarQuadraticFunction(
            MOI.ScalarQuadraticTerm.(
                T[2, 1//4, 1//4, 1//2, 2],
                [x[1], x[1], x[2], x[1], x[2]],
                [x[1], x[2], x[1], x[2], x[2]],
            ),  # quad
            MOI.ScalarAffineTerm{T}[],  # affine terms
            T(0),  # constant
        ),
    )
    MOI.optimize!(model)
    println(opt.problem.P)
    variable_primal = [(x[1], T(1)), (x[2], T(2))]
    for (index, solution_value) in variable_primal
        @test isapprox(
            MOI.get(model, MOI.VariablePrimal(), index),
            solution_value
        )
    end
    @test isapprox(
            MOI.get(model, MOI.ObjectiveValue()),
            T(7),
        )
end

    const OPTIMIZER = MOI.instantiate(
        MOI.OptimizerWithAttributes(PIPG.Optimizer, MOI.Silent() => true),
    )

    const BRIDGED = MOI.instantiate(
        MOI.OptimizerWithAttributes(PIPG.Optimizer, MOI.Silent() => true),
        with_bridge_type = Float64,
    )
    const CONFIG = MOI.Test.Config(
        atol=1e-5,
        rtol=1e-5,
        optimal_status=MOI.OPTIMAL,
        exclude=Any[MOI.VariableName, MOI.delete,
            # not an integer solver
            MOI.ObjectiveBound, MOI.RelativeGap,
            # not simplex
            MOI.ConstraintBasisStatus, MOI.SimplexIterations,
            # not barrier, nor branch-and-cut
            MOI.BarrierIterations, MOI.NodeCount])

    MOI.Test.runtests(
        BRIDGED,
        CONFIG,
        exclude = [
            "test_attribute_NumberOfThreads",
            "test_linear_Indicator_",
            "test_model_UpperBoundAlreadySet"
        ],
        # This argument is useful to prevent tests from failing on future
        # releases of MOI that add new tests. Don't let this number get too far
        # behind the current MOI release though! You should periodically check
        # for new tests in order to fix bugs and implement new features.
        exclude_tests_after = v"0.10.5",
    )


using Test, YaoBlocks, YaoArrayRegister

reg = rand_state(10)
block = kron(4, 2=>X)
c = concentrate(10, block, [1, 3, 9, 2])
@test nqubits(c) == 10
@test nactive(c) == 4
@test isunitary(c) == true
@test isreflexive(c) == true
@test ishermitian(c) == true
blk = kron(4, 2=>Rx(0.3))
@test chsubblocks(c, [blk]) |> subblocks |> first == blk

@test apply!(copy(reg), c) == apply!(copy(reg), kron(10, 3=>X))
@test apply!(rand_state(12, nbatch=10) |> focus!(Tuple(1:10)...), c) |> nactive == 10

@testset "test repeat" begin
    c = concentrate(8, repeat(5, H), 1:5)
    r = rand_state(8)
    r1 = copy(r) |> c
    r2 = copy(r) |> repeat(8, H, 1:5)
    @test r1 ≈ r2
end

@testset "mat" begin
    cc = concentrate(5, kron(X,X), (1,3))
    @test applymatrix(cc) ≈ mat(cc)
end


s = BenchmarkGroup()
suite["flow"] = s

edgs = [
  (1,2,10),(1,3,5),(1,4,15),(2,3,4),(2,5,9),
  (2,6,15),(3,4,4),(3,6,8),(4,7,16),(5,6,15),
  (5,8,10),(6,7,15),(6,8,10),(7,3,6),(7,8,10)
]

flow_graph = DiGraph(8)
capacity_matrix = zeros(Int,8,8)
for e in edgs
    u,v,f = e
    add_edge!(flow_graph,u,v)
    capacity_matrix[u,v] = f
end

s["push_relabel","$flow_graph"] = @benchmarkable maximum_flow($flow_graph, 1, 8
                    , $capacity_matrix, algorithm=PushRelabelAlgorithm())
s["dinic","$flow_graph"] = @benchmarkable maximum_flow($flow_graph, 1, 8
                    , $capacity_matrix, algorithm=DinicAlgorithm())

s["boykov","$flow_graph"] = @benchmarkable maximum_flow($flow_graph, 1, 8
                    , $capacity_matrix, algorithm=BoykovKolmogorovAlgorithm())

# from LittleScienceTools.RFIM
Random.seed!(17)
N = 200
g = random_regular_graph(N, 3, seed=1)
h = randn(N)
J = 1.

function net_capacity(g::AGraph, h::Vector{T}, J::AbstractFloat) where T
    N = nv(g)
    dg = digraph(g)
    add_vertices!(dg, 2)
    source = N+1
    target = N+2
    for i=1:N
        if h[i] > 0
            add_edge!(dg, source, i)
            add_edge!(dg, i, source)
        else
            add_edge!(dg, i, target)
            add_edge!(dg, target, i)
        end
    end

    c = spzeros(N+2, N+2)
    for i=1:N
        neigs = neighbors(dg, i)
        for j in neigs
            if j <= N
                c[i,j] = J
            elseif j == target
                c[i,j] = abs(h[i])
            end

        end
    end
    for j in neighbors(dg, source)
        c[source,j] = abs(h[j])
    end

    return dg, c
end

dg, c = net_capacity(g, h, J)

s["push_relabel","$dg"] = @benchmarkable maximum_flow($dg, $(N+1), $(N+2)
                    , $c, algorithm=PushRelabelAlgorithm())
s["dinic","$dg"] = @benchmarkable maximum_flow($dg, $(N+1), $(N+2)
                    , $c, algorithm=DinicAlgorithm())
s["boykov","$dg"] = @benchmarkable maximum_flow($dg, $(N+1),$(N+2)
                    , $c, algorithm=BoykovKolmogorovAlgorithm())


"""
Provides unofficial implementations (in the form of `ODESystem` and
`ODEFunction` instances) of a polynomial approximation for longitudinal
aircraft dynamics. Specifically, these models approximate NASA's
Generic Transport Model – a radio-controlled, sub-scale model aircraft which
is used flight control research. These publicly available equations were published by
[Chakraborty et al](https://www.sciencedirect.com/science/article/abs/pii/S0967066110002595).

# Extended help

## License

$(LICENSE)

## Exports

$(EXPORTS)

## Imports

$(IMPORTS)
"""
module PolynomialGTM

using Memoize
using Symbolics
using LinearAlgebra
using ModelingToolkit
using DocStringExtensions

@template (FUNCTIONS, METHODS, MACROS) =
    """
    $(SIGNATURES)

    $(DOCSTRING)
    """

@template (TYPES, CONSTANTS) =
    """
    $(TYPEDEF)

    $(DOCSTRING)
    """

export GTM, GTMFunction

"""
NASA's Generic Transport Model can be approximated
(near select flight conditions) as a polynomial model.
Here, `GTM` is a `ModelingToolkit.ODESystem` which
provides publicly available polynomial approximations
for GTM longitudinal flight dynamics.

# Extended Help

## Initial Conditions

The default initial conditions are one trim condition.
Two trim conditions for these polynomial-approximated
dynamics are shown below.

```julia
trim₁ = [[29.6, deg2rad(9), 0.0, deg2rad(9)],   [deg2rad(0.68), 12.7]]
trim₂ = [[25.0, deg2rad(18), 0.0, deg2rad(18)], [deg2rad(-7.2), 59]]
```

## Usage

```julia
model = GTM()
```
## References:

- [Chakraborty et al](https://www.sciencedirect.com/science/article/abs/pii/S0967066110002595)
- [Joe Carpinelli](https://github.com/cadojo/Replicated-ROA-Analysis)
"""
@memoize function GTM(; stm=false, structural_simplify=true, name=:GTM)

    # First, let's define the necessary parameters
    # and variables for modeling the longitudinal
    # dynamics for NASA's Generic Transport Model

    @parameters t δₑ δₜ
    @variables V(t) α(t) q(t) θ(t)
    δ = Differential(t)
    x = [V, α, q, θ]
    p = [δₑ, δₜ]

    # Great! Now, we'll need to hard-code in the polynomial
    # approximations the longitudinal dynamics.

    eqs = [

        # Damn Unicode and its variable character font width!!!

        δ(V) ~  1.233e-8*V^4*q^2 + 4.853e-9*α^3*δₜ^3 	  	+
                3.705e-5*V^3*α*q - 2.184e-6*V^3*q^2 		+
                2.203e-2*V^2*α^3 - 2.836e-6*α^3*δₜ^2 	  	+
                3.885e-7*α^2*δₜ^3 - 1.069e-6*V^3*q 	  		-
                4.517e-2*V^2*α^2 - 2.140e-3*V^2*α*δₑ 		-
                3.282e-3*V^2*α*q - 8.901e-4*V^2*δₑ^2 		+
                9.677e-5*V^2*q^2 - 2.037e-4*α^3*δₜ 	  	 	 -
                2.270e-4*α^2*δₜ^2 - 2.912e-8*α*δₜ^3 	     +
                1.591e-3*V^2*α - 4.077e-4*V^2*δₑ 		  	+
                9.475e-5*V^2*q - 1.637*α^3 				  	-
                1.631e-2*α^2*δₜ + 4.903*α^2*θ 		  		-
                4.903*α*θ^2 + 1.702e-5*α*δₜ^2 		  		-
                7.771e-7*δₜ^3 + 1.634*θ^3 - 4.319e-4*V^2    -
                2.142e-1*α^2 + 1.222e-3*α*δₜ 			  	+
                4.541e-4*δₜ^2 + 9.823*α + 3.261e-2*δₜ 	    -
                9.807*θ + 4.282e-1,
        δ(α) ~ -3.709e-11*V^5*q^2 + 6.869e-11*V*α^3*δₜ^3 	+
                7.957e-10*V^4*α*q + 9.860e-9*V^4*q^2 		+
                1.694e-5*V^3*α^3 - 4.015e-8*V*α^3*δₜ^2 		 -
                7.722e-12*V*α^2*δₜ^3 - 6.086e-9*α^3*δₜ^3 	 -
                2.013e-8*V^4*q - 5.180e-5*V^3*α^2 			-
                2.720e-6*V^3*α*δₑ - 1.410e-7*V^3*α*q 		+
                7.352e-7*V^3*δₑ^2 - 8.736e-7*V^3*q^2 		-
                1.501e-3*V^2*α^3 - 2.883e-6*V*α^3*δₜ 		+
                4.513e-9*V*α^2*δₜ^2 - 4.121e-10*V*α*δₜ^3 	+
                3.557e-6*α^3*δₜ^2 + 6.841e-10*α^2*δₜ^3 		+
                4.151e-5*V^3*α + 3.648e-6*V^3*δₑ 		    +
                3.566e-6*V^3*q + 6.246e-6*V^2*α*q 			+
                4.589e-3*V^2*α^2 + 2.410e-74*V^2*α*δₑ 		-
                6.514e-5*V^2*δₑ^2 + 2.580e-5*V^2*q^2 		-
                3.787e-5*V*α^3 + 3.241e-7*V*α^2*δₜ 			+
                2.409e-7*V*α*δₜ^2 + 1.544e-11*V*δₜ^3 		+
                2.554e-4*α^3*δₜ - 3.998e-7*α^2*δₜ^2 		+
                3.651e-8*α*δₜ^3 + 4.716e-7*V^3 				-
                3.677e-3*V^2*α - 3.231e-4*V^2*δₑ 			-
                1.579e-4*V^2*q + 2.605e-3*V*α^2 			+
                1.730e-5*V*α*δₜ - 5.201e-3*V*α*θ 			-
                9.026e-9*V*δₜ^2 + 2.601e-3*V*θ^2 			+
                3.355e-3*α^3 - 2.872e-5*α^2*δₜ 				-
                2.134e-5*α*δₜ^2 - 1.368e-9*δₜ^3 		    -
                4.178e-5*V^2 + 2.272e-4*V*α 				-
                6.483e-7*V*δₜ - 2.308e-1*α^2 				-
                1.532e-3*α*δₜ + 4.608e-1*α*θ - 2.304e-1*θ^2 +
                7.997e-7*δₜ^2 - 5.210e-3*V - 2.013e-2*α 	+
                5.744e-5*δₜ + q + 4.616e-1,
        δ(q) ~ -6.573e-9*V^5*q^3 + 1.747e-6*V^4*q^3 	    -
                1.548e-4*V^3*q^3 - 3.569e-3*V^2*α^3 	    +
                4.571e-3*V^2*q^3 + 4.953e-5*V^3*q 		    +
                9.596e-3*V^2*α^2 + 2.049e-2*V^2*α*δₑ 	    -
                2.431e-2*V^2*α - 3.063e-2*V^2*δₑ 		    -
                4.388e-3*V^2*q - 2.594e-7*δₜ^3 			    +
                2.461e-3*V^2 + 1.516e-4*δₜ^2 + 1.089e-2*δₜ 	 +
                1.430e-1,
        δ(θ) ~ q
    ]

    # Let's set some default values, so we can plug this system directly
    # into `ODEProblem` and perform quick analysis. Note that what we're
    # doing here is setting one equilibrium position as a default
    # initial condition if no arguments are provided to `ODEProblem`!
    defaults = Dict(
        V  => 29.6,
        α  => deg2rad(9),
        q  => 0.0,
        θ  => deg2rad(0),
        δₑ => deg2rad(0.68),
        δₜ => 12.7
    )

    # If state transition matrix dynamics are enabled, append
    # the dynamics to our equations of motion, append the
    # state variables to our state vector, and append default
    # values (the identity matrix) to the defaults field!
    if stm
        @variables Φ[1:4,1:4](t)
        Φ = Symbolics.scalarize(Φ)
        A = Symbolics.jacobian(map(el -> el.rhs, eqs), x)

        LHS = map(δ, Φ)
        RHS = map(simplify, A * Φ)

        eqs = vcat(eqs, [LHS[i] ~ RHS[i] for i in 1:length(LHS)])

        for ϕ in vec(Φ)
            push!(x, ϕ)
        end

        for (ϕ, i) in zip(vec(Φ), vec(Matrix(I(4))))
            defaults[ϕ] = i
        end
    end

    # Model name
    if string(name) == "GTM" && stm
        modelname = Symbol("GTMWithSTM")
    else
        modelname = name
    end

    # Make the model and return
    model = ODESystem(eqs, t, x, p; defaults=defaults, name=Symbol(modelname))
    return structural_simplify ? ModelingToolkit.structural_simplify(model) : model
end

"""
Returns a `DifferentialEquations`-compatible `ODEFunction` for GTM dynamics.
The `stm`, `structural_simplify`, and `name` keyword arguments are passed
to `GTM`. All other keyword arguments are passed directly to `ODEFunction`.

# Extended Help

## Usage

Note that this `ODEFunction` output has several methods, including an in-place
method! Function signatures follow `ModelingToolkit` and `DifferentialEquations`
conventions.

```julia
f = GTMFunction()
let u = randn(4), p = randn(2), t = rand()
    f(u,p,t)
end
```
"""
@memoize function GTMFunction(; stm=false, structural_simplify=true, name=:GTM, kwargs...)
    defaults = (; jac = true)
    options  = merge(defaults, kwargs)

    model = GTM(; stm=stm, structural_simplify=structural_simplify, name=name)
    return ODEFunction(model; options...)
end

end # module


import RData
using RNGPool
using CoupledConditionalSMC

# data generated with A = 0.95 (not 0.9)
dataJLS = RData.load("demo/ar1data.RData")
ys = vec(dataJLS["observations"])

include("../demo/lgModel.jl")

# they start at time 0 with observations from time 1,2,...
# so we have an initial distribution of N(0,1+0.9^2)
theta = LGTheta(0.9, 1.0, 1.0, 1.0, 0.0, 1.81)
model = makeLGModel(theta, ys)
lM = LinearGaussian.makelM(theta)
ko = kalman(theta, ys)

function checkBoundary(ref1::Vector{Float64Particle}, ref2::Vector{Float64Particle}) :: Int64
    b = 0
    for i in 1:length(ref1)
        (ref1[i].x != ref2[i].x) && break
        b = i
    end
    b
end

function couplingBoundary(model, lM, algorithm=:AS, N=2000, n=500, maxit=500,
    rngCouple = true, independentInitialisation = true)
    io = CCSMCIO{model.particle, model.pScratch}(N, n)
    initializeCCSMC(model, lM, io, algorithm, independentInitialisation)
    bd = fill(n, maxit)
    for i in 1:maxit
        ccXpf!(model, lM, io, algorithm, rngCouple)
        bd_ = checkBoundary(io.ref1, io.ref2)
        bd[i] = bd_
        (bd_ == n) && break
    end
    bd
end

setRNGs(12345)
N = 64    # Number of particles
n = 800     # Length of time series
maxit = 500 # Maximum number of iterations
bdAT = couplingBoundary(model, lM, :AT, N, n, maxit)
bdAS = couplingBoundary(model, lM, :AS, N, n, maxit)
bdBS = couplingBoundary(model, lM, :BS, N, n, maxit)

using GR
figure(figsize=(15,5))
beginprint("paper/output/lgDemoPaths.pdf")
GR.legend("BS", "AS", "ATT"; loc=2)
GR.xlabel("Iteration")
GR.ylabel("Coupling boundary")
GR.ylim([-30, n+30])
GR.xlim([0, maxit])
GR.plot(1:maxit, bdBS, "k-.", 1:maxit, bdAS, "k--", 1:maxit, bdAT, "k-")
endprint()


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ 0edbe34a-3714-11eb-1350-af72689d8855
md" ## _Data Structures_
In this lesson we will study how data can be collected and stored in memory. In particular we will deal with vectors, matrices, n-dimensional arrays, tuples and dictionaries.
"

# ╔═╡ 4ef2e454-3716-11eb-3014-774db1fb57d0
md" ## Vectors
A vector is a list of ordered data which share a common type (be it `Int`, `Float` or `Any`). Furthermore a vector is a one-dimensional array, and often “vector” and “array” are used a synonyms.

Contrarily to the mathematical definition of a vector, in programming a vector is simply a list of values and has no a priori geometrical meaning.

In Julia, to create a vector we use the following syntax:
"

# ╔═╡ 9833f658-3716-11eb-300e-d3b5d02b70e5
a = [1,2,3,4,5]

# ╔═╡ b2cbd7d8-3716-11eb-15cc-03c3fec2c981
b = [1.2, 3,4,5]

# ╔═╡ b2d549da-3716-11eb-0cce-b703452d7f34
c = ["Hello", "it's me", "TechyTok"]

# ╔═╡ bbe66bfa-3716-11eb-1929-ef53e271e354
md" We can access the members of an array using the indexing syntax: `array_name[element_number]`, for example, if we want to retrieve the third element of c we type `c[3]`. Contrarily to other programming languages, in **Julia vectors start at 1**, there is not much to say, it is just like that.
"

# ╔═╡ d5c03624-3716-11eb-38ed-3f395efc01f6
c[2], a[3], b[1]

# ╔═╡ e7768e7e-3716-11eb-366a-2f81166da844
md"If you want to append an element to an array, we can do it by using the `append!` function. Notice the !, this is a Julia convention to say that the function will modify the first argument given to the function, in fact `append!(a, 6)` results in `a = [1, 2, 3, 4, 5, 6]`.

Previously I have said that vector elements must share the same type, if we try to append (or add) a value to an array with a different type it will result into an error:
"

# ╔═╡ 5268b19c-3717-11eb-3854-ddd093a61057
append!(a, 3.14)

# ╔═╡ 6ced58c6-3717-11eb-1009-7909ebb6866b
md"In this case, we cannot convert a Float64 to an Int64 without losing precision, thus the error.

To check the type of an array we can type:
"

# ╔═╡ f2cb5918-3717-11eb-2b4b-77642fbeaf20
typeof(a)

# ╔═╡ ff24fcbc-3717-11eb-2ac3-6f18fde9484f
md"As we can see, `a` can only store `Int64` values or values that can be safely converted to `Int64` (such as `Int32` for example).

An alternative way to define a Vector, if we want to specify the type of the elements included, is by using a type name followed by a square bracket and the desired elements:
"

# ╔═╡ 21342288-3718-11eb-1b86-8d438637ac3a
d = Int[1,2,3,4,5]

# ╔═╡ 3fb54d2c-3718-11eb-305f-e7159394ac7b
md" ## Matrices
Matrices are two dimensional arrays. We can define a matrix with the following syntax:
"

# ╔═╡ 4c2a11c8-3718-11eb-0f59-8382e767058f
mat1 = [1 2 3; 4 5 6]

# ╔═╡ 538463b0-3718-11eb-03e7-171beded2a41
md"Rows are separated by semi colons ; and columns are separated by spaces.

When we want to access the elements of a matrix, we use the following notation `mat1[row_index, column_index]`. If we want to access the 2nd element of the first row we shall write `mat1[1, 2]`.
"

# ╔═╡ 6ed4e510-3718-11eb-3d00-9f60abc37333
md" ## N-Dimensional Arrays
Sometimes we need to create tables with more than 2 dimensions. In this case usually the tables tend to be big, so there is no explicit way to create an n-dimensional array. The suggested practice is to create an empty array first, using a new function called `zeros`, and then fill it either manually of using a loop.

For example, let’s suppose we want to create a 2x3x4 table, we would do it with `zeros(2,3,4)`. Let’s suppose we want to fill it with the product of the indexes, we can do it in the following way:

"

# ╔═╡ 9a27ee68-3718-11eb-1c7f-33eaf6e7cb3f
table = zeros(2,3,4)

# ╔═╡ ab83ddde-3718-11eb-0847-9fdeaca6d470
for k in 1:4
    for j in 1:3
        for i in 1:2
            table[i,j,k] = i*j*k
        end
    end
end

# ╔═╡ bbbc6720-3718-11eb-2c1e-978871157853
table

# ╔═╡ d3ce8154-3718-11eb-08ff-7b29208c5099
md" Please not that Julia stores values in memory differently from Python: in Julia to obtain fast loops we need to iter first over columns (which means that the first index must vary first and so on). For this reason if we plan to store, for example, 42 2x2 matrices, we need to create an array of size 2x2x42 (while in Python we would have created a 42x2x2 table).
"

# ╔═╡ df969238-3718-11eb-2196-a7a68b5a69f2
md" ## Slices
There is a convenient notation (called **slicing**) to access a subset of elements from an array. Let’s suppose we want to access the 2nd to 5th elements of an array of length 6, we can do it in the following way:

"

# ╔═╡ f82b9558-3718-11eb-0ef8-5fa509ba5b6f
begin
	e = [1,2,3,4,5,6]
	f = e[2:5]
end

# ╔═╡ 6ffa0c1a-3719-11eb-34d2-d7fc89a62e24
md"We can also use this notation to access a subset of a matrix, for example:
"

# ╔═╡ 7edb458c-3719-11eb-2bb3-4bd9173b530c
mat3 = reshape([i for i in 1:16],4,4)

# ╔═╡ 8bae5826-3719-11eb-084e-fbf852773ffd
mat4 = mat3[2:3, 2:3]

# ╔═╡ a17c0086-3719-11eb-1ab0-2381decb5422
md"On line 1 we have used a handy notation called **list comprehension**. `[i for i in 1:16]` means “create an array containing each i comprised from 1 to 16”. We then reshape it to have a size of 4x4 and store the result in `mat1`.

It is also possible to have nested list comprehensions like this one:
"

# ╔═╡ c62a2930-3719-11eb-3c28-750300385ad8
[i+j for i in 1:10 for j in 1:5]

# ╔═╡ e07fe63a-3719-11eb-05eb-81512953fb4b
md"This will create an array of length 50 containing all the sums of i+j.
"

# ╔═╡ 2655af78-371a-11eb-331c-310db7a9ccb1
md" ## Views
As in other programming languages, arrays are pointers to location in memory, thus we need to pay attention when we handle them.

If we create an array `g` and we assign a to `h` with `h=g`, the elements of `g` be modified by accessing `h`:
"

# ╔═╡ 3cdd35ea-371a-11eb-0f7f-135c3a766ecd
g=[1,2,3]

# ╔═╡ 67a3d7f2-371a-11eb-07cb-e3e233e2aab1
h=g

# ╔═╡ 67ad9e04-371a-11eb-36ad-9710ff0b0db3
g[2] = 42

# ╔═╡ 67bdc450-371a-11eb-08b2-55230f42c0e8
g

# ╔═╡ b2873584-371a-11eb-328d-37a39540905c
md"
This is particularly useful because it lets us save memory, but may have undesirable effects. If we want to make a copy of an array we need to use the function `copy`.
"

# ╔═╡ c17d2258-371a-11eb-25b5-933f413cd134
i=[1,2,3]

# ╔═╡ d902ab08-371a-11eb-31ce-8d50b3388b2b
j=copy(i)

# ╔═╡ d90d570e-371a-11eb-0e4b-075841a6a73a
j[2] = 42

# ╔═╡ dbfd0a06-371a-11eb-1bf5-e9112227d8b0
i

# ╔═╡ de4b38a0-371a-11eb-39b1-7575188e8063
j

# ╔═╡ ee1a296c-371a-11eb-0036-3b8ad05c6804
md"In some cases, when there are arrays containing other arrays, if we want to make a full copy of all the contents we need to use `deepcopy` instead."

# ╔═╡ f6fa2e36-371a-11eb-36e7-a7de898b78c5
md" ## Tuples
A tuple is a fixed size group of variables which may share a common type but don’t need to.

Unlike arrays, you cannot increase the size of a tuple once it has been created. Tuples are created using the following syntax:
"

# ╔═╡ 0e12b2be-371b-11eb-0e56-99ec2e4a224d
k =(1,2,3)

# ╔═╡ 187b73f0-371b-11eb-34e6-d7a2aeb00b58
l= 1,2,3

# ╔═╡ 1eb587ec-371b-11eb-2901-a1ebf0197aa7
md"
So tuples can be created by using regular brackets or no brackets at all! Tuples are really handy, as it is possible to “unpack” a tuple over many values:
"

# ╔═╡ 2bc6ba1e-371b-11eb-12d0-a3665fe9e1d0
tuple1 = (1, 2, 3)

# ╔═╡ 3985145c-371b-11eb-2dcc-5308e23dc5f0
m, n, o = tuple1

# ╔═╡ 39a4874c-371b-11eb-164c-692fd38680a9
("$m $n $o")

# ╔═╡ 58abede2-371b-11eb-24d3-2b16686937d5
md"
It is also possible to use tuples to emulate multiple return values from functions:
"

# ╔═╡ 6205bad0-371b-11eb-1cfb-3f1ff43c255e
function return_multiple()
    return 42, 43, 44
end

# ╔═╡ 64cd1682-371b-11eb-2097-a9267c1e63dc
p, q, r = return_multiple()

# ╔═╡ 75e813a2-371b-11eb-21d8-3d2bf2e80cac
("$p $r $q")

# ╔═╡ 827ebb9a-371b-11eb-006b-6107854f0a9c
md" ## Splatting
It is possible to “unpack” a tuple and pass its arguments to a function with the following syntax:
"

# ╔═╡ 04ea7b0a-371c-11eb-2827-7d192535d804
function splat_me(s, t, u)
    return s*t*u
end

# ╔═╡ 0e3a8240-371c-11eb-1720-7f7de97d90aa
tuple2 = (1,2,3)

# ╔═╡ 266b7540-371c-11eb-06f2-553a12a3314b
splat_me(tuple2...)

# ╔═╡ 4417da34-371c-11eb-155e-77b5eacbfc88
md" So the `...` after a tuple will unpack it! This is useful but addictive, use it only if needed as it is better for clarity (and to avoid multiple dispatch errors) to call a function with its single parameters.
"

# ╔═╡ 55d5c8a8-371c-11eb-00cf-91fa36d66a01
md" ## Named Tuples
Named tuples are like tuples but with a name identifier for a single value, for example:
"

# ╔═╡ 68bde3ec-371c-11eb-0aa7-8161a51fdbf5
namedTuple1 = (a = 1, b = "hello")

# ╔═╡ 6ba908ca-371c-11eb-1126-cfcbeb01cb0d
namedTuple1[:a]

# ╔═╡ 9b678d86-371c-11eb-0b34-2917771de661
md" or in alternative:
"

# ╔═╡ 83a9ac22-371c-11eb-09d9-c75f570c9b31
namedTuple2 = NamedTuple{(:a, :b)}((2,"hello2"))

# ╔═╡ 89ade9bc-371c-11eb-21aa-8f0831a9fff5
namedTuple2[:b]

# ╔═╡ 970399ea-371c-11eb-0cee-b54e5f3f5498
md"## Dictionaries
A dictionary is a collection of keys and values. They are unordered (which means that the order of the keys is random) and are really useful when you need to organise, for example, a dataset.

Let’s suppose we want to create an address book. A single entry should be able to store all the fundamental characteristics needed to identify a friend: the name of the contact, the phone number and the shoe size!
"

# ╔═╡ 706e79fc-371d-11eb-0894-ebb085103956
person1 = Dict("Name" => "Peppa", "Phone" => 123456789, "Shoe-size" => 40)

# ╔═╡ 7f46d19a-371d-11eb-2971-ef0eda3d5d18
person2 = Dict("Name" => "George", "Phone" => 123456789, "Shoe-size" => 36)

# ╔═╡ 9648198a-371d-11eb-3c9f-9338602af4af
md"I can even make a dictionary containing other dictionaries:
"

# ╔═╡ a25da512-371d-11eb-33ea-0f0740b919e3
addressBook = Dict("Peppa" => person1, "George" => person2)

# ╔═╡ ee8143ec-371d-11eb-3c9b-754027e67304
md"
We can add another friend to the addressBook , once it has been created, with the following syntax:
"

# ╔═╡ f94e1d2c-371d-11eb-0e5c-819b938a8d00
person3 = Dict("Name" => "Vittorio", "Phone" => 123456789, "Shoe-size" => 42)

# ╔═╡ 0560f292-371e-11eb-07de-c55444c0a3d5
addressBook["Vittorio"] = person3

# ╔═╡ 0ac0ae12-371e-11eb-3d9a-a7ef8aa4f5d9
addressBook

# ╔═╡ 19e817c2-371e-11eb-15ce-5bee1fd2f891
md" ## Conclusion
In this lesson we have learned how to operate on arrays, tuples and dictionaries. Those structures are the basics of in-memory storage for any data. Arrays are lightweight and useful solutions to pass blocks of data, so use them whenever needed! Contrarily to C++, Julia has a built in garbage collector, so you don’t have to care about freeing memory and deleting pointers, as Julia will take care of it!
"

# ╔═╡ Cell order:
# ╟─0edbe34a-3714-11eb-1350-af72689d8855
# ╟─4ef2e454-3716-11eb-3014-774db1fb57d0
# ╠═9833f658-3716-11eb-300e-d3b5d02b70e5
# ╠═b2cbd7d8-3716-11eb-15cc-03c3fec2c981
# ╠═b2d549da-3716-11eb-0cce-b703452d7f34
# ╟─bbe66bfa-3716-11eb-1929-ef53e271e354
# ╠═d5c03624-3716-11eb-38ed-3f395efc01f6
# ╟─e7768e7e-3716-11eb-366a-2f81166da844
# ╠═5268b19c-3717-11eb-3854-ddd093a61057
# ╟─6ced58c6-3717-11eb-1009-7909ebb6866b
# ╠═f2cb5918-3717-11eb-2b4b-77642fbeaf20
# ╟─ff24fcbc-3717-11eb-2ac3-6f18fde9484f
# ╠═21342288-3718-11eb-1b86-8d438637ac3a
# ╟─3fb54d2c-3718-11eb-305f-e7159394ac7b
# ╠═4c2a11c8-3718-11eb-0f59-8382e767058f
# ╟─538463b0-3718-11eb-03e7-171beded2a41
# ╟─6ed4e510-3718-11eb-3d00-9f60abc37333
# ╠═9a27ee68-3718-11eb-1c7f-33eaf6e7cb3f
# ╠═ab83ddde-3718-11eb-0847-9fdeaca6d470
# ╠═bbbc6720-3718-11eb-2c1e-978871157853
# ╟─d3ce8154-3718-11eb-08ff-7b29208c5099
# ╟─df969238-3718-11eb-2196-a7a68b5a69f2
# ╠═f82b9558-3718-11eb-0ef8-5fa509ba5b6f
# ╟─6ffa0c1a-3719-11eb-34d2-d7fc89a62e24
# ╠═7edb458c-3719-11eb-2bb3-4bd9173b530c
# ╠═8bae5826-3719-11eb-084e-fbf852773ffd
# ╟─a17c0086-3719-11eb-1ab0-2381decb5422
# ╠═c62a2930-3719-11eb-3c28-750300385ad8
# ╟─e07fe63a-3719-11eb-05eb-81512953fb4b
# ╟─2655af78-371a-11eb-331c-310db7a9ccb1
# ╠═3cdd35ea-371a-11eb-0f7f-135c3a766ecd
# ╠═67a3d7f2-371a-11eb-07cb-e3e233e2aab1
# ╠═67ad9e04-371a-11eb-36ad-9710ff0b0db3
# ╠═67bdc450-371a-11eb-08b2-55230f42c0e8
# ╟─b2873584-371a-11eb-328d-37a39540905c
# ╠═c17d2258-371a-11eb-25b5-933f413cd134
# ╠═d902ab08-371a-11eb-31ce-8d50b3388b2b
# ╠═d90d570e-371a-11eb-0e4b-075841a6a73a
# ╠═dbfd0a06-371a-11eb-1bf5-e9112227d8b0
# ╠═de4b38a0-371a-11eb-39b1-7575188e8063
# ╟─ee1a296c-371a-11eb-0036-3b8ad05c6804
# ╟─f6fa2e36-371a-11eb-36e7-a7de898b78c5
# ╠═0e12b2be-371b-11eb-0e56-99ec2e4a224d
# ╠═187b73f0-371b-11eb-34e6-d7a2aeb00b58
# ╟─1eb587ec-371b-11eb-2901-a1ebf0197aa7
# ╠═2bc6ba1e-371b-11eb-12d0-a3665fe9e1d0
# ╠═3985145c-371b-11eb-2dcc-5308e23dc5f0
# ╠═39a4874c-371b-11eb-164c-692fd38680a9
# ╟─58abede2-371b-11eb-24d3-2b16686937d5
# ╠═6205bad0-371b-11eb-1cfb-3f1ff43c255e
# ╠═64cd1682-371b-11eb-2097-a9267c1e63dc
# ╠═75e813a2-371b-11eb-21d8-3d2bf2e80cac
# ╟─827ebb9a-371b-11eb-006b-6107854f0a9c
# ╠═04ea7b0a-371c-11eb-2827-7d192535d804
# ╠═0e3a8240-371c-11eb-1720-7f7de97d90aa
# ╠═266b7540-371c-11eb-06f2-553a12a3314b
# ╟─4417da34-371c-11eb-155e-77b5eacbfc88
# ╟─55d5c8a8-371c-11eb-00cf-91fa36d66a01
# ╠═68bde3ec-371c-11eb-0aa7-8161a51fdbf5
# ╠═6ba908ca-371c-11eb-1126-cfcbeb01cb0d
# ╠═9b678d86-371c-11eb-0b34-2917771de661
# ╠═83a9ac22-371c-11eb-09d9-c75f570c9b31
# ╠═89ade9bc-371c-11eb-21aa-8f0831a9fff5
# ╟─970399ea-371c-11eb-0cee-b54e5f3f5498
# ╠═706e79fc-371d-11eb-0894-ebb085103956
# ╠═7f46d19a-371d-11eb-2971-ef0eda3d5d18
# ╟─9648198a-371d-11eb-3c9f-9338602af4af
# ╠═a25da512-371d-11eb-33ea-0f0740b919e3
# ╠═ee8143ec-371d-11eb-3c9b-754027e67304
# ╠═f94e1d2c-371d-11eb-0e5c-819b938a8d00
# ╠═0560f292-371e-11eb-07de-c55444c0a3d5
# ╠═0ac0ae12-371e-11eb-3d9a-a7ef8aa4f5d9
# ╟─19e817c2-371e-11eb-15ce-5bee1fd2f891


# ---------------------------------------------------------------------------- #
#
#   basisFuncLineLag.jl
#
#   Get Lagrangian basis functions for Line
#
#   λυτέος
#   Fall 2017
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

function basisFuncLineLag( P::P1, s::Vector{Float64} )

  nx  = length( s )

  # Linear
  nphi = 2
  phi  = Array{Float64}(nphi, nx)
  dphi = Array{Float64}(nphi, nx)

  t = 1 - s

  phi[1,:] = t
  phi[2,:] = s

  dphi[1,:] =  -1.0
  dphi[2,:] =   1.0

  return phi, dphi

end

function basisFuncLineLag( P::P2, s::Vector{Float64} )

  nx  = length( s )

  # Quadratric
  nphi = 3
  phi  = Array{Float64}(nphi, nx)
  dphi = Array{Float64}(nphi, nx)

  t = 1 - s

  phi[1,:] =  s.*-3.0+(s.*s).*2.0+1.0
  phi[2,:] =  -s+(s.*s).*2.0
  phi[3,:] =  s.*4.0-(s.*s).*4.0

  dphi[1,:,1] =  s.*4.0-3.0
  dphi[2,:,1] =  s.*4.0-1.0
  dphi[3,:,1] =  s.*-8.0+4.0

  return phi, dphi

end

function basisFuncLineLag( P::P3, s::Vector{Float64} )

  nx  = length( s )

  # Cubic
  nphi = 4
  phi  = Array{Float64}(nphi, nx)
  dphi = Array{Float64}(nphi, nx)

  t = 1 - s

  phi[1,:] =  s.*(-1.1E1/2.0)+(s.*s).*9.0-(s.*s.*s).*(9.0/2.0)+1.0
  phi[2,:] =  s-(s.*s).*(9.0/2.0)+(s.*s.*s).*(9.0/2.0)
  phi[3,:] =  s.*9.0-(s.*s).*(4.5E1/2.0)+(s.*s.*s).*(2.7E1/2.0)
  phi[4,:] =  s.*(-9.0/2.0)+(s.*s).*1.8E1-(s.*s.*s).*(2.7E1/2.0)

  dphi[1,:,1] =  s.*1.8E1-(s.*s).*(2.7E1/2.0)-1.1E1/2.0
  dphi[2,:,1] =  s.*-9.0+(s.*s).*(2.7E1/2.0)+1.0
  dphi[3,:,1] =  s.*-4.5E1+(s.*s).*(8.1E1/2.0)+9.0
  dphi[4,:,1] =  s.*3.6E1-(s.*s).*(8.1E1/2.0)-9.0/2.0

  return phi, dphi

end

function basisFuncLineLag( P::P4, s::Vector{Float64} )

  nx  = length( s )

  # Quartic
  nphi = 5
  phi  = Array{Float64}(nphi, nx)
  dphi = Array{Float64}(nphi, nx)

  t = 1 - s

  phi[1,:] =  (1./3).*(s.*-2.5E1+(s.*s).*7.0E1-(s.*s.*s).*8.0E1+(s.*s.*s.*s).*3.2E1+3.0)
  phi[2,:] =  -(1./3).*(s.*3.0-(s.*s).*2.2E1+(s.*s.*s).*4.8E1-(s.*s.*s.*s).*3.2E1)
  phi[3,:] =  (1./3).*(s.*4.8E1-(s.*s).*2.08E2+(s.*s.*s).*2.88E2-(s.*s.*s.*s).*1.28E2)
  phi[4,:] =  -(1./3).*(s.*3.6E1-(s.*s).*2.28E2+(s.*s.*s).*3.84E2-(s.*s.*s.*s).*1.92E2)
  phi[5,:] =  (1./3).*(s.*1.6E1-(s.*s).*1.12E2+(s.*s.*s).*2.24E2-(s.*s.*s.*s).*1.28E2)

  dphi[1,:,1] =  (1./3).*(s.*1.4E2-(s.*s).*2.4E2+(s.*s.*s).*1.28E2-2.5E1)
  dphi[2,:,1] =  (1./3).*(s.*4.4E1-(s.*s).*1.44E2+(s.*s.*s).*1.28E2-3.0)
  dphi[3,:,1] =  -(1./3).*(s.*4.16E2-(s.*s).*8.64E2+(s.*s.*s).*5.12E2-4.8E1)
  dphi[4,:,1] =  (1./3).*(s.*4.56E2-(s.*s).*1.152E3+(s.*s.*s).*7.68E2-3.6E1)
  dphi[5,:,1] =  -(1./3).*(s.*2.24E2-(s.*s).*6.72E2+(s.*s.*s).*5.12E2-1.6E1)

  return phi, dphi

end

# function basisFuncLineLag( P::P5, s::Vector{Float64} )
#
#   nx  = length( s )
#
#   # Quintic
#   nphi = 6
#   phi  = Array{Float64}(nphi, nx)
#   dphi = Array{Float64}(nphi, nx)
#
#   t = 1 - s
#
#   phi[1,:] =  -(1./24).*(s.*2.74E2-(s.*s).*1.125E3+(s.*s.*s).*2.125E3-(s.*s.*s.*s).*1.875E3+(s.*s.*s.*s.*s).*6.25E2-2.4E1)
#   phi[2,:] =  (1./24).*(s.*2.4E1-(s.*s).*2.5E2+(s.*s.*s).*8.75E2-(s.*s.*s.*s).*1.25E3+(s.*s.*s.*s.*s).*6.25E2)
#   phi[3,:] =  (1./24).*(s.*6.0E2-(s.*s).*3.85E3+(s.*s.*s).*8.875E3-(s.*s.*s.*s).*8.75E3+(s.*s.*s.*s.*s).*3.125E3)
#   phi[4,:] =  -(1./24).*(s.*6.0E2-(s.*s).*5.35E3+(s.*s.*s).*1.475E4-(s.*s.*s.*s).*1.625E4+(s.*s.*s.*s.*s).*6.25E3)
#   phi[5,:] =  (1./24).*(s.*4.0E2-(s.*s).*3.9E3+(s.*s.*s).*1.225E4-(s.*s.*s.*s).*1.5E4+(s.*s.*s.*s.*s).*6.25E3)
#   phi[6,:] =  -(1./24).*(s.*1.5E2-(s.*s).*1.525E3+(s.*s.*s).*5.125E3-(s.*s.*s.*s).*6.875E3+(s.*s.*s.*s.*s).*3.125E3)
#
#
#   dphi[1,:,1] =  -(1./24).*(s.*-2.25E3+(s.*s).*6.375E3-(s.*s.*s).*7.5E3+(s.*s.*s.*s).*3.125E3+2.74E2)
#   dphi[2,:,1] =  (1./24).*(s.*-5.0E2+(s.*s).*2.625E3-(s.*s.*s).*5.0E3+(s.*s.*s.*s).*3.125E3+2.4E1)
#   dphi[3,:,1] =  (1./24).*(s.*-7.7E3+(s.*s).*2.6625E4-(s.*s.*s).*3.5E4+(s.*s.*s.*s).*1.5625E4+6.0E2)
#   dphi[4,:,1] =  -(1./24).*(s.*-1.07E4+(s.*s).*4.425E4-(s.*s.*s).*6.5E4+(s.*s.*s.*s).*3.125E4+6.0E2)
#   dphi[5,:,1] =  (1./24).*(s.*-7.8E3+(s.*s).*3.675E4-(s.*s.*s).*6.0E4+(s.*s.*s.*s).*3.125E4+4.0E2)
#   dphi[6,:,1] =  -(1./24).*(s.*-3.05E3+(s.*s).*1.5375E4-(s.*s.*s).*2.75E4+(s.*s.*s.*s).*1.5625E4+1.5E2)
#
#   return phi, dphi
#
# end


#
# Jeffrey D. Scargle et.al., 2013,
# "Studies In Astronomical Time Series Analysis. VI. Bayesian Block Representations",
# http://iopscience.iop.org/article/10.1088/0004-637X/764/2/167/meta
#
# To be used for adaptive ADWIN
#
# Bayesian blocks method joins bins in the regular histogram in an optimal way.
#
using Base.Test

"""
Eq. 19 from Scargle 2012
:param n_k: number of elements in each block
:param t_k: width of the block
:return: fitness value.

NOTE: Looks similar to the entropy function.
"""
function fitness_eq19(n_k, t_k)
    n_k * (log(n_k) - log(t_k))
end

function fit_prior(n, p0=0.05)
    4 - log(75.33 * p0 * n ^(-0.478))
end

"""
In this method the regular histogram is computed first and after that its bins
are joined using dynamic programming.

It is like a normal changedetection but along the Y-axis.
That is why last[] and best[] have lengths of the number of unique elements
(defining edges boundaries).

About the statement
 n_k = n_k = cumsum(x[end:-1:1])[end:-1:1]
 Here the twice reversing is used because e.g for the input
 [1, 1, 1, 2, 1]
the output should be
 [6, 5, 4, 3, 1]
which is a possible number of elements in the last block prior to r.
"""
function bayesian_blocks(t_input, p0_prior=0.05, delta=0.0)
    uniq_t = sort(unique(t_input))
    #... Start Calculate histogram
    edges = [uniq_t[1]-delta;
             0.5 * (uniq_t[2:end] + uniq_t[1:end-1]);
             uniq_t[end]+delta]

    # Count elements in blocks.
    x::Array{Int,1} = []
    for e in uniq_t
        push!(x, count(x -> x == e, t_input))
    end
    @assert length(x) == length(edges) - 1
    #... End Calculate histogram
    # println("...1 edges, counts:", edges, x)
    block_lengths = edges[end] - edges[1:end-1]
    last::Array{Int,1} = zeros(length(uniq_t))
    best::Array{Float64,1} = zeros(length(uniq_t))
    n = length(uniq_t)
    ## Construct a vector of possible block lengths
    ## and count the number of elements in each block.
    for r = 1:n
        if r==1
            last[r] = 1
            best[r] = 0.0
        else
            ## Possible blocks widths prior to r
            t_k = block_lengths[1:r-1] - block_lengths[r] # ??? + delta
            ## Possible numbers of elements in blocks prior to r.
            n_k = cumsum(x[1:r-1][end:-1:1])[end:-1:1]
            # if r==5 @assert length(t_k) == length(n_k) end
            fit_vec = map(x -> fitness_eq19(x[1], x[2]), zip(n_k, t_k))
            a_r = zeros(r) ## a_r[1]==0 automatically; to make sure that the size is correct
            a_r[2:r] = fit_vec - fit_prior(n, p0_prior)
            a_r[2:r] += best[1:r-1]
            v_max, i_max = findmax(a_r)
            last[r] = i_max
            best[r] = v_max
        end
    end
    # last, best
    chps = extract_changepoints(last)
    edges_c = [edges[c] for c in chps]
    if edges_c[end]!=edges[end]
        edges_c = [edges_c; edges[end]]
    end
    edges_c
end

"""
    In this method an input is a is a pre-calculated histogram.
"""
function bayesian_blocks(edges::Array{Float64,1},
                         counts::Array{Int,1},
                         p0_prior=0.05)
    x = counts
    # println("... 2 edges, counts:", edges, x)
    @assert length(edges) - 1 == length(x)
    block_lengths = edges[end] - edges[1:end-1]
    uniq_t = zeros(length(edges)-1) ## to be readable along other method
    last::Array{Int,1} = zeros(length(uniq_t))
    best::Array{Float64,1} = zeros(length(uniq_t))
    n=length(uniq_t)
    # println("...2: n=$n")
    ## Construct a vector of possible block lengths
    ## and count the number of elements in each block.
    for r = 1:n
        if r==1
            last[r] = 1
            best[r] = 0.0
        else
            ## Possible blocks widths prior to r
            t_k = block_lengths[1:r-1] - block_lengths[r] # ??? + delta
            ## Possible numbers of elements in blocks prior to r.
            n_k = cumsum(x[1:r-1][end:-1:1])[end:-1:1]
            if r==5 @assert length(t_k) == length(n_k) end
            fit_vec = map(x -> fitness_eq19(x[1], x[2]), zip(n_k, t_k))
            a_r = zeros(r) ## a_r[1]==0 automatically; to make sure that the size is correct
            a_r[2:r] = fit_vec - fit_prior(n, p0_prior)
            a_r[2:r] += best[1:r-1]
            v_max, i_max = findmax(a_r)
            last[r] = i_max
            best[r] = v_max
        end
    end
    chps=extract_changepoints(last)
    edges_c = [edges[c] for c in chps]
    if edges_c[end]!=edges[end]
        edges_c = [edges_c; edges[end]]
    end
    edges_c
    # last
end

function extract_changepoints(last::Array{Int,1})
    # cp1 = last[end]
    # cp2 = last[cp1]
    chps::Array{Int,1} = []
    next_chps = last[end]
    while next_chps != 1
        push!(chps, next_chps)
        if next_chps != last[next_chps]
            next_chps = last[next_chps]
        else
            next_chps = last[next_chps-1]
        end
    end
    push!(chps,1)
    sort(chps)
end

function runtests()
    println("... Start tests()")
    ret1 = bayesian_blocks([1, 2, 3, 4, 5, 5, 1.1, 1.1], 0.05)

    ret2 = bayesian_blocks([1.0, 1.05, 1.55, 2.5, 3.5, 4.5, 5.0],
                           [1, 2, 1, 1, 1, 2], 0.05)
    @test ret1 == ret2
    # test extract_changepoints
    @test extract_changepoints([1,1,1,2,2,2,5]) == [1,2,5]
    @test extract_changepoints([1,1,1,2,2,2,5,3]) == [1,3]
    println("... End tests()")
end
# runtests()


export tod2map, combinemaps!

@doc raw"""
    tod2map{T,O}(pixidx, tod::Array{T}; nside=128) :: (map, hits)

Create a binned map for a TOD and return a tuple containing the map
itself and the hit map.
"""
function tod2map(pixidx, tod::Array{T}; nside=128, ordering=Healpix.Ring) where {T}
    @assert length(pixidx) == length(tod)

    binnedmap = Map{T,ordering}(nside)
    hitmap = Map{Int,ordering}(nside)

    @inbounds for i in eachindex(pixidx)
        if !ismissing(tod[i])
            binnedmap.pixels[pixidx[i]] += tod[i]
            hitmap.pixels[pixidx[i]] += 1
        end
    end

    @inbounds for i in eachindex(binnedmap.pixels)
        if hitmap.pixels[i] > 0
            binnedmap.pixels[i] /= hitmap.pixels[i]
        end
    end

    (binnedmap, hitmap)
end

@doc raw"""
    combinemaps{T, O, H}(destmap::Map{T, O}, desthitmap::Map{H, O}, othermap::Map{T, O}, otherhitmap::Map{H, O})

Sum "othermap" to "destmap", assuming that both maps have been
produced by binning TODs. The parameters `desthitmap` and
`otherhitmap` are the two hit maps. At the end of the call, `destmap`
and `desthitmap` are updated.
"""
function combinemaps!(destmap::Map{T,O},
                      desthitmap::Map{H,O},
                      othermap::Map{T,O},
                      otherhitmap::Map{H,O}) where {T,O,H}
    (conformables(destmap, othermap)
     && conformables(desthitmap, otherhitmap)
     && conformables(destmap, desthitmap)) || throw(DomainError())

    for i in eachindex(destmap.pixels)
        (pix1, pix2) = (destmap.pixels[i], othermap.pixels[i])
        (hit1, hit2) = (desthitmap.pixels[i], otherhitmap.pixels[i])
        hitsum = hit1 + hit2
        if hitsum > 0
            destmap.pixels[i] = (pix1 * hit1 + pix2 * hit2) / hitsum
        end
        desthitmap.pixels[i] = hitsum
    end
end


module MyMethods

# Dependencies =================================================================
using LinearAlgebra
using Random
using Distributions
using DataFrames

# Module export ================================================================

export myEstimator, myLS, myTSLS, mySieve
export coef, predict, inference # my methods
export get_basis # helper functions

# Module content ===============================================================

include("FUN/myEstimator.jl") # myEstimators

include("FUN/myLS.jl") # myLS

include("FUN/myTSLS.jl") # myLS

include("FUN/mySieve.jl") # mySieve

end


module PolrfitTest

using Test, OrdinalMultinomialModels, RDatasets, MathOptInterface
const MOI = MathOptInterface

housing = dataset("MASS", "housing")
@info "Housing example for `polr` function in R package MASS"

@testset "logit link" begin
    Ipopt_solver = Ipopt.Optimizer()
    MOI.set(Ipopt_solver, MOI.RawOptimizerAttribute("print_level"), 0)
    NLopt_solver = NLopt.Optimizer()
    MOI.set(NLopt_solver, MOI.RawOptimizerAttribute("algorithm"), :LD_SLSQP)
    for solver in [Ipopt_solver, NLopt_solver]
        houseplr = polr(@formula(Sat ~ Infl + Type + Cont), housing,
            LogitLink(), solver; wts = housing[!, :Freq])
        @test nobs(houseplr) == 72
        @test dof(houseplr) == 8
        @test isapprox(deviance(houseplr), 3479.149; rtol=1e-4)
        @test isapprox(coef(houseplr), [-0.4961353, 0.6907083, 0.5663937, 1.2888191, -0.5723501, -0.3661866, -1.0910149, 0.3602841]; rtol=1e-4)
        @test isapprox(stderror(houseplr), [0.124541, 0.125212, 0.104963, 0.126705, 0.118747, 0.156766, 0.151514, 0.0953575]; rtol=1e-4)
        @test loglikelihood(houseplr.model) ≈ -1739.5746495294738
        @test fitted(houseplr.model) == nothing # TODO
        @test predict(houseplr.model) == nothing # TODO
        @test response(houseplr.model) === houseplr.model.Y
        @test modelmatrix(houseplr.model) === houseplr.model.X
        @test score(houseplr.model) === houseplr.model.∇
        @test vcov(houseplr.model) === houseplr.model.vcov
        @test all(abs.(cor(houseplr.model)) .≤ 1 + eps(Float64))
        @test all(weights(houseplr.model) .== housing[!, :Freq])
        @show coeftable(houseplr)
        @show confint(houseplr.model)
    end
end

@testset "probit link" begin
    Ipopt_solver = Ipopt.Optimizer()
    MOI.set(Ipopt_solver, MOI.RawOptimizerAttribute("print_level"), 0)
    NLopt_solver = NLopt.Optimizer()
    MOI.set(NLopt_solver, MOI.RawOptimizerAttribute("algorithm"), :LD_SLSQP)
    for solver in [Ipopt_solver, NLopt_solver]
        houseplr = polr(@formula(Sat ~ Infl + Type + Cont), housing,
            ProbitLink(), solver; wts = housing[!, :Freq])
        @test nobs(houseplr) == 72
        @test dof(houseplr) == 8
        @test isapprox(deviance(houseplr), 3479.689; rtol=1e-4)
        @test isapprox(coef(houseplr), [-0.29983, 0.42672, 0.34642, 0.78291, -0.34754, -0.21789, -0.66417, 0.22239]; rtol=1e-4)
        @test isapprox(stderror(houseplr), [0.0761614, 0.0763991, 0.0641796, 0.0762645, 0.0722116, 0.0955741, 0.0919294, 0.0581214]; rtol=1e-4)
    end
end

@testset "cloglog link" begin
    Ipopt_solver = Ipopt.Optimizer()
    MOI.set(Ipopt_solver, MOI.RawOptimizerAttribute("print_level"), 0)
    NLopt_solver = NLopt.Optimizer()
    MOI.set(NLopt_solver, MOI.RawOptimizerAttribute("algorithm"), :LD_SLSQP)
    for solver in [Ipopt_solver, NLopt_solver]
        houseplr = polr(@formula(Sat ~ Infl + Type + Cont), housing,
            CloglogLink(), solver; wts = housing[!, :Freq])
        @test nobs(houseplr) == 72
        @test dof(houseplr) == 8
        @test isapprox(deviance(houseplr), 3484.053; rtol=1e-4)
        @test isapprox(coef(houseplr), [-0.79622, 0.05537, 0.38204, 0.91536, -0.40720, -0.28053, -0.74245, 0.20922]; rtol=1e-3)
        @test isapprox(stderror(houseplr), [0.0904734, 0.0866657, 0.0701218, 0.0924964, 0.0861325, 0.112948, 0.102084, 0.0653755]; rtol=1e-4)
    end
end

end


"""
    index_map, num_segments = felzenszwalb(edges, num_vertices, k, min_size=0)

Segment an image represented as Region Adjacency Graph(RAG) using Felzenszwalb's segmentation algorithm. Each pixel/region
corresponds to a node in the graph and weights on each edge measure the dissimilarity between pixels.
The function returns the number of segments and index mapping from nodes of the RAG to segments.

Parameters:
- `edges`:        Array of edges in RAG. Each edge is represented as `ImageEdge`.
- `num_vertices`: Number of vertices in RAG
- `k`:            Threshold for region merging step. Larger threshold will result in bigger segments.
- `min_size`:     Minimum segment size (in # pixels)
"""
function felzenszwalb(edges::Array{ImageEdge}, num_vertices::Int, k::Float64, min_size::Int = 0)

    num_edges = length(edges)
    G = IntDisjointSets(num_vertices)
    set_size = ones(num_vertices)
    threshold = fill(convert(Float64,k), num_vertices)

    sort!(edges, lt = (x,y)->(x.weight<y.weight))

    for edge in edges
        w = edge.weight
        a = find_root!(G, edge.index1)
        b = find_root!(G, edge.index2)
        if a!=b
            if w <= min(threshold[a], threshold[b])
                merged_root = union!(G, a, b)
                set_size[merged_root] = set_size[a] + set_size[b]
                threshold[merged_root] = w + k/set_size[merged_root]
            end
        end
    end

    #merge small segments
    for edge in edges
        a = find_root!(G, edge.index1)
        b = find_root!(G, edge.index2)
        if a!=b && (set_size[a] < min_size || set_size[b] < min_size)
            union!(G, a, b)
        end
    end

    segments = OrderedSet{Int}()
    for i in 1:num_vertices
        push!(segments, find_root!(G, i))
    end

    num_sets = length(segments)
    segments2index = Dict{Int, Int}()
    for (i, s) in enumerate(segments)
        segments2index[s] = i
    end

    index_map = Vector{Int}(undef, num_vertices)
    for i in 1:num_vertices
        index_map[i] = segments2index[find_root!(G, i)]
    end

    return index_map, num_sets
end
felzenszwalb(edges::Array{ImageEdge}, num_vertices::Integer, k::Real, min_size::Integer = 0) =
    felzenszwalb(edges, convert(Int, num_vertices)::Int, convert(Float64, k)::Float64, convert(Int, min_size)::Int)

meantype(::Type{T}) where T = typeof(zero(accum_type(T))/2)

"""
    segments = felzenszwalb(img, k, [min_size])

Segment an image using Felzenszwalb's segmentation algorithm and returns the result as `SegmentedImage`.
The algorithm uses euclidean distance in color space as edge weights for the region adjacency graph.

Parameters:
- `img`:      input image
- `k`:        Threshold for region merging step. Larger threshold will result in bigger segments.
- `min_size`: Minimum segment size (in # pixels)
"""
function felzenszwalb(img::AbstractArray{T}, k::Real, min_size::Int = 0) where T<:Union{Real,Color}

    sz = size(img)
    num_vertices = prod(sz)

    R = CartesianIndices(img)
    L = LinearIndices(img)
    Ibegin, Iend = first(R), last(R)
    I1 = _oneunit(Ibegin)

    # Compute the number of entries per pixel (other than at the image edges)
    num_edges = 0
    for I in _colon(-I1, I1)
        I >= zero(I1) && continue
        num_edges += 1
    end
    num_edges *= num_vertices   # now the number for the whole image
    edges = Vector{ImageEdge}(undef, num_edges)

    num = 0
    for I in R
        imgI = img[I]
        for J in _colon(max(Ibegin, I-I1), min(Iend, I+I1))
            if I >= J
                continue
            end
            edges[num+=1] = ImageEdge(L[I], L[J], sqrt(_abs2(imgI-meantype(T)(img[J]))))
        end
    end
    deleteat!(edges, num+1:num_edges)   # compensate for the ones we were missing at the image edges

    index_map, num_segments = felzenszwalb(edges, num_vertices, k, min_size)

    result              = similar(img, Int)
    labels              = Array(1:num_segments)
    region_means        = Dict{Int, meantype(T)}()
    region_pix_count    = Dict{Int, Int}()

    for I in R
        result[I] = index_map[L[I]]
        region_pix_count[result[I]] = get(region_pix_count, result[I], 0) + 1
        region_means[result[I]] = get(region_means, result[I], zero(meantype(T))) + (img[I] - get(region_means, result[I], zero(meantype(T))))/region_pix_count[result[I]]
    end

    return SegmentedImage(result, labels, region_means, region_pix_count)
end


struct NewtonSolver{T, FT, TJ, TL} <: AbstractNewtonSolver{T}
    @newton_solver_variables
    function NewtonSolver{T,FT,TJ,TL}(x, F!, Jparams, linear_solver) where {T,FT,TJ,TL}
        J  = zero(linear_solver.A)
        x₀ = zero(x)
        x₁ = zero(x)
        y₁ = zero(x)
        y₀ = zero(x)
        δx = zero(x)
        δy = zero(x)

        nls_params = NonlinearSolverParameters(T)
        nls_status = NonlinearSolverStatus{T}(length(x))

        new(x, J, x₀, x₁, y₀, y₁, δx, δy, F!, Jparams, linear_solver,
            nls_params, nls_status)
    end
end

function NewtonSolver(x::AbstractVector{T}, F!::Function; J!::Union{Function,Nothing}=nothing) where {T}
    n = length(x)
    Jparams = getJacobianParameters(J!, F!, T, n)
    linear_solver = getLinearSolver(x)
    NewtonSolver{T, typeof(F!), typeof(Jparams), typeof(linear_solver)}(x, F!, Jparams, linear_solver)
end

function solve!(s::NewtonSolver{T}; n::Int=0) where {T}
    local nmax::Int = n > 0 ? nmax = n : s.params.nmax

    s.F!(s.x, s.y₀)
    residual_initial!(s.status, s.x, s.y₀)
    s.status.i  = 0

    if s.status.rₐ ≥ s.params.atol²
        for s.status.i = 1:nmax
            computeJacobian(s)
            copyto!(s.linear.A, s.J)
            s.linear.b .= -one(T) .* s.y₀
            factorize!(s.linear)
            solve!(s.linear)
            copyto!(s.δx, s.linear.b)
            s.x .+= s.δx
            s.F!(s.x, s.y₀)
            residual!(s.status, s.δx, s.x, s.y₀)

            if check_solver_converged(s.status, s.params) && s.status.i ≥ s.params.nmin && !(n > 0)
                break
            end
        end
    end
end


using Test, SegmentIntersections

@testset "Test results" begin
    segments = [
        Segment(4.9, 3.56, 2.86, 1.62),
        Segment(3.63, 3.88, 5.45, 2.32),
        Segment(2.99, 3.29, 4.47, 1.8),
        Segment(2.66, 2.85, 3.81, 1.44),
    ];
    expected_intersections = [Point(4.48, 3.16), Point(3.78, 2.5), Point(3.31, 2.05)];
    intersections = find_intersections_brute(segments)
    for expected in expected_intersections
        match = false
        for intersection in intersections
            if isapprox(intersection, expected, rtol = 1e-2)
                match = true
            end
        end
        @test match == true
    end
end


@testset "multilevel_PARSDMM" begin

Random.seed!(123)

#test serial multilevel PARSDMM
    TF = Float32
    options = PARSDMM_options()
    default_PARSDMM_options(options,TF)
    options.parallel     = false
    options.evol_rel_tol = 100*eps(TF)
    options.maxit        = 5000

    n_levels          = 2
    coarsening_factor = 1.87

    x         = randn(TF,100,201)
    comp_grid = compgrid((1.0, 1.0),(100, 201))
    x         = vec(x)
    m         = deepcopy(x)

    constraint = Vector{SetIntersectionProjection.set_definitions}()

    #total variation
    (TD_OP, AtA_diag, dense, TD_n)=get_TD_operator(comp_grid,"TV",options.FL)
    m_min     = 0.0
    m_max     = 0.5f0*norm(TD_OP*x,1)
    set_type  = "l1"
    TD_OP     = "TV"
    app_mode  = ("matrix","")
    custom_TD_OP = ([],false)
    push!(constraint, set_definitions(set_type,TD_OP,m_min,m_max,app_mode,custom_TD_OP))

    #(P_sub,TD_OP,set_Prop) = setup_constraints(constraint,comp_grid,options.FL)
    #(TD_OP,AtA,l,y)        = PARSDMM_precompute_distribute(TD_OP,set_Prop,comp_grid,options)
    (TD_OP_levels,AtA_levels,P_sub_levels,set_Prop_levels,comp_grid_levels)=setup_multi_level_PARSDMM(m,n_levels,coarsening_factor,comp_grid,constraint,options)

    (x1,log_PARSDMM) = PARSDMM(m,AtA,TD_OP,set_Prop,P_sub,comp_grid,options);

    result = deepcopy(x1)
    for i=1:length(TD_OP)-1
      copy!(result,x1)
      @test norm(P_sub[i](TD_OP[i]*result)-(TD_OP[i]*result))/norm((TD_OP[i]*result)) <= 1.5*options.feas_tol
    end

#test parallel multilevel PARSDMM
        options=PARSDMM_options()
        default_PARSDMM_options(options,Float64)
        options.parallel = true

            (TD_OP_levels,AtA_levels,P_sub_levels,set_Prop_levels,comp_grid_levels,constraint_level)=setup_multi_level_PARSDMM(m,n_levels,coarsening_factor,comp_grid,constraint,options)

            (x2,log_PARSDMM) = PARSDMM_multi_level(m,TD_OP_levels,AtA_levels,P_sub_levels,set_Prop_levels,comp_grid_levels,options);
            result=Vector{typeof(x[1])}(length(x2))
            for i=1:length(TD_OP_levels[1])-1
              copy!(result,x2)
              @test norm(P_sub_levels[1][i](TD_OP_levels[1][i]*result)-(TD_OP_levels[1][i]*result))/norm((TD_OP_levels[1][i]*result)) <= options.feas_tol
            end
@test isapprox(x1,x2,rtol=options.feas_tol)

#run serial PARSDMM and compare results
default_PARSDMM_options(options,Float64)
options.parallel    = false

(P_sub,TD_OP,set_Prop) = setup_constraints(constraint,comp_grid,options.FL);
(TD_OP,AtA,l,y) = PARSDMM_precompute_distribute(TD_OP,set_Prop,comp_grid,options)

(x3,log_PARSDMM) = PARSDMM(m,AtA,TD_OP,set_Prop,P_sub,comp_grid,options);
result=Vector{typeof(x[1])}(length(x3))
for i=1:length(TD_OP)-1
  copy!(result,x3)
  @test norm(P_sub[i](TD_OP[i]*result)-(TD_OP[i]*result))/norm((TD_OP[i]*result)) <= options.feas_tol
end

@test isapprox(x1,x3,rtol=options.feas_tol)

end


"""
Check global mean temperature at every vertical level from the m=n=1 coefficient
of the associated Legendre polynomials.
"""
function check_global_mean_temperature( Tabs::Array{Complex{T},3},
                                        P::Parameters) where T
    @unpack print_dialog = P

    # The m=n=1 coefficent in spectral space is the mean value
    # √2 is due to the normalization of the associated Legendre polynomials
    gmt = sqrt(0.5)*Float64.(real.(Tabs[1,1,:]))

    if print_dialog
        println(gmt)
    end
end

# # Prints global means of eddy kinetic energy and temperature.
# # Also stops the integration if the computed diagnostics are outside of
# # allowable ranges.
# function check_diagnostics(geometry::Geometry, spectral_trans::SpectralTrans, ξ, D, Tₐ, step,
#                            print_diag=true)
#     @unpack nlev, mx, nx = geometry
#
#     diagnostics = zeros(Float64, nlev, 3)
#
#     # 1. Get global-mean temperature and compute eddy kinetic energy
#     for k in 1:nlev
#         diagnostics[k,3] = √(0.5)*real(Tₐ[1,1,k])
#
#         temp = ∇⁻²(spectral_trans, ξ[:,:,k])
#
#         for m in 2:mx
#             for n in 1:nx
#                 diagnostics[k,1] = diagnostics[k,1] - real(temp[m,n]*conj(ξ[m,n,k]))
#             end
#         end
#
#         temp = ∇⁻²(spectral_trans, D[:,:,k])
#
#         for m in 2:mx
#             for n in 1:nx
#                 diagnostics[k,2] = diagnostics[k,2] - real(temp[m,n]*conj(D[m,n,k]))
#             end
#         end
#     end
#
#     # 2. Print results to screen
#     if print_diag
#         println("Step = $step")
#         println(diagnostics[:,1])
#         println(diagnostics[:,2])
#         println(diagnostics[:,3])
#     end
#
#     # 3. Stop integration if model variables are out of range
#     for k in 1:nlev
#         if diagnostics[k,1] > 500.0 || diagnostics[k,2] > 500.0 ||
#             diagnostics[k,3] < 180.0 || diagnostics[k,3] > 320.0
#
#             println(diagnostics[:,1])
#             println(diagnostics[:,2])
#             println(diagnostics[:,3])
#
#             throw("Model variables out of accepted range")
#         end
#     end
# end


import Base: rand,
             ∈

export Ball2,
       sample

"""
    Ball2{N<:AbstractFloat} <: AbstractCentrallySymmetric{N}

Type that represents a ball in the 2-norm.

### Fields

- `center` -- center of the ball as a real vector
- `radius` -- radius of the ball as a real scalar (``≥ 0``)

### Notes

Mathematically, a ball in the 2-norm is defined as the set

```math
\\mathcal{B}_2^n(c, r) = \\{ x ∈ \\mathbb{R}^n : ‖ x - c ‖_2 ≤ r \\},
```
where ``c ∈ \\mathbb{R}^n`` is its center and ``r ∈ \\mathbb{R}_+`` its radius.
Here ``‖ ⋅ ‖_2`` denotes the Euclidean norm (also known as 2-norm), defined as
``‖ x ‖_2 = \\left( \\sum\\limits_{i=1}^n |x_i|^2 \\right)^{1/2}`` for any
``x ∈ \\mathbb{R}^n``.

### Examples

Create a five-dimensional ball `B` in the 2-norm centered at the origin with
radius 0.5:

```jldoctest ball2_label
julia> B = Ball2(zeros(5), 0.5)
Ball2{Float64}([0.0, 0.0, 0.0, 0.0, 0.0], 0.5)
julia> dim(B)
5
```

Evaluate `B`'s support vector in the direction ``[1,2,3,4,5]``:

```jldoctest ball2_label
julia> σ([1.,2.,3.,4.,5.], B)
5-element Array{Float64,1}:
 0.06741998624632421
 0.13483997249264842
 0.20225995873897262
 0.26967994498529685
 0.3370999312316211
```
"""
struct Ball2{N<:AbstractFloat} <: AbstractCentrallySymmetric{N}
    center::Vector{N}
    radius::N

    # default constructor with domain constraint for radius
    function Ball2{N}(center::Vector{N}, radius::N) where {N<:AbstractFloat}
        @assert radius >= zero(N) "radius must not be negative"
        return new{N}(center, radius)
    end
end

# convenience constructor without type parameter
Ball2(center::Vector{N}, radius::N) where {N<:AbstractFloat} =
    Ball2{N}(center, radius)

# --- AbstractCentrallySymmetric interface functions ---

"""
    center(B::Ball2{N})::Vector{N} where {N<:AbstractFloat}

Return the center of a ball in the 2-norm.

### Input

- `B` -- ball in the 2-norm

### Output

The center of the ball in the 2-norm.
"""
function center(B::Ball2{N})::Vector{N} where {N<:AbstractFloat}
    return B.center
end

# --- LazySet interface functions ---

"""
    σ(d::AbstractVector{N}, B::Ball2{N}) where {N<:AbstractFloat}

Return the support vector of a 2-norm ball in a given direction.

### Input

- `d` -- direction
- `B` -- ball in the 2-norm

### Output

The support vector in the given direction.
If the direction has norm zero, the origin is returned.

### Notes

Let ``c`` and ``r`` be the center and radius of a ball ``B`` in the 2-norm,
respectively.
For nonzero direction ``d`` we have

```math
σ_B(d) = c + r \\frac{d}{‖d‖_2}.
```

This function requires computing the 2-norm of the input direction, which is
performed in the given precision of the numeric datatype of both the direction
and the set.
Exact inputs are not supported.
"""
function σ(d::AbstractVector{N}, B::Ball2{N}) where {N<:AbstractFloat}
    dnorm = norm(d, 2)
    if dnorm <= zero(N)
        return zeros(N, length(d))
    else
        return @. B.center + d * (B.radius / dnorm)
    end
end

"""
    ∈(x::AbstractVector{N}, B::Ball2{N})::Bool where {N<:AbstractFloat}

Check whether a given point is contained in a ball in the 2-norm.

### Input

- `x` -- point/vector
- `B` -- ball in the 2-norm

### Output

`true` iff ``x ∈ B``.

### Notes

This implementation is worst-case optimized, i.e., it is optimistic and first
computes (see below) the whole sum before comparing to the radius.
In applications where the point is typically far away from the ball, a fail-fast
implementation with interleaved comparisons could be more efficient.

### Algorithm

Let ``B`` be an ``n``-dimensional ball in the 2-norm with radius ``r`` and let
``c_i`` and ``x_i`` be the ball's center and the vector ``x`` in dimension
``i``, respectively.
Then ``x ∈ B`` iff ``\\left( ∑_{i=1}^n |c_i - x_i|^2 \\right)^{1/2} ≤ r``.

### Examples

```jldoctest
julia> B = Ball2([1., 1.], sqrt(0.5))
Ball2{Float64}([1.0, 1.0], 0.7071067811865476)
julia> ∈([.5, 1.6], B)
false
julia> ∈([.5, 1.5], B)
true
```
"""
function ∈(x::AbstractVector{N}, B::Ball2{N})::Bool where {N<:AbstractFloat}
    @assert length(x) == dim(B)
    sum = zero(N)
    for i in eachindex(x)
        sum += (B.center[i] - x[i])^2
    end
    return sqrt(sum) <= B.radius
end

"""
    rand(::Type{Ball2}; [N]::Type{<:Real}=Float64, [dim]::Int=2,
         [rng]::AbstractRNG=GLOBAL_RNG, [seed]::Union{Int, Nothing}=nothing
        )::Ball2{N}

Create a random ball in the 2-norm.

### Input

- `Ball2` -- type for dispatch
- `N`     -- (optional, default: `Float64`) numeric type
- `dim`   -- (optional, default: 2) dimension
- `rng`   -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`  -- (optional, default: `nothing`) seed for reseeding

### Output

A random ball in the 2-norm.

### Algorithm

All numbers are normally distributed with mean 0 and standard deviation 1.
Additionally, the radius is nonnegative.
"""
function rand(::Type{Ball2};
              N::Type{<:Real}=Float64,
              dim::Int=2,
              rng::AbstractRNG=GLOBAL_RNG,
              seed::Union{Int, Nothing}=nothing
             )::Ball2{N}
    rng = reseed(rng, seed)
    center = randn(rng, N, dim)
    radius = abs(randn(rng, N))
    return Ball2(center, radius)
end

"""
    translate(B::Ball2{N}, v::AbstractVector{N}) where {N<:AbstractFloat}

Translate (i.e., shift) a ball in the 2-norm by a given vector.

### Input

- `B` -- ball in the 2-norm
- `v` -- translation vector

### Output

A translated ball in the 2-norm.

### Algorithm

We add the vector to the center of the ball.
"""
function translate(B::Ball2{N}, v::AbstractVector{N}) where {N<:AbstractFloat}
    @assert length(v) == dim(B) "cannot translate a $(dim(B))-dimensional " *
                                "set by a $(length(v))-dimensional vector"
    return Ball2(center(B) + v, B.radius)
end

"""
    sample(B::Ball2{N}, nsamples::Int=1;
           [rng]::AbstractRNG=GLOBAL_RNG,
           [seed]::Union{Int, Nothing}=nothing) where {N<:AbstractFloat}

Return samples from a uniform distribution on the given ball in the 2-norm.

### Input

- `B`        -- ball in the 2-norm
- `nsamples` -- (optional, default: `1`) number of samples
- `rng`      -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`     -- (optional, default: `nothing`) seed for reseeding

### Output

A linear array of `nsamples` elements drawn from a uniform distribution in `B`.

### Algorithm

Random sampling with uniform distribution in `B` is computed using Muller's method
of normalized Gaussians. This function requires the package `Distributions`.
See `_sample_unit_nball_muller!` for implementation details.
"""
function sample(B::Ball2{N}, nsamples::Int=1;
                rng::AbstractRNG=GLOBAL_RNG,
                seed::Union{Int, Nothing}=nothing) where {N<:AbstractFloat}
    require(:Distributions; fun_name="sample")
    n = dim(B)
    D = Vector{Vector{N}}(undef, nsamples) # preallocate output
    _sample_unit_nball_muller!(D, n, nsamples, rng=rng, seed=seed)

    # customize for the given ball
    r, c = B.radius, B.center
    @inbounds for i in 1:nsamples
        axpby!(one(N), c, r, D[i])
    end
    return D
end


"""
    LeakyReLU{T, N}(α::T) <: Bijector{N}

Defines the invertible mapping

    x ↦ x if x ≥ 0 else αx

where α > 0.
"""
struct LeakyReLU{T, N} <: Bijector{N}
    α::T
end

LeakyReLU(α::T; dim::Val{N} = Val(0)) where {T<:Real, N} = LeakyReLU{T, N}(α)
LeakyReLU(α::T; dim::Val{N} = Val(D)) where {D, T<:AbstractArray{<:Real, D}, N} = LeakyReLU{T, N}(α)

up1(b::LeakyReLU{T, N}) where {T, N} = LeakyReLU{T, N + 1}(b.α)

# (N=0) Univariate case
function (b::LeakyReLU{<:Any, 0})(x::Real)
    mask = x < zero(x)
    return mask * b.α * x + !mask * x
end
(b::LeakyReLU{<:Any, 0})(x::AbstractVector{<:Real}) = map(b, x)

function Base.inv(b::LeakyReLU{<:Any,N}) where N
    invα = inv.(b.α)
    return LeakyReLU{typeof(invα),N}(invα)
end

function logabsdetjac(b::LeakyReLU{<:Any, 0}, x::Real)
    mask = x < zero(x)
    J = mask * b.α + (1 - mask) * one(x)
    return log(abs(J))
end
logabsdetjac(b::LeakyReLU{<:Real, 0}, x::AbstractVector{<:Real}) = map(x -> logabsdetjac(b, x), x)

# We implement `forward` by hand since we can re-use the computation of
# the Jacobian of the transformation. This will lead to faster sampling
# when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
function forward(b::LeakyReLU{<:Any, 0}, x::Real)
    mask = x < zero(x)
    J = mask * b.α + !mask * one(x)
    return (rv=J * x, logabsdetjac=log(abs(J)))
end

# Batched version
function forward(b::LeakyReLU{<:Any, 0}, x::AbstractVector)
    J = let T = eltype(x), z = zero(T), o = one(T)
        @. (x < z) * b.α + (x > z) * o
    end
    return (rv=J .* x, logabsdetjac=log.(abs.(J)))
end

# (N=1) Multivariate case
function (b::LeakyReLU{<:Any, 1})(x::AbstractVecOrMat)
    return let z = zero(eltype(x))
        @. (x < z) * b.α * x + (x > z) * x
    end
end

function logabsdetjac(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
    # Is really diagonal of jacobian
    J = let T = eltype(x), z = zero(T), o = one(T)
        @. (x < z) * b.α + (x > z) * o
    end

    if x isa AbstractVector
        return sum(log.(abs.(J)))
    elseif x isa AbstractMatrix
        return vec(sum(log.(abs.(J)); dims = 1))  # sum along column
    end
end

# We implement `forward` by hand since we can re-use the computation of
# the Jacobian of the transformation. This will lead to faster sampling
# when using `rand` on a `TransformedDistribution` making use of `LeakyReLU`.
function forward(b::LeakyReLU{<:Any, 1}, x::AbstractVecOrMat)
    # Is really diagonal of jacobian
    J = let T = eltype(x), z = zero(T), o = one(T)
        @. (x < z) * b.α + (x > z) * o
    end

    if x isa AbstractVector
        logjac = sum(log.(abs.(J)))
    elseif x isa AbstractMatrix
        logjac = vec(sum(log.(abs.(J)); dims = 1))  # sum along column
    end

    y = J .* x
    return (rv=y, logabsdetjac=logjac)
end


"""
nstrides1d(l, k, p, s)

nstrides1d is a function to calculate the number of kernel strides that fit in a given
length. This equals the size of a dimension after a convolution.
"""
nstrides1d(l, k, p, s) = ((l - k + p) / s + 1)

"""
nstrides2d(l, k, p, s)

nstrides2d is a convenience function that accepts tuple inputs in the manner expected for
the pad1d and pad2d functions and outpus 2-element tuple.
"""
nstrides2d(l, k, p, s) = (nstrides1d(l[1], k[1], p[1] + p[2], s[1]),
                          nstrides1d(l[2], k[2], p[3] + p[4], s[2]))

"""
pad1d(l, k, s)

Calculate the amount of padding needed to fit an integer number of kernel strides
across a given length. While pad_1d works for 2-D inputs, it is suggested to use
the helper function pad2d.

Inputs:
-------
l - length\n
k - kernel size\n
s - stride\n

Output:
----
A 2-element tuple with the number of elements needed to pad each side of a 1-D length.
"""
function pad1d(l, k, s)
    n = ceil((l - k)/s + 1)  # Number of whole kernels needed to cross l
    pad_tot = s*(n - 1) + k - l  # Padding needed to get to n kernels
    # If padding still does not give an inter number of kernel strides
    # Then print an error
    @info "$l  $k  $s"
    if ((l - k + pad_tot) / s + 1) % 1 != 0
        @error "Could not pad array properly"
    end
    # Try to evenly space padding between sides
    # left, right, top, and bottom sides.
    # The following pads bottom/right sides first
    @info "pad_tot = $pad_tot  |  div = $(pad_tot ÷ 2)"
    p1 = convert(Int, pad_tot ÷ 2)
    p2 = convert(Int, pad_tot) - p1

    return (p1, p2)
end

"""
pad2d(l, k, s)

Convience function for pad_1d. This function splats the output into a single,
N-element tuple rather than using pad_1d.(x, y, z), which outputs a tuple of tuples

Inputs:
-------
l - length
k - kernel size
s - stride

Output:
----
A 4-element tuple with the number of elements needed to pad each side of a 2-D matrix.
"""
pad2d(l, k, s) = (pad1d(l[1], k[1], s[1])...,  pad1d(l[2], k[2], s[2])...)

"""
padchain(N, l, k, s)

A function for outputting an array of padding tuples for use in Flux Chain models
"""
function padchain(N, l, k, s)
    ls = Array{Tuple}(undef, N)
    ps = Array{Tuple}(undef, N)

    ls[1] = l
    ps[1] = pad2d(ls[1], k[1], s[1])
    for i = 2:N
        ls[i] = nstrides2d(ls[i-1], k[i-1], ps[i-1], s[i-1])
        ps[i] = pad2d(ls[i], k[i], s[i])
    end

    return ps
end

"""
image2tensor(image)

Converts an input greyscale or RGB image into a tensor suitable for use by Flux.
"""
function image2tensor(image)
    # Using views is 2.5x faster and makes 2x fewer allocations
    imsize = size(image)
    # In one line, turn a an image type into array, convert to Float32, and rearrange
    array = permuteddimsview(Float32.(channelview(image)), (3, 2, 1))
    tensor = reshape(array, (size(array)..., 1))

    return tensor
end

"""
tensor2image(tensor)

Converts an input tensor suitable into greyscale or RGB image.
"""
function tensor2image(tensor)
    if size(tensor)[3] == 1  # Grayscale
        image = colorview(Gray, permuteddimsview(reshape(tensor, size(tensor)[1:3]), (3, 2, 1)))
    elseif size(tensor)[3] == 3
        image = colorview(RGB, permuteddimsview(reshape(tensor, size(tensor)[1:3]), (3, 2, 1)))
    else
        @error "Number of image channels ($(size(image)[3])) not recognized (1 = grayscale, 3 = RGB)"
    end

    return image
end


struct ExponentialSampler <: Sampleable{Univariate,Continuous}
    scale::AbstractFloat
end

rand(rng::AbstractRNG, s::ExponentialSampler) = s.scale * randexp(rng)

struct ExponentialLogUSampler <: Sampleable{Univariate,Continuous}
    scale::AbstractFloat
end

rand(rng::AbstractRNG, s::ExponentialLogUSampler) = -s.scale * log(rand(rng))


using FFTW, Plots, SampledSignals

function plotwave(wave)
    t = 2*pi*(1:100)/100
    # Broadcast map to support plotting a list of functions
    plot(t, map.(wave, Ref(t)))
end

function plotfreqs!(plt, sound, fs; kargs...)
    fs = sound.samplerate
    len = length(sound)
    domain = fftfreq(len, fs) |> fftshift
    freqs = fft(sound) |> fftshift
    # shifted slightly to avoid trying to compute log10(0)
    decibels = 10*log10.((abs.(freqs) .+ 1)/len)
    plot!(plt, domain, decibels; xlimit=(0,fs/2) , ylimit=(min(decibels...), 0), kargs...)
end

function plotfreqs!(sound, fs; kargs...)
	plotfreqs!(Plots.current(), sound, fs; kargs...)
end

function plotfreqs(sound, fs; kargs...)
    plotfreqs!(plot(), sound, fs; kargs...)
end