# structs
export mjContact, mjWarningStat, mjTimerStat, mjSolverStat, mjData, mjLROpt, mjVFS, mjOption, mjVisual, mjStatistic, mjModel, mjrRect, mjrContext, mjuiState, mjuiThemeSpacing, mjuiThemeColor, mjuiItem, mjuiSection, mjUI, mjuiDef, mjvPerturb, mjvCamera, mjvGLCamera, mjvGeom, mjvLight, mjvOption, mjvScene, mjvFigure
#names,
#_global,
#quality,
#headlight,
#map,
#scale,
#rgba,

export mj_activate, mj_deactivate, mj_certQuestion, mj_certAnswer, mj_certCheck, mj_defaultVFS, mj_addFileVFS, mj_makeEmptyFileVFS, mj_findFileVFS, mj_deleteFileVFS, mj_deleteVFS, mj_loadXML, mj_saveLastXML, mj_freeLastXML, mj_printSchema, mj_step, mj_step1, mj_step2, mj_forward, mj_inverse, mj_forwardSkip, mj_inverseSkip, mj_defaultLROpt, mj_defaultSolRefImp, mj_defaultOption, mj_defaultVisual, mj_copyModel, mj_saveModel, mj_loadModel, mj_deleteModel, mj_sizeModel, mj_makeData, mj_copyData, mj_resetData, mj_resetDataDebug, mj_resetDataKeyframe, mj_stackAlloc, mj_deleteData, mj_resetCallbacks, mj_setConst, mj_setLengthRange, mj_printModel, mj_printData, mj_fwdPosition, mj_fwdVelocity, mj_fwdActuation, mj_fwdAcceleration, mj_fwdConstraint, mj_Euler, mj_RungeKutta, mj_invPosition, mj_invVelocity, mj_invConstraint, mj_compareFwdInv, mj_sensorPos, mj_sensorVel, mj_sensorAcc, mj_energyPos, mj_energyVel, mj_checkPos, mj_checkVel, mj_checkAcc, mj_kinematics, mj_comPos, mj_camlight, mj_tendon, mj_transmission, mj_crb, mj_factorM, mj_solveM, mj_solveM2, mj_comVel, mj_passive, mj_subtreeVel, mj_rne, mj_rnePostConstraint, mj_collision, mj_makeConstraint, mj_projectConstraint, mj_referenceConstraint, mj_constraintUpdate, mj_addContact, mj_isPyramidal, mj_isSparse, mj_isDual, mj_mulJacVec, mj_mulJacTVec, mj_jac, mj_jacBody, mj_jacBodyCom, mj_jacGeom, mj_jacSite, mj_jacPointAxis, mj_name2id, mj_id2name, mj_fullM, mj_mulM, mj_mulM2, mj_addM, mj_applyFT, mj_objectVelocity, mj_objectAcceleration, mj_contactForce, mj_differentiatePos, mj_integratePos, mj_normalizeQuat, mj_local2Global, mj_getTotalmass, mj_setTotalmass, mj_version, mj_ray, mj_rayHfield, mj_rayMesh, mjv_defaultCamera, mjv_defaultPerturb, mjv_room2model, mjv_model2room, mjv_cameraInModel, mjv_cameraInRoom, mjv_frustumHeight, mjv_alignToCamera, mjv_moveCamera, mjv_movePerturb, mjv_moveModel, mjv_initPerturb, mjv_applyPerturbPose, mjv_applyPerturbForce, mjv_averageCamera, mjv_select, mjv_defaultOption, mjv_defaultFigure, mjv_initGeom, mjv_makeConnector, mjv_defaultScene, mjv_makeScene, mjv_freeScene, mjv_updateScene, mjv_addGeoms, mjv_makeLights, mjv_updateCamera, mjv_updateSkin, mjr_defaultContext, mjr_makeContext, mjr_changeFont, mjr_addAux, mjr_freeContext, mjr_uploadTexture, mjr_uploadMesh, mjr_uploadHField, mjr_restoreBuffer, mjr_setBuffer, mjr_readPixels, mjr_drawPixels, mjr_blitBuffer, mjr_setAux, mjr_blitAux, mjr_text, mjr_overlay, mjr_maxViewport, mjr_rectangle, mjr_figure, mjr_render, mjr_finish, mjr_getError, mjr_findRect, mju_printMat, mju_printMatSparse, mju_rayGeom, mju_raySkin, mju_error, mju_error_i, mju_error_s, mju_warning, mju_warning_i, mju_warning_s, mju_clearHandlers, mju_malloc, mju_free, mj_warning, mju_writeLog, mju_zero3, mju_copy3, mju_scl3, mju_add3, mju_sub3, mju_addTo3, mju_subFrom3, mju_addToScl3, mju_addScl3, mju_normalize3, mju_norm3, mju_dot3, mju_dist3, mju_rotVecMat, mju_rotVecMatT, mju_cross, mju_zero4, mju_unit4, mju_copy4, mju_normalize4, mju_zero, mju_copy, mju_sum, mju_L1, mju_scl, mju_add, mju_sub, mju_addTo, mju_subFrom, mju_addToScl, mju_addScl, mju_normalize, mju_norm, mju_dot, mju_mulMatVec, mju_mulMatTVec, mju_transpose, mju_mulMatMat, mju_mulMatMatT, mju_mulMatTMat, mju_sqrMatTD, mju_transformSpatial, mju_dotSparse, mju_dotSparse2, mju_dense2sparse, mju_sparse2dense, mju_mulMatVecSparse, mju_compressSparse, mju_combineSparse, mju_sqrMatTDSparse, mju_transposeSparse, mju_rotVecQuat, mju_negQuat, mju_mulQuat, mju_mulQuatAxis, mju_axisAngle2Quat, mju_quat2Vel, mju_subQuat, mju_quat2Mat, mju_mat2Quat, mju_derivQuat, mju_quatIntegrate, mju_quatZ2Vec, mju_mulPose, mju_negPose, mju_trnVecPose, mju_cholFactor, mju_cholBacksub, mju_cholSolve, mju_cholUpdate, mju_cholFactorSparse, mju_cholSolveSparse, mju_cholUpdateSparse, mju_eig3, mju_muscleFVL, mju_musclePassive, mju_pneumatic, mju_muscleGain, mju_muscleBias, mju_muscleDynamics, mju_encodePyramid, mju_decodePyramid, mju_springDamper, mju_min, mju_max, mju_sign, mju_round, mju_type2Str, mju_str2Type, mju_warningText, mju_isBad, mju_isZero, mju_standardNormal, mju_f2n, mju_n2f, mju_d2n, mju_n2d, mju_insertionSort, mju_Halton, mju_strncpy, mjui_themeSpacing, mjui_themeColor, mjui_add, mjui_resize, mjui_update, mjui_event, mjui_render


module Library

using BioSequences
using ..RNASplicing: OpaqueTransElement, Interaction, Spliceosome, only
# The core splicing machinery can be expressed in this framework

# # Core Machinery

"5’ splice site - 9-nucleotide motif that generally conforms to the sequence"
donorsite = biore"(C|A)AG|GU(A|G)AGU"rna
# donorsite(seq) = occursin(biore"(C|A)AG|GU(A|G)AGU"rna, seq) ? 1.0 : 0.0
u1 = OpaqueTransElement(:u1)
u2 = OpaqueTransElement(:u2)
u2af = OpaqueTransElement(:u2af)
e1 = Interaction(donorsite, only(u1), 0, 10.0)
coremachinery = [e1]

acceptsite = biore"A(C|U)AG"rna

# # SREs

# `seq` have a sequence of at least three Gs?
grun = biore"GGG+"rna

# Major splicesosome
majorsos = Spliceosome(Dict(u1 => 1.0))
end


#!/usr/bin/env julia

println("Loading dependencies...")

using Plots
using LibSerialPort

# Plot backend
gr(size=(900, 500))

address = length(ARGS) == 1 ? ARGS[1] : "/dev/cu.usbmodem1421"

mcu = open(address, 115200)

# Number of points to display
n = 500

# Values for horizontal (time) axis and data
x = Float64[]
y, s = Float64[], Float64[]

function readmcu(sp::SerialPort)
    # Get message as a String, split to array
    msg = readuntil(mcu, '\n', 100)
    data = zeros(Float64, 3)
    try
        data = [parse(Float64, x) for x in split(msg)]
    catch
        println("Skipping: $msg")
        return
    end
    length(data) == 3 || return
    time, y1, y2 = data

    # println("$time, $y1, $y2")

    # Manage arrays as FIFO queues
    push!(x, time/1000) # ms to seconds
    push!(y, y1/1023)
    push!(s, y2/1023)
    if length(x) > n
        for v in (x, y, s)
            shift!(v)
        end
    end
end

function rtplot(x, y, s)
    plot(x, [y,s],
         title="Data stream", # Not displaying with GR backend (?)
         xaxis="time [s]",
         yaxis=("example [units]", (-0.2, 1.2)), # title, limits
         labels=(["Measured","Filtered"]), # legend
         line=([1 2], [:steppre :path], ["gray" "crimson"], [1.0 0.5]), # lineweights, linetypes, colors, alphas
         fill=(0, [0.1 0], "dodgerblue"), # y-origin, alpha, colors
         )
    # Update the plot window
    gui()
end

println("Creating plot...")
readmcu(mcu)
rtplot(x,y,s)

println("Ready")

while true

    @async begin
        a = readline(STDIN)
        a == "q\n" && quit()
        write(mcu, "$a")
    end

    @async begin
        readmcu(mcu)
        rtplot(x,y,s)
    end

    # Give the queued tasks a chance to run
    sleep(0.001)
end


#!/usr/bin/env julia -p 4

@everywhere using Params, Rasim

pmap(Rasim.run_whole_simulation, genparams())


using StanRun, Test

@testset "cmdstan run and results" begin
    # setup environment
    MODELDIR = mktempdir()
    SAMPLEDIR = mktempdir()
    SRC = joinpath(MODELDIR, "test.stan")
    cp(joinpath(@__DIR__, "test.stan"), SRC)

    # run a model
    model = StanModel(SRC)
    @test repr(model) ==
        "Stan model at $(SRC)\n    (CmdStan home: $(StanRun.get_cmdstan_home()))"
    exec_path = StanRun.ensure_executable(model)
    @test isfile(exec_path)
    n_chains = 5
    OUTPUT_BASE = joinpath(SAMPLEDIR, "test")
    time_before = time()
    @test stan_compile(model) ≡ nothing
    chains = stan_sample(model, (N = 100, x = randn(100)), 5; output_base = OUTPUT_BASE)
    for (sample, logfile) in chains
        @test ctime(sample) ≥ time_before
        @test ctime(sample) ≥ time_before
    end
    @test first.(chains) == sort(StanRun.find_samples(OUTPUT_BASE)) ==
        [joinpath(SAMPLEDIR, "test_chain_$(i).csv") for i in 1:n_chains]
end

@testset "unset cmdstan environment" begin
    withenv("JULIA_CMDSTAN_HOME" => nothing) do
        @test_throws ErrorException StanRun.get_cmdstan_home()
    end
end

@testset "model error and message" begin
    model = StanModel(joinpath(@__DIR__, "test_incorrect.stan"))
    try
        stan_compile(model)
    catch e
        @test e isa StanRun.StanModelError
        @test occursin("Variable \"x\" does not exist", e.message)
        io = IOBuffer()
        showerror(io, e)
        e_repr = String(take!(io))
        @test occursin("error when compiling", e_repr)
    end
end


using Mimi
using Distributions
using Test

include("test-model/test-model.jl")
using .TestModel

m = create_model()
mm1 = create_marginal_model(create_model())
mm2 = create_marginal_model(create_model())

simdef = @defsim begin
    share = Uniform(0, 1)
    save(emissions.E_Global)
end

# Test running one MarginalModel

function post_trial1(si, trialnum, ntimesteps, tup)
    @test length(si.models) == 1
    @test si.models[1] isa MarginalModel
end

si = run(simdef, mm1, 2, post_trial_func = post_trial1, results_in_memory = true)
@test all(iszero, si.results[1][:emissions, :E_Global][!, :E_Global])   # Test that the marginal emission saved from the MarginalModel are zeros (because there's no difference between mm.base and mm.marginal)

# Test running a vector of MarginalModels

function post_trial2(si, trialnum, ntimesteps, tup)
    @test length(si.models) == 2
    @test si.models[1] isa MarginalModel
    @test si.models[2] isa MarginalModel
end

si = run(simdef, [mm1, mm2], 2, post_trial_func = post_trial2, results_in_memory = true)
@test all(iszero, si.results[1][:emissions, :E_Global][!, :E_Global])  # Test that the regular model has non-zero emissions
@test all(iszero, si.results[2][:emissions, :E_Global][!, :E_Global])   # Test that the marginal emission saved from the MarginalModel are zeros (because there's no difference between mm.base and mm.marginal)

# Test running a vector of a Model and a MarginalModel

function post_trial3(si, trialnum, ntimesteps, tup)
    @test length(si.models) == 2
    @test si.models[1] isa Model
    @test si.models[2] isa MarginalModel
end

si = run(simdef, [m, mm1], 2, post_trial_func = post_trial3, results_in_memory = true)
@test all(!iszero, si.results[1][:emissions, :E_Global][!, :E_Global])  # Test that the regular model has non-zero emissions
@test all(iszero, si.results[2][:emissions, :E_Global][!, :E_Global])   # Test that the marginal emission saved from the MarginalModel are zeros (because there's no difference between mm.base and mm.marginal)


# Various constants used to determine problem type

const RASTER = ["raster", "Raster"]
const PAIRWISE = ["pairwise", "Pairwise"]
const ADVANCED = ["advanced", "Advanced"]
const ONETOALL = ["one-to-all", "one_to_all"]
const ALLTOONE = ["all-to-one", "all_to_one"]
const SINGLE = ["single", "Single"]
const DOUBLE = ["double", "Double"]

# Solver constants
const AMG = ["cg+amg", "amg+cg"]
const CHOLMOD = ["cholmod", "cholesky", "cholfact"]
const MKLPARDISO = ["mklpardiso", "MKLPardiso", "MKLPARDISO", "pardiso"]

# Constants used in IO
const AAGRID = 2
const TXTLIST = 3
const PAIRS_AAGRID = 4
const PAIRS_LIST = 5
const TRUELIST = ["True", "true", "1"]

const FILE_TYPE_NPY = 1
const FILE_TYPE_AAGRID = 2
const FILE_TYPE_TXTLIST = 3
const FILE_TYPE_INCL_PAIRS_AAGRID = 4
const FILE_TYPE_INCL_PAIRS = 5
const FILE_TYPE_GEOTIFF = 6

const FILE_HDR_GZIP = "\x1f\x8b\x08"
const FILE_HDR_NPY = "\x93NUMPY"
const FILE_HDR_AAGRID = "ncols"
const FILE_HDR_INCL_PAIRS_AAGRID = "min"
const FILE_HDR_INCL_PAIRS = "mode"

# Constants for logging
const NONE = ["NONE", "None", "none"]
const INFO = ["INFO", "info", "Info"]
const DEBUG = ["DEBUG", "debug", "Debug"]
const WARNING = ["WARNING", "WARN", "warning", "warn", "Warning", "Warn"]
const CRITICAL = ["CRITICAL", "critical", "Critical"]

# Constants for tolerance when checking matrix norms
const TOL_SINGLE = 1e-3
const TOL_DOUBLE = 1e-5


using GeometryBasics, RPRMakie
using Colors, FileIO, ImageShow
using Colors: N0f8

RPRMakie.activate!()
fig = Figure(; resolution=(1200, 800), fontsize=26)
radiance = 10000
lights = [EnvironmentLight(0.5, load(RPR.assetpath("studio026.exr"))),
          PointLight(Vec3f(0, 0, 20), RGBf(radiance, radiance, radiance))]

ax = LScene(fig[1, 1]; scenekw=(lights=lights,))
Makie.axis3d!(ax.scene, Rect3f(Vec3f(0), Vec3f(1, 1, 1.2)))
screen = RPRScreen(ax.scene; iterations=200, plugin=RPR.Northstar)
matsys = screen.matsys
rectMesh = FRect3D(Vec3f0(-0.5, -0.5, 0), Vec3f0(1))
recmesh = GeometryBasics.normal_mesh(rectMesh)
n = 100
pos = [Point3f(i, j, 0) ./ 10 for i in 1:n for j in 1:n]
z = rand(n, n)
mat = RPR.Chrome(matsys)
mat.roughness = Vec4f(0.3)
meshscatter!(ax, pos; marker=recmesh, markersize=Vec3f.(0.1, 0.1, z[:]), material=mat, color=vec(z))

cam = cameracontrols(ax.scene)
cam.eyeposition[] = Float32[5, 22, 12]
cam.lookat[] = Float32[5, 5, -0.5]
cam.upvector[] = Float32[0.0, 0.0, 1.0]
cam.fov[] = 14.0

# show the screen directly in e.g. VSCode (display unecessary, )
# display(screen)
# alternatively, get the Image directly from the screen

image = colorbuffer(screen)::Matrix{RGB{N0f8}}

save("bars.png", image)


abstract type AbstractContact{T} <: AbstractObject{T} end

"""
    mutable struct Contact{T, G, MT} <: AbstractContact{T}

Contact of a [`SolidStateDetector`](@ref).

For the simulation of the [`ElectricPotential`](@ref), all contacts are fixed to a constant potential value.

## Parametric types:
* `T`: Precision type.
* `G`: Type of `geometry`.
* `MT`: Type of `material`.

## Fields
* `potential::T`: Potential (in V) to which the contact will be fixed during the calculation of the [`ElectricPotential`](@ref).
* `material::MT`: Material of the contact.
* `id::Int`: Unique id that will unambiguously identify the contact.
* `name::String`: Custom name for the contact, relevant for plotting.
* `geometry::G`: Geometry of the contact, see [Constructive Solid Geometry (CSG)](@ref).

## Definition in Configuration File

A `Contact` is defined in the configuration file through an entry in the `contacts` array of a detector.
It needs `id`, `potential` and `geometry` and can optionally be given a `name` and `material`.

An example definition of contacts looks like this:
```yaml
contacts:
  - name: "n+ contact"
    id: 1
    potential: 5000V
    material: HPGe # optional
    geometry: # ....
  - name: "p+ contact"
    id: 2
    potential: 0
    material: HPGe #optional
    geometry: # ....
```
"""
struct Contact{T,G,MT} <: AbstractContact{T}
    potential::T
    material::MT
    id::Int
    name::String
    geometry::G
end

function Contact{T}(dict::Union{Dict{String,Any}, Dict{Any, Any}}, input_units::NamedTuple, outer_transformations)::Contact{T} where {T <: SSDFloat}
    id::Int = haskey(dict, "id") ? dict["id"] : -1
    material = haskey(dict, "material") ? material_properties[materials[dict["material"]]] : material_properties[materials["HPGe"]]
    name = haskey(dict,"name") ? dict["name"] : ""
    inner_transformations = parse_CSG_transformation(T, dict, input_units)
    transformations = combine_transformations(inner_transformations, outer_transformations)
    geometry = Geometry(T, dict["geometry"], input_units, transformations)
    return Contact( T(dict["potential"]), material, id, name, geometry )
end

function println(io::IO, d::Contact)
    println("\t________"*"Contact $(d.id)"*"________\n")
    println("\t---General Properties---")
    println("\t-Potential: \t\t $(d.potential) V")
    println("\t-Contact Material: \t $(d.material.name)")
    println()
end
print(io::IO, d::Contact{T}) where {T} = print(io, "Contact $(d.id) - $(d.potential) V - $(d.material.name)")
show(io::IO, d::Contact) = print(io, d)
show(io::IO,::MIME"text/plain", d::Contact) = show(io, d)


@testset "DeviceModel Tests" begin
    @test_throws ArgumentError DeviceModel(PSY.ThermalGen, PSI.ThermalStandardUnitCommitment)
    @test_throws ArgumentError DeviceModel(PSY.ThermalStandard, PSI.AbstractDeviceFormulation)
end

@testset "OperationsProblem Tests" begin
    for p in [true, false]
        t = OperationsProblem(TestOpProblem, CopperPlatePowerModel, c_sys5; use_parameters = p)
        @test PSI.model_has_parameters(t.canonical) == p
    end
end


#!/cvmfs/soft.computecanada.ca/easybuild/software/2017/avx512/Compiler/gcc7.3/julia/1.1.0/bin/julia
#SBATCH -o four_rooms_var_red.out # Standard output
#SBATCH -e four_rooms_var_red.err # Standard error
#SBATCH --mem-per-cpu=2000M # Memory request of 2 GB
#SBATCH --time=12:00:00 # Running time of 6 hours
#SBATCH --ntasks=64
#SBATCH --account=def-whitem

using Pkg
Pkg.activate(".")

using Reproduce
using Logging

const save_loc = "four_rooms_cont_exp_int_study_rmsprop"
const exp_file = "experiment/four_rooms_cont_int.jl"
const exp_module_name = :FourRoomsContIntExperiment
const exp_func_name = :main_experiment
const alphas = [0.0, 0.00001, 0.00005, 0.0001, 0.0005, 0.001, 0.005, 0.01]
const betas = collect(0.0:0.1:1.0)
const policies = ["random_state_variant", "random_state_weight_variant", "uniform"]
const gvfs = ["collide_down", "favored_down"]
const batchsizes = [16]
const train_gaps = [1, 2, 3, 4, 8, 16, 32, 48, 64, 128, 192, 256]
const warm_up = 1000
const buffersize = 15000
const numsteps = 250000

function make_arguments(args::Dict{String, String})
    new_args=["--policy", args["policy"],
              "--gvf", args["gvf"],
              "--train_gap", args["train_gap"],
              "--batchsize", args["batchsize"],
              "--beta", args["beta"],
              "--run", args["run"],
              ]
    return new_args
end

function main()

    as = ArgParseSettings()
    @add_arg_table as begin
        "numworkers"
        arg_type=Int64
        default=1
        "--numruns"
        arg_type=Int64
        default=10
        "--saveloc"
        arg_type=String
        default=save_loc
        "--numjobs"
        action=:store_true
    end
    parsed = parse_args(as)
    num_workers = parsed["numworkers"]

    arg_dict = Dict([
        "policy"=>policies,
        "gvf"=>gvfs,
        "train_gap"=>train_gaps,
        "batchsize"=>batchsizes,
        "beta"=>betas,
        "run"=>1:parsed["numruns"]
    ])
    arg_list = ["policy", "gvf", "beta", "train_gap", "batchsize", "run"]

    static_args = [["--exp_loc", parsed["saveloc"],
                    "--warm_up", string(warm_up),
                    "--buffersize", string(buffersize),
                    "--seed", "0",
                    "--numinter", string(numsteps),
                    "--eval_points", "100",
                    "--eval_steps", "100",
                    "--compress",
                    "--opt", "RMSProp"
                    "--alphas"];
                   string.(alphas)]
    args_iterator = ArgIterator(arg_dict, static_args; arg_list=arg_list, make_args=make_arguments)

    if parsed["numjobs"]
        @info "This experiment has $(length(collect(args_iterator))) jobs."
        println(collect(args_iterator)[num_workers])
        exit(0)
    end

    experiment = Experiment(save_loc,
                            exp_file,
                            exp_module_name,
                            exp_func_name,
                            args_iterator)

    create_experiment_dir(experiment)
    add_experiment(experiment; settings_dir="settings")
    ret = job(experiment; num_workers=num_workers)
    post_experiment(experiment, ret)
end

main()


abstract MonteCarloScenarios
immutable FixedMonteCarloScenarios <: MonteCarloScenarios
  scenarios::Reservoir
end
immutable GeneratedMonteCarloScenarios <: MonteCarloScenarios
  method::Symbol
end

abstract ReportEvaluator
immutable CrossValidatedFeasibilityProbability <: ReportEvaluator end
immutable InverseModel <: ReportEvaluator end
immutable MonteCarloFeasibilityProbability <: ReportEvaluator
  scenarios::MonteCarloScenarios
end

"""
Prints a complete solution report for a reservoir, with multiple solutions and their evaluation.

Mandatory arguments:

  * `reservoir`: the reservoir to study
  * `ruleDuration`: the duration of the rule curves to produce
  * `guarantee`: the duration to guarantee for those rule curves

Optional arguments for input data:

  * `currentRuleCurve`: the existing rule curve, used as a benchmark for the newly computed ones.
    Its periodicity must be that of `reservoir`.
  * `climateChangeScenarios`: a set of inflow scenarios for some climate change, associating a description
    to the inflow. There is no distinction between the rivers in this argument!

Optional arguments to select the solvers and evaluators to use:

  * `scenarioGenerationAlgorithms`: algorithms used to generate scenarios (see Scenario generation module, the symbols
    are the argument to `scenarioGeneration`).
    Default: `[:Merge, :Mix]`
  * `ciAlgorithmsLevels`: algorithms used to compute the confidence intervals, associated with their confidence values
    (see Statistical module, the symbols are the argument to `generateScenarios`).
    Default: `Dict(:TStudent => [.95])`
  * `solvers`: solvers that are used to generate solutions (see Solver module, the symbols correspond to the function
    names of the solvers).
    Default: `[:minimumRuleCurve, :safeMinimumRuleCurve]`
  * `evaluators`: algorithms that are used to evaluate the new rule curves (see Solver module, the symbols correspond
    to the function names of the solvers).
    Default: `[]`
    TODO: not really flexible (think of parameters).

    The current list of possible evaluators is:

    * `:feasibilityProbabilitySyntheticScenarios`: computes the feasibility probability on synthetic scenarios
      (generated based on the statistical analysis of the data sets)
    * `:feasibilityProbabilityAllDatasets`: computes the feasibility probability on historical scenarios (all the given
       data sets), i.e. resubstitution
    * `:feasibilityProbabilityClimateChange`: computes the feasibility probability on provided climate change scenarios
      (where all rivers are merged)
    * `:hydropower`: evaluates the hydropower potential (with a generated price scenario)
      TODO: price scenario from outside.
    * `:worstCaseFunctions`:
    * `:maximumDeviation`:

  * `globalEvaluators`: algorithms that are used to evaluate algorithms that generate the rule curves, and not the
    rule curves themselves (see Solver module, the symbols correspond to the function names of the solvers).
    Default: `[]`
    TODO: not really flexible (think of parameters).
  * `disabledCombinations`: combinations of the above parameters that should *not* be run. This list contains pairs
    of a scenario generation algorithm and a solver. It is not possible to disable confidence intervals this way.
    Default: `[]`.
    Recommended value: `[(:Mix, :safeMinimumRuleCurve)]`, as this combination takes a lot of time to compute.
  * `verbosity`: passed as is to the solvers (`-1` to disable all output)
  * `saveFolder`: folder name to use with the following options
  * `savePlots`: whether the script should prepare some plots
  * `saveIntermediate`: whether the script should save the outputs for future reuse

TODO: how to add new solvers/evaluators? Work with a hash table and store lambdas with a fixed signature?
TODO: how to deal with solvers/evaluators parameters? Still via lambdas? Within the hash table?
TODO: provide more types to avoid typing Tuple{Tuple} too often. Use the defined data structures.
"""
function report(reservoir::Reservoir;
                ruleDuration::Period=Year(1), guarantee::Period=Year(2), ciPeriod::Period=Week(1),
                currentRuleCurve::Array{Float64, 1}=zeros(Float64, 0),
                climateChangeScenarios::Dict{Symbol, Array{TimeSeries.TimeArray{Float64,1,DateTime,Array{Float64,1}}, 1}}=Dict{Symbol, Array{TimeSeries.TimeArray{Float64,1,DateTime,Array{Float64,1}}, 1}}(),

                scenarioGenerationAlgorithms::Array{Symbol, 1}=Symbol[:Merge, :Mix],
                ciAlgorithmsLevels::Dict{Symbol, Array{Float64, 1}}=Dict{Symbol, Array{Float64, 1}}(:TStudent => [.95, .99]),
                solvers::Array{Symbol, 1}=[:minimumRuleCurve, :safeMinimumRuleCurve],
                evaluators::Array{Symbol, 1}=Symbol[],#:feasibilityProbabilityAllDatasets,:feasibilityProbabilityClimateChange], # :feasibilityProbability, # , :hydropower # , :worstCaseFunctions, :maximumDeviation
                globalEvaluators::Array{Symbol, 1}=Symbol[],#:crossValidatedFeasibilityProbability],
                disabledCombinations::Array{Tuple{Symbol, Symbol}, 1}=Tuple{Symbol, Symbol}[],#(:Mix, :safeMinimumRuleCurve)],
                saveFolder::AbstractString="", savePlots::Bool=false, saveIntermediate::Bool=false,
                verbosity::Int=1)
  # Some assumptions on the durations.
  @assert typeof(ruleDuration) <: Year
  @assert typeof(guarantee) <: Year

  # Ensure the output folder is ready.
  if length(saveFolder) > 0 && ! isdir(saveFolder)
    mkdir(saveFolder)
  end

  ### Generate parameters.
  hasCurrentRuleCurve = ! isempty(currentRuleCurve)

  generateDefaultLengthScenarios = in(:minimumRuleCurve, solvers) # Default length scenarios (i.e. guarantee).
  generateExtraLengthScenarios = in(:safeMinimumRuleCurve, solvers) # Longer scenarios (i.e. ruleDuration + guarantee).
  scenariosDefaultLength = Int(guarantee)
  scenariosExtraLength = Int(ruleDuration) + Int(guarantee)

  ruleLengthPeriods = Int(floor(Dates.days(ruleDuration) / Dates.days(getPeriod(reservoir))))
  guaranteePeriods = Int(floor(Dates.days(guarantee) / Dates.days(getPeriod(reservoir))))

  generateCIScenarios = length(ciAlgorithmsLevels) > 0
  ciAlgorithms = [algo[1] for algo in ciAlgorithmsLevels]

  # TODO: COMPLETE REFACTORING OF THE WHOLE FUNCTION! TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO
  # A complete data structure of available solvers, presented with a uniform syntax:
  #     (reservoir, verbosity) -> f(...)::ReservoirSolution
  # TODO: Not yet useable throughout this script (just for cross-validation): for the solving process, needs information for
  # TODO:   the solution name (could just currify the functions for cross-validation and let that string as a parameter)
  # TODO:       (reservoir, verbosity, name) -> f(...)::ReservoirSolution
  # TODO:   however, this is still not enough, as minimumRuleCurve and safeMinimumRuleCurve do not take their scenarios from
  # TODO:   the same place (scgen_default, scgen_extra). Would need a more complex data structure to hold the solvers?
  # solversF = Dict{Symbol, Function}
  # for solver in solvers
  #   if solver == :minimumRuleCurve
  #     solversF[solver] = (sc, verbosity) -> minimumRuleCurve(sc, "", outputLength=ruleLengthPeriods, verbosity=verbosity)
  #   elseif solver == :safeMinimumRuleCurve
  #     solversF[solver] = (sc, verbosity) -> safeMinimumRuleCurve(sc, ruleLengthPeriods, guaranteePeriods, "", verbosity=verbosity)
  #   else
  #     error("Unrecognised solver: " * string(solver))
  #   end
  # end
  # TODO: COMPLETE REFACTORING OF THE WHOLE FUNCTION! TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO TODO

  ### Handle scenario generation.
  if verbosity >= 0; println("[=*=] Scenario generation: started... "); end
  cis = Dict{Tuple{Symbol, Float64}, Tuple{Reservoir, Reservoir, Reservoir}}()
  scgen_default = Dict{Symbol, Reservoir}()
  scgen_extra = Dict{Symbol, Reservoir}()

  # Confidence interval computations if needed.
  for algo in ciAlgorithms
    for level in ciAlgorithmsLevels[algo]
      cis[(algo, level)] = computeCI(reservoir, level, ci_period=ciPeriod, method=algo)
    end
  end

  # Actual scenario generation.
  if generateDefaultLengthScenarios
    for algo in scenarioGenerationAlgorithms
      scgen_default[algo] = scenarioGeneration(reservoir, algo, scenariosDefaultLength)
    end
  end

  if generateExtraLengthScenarios
    for algo in scenarioGenerationAlgorithms
      scgen_extra[algo] = scenarioGeneration(reservoir, algo, scenariosExtraLength)
    end
  end

  if generateCIScenarios
    for algo in ciAlgorithms
      for level in ciAlgorithmsLevels[algo]
        s = symbol("ci_", algo, "_", level)
        if generateDefaultLengthScenarios
          scgen_default[s] = scenarioGeneration(cis[(algo, level)][2], :Duplicate, scenariosDefaultLength)
        end
        if generateExtraLengthScenarios
          scgen_extra[s] = scenarioGeneration(cis[(algo, level)][2], :Duplicate, scenariosExtraLength)
        end
      end
    end
  end

  if verbosity >= 0; println("[=*=] Scenario generation: done! "); end

  ### Handle solving.
  if verbosity >= 0; println("[=*=] Solving: started... "); end
  solutions = Dict{Tuple{Symbol, Symbol}, ReservoirSolution}()
  if hasCurrentRuleCurve
    solutions[(:currentRuleCurve, :currentRuleCurve)] = wrapExistingSolution(reservoir, currentRuleCurve)
  end

  for solver in solvers
    if verbosity >= 0; println("[=*=]   Solving with the algorithm: " * string(solver)); end

    if solver == :minimumRuleCurve
      for (lbl, sc) in scgen_default
        if ! in((lbl, solver), disabledCombinations)
          if verbosity >= 0; println("[=*=]     Solving with the data set: " * string(lbl)); end
          sol_lbl = "stochastic, default solver, " * string(guarantee) * ", " * string(lbl)
          solutions[(solver, lbl)] = minimumRuleCurve(sc, sol_lbl, outputLength=ruleLengthPeriods, useOutputs=getOutputIds(sc), verbosity=verbosity)
        end
      end
    elseif solver == :safeMinimumRuleCurve
      for (lbl, sc) in scgen_extra
        if ! in((lbl, solver), disabledCombinations)
          if verbosity >= 0; println("[=*=]     Solving with the data set: " * string(lbl)); end
          sol_lbl = "stochastic, safety solver, " * string(guarantee) * " + " * string(ruleDuration) * ", " * string(lbl)
          solutions[(solver, lbl)] = safeMinimumRuleCurve(sc, ruleLengthPeriods, guaranteePeriods, sol_lbl, verbosity=verbosity)
        end
      end
    else
      error("Unrecognised solver: " * string(solver))
    end
  end

  if verbosity >= 0; println("[=*=] Solving: done! "); end

  ## Check the consistency of the solutions.
  # All must have the same number of timesteps.
  sol_n_ts = -1 # Will be initialised to a positive value when the first feasible solution is found.
  sol_ts_pb = Dict{Any, Int}()
  for (desc, solution) in solutions
    if solution.feasible
      l = countTimeSteps(solution)
      if sol_n_ts < 0
        sol_n_ts = l
      elseif sol_n_ts != l
        sol_ts_pb[desc] = l
      end
    end
  end
  if sol_n_ts > 0 && length(sol_ts_pb) > 0
    txt_sol = "A solution has"
    if length(sol_ts_pb) > 1
      txt_sol = string(length(sol_ts_pb)) * " solutions have"
    end
    error(txt_sol * " an inconsistent number of time steps: " * string(sol_ts_pb))
  end

  # Other tests to perform?

  if verbosity >= 1; println("[=*=] Solving: all solutions are consistent! "); end

  ### Handle evaluation.
  if verbosity >= 0; println("[=*=] Evaluating: started... "); end

  evaluations = Dict{Tuple{Tuple{Symbol, Symbol}, Symbol}, ReservoirSolution}() # ((solver, scenarioGenerator), evaluator) -> evaluation
  # Associates (solutionAlgorithm, evaluator) to the result of an evaluation.
  # A solutionAlgorithm is a pair (solver, scenarioGenerator).
  for evaluator in evaluators
    if verbosity >= 0; println("[=*=]   Evaluating with the algorithm: " * string(evaluator)); end

    for (desc, solution) in solutions
      if ! solution.feasible
        continue
      end

      solver = desc[1]
      sc_gen = desc[2]

      if verbosity >= 0; println("[=*=]     Evaluating the solution: solver " * string(solver) * ", scenario generation " * string(sc_gen)); end

      if evaluator == :feasibilityProbabilitySyntheticScenarios
        evaluations[(desc, evaluator)] = feasibilityProbability(solution, numberScenarios=1000, scenarioDuration=guarantee, useOutputs=getOutputIds(getReservoir(solution)), verbosity=verbosity)
      elseif evaluator == :feasibilityProbabilityAllDatasets
        if verbosity >= 0; println("[=*=]     Evaluating with default length data (" * string(guarantee) * ")."); end
        for (lbl, sc) in scgen_default
         evaluations[(desc, symbol(evaluator, "_default_", lbl))] = feasibilityProbability(solution, scenarios=sc, useOutputs=getOutputIds(getReservoir(solution)), verbosity=verbosity)
        end
        if verbosity >= 0; println("[=*=]     Evaluating with extra length data (" * string(ruleDuration + guarantee) * ")."); end
        for (lbl, sc) in scgen_extra
          evaluations[(desc, symbol(evaluator, "_extra_", lbl))] = feasibilityProbability(solution, scenarios=sc, useOutputs=getOutputIds(getReservoir(solution)), verbosity=verbosity)
        end
      elseif evaluator == :feasibilityProbabilityClimateChange
        if verbosity >= 0; println("[=*=]     Evaluating with climate change scenarios."); end
        if isempty(climateChangeScenarios); error("No climate change scenarios given, while the script was asked to test against them!"); end;
        for (lbl, sc) in climateChangeScenarios
          evaluations[(desc, symbol(evaluator, "_", lbl))] = feasibilityProbability(solution, scenarios=sc, useOutputs=getOutputIds(getReservoir(solution)), verbosity=verbosity)
        end
      elseif evaluator == :crossValidatedFeasibilityProbability
        error("Cross-validation is not a solution evaluator! ")
      elseif evaluator == :worstCaseFunctions
        evaluations[(desc, evaluator)] = worstCaseFunctions(solution, useOutputs=getOutputIds(getReservoir(solution)), verbosity=verbosity)
      elseif evaluator == :maximumDeviation
        inflow = vec(getSolution(worstCaseFunctions(solutions[(:currentRuleCurve, :currentRuleCurve)], useOutputs=getOutputIds(getReservoir(solution)), verbosity=verbosity)))
        evaluations[(desc, evaluator)] = maximumDeviation(solution, inflow, verbosity=verbosity)
      elseif evaluator == :hydropower
        if solver == :minimumRuleCurve
          price_scenario = [sin(2 * π / 20 * t) for t in 1:scenariosDefaultLength]
        elseif solver == :safeMinimumRuleCurve
          price_scenario = [sin(2 * π / 30 * t) for t in 1:scenariosExtraLength]
        else
          error("Unrecognised solver for generating price scenarios: " * string(solver))
        end
        evaluations[(desc, evaluator)] = hydropower(solution, useOutputs=getOutputIds(getReservoir(solution)), electricityPrices=price_scenario)
      else
        error("Unrecognised evaluator: " * string(evaluator))
      end
    end
  end

  globalEvaluations = Dict{Tuple{Symbol, Symbol, Symbol}, ReservoirSolution}() # (solver, sc_gen, globalEvaluator) -> evaluation result
  for globalEvaluator in globalEvaluators
    if verbosity >= 0; println("[=*=]   Globally evaluating with the algorithm: " * string(globalEvaluator)); end

    for solver in solvers
      if verbosity >= 0; println("[=*=]     Evaluating solver with " * string(globalEvaluator) * ": " * string(solver)); end

      if globalEvaluator == :crossValidatedFeasibilityProbability
        if solver == :minimumRuleCurve
          for sc_gen in scenarioGenerationAlgorithms
            if in((sc_gen, solver), disabledCombinations)
              continue
            end

            if verbosity >= 0; println("[=*=]     Evaluating solver with " * string(globalEvaluator) * ": " * string(solver) * ", scenario generation " * string(sc_gen)); end

            f = function (sc, verbosity)
              scgen = scenarioGeneration(sc, sc_gen, scenariosDefaultLength)
              return minimumRuleCurve(scgen, "", useOutputs=getOutputIds(sc), outputLength=ruleLengthPeriods, verbosity=verbosity)
            end
            globalEvaluations[(solver, sc_gen, globalEvaluator)] = crossValidatedFeasibilityProbability(reservoir, f, verbosity=verbosity, folds=10)
          end
          for algo in ciAlgorithms
            for level in ciAlgorithmsLevels[algo]
              sc_gen = symbol("ci_", algo, "_", level)

              if in((sc_gen, solver), disabledCombinations)
                continue
              end

              if verbosity >= 0; println("[=*=]     Evaluating solver with " * string(globalEvaluator) * ": " * string(solver) * ", scenario generation " * string(sc_gen)); end

              f = function (sc, verbosity)
                cis = computeCI(sc, level, ci_period=ciPeriod, method=algo)[2]
                scgen = scenarioGeneration(cis, :Duplicate, scenariosDefaultLength)
                return minimumRuleCurve(scgen, "", useOutputs=getOutputIds(sc), outputLength=ruleLengthPeriods, verbosity=verbosity)
              end
              globalEvaluations[(solver, sc_gen, globalEvaluator)] = crossValidatedFeasibilityProbability(reservoir, f, verbosity=verbosity, folds=10)
            end
          end
        elseif solver == :safeMinimumRuleCurve
          for sc_gen in scenarioGenerationAlgorithms
            if in((sc_gen, solver), disabledCombinations)
              continue
            end

            if verbosity >= 0; println("[=*=]     Evaluating solver with " * string(globalEvaluator) * ": " * string(solver) * ", scenario generation " * string(sc_gen)); end

            f = function (sc, verbosity)
              scgen = scenarioGeneration(sc, sc_gen, scenariosExtraLength)
              return safeMinimumRuleCurve(scgen, ruleLengthPeriods, guaranteePeriods, "", verbosity=verbosity)
            end
            globalEvaluations[(solver, sc_gen, globalEvaluator)] = crossValidatedFeasibilityProbability(reservoir, f, verbosity=verbosity, folds=10)
          end
          for algo in ciAlgorithms
            for level in ciAlgorithmsLevels[algo]
              sc_gen = symbol("ci_", algo, "_", level)

              if in((sc_gen, solver), disabledCombinations)
                continue
              end

              if verbosity >= 0; println("[=*=]     Evaluating solver with " * string(globalEvaluator) * ": " * string(solver) * ", scenario generation " * string(sc_gen)); end

              f = function (sc, verbosity)
                cis = computeCI(sc, level, ci_period=ciPeriod, method=algo)[2]
                scgen = scenarioGeneration(cis, :Duplicate, scenariosExtraLength)
                return safeMinimumRuleCurve(scgen, ruleLengthPeriods, guaranteePeriods, "", verbosity=verbosity)
              end
              globalEvaluations[(solver, sc_gen, globalEvaluator)] = crossValidatedFeasibilityProbability(reservoir, f, verbosity=verbosity, folds=10)
            end
          end
        else
          error("Unrecognised solver: " * string(solver))
        end
      else
        error("Unrecognised global evaluator: " * string(solver))
      end
    end
  end

  if verbosity >= 0; println("[=*=] Evaluating: done! "); end

  ### Finally, display report, generate figures, and save what should be saved.
  if verbosity >= 1
    for desc in keys(solutions)
      println("**************************************************************************************************************")
      println("** Solver: " * string(desc[1]))
      println("** Scenario generation: " * string(desc[2]))
      println("**")
      for evaluator in evaluators
        # The generated keys for those solvers are not detected by the following condition!
        if evaluator == :feasibilityProbabilityAllDatasets
          for (eval_desc, evaluation) in evaluations
            if desc == eval_desc[1] && startswith(string(eval_desc[2]), "feasibilityProbabilityAllDatasets")
              ev = evaluation.solution
              first_part_length = length("feasibilityProbabilityAllDatasets_")
              dataset = string(eval_desc[2])[first_part_length + 1:end]
              println("** Feasibility analysis on data set " * dataset * ": probability of being infeasible: " * string(ev.probabilityInfeasible * 100) * "%")
              println("** Feasibility analysis on data set " * dataset * ": probability of being infeasible due to water surplus: " * string(ev.probabilityInfeasibleSurplus * 100) * "%")
              println("** Feasibility analysis on data set " * dataset * ": probability of being infeasible due to water shortage: " * string(ev.probabilityInfeasibleShortage * 100) * "%")
              println("** Feasibility analysis on data set " * dataset * ": average water shortage: " * string(ev.averageShortage) * " 10^6 m^3")
              println("** Feasibility analysis on data set " * dataset * ": average shortage index: " * string(ev.averageShortageIndex))
              println("** Feasibility analysis on data set " * dataset * ": average shortage vulnerability index: " * string(ev.averageShortageVulnerability))
            end
          end
        end
        if evaluator == :feasibilityProbabilityClimateChange
          for (eval_desc, evaluation) in evaluations
            if desc == eval_desc[1] && startswith(string(eval_desc[2]), "feasibilityProbabilityClimateChange")
              ev = evaluation.solution
              first_part_length = length("feasibilityProbabilityClimateChange_")
              dataset = string(eval_desc[2])[first_part_length + 1:end]
              println("** Feasibility analysis on the climate change scenario " * dataset * ": probability of being infeasible: " * string(ev.probabilityInfeasible * 100) * "%")
              println("** Feasibility analysis on the climate change scenario " * dataset * ": probability of being infeasible due to water surplus: " * string(ev.probabilityInfeasibleSurplus * 100) * "%")
              println("** Feasibility analysis on the climate change scenario " * dataset * ": probability of being infeasible due to water shortage: " * string(ev.probabilityInfeasibleShortage * 100) * "%")
              println("** Feasibility analysis on the climate change scenario " * dataset * ": average water shortage: " * string(ev.averageShortage) * " 10^6 m^3")
              println("** Feasibility analysis on the climate change scenario " * dataset * ": average shortage index: " * string(ev.averageShortageIndex))
              println("** Feasibility analysis on the climate change scenario " * dataset * ": average shortage vulnerability index: " * string(ev.averageShortageVulnerability))
            end
          end
        end

        if haskey(evaluations, (desc, evaluator)) # Not all evaluators could work previously (e.g. cross-validation needs more scenarios than available with CIs, with more parameters than a sheer symbol).
          ev = evaluations[(desc, evaluator)].solution
          if evaluator == :feasibilityProbabilitySyntheticScenarios
            println("** Monte-Carlo feasibility analysis: probability of being infeasible: " * string(ev.probabilityInfeasible * 100) * "%")
            println("** Monte-Carlo feasibility analysis: probability of being infeasible due to water surplus: " * string(ev.probabilityInfeasibleSurplus * 100) * "%")
            println("** Monte-Carlo feasibility analysis: probability of being infeasible due to water shortage: " * string(ev.probabilityInfeasibleShortage * 100) * "%")
            println("** Monte-Carlo feasibility analysis: average water shortage: " * string(ev.averageShortage) * " 10^6 m^3")
            println("** Monte-Carlo feasibility analysis: average shortage index: " * string(ev.averageShortageIndex))
            println("** Monte-Carlo feasibility analysis: average shortage vulnerability index: " * string(ev.averageShortageVulnerability))
          elseif evaluator == :worstCaseFunctions
            println("** Inverse model analysis: annual volume below which infeasibility is guaranteed: " * string(ev.total) * " 10^6 m^3")
          elseif evaluator == :maximumDeviation
            println("** Maximum deviation: from the inverse solution for the current rule curve: " * string(ev.deviation) * " 10^6 m^3")
          else
            error("Unrecognised evaluator: " * string(evaluator))
          end
        end
      end
      println("**************************************************************************************************************")
    end

    for globalEvaluator in globalEvaluators
      if globalEvaluator == :crossValidatedFeasibilityProbability
        for (globalEvaluationDescription, globalEvaluation) in globalEvaluations
          solver = globalEvaluationDescription[1]
          sc_gen = globalEvaluationDescription[2]
          evaluator = globalEvaluationDescription[3]
          if evaluator == :crossValidatedFeasibilityProbability
            ev = globalEvaluation.solution

            println("**************************************************************************************************************")
            println("** Solver: ", solver)
            println("** Scenario generation: ", sc_gen)
            println("** ")
            println("** Cross-validation feasibility analysis: probability of being infeasible: " * string(ev.probabilityInfeasible * 100) * "%")
            println("** Cross-validation feasibility analysis: probability of being infeasible due to water surplus: " * string(ev.probabilityInfeasibleSurplus * 100) * "%")
            println("** Cross-validation feasibility analysis: probability of being infeasible due to water shortage: " * string(ev.probabilityInfeasibleShortage * 100) * "%")
            println("** Cross-validation feasibility analysis: probability of having no rule curve: " * string(ev.probabilityNoRuleCurve * 100) * "%")
            println("** Cross-validation feasibility analysis: average water shortage: " * string(ev.averageShortage) * " 10^6 m^3")
            println("** Cross-validation feasibility analysis: average shortage index: " * string(ev.averageShortageIndex))
            println("** Cross-validation feasibility analysis: average shortage vulnerability index: " * string(ev.averageShortageVulnerability))
            println("**************************************************************************************************************")
          end
        end
      end
    end
  end

  if saveIntermediate
    if verbosity >= 0; println("[=*=] Serialising: started... "); end

    folder_intermediate = saveFolder * "/intermediate/"
    if ! isdir(folder_intermediate)
      mkdir(folder_intermediate)
    end
    serialise_own(filename, o) = open(folder_intermediate * filename, "w") do stream; serialize(stream, o); end

    serialise_own("cis.jls", cis)
    serialise_own("scgen_default.jls", scgen_default)
    serialise_own("scgen_extra.jls", scgen_extra)
    serialise_own("solutions.jls", solutions)
    serialise_own("evaluations.jls", evaluations)
    serialise_own("global_evaluations.jls", globalEvaluations)

    if verbosity >= 0; println("[=*=] Serialising: done! "); end
  end

  if savePlots
    if verbosity >= 0; println("[=*=] Plotting: started... "); end

    folder_plots = saveFolder * "/plots/"
    if ! isdir(folder_plots)
      mkdir(folder_plots)
    end

    function exportPlot(plot::AbstractPlot, name::AbstractString)
      savefig(plot, folder_plots * name * ".png")
      savefig(plot, folder_plots * name * ".svg")
      nothing
    end

    function exportData(data::Array{Any, 2}, name::AbstractString)
      writedlm(folder_plots * name * ".tsv", data, '\t')
      nothing
    end

    period = getPeriod(reservoir)

    ## Data treatment.
    # Confidence intervals. First plots for each algorithm, then a global plot.
    # Indices in the code: once the average [1], then for each pair upper/lower bounds [2] and [3].
    if ! isempty(currentRuleCurve) && length(ciAlgorithms) > 0
      # Plot each algorithm independently (if at least two algorithms), all confidence levels at once.
      if length(ciAlgorithms) > 1
        for algo in ciAlgorithms
          # All confidence levels at once.
          data = vcat(cis[(algo, ciAlgorithmsLevels[algo][1])][1], [let lcis = cis[(algo, level)]; [lcis[2], lcis[3]]; end for level in ciAlgorithmsLevels[algo]]...)
          plt, dat = plotRiverAlongReservoirs(data, period, 1:ruleLengthPeriods, getName(reservoir), exportPlotData=true)
          exportPlot(plt, "confidence_intervals_" * string(algo))
          exportData(dat, "confidence_intervals_" * string(algo))
        end
      end

      # Each confidence level in its own plot (if more than one level requested), for each algorithm.
      for algo in ciAlgorithms
        if length(ciAlgorithmsLevels[algo]) > 1
          for level in ciAlgorithmsLevels[algo]
            data = vcat(cis[(algo, ciAlgorithmsLevels[algo][1])][1], cis[(algo, level)][2], cis[(algo, level)][3])
            plt, dat = plotRiverAlongReservoirs(data, period, 1:ruleLengthPeriods, getName(reservoir), exportPlotData=true)
            exportPlot(plt, "confidence_intervals_" * string(algo) * "_cl" * string(level))
            exportData(dat, "confidence_intervals_" * string(algo) * "_cl" * string(level))
          end
        end
      end

      # TODO: plot for a given confidence level (irrespectively of algorithms).

      # Global plot for all algorithms and confidence levels at once.
      data = vcat([[cis[(algo, ciAlgorithmsLevels[algo][1])][1]; vcat([let lcis = cis[(algo, level)]; [lcis[2]; lcis[3]]; end for level in ciAlgorithmsLevels[algo]]...)] for algo in ciAlgorithms]...)
      plt, dat = plotRiverAlongReservoirs(data, period, 1:ruleLengthPeriods, getName(reservoir), exportPlotData=true)
      exportPlot(plt, "confidence_intervals")
      exportData(dat, "confidence_intervals")
    end

    # Skip scenario generation output.

    ## Solutions.
    # Existing target, alone.
    if ! isempty(currentRuleCurve)
      plt, dat = plotRule(currentRuleCurve, period, title="Current rule curve", exportPlotData=true)
      exportPlot(plt, "current_rule")
      exportData(dat, "current_rule")
    end

    # Complete set of scenarios when merging them and compare them to the solution.
    if in(:Merge, scenarioGenerationAlgorithms) && in(:minimumRuleCurve, solvers) && ! isempty(currentRuleCurve)
      merge_sol = solutions[(:minimumRuleCurve, :Merge)]
      colourings = [:YearlyAverage, :WetSeasonAverage, :DrySeasonAverage, :DriestMonth, :DriestThreeMonths, :DriestSixMonths]

      for colouring in colourings
        plt, dat = plotSolutionVsScenarios(merge_sol, currentRuleCurve, colouring=colouring, exportPlotData=true)
        exportPlot(plt, "solution_safety_stochastic_merge_splitperscenario_" * string(colouring))
        exportData(dat, "solution_safety_stochastic_merge_splitperscenario_" * string(colouring))
      end
    end

    # Solutions for the various solvers: one plot for each solver first, then one plot with each and every solution.
    if ! isempty(currentRuleCurve)
      for solver in solvers
        solver_solution = filter((k, v) -> k[1] == solver, solutions)
        plt, dat = plotSolutions(collect(values(solver_solution)), currentRuleCurve, exportPlotData=true)
        exportPlot(plt, "solutions_" * string(solver))
        exportData(dat, "solutions_" * string(solver))
      end

      plt, dat = plotSolutions(collect(values(solutions)), currentRuleCurve, exportPlotData=true)
      exportPlot(plt, "solutions")
      exportData(dat, "solutions")
    end

    ## Evaluation.
    # Nothing for solution evaluators.
    # Nothing for global evaluators.

    if verbosity >= 0; println("[=*=] Plotting: done! "); end
  end

  return cis, scgen_default, scgen_extra, solutions, evaluations
end


#### EDMF model kernels

using CLIMAParameters.Planet: e_int_v0, grav, day, R_d, R_v, molmass_ratio
using Printf
using ClimateMachine.Atmos: nodal_update_auxiliary_state!

using ClimateMachine.BalanceLaws: number_states

using ClimateMachine.MPIStateArrays: MPIStateArray
using ClimateMachine.DGMethods: LocalGeometry, DGModel

import ClimateMachine.BalanceLaws:
    vars_state,
    update_auxiliary_state!,
    init_state_prognostic!,
    flux_first_order!,
    flux_second_order!,
    compute_gradient_argument!,
    compute_gradient_flux!

import ClimateMachine.TurbulenceConvection:
    init_aux_turbconv!,
    turbconv_nodal_update_auxiliary_state!,
    turbconv_boundary_state!,
    turbconv_normal_boundary_flux_second_order!

using ClimateMachine.Thermodynamics: air_pressure, air_density

include(joinpath("helper_funcs", "nondimensional_exchange_functions.jl"))
include(joinpath("helper_funcs", "lamb_smooth_minimum.jl"))
include(joinpath("helper_funcs", "utility_funcs.jl"))
include(joinpath("helper_funcs", "subdomain_statistics.jl"))
include(joinpath("helper_funcs", "diagnose_environment.jl"))
include(joinpath("helper_funcs", "subdomain_thermo_states.jl"))
include(joinpath("helper_funcs", "save_subdomain_temperature.jl"))
include(joinpath("closures", "entr_detr.jl"))
include(joinpath("closures", "pressure.jl"))
include(joinpath("closures", "mixing_length.jl"))
include(joinpath("closures", "turbulence_functions.jl"))
include(joinpath("closures", "surface_functions.jl"))

function vars_state(m::NTuple{N, Updraft}, st::Auxiliary, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, ::Auxiliary, FT)
    @vars(
        buoyancy::FT,
        a::FT,
        ε_dyn::FT,
        δ_dyn::FT,
        ε_trb::FT,
        T::FT,
        θ_liq::FT,
        q_tot::FT,
        w::FT,
    )
end

function vars_state(::Environment, ::Auxiliary, FT)
    @vars(T::FT, cld_frac::FT, buoyancy::FT)
end

function vars_state(m::EDMF, st::Auxiliary, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Prognostic, FT)
    @vars(ρa::FT, ρaw::FT, ρaθ_liq::FT, ρaq_tot::FT,)
end

function vars_state(::Environment, ::Prognostic, FT)
    @vars(ρatke::FT, ρaθ_liq_cv::FT, ρaq_tot_cv::FT, ρaθ_liq_q_tot_cv::FT,)
end

function vars_state(m::NTuple{N, Updraft}, st::Prognostic, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Prognostic, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(::Updraft, ::Gradient, FT)
    @vars(w::FT,)
end

function vars_state(::Environment, ::Gradient, FT)
    @vars(
        θ_liq::FT,
        q_tot::FT,
        w::FT,
        tke::FT,
        θ_liq_cv::FT,
        q_tot_cv::FT,
        θ_liq_q_tot_cv::FT,
        θv::FT,
        e::FT,
    )
end

function vars_state(m::NTuple{N, Updraft}, st::Gradient, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(m::EDMF, st::Gradient, FT)
    @vars(
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

function vars_state(m::NTuple{N, Updraft}, st::GradientFlux, FT) where {N}
    return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}
end

function vars_state(::Updraft, st::GradientFlux, FT)
    @vars(∇w::SVector{3, FT},)
end

function vars_state(::Environment, ::GradientFlux, FT)
    @vars(
        ∇θ_liq::SVector{3, FT},
        ∇q_tot::SVector{3, FT},
        ∇w::SVector{3, FT},
        ∇tke::SVector{3, FT},
        ∇θ_liq_cv::SVector{3, FT},
        ∇q_tot_cv::SVector{3, FT},
        ∇θ_liq_q_tot_cv::SVector{3, FT},
        ∇θv::SVector{3, FT},
        ∇e::SVector{3, FT},
    )

end

function vars_state(m::EDMF, st::GradientFlux, FT)
    @vars(
        S²::FT, # should be conditionally grabbed from atmos.turbulence
        environment::vars_state(m.environment, st, FT),
        updraft::vars_state(m.updraft, st, FT)
    )
end

"""
    init_aux_turbconv!(
        turbconv::EDMF{FT},
        m::AtmosModel{FT},
        aux::Vars,
        geom::LocalGeometry,
    ) where {FT}

Initialize EDMF auxiliary variables.
"""
function init_aux_turbconv!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    aux::Vars,
    geom::LocalGeometry,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    en_aux = aux.turbconv.environment
    up_aux = aux.turbconv.updraft

    en_aux.cld_frac = FT(0)
    en_aux.buoyancy = FT(0)

    @unroll_map(N_up) do i
        up_aux[i].buoyancy = FT(0)
        up_aux[i].θ_liq = FT(0)
        up_aux[i].q_tot = FT(0)
        up_aux[i].w = FT(0)
    end
end;

function turbconv_nodal_update_auxiliary_state!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    save_subdomain_temperature!(m, state, aux)

    en_aux = aux.turbconv.environment
    up_aux = aux.turbconv.updraft
    gm = state
    en = state.turbconv.environment
    up = state.turbconv.updraft

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    # Compute buoyancies of subdomains
    ρ_inv = 1 / gm.ρ
    _grav::FT = grav(m.param_set)

    z = altitude(m, aux)

    ρ_en = air_density(ts.en)
    en_aux.buoyancy = -_grav * (ρ_en - aux.ref_state.ρ) * ρ_inv

    @unroll_map(N_up) do i
        ρ_i = air_density(ts.up[i])
        up_aux[i].buoyancy = -_grav * (ρ_i - aux.ref_state.ρ) * ρ_inv
        up_aux[i].a = up[i].ρa * ρ_inv
        up_aux[i].θ_liq = up[i].ρaθ_liq / up[i].ρa
        up_aux[i].q_tot = up[i].ρaq_tot / up[i].ρa
        up_aux[i].w = up[i].ρaw / up[i].ρa
    end
    b_gm = grid_mean_b(state, aux, N_up)

    # remove the gm_b from all subdomains
    @unroll_map(N_up) do i
        up_aux[i].buoyancy -= b_gm
    end
    en_aux.buoyancy -= b_gm

    εδ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end

    ε_dyn, δ_dyn, ε_trb = ntuple(i -> map(x -> x[i], εδ_up), 3)

    @unroll_map(N_up) do i
        up_aux[i].ε_dyn = ε_dyn[i]
        up_aux[i].δ_dyn = δ_dyn[i]
        up_aux[i].ε_trb = ε_trb[i]
    end

end;

function compute_gradient_argument!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)
    z = altitude(m, aux)

    # Aliases:
    up_tf = transform.turbconv.updraft
    en_tf = transform.turbconv.environment
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    @unroll_map(N_up) do i
        up_tf[i].w = up[i].ρaw / up[i].ρa
    end
    _grav::FT = grav(m.param_set)

    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)
    q_tot_en = total_specific_humidity(ts.en)

    # populate gradient arguments
    en_tf.θ_liq = θ_liq_en
    en_tf.q_tot = q_tot_en
    en_tf.w = env.w

    en_tf.tke = enforce_positivity(en.ρatke) / (env.a * gm.ρ)
    en_tf.θ_liq_cv = enforce_positivity(en.ρaθ_liq_cv) / (env.a * gm.ρ)
    en_tf.q_tot_cv = enforce_positivity(en.ρaq_tot_cv) / (env.a * gm.ρ)
    en_tf.θ_liq_q_tot_cv = en.ρaθ_liq_q_tot_cv / (env.a * gm.ρ)

    # TODO: is this supposed to be grabbed from grid mean?
    en_tf.θv = virtual_pottemp(ts.gm)
    e_kin = FT(1 // 2) * ((gm.ρu[1] * ρ_inv)^2 + (gm.ρu[2] * ρ_inv)^2 + env.w^2)
    en_tf.e = total_energy(e_kin, _grav * z, ts.en)
end;

function compute_gradient_flux!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    up_dif = diffusive.turbconv.updraft
    tc_dif = diffusive.turbconv
    up_∇tf = ∇transform.turbconv.updraft
    en_dif = diffusive.turbconv.environment
    en_∇tf = ∇transform.turbconv.environment

    @unroll_map(N_up) do i
        up_dif[i].∇w = up_∇tf[i].w
    end

    ρ_inv = 1 / gm.ρ
    # first moment grid mean coming from environment gradients only
    en_dif.∇θ_liq = en_∇tf.θ_liq
    en_dif.∇q_tot = en_∇tf.q_tot
    en_dif.∇w = en_∇tf.w
    # second moment env cov
    en_dif.∇tke = en_∇tf.tke
    en_dif.∇θ_liq_cv = en_∇tf.θ_liq_cv
    en_dif.∇q_tot_cv = en_∇tf.q_tot_cv
    en_dif.∇θ_liq_q_tot_cv = en_∇tf.θ_liq_q_tot_cv

    en_dif.∇θv = en_∇tf.θv
    en_dif.∇e = en_∇tf.e

    tc_dif.S² = ∇transform.u[3, 1]^2 + ∇transform.u[3, 2]^2 + en_dif.∇w[3]^2
end;

function turbconv_source!(
    m::AtmosModel{FT},
    source::Vars,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
    direction,
) where {FT}
    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    en = state.turbconv.environment
    up = state.turbconv.updraft
    en_src = source.turbconv.environment
    up_src = source.turbconv.updraft
    en_dif = diffusive.turbconv.environment
    up_aux = aux.turbconv.updraft

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    εδ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end
    ε_dyn, δ_dyn, ε_trb = ntuple(i -> map(x -> x[i], εδ_up), 3)

    # get environment values
    _grav::FT = grav(m.param_set)
    ρ_inv = 1 / gm.ρ
    θ_liq_en = liquid_ice_pottemp(ts.en)
    q_tot_en = total_specific_humidity(ts.en)
    tke_en = enforce_positivity(en.ρatke) * ρ_inv / env.a
    θ_liq = liquid_ice_pottemp(ts.gm)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max

    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end

    @unroll_map(N_up) do i

        ρa_up_i = ρa_up[i]
        w_up_i = up[i].ρaw / ρa_up_i
        ρa_up_i_inv = FT(1) / ρa_up_i

        # first moment sources - for now we compute these as aux variable
        dpdz = perturbation_pressure(
            m,
            m.turbconv.pressure,
            state,
            diffusive,
            aux,
            t,
            env,
            i,
        )

        # entrainment and detrainment
        up_src[i].ρa += up[i].ρaw * (ε_dyn[i] - δ_dyn[i])
        up_src[i].ρaw +=
            up[i].ρaw *
            ((ε_dyn[i] + ε_trb[i]) * env.w - (δ_dyn[i] + ε_trb[i]) * w_up_i)
        up_src[i].ρaθ_liq +=
            up[i].ρaw * (
                (ε_dyn[i] + ε_trb[i]) * θ_liq_en -
                (δ_dyn[i] + ε_trb[i]) * up[i].ρaθ_liq * ρa_up_i_inv
            )
        up_src[i].ρaq_tot +=
            up[i].ρaw * (
                (ε_dyn[i] + ε_trb[i]) * q_tot_en -
                (δ_dyn[i] + ε_trb[i]) * up[i].ρaq_tot * ρa_up_i_inv
            )

        # add buoyancy and perturbation pressure in subdomain w equation
        up_src[i].ρaw += up[i].ρa * (up_aux[i].buoyancy - dpdz)
        # microphysics sources should be applied here

        # environment second moments:
        en_src.ρatke += (
            up[i].ρaw *
            δ_dyn[i] *
            (w_up_i - env.w) *
            (w_up_i - env.w) *
            FT(0.5) +
            up[i].ρaw *
            ε_trb[i] *
            (env.w - gm.ρu[3] * ρ_inv) *
            (env.w - w_up_i) - up[i].ρaw * (ε_dyn[i] + ε_trb[i]) * tke_en
        )

        en_src.ρaθ_liq_cv += (
            up[i].ρaw *
            δ_dyn[i] *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) +
            up[i].ρaw *
            ε_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) +
            up[i].ρaw *
            ε_trb[i] *
            (θ_liq_en - θ_liq) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) -
            up[i].ρaw * (ε_dyn[i] + ε_trb[i]) * en.ρaθ_liq_cv
        )

        en_src.ρaq_tot_cv += (
            up[i].ρaw *
            δ_dyn[i] *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) +
            up[i].ρaw *
            ε_trb[i] *
            (q_tot_en - gm.moisture.ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) +
            up[i].ρaw *
            ε_trb[i] *
            (q_tot_en - gm.moisture.ρq_tot * ρ_inv) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) -
            up[i].ρaw * (ε_dyn[i] + ε_trb[i]) * en.ρaq_tot_cv
        )

        en_src.ρaθ_liq_q_tot_cv += (
            up[i].ρaw *
            δ_dyn[i] *
            (up[i].ρaθ_liq * ρa_up_i_inv - θ_liq_en) *
            (up[i].ρaq_tot * ρa_up_i_inv - q_tot_en) +
            up[i].ρaw *
            ε_trb[i] *
            (θ_liq_en - θ_liq) *
            (q_tot_en - up[i].ρaq_tot * ρa_up_i_inv) +
            up[i].ρaw *
            ε_trb[i] *
            (q_tot_en - gm.moisture.ρq_tot * ρ_inv) *
            (θ_liq_en - up[i].ρaθ_liq * ρa_up_i_inv) -
            up[i].ρaw * (ε_dyn[i] + ε_trb[i]) * en.ρaθ_liq_q_tot_cv
        )

        # pressure tke source from the i'th updraft
        en_src.ρatke += ρa_up_i * (w_up_i - env.w) * dpdz
    end
    l_mix = mixing_length(
        m,
        m.turbconv.mix_len,
        state,
        diffusive,
        aux,
        t,
        δ_dyn,
        ε_trb,
        ts,
        env,
    )
    K_eddy = m.turbconv.mix_len.c_m * l_mix * sqrt(tke_en)
    Shear² = diffusive.turbconv.S²
    ρa₀ = gm.ρ * env.a
    Diss₀ = m.turbconv.mix_len.c_d * sqrt(tke_en) / l_mix

    # production from mean gradient and Dissipation
    en_src.ρatke += ρa₀ * (K_eddy * Shear² - Diss₀ * tke_en)
    en_src.ρaθ_liq_cv +=
        ρa₀ *
        (K_eddy * en_dif.∇θ_liq[3] * en_dif.∇θ_liq[3] - Diss₀ * en.ρaθ_liq_cv)
    en_src.ρaq_tot_cv +=
        ρa₀ *
        (K_eddy * en_dif.∇q_tot[3] * en_dif.∇q_tot[3] - Diss₀ * en.ρaq_tot_cv)
    en_src.ρaθ_liq_q_tot_cv +=
        ρa₀ * (
            K_eddy * en_dif.∇θ_liq[3] * en_dif.∇q_tot[3] -
            Diss₀ * en.ρaθ_liq_q_tot_cv
        )
    # covariance microphysics sources should be applied here
end;

# # in the EDMF first order (advective) fluxes exist only in the grid mean (if <w> is nonzero) and the uprdafts
function flux_first_order!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    flux::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    # Aliases:
    gm = state
    up = state.turbconv.updraft
    up_flx = flux.turbconv.updraft
    N_up = n_updrafts(turbconv)

    ρ_inv = 1 / gm.ρ
    ẑ = vertical_unit_vector(m, aux)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max
    # in future GCM implementations we need to think about grid mean advection

    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end

    @unroll_map(N_up) do i
        ρa_i = ρa_up[i]
        up_flx[i].ρa = up[i].ρaw * ẑ
        w_up_i = up[i].ρaw / ρa_i
        up_flx[i].ρaw = up[i].ρaw * w_up_i * ẑ
        up_flx[i].ρaθ_liq = w_up_i * up[i].ρaθ_liq * ẑ
        up_flx[i].ρaq_tot = w_up_i * up[i].ρaq_tot * ẑ
    end
end;

# in the EDMF second order (diffusive) fluxes
# exist only in the grid mean and the environment
function flux_second_order!(
    turbconv::EDMF{FT},
    m::AtmosModel{FT},
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    aux::Vars,
    t::Real,
) where {FT}
    N_up = n_updrafts(turbconv)

    # Aliases:
    gm = state
    up = state.turbconv.updraft
    en = state.turbconv.environment
    gm_flx = flux
    en_flx = flux.turbconv.environment
    en_dif = diffusive.turbconv.environment

    # Recover thermo states
    ts = recover_thermo_state_all(m, state, aux)

    # Get environment variables
    env = environment_vars(state, aux, N_up)

    ρ_inv = FT(1) / gm.ρ
    _grav::FT = grav(m.param_set)
    z = altitude(m, aux)
    a_min = turbconv.subdomains.a_min
    a_max = turbconv.subdomains.a_max

    εδ_up = ntuple(N_up) do i
        entr_detr(m, m.turbconv.entr_detr, state, aux, t, ts, env, i)
    end

    ε_dyn, δ_dyn, ε_trb = ntuple(i -> map(x -> x[i], εδ_up), 3)

    l_mix = mixing_length(
        m,
        turbconv.mix_len,
        state,
        diffusive,
        aux,
        t,
        δ_dyn,
        ε_trb,
        ts,
        env,
    )
    tke_en = enforce_positivity(en.ρatke) / env.a * ρ_inv
    K_eddy = m.turbconv.mix_len.c_m * l_mix * sqrt(tke_en)

    #TotalFlux(ϕ) = Eddy_Diffusivity(ϕ) + MassFlux(ϕ)
    e_int = internal_energy(m, state, aux)

    e_kin = vuntuple(N_up) do i
        FT(1 // 2) * (
            (gm.ρu[1] * ρ_inv)^2 +
            (gm.ρu[2] * ρ_inv)^2 +
            (up[i].ρaw / up[i].ρa)^2
        )
    end
    e_tot_up = ntuple(i -> total_energy(e_kin[i], _grav * z, ts.up[i]), N_up)
    ρa_up = vuntuple(N_up) do i
        gm.ρ * enforce_unit_bounds(up[i].ρa * ρ_inv, a_min, a_max)
    end

    massflux_e = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.ρe * ρ_inv - e_tot_up[i]) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    massflux_q_tot = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.moisture.ρq_tot * ρ_inv - up[i].ρaq_tot / up[i].ρa) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    massflux_w = sum(
        vuntuple(N_up) do i
            up[i].ρa *
            ρ_inv *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / up[i].ρa) *
            (gm.ρu[3] * ρ_inv - up[i].ρaw / ρa_up[i])
        end,
    )

    # update grid mean flux_second_order
    ρe_sgs_flux = -gm.ρ * env.a * K_eddy * en_dif.∇e[3] + massflux_e
    ρq_tot_sgs_flux = -gm.ρ * env.a * K_eddy * en_dif.∇q_tot[3] + massflux_q_tot
    ρu_sgs_flux = -gm.ρ * env.a * K_eddy * en_dif.∇w[3] + massflux_w

    # for now the coupling to the dycore is commented out

    # gm_flx.ρe              += SVector{3,FT}(0,0,ρe_sgs_flux)
    # gm_flx.moisture.ρq_tot += SVector{3,FT}(0,0,ρq_tot_sgs_flux)
    # gm_flx.ρu              += SMatrix{3, 3, FT, 9}(
    #     0, 0, 0,
    #     0, 0, 0,
    #     0, 0, ρu_sgs_flux,
    # )

    ẑ = vertical_unit_vector(m, aux)
    # env second moment flux_second_order
    en_flx.ρatke = -gm.ρ * env.a * K_eddy * en_dif.∇tke[3] * ẑ
    en_flx.ρaθ_liq_cv = -gm.ρ * env.a * K_eddy * en_dif.∇θ_liq_cv[3] * ẑ
    en_flx.ρaq_tot_cv = -gm.ρ * env.a * K_eddy * en_dif.∇q_tot_cv[3] * ẑ
    en_flx.ρaθ_liq_q_tot_cv =
        -gm.ρ * env.a * K_eddy * en_dif.∇θ_liq_q_tot_cv[3] * ẑ
end;

# First order boundary conditions
function turbconv_boundary_state!(
    nf,
    bc::EDMFBCs,
    m::AtmosModel{FT},
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    aux⁻::Vars,
    bctype,
    t,
    state_int::Vars,
    aux_int::Vars,
) where {FT}

    turbconv = m.turbconv
    N_up = n_updrafts(turbconv)
    up = state⁺.turbconv.updraft
    en = state⁺.turbconv.environment
    gm = state⁺
    gm_a = aux⁺
    if bctype == 1 # bottom
        zLL = altitude(m, aux_int)
        a_up_surf,
        θ_liq_up_surf,
        q_tot_up_surf,
        θ_liq_cv,
        q_tot_cv,
        θ_liq_q_tot_cv,
        tke = subdomain_surface_values(
            turbconv.surface,
            turbconv,
            m,
            gm,
            gm_a,
            zLL,
        )

        @unroll_map(N_up) do i
            up[i].ρaw = FT(0)
            up[i].ρa = a_up_surf[i] * gm.ρ
            up[i].ρaθ_liq = up[i].ρa * θ_liq_up_surf[i]
            up[i].ρaq_tot = up[i].ρa * q_tot_up_surf[i]
        end
        a_en = environment_area(gm, gm_a, N_up)
        en.ρatke = gm.ρ * a_en * tke
        en.ρaθ_liq_cv = gm.ρ * a_en * θ_liq_cv
        en.ρaq_tot_cv = gm.ρ * a_en * q_tot_cv
        en.ρaθ_liq_q_tot_cv = gm.ρ * a_en * θ_liq_q_tot_cv
    end
end;

# The boundary conditions for second-order unknowns
function turbconv_normal_boundary_flux_second_order!(
    nf,
    bc::EDMFBCs,
    m::AtmosModel{FT},
    fluxᵀn::Vars,
    n⁻,
    state⁻::Vars,
    diff⁻::Vars,
    hyperdiff⁻::Vars,
    aux⁻::Vars,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    bctype,
    t,
    _...,
) where {FT}
    turbconv = m.turbconv
    N = n_updrafts(turbconv)
    up_flx = fluxᵀn.turbconv.updraft
    en_flx = fluxᵀn.turbconv.environment
    if bctype == 2 # top
        @unroll_map(N_up) do i
            up_flx[i].ρaw = -n⁻ * FT(0)
            up_flx[i].ρa = -n⁻ * FT(0)
            up_flx[i].ρaθ_liq = -n⁻ * FT(0)
            up_flx[i].ρaq_tot = -n⁻ * FT(0)
        end
        en_flx.∇tke = -n⁻ * FT(0)
        en_flx.∇e_int_cv = -n⁻ * FT(0)
        en_flx.∇q_tot_cv = -n⁻ * FT(0)
        en_flx.∇e_int_q_tot_cv = -n⁻ * FT(0)
    end
end;


m = rand(3, 3)
axs = (1:3, 1:3)
ty = Interpolations.BSplineInterpolation{Float64, ndims(m), typeof(m), BSpline{Linear{Throw{OnGrid}}}, typeof(axs)}
etty = Interpolations.Extrapolation{Float64, 2, Interpolations.BSplineInterpolation{Float64, 2, Matrix{Float64}, BSpline{Linear{Throw{OnGrid}}}, Tuple{UnitRange{Int64}, UnitRange{Int64}}}, BSpline{Linear{Throw{OnGrid}}}, Flat{Nothing}}
fieldsum(x) = mapreduce(y->getfield(x, y), +, ntuple(identity, nfields(x)))

@testset "Interpolations.BSplineInterpolation constructor gradient" begin
    zg = Zygote.gradient((x,y,z)->sum(ty(x,y,z)), m, axs, BSpline(Linear()))
    fd = grad(central_fdm(5,1),
              x->sum(ty(x, axs, BSpline(Linear()))),
              m)
    @test isapprox(zg[1], fd[1])
    @test zg[2:3] == (nothing, nothing)
end

@testset "Interpolations.BSplineInterpolation functor gradient" begin
    zg = Zygote.gradient((x,y,z,w)->sum(Interpolations.BSplineInterpolation(x,y,z,w)), Float64, m, BSpline(Linear()), axs)
    fd = grad(central_fdm(5,1),
              x->sum(Interpolations.BSplineInterpolation(Float64, x, BSpline(Linear()), axs)),
              m)
    @test isapprox(zg[2], fd[1])
    @test findall(==(nothing), zg) == [1, 3, 4]
end

@testset "Interpolations.Extrapolation constructor gradient" begin
    itp = ty(m, axs, BSpline(Linear()))
    zg = Zygote.gradient((x,y)->sum(etty(x,y)), itp, Flat())
    fd = grad(central_fdm(5,1),
              x->sum(etty(itp, Flat())),
              m)
    @test isapprox(zg[1], fd[1])
    @test findall(==(nothing), zg) == [2]
end

@testset "Interpolations.copy_with_padding function gradient" begin
    zg = Zygote.gradient((x,y,z)->sum(Interpolations.copy_with_padding(x,y,z)), Float64, m, BSpline(Linear()))
    fd = grad(central_fdm(5,1),
              x->sum(Interpolations.copy_with_padding(Float64, x, BSpline(Linear()))),
              m)
    @test isapprox(zg[2], fd[1])
    @test findall(==(nothing), zg) == [1, 3]
end

@testset "Interpolations.FilledExtrapolation constructor gradient" begin
    @test_broken begin
        itp = ty(m, axs, BSpline(Linear()))
        zg = Zygote.gradient((x,y)->sum(Interpolations.FilledExtrapolation(x,y)), itp, 0.0)
        fd = grad(central_fdm(5,1),
                x->sum(Interpolations.FilledExtrapolation(x, 0.0)),
                itp)
        @test_broken isapprox(zg[1], fd[1])
        @test_broken findall(==(nothing), zg) == [2]
    end # TODO: Remove the @test_broken after tests work.
end

@testset "SVector and SMatrix gradient" begin
    for t in (Float32, Float64, RGB{Float32}, RGB{Float64})
        inp = rand(t, 2)
        inp_mat = rand(t, 3, 3)
        _abs(c)  = mapreducec(v->abs(float(v)), +, 0, c)
        @test Zygote.gradient(x->_abs(sum(SVector{2, t}(x))), inp)[1] == ones(t, 2)
        @test Zygote.gradient(x->_abs(sum(SMatrix{3, 3, t, 9}(x))), inp_mat)[1] == ones(t, 3, 3)
    end
end

@testset "DiffImages.Homography method gradient" begin
    h = DiffImages.Homography{Float64}()
    v = rand(SVector{2, Float64})
    zy = Zygote.gradient((x,y)->sum(y(x)), v, h)
    fd = grad(central_fdm(5,1), (x,y)->sum(y(x)), v, h)
    @test (zy[1] .≈ fd[1]) == Bool[1, 1]
    @test zy[2].H ≈ fd[2].H
    @test typeof(zy[2].H) == typeof(fd[2].H)
end

@testset "CoordinateTransformations.LinearMap gradient" begin
    m = CoordinateTransformations.LinearMap(rand(3, 3))
    sm = CoordinateTransformations.LinearMap(rand(SMatrix{3, 3, Float64}))
    v = rand(3)
    sv = rand(SVector{3, Float64})
    zy = Zygote.gradient((x, y)->sum(x(y)), m, v)
    fd = grad(central_fdm(5,1), (x, y)->sum(x(y)), m, v)
    zy_s = Zygote.gradient((x, y)->sum(x(y)), sm, sv)
    fd_s = grad(central_fdm(5,1), (x, y)->sum(x(y)), sm, sv)

    @test zy[1].linear ≈ fd[1].linear
    @test zy[2] ≈ fd[2]
    @test typeof(zy[1].linear) == typeof(fd[1].linear)
    @test zy_s[1].linear ≈ fd_s[1].linear
    @test zy_s[2] ≈ fd_s[2]
    @test_broken typeof(zy_s[1].linear) == typeof(fd_s[1].linear) # TODO: fix in another issue.
end

@testset "Rotations.RotMatrix{N, T, L} gradient" begin
    rm = RotMatrix(π/4)
    v = rand(2)
    zy = Zygote.gradient((x, y)->sum(LinearMap(x)(y)), rm, v)
    fd = grad(central_fdm(5,1), (x, y)->sum(LinearMap(x)(y)), rm, v)
    @test zy[1] ≈ fd[1]
    @test zy[2] ≈ fd[2]
end

@testset "ImageTransformations._getindex gradient" begin
    _sep(x) = x.r + x.g + x.b
    for t in (Float32, Float64, RGB{Float32}, RGB{Float64})
        itp = extrapolate(interpolate(rand(t, 3, 3), BSpline(Linear())), zero(t))
        for ind in ((2.5, 2.5), (5, 5))
            if t <: Colorant
                zy = Zygote.gradient((x,y)->_sep(ImageTransformations._getindex(x,y)), itp, ind)
                @test zy[2] ≈ fieldsum.(Interpolations.gradient(itp, ind...))
            else
                zy = Zygote.gradient((x,y)->ImageTransformations._getindex(x,y), itp, ind)
                @test zy[2] ≈ Interpolations.gradient(itp, ind...)
            end
        end
    end
end


"""
`prun_we`: Run a parallel WE simulation, optionally returning the ensemble at
each step. This performs the mutation steps in parallel, and assumes a worker
pool has already been created.

### Arguments
* `E₀` - initial particle ensemble
* `B₀` - initial bin data structure
* `sampler` - WE sampler functions data structure
* `n_we_steps` - number of steps in the WE run
### Optional Arguments
* `n_save_iters = 1` - save the ensemble and bins every `n_save_iters` iterations.  Set `n_save_iters=n_we_steps` to only save the final values.
"""
function prun_we(E₀::TE, B₀::TB, sampler::TWES, n_we_steps::Int; n_save_iters = 1) where
   {TE<:Ensemble, TB<:Bins, TWES<:DistributedWEsampler}

   E = deepcopy(E₀);
   B = deepcopy(B₀);
   E_trajectory = TE[];
   B_trajectory = TB[];

   for t in 0:n_we_steps-1
      # first selection is at t = 0
      sampler.selection!(E, B, t);
      copy!(E.ω, E.ω̂);
      @. E.d = deepcopy(E.d̂);
      E.ξ .= pmap(sampler.mutation, E.ξ̂);
      # after mutation, time is t ↦ t+1
      sampler.rebin!(E, B, t+1);
      # analysis
      sampler.analysis!(E, B, t+1);

      if(mod(t+1, n_save_iters)==0)
         push!(E_trajectory, deepcopy(E))
         push!(B_trajectory, deepcopy(B))
      end
   end

   return E_trajectory, B_trajectory

end

"""
`prun_we_observables`: Run a parallel WE simulation, optionally returning the ensemble at
each step. This performs the mutation steps in parallel, and assumes a worker
pool has already been created.

### Arguments
* `E₀` - initial particle ensemble
* `B₀` - initial bin data structure
* `sampler` - WE sampler functions data structure
* `n_we_steps` - number of steps in the WE run
* `observables` - Tuple of scalar observable functions for the ergodic average
"""
function prun_we_observables(E₀::TE, B₀::TB, sampler::TWES, n_we_steps::Int, observables::Tuple{Vararg{<:Function,NO}}) where
   {TE<:Ensemble, TB<:Bins, TWES<:DistributedWEsampler, NO}

   E = deepcopy(E₀);
   B = deepcopy(B₀);
   observables_trajectory = zeros(NO, n_we_steps);

   for t in 0:n_we_steps-1
      # first selection is at t = 0
      sampler.selection!(E, B, t);
      copy!(E.ω, E.ω̂);
      @. E.d = deepcopy(E.d̂);
      E.ξ .= pmap(sampler.mutation, E.ξ̂);
      # after mutation, time is t ↦ t+1
      sampler.rebin!(E, B, t+1);
      # analysis
      sampler.analysis!(E, B, t+1);

      ntuple(k-> observables_trajectory[k,t+1] =(observables[k]).(E.ξ) ⋅ E.ω, NO)
   end

   return observables_trajectory
end

"""
`prun_we!`: Run an in place parallel WE simulation.  This performs the mutation
steps in parallel, and assumes a worker pool has already been created.

### Arguments
* `E` - particle ensemble
* `B` - bin data structure
* `sampler` - WE sampler functions data structure
* `n_we_steps` - number of steps in the WE run
"""
function prun_we!(E::TE, B::TB, sampler::TWES, n_we_steps::Int) where
   {TE<:Ensemble, TB<:Bins, TWES<:DistributedWEsampler}

   for t in 0:n_we_steps-1
      # first selection is at t = 0
      sampler.selection!(E, B, t);
      copy!(E.ω, E.ω̂);
      @. E.d = deepcopy(E.d̂);
      E.ξ .= pmap(sampler.mutation, E.ξ̂);
      # after mutation, time is t ↦ t+1
      sampler.rebin!(E, B, t+1);
      # analysis
      sampler.analysis!(E, B, t+1);

   end
   E, B

end


get_connection_type(s) = getmetadata(unwrap(s), VariableConnectType, Equality)

function with_connector_type(expr)
    @assert expr isa Expr && (expr.head == :function || (expr.head == :(=) &&
                                       expr.args[1] isa Expr &&
                                       expr.args[1].head == :call))

    sig = expr.args[1]
    body = expr.args[2]

    fname = sig.args[1]
    args = sig.args[2:end]

    quote
        function $fname($(args...))
            function f()
                $body
            end
            res = f()
            $isdefined(res, :connector_type) && $getfield(res, :connector_type) === nothing ? $Setfield.@set!(res.connector_type = $connector_type(res)) : res
        end
    end
end

macro connector(expr)
    esc(with_connector_type(expr))
end

abstract type AbstractConnectorType end
struct StreamConnector <: AbstractConnectorType end
struct RegularConnector <: AbstractConnectorType end

function connector_type(sys::AbstractSystem)
    sts = get_states(sys)
    n_stream = 0
    n_flow = 0
    n_regular = 0 # state that is not input, output, stream, or flow.
    for s in sts
        vtype = get_connection_type(s)
        if vtype === Stream
            isarray(s) && error("Array stream variables are not supported. Got $s.")
            n_stream += 1
        elseif vtype === Flow
            n_flow += 1
        elseif !(isinput(s) || isoutput(s))
            n_regular += 1
        end
    end
    (n_stream > 0 && n_flow > 1) && error("There are multiple flow variables in $(nameof(sys))!")
    if n_flow != n_regular
        @warn "$(nameof(sys)) contains $n_flow variables, yet $n_regular regular " *
        "(non-flow, non-stream, non-input, non-output) variables." *
        "This could lead to imbalanced model that are difficult to debug." *
        "Consider marking some of the regular variables as input/output variables."
    end
    n_stream > 0 ? StreamConnector() : RegularConnector()
end

get_systems(c::Connection) = c.systems

instream(a) = term(instream, unwrap(a), type=symtype(a))
SymbolicUtils.promote_symtype(::typeof(instream), _) = Real

isconnector(s::AbstractSystem) = has_connector_type(s) && get_connector_type(s) !== nothing

function flowvar(sys::AbstractSystem)
    sts = get_states(sys)
    for s in sts
        vtype = get_connection_type(s)
        vtype === Flow && return s
    end
    error("There in no flow variable in $(nameof(sys))")
end

collect_instream!(set, eq::Equation) = collect_instream!(set, eq.lhs) | collect_instream!(set, eq.rhs)

function collect_instream!(set, expr, occurs=false)
    istree(expr) || return occurs
    op = operation(expr)
    op === instream && (push!(set, expr); occurs = true)
    for a in SymbolicUtils.unsorted_arguments(expr)
        occurs |= collect_instream!(set, a, occurs)
    end
    return occurs
end

#positivemax(m, ::Any; tol=nothing)= max(m, something(tol, 1e-8))
#_positivemax(m, tol) = ifelse((-tol <= m) & (m <= tol), ((3 * tol - m) * (tol + m)^3)/(16 * tol^3) + tol, max(m, tol))
function _positivemax(m, si)
    T = typeof(m)
    relativeTolerance = 1e-4
    nominal = one(T)
    eps = relativeTolerance * nominal
    alpha = if si > eps
        one(T)
    else
        if si > 0
            (si/eps)^2*(3-2* si/eps)
        else
            zero(T)
        end
    end
    alpha * max(m, 0) + (1-alpha)*eps
end
@register _positivemax(m, tol)
positivemax(m, ::Any; tol=nothing) = _positivemax(m, tol)
mydiv(num, den) = if den == 0
    error()
else
    num / den
end
@register mydiv(n, d)

function generate_isouter(sys::AbstractSystem)
    outer_connectors = Symbol[]
    for s in get_systems(sys)
        n = nameof(s)
        isconnector(s) && push!(outer_connectors, n)
    end
    let outer_connectors=outer_connectors
        function isouter(sys)::Bool
            s = string(nameof(sys))
            isconnector(sys) || error("$s is not a connector!")
            idx = findfirst(isequal('₊'), s)
            parent_name = Symbol(idx === nothing ? s : s[1:prevind(s, idx)])
            parent_name in outer_connectors
        end
    end
end

struct LazyNamespace
    namespace::Union{Nothing,Symbol}
    sys
end

Base.copy(l::LazyNamespace) = renamespace(l.namespace, l.sys)
Base.nameof(l::LazyNamespace) = renamespace(l.namespace, nameof(l.sys))

struct ConnectionElement
    sys::LazyNamespace
    v
    isouter::Bool
end
Base.hash(l::ConnectionElement, salt::UInt) = hash(nameof(l.sys)) ⊻ hash(l.v) ⊻ hash(l.isouter) ⊻ salt
Base.isequal(l1::ConnectionElement, l2::ConnectionElement) = l1 == l2
Base.:(==)(l1::ConnectionElement, l2::ConnectionElement) = nameof(l1.sys) == nameof(l2.sys) && isequal(l1.v, l2.v) && l1.isouter == l2.isouter
namespaced_var(l::ConnectionElement) = states(l, l.v)
states(l::ConnectionElement, v) = states(copy(l.sys), v)

struct ConnectionSet
    set::Vector{ConnectionElement} # namespace.sys, var, isouter
end

function Base.show(io::IO, c::ConnectionSet)
    print(io, "<")
    for i in 1:length(c.set)-1
        @unpack sys, v, isouter = c.set[i]
        print(io, nameof(sys), ".", v, "::", isouter ? "outer" : "inner", ", ")
    end
    @unpack sys, v, isouter = last(c.set)
    print(io, nameof(sys), ".", v, "::", isouter ? "outer" : "inner", ">")
end

@noinline connection_error(ss) = error("Different types of connectors are in one conenction statement: <$(map(nameof, ss))>")

function connection2set!(connectionsets, namespace, ss, isouter)
    nn = map(nameof, ss)
    sts1 = Set(states(first(ss)))
    T = ConnectionElement
    csets = [T[] for _ in 1:length(sts1)]
    for (i, s) in enumerate(ss)
        sts = states(s)
        i != 1 && ((length(sts1) == length(sts) && all(Base.Fix2(in, sts1), sts)) || connection_error(ss))
        io = isouter(s)
        for (j, v) in enumerate(sts)
            push!(csets[j], T(LazyNamespace(namespace, s), v, io))
        end
    end
    for cset in csets
        vtype = get_connection_type(first(cset).v)
        for k in 2:length(cset)
            vtype === get_connection_type(cset[k].v) || connection_error(ss)
        end
        push!(connectionsets, ConnectionSet(cset))
    end
end

function generate_connection_set(sys::AbstractSystem)
    connectionsets = ConnectionSet[]
    sys = generate_connection_set!(connectionsets, sys)
    sys, merge(connectionsets)
end

function generate_connection_set!(connectionsets, sys::AbstractSystem, namespace=nothing)
    subsys = get_systems(sys)

    isouter = generate_isouter(sys)
    eqs′ = get_eqs(sys)
    eqs = Equation[]

    cts = [] # connections
    for eq in eqs′
        if eq.lhs isa Connection
            push!(cts, get_systems(eq.rhs))
        else
            push!(eqs, eq) # split connections and equations
        end
    end

    # all connectors are eventually inside connectors.
    T = ConnectionElement
    for s in subsys
        isconnector(s) || continue
        for v in states(s)
            Flow === get_connection_type(v) || continue
            push!(connectionsets, ConnectionSet([T(LazyNamespace(namespace, s), v, false)]))
        end
    end

    for ct in cts
        connection2set!(connectionsets, namespace, ct, isouter)
    end

    # pre order traversal
    @set! sys.systems = map(s->generate_connection_set!(connectionsets, s, renamespace(namespace, nameof(s))), subsys)
    @set! sys.eqs = eqs
end

function Base.merge(csets::AbstractVector{<:ConnectionSet})
    mcsets = ConnectionSet[]
    ele2idx = Dict{ConnectionElement,Int}()
    cacheset = Set{ConnectionElement}()
    for cset in csets
        idx = nothing
        for e in cset.set
            idx = get(ele2idx, e, nothing)
            idx !== nothing && break
        end
        if idx === nothing
            push!(mcsets, cset)
            for e in cset.set
                ele2idx[e] = length(mcsets)
            end
        else
            for e in mcsets[idx].set
                push!(cacheset, e)
            end
            for e in cset.set
                push!(cacheset, e)
            end
            empty!(mcsets[idx].set)
            for e in cacheset
                ele2idx[e] = idx
                push!(mcsets[idx].set, e)
            end
            empty!(cacheset)
        end
    end
    mcsets
end

function generate_connection_equations_and_stream_connections(csets::AbstractVector{<:ConnectionSet})
    eqs = Equation[]
    stream_connections = ConnectionSet[]

    for cset in csets
        v = cset.set[1].v
        if hasmetadata(v, Symbolics.GetindexParent)
            v = getparent(v)
        end
        vtype = get_connection_type(v)
        if vtype === Stream
            push!(stream_connections, cset)
            continue
        elseif vtype === Flow
            rhs = 0
            for ele in cset.set
                v = namespaced_var(ele)
                rhs += ele.isouter ? -v : v
            end
            push!(eqs, 0 ~ rhs)
        else # Equality
            base = namespaced_var(cset.set[1])
            for i in 2:length(cset.set)
                v = namespaced_var(cset.set[i])
                push!(eqs, base ~ v)
            end
        end
    end
    eqs, stream_connections
end

function expand_connections(sys::AbstractSystem; debug=false, tol=1e-10)
    sys, csets = generate_connection_set(sys)
    ceqs, instream_csets = generate_connection_equations_and_stream_connections(csets)
    _sys = expand_instream(instream_csets, sys; debug=debug, tol=tol)
    sys = flatten(sys, true)
    @set! sys.eqs = [equations(_sys); ceqs]
end

function unnamespace(root, namespace)
    root === nothing && return namespace
    root = string(root)
    namespace = string(namespace)
    if length(namespace) > length(root)
        @assert root == namespace[1:length(root)]
        Symbol(namespace[nextind(namespace, length(root)):end])
    else
        @assert root == namespace
        nothing
    end
end

function expand_instream(csets::AbstractVector{<:ConnectionSet}, sys::AbstractSystem, namespace=nothing, prevnamespace=nothing; debug=false, tol=1e-8)
    subsys = get_systems(sys)
    # post order traversal
    @set! sys.systems = map(s->expand_instream(csets, s, renamespace(namespace, nameof(s)), namespace; debug, tol), subsys)
    subsys = get_systems(sys)

    if debug
        @info "Expanding" namespace
    end

    sub = Dict()
    eqs = Equation[]
    instream_eqs = Equation[]
    instream_exprs = Set()
    for s in subsys
        for eq in get_eqs(s)
            eq = namespace_equation(eq, s)
            if collect_instream!(instream_exprs, eq)
                push!(instream_eqs, eq)
            else
                push!(eqs, eq)
            end
        end

    end

    for ex in instream_exprs
        cset, idx_in_set, sv = get_cset_sv(namespace, ex, csets)

        n_inners = n_outers = 0
        for (i, e) in enumerate(cset)
            if e.isouter
                n_outers += 1
            else
                n_inners += 1
            end
        end
        if debug
            @info "Expanding at [$idx_in_set]" ex ConnectionSet(cset)
            @show n_inners, n_outers
        end
        if n_inners == 1 && n_outers == 0
            sub[ex] = sv
        elseif n_inners == 2 && n_outers == 0
            other = idx_in_set == 1 ? 2 : 1
            sub[ex] = get_current_var(namespace, cset[other], sv)
        elseif n_inners == 1 && n_outers == 1
            if !cset[idx_in_set].isouter
                other = idx_in_set == 1 ? 2 : 1
                outerstream = get_current_var(namespace, cset[other], sv)
                sub[ex] = instream(outerstream)
            end
        else
            if !cset[idx_in_set].isouter
                fv = flowvar(first(cset).sys.sys)
                # mj.c.m_flow
                innerfvs = [get_current_var(namespace, s, fv) for (j, s) in enumerate(cset) if j != idx_in_set && !s.isouter]
                innersvs = [get_current_var(namespace, s, sv) for (j, s) in enumerate(cset) if j != idx_in_set && !s.isouter]
                # ck.m_flow
                outerfvs = [get_current_var(namespace, s, fv) for s in cset if s.isouter]
                outersvs = [get_current_var(namespace, s, sv) for s in cset if s.isouter]

                sub[ex] = term(instream_rt, Val(length(innerfvs)), Val(length(outerfvs)), innerfvs..., innersvs..., outerfvs..., outersvs...)
            end
        end
    end

    # additional equations
    additional_eqs = Equation[]
    csets = filter(cset->any(e->e.sys.namespace === namespace, cset.set), csets)
    for cset′ in csets
        cset = cset′.set
        connectors = Vector{Any}(undef, length(cset))
        n_inners = n_outers = 0
        for (i, e) in enumerate(cset)
            connectors[i] = e.sys.sys
            if e.isouter
                n_outers += 1
            else
                n_inners += 1
            end
        end
        iszero(n_outers) && continue
        connector_representative = first(cset).sys.sys
        fv = flowvar(connector_representative)
        sv = first(cset).v
        vtype = get_connection_type(sv)
        vtype === Stream || continue
        if n_inners == 1 && n_outers == 1
            push!(additional_eqs, states(cset[1].sys.sys, sv) ~ states(cset[2].sys.sys, sv))
        elseif n_inners == 0 && n_outers == 2
            # we don't expand `instream` in this case.
            v1 = states(cset[1].sys.sys, sv)
            v2 = states(cset[2].sys.sys, sv)
            push!(additional_eqs, v1 ~ instream(v2))
            push!(additional_eqs, v2 ~ instream(v1))
        else
            sq = 0
            s_inners = (s for s in cset if !s.isouter)
            s_outers = (s for s in cset if s.isouter)
            for (q, oscq) in enumerate(s_outers)
                sq += sum(s->max(-states(s, fv), 0), s_inners)
                for (k, s) in enumerate(s_outers); k == q && continue
                    f = states(s.sys.sys, fv)
                    sq += max(f, 0)
                end

                num = 0
                den = 0
                for s in s_inners
                    f = states(s.sys.sys, fv)
                    tmp = positivemax(-f, sq; tol=tol)
                    den += tmp
                    num += tmp * states(s.sys.sys, sv)
                end
                for (k, s) in enumerate(s_outers); k == q && continue
                    f = states(s.sys.sys, fv)
                    tmp = positivemax(f, sq; tol=tol)
                    den += tmp
                    num += tmp * instream(states(s.sys.sys, sv))
                end
                push!(additional_eqs, states(oscq.sys.sys, sv) ~ num / den)
            end
        end
    end

    subed_eqs = substitute(instream_eqs, sub)
    if debug && !(isempty(csets) && isempty(additional_eqs) && isempty(instream_eqs))
        println("======================================")
        @info "Additional equations" csets
        display(additional_eqs)
        println("======================================")
        println("Substitutions")
        display(sub)
        println("======================================")
        println("Substituted equations")
        foreach(i->println(instream_eqs[i] => subed_eqs[i]), eachindex(subed_eqs))
        println("======================================")
    end

    @set! sys.systems = []
    @set! sys.eqs = [get_eqs(sys); eqs; subed_eqs; additional_eqs]
    sys
end

function get_current_var(namespace, cele, sv)
    states(renamespace(unnamespace(namespace, cele.sys.namespace), cele.sys.sys), sv)
end

function get_cset_sv(namespace, ex, csets)
    ns_sv = only(arguments(ex))
    full_name_sv = renamespace(namespace, ns_sv)

    cidx = -1
    idx_in_set = -1
    sv = ns_sv
    for (i, c) in enumerate(csets)
        crep = first(c.set)
        current = namespace == crep.sys.namespace
        for (j, v) in enumerate(c.set)
            if isequal(namespaced_var(v), full_name_sv) && (current || !v.isouter)
                cidx = i
                idx_in_set = j
                sv = v.v
            end
        end
    end
    cidx < 0 && error("$ns_sv is not a variable inside stream connectors")
    cset = csets[cidx].set
    #if namespace != first(cset).sys.namespace
    #    cset = map(c->@set(c.isouter = false), cset)
    #end
    cset, idx_in_set, sv
end

# instream runtime
@generated function _instream_split(::Val{inner_n}, ::Val{outer_n}, vars::NTuple{N,Any}) where {inner_n, outer_n, N}
    #instream_rt(innerfvs..., innersvs..., outerfvs..., outersvs...)
    ret = Expr(:tuple)
    # mj.c.m_flow
    inner_f = :(Base.@ntuple $inner_n i -> vars[i])
    offset = inner_n
    inner_s = :(Base.@ntuple $inner_n i -> vars[$offset+i])
    offset += inner_n
    # ck.m_flow
    outer_f = :(Base.@ntuple $outer_n i -> vars[$offset+i])
    offset += outer_n
    outer_s = :(Base.@ntuple $outer_n i -> vars[$offset+i])
    Expr(:tuple, inner_f, inner_s, outer_f, outer_s)
end

function instream_rt(ins::Val{inner_n}, outs::Val{outer_n}, vars::Vararg{Any,N}) where {inner_n, outer_n, N}
    @assert N == 2*(inner_n + outer_n)

    # inner: mj.c.m_flow
    # outer: ck.m_flow
    inner_f, inner_s, outer_f, outer_s = _instream_split(ins, outs, vars)

    T = float(first(inner_f))
    si = zero(T)
    num = den = zero(T)
    for f in inner_f
        si += max(-f, 0)
    end
    for f in outer_f
        si += max(f, 0)
    end
    #for (f, s) in zip(inner_f, inner_s)
    for j in 1:inner_n
        @inbounds f = inner_f[j]
        @inbounds s = inner_s[j]
        num += _positivemax(-f, si) * s
        den += _positivemax(-f, si)
    end
    #for (f, s) in zip(outer_f, outer_s)
    for j in 1:outer_n
        @inbounds f = outer_f[j]
        @inbounds s = outer_s[j]
        num += _positivemax(-f, si) * s
        den += _positivemax(-f, si)
    end
    return num / den
    #=
    si = sum(max(-mj.c.m_flow,0) for j in cat(1,1:i-1, i+1:N)) +
            sum(max(ck.m_flow ,0) for k  in 1:M)

    inStream(mi.c.h_outflow) =
       (sum(positiveMax(-mj.c.m_flow,si)*mj.c.h_outflow)
      +  sum(positiveMax(ck.m_flow,s_i)*inStream(ck.h_outflow)))/
     (sum(positiveMax(-mj.c.m_flow,s_i))
        +  sum(positiveMax(ck.m_flow,s_i)))
                  for j in 1:N and i <> j and mj.c.m_flow.min < 0,
                  for k in 1:M and ck.m_flow.max > 0
    =#
end
SymbolicUtils.promote_symtype(::typeof(instream_rt), ::Vararg) = Real