#
# This file is part of the DiscreteEvents.jl Julia package, MIT license
#
# Paul Bayer, 2019
#
# This is a Julia package for discrete event simulation
#
# this implements the event handling
#

# Return the next scheduled event.
_nextevent(c::Clock) = peek(c.sc.events)[1]

# Return the internal time (unitless) of next scheduled event.
_nextevtime(c::Clock) = peek(c.sc.events)[2]

# Execute or evaluate the next timed event on a clock c.
@inline function _event!(c::Clock)
    c.time = c.tev
    ev = dequeue!(c.sc.events)
    _evaluate(ev.ex)
    c.evcount += 1
    ev.Δt > 0.0 && event!(c, ev.ex, c.time + ev.Δt, cycle=ev.Δt)
end

# First execute all sampling expressions in a schedule, then evaluate all
# conditional events and if conditions are met, execute them.
function _tick!(c::Clock)
    c.time = c.tn
    foreach(x -> _evaluate(x.ex), c.sc.samples)  # exec sampling
    # then lookup conditional events
    ix = findfirst(x->all(_evaluate(x.cond)), c.sc.cevents)
    while ix !== nothing
        _evaluate(splice!(c.sc.cevents, ix).ex)
        if isempty(c.sc.cevents)
            if isempty(c.sc.samples) && isempty(c.ac) # no sampling and active clocks
                c.Δt = 0.0 # delete sample rate
            end
            break
        end
        ix = findfirst(x->all(_evaluate(x.cond)), c.sc.cevents)
    end
    c.scount +=1
end

# ------------------------------------------------------
# step forward to next tick or scheduled event. At a tick evaluate
# 1) all sampling functions or expressions,
# 2) all conditional events, then
# 3) if an event is encountered, trigger the event.
# ------------------------------------------------------
# If sampling rate Δt==0, c.tn is set to 0
# If no events are present, c.tev is set to c.end_time
# -------------------------------------------------------
@inline function _step!(c::Clock)
    c.state = Busy()
    if (c.tev ≤ c.time) && (length(c.sc.events) > 0)
        c.tev = _nextevtime(c)
    end

    if !isempty(c.sc.events)
        if c.Δt > 0.0
            if c.tn <= c.tev     # if t_next_tick  ≤ t_next_event
                _tick!(c)
                if c.tn == c.tev # an event is scheduled at the same time
                    _event!(c)
                end
                c.tn += c.Δt
            else
                _event!(c)
            end
        else
            _event!(c)
        end
    elseif c.Δt > 0.0
        _tick!(c)
        c.tn += c.Δt
    else
        (c.state == Busy()) && (c.state = Idle())
        return 99
        # error("_step!: nothing to evaluate")
    end
    !isempty(c.processes) && yield() # let processes run
    c.tev = !isempty(c.sc.events) ? _nextevtime(c) : c.end_time
    (c.state == Busy()) && (c.state = Idle())
    return 0
end

# ----------------------------------------------------
# execute all events in a clock cycle, then do the periodic actions
# ----------------------------------------------------
function _cycle!(clk::Clock, Δt::Float64, sync::Bool=false)
    clk.state = Busy()
    tcyc = clk.time + Δt
    clk.tev = length(clk.sc.events) ≥ 1 ? _nextevtime(clk) : clk.time
    while clk.time ≤ tcyc
        if (clk.tev ≤ tcyc) && length(clk.sc.events) ≥ 1
            _event!(clk)
        else
            clk.time = tcyc
            break
        end
        length(clk.processes) == 0 || yield() # let processes run
    end
    if !sync
        clk.tn = tcyc
        _tick!(clk)
    end
    (clk.state == Busy()) && (clk.state = Idle())
end


#############################################################
# This file is primarily for the purpose of demostrating
# various features:
#
# * DecayOnValidation learning rate policy scheduling
# * User defined Coffee
#
#############################################################
using Mocha
srand(12345678)

data_layer  = AsyncHDF5DataLayer(name="train-data", source="data/train.txt", batch_size=64, shuffle=true)
conv_layer  = ConvolutionLayer(name="conv1", n_filter=20, kernel=(5,5), bottoms=[:data], tops=[:conv])
pool_layer  = PoolingLayer(name="pool1", kernel=(2,2), stride=(2,2), bottoms=[:conv], tops=[:pool])
conv2_layer = ConvolutionLayer(name="conv2", n_filter=50, kernel=(5,5), bottoms=[:pool], tops=[:conv2])
pool2_layer = PoolingLayer(name="pool2", kernel=(2,2), stride=(2,2), bottoms=[:conv2], tops=[:pool2])
fc1_layer   = InnerProductLayer(name="ip1", output_dim=500, neuron=Neurons.ReLU(), bottoms=[:pool2], tops=[:ip1])
fc2_layer   = InnerProductLayer(name="ip2", output_dim=10, bottoms=[:ip1], tops=[:ip2])
loss_layer  = SoftmaxLossLayer(name="loss", bottoms=[:ip2,:label])

backend = DefaultBackend()
init(backend)

common_layers = [conv_layer, pool_layer, conv2_layer, pool2_layer, fc1_layer, fc2_layer]
net = Net("MNIST-train", backend, [data_layer, common_layers..., loss_layer])

exp_dir = "snapshots-$(Mocha.default_backend_type)"

data_layer_test = HDF5DataLayer(name="test-data", source="data/test.txt", batch_size=100)
acc_layer = AccuracyLayer(name="test-accuracy", bottoms=[:ip2, :label])
test_net = Net("MNIST-test", backend, [data_layer_test, common_layers..., acc_layer])
validate_performance = ValidationPerformance(test_net)
lr_policy = LRPolicy.DecayOnValidation(1.e-2, "test-accuracy-accuracy", 0.5)

method = SGD()
params = make_solver_parameters(method, max_iter=10000, regu_coef=0.0005,
                                mom_policy=MomPolicy.Fixed(0.9),
                                lr_policy=lr_policy,
                                load_from=exp_dir)
solver = Solver(method, params)

setup_coffee_lounge(solver, save_into="$exp_dir/statistics.jld", every_n_iter=1000)

# a dummy demo of user-defined coffe break
type DummyCoffee <: Coffee
end
function Mocha.enjoy(lounge::CoffeeLounge, coffee::DummyCoffee, ::Net, ::SolverState)
  println("See also src/coffee/*.jl for implementation of built-in coffee breaks")
end
add_coffee_break(solver, DummyCoffee(), every_n_iter=500)

# report training progress every 100 iterations
add_coffee_break(solver, TrainingSummary(), every_n_iter=100)

# save snapshots every 5000 iterations
add_coffee_break(solver, Snapshot(exp_dir), every_n_iter=5000)

# show performance on test data every 1000 iterations
add_coffee_break(solver, validate_performance, every_n_iter=1000)

# link the decay-on-validation policy with the actual performance validator
setup(lr_policy, validate_performance, solver)

solve(solver, net)

#Profile.init(int(1e8), 0.001)
#@profile solve(solver, net)
#open("profile.txt", "w") do out
#  Profile.print(out)
#end

destroy(net)
destroy(test_net)
shutdown(backend)


using PSCOPF

using Test
using Dates
using DataStructures

@testset verbose=true "example_usecase_from_folder" begin

    @testset "example_2bus_mode1" begin
        data_path = joinpath(@__DIR__, "..", "..", "data", "2buses_usecase", "instance")
        out_path = joinpath(@__DIR__, "..", "..", "default_out", "example_2bus_mode1")
        rm(out_path, recursive=true, force=true)
        mode = PSCOPF.PSCOPF_MODE_1

        # load network
        network = PSCOPF.Data.pscopfdata2network(data_path)
        uncertainties = PSCOPF.PSCOPFio.read_uncertainties(data_path)
        gen_init_state = PSCOPF.PSCOPFio.read_initial_state(data_path)
        println("init_state: ", gen_init_state)

        ts1 = Dates.DateTime("2015-01-01T11:00:00")
        TS = PSCOPF.create_target_timepoints(ts1) #11h, 11h15, 11h30, 11h45

        ECH = PSCOPF.generate_ech(network, TS, mode)

        sequence = PSCOPF.generate_sequence(network, TS, ECH, mode)

        exec_context = PSCOPF.PSCOPFContext(network, TS, mode, gen_init_state, uncertainties,
                                            nothing,
                                            out_path)
        PSCOPF.run!(exec_context, sequence)
    end

end


module TestModelAPI

using Test
using MLJBase
import MLJModelInterface
using ..Models
using Distributions
using StableRNGs
using JLSO

rng = StableRNG(661)

@testset "predict_*" begin
    X = (x = rand(rng, 5),)
    yfinite = categorical(collect("abaaa"))
    ycont = float.(1:5)

    clf = ConstantClassifier()
    fitresult, _, _ = MLJBase.fit(clf, 1, X, yfinite)
    @test predict_mode(clf, fitresult, X)[1] == 'a'
    @test_throws ArgumentError predict_mean(clf, fitresult, X)
    @test_throws ArgumentError predict_median(clf, fitresult, X)

    rgs = ConstantRegressor()
    fitresult, _, _ = MLJBase.fit(rgs, 1, X, ycont)
    @test predict_mean(rgs, fitresult, X)[1] == 3
    @test predict_median(rgs, fitresult, X)[1] == 3
    @test_throws ArgumentError predict_mode(rgs, fitresult, X)
end

end
true


using SIIPExamples
pkgpath = dirname(dirname(pathof(SIIPExamples)))
include(
    joinpath(pkgpath, "test", "3_PowerSimulations_examples", "01_operations_problems.jl"),
)

sys_RT = build_system(PSITestSystems, "modified_RTS_GMLC_RT_sys")

#set_device_model!(template_ed, GenericBattery, BookKeeping)
template_uc

template_ed = template_economic_dispatch()

problems = SimulationProblems(
    UC = OperationsProblem(template_uc, sys, optimizer = solver),
    ED = OperationsProblem(
        template_ed,
        sys_RT,
        optimizer = solver,
        balance_slack_variables = true,
    ),
)

feedforward_chronologies = Dict(("UC" => "ED") => Synchronize(periods = 24))

feedforward = Dict(
    ("ED", :devices, :ThermalStandard) => SemiContinuousFF(
        binary_source_problem = PSI.ON,
        affected_variables = [PSI.ACTIVE_POWER],
    ),
)

intervals = Dict("UC" => (Hour(24), Consecutive()), "ED" => (Minute(15), Consecutive()))

DA_RT_sequence = SimulationSequence(
    problems = problems,
    intervals = intervals,
    ini_cond_chronology = InterProblemChronology(),
    feedforward_chronologies = feedforward_chronologies,
    feedforward = feedforward,
)

sim = Simulation(
    name = "rts-test",
    steps = 1,
    problems = problems,
    sequence = DA_RT_sequence,
    simulation_folder = dirname(dirname(pathof(SIIPExamples))),
)

build!(sim)

execute!(sim, enable_progress_bar = false)

results = SimulationResults(sim);
uc_results = get_problem_results(results, "UC"); # UC stage result metadata
ed_results = get_problem_results(results, "ED"); # ED stage result metadata

read_variables(uc_results, names = [:P__ThermalStandard, :P__RenewableDispatch])

read_parameter(
    ed_results,
    :P__max_active_power__RenewableFix,
    initial_time = DateTime("2020-01-01T06:00:00"),
    count = 5,
)

read_realized_variables(uc_results, names = [:P__ThermalStandard, :P__RenewableDispatch])

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


include("delta.jl")
include("heaviside.jl")
include("angle_of_incidence.jl")
include("bulk.jl")
include("holtslag.jl")
include("large.jl")
include("orad.jl")
include("grid.jl")
include("parameters.jl")

include("ForwardStep.jl")
include("SaveOutput.jl")

fname="out/test.nc"
ReadOutput(fname)

mtime=time

SST=TO[1]
ForwardStep(mtime)


using Random
using ProgressMeter
include("./utils.jl")

using Distributed
addpross(4)

@everywhere begin
    include("./mcts.jl")
end

@everywhere function playonce()
    game = Game()
    randomPlace(game, 2)
    playerX = MCTS()
    playerO = MCTS()
    chunk = Array{Array{Float32, 3}, Array{Float32, 1}, Integer}[]
    while iszero(game.winner)
        move, p = playerX(game, Clients)
        data = getAugmentation(game, p)
        append!(chunk, data)
        place(game, move)

        !iszero(game.winner) && break
        move, p = playerO(game, Clients)
        data = getAugmentation(game, p)
        append!(chunk, data)
        place(game, move)
    end
    @async closeClients(playerX.Clients)
    @async closeClients(playerO.Clients)
    for e in chunk
        @inbounds e[3] .*= game.winner
    end
    return chunk
end

function playNgames(N::Integer=64, bs::Integer=16)
    results = @showprogress pmap(1:N, batch_size=bs) do x
        playonce()
    end
    save_data(shuffle!(reduce(append!, results)))
    nothing
end


@testset "Arithmetic" begin
    x = DNAMer("AA")
    @test x - 1 == x + (-1) == x - 0x01 == DNAMer("TT")
    @test x + 1 == x - (-1) == x + 0x01 == DNAMer("AC")

    x = BigDNAMer("AA")
    @test x - 1 == x + (-1) == x - 0x01 == BigDNAMer("TT")
    @test x + 1 == x - (-1) == x + 0x01 == BigDNAMer("AC")

    x = DNAMer("TT")
    @test x - 1 == x + (-1) == x - 0x01 == DNAMer("TG")
    @test x + 1 == x - (-1) == x + 0x01 == DNAMer("AA")

    x = BigDNAMer("TT")
    @test x - 1 == x + (-1) == x - 0x01 == BigDNAMer("TG")
    @test x + 1 == x - (-1) == x + 0x01 == BigDNAMer("AA")

    base = DNAMer("AAA")
    offset = 0
    nucs = "ACGT"
    for a in nucs, b in nucs, c in nucs
        @test base + offset == DNAMer(string(a, b, c))
        offset += 1
    end

    base = BigDNAMer("AAA")
    offset = 0
    nucs = "ACGT"
    for a in nucs, b in nucs, c in nucs
        @test base + offset == BigDNAMer(string(a, b, c))
        offset += 1
    end

    @testset "typemin" begin
        for k in 1:64
            if k <= 32
                @test typemin(DNAMer{k}) === DNAMer("A"^k)
                @test typemin(RNAMer{k}) === RNAMer("A"^k)
            end

            @test typemin(BigDNAMer{k}) === BigDNAMer("A"^k)
            @test typemin(BigRNAMer{k}) === BigRNAMer("A"^k)
        end
        @test_throws ArgumentError typemin(DNAMer{0})
        @test_throws ArgumentError typemin(DNAMer{33})

        @test_throws ArgumentError typemin(BigDNAMer{0})
        @test_throws ArgumentError typemin(BigDNAMer{65})
    end

    @testset "typemax" begin
        for k in 1:64
            if k <= 32
                @test typemax(DNAMer{k}) === DNAMer("T"^k)
                @test typemax(RNAMer{k}) === RNAMer("U"^k)
            end

            @test typemax(BigDNAMer{k}) === BigDNAMer("T"^k)
            @test typemax(BigRNAMer{k}) === BigRNAMer("U"^k)
        end
        @test_throws ArgumentError typemax(DNAMer{0})
        @test_throws ArgumentError typemax(DNAMer{33})

        @test_throws ArgumentError typemax(BigDNAMer{0})
        @test_throws ArgumentError typemax(BigDNAMer{65})
    end
end


module SquareGreedImages

using Images, Reel

using ...PCG.Geometry: Size, Point
using ...PCG.Types: Recorder
using ...PCG.Storages
using ...PCG.Storages.LinearStorages
using ...PCG.Topologies
using ...PCG.Topologies.SquareGreedTopologies
using ...PCG.Universes
using ...PCG.Universes: Universe
using ...PCG.Operations
using ...PCG.Recorders.GreedImages

export SquareSprite

SquareSprite(color, cell_size::Size) = Sprite(color, fill(color, (Int64(cell_size.y), Int64(cell_size.x))))

function GreedImages.node_position(recorder::GreedImageRecorder, index::SquareGreedIndex, canvas_size::Size)
    point = (Point(index) - Point(1.0, 1.0)) * recorder.cell_size
    return ceil(Int64, point) + Point(1, 1)
end

function GreedImages.canvas_size(topology::SquareGreedTopology, recorder::GreedImageRecorder)::Size
    return ceil(Int64, Size(topology.width, topology.height) * recorder.cell_size)
end

end


import ...BalanceLaws:
    tendency_from_slow_to_fast!,
    cummulate_fast_solution!,
    reconcile_from_fast_to_slow!

#using Printf

@inline function initialize_fast_state!(
    slow::OceanModel,
    fast::BarotropicModel,
    dgSlow,
    dgFast,
    Qslow,
    Qfast,
    slow_dt,
    fast_time_rec,
    fast_steps,
)

    #- inverse ratio of additional fast time steps (for weighted average)
    #  --> do 1/add more time-steps and average from: 1 - 1/add up to: 1 + 1/add
    add = slow.add_fast_substeps

    #- set time-step and number of sub-steps we need
    #  will time-average fast over: fast_steps[1] , fast_steps[3]
    #  centered on fast_steps[2] which corresponds to advance in time of slow
    fast_dt = fast_time_rec[1]
    if add == 0
        steps = fast_dt > 0 ? ceil(Int, slow_dt / fast_dt) : 1
        fast_steps[1:3] = [1 1 1] * steps
    else
        steps = fast_dt > 0 ? ceil(Int, slow_dt / fast_dt / add) : 1
        fast_steps[2] = add * steps
        fast_steps[1] = (add - 1) * steps
        fast_steps[3] = (add + 1) * steps
    end
    fast_time_rec[1] = slow_dt / fast_steps[2]
    fast_time_rec[2] = 0.0
    # @printf("Update: frac_dt = %.1f , dt_fast = %.1f , nsubsteps= %i\n",
    #          slow_dt,fast_time_rec[1],fast_steps[3])
    # println(" fast_steps = ",fast_steps)

    dgFast.state_auxiliary.η_c .= -0
    dgFast.state_auxiliary.U_c .= (@SVector [-0, -0])'

    # preliminary test: no average
    Qfast.η .= dgFast.state_auxiliary.η_s
    Qfast.U .= dgFast.state_auxiliary.U_s

    return nothing
end

@inline function initialize_adjustment!(
    slow::OceanModel,
    fast::BarotropicModel,
    dgSlow,
    dgFast,
    Qslow,
    Qfast,
)
    ## reset tendency adjustment before calling Baroclinic Model
    dgSlow.state_auxiliary.ΔGu .= 0

    return nothing
end

@inline function tendency_from_slow_to_fast!(
    slow::OceanModel,
    fast::BarotropicModel,
    dgSlow,
    dgFast,
    Qslow,
    Qfast,
    dQslow2fast,
)
    FT = eltype(Qslow)

    # integrate the tendency
    tendency_dg = dgSlow.modeldata.tendency_dg
    tend = tendency_dg.balance_law
    grid = dgSlow.grid
    elems = grid.topology.elems
    update_auxiliary_state!(tendency_dg, tend, dQslow2fast, 0, elems)

    Nq, Nqk, _, _, nelemv, nelemh, nrealelemh, _ = basic_grid_info(dgSlow)

    ## get top value (=integral over full depth) of ∫du
    nb_aux_tnd = number_states(tend, Auxiliary())
    data_tnd = reshape(
        tendency_dg.state_auxiliary.data,
        Nq^2,
        Nqk,
        nb_aux_tnd,
        nelemv,
        nelemh,
    )
    index_∫du = varsindex(vars_state(tend, Auxiliary(), FT), :∫du)
    flat_∫du = @view data_tnd[:, end, index_∫du, end, 1:nrealelemh]

    ## copy into Gᵁ of dgFast
    nb_aux_fst = number_states(fast, Auxiliary())
    data_fst = reshape(dgFast.state_auxiliary.data, Nq^2, nb_aux_fst, nelemh)
    index_Gᵁ = varsindex(vars_state(fast, Auxiliary(), FT), :Gᵁ)
    boxy_Gᵁ = @view data_fst[:, index_Gᵁ, 1:nrealelemh]
    boxy_Gᵁ .= flat_∫du

    ## scale by -1/H and copy back to ΔGu
    # note: since tendency_dg.state_auxiliary.∫du is not used after this, could be
    #   re-used to store a 3-D copy of "-Gu"
    nb_aux_slw = number_states(slow, Auxiliary())
    data_slw = reshape(
        dgSlow.state_auxiliary.data,
        Nq^2,
        Nqk,
        nb_aux_slw,
        nelemv,
        nelemh,
    )
    index_ΔGu = varsindex(vars_state(slow, Auxiliary(), FT), :ΔGu)
    boxy_ΔGu = @view data_slw[:, :, index_ΔGu, :, 1:nrealelemh]
    boxy_∫du = @view data_tnd[:, end:end, index_∫du, end:end, 1:nrealelemh]
    boxy_ΔGu .= -boxy_∫du / slow.problem.H

    return nothing
end

@inline function cummulate_fast_solution!(
    fast::BarotropicModel,
    dgFast,
    Qfast,
    fast_time,
    fast_dt,
    substep,
    fast_steps,
    fast_time_rec,
)
    #- might want to use some of the weighting factors: weights_η & weights_U
    #- should account for case where fast_dt < fast.param.dt

    # cumulate Fast solution:
    if substep >= fast_steps[1]
        dgFast.state_auxiliary.U_c .+= Qfast.U
        dgFast.state_auxiliary.η_c .+= Qfast.η
        fast_time_rec[2] += 1.0
    end

    # save mid-point solution to start from the next time-step
    if substep == fast_steps[2]
        dgFast.state_auxiliary.U_s .= Qfast.U
        dgFast.state_auxiliary.η_s .= Qfast.η
    end

    return nothing
end

@inline function reconcile_from_fast_to_slow!(
    slow::OceanModel,
    fast::BarotropicModel,
    dgSlow,
    dgFast,
    Qslow,
    Qfast,
    fast_time_rec,
)
    FT = eltype(Qslow)
    Nq, Nqk, _, _, nelemv, nelemh, nrealelemh, _ = basic_grid_info(dgSlow)
    grid = dgSlow.grid
    elems = grid.topology.elems

    # need to calculate int_u using integral kernels
    # u_slow := u_slow + (1/H) * (u_fast - \int_{-H}^{0} u_slow)

    ## get time weighted averaged out of cumulative arrays
    dgFast.state_auxiliary.U_c .*= 1 / fast_time_rec[2]
    dgFast.state_auxiliary.η_c .*= 1 / fast_time_rec[2]

    ## Compute: \int_{-H}^{0} u_slow)

    # integrate vertically horizontal velocity
    flowintegral_dg = dgSlow.modeldata.flowintegral_dg
    flowint = flowintegral_dg.balance_law
    update_auxiliary_state!(flowintegral_dg, flowint, Qslow, 0, elems)

    # get top value (=integral over full depth)
    nb_aux_flw = number_states(flowint, Auxiliary())
    data_flw = reshape(
        flowintegral_dg.state_auxiliary.data,
        Nq^2,
        Nqk,
        nb_aux_flw,
        nelemv,
        nelemh,
    )
    index_∫u = varsindex(vars_state(flowint, Auxiliary(), FT), :∫u)
    flat_∫u = @view data_flw[:, end, index_∫u, end, 1:nrealelemh]

    ## substract ∫u from U and divide by H

    # Δu is a place holder for 1/H * (Ū - ∫u)
    Δu = dgFast.state_auxiliary.Δu
    Δu .= dgFast.state_auxiliary.U_c

    nb_aux_fst = number_states(fast, Auxiliary())
    data_fst = reshape(dgFast.state_auxiliary.data, Nq^2, nb_aux_fst, nelemh)
    index_Δu = varsindex(vars_state(fast, Auxiliary(), FT), :Δu)
    boxy_Δu = @view data_fst[:, index_Δu, 1:nrealelemh]
    boxy_Δu .-= flat_∫u
    boxy_Δu ./= slow.problem.H

    ## apply the 2D correction to the 3D solution
    nb_cons_slw = number_states(slow, Prognostic())
    data_slw = reshape(Qslow.data, Nq^2, Nqk, nb_cons_slw, nelemv, nelemh)
    index_u = varsindex(vars_state(slow, Prognostic(), FT), :u)
    boxy_u = @view data_slw[:, :, index_u, :, 1:nrealelemh]
    boxy_u .+= reshape(boxy_Δu, Nq^2, 1, 2, 1, nrealelemh)

    ## save Eta from 3D model into η_diag (aux var of 2D model)
    ## and store difference between η from Barotropic Model and η_diag
    index_η = varsindex(vars_state(slow, Prognostic(), FT), :η)
    boxy_η_3D = @view data_slw[:, :, index_η, :, 1:nrealelemh]
    flat_η = @view data_slw[:, end, index_η, end, 1:nrealelemh]
    index_η_diag = varsindex(vars_state(fast, Auxiliary(), FT), :η_diag)
    boxy_η_diag = @view data_fst[:, index_η_diag, 1:nrealelemh]
    boxy_η_diag .= flat_η
    dgFast.state_auxiliary.Δη .=
        dgFast.state_auxiliary.η_c - dgFast.state_auxiliary.η_diag

    ## copy 2D model Eta over to 3D model
    index_η_c = varsindex(vars_state(fast, Auxiliary(), FT), :η_c)
    boxy_η_2D = @view data_fst[:, index_η_c, 1:nrealelemh]
    boxy_η_3D .= reshape(boxy_η_2D, Nq^2, 1, 1, 1, nrealelemh)

    return nothing
end


export IDAPBCProblem, PDELossKinetic, PDELossPotential
export kineticpde, potentialpde
export trainable, unstack, paramstack
export hamiltonian, hamiltoniand, interconnection, controller
export gradient

const UFA = Union{T,C} where {T<:Function, C<:Chain}
const J2Chain = NTuple{N,<:Chain} where {N}
const J2FChain = NTuple{N,<:FastChain} where {N}

struct IDAPBCProblem{TJ2,M,MD,V,VD,GT,GP}
    N::Int
    M⁻¹::M
    Md⁻¹::MD
    V::V
    Vd::VD
    J2::TJ2
    G::GT
    G⊥::GP
    function IDAPBCProblem(N,M⁻¹,Md⁻¹,V,Vd,J2,G,G⊥)
        if isa(M⁻¹,Matrix) && (isa(Md⁻¹,Matrix) || isa(Md⁻¹,PSDMatrix)) && !isnothing(J2)
            error("Constant M and Md requires J2=nothing.")
        end
        if isa(M⁻¹,Function) && isa(Md⁻¹,Matrix)
            error("M is a function, so Md cannot be constant.")
        end
        for f in (isnothing(J2) ? (Md⁻¹, Vd) : (Md⁻¹, Vd, J2...))
            if isa(f, Chain)
                valid = (
                    isa(f.layers, DenseLayers),
                    isa(front(f.layers), DenseLayers),
                    isa(front(tail(f.layers)), DenseLayers)
                )
                if !any(valid)
                    error("Unsupported type $(f)")
                end
            elseif isa(f, FastChain)
                valid = (
                    isa(f.layers, FastDenseLayers),
                    isa(front(f.layers), FastDenseLayers),
                    isa(front(tail(f.layers)), FastDenseLayers)
                )
                if !any(valid)
                    error("Unsupported type $(f)")
                end
            end
        end
        new{typeof(J2),typeof(M⁻¹),typeof(Md⁻¹),typeof(V),typeof(Vd),typeof(G),typeof(G⊥)}(
            N,M⁻¹,Md⁻¹,V,Vd,J2,G,G⊥
        )
    end
end

function IDAPBCProblem(N,M⁻¹,Md⁻¹,V,Vd,G,G⊥)
    IDAPBCProblem(N,M⁻¹,Md⁻¹,V,Vd,nothing,G,G⊥)
end

function Base.show(io::IO, p::IDAPBCProblem)
    println(io, "IDAPBCProblem [q ∈ ℜ^$(p.N)]");
    println(io, "M⁻¹  => $(typeof(p.M⁻¹).name.name)");
    println(io, "Md⁻¹ => $(typeof(p.Md⁻¹).name.name)");
    println(io, "Vd   => $(typeof(p.Vd).name.name)");
    println(io, "J2   => $(typeof(p.J2).name.name)");
end
function Base.getindex(P::IDAPBCProblem, sym::Symbol)
    if sym === :Md
        if isimplicit(P)
            return (q,::Any=missing)->inv(_Md⁻¹(P,q))
        else
            return (q,ps)->inv(_Md⁻¹(P,q,ps))
        end
    elseif sym === :Mdinv
        if isimplicit(P)
            return (q,::Any=missing)->_Md⁻¹(P,q)
        else
            return (q,ps)->_Md⁻¹(P,q,ps)
        end
    elseif sym === :Vd
        if isimplicit(P)
            return (q,::Any=missing)->_Vd(P,q)
        else
            return (q,ps)->_Vd(P,q,ps)
        end
    elseif sym === :Hd
        if isimplicit(P)
            function (x,::Any=missing)
                q, qdot = x[1:P.N], x[P.N+1:end]
                momentum = _M⁻¹(P,q) \ qdot
                hamiltoniand(P, [q; momentum])
            end
        else
            function (x,ps)
                q, qdot = x[1:P.N], x[P.N+1:end]
                momentum = _M⁻¹(P,q) \ qdot
                hamiltoniand(P, [q; momentum], ps)
            end
        end
    elseif sym === :J2
        isnothing(P.J2) && return nothing
        if isimplicit(P)
            return (q,::Any=missing)->interconnection(P,q)
        else
            return (q,ps)->interconnection(P,q,ps)
        end
    end
end

hasfreevariables(::IDAPBCProblem{J2}) where {J2} = J2===Nothing
isimplicit(P::IDAPBCProblem) = any(f->isa(f,Chain), trainable(P))

function kineticpde(M⁻¹, Md⁻¹, ∇M⁻¹, ∇Md⁻¹, G⊥, J2=0)
    return G⊥ * (∇M⁻¹' - (Md⁻¹\M⁻¹)*∇Md⁻¹' + J2*Md⁻¹)
end

function potentialpde(M⁻¹, Md⁻¹, ∇V, ∇Vd, G⊥)
    return first(G⊥ * (∇V - (Md⁻¹\M⁻¹)*∇Vd))
end

trainable(p::IDAPBCProblem{Nothing,M,MD,V,VD}) where
    {M,V,MD<:UFA,VD<:UFA} = (p.Md⁻¹, p.Vd)
trainable(p::IDAPBCProblem{Nothing,M,MD,V,VD}) where
    {M,V,MD<:Matrix,VD<:UFA} = (p.Vd,)
trainable(p::IDAPBCProblem{J,M,MD,V,VD}) where
    {M,V,MD<:UFA,VD<:UFA,N,J<:NTuple{N,<:UFA}} = (p.Md⁻¹, p.Vd, p.J2...)
function unstack(ts::NTuple{N,Function}, ps) where N
    n = DiffEqFlux.paramlength(first(ts))
    res = ps[1:n]
    return (res, unstack(tail(ts), ps[n+1:end])...)
end
unstack(::Tuple{}, ::Any) = ()
unstack(p::IDAPBCProblem, ps) = unstack(trainable(p), ps)
unstack(::IDAPBCProblem, ::Nothing) = nothing
paramstack(p::IDAPBCProblem) = vcat(DiffEqFlux.initial_params.(trainable(p))...)

_M⁻¹(P::IDAPBCProblem{J2,M}, ::Any) where {J2,M<:Matrix} = P.M⁻¹
_M⁻¹(P::IDAPBCProblem{J2,M}, q) where {J2,M<:Function} = P.M⁻¹(q)
_Md⁻¹(P::IDAPBCProblem{J2,M,MD}, q, ::Any=nothing) where {J2,M,MD<:Matrix} = P.Md⁻¹
_Md⁻¹(P::IDAPBCProblem{J2,M,MD}, q) where {J2,M,MD<:Chain} = P.Md⁻¹(q)
function _Md⁻¹(P::IDAPBCProblem{J2,M,MD}, q) where {J2,M,MD<:Function}
    P.Md⁻¹(q)
end
function _Md⁻¹(P::IDAPBCProblem{J2,M,MD}, q, ps) where {J2,M,MD<:Function}
    θMd = first(unstack(P, ps))
    P.Md⁻¹(q, θMd)
end
_∇M⁻¹(P::IDAPBCProblem{J2,M}, q) where {J2,M<:Matrix} = collect(0P.M⁻¹ for _=1:P.N)
_∇M⁻¹(P::IDAPBCProblem{J2,M}, q) where {J2,M<:Function} = jacobian(P.M⁻¹,q)
_∇Md⁻¹(P::IDAPBCProblem{J2,M,MD}, q, ::Any=nothing) where {J2,M,MD<:Matrix} = collect(0P.M⁻¹ for _=1:P.N)
_∇Md⁻¹(P::IDAPBCProblem{J2,M,MD}, q) where {J2,M,MD<:Chain} = jacobian(P.Md⁻¹,q)
function _∇Md⁻¹(P::IDAPBCProblem{J2,M,MD}, q, ps) where {J2,M,MD<:Function}
    θMd = first(unstack(P, ps))
    jacobian(P.Md⁻¹, q, θMd)
end
function _∇Md⁻¹(P::IDAPBCProblem{J2,M,MD}, q) where {J2,M,MD<:Function}
    jacobian(P.Md⁻¹, q)
end
_Vd(P::IDAPBCProblem{J2,M,MD,V,VD}, q) where {J2,M,MD,V,VD<:Chain} = P.Vd(q)
function _Vd(P::IDAPBCProblem{J2,M,MD,V,VD}, q, ps) where {J2,M,MD<:Matrix,V,VD<:Function}
    P.Vd(q, ps)
end
function _Vd(P::IDAPBCProblem{J2,M,MD,V,VD}, q, ps) where {J2,M,MD<:Function,V,VD<:Function}
    _, θVd = unstack(P, ps)
    P.Vd(q, θVd)
end
function _∇Vd(P::IDAPBCProblem{J2,M,MD,V,VD}, q) where {J2,M,MD,V,VD<:Chain}
    if isa(last(P.Vd.layers), Flux.Dense) && isequal(size(last(P.Vd.layers).weight,1), P.N)
        return P.Vd(q)
    else
        return jacobian(P.Vd, q)[:]
    end
end
function _∇Vd(P::IDAPBCProblem{J2,M,MD,V,VD}, q, ps) where {J2,M,V,MD<:Matrix,VD<:FastChain}
    if isa(last(P.Vd.layers), DiffEqFlux.FastDense) && isequal(last(P.Vd.layers).out, P.N)
        return P.Vd(q, ps)
    else
        return jacobian(P.Vd, q, ps)[:]
    end
end
function _∇Vd(P::IDAPBCProblem{J2,M,MD,V,VD}, q, ps) where {J2,M,V,MD<:Function,VD<:FastChain}
    _, θVd = unstack(P,ps)
    if isa(last(P.Vd.layers), DiffEqFlux.FastDense) && isequal(last(P.Vd.layers).out, P.N)
        return P.Vd(q, θVd)
    else
        return jacobian(P.Vd, q, θVd)[:]
    end
end
function _∇Vd(P::IDAPBCProblem{J2,M,MD,V,VD}, q, ps) where {J2,M,V,MD<:Matrix,VD<:Function}
    return jacobian(P.Vd, q, ps)[:]
end
function _∇Vd(P::IDAPBCProblem{J2,M,MD,V,VD}, q, ps) where {J2,M,V,MD<:Function,VD<:Function}
    _, θVd = unstack(P,ps)
    return jacobian(P.Vd, q, θVd)[:]
end

interconnection(P::IDAPBCProblem{Nothing}, x, ::Any=nothing) = 0
function interconnection(P::IDAPBCProblem{<:J2Chain}, x)
    N = P.N
    q, p = x[1:N], x[N+1:2N]
    mapreduce(+, P.J2, p) do Uk, pk
        Uk(q)*pk/2
    end
end
function interconnection(P::IDAPBCProblem{<:J2FChain}, x, ps)
    N = P.N
    q, p = x[1:N], x[N+1:2N]
    θ = last(unstack(P, ps), 2)
    mapreduce(+, P.J2, p, θ) do Uk, pk, θk
        Uk(q,θk)*pk/2
    end
end

function hamiltonian(P::IDAPBCProblem, qp)
    N = P.N
    q, p = qp[1:N], qp[N+1:2N]
    dot(p, _M⁻¹(P, q)*p)/2 + first(P.V(q))
end
function ∇H(P::IDAPBCProblem, qp)
    N = P.N
    q, p = qp[1:N], qp[N+1:2N]
    ∇M⁻¹ = _∇M⁻¹(P, q)
    (mapreduce(*,+,∇M⁻¹,p)'*p)/2 + jacobian(P.V, q)[:]
end

function hamiltoniand(P::IDAPBCProblem, qp)
    N = P.N
    q, p = qp[1:N], qp[N+1:2N]
    dot(p, _Md⁻¹(P, q)*p)/2 + first(_Vd(P,q))
end
function hamiltoniand(P::IDAPBCProblem, qp, ps)
    N = P.N
    q, p = qp[1:N], qp[N+1:2N]
    dot(p, _Md⁻¹(P, q, ps)*p)/2 + first(_Vd(P,q,ps))
end
function ∇Hd(P::IDAPBCProblem, qp)
    N = P.N
    q, p = qp[1:N], qp[N+1:2N]
    ∇Md⁻¹ = _∇Md⁻¹(P, q)
    (mapreduce(*,+,∇Md⁻¹,p)'*p)/2 + _∇Vd(P, q)
end
function ∇Hd(P::IDAPBCProblem, qp, ps)
    N = P.N
    q, p = qp[1:N], qp[N+1:2N]
    ∇Md⁻¹ = _∇Md⁻¹(P, q, ps)
    (mapreduce(*,+,∇Md⁻¹,p)'*p)/2 + _∇Vd(P, q, ps)
end

function controller(P::IDAPBCProblem, x; kv=1)
    N = P.N
    q, qdot = x[1:N], x[N+1:2N]
    G⁻¹ = (P.G'*P.G)\(P.G')
    M⁻¹ = _M⁻¹(P,q)
    Md⁻¹ = _Md⁻¹(P,q)
    MdM⁻¹ = Md⁻¹ \ M⁻¹
    momentum = M⁻¹ \ qdot
    qp = [q; momentum]
    u_es = G⁻¹ * ( ∇H(P,qp) - MdM⁻¹*∇Hd(P,qp) + interconnection(P,qp)*Md⁻¹*momentum )
    u_di = -kv * dot(P.G, Md⁻¹*momentum)
    return u_es + u_di
end
function controller(P::IDAPBCProblem, x, ps; kv=1)
    N = P.N
    q, qdot = x[1:N], x[N+1:2N]
    G⁻¹ = (P.G'*P.G)\(P.G')
    M⁻¹ = _M⁻¹(P,q)
    Md⁻¹ = _Md⁻¹(P,q,ps)
    MdM⁻¹ = Md⁻¹ \ M⁻¹
    momentum = M⁻¹ \ qdot
    qp = [q; momentum]
    u_es = G⁻¹ * ( ∇H(P,qp) - MdM⁻¹*∇Hd(P,qp,ps) + interconnection(P,qp,ps)*Md⁻¹*momentum )
    u_di = -kv * dot(P.G, Md⁻¹*momentum)
    return u_es + u_di
end

abstract type IDAPBCLoss{P} end

function gradient(l::IDAPBCLoss{P}, q, ps::AbstractVector) where
    {J2,M,MD,V,VD<:Function,P<:IDAPBCProblem{J2,M,MD,V,VD}}
    ReverseDiff.gradient(_2->l(q,_2), ps)
end
function gradient(l::IDAPBCLoss{P}, q, ps::Flux.Params) where
    {J2,M,MD,V,VD<:Chain,P<:IDAPBCProblem{J2,M,MD,V,VD}}
    Flux.gradient(()->l(q), ps)
end

struct PDELossKinetic{hasJ2,P} <: IDAPBCLoss{P}
    prob::P
    function PDELossKinetic(prob::P) where {J2,P<:IDAPBCProblem{J2}}
        new{J2!==Nothing,P}(prob)
    end
end
function (L::PDELossKinetic{false})(q)
    p = L.prob
    M⁻¹ = _M⁻¹(p, q)
    ∇M⁻¹ = _∇M⁻¹(p, q)
    Md⁻¹ = _Md⁻¹(p, q)
    ∇Md⁻¹ = _∇Md⁻¹(p, q)
    r = map(∇M⁻¹, ∇Md⁻¹) do _1, _2
        sum(abs, kineticpde(M⁻¹, Md⁻¹, _1, _2, p.G⊥))
    end
    sum(r)
end
function (L::PDELossKinetic{true})(q)
    p = L.prob
    M⁻¹ = _M⁻¹(p, q)
    ∇M⁻¹ = _∇M⁻¹(p, q)
    Md⁻¹ = _Md⁻¹(p, q)
    ∇Md⁻¹ = _∇Md⁻¹(p, q)
    r = map(∇M⁻¹, ∇Md⁻¹, p.J2) do _1, _2, Uk
        sum(abs, kineticpde(M⁻¹, Md⁻¹, _1, _2, p.G⊥, Uk(q)))
    end
    sum(r)
end
function (L::PDELossKinetic{false})(q,ps)
    p = L.prob
    M⁻¹ = _M⁻¹(p, q)
    ∇M⁻¹ = _∇M⁻¹(p, q)
    Md⁻¹ = _Md⁻¹(p, q, ps)
    ∇Md⁻¹ = _∇Md⁻¹(p, q, ps)
    r = map(∇M⁻¹, ∇Md⁻¹) do _1, _2
        sum(abs, kineticpde(M⁻¹, Md⁻¹, _1, _2, p.G⊥))
    end
    sum(r)
end
function (L::PDELossKinetic{true})(q,ps)
    p = L.prob
    M⁻¹ = _M⁻¹(p, q)
    ∇M⁻¹ = _∇M⁻¹(p, q)
    Md⁻¹ = _Md⁻¹(p, q, ps)
    ∇Md⁻¹ = _∇Md⁻¹(p, q, ps)
    J2ps = last(unstack(p, ps), 2)
    r = map(∇M⁻¹, ∇Md⁻¹, p.J2, J2ps) do _1, _2, Uk, Ukps
        sum(abs, kineticpde(M⁻¹, Md⁻¹, _1, _2, p.G⊥, Uk(q,Ukps)))
    end
    sum(r)
end

struct PDELossPotential{P} <: IDAPBCLoss{P}
    prob::P
end
function (L::PDELossPotential)(q)
    p = L.prob
    ∇V = jacobian(p.V, q)[:]
    ∇Vd = _∇Vd(p, q)
    M⁻¹ = _M⁻¹(p, q)
    Md⁻¹ = _Md⁻¹(p, q)
    abs(potentialpde(M⁻¹, Md⁻¹, ∇V, ∇Vd, p.G⊥))
end
function (L::PDELossPotential)(q,ps)
    p = L.prob
    ∇V = jacobian(p.V, q)[:]
    ∇Vd = _∇Vd(p, q, ps)
    M⁻¹ = _M⁻¹(p, q)
    Md⁻¹ = _Md⁻¹(p, q, ps)
    abs(potentialpde(M⁻¹, Md⁻¹, ∇V, ∇Vd, p.G⊥))
end

struct PotentialHessianSymLoss{P} <: IDAPBCLoss{P}
    prob::P
    function PotentialHessianSymLoss(p::IDAPBCProblem{J,M,MD,V,VD}) where
        {J,M,MD,V,VD<:Chain}
        if isequal(length(last(p.Vd.layers).bias), p.N)
            new{typeof(p)}(p)
        else
            error("Not applicable for this type of IDAPBCProblem")
        end
    end
    function PotentialHessianSymLoss(p::IDAPBCProblem{J,M,MD,V,VD}) where
        {J,M,MD,V,VD<:FastChain}
        if isequal(last(p.Vd.layers).out, p.N)
            new{typeof(p)}(p)
        else
            error("Not applicable for this type of IDAPBCProblem")
        end
    end
end
function (L::PotentialHessianSymLoss)(q)
    J = jacobian(L.prob.Vd, q)
    mapreduce(abs, +, J - J')
end
function (L::PotentialHessianSymLoss{P})(q, ps) where {J2,M,MD<:Matrix,P<:IDAPBCProblem{J2,M,MD}}
    J = jacobian(L.prob.Vd, q, ps)
    mapreduce(abs, +, J - J')
end
function (L::PotentialHessianSymLoss{P})(q, ps) where {J2,M,MD<:Function,P<:IDAPBCProblem{J2,M,MD}}
    _, θVd = unstack(P,ps)
    J = jacobian(L.prob.Vd, q, θVd)
    mapreduce(abs, +, J - J')
end


@inline function mulsizeof(::Type{T}, x::Number) where {T}
  (Base.allocatedinline(T) ? sizeof(T) : sizeof(Int)) * x
end
@generated function mulsizeof(::Type{T}, ::StaticInt{N}) where {T,N}
  st = Base.allocatedinline(T) ? sizeof(T) : sizeof(Int)
  Expr(:block, Expr(:meta,:inline), Expr(:call, Expr(:curly, :StaticInt, N*st)))
end
@inline mulsizeof(::Type{T}, ::Tuple{}) where {T} = ()
@inline mulsizeof(::Type{T}, x::Tuple{X}) where {T,X} = (mulsizeof(T, getfield(x,1,false)), )
@inline mulsizeof(::Type{T}, x::Tuple{X1,X2,Vararg}) where {T,X1,X2} = (mulsizeof(T, getfield(x,1,false)), mulsizeof(T, Base.tail(x))...)

@inline bytestrides(A::AbstractArray{T}) where {T} = mulsizeof(T, ArrayInterface.strides(A))

@inline memory_reference(A::AbstractArray) = memory_reference(ArrayInterface.device(A), A)
@inline memory_reference(A::BitArray) = Base.unsafe_convert(Ptr{Bit}, A.chunks), A.chunks
@inline memory_reference(::CPUPointer, A) = pointer(A), preserve_buffer(A)
@inline memory_reference(::CPUPointer, A::Union{LinearAlgebra.Adjoint, Base.ReshapedArray, Base.PermutedDimsArray, LinearAlgebra.Transpose}) = memory_reference(CPUPointer(), parent(A))
@inline function memory_reference(::CPUPointer, A::Base.ReinterpretArray{T}) where {T}
  p, m = memory_reference(CPUPointer(), parent(A))
  reinterpret(Ptr{T}, p), m
end
@inline ind_diff(::Base.Slice, ::Any) = Zero()
@inline ind_diff(x::AbstractRange, o) = static_first(x) - o
@inline ind_diff(x::Integer, o) = x - o
@inline memory_reference(::CPUPointer, A::SubArray) = memory_reference_subarray(CPUPointer(), A)
@inline memory_reference(::CPUTuple, A::SubArray) = memory_reference_subarray(CPUTuple(), A)
@inline function memory_reference_subarray(::PT, A::SubArray) where {PT}
  p, m = memory_reference(PT(), parent(A))
  pA = parent(A)
  offset = ArrayInterface.reduce_tup(+, _map(*, _map(ind_diff, A.indices, offsets(pA)), strides(pA)))
  p + sizeof(eltype(A))*offset, m
end
@inline function memory_reference(::CPUTuple, A)
  r = Ref(A)
  Base.unsafe_convert(Ptr{eltype(A)}, Base.pointer_from_objref(r)), r
end
@inline function memory_reference(::ArrayInterface.CheckParent, A)
  P = parent(A)
  if P === A
    memory_reference(ArrayInterface.CPUIndex(), A)
  else
    memory_reference(ArrayInterface.device(P), P)
  end
end
@inline memory_reference(::ArrayInterface.CPUIndex, A) = throw("Memory access for $(typeof(A)) not implemented yet.")

@inline ArrayInterface.contiguous_axis(::Type{A}) where {T,N,C,A<:AbstractStridedPointer{T,N,C}} = StaticInt{C}()
@inline ArrayInterface.contiguous_batch_size(::Type{A}) where {T,N,C,B,A<:AbstractStridedPointer{T,N,C,B}} = StaticInt{B}()
@inline ArrayInterface.stride_rank(::Type{A}) where {T,N,C,B,R,A<:AbstractStridedPointer{T,N,C,B,R}} = _map(StaticInt, R)
@inline memory_reference(A::AbstractStridedPointer) = pointer(A), nothing

@inline Base.eltype(::AbstractStridedPointer{T}) where {T} = T

struct StridedPointer{T,N,C,B,R,X,O} <: AbstractStridedPointer{T,N,C,B,R,X,O}
  p::Ptr{T}
  si::StrideIndex{N,R,C,X,O}
end
@inline stridedpointer(p::Ptr{T}, si::StrideIndex{N,R,C,X,O}, ::StaticInt{B}) where {T,N,R,C,B,X,O} = StridedPointer{T,N,C,B,R,X,O}(p, si)
@inline stridedpointer(p::Ptr{T}, si::StrideIndex{N,R,C,X,O}) where {T,N,R,C,X,O} = StridedPointer{T,N,C,0,R,X,O}(p, si)

@inline bytestrideindex(A::AbstractArray{T}) where {T} = mulstrides(T, StrideIndex(A))

@inline function mulstrides(::Type{T}, si::StrideIndex{N,R,C,X,O}) where {N,R,C,X,O,T}
  StrideIndex{N,R,C}(mulsizeof(T, si.strides), si.offsets)
end
@inline function stridedpointer(A::AbstractArray)
  p, r = memory_reference(A)
  stridedpointer(p, bytestrideindex(A), ArrayInterface.contiguous_batch_size(A))
end
@inline function stridedpointer_preserve(A::AbstractArray)
  p, r = memory_reference(A)

  stridedpointer(p, bytestrideindex(A), ArrayInterface.contiguous_batch_size(A)), r
end
@inline val_stride_rank(::AbstractStridedPointer{T,N,C,B,R}) where {T,N,C,B,R} = Val{R}()
@generated val_dense_dims(::AbstractStridedPointer{T,N}) where {T,N} = Val{ntuple(==(0), Val(N))}()
@inline val_stride_rank(A) = Val(known(stride_rank(A)))
@inline val_dense_dims(A) = Val(known(ArrayInterface.dense_dims(A)))

function zerotupleexpr(N::Int)
  t = Expr(:tuple);
  for n in 1:N
    push!(t.args, :(Zero()))
  end
  Expr(:block, Expr(:meta,:inline), t)
end
@generated zerotuple(::Val{N}) where {N} = zerotupleexpr(N)
# @inline center(p::AbstractStridedPointer{T,N}) where {T,N} = gesp(p, zerotuple(Val(N)))
@inline zero_offsets(si::StrideIndex{N,R,C}) where {N,R,C} = StrideIndex{N,R,C}(getfield(si,:strides), zerotuple(Val(N)))
@inline zero_offsets(sptr::AbstractStridedPointer) = stridedpointer(pointer(sptr), zero_offsets(StrideIndex(sptr)), contiguous_batch_size(sptr))
@inline zstridedpointer(A) = zero_offsets(stridedpointer(A))
@inline function zstridedpointer_preserve(A::AbstractArray{T,N}) where {T,N}
  strd = mulsizeof(T, ArrayInterface.strides(A))
  si = StrideIndex{N,known(stride_rank(A)),Int(contiguous_axis(A))}(strd, zerotuple(Val(N)), Zero())
  p, r = memory_reference(A)
  stridedpointer(p, si, contiguous_batch_size(A)), r
end

@inline Base.similar(sptr::StridedPointer, ptr::Ptr) = stridedpointer(ptr, StrideIndex(sptr), contiguous_batch_size(sptr))

Base.unsafe_convert(::Type{Ptr{T}}, ptr::AbstractStridedPointer{T}) where {T} = pointer(ptr)
# Shouldn't need to special case Array
@generated function nopromote_axis_indicator(::AbstractStridedPointer{<:Any,N}) where {N}
  t = Expr(:tuple); foreach(n -> push!(t.args, True()), 1:N)
  Expr(:block, Expr(:meta, :inline), t)
end

@inline dynamic_offsets(si::StrideIndex{N,R,C}) where {N,R,C} = StrideIndex{N,R,C}(strides(si), _map(Int, offsets(si)))
struct StridedBitPointer{N,C,B,R,X,O} <: AbstractStridedPointer{Bit,N,C,B,R,X,O}
  p::Ptr{Bit}
  si::StrideIndex{N,R,C,X,O}
end
@inline stridedpointer(p::Ptr{Bit}, si::StrideIndex{N,R,C,X,O}, ::StaticInt{B}) where {N,R,C,B,X,O} = StridedBitPointer{N,C,B,R,X,NTuple{N,Int}}(p, dynamic_offsets(si))
@inline stridedpointer(p::Ptr{Bit}, si::StrideIndex{N,R,C,X,O}) where {N,R,C,X,O} = StridedBitPointer{N,C,0,R,X,NTuple{N,Int}}(p, dynamic_offsets(si))

@inline Base.pointer(p::Union{StridedPointer,StridedBitPointer}) = getfield(p, :p)
@inline ArrayInterface.StrideIndex(sptr::Union{StridedPointer,StridedBitPointer}) = getfield(sptr,:si)
@inline bytestrideindex(sptr::AbstractStridedPointer) = StrideIndex(sptr)

@inline bytestrides(si::StrideIndex) = _map(Base.Fix2(*,StaticInt{8}()), getfield(si,:strides))
@inline bytestrides(ptr::AbstractStridedPointer) = getfield(StrideIndex(ptr), :strides)
@inline Base.strides(ptr::AbstractStridedPointer) = getfield(StrideIndex(ptr), :strides)
@inline ArrayInterface.strides(ptr::AbstractStridedPointer) = getfield(StrideIndex(ptr), :strides)
@inline ArrayInterface.offsets(ptr::AbstractStridedPointer) = offsets(StrideIndex(ptr))
@inline ArrayInterface.contiguous_axis_indicator(ptr::AbstractStridedPointer{T,N,C}) where {T,N,C} = contiguous_axis_indicator(StaticInt{C}(), Val{N}())

@inline function similar_with_offset(sptr::AbstractStridedPointer{T,N,C,B,R}, ptr::Ptr, offset::Tuple) where {T,N,C,B,R}
  si = StrideIndex{N,R,C}(strides(sptr), offset)
  stridedpointer(ptr, si, contiguous_batch_size(sptr))
end
@inline function similar_no_offset(sptr::AbstractStridedPointer{T,N,C,B,R}, ptr::Ptr) where {T,N,C,B,R}
  si = StrideIndex{N,R,C}(strides(sptr), zerotuple(Val(N)))
  stridedpointer(ptr, si, contiguous_batch_size(sptr))
end

# There is probably a smarter way to do indexing adjustment here.
# The reasoning for the current approach of geping for Zero() on extracted inds
# and for offsets otherwise is best demonstrated witht his motivational example:
#
# A = OffsetArray(rand(10,11), 6:15, 5:15);
# for i in 6:15
    # s += A[i,i]
# end
# first access is at zero-based index
# (first(6:16) - ArrayInterface.offsets(a)[1]) * ArrayInterface.strides(A)[1] + (first(6:16) - ArrayInterface.offsets(a)[2]) * ArrayInterface.strides(A)[2]
# equal to
#  (6 - 6)*1 + (6 - 5)*10 = 10
# i.e., the 1-based index 11.
# So now we want to adjust the offsets and pointer's value
# so that indexing it with a single `i` (after summing strides) is correct.
# Moving to 0-based indexing is easiest for combining strides. So we gep by 0 on these inds.
# E.g., gep(stridedpointer(A), (0,0))
# ptr += (0 - 6)*1 + (0 - 5)*10 = -56
# new stride = 1 + 10 = 11
# new_offset = 0
# now if we access the 6th element
# (6 - new_offse) * new_stride
# ptr += (6 - 0) * 11 = 66
# cumulative:
# ptr = pointer(A) + 66 - 56  = pointer(A) + 10
# so initial load is of pointer(A) + 10 -> the 11th element w/ 1-based indexing

@inline stridedpointer(ptr::AbstractStridedPointer) = ptr
@inline stridedpointer_preserve(ptr::AbstractStridedPointer) = (ptr,nothing)

struct FastRange{T,F,S,O}# <: AbstractStridedPointer{T,1,1,0,(1,),Tuple{S},Tuple{O}}# <: AbstractRange{T}
  f::F
  s::S
  o::O
end
FastRange{T}(f::F,s::S) where {T<:Integer,F,S} = FastRange{T,Zero,S,F}(Zero(),s,f)
FastRange{T}(f::F,s::S,o::O) where {T,F,S,O} = FastRange{T,F,S,O}(f,s,o)

FastRange{T}(f::F,s::S) where {T<:FloatingTypes,F,S} = FastRange{T,F,S,Int}(f,s,0)
# FastRange{T}(f,s) where {T<:FloatingTypes} = FastRange{T}(f,s,fast_int64_to_double())
# FastRange{T}(f::F,s::S,::True) where {T<:FloatingTypes,F,S} = FastRange{T,F,S,Int}(f,s,0)
# FastRange{T}(f::F,s::S,::False) where {T<:FloatingTypes,F,S} = FastRange{T,F,S,Int32}(f,s,zero(Int32))

@inline function memory_reference(r::AbstractRange{T}) where {T}
  s = ArrayInterface.static_step(r)
  FastRange{T}(ArrayInterface.static_first(r) - s, s), nothing
end
@inline memory_reference(r::FastRange) = (r,nothing)
@inline bytestrides(::FastRange{T}) where {T} = (StaticInt(sizeof(T)),)
@inline ArrayInterface.offsets(::FastRange) = (One(),)
@inline val_stride_rank(::FastRange) = Val{(1,)}()
@inline val_dense_dims(::FastRange) = Val{(true,)}()
@inline ArrayInterface.contiguous_axis(::FastRange) = One()
@inline ArrayInterface.contiguous_batch_size(::FastRange) = Zero()

@inline stridedpointer(fr::FastRange, ::StrideIndex, ::StaticInt{0}) = fr
struct NoStrides end
@inline bytestrideindex(::FastRange) = NoStrides()
@inline ArrayInterface.offsets(::NoStrides) = NoStrides()
@inline reconstruct_ptr(r::FastRange{T}, o) where {T} = FastRange{T}(getfield(r,:f), getfield(r, :s), o)

@inline function zero_offsets(fr::FastRange{T,Zero}) where {T<:Integer}
  s = getfield(fr,:s)
  FastRange{T}(Zero(), s, getfield(fr,:o) + s)
end
@inline function zero_offsets(fr::FastRange{T}) where {T<:FloatingTypes}
  FastRange{T}(getfield(fr,:f), getfield(fr,:s), getfield(fr,:o) + 1)
end

# `FastRange{<:FloatingTypes}` must use an integer offset because `ptrforcomparison` needs to be exact/integral.

# discard unnueeded align/reg size info
@inline Base.eltype(::FastRange{T}) where {T} = T

# @inline Base.pointer(p::AbstractStridedPointer, i::Tuple) = pointer(p) + StrideIndex(p)[i...]
function reinterpret_strideindex_quote(
  sz_new::Int, sz_old::Int, C::Int, N::Int, curly::Expr
)
  if sz_old == sz_new
    # size_expr = :size_A
    bx_expr = :bx
  else
    @assert 1 ≤ C ≤ N
    # size_expr = Expr(:tuple)
    bx_expr = Expr(:tuple)
    for n ∈ 1:N
      # sz_n = Expr(:call, GlobalRef(Core,:getfield), :size_A, n, false)
      bx_n = Expr(:call, GlobalRef(Core,:getfield), :bx, n, false)
      if n ≠ C
        # push!(size_expr.args, sz_n)
        push!(bx_expr.args, bx_n)
      elseif sz_old > sz_new
        si = :(StaticInt{$(sz_old ÷ sz_new)}())
        # push!(size_expr.args, Expr(:call, :*, sz_n, si))
        push!(bx_expr.args, Expr(:call, :÷, bx_n, si))
      else
        si = :(StaticInt{$(sz_new ÷ sz_old)}())
        # push!(size_expr.args, Expr(:call, :÷, sz_n, si))
        push!(bx_expr.args, Expr(:call, :*, bx_n, si))
      end
    end
  end
  quote
    $(Expr(:meta,:inline))
    bx = si.strides
    $curly($bx_expr, si.offsets)
  end
end
@generated function reinterpret_strideindex(::Type{Tnew}, ::Type{Told}, si::StrideIndex{N,R,C}) where {Tnew,Told,N,R,C}
  # 1+2
  reinterpret_strideindex_quote(offsetsize(Tnew), offsetsize(Told), C, N, :(StrideIndex{$N,$R,$C}))
end
@generated function Base.reinterpret(::Type{Tnew}, sp::AbstractStridedPointer{Told,N,C,B,R}) where {Tnew,Told,N,C,B,R}
  # 1+2
  sz_old = offsetsize(Told)
  sz_new = offsetsize(Tnew)
  Bnew = ((B > 0) & (sz_new ≠ sz_old)) ? ((B * sz_old) ÷ sz_new) : B
  quote
    $(Expr(:meta,:inline))
    si = reinterpret_strideindex(Tnew, Told, StrideIndex(sp))
    stridedpointer(reinterpret(Ptr{$Tnew}, pointer(sp)), si, StaticInt{$Bnew}())
  end
end

@inline vload(::StaticInt{N}, args...) where {N} = StaticInt{N}()
@inline stridedpointer(::StaticInt{N}) where {N} = StaticInt{N}()
@inline zero_offsets(::StaticInt{N}) where {N} = StaticInt{N}()


import NPZ

struct OutputConfig
    "where to write output"
    directory
    "times at which to output"
    output_times::Array{Real}

    "whether to output boxes"
    box::Bool
    "whether to output slices"
    slice::Bool
    "whether to output ψ"
    psi::Bool
    "whether to output ρ"
    rho::Bool
end

function OutputConfig(
    directory, output_times;
    box=true, slice=false, psi=true, rho=true
)

    OutputConfig(directory, output_times, box, slice, psi, rho)
end

"""
    output_grids(grids, output_config, step)

Write output from `grids` as specified in `output_config`
"""
function output_grids(grids, output_config, step)

    if output_config.box
        if output_config.psi
            NPZ.npzwrite(
                joinpath(output_config.directory, "psi_$step.npy"),
                grids.ψx
            )
        end
        if output_config.rho
            NPZ.npzwrite(
                joinpath(output_config.directory, "rho_$step.npy"),
                grids.ρx
            )
        end
    end

    if output_config.slice
        if output_config.psi
            NPZ.npzwrite(
                joinpath(output_config.directory, "psi_slice_$step.npy"),
                grids.ψx[1, :, :]
            )
        end
        if output_config.rho
            NPZ.npzwrite(
                joinpath(output_config.directory, "rho_slice_$step.npy"),
                grids.ρx[1, :, :]
            )
        end
    end
end

function output_grids(grids::PencilGrids, output_config, step)

    #TODO: don't use gather.  This sends all data to one node.  Should instead use multithreaded HDF5 output
    if output_config.box
        if output_config.psi
            NPZ.npzwrite(
                joinpath(output_config.directory, "psi_$step.npy"),
                gather(grids.ψx)
            )
        end
        if output_config.rho
            NPZ.npzwrite(
                joinpath(output_config.directory, "rho_$step.npy"),
                gather(grids.ρx)
            )
        end
    end

    if output_config.slice
        if output_config.psi
            NPZ.npzwrite(
                joinpath(output_config.directory, "psi_slice_$step.npy"),
                gather(grids.ψx[1, :, :])
            )
        end
        if output_config.rho
            NPZ.npzwrite(
                joinpath(output_config.directory, "rho_slice_$step.npy"),
                gather(grids.ρx[1, :, :])
            )
        end
    end
end

"""
    output_summary_header(output_config)

Write a header for a summary file

The header contains labels for each column of the summary CSV file.
This function overwrites the current contents of the file.
"""
function output_summary_header(output_config)

    open(joinpath(output_config.directory, "summary.csv"), "w") do file
        write(file, "t,a,Δt,ρx_mean,δx_rms\n")
    end
end

"""
    output_summary_row(grids, output_config, t, a, Δt)

Write a new row to the summary file
"""
function output_summary_row(grids, output_config, t, a, Δt)
    open(joinpath(output_config.directory, "summary.csv"), "a") do file
        ρx_mean = mean(grids.ρx)
        δx_rms = mean(((grids.ρx .- ρx_mean).^2))^0.5
        write(file, "$t, $a, $Δt, $ρx_mean, $δx_rms\n")
    end
end


# ---
# title: Composition utilities
# cover: utilities.gif
# description: a set of helper operations that may be useful when compositing more complex augmentation workflow
# ---

# Aside from "true" operations that specify some kind of transformation, there are also a couple of
# special utility operations used for functionality such as stochastic branching.

using Augmentor
using Random
Random.seed!(1337)

img_in = testpattern(RGB, ratio=0.5)
img_out = augment(img_in, Either(0.5=>NoOp(), 0.25=>FlipX(), 0.25=>FlipY()))
#md nothing #hide

# ![](utilities.gif)

# ## References

#md # ```@docs
#md # NoOp
#md # Either
#md # CacheImage
#md # ```

## save covers #src
using ImageMagick #src
include(joinpath("..", "assets", "utilities.jl")) #src
pl = Either(ShearX(-5:5), ShearY(-5:5)) |> Rotate(-10:10) |> Either(NoOp(), FlipX(), FlipY()) |> CropNative(axes(img_in)) #src
cover = make_gif(img_in, pl, 10) #src
ImageMagick.save("utilities.gif", cover; fps=2) #src


import Base: ==, <, >

"""
A record of a transition and the time.
It's sortable by time. Immutable.
"""
struct NRTransition
	key
	time::Float64
end

function Base.:<(a::NRTransition, b::NRTransition)
	a.time < b.time
end

function Base.:>(a::NRTransition, b::NRTransition)
    a.time > b.time
end

function ==(a::NRTransition, b::NRTransition)
    a.time == b.time
end

"""
Classic First Reaction method
"""
struct FirstReaction
end

function Next(fr::FirstReaction, system, rng)
	least=NRTransition(nothing, Inf)
	Hazards(system, rng) do clock, now, enabled, rng2
	  trial_time = Sample(clock.intensity, now, rng2)
	  @assert(trial_time >= now)
	  if trial_time < least.time
	  	least = NRTransition(clock, trial_time)
	  end
    end
    (least.time, least.key)
end

Observer(fr::FirstReaction) = (hazard, time, updated, rng) -> nothing


################################################################################
function init(::Type{Metropolis}, config::ConfigMetropolis, objective::Objective, difftune)
    return Metropolis(ModelWrappers.ℓDensityResult(objective))
end

function infer(
    _rng::Random.AbstractRNG,
    diagnostics::Type{AbstractDiagnostics},
    kernel::Metropolis,
    model::ModelWrapper,
    data::D,
) where {D}
    return DiagnosticsMetropolis{model.info.reconstruct.default.output}
end

################################################################################
# Export
export infer, init


using MCIntegration
using Test

@testset "MCIntegration.jl" begin
    # Write your tests here.
    include("montecarlo.jl")
end


include("basic/gamsworld.jl")

@testset "basic tests" begin

@testset "no objective and start value" begin
    println("==================================")
    println("No objective and start value")
    println("==================================")
    juniper = DefaultTestSolver(log_levels=[:Table])

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper...)
    )

    @variable(m, x, Int, start=3)

    @constraint(m, x >= 0)
    @constraint(m, x <= 5)
    @NLconstraint(m, x^2 >= 17)

    status = solve(m)
    rand_num = rand()
    status = solve(m)
    @test rand() != rand_num

    inner = internalmodel(m)
    @test JuMP.termination_status(m) == MOI.LOCALLY_SOLVED
    @test JuMP.primal_status(m) == MOI.FEASIBLE_POINT
    @test JuMP.dual_status(m) == MOI.FEASIBLE_POINT
    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.value(x), 5, atol=sol_atol)
    @test Juniper.getnsolutions(inner) == 1
    @test inner.primal_start[1] == 3
end

@testset "bruteforce" begin
    println("==================================")
    println("Bruteforce")
    println("==================================")

    m = Model()

    @variable(m, 1 <= x[1:4] <= 5, Int)

    special_minimizer_fct(x) = x
    grad(x) = 1.0
    grad2(x) = 0.0

    register_args = [:special_minimizer_fct, 1, special_minimizer_fct, grad, grad2]
    JuMP.register(m, register_args...)

    @NLobjective(m,Min,special_minimizer_fct(x[1]))

    juniper_all_solutions = DefaultTestSolver(
        branch_strategy=:StrongPseudoCost,
        list_of_solutions = true,
        strong_restart = true,
        debug = true,
    )

    JuMP.set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_all_solutions...
        )
    )

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    JuMP.set_optimizer_attribute(m, "all_solutions", true)
    JuMP.optimize!(m)
    bm = JuMP.backend(m)
    status = MOI.get(bm, MOI.TerminationStatus())

    innermodel = bm.optimizer.model.inner
    debugDict = innermodel.debugDict
    @test getnstate(debugDict,:Integral) == 24
    @test different_hashes(debugDict) == true
    counter_test(debugDict,Juniper.getnbranches(innermodel))

    list_of_solutions = Juniper.getsolutions(innermodel)
    @test length(unique(list_of_solutions)) == Juniper.getnsolutions(innermodel)

    @test status == MOI.LOCALLY_SOLVED
    @test Juniper.getnsolutions(innermodel) == 24
end

@testset "bruteforce optimizer without attributes" begin
    m = Model(Juniper.Optimizer)
    set_optimizer_attribute(m, "nl_solver", Ipopt.Optimizer)
    set_optimizer_attribute(m, "mip_solver", Cbc.Optimizer)

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    mktemp() do path,io
        out = stdout
        err = stderr
        redirect_stdout(io)
        redirect_stderr(io)

        status = try
            optimize!(m)
        catch e
            e
        end

        flush(io)
        redirect_stdout(out)
        redirect_stderr(err)
    end
    @test length(sort(unique(convert.(Int, JuMP.value.(x))))) == 4
    @test all(1 .<= JuMP.value.(x) .<= 5)
    @test JuMP.value(x[1]) ≈ 1.0
end

@testset "bruteforce full strong w/o restart" begin
    println("==================================")
    println("Bruteforce full strong w/o restart")
    println("==================================")
    juniper_all_solutions = DefaultTestSolver(
        branch_strategy=:StrongPseudoCost,
        all_solutions = true,
        list_of_solutions = true,
        strong_branching_perc = 100,
        strong_branching_nsteps = 100,
        strong_restart = false
    )

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_all_solutions...)
    )

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    status = solve(m)

    list_of_solutions = Juniper.getsolutions(internalmodel(m))
    @test length(unique(list_of_solutions)) == Juniper.getnsolutions(internalmodel(m))

    @test status == MOI.LOCALLY_SOLVED
    @test Juniper.getnsolutions(internalmodel(m)) == 24
end

@testset "bruteforce approx time limit" begin
    println("==================================")
    println("Bruteforce  approx time limit")
    println("==================================")
    juniper_all_solutions = DefaultTestSolver(
        branch_strategy=:StrongPseudoCost,
        strong_branching_approx_time_limit=2,
        all_solutions = true,
        list_of_solutions = true,
        strong_restart = true
    )

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_all_solutions...)
    )

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    status = solve(m)

    list_of_solutions = Juniper.getsolutions(internalmodel(m))
    @test length(unique(list_of_solutions)) == Juniper.getnsolutions(internalmodel(m))

    @test status == MOI.LOCALLY_SOLVED
    @test Juniper.getnsolutions(internalmodel(m)) == 24
end

@testset "bruteforce time limit reliable" begin
    println("==================================")
    println("Bruteforce  time reliable")
    println("==================================")
    juniper_all_solutions = DefaultTestSolver(
        branch_strategy=:Reliability,
        strong_branching_time_limit=0.02,
        reliability_branching_perc=100,
        all_solutions = true,
        list_of_solutions = true,
        strong_restart = true
    )

    optimizer = optimizer_with_attributes(Juniper.Optimizer, juniper_all_solutions...)

    m = Model(optimizer)
    moi_optimizer = JuMP.backend(m).optimizer.model

    # default => 1
    @test MOI.supports(moi_optimizer, MOI.NumberOfThreads()) == true
    MOI.set(moi_optimizer, MOI.NumberOfThreads(), nothing)

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    status = solve(m)

    list_of_solutions = Juniper.getsolutions(internalmodel(m))
    @test length(unique(list_of_solutions)) == Juniper.getnsolutions(internalmodel(m))
    @test MOI.get(moi_optimizer, MOI.NumberOfThreads()) == 1
    # all solutions are saved => nsolutions should equal length(solutions)
    @test MOI.get(moi_optimizer, MOI.ResultCount()) == Juniper.getnsolutions(internalmodel(m))

    @test status == MOI.LOCALLY_SOLVED
    @test Juniper.getnsolutions(internalmodel(m)) == 24
end

@testset "bruteforce PseudoCost" begin
    println("==================================")
    println("Bruteforce PseudoCost")
    println("==================================")
    juniper_all_solutions = DefaultTestSolver(
        branch_strategy=:PseudoCost,
        all_solutions = true,
        list_of_solutions = true,
        strong_restart = true
    )

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_all_solutions...)
    )

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    status = solve(m)
    println("Status: ", status)
    list_of_solutions = Juniper.getsolutions(internalmodel(m))
    @test length(unique(list_of_solutions)) == Juniper.getnsolutions(internalmodel(m))

    @test status == MOI.LOCALLY_SOLVED
    @test Juniper.getnsolutions(internalmodel(m)) == 24
end

@testset "bruteforce Reliability" begin
    println("==================================")
    println("Bruteforce Reliability")
    println("==================================")
    juniper_all_solutions = DefaultTestSolver(
        branch_strategy=:Reliability,
        all_solutions = true,
        list_of_solutions = true,
    )

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_all_solutions...)
    )

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    status = solve(m)
    println("Status: ", status)
    list_of_solutions = Juniper.getsolutions(internalmodel(m))
    @test length(unique(list_of_solutions)) == Juniper.getnsolutions(internalmodel(m))

    @test status == MOI.LOCALLY_SOLVED
    @test Juniper.getnsolutions(internalmodel(m)) == 24
end

@testset "no integer" begin
    println("==================================")
    println("no integer")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_restart...)
    )

    println("Create variables/constr/obj")
    @variable(m, 1 <= x <= 5, start = 2.7)
    @variable(m, -2 <= y <= 2)

    @objective(m, Min, -x-y)

    @NLconstraint(m, y==2*cos(2*x))
    println("before solve")
    status = solve(m)
    println("Status: ", status)

    @test status == MOI.LOCALLY_SOLVED
end

@testset "infeasible cos" begin
    println("==================================")
    println("Infeasible cos")
    println("==================================")
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_restart...)
    )

    @variable(m, 1 <= x <= 5, Int)
    @variable(m, -2 <= y <= 2, Int)

    @objective(m, Min, -x-y)

    @NLconstraint(m, y==2*cos(2*x))

    status = solve(m)

    println("Status: ", status)

    @test status == MOI.LOCALLY_INFEASIBLE
    @test JuMP.termination_status(m) == MOI.LOCALLY_INFEASIBLE
    @test JuMP.primal_status(m) == MOI.INFEASIBLE_POINT
    @test JuMP.dual_status(m) == MOI.INFEASIBLE_POINT
    @test isnan(getobjgap(m))
end

@testset "infeasible int reliable" begin
    println("==================================")
    println("Infeasible int reliable")
    println("==================================")
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_reliable_restart...)
    )

    @variable(m, 1 <= x <= 5, Int)
    @variable(m, -2 <= y <= 2, Int)

    @objective(m, Min, -x-y)

    @NLconstraint(m, y >= sqrt(2))
    @NLconstraint(m, y <= sqrt(3))

    status = solve(m)
    println("Status: ", status)

    @test status == MOI.LOCALLY_INFEASIBLE
    @test isnan(getobjgap(m))
end

@testset "infeasible sin with different bounds" begin
    println("==================================")
    println("Infeasible  sin with different bounds")
    println("==================================")
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(
            branch_strategy=:MostInfeasible,
            feasibility_pump = true,
            time_limit = 1,
            mip_solver=optimizer_with_attributes(Cbc.Optimizer)
        )...)
    )

    @variable(m, x, Int, start=3)
    @variable(m, y >= 2, Int)

    @constraint(m, 0 <= x <= 5)

    @objective(m, Min, -x-y)

    @NLconstraint(m, y==sin(x))

    status = solve(m)

    @test status == MOI.LOCALLY_INFEASIBLE
end

@testset "infeasible relaxation" begin
    println("==================================")
    println("Infeasible relaxation")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;debug=true)...)
    )

    @variable(m, 0 <= x[1:10] <= 2, Int)

    @objective(m, Min, sum(x))

    @constraint(m, sum(x[1:5]) <= 20)
    @NLconstraint(m, x[1]*x[2]*x[3] >= 10)

    status = solve(m)

    debug1 = internalmodel(m).debugDict

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;debug=true)...)
    )

    @variable(m, 0 <= x[1:10] <= 2, Int)

    @objective(m, Min, sum(x))

    @constraint(m, sum(x[1:5]) <= 20)
    @NLconstraint(m, x[1]*x[2]*x[3] >= 10)

    status = solve(m)

    debug2 = internalmodel(m).debugDict
    opts = internalmodel(m).options

    # should be deterministic
    @test debug1[:relaxation][:nrestarts] == debug2[:relaxation][:nrestarts] == opts.num_resolve_root_relaxation
    for i=1:debug1[:relaxation][:nrestarts]
        @test debug1[:relaxation][:restarts][i] == debug2[:relaxation][:restarts][i]
    end

    println("Status: ", status)

    @test status == MOI.LOCALLY_INFEASIBLE
end

@testset "infeasible relaxation 2" begin
    println("==================================")
    println("Infeasible relaxation 2")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_no_restart...)
    )

    @variable(m, x[1:3], Int)
    @variable(m, y)

    @objective(m, Max, sum(x))

    @NLconstraint(m, x[1]^2+x[2]^2+x[3]^2+y^2 <= 3)
    @NLconstraint(m, x[1]^2*x[2]^2*x[3]^2*y^2 >= 10)

    status = solve(m)
    println("Status: ", status)

    @test status == MOI.LOCALLY_INFEASIBLE
end

@testset "infeasible integer" begin
    println("==================================")
    println("Infeasible integer")
    println("==================================")
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_no_restart...)
    )

    @variable(m, 0 <= x[1:10] <= 2, Int)

    @objective(m, Min, sum(x))

    @constraint(m, sum(x[1:5]) <= 20)
    @NLconstraint(m, x[1]*x[2]*x[3] >= 7)
    @NLconstraint(m, x[1]*x[2]*x[3] <= 7.5)

    status = solve(m)
    println("Status: ", status)

    @test status == MOI.LOCALLY_INFEASIBLE
end

@testset "infeasible in strong" begin
    println("==================================")
    println("Infeasible in strong")
    println("==================================")
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_no_restart...)
    )

    @variable(m, 0 <= x[1:5] <= 2, Int)

    @objective(m, Min, sum(x))

    @NLconstraint(m, x[3]^2 <= 2)
    @NLconstraint(m, x[3]^2 >= 1.2)

    status = solve(m)
    println("Status: ", status)

    @test status == MOI.LOCALLY_INFEASIBLE
end

@testset "One Integer small Reliable" begin
    println("==================================")
    println("One Integer small Reliable")
    println("==================================")
    m = Model()

    @variable(m, x >= 0, Int, start=2)
    @variable(m, y >= 0)
    @variable(m, 0 <= u <= 10, Int)
    @variable(m, w == 1)

    myf(x,y) = -3x-y
    function ∇f(g,x,y)
        g[1] = -3
        g[2] = -1
    end

    register_args = [:myf, 2, myf, ∇f]
    JuMP.register(m, register_args...)

    @NLobjective(m,Min,myf(x,y))

    JuMP.set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_reliable_restart...
        )
    )

    @constraint(m, 3x + 10 <= 20)
    @NLconstraint(m, y^2 <= u*w)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))
    println("x: ", JuMP.value(x))
    println("y: ", JuMP.value(y))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), -12.162277, atol=opt_atol)
    @test isapprox(JuMP.value(x), 3, atol=sol_atol)
    @test isapprox(JuMP.value(y), 3.162277, atol=sol_atol)
end

@testset "One Integer small Strong" begin
    println("==================================")
    println("One Integer small Strong")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_no_restart...)
    )

    @variable(m, x >= 0, Int)
    @variable(m, y >= 0)
    @variable(m, 0 <= u <= 10, Int)
    @variable(m, w == 1)

    @objective(m, Min, -3x - y)

    @constraint(m, 3x + 10 <= 20)
    @NLconstraint(m, y^2 <= u*w)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))
    println("x: ", JuMP.value(x))
    println("y: ", JuMP.value(y))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), -12.162277, atol=opt_atol)
    @test isapprox(JuMP.value(x), 3, atol=sol_atol)
    @test isapprox(JuMP.value(y), 3.162277, atol=sol_atol)
end

@testset "One Integer small MostInfeasible" begin
    println("==================================")
    println("One Integer small MostInfeasible")
    println("==================================")

    m = Model()

    @variable(m, x >= 0, Int)
    @variable(m, y >= 0)
    @variable(m, 0 <= u <= 10, Int)
    @variable(m, w == 1)

    special_minimizer_fct(x,y) = -3x-y
    register_args = [:special_minimizer_fct, 2, special_minimizer_fct]
    JuMP.register(m, register_args...; autodiff=true)

    @NLobjective(m,Min,special_minimizer_fct(x, y))

    @constraint(m, 3x + 10 <= 20)
    @NLconstraint(m, y^2 <= u*w)

    JuMP.set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(
            branch_strategy=:MostInfeasible,
        )...)
    )

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))
    println("x: ", JuMP.value(x))
    println("y: ", JuMP.value(y))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), -12.162277, atol=opt_atol)
    @test isapprox(JuMP.value(x), 3, atol=sol_atol)
    @test isapprox(JuMP.value(y), 3.162277, atol=sol_atol)
end

@testset "One Integer small PseudoCost" begin
    println("==================================")
    println("One Integer small PseudoCost")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_pseudo...)
    )

    @variable(m, x >= 0, Int)
    @variable(m, y >= 0)
    @variable(m, 0 <= u <= 10, Int)
    @variable(m, w == 1)

    @objective(m, Min, -3x - y)

    @constraint(m, 3x + 10 <= 20)
    @NLconstraint(m, y^2 <= u*w)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))
    println("x: ", JuMP.value(x))
    println("y: ", JuMP.value(y))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), -12.162277, atol=opt_atol)
    @test isapprox(JuMP.value(x), 3, atol=sol_atol)
    @test isapprox(JuMP.value(y), 3.162277, atol=sol_atol)
end

@testset "Three Integers Small Strong" begin
    println("==================================")
    println("Three Integers Small Strong")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_no_restart...)
    )

    @variable(m, x >= 0, Int)
    @variable(m, y >= 0, Int)
    @variable(m, 0 <= u <= 10, Int)
    @variable(m, w == 1)

    @objective(m, Min, -3x - y)

    @constraint(m, 3x + 10 <= 20)
    @NLconstraint(m, y^2 <= u*w)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), -12, atol=opt_atol)
    @test isapprox(JuMP.value(x), 3, atol=sol_atol)
    @test isapprox(JuMP.value(y), 3, atol=sol_atol)
end

@testset "Three Integers Small MostInfeasible" begin
    println("==================================")
    println("Three Integers Small MostInfeasible")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_mosti...)
    )

    @variable(m, x >= 0, Int)
    @variable(m, y >= 0, Int)
    @variable(m, 0 <= u <= 10, Int)
    @variable(m, w == 1)

    @objective(m, Min, -3x - y)

    @constraint(m, 3x + 10 <= 20)
    @NLconstraint(m, y^2 <= u*w)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), -12, atol=opt_atol)
    @test isapprox(JuMP.value(x), 3, atol=sol_atol)
    @test isapprox(JuMP.value(y), 3, atol=sol_atol)
end

@testset "Three Integers Small PseudoCost" begin
    println("==================================")
    println("Three Integers Small PseudoCost")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_pseudo...)
    )

    @variable(m, x >= 0, Int)
    @variable(m, y >= 0, Int)
    @variable(m, 0 <= u <= 10, Int)
    @variable(m, w == 1)

    @objective(m, Min, -3x - y)

    @constraint(m, 3x + 10 <= 20)
    @NLconstraint(m, y^2 <= u*w)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), -12, atol=opt_atol)
    @test isapprox(JuMP.value(x), 3, atol=sol_atol)
    @test isapprox(JuMP.value(y), 3, atol=sol_atol)
end

@testset "Knapsack Max" begin
    println("==================================")
    println("KNAPSACK")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;traverse_strategy=:DBFS,
            incumbent_constr=true,mip_solver=optimizer_with_attributes(Cbc.Optimizer),
            strong_branching_time_limit=1)...)
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Bin)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), 65, atol=opt_atol)
    @test isapprox(JuMP.objective_bound(m), 65, atol=opt_atol)
    @test isapprox(JuMP.value.(x), [0,0,0,1,1], atol=sol_atol)
end

@testset "Knapsack Max Reliable" begin
    println("==================================")
    println("KNAPSACK Reliable no restart")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;branch_strategy=:Reliability,
              strong_restart=false,strong_branching_time_limit=1,gain_mu=0.5)...)
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Bin)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), 65, atol=opt_atol)
    @test isapprox(JuMP.objective_bound(m), 65, atol=opt_atol)
    @test isapprox(JuMP.value.(x), [0,0,0,1,1], atol=sol_atol)
end

@testset "Integer at root" begin
    println("==================================")
    println("INTEGER AT ROOT")
    println("==================================")
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver()...)
    )

    @variable(m, x[1:6] <= 1, Int)
    @constraint(m, x[1:6] .== 1)
    @objective(m, Max, sum(x))
    @NLconstraint(m, x[1]*x[2]*x[3]+x[4]*x[5]*x[6] <= 100)
    status = solve(m)
    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), 6, atol=opt_atol)
    # objective bound doesn't exists anymore in Ipopt
    #@test isapprox(JuMP.objective_bound(m), 0, atol=opt_atol) # Ipopt return 0
end

@testset "Knapsack Max with epsilon" begin
    println("==================================")
    println("KNAPSACK with epsilon")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;traverse_strategy=:DBFS,obj_epsilon=0.5)...)
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Bin)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    status = solve(m)
    println("Obj: ", JuMP.objective_value(m))

    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(JuMP.objective_value(m), 65, atol=opt_atol)
    @test isapprox(JuMP.value.(x), [0,0,0,1,1], atol=sol_atol)
end

@testset "Knapsack Max with epsilon too strong" begin
    println("==================================")
    println("KNAPSACK with epsilon too strong")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(;traverse_strategy=:DBFS,obj_epsilon=0.1)...)
    )

    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    # set all to 1 which is infeasible otherwise incumbent solution would be found
    @variable(m, x[1:5], Bin, start=1)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    status = solve(m)

    @test status == MOI.LOCALLY_INFEASIBLE
end

@testset "Batch.mod Restart" begin
    println("==================================")
    println("BATCH.MOD RESTART")
    println("==================================")

    m = batch_problem()

    JuMP.set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_restart...)
    )

    status = solve(m)
    @test status == MOI.LOCALLY_SOLVED

    juniper_val = JuMP.objective_value(m)

    println("Solution by Juniper")
    println("obj: ", juniper_val)

    @test isapprox(juniper_val, 285506.5082, atol=opt_atol, rtol=opt_rtol)
end

@testset "Batch.mod Restart 2 Levels" begin
    println("==================================")
    println("BATCH.MOD RESTART 2 LEVELS")
    println("==================================")

    m = batch_problem()

    JuMP.set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_restart_2...)
    )

    status = solve(m)
    @test status == MOI.LOCALLY_SOLVED

    juniper_val = JuMP.objective_value(m)

    println("Solution by Juniper")
    println("obj: ", juniper_val)

    @test isapprox(juniper_val, 285506.5082, atol=opt_atol, rtol=opt_rtol)
end

@testset "cvxnonsep_nsig20r.mod restart" begin
    println("==================================")
    println("cvxnonsep_nsig20r.MOD RESTART")
    println("==================================")

    m = cvxnonsep_nsig20r_problem()

    JuMP.set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_restart...)
    )

    status = solve(m)
    @test status == MOI.LOCALLY_SOLVED

    juniper_val = JuMP.objective_value(m)

    println("Solution by Juniper")
    println("obj: ", juniper_val)

    @test isapprox(juniper_val, 80.9493, atol=opt_atol, rtol=opt_rtol)
end

@testset "cvxnonsep_nsig20r.mod no restart" begin
    println("==================================")
    println("cvxnonsep_nsig20r.MOD NO RESTART")
    println("==================================")

    m = cvxnonsep_nsig20r_problem()

    JuMP.set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_strong_no_restart...)
    )

    status = solve(m)
    @test status == MOI.LOCALLY_SOLVED

    juniper_val = JuMP.objective_value(m)

    println("Solution by Juniper")
    println("obj: ", juniper_val)

    @test isapprox(juniper_val, 80.9493, atol=opt_atol, rtol=opt_rtol)
end

@testset "Knapsack solution limit and table print test" begin
    println("==================================")
    println("Knapsack solution limit and table print test")
    println("==================================")

    juniper_one_solution = DefaultTestSolver(
        log_levels=[:Table],
        branch_strategy=:MostInfeasible,
        solution_limit=1
    )

    m = Model()
    v = [10,20,12,23,42]
    w = [12,45,12,22,21]
    @variable(m, x[1:5], Bin)

    @objective(m, Max, dot(v,x))

    @NLconstraint(m, sum(w[i]*x[i]^2 for i=1:5) <= 45)

    JuMP.set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_one_solution...)
    )

    status = solve(m)
    @test status == MOI.SOLUTION_LIMIT

    # maybe 2 found at the same time
    @test Juniper.getnsolutions(internalmodel(m)) <= 2
    @test Juniper.getnsolutions(internalmodel(m)) >= 1
end

@testset "bruteforce obj_epsilon" begin
    println("==================================")
    println("Bruteforce obj_epsilon")
    println("==================================")
    juniper_obj_eps = DefaultTestSolver(
        branch_strategy=:StrongPseudoCost,
        obj_epsilon=0.4
    )

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_obj_eps...)
    )

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    status = solve(m)

    # maybe 2 found at the same time
    @test status == MOI.LOCALLY_SOLVED
end

@testset "bruteforce best_obj_stop not reachable" begin
    println("==================================")
    println("Bruteforce best_obj_stop not reachable")
    println("==================================")
    juniper_best_obj_stop = DefaultTestSolver(
        branch_strategy=:StrongPseudoCost,
        best_obj_stop=0.8
    )
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_best_obj_stop...)
    )

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    status = solve(m)

    # not possible to reach but should be solved anyway
    @test status == MOI.LOCALLY_SOLVED
end

@testset "bruteforce best_obj_stop reachable" begin
    println("==================================")
    println("Bruteforce best_obj_stop reachable")
    println("==================================")
    juniper_one_solution = DefaultTestSolver(
        branch_strategy=:StrongPseudoCost,
        best_obj_stop=1
    )

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_one_solution...)
    )

    @variable(m, 1 <= x[1:4] <= 5, Int)

    @objective(m, Min, x[1])

    @constraint(m, x[1] >= 0.9)
    @constraint(m, x[1] <= 1.1)
    @NLconstraint(m, (x[1]-x[2])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[3])^2 >= 0.1)
    @NLconstraint(m, (x[1]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[2]-x[4])^2 >= 0.1)
    @NLconstraint(m, (x[3]-x[4])^2 >= 0.1)

    status = solve(m)

    # reachable and should break
    @test status == MOI.OBJECTIVE_LIMIT
    # maybe 2 found at the same time
    @test Juniper.getnsolutions(internalmodel(m)) <= 2
    @test Juniper.getnsolutions(internalmodel(m)) >= 1
end

# this test has a lot "Only almost locally solved" warnings (in mumps at least)
@testset "Sum 1/x = 2" begin
    println("==================================")
    println("Sum 1/x = 2")
    println("==================================")
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        juniper_pseudo...)
    )

    @variable(m, 1 <= x[1:11], Int)
    @constraint(m, [i=2:11], x[i-1] <= x[i] - 1)
    @NLconstraint(m, sum(1 / x[i] for i in 1:11) == 2)

    status = solve(m)

    # in mumps/on travis this returns ALMOST_LOCALLY_SOLVED but with ma27/locally it is LOCALLY_SOLVED
    @test Juniper.state_is_optimal(status; allow_almost=true)
    @test isapprox(2, sum(1/v for v in JuMP.value.(x)), atol=opt_atol)
end

# this test has a lot "Only almost locally solved" warnings
@testset "Sum 1/x = 2 don't allow almost" begin
    println("==================================")
    println("Sum 1/x = 2 don't allow almost")
    println("==================================")
    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(
                branch_strategy=:PseudoCost,
                allow_almost_solved = false
            )...)
    )

    @variable(m, 1 <= x[1:11], Int)
    @constraint(m, [i=2:11], x[i-1] <= x[i] - 1)
    @NLconstraint(m, sum(1 / x[i] for i in 1:11) == 2)

    status = solve(m)

    # even without almost this should still be solveable
    @test status == MOI.LOCALLY_SOLVED
    @test isapprox(2, sum(1/v for v in JuMP.value.(x)), atol=opt_atol)
end

#this test has an expression where a variable will be dereferenced twice
@testset "Nested variable reference" begin
    println("==================================")
    println("Nested variable reference")
    println("==================================")

    m = Model(optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver()...)
    )

    x = @variable(m, x[i=1:5,j=1:5], Bin)
    xrowsum = @NLexpression(m, xrowsum[i=1:5], sum(x[i,j] for j in 1:5))
    @NLobjective(m, Max, sum(sum(x[i,j] + xrowsum[i] for i=1:5) for j = 1:5))

    optimize!(m)

    @test JuMP.termination_status(m) == MOI.LOCALLY_SOLVED
end

@testset "Expr dereferencing for @NLexpression" begin
    # See issue 184
    m = Model()
    set_optimizer(m, optimizer_with_attributes(
        Juniper.Optimizer,
        DefaultTestSolver(
            mip_solver=optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0))...
    ))

    @variable(m, 0 <= a_var <= 1)
    @variable(m, bin_var, Int)
    b_expr = @NLexpression(m, bin_var/1)
    @NLconstraint(m, 1.1 >= b_expr)
    an_expr = @NLexpression(m, a_var / 1)

    @NLobjective(m, Max, bin_var + an_expr)

    optimize!(m)
    @test JuMP.objective_value(m) ≈ 2.0
    @test JuMP.value(bin_var) ≈ 1
    @test JuMP.value(a_var) ≈ 1
end

end


using DataFrames
using CSV
using Statistics

filenames = [

    ### Put result table files here for reviewing the results, e.g.

    # LD 10
    # "Experiments/results/ld10_2020-02-19_003851.csv"

    # LD 3
    # "Experiments/results/ld3_2020-02-19_050959.csv"

    # CLD 10
    # "Experiments/results/cld10_2020-02-19_015442.csv"

    # CLD 3
    # "Experiments/results/cld3_2020-02-19_030803.csv"

    ]

data = vcat([CSV.read(fname) for fname in filenames]...)

r = by(data, :solver, N = :reward => length, mean = :reward => mean, std = :reward => r->std(r)/sqrt(length(r)))
sort!(r, :mean, rev=true)
display(filenames)
@show r


mutable struct GPRealisations{KEY}
    groupKeys::Vector{KEY}
    mgp::Dict{KEY,GPE}
    nobs::Int
    logNoise::Float64
    mean::Mean
    kernel::Kernel
    mll::Float64            # Marginal log-likelihood
    dmll::Vector{Float64}   # Gradient marginal log-likelihood
    function GPRealisations{KEY}(
            groupKeys::Vector{KEY}, mgp::Dict{KEY,GPE}, nobs::Int, logNoise::Float64,
            # groupKeys::Vector{KEY}, mgp::Dict{KEY,GPE}, nobs::Int, logNoise::Scalar,
            mean::Mean, kernel::Kernel) where {KEY}
        length(groupKeys) == length(mgp) || throw("groupKeys and mgp should have same length")
        gpreals = new(groupKeys, mgp, nobs, logNoise, mean, kernel)
        propagate_params!(gpreals)
        update_mll!(gpreals)
        return gpreals
    end
end
function GPRealisations(gpList::Vector{GPE}, groupKeys::Vector{KEY}) where {KEY}
    total_nobs = sum([gp.nobs for gp in gpList])
    # get kernel and parameters from first GP in list
    first_gp = gpList[1]
    dim = first_gp.dim
    logNoise = first_gp.logNoise
    kern = first_gp.kernel
    mean = first_gp.mean
    ngroups = length(groupKeys)
    @assert ngroups == length(gpList)
    gpDict = Dict{KEY, GPE}()
    istart = 1
    for j in 1:ngroups
        gp = gpList[j]
        key = groupKeys[j]
        nobs = gp.nobs
        gp.nobs > 0 || throw("empty group")
        gpDict[key] = gp
        istart += nobs
    end
    gpreals = GPRealisations{KEY}(
                groupKeys, gpDict,
                total_nobs, logNoise, mean, kern)
    return gpreals
end
function GPRealisations{KEY}(groupKeys, mgp, nobs,
        logNoise::Scalar, mean, kernel) where {KEY}
    logNoise_float = convert(Float64, logNoise)
    return GPRealisations{KEY}(groupKeys, mgp, nobs,
            logNoise_float, mean, kernel)
end

Base.getindex(gpreals::GPRealisations{KEY}, key::KEY) where {KEY} = gpreals.mgp[key]
Base.keys(gpreals::GPRealisations) = gpreals.groupKeys
Base.values(gpreals::GPRealisations) = values(gpreals.mgp)

function get_params(gpreals::GPRealisations;
                    noise::Bool=true, domean::Bool=true, kern::Bool=true)
    params = Float64[]
    if noise; push!(params, gpreals.logNoise); end
    if domean
        for gp in values(gpreals.mgp)
            append!(params, get_params(gp.mean))
        end
    end
    if kern; append!(params, get_params(gpreals.kernel)); end
    return params
end
function propagate_params!(gpreals::GPRealisations; noise::Bool=true, kern::Bool=true)
    for gp in values(gpreals.mgp)
        # harmonize parameters
        if kern
            gp.kernel = gpreals.kernel
        end
        if noise
            gp.logNoise = gpreals.logNoise
        end
    end
end
function set_params!(gpreals::GPRealisations, hyp::Vector{Float64};
                     noise::Bool=true, domean::Bool=true, kern::Bool=true)
    # println("mean=$(mean)")
    istart=1
    if noise
        gpreals.logNoise = hyp[istart]
        istart += 1
    end
    if domean
        for key in gpreals.groupKeys
            gp = gpreals.mgp[key]
            set_params!(gp.mean, hyp[istart:istart+num_params(gp.mean)-1])
            istart += num_params(gp.mean)
        end
    end
    if kern
        set_params!(gpreals.kernel, hyp[istart:end])
    end
    propagate_params!(gpreals, noise=noise, kern=kern)
end

function update_mll!(gpreals::GPRealisations)
    mll = 0.0
    for gp in values(gpreals.mgp)
        update_mll!(gp)
        mll += gp.mll
    end
    gpreals.mll = mll
    return mll
end
function update_mll_and_dmll!(gpreals::GPRealisations,
                              ααinvcKIs::Dict{Int,Matrix};
                              noise::Bool=true, domean::Bool=true, kern::Bool=true)
    gpreals.mll = 0.0
    gpreals.dmll = zeros(get_params(gpreals; noise=noise, domean=domean, kern=kern))
    imean=2
    ikern=collect((length(gpreals.dmll)-num_params(gpreals.kernel)+1):length(gpreals.dmll))
    for gp in values(gpreals.mgp)
        update_mll_and_dmll!(gp, ααinvcKIs[gp.nobs];
            noise=noise,  domean=domean, kern=kern)
        gpreals.mll += gp.mll
        dmll_indices=Int[]
        if noise
            push!(dmll_indices, 1)
        end
        if domean
            append!(dmll_indices, imean:imean+num_params(gp.mean)-1)
            imean+=num_params(gp.mean)
        end
        if kern
            append!(dmll_indices, ikern)
        end
        gpreals.dmll[dmll_indices] .+= gp.dmll
    end
    return gpreals.dmll
end
function update_mll_and_dmll!(gpreals::GPRealisations; kwargs...)
    ααinvcKIs = Dict{Int,Matrix}()
    for gp in values(gpreals.mgp)
        if haskey(ααinvcKIs, gp.nobs)
            continue
        end
        ααinvcKIs[gp.nobs] = Array{Float64}(gp.nobs, gp.nobs)
    end
    return update_mll_and_dmll!(gpreals, ααinvcKIs)
end

function get_optim_target(gpreals::GPRealisations;
                          noise::Bool=true, domean::Bool=true, kern::Bool=true)
    ααinvcKIs = Dict{Int,Matrix}()
    for gp in values(gpreals.mgp)
        if haskey(ααinvcKIs, gp.nobs)
            continue
        end
        ααinvcKIs[gp.nobs] = Array{Float64}(gp.nobs, gp.nobs)
    end
    function mll(hyp::Vector{Float64})
        try
            set_params!(gpreals, hyp; noise=noise, domean=domean, kern=kern)
            update_mll!(gpreals)
            return -gpreals.mll
        catch err
             if !all(isfinite(hyp))
                println(err)
                return Inf
            elseif isa(err, ArgumentError)
                println(err)
                return Inf
            elseif isa(err, Base.LinAlg.PosDefException)
                println(err)
                return Inf
            else
                throw(err)
            end
        end
    end

    function mll_and_dmll!(grad::Vector{Float64}, hyp::Vector{Float64})
        set_params!(gpreals, hyp; noise=noise, domean=domean, kern=kern)
        update_mll_and_dmll!(gpreals, ααinvcKIs; noise=noise, domean=domean, kern=kern)
        grad[:] = -gpreals.dmll
        return -gpreals.mll
    end
    function dmll!(grad::Vector{Float64}, hyp::Vector{Float64})
        mll_and_dmll!(grad, hyp)
    end

    func = OnceDifferentiable(mll, dmll!, mll_and_dmll!,
        get_params(gpreals;noise=noise,domean=domean,kern=kern))
    return func
end
function optimize!(gpreals::GPRealisations;
                   noise::Bool=true, domean::Bool=true, kern::Bool=true)
    func = get_optim_target(gpreals, noise=noise, domean=domean, kern=kern)
    init = get_params(gpreals;  noise=noise, domean=domean, kern=kern)  # Initial hyperparameter values
    nparams = length(init)
    lower = -Inf*ones(nparams)
    upper = Inf*ones(nparams)
    ikernstart=1
    if noise
        ikernstart+=1
        lower[1]=-10.0
    end
    upper[ikernstart:ikernstart+num_params(gpreals.kernel)-1]=20.0
    best_x = copy(init)
    best_y = Inf
    count = 0
    function myfunc(x::Vector, grad::Vector)
        count += 1
        if length(grad) > 0
            y = func.fg!(grad, x)
            if y < best_y
                best_x[:] = x
            end
            return y
        else
            try
                y = func.f(x)
                if y < best_y
                    best_x[:] = x
                end
                return y
            catch
                return Inf
            end
        end
    end

    opt = Opt(:LD_LBFGS, nparams)
    lower_bounds!(opt, lower)
    upper_bounds!(opt, upper)
    min_objective!(opt, myfunc)
    xtol_rel!(opt,1e-4)
    ftol_rel!(opt, 1e-20)

    (minf,minx,ret) = NLopt.optimize(opt, init)
    myfunc(best_x, [])
    return best_x, count
end


module MCMCChainsStorage

import Base: read, write

using HDF5
using MCMCChains

"""
    write(f::Union{HDF5.File, HDF5.Group}, chains::Chains)

Write MCMCChains object to the HDF5 file or group.
"""
function Base.write(f::Union{HDF5.File, HDF5.Group}, c::Chains)
    for s in sections(c)
        g = create_group(f, string(s))
        for n in names(c, s)
            g[string(n), shuffle=(), compress=3] = Array(c[n])
        end
    end
end

"""
    read(f::Union{HDF5.File, HDF5.Group}, ::Type{Chains})

Read a chain object from the given HDF5 file our group.

"""
function Base.read(f::Union{HDF5.File, HDF5.Group}, ::Type{Chains})
    secs = keys(f)
    pns = []
    datas = []
    name_map = Dict()
    for s in secs
        ns = keys(f[s])
        name_map[Symbol(s)] = ns
        for n in ns
            push!(pns, n)
            push!(datas, read(f[s], n))
        end
    end

    nc, ns = size(datas[1])
    np = size(datas,1)

    a = zeros(nc, np, ns)

    for i in 1:np
        a[:,i,:] = datas[i]
    end

    Chains(a, pns, name_map)
end

end # module


# Having Plot as an argument is a lazy way to avoid using RecipesBase
function gray_square(xs,ys,Plot; inside = true)
  if inside == true
    dx = abs(xs[2]- xs[1])/160.0;
    dy = abs(ys[2]- ys[1])/160.0;
    ys[1] += dy
    ys[2] -= dy
    xs[1] += dx
    xs[2] -= dx
  end
    Plot([xs[1],xs[2]],[ys[2],ys[2]], linecolor = :grey, linestyle=:dash, lab="",border=false);
    Plot([xs[2],xs[2]],[ys[2],ys[1]], linecolor = :grey, linestyle=:dash, lab="",border=false);
    Plot([xs[2],xs[1]],[ys[1],ys[1]], linecolor = :grey, linestyle=:dash, lab="",border=false);
    Plot([xs[1],xs[1]],[ys[1],ys[2]], linecolor = :grey, linestyle=:dash, lab="",border=false)
end


export TR_step_computation

function TR_step_computation(h  :: Any,
                             g  :: Any,
                             dN :: Any,
                             Δ  :: Any)
  if true in (h .> 0.0)
    if true in (g .> 0.0)
      d = max.(-Δ, dN)
    else
      d = min.(dN, Δ)
    end
  else
    if true in (g .> 0.0)
      d = .- Δ
    else
      d = Δ
    end
  end

  d = d[1]

  return [d]
end