module ConversionWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const AbstractIntegrator = @jimport org.hipparchus.ode.AbstractIntegrator
const AbstractPropagatorBuilder = @jimport org.orekit.propagation.conversion.AbstractPropagatorBuilder
const AbstractPropagatorConverter = @jimport org.orekit.propagation.conversion.AbstractPropagatorConverter
const AdamsBashforthIntegratorBuilder = @jimport org.orekit.propagation.conversion.AdamsBashforthIntegratorBuilder
const AdamsMoultonIntegratorBuilder = @jimport org.orekit.propagation.conversion.AdamsMoultonIntegratorBuilder
const AdditionalEquations = @jimport org.orekit.propagation.integration.AdditionalEquations
const AttitudeProvider = @jimport org.orekit.attitudes.AttitudeProvider
const BatchLSModel = @jimport org.orekit.estimation.leastsquares.BatchLSModel
const BatchLSODModel = @jimport org.orekit.estimation.leastsquares.BatchLSODModel
const Class = @jimport java.lang.Class
const ClassicalRungeKuttaIntegratorBuilder = @jimport org.orekit.propagation.conversion.ClassicalRungeKuttaIntegratorBuilder
const DSSTBatchLSModel = @jimport org.orekit.estimation.leastsquares.DSSTBatchLSModel
const DSSTForceModel = @jimport org.orekit.propagation.semianalytical.dsst.forces.DSSTForceModel
const DSSTKalmanModel = @jimport org.orekit.estimation.sequential.DSSTKalmanModel
const DSSTPropagator = @jimport org.orekit.propagation.semianalytical.dsst.DSSTPropagator
const DSSTPropagatorBuilder = @jimport org.orekit.propagation.conversion.DSSTPropagatorBuilder
const DataContext = @jimport org.orekit.data.DataContext
const DormandPrince54IntegratorBuilder = @jimport org.orekit.propagation.conversion.DormandPrince54IntegratorBuilder
const DormandPrince853IntegratorBuilder = @jimport org.orekit.propagation.conversion.DormandPrince853IntegratorBuilder
const EcksteinHechlerPropagatorBuilder = @jimport org.orekit.propagation.conversion.EcksteinHechlerPropagatorBuilder
const EulerIntegratorBuilder = @jimport org.orekit.propagation.conversion.EulerIntegratorBuilder
const FiniteDifferencePropagatorConverter = @jimport org.orekit.propagation.conversion.FiniteDifferencePropagatorConverter
const ForceModel = @jimport org.orekit.forces.ForceModel
const Frame = @jimport org.orekit.frames.Frame
const GillIntegratorBuilder = @jimport org.orekit.propagation.conversion.GillIntegratorBuilder
const GraggBulirschStoerIntegratorBuilder = @jimport org.orekit.propagation.conversion.GraggBulirschStoerIntegratorBuilder
const HighamHall54IntegratorBuilder = @jimport org.orekit.propagation.conversion.HighamHall54IntegratorBuilder
const IntegratedPropagatorBuilder = @jimport org.orekit.propagation.conversion.IntegratedPropagatorBuilder
const JString = @jimport java.lang.String
const JacobianPropagatorConverter = @jimport org.orekit.propagation.conversion.JacobianPropagatorConverter
const KalmanModel = @jimport org.orekit.estimation.sequential.KalmanModel
const KalmanODModel = @jimport org.orekit.estimation.sequential.KalmanODModel
const KeplerianPropagatorBuilder = @jimport org.orekit.propagation.conversion.KeplerianPropagatorBuilder
const List = @jimport java.util.List
const LutherIntegratorBuilder = @jimport org.orekit.propagation.conversion.LutherIntegratorBuilder
const MidpointIntegratorBuilder = @jimport org.orekit.propagation.conversion.MidpointIntegratorBuilder
const ModelObserver = @jimport org.orekit.estimation.leastsquares.ModelObserver
const NumericalPropagator = @jimport org.orekit.propagation.numerical.NumericalPropagator
const NumericalPropagatorBuilder = @jimport org.orekit.propagation.conversion.NumericalPropagatorBuilder
const ODEIntegratorBuilder = @jimport org.orekit.propagation.conversion.ODEIntegratorBuilder
const Object = @jimport java.lang.Object
const Orbit = @jimport org.orekit.orbits.Orbit
const OrbitType = @jimport org.orekit.orbits.OrbitType
const OsculatingToMeanElementsConverter = @jimport org.orekit.propagation.conversion.OsculatingToMeanElementsConverter
const ParameterDriversList = @jimport org.orekit.utils.ParameterDriversList
const PositionAngle = @jimport org.orekit.orbits.PositionAngle
const PropagationType = @jimport org.orekit.propagation.PropagationType
const Propagator = @jimport org.orekit.propagation.Propagator
const PropagatorBuilder = @jimport org.orekit.propagation.conversion.PropagatorBuilder
const PropagatorConverter = @jimport org.orekit.propagation.conversion.PropagatorConverter
const SpacecraftState = @jimport org.orekit.propagation.SpacecraftState
const TLE = @jimport org.orekit.propagation.analytical.tle.TLE
const TLEPropagatorBuilder = @jimport org.orekit.propagation.conversion.TLEPropagatorBuilder
const ThreeEighthesIntegratorBuilder = @jimport org.orekit.propagation.conversion.ThreeEighthesIntegratorBuilder
const TideSystem = @jimport org.orekit.forces.gravity.potential.TideSystem
const UnnormalizedSphericalHarmonicsProvider = @jimport org.orekit.forces.gravity.potential.UnnormalizedSphericalHarmonicsProvider

export AbstractPropagatorBuilder
export AbstractPropagatorConverter
export AdamsBashforthIntegratorBuilder
export AdamsMoultonIntegratorBuilder
export ClassicalRungeKuttaIntegratorBuilder
export DSSTPropagatorBuilder
export DormandPrince54IntegratorBuilder
export DormandPrince853IntegratorBuilder
export EcksteinHechlerPropagatorBuilder
export EulerIntegratorBuilder
export FiniteDifferencePropagatorConverter
export GillIntegratorBuilder
export GraggBulirschStoerIntegratorBuilder
export HighamHall54IntegratorBuilder
export IntegratedPropagatorBuilder
export JacobianPropagatorConverter
export KeplerianPropagatorBuilder
export LutherIntegratorBuilder
export MidpointIntegratorBuilder
export NumericalPropagatorBuilder
export ODEIntegratorBuilder
export OsculatingToMeanElementsConverter
export PropagatorBuilder
export PropagatorConverter
export TLEPropagatorBuilder
export ThreeEighthesIntegratorBuilder
export add_additional_equations
export add_force_model
export build_integrator
export build_kalman_model
export build_ls_model
export build_propagator
export convert
export copy
export equals
export get_adapted_propagator
export get_all_force_models
export get_attitude_provider
export get_class
export get_evaluations
export get_frame
export get_initial_orbit_date
export get_integrator_builder
export get_mass
export get_mu
export get_orbit_type
export get_orbital_parameters_drivers
export get_position_angle
export get_position_scale
export get_propagation_parameters_drivers
export get_rms
export get_selected_normalized_parameters
export hash_code
export notify
export notify_all
export reset_orbit
export set_attitude_provider
export set_mass
export to_string
export wait

include("abstract_propagator_builder.jl")
include("abstract_propagator_converter.jl")
include("adams_bashforth_integrator_builder.jl")
include("adams_moulton_integrator_builder.jl")
include("classical_runge_kutta_integrator_builder.jl")
include("dormand_prince54_integrator_builder.jl")
include("dormand_prince853_integrator_builder.jl")
include("dsst_propagator_builder.jl")
include("eckstein_hechler_propagator_builder.jl")
include("euler_integrator_builder.jl")
include("finite_difference_propagator_converter.jl")
include("gill_integrator_builder.jl")
include("gragg_bulirsch_stoer_integrator_builder.jl")
include("higham_hall54_integrator_builder.jl")
include("integrated_propagator_builder.jl")
include("jacobian_propagator_converter.jl")
include("keplerian_propagator_builder.jl")
include("luther_integrator_builder.jl")
include("midpoint_integrator_builder.jl")
include("numerical_propagator_builder.jl")
include("ode_integrator_builder.jl")
include("osculating_to_mean_elements_converter.jl")
include("propagator_builder.jl")
include("propagator_converter.jl")
include("three_eighthes_integrator_builder.jl")
include("tle_propagator_builder.jl")

end


# Turn off precompilation during development
__precompile__(false)

#-------------------------------------------------------------------------------
# Main file implementing LOWDER in Julia
#-------------------------------------------------------------------------------

module LOWDER

    # Load dependencies

    using LinearAlgebra
    using Printf

    # Load code
    include("lowder_types.jl")
    include("linear_models.jl")
    include("lovo_utils.jl")
    include("utils.jl")
    include("lowder_output.jl")

    function lowder(
                    func_list::Array{Function,1},
                    x::Vector{Float64},
                    a::Vector{Float64},
                    b::Vector{Float64};
                    δ::Float64=predefined_sample_radius(a, b),
                    Δ::Float64=(1.5 * δ),
                    m::Int64=(2 * length(x) + 1),
                    maxit::Int64=(1000 * length(x)),
                    maxfun::Int64=(1000 * length(func_list) * length(x)),
                    maxcrit::Int64=(m - 1),
                    maxalt::Int64=(m - 1),
                    nρmax::Int64=3,
                    Γmax::Int64=1,
                    verbose::Int64=0,
                    δmin::Float64=1.0e-4,
                    β::Float64=1.0,
                    τ1::Float64=0.6,
                    τ2::Float64=1.5,
                    τ3::Float64=2.0,
                    η::Float64=0.1,
                    η1::Float64=0.3,
                    η2::Float64=0.6,
                    filename::String=""
                    )

        # Calculates the search space dimension.
        n = length(x)

        # Calculates the number of functions that make up the objective function fmin
        r = length(func_list)

        # Verify algorithm initialization conditions.
        if m != ( n + 1 )
            m_min = n + 2
            m_max = convert(Int64, ( n + 1 ) * ( n + 2 ) / 2)
            @assert m_min ≤ m ≤ m_max "The number of interpolation points 'm' must satisfy m ∈ [$(m_min), $(m_max)] for quadratic models. In the case of linear models, 'm' must be defined as $(n + 1)."
        end
        @assert length(a) == n "The vector 'a' must have dimension $(n)."
        @assert length(b) == n "The vector 'b' must have dimension $(n)."
        @assert 0.0 < δ ≤ Δ "The radius of the sample set 'δ' must be positive and less or equal to the trust-region radius 'Δ'."
        @assert verify_initial_room(n, δ, a, b) "The radius of the initial sample set 'δ' is not suitable, it must satisfy a[i]- b[i] >= 2δ, for i = 1, ..., $(n)."
        @assert maxit > 0 "The parameter 'maxit' must be positive."
        @assert maxfun ≥ r "The parameter 'maxfun' must be greater than ou equal to $(r)."
        @assert maxcrit > 0 "The parameter 'maxcrit' must be positive."
        @assert maxalt > 0 "The parameter 'maxalt' must be positive."
        @assert nρmax > 0 "The parameter 'nρmax' must be positive."
        @assert Γmax > 0 "The parameter 'Γmax' must be positive."
        @assert δmin > 0.0 "The parameter 'δmin' must be positive."
        @assert β > 0.0 "The parameter 'β' must be positive."
        @assert 0.0 < τ1 < 1.0 "The parameter 'τ1' must be positive and less than one."
        @assert 1.0 ≤ τ2 "The parameter 'τ2' must be greater than or equal to one."
        @assert τ2 ≤ τ3 "The parameter 'τ3' must be greater than or equal to 'τ2'."
        @assert 0.0 < η1 < 1.0 "The parameter 'η1' must be positive and less than one."
        @assert 0 ≤ η < η1 "The parameter 'η' must be nonnegative and less than 'η1'."
        @assert η1 ≤ η2 "The parameter 'η2' must be greater than or equal to 'η1'."
        @assert 0 ≤ verbose ≤ 3 "The parameter 'verbose' must be 0, 1, 2 or 3."

        # Verifies if the usuary wants to save a file with information of the iterations.
        if filename == ""

            saveinfo = false

        else

            saveinfo = true

        end

        # Sets some useful constants.
        Δinit = Δ
        δcrit = exp( log( ( 1.0 - sqrt( eps( Float64 ) ) ) * δmin ) - maxcrit * log( τ1 ) )

        # Initializes useful variables, vectors, and matrices.
        nρ = 0                        # Auxiliary counter for simplified 'ρ' calculations.
        nΓ = 0                        # Auxiliary counter for radii adjustments phase.
        nit = 0                       # Counts the number of iterations.
        nf = 0                        # Counts the number of 'f_{i}' function evaluations.
        ncrit = 0                     # Counts the number of consecutive criticality iterations.
        nalt = 0                      # Counts the number of consecutive iterations of type ALTMOV, ignoring criticality iterations.
        fi_x = NaN                    # New function value.
        pred_red = NaN                # Predicted reduction.
        real_red = NaN                # Real reduction.
        ρ = NaN                       # Relative reduction.
        π = NaN                       # Stationarity measure.
        norm_d = NaN                  # Norm of the computed direction 'd'.
        δold = NaN                    # Radius of the sample set at the beginning of the iteration.
        Δold = NaN                    # Radius of the trust-region at the beginning of the iteration.
        full_calc = false             # Indicates that ρ was calculated using fmin.
        verif_new_info = true         # Indicates the need to verify the information obtained in the last iteration.
        ao = zeros(Float64, n)        # Difference between the lower bounds 'a' and the center of the sample set, given by 'xbase'.
        bo = zeros(Float64, n)        # Difference between the upper bounds 'b' and the center of the sample set, given by 'xbase'.
        d = zeros(Float64, n)         # TRSBOX or ALTMOV direction.
        aux_v = zeros(Float64, n)     # Auxiliar vector for workspace.
        aux_w = zeros(Float64, n)     # Auxiliar vector for workspace.
        active_set = zeros(Bool, n)   # Set of active constraints.
        imin_set = zeros(Bool, r)     # Set I_{min}(x)
        it_flag = :nonspecified       # Initializes the iteration flag.
        exit_flag = :nonspecified     # Initializes the exit flag.
        status_flag = :nonspecified   # Initializes the status flag.

        #-------------------- Preparations for the first iteration ---------------------

        # Create an empty structure called 'model', which can hold a linear or quadratic model.
        if n == (m - 1)

            model = LinearModel(n)

        end

        # Modifies the initial estimate 'x' to be suitable for building the first model.
        # Modifies 'ao' and 'bo' to store the 'a-x' and 'b-x' differences, respectively.
        correct_guess_bounds!(n, δ, a, b, x, ao, bo)

        # Computes the value of f_{min}('x') and an index 'imin_idx' ∈ I_{min}('x').
        fi_x, imin_idx = fmin_eval!(func_list, r, x, imin_set)

        # Updates de function call counter.
        nf += r

        # Creates the initial model by modifying the structure of the previous 'model'
        # and calculates the QR factorization of the matrix M associated with the model definition.
        construct_model!(func_list, imin_idx, δ, fi_x, x, ao, bo, model)

        # Updates de function call counter.
        nf += n - 1

        if saveinfo

            data_file = open(filename, "w")

        end

        # Returns if 'nf' exceeds 'maxfun'.
        if nf ≥ maxfun

            it_flag = :nonspecified
            exit_flag = :max_evaluations

            # Creates the LOWDEROutput.
            output = create_output(model, exit_flag, full_calc, nit, nf, δold, Δold, π)

            # Prints information about the iteration, exit flag and LOWDEROutput.
            print_info(model, output, it_flag, status_flag, full_calc, verbose, δ, Δ, π, ρ, pred_red, real_red, d)

            if saveinfo

                save_info!(model, it_flag, status_flag, full_calc, nit, nf, δ, Δ, π, ρ, pred_red, real_red, d, data_file)
                println(data_file, exit_flag)
                close(data_file)

            end

            return output

        end

        # Main loop
        while true

            it_flag = :nonspecified
            status_flag = :nonspecified

            # Saves old information about radii
            δold = δ
            Δold = Δ

            # Computes the stationarity measure π = ||P_{Ω}(xopt - g) - xopt||
            π = stationarity_measure(model, a, b)

            # Verifies if 'δ' and 'π' are less than or equal to 'δmin' and 'δmin / β', respectively.
            if ( ( δ ≤ δmin ) && ( β * π ≤ δmin ) )

                # Sets exit flags
                exit_flag = :success
                break

            elseif ( ( δ ≤ δmin ) && ( Δ ≤ δmin ) && ( nalt > maxalt ) )

                # Sets exit flags
                exit_flag = :stalled
                break

            end

            if δ > β * π

                #------------------------------ Criticality phase ------------------------------

                # Sets iteration flag
                it_flag = :criticality

                # Update parameters and counter
                ncrit += 1
                nalt = 0
                ρ = NaN
                pred_red = NaN
                real_red = NaN
                norm_d = NaN

                if ncrit ≥ maxcrit

                    # It allows performing at most 'maxcrit' criticality iterations until the end of the algorithm.
                    δ = min( δ, δcrit )
                    ncrit = 0

                else

                    δ = max( τ1 * δ, δmin )

                end

            else

                # Update counter
                ncrit = 0

                #------------------------------- Step calculation ------------------------------

                status_flag = trsbox!(model, Δ, a, b, active_set, x, d, aux_v)
                norm_d = norm( d )

                # Verifies if 'trsbox' returns a unexpected direction.
                if norm_d == NaN

                    # Sets iteration and exit flags
                    it_flag = :trust_region
                    exit_flag = :numerical_error
                    break

                elseif norm_d ≥ 0.5 * Δ

                    #------------------------------- Step acceptance -------------------------------

                    mnew = model(x)
                    pred_red = model.fval[ model.kopt[] + 1 ] - mnew
                    full_calc = false

                    if pred_red < 0

                        it_flag = :nonspecified
                        exit_flag = :nondescent
                        break

                    else

                        if nρ < nρmax

                            ρ, real_red, fi_x, imin_idx = relative_reduction(model, func_list, pred_red, x)
                            nf +=1

                        else

                            full_calc = true
                            ρ, real_red, fi_x, imin_idx = relative_reduction!(model, func_list, r, pred_red, x, imin_set)
                            nf += r - 1
                            nρ = 0

                        end

                    end

                else

                    # Update parameters
                    ρ = 0.0
                    pred_red = 0.0
                    real_red = 0.0

                end

                #------------------------------- Radii updates ---------------------------------

                if ρ < η1

                    if ρ > 0.0

                        # Iterations of type 'trust_region' and 'bad_trust_region'.

                        Δ = max( τ1 * Δ ,  δ )

                    else

                        # Iterations of type 'altmov'.

                        if nalt > maxalt

                            δ = max( τ1 * δ, δmin )
                            Δ = max( τ1 * Δ, δmin )

                        else

                            Δ = max( τ1 * Δ ,  δ )

                        end

                    end

                elseif ( ρ > η2 ) && ( norm_d ≈ Δ )

                    δ *= τ2
                    Δ *= τ2

                end

            end

            # Chooses the point 't' that must leave the interpolation set. In the case of an ALTMOV call,
            # it also computes the new point 'x' and the direction 'd'.
            if ρ ≥ η

                # Sets iteration flag
                it_flag = :trust_region

                # Chooses the point that must leave the interpolation set 'model.Y'.
                t = choose_index_trsbox(model, x, aux_v)

                # Renitializes the counter
                nρ = 0
                nalt = 0

            elseif ρ > 0.0

                # Sets iteration flag
                it_flag = :bad_trust_region

                # Chooses the point that must leave the interpolation set 'model.Y'.
                t = choose_index_trsbox(model, x, aux_v)

                # Updates the counters.
                nρ += 1
                nalt = 0

            else

                # Sets iteration flag
                if it_flag != :criticality

                    it_flag = :altmov

                    if nalt > maxalt

                        nalt = 0

                    else

                        nalt += 1

                    end

                end

                # Chooses the point that must leave the interpolation set 'model.Y'.
                t = choose_index_altmov(model)

                # Computes the new interpolation point.
                status_flag = altmov!(model, t, δold, a, b, x, d, aux_v, aux_w, active_set)

                # Verifies if 'altmov' returns a unexpected direction.
                if norm( d ) == NaN

                    exit_flag = :numerical_error
                    break

                end

                # Computes the new function value.
                fi_x = fi_eval(func_list, model.imin[], x)

                # Updates the counters.
                nf += 1

            end

            #------------------------- Verifies output conditions --------------------------

            # Verifies if 'nit' exceeds 'maxit', or if 'nf' exceeds 'maxfun'.
            if nit ≥ maxit

                exit_flag = :max_iterations
                break

            elseif nf ≥ maxfun

                exit_flag = :max_evaluations
                break

            else

                if verbose ≥ 1

                    print_iteration(it_flag, status_flag, full_calc, nit, nf, model.imin[], δold, Δold, π, ρ, pred_red, real_red, model.fval[ model.kopt[] + 1 ], model.xopt, d)

                end

                if saveinfo

                    save_info!(model, it_flag, status_flag, full_calc, nit, nf, δold, Δold, π, ρ, pred_red, real_red, d, data_file)

                end

            end

            #--------------------------- Radii adjustments phase ---------------------------

            if ( ρ > 0.0 && full_calc ) || ( ρ ≥ η )

                if full_calc

                    if imin_set[ model.imin[] ]

                        imin_idx = model.imin[]

                    end

                else

                    fi_x, imin_idx = fmin_partial_eval!( func_list, r, model.imin[], fi_x, x, imin_set )
                    nf += r - 1

                    if imin_set[ model.imin[] ]

                        imin_idx = model.imin[]

                    end

                end

                if imin_idx != model.imin[]

                    if ρ ≥ η1

                        nΓ = 0

                    end

                    if nΓ ≤ Γmax

                        # Adjusts the radii to create a new model.
                        δ = τ3 * δold
                        Δ = max( τ3 * Δold, Δinit )
                        nΓ += 1

                    end

                    # Constructs a new model and set the new relative bounds 'ao' and 'bo'.
                    construct_new_model!( func_list, imin_idx, δ, fi_x, x, a, b, ao, bo, model )

                else

                    # Updates the model with TRSBOX information
                    update_model!( t, fi_x, x, model; trsbox_step = true)

                end

            elseif ρ > 0.0

                # Updates the model with TRSBOX information
                update_model!( t, fi_x, x, model; trsbox_step = true)

            else

                # Updates the model with ALTMOV information
                update_model!( t, fi_x, x, model)

            end

            # Increases iteration counter
            nit += 1

        end

        #---------------------- Preparations to finish execution  ----------------------

        if ( it_flag == :nonspecified ) || ( it_flag == :criticality ) || ( exit_flag == :numerical_error )

            verif_new_info = false

        end

        # Verifies the information obtained in the last iteration.
        if verif_new_info && ( fi_x < model.fval[ model.kopt[] + 1 ] )

            # Calculates the value of fmin if there is enough computational budget and it has not already been calculated.
            if ( ( maxfun - nf ) ≥ ( r - 1 ) ) && full_calc == false

                fi_x, imin_idx = fmin_partial_eval!( func_list, r, model.imin[], fi_x, x, imin_set )
                nf += r - 1
                full_calc = true

                # Verifies if 'model.imin' belongs to the 'imin_set' of the new point.
                if !( imin_set[ model.imin[] ] )

                    model.imin[] = imin_idx

                end

                # Updates 'model.xopt' and its function value.
                @. model.xopt = x
                model.fval[ model.kopt[] + 1 ] = fi_x

            else

                # Updates 'model.xopt' and its function value.
                @. model.xopt = x
                model.fval[ model.kopt[] + 1 ] = fi_x

            end

        end

        # Creates the LOWDEROutput.
        output = create_output(model, exit_flag, full_calc, nit, nf, δold, Δold, π)

        # Prints information about the iteration, exit flag and LOWDEROutput.
        print_info(model, output, it_flag, status_flag, full_calc, verbose, δold, Δold, π, ρ, pred_red, real_red, d)

        if saveinfo

            save_info!(model, it_flag, status_flag, full_calc, nit, nf, δold, Δold, π, ρ, pred_red, real_red, d, data_file)
            println(data_file, exit_flag)
            close(data_file)

        end

        return output

    end

end


# ----------------------------------------------------------------------------------- #
# Copyright (c) 2021 Varnerlab
# Robert Frederick Smith School of Chemical and Biomolecular Engineering
# Cornell University, Ithaca NY 14850
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED AS IS, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# ----------------------------------------------------------------------------------- #

function update_parameter_dictionary(parameter_dictionary::Dict{String,Any}, varargs...)::Union{Nothing, Dict{String,Any}}

    # Default: return the unmodified dictionary -
    @info "Default implementation of update_parameter_dictionary returns unmodified dictionary"

    # return parameter dictionary -
    return parameter_dictionary
end

function generate_default_parameter_dictionary()::Dict{String,Any}

    # initialize -
    parameter_dictionary = Dict{String,Any}()
    system_type_flag = 	:CF_TXTL

    try

        # load the stoichiometric_matrix (SM) and degradation_dilution_matrix (DM) -
        SM = readdlm(joinpath(_PATH_TO_NETWORK,"Network.dat"))
        DM = readdlm(joinpath(_PATH_TO_NETWORK,"Degradation.dat"))

        # build the species initial condition array -
        initial_condition_array = [
			5.0	;	#	1	gene_venus	units: nM
			0.0	;	#	2	mRNA_venus	units: nM
			0.0	;	#	3	P_venus	units: nM

			# translation capacity -
			100.0	;	#	4	translation capacity	units: NA
		]

        # setup the system dimension -

		number_of_transcription_processes = 1
		number_of_translation_processes = 1

        # build the system species concentration array -
        system_concentration_array = [
			0.07	;	#	RNAP	units: µM
			2.3	;	#	RIBOSOME	units: µM
			0.035	;	#	P_σ70	units: µM
		]

        # get the biophysical parameters for this system type -
        biophysical_parameters_dictionary = get_biophysical_parameters_dictionary_for_system(system_type_flag)

        # setup the model parameter array -
        model_parameter_array = [

			# control parameters: gene_venus
			0.001	;	#	1	W_gene_venus
			1.0		;	#	2	W_gene_venus_P_σ70
			1.0		;	#	3	K_gene_venus_P_σ70
			1.0		;	#	4	n_gene_venus_P_σ70

			# transcription kinetic limit parameters: gene_venus
			1.0		;	#	5	K_gene_venus
			1.0		;	#	6	𝛕_gene_venus

			# translation kinetic limit parameters: P_venus
			1.0		;	#	7	K_P_venus
			1.0		;	#	8	𝛕_P_venus

			# species degradation constants -
			0.001		;	#	9	𝛳_mRNA_venus
			0.0001		;	#	10	𝛳_P_venus
		]

        # setup the parameter symbol index map -
        model_parameter_symbol_index_map = Dict{Symbol,Int}()
		model_parameter_symbol_index_map[:W_gene_venus] = 1
		model_parameter_symbol_index_map[:W_gene_venus_P_σ70] = 2
		model_parameter_symbol_index_map[:K_gene_venus_P_σ70] = 3
		model_parameter_symbol_index_map[:n_gene_venus_P_σ70] = 4
		model_parameter_symbol_index_map[:K_gene_venus] = 5
		model_parameter_symbol_index_map[:𝛕_gene_venus] = 6
		model_parameter_symbol_index_map[:K_P_venus] = 7
		model_parameter_symbol_index_map[:𝛕_P_venus] = 8
		model_parameter_symbol_index_map[:𝛳_mRNA_venus] = 9
		model_parameter_symbol_index_map[:𝛳_P_venus] = 10

        # setup the inverse parameter symbol index map -
        inverse_model_parameter_symbol_index_map = Dict{Int, Symbol}()
		inverse_model_parameter_symbol_index_map[1] = :W_gene_venus
		inverse_model_parameter_symbol_index_map[2] = :W_gene_venus_P_σ70
		inverse_model_parameter_symbol_index_map[3] = :K_gene_venus_P_σ70
		inverse_model_parameter_symbol_index_map[4] = :n_gene_venus_P_σ70
		inverse_model_parameter_symbol_index_map[5] = :K_gene_venus
		inverse_model_parameter_symbol_index_map[6] = :𝛕_gene_venus
		inverse_model_parameter_symbol_index_map[7] = :K_P_venus
		inverse_model_parameter_symbol_index_map[8] = :𝛕_P_venus
		inverse_model_parameter_symbol_index_map[9] = :𝛳_mRNA_venus
		inverse_model_parameter_symbol_index_map[10] = :𝛳_P_venus

        # specific growth rate (default units: h^-1)
        μ = 0.0 # default units: h^-1

        # translation capacity half-life (default units: h^-1) -
		half_life_translation_capacity=value(biophysical_parameters_dictionary,string(:half_life_translation_capacity))	# units: h^-1
		half_life_translation_capacity = half_life_translation_capacity*(3600)

        # == DO NOT EDIT BELOW THIS LINE ========================================================== #
        parameter_dictionary["initial_condition_array"] = initial_condition_array
        parameter_dictionary["number_of_states"] = length(initial_condition_array)
        parameter_dictionary["number_of_transcription_processes"] = number_of_transcription_processes
        parameter_dictionary["number_of_translation_processes"] = number_of_translation_processes
        parameter_dictionary["system_concentration_array"] = system_concentration_array
        parameter_dictionary["biophysical_parameters_dictionary"] = biophysical_parameters_dictionary
        parameter_dictionary["model_parameter_array"] = model_parameter_array
        parameter_dictionary["model_parameter_symbol_index_map"] = model_parameter_symbol_index_map
        parameter_dictionary["inverse_model_parameter_symbol_index_map"] = inverse_model_parameter_symbol_index_map
        parameter_dictionary["stoichiometric_matrix"] = SM
        parameter_dictionary["dilution_degradation_matrix"] = DM
        parameter_dictionary["specific_growth_rate"] = μ
        parameter_dictionary["half_life_translation_capacity"] = half_life_translation_capacity

        # return -
        return parameter_dictionary
        # ========================================================================================= #
    catch error
        throw(error)
    end
end


using PowerModels
using Graphs

type Graph
    graph               ::SimpleGraph
    vertex_map          ::Dict{Int,Any}
    edge_map            ::Dict{Int,Any}
    bus_to_vertex_map   ::Dict{Int,Int}
    data                ::Dict{AbstractString,Any}
    ref                 ::Dict{Symbol,Any}
    function Graph()
        gm = new()
        gm.graph = simple_graph(0, is_directed=false)
        gm.vertex_map = Dict{Int,Any}()
        gm.edge_map = Dict{Int,Any}()
        gm.bus_to_vertex_map = Dict{Int,Int}()
        gm.data = Dict{AbstractString,Any}()
        return gm
    end
end

type VertexNeighbors
    id          ::Int
    edge_ids    ::Vector{Int}
    function VertexNeighbors()
        vn = new()
        vn.id = 0
        vn.edge_ids = Int[]
        return vn
    end
end

type Attack
    bus_ids             ::Vector{Int}
    branch_ids          ::Vector{Int}
    total_load          ::Float64
    function Attack()
        a = new()
        a.bus_ids = Int[]
        a.branch_ids = Int[]
        a.total_load = 0.0
        return a
    end
end

function populate_gm(gm::Graph, buses::Dict{Int,Any}, branches::Dict{Int,Any})
    i = 1
    for bus in buses
        add_vertex!(gm.graph)
        key = bus[1]
        value = bus[2]
        gm.vertex_map[i] = value
        gm.bus_to_vertex_map[key] = i
        i += 1
    end

    i = 1
    for (key,value) in branches
        add_edge!(gm.graph, gm.bus_to_vertex_map[value["f_bus"]], gm.bus_to_vertex_map[value["t_bus"]])
        gm.edge_map[i] = value
        i += 1
    end

    for edge in edges(gm.graph)
        @assert gm.vertex_map[edge.source]["index"] == gm.edge_map[edge.index]["f_bus"]
        @assert gm.vertex_map[edge.target]["index"] == gm.edge_map[edge.index]["t_bus"]
    end
    return
end

function get_vertices(gm::Graph; degree = 1)
    vertices_with_degree = Dict(v => out_degree(v, gm.graph) for v in vertices(gm.graph))
    vertices_with_degree = filter((v,deg) -> deg == degree, vertices_with_degree)
    neighbors = Dict(v => collect(out_edges(v, gm.graph)) for v in collect(keys(vertices_with_degree)))

    vertex_neighbors = VertexNeighbors[]
    for (v,n) in neighbors
        if gm.vertex_map[v]["bus_type"] == 2 || gm.vertex_map[v]["bus_type"] == 3
            continue
        end
        vn = VertexNeighbors()
        vn.id = gm.vertex_map[v]["index"]
        for edge in n
            push!(vn.edge_ids, gm.edge_map[edge.index]["index"])
        end
        push!(vertex_neighbors, vn)
    end

    return vertex_neighbors
end

function create_n_1_graph(gm::Graph, buses, branches, line)
    i = 1
    for bus in buses
        add_vertex!(gm.graph)
        key = bus[1]
        value = bus[2]
        gm.vertex_map[i] = value
        gm.bus_to_vertex_map[key] = i
        i += 1
    end

    i = 1
    for branch in branches
        if branch == line
            continue
        end
        key = branch[1]
        value = branch[2]
        add_edge!(gm.graph, gm.bus_to_vertex_map[value["f_bus"]], gm.bus_to_vertex_map[value["t_bus"]])
        gm.edge_map[i] = value
        i += 1
    end

    for edge in edges(gm.graph)
        @assert gm.vertex_map[edge.source]["index"] == gm.edge_map[edge.index]["f_bus"]
        @assert gm.vertex_map[edge.target]["index"] == gm.edge_map[edge.index]["t_bus"]
    end

    return
end

function check_graph_for_isolation(gm::Graph, buses, branches)
    components = connected_components(gm.graph)
    isolated_load_vertices = Attack[]
    if length(components) == 1
        return Attack[]
    end
    ref = gm.ref
    for component in components
        is_gen = false
        for v in component
            i = gm.vertex_map[v]["index"]
            if buses[i]["bus_type"] == 2 || buses[i]["bus_type"] == 3
                total_pmin = sum([abs(ref[:gen][g]["pmin"]) for g in ref[:bus_gens][i]])
                total_pmax = sum([abs(ref[:gen][g]["pmax"]) for g in ref[:bus_gens][i]])
                is_condensor = (total_pmin == 0 && total_pmax == 0) ? true : false
                if is_condensor
                    continue
                else
                    is_gen = true
                    break
                end
            end
        end
        if !(is_gen)
            load = Attack()
            for v in component
                push!(load.bus_ids, gm.vertex_map[v]["index"])
                load.total_load += abs(buses[gm.vertex_map[v]["index"]]["pd"])
            end
            push!(isolated_load_vertices, load)

        end
    end

    return isolated_load_vertices
end

function remove_redundant_attacks(attacks::Vector{Attack})
    redundant_attack_indexes = Set()
    num_attacks = length(attacks)
    for i in 1:num_attacks
        for j in 1:num_attacks
            if j == i
                continue
            end
            num_elements_i = length(attacks[i].bus_ids)
            num_elements_j = length(attacks[j].bus_ids)
            attack_set_i = Set(attacks[i].bus_ids)
            attack_set_j = Set(attacks[j].bus_ids)
            intersection = intersect(attack_set_i, attack_set_j)
            if attack_set_i == intersection
                push!(redundant_attack_indexes, i)
            end
            if attack_set_j == intersection
                push!(redundant_attack_indexes, j)
            end
        end
    end
    deleteat!(attacks, sort(collect(redundant_attack_indexes)))
    return
end


const EMPTY = 0x0000000000000000
const FULL = 0xffffffffffffffff

const NIGHT_CLEAR_FILES = [0x3f3f3f3f3f3f3f3f, 0x7f7f7f7f7f7f7f7f,
                           0xfefefefefefefefe, 0xfcfcfcfcfcfcfcfc,
                           0xfcfcfcfcfcfcfcfc, 0xfefefefefefefefe,
                           0x7f7f7f7f7f7f7f7f, 0x3f3f3f3f3f3f3f3f]

const NIGHT_JUMPS = [-10, -17, -15, -6, 10, 17, 15, 6]

const FILE_SHIFTS = [56, 48, 40, 32, 24, 16, 8, 0]
const MASK_FILE_A = 0x8080808080808080
const MASK_FILE_B = 0x4040404040404040
const MASK_FILE_C = 0x2020202020202020
const MASK_FILE_D = 0x1010101010101010
const MASK_FILE_E = 0x0808080808080808
const MASK_FILE_F = 0x0404040404040404
const MASK_FILE_G = 0x0202020202020202
const MASK_FILE_H = 0x0101010101010101
const MASK_FILES = [0x8080808080808080,
					0x4040404040404040,
					0x2020202020202020,
					0x1010101010101010,
					0x0808080808080808,
					0x0404040404040404,
					0x0202020202020202,
					0x0101010101010101]

const RANK_SHIFTS = [0, 8, 16, 24, 32, 40, 48, 56]
const MASK_RANK_1 = 0x00000000000000ff
const MASK_RANK_2 = 0x000000000000ff00
const MASK_RANK_3 = 0x0000000000ff0000
const MASK_RANK_4 = 0x00000000ff000000
const MASK_RANK_5 = 0x000000ff00000000
const MASK_RANK_6 = 0x0000ff0000000000
const MASK_RANK_7 = 0x00ff000000000000
const MASK_RANK_8 = 0xff00000000000000
const MASK_RANKS = [0x00000000000000ff,
					0x000000000000ff00,
					0x0000000000ff0000,
					0x00000000ff000000,
					0x000000ff00000000,
					0x0000ff0000000000,
					0x00ff000000000000,
					0xff00000000000000]

const FRAME = MASK_RANK_1 | MASK_RANK_8 | MASK_FILE_A | MASK_FILE_H
const CORNERS = (0x8000000000000000 | 0x0100000000000000 |
				 0x0000000000000080 | 0x0000000000000001)

const CLEAR_FILE_A = 0x7f7f7f7f7f7f7f7f
const CLEAR_FILE_B = 0xbfbfbfbfbfbfbfbf
const CLEAR_FILE_C = 0xdfdfdfdfdfdfdfdf
const CLEAR_FILE_D = 0xefefefefefefefef
const CLEAR_FILE_E = 0xf7f7f7f7f7f7f7f7
const CLEAR_FILE_F = 0xfbfbfbfbfbfbfbfb
const CLEAR_FILE_G = 0xfdfdfdfdfdfdfdfd
const CLEAR_FILE_H = 0xfefefefefefefefe

const CLEAR_RANK_1 = 0xffffffffffffff00
const CLEAR_RANK_2 = 0xffffffffffff00ff
const CLEAR_RANK_3 = 0xffffffffff00ffff
const CLEAR_RANK_4 = 0xffffffff00ffffff
const CLEAR_RANK_5 = 0xffffff00ffffffff
const CLEAR_RANK_6 = 0xffff00ffffffffff
const CLEAR_RANK_7 = 0xff00ffffffffffff
const CLEAR_RANK_8 = 0x00ffffffffffffff

const KING_SHIFTS = [9, 8, 7, -1, -9, -8, -7, 1]
const KING_CLEAR_FILES = [0x7f7f7f7f7f7f7f7f, 0xffffffffffffffff,
						  0xfefefefefefefefe, 0xfefefefefefefefe,
						  0xfefefefefefefefe, 0xffffffffffffffff,
						  0x7f7f7f7f7f7f7f7f, 0x7f7f7f7f7f7f7f7f]
const KING_MASK_FILES = [0x8080808080808080, 0xffffffffffffffff,
						 0x0101010101010101, 0x0101010101010101,
						 0x0101010101010101, 0xffffffffffffffff,
						 0x8080808080808080, 0x8080808080808080]

const PAWN_SQUARES = 0x00ffffffffffff00

const ANTIDIAGONALS = [0x0102040810204080,
                 	   0x0001020408102040,
                 	   0x0000010204081020,
                 	   0x0000000102040810,
                 	   0x0000000001020408,
                 	   0x0000000000010204,
                 	   0x0000000000000102,
                 	   0x0000000000000001,
                 	   0x0001020408102040<<9,
                 	   0x0000010204081020<<18,
                 	   0x0000000102040810<<27,
                 	   0x0000000001020408<<36,
                 	   0x0000000000010204<<45,
                 	   0x0000000000000102<<54,
                 	   0x0000000000000001<<63]

const DIAGONALS = [0x8040201008040201,
             	   0x0080402010080402,
             	   0x0000804020100804,
             	   0x0000008040201008,
             	   0x0000000080402010,
             	   0x0000000000804020,
             	   0x0000000000008040,
             	   0x0000000000000080,
             	   0x0080402010080402<<7,
             	   0x0000804020100804<<14,
             	   0x0000008040201008<<21,
             	   0x0000000080402010<<28,
             	   0x0000000000804020<<35,
             	   0x0000000000008040<<42,
             	   0x0000000000000080<<49]


# the number of masked elements in each set of mosaic
function nmasked_perset_mask!(nmasked::AbstractVector{Int}, mosaic::SetMosaic, elmask::AbstractVector{Bool},
                              set2index::Union{Nothing,Dict{Int,Int}} = nothing)
    @inbounds for elix in eachindex(elmask)
        elmask[elix] || continue
        for setix in view(mosaic.setXelm, :, elix)
            if set2index === nothing
                nmasked[setix] += 1
            else
                ix = get(set2index, setix, 0)
                if ix > 0
                    nmasked[ix] += 1
                end
            end
        end
    end
    return nmasked
end

# the number of masked elements in each set of mosaic
function nmasked_perset(mosaic::SetMosaic, elmasks::AbstractVector,
                        set2index::Union{Nothing,Dict{Int,Int}} = nothing)
    res = zeros(Int, set2index === nothing ? nsets(mosaic) : length(set2index), length(elmasks))
    @inbounds for (expix, elmask) in enumerate(elmasks)
        @assert length(elmask) == nelements(mosaic)
        nmasked_perset_mask!(view(res, :, expix), mosaic, elmask, set2index)
    end
    return res
end

# the number of masked elements in each set of mosaic
function nmasked_perset(mosaic::SetMosaic, elmasks::AbstractMatrix{Bool},
                        set2index::Union{Nothing,Dict{Int,Int}} = nothing)
    @assert size(elmasks, 1) == nelements(mosaic)
    res = zeros(Int, set2index === nothing ? nsets(mosaic) : length(set2index), size(elmasks, 2))
    @inbounds for expix in 1:size(elmasks, 2)
        nmasked_perset_mask!(view(res, :, expix), mosaic, view(elmasks, :, expix), set2index)
    end
    return res
end

# info for the set × mask overlap for the MaskedSetMosaic
struct MaskOverlap
    nmasked::Int    # number of masked elements
    nunmasked::Int  # number of unmasked elements
    used::Bool      # true if used for the cover problem
end

"""
[`SetMosaic`](@ref) with the elements masks (selections) on top.
Sets that are not overlapping with the masks are excluded(skipped) from `MaskedSetMosaic`.
Optionally, the filtering can include testing for the minimal overlap significance P-value.

The tiles of non-overlapped sets are removed, the tiles that have identical membership
for all the masked sets are squashed into a single tile.

## Type parameters
* `T`: the type of elements
* `S`: the type of set ids
* `E`: the type of experiment ids
"""
mutable struct MaskedSetMosaic{T, S, E} <: AbstractWeightedSetMosaic{T, S, E, Float64}
    original::SetMosaic{T, S}       # original mosaic
    elmasks::BitMatrix              # elements×masks
    elunionmask::BitVector          # all masked elements
    total_masked::Vector{Int}       # total number of masked elements for each mask

    # info for overlaps with all experiment masks
    # experiments are ordered by experiment index
    # if there's no overlap, the value is "missing"
    setXexp_olaps::Matrix{MaskOverlap}

    loc2glob_setix::Vector{Int}         # local set index to global set index in original mosaic
    glob2loc_setix::SparseVector{Int}   # global set index (in orig. mosaic) to local set index

    ix2experiment::Vector{E}        # experiment index to the ID of the mask
    experiment2ix::Dict{E, Int}     # experiment Id to index
end

function MaskedSetMosaic(mosaic::SetMosaic{T, S}, elmasks::AbstractMatrix{Bool},
                         setXexp_olaps::Matrix{MaskOverlap},
                         set_indices::Vector{Int},
                         experiment_ids::Union{AbstractVector{E}, Nothing} = nothing
) where {T, S, E}
    size(elmasks, 1) == nelements(mosaic) ||
        throw(ArgumentError("Elements mask length ($(size(elmasks, 1))) should match the number of elements ($(nelements(mosaic)))"))
    nmasks = size(elmasks, 2)
    (experiment_ids === nothing) || (nmasks == length(experiment_ids)) ||
        throw(ArgumentError("Number of masks ($nmasks) should match the number of mask IDs ($(length(masks_ids)))"))
    bit_elmasks = convert(BitMatrix, elmasks)
    bit_elunionmask = dropdims(any(bit_elmasks, dims=2), dims=2)
    @assert length(bit_elunionmask) == size(bit_elmasks, 1)
    MaskedSetMosaic(
        mosaic, bit_elmasks, bit_elunionmask,
        dropdims(sum(elmasks, dims=1), dims=1),
        setXexp_olaps,
        set_indices, SparseVector(nsets(mosaic), set_indices, collect(eachindex(set_indices))),
        collect(experiment_ids !== nothing ? experiment_ids : 1:nmasks),
        # FIXME use vector container if experiment_ids === nothing or if ids are integers
        Dict(id => ix for (ix, id) in enumerate(experiment_ids !== nothing ? experiment_ids : 1:nmasks)))
end

"""
    mask(mosaic::SetMosaic, elmasks::AbstractMatrix{Bool};
         [experiment_ids::Union{AbstractVector, AbstractSet, Nothing} = nothing],
         [min_nmasked=1], [max_setsize=nothing],
         [max_overlap_logpvalue=0.0]) -> MaskedSetMosaic

Construct [`MaskedSetMosaic`](@ref) from the [`SetMosaic`](@ref) and the collection of element masks.

## Arguments
* `min_nmasked`: the minimal number of masked elements in a set to include in the mosaic
* `max_setsize` (optional): ignore the annotation sets bigger than the specified size
* `max_overlap_logpvalue`: the threshold of Fisher's Exact Test log P-value of the overlap
   between the set and the mask for the inclusion of the set into the mosaic.
   *0* accepts all sets.
"""
function mask(mosaic::SetMosaic, elmasks::AbstractMatrix{Bool};
              experiment_ids::Union{AbstractVector, Nothing} = nothing,
              min_nmasked::Integer=1, max_setsize::Union{Integer, Nothing} = nothing,
              max_overlap_logpvalue::Number=0.0, # 0.0 would accept any overlap (as log(Fisher Exact Test P-value))
              max_min_overlap_logpvalue::Number=max_overlap_logpvalue
)
    size(elmasks, 1) == nelements(mosaic) ||
        throw(ArgumentError("Elements mask rows ($(size(elmasks, 1))) should match the number of elements ($(nelements(mosaic)))"))

    # calculate overlap sizes for each set and each mask
    nmasked_orgsets = nmasked_perset(mosaic, elmasks)

    # get the sets that pass the overlap filters
    setixs = Vector{Int}()
    ntotal = nelements(mosaic)
    ntotal_masked = dropdims(sum(elmasks, dims=1), dims=1)
    @inbounds for (setix, nmasked_permask) in enumerate(eachslice(nmasked_orgsets, dims=1))
        nset = setsize(mosaic, setix)
        min_logpvalue = 0.0
        for (expix, nmasked) in enumerate(nmasked_permask)
            (nmasked < min_nmasked) && continue # skip small overlap
            (max_setsize !== nothing) && (nset > max_setsize) && continue # skip very big and generic sets
            overlap_logpvalue = logpvalue(nmasked, nset, ntotal_masked[expix], ntotal)
            (overlap_logpvalue > max_overlap_logpvalue) && continue # skip non-signif overlaps
            min_logpvalue = min(min_logpvalue, overlap_logpvalue)
            if min_logpvalue <= max_min_overlap_logpvalue
                # set has passed the criteria
                push!(setixs, setix)
                break
            end
        end
    end

    # fill the overlaps matrix
    setXexp_olaps = Matrix{MaskOverlap}(undef, length(setixs), size(elmasks, 2))
    @inbounds for (expix, nmasked_perset) in enumerate(eachslice(nmasked_orgsets, dims=2))
        mask_olaps = view(setXexp_olaps, :, expix)
        masksize = ntotal_masked[expix]
        for (i, setix) in enumerate(setixs)
            nset = setsize(mosaic, setix)
            nmasked = nmasked_perset[setix]
            overlap_pvalue = logpvalue(nmasked, nset, masksize, ntotal)
            mask_olaps[i] = MaskOverlap(nmasked, nset - nmasked,
                                        overlap_pvalue <= max_overlap_logpvalue)
        end
    end

    return MaskedSetMosaic(mosaic, elmasks, setXexp_olaps, setixs, experiment_ids)
end

function mask(mosaic::SetMosaic{T}, elmasks::AbstractVector #= iterable with eltype()==Set{T} =#;
              kwargs...) where T
    @assert eltype(elmasks) === Set{T}
    mask(mosaic, Bool[in(e, elmask::Set{T}) for e in mosaic.ix2elm, elmask in elmasks]; kwargs...)
end

function mask(mosaic::SetMosaic{T}, elmasks::AbstractDict #= iterable with eltype()==Set{T} =#;
              kwargs...) where T
    @assert valtype(elmasks) === Set{T}
    mask(mosaic, Bool[in(e, elmask::Set{T}) for e in mosaic.ix2elm, elmask in values(elmasks)];
         experiment_ids = collect(keys(elmasks)), kwargs...)
end

nelements(mosaic::MaskedSetMosaic) = nelements(mosaic.original)
nsets(mosaic::MaskedSetMosaic) = size(mosaic.setXexp_olaps, 1) # only sets overlapping with masks
nsets(mosaic::MaskedSetMosaic, expix::Integer) = # number of sets overlapping with given mask
    sum(olap -> olap.used && olap.nmasked > 0,
        view(mosaic.setXexp_olaps, :, expix))
nmasks(mosaic::MaskedSetMosaic) = length(mosaic.total_masked)
nexperiments(mosaic::MaskedSetMosaic) = nmasks(mosaic)

_nmasked(mosaic::MaskedSetMosaic, expix::Int) = mosaic.total_masked[expix]
_nunmasked(mosaic::MaskedSetMosaic, expix::Int) = nelements(mosaic) - mosaic.total_masked[expix]

nmasked(mosaic::MaskedSetMosaic, mask::Any) = _nmasked(mosaic, mosaic.experiment2ix[mask])
nunmasked(mosaic::MaskedSetMosaic, mask::Any) = _nunmasked(mosaic, mosaic.experiment2ix[mask])

# FIXME swap exp and set args order
function overlap(mosaic::MaskedSetMosaic, exp::Any, set::Any)
    expix = mosaic.experiment2ix[exp] # throws if mask is not defined
    globsetix = mosaic.original.set2ix[set] # throws if set is not defined
    setix = mosaic.glob2loc_setix[globsetix]
    (setix == 0) && return missing # no overlap data for given set
    return mosaic.setXexp_olaps[setix, expix]
end

# FIXME swap mask and set args order
function logpvalue(mosaic::MaskedSetMosaic, exp::Any, set::Any)
    olap = overlap(mosaic, exp, set)
    ismissing(olap) && return missing
    return logpvalue(olap.nmasked, olap.nmasked + olap.nunmasked,
                     nmasked(mosaic, exp), nelements(mosaic))
end

function _setweight(mosaic::MaskedSetMosaic, setix::Integer, expix::Integer;
                    filter::Bool = true)
    olap = mosaic.setXexp_olaps[setix, expix]
    (ismissing(olap) || (filter && !olap.used)) && return missing
    return logpvalue(olap.nmasked, olap.nmasked + olap.nunmasked,
                     _nmasked(mosaic, expix), nelements(mosaic))
end

# copy everything, except the original mosaic (leave the reference to the same object)
Base.copy(mosaic::MaskedSetMosaic) =
    MaskedSetMosaic(mosaic.original, copy(mosaic.elmasks), copy(mosaic.elunionmask),
                    copy(mosaic.total_masked),
                    copy(mosaic.setXexp_olaps),
                    copy(mosaic.loc2glob_setix), copy(mosaic.glob2loc_setix),
                    copy(mosaic.ix2experiment), copy(mosaic.experiment2ix))


export TxDAQControllerParams, TxDAQController, controlTx

Base.@kwdef mutable struct TxDAQControllerParams <: DeviceParams
  phaseAccuracy::Float64
  amplitudeAccuracy::Float64
  controlPause::Float64
  maxControlSteps::Int64 = 20
  correctCrossCoupling::Bool = false
end
TxDAQControllerParams(dict::Dict) = params_from_dict(TxDAQControllerParams, dict)

struct ControlledChannel
  seqChannel::PeriodicElectricalChannel
  daqChannel::TxChannelParams
end

Base.@kwdef mutable struct TxDAQController <: VirtualDevice
  "Unique device ID for this device as defined in the configuration."
  deviceID::String
  "Parameter struct for this devices read from the configuration."
  params::TxDAQControllerParams
  "Flag if the device is optional."
	optional::Bool = false
  "Flag if the device is present."
  present::Bool = false
  "Vector of dependencies for this device."
  dependencies::Dict{String, Union{Device, Missing}}

  currTx::Union{Matrix{ComplexF64}, Nothing} = nothing
  desiredTx::Union{Matrix{ComplexF64}, Nothing} = nothing
  sinLUT::Union{Matrix{Float64}, Nothing} = nothing
  cosLUT::Union{Matrix{Float64}, Nothing} = nothing
  controlledChannels::Vector{ControlledChannel} = []
end

function init(tx::TxDAQController)
  @info "Initializing TxDAQController with ID `$(tx.deviceID)`."
  tx.present = true
end

neededDependencies(::TxDAQController) = [AbstractDAQ]
optionalDependencies(::TxDAQController) = []

function controlTx(txCont::TxDAQController, seq::Sequence, initTx::Union{Matrix{ComplexF64}, Nothing} = nothing)
  # Prepare and check channel under control
  daq = dependency(txCont, AbstractDAQ)
  seqControlledChannel = getControlledChannel(seq)
  missingControlDef = []
  txCont.controlledChannels = []
  for seqChannel in seqControlledChannel
    name = id(seqChannel)
    daqChannel = get(daq.params.channels, name, nothing)
    if isnothing(daqChannel) || isnothing(daqChannel.feedback) || !in(daqChannel.feedback.channelID, daq.refChanIDs)
      push!(missingControlDef, name)
    else
      push!(txCont.controlledChannels, ControlledChannel(seqChannel, daqChannel))
    end
  end
  if length(missingControlDef) > 0
    message = "The sequence requires control for the following channel " * join(string.(missingControlDef), ", ", " and") * ", but either the channel was not defined or had no defined feedback channel."
    throw(IllegalStateException(message))
  end

  # Check that channels only have one component
  if any(x -> length(x.components) > 1, seqControlledChannel)
    throw(IllegalStateException("Sequence has channel with more than one component. Such a channel cannot be controlled by this controller"))
  end

  if !isnothing(initTx)
    s = size(initTx)
    # Not square or does not fit controlled channel matrix
    if !(length(s) == 0 || all(isequal(s[1]), s))
      throw(IllegalStateException("Given initTx for control tx has dimenions $s that is either not square or does not match the amount of controlled channel"))
    end
  end

  # Prepare values
  if isnothing(initTx)
    txCont.currTx = convert(Matrix{ComplexF64}, diagm(ustrip.(u"V", [channel.daqChannel.limitPeak/10 for channel in txCont.controlledChannels])))
  else
    txCont.currTx = initTx
  end
  sinLUT, cosLUt = createLUTs(seqControlledChannel, seq::Sequence)
  txCont.sinLUT = sinLUT
  txCont.cosLUT = cosLUt
  Ω = calcDesiredField(seqControlledChannel)

  # Start Tx
  setTxParams(daq, txFromMatrix(txCont, txCont.currTx)...)
  startTx(daq)

  controlPhaseDone = false
  i = 1
  while !controlPhaseDone && i <= txCont.params.maxControlSteps
    @info "CONTROL STEP $i"
    period = currentPeriod(daq)
    uMeas, uRef = readDataPeriods(daq, 1, period + 1, acqNumAverages(seq))

    controlPhaseDone = doControlStep(txCont, seq, uRef, Ω)

    sleep(txCont.params.controlPause)
    i += 1
  end
  stopTx(daq)
  setTxParams(daq, txFromMatrix(txCont, txCont.currTx)...)
  return txFromMatrix(txCont, txCont.currTx)
end

getControlledChannel(seq::Sequence) = [channel for field in seq.fields if field.control for channel in field.channels if typeof(channel) <: PeriodicElectricalChannel]

function txFromMatrix(txCont::TxDAQController, Γ::Matrix{ComplexF64})
  amplitudes = Dict{String, Vector{Union{Float32, Nothing}}}()
  phases = Dict{String, Vector{Union{Float32, Nothing}}}()
  for (d, channel) in enumerate(txCont.controlledChannels)
    amps = []
    phs = []
    for (e, channel) in enumerate(txCont.controlledChannels)
      push!(amps, abs(Γ[d, e]))
      push!(phs, angle(Γ[d, e]))
    end
    amplitudes[id(channel.seqChannel)] = amps
    phases[id(channel.seqChannel)] = phs
  end
  return amplitudes, phases
end

function createLUTs(seqChannel::Vector{PeriodicElectricalChannel}, seq::Sequence)
  N = rxNumSamplingPoints(seq)
  D = length(seqChannel)
  cycle = ustrip(dfCycle(seq))
  base = ustrip(dfBaseFrequency(seq))
  dfFreq = [base/x.components[1].divider for x in seqChannel]

  sinLUT = zeros(N,D)
  cosLUT = zeros(N,D)
  for d=1:D
    Y = round(Int64, cycle*dfFreq[d] )
    for n=1:N
      sinLUT[n,d] = sin(2 * pi * (n-1) * Y / N) / N #sqrt(N)*2
      cosLUT[n,d] = cos(2 * pi * (n-1) * Y / N) / N #sqrt(N)*2
    end
  end
  return sinLUT, cosLUT
end

function doControlStep(txCont::TxDAQController, seq::Sequence, uRef, Ω::Matrix)

  Γ = calcFieldFromRef(txCont,seq, uRef)
  daq = dependency(txCont, AbstractDAQ)

  @info "reference Γ=" Γ

  if controlStepSuccessful(Γ, Ω, txCont)
    @info "Could control"
    return true
  else
    newTx = newDFValues(Γ, Ω, txCont)
    oldTx = txCont.currTx
    @info "oldTx=" oldTx
    @info "newTx=" newTx

    if checkDFValues(newTx, oldTx, Γ,txCont)
      txCont.currTx[:] = newTx
      setTxParams(daq, txFromMatrix(txCont, txCont.currTx)...)
    else
      @warn "New values are above voltage limits or are different than expected!"
    end

    @info "Could not control !"
    return false
  end
end

function calcFieldFromRef(txCont::TxDAQController, seq::Sequence, uRef)
  len = length(txCont.controlledChannels)
  Γ = zeros(ComplexF64, len, len)

  for d=1:len
    for e=1:len
      c = ustrip(u"T/V", txCont.controlledChannels[d].daqChannel.feedback.calibration)

      uVolt = float(uRef[1:rxNumSamplingPoints(seq),d,1])

      a = 2*sum(uVolt.*txCont.cosLUT[:,e])
      b = 2*sum(uVolt.*txCont.sinLUT[:,e])

      Γ[d,e] = c*(b+im*a)
    end
  end
  return Γ
end

function calcDesiredField(seqChannel::Vector{PeriodicElectricalChannel})
  temp = [ustrip(amplitude(ch.components[1])) * exp(im*ustrip(phase(ch.components[1]))) for ch in seqChannel]
  return convert(Matrix{ComplexF64}, diagm(temp))
end

function controlStepSuccessful(Γ::Matrix, Ω::Matrix, txCont::TxDAQController)

  if txCont.params.correctCrossCoupling
    diff = Ω - Γ
  else
    diff = diagm(diag(Ω)) - diagm(diag(Γ))
  end
  deviation = maximum(abs.(diff)) / maximum(abs.(Ω))
  @info "Ω = " Ω
  @info "Γ = " Γ
  @info "Ω - Γ = " diff
  @info "deviation = $(deviation)   allowed= $(txCont.params.amplitudeAccuracy)"
  return deviation < txCont.params.amplitudeAccuracy
end

function newDFValues(Γ::Matrix, Ω::Matrix, txCont::TxDAQController)

  κ = txCont.currTx
  if txCont.params.correctCrossCoupling
    β = Γ*inv(κ)
  else
    @show size(Γ), size(κ)
    β = diagm(diag(Γ))*inv(diagm(diag(κ)))
  end
  newTx = inv(β)*Ω

  @warn "here are the values"
  @show κ
  @show Γ
  @show Ω

  return newTx
end

function checkDFValues(newTx, oldTx, Γ, txCont::TxDAQController)

  calibFieldToVoltEstimate = [ustrip(u"V/T", ch.daqChannel.calibration) for ch in txCont.controlledChannels]
  calibFieldToVoltMeasured = abs.(diag(oldTx) ./ diag(Γ))

  @info "" calibFieldToVoltEstimate[1] calibFieldToVoltMeasured[1]

  deviation = abs.(1.0 .- calibFieldToVoltMeasured./calibFieldToVoltEstimate)

  @info "We expected $(calibFieldToVoltEstimate) and got $(calibFieldToVoltMeasured), deviation: $deviation"

  return all( abs.(newTx) .<  ustrip.(u"V", [channel.daqChannel.limitPeak for channel in txCont.controlledChannels]) ) && maximum( deviation ) < 0.2
end


using DrWatson
@quickactivate
using ArgParse
using GenerativeAD
import StatsBase: fit!, predict
using StatsBase
using BSON

s = ArgParseSettings()
@add_arg_table! s begin
   "max_seed"
        required = true
        arg_type = Int
        help = "seed"
    "dataset"
        required = true
        arg_type = String
        help = "dataset"
    "contamination"
    	arg_type = Float64
    	help = "contamination rate of training data"
    	default = 0.0
end
parsed_args = parse_args(ARGS, s)
@unpack dataset, max_seed, contamination = parsed_args

#######################################################################################
################ THIS PART IS TO BE PROVIDED FOR EACH MODEL SEPARATELY ################
modelname = "if"
# sample parameters, should return a Dict of model kwargs

function set_params(data)
	return (n_estimators=100, max_samples="auto",max_features=1.0)
end
function fit(data, parameters)
	# construct model - constructor should only accept kwargs
	model = GenerativeAD.Models.IForest(;parameters...)

	# fit train data
	try
		global info, fit_t, _, _, _ = @timed fit!(model, data[1][1])
	catch e
		# return an empty array if fit fails so nothing is computed
		return (fit_t = NaN,), []
	end

	# construct return information - put e.g. the model structure here for generative models
	training_info = (
		fit_t = fit_t,
		model = nothing
		)

	# now return the different scoring functions
	training_info, [(x->predict(model, x), parameters)]
end

####################################################################
################ THIS PART IS COMMON FOR ALL MODELS ################
# set a maximum for parameter sampling retries
try_counter = 0
max_tries = 10*max_seed
cont_string = (contamination == 0.0) ? "" : "_contamination-$contamination"
while try_counter < max_tries
    for seed in 1:max_seed
		savepath = datadir("experiments/tabular_clean_val_default/$(modelname)/$(dataset)/seed=$(seed)")
		mkpath(savepath)

		# get data
		data = GenerativeAD.load_data(dataset, seed=seed, contamination=contamination)

		# parameters
		parameters = set_params(data)

		@info "Trying to fit $modelname on $dataset with parameters $(parameters)..."
		# check if a combination of parameters and seed alread exists
		if GenerativeAD.check_params(savepath, parameters)
			# fit
			training_info, results = fit(data, parameters)
			# here define what additional info should be saved together with parameters, scores, labels and predict times
			save_entries = merge(training_info, (modelname = modelname, seed = seed, dataset = dataset, contamination = contamination))

			# now loop over all anomaly score funs
			for result in results
				GenerativeAD.experiment(result..., data, savepath; save_entries...)
			end
			global try_counter = max_tries + 1
		else
			@info "Model already present, trying new hyperparameters..."
			global try_counter += 1
		end
	end
end
(try_counter == max_tries) ? (@info "Reached $(max_tries) tries, giving up.") : nothing


function optimize!(structure::DeterministicEquivalent, optimizer::MOI.AbstractOptimizer, x₀::AbstractVector)
    # Sanity check
    backend(structure.model) === optimizer || error("Stochastic program optimizer has not been connected to the deterministically equivalent problem.")
    # Crash if supported
    if MOI.supports(optimizer, MOI.VariablePrimalStart(), MOI.VariableIndex)
        for (i, idx) in enumerate(all_decisions(structure))
            MOI.set(optimizer, MOI.VariablePrimalStart(), idx, x₀[i])
        end
    end
    # Run standard MOI optimization procedure
    MOI.optimize!(optimizer)
    return nothing
end

function supports_structure(::MOI.AbstractOptimizer, ::DeterministicEquivalent)
    return true
end

function check_loadable(::MOI.AbstractOptimizer, ::DeterministicEquivalent)
    return nothing
end

function set_master_optimizer!(structure::DeterministicEquivalent, optimizer)
    # Ensure decision bridges are added
    for bridge_type in structure.model.bridge_types
        JuMP._moi_add_bridge(structure.model.moi_backend, bridge_type)
    end
    return nothing
end

function set_master_optimizer_attribute!(structure::DeterministicEquivalent, attr::MOI.AbstractOptimizerAttribute, value)
    MOI.set(backend(structure.model), attr, value)
    return nothing
end

function set_subproblem_optimizer!(structure::DeterministicEquivalent, optimizer)
    return nothing
end

function set_subproblem_optimizer_attribute!(::DeterministicEquivalent, ::MOI.AbstractOptimizerAttribute, value)
    return nothing
end

function optimize!(structure::DeterministicEquivalent, optimizer::AbstractStructuredOptimizer)
    # Sanity check
    supports_structure(optimizer, typeof(structure)) || throw(UnsupportedStructure{typeof(optimizer), typeof(structure)}())
    # Ensure that no decisions are fixed in the first stage
    untake_decisions!(structure.model, structure.decision_variables[1])
    # Load structure
    load_structure!(optimizer, structure)
    # Run structure-exploiting optimization procedure
    MOI.optimize!(optimizer)
    return nothing
end


# ============================================
# Flowpipes not contiguous in time
# ============================================

"""
    MixedFlowpipe{N, D, FT<:AbstractFlowpipe, VOA<:VectorOfArray{N, D, Vector{FT}}} <: AbstractFlowpipe

Type that wraps a vector of flowpipes of the same time, such that they are
not necessarily contiguous in time.

### Fields

- `Fk`  -- vector of flowpipes
- `ext` -- (optional, default: empty) dictionary for extensions

### Notes

This type does not assume that the flowpipes are contiguous in time.
"""
struct MixedFlowpipe{N, RT<:AbstractReachSet{N}, FT<:AbstractFlowpipe} <: AbstractFlowpipe
    Fk::VectorOfArray{RT, 2, Vector{FT}}
    ext::Dict{Symbol, Any}
end

function MixedFlowpipe(Fk::Vector{FT}) where {N, RT<:AbstractReachSet{N}, FT<:Flowpipe{N, RT}}
    voa = VectorOfArray{RT, 2, Vector{FT}}(Fk)
    ext = Dict{Symbol, Any}()
    return MixedFlowpipe{N, RT, FT}(voa, ext)
end

function MixedFlowpipe(Fk::Vector{FT}, ext::Dict{Symbol, Any}) where {N, RT<:AbstractReachSet{N}, FT<:Flowpipe{N, RT}}
    voa = VectorOfArray{RT, 2, Vector{FT}}(Fk)
    return MixedFlowpipe{N, RT, FT}(voa, ext)
end

# interface functions
array(fp::MixedFlowpipe) = fp.Fk
flowpipe(fp::MixedFlowpipe) = fp
setrep(::Type{MixedFlowpipe{N, RT, FT}}) where {N, RT, FT} = RT
numrsets(fp::MixedFlowpipe) = mapreduce(length, +, fp)

# indexing: fp[j, i] returning the j-th reach-set of the i-th flowpipe
Base.getindex(fp::MixedFlowpipe, I::Int...) = getindex(fp.Fk, I...)

@inline tspan(fp::MixedFlowpipe) = error("for mixed flowpipes you should specify the index, as in `tspan(fp, i)`")
@inline tspan(fp::MixedFlowpipe, i::Int) = tspan(fp[i])

function Base.similar(fp::MixedFlowpipe{N, RT, FT}) where {N, RT, FT}
    return MixedFlowpipe(Vector{FT}())
end

function (fp::MixedFlowpipe)(t::Number)
    error("not implemented yet")
end

function (fp::MixedFlowpipe)(dt::TimeInterval)
    error("not implemented yet")
end

function overapproximate(fp::MixedFlowpipe, args...)
    return MixedFlowpipe([overapproximate(Fi, args...) for Fi in fp], fp.ext)
end

function project(fp::MixedFlowpipe, args...)
    return [project(F, args...) for F in fp]
end

# LazySets interface
function LazySets.ρ(d::AbstractVector, fp::MixedFlowpipe)
    return maximum(ρ(d, F) for F in array(fp))
end

function LazySets.σ(d::AbstractVector, fp::MixedFlowpipe)
    error("not implemented")
end

function LazySets.linear_map(M, fp::MixedFlowpipe)
    out = [linear_map(M, F) for F in fp]
    return MixedFlowpipe(out, fp.ext)
end

function LazySets.affine_map(M, b, fp::MixedFlowpipe)
    out = [affine_map(M, b, F) for F in fp]
    return MixedFlowpipe(out, fp.ext)
end


#!/usr/bin/julia

using Pkg
#Pkg.add("StatsModels")
using ExperimentalDesign, StatsModels, GLM, DataFrames, Distributions, Random
include("gcc.jl")

function generate_file_distribution()
	prefix = "processed"

	files = readdir(prefix)
	maximum = 0

	for file in files
		splitted = split(file, ".")
		base = splitted[1]
		ext = splitted[2]

		base_split = split(base, "-")
		num_func_str = base_split[1]
		num_funcs = parse(Int64, num_func_str)

		maximum = max(maximum, num_funcs)
	end

	names = fill(String[], maximum)

	for i in 1:maximum
		names[i] = String[]
	end

	for file in files
		splitted = split(file, ".")
		base = splitted[1]
		ext = splitted[2]

		base_split = split(base, "-")
		num_func_str = base_split[1]
		size_str = base_split[2]
		ret_str = base_split[3]

		num_funcs = parse(Int64, num_func_str)

		if num_funcs == 0
			continue
		end

		push!(names[num_funcs], file)
	end

	return names
end

function compute_fill_rate(names)
	len = length(names)
	acc = 0

	for i in 1:len
		if names[i] != []
			acc = acc + 1
		end
	end

	print(acc/len)
end

function y(x)
    return run_gcc(x)
end

# Used for DiscreteNonParametric(fp, prob), but use DiscreteUniform for now.
# fp   = [20 + 14*n for n in 0:14]
# prob = [1/15 for n in 0:14]

#random_design_generator = RandomDesign((
#          Distributions.DiscreteUniform(1, 300),     # File size, in number of functions
#          Distributions.DiscreteUniform(1, 10000),   # Min lto partition
#          Distributions.Bernoulli(0.5)))             # Run in parallel?
#
#println(random_design_generator)
#
#random_design = rand(random_design_generator, 8)
#random_design[!, :response] = y.(eachrow(random_design[:, :]))
#
#println(random_design)

names = generate_file_distribution()
compute_fill_rate(names)


module TestCallbacks

using Gurobi
using Random
using Test

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

const MOI = Gurobi.MOI

const GRB_ENV =
    isdefined(Main, :GRB_ENV) ? Main.GRB_ENV : Gurobi.Env(output_flag = 0)

function callback_simple_model()
    model = Gurobi.Optimizer(GRB_ENV)
    MOI.set(model, MOI.RawOptimizerAttribute("OutputFlag"), 0)
    MOI.set(model, MOI.RawOptimizerAttribute("Cuts"), 0)
    MOI.set(model, MOI.RawOptimizerAttribute("Presolve"), 0)
    MOI.set(model, MOI.RawOptimizerAttribute("Heuristics"), 0)
    MOI.Utilities.loadfromstring!(
        model,
        """
    variables: x, y
    maxobjective: y
    c1: x in Integer()
    c2: y in Integer()
    c3: x in Interval(0.0, 2.5)
    c4: y in Interval(0.0, 2.5)
""",
    )
    x = MOI.get(model, MOI.VariableIndex, "x")
    y = MOI.get(model, MOI.VariableIndex, "y")
    return model, x, y
end

function callback_knapsack_model()
    model = Gurobi.Optimizer(GRB_ENV)
    MOI.set(model, MOI.RawOptimizerAttribute("OutputFlag"), 0)
    MOI.set(model, MOI.RawOptimizerAttribute("Cuts"), 0)
    MOI.set(model, MOI.RawOptimizerAttribute("Presolve"), 0)
    MOI.set(model, MOI.RawOptimizerAttribute("PreCrush"), 0)
    MOI.set(model, MOI.RawOptimizerAttribute("Heuristics"), 0)
    N = 30
    x = MOI.add_variables(model, N)
    MOI.add_constraints(model, x, MOI.ZeroOne())
    MOI.set.(model, MOI.VariablePrimalStart(), x, 0.0)
    Random.seed!(1)
    item_weights, item_values = rand(N), rand(N)
    MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(item_weights, x), 0.0),
        MOI.LessThan(10.0),
    )
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(item_values, x), 0.0),
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MAX_SENSE)
    return model, x, item_weights
end

function test_lazy_constraint_callback()
    model, x, y = callback_simple_model()
    lazy_called = false
    MOI.set(
        model,
        MOI.LazyConstraintCallback(),
        cb_data -> begin
            lazy_called = true
            x_val = MOI.get(model, MOI.CallbackVariablePrimal(cb_data), x)
            y_val = MOI.get(model, MOI.CallbackVariablePrimal(cb_data), y)
            status = MOI.get(
                model,
                MOI.CallbackNodeStatus(cb_data),
            )::MOI.CallbackNodeStatusCode
            int_sol = round.(Int, [x_val, y_val])
            if status == MOI.CALLBACK_NODE_STATUS_INTEGER
                @test isapprox(int_sol, [x_val, y_val], atol = 1e-6)
            end
            @test MOI.supports(model, MOI.LazyConstraint(cb_data))
            if y_val - x_val > 1 + 1e-6
                @test MOI.submit(
                    model,
                    MOI.LazyConstraint(cb_data),
                    MOI.ScalarAffineFunction{Float64}(
                        MOI.ScalarAffineTerm.([-1.0, 1.0], [x, y]),
                        0.0,
                    ),
                    MOI.LessThan{Float64}(1.0),
                ) === nothing
            elseif y_val + x_val > 3 + 1e-6
                @test MOI.submit(
                    model,
                    MOI.LazyConstraint(cb_data),
                    MOI.ScalarAffineFunction{Float64}(
                        MOI.ScalarAffineTerm.([1.0, 1.0], [x, y]),
                        0.0,
                    ),
                    MOI.LessThan{Float64}(3.0),
                ) === nothing
            end
        end,
    )
    @test MOI.supports(model, MOI.LazyConstraintCallback())
    MOI.optimize!(model)
    @test lazy_called
    @test MOI.get(model, MOI.VariablePrimal(), x) == 1
    @test MOI.get(model, MOI.VariablePrimal(), y) == 2
end

function test_lazy_constraint_callback_fractional()
    # callback_simple_model() is not large enough to see MIPNODE callbacks.
    model, x, item_weights = callback_knapsack_model()
    lazy_called_integer = false
    lazy_called_fractional = false
    MOI.set(
        model,
        MOI.LazyConstraintCallback(),
        cb_data -> begin
            status = MOI.get(
                model,
                MOI.CallbackNodeStatus(cb_data),
            )::MOI.CallbackNodeStatusCode
            if status == MOI.CALLBACK_NODE_STATUS_INTEGER
                lazy_called_integer = true
            elseif status == MOI.CALLBACK_NODE_STATUS_FRACTIONAL
                lazy_called_fractional = true
            end
        end,
    )
    @test MOI.supports(model, MOI.LazyConstraintCallback())
    MOI.optimize!(model)
    @test lazy_called_integer
    @test lazy_called_fractional
end

function test_lazy_constraint_callback_OptimizeInProgress()
    model, x, y = callback_simple_model()
    MOI.set(
        model,
        MOI.LazyConstraintCallback(),
        cb_data -> begin
            @test_throws(
                MOI.OptimizeInProgress(MOI.VariablePrimal()),
                MOI.get(model, MOI.VariablePrimal(), x)
            )
            @test_throws(
                MOI.OptimizeInProgress(MOI.ObjectiveValue()),
                MOI.get(model, MOI.ObjectiveValue())
            )
            @test_throws(
                MOI.OptimizeInProgress(MOI.ObjectiveBound()),
                MOI.get(model, MOI.ObjectiveBound())
            )
        end,
    )
    return MOI.optimize!(model)
end

function test_lazy_constraint_callback_UserCut()
    model, x, y = callback_simple_model()
    cb = nothing
    MOI.set(
        model,
        MOI.LazyConstraintCallback(),
        cb_data -> begin
            cb = cb_data
            MOI.submit(
                model,
                MOI.UserCut(cb_data),
                MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.0),
                MOI.LessThan(2.0),
            )
        end,
    )
    @test_throws(
        MOI.InvalidCallbackUsage(MOI.LazyConstraintCallback(), MOI.UserCut(cb)),
        MOI.optimize!(model)
    )
end

function test_lazy_constraint_callback_HeuristicSolution()
    model, x, y = callback_simple_model()
    cb = nothing
    MOI.set(
        model,
        MOI.LazyConstraintCallback(),
        cb_data -> begin
            cb = cb_data
            MOI.submit(model, MOI.HeuristicSolution(cb_data), [x], [2.0])
        end,
    )
    @test_throws(
        MOI.InvalidCallbackUsage(
            MOI.LazyConstraintCallback(),
            MOI.HeuristicSolution(cb),
        ),
        MOI.optimize!(model)
    )
end

function test_user_cut_callback()
    model, x, item_weights = callback_knapsack_model()
    user_cut_submitted = false
    MOI.set(
        model,
        MOI.UserCutCallback(),
        cb_data -> begin
            terms = MOI.ScalarAffineTerm{Float64}[]
            accumulated = 0.0
            for (i, xi) in enumerate(x)
                if MOI.get(model, MOI.CallbackVariablePrimal(cb_data), xi) > 0.0
                    push!(terms, MOI.ScalarAffineTerm(1.0, xi))
                    accumulated += item_weights[i]
                end
            end
            @test MOI.supports(model, MOI.UserCut(cb_data))
            if accumulated > 10.0
                @test MOI.submit(
                    model,
                    MOI.UserCut(cb_data),
                    MOI.ScalarAffineFunction{Float64}(terms, 0.0),
                    MOI.LessThan{Float64}(length(terms) - 1),
                ) === nothing
                user_cut_submitted = true
            end
        end,
    )
    @test MOI.supports(model, MOI.UserCutCallback())
    MOI.optimize!(model)
    @test user_cut_submitted
end

function test_user_cut_callback_LazyConstraint()
    model, x, item_weights = callback_knapsack_model()
    cb = nothing
    MOI.set(
        model,
        MOI.UserCutCallback(),
        cb_data -> begin
            cb = cb_data
            MOI.submit(
                model,
                MOI.LazyConstraint(cb_data),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.0),
                MOI.LessThan(5.0),
            )
        end,
    )
    @test_throws(
        MOI.InvalidCallbackUsage(MOI.UserCutCallback(), MOI.LazyConstraint(cb)),
        MOI.optimize!(model)
    )
end

function test_user_cut_callback_HeuristicSolution()
    model, x, item_weights = callback_knapsack_model()
    cb = nothing
    MOI.set(
        model,
        MOI.UserCutCallback(),
        cb_data -> begin
            cb = cb_data
            MOI.submit(model, MOI.HeuristicSolution(cb_data), [x[1]], [0.0])
        end,
    )
    @test_throws(
        MOI.InvalidCallbackUsage(
            MOI.UserCutCallback(),
            MOI.HeuristicSolution(cb),
        ),
        MOI.optimize!(model)
    )
end

function test_heuristic_callback()
    model, x, item_weights = callback_knapsack_model()
    solution_accepted = false
    solution_rejected = false
    MOI.set(
        model,
        MOI.HeuristicCallback(),
        cb_data -> begin
            x_vals = MOI.get.(model, MOI.CallbackVariablePrimal(cb_data), x)
            @test MOI.supports(model, MOI.HeuristicSolution(cb_data))
            status = MOI.get(
                model,
                MOI.CallbackNodeStatus(cb_data),
            )::MOI.CallbackNodeStatusCode
            int_sol = round.(Int, x_vals)
            if status == MOI.CALLBACK_NODE_STATUS_INTEGER
                @test isapprox(int_sol, x_vals, atol = 1e-6)
            end
            if MOI.submit(
                model,
                MOI.HeuristicSolution(cb_data),
                x,
                floor.(x_vals),
            ) == MOI.HEURISTIC_SOLUTION_ACCEPTED
                solution_accepted = true
            end
            if MOI.submit(
                model,
                MOI.HeuristicSolution(cb_data),
                x,
                ceil.(x_vals),
            ) == MOI.HEURISTIC_SOLUTION_REJECTED
                solution_rejected = true
            end
        end,
    )
    @test MOI.supports(model, MOI.HeuristicCallback())
    MOI.optimize!(model)
    @test solution_accepted
    @test solution_rejected
end

function test_heuristic_callback_LazyConstraint()
    model, x, item_weights = callback_knapsack_model()
    cb = nothing
    MOI.set(
        model,
        MOI.HeuristicCallback(),
        cb_data -> begin
            cb = cb_data
            MOI.submit(
                model,
                MOI.LazyConstraint(cb_data),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.0),
                MOI.LessThan(5.0),
            )
        end,
    )
    @test_throws(
        MOI.InvalidCallbackUsage(
            MOI.HeuristicCallback(),
            MOI.LazyConstraint(cb),
        ),
        MOI.optimize!(model)
    )
end

function test_heuristic_callback_UserCut()
    model, x, item_weights = callback_knapsack_model()
    cb = nothing
    MOI.set(
        model,
        MOI.HeuristicCallback(),
        cb_data -> begin
            cb = cb_data
            MOI.submit(
                model,
                MOI.UserCut(cb_data),
                MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.0),
                MOI.LessThan(5.0),
            )
        end,
    )
    @test_throws(
        MOI.InvalidCallbackUsage(MOI.HeuristicCallback(), MOI.UserCut(cb)),
        MOI.optimize!(model)
    )
end

function test_CallbackFunction_callback_OptimizeInProgress()
    model, x, y = callback_simple_model()
    MOI.set(
        model,
        Gurobi.CallbackFunction(),
        (cb_data, cb_where) -> begin
            @test_throws(
                MOI.OptimizeInProgress(MOI.VariablePrimal()),
                MOI.get(model, MOI.VariablePrimal(), x)
            )
            @test_throws(
                MOI.OptimizeInProgress(MOI.ObjectiveValue()),
                MOI.get(model, MOI.ObjectiveValue())
            )
            @test_throws(
                MOI.OptimizeInProgress(MOI.ObjectiveBound()),
                MOI.get(model, MOI.ObjectiveBound())
            )
        end,
    )
    @test MOI.supports(model, Gurobi.CallbackFunction())
    return MOI.optimize!(model)
end

function test_CallbackFunction_callback_LazyConstraint()
    model, x, y = callback_simple_model()
    cb_calls = Int32[]
    function callback_function(cb_data::Gurobi.CallbackData, cb_where::Cint)
        push!(cb_calls, cb_where)
        if cb_where == Gurobi.GRB_CB_MIPSOL
            Gurobi.load_callback_variable_primal(cb_data, cb_where)
            x_val = MOI.get(model, MOI.CallbackVariablePrimal(cb_data), x)
            y_val = MOI.get(model, MOI.CallbackVariablePrimal(cb_data), y)
            if y_val - x_val > 1 + 1e-6
                MOI.submit(
                    model,
                    MOI.LazyConstraint(cb_data),
                    MOI.ScalarAffineFunction{Float64}(
                        MOI.ScalarAffineTerm.([-1.0, 1.0], [x, y]),
                        0.0,
                    ),
                    MOI.LessThan{Float64}(1.0),
                )
            elseif y_val + x_val > 3 + 1e-6
                MOI.submit(
                    model,
                    MOI.LazyConstraint(cb_data),
                    MOI.ScalarAffineFunction{Float64}(
                        MOI.ScalarAffineTerm.([1.0, 1.0], [x, y]),
                        0.0,
                    ),
                    MOI.LessThan{Float64}(3.0),
                )
            end
        end
    end
    MOI.set(model, MOI.RawOptimizerAttribute("LazyConstraints"), 1)
    MOI.set(model, Gurobi.CallbackFunction(), callback_function)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.VariablePrimal(), x) == 1
    @test MOI.get(model, MOI.VariablePrimal(), y) == 2
    @test length(cb_calls) > 0
    @test Gurobi.GRB_CB_MESSAGE in cb_calls
    @test Gurobi.GRB_CB_PRESOLVE in cb_calls
    @test Gurobi.GRB_CB_MIPSOL in cb_calls
end

function test_CallbackFunction_callback_UserCut()
    model, x, item_weights = callback_knapsack_model()
    user_cut_submitted = false
    cb_calls = Int32[]
    MOI.set(
        model,
        Gurobi.CallbackFunction(),
        (cb_data, cb_where) -> begin
            push!(cb_calls, cb_where)
            if cb_where != Gurobi.GRB_CB_MIPNODE
                return
            end
            status = Ref{Cint}()
            Gurobi.GRBcbget(
                cb_data,
                cb_where,
                Gurobi.GRB_CB_MIPNODE_STATUS,
                status,
            )
            if status[] != 2
                return  # Not optimal.
            end
            Gurobi.load_callback_variable_primal(cb_data, cb_where)
            terms = MOI.ScalarAffineTerm{Float64}[]
            accumulated = 0.0
            for (i, xi) in enumerate(x)
                if MOI.get(model, MOI.CallbackVariablePrimal(cb_data), xi) > 0.0
                    push!(terms, MOI.ScalarAffineTerm(1.0, xi))
                    accumulated += item_weights[i]
                end
            end
            if accumulated > 10.0
                MOI.submit(
                    model,
                    MOI.UserCut(cb_data),
                    MOI.ScalarAffineFunction{Float64}(terms, 0.0),
                    MOI.LessThan{Float64}(length(terms) - 1),
                )
                user_cut_submitted = true
            end
        end,
    )
    MOI.optimize!(model)
    @test user_cut_submitted
    @test Gurobi.GRB_CB_MIPNODE in cb_calls
end

function test_CallbackFunction_callback_HeuristicSolution()
    model, x, item_weights = callback_knapsack_model()
    solution_accepted = false
    solution_rejected = false
    cb_calls = Int32[]
    MOI.set(
        model,
        Gurobi.CallbackFunction(),
        (cb_data, cb_where) -> begin
            push!(cb_calls, cb_where)
            if cb_where != Gurobi.GRB_CB_MIPNODE
                return
            end
            status = Ref{Cint}()
            Gurobi.GRBcbget(
                cb_data,
                cb_where,
                Gurobi.GRB_CB_MIPNODE_STATUS,
                status,
            )
            if status[] != 2
                return  # Not optimal.
            end
            Gurobi.load_callback_variable_primal(cb_data, cb_where)
            x_vals =
                MOI.get.(model, MOI.CallbackVariablePrimal(cb_data), x)
            if MOI.submit(
                model,
                MOI.HeuristicSolution(cb_data),
                x,
                floor.(x_vals),
            ) == MOI.HEURISTIC_SOLUTION_ACCEPTED
                solution_accepted = true
            end
            if MOI.submit(
                model,
                MOI.HeuristicSolution(cb_data),
                x,
                ceil.(x_vals),
            ) == MOI.HEURISTIC_SOLUTION_REJECTED
                solution_rejected = true
            end
        end,
    )
    MOI.optimize!(model)
    @test solution_accepted
    @test solution_rejected
    @test Gurobi.GRB_CB_MIPNODE in cb_calls
end

function test_CallbackFunction_CallbackNodeStatus()
    model, x, item_weights = callback_knapsack_model()
    unknown_reached = false
    MOI.set(
        model,
        Gurobi.CallbackFunction(),
        (cb_data, cb_where) -> begin
            if MOI.get(model, MOI.CallbackNodeStatus(cb_data)) ==
               MOI.CALLBACK_NODE_STATUS_UNKNOWN
                unknown_reached = true
            end
        end,
    )
    MOI.optimize!(model)
    @test unknown_reached
end

function test_CallbackFunction_broadcast()
    model, x, _ = callback_knapsack_model()
    f(cb_data, x) = MOI.get(model, MOI.CallbackVariablePrimal(cb_data), x)
    solutions = Vector{Float64}[]
    MOI.set(
        model,
        Gurobi.CallbackFunction(),
        (cb_data, cb_where) -> begin
            if cb_where == Gurobi.GRB_CB_MIPSOL
                Gurobi.load_callback_variable_primal(cb_data, cb_where)
                push!(solutions, f.(cb_data, x))
            end
        end,
    )
    MOI.optimize!(model)
    @test length(solutions) > 0
    @test length(solutions[1]) == length(x)
end

end  # module TestCallbacks

TestCallbacks.runtests()


const MKL_DSS_DEFAULTS = 0
const MKL_DSS_NON_SYMMETRIC = 536871104
const MKL_DSS_SUCCESS = 0
const MKL_DSS_AUTO_ORDER = 268435520
const MKL_DSS_POSITIVE_DEFINITE = 134217792

function dss_analyze(A::SparseMatrixCSC)
  handle = Int[0]
  opt = MKL_DSS_DEFAULTS
  error = ccall((:dss_create, Sparso.LIB_PATH), Cint,
        (Ptr{Void}, Ptr{Cint}),
        handle, &opt)
  if error != MKL_DSS_SUCCESS
    println("dss_create returned error code $error")
  end

  opt = MKL_DSS_NON_SYMMETRIC
  m = size(A, 1)
  nnz = A.colptr[m + 1] - 1
  error = ccall((:dss_define_structure, Sparso.LIB_PATH), Cint,
        (Ptr{Void}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}),
        handle, &opt, A.colptr, &m, &m, A.rowval, &nnz)
  if error != MKL_DSS_SUCCESS
    println("dss_define_structure returned error code $error")
  end

  opt = MKL_DSS_AUTO_ORDER
  error = ccall((:dss_reorder, Sparso.LIB_PATH), Cint,
                (Ptr{Void}, Ptr{Cint}, Ptr{Cint}),
                handle, &opt, C_NULL)
  if error != MKL_DSS_SUCCESS
    println("dss_reorder returned error code $error")
  end

  handle
end

function dss_factor(handle, A::SparseMatrixCSC)
  opt = MKL_DSS_DEFAULTS
  error = ccall((:dss_factor_real, Sparso.LIB_PATH), Cint,
                (Ptr{Void}, Ptr{Cint}, Ptr{Cdouble}),
                handle, &opt, A.nzval)
  if error != MKL_DSS_SUCCESS
    println("dss_factor_real returned error code $error")
  end
end

# solve A*sol = rhs
function dss_solve!(handle, rhs::Vector, sol::Vector)
  opt = MKL_DSS_DEFAULTS
  nrhs = 1
  error = ccall((:dss_solve_real, Sparso.LIB_PATH), Cint,
                (Ptr{Void}, Ptr{Cint}, Ptr{Cdouble}, Ptr{Cint}, Ptr{Cdouble}),
                handle, &opt, rhs, &nrhs, sol)
  ccall((:dss_delete, Sparso.LIB_PATH), Cint,
        (Ptr{Void}, Ptr{Cint}), handle, &opt)
end


include("../models/baseline.jl")

model = Baseline()

function run_step(img, speed, command)
    img = convert.(Float32, img[:,:,1:3]) / 255
    img = reshape(img, (size(img)..., 1))
    img = KnetArray{Float32, 4}(img)

    speed = KnetArray{Float32}([speed])

    preds = model(img)
    throttle = preds[1,1]
    steer = preds[2,1]
    return return throttle, steer
end


@testset "HMC Sampling" begin
    seed!(42)
    L = 3
    D = 10
    N = 20
    nSamples = 10
    b = 5
    i = HMCSampling(; nBurnin=0, thinning=1)

    @test repr(i) == "Hamilton Monte Carlo Sampler"

    @test AGP.ρ(i) == 1
    @test AGP.is_stochastic(i) == false

    i = AGP.init_sampler!(i, L, N, nSamples, false)

    @test i.sample_store == zeros(Float64, nSamples, N, L)

    i.n_iter = nSamples
    i = AGP.init_sampler!(i, L, N, nSamples, true)

    @test i.sample_store == zeros(Float64, 2 * nSamples, N, L)

    i.n_iter = 2 * nSamples
    i = AGP.init_sampler!(i, L, N, 2, false)

    @test i.sample_store == zeros(Float64, 2, N, L)

    i.n_iter = 1
    fs = [rand(N) for _ in 1:L]
    AGP.store_variables!(i, fs)
    @test i.sample_store[1, :, 1] == fs[1]
end