@everywhere module TestDecisionObjective

using StochasticPrograms
using Test

function test_SingleDecision_objective(Structure)
    ξ₁ = @scenario a = 1. probability = 0.5
    ξ₂ = @scenario a = 2. probability = 0.5
    sp = StochasticProgram([ξ₁,ξ₂], Structure...)
    @first_stage sp = begin
        @decision(model, x)
        @objective(model, Min, x)
    end
    @second_stage sp = begin
        @recourse(model, y)
        @objective(model, Min, y)
    end
    # First-stage
    x = DecisionRef(sp[1,:x])
    @test MOI.MIN_SENSE == @inferred JuMP.objective_sense(sp, 1)
    @test JuMP.objective_function_type(sp, 1) == DecisionRef
    @test JuMP.objective_function(sp, 1) == x
    @test x == @inferred JuMP.objective_function(sp, 1, DecisionRef)
    # Second-stage
    y1 = DecisionRef(sp[2,:y], 1)
    y2 = DecisionRef(sp[2,:y], 2)
    @test MOI.MIN_SENSE == @inferred JuMP.objective_sense(sp, 2, 1)
    # Structure specific
    if sp.structure isa StochasticPrograms.DeterministicEquivalent
        @test JuMP.objective_function_type(sp, 2, 1) == DecisionRef
        @test JuMP.objective_function_type(sp, 2, 2) == DecisionRef
        @test JuMP.objective_function(sp, 2, 1) == y1
        @test JuMP.objective_function(sp, 2, 2) == y2
        @test y1 == @inferred JuMP.objective_function(sp, 2, 1, DecisionRef)
        @test y2 == @inferred JuMP.objective_function(sp, 2, 2, DecisionRef)
        @test JuMP.isequal_canonical(x + 0.5*y1 + 0.5*y2, JuMP.objective_function(sp))
        @test JuMP.isequal_canonical(x + 0.5*y1 + 0.5*y2,
                                     @inferred JuMP.objective_function(sp, DecisionAffExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.VerticalStructure
        @test JuMP.objective_function_type(sp, 2, 1) == DecisionRef
        @test JuMP.objective_function_type(sp, 2, 2) == DecisionRef
        @test JuMP.objective_function(sp, 2, 1) == y1
        @test JuMP.objective_function(sp, 2, 2) == y2
        @test y1 == @inferred JuMP.objective_function(sp, 2, 1, DecisionRef)
        @test y2 == @inferred JuMP.objective_function(sp, 2, 2, DecisionRef)
        @test JuMP.objective_function(sp) == x
        @test x == @inferred JuMP.objective_function(sp, DecisionRef)
    end
    if sp.structure isa StochasticPrograms.HorizontalStructure
        x1 = DecisionRef(sp[1,:x], 2, 1)
        x2 = DecisionRef(sp[1,:x], 2, 2)
        @test JuMP.objective_function_type(sp, 2, 1) == DecisionAffExpr{Float64}
        @test JuMP.objective_function_type(sp, 2, 2) == DecisionAffExpr{Float64}
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), x1 + y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), x2 + y2)
        @test JuMP.isequal_canonical(x1 + y1,
                                     @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(x2 + y2,
                                     @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        @test_throws ErrorException JuMP.objective_function(sp)
    end
    @first_stage sp = begin
        @decision(model, x)
        @objective(model, Max, x)
    end
    @second_stage sp = begin
        @recourse(model, y)
        @objective(model, Max, y)
    end
    # First-stage
    x = DecisionRef(sp[1,:x])
    @test MOI.MAX_SENSE == @inferred JuMP.objective_sense(sp, 1)
    @test JuMP.objective_function_type(sp, 1) == DecisionRef
    @test JuMP.objective_function(sp, 1) == x
    @test x == @inferred JuMP.objective_function(sp, 1, DecisionRef)
    # Second-stage
    y1 = DecisionRef(sp[2,:y], 1)
    y2 = DecisionRef(sp[2,:y], 2)
    @test MOI.MAX_SENSE == @inferred JuMP.objective_sense(sp, 2, 1)
    # Structure specific
    if sp.structure isa StochasticPrograms.DeterministicEquivalent
        @test JuMP.objective_function_type(sp, 2, 1) == DecisionRef
        @test JuMP.objective_function_type(sp, 2, 2) == DecisionRef
        @test JuMP.objective_function(sp, 2, 1) == y1
        @test JuMP.objective_function(sp, 2, 2) == y2
        @test y1 == @inferred JuMP.objective_function(sp, 2, 1, DecisionRef)
        @test y2 == @inferred JuMP.objective_function(sp, 2, 2, DecisionRef)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), x + 0.5*y1 + 0.5*y2)
        @test JuMP.isequal_canonical(x + 0.5*y1 + 0.5*y2,
                                     @inferred JuMP.objective_function(sp, DecisionAffExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.VerticalStructure
        @test JuMP.objective_function_type(sp, 2, 1) == DecisionRef
        @test JuMP.objective_function_type(sp, 2, 2) == DecisionRef
        @test JuMP.objective_function(sp, 2, 1) == y1
        @test JuMP.objective_function(sp, 2, 1) == y2
        @test y1 == @inferred JuMP.objective_function(sp, 2, 1, DecisionRef)
        @test y2 == @inferred JuMP.objective_function(sp, 2, 1, DecisionRef)
        @test JuMP.objective_function(sp) == x
        @test x == @inferred JuMP.objective_function(sp, DecisionRef)
    end
    if sp.structure isa StochasticPrograms.HorizontalStructure
        x1 = DecisionRef(sp[1,:x], 2, 1)
        x2 = DecisionRef(sp[1,:x], 2, 2)
        @test JuMP.objective_function_type(sp, 2, 1) == DecisionAffExpr{Float64}
        @test JuMP.objective_function_type(sp, 2, 2) == DecisionAffExpr{Float64}
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), x1 + y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), x2 + y2)
        @test JuMP.isequal_canonical(x1 + y1,
                                     @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(x2 + y2,
                                     @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        @test_throws ErrorException JuMP.objective_function(sp)
    end
end

function test_DecisionAffExpr_objective(Structure)
    ξ₁ = @scenario a = 2. probability = 0.5
    ξ₂ = @scenario a = 4 probability = 0.5
    sp = StochasticProgram([ξ₁,ξ₂], Structure...)
    @first_stage sp = begin
        @decision(model, x)
        @objective(model, Min, 2x)
    end
    @second_stage sp = begin
        @uncertain a
        @recourse(model, y)
        @objective(model, Min, a*y)
    end
    # First-stage
    x = DecisionRef(sp[1,:x])
    @test MOI.MIN_SENSE == @inferred JuMP.objective_sense(sp, 1)
    @test JuMP.objective_function_type(sp, 1) == DecisionAffExpr{Float64}
    @test JuMP.isequal_canonical(JuMP.objective_function(sp, 1), 2x)
    @test JuMP.isequal_canonical(
        2x, @inferred JuMP.objective_function(sp, 1, DecisionAffExpr{Float64}))
    # Second-stage
    y1 = DecisionRef(sp[2,:y], 1)
    y2 = DecisionRef(sp[2,:y], 2)
    @test MOI.MIN_SENSE == @inferred JuMP.objective_sense(sp, 2, 1)
    @test JuMP.objective_function_type(sp, 2, 1) == DecisionAffExpr{Float64}
    @test JuMP.objective_function_type(sp, 2, 2) == DecisionAffExpr{Float64}
    # Structure specific
    if sp.structure isa StochasticPrograms.DeterministicEquivalent
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 2y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 4y2)
        @test JuMP.isequal_canonical(
            2y1, @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(
            4y2, @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), 2x + y1 + 2y2)
        @test JuMP.isequal_canonical(2x + y1 + 2y2,
                                     @inferred JuMP.objective_function(sp, DecisionAffExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.VerticalStructure
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 2y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 4y2)
        @test JuMP.isequal_canonical(
            2y1, @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(
            4y2, @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), 2x)
        @test JuMP.isequal_canonical(2x,
                                     @inferred JuMP.objective_function(sp, DecisionAffExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.HorizontalStructure
        x1 = DecisionRef(sp[1,:x], 2, 1)
        x2 = DecisionRef(sp[1,:x], 2, 2)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 2x1 + 2y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 2x2 + 4y2)
        @test JuMP.isequal_canonical(
            2x1 + 2y1, @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(
            2x2 + 4y2, @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        @test_throws ErrorException JuMP.objective_function(sp)
    end
    @first_stage sp = begin
        @decision(model, x)
        @objective(model, Max, x + 3x + 1)
    end
    @second_stage sp = begin
        @known x
        @uncertain a
        @recourse(model, y)
        @objective(model, Max, x + a*x + 1 + y + 3y + 1)
    end
    # First-stage
    x = DecisionRef(sp[1,:x])
    @test MOI.MAX_SENSE == @inferred JuMP.objective_sense(sp, 1)
    @test JuMP.objective_function_type(sp, 1) == DecisionAffExpr{Float64}
    @test JuMP.isequal_canonical(JuMP.objective_function(sp, 1), 4x + 1)
    @test JuMP.isequal_canonical(
        4x + 1, @inferred JuMP.objective_function(sp, 1, DecisionAffExpr{Float64}))
    # Second-stage
    y1 = DecisionRef(sp[2,:y], 1)
    y2 = DecisionRef(sp[2,:y], 2)
    @test MOI.MAX_SENSE == @inferred JuMP.objective_sense(sp, 2, 1)
    @test JuMP.objective_function_type(sp, 2, 1) == DecisionAffExpr{Float64}
    @test JuMP.objective_function_type(sp, 2, 2) == DecisionAffExpr{Float64}
    if sp.structure isa StochasticPrograms.DeterministicEquivalent
        x = DecisionRef(sp[1,:x], 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 3x + 4y1 + 2)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 5x + 4y2 + 2)
        @test JuMP.isequal_canonical(
            3x + 4y1 + 2, @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(
            5x + 4y2 + 2, @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        x = DecisionRef(sp[1,:x])
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), 8x + 2y1 + 2y2 + 3)
        @test JuMP.isequal_canonical(8x + 2y1 + 2y2 + 3,
                                     @inferred JuMP.objective_function(sp, DecisionAffExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.VerticalStructure
        x = DecisionRef(sp[1,:x], 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 3x + 4y1 + 2)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 5x + 4y2 + 2)
        @test JuMP.isequal_canonical(
            3x + 4y1 + 2, @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(
            5x + 4y2 + 2, @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        x = DecisionRef(sp[1,:x])
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), 4x + 1)
        @test JuMP.isequal_canonical(4x + 1,
                                     @inferred JuMP.objective_function(sp, DecisionAffExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.HorizontalStructure
        x1 = DecisionRef(sp[1,:x], 2, 1)
        x2 = DecisionRef(sp[1,:x], 2, 2)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 7x1 + 4y1 + 3)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 9x2 + 4y2 + 3)
        @test JuMP.isequal_canonical(
            7x1 + 4y1 + 3, @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(
            9x2 + 4y2 + 3, @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        @test_throws ErrorException JuMP.objective_function(sp)
    end
end

function test_DecisionQuadExpr_objective(Structure)
    ξ₁ = @scenario a = 2. probability = 0.5
    ξ₂ = @scenario a = 4 probability = 0.5
    sp = StochasticProgram([ξ₁,ξ₂], Structure...)
    @first_stage sp = begin
        @decision(model, x)
        @objective(model, Min, x^2 + 2x)
    end
    @second_stage sp = begin
        @known x
        @uncertain a
        @recourse(model, y)
        @objective(model, Min, y^2 + a*y + x^2 + 2x)
    end
    # First-stage
    x = DecisionRef(sp[1,:x])
    @test MOI.MIN_SENSE == @inferred JuMP.objective_sense(sp, 1)
    @test JuMP.objective_function_type(sp, 1) == DecisionQuadExpr{Float64}
    @test JuMP.isequal_canonical(JuMP.objective_function(sp, 1), x^2 + 2x)
    @test JuMP.isequal_canonical(
        x^2 + 2x, @inferred JuMP.objective_function(sp, 1, DecisionQuadExpr{Float64}))
    # Second-stage
    y1 = DecisionRef(sp[2,:y], 1)
    y2 = DecisionRef(sp[2,:y], 2)
    @test MOI.MIN_SENSE == @inferred JuMP.objective_sense(sp, 2, 1)
    @test JuMP.objective_function_type(sp, 2, 1) == DecisionQuadExpr{Float64}
    # Structure specific
    x = DecisionRef(sp[1,:x])
    if sp.structure isa StochasticPrograms.DeterministicEquivalent
        x = DecisionRef(sp[1,:x], 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), y1^2 + 2y1 + x^2 + 2x)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), y2^2 + 4y2 + x^2 + 2x)
        @test JuMP.isequal_canonical(
            y1^2 + 2y1 + x^2 + 2x, @inferred JuMP.objective_function(sp, 2, 1, DecisionQuadExpr{Float64}))
        @test JuMP.isequal_canonical(
            y2^2 + 4y2 + x^2 + 2x, @inferred JuMP.objective_function(sp, 2, 2, DecisionQuadExpr{Float64}))
        x = DecisionRef(sp[1,:x])
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), 2x^2 + 4x + 0.5*y1^2 + y1 + 0.5*y2^2 + 2y2)
        @test JuMP.isequal_canonical(2x^2 + 4x + 0.5*y1^2 + y1 + 0.5*y2^2 + 2y2,
                                     @inferred JuMP.objective_function(sp, DecisionQuadExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.VerticalStructure
        x = DecisionRef(sp[1,:x], 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), y1^2 + 2y1 + x^2 + 2x)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), y2^2 + 4y2 + x^2 + 2x)
        @test JuMP.isequal_canonical(
            y1^2 + 2y1 + x^2 + 2x, @inferred JuMP.objective_function(sp, 2, 1, DecisionQuadExpr{Float64}))
        @test JuMP.isequal_canonical(
            y2^2 + 4y2 + x^2 + 2x, @inferred JuMP.objective_function(sp, 2, 2, DecisionQuadExpr{Float64}))
        x = DecisionRef(sp[1,:x])
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), x^2 + 2x)
        @test JuMP.isequal_canonical(x^2 + 2x,
                                     @inferred JuMP.objective_function(sp, DecisionQuadExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.HorizontalStructure
        x1 = DecisionRef(sp[1,:x], 2, 1)
        x2 = DecisionRef(sp[1,:x], 2, 2)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 2x1^2 + 4x1 + y1^2 + 2y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 2x2^2 + 4x2 + y2^2 + 4y2)
        @test JuMP.isequal_canonical(
            2x1^2 + 4x1 + y1^2 + 2y1, @inferred JuMP.objective_function(sp, 2, 1, DecisionQuadExpr{Float64}))
        @test JuMP.isequal_canonical(
            2x2^2 + 4x2 + y2^2 + 4y2, @inferred JuMP.objective_function(sp, 2, 2, DecisionQuadExpr{Float64}))
        @test_throws ErrorException JuMP.objective_function(sp)
    end
end

function test_decision_objective_coefficient_modifiction(Structure)
    ξ₁ = @scenario a = 2. probability = 0.5
    ξ₂ = @scenario a = 4 probability = 0.5
    sp = StochasticProgram([ξ₁,ξ₂], Structure...)
    @first_stage sp = begin
        @decision(model, x)
        @objective(model, Min, x)
    end
    @second_stage sp = begin
        @known x
        @recourse(model, y)
        @objective(model, Min, y)
    end
    # First-stage
    x = sp[1,:x]
    JuMP.set_objective_coefficient(sp, x, 1, 4.0)
    x = DecisionRef(x)
    @test JuMP.isequal_canonical(JuMP.objective_function(sp, 1), 4x)
    # Second-stage
    x = sp[1,:x]
    y = sp[2,:y]
    JuMP.set_objective_coefficient(sp, y, 2, 1, 4.0)
    JuMP.set_objective_coefficient(sp, y, 2, 2, 4.0)
    y1 = DecisionRef(y, 1)
    y2 = DecisionRef(y, 2)
    # Structure specific
    if sp.structure isa StochasticPrograms.DeterministicEquivalent
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 4y2)
        JuMP.set_objective_coefficient(sp, x, 2, 1, 2.0)
        JuMP.set_objective_coefficient(sp, x, 2, 2, 4.0)
        x = DecisionRef(sp[1,:x], 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 2x + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 4x + 4y2)
        x = DecisionRef(sp[1,:x])
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), 7x + 2y1 + 2y2)
        @test JuMP.isequal_canonical(7x + 2y1 + 2y2,
                                     @inferred JuMP.objective_function(sp, DecisionAffExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.VerticalStructure
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 4y2)
        JuMP.set_objective_coefficient(sp, x, 2, 1, 2.0)
        JuMP.set_objective_coefficient(sp, x, 2, 2, 4.0)
        x = DecisionRef(sp[1,:x], 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 2x + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 4x + 4y2)
        x = DecisionRef(sp[1,:x])
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), 4x)
        @test JuMP.isequal_canonical(4x,
                                     @inferred JuMP.objective_function(sp, DecisionAffExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.HorizontalStructure
        x = DecisionRef(sp[1,:x], 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 4x + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 4x + 4y2)
        @test JuMP.isequal_canonical(4x + 4y1,
                                     @inferred JuMP.objective_function(sp, 2, 1, DecisionAffExpr{Float64}))
        @test JuMP.isequal_canonical(4x + 4y2,
                                     @inferred JuMP.objective_function(sp, 2, 2, DecisionAffExpr{Float64}))
        @test_throws ErrorException JuMP.objective_function(sp)
    end
    @first_stage sp = begin
        @decision(model, x)
        @objective(model, Min, x^2 + x)
    end
    @second_stage sp = begin
        @known x
        @recourse(model, y)
        @objective(model, Min, x^2 + y^2 + y)
    end
    # First-stage
    x = sp[1,:x]
    JuMP.set_objective_coefficient(sp, x, 1, 4.0)
    x = DecisionRef(x)
    @test JuMP.isequal_canonical(JuMP.objective_function(sp, 1), x^2 + 4x)
    # Second-stage
    x = sp[1,:x]
    y = sp[2,:y]
    JuMP.set_objective_coefficient(sp, y, 2, 1, 4.0)
    JuMP.set_objective_coefficient(sp, y, 2, 2, 4.0)
    # Structure specific
    if sp.structure isa StochasticPrograms.DeterministicEquivalent
        x = DecisionRef(x, 2, 1)
        y1 = DecisionRef(y, 1)
        y2 = DecisionRef(y, 2)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), x^2 + y1^2 + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), x^2 + y2^2 + 4y2)
        x = sp[1,:x]
        JuMP.set_objective_coefficient(sp, x, 2, 1, 2.0)
        JuMP.set_objective_coefficient(sp, x, 2, 2, 4.0)
        x = DecisionRef(x, 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), x^2 + y1^2 + 2x + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), x^2 + y2^2 + 4x + 4y2)
        x = DecisionRef(sp[1,:x])
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), 2x^2 + 7x + 0.5*y1^2 + 2y1 + 0.5*y2^2 + 2y2)
        @test JuMP.isequal_canonical(2x^2 + 7x + 0.5*y1^2 + 2y1 + 0.5*y2^2 + 2y2,
                                     @inferred JuMP.objective_function(sp, DecisionQuadExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.VerticalStructure
        x = DecisionRef(x, 2, 1)
        y1 = DecisionRef(y, 1)
        y2 = DecisionRef(y, 2)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), x^2 + y1^2 + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), x^2 + y2^2 + 4y2)
        x = sp[1,:x]
        JuMP.set_objective_coefficient(sp, x, 2, 1, 2.0)
        JuMP.set_objective_coefficient(sp, x, 2, 2, 4.0)
        x = DecisionRef(x, 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), x^2 + y1^2 + 2x + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), x^2 + y2^2 + 4x + 4y2)
        x = DecisionRef(sp[1,:x])
        @test JuMP.isequal_canonical(JuMP.objective_function(sp), x^2 + 4x)
        @test JuMP.isequal_canonical(x^2 + 4x,
                                     @inferred JuMP.objective_function(sp, DecisionQuadExpr{Float64}))
    end
    if sp.structure isa StochasticPrograms.HorizontalStructure
        x = DecisionRef(x, 2, 1)
        y1 = DecisionRef(y, 1)
        y2 = DecisionRef(y, 2)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 2x^2 + y1^2 + 4x + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 2x^2 + y2^2 + 4x + 4y2)
        @test JuMP.isequal_canonical(2x^2 + y1^2 + 4x + 4y1,
                                     @inferred JuMP.objective_function(sp, 2, 1, DecisionQuadExpr{Float64}))
        @test JuMP.isequal_canonical(2x^2 + y2^2 + 4x + 4y2,
                                     @inferred JuMP.objective_function(sp, 2, 2, DecisionQuadExpr{Float64}))
        x = sp[1,:x]
        JuMP.set_objective_coefficient(sp, x, 2, 1, 2.0)
        JuMP.set_objective_coefficient(sp, x, 2, 2, 4.0)
        x = DecisionRef(x, 2, 1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 1), 2x^2 + y1^2 + 2x + 4y1)
        @test JuMP.isequal_canonical(JuMP.objective_function(sp, 2, 2), 2x^2 + y2^2 + 4x + 4y2)
        @test JuMP.isequal_canonical(2x^2 + y1^2 + 2x + 4y1,
                                     @inferred JuMP.objective_function(sp, 2, 1, DecisionQuadExpr{Float64}))
        @test JuMP.isequal_canonical(2x^2 + y2^2 + 4x + 4y2,
                                     @inferred JuMP.objective_function(sp, 2, 2, DecisionQuadExpr{Float64}))
        @test_throws ErrorException JuMP.objective_function(sp)
    end
end

function runtests()
    @testset "DecisionObjective" begin
        for structure in [(Deterministic(),),
                          (Vertical(),),
                          (Horizontal(),),
                          (Deterministic(), () -> MOIU.MockOptimizer(MOIU.Model{Float64}()))]
            name = length(structure) == 1 ? "$(structure[1])" : "$(structure[1]) with decision bridges"
            @testset "$name" begin
                for name in names(@__MODULE__; all = true)
                    if !startswith("$(name)", "test_")
                        continue
                    end
                    f = getfield(@__MODULE__, name)
                    @testset "$(name)" begin
                        f(structure)
                    end
                end
            end
        end
    end
end

function run_dtests()
    @testset "DecisionObjective" begin
        for structure in [(DistributedVertical(),), (DistributedHorizontal(),)]
            @testset "$(structure[1])" begin
                for name in names(@__MODULE__; all = true)
                    if !startswith("$(name)", "test_")
                        continue
                    end
                    f = getfield(@__MODULE__, name)
                    @testset "$(name)" begin
                        f(structure)
                    end
                end
            end
        end
    end
end

end


using Flux
using Flux: onehotbatch, onecold, crossentropy
using JLD2, FileIO
using Statistics: mean

include("modulos/dataset_DL.jl");
include("modulos/graphics.jl");

#===============================================================================
COSAS:
    · en toFloatArray no me deja hacer un assert de una comp de ints:  MethodError: objects of type Int64 are not callable
    · ¿por qué filtros de 3x3?
    · ¿Están bien las capas convolucionales?

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

interval = 8;
(train_imgs, train_labels,
    test_imgs, test_labels) = getInputs("datasets", interval);

#=
for image in train_imgs
    display(image);
end;
for image in test_imgs
    display(image);
end;
=#

train_imgs = toFloatArray(train_imgs);
test_imgs = toFloatArray(test_imgs);

println("Tamaño de la matriz de entrenamiento: ", size(train_imgs))
println("Tamaño de la matriz de test:          ", size(test_imgs))

batch_size = 64;
gruposIndicesBatch = Iterators.partition(1:size(train_imgs,4), batch_size);
println("He creado ", length(gruposIndicesBatch), " grupos de indices para distribuir los patrones en batches");

train_set = [( train_imgs[:,:,:,indicesBatch], Array(onehotbatch(train_labels[indicesBatch], 0:2)) ) for indicesBatch in gruposIndicesBatch]

test_set = (test_imgs, onehotbatch(test_labels, 0:2));

train_imgs = nothing;
test_imgs = nothing;
GC.gc();

funcionTransferenciaCapasConvolucionales = relu;

# Definimos la red con la funcion Chain, que concatena distintas capas

modelo = Chain(
    Conv((5, 5), 3=>16, pad=(1,1), funcionTransferenciaCapasConvolucionales),
    MaxPool((2,2)),
    Conv((5, 5), 16=>32, pad=(1,1), funcionTransferenciaCapasConvolucionales),
    MaxPool((2,2)),
    Conv((5, 5), 32=>32, pad=(1,1), funcionTransferenciaCapasConvolucionales),
    MaxPool((2,2)),
    x -> reshape(x, :, size(x, 4)),
    Dense(9248, 3),
    softmax
);

# Vamos a probar la RNA y poner algunos datos de cada capa
# Usaremos como entrada varios patrones de un batch
numBatchCoger = 1; numImagenEnEseBatch = [12, 6];

entradaCapa = train_set[numBatchCoger][1][:,:,:,numImagenEnEseBatch];
numCapas = length(params(modelo));
println("La RNA tiene ", numCapas, " capas:");
for numCapa in 1:numCapas
    println("   Capa ", numCapa, ": ", modelo[numCapa]);
    # Le pasamos la entrada a esta capa
    global entradaCapa # Esta linea es necesaria porque la variable entradaCapa es global y se modifica en este bucle
    capa = modelo[numCapa];
    salidaCapa = capa(entradaCapa);
    println("      La salida de esta capa tiene dimension ", size(salidaCapa));
    entradaCapa = salidaCapa;
end;

modelo(train_set[numBatchCoger][1][:,:,:,numImagenEnEseBatch]);

loss(x, y) = crossentropy(modelo(x), y)
accuracy(batch) = mean(onecold(modelo(batch[1])) .== onecold(batch[2]))

function calculate_loss(array_of_tuples)
    return loss(array_of_tuples[1], array_of_tuples[2]);
end;

println("Ciclo 0: Precision en el conjunto de entrenamiento: ", 100*mean(accuracy.(train_set)), " %");

opt = ADAM(0.001);

println("Comenzando entrenamiento...")
mejorPrecision = -Inf;
criterioFin = false;
numCiclo = 0;
numCicloUltimaMejora = 0;
mejorModelo = nothing;
num_ciclos = [];
trainingLosses = [];
testLosses = [];
trainingAccuracies = [];
testAccuracies = [];

while (!criterioFin)

    global numCicloUltimaMejora, numCiclo, mejorPrecision, mejorModelo,
        criterioFin, num_ciclos, trainingLosses, testLosses,
        trainingAccuracies, testAccuracies;

    Flux.train!(loss, params(modelo), train_set, opt);

    numCiclo += 1;

    precisionEntrenamiento = mean(accuracy.(train_set));
    println("Ciclo ", numCiclo, ": Precision en el conjunto de entrenamiento: ", 100*precisionEntrenamiento, " %");
    precisionTest = accuracy(test_set);
    println("   Precision en el conjunto de test: ", 100*precisionTest, " %");

    # Calculamos las métricas
    push!(num_ciclos, numCiclo);
    push!(trainingAccuracies, precisionEntrenamiento);
    push!(testAccuracies, precisionTest);
    push!(trainingLosses, mean(calculate_loss.(train_set)));
    push!(testLosses, mean(calculate_loss(test_set)));

    if (precisionEntrenamiento > mejorPrecision)
        mejorPrecision = precisionEntrenamiento;
        println("   Mejora en el conjunto de entrenamiento -> Precision en el conjunto de test: ", 100*precisionTest, " %");
        mejorModelo = deepcopy(modelo);
        numCicloUltimaMejora = numCiclo;
    end;
#=
    if (numCiclo - numCicloUltimaMejora >= 5) && (opt.eta > 1e-6)
        opt.eta /= 10.0
        println("   No se ha mejorado en 5 ciclos, se baja la tasa de aprendizaje a ", opt.eta);
        numCicloUltimaMejora = numCiclo;
    end;
=#
    if (precisionEntrenamiento >= 0.999)
        println("   Se para el entenamiento por haber llegado a una precision de 99.9%")
        criterioFin = true;
    end;

    if (numCiclo - numCicloUltimaMejora >= 10)
        println("   Se para el entrenamiento por no haber mejorado la precision en el conjunto de entrenamiento durante 10 ciclos")
        criterioFin = true;
    end;
end;

print_train_results(num_ciclos, trainingLosses, testLosses, trainingAccuracies, testAccuracies);


using MultiStochGrad, Test

using Logging
using Base.CoreLogging

logger = ConsoleLogger(stdout, CoreLogging.Info)
global_logger(logger)

include("TestLinearRegression.jl")

@testset "linreg" begin
@test test_linear_regression_scsg()
@test test_linear_regression_svrg()
end

# possibly  after download of mnist files
include("TestLogisticRegression.jl")

@testset "logistic_reg" begin
@test mnist_logistic_regression_svrg()
@test mnist_logistic_regression_scsg()
end


import StatsBase

pointestimate_funcs = [
    StatsBase.var,
    StatsBase.std,
    StatsBase.middle,
    StatsBase.median,
    StatsBase.mean,
    StatsBase.genmean,
    StatsBase.genvar,
    StatsBase.harmmean,
    StatsBase.mode,
    StatsBase.percentile,
    StatsBase.quantile,
    StatsBase.rle,
    StatsBase.sem,
    StatsBase.span,
    StatsBase.summarystats,
    StatsBase.totalvar,
    StatsBase.kurtosis,
    StatsBase.moment,
    StatsBase.skewness,
    StatsBase.renyientropy
]

n = 10

@testset "Point-estimate statistics" begin
    @testset "point-estimate statistic: $(pointestimate_funcs[i])" for i = 1:length(pointestimate_funcs)
        f = pointestimate_funcs[i]
        @testset for (i, uval) in enumerate(example_uvals)
            if f == StatsBase.summarystats
                @test f(uval, n) isa StatsBase.SummaryStats{T} where T
            elseif f == StatsBase.percentile
                @test f(uval, 10, n) isa T where T <: Real
                @test f(uval, [10, 20], n) isa Vector{T} where T <: Real
            elseif f == StatsBase.quantile
                @test f(uval, 0.1, n) isa T where T <: Real
                @test f(uval, [0.2, 0.5], n) isa Vector{T} where T <: Real
            elseif f == StatsBase.moment
                @test f(uval, 1, n) isa T where T <: Real
                @test f(uval, 2, n) isa T where T <: Real
            elseif f == StatsBase.genmean
                @test f(uval, 4, n) isa T where T <: Real
            elseif f == StatsBase.rle
                @test rle(uval, n) isa Tuple{Vector{T} where T, Vector{T2} where T2}
            elseif f == StatsBase.span
                @test f(uval, n) isa AbstractRange{T} where T <: Real
            elseif f == StatsBase.renyientropy
                @test f(uval, 1, n) isa T where T <: Real
            else
                @test f(uval, n) isa T where T <: Real
            end
        end
    end;
end


@testset "PLL discriminator" begin

    correlator_minus60off = EarlyPromptLateCorrelator(
        -0.5 + sqrt(3) / 2im,
        -1 + sqrt(3) * 1im,
        -0.5 + sqrt(3) / 2im
    )
    correlator_0off = EarlyPromptLateCorrelator(
        0.5 + 0.0im,
        1 + 0.0im,
        0.5 + 0.0im
    )
    correlator_plus60off = EarlyPromptLateCorrelator(
        0.5 + sqrt(3) / 2im,
        1 + sqrt(3) * 1im,
        0.5 + sqrt(3) / 2im
    )

    @test @inferred(Tracking.pll_disc(GPSL1, correlator_minus60off)) == -π / 3  #-60°
    @test @inferred(Tracking.pll_disc(GPSL1, correlator_0off)) == 0
    @test @inferred(Tracking.pll_disc(GPSL1, correlator_plus60off)) == π / 3  #+60°
end

@testset "DLL discriminator" begin

    very_early_correlator = EarlyPromptLateCorrelator(0.0 + 0.0im, 0.5 + 0.0im, 1.0 + 0.0im)
    early_correlator = EarlyPromptLateCorrelator(0.25 + 0.0im, 0.75 + 0.0im, 0.75 + 0.0im)
    prompt_correlator = EarlyPromptLateCorrelator(0.5 + 0.0im, 1 + 0.0im, 0.5 + 0.0im)
    late_correlator = EarlyPromptLateCorrelator(0.75 + 0.0im, 0.75 + 0.0im, 0.25 + 0.0im)
    very_late_correlator = EarlyPromptLateCorrelator(1.0 + 0.0im, 0.5 + 0.0im, 0.0 + 0.0im)

    sample_shift = 2
    delta = 0.25

    @test @inferred(
        Tracking.dll_disc(GPSL1, very_early_correlator, sample_shift, delta)
    ) == -0.5
    @test @inferred(
        Tracking.dll_disc(GPSL1, early_correlator, sample_shift, delta)
    ) == -0.25
    @test @inferred(
        Tracking.dll_disc(GPSL1, prompt_correlator, sample_shift, delta)
    ) == 0
    @test @inferred(
        Tracking.dll_disc(GPSL1, late_correlator, sample_shift, delta)
    ) == 0.25
    @test @inferred(
        Tracking.dll_disc(GPSL1, very_late_correlator, sample_shift, delta)
    ) == 0.5
end


###################################
## Example code
## for Bayesian Poisson HMM

using PyPlot, PyCall
using HDF5, JLD
@pyimport matplotlib.gridspec as gspec

push!(LOAD_PATH,".")
import PoissonHMM
import PoissonMixtureModel

"""
Simple comparison between HMM and mixture model.
"""
function test_comparison()
    #########################
    ## load data
    file_name = "../data/timeseries.jld"
    X = load(file_name)["obs"]
    N = length(X)

    #########################
    ## Poison HMM

    ## set model
    K = 2 #  number of mixture components
    alpha_phi = 10.0 * ones(K)
    alpha_A = 100.0 * eye(K) + 1.0*ones(K, K)
    cmp = [PoissonHMM.Gam(1.0, 0.01), PoissonHMM.Gam(1.0, 0.01)]
    bhmm = PoissonHMM.BHMM(K, alpha_phi, alpha_A, cmp)

    ## inference
    max_iter = 100
    tic()
    Z_est_hmm, post_bhmm = PoissonHMM.learn_VI(X, bhmm, max_iter)
    toc()

    #########################
    ## Poison Mixture Model

    ## set model
    K = 2 #  number of mixture components
    alpha_phi = 10.0 * ones(K)
    cmp = [PoissonMixtureModel.Gam([1.0], 0.01), PoissonMixtureModel.Gam([1.0], 0.01)]
    bpmm = PoissonMixtureModel.BPMM(1, K, alpha_phi, cmp)

    ## inference
    max_iter = 100
    tic()
    Z_est_pmm, post_bpmm = PoissonMixtureModel.learn_VI(reshape(X, 1, N), bpmm, max_iter)
    toc()

    #########################
    ## Compare results
    figure("Hidden Markov Model vs Mixture Model")
    subplot(3,1,1);plot(X);ylabel("data")
    subplot(3,1,2);fill_between(1:N, reshape(Z_est_hmm[1,:]', N), zeros(N));ylim([0.0, 1.0]);ylabel("S (PHMM)")
    subplot(3,1,3);fill_between(1:N, reshape(Z_est_pmm[1,:]', N), zeros(N));ylim([0.0, 1.0]);ylabel("S (PMM)")
    show()
end

test_comparison()


using LinearAlgebra
"""
ExKIObj{FT<:AbstractFloat, IT<:Int}
Extended Kalman Inversion struct (ExKI)
For solving the inverse problem
    y = G(θ) + η

"""

mutable struct ExKIObj{FT<:AbstractFloat, IT<:Int}
     method::String
    "a vector of arrays of size  N_parameters containing the mean of the parameters (in each exki iteration a new array of mean is added)"
     θ_mean::Vector{Array{FT, 1}}
     "a vector of arrays of size (N_parameters x N_parameters) containing the covariance of the parameters (in each exki iteration a new array of cov is added)"
     θθ_cov::Vector{Array{FT, 2}}
     "a vector of arrays of size N_y containing the predicted observation (in each exki iteration a new array of predicted observation is added)"
     y_pred::Vector{Array{FT, 1}}
     "vector of observations (length: N_y)"
     y::Array{FT, 1}
     "covariance of the observational noise"
     Σ_η::Array{FT, 2}
     "size of θ"
     N_θ::IT
     "size of y"
     N_y::IT
     "prior mean"
     r_0::Array{FT, 1}
     "prior covariance"
     Σ_0::Array{FT, 2}
     "hyperparameter"
     γ::FT
     "current iteration number"
     iter::IT
     N_iter::IT
end

"""
ExKIObj Constructor
N_iter: number of iterations
y::Array{FT,1} : observation
Σ_η::Array{FT, 2} : observation error covariance
γ::FT : hyperparameter
r_0::Array{FT, 1} : prior mean
Σ_0::Array{FT, 2} : prior covariance
θ0_mean::Array{FT, 1} : initial mean
θθ0_cov::Array{FT, 2} : initial covariance
"""
function ExKIObj(method::String,
                 N_iter::IT,
                 y::Array{FT,1},
                 Σ_η::Array{FT, 2},
                 γ::FT,
                 r_0::Array{FT, 1},
                 Σ_0::Array{FT, 2},
                 θ0_mean::Array{FT,1} = r_0,
                 θθ0_cov::Array{FT, 2} = Σ_0) where {FT<:AbstractFloat, IT<:Int}

    N_θ = size(θ0_mean,1)
    N_y = size(y, 1)

    θ_mean = Array{FT,1}[]  # array of Array{FT, 2}'s
    push!(θ_mean, θ0_mean) # insert parameters at end of array (in this case just 1st entry)
    θθ_cov = Array{FT,2}[] # array of Array{FT, 2}'s
    push!(θθ_cov, θθ0_cov) # insert parameters at end of array (in this case just 1st entry)

    y_pred = Array{FT, 1}[]  # array of Array{FT, 2}'s

    iter = IT(0)

    ExKIObj{FT,IT}(method,
                  θ_mean, θθ_cov, y_pred,
                  y,   Σ_η,
                  N_θ, N_y,
                  r_0, Σ_0,
                  γ, iter, N_iter)

end

"""
update exki struct
forward(θ) = G(θ) dG(θ)
"""
function update_ensemble!(exki::ExKIObj{FT, IT}, forward::Function) where {FT<:AbstractFloat, IT<:Int}

    exki.iter += 1

    θ_mean  = exki.θ_mean[end]
    θθ_cov = exki.θθ_cov[end]
    y = exki.y
    γ = exki.γ
    r_0, Σ_0 = exki.r_0, exki.Σ_0
    N_θ, N_y = exki.N_θ, exki.N_y
    ############# Prediction step:

    if exki.method == "ExKI"
        θ_p_mean  = θ_mean
        θθ_p_cov =  (γ + 1)*θθ_cov

        g_mean, dg = forward(θ_p_mean)
        θg_cov = θθ_p_cov * dg'

        # extended system
        ff_cov = [ dg*θg_cov+(1 + 1/γ)*Σ_η  θg_cov';
                    θg_cov     θθ_p_cov+(1 + 1/γ)*Σ_0]

        θf_cov = [θg_cov   θθ_p_cov]

        tmp = θf_cov/ff_cov

        θ_mean =  θ_p_mean + tmp*([y ; r_0] - [g_mean; θ_p_mean])

        θθ_cov =  θθ_p_cov - tmp*θf_cov'
    elseif exki.method == "ExKF"
        Δt = 1.0/exki.N_iter

        θ_p_mean  = θ_mean
        θθ_p_cov =   θθ_cov

        g_mean, dg = forward(θ_p_mean)
        θg_cov = θθ_p_cov * dg'

        # extended system

        gg_cov = dg*θg_cov + Σ_η/Δt

        tmp = θg_cov/gg_cov

        θ_mean =  θ_p_mean + tmp*(y - g_mean)

        θθ_cov =  θθ_p_cov - tmp*θg_cov'
    end

    ########### Save resutls
    push!(exki.y_pred, g_mean) # N_ens x N_data
    push!(exki.θ_mean, θ_mean) # N_ens x N_params
    push!(exki.θθ_cov, θθ_cov) # N_ens x N_data
end

function ExKI_Run(s_param, forward::Function,
    method,
    y, Σ_η,
    r_0,
    Σ_0,
    θ0_mean, θθ0_cov,
    N_iter,
    γ = 1.0)

    exkiobj = ExKIObj(
        method,
        N_iter,
        y,
        Σ_η,
        γ,
        r_0,
        Σ_0,
        θ0_mean,
        θθ0_cov)

    func(θ) = forward(s_param, θ)

    for i in 1:N_iter

        update_ensemble!(exkiobj, func)

        @info "optimization error at iter $(i) = ", 0.5*(exkiobj.y_pred[i] - exkiobj.y)'*(exkiobj.Σ_η\(exkiobj.y_pred[i] - exkiobj.y))
        @info "Frobenius norm of the covariance at iter $(i) = ", norm(exkiobj.θθ_cov[i])
    end

    return exkiobj

end

function plot_param_iter(exkiobj::ExKIObj{FT, IT}, θ_ref::Array{FT,1}, θ_ref_names::Array{String}) where {FT<:AbstractFloat, IT<:Int}

    θ_mean = exkiobj.θ_mean
    θθ_cov = exkiobj.θθ_cov

    N_iter = length(θ_mean) - 1
    ites = Array(LinRange(1, N_iter+1, N_iter+1))

    θ_mean_arr = abs.(hcat(θ_mean...))

    N_θ = length(θ_ref)
    θθ_std_arr = zeros(Float64, (N_θ, N_iter+1))
    for i = 1:N_iter+1
        for j = 1:N_θ
            θθ_std_arr[j, i] = sqrt(θθ_cov[i][j,j])
        end
    end

    for i = 1:N_θ
        errorbar(ites, θ_mean_arr[i,:], yerr=3.0*θθ_std_arr[i,:], fmt="--o",fillstyle="none", label=θ_ref_names[i])
        plot(ites, fill(θ_ref[i], N_iter+1), "--", color="gray")
    end

    xlabel("Iterations")
    legend()
    tight_layout()
end

function plot_opt_errors(exkiobj::ExKIObj{FT, IT},
    θ_ref::Union{Array{FT,1}, Nothing} = nothing,
    transform_func::Union{Function, Nothing} = nothing) where {FT<:AbstractFloat, IT<:Int}

    θ_mean = exkiobj.θ_mean
    θθ_cov = exkiobj.θθ_cov
    y_pred = exkiobj.y_pred
    Σ_η = exkiobj.Σ_η
    y = exkiobj.y

    N_iter = length(θ_mean) - 1

    ites = Array(LinRange(1, N_iter, N_iter))
    N_subfigs = (θ_ref === nothing ? 2 : 3)

    errors = zeros(Float64, N_subfigs, N_iter)
    fig, ax = PyPlot.subplots(ncols=N_subfigs, figsize=(N_subfigs*6,6))
    for i = 1:N_iter
        errors[N_subfigs - 1, i] = 0.5*(y - y_pred[i])'*(Σ_η\(y - y_pred[i]))
        errors[N_subfigs, i]     = norm(θθ_cov[i])

        if N_subfigs == 3
            errors[1, i] = norm(θ_ref - (transform_func === nothing ? θ_mean[i] : transform_func(θ_mean[i])))/norm(θ_ref)
        end

    end

    markevery = max(div(N_iter, 10), 1)
    ax[N_subfigs - 1].plot(ites, errors[N_subfigs - 1, :], linestyle="--", marker="o", fillstyle="none", markevery=markevery)
    ax[N_subfigs - 1].set_xlabel("Iterations")
    ax[N_subfigs - 1].set_ylabel("Optimization error")
    ax[N_subfigs - 1].grid()

    ax[N_subfigs].plot(ites, errors[N_subfigs, :], linestyle="--", marker="o", fillstyle="none", markevery=markevery)
    ax[N_subfigs].set_xlabel("Iterations")
    ax[N_subfigs].set_ylabel("Frobenius norm of the covariance")
    ax[N_subfigs].grid()
    if N_subfigs == 3
        ax[1].set_xlabel("Iterations")
        ax[1].plot(ites, errors[1, :], linestyle="--", marker="o", fillstyle="none", markevery=markevery)
        ax[1].set_ylabel("L₂ norm error")
        ax[1].grid()
    end

end

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

# function Two_Param_Linear_Test(problem_type::String, θ0_bar, θθ0_cov)

#     N_θ = length(θ0_bar)

#     if problem_type == "under-determined"
#         # under-determined case
#         θ_ref = [0.6, 1.2]
#         G = [1.0 2.0;]

#         y = [3.0;]
#         Σ_η = Array(Diagonal(fill(0.1^2, size(y))))

#     elseif problem_type == "over-determined"
#         # over-determined case
#         θ_ref = [1/3, 8.5/6]
#         G = [1.0 2.0; 3.0 4.0; 5.0 6.0]

#         y = [3.0;7.0;10.0]
#         Σ_η = Array(Diagonal(fill(0.1^2, size(y))))

#     elseif problem_type == "Hilbert"

#         G = zeros(N_θ, N_θ)
#         for i = 1:N_θ
#             for j = 1:N_θ
#                 G[i,j] = 1/(i + j - 1)
#             end
#         end

#         θ_ref = fill(1.0, N_θ)
#         y   = G*θ_ref
#         Σ_η = Array(Diagonal(fill(0.5^2, N_θ)))

#     else
#         error("Problem type : ", problem_type, " has not implemented!")
#     end

#     Σ_post = inv(G'*(Σ_η\G) + inv(θθ0_cov))
#     θ_post = θ0_bar + Σ_post*(G'*(Σ_η\(y - G*θ0_bar)))

#     return θ_post, Σ_post, G, y, Σ_η
# end

# struct Setup_Param{MAT, IT<:Int}
#     G::MAT
#     N_θ::IT
#     N_y::IT
# end

# function forward(s_param::Setup_Param, θ::Array{FT, 1}) where {FT<:AbstractFloat}
#     G = s_param.G
#     return G * θ, G
# end

# N_iter = 30
# N_θ = 2

# r_0, Σ_0 = zeros(Float64, N_θ), Array(Diagonal(fill(1.0^2, N_θ)))
# θ0_mean, θθ0_cov = r_0, Σ_0

# problem_type = "under-determined"

# θ_post, Σ_post, G, y, Σ_η = Two_Param_Linear_Test(problem_type, r_0, Σ_0)

# N_y = length(y)

# s_param = Setup_Param(G, N_θ, N_y)

# method = "ExKI"
# # UKI-1
# exki_obj = ExKI_Run(s_param, forward,
# method,
# y, Σ_η,
# r_0,
# Σ_0,
# θ0_mean, θθ0_cov,
# N_iter,)

# @info "ERRORs are :" , norm(exki_obj.θ_mean[end] - θ_post), norm(exki_obj.θθ_cov[end] - Σ_post)


using LinearAlgebra

function get_module_3_excersise_4_author_name()
    return "YOUR NAME HERE"
end

function problem_4(a)

    A = [
    a 1 0 0 0 1
    1 a 1 0 0 0
    0 1 a 1 0 0
    0 0 1 a 1 0
    0 0 0 1 a 1
    1 0 0 0 1 a
    ]

    eigen_vals = []
    eigen_vectors = []
    return eigen_vals, eigen_vectors
end


abstract type SVAE end

normalizecolumns(m::AbstractArray{T, 2}) where {T<:Number} = m ./ sqrt.(sum(m .^ 2, dims = 1) .+ eps(T))
function normalizecolumns!(m::AbstractArray{T, 2}) where {T<:Number}
	m .= normalizecolumns(m)
end

"""
	vmfentropy(m, κ)

	Entropy of Von Mises-Fisher distribution
"""
vmfentropy(m, κ) = .-κ .* besselix(m / 2, κ) ./ besselix(m / 2 - 1, κ) .- ((m ./ 2 .- 1) .* log.(κ) .- (m ./ 2) .* log(2π) .- (κ .+ log.(besselix(m / 2 - 1, κ))))

"""
	huentropy(m)

	Entropy of Hyperspherical Uniform distribution
"""
huentropy(m) = m / 2 * log(π) + log(2) - lgamma(m / 2)

"""
	kldiv(model::SVAE, κ)

	KL divergence between Von Mises-Fisher and Hyperspherical Uniform distributions
"""
kldiv(model::SVAE, κ) = .- vmfentropy(model.zdim, κ) .+ model.hue

log_normal(x) = - sum((x .^ 2), dims = 1) ./ 2 .- size(x, 1) .* log(2π) ./ 2
log_normal(x, μ) = log_normal(x - μ)
log_normal(x,μ, σ2::AbstractArray{T}) where {T<:Number} = - sum((x - μ) .^ 2 ./ σ2 .+ log.(σ2 .* 2π), dims = 1) / 2

# Likelihood estimation of a sample x under VMF with given parameters taken from https://pdfs.semanticscholar.org/2b5b/724fb175f592c1ff919cc61499adb26996b1.pdf
# normalizing constant for density function of VMF
c(p, κ) = κ ^ (p / 2 - 1) / ((2π) ^ (p / 2) * besseli(p / 2 - 1, κ))

# log likelihood of one sample under the VMF dist with given parameters
log_vmf_c(x, μ, κ) = κ * μ' * x .+ log(c(length(μ), κ))
log_vmf_wo_c(x, μ, κ) = κ * μ' * x

pairwisecos(x, y) = max.(1 .- (x' * y), 0) # This is a bit of a hack to avoid the distance being negative due to float
pairwisecos(x) = pairwisecos(x, x)

k_imq(x, y, c) = sum( c./ (c .+ pairwisecos(x, y))) / (size(x, 2) * size(y, 2))
k_imq(x::T, c) where {T <: AbstractMatrix} = sum(c ./ (c .+ pairwisecos(x))) / (size(x, 2) * (size(x, 2) - 1))
k_imq(x::T, c) where {T <: AbstractVector} = zero(eltype(x))

mmd_imq(x,y,c) = k_imq(x,c) + k_imq(y,c) - 2 * k_imq(x,y,c)

function samplehsuniform(size...)
	v = randn(size...)
	v = normalizecolumns(v)
end

function log_pxexpectedz(m::SVAE, x)
	xgivenz = m.g(zparams(m, x)[1])
	Flux.Tracker.data(log_normal(x, xgivenz))
end

cart3dtoangular(x) = [acos(1.0 * x[3] / norm(x)), 1.0 * atan(x[2] / x[1])]

function cart2spherical(x)
	scoords = vcat([acot(x[i] / norm(x[(i + 1):end])) for i in 1:(length(x) - 2)]...)
	scoords = vcat(scoords, 2acot((x[end - 1] + norm(x[(end-1):end])) / x[end]))
end

jacobian_encoder(m::SVAE, x) = Flux.Tracker.jacobian(a -> Chain(m.q, m.μzfromhidden)(a), x)

function log_pz_jacobian_encoder(m::SVAE, x)
	if size(x, 2) > 1
		xs = [x[:, i] for i in 1:size(x, 2)]
		return map(x -> log_pz_jacobian_encoder_singleinstance(m, x), xs)
	else
		return log_pz_jacobian_encoder_singleinstance(m, x)
	end
end

function log_pz_jacobian_encoder_singleinstance(m::SVAE, x)
	@assert size(x, 2) == 1
	s = svd(jacobian_encoder(m, x).data)
	d = reduce(+, log.(abs.(s.S))) * 2
	d + log_pz(m, x)
end

jacobian_decoder(m::SVAE, z) = Flux.Tracker.jacobian(a -> m.g(a), z)

function log_pz_jacobian_decoder(m::SVAE, z)
	if size(z, 2) > 1
		zs = [z[:, i] for i in 1:size(z, 2)]
		return map(z -> log_pz_jacobian_decoder_singleinstance(m, z), zs)
	else
		return log_pz_jacobian_decoder_singleinstance(m, z)
	end
end

function log_pz_jacobian_decoder_singleinstance(m::SVAE, z)
	@assert size(z, 2) == 1
	s = svd(jacobian_decoder(m, z).data)
	d = reduce(+, log.(abs.(s.S))) * 2
	-d + log_pz_from_z(m, z)
end

# function log_pz_jacobian(m::SVAE, x, z)
# 	if size(x, 2) > 1
# 		xs = [x[:, i] for i in 1:size(x, 2)]
# 		zs = [z[:, i] for i in 1:size(z, 2)]
# 		return map((x, z) -> log_pz_jacobian_singleinstance(m, x, z), xs, zs)
# 	else
# 		return log_pz_jacobian_singleinstance(m, x, z)
# 	end
# end

# function log_pz_jacobian_singleinstance(m::SVAE, x, z)
# 	@assert size(x, 2) == 1
# 	s = svd(jacobian(m, x).data)
# 	d = reduce(+, log.(abs.(s.S)))
# 	d + log(pz_from_z(m, z))
# end

"""
	infer(m::SVAE, x)

	infer latent variables and sample output x
"""
function infer(m::SVAE, x)
	(μz, κz) = zparams(m, x)
	z = samplez(m, μz, κz)
	xgivenz = m.g(z)
	return xgivenz
end

"""
	zparams(model::SVAE, x)

	Computes μ and κ from the hidden layer
"""

function zparams(model::SVAE, x)
	hidden = model.q(x)
	# return model.μzfromhidden(hidden), max.(model.κzfromhidden(hidden), 100)
	return model.μzfromhidden(hidden), model.κzfromhidden(hidden)
end

"""
	zfromx(m::SVAE, x)

	convenience function that returns latent layer based on the input `x`
"""
zfromx(m::SVAE, x) = samplez(m, zparams(m, x)...)

"""
	samplez(m::SVAE, μz, κz)

	samples z layer based on its parameters
"""
function samplez(m::SVAE, μz, κz)
	ω = sampleω(m, κz)
	normal = Normal()
	v = Adapt.adapt(eltype(Flux.Tracker.data(κz)), rand(normal, size(μz, 1) - 1, size(μz, 2)))
	v = normalizecolumns(v)
	z = householderrotation(vcat(ω, sqrt.(1 .- ω .^ 2 .+ eps(Float32)) .* v), μz)
	return z
end

"""
	householderrotation(zprime, μ)

	uses Householder reflection to rotate the `zprime` vector according to mapping of e1 to μ
"""
function householderrotation(zprime, μ)
	e1 = similar(μ) .= 0
	e1[1, :] .= 1
	u = e1 .- μ
	normalizedu = normalizecolumns(u)
	return zprime .- 2 .* sum(zprime .* normalizedu, dims = 1) .* normalizedu
end

function sampleω(model::SVAE, κ)
	m = model.zdim
	c = @. sqrt(4κ ^ 2 + (m - 1) ^ 2)
	b = @. (-2κ + c) / (m - 1)
	a = @. (m - 1 + 2κ + c) / 4
	d = @. (4 * a * b) / (1 + b) - (m - 1) * log(m - 1)
	ω = rejectionsampling(m, a, b, d, κ)
	return ω
end

function rejectionsampling(m, a, b, d, κ)
	beta = Beta((m - 1) / 2, (m - 1) / 2)
	T = eltype(Flux.Tracker.data(a))
	ϵ, u = Adapt.adapt(T, rand(beta, size(a)...)), rand(T, size(a)...)

	accepted = isaccepted(ϵ, u, m, Flux.data(a), Flux.Tracker.data(b), Flux.data(d))
	it = 0
	while (!all(accepted)) & (it < 10000)
		mask = .! accepted
		ϵ[mask] = Adapt.adapt(T, rand(beta, sum(mask)))
		u[mask] = rand(T, sum(mask))
		accepted[mask] = isaccepted(mask, ϵ, u, m, Flux.data(a), Flux.data(b), Flux.data(d))
		it += 1
	end
	if it >= 10000
		println("Warning - sampler was stopped by 10000 iterations - it did not accept the sample!")
		println(m)
		println(a)
		println(b)
		println(d)
		println(κ)
		println(accepted)
	end
	return @. (1 - (1 + b) * ϵ) / (1 - (1 - b) * ϵ)
end

# function rejectionsampling(m, a, b, d)
# 	beta = Beta((m - 1) / 2, (m - 1) / 2)
# 	T = eltype(Flux.Tracker.data(a))
# 	ϵ, u = Adapt.adapt(T, rand(beta, size(a)...)), rand(T, size(a)...)
#
# 	accepted = isaccepted(ϵ, u, m, Flux.data(a), Flux.Tracker.data(b), Flux.data(d))
# 	while !all(accepted)
# 		mask = .! accepted
# 		ϵ[mask] = Adapt.adapt(T, rand(beta, sum(mask)))
# 		u[mask] = rand(T, sum(mask))
# 		ia = isaccepted(mask, ϵ, u, m, Flux.data(a), Flux.data(b), Flux.data(d))
# 		accepted[mask] = ia
# 	end
# 	return @. (1 - (1 + b) * ϵ) / (1 - (1 - b) * ϵ)
# end

isaccepted(mask, ϵ, u, m:: Int, a, b, d) = isaccepted(ϵ[mask], u[mask], m, a[mask], b[mask], d[mask]);
function isaccepted(ϵ, u, m:: Int, a, b, d)
	ω = @. (1 - (1 + b) * ϵ) / (1 - (1 - b) * ϵ)
	t = @. 2 * a * b / (1 - (1 - b) * ϵ)
	@. (m - 1) * log(t) - t + d >= log(u)
end


using BenchmarkTools
using LoopVectorization
using Flux
using ExplainableAI
import ExplainableAI: modify_layer, lrp!

on_CI = haskey(ENV, "GITHUB_ACTIONS")

include("../test/vgg11.jl")
vgg11 = VGG11(; pretrain=false)
model = flatten_model(strip_softmax(vgg11.layers))

T = Float32
img = rand(MersenneTwister(123), T, (224, 224, 3, 1))

# Benchmark custom LRP composite
function LRPCustom(model::Chain)
    return LRP(
        model,
        [ZBoxRule(zero(T), oneunit(T)), repeat([GammaRule()], length(model.layers) - 1)...],
    )
end

# Use one representative algorithm of each type
algs = Dict(
    "Gradient" => Gradient,
    "InputTimesGradient" => InputTimesGradient,
    "LRPZero" => LRPZero,
    "LRPCustom" => LRPCustom, #modifies weights
    "SmoothGrad" => model -> SmoothGrad(model, 10),
    "IntegratedGradients" => model -> IntegratedGradients(model, 10),
)

# Define benchmark
contruct_analyzer(alg, model) = alg(model) # for use with @benchmarkable macro

SUITE = BenchmarkGroup()
SUITE["VGG"] = BenchmarkGroup([k for k in keys(algs)])
for (name, alg) in algs
    SUITE["VGG"][name] = BenchmarkGroup(["construct analyzer", "analyze"])
    SUITE["VGG"][name]["construct analyzer"] = @benchmarkable contruct_analyzer(
        $(alg), $(model)
    )

    analyzer = alg(model)
    SUITE["VGG"][name]["analyze"] = @benchmarkable analyze($(img), $(analyzer))
end

# Rules benchmarks – use wrapper to trigger AD fallback
struct TestWrapper{T}
    layer::T
end
(w::TestWrapper)(x) = w.layer(x)
modify_layer(r::AbstractLRPRule, w::TestWrapper) = modify_layer(r, w.layer)
lrp!(Rₖ, rule::ZBoxRule, w::TestWrapper, aₖ, Rₖ₊₁) = lrp!(Rₖ, rule, w.layer, aₖ, Rₖ₊₁)

# generate input for conv layers
insize = (64, 64, 3, 1)
in_dense = 500
out_dense = 100
aₖ = randn(T, insize)

layers = Dict(
    "MaxPool" => (MaxPool((3, 3); pad=0), aₖ),
    "Conv" => (Conv((3, 3), 3 => 2), aₖ),
    "Dense" => (Dense(in_dense, out_dense, relu), randn(T, in_dense, 1)),
    "WrappedDense" =>
        (TestWrapper(Dense(in_dense, out_dense, relu)), randn(T, in_dense, 1)),
)
rules = Dict(
    "ZeroRule" => ZeroRule(),
    "EpsilonRule" => EpsilonRule(),
    "GammaRule" => GammaRule(),
    "ZBoxRule" => ZBoxRule(zero(T), oneunit(T)),
)

SUITE["Layer"] = BenchmarkGroup([k for k in keys(layers)])
for (layername, (layer, aₖ)) in layers
    SUITE["Layer"][layername] = BenchmarkGroup([k for k in keys(rules)])
    Rₖ = similar(aₖ)
    Rₖ₊₁ = layer(aₖ)
    for (rulename, rule) in rules
        SUITE["Layer"][layername][rulename] = @benchmarkable lrp!(
            $(Rₖ), $(rule), $(layer), $(aₖ), $(Rₖ₊₁)
        )
    end
end


# -----------------------------------------------------------------------------
# Diagonally-dominant bridge [MOI]
# -----------------------------------------------------------------------------

"""
Converts a Semidefinite Program (SDP) into a Diagonally-Dominant Program (DDP).
Positive-semidefiniteness constraints are relaxed into Diagonal-dominance
constraints.
"""
struct DDPBridge{T,F,G} <: MOIB.Constraint.AbstractBridge
        dominance::Vector{MOI.ConstraintIndex{F,MOI.EqualTo{T}}}
end

function MOIB.Constraint.bridge_constraint(
        ::Type{DDPBridge{T,F,G}},
        m::MOI.ModelLike,
        f::MOI.AbstractVectorFunction,
        s::MOI.PositiveSemidefiniteConeTriangle,
) where {T,F,G}
        n = s.side_dimension
        z = MOI.add_variables(m, n)
        v = MOI.SingleVariable.(z)
        matrix = square_a_triangle(f, n)

        ds = Vector{MOI.ConstraintIndex{F,MOI.EqualTo{T}}}(undef, n)
        for i in 1:n
                zs = MOIU.operate(vcat, T, v[i], matrix[i, :]...)
                dominance = T(2) * matrix[i, i] - v[i]
                MOI.add_constraint(m, zs, MOI.NormOneCone(n + 1))
                ds[i] = MOI.add_constraint(m, dominance, MOI.EqualTo(zero(T)))
        end
        return DDPBridge{T,F,G}(ds)
end

function MOI.supports_constraint(
        ::Type{<:DDPBridge},
        ::Type{<:MOI.AbstractVectorFunction},
        ::Type{<:MOI.PositiveSemidefiniteConeTriangle},
)
        return true
end
function MOIB.added_constrained_variable_types(::Type{<:DDPBridge})
        return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:DDPBridge{T,F,G}}) where {T,F,G}
        return [(F, MOI.EqualTo{T}); (G, MOI.NormOneCone)]
end
function MOIB.Constraint.concrete_bridge_type(
        ::Type{<:DDPBridge{T}},
        G::Type{<:MOI.AbstractVectorFunction},
        ::Type{MOI.PositiveSemidefiniteConeTriangle},
) where {T}
        S = MOIU.scalar_type(G)
        F = MOIU.promote_operation(-, T, S, MOI.SingleVariable)
        return DDPBridge{T,F,G}
end

function MOI.get(
        model::MOI.ModelLike,
        attr::MOI.ConstraintDual,
        bridge::DDPBridge{T},
) where {T}
        diag_dual = MOI.get(model, attr, bridge.dominance)
        side = length(diag_dual)
        n = MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(side))

        dd_dual = Vector{T}(undef, n)
        k = 0
        for j in 1:side
                for i in 1:(j-1)
                        k += 1
                        dd_dual[k] = -(diag_dual[i] + diag_dual[j])
                end
                k += 1
                dd_dual[k] = diag_dual[j] * 2
        end
        return dd_dual
end


using Distributions: Dirichlet
using SpecialFunctions: lgamma, digamma, lbeta, beta
using Base.MathConstants: γ

struct BatchKumaraswamy{T}
    a::T
    b::T
end

function _u2kumaraswamysample(u, kuma::BatchKumaraswamy)
    _one = one(eltype(u))
    return (_one .- u.^(_one ./ kuma.b)).^(_one ./ kuma.a)
end

function _u2logkumaraswamysample(u, kuma::BatchKumaraswamy)
    _one = one(eltype(u))
    _eps = eps(eltype(u))
    return log1p.(-exp.(log.(u .+ _eps) ./ kuma.b)) ./ kuma.a
end

"""
    rand(kuma::BatchKumaraswamy) where

Sample from Kumaraswamy distribution.

NOTE: `k.a` and `k.b` are assumed to be in batch

Ref: https://arxiv.org/abs/1605.06197
"""
function rand(kuma::BatchKumaraswamy)
    u = randsimilar(kuma.a)
    x = _u2kumaraswamysample(u, kuma)
    return x
end

function rand(kuma::BatchKumaraswamy{T}, dims::Int...) where {T<:Real}
    @assert length(kuma.a) == 1 "`rand` for multiple samples only supports for univariate case"
    @assert length(kuma.b) == 1 "`rand` for multiple samples only supports for univariate case"
    u = length(dims) == 0 ? rand(T) : randarr(dims)
    x = _u2kumaraswamysample(u, kuma)
    return x
end

function logrand(kuma::BatchKumaraswamy)
    u = randsimilar(kuma.a)
    logx = _u2logkumaraswamysample(u, kuma)
    return logx
end

function logrand(kuma::BatchKumaraswamy{T}, dims::Int...) where {T<:Real}
    @assert length(kuma.a) == 1 "`rand` for multiple samples only supports for univariate case"
    @assert length(kuma.b) == 1 "`rand` for multiple samples only supports for univariate case"
    u = length(dims) == 0 ? rand(T) : randarr(dims)
    logx = _u2logkumaraswamysample(u, kuma)
    return logx
end

"""
    logpdf(kuma::BatchKumaraswamy, x)

Compute ``Kumaraswamy(x; a, b)``.

WARN: this function is not tested.
"""
function logpdf(kuma::BatchKumaraswamy, x)
    _one = one(eltype(x))
    _eps = eps(eltype(x))
    lp = log.(kuma.a) .+
         log.(kuma.b) .+
         (kuma.a .- _one) .* log.(x .+ _eps) .+
         (kuma.b .- _one) .* log.(_one + _eps .- x.^kuma.a)
    return lp
end

"""
    kldiv(d1::Dirichlet, d2::Dirichlet)

Compute ``KL(Dir(α)||Dir(β))`` where ``α = [α_1, \\dots, α_K]`` and ``β = [β_1, \\dots, β_K]``.

Ref: http://bariskurt.com/kullback-leibler-divergence-between-two-dirichlet-and-beta-distributions/
"""
function kldiv(d1::Dirichlet, d2::Dirichlet)
    α = d1.alpha
    β = d2.alpha

    α0 = sum(α)
    β0 = sum(β)

    kl = (lgamma.(α0) .- sum(lgamma.(α))) .-
         (lgamma.(β0) - sum(lgamma.(β))) .+
         sum((α .- β) .* (digamma.(α) .- digamma(α0)))
    return kl
end

struct BatchBeta{T}
    α::T
    β::T
end

function logpdf(bb::BatchBeta, x)
    _one = one(eltype(x))
    _eps = eps(eltype(x))
    lp = (bb.α .- _one) .* log.(x .+ _eps) .+
         (bb.β .- _one) .* log.(_one + _eps .- x) .-
         lbeta(bb.α, bb.β)
    return lp
end

"""
    kldiv(bb1::BatchBeta, bb2::BatchBeta)

Compute ``KL(Beta(α1, β1)||Beta(α2, β2))``.
"""
function kldiv(bb1::BatchBeta, bb2::BatchBeta)
    α1 = bb1.α
    β1 = bb1.β
    α2 = bb2.α
    β2 = bb2.β

    α0 = α1 .+ β1
    β0 = α2 .+ β2

    kl = (lgamma.(α0) .- (lgamma.(α1) .+ lgamma.(β1))) .-
         (lgamma.(β0) .- (lgamma.(α2) .+ lgamma.(β2))) .+
         (α1 - α2) .* (digamma.(α1) - digamma.(α0)) .+
         (β1 - β2) .* (digamma.(β1) - digamma.(α0))
    return kl
end

"""
    kldiv(kuma::BatchKumaraswamy, bb::BatchBeta; M=10)

Compute ``KL(Kumaraswamy(a, b)||Beta(α, β))``.

NOTE: only `a` and `b` are assumed to be in batch
"""
function kldiv(kuma::BatchKumaraswamy, bb::BatchBeta; M::Int=11)
    a = kuma.a
    b = kuma.b
    α = bb.α
    β = bb.β
    T = eltype(a)
    _one = one(eltype(a))

    @assert M > 0
    a_times_b = a .* b
    acc = _one ./ (_one .+ a_times_b) .* beta(_one ./ a, b)
    for m = 2:M
        acc = acc .+ _one ./ (T(m) .+ a_times_b) .* beta(T(m) ./ a, b)
    end

    kl = (a .- α) ./ a .* (-T(γ) .- digamma.(b) .- _one ./ b) .+ log.(a_times_b) .+
         lbeta(α, β) .- (b .- _one) ./ b .+ (β - _one) .* b .* acc
    return kl
end


struct MatrixLogNormal{T<:Real,M<:AbstractMatrix,C<:AbstractPDMat} <:
    Distribution{Matrixvariate,Continuous}
    normal::MatrixNormal{T, M, C}

    function MatrixLogNormal(normal::MatrixNormal{T, M, C}) where {T, M, C}
        new{T, M, C}(normal)
    end
end

function MatrixLogNormal(M, U, V)
    MatrixLogNormal(MatrixNormal(M, U, V))
end

rand(MLN::MatrixLogNormal) = exp.(rand(MLN.normal))

function location(::Type{D}, s::Symbol,
                  M::AbstractMatrix,
                  U::AbstractMatrix, V::AbstractMatrix) where {D<:MatrixLogNormal}
    loc = similar(M)
    @inbounds for jdx in size(loc, 2)
        M[:, jdx] .= location(MvLogNormal, s,
                              M[:, jdx],
                              diagm(V[jdx, jdx] .* diag(U)))
    end
    loc
end

function location(::Type{D}, s::Symbol,
                  M::AbstractMatrix,
                  U::AbstractPDMat, V::AbstractPDMat) where {D<:MatrixLogNormal}
    location(D, s, M, U.mat, V.mat)
end


using FunctionZeros
using BenchmarkTools
import SpecialFunctions

@btime besselj_zero(3,10)
@btime besselj_zero(3,10)

const z = besselj_zero(3,10)
zer = SpecialFunctions.besselj(3,z)


# Computes TTVs with TTVFaster for N planets with pairwise TTV calculation.
include("compute_ttv.jl")

function ttv_nplanet(nplanet::Int64,jmax::Int64,ntrans::Vector{Int64},data::Vector{T}) where T<:Real
  # Need at least two planets!
  @assert(nplanet>=2)
  # The ntrans vectors should have length nplanet:
  @assert(length(ntrans)==nplanet)
  # Define type of ttv array:
  ttv_el_type = eltype(data) == Float64 ? Float64 : Number
  # Need to create an array to store TTVs with maximum length equal to maximum number
  # of transit times of any planet:
  ntransmax = maximum(ntrans)
  #ttv = zeros(ttv_el_type,nplanet,ntransmax)
  ttv = zeros(T,nplanet,ntransmax) #check memory allocation <<<<<<<<<<<<<<<<<<<
  # Each planet requires 5 elements in data: mass_ratio,period,trans0,ecosw,esinw:
  @assert(length(data)==5*nplanet)
  @assert(jmax>=1)  # Should there be a larger minimum?
  for iplanet=1:nplanet
  # Each planet should have at least 2 transits:
    @assert(ntrans[iplanet]>=2)
  end
  for iplanet=1:nplanet-1
  # The periods of the planets should be ordered from least to greatest:
    if (data[(iplanet-1)*5+2] >= data[iplanet*5+2])
      return ttv
    end
  end
  # Set up planets planar-planet types for all of the planets:
  #planet = Array{Planet_plane_hk}(nplanet)
  #planet = Array{Any}(nplanet)
  # Loop over pairs of planets to compute pairwise TTVs
  # Loop over inner planets:
  #println("Looping over planets in ttv_nplanet:")
  for iplanet=1:nplanet-1
    # Create a Planet_plane_hk type for the inner planet:
    p1=Planet_plane_hk(data[(iplanet-1)*5+1],data[(iplanet-1)*5+2],data[(iplanet-1)*5+3],data[(iplanet-1)*5+4],data[(iplanet-1)*5+5])
    # Create an array of times for the inner planet:
    n1 = ntrans[iplanet]
    time1 = collect(p1.trans0 .+ range(0,stop=n1-1,length=n1) .* p1.period)
    # Loop over outer planets:
    for jplanet=iplanet+1:nplanet
      # Create a Planet_plane_hk type for the outer planet:
      p2=Planet_plane_hk(data[(jplanet-1)*5+1],data[(jplanet-1)*5+2],data[(jplanet-1)*5+3],data[(jplanet-1)*5+4],data[(jplanet-1)*5+5])
      # Create an array of times for the outer planet:
      n2 = ntrans[jplanet]
      time2 = collect(p2.trans0 .+ range(0,stop=n2-1,length=n2) .* p2.period)
      # Define arrays to hold the TTVs:
      ttv1=zeros(T,n1)
      ttv2=zeros(T,n2)
      # Call the compute_ttv code which implements equation (33) from Agol & Deck (2016):
      #    println("Calling compute_ttv")
      compute_ttv!(jmax,p1,p2,time1,time2,ttv1,ttv2)
      #    println("Finished compute_ttv")
      for i=1:n1
        ttv[iplanet,i] += ttv1[i]
      end
      for i=1:n2
        ttv[jplanet,i] += ttv2[i]
      end
    end
  end
  return ttv
end


mutable struct Point
    x::Float64
    y::Float64
end

mutable struct mySurface
    P::Point
    Q::Point
end

function mySurface(Parr::Vector{Float64},Qarr::Vector{Float64})
  P = Point(Parr[1],Parr[2])
  Q = Point(Qarr[1],Qarr[2])
  return mySurface(P,Q)
end

"""
  DomainQuad(mesh,u)

  Computes the 2D integral over the entire domain defined in mesh of the
  solution u. Uses Gaussian Quadrature.
"""
function DomainQuad(mesh::M,u::Vector{Float64}) where {M<:AbstractMesh}
  Nel   = length(mesh.cm)
  ngpt  = length(mesh.cm[1].NodeList)
  order = (ngpt==4 ? 1 : 2)
  quad  = 0.0

  # generate basis values at gauss points
  w,s,t = GaussQuadPoints2D(order+1)

  xN     = zeros(Float64,ngpt)
  yN     = zeros(Float64,ngpt)
  ucoefs = zeros(Float64,ngpt)

  for el=1:Nel
    for ti=1:ngpt
      c          = mesh.cm[el].NodeList[ti]
      xN[ti]     = mesh.xy[c].x
      yN[ti]     = mesh.xy[c].y
      ucoefs[ti] = u[c]
    end

    for gpt=1:ngpt
      phi,_,_,jac = derivShape2D(s[gpt],t[gpt],xN,yN,order)
      ugpt = shapeEval(ucoefs,phi)

      quad += w[gpt]*abs(ugpt)*jac
    end
  end

  return quad
end
function DomainQuad(mesh,u::S) where S<:AbstractSolution
  uarr = u.u
  return DomainQuad(mesh,uarr)
end

"""
  SurfaceQuad(mesh,u,surf,N)

  Computes N-discretized line integral of u on 'surface' (aka line) surf.
  Quadrature uses trapezoid rule
"""
function SurfaceQuad(mesh::M,u::Vector{Float64},surf::mySurface,N::Int) where
                    {M<:AbstractMesh}
  (mesh.order==:Linear ? order=1 : order=2)
  (order==1 ? nGpt=2 : nGpt=3)

  # Define 2 points
  P = surf.P; Q = surf.Q
  rlen = sqrt((P.x-Q.x)^2 + (P.y-Q.y)^2)

  # Interpolate N segments (N+1 points)
  Δt = rlen/N
  xN = zeros(Float64,N+1)
  yN = zeros(Float64,N+1)
  uN = zeros(Float64,N+1)
  tx,ty = ComputeTangent(P,Q)

  xN[1] = P.x; yN[1] = P.y
  for i=2:(N+1)
    xN[i] = xN[1] + (i-1)*Δt*tx
    yN[i] = yN[1] + (i-1)*Δt*ty
  end

  # Interpolate solution at N+1 points
  for i=1:(N+1)
    point = [xN[i];yN[i]]
    uN[i] = pointTransform(mesh,u,point)
  end

  # For each segment, do trapezoid rule quadrature
  quad = 0.0
  for ti=2:N
    quad += uN[ti]*Δt
  end
  quad += 0.5*Δt*(uN[1] + uN[N+1])

  # return quad
  return quad
end
function SurfaceQuad(mesh::M,u::S,surf::mySurface,N::Int) where
  {M<:AbstractMesh,S<:AbstractSolution}
  uarr = u.u
  return SurfaceQuad(mesh,uarr,surf,N)
end

"""
  Performs interpolation along a surface (line)
"""
function SurfaceInterp(mesh::M,u::Vector{Float64},surf::mySurface,N::Int) where
  {M<:AbstractMesh}
  (mesh.order==:Linear ? order=1 : order=2)
  (order==1 ? nGpt=2 : nGpt=3)

  # Define 2 points
  P = surf.P; Q = surf.Q
  rlen = sqrt((P.x-Q.x)^2 + (P.y-Q.y)^2)

  # Interpolate N segments (N+1 points)
  Δt = rlen/N
  xN = zeros(Float64,N+1)
  yN = zeros(Float64,N+1)
  uN = zeros(Float64,N+1)
  tx,ty = ComputeTangent(P,Q)

  xN[1] = P.x; yN[1] = P.y
  for i=2:(N+1)
    xN[i] = xN[1] + (i-1)*Δt*tx
    yN[i] = yN[1] + (i-1)*Δt*ty
  end

  # Interpolate solution at N+1 points
  for i=1:(N+1)
    point = [xN[i];yN[i]]
    uN[i] = pointTransform(mesh,u,point)
  end

  # return quad
  return uN
end

"""
  for some scalar field c, computes grad(c)*n across a surface
"""
function SurfaceFlux(mesh::M,u::Vector{Float64},surf::mySurface,Nquad::Int) where
                    {M<:AbstractMesh}
  Npts = length(mesh.xy)
  Nel  = length(mesh.cm)
  (mesh.order==:Linear ? order=1 : order=2)
  (order==1 ? nGpt=2 : nGpt=3)
  NnodesPerElm = (order==1 ? 4 : 9)

  # Define 2 points
  P = surf.P; Q = surf.Q
  rlen = sqrt((P.x-Q.x)^2 + (P.y-Q.y)^2)

  # Initialize arrays
  Δt = rlen/Nquad
  xN = zeros(Float64,Nquad+1)
  yN = zeros(Float64,Nquad+1)
  uN = zeros(Float64,Nquad+1)

  dudx  = zeros(Float64,Npts)
  dudy  = zeros(Float64,Npts)
  dudxN = zeros(Float64,Nquad+1)
  dudyN = zeros(Float64,Nquad+1)

  tx,ty = ComputeTangent(P,Q)
  nx,ny = ComputeNormal(P,Q)

  # Interpolate N segments (N+1 points)
  xN[1] = P.x; yN[1] = P.y
  for i=2:(Nquad+1)
    xN[i] = xN[1] + (i-1)*Δt*tx
    yN[i] = yN[1] + (i-1)*Δt*ty
  end

  # compute derivatives at mesh points
  s = [-1.0,1.0,1.0,-1.0,0.0,1.0,0.0,-1.0,0.0]
  t = [-1.0,-1.0,1.0,1.0,-1.0,0.0,1.0,0.0,0.0]
  for el=1:Nel
    cv = mesh.cm[el].NodeList
    xNMesh = [mesh.xy[i].x for i in cv]
    yNMesh = [mesh.xy[i].y for i in cv]
    utemp  = [u[i] for i in cv]

    for i=1:NnodesPerElm
      # compute basis at node points
      _,dphidx,dphidy,jac = derivShape2D(s[i],t[i],xNMesh,yNMesh,order)
      dudx[mesh.cm[el].NodeList[i]] = shapeEval(utemp,dphidx)/jac
      dudy[mesh.cm[el].NodeList[i]] = shapeEval(utemp,dphidy)/jac
    end
  end

  # Interpolate derivative solution at N+1 points
  for i=1:(Nquad+1)
    point  = [xN[i];yN[i]]
    dudxN[i]  = pointTransform(mesh,dudx,point)
    dudyN[i]  = pointTransform(mesh,dudy,point)
  end

  # For each segment, do trapezoid rule quadrature
  quad = 0.0
  for ti=2:Nquad
    quad += (dudxN[ti]*nx + dudyN[ti]*ny)*Δt
  end
  quad += 0.5*Δt*(dudxN[1]*nx + dudyN[1]*ny + dudxN[Nquad+1]*nx + dudyN[Nquad+1]*ny)

  # return quad
  return quad
end
function SurfaceFlux(mesh::M,u::S,surf::mySurface,N::Int) where
  {M<:AbstractMesh,S<:AbstractSolution}
  uarr = u.u
  return SurfaceFlux(mesh,uarr,surf,N)
end

function VelocityFlux(mesh::M,u::Vector{Float64},v::Vector{Float64},
                surf::mySurface,N::Int) where {M<:AbstractMesh}
  # compute normal
  nx,ny = ComputeNormal(surf.P,surf.Q)

  # copmute u-flux
  uFlux = SurfaceQuad(mesh,u,surf,N)

  # scale u-flux by horizontal normal
  uval = nx*uFlux

  # compute v-flux
  vFlux = SurfaceQuad(mesh,v,surf,N)

  # scale v-flux by vertical normal
  vval = ny*vFlux

  return uval + vval
end

#############################################
# NECESSARY FUNCTIONS FOR SURFACE Quadrature
#############################################

"""
  Interp2Dline(surf,N)

  Interpolates the points defining a surface in 2D (aka a line) into N segments.
"""
function Interp2Dline(surf::mySurface,N::Int)
  vt = ComputeTangent(surf.P,surf.Q)
  vn = ComputeNormal(surf.P,surf.Q)

  return vt,vn
end

# computes the normal at 2 points p1=[x1,y1] and p2=[x2,y2]
function ComputeTangent(p1,p2)
  rx = p2.x - p1.x
  ry = p2.y - p1.y

  rmag = sqrt(rx^2 + ry^2)

  tx = rx/rmag
  ty = ry/rmag

  return tx,ty
end

# computes the normal at 2 points p1=[x1,y1] and p2=[x2,y2]
# Normal is defined if you are standing on p1 and looking at p2, the normal goes to
# the right
function ComputeNormal(p1,p2)
  tx,ty = ComputeTangent(p1,p2)

  nx = ty
  ny = -tx

  return nx,ny
end


import Distributions
import LinearAlgebra
import Statistics
import StatsBase

ncol(df::DataFrames.AbstractDataFrame) = size(df, 2)

ncol(m::AbstractMatrix) = size(m, 2)

function numeric_columns(
        dat,
        )
    p = ncol(dat)
    is_numeric = falses(p)
    for j = 1:p
        typ = eltype(dat[:, j])
        if typ <: Real && typ ≠ Bool
            is_numeric[j] = true
        end
    end
    res = findall(is_numeric)
    res
end

function fill_diagonal!(
        X,
        diag_elems,
        )
    p = size(X, 2)
    for i = 1:p
        X[i, i] = diag_elems[i]
    end
end

function rose_real(
        X,
        n,
        ids_class,
        ids_generation,
        h_mult = 1,
        )
    X_array = convert(Array, X)
    p = size(X_array, 2)
    n_new = length(ids_generation)
    cons_kernel = (4/((p+2) * n))^(1/(p+4))

    if p ≠ 1
        # sd_mat = eye(p)
        sd_mat = Matrix(1.0LinearAlgebra.I, p, p)
        sd_vect = Statistics.std(X_array[ids_class, :];dims = 2,)
        fill_diagonal!(sd_mat, sd_vect)
        H = h_mult * cons_kernel * sd_mat
    else
        H = h_mult * cons_kernel * Statistics.std(X_array[ids_class, :])
    end
    X_new_num = randn(n_new, p) * H
    X_new_num = X_new_num + X_array[ids_generation, :]
    return X_new_num
end

function rose_sampling(
        X,
        y,
        prop,
        indcs_maj,
        indcs_min,
        y_majority,
        y_minority,
        h_mult_maj,
        h_mult_min,
        )
    n = size(X, 1)
    n_minority = sum(rand(Distributions.Binomial(1, prop), n))
    n_majority = n - n_minority

    indcs_maj_new = StatsBase.sample(indcs_maj, n_majority, replace = true)
    indcs_min_new = StatsBase.sample(indcs_min, n_minority, replace = true)

    numeric_cols = numeric_columns(X)

    # Create  X
    indcs = vcat(indcs_maj_new, indcs_min_new)
    X_new = X[indcs, :]
    if length(numeric_cols) > 0
        majority_rose_real_results = rose_real(
            X[:, numeric_cols],
            length(indcs_maj),
            indcs_maj,
            indcs_maj_new,
            h_mult_maj,
            )
        for k = 1:length(numeric_cols)
            col = numeric_cols[k]
            X_new[1:n_majority, col] = majority_rose_real_results[:, k]
        end
        minority_rose_real_results = rose_real(
            X[:, numeric_cols],
            length(indcs_min),
            indcs_min,
            indcs_min_new,
            h_mult_min,
            )
        for k = 1:length(numeric_cols)
            col = numeric_cols[k]
            X_new[(n_majority + 1):n, col] = minority_rose_real_results[:, k]
        end
    end

    # Create y
    y_new = similar(y)
    y_new[1:n_majority] .= y_majority
    y_new[(n_majority + 1):n] .= y_minority

    result = (X_new, y_new,)
    return result
end

"""
    classlabel(y)
Given a column from a DataFrames.DataFrame, this function returns the majority/minority class label.
"""
function classlabel(
        y::Array{T, 1},
        labeltype = :minority,
        ) where T
    count_dict = StatsBase.countmap(y)
    labels = collect(keys(count_dict))
    counts = collect(values(count_dict))
    if length(labels) > 2
        error("There are more than two classes in the target variable.")
    elseif length(labels) < 2
        error("There is only one class in the target variable.")
    end
    func = (labeltype == :majority) ? :argmax : :argmin
    indx = eval(func)(counts)
    res = collect(labels)[indx]
    res
end

function rose(
        dat::DataFrames.DataFrame,
        y_column::Symbol,
        prop::Float64 = 0.5,
        h_mult_maj = 1,
        h_mult_min = 1,
        )
    majority_label = classlabel(dat[y_column], :majority)
    minority_label = classlabel(dat[y_column], :minority)

    indcs_maj = findall(dat[y_column] .== majority_label)
    indcs_min = findall(dat[y_column] .== minority_label)
    p = size(dat, 2)
    y_indx = findfirst(names(dat) .== y_column)
    X_indcs = setdiff(1:p, y_indx)
    X = dat[:, X_indcs]
    y = dat[y_column]
    y_majority = majority_label
    y_minority = minority_label
    X_new, y_new = rose_sampling(
        X,
        y,
        prop,
        indcs_maj,
        indcs_min,
        y_majority,
        y_minority,
        h_mult_maj,
        h_mult_min,
        )
    result = X_new
    result[y_column] = y_new
    return result
end


@testset "PolyProducts" begin

@info("--------- Testing Products of Polynomials ----------")

##
using ACE, Test, ForwardDiff, JuLIP, JuLIP.Testing

using LinearAlgebra: norm, cond
using ACE.OrthPolys: OrthPolyBasis
using JuLIP: evaluate, evaluate_d

##

maxn = 30
tdf = rand(200)
ww = 1.0 .+ rand(200)
Jd = OrthPolyBasis(maxn, 2, 1.0, 1, -1.0, tdf, ww)

##

@info("Testing the Expansion of products of radial polynomials")
coeffs = ACE.OrthPolys.OrthPolyProdCoeffs(Jd)

for n1 = 1:5, n2 = 1:5
   # product
   f1(x) = evaluate(Jd, x)[n1] * evaluate(Jd, x)[n2]
   # expansion
   f2(x) = ( P = coeffs(n1, n2);
             sum( evaluate(Jd, x)[ν] * P[ν]  for ν = 1:length(P) ) )
   for ntest = 1:10
      x = 2 * rand() - 1
      print_tf(@test f1(x) ≈ f2(x))
   end
end

##

@info("Testing Expansion of Products of Spherical Harmonics")
using ACE.Orth: SHProdCoeffs
using ACE.SphericalHarmonics: SHBasis, index_y

coeffs = SHProdCoeffs()
for l1 = 0:4, l2=l1:4, m1 = -l1:l1, m2 = -l2:l2
   f1 = let SH = SHBasis(10)
      x -> ( Y = evaluate(SH, x);
             Y[index_y(l1, m1)] * Y[index_y(l2, m2)] )
   end
   f2 = let P = coeffs(l1, m1, l2, m2), SH = SHBasis(10)
      x -> sum( p * evaluate(SH, x)[index_y(L, M)]
                for (L, M, p) in P )
   end
   nfail = 0
   for ntest = 1:10
      x = randn(JVec); x /= norm(x)
      nfail += (abs(f1(x) - f2(x)) > 1e-12)
   end
   print_tf(@test nfail == 0)
end

##

end


using Test, AdmittanceModels, LinearAlgebra
import .PSOModelExample, .LRCExample, .CircuitCascadeExample
const ae = PSOModelExample
const lrce = LRCExample
const cce = CircuitCascadeExample

@testset "PSOModel to Blackbox" begin
    # test 1
    ω = collect(range(4, stop=6, length=1000)) * 1e9 * 2π
    bbox = Blackbox(ω, ae.pso)
    @test get_ports(bbox) == get_ports(ae.pso)
    @test length(get_Y(bbox)) == length(ω)
    @test size(get_Y(bbox)[1]) == size(ae.pso.K)
    bbox = canonical_gauge(bbox)
    @test size(get_Y(bbox)[1]) == (2,2)
    # test 2
    l, r, c, ω₀ = lrce.l, lrce.r, lrce.c, lrce.ω
    ω = collect(range(.8, stop=1.2, length=1000)) * ω₀
    bbox = Blackbox(ω, lrce.pso)
    z = 1 ./(1 ./(l * 1im * ω) .+ 1/r .+ c * 1im * ω)
    @test length(get_Y(bbox)) == length(ω)
    impedance = impedance_matrices(bbox)
    @test all([size(x) == (1,1) for x in impedance])
    @test [x[1,1] for x in impedance] ≈ z
end

@testset "cascade, unite_ports, open_ports, short_ports" begin
    # cascade
    pso0 = PSOModel(cce.circuit0, [(0,1), (0,2)], [1, 2])
    pso1 = PSOModel(cce.circuit1, [(4,5), (4,6)], [5, 6])
    pso2 = PSOModel(cce.circuit2, [(8,9), (8,10)], [9, 10])
    pso = cascade(pso0, pso1, pso2)
    ω = collect(range(.1, stop=100, length=1000)) * 2π
    bbox_from_pso = Blackbox(ω, pso)
    @test get_ports(bbox_from_pso) == get_ports(pso)
    bboxes0 = [Blackbox(ω, a) for a in [pso0, pso1, pso2]]
    bbox0 = cascade(bboxes0)
    @test bbox0 ≈ bbox_from_pso
    bboxes1 = [canonical_gauge(b) for b in bboxes0]
    bbox1 = cascade(bboxes1)
    @test bbox1 ≈ canonical_gauge(bbox_from_pso)
    # unite_ports
    united_pso = unite_ports(pso, 1, 5, 9)
    bbox_from_pso = Blackbox(ω, united_pso)
    @test get_ports(bbox_from_pso) == get_ports(united_pso)
    united_bbox0 = unite_ports(bbox0, 1, 5, 9)
    @test united_bbox0 ≈ bbox_from_pso
    united_bbox1 = unite_ports(bbox1, 1, 5, 9)
    @test canonical_gauge(united_bbox1) ≈ canonical_gauge(bbox_from_pso)
    # open_ports
    open_pso = open_ports(united_pso, 2, 6)
    bbox_from_pso = Blackbox(ω, open_pso)
    open_bbox = open_ports(united_bbox1, 2, 6)
    @test canonical_gauge(open_bbox) ≈ canonical_gauge(bbox_from_pso)
    # short_ports
    short_pso = short_ports(united_pso, 2, 6)
    bbox_from_pso = Blackbox(ω, short_pso)
    short_bbox = short_ports(united_bbox1, 2, 6)
    @test canonical_gauge(short_bbox) ≈ canonical_gauge(bbox_from_pso)
end

# Pozar p 192
@testset "transfer function conversions" begin
    a, b, c, d = 1.0, 2.0, 3.0, 4.0
    Z0 = 10.0
    Y = [[d, -1] [b * c - a * d, a]] ./ b
    Z = [[a, 1] [a * d - b * c, d]] ./ c
    denom = a + b/Z0 + c * Z0 + d
    S = [[a + b/Z0 - c * Z0 - d, 2] [2 * (a * d - b * c), -a + b/Z0 - c * Z0 + d]] ./ denom
    @test inv(Y) ≈ Z
    Z0s = [Z0, Z0]
    @test AdmittanceModels.admittance_to_scattering(Y, Z0s) ≈ S
    @test AdmittanceModels.scattering_to_admittance(S, Z0s) ≈ Y
    @test AdmittanceModels.impedance_to_scattering(Z, Z0s) ≈ S
    @test AdmittanceModels.scattering_to_impedance(S, Z0s) ≈ Z
end


using Revise, InfiniteLinearAlgebra, MatrixFactorizations
import IntervalArithmetic
import IntervalArithmetic: Interval, emptyinterval
import MatrixFactorizations: add_mul_signed, reflector!

Δ = SymTridiagonal(Vcat([3.0,4.0],Fill(0.0,∞)), Vcat([0.5],Fill(0.5,∞)))
λ = Interval(4.1,4.2); ql(Δ - λ*I).L[1,1]
m = (λ.hi + λ.lo)/2
λ1 = Interval(λ.lo,m)
λ = 0 in ql(Δ - λ1*I).L[1,1] ? λ1 : Interval(m,λ.hi)
ql(Δ - λ*I).L[1,1]

@inline function add_mul_signed(normu::Interval, ξ1::Interval)
    ν = copysign(normu, real(ξ1))
    ξ1 + ν, ν
end

λ2 = Interval(m,λ.hi)
ql(Δ - λ1.hi*I).L[1,1]

ql(Δ - Interval(4.2)I).L[1,1]

rig_qltail(0.5,-1.98,2.0)

A = (Δ-λ*I)
import InfiniteLinearAlgebra: _qlfactUnblocked!
InfiniteLinearAlgebra.qltail(0.5,Interval(-4.1,-4.0), 0.5)

InfiniteLinearAlgebra.qltail(0.5,-4.1, 0.5)

InfiniteLinearAlgebra.qltail(0.5,-4.0, 0.5)

Z,A,B = 0.5,-4.0,0.5

(Interval(0,1) ∩ Interval(2,3)).lo

A = 4
rig_qltail(-2,10,B)

d,e = d2,e2

rig_qltail(Z,A,B,d,e)

InfiniteLinearAlgebra.qltail(Z,A,B)

@which reflector!(X[2:-1:1,end])

@which ql!(X)
A = X

import MatrixFactorizations: reflector!, reflectorApply!

m, n = size(A)
τ = zeros(eltype(A), min(m,n))
for k = n:-1:max((n - m + 1 + (T<:Real)),1)
    k = n
    μ = m+k-n
    x = view(A, μ:-1:1, k)
    τk = reflector!(x)
    τ[k-n+min(m,n)] = τk
    reflectorApply!(x, τk, view(A, μ:-1:1, 1:k-1))
end
x

@which reflector!(x)
QL(A, τ)

sign(Interval(-10,10))

length(d)

InfiniteLinearAlgebra.qltail(Z,A,B)

Z,A,B = 0.5,Interval(-4.1,-4.0),0.5
T = promote_type(eltype(Z),eltype(A),eltype(B))
ñ1 = (A + sqrt(A^2-4B*Z))/2
ñ2 = (A - sqrt(A^2-4B*Z))/2
ñ = abs(ñ1) > abs(ñ2) ? ñ1 : ñ2
(n,σ) = (abs(ñ),conj(sign(ñ)))
if n^2 < abs2(B)
    throw(ContinuousSpectrumError())
end

e = sqrt(n^2 - abs2(B))
d = σ*e*Z/n

X = [Z A B;
     0 d e]
QL = _qlfactUnblocked!(X)

# two iterations to correct for sign
X[2,:] .= (zero(T), X[1,1], X[1,2]);
X[1,:] .= (Z,A,B);
QL = _qlfactUnblocked!(X)

X, QL.τ[end]


using NMF: nnmf, solve!, CoordinateDescent, nndsvd

"""
    NMF(;ncomps=nothing) <: LinearAlgorithm
    NMF(ncomps)

Use non-negative matrix factorization (NMF, NNMF) to form a non-negative low-rank orthonormal basis of the input. The implementation of the underlying NMF is provided by [NMF.jl](https://github.com/JuliaStats/NMF.jl). The implementation uses a non-negative SVD for initialization and a coordinate-descent solver to fit.

If `ncomps` is `nothing`, it will be set to the number of frames in the reference cube when processed.

# References
* [Ren et al. 2018](http://adsabs.harvard.edu/abs/2018ApJ...852..104R) Non-negative Matrix Factorization: Robust Extraction of Extended Structures
"""
@with_kw struct NMF <: LinearAlgorithm
    ncomps::Union{Int,Nothing} = nothing
end

function ADI.decompose(alg::NMF, cube, angles, cube_ref; kwargs...)
    X = flatten(cube)
    X_ref = flatten(cube_ref)

    k = isnothing(alg.ncomps) ? size(cube, 1) : alg.ncomps
    k > size(cube, 1) && error("ncomps ($k) cannot be greater than the number of frames ($(size(cube, 1)))")
    # manually initialize W and H
    _, H = nndsvd(X_ref, k)
    W = X * H'
    solve!(CoordinateDescent{eltype(X)}(maxiter=100, α=0), X, W, H)
    return H, W
end

function ADI.decompose(alg::NMF, cube, angles; kwargs...)
    # X = clamp.(flatten(cube), 0, Inf) # clamp non-negative values
    X = flatten(cube)

    k = isnothing(alg.ncomps) ? size(cube, 1) : alg.ncomps
    k > size(cube, 1) && error("ncomps ($k) cannot be greater than the number of frames ($(size(cube, 1)))")

    W, H = nndsvd(X, k)
    solve!(CoordinateDescent{eltype(X)}(maxiter=100, α=0), X, W, H)
    return H, W
end


using MathProgComplex, OPFInstances
using Base.Test

brokentestset = Set([])

testcases = OrderedSet([("WB2", 1),
                        ("LMBM3", 1),
                        ("WB5", 1),
                        ("case6ww", 1),
                        ("case9mod", 1),
                        ("case9", 1),
                        ("case14", 1)])
                        # ("case24_ieee_rts", 1),
                        # ("case_ieee30", 1),
                        # ("case30pwl", 1),
                        # ("case30", 1)])

@testset "SDP hierarchy ACOPFs relaxation - rank 1" begin

    @testset "$pbsolved tests" for pbsolved in [:SOSRelaxation, :MomentRelaxation]
        @testset "$instance problem - rank $d" for (instance, d) in testcases
            problem_c, point = import_from_dat(getinstancepath("Matpower", "QCQP", instance))
            problem = pb_cplx2real(problem_c)

            relax_ctx = set_relaxation(problem; hierarchykind=:Real,
                                                d = d,
                                                params = Dict(:opt_outlev=>0,
                                                                :opt_solver=>testsolver,
                                                                :opt_pbsolved=>pbsolved))

            primobj, dualobj = run_hierarchy(problem, relax_ctx);

            ε_rel = OPFsols[(instance, 1)].rel_opt_gap
            ε_abs = max(primobj, dualobj) * ε_rel

            @show testsolver, pbsolved, instance, d
            @show primobj, dualobj,  OPFsols[(instance, 1)].primal_solvalue, ε_abs, instance

            @test primobj ≈ OPFsols[(instance, 1)].primal_solvalue atol=ε_abs
            @test dualobj ≈ primobj atol=ε_abs
        end
    end
end


export pradius, pradiusb

function pradiusk(As, p, k, pnorm)
    mean(map(A->opnorm(A, pnorm)^p, As))^(1/(p*k))
end

function pradius(s::DiscreteSwitchedLinearSystem, p, lift::Function; pnorm=Inf, ɛ=1e-2, forceub=false)
    ρpmpp = ρ(sum((rm -> lift(rm.A, p)).(s.resetmaps)))
    if forceub
        ρpmpp^(1/p)
    else
        (ρpmpp / ρA(s))^(1/p)
    end
end

function kozyakinlift(s, t, As)
    n = nstates(s)
    A = As[symbol(s, t)]
    kron(sparse([source(s, t)], [target(s, t)], [one(eltype(A))], n, n), A)
end
function pradius(s::ConstrainedDiscreteSwitchedLinearSystem, p, lift::Function; pnorm=Inf, ɛ=1e-2, forceub=false)
    Alifted = (rm -> lift(rm.A, p)).(s.resetmaps)
    ρpmpp = ρ(sum(kozyakinlift(s, t, Alifted) for t in transitions(s)))
    if forceub
        ρpmpp^(1/p)
    else
        (ρpmpp / ρA(s))^(1/p)
    end
end

abstract type PRadiusAlgorithm end
abstract type ExactPRadiusAlgorithm <: PRadiusAlgorithm end
function checkeven(p, algo::ExactPRadiusAlgorithm)
    if !iseven(p)
        throw(ArgumentError("Odd p is not supported yet for pradius computation with $algo algorithm"))
    end
end

export VeroneseLift, KroneckerLift, BruteForce
struct VeroneseLift <: ExactPRadiusAlgorithm end
struct KroneckerLift <: ExactPRadiusAlgorithm end
struct BruteForce <: PRadiusAlgorithm end

liftfunction(::VeroneseLift) = veroneselift
liftfunction(::KroneckerLift) = kroneckerlift
function kroneckerlift(A::AbstractMatrix, p::Integer)
    @assert iseven(p)
    kronpow(veroneselift(A, 2), div(p, 2))
end

function pradius(s::AbstractDiscreteSwitchedSystem, p, algo::ExactPRadiusAlgorithm=VeroneseLift(); kws...)
    checkeven(p, algo)
    pradius(s, p, liftfunction(algo); kws...)
end
# The Inf- and 1-norm of a matrices are easier to compute than the 2-norm
function pradius(s::DiscreteSwitchedLinearSystem, p, ::BruteForce; pnorm=Inf, ɛ=1e-2, forceub=false)
    As = map(rm -> rm.A, s.resetmaps)
    Ascur = As
    ρp = Float64[]
    m = ntransitions(s)
    k = 0
    while k < 2 || !isapprox(ρp[end-1], ρp[end], rtol=ɛ)
        k += 1
        l = length(Ascur)
        Asnew = similar(Ascur, l * m)
        for (i,A) in enumerate(As)
            for (j,B) in enumerate(Ascur)
                Asnew[(i-1)*l+j] = A*B
            end
        end
        Ascur = Asnew
        push!(ρp, pradiusk(Ascur, p, k, pnorm))
    end
    if forceub
        ρp[end] * length(As)^(1/p)
    else
        ρp[end]
    end
end

# p-radius algo needs to be exact to compute bounds on JSR
function pradiusb(s::AbstractDiscreteSwitchedSystem, p, algo::ExactPRadiusAlgorithm=VeroneseLift())
    ρpmp = pradius(s, p, algo, forceub=true)
    ρp = ρpmp / ρA(s)^(1/p)
    updateb!(s, ρp, ρpmp)
end


module ModuleDataNormalization
using MultipleViewGeometry.ModuleTypes, MultipleViewGeometry.ModuleOperators, MultipleViewGeometry.ModuleMathAliases
using StaticArrays, LinearAlgebra
export hartley_transformation, hartley_normalization!, hartley_normalization
include("hartley_transformation.jl")
end


include("triangle.jl");
include("genmesh.jl");
include("assemble_BEM.jl");
import .NEPCore.compute_Mder
import .NEPCore.size
import .NEPCore.compute_Mlincomb

using LinearAlgebra
"""
    struct BEM_NEP <: NEP

Represents a boundary element discretization. Create with a `mesh::Vector{Triangle}` using `nep=BEM_NEP(mesh)`.

"""
struct BEM_NEP <: NEP
    mesh::Vector{Triangle}
    n::Int
    gauss_order::Int
end

function BEM_NEP(mesh::Vector,gauss_order=3)
    return BEM_NEP(mesh,size(mesh,1),gauss_order);
end

function size(nep::BEM_NEP)
    return (nep.n,nep.n);
end
function size(nep::BEM_NEP,dim)
    return nep.n;
end
#
function compute_Mder(nep::BEM_NEP,λ::Number,der::Int=0)
    return assemble_BEM(λ, nep.mesh, nep.gauss_order,der)[:,:,1];
end
# Delegate the compute Mlincomb functions. This is quite inefficient.
compute_Mlincomb(nep::BEM_NEP,λ::Number,V::Union{AbstractMatrix,AbstractVector}, a::Vector) = compute_Mlincomb_from_Mder(nep,λ,V,a)
compute_Mlincomb(nep::BEM_NEP,λ::Number,V::Union{AbstractMatrix,AbstractVector}) = compute_Mlincomb(nep,λ,V, ones(eltype(V),size(V,2)))

function bem_fichera(N::Int=3)
    mesh=gen_ficheramesh(N)
    nep=BEM_NEP(mesh);
    precompute_quad!(nep.mesh,nep.gauss_order)
    return nep;
end


const Rmath = "libRmath-julia"

function rbeta(alpha::Float64, beta::Float64)
  ccall( (:rbeta, Rmath), Float64, (Float64, Float64), alpha, beta )
end

function rbeta(n::Int, alpha::Float64, beta::Float64)
  res = zeros(Float64, n)
  for i in 1:n
    res[i] = rbeta(alpha, beta)
  end
  res
end

#function pbeta(x::Float64, alpha::Float64, beta::Float64)
#  ccall( (:pbeta, Rmath), Float64, (Float64, Float64, Float64, Int64, Int64), x, alpha, beta, 0, 0)
#end

# Incomplete beta function.
#function betainc(x, a, b)
#  pbeta(float(x), float(a), float(b))
#end


module AlgebraicTraits

# useful because Complex is a leaf type in Julia
typealias AbstractComplex Union{Real, Complex}

include("associativity.jl")
include("closure.jl")
include("identity.jl")

include("monoid.jl")

end  # module AlgebraicTraits


function Pauli(i)

if i==0

return [1.0 0.0; 0.0 1.0]

elseif i==1

return [0.0 1.0; 1.0 0.0]

elseif i==2

return [0. -1.0im; 1.0im 0.0]

elseif i==3

return [1.0 0.0 ; 0.0 -1.0]

else throw(DomainError())

end

end


function LinearAlgebra.pinv(A::Array{T,2}) where T<:ForwardDiff.Dual
    m = size(A,1)
    n = size(A,2)

    Av = map(ForwardDiff.value, A)
    Ap = map(ForwardDiff.partials, A)
    Apt = Ap[reshape(1:m*n,(m,n))'] # taking the transpose

    Pv = pinv(Av)
    Pp = Pv*Pv'*Apt*(I - Av*Pv) + (I - Pv*Av)*Apt*Pv'*Pv - Pv*Ap*Pv

    P = map(T,Pv,Pp)

    return P
end


using VQE4NPP

s = 100 # Define the sum of the elements inside the target set.
n = 5 # Define the number of elements inside the target set.
d = 6 # Define the depth of the differentiable circuit.

set = RandIntNumSet(n,s) # Create the target number set based on configurartions
VQEtrain(set, depth = d) # Train the VQE for the corresponding Hamiltonian of the target set.


using Sigma

## 1D Ising Model
## ==============
nbits = 10
bits = Sigma.RandArray([Sigma.flip() for i = 1:nbits])
noisyeq(a,b) = Sigma.flip(ifelse(a == b, 1.0, 0.2))

ising_condition = (&)(map(noisyeq, bits[2:end], bits[1:end-1])...)
rand(bits, ising_condition)


function FiniteDifferenceAtParameters(P,ϵs)
    y = CF(P)
    ωᵢ = P.ω
    Almᵢ = P.Alm

    P.ω = ωᵢ + ϵs+im*0
    P.Alm,_,_ = ComputeAₗₘ(P.s,P.m,P.a*(P.ω), Almᵢ,P.lmax)
    yʼr = (CF(P) - y)/ϵs

    P.ω = ωᵢ + im*ϵs;
    P.Alm,_,_ = ComputeAₗₘ(P.s,P.m,P.a*(P.ω), Almᵢ,P.lmax)
    yʼi = (CF(P) - y)/ϵs

    P.ω = ωᵢ
    P.Alm = Almᵢ
    return yʼr,yʼi,y
end

function Newton!(P,ϵs)
    ω₀ = P.ω
    Alm₀ = P.Alm
    function f!(F,x)
        ωᵢ = P.ω
        Almᵢ = P.Alm
        P.ω = x[1]+im*x[2]
        P.Alm,_,_ = ComputeAₗₘ(P.s,P.m,P.a*(P.ω), P.Alm,P.lmax)
        val = CF(P)
        P.ω = ωᵢ
        P.Alm = Almᵢ
        F[1] = real(val)
        F[2] = imag(val)
    end

    function j!(J,x)
        ω = x[1]+im*x[2]
        P.ω = ω
        ydr,ydi,y = FiniteDifferenceAtParameters(P,ϵs)
        J[1,1] = real(ydr)
        J[1,2] = real(ydi)
        J[2,1] = imag(ydr)
        J[2,2] = imag(ydi)
    end

    function fj!(F,J,x)
        ωᵢ = P.ω
        Almᵢ = P.Alm
        P.ω = x[1]+im*x[2]
        P.Alm,_,_ = ComputeAₗₘ(P.s,P.m,P.a*(P.ω), Almᵢ,P.lmax)
        val = CF(P)
        F[1] = real(val)
        F[2] = imag(val)

        P.ω = x[1]+im*x[2] + ϵs+im*0
        P.Alm,_,_ = ComputeAₗₘ(P.s,P.m,P.a*(P.ω), Almᵢ,P.lmax)
        yʼr = (CF(P) - val)/ϵs

        P.ω = x[1]+im*x[2] + im*ϵs
        P.Alm,_,_ = ComputeAₗₘ(P.s,P.m,P.a*(P.ω), Almᵢ,P.lmax)
        yʼi = (CF(P) - val)/ϵs

        P.ω = ωᵢ
        P.Alm = Almᵢ

        J[1,1] = real(yʼr )
        J[1,2] = real(yʼi)
        J[2,1] = imag(yʼr )
        J[2,2] = imag(yʼi)
    end

    x0 = [real(ω₀),imag(ω₀)]
    df = OnceDifferentiable(f!, j!, fj!,x0,similar(x0))
    sol = nlsolve(df, x0, ftol = 1e-14)

    ωnew = sol.zero
    P.ω = ωnew[1] + im*ωnew[2]
    Almnew,_,_ = ComputeAₗₘ(P.s,P.m,P.a*(P.ω), P.Alm,P.lmax)
    P.Alm = Almnew
    P
end


######################################################################
# observations.jl
#
# Contains everything needed for the observation function.
######################################################################

"""
`O(m::SearchDomain, x::Vehicle, theta, o::Obs)`

Arguments:

 * `m` is a `SearchDomain`
 * `x` is a `Vehicle`
 * `theta` is a possible jammer location
 * `o` is an observation, 0 to 35

Returns probability of observing `o` from `(xp, theta)` in domain `m`.
"""
function O(m::SearchDomain, x::Vehicle, theta, o::Obs)

	# Calculate true bearing, and find distance to bin edges
	xp = (x.x, x.y)
	ang_deg = true_bearing(xp, theta)
	rel_start, rel_end = rel_bin_edges(ang_deg, o)

	# now look at probability
	#p = cdf(m.d, deg2rad(rel_end)) - cdf(m.d, deg2rad(rel_start))
	p = cdf(m.d, rel_end) - cdf(m.d, rel_start)
	return p
end

function O(m::SearchDomain, xv::Float64, yv::Float64, theta, o::Obs)

	# Calculate true bearing, and find distance to bin edges
	xp = (xv, yv)
	ang_deg = true_bearing(xp, theta)
	rel_start, rel_end = rel_bin_edges(ang_deg, o)

	# now look at probability
	#p = cdf(m.d, deg2rad(rel_end)) - cdf(m.d, deg2rad(rel_start))
	p = cdf(m.d, rel_end) - cdf(m.d, rel_start)
	return p
end

"""
`observe(m::SearchDomain, X::VehicleSet)`

Returns a set of possible observations for vehicles in set `X`.
	The output is a vector of observations, one per vehicle.
"""
function observe(m::SearchDomain, X::VehicleSet)
	n = length(X)
	Z = Array(Obs, 0)
	for xi in X
		obs_probs = Array(Float64, 0)
		for o in 0:35
			push!(obs_probs, O(m, xi, m.theta, o))
		end
		push!(Z, sample(0:35, WeightVec(obs_probs)))
	end
	return Z
end

# Find the relative offset
function rel_bin_edges(bearing_deg, o::Obs)

	# calculate start, end degrees of bin
	start_deg, end_deg = bin2deg(o)

	# compute relative distance to true bearing
	rel_start = fit_180(bearing_deg - start_deg)
	rel_end = fit_180(bearing_deg - end_deg)

	# Make sure start is further left on number line
	if rel_end < rel_start
		temp = rel_start
		rel_start = rel_end
		rel_end = temp
	end

	# Handle odd case where we stradle 180 degrees off
	if (rel_start < 0) && (rel_end > 100)
		rel_start = rel_end
		rel_end += 10
	end

	return rel_start, rel_end
end

# Fits an angle into -180 to 180
# Assume the angle is within -360 to 360
function fit_180(angle)
	if angle > 180
		angle = 360 - angle
	elseif angle < -180
		angle += 360
	end
	return angle
end

# Find the true angle between UAV and jammer.
#
# Parameters:
#  xp - location of vehicle
#  theta - location of jammer
#
# Returns true angle, measured from north, in degrees.
function true_bearing(xp, theta)

	xr = theta[1] - xp[1]
	yr = theta[2] - xp[2]
	return mod(rad2deg(atan2(xr,yr)), 360)
end

# What bin does degree fall into?
function findbin(degree)
	deg_bin = round(Int, degree / 10, RoundNearestTiesAway)
	if deg_bin == 36
		deg_bin = 0
	end
	return deg_bin
end

# returns (start_deg, end_deg) integer tuple
function bin2deg(bin_deg::Int64)
	if bin_deg == 0
		start_val = -5
		end_val = 5
	else
		start_val = 10bin_deg - 5
		end_val  = 10bin_deg + 5
	end
	return (start_val, end_val)
end