# Load in the dependencies
using InfiniteOpt, JuMP, MathOptInterface, Distributions, Random,
FastGaussQuadrature, DataStructures

# load the test module
using Test

# Define convenient aliases
const IC = InfiniteOpt.Collections
const MT = InfiniteOpt.MeasureToolbox
const IOTO = InfiniteOpt.TranscriptionOpt
const JuMPC = JuMP.Containers
const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const MOIUC = MathOptInterface.Utilities.CleverDicts
const FGQ = FastGaussQuadrature
const IOMT = InfiniteOpt.MeasureToolbox

# Load in testing utilities
include("utilities.jl")

# Run unit tests
println("----------------------------------------------------------------------------")
println("---------------------------------UNIT TESTS---------------------------------")
println("----------------------------------------------------------------------------")
@time @testset "Collections" begin
    include("Collections/VectorTuple.jl")
    include("Collections/DualDict.jl")
end
println("")
@time @testset "Datatypes" begin include("datatypes.jl") end
println("")
@time @testset "Utilities" begin include("utility_tests.jl") end
println("")
@time @testset "Infinite Set Methods" begin include("infinite_sets.jl") end
println("")
@time @testset "General Variable Methods" begin
    include("general_variables.jl")
end
println("")
@time @testset "Optimizer Setup Methods" begin include("optimizer_setup.jl") end
println("")
@time @testset "Macro Utilities" begin include("macro_utilities.jl") end
println("")
@time @testset "Parameter Methods" begin
   @testset "Scalar" begin include("scalar_parameters.jl") end
   @testset "Array" begin include("array_parameters.jl") end
end
println("")
@time @testset "Variable Methods" begin
   @testset "Infinite Variables" begin include("infinite_variables.jl") end
   @testset "Point Variables" begin include("point_variables.jl") end
   @testset "Hold Variables" begin include("hold_variables.jl") end
   @testset "Info Constraints" begin include("variable_info.jl") end
   @testset "Reduced Variables" begin include("reduced_variables.jl") end
end
println("")
@time @testset "Derivative Methods" begin include("derivatives.jl") end
println("")
@time @testset "Operators" begin include("operators.jl") end
println("")
@time @testset "Expression Methods" begin include("expressions.jl") end
println("")
@time @testset "Macro Expressions" begin include("macro_expressions.jl") end
println("")
@time @testset "Measure Methods" begin include("measures.jl") end
println("")
@time @testset "Measure Toolbox Methods" begin
    @testset "Integrals" begin include("MeasureToolbox/integrals.jl") end
    @testset "Expectations" begin include("MeasureToolbox/expectations.jl") end
    @testset "Support Sums" begin include("MeasureToolbox/support_sums.jl") end
end
println("")
@time @testset "Objective Methods" begin include("objective.jl") end
println("")
@time @testset "Constraint Methods" begin include("constraints.jl") end
println("")
@time @testset "Printing Methods" begin include("show.jl") end
println("")
@time @testset "Deletion Methods" begin include("deletion.jl") end
println("")
@time @testset "Expansion Methods" begin include("measure_expansions.jl") end
println("")
@time @testset "Derivative Evaluation" begin include("derivative_evaluation.jl") end
println("")
@time @testset "TranscriptionOpt" begin
    @testset "Model" begin include("TranscriptionOpt/model.jl") end
    @testset "Measures" begin include("TranscriptionOpt/measure.jl") end
    @testset "Transcribe" begin include("TranscriptionOpt/transcribe.jl") end
    @testset "Optimize" begin include("TranscriptionOpt/optimize.jl") end
end
println("")
@time @testset "Solution Methods" begin include("optimizer.jl") end
println("")
@time @testset "Solution Queries" begin include("results.jl") end
println("")
@time @testset "Extensions" begin include("extensions.jl") end
println("")
println("----------------------------------------------------------------------------")
println("-----------------------------TESTING COMPLETE!------------------------------")
println("----------------------------------------------------------------------------")


include("src/Blake3Hash.jl")

using BenchmarkTools
using .Blake3Hash

@benchmark begin
    Blake3Hash.Simd128.update!(h, d)
    Blake3Hash.Simd128.digest(h, o)
end setup=begin
    h = Blake3Hash.Simd128.Blake3Ctx()
    d = rand(UInt8, 2*8192)
    o = Vector{UInt8}(undef, 32)
end


using Test

@testset "ITensors.jl" begin
  @testset "$filename" for filename in [
    "tagset.jl",
    "smallstring.jl",
    "symmetrystyle.jl",
    "index.jl",
    "indexset.jl",
    "not.jl",
    "itensor.jl",
    "itensor_slice.jl",
    "itensor_scalar_contract.jl",
    "itensor_combine_contract.jl",
    "itensor_contraction_sequence_optimization.jl",
    "broadcast.jl",
    "diagitensor.jl",
    "contract.jl",
    "combiner.jl",
    "debug_checks.jl",
    "trg.jl",
    "ctmrg.jl",
    "iterativesolvers.jl",
    "dmrg.jl",
    "sitetype.jl",
    "phys_site_types.jl",
    "decomp.jl",
    "lattices.jl",
    "mps.jl",
    "mpo.jl",
    "sweeps.jl",
    "sweepnext.jl",
    "autompo.jl",
    "svd.jl",
    "qn.jl",
    "qnindex.jl",
    "qnitensor.jl",
    "qncombiner.jl",
    "qndiagitensor.jl",
    "empty.jl",
    "qnmpo.jl",
    "readwrite.jl",
    "readme.jl",
    "examples.jl",
  ]
    println("Running $filename")
    include(filename)
  end
end


using Minibatch
using Revtok
using DataStructures

struct Vocab
    itos::Vector
    stoi::Associative
end
Base.length(vocab::Vocab) = length(vocab.itos)

function process(corpus, vocabsize=200, batchsize=4)
    corpus = tokenize.(split(corpus, '\n'))
    segments = buildvocab(counter(Iterators.flatten(corpus)), vocabsize)
    corpus = segment.(corpus, segments)
    itos = keys(segments)
    stoi = Dict((reverse(p) for p in enumerate(itos)))
    vocab = Vocab(collect(itos), stoi)
    corpus = [[vocab.stoi[x] for x in ex] for ex in corpus]
    corpus = Iterators.partition(corpus, batchsize)
    corpus = map(b -> MaskedBatch(b, (true,)), corpus)
    return corpus, vocab
end

english = """This is an account of how magical thinking made us modern.
When people talk about magical thinking, it is usually as a cognitive feature of children, uneducated people, the mushy-minded, or the mentally ill.
If we notice magical thinking in ourselves, it is with a pang of shame: literate adults are supposed to be more sophisticated than that.
At the same time, magical thinking is obviously rampant in the world.
It’s hard not to be fascinated, even if it’s a horrified fascination.
I think that the reason that it has been so difficult to precisely define “magical thinking” is that what we call “magical thinking” is a collection of stigmatized examples of a more general, and generally useful, cognitive capacity.
This is the ability to think in “as if” mode: “as if” inanimate objects had minds, “as if” thoughts could affect reality, “as if” symbols had power over their referents."""

spanish = """Este es un relato de cómo el pensamiento mágico nos hizo modernos.
Cuando la gente habla de pensamiento mágico, generalmente es como un rasgo cognitivo de los niños, de las personas sin educación, de los que tienen la mente blanda o de los enfermos mentales.
Si nos damos cuenta del pensamiento mágico en nosotros mismos, es con una espiral de vergüenza: se supone que los adultos alfabetizados son más sofisticados que eso.
Al mismo tiempo, el pensamiento mágico es obviamente desenfrenado en el mundo.
Es difícil no estar fascinado, aunque sea una fascinación horrorizada.
Creo que la razón por la que ha sido tan difícil definir con precisión el “pensamiento mágico” es que lo que llamamos "pensamiento mágico" es una colección de ejemplos estigmatizados de una capacidad cognitiva más general, y en general útil.
Esta es la capacidad de pensar en modo “como si”: “como si” los objetos inanimados tuvieran mente, “como si” los pensamientos pudieran afectar a la realidad, “como si” los símbolos tuvieran poder sobre sus referentes."""

#corpus = zip(english, spanish)
en, vocab_en = process(english)
es, vocab_es = process(spanish)


@testset "RockPaperScissorsEnv" begin

    rng = StableRNG(123)
    env = RockPaperScissorsEnv()

    RLBase.test_interfaces!(env)
    RLBase.test_runnable!(env)

    rewards = [[], []]
    for _ in 1:50_000
        while !is_terminated(env)
            env(rand(rng, legal_action_space(env)))
        end
        push!(rewards[1], reward(env, 1))
        push!(rewards[2], reward(env, 2))
        reset!(env)
    end

    @test all(rewards[1] .+ rewards[2] .== 0)

end


function get_branch_to_pm(ix::Int64, branch::PowerSystems.Transformer2W)
    PM_branch = Dict{String,Any}(
        "br_r"        => branch.r,
        "rate_a"      => branch.rate,
        "shift"       => 0.0,
        "rate_b"      => branch.rate,
        "br_x"        => branch.r,
        "rate_c"      => branch.rate,
        "g_to"        => 0.0,
        "g_fr"        => 0.0,
        "b_fr"        => branch.primaryshunt/2,
        "f_bus"       => branch.connectionpoints.from.number,
        "br_status"   => Float64(branch.available),
        "t_bus"       => branch.connectionpoints.to.number,
        "b_to"        => branch.primaryshunt/2,
        "index"       => ix,
        "angmin"      => -1.50,
        "angmax"      =>  1.50,
        "transformer" => true,
        "tap"         => 1.0,
    )
    return PM_branch
end

function get_branch_to_pm(ix::Int64, branch::PowerSystems.TapTransformer)
    PM_branch = Dict{String,Any}(
        "br_r"        => branch.r,
        "rate_a"      => branch.rate,
        "shift"       => 0.0,
        "rate_b"      => branch.rate,
        "br_x"        => branch.r,
        "rate_c"      => branch.rate,
        "g_to"        => 0.0,
        "g_fr"        => 0.0,
        "b_fr"        => branch.primaryshunt/2,
        "f_bus"       => branch.connectionpoints.from.number,
        "br_status"   => Float64(branch.available),
        "t_bus"       => branch.connectionpoints.to.number,
        "b_to"        => branch.primaryshunt/2,
        "index"       => ix,
        "angmin"      => -1.50,
        "angmax"      =>  1.50,
        "transformer" => true,
        "tap"         => branch.tap
    )
    return PM_branch
end

function get_branch_to_pm(ix::Int64, branch::PowerSystems.Line)
    PM_branch = Dict{String,Any}(
        "br_r"        => branch.r,
        "rate_a"      => branch.rate.from_to,
        "shift"       => 0.0,
        "rate_b"      => branch.rate.from_to,
        "br_x"        => branch.r,
        "rate_c"      => branch.rate.from_to,
        "g_to"        => 0.0,
        "g_fr"        => 0.0,
        "b_fr"        => branch.b.from,
        "f_bus"       => branch.connectionpoints.from.number,
        "br_status"   => Float64(branch.available),
        "t_bus"       => branch.connectionpoints.to.number,
        "b_to"        => branch.b.to,
        "index"       => ix,
        "angmin"      => branch.anglelimits.min,
        "angmax"      => branch.anglelimits.max,
        "transformer" => false,
        "tap"         => 1.0,
    )
    return PM_branch
end

function get_branches_to_pm(branches::Array{T}) where {T <: PowerSystems.Branch}
        PM_branches = Dict{String,Any}()

        for (ix, branch) in enumerate(branches)
            PM_branches["$(ix)"] = get_branch_to_pm(ix, branch)
        end

    return PM_branches
end

function get_buses_to_pm(buses::Array{PowerSystems.Bus})
    PM_buses = Dict{String,Any}()
    for bus in buses
        PM_bus = Dict{String,Any}(
        "zone"     => 1,
        "bus_i"    => bus.number,
        "bus_type" => 2,
        "vmax"     => bus.voltagelimits.max,
        "area"     => 1,
        "vmin"     => bus.voltagelimits.min,
        "index"    => bus.number,
        "va"       => bus.angle,
        "vm"       => bus.voltage,
        "base_kv"  => bus.basevoltage,
        "pni"      => 0.0,
        "qni"      => 0.0,
        )
        PM_buses["$(bus.number)"] = PM_bus
    end
    PM_buses["$(buses[1].number)"]["bus_type"] = 3
    return PM_buses
end

#=
function expression_to_pm_active(PM_dict::Dict{String,Any}, netinjection::BalanceNamedTuple, sys::PowerSystems.PowerSystem)

    for bus in sys.buses, time in 1:sys.time_periods

        PM_dict["nw"]["$(time)"]["bus"]["$(bus.number)"]["pni"] = netinjection.var_active[bus.number,time]

    end

end

function expression_to_pm_reactive(PM_dict::Dict{String,Any}, netinjection::BalanceNamedTuple, sys::PowerSystems.PowerSystem)

    for bus in sys.buses, time in 1:sys.time_periods

        PM_dict["nw"]["$(time)"]["bus"]["$(bus.number)"]["qni"] = netinjection.var_reactive[bus.number,time]

    end

end
=#

function pass_to_pm(sys::PowerSystems.PowerSystem, netinjection::BalanceNamedTuple)

    PM_translation = Dict{String,Any}(
    "bus" => get_buses_to_pm(sys.buses),
    "branch" => get_branches_to_pm(sys.branches),
    "baseMVA" => sys.basepower,
    "per_unit" => true,
    "dcline"         => Dict{String,Any}(),
    "gen"            => Dict{String,Any}(),
    "shunt"          => Dict{String,Any}(),
    "load"           => Dict{String,Any}(),
    )

    # TODO: this function adds overhead in large number of time_steps
    # We can do better later.

    PM_translation = IM.replicate(PM_translation,sys.time_periods)

    return PM_translation

end


abstract type AbstractImmutableSimplex{D, T} <: AbstractSimplex{D, T} end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    GenericDensity{F<:Function} <: AbstractDensity

Constructors:

    GenericDensity(f)

Turns the density function `f` into a BAT-compatible [`AbstractDensity`](@ref).
The return type of `f(v)` must supported by [`logvalof`](@ref).

It must be safe to execute `f` in parallel on multiple threads and
processes.
"""
struct GenericDensity{F<:Function} <: AbstractDensity
    f::F
end

Base.convert(::Type{GenericDensity}, f::Function) = GenericDensity(f)
Base.convert(::Type{AbstractDensity}, f::Function) = GenericDensity(f)

Base.parent(density::GenericDensity) = density.f

function eval_logval_unchecked(density::GenericDensity, v::Any)
    logvalof(density.f(v))
end

"""
    BAT.LogFuncDensity{F<:Function} <: AbstractDensity

*BAT-internal, not part of stable public API.*

Constructors:

    LogFuncDensity(logf::Function)

A density defined by a function that computes it's logarithmic value at given
points.

It must be safe to execute `f` in parallel on multiple threads and
processes.
"""
struct LogFuncDensity{F<:Function} <: AbstractDensity
    logf::F
end

Base.convert(::Type{LogFuncDensity}, nt::NamedTuple{(:logdensity,)}) = LogFuncDensity(nt.logdensity)
Base.convert(::Type{AbstractDensity}, nt::NamedTuple{(:logdensity,)}) = convert(LogFuncDensity, nt)

function eval_logval_unchecked(density::LogFuncDensity, v::Any)
    density.logf(v)
end

"""
    BAT.LogFuncDensityWithGrad{F<:Function} <: AbstractDensity

*BAT-internal, not part of stable public API.*

Constructors:

    LogFuncDensityWithGrad(logf::Function, valgradlogf::Function)

A density defined by a function that computes it's logarithmic value at given
points, as well as a function that computes both the value and the gradient.

It must be safe to execute both functions in parallel on multiple threads and
processes.
"""
struct LogFuncDensityWithGrad{F<:Function,G<:Function} <: AbstractDensity
    logf::F
    valgradlogf::G
end

function eval_logval_unchecked(density::LogFuncDensityWithGrad, v::Any)
    density.logf(v)
end

function ChainRulesCore.rrule(::typeof(eval_logval_unchecked), density::LogFuncDensityWithGrad, v)
    value, gradient = density.valgradlogf(v)
    @assert value isa Real
    function lfdwg_pullback(thunked_ΔΩ)
        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
        @assert ΔΩ isa Real
        tangent = gradient * ΔΩ
        (NoTangent(), ZeroTangent(), tangent)
    end
    return value, lfdwg_pullback
end

vjp_algorithm(density::LogFuncDensityWithGrad) = ZygoteAD()


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

VERSION >= v"0.4.0" && __precompile__()

module Spectra

using StatsBase
using PyPlot
using LsqFit
using PyCall
using Dierckx
using Ipopt
using JuMP

# For PyCall modules
const preprocessing = PyNULL()
const grid_search = PyNULL()
const cross_validation = PyNULL()
const decomposition = PyNULL()
const kernel_ridge = PyNULL()
const svm = PyNULL()
const gaussian_process = PyNULL()
const linear_model = PyNULL()

# some initial setup for calling the GCVSPL.f library
unixpath = "../deps/src/gcvspline/libgcvspl"
winpath = "../deps/bin$WORD_SIZE/libgcvspl" # let it there as an example but I did not tried yet any build on Windows... TODO
const gcvspl = joinpath(dirname(@__FILE__), @unix? unixpath : winpath)

function __init__()
    # Ensure library is available.
    if (Libdl.dlopen_e(gcvspl) == C_NULL)
        error("GCVSPL not properly installed. Run Pkg.build(\"Spectra\"). Windows auto-build is not setup, you might want to build the library manually.")
    end

	copy!(preprocessing, pyimport_conda("sklearn.preprocessing", "scikit-learn"))
	copy!(grid_search, pyimport_conda("sklearn.grid_search", "scikit-learn"))
	copy!(cross_validation, pyimport_conda("sklearn.cross_validation", "scikit-learn"))
	copy!(decomposition, pyimport_conda("sklearn.decomposition", "scikit-learn"))
	copy!(kernel_ridge, pyimport_conda("sklearn.kernel_ridge", "scikit-learn"))
	copy!(svm, pyimport_conda("sklearn.svm","scikit-learn"))
	copy!(gaussian_process, pyimport_conda("sklearn.gaussian_process", "scikit-learn"))
	copy!(linear_model, pyimport_conda("sklearn.linear_model", "scikit-learn"))

end

include("diffusion.jl")
include("integrale.jl")
include("functions.jl")
include("baseline.jl")
include("bootstrap.jl")
include("tlcorrection.jl")
include("rameau.jl")
include("deprecated.jl")
include("ml_regressor.jl")
include("ctxremoval.jl")

#From integrale.jl
export trapz, gaussianarea

#From diffusion.jl
export peak_diffusion, model, IRdataprep

#From functions.jl
export poly, polyfit, gaussiennes, lorentziennes, pseudovoigts, pearson7, normal_dist, xshift_inversion, xshift_direct,xshift_correction

#From baseline.jl
export baseline

#From gcvspl.jl
export gcvspl, splderivative

#From bootstrap
export bootsample, bootperf

#From tlcorrection
export tlcorrection

#From rameau
export rameau

#From ml_regressor
export mlregressor

#From ctxremoval
export ctxremoval

end # module


# This script has a heavy dependency on PyCall and PyTorch,
# so we should re-distribute the model and make this unnecessary
using PyCall
@pyimport torchvision.models as models

const model_dict = Dict([
        101=>models.resnet101,
        50=>models.resnet50,
        152=>models.resnet152,
        18=>models.resnet18,
        34=>models.resnet34
])

function load_resnet!(model, depth=101)
    conv_mode!(model, 1)
    stride3!()
    ptdict = state_dict(models, depth)
    #println(keys(ptdict))
    parse_conv!(model.conv1, ptdict, "conv1")
    parse_bn!(model.bn1, ptdict, "bn1")
    # Load the layers
    for i = 1:length(model.layers)
        seq = model.layers[i]
        for j = 1:length(seq.layers)
            bneck = seq.layers[j]
            key = string("layer",i, ".", j-1)
            for k = 1:(isa(bneck, Bottleneck) ? 3 : 2) #convs and bnorms
                ckey = string(key, ".conv",k)
                bkey = string(key, ".bn", k)
                parse_conv!(
                    getfield(bneck, Symbol(string("conv", k))),
                    ptdict, ckey)
                parse_bn!(
                    getfield(bneck, Symbol(string("bn", k))),
                    ptdict, bkey)
            end
            if bneck.downsample !== nothing
                conv = bneck.downsample.layers[1]
                bn = bneck.downsample.layers[2]
                parse_conv!(conv, ptdict, string(
                    key, ".", "downsample.0"
                ))
                parse_bn!(bn, ptdict, string(
                    key, ".", "downsample.1"
                ))
            end
        end
    end
    if model.output !== nothing
        parse_linear!(model.output, ptdict, "fc")
    end
end

state_dict(models, depth) =
    model_dict[depth](pretrained=true)[:state_dict]()

function parse_linear!(lin, ptdict, key)
    setval!(lin.w, reshape(
        ptdict[string(key, ".weight")][:numpy](),
        size(aval(lin.w))))

    setval!(lin.b, reshape(
        ptdict[string(key, ".bias")][:numpy](),
        size(aval(lin.b))))
end

function parse_conv!(conv, ptdict, key)
    setval!(conv.w, permutedims(
        ptdict[string(key, ".weight")][:numpy](), (4, 3, 2, 1)))
end

function parse_bn!(bn, ptdict, key)
    setval!(bn.w, vcat(ptdict[string(key, ".weight")][:numpy]()[:],
                       ptdict[string(key, ".bias")][:numpy]()[:]))
    sz = (1, 1, div(length(aval(bn.w)), 2), 1)
    bn.moments.mean =
        reshape(ptdict[string(key, ".running_mean")][:numpy](), sz)
    bn.moments.var =
        reshape(ptdict[string(key, ".running_var")][:numpy](), sz)
    if bn.remove_initfns
        bn.moments.meaninit = nothing
        bn.moments.varinit  = nothing
    end
end


decompose_trace(trace::Optim.OptimizationTrace) = last(trace)

function __map_optimizer_args(prob::OptimizationProblem, opt::Union{Optim.AbstractOptimizer, Optim.Fminbox,Optim.SAMIN, Optim.ConstrainedOptimizer};
    cb=nothing,
    maxiters::Union{Number, Nothing}=nothing,
    maxtime::Union{Number, Nothing}=nothing,
    abstol::Union{Number, Nothing}=nothing,
    reltol::Union{Number, Nothing}=nothing,
    kwargs...)
    if !isnothing(abstol)
        @warn "common abstol is currently not used by $(opt)"
    end

    mapped_args = (;extended_trace=true, kwargs...)

    if !isnothing(cb)
        mapped_args = (; mapped_args...,  callback = cb)
    end

    if !isnothing(maxiters)
        mapped_args = (; mapped_args..., iterations=maxiters)
    end

    if !isnothing(maxtime)
        mapped_args = (; mapped_args..., time_limit=maxtime)
    end

    if !isnothing(reltol)
        mapped_args = (; mapped_args..., f_tol=reltol)
    end

    return Optim.Options(;mapped_args...)
end

function __solve(prob::OptimizationProblem, opt::Optim.AbstractOptimizer, data = DEFAULT_DATA;
    kwargs...)
    if !isnothing(prob.lb) | !isnothing(prob.ub)
        if !(opt isa Union{Optim.Fminbox, Optim.SAMIN, Optim.AbstractConstrainedOptimizer})
            if opt isa Optim.ParticleSwarm
                opt = Optim.ParticleSwarm(;lower=prob.lb, upper=prob.ub, n_particles=opt.n_particles)
            elseif opt isa Optim.SimulatedAnnealing
                @warn "$(opt) can currently not be wrapped in Fminbox(). The lower and upper bounds thus will be ignored. Consider using a different optimizer or open an issue with Optim.jl"
            else
                opt = Optim.Fminbox(opt)
            end
        end
    end

    return ___solve(prob, opt, data; kwargs...)
end

function ___solve(prob::OptimizationProblem, opt::Optim.AbstractOptimizer,
                 data = DEFAULT_DATA;
                 cb = (args...) -> (false),
                 maxiters::Union{Number, Nothing} = nothing,
                 maxtime::Union{Number, Nothing} = nothing,
                 abstol::Union{Number, Nothing}=nothing,
                 reltol::Union{Number, Nothing}=nothing,
                 progress = false,
                 kwargs...)
    local x, cur, state

    if data != DEFAULT_DATA
        maxiters = length(data)
    end

    cur, state = iterate(data)

    function _cb(trace)
        cb_call = opt == Optim.NelderMead() ? cb(decompose_trace(trace).metadata["centroid"],x...) : cb(decompose_trace(trace).metadata["x"],x...)
        if !(typeof(cb_call) <: Bool)
            error("The callback should return a boolean `halt` for whether to stop the optimization process.")
        end
        cur, state = iterate(data, state)
        cb_call
    end

    maxiters = _check_and_convert_maxiters(maxiters)
    maxtime = _check_and_convert_maxtime(maxtime)

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p)

    !(opt isa Optim.ZerothOrderOptimizer) && f.grad === nothing && error("Use OptimizationFunction to pass the derivatives or automatically generate them with one of the autodiff backends")

    _loss = function(θ)
        x = f.f(θ, prob.p, cur...)
        __x = first(x)
        return prob.sense === MaxSense ? -__x : __x
    end

    fg! = function (G,θ)
        if G !== nothing
            f.grad(G, θ, cur...)
            if prob.sense === MaxSense
                G .*= false
            end
        end
        return _loss(θ)
    end

    if opt isa Optim.KrylovTrustRegion
        hv = function (H,θ,v)
            f.hv(H,θ,v,cur...)
            if prob.sense === MaxSense
                H .*= false
            end
        end
        optim_f = Optim.TwiceDifferentiableHV(_loss, fg!, hv, prob.u0)
    else
        gg = function (G, θ)
            f.grad(G, θ, cur...)
            if prob.sense === MaxSense
                G .*= false
            end
        end

        hh = function (H,θ)
            f.hess(H,θ,cur...)
            if prob.sense === MaxSense
                H .*= false
            end
        end
        optim_f = Optim.TwiceDifferentiable(_loss, gg, fg!, hh, prob.u0)
    end

    opt_args = _map_optimizer_args(prob, opt, cb=_cb, maxiters=maxiters, maxtime=maxtime,abstol=abstol, reltol=reltol; kwargs...)

    t0 = time()
    opt_res = Optim.optimize(optim_f, prob.u0, opt, opt_args)
    t1 = time()
    opt_ret = Symbol(Optim.converged(opt_res))

    SciMLBase.build_solution(prob, opt, opt_res.minimizer, prob.sense === MaxSense ? -opt_res.minimum : opt_res.minimum; original=opt_res, retcode=opt_ret)
end

function ___solve(prob::OptimizationProblem, opt::Union{Optim.Fminbox,Optim.SAMIN},
                 data = DEFAULT_DATA;
                 cb = (args...) -> (false),
                 maxiters::Union{Number, Nothing} = nothing,
                 maxtime::Union{Number, Nothing} = nothing,
                 abstol::Union{Number, Nothing}=nothing,
                 reltol::Union{Number, Nothing}=nothing,
                 progress = false,
                 kwargs...)

    local x, cur, state

    if data != DEFAULT_DATA
        maxiters = length(data)
    end

    cur, state = iterate(data)

    function _cb(trace)
        cb_call = !(opt isa Optim.SAMIN) && opt.method == Optim.NelderMead() ? cb(decompose_trace(trace).metadata["centroid"],x...) : cb(decompose_trace(trace).metadata["x"],x...)
        if !(typeof(cb_call) <: Bool)
            error("The callback should return a boolean `halt` for whether to stop the optimization process.")
        end
        cur, state = iterate(data, state)
        cb_call
    end

    maxiters = _check_and_convert_maxiters(maxiters)
    maxtime = _check_and_convert_maxtime(maxtime)

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p)

    !(opt isa Optim.ZerothOrderOptimizer) && f.grad === nothing && error("Use OptimizationFunction to pass the derivatives or automatically generate them with one of the autodiff backends")

    _loss = function(θ)
        x = f.f(θ, prob.p, cur...)
        __x = first(x)
        return prob.sense === MaxSense ? -__x : __x
    end
    fg! = function (G,θ)
        if G !== nothing
            f.grad(G, θ, cur...)
            if prob.sense === MaxSense
                G .*= false
            end
        end
        return _loss(θ)
    end

    gg = function (G, θ)
        f.grad(G, θ, cur...)
        if prob.sense === MaxSense
            G .*= false
        end
    end
    optim_f = Optim.OnceDifferentiable(_loss, gg, fg!, prob.u0)

    opt_args = _map_optimizer_args(prob, opt, cb=_cb, maxiters=maxiters, maxtime=maxtime,abstol=abstol, reltol=reltol; kwargs...)

    t0 = time()
    opt_res = Optim.optimize(optim_f, prob.lb, prob.ub, prob.u0, opt, opt_args)
    t1 = time()
    opt_ret = Symbol(Optim.converged(opt_res))

    SciMLBase.build_solution(prob, opt, opt_res.minimizer, opt_res.minimum; original=opt_res, retcode=opt_ret)
end

function ___solve(prob::OptimizationProblem, opt::Optim.ConstrainedOptimizer,
                 data = DEFAULT_DATA;
                 cb = (args...) -> (false),
                 maxiters::Union{Number, Nothing} = nothing,
                 maxtime::Union{Number, Nothing} = nothing,
                 abstol::Union{Number, Nothing}=nothing,
                 reltol::Union{Number, Nothing}=nothing,
                 progress = false,
                 kwargs...)

    local x, cur, state

    if data != DEFAULT_DATA
        maxiters = length(data)
    end

    cur, state = iterate(data)

      function _cb(trace)
      cb_call = cb(decompose_trace(trace).metadata["x"],x...)
      if !(typeof(cb_call) <: Bool)
          error("The callback should return a boolean `halt` for whether to stop the optimization process.")
      end
      cur, state = iterate(data, state)
      cb_call
    end

    maxiters = _check_and_convert_maxiters(maxiters)
    maxtime = _check_and_convert_maxtime(maxtime)

    f = instantiate_function(prob.f,prob.u0,prob.f.adtype,prob.p,prob.ucons === nothing ? 0 : length(prob.ucons))

    f.cons_j ===nothing && error("This optimizer requires derivative definitions for nonlinear constraints. If the problem does not have nonlinear constraints, choose a different optimizer. Otherwise define the derivative for cons using OptimizationFunction either directly or automatically generate them with one of the autodiff backends")

    _loss = function(θ)
        x = f.f(θ, prob.p, cur...)
        __x = first(x)
        return prob.sense === MaxSense ? -__x : __x
    end
    fg! = function (G,θ)
        if G !== nothing
            f.grad(G, θ, cur...)
            if prob.sense === MaxSense
                G .*= false
            end
        end
        return _loss(θ)
    end
    gg = function (G, θ)
        f.grad(G, θ, cur...)
        if prob.sense === MaxSense
            G .*= false
        end
    end

    hh = function (H,θ)
        f.hess(H,θ,cur...)
        if prob.sense === MaxSense
            H .*= false
        end
    end
    optim_f = Optim.TwiceDifferentiable(_loss, gg, fg!, hh, prob.u0)

    cons! = (res, θ) -> res .= f.cons(θ);

    cons_j! = function(J, x)
        f.cons_j(J, x)
    end

    cons_hl! = function (h, θ, λ)
        res = [similar(h) for i in 1:length(λ)]
        f.cons_h(res, θ)
        for i in 1:length(λ)
            h .+= λ[i]*res[i]
        end
    end

    lb = prob.lb === nothing ? [] : prob.lb
    ub = prob.ub === nothing ? [] : prob.ub
    optim_fc = Optim.TwiceDifferentiableConstraints(cons!, cons_j!, cons_hl!, lb, ub, prob.lcons, prob.ucons)

    opt_args = _map_optimizer_args(prob, opt, cb=_cb, maxiters=maxiters, maxtime=maxtime,abstol=abstol, reltol=reltol; kwargs...)

    t0 = time()
    opt_res = Optim.optimize(optim_f, optim_fc, prob.u0, opt, opt_args)
    t1 = time()
    opt_ret = Symbol(Optim.converged(opt_res))

    SciMLBase.build_solution(prob, opt, opt_res.minimizer, opt_res.minimum; original=opt_res, retcode=opt_ret)
end


# ======================
# Control postprocessing
# ======================

abstract type ControlPostprocessing end

struct NoPostprocessing <: ControlPostprocessing
end

apply(::NoPostprocessing, x) = x

struct UniformAdditivePostprocessing{N<:Number} <: ControlPostprocessing
    shift::N

    function UniformAdditivePostprocessing(shift::N) where {N<:Number}
        if shift == zero(N)
            @warn("zero additive control postprocessing is discouraged; " *
                  "use `NoPostprocessing` instead")
        end
        return new{N}(shift)
    end
end

function apply(postprocessing::UniformAdditivePostprocessing, x)
    return x .+ postprocessing.shift
end

function apply(postprocessing::UniformAdditivePostprocessing, X::LazySet)
    return translate(X, fill(postprocessing.shift, dim(X)))
end

function apply(postprocessing::UniformAdditivePostprocessing, X::UnionSetArray)
    return UnionSetArray([apply(postprocessing, Xi) for Xi in array(X)])
end

# =====================
# Control preprocessing
# =====================

abstract type ControlPreprocessing end

struct NoPreprocessing <: ControlPreprocessing
end

apply(::NoPreprocessing, x) = x

struct FunctionPreprocessing{F<:Function} <: ControlPreprocessing
    f::F

    function FunctionPreprocessing(f::F) where {F<:Function}
        return new{F}(f)
    end
end

function apply(funct::FunctionPreprocessing, x)
    return funct.f(x)
end

# ================
# Control problems
# ================

abstract type AbstractControlProblem end

"""
    ControlledPlant{ST, CT, XT, DT, PT, CPRT, CPST} <: AbstractControlProblem

Struct representing a closed-loop controlled system.

### Fields

- `ivp`            -- initial-value problem
- `controller`     -- controller
- `vars`           -- dictionary storing state variables, input variables and
                      control variables
- `period`         -- control period
- `postprocessing` -- postprocessing of the controller output
- `preprocessing`  -- preprocessing of the controller input

### Parameters

- `ST`:  type of system
- `CT`:  type of controller
- `XT`:  type of initial condition
- `DT`:  type of variables
- `PT`:  type of period
- `CPRT`: type of control preprocessing
- `CPST`: type of control postprocessing

### Notes

While typically the `controller` is a neural network, this struct does not
prescribe the type.
"""
struct ControlledPlant{ST, CT, XT, DT, PT, CPRT, CPST} <: AbstractControlProblem
    ivp::InitialValueProblem{ST, XT}
    controller::CT
    vars::Dict{Symbol, DT}
    period::PT
    preprocessing::CPRT
    postprocessing::CPST

    function ControlledPlant(ivp::InitialValueProblem{ST, XT},
                             controller::CT,
                             vars::Dict{Symbol, DT},
                             period::PT;
                             preprocessing::CPRT=NoPreprocessing(),
                             postprocessing::CPST=NoPostprocessing()) where {ST, CT, XT, DT, PT, CPRT, CPST}
        return new{ST, CT, XT, DT, PT, CPRT, CPST}(ivp, controller, vars, period, preprocessing, postprocessing)
    end
end

plant(cp::ControlledPlant) = cp.ivp
MathematicalSystems.initial_state(cp::ControlledPlant) = initial_state(cp.ivp)
system(cp::ControlledPlant) = cp.ivp.s
controller(cp::ControlledPlant) = cp.controller
period(cp::ControlledPlant) = cp.period
control_postprocessing(cp::ControlledPlant) = cp.postprocessing
control_preprocessing(cp::ControlledPlant) = cp.preprocessing

state_vars(cp::ControlledPlant) = get(cp.vars, :state_vars, Int[])
input_vars(cp::ControlledPlant) = get(cp.vars, :input_vars, Int[])
control_vars(cp::ControlledPlant) = get(cp.vars, :control_vars, Int[])

# =============
# Specification
# =============

@with_kw struct Specification{NT, PT, ET}
    T::NT
    predicate::PT
    ext::ET = nothing
end

predicate(spec::Specification) = spec.predicate
time_horizon(spec::Specification) = spec.T
ext(spec::Specification) = spec.ext


@testset "Quick" begin
    qm = QuickMDP(statetype=Int, actiontype=Int)
    @test @inferred(statetype(qm)) == Int
    @test @inferred(actiontype(qm)) == Int
    @test @inferred(discount(qm)) == 1.0
    @test @inferred(isterminal(qm, 1)) == false
    @test_throws MissingQuickArgument transition(qm, 1, 1)
    @test_throws MissingQuickArgument initialstate(qm)
    @test_throws MissingQuickArgument reward(qm, 1, 1, 1)
    @test_throws MissingQuickArgument initialstate(qm)
    @test_throws MissingQuickArgument states(qm)
    @test_throws MissingQuickArgument actions(qm)
    @test_throws MissingQuickArgument stateindex(qm, 1)
    @test_throws MissingQuickArgument actionindex(qm, 1)

    qp = QuickPOMDP(obstype=Int)
    @test @inferred(obstype(qp)) == Int
    @test @inferred(statetype(qp)) == Any
    @test @inferred(actiontype(qp)) == Any
    @test @inferred(discount(qp)) == 1.0
    @test @inferred(isterminal(qp, 1)) == false
    @test_throws MissingQuickArgument transition(qp, 1, 1)
    @test_throws MissingQuickArgument observation(qp, 1, 1)
    @test_throws MissingQuickArgument initialstate(qp)
    @test_throws MissingQuickArgument reward(qp, 1, 1, 1)
    @test_throws MissingQuickArgument initialstate(qp)
    @test_throws MissingQuickArgument initialobs(qp, 1)
    @test_throws MissingQuickArgument states(qp)
    @test_throws MissingQuickArgument actions(qp)
    @test_throws MissingQuickArgument observations(qp)
    @test_throws MissingQuickArgument stateindex(qp, 1)
    @test_throws MissingQuickArgument actionindex(qp, 1)
    @test_throws MissingQuickArgument obsindex(qp, 1)

    @test_throws MissingQuickArgument @gen(:sp,:o,:r)(qp, 1, 1, MersenneTwister(2))

    qp = QuickMDP(states=[3,2,1], actions=[1,2,3])
    @test ordered_states(qp) == [3,2,1]
    @test ordered_actions(qp) == [1,2,3]
end

@testset "Mountan Car" begin
    mountaincar = begin
        include("../examples/mountaincar.jl")
    end

    @show statetype(mountaincar)

    energize = FunctionPolicy(
        function (s)
            if s[2] < 0.0
                return minimum(actions(mountaincar))
            else
                return maximum(actions(mountaincar))
            end
        end
    )
    sim = HistoryRecorder(max_steps=1000)
    # below should work with a new version of POMDPSimulators
    @test_skip simulate(sim, mountaincar, energize) != nothing
    hist = simulate(sim, mountaincar, energize, rand(initialstate(mountaincar)))
    @test maximum(hist[:r]) == 100.0
    @test all(isa.(hist[:s], statetype(mountaincar)))
    @test last(collect(hist[:sp]))[1] > 0.5
end

@testset "Mountan Car Visualization" begin
    m = begin
        include("../examples/mountaincar_with_visualization.jl")
    end

    @inferred @gen(:sp,:r)(m, (0.2, 0.0), 0.0, MersenneTwister(2))

    energize = FunctionPolicy(
        function (s)
            if s[2] < 0.0
                return minimum(actions(m))
            else
                return maximum(actions(m))
            end
        end
    )
    sim = HistoryRecorder(max_steps=1000)
    # below should work when POMDPSimulators is fixed, but it can be deleted then
    @test_skip simulate(sim, m, energize) != nothing
    hist = simulate(sim, m, energize, rand(initialstate(m)))
    @test maximum(hist[:r]) == 100.0
    @test all(isa.(hist[:s], statetype(mountaincar)))
    @test last(collect(hist[:sp]))[1] > 0.5
    draw(SVG("test_render.svg", 5cm, 4cm), render(m, (s=(0.2, 0.0),)))
    @test isfile("test_render.svg")
end

@testset "Simple Light-Dark" begin
    m = begin
        include("../examples/lightdark.jl")
    end
    p = RandomPolicy(m, rng = MersenneTwister(2))
    sim = HistoryRecorder(max_steps=1000)
    up = DiscreteUpdater(m)
    hist = simulate(sim, m, p, up)
    @test all(-100 .<= hist[:r] .<= 100)
    @test all(isa.(hist[:s], statetype(m)))
    @test last(collect(hist[:sp])) == 61
end

struct A end
@testset "preprocess" begin
    QuickPOMDPs.preprocess(::Type{A}) = Char
    qm = QuickMDP(statetype=A, actiontype=Int)
    @test statetype(qm) == Char
    qp = QuickPOMDP(statetype=A, actiontype=Int, obstype=Int)
    @test statetype(qp) == Char
end

@testset "state action reward" begin
    qm = QuickMDP(states=1:3, actions=1:3, reward=(s, a)->0.0)
    @test StateActionReward(qm) isa FunctionSAR
    qm = QuickMDP(states=1:3, actions=1:3, reward=(s, a, sp)->0.0)
    @test StateActionReward(qm) isa LazyCachedSAR

    qp = QuickPOMDP(states=1:3, actions=1:3, observations=1:3, reward=(s, a)->0.0)
    @test StateActionReward(qp) isa FunctionSAR
    qp = QuickPOMDP(states=1:3, actions=1:3, observations=1:3, reward=(s, a, sp, o)->0.0)
    @test StateActionReward(qp) isa LazyCachedSAR
end

@testset "variable args" begin
    m = QuickMDP(states=1:3, actions=1:3, reward=(s, a)->0.0)
    @test @inferred(reward(m, 1, 2)) == 0.0
    @test @inferred(reward(m, 1, 2, 3)) == 0.0
    m = QuickPOMDP(states=1:3, actions=1:3, observations=1:3, reward=(s, a)->0.0)
    @test @inferred(reward(m, 1, 2)) == 0.0
    @test @inferred(reward(m, 1, 2, 3)) == 0.0
    @test @inferred(reward(m, 1, 2, 3, 1)) == 0.0
    m = QuickPOMDP(states=1:3, actions=1:3, observations=1:3, observation=(a, sp)->Deterministic(sp))
    @test rand(@inferred(observation(m, 2, 3))) == 3
    @test rand(@inferred(observation(m, 1, 2, 3))) == 3
end

@testset "#23" begin
    @test_throws ArgumentError QuickPOMDP(initialstate=:test)
end


#!/usr/bin/env julia
using CSV, DataFrames, PGFPlotsX
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage{fontspec}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\setmainfont{Asana-Math}")
seflu1617 = CSV.File("data/Rpt_LabSurveillanceDataAnyWeekByCtry_Chart_se_201640-201720.tsv") |> DataFrame
seflu21 = CSV.File("data/Rpt_LabSurveillanceDataAnyWeekByCtry_Chart_se_202120-202149.tsv") |> DataFrame
seflu1617[!,:AH3_NS] = seflu1617[!,:AH3] + seflu1617[!,:ANOTSUBTYPED]
seflu21[!,:AH3_NS] = seflu21[!,:AH3] + seflu21[!,:ANOTSUBTYPED]
p = @pgf Axis({ymode="log",xlabel="vecka", ylabel="antal labbfall", xmajorgrids,
    title="Fall influensa A(H3N2) och A ej subtypad Sverige",
    xtick="{1,10,13,33}",xticklabels="{40,49,52,20}"})
@pgf push!(p, PlotInc({"mark=x", "color=blue"}, Table([1:33,seflu1617[!,:AH3_NS]])), LegendEntry("2016/17"))
@pgf push!(p, PlotInc({"mark=*", "color=red"}, Table([1:10,seflu21[21:30,:AH3_NS]])), LegendEntry("2021/22"))
inflpl = "se_flu_1617_4020_21_4049"
inflplpath = "../../images/$(inflpl).svg"
print(*("![Labbfall influensa A(H3N2) och A ej subtypad Sverige 2016/17 och 2021/22.]",
    "($(inflplpath)){#fig:$(inflpl) width=100%}\n\n"))
pgfsave(inflplpath, p)


export GaussianMeanVariance, GaussianMeanPrecision, GaussianWeighteMeanPrecision
export MvGaussianMeanCovariance, MvGaussianMeanPrecision, MvGaussianWeightedMeanPrecision
export UnivariateNormalDistributionsFamily, MultivariateNormalDistributionsFamily, NormalDistributionsFamily
export UnivariateGaussianDistributionsFamily, MultivariateGaussianDistributionsFamily, GaussianDistributionsFamily

const GaussianMeanVariance            = NormalMeanVariance
const GaussianMeanPrecision           = NormalMeanPrecision
const GaussianWeighteMeanPrecision    = NormalWeightedMeanPrecision
const MvGaussianMeanCovariance        = MvNormalMeanCovariance
const MvGaussianMeanPrecision         = MvNormalMeanPrecision
const MvGaussianWeightedMeanPrecision = MvNormalWeightedMeanPrecision

const UnivariateNormalDistributionsFamily{T}   = Union{NormalMeanPrecision{T}, NormalMeanVariance{T}, NormalWeightedMeanPrecision{T}}
const MultivariateNormalDistributionsFamily{T} = Union{MvNormalMeanPrecision{T}, MvNormalMeanCovariance{T}, MvNormalWeightedMeanPrecision{T}}
const NormalDistributionsFamily{T}             = Union{UnivariateNormalDistributionsFamily{T}, MultivariateNormalDistributionsFamily{T}}

const UnivariateGaussianDistributionsFamily   = UnivariateNormalDistributionsFamily
const MultivariateGaussianDistributionsFamily = MultivariateNormalDistributionsFamily
const GaussianDistributionsFamily             = NormalDistributionsFamily

import Base: prod, convert

# Variate forms promotion

promote_variate_type(::Type{ Univariate },   ::Type{ F }) where { F <: UnivariateNormalDistributionsFamily }   = F
promote_variate_type(::Type{ Multivariate }, ::Type{ F }) where { F <: MultivariateNormalDistributionsFamily } = F

promote_variate_type(::Type{ Univariate }, ::Type{ <: MvNormalMeanCovariance })        = NormalMeanVariance
promote_variate_type(::Type{ Univariate }, ::Type{ <: MvNormalMeanPrecision })         = NormalMeanPrecision
promote_variate_type(::Type{ Univariate }, ::Type{ <: MvNormalWeightedMeanPrecision }) = NormalWeightedMeanPrecision

promote_variate_type(::Type{ Multivariate }, ::Type{ <: NormalMeanVariance })          = MvNormalMeanCovariance
promote_variate_type(::Type{ Multivariate }, ::Type{ <: NormalMeanPrecision })         = MvNormalMeanPrecision
promote_variate_type(::Type{ Multivariate }, ::Type{ <: NormalWeightedMeanPrecision }) = MvNormalWeightedMeanPrecision

# Conversion to mean - variance parametrisation

function Base.convert(::Type{ NormalMeanVariance{T} }, dist::UnivariateNormalDistributionsFamily) where { T <: Real }
    return NormalMeanVariance(convert(T, mean(dist)), convert(T, var(dist)))
end

function Base.convert(::Type{ MvNormalMeanCovariance{T} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real }
    return convert(MvNormalMeanCovariance{T, AbstractArray{T, 1}}, dist)
end

function Base.convert(::Type{ MvNormalMeanCovariance{T, M} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real, M <: AbstractArray{T} }
    return convert(MvNormalMeanCovariance{T, AbstractArray{T, 1}, AbstractArray{T, 2}}, dist)
end

function Base.convert(::Type{ MvNormalMeanCovariance{T, M, P} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real, M <: AbstractArray{T}, P <: AbstractArray{T} }
    return MvNormalMeanCovariance(convert(M, mean(dist)), convert(P, cov(dist)))
end

function Base.convert(::Type{ NormalMeanVariance }, dist::UnivariateNormalDistributionsFamily{T}) where { T <: Real }
    return convert(NormalMeanVariance{T}, dist)
end

function Base.convert(::Type{ MvNormalMeanCovariance }, dist::MultivariateNormalDistributionsFamily{T}) where { T <: Real }
    return convert(MvNormalMeanCovariance{T}, dist)
end

# Conversion to mean - precision parametrisation

function Base.convert(::Type{ NormalMeanPrecision{T} }, dist::UnivariateNormalDistributionsFamily) where { T <: Real }
    return NormalMeanPrecision(convert(T, mean(dist)), convert(T, precision(dist)))
end

function Base.convert(::Type{ MvNormalMeanPrecision{T} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real }
    return convert(MvNormalMeanPrecision{T, AbstractArray{T, 1}}, dist)
end

function Base.convert(::Type{ MvNormalMeanPrecision{T, M} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real, M <: AbstractArray{T} }
    return convert(MvNormalMeanPrecision{T, AbstractArray{T, 1}, AbstractArray{T, 2}}, dist)
end

function Base.convert(::Type{ MvNormalMeanPrecision{T, M, P} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real, M <: AbstractArray{T}, P <: AbstractArray{T} }
    return MvNormalMeanPrecision(convert(M, mean(dist)), convert(P, precision(dist)))
end

function Base.convert(::Type{ NormalMeanPrecision }, dist::UnivariateNormalDistributionsFamily{T}) where { T <: Real }
    return convert(NormalMeanPrecision{T}, dist)
end

function Base.convert(::Type{ MvNormalMeanPrecision }, dist::MultivariateNormalDistributionsFamily{T}) where { T <: Real }
    return convert(MvNormalMeanPrecision{T}, dist)
end

# Conversion to weighted mean - precision parametrisation

function Base.convert(::Type{ NormalWeightedMeanPrecision{T} }, dist::UnivariateNormalDistributionsFamily) where { T <: Real }
    return NormalWeightedMeanPrecision(convert(T, weightedmean(dist)), convert(T, precision(dist)))
end

function Base.convert(::Type{ MvNormalWeightedMeanPrecision{T} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real }
    return convert(MvNormalWeightedMeanPrecision{T, AbstractArray{T, 1}}, dist)
end

function Base.convert(::Type{ MvNormalWeightedMeanPrecision{T, M} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real, M <: AbstractArray{T} }
    return convert(MvNormalWeightedMeanPrecision{T, AbstractArray{T, 1}, AbstractArray{T, 2}}, dist)
end

function Base.convert(::Type{ MvNormalWeightedMeanPrecision{T, M, P} }, dist::MultivariateNormalDistributionsFamily) where { T <: Real, M <: AbstractArray{T}, P <: AbstractArray{T} }
    return MvNormalWeightedMeanPrecision(convert(M, weightedmean(dist)), convert(P, precision(dist)))
end

function Base.convert(::Type{ NormalWeightedMeanPrecision }, dist::UnivariateNormalDistributionsFamily{T}) where { T <: Real }
    return convert(NormalWeightedMeanPrecision{T}, dist)
end

function Base.convert(::Type{ MvNormalWeightedMeanPrecision }, dist::MultivariateNormalDistributionsFamily{T}) where { T <: Real }
    return convert(MvNormalWeightedMeanPrecision{T}, dist)
end

# Basic prod fallbacks to weighted mean precision and converts first argument back

prod_analytical_rule(::Type{ <: UnivariateNormalDistributionsFamily }, ::Type{ <: UnivariateNormalDistributionsFamily }) = ProdAnalyticalRuleAvailable()

function Base.prod(::ProdAnalytical, left::L, right::R) where { L <: UnivariateNormalDistributionsFamily, R <: UnivariateNormalDistributionsFamily }
    wleft  = convert(NormalWeightedMeanPrecision, left)
    wright = convert(NormalWeightedMeanPrecision, right)
    return prod(ProdAnalytical(), wleft, wright)
end

prod_analytical_rule(::Type{ <: MultivariateNormalDistributionsFamily }, ::Type{ <: MultivariateNormalDistributionsFamily }) = ProdAnalyticalRuleAvailable()

function Base.prod(::ProdAnalytical, left::L, right::R) where { L <: MultivariateNormalDistributionsFamily, R <: MultivariateNormalDistributionsFamily }
    wleft  = convert(MvNormalWeightedMeanPrecision, left)
    wright = convert(MvNormalWeightedMeanPrecision, right)
    return prod(ProdAnalytical(), wleft, wright)
end


# This file is a part of Julia. License is MIT: https://julialang.org/license

# Support for @fastmath

# This module provides versions of math functions that may violate
# strict IEEE semantics.

# This allows the following transformations. For more information see
# http://llvm.org/docs/LangRef.html#fast-math-flags:
# nnan: No NaNs - Allow optimizations to assume the arguments and
#       result are not NaN. Such optimizations are required to retain
#       defined behavior over NaNs, but the value of the result is
#       undefined.
# ninf: No Infs - Allow optimizations to assume the arguments and
#       result are not +/-Inf. Such optimizations are required to
#       retain defined behavior over +/-Inf, but the value of the
#       result is undefined.
# nsz:  No Signed Zeros - Allow optimizations to treat the sign of a
#       zero argument or result as insignificant.
# arcp: Allow Reciprocal - Allow optimizations to use the reciprocal
#       of an argument rather than perform division.

module FastMath

export @fastmath

import Core.Intrinsics: sqrt_llvm, neg_float_fast,
    add_float_fast, sub_float_fast, mul_float_fast, div_float_fast, rem_float_fast,
    eq_float_fast, ne_float_fast, lt_float_fast, le_float_fast

const fast_op =
    Dict(# basic arithmetic
         :+ => :add_fast,
         :- => :sub_fast,
         :* => :mul_fast,
         :/ => :div_fast,
         :(==) => :eq_fast,
         :!= => :ne_fast,
         :< => :lt_fast,
         :<= => :le_fast,
         :abs => :abs_fast,
         :abs2 => :abs2_fast,
         :cmp => :cmp_fast,
         :conj => :conj_fast,
         :inv => :inv_fast,
         :rem => :rem_fast,
         :sign => :sign_fast,
         :isfinite => :isfinite_fast,
         :isinf => :isinf_fast,
         :isnan => :isnan_fast,
         :issubnormal => :issubnormal_fast,
         # math functions
         :^ => :pow_fast,
         :acos => :acos_fast,
         :acosh => :acosh_fast,
         :angle => :angle_fast,
         :asin => :asin_fast,
         :asinh => :asinh_fast,
         :atan => :atan_fast,
         :atan2 => :atan2_fast,
         :atanh => :atanh_fast,
         :cbrt => :cbrt_fast,
         :cis => :cis_fast,
         :cos => :cos_fast,
         :cosh => :cosh_fast,
         :exp10 => :exp10_fast,
         :exp2 => :exp2_fast,
         :exp => :exp_fast,
         :expm1 => :expm1_fast,
         :hypot => :hypot_fast,
         :lgamma => :lgamma_fast,
         :log10 => :log10_fast,
         :log1p => :log1p_fast,
         :log2 => :log2_fast,
         :log => :log_fast,
         :max => :max_fast,
         :min => :min_fast,
         :minmax => :minmax_fast,
         :sin => :sin_fast,
         :sincos => :sincos_fast,
         :sinh => :sinh_fast,
         :sqrt => :sqrt_fast,
         :tan => :tan_fast,
         :tanh => :tanh_fast)

const rewrite_op =
    Dict(:+= => :+,
         :-= => :-,
         :*= => :*,
         :/= => :/,
         :^= => :^)

function make_fastmath(expr::Expr)
    if expr.head === :quote
        return expr
    end
    op = get(rewrite_op, expr.head, :nothing)
    if op !== :nothing
        var = expr.args[1]
        rhs = expr.args[2]
        if isa(var, Symbol)
            # simple assignment
            expr = :($var = $op($var, $rhs))
        elseif isa(var, Expr) && var.head === :ref
            # array reference
            arr = var.args[1]
            inds = tuple(var.args[2:end]...)
            arrvar = gensym()
            indvars = tuple([gensym() for i in inds]...)
            expr = quote
                $(Expr(:(=), arrvar, arr))
                $(Expr(:(=), Expr(:tuple, indvars...), Expr(:tuple, inds...)))
                $arrvar[$(indvars...)] = $op($arrvar[$(indvars...)], $rhs)
            end
        end
    end
    Expr(make_fastmath(expr.head), map(make_fastmath, expr.args)...)
end
function make_fastmath(symb::Symbol)
    fast_symb = get(fast_op, symb, :nothing)
    if fast_symb === :nothing
        return symb
    end
    :(Base.FastMath.$fast_symb)
end
make_fastmath(expr) = expr

"""
    @fastmath expr

Execute a transformed version of the expression, which calls functions that
may violate strict IEEE semantics. This allows the fastest possible operation,
but results are undefined -- be careful when doing this, as it may change numerical
results.

This sets the [LLVM Fast-Math flags](http://llvm.org/docs/LangRef.html#fast-math-flags),
and corresponds to the `-ffast-math` option in clang. See [the notes on performance
annotations](@ref man-performance-annotations) for more details.

# Examples
```jldoctest
julia> @fastmath 1+2
3

julia> @fastmath(sin(3))
0.1411200080598672
```
"""
macro fastmath(expr)
    make_fastmath(esc(expr))
end

# Basic arithmetic

const FloatTypes = Union{Float32,Float64}

sub_fast(x::FloatTypes) = neg_float_fast(x)

add_fast(x::T, y::T) where {T<:FloatTypes} = add_float_fast(x, y)
sub_fast(x::T, y::T) where {T<:FloatTypes} = sub_float_fast(x, y)
mul_fast(x::T, y::T) where {T<:FloatTypes} = mul_float_fast(x, y)
div_fast(x::T, y::T) where {T<:FloatTypes} = div_float_fast(x, y)
rem_fast(x::T, y::T) where {T<:FloatTypes} = rem_float_fast(x, y)

add_fast(x::T, y::T, zs::T...) where {T<:FloatTypes} =
    add_fast(add_fast(x, y), zs...)
mul_fast(x::T, y::T, zs::T...) where {T<:FloatTypes} =
    mul_fast(mul_fast(x, y), zs...)

@fastmath begin
    cmp_fast(x::T, y::T) where {T<:FloatTypes} = ifelse(x==y, 0, ifelse(x<y, -1, +1))
end

eq_fast(x::T, y::T) where {T<:FloatTypes} = eq_float_fast(x, y)
ne_fast(x::T, y::T) where {T<:FloatTypes} = ne_float_fast(x, y)
lt_fast(x::T, y::T) where {T<:FloatTypes} = lt_float_fast(x, y)
le_fast(x::T, y::T) where {T<:FloatTypes} = le_float_fast(x, y)

isinf_fast(x) = false
isfinite_fast(x) = true
isnan_fast(x) = false
issubnormal_fast(x) = false

# complex numbers

ComplexTypes = Union{Complex64, Complex128}

@fastmath begin
    abs_fast(x::ComplexTypes) = hypot(real(x), imag(x))
    abs2_fast(x::ComplexTypes) = real(x)*real(x) + imag(x)*imag(x)
    conj_fast(x::T) where {T<:ComplexTypes} = T(real(x), -imag(x))
    inv_fast(x::ComplexTypes) = conj(x) / abs2(x)
    sign_fast(x::ComplexTypes) = x == 0 ? float(zero(x)) : x/abs(x)

    add_fast(x::T, y::T) where {T<:ComplexTypes} =
        T(real(x)+real(y), imag(x)+imag(y))
    add_fast(x::Complex{T}, b::T) where {T<:FloatTypes} =
        Complex{T}(real(x)+b, imag(x))
    add_fast(a::T, y::Complex{T}) where {T<:FloatTypes} =
        Complex{T}(a+real(y), imag(y))

    sub_fast(x::T, y::T) where {T<:ComplexTypes} =
        T(real(x)-real(y), imag(x)-imag(y))
    sub_fast(x::Complex{T}, b::T) where {T<:FloatTypes} =
        Complex{T}(real(x)-b, imag(x))
    sub_fast(a::T, y::Complex{T}) where {T<:FloatTypes} =
        Complex{T}(a-real(y), -imag(y))

    mul_fast(x::T, y::T) where {T<:ComplexTypes} =
        T(real(x)*real(y) - imag(x)*imag(y),
          real(x)*imag(y) + imag(x)*real(y))
    mul_fast(x::Complex{T}, b::T) where {T<:FloatTypes} =
        Complex{T}(real(x)*b, imag(x)*b)
    mul_fast(a::T, y::Complex{T}) where {T<:FloatTypes} =
        Complex{T}(a*real(y), a*imag(y))

    @inline div_fast(x::T, y::T) where {T<:ComplexTypes} =
        T(real(x)*real(y) + imag(x)*imag(y),
          imag(x)*real(y) - real(x)*imag(y)) / abs2(y)
    div_fast(x::Complex{T}, b::T) where {T<:FloatTypes} =
        Complex{T}(real(x)/b, imag(x)/b)
    div_fast(a::T, y::Complex{T}) where {T<:FloatTypes} =
        Complex{T}(a*real(y), -a*imag(y)) / abs2(y)

    eq_fast(x::T, y::T) where {T<:ComplexTypes} =
        (real(x)==real(y)) & (imag(x)==imag(y))
    eq_fast(x::Complex{T}, b::T) where {T<:FloatTypes} =
        (real(x)==b) & (imag(x)==T(0))
    eq_fast(a::T, y::Complex{T}) where {T<:FloatTypes} =
        (a==real(y)) & (T(0)==imag(y))

    ne_fast(x::T, y::T) where {T<:ComplexTypes} = !(x==y)
end

# fall-back implementations and type promotion

for op in (:abs, :abs2, :conj, :inv, :sign)
    op_fast = fast_op[op]
    @eval begin
        # fall-back implementation for non-numeric types
        $op_fast(xs...) = $op(xs...)
    end
end

for op in (:+, :-, :*, :/, :(==), :!=, :<, :<=, :cmp, :rem)
    op_fast = fast_op[op]
    @eval begin
        # fall-back implementation for non-numeric types
        $op_fast(xs...) = $op(xs...)
        # type promotion
        $op_fast(x::Number, y::Number, zs::Number...) =
            $op_fast(promote(x,y,zs...)...)
        # fall-back implementation that applies after promotion
        $op_fast(x::T,ys::T...) where {T<:Number} = $op(x,ys...)
    end
end

# Math functions

# builtins

pow_fast(x::Float32, y::Integer) = ccall("llvm.powi.f32", llvmcall, Float32, (Float32, Int32), x, y)
pow_fast(x::Float64, y::Integer) = ccall("llvm.powi.f64", llvmcall, Float64, (Float64, Int32), x, y)

sqrt_fast(x::FloatTypes) = sqrt_llvm(x)

# libm

const libm = Base.libm_name

for f in (:acos, :acosh, :asin, :asinh, :atan, :atanh, :cbrt, :cos,
          :cosh, :exp2, :expm1, :lgamma, :log10, :log1p, :log2,
          :log, :sin, :sinh, :tan, :tanh)
    f_fast = fast_op[f]
    @eval begin
        $f_fast(x::Float32) =
            ccall(($(string(f,"f")),libm), Float32, (Float32,), x)
        $f_fast(x::Float64) =
            ccall(($(string(f)),libm), Float64, (Float64,), x)
    end
end

pow_fast(x::Float32, y::Float32) =
    ccall(("powf",libm), Float32, (Float32,Float32), x, y)
pow_fast(x::Float64, y::Float64) =
    ccall(("pow",libm), Float64, (Float64,Float64), x, y)

atan2_fast(x::Float32, y::Float32) =
    ccall(("atan2f",libm), Float32, (Float32,Float32), x, y)
atan2_fast(x::Float64, y::Float64) =
    ccall(("atan2",libm), Float64, (Float64,Float64), x, y)

# explicit implementations

# FIXME: Change to `ccall((:sincos, libm))` when `Ref` calling convention can be
#        stack allocated.
@inline function sincos_fast(v::Float64)
    return Base.llvmcall("""
    %f = bitcast i8 *%1 to void (double, double *, double *)*
    %ps = alloca double
    %pc = alloca double
    call void %f(double %0, double *%ps, double *%pc)
    %s = load double, double* %ps
    %c = load double, double* %pc
    %res0 = insertvalue [2 x double] undef, double %s, 0
    %res = insertvalue [2 x double] %res0, double %c, 1
    ret [2 x double] %res
    """, Tuple{Float64,Float64}, Tuple{Float64,Ptr{Void}}, v, cglobal((:sincos, libm)))
end

@inline function sincos_fast(v::Float32)
    return Base.llvmcall("""
    %f = bitcast i8 *%1 to void (float, float *, float *)*
    %ps = alloca float
    %pc = alloca float
    call void %f(float %0, float *%ps, float *%pc)
    %s = load float, float* %ps
    %c = load float, float* %pc
    %res0 = insertvalue [2 x float] undef, float %s, 0
    %res = insertvalue [2 x float] %res0, float %c, 1
    ret [2 x float] %res
    """, Tuple{Float32,Float32}, Tuple{Float32,Ptr{Void}}, v, cglobal((:sincosf, libm)))
end

@inline function sincos_fast(v::Float16)
    s, c = sincos_fast(Float32(v))
    return Float16(s), Float16(c)
end

sincos_fast(v::AbstractFloat) = (sin_fast(v), cos_fast(v))
sincos_fast(v::Real) = sincos_fast(float(v)::AbstractFloat)
sincos_fast(v) = (sin_fast(v), cos_fast(v))

@fastmath begin
    hypot_fast(x::T, y::T) where {T<:FloatTypes} = sqrt(x*x + y*y)

    # Note: we use the same comparison for min, max, and minmax, so
    # that the compiler can convert between them
    max_fast(x::T, y::T) where {T<:FloatTypes} = ifelse(y > x, y, x)
    min_fast(x::T, y::T) where {T<:FloatTypes} = ifelse(y > x, x, y)
    minmax_fast(x::T, y::T) where {T<:FloatTypes} = ifelse(y > x, (x,y), (y,x))

    # complex numbers

    function cis_fast(x::T) where {T<:FloatTypes}
        s, c = sincos_fast(x)
        Complex{T}(c, s)
    end

    # See <http://en.cppreference.com/w/cpp/numeric/complex>
    pow_fast(x::T, y::T) where {T<:ComplexTypes} = exp(y*log(x))
    pow_fast(x::T, y::Complex{T}) where {T<:FloatTypes} = exp(y*log(x))
    pow_fast(x::Complex{T}, y::T) where {T<:FloatTypes} = exp(y*log(x))
    acos_fast(x::T) where {T<:ComplexTypes} =
        convert(T,π)/2 + im*log(im*x + sqrt(1-x*x))
    acosh_fast(x::ComplexTypes) = log(x + sqrt(x+1) * sqrt(x-1))
    angle_fast(x::ComplexTypes) = atan2(imag(x), real(x))
    asin_fast(x::ComplexTypes) = -im*asinh(im*x)
    asinh_fast(x::ComplexTypes) = log(x + sqrt(1+x*x))
    atan_fast(x::ComplexTypes) = -im*atanh(im*x)
    atanh_fast(x::T) where {T<:ComplexTypes} = convert(T,1)/2*(log(1+x) - log(1-x))
    cis_fast(x::ComplexTypes) = exp(-imag(x)) * cis(real(x))
    cos_fast(x::ComplexTypes) = cosh(im*x)
    cosh_fast(x::T) where {T<:ComplexTypes} = convert(T,1)/2*(exp(x) + exp(-x))
    exp10_fast(x::T) where {T<:ComplexTypes} =
        exp10(real(x)) * cis(imag(x)*log(convert(T,10)))
    exp2_fast(x::T) where {T<:ComplexTypes} =
        exp2(real(x)) * cis(imag(x)*log(convert(T,2)))
    exp_fast(x::ComplexTypes) = exp(real(x)) * cis(imag(x))
    expm1_fast(x::ComplexTypes) = exp(x)-1
    log10_fast(x::T) where {T<:ComplexTypes} = log(x) / log(convert(T,10))
    log1p_fast(x::ComplexTypes) = log(1+x)
    log2_fast(x::T) where {T<:ComplexTypes} = log(x) / log(convert(T,2))
    log_fast(x::T) where {T<:ComplexTypes} = T(log(abs2(x))/2, angle(x))
    sin_fast(x::ComplexTypes) = -im*sinh(im*x)
    sinh_fast(x::T) where {T<:ComplexTypes} = convert(T,1)/2*(exp(x) - exp(-x))
    sqrt_fast(x::ComplexTypes) = sqrt(abs(x)) * cis(angle(x)/2)
    tan_fast(x::ComplexTypes) = -im*tanh(im*x)
    tanh_fast(x::ComplexTypes) = (a=exp(x); b=exp(-x); (a-b)/(a+b))
end

# fall-back implementations and type promotion

for f in (:acos, :acosh, :angle, :asin, :asinh, :atan, :atanh, :cbrt,
          :cis, :cos, :cosh, :exp10, :exp2, :exp, :expm1, :lgamma,
          :log10, :log1p, :log2, :log, :sin, :sinh, :sqrt, :tan,
          :tanh)
    f_fast = fast_op[f]
    @eval begin
        $f_fast(x) = $f(x)
    end
end

for f in (:^, :atan2, :hypot, :max, :min, :minmax)
    f_fast = fast_op[f]
    @eval begin
        # fall-back implementation for non-numeric types
        $f_fast(x, y) = $f(x, y)
        # type promotion
        $f_fast(x::Number, y::Number) = $f_fast(promote(x, y)...)
        # fall-back implementation that applies after promotion
        $f_fast(x::T, y::T) where {T<:Number} = $f(x, y)
    end
end

end


#=
using PyCall

@pyimport pickle

f = open(raw"D:\projets\PlanHeat\tmp\instances\not_dag_2\flows.pickle", "r")
flows=pickle.load(f)
close(f)

include("local_optimizer.jl")

flows = Dict{Tuple, Float32}((0, 1, 0) => 1,
             (1, 2, 0) => 1,
             (2, 3, 0) => 1,
             (1, 4, 0) => 1,
             (5, 6, 0) => 1,
             (6, 4, 0) => 1,
             (6, 7, 0) => 1,
             (7, 8, 0) => 1,
             (8, 10, 0) => 1,
             (10, 9, 0) => 1,
             (9, 10, 0) => 1,
             (9, 10, 1) => 1,
             (10, 3, 0) => 1)

flow_dag = get_flow_dag(flows)

edges = Set{Tuple}()
for (u, ns) in flow_dag
    for (v, ks) in ns
        for k in ks
            e = (u, v, k)
            push!(edges, e)
        end
    end
end
for e in edges
    re = (e[2], e[1], e[3])
    if re in edges
        println(e)
    end
end
start_node= "27307035"

k = find_cycle(flow_dag)

i_flows = copy(flows)

clean_flow_two_cycles(flows)

break_cycles(flows)

flow_dag = get_flow_dag(flows)

get_nodes(flow_dag)
descendants(flow_dag, "S_0")
descendants(flow_dag, start_node)

costs = Dict(e => 1.0 for e in keys(flows))
G = get_adj_dict(keys(flows))
dijkstra_multisource(G, Set(["S_0"]), start_node, costs)

optimization_graph = get_adj_dict(keys(flows))
=#

#=   ======================================================================== =#
include("FCFP.jl")
import FCFP.FCFPInstance
include("DSSP.jl")
using DSSP
using PyCall
using JSON
using DataStructures

@pyimport pickle

p = raw"D:\projets\PlanHeat\tmp\master_plugin\debug_local_optimizer2/"
f = open(p * "flows.pickle", "r")
flows=pickle.load(f)
close(f)

f = open(p * "costs.pickle", "r")
costs=pickle.load(f)
close(f)

f = open(p * "heat_demand.json", "r")
head_demand=JSON.parse(read(f, String))
close(f)

f = open(p * "production.json", "r")
productions = JSON.parse(read(f, String))
close(f)

include("local_optimizer.jl")
import LocalOptimizer:optimize_locally, MultiDiGraph, get_nodes, get_flow_dag, has_no_parent, add_edge!, get_edges
import LocalOptimizer:descendants, get_sub_graph, break_cycles, clean_flow_two_cycles, validate_move, LocalOptimizerInstance, get_flow_arteries, remove_node!, cascade_union, add_zero_weight_edges_from_nodes, dijkstra_multisource, get_edge_like_path, reroute_artery, lowest_common_ancestor
import LocalOptimizer:check_solution, find_cycle, VertexType, KeyType, EdgeType, MultiEdgeType, AdjencyView
import LocalOptimizer:FlowDictType, replace_arteries, WeightType, out_edges, in_edges, has_one_parent

optimization_graph = MultiDiGraph(keys(flows))
V = get_nodes(optimization_graph)
old_buildings = Set()
capacities = Dict()
fcfp = FCFPInstance(0.0, V, costs, head_demand, productions, capacities, old_buildings)
logger = println

using StatProfilerHTML
Profile.clear()
@profile new_flows = optimize_locally(flows, fcfp, logger)
statprofilehtml()

flows2 = FlowDictType(flows)
loi = LocalOptimizerInstance(flows, fcfp, logger)
initial_cost = sum(loi.fcfp.fixed_costs[e] for (e, f) in loi.flows if f > 0)
check_solution(loi, loi.flows)

start = time()

loi = LocalOptimizerInstance(flows, fcfp, logger)
old_flows = copy(loi.flows)
initial_cost = sum(loi.fcfp.fixed_costs[e] for (e, f) in loi.flows if f > 0)
clean_flow_two_cycles(loi.working_flows)
break_cycles(loi, loi.working_flows)
validate_move(loi)
while replace_arteries(loi)
    a = check_solution(loi, loi.flows)
    if a
        old_flows = copy(loi.flows)
    else
        println("now")
        break
    end
    gain = initial_cost - sum(loi.fcfp.fixed_costs[e] for (e, f) in loi.flows if f > 0)
    loi.logger("Cost gain: $(gain)")
end

check_solution(loi, old_flows)
check_solution(loi, loi.flows)

f = open(p * "previous_correct_flows.pickle", "w")
pickle.dump(old_flows, f)
close(f)

################################################################################
################### Understanding while check is false #########################
################################################################################

flows = loi.flows
flow_dag = get_flow_dag(loi.optimization_graph, loi.flows)
c = find_cycle(flow_dag)
tol = 1e-5
nodes = get_nodes(flow_dag)
for n in nodes
    out_flow = sum(Array{WeightType}([flows[e] for e in out_edges(flow_dag, n)]))
    in_flow = sum(Array{WeightType}([flows[e] for e in in_edges(flow_dag, n)]))
    if n in loi.fcfp.building_nodes
        continue
    elseif n in loi.fcfp.production_nodes
        if out_flow > (loi.fcfp.production[n] + tol)
            println("here")
            break
        end
    else
        if abs(out_flow - in_flow) > tol
            println("or here")
            println(n)
            break
        end
    end
end

n= "intersect_S_2"
abs(sum([loi.flows[e] for e in out_edges(flow_dag, n)]) - sum([loi.flows[e] for e in in_edges(flow_dag, n)]))
v = "6341231519"
minimum(values(loi.fcfp.heat_demand))

ns = [u for u in get_nodes(flow_dag)]
v in ns
loi.flows[(v, n , 0)]

out_flow = sum(Array{WeightType}([flows[e] for e in out_edges(flow_dag, n)]))
in_flow = sum(Array{WeightType}([flows[e] for e in in_edges(flow_dag, n)]))

previous_flow = loi.flows
f = open(p * "flows_sources_move_pr.pickle", "w")
pickle.dump(previous_flow, f)
close(f)

################################################################################
################### Finding the artery that is involved ########################
################################################################################

loi.flows = copy(old_flows)
loi.working_flows = copy(old_flows)

loi.roots
[loi.fcfp.production[r] for r in loi.roots]

flow_dag = get_flow_dag(loi.optimization_graph, loi.flows)
arteries = get_flow_arteries(loi, flow_dag)
roots_descendants = Array{Set{VertexType}, 1}([descendants(flow_dag, r, false) for r in loi.roots])
roots_marginal_capacity = [loi.fcfp.production[r] - sum(get(loi.fcfp.heat_demand, n, 0.0)
                                                        for n in roots_descendants[r_index])
                                                                for (r_index, r) in enumerate(loi.roots)]
a = []
edge_like_path = []
while !isempty(arteries)
    a = dequeue!(arteries)
    # determine sub tree
    end_node = last(a)[2]
    sub_tree_nodes = descendants(flow_dag, end_node, false)
    if isempty(sub_tree_nodes)
       print("empty here")
    end
    sub_tree_demand = sum(get(loi.fcfp.heat_demand, n, 0.0) for n in sub_tree_nodes)
    # determine parent sub tree
    start_node = first(a)[1]
    current_roots_indexes = Set{Int64}()
    other_roots_indexes = Set{Int64}()
    for (d_index, d) in enumerate(roots_descendants)
       if (start_node in d)
           push!(current_roots_indexes, d_index)
       elseif (roots_marginal_capacity[d_index] >= sub_tree_demand)
           push!(other_roots_indexes, d_index)
       end
    end
    concerned_sub_trees_indexes = union(current_roots_indexes, other_roots_indexes)
    parent_dag_trees_nodes = cascade_union([roots_descendants[d_index] for d_index in concerned_sub_trees_indexes])
    parent_dag_trees_nodes = setdiff(parent_dag_trees_nodes, union(sub_tree_nodes,
                                                                  Set(e[2] for e in a)))
    if isempty(parent_dag_trees_nodes)
       print("empty here2")
    end
    # Add 0-weight edges in the sub-graph
    modified_costs = add_zero_weight_edges_from_nodes(sub_tree_nodes, loi.costs)
    path, cost = dijkstra_multisource(loi.optimization_graph,
                                     parent_dag_trees_nodes,
                                     end_node,
                                     modified_costs)
    edge_like_path = get_edge_like_path(path, loi.optimization_graph, loi.fcfp.fixed_costs)
    if edge_like_path != a
       print("hey")
       break
    end
end
edge_like_path
a
[loi.flows[e] for e in a]
[has_one_parent(flow_dag, e[2]) for e in a]
flow_dag = get_flow_dag(loi.optimization_graph, loi.working_flows)

potential_cycle_creation = false
if any((n != end_node) && (n in sub_tree_nodes) for n in path)
    potential_cycle_creation = true
end
new_start = first(path)
any(new_start in roots_descendants[d_index] for d_index in current_roots_indexes)
if any(new_start in roots_descendants[d_index] for d_index in current_roots_indexes)
    loi.logger("Moving flow inside current dag")
    reroute_artery(loi, a, edge_like_path, flow_dag, potential_cycle_creation)
else
    loi.logger("Moving flow to another dag")
    new_roots_candidates = Set{VertexType}([loi.roots[r] for r in other_roots_indexes if new_start in roots_descendants[r]])
    current_roots = Set{VertexType}([loi.roots[r] for r in current_roots_indexes])
    move_artery_to_other_dag(loi, a, edge_like_path, flow_dag, current_roots, new_roots_candidates, potential_cycle_creation)
end

###### debug : move_artery_to_other_dag ################################
new_a = edge_like_path
artery_flow = loi.working_flows[first(a)]
for e in a
    loi.working_flows[e] = zero(WeightType)
end
# Find path from old root to old start and remove flow
old_start = first(a)[1]
residual_artery_flow = artery_flow
r = first(current_roots)
for r in current_roots  # Removing artery flow taking into account multiple roots
    root_flow = sum(loi.working_flows[e] for e in out_edges(flow_dag, r))
    flow_to_remove = min(root_flow, residual_artery_flow)
    (back_path, _) = dijkstra_multisource(flow_dag, Set{VertexType}([r]), old_start, loi.costs)
    edge_back_path = get_edge_like_path(back_path, flow_dag, loi.fcfp.fixed_costs)
    for e in edge_back_path
        loi.working_flows[e] -= flow_to_remove
    end
    residual_artery_flow -= flow_to_remove
    if residual_artery_flow <= 0.0
        break
    end
end
# Find path from new root to new start and add flow
new_start = first(new_a)[1]
(down_path, _) = dijkstra_multisource(flow_dag, new_roots_candidates, new_start, loi.costs)
edge_down_path = get_edge_like_path(down_path, flow_dag, loi.fcfp.fixed_costs)
e = edge_down_path[2]
for e in edge_down_path
    loi.working_flows[e] += artery_flow
end
n = "intersect_S_2"

e2 = edge_down_path[2]
v1 = Float64(loi.flows[e2]) + Float64(artery_flow)
e1 = edge_down_path[1]
e3 = ("intersect_S_2", "intersect_5321182.0", 0)
v2 = Float64(loi.flows[e1]) + Float64(artery_flow) - Float64(loi.flows[e3])
v1 - v2

sum([loi.working_flows[e] for e in out_edges(flow_dag, n)]) - (sum([loi.working_flows[e] for e in in_edges(flow_dag, n)]))
[e for e in in_edges(flow_dag, n)]
# Set new artery flow
for e in new_a
    loi.working_flows[e] += artery_flow  # Adding flow because their may be already a flow on it
end
# Clean flows
clean_flow_two_cycles(loi.working_flows)
if potential_cycle_creation
    break_cycles(loi, loi.working_flows, Set{VertexType}([new_start]))
end

validate_move(loi)

new_a = edge_like_path
artery_flow = loi.working_flows[first(a)]
for e in a
    loi.working_flows[e] = zero(WeightType)
end
new_start = first(new_a)[1]
start_node = first(a)[1]
if new_start != start_node
    lca = lowest_common_ancestor(flow_dag, start_node, new_start)
    if lca === nothing
        throw(ErrorException("Unable to find common ancestor"))
    end
    # 1: Browse back the flow to not show flow in this direction
    if lca != start_node
        back_path, _ = dijkstra_multisource(flow_dag, Set{VertexType}([lca]), start_node, loi.costs)
        edge_back_path = get_edge_like_path(back_path, flow_dag, loi.fcfp.fixed_costs)
        for e in edge_back_path
            if artery_flow <= loi.working_flows[e]
                loi.working_flows[e] -= artery_flow
            else
                potential_cycle_creation = true
                loi.working_flows[(e[2], e[1], e[3])] = artery_flow - loi.working_flows[e]
                loi.working_flows[e] = zero(WeightType)
            end
        end
    end
    # 2: Browse down the flow to the start of the new artery
    if lca != new_start
        (down_path, _) = dijkstra_multisource(flow_dag, Set{VertexType}([lca]), new_start, loi.costs)
        edge_down_path = get_edge_like_path(down_path, flow_dag, loi.fcfp.fixed_costs)
        for e in edge_down_path
            loi.working_flows[e] += artery_flow
        end
    end
end
for e in new_a
    loi.working_flows[e] += artery_flow  # Adding flow because their may be already a flow on it
end
clean_flow_two_cycles(loi.working_flows)
if potential_cycle_creation
    break_cycles(loi, loi.working_flows, Set{VertexType}([new_start]))
end

gain = initial_cost - sum(loi.fcfp.fixed_costs[e] for (e, f) in loi.flows if f > 0)

optimization_time = time() - start
loi.logger("Arteries flows rerouting postprocess: gain=$(gain),  time: $(optimization_time)")
if check_solution(loi, loi.flows)
    return loi.flows
else
    loi.logger("Local optimization failed")
    return loi.init_flows
end

replace_arteries(loi)

import LocalOptimizer:replace_arteries, check_solution
check_solution(loi, loi.working_flows)

flow_dag = get_flow_dag(loi.optimization_graph, loi.flows)
arteries = get_flow_arteries(loi, flow_dag)
roots_descendants = Array{Set{VertexType}, 1}([descendants(flow_dag, r, false) for r in loi.roots])
roots_marginal_capacity = [loi.fcfp.production[r] - sum(get(loi.fcfp.heat_demand, n, 0.0)
                                                        for n in roots_descendants[r_index])
                           for (r_index, r) in enumerate(loi.roots)]

new_start = nothing
path = nothing
edge_like_path = nothing
a = nothing
potential_cycle_creation = nothing
using DataStructures

while !isempty(arteries)
    a = dequeue!(arteries)
    # determine sub tree
    end_node = last(a)[2]
    sub_tree_nodes = descendants(flow_dag, end_node, false)
    sub_tree_demand = sum(get(loi.fcfp.heat_demand, n, 0.0) for n in sub_tree_nodes)
    # determine parent sub tree
    start_node = first(a)[1]
    current_roots_indexes = Set{Int64}()
    other_roots_indexes = Set{Int64}()
    for (d_index, d) in enumerate(roots_descendants)
        if (start_node in d)
            push!(current_roots_indexes, d_index)
        elseif (roots_marginal_capacity[d_index] >= sub_tree_demand)
            push!(other_roots_indexes, d_index)
        end
    end
    concerned_sub_trees_indexes = union(current_roots_indexes, other_roots_indexes)
    parent_dag_trees_nodes = cascade_union([roots_descendants[d_index] for d_index in concerned_sub_trees_indexes])
    parent_dag_trees_nodes = setdiff(parent_dag_trees_nodes, union(sub_tree_nodes,
                                                                   Set(e[2] for e in a)))
    # Add 0-weight edges in the sub-graph
    modified_costs = add_zero_weight_edges_from_nodes(sub_tree_nodes, loi.costs)
    path, cost = dijkstra_multisource(loi.optimization_graph,
                                      parent_dag_trees_nodes,
                                      end_node,
                                      modified_costs)
    edge_like_path = get_edge_like_path(path, loi.optimization_graph, loi.fcfp.fixed_costs)
    if edge_like_path != a
        potential_cycle_creation = false
        if any((n != end_node) && (n in sub_tree_nodes) for n in path)
            potential_cycle_creation = true
        end
        new_start = first(path)
        if any(new_start in roots_descendants[d_index] for d_index in current_roots_indexes)
            loi.logger("Moving flow inside current dag")
            reroute_artery(loi, a, edge_like_path, flow_dag, potential_cycle_creation)
            break
        else
            loi.logger("Moving flow to another dag")
            println("NOWWW")
            break
        end
    end
end

edge_like_path

a

check_solution(loi, loi.working_flows)
validate_move(loi)

import LocalOptimizer: move_artery_to_other_dag, WeightType, VertexType, MultiEdgeType
other_roots_indexes = Set{Int64}(1)
new_roots_candidates = Set{VertexType}([loi.roots[r] for r in other_roots_indexes if new_start in roots_descendants[r]])
loi.working_flows = copy(loi.flows)

new_a = edge_like_path
artery_flow = loi.working_flows[first(a)]
for e in a
    loi.working_flows[e] = zero(WeightType)
end
new_start = first(new_a)[1]
# Find path from root to new start
(down_path, _) = dijkstra_multisource(flow_dag, new_roots_candidates, new_start, loi.costs)
edge_down_path = get_edge_like_path(down_path, flow_dag, loi.fcfp.fixed_costs)
for e in edge_down_path
    loi.working_flows[e] += artery_flow
end
# Set new artery flow
for e in new_a
    loi.working_flows[e] += artery_flow  # Adding flow because their may be already a flow on it
end
# Clean flows
clean_flow_two_cycles(loi.working_flows)
if potential_cycle_creation
    break_cycles(loi, loi.working_flows)
end

check_solution(loi, loi.working_flows)

flows = loi.working_flows
# Check that flow graph is a DAG
flow_dag = get_flow_dag(loi.optimization_graph, flows)
c = find_cycle(flow_dag)
if c !== nothing
    return false
end
# Check flow conservation
tol = 1e-5
nodes = get_nodes(flow_dag)
for n in nodes
    out_flow = sum(Array{WeightType}([flows[e] for e in out_edges(flow_dag, n)]))
    in_flow = sum(Array{WeightType}([flows[e] for e in in_edges(flow_dag, n)]))
    if n in loi.fcfp.building_nodes
        continue
    elseif n in loi.fcfp.production_nodes
        if out_flow > (loi.fcfp.production[n] + tol)
            println("$(n) : $(out_flow) > $((loi.fcfp.production[n] + tol))")
        end
    else
        if abs(out_flow - in_flow) > tol
            println("$(n) : $(out_flow) != $(in_flow)")
        end
    end
end
return true

import LocalOptimizer: out_edges, in_edges

previous_flow = loi.flows
f = open(raw"D:\projets\PlanHeat\tmp\instances\not_dag_3\flows_sources_move_pr.pickle", "w")
pickle.dump(previous_flow, f)
close(f)

f = open(raw"D:\projets\PlanHeat\tmp\instances\not_dag_3\flows_sources_move.pickle", "w")
pickle.dump(loi.working_flows, f)
close(f)
=#

[has_one_parent(flow_dag, e[2]) for e in current_artery]

a = dequeue!(arteries)

arteries = PriorityQueue{Array{MultiEdgeType, 1}, WeightType}()
explored_nodes = Set{VertexType}()
r = loi.roots[3]
nodes_stack = Array{VertexType, 1}([r])
while !isempty(nodes_stack)
    n = pop!(nodes_stack)
    if n in explored_nodes
        continue
    end
    push!(explored_nodes, n)
    out_flow_edges = collect(out_edges(flow_dag, n))
    println(n)
    if (length(out_flow_edges) != 1) || (! has_one_parent(flow_dag, n))
        for e in out_flow_edges
            push!(nodes_stack, e[2])
        end
        continue
    end
    println("here")
    println(n)
    break
end

n = "intersect_S_0_2"
out_flow_edges = collect(out_edges(flow_dag, n))
# A new artery is found, explore it in depth
current_artery = [first(in_edges(flow_dag, n)), out_flow_edges[1]]
sub_nodes_stack = [out_flow_edges[1][2]]
while !isempty(sub_nodes_stack)
    sub_node = pop!(sub_nodes_stack)
    push!(explored_nodes, sub_node)
    out_flow_edges = collect(out_edges(flow_dag, sub_node))
    if (length(out_flow_edges) != 1) || (! has_one_parent(flow_dag, n))
        # End of the artery
        artery_cost = - sum(loi.fcfp.fixed_costs[e] for e in current_artery)
        if current_artery in loi.treated_arteries
            artery_cost += 1e5
        end
        arteries[current_artery] = artery_cost
        for e in out_flow_edges
            push!(nodes_stack, e[2])
        end
        break
    else
        push!(current_artery, out_flow_edges[1])
        push!(sub_nodes_stack, out_flow_edges[1][2])
    end
end


###############################################################################
#
# Canopy radiation
#
###############################################################################
"""
    mutable struct CanopyRads{FT}

A struct for canopy radiation information

# Fields
$(TYPEDFIELDS)
"""
Base.@kwdef mutable struct CanopyRads{FT}
    # local storage of dimension information
    "Number of azimuth angles"
    nAzi  ::Int = 36
    "Number of inclination agles"
    nIncl ::Int = 9
    "Number of canopy layers"
    nLayer::Int = 5
    "Number of canopy levels"
    nLevel::Int = nLayer+1
    "Number of wave lengths"
    nWL   ::Int = 10
    "Number of wave lengths for SIF"
    nWLF  ::Int = 10

    # Scalars
    "Integrated TOC outgoing flux `[W m⁻²]`"
    intEout            ::FT = 0
    "Incident spectrally integrated total PAR `[mol m⁻² s⁻¹]`"
    incomingPAR        ::FT = 0
    "Incident spectrally integrated direct PAR `[mol m⁻² s⁻¹]`"
    incomingPAR_direct ::FT = 0
    "Incident spectrally integrated diffuse PAR `[mol m⁻² s⁻¹]`"
    incomingPAR_diffuse::FT = 0
    "Net radiation of shaded soil `[W m⁻²]`"
    RnSoil_diffuse     ::FT = 0
    "Net Short-wave radiation of sunlit soil `[W m⁻²]`"
    RnSoil_direct      ::FT = 0
    "Net Short-wave radiation of soil (shaded + sunlit) `[W m⁻²]`"
    RnSoil             ::FT = 0
    "Net long-wave radiation of soil (shaded + sunlit) `[W m⁻²]`"
    RnSoilLW           ::FT = 0

    # Dim of nLayer
    "Net PAR of shaded leaves `[mol m⁻² s⁻¹]`"
    absPAR_shade   ::Array{FT,1} = zeros(FT, nLayer)
    "Net PAR by Cab+Car of shaded leaves `[moles m⁻² s⁻¹]`"
    absPAR_shadeCab::Array{FT,1} = zeros(FT, nLayer)
    "Spectrally integrated net absorbed direct radiation in each layer per leaf area `[W m⁻²]`"
    intNetSW_sunlit::Array{FT,1} = zeros(FT, nLayer)
    "Spectrally integrated net absorbed diffuse radiation in each layer per leaf area `[W m⁻²]`"
    intNetSW_shade ::Array{FT,1} = zeros(FT, nLayer)
    "Spectrally integrated net absorbed direct radiation in each layer per leaf area `[W m⁻²]`"
    intNetLW_sunlit::Array{FT,1} = zeros(FT, nLayer)
    "Spectrally integrated net absorbed diffuse radiation in each layer per leaf area `[W m⁻²]`"
    intNetLW_shade ::Array{FT,1} = zeros(FT, nLayer)
    "Leaf temperature (sunlit) `[K]`"
    T_sun          ::Array{FT,1} = zeros(FT, nLayer) .+ FT(298.15)
    "Leaf temperature (shaded) `[K]`"
    T_shade        ::Array{FT,1} = zeros(FT, nLayer) .+ FT(298.15)
    "Fluorescence yield for shaded leaves"
    ϕ_shade        ::Array{FT,1} =  ones(FT, nLayer) .* FT(0.01)
    "Sensible Heat flux H of shaded leaves `[W m⁻²]`"
    H_shade        ::Array{FT,1} = zeros(FT, nLayer)
    "Latent Heat flux LE of shaded leaves `[W m⁻²]`"
    LE_shade       ::Array{FT,1} = zeros(FT, nLayer)
    "NPQ of shaded leaves"
    NPQ_shade      ::Array{FT,1} = zeros(FT, nLayer)
    # TODO remove these?
    "GPP of shaded leaves `[μmol m⁻² s⁻¹]`"
    GPP_shade      ::Array{FT,1} = zeros(FT, nLayer)
    "gs of shaded leaves `[mol m⁻² s⁻¹]`"
    gs_shade       ::Array{FT,1} = zeros(FT, nLayer)
    "Leaf water potential of shaded leaves `[MPa]`"
    ψl_shade       ::Array{FT,1} = zeros(FT, nLayer)
    "Cc of shaded leaves `[µmol/mol]`"
    Cc_shade       ::Array{FT,1} = zeros(FT, nLayer)
    "internal CO₂ concentration of shaded leaves `[µmol/mol]`"
    Pi_shade       ::Array{FT,1} = zeros(FT, nLayer)

    # Dimension of wavelength
    "Short-wave TOC outgoing radiance in observation direction `[mW m⁻² nm⁻¹ sr⁻¹]`"
    Lo         ::Array{FT,1} = zeros(FT, nWL)
    "Short-wave TOC outgoing radiation `[mW m⁻² nm⁻¹]`"
    Eout       ::Array{FT,1} = zeros(FT, nWL)
    "Short-wave Albedo in viewing direction"
    alb_obs    ::Array{FT,1} = zeros(FT, nWL)
    "Short-wave Albedo for direct incoming radiation"
    alb_direct ::Array{FT,1} = zeros(FT, nWL)
    "Short-wave Albedo for diffuse incoming radiation"
    alb_diffuse::Array{FT,1} = zeros(FT, nWL)

    # Dimension of nLevel * nWavelengths
    "Upwelling diffuse short-wave radiation within canopy `[mW m⁻² nm⁻¹]`"
    E_up  ::Array{FT,2} = zeros(FT, (nWL,nLevel))
    "Downwelling diffuse short-wave radiation within canopy `[mW m⁻² nm⁻¹]`"
    E_down::Array{FT,2} = zeros(FT, (nWL,nLevel))

    # Dimension of nLayer * nWavelengths
    "Net absorbed direct radiation in each layer `[mW m⁻² nm⁻¹]`"
    netSW_sunlit   ::Array{FT,2} = zeros(FT, (nWL,nLayer))
    "net absorbed diffuse radiation in each layer `[mW m⁻² nm⁻¹]`"
    netSW_shade    ::Array{FT,2} = zeros(FT, (nWL,nLayer))

    # Dimension of nLeafInclination * nLeafAzimuth * nLayer
    "net PAR of sunlit leaves `[mol m⁻² s⁻¹]`"
    absPAR_sun   ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    "net PAR by Cab+Car of sunlit leaves `[mol m⁻² s⁻¹]`"
    absPAR_sunCab::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    "Leaf temperature (sunlit) `[K]`"
    T_sun3D      ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer)) .+ FT(298.15)
    "Fluorescence yield for sunlit leaves"
    ϕ_sun        ::Array{FT,3} =  ones(FT, (nIncl,nAzi,nLayer)) .* FT(0.01)
    "Sensible Heat flux H of sunlit leaves `[W m⁻²]`"
    H_sun        ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    "Latent Heat flux LE of sunlit leaves `[W m⁻²]`"
    LE_sun       ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    "NPQ of sunlit leaves"
    NPQ_sun      ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    # TODO remove these?
    "GPP of sunlit leaves `[μmol m⁻² s⁻¹]`"
    GPP_sun      ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    "gs of sunlit leaves `[mol m⁻² s⁻¹]`"
    gs_sun       ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    "Leaf water potential of sunlit leaves `[MPa]`"
    ψl_sun       ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    "Cc of sunlit leaves `[µmol/mol]`"
    Cc_sun       ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))
    "Internal CO₂ concentration of sunlit leaves `[µmol/mol]`"
    Pi_sun       ::Array{FT,3} = zeros(FT, (nIncl,nAzi,nLayer))

    # Fluorescence Output:
    "Hemispheric total outgoing SIF flux `[mW m⁻² nm⁻¹]`)"
    SIF_hemi         ::Array{FT,1} = zeros(FT, nWLF)
    "Observer-direction outgoing SIF radiance (mW m⁻² nm⁻¹ sr⁻¹))"
    SIF_obs          ::Array{FT,1} = zeros(FT, nWLF)
    "Observer-direction outgoing SIF radiance, sunlit leaves (mW m⁻² nm⁻¹ sr⁻¹)"
    SIF_obs_sunlit   ::Array{FT,1} = zeros(FT, nWLF)
    "Observer-direction outgoing SIF radiance, shaded leaves (mW m⁻² nm⁻¹ sr⁻¹)"
    SIF_obs_shaded   ::Array{FT,1} = zeros(FT, nWLF)
    "Observer-direction outgoing SIF radiance, scattered (mW m⁻² nm⁻¹ sr⁻¹)"
    SIF_obs_scattered::Array{FT,1} = zeros(FT, nWLF)
    "Observer-direction outgoing SIF radiance, soil-reflected (mW m⁻² nm⁻¹ sr⁻¹)"
    SIF_obs_soil     ::Array{FT,1} = zeros(FT, nWLF)
    "Total SIF sum of layer sources  `[mW m⁻² nm⁻¹]`)"
    SIF_sum          ::Array{FT,1} = zeros(FT, nWLF)
end


module SimpleFluid
using HoJMTK

@parameters t
const D = Differential(t)

include("core.jl")
include("sources.jl")
include("components.jl")

export Ground,
    connect_series,
    connect_parallel,
    connect_star,
    Pump,
    FluidResistor,
    VesselCompliance,
    FluidInertia,
    FluidGround
end


export BasicDQNLearner

"""
    BasicDQNLearner(;kwargs...)

See paper: [Playing Atari with Deep Reinforcement Learning](https://www.cs.toronto.edu/~vmnih/docs/dqn.pdf)

This is the very basic implementation of DQN. Compared to the traditional Q learning, the only difference is that,
in the updating step it uses a batch of transitions sampled from an experience buffer instead of current transition.
And the `approximator` is usually a [`NeuralNetworkApproximator`](@ref).
You can start from this implementation to understand how everything is organized and how to write your own customized algorithm.

# Keywords

- `approximator`::[`AbstractApproximator`](@ref): used to get Q-values of a state.
- `loss_func`: the loss function to use.
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `rng=Random.GLOBAL_RNG`
"""
mutable struct BasicDQNLearner{Q,F,R} <: AbstractLearner
    approximator::Q
    loss_func::F
    γ::Float32
    sampler::BatchSampler
    min_replay_history::Int
    rng::R
    # for debugging
    loss::Float32
end

Flux.functor(x::BasicDQNLearner) = (Q = x.approximator,), y -> begin
    x = @set x.approximator = y.Q
    x
end

(learner::BasicDQNLearner)(env) =
    env |>
    get_state |>
    x -> send_to_device(device(learner), x) |>
    learner.approximator |>
    send_to_host

function BasicDQNLearner(;
    approximator::Q,
    loss_func::F = huber_loss,
    γ = 0.99f0,
    batch_size = 32,
    min_replay_history = 32,
    rng = Random.GLOBAL_RNG,
) where {Q,F}
    BasicDQNLearner{Q,F,typeof(rng)}(
        approximator,
        loss_func,
        γ,
        BatchSampler{SARTS}(batch_size),
        min_replay_history,
        rng,
        0.0,
    )
end

function RLBase.update!(learner::BasicDQNLearner, traj::AbstractTrajectory)
    if length(traj) >= learner.min_replay_history
        inds, batch = sample(learner.rng, traj, learner.sampler)
        update!(learner, batch)
    end
end

function RLBase.update!(learner::BasicDQNLearner, batch::NamedTuple{SARTS})

    Q = learner.approximator
    γ = learner.γ
    loss_func = learner.loss_func

    s, a, r, t, s′ = send_to_device(device(Q), batch)
    a = CartesianIndex.(a, 1:length(a))

    gs = gradient(params(Q)) do
        q = Q(s)[a]
        q′ = vec(maximum(Q(s′); dims = 1))
        G = r .+ γ .* (1 .- t) .* q′
        loss = loss_func(G, q)
        ignore() do
            learner.loss = loss
        end
        loss
    end

    update!(Q, gs)
end


# This module contains the necessary methods to use Fermi Arrays within TensorOperations

import TensorOperations: contract!
import TensorOperations: IndexTuple
import TensorOperations: IndexError
import TensorOperations: similarstructure_from_indices
import TensorOperations: scalar
import TBLIS

function tblis_set_num_threads(N)
    TBLIS.set_num_threads(N)
end

function set_num_threads(N = Threads.nthreads())
    tblis_set_num_threads(N)
    LinearAlgebra.BLAS.set_num_threads(N)
end

function oind2eins(oindA::NTuple{NAo}, cindA::NTuple{NAc},
          oindB::NTuple{NBo}, cindB::NTuple{NBc},
          tindC::NTuple{NCt}) where {NAo, NAc, NBo, NBc, NCt}

    # This function converts the numerical scheme used in TensorOperations
    # to represent contractions into an Einsum scheme.

    # Check contraction consistency.
    NAo + NBo == NCt || throw(IndexError("number of outer index not consistent."))
    NAc == NBc || throw(IndexError("number of contracted index not consistent."))

    cPadding = 'a' - 'A'
    einA = zeros(Int8, NAo+NAc)
    einB = zeros(Int8, NBo+NBc)

    # Outer indices.
    for i = 1:NAo
        einA[oindA[i]] = i
    end
    for i = 1:NBo
        einB[oindB[i]] = i + NAo
    end

    # Contracted indices.
    for i = 1:NAc
        einA[cindA[i]] = i + cPadding
        einB[cindB[i]] = i + cPadding
    end

    einA = string((einA .+'A')...)
    einB = string((einB .+'A')...)
    einC = string((tindC.+'A')...) # C has direct conversion relations.
    einA, einB, einC
end

function contract!(α,
          A::FermiMDArray, conjA::Symbol,
          B::FermiMDArray, conjB::Symbol,
          β,
          C::FermiMDArray,
          oindA::IndexTuple, cindA::IndexTuple,
          oindB::IndexTuple, cindB::IndexTuple,
          tindC::IndexTuple, syms::Union{Nothing, NTuple{3,Symbol}} = nothing)

    if Fermi.Options.get("tblis")
        # Check permutation consistency.
        # This check is copied from stridedarray.jl
        pA = (oindA...,cindA...)
        (length(pA) == ndims(A) && isperm(pA)) ||
            throw(IndexError("invalid permutation of length $(ndims(A)): $pA"))
        pB = (oindB...,cindB...)
        (length(pB) == ndims(B) && isperm(pB)) ||
            throw(IndexError("invalid permutation of length $(ndims(B)): $pB"))
        (length(tindC) == ndims(C) && isperm(tindC)) ||
            throw(IndexError("invalid permutation of length $(ndims(C)): $tindC"))

        # Complex conjugate is not supported by TBLIS yet.
        # Here we do a intermediate copying-conversion.
        if conjA == :C || conjA == :A || conjB == :C || conjB == :A
            throw(Fermi.Options.FermiException("conj operations not suported for TBLIS"))
        end
        einA, einB, einC = oind2eins(oindA, cindA,
                                     oindB, cindB,
                                     tindC)

        CType = eltype(C.data)
        _A = TBLIS.TTensor{CType}(A.data,CType(α))
        _B = TBLIS.TTensor{CType}(B.data)
        _C = TBLIS.TTensor{CType}(C.data,CType(β))

        TBLIS.mul!(_C, _A, _B, einA, einB, einC)
        return C
    else
        contract!(α, A.data, conjA, B.data, conjB, β, C.data, oindA, cindA, oindB, cindB, tindC)
        return C
    end
end

function contract!(α,
          A::FermiMDArray, conjA::Symbol,
          B::FermiMDArray, conjB::Symbol,
          β,
          C::AbstractArray,
          oindA::IndexTuple, cindA::IndexTuple,
          oindB::IndexTuple, cindB::IndexTuple,
          tindC::IndexTuple, syms::Union{Nothing, NTuple{3,Symbol}} = nothing)

    contract!(α, A, conjA, B, conjB, β, FermiMDArray(C), oindA, cindA, oindB, cindB, tindC, syms)
end

function contract!(α,
          A::FermiMDArray, conjA::Symbol,
          B::AbstractArray, conjB::Symbol,
          β,
          C::AbstractArray,
          oindA::IndexTuple, cindA::IndexTuple,
          oindB::IndexTuple, cindB::IndexTuple,
          tindC::IndexTuple, syms::Union{Nothing, NTuple{3,Symbol}} = nothing)

    contract!(α, A, conjA, FermiMDArray(B), conjB, β, FermiMDArray(C), oindA, cindA, oindB, cindB, tindC, syms)
end

function contract!(α,
          A::AbstractArray, conjA::Symbol,
          B::FermiMDArray, conjB::Symbol,
          β,
          C::AbstractArray,
          oindA::IndexTuple, cindA::IndexTuple,
          oindB::IndexTuple, cindB::IndexTuple,
          tindC::IndexTuple, syms::Union{Nothing, NTuple{3,Symbol}} = nothing)

    contract!(α, FermiMDArray(A), conjA, B, conjB, β, FermiMDArray(C), oindA, cindA, oindB, cindB, tindC, syms)
end