module OBJMeshes

export read_obj, write_obj, tri_vertex_str
export cube_mesh, cylinder_mesh, sphere_mesh, estimate_normals, portal_mesh, portal_background_mesh
export OBJTriangle, OBJMesh

using FileIO
using LinearAlgebra

#push!(LOAD_PATH, pwd())
using ..GfxBase

""" OBJTriangle
A struct that represents a single triangle in a mesh. """
mutable struct OBJTriangle
    positions::Array{Int, 1} # vertex position indices
    uvs::Array{Int, 1} # vertex texture coordinate indices
    normals::Array{Int, 1} # normal vector indices
end

""" OBJMesh
A struct that represents an indexed triangle mesh for reading from
or writing to OBJ format. """
mutable struct OBJMesh
    positions::Array{Vec3, 1} # all vertex positions
    uvs::Array{Vec2, 1} # all texture coordinates
    normals::Array{Vec3, 1} # all vertex normals
    triangles::Array{OBJTriangle, 1} # the OBJTriangles belonging to the mesh
end

""" read_obj(obj_filename)
Read a mesh in OBJ format from file obj_filename."""
function read_obj(obj_filename)
    m = OBJMesh([], [], [], []) # create a mesh
    open(obj_filename) do f
        for (line_number, line) in enumerate(eachline(f))
            if line == "" || line[1] == "#"
                continue # skip comments
            end
            # Read the line and add its contents to the correct field of m:
            tokens = split(strip(line))
            if tokens[1] == "v" # vertex
                push!(m.positions, Vec3([parse(Float64, x) for x in tokens[2:end]]...))
            elseif tokens[1] == "vt" # vertex texture
                push!(m.uvs, Vec2([parse(Float64, x) for x in tokens[2:end]]...))
            elseif tokens[1] == "vn" # vertex normal
                push!(m.normals, Vec3([parse(Float64, x) for x in tokens[2:end]]...))
            elseif tokens[1] == "f"
                # create a OBJTriangle face:
                points = []
                uvs = []
                normals = []
                # handle faces with no texture and/or normals
                for corner in tokens[2:end]
                    indices = split(corner, '/')
                    if length(indices) == 3 # all 3 present, third is a normal
                        push!(normals, parse(Int, indices[3]))
                    end
                    if length(indices) >= 2 && indices[2] != ""
                        # if there are 2 or more and the second isn't blank, it's a texture
                        push!(uvs, parse(Int, indices[2]))
                    end
                    if length(indices) >= 1 # first value is the position
                        push!(points, parse(Int, indices[1]))
                    else # unless it has none, in which case it's not valid
                        error("in line $line_number: face vertex $corner could not be parsed")
                    end
                end
                # create the triangle and add it to the triangles array
                push!(m.triangles, OBJTriangle(points, uvs, normals))
            end
        end
    end
    return m
end

""" write_obj(obj_filename)
Write the given mesh in OBJ format to file obj_filename."""
function write_obj(obj_filename, mesh::OBJMesh)
    open(obj_filename, "w") do f
        # write all positions:
        for v in mesh.positions
            write(f, "v $(v[1]) $(v[2]) $(v[3])\n")
        end

        # write all texture coords:
        for v in mesh.uvs
            write(f, "vt $(v[1]) $(v[2])\n")
        end
        # write all normals:
        for v in mesh.normals
            write(f, "vn $(v[1]) $(v[2]) $(v[3])\n")
        end

        # write all triangles:
        for tri in mesh.triangles
            write(f, "f $(tri_vertex_str(tri))\n")
        end

    end

end

""" tri_vertex_str(triangle)
Return a string with the indices of applicable positions, texture coordinates,
and normals for a given triangle according to the OBJ specification.
In particular, if p, u, and n are position, vertex and normal, each corner
of the triangle is represented as one of the following:
    p       (position only)
    p/u     (position and texture)
    p//n    (position and normal)
    p/u/n   (position, texture, and normal)
"""
function tri_vertex_str(triangle::OBJTriangle)
    # determine whether textures and normals are present:
    write_uv = length(triangle.uvs) == length(triangle.positions)
    write_normals = length(triangle.normals) == length(triangle.positions)
    corners = []
    for i = 1:3
        output = "$(triangle.positions[i])"
        if write_uv && !write_normals
            output = output * "/$(triangle.uvs[i])" # * does concatenation(!)
        elseif !write_uv && write_normals
            output = output * "//$(triangle.normals[i])"
        elseif write_uv && write_normals
            output = output * "/$(triangle.uvs[i])/$(triangle.normals[i])"
        end
        push!(corners, output)
    end
    join(corners, " ")
end

""" cube_mesh()
Return a new OBJMesh representing a 2x2x2 cube centered at the origin and
axis-aligned. """
function cube_mesh()
    positions = []
    uvs = []
    normals = []
    triangles = []
    # key to comments:
    # L/R = x = right/left
    # B/T = y = top/bottom
    # C/F = z = close/far
    push!(positions, Vec3( 1, -1, -1)) # 1 RBC
    push!(positions, Vec3( 1, -1,  1)) # 2 RBF
    push!(positions, Vec3(-1, -1,  1)) # 3 LBF
    push!(positions, Vec3(-1, -1, -1)) # 4 LBC
    push!(positions, Vec3( 1,  1, -1)) # 5 RTC
    push!(positions, Vec3( 1,  1,  1)) # 6 RTF
    push!(positions, Vec3(-1,  1,  1)) # 7 LTF
    push!(positions, Vec3(-1,  1, -1)) # 8 LTC

    # texture coordinates:
    push!(uvs, Vec2(1, 1)) # TR
    push!(uvs, Vec2(0, 1)) # TL
    push!(uvs, Vec2(0, 0)) # BL
    push!(uvs, Vec2(1, 0)) # BR

    #by removing pre-made normals from our objects, we're able to rotate them and calculate new normals (gets rid of black surfaces!)

    # normals:
    #push!(normals, Vec3( 1, 0, 0)) # R
    #push!(normals, Vec3(-1, 0, 0)) # L
    #push!(normals, Vec3( 0, 1, 0)) # U
    #push!(normals, Vec3( 0,-1, 0)) # D
    #push!(normals, Vec3( 0, 0, 1)) # C
    #push!(normals, Vec3( 0, 0,-1)) # F

    # 8 faces, 2 triangles each
    push!(triangles, OBJTriangle([1,2,3], [1,2,3], [])) # bottom face 1
    push!(triangles, OBJTriangle([1,3,4], [1,3,4], [])) # bottom face 2
    push!(triangles, OBJTriangle([1,5,6], [4,1,2], [])) # right face 1
    push!(triangles, OBJTriangle([1,6,2], [4,2,3], [])) # right face 2
    push!(triangles, OBJTriangle([2,6,7], [4,1,2], [])) # far face 1
    push!(triangles, OBJTriangle([2,7,3], [4,2,3], [])) # far face 2
    push!(triangles, OBJTriangle([3,7,8], [2,3,4], [])) # left face 1
    push!(triangles, OBJTriangle([3,8,4], [2,4,1], [])) # left face 2
    push!(triangles, OBJTriangle([4,8,5], [2,3,4], [])) # far face 1
    push!(triangles, OBJTriangle([4,5,1], [2,4,1], [])) # far face 2
    push!(triangles, OBJTriangle([5,8,7], [1,2,3], [])) # top face 1
    push!(triangles, OBJTriangle([5,7,6], [1,3,4], [])) # top face 2

    # julia automatically returns the last value in the function:
    OBJMesh(positions, uvs, normals, triangles)

end

function portal_mesh(x, y)

    positions = []
    uvs = []
    normals = []
    triangles = []
    push!(positions, Vec3(0,0,1)) # top of cap
    #push!(normals, Vec3(0,0,1))
    n = 20
    for i = 1:n+1
        ratio = i/n
        theta = 2*pi*ratio
        push!(positions, Vec3(-sin(theta)*x, -cos(theta)*y, 1))
        #push!(normals, Vec3(-sin(theta),  -cos(theta), 1))
    end
	length_pos = length(positions)
     for i=1:length_pos-1
       push!(triangles, OBJTriangle([i, 1,i+1],[i+length_pos,1,i+length_pos],[]))
     end
     OBJMesh(positions, uvs, normals, triangles)

end

#a copy of the portal mesh function that (hopefully) will be recognized as not-a-portal
function portal_background_mesh(x, y)
    positions = []
    uvs = []
    normals = []
    triangles = []
    push!(positions, Vec3(0,0,1)) # top of cap
    n = 20
    for i = 1:n+1
        ratio = i/n
        theta = 2*pi*ratio
        push!(positions, Vec3(-sin(theta)*x, -cos(theta)*y, 1))
    end
	length_pos = length(positions)
     for i=1:length_pos-1
       push!(triangles, OBJTriangle([i, 1,i+1],[i+length_pos,1,i+length_pos],[]))
     end
     OBJMesh(positions, uvs, normals, triangles)
end

""" cylinder_mesh(n)
Return a new OBJMesh object approximation of a cylinder with radius 1 and
height 2, centered at the origin. The logitudinal axis is aligned with y, and
it is tesselated with n divisions arranged radially around the outer surface.
The ends of the cylinder are disc-shaped caps parallel to the xz plane.
"""
function cylinder_mesh(divisionsU)
  positions = []
  uvs = []
  normals = []
  triangles = []

  n = divisionsU # change to n for ease of typing
  push!(positions, Vec3(0,1,0)) # top of cap
  push!(normals, Vec3(0,1,0))

  push!(positions, Vec3(0,-1,0)) # bottom of cap
  push!(normals, Vec3(0,-1,0))

  push!(positions, Vec3(0,1,-1)) # start of shell
  push!(positions, Vec3(0,-1,-1))
  push!(normals, Vec3(0,0,-1))

  push!(uvs, Vec2(.75, .75)) # initial uv coordinates
  push!(uvs, Vec2(.25, .75))
  push!(uvs, Vec2(0, .5))
  push!(uvs, Vec2(0, 0))
  # positions
  for i = 1:n-1
    ratio = i/n
    theta = 2*pi*ratio
    push!(positions, Vec3(-sin(theta), 1, -cos(theta)))
    push!(positions, Vec3(-sin(theta), -1, -cos(theta)))
    push!(normals, Vec3(-sin(theta), 0, -cos(theta)))
    push!(uvs, Vec2(ratio, .5)) # uvs for shell
    push!(uvs, Vec2(ratio, 0))
  end

  length_pos = length(positions) # for later use
  length_norm = length(normals)

  push!(uvs, Vec2(1, .5)) # missing coordinates
  push!(uvs, Vec2(1, 0))
  push!(uvs, Vec2(.75, 1))
  push!(uvs, Vec2(.25, .5))

  # uvs for caps
  for i=5:length_pos
      x = positions[i][1] # -sin(θ)
      y = positions[i][2] # y ∈ (-1,1)
      z = positions[i][3] # -cos(θ)
    if y == 1.0 # top cap
      push!(uvs, Vec2((x+1)/4+.5, (1-z)/4+.5))
    else # bottom cap
      push!(uvs, Vec2((x+1)/4, (z+1)/4+.5))
    end
  end

  n1 = 3 # normal counters
  n2 = 3

  # create the shell ignoring the top and bottom coordinates
  for i =3:2:length_pos-3
    # create triangles in a ccw order, texture coords are same index
    push!(triangles, OBJTriangle([i,i+1, i+2],[i,i+1,i+2],
                                 [n1,n1,n1+1]))
    push!(triangles, OBJTriangle([i+1,i+3, i+2],[i+1,i+3,i+2]
                                 ,[n2, n2+1, n2+1]))
    n1 += 1 # normals increment half as fast as the shell
    n2 += 1
  end

  # last bit of the shell
  push!(triangles, OBJTriangle([length_pos-1,length_pos,3],[length_pos-1,length_pos,length_pos+1],
                               [length_norm, length_norm, 3]))
  push!(triangles, OBJTriangle([length_pos,4,3],[length_pos, length_pos+2, length_pos+1],
                               [length_norm,3 ,3 ]))

  # create the caps for the cylinder
  for i=3:2:length_pos-2
    push!(triangles, OBJTriangle([i, i+2, 1],[i+length_pos,i+length_pos+2,1],[1,1,1]))
    push!(triangles, OBJTriangle([i+1, 2, i+3],[i+length_pos+1,2,i+length_pos+3],[2,2,2]))
  end

  # last bit of the caps
  push!(triangles, OBJTriangle([length_pos-1, 3, 1], [2*length_pos-1,length_pos+3,1], [1,1,1]))
  push!(triangles, OBJTriangle([length_pos, 2, 4], [2*length_pos,2,length_pos+4], [2,2,2]))

  OBJMesh(positions,uvs ,normals,triangles)
end

""" sphere_mesh(n, m
Create a Latitude-Longitude-tesselated approximation of a sphere with radius 1
centered at the origin. There are n divisions around the equator and m
divisions from pole to pole along each line of longitude. The North pole is at
(0,1,0), the South pole at (0,-1,0)((-1, and points on the Greenwich meridian are
in the x = 0 plane with z > 0. The u texture coordinate depends on longitude,
with u=0 at 180 degrees West and u=1 at 180 degrees East. The v texture
coordinate varies with latitude with v=0 at the South pole and v=1 at the North
pole. Normals should be normal to the ideal sphere surface. See the assignment
for a diagram and further details. """
function sphere_mesh(n, m)
    positions = []
    uvs = []
    normals = []
    triangles = []

    u = 2 * pi / n # normalize
    v = pi / m # normalize

    push!(positions, Vec3(0, -1, 0)) # bottom pole
    push!(normals, Vec3(0,-1,0))

    push!(positions, Vec3(0, 1, 0)) # top pole
    push!(normals, Vec3(0,1,0))

    # In this next block we convert x,y,z to spherical coordinates and
    # add the vectors to the positions array
    for i=1:m-1
      theta = (pi / 2) - (i * v) # π/2 to -π/2
      y = sin(theta) # Spherical y coord
      for j=0:n-1
        phi = j * u
        x = cos(theta) * sin(phi) # Spherical x coord
        z = cos(theta) * cos(phi) # spherical y coord
        push!(positions, Vec3(-x, -y, -z)) # negative as we start from the bottom
        push!(normals, Vec3(-x,-y,-z))
      end
    end
    length_pos = length(positions)
    # uv coords
    for i=0:m
      for j=0:n
        push!(uvs, Vec2(j/n,i/m))
      end
    end

    # create the bottom "ring"
    for i=1:n-1
      push!(triangles, OBJTriangle([1, i+3,i+2],[i,i+n+1,i+n],[1,i+3,i+2]))
    end

    # missing piece of bottom ring
    push!(triangles, OBJTriangle([1,3,n+2],[n, 2*n+1,2*n],[1,3,n+2]))

    t1, t2 = n+1,2*n+2 # texture counters
    for i=3:length_pos-n
      # if i is a multiple of m we subtract m from the first and second coordinates, respectively.
      if ((i-2)%n) == 0
       push!(triangles,OBJTriangle([i,i-n+1,i+n],[t1,t1+1,t2],[i,i-n+1,i+n]))
       push!(triangles,OBJTriangle([i-n+1,i+1,i+n],[t1+1,t2+1,t2],[i-n+1,i+1,i+n]))
       t1 +=2 # textures increment twice as fast
       t2 += 2
       continue
      end

     # for most of the triangles, i is not a multiple of m, add m to 3rd coord.
     push!(triangles, OBJTriangle([i, i+1,i+n],[t1,t1+1,t2],[i, i+1,i+n]))
     push!(triangles, OBJTriangle([i+1,i+1+n,i+n],[t1+1,t2+1,t2],[i+1,i+1+n,i+n]))
     t1 += 1 # textures increment normally here
     t2 += 1
    end

    # top "ring"
    for i =length_pos-n+1:length_pos-1
      push!(triangles,OBJTriangle([i,i+1,2],[t1,t1+1,t1+1+n],[i,i+1,2]))
      t1 += 1
      t2 += 1
    end
    # missing bit of the top "ring"
    push!(triangles,OBJTriangle([length_pos,length_pos-n+1,2],[t1,t1+1,t1+1+n],
                                [length_pos, length_pos-n+1,2]))

   deleteat!(uvs, 1) # delete the first/last index of texture coords as it is unneeded
   deleteat!(uvs, length(uvs))
   OBJMesh(positions,uvs, normals, triangles)
end

"""
    estimate_normals(mesh::OBJMesh)
Estimates normals for the given mesh. Overwrites any existing normals and returns a new OBJMesh object.
"""
function estimate_normals(mesh::OBJMesh)
  normals = []
  # cerate a normal for every position
  for  x in mesh.positions
    push!(normals, Vec3(0,0,0))
  end
  # average all the normals for every triangle
  for x in mesh.triangles
    a = mesh.positions[x.positions[1]]
    b = mesh.positions[x.positions[2]]
    c = mesh.positions[x.positions[3]]
    # definition of computing a normal
    cr = cross(b-a, c-a)
    normals[x.positions[1]] += cr/norm(cr)
    normals[x.positions[2]] += cr/norm(cr)
    normals[x.positions[3]] += cr/norm(cr)
  end

  tmp = [] # create another array to normalize the normals
  for i in normals
    i = i/norm(i) # normalize
    push!(tmp, i)
  end

  normals = tmp # for the return value
  for x in mesh.triangles
    x.normals = x.positions # normals have same index as positions per the reference OBJs
  end
  OBJMesh(mesh.positions, mesh.uvs, normals, mesh.triangles)
end

end # module OBJMeshes


const SVF = MOI.SingleVariable
const VVF = MOI.VectorOfVariables
const SAF{T} = MOI.ScalarAffineFunction{T}
const VAF{T} = MOI.VectorAffineFunction{T}

const VI = MOI.VariableIndex
const CI = MOI.ConstraintIndex

const GT = MOI.GreaterThan
const LT = MOI.LessThan
const ET = MOI.EqualTo

const SCALAR_SETS = Union{
    MOI.GreaterThan{Float64},
    MOI.LessThan{Float64},
    MOI.EqualTo{Float64},
}

const VECTOR_SETS = Union{
    MOI.SecondOrderCone,
    MOI.RotatedSecondOrderCone,
}

# dual variable info
mutable struct ConstraintInfo{T<:Union{Float64, Vector{Float64}}}
    start::T
    upper::T
    lower::T
    function ConstraintInfo{Float64}()
        new(NaN, NaN, NaN)
    end
    function ConstraintInfo{Vector{Float64}}(N::Integer)
        new(fill(NaN, N), fill(NaN, N), fill(NaN, N))
    end
end

mutable struct VariableInfo{T<:Union{Float64, Vector{Float64}}}
    upper::T
    lower::T
    function VariableInfo()
        new{Float64}(NaN, NaN)
    end
    function VariableInfo(N::Integer)
        new{Vector{Float64}}(fill(NaN, N), fill(NaN, N))
    end
end
function VariableInfo(_info::ConstraintInfo{T}) where T
    if isa(_info.upper, Number)
        info = VariableInfo()
        info.lower = _info.lower
        info.upper = _info.upper
    else
        info = VariableInfo(length(_info.upper))
        info.lower .= _info.lower
        info.upper .= _info.upper
    end
    return info
end

struct Complement#{M1 <: MOI.ModelLike, M2 <: MOI.ModelLike, F, S}
    is_vec
    # primal::M1
    constraint
    func_w_cte#::F
    set_w_zero#::S
    # dual::M2
    variable#::VI
    # var_set#::S2
end

abstract type AbstractBilevelSolverMode{T} end

mutable struct NoMode{T} <: AbstractBilevelSolverMode{T}
end

mutable struct SOS1Mode{T} <: AbstractBilevelSolverMode{T}
    function SOS1Mode()
        return new{Float64}()
    end
end

mutable struct ComplementMode{T} <: AbstractBilevelSolverMode{T}
    with_slack::Bool
    function ComplementMode(;with_slack = false)
        return new{Float64}(with_slack)
    end
end

mutable struct ProductMode{T} <: AbstractBilevelSolverMode{T}
    epsilon::T
    with_slack::Bool
    function ProductMode(eps::T=zero(Float64); with_slack = false) where T
        return new{Float64}(eps, with_slack)
    end
end

@enum IndicatorSetting ZERO_ONE ZERO_ZERO ONE_ONE

mutable struct IndicatorMode{T} <: AbstractBilevelSolverMode{T}
    mode::IndicatorSetting
    function IndicatorMode()
        return new{Float64}(ONE_ONE)
    end
end

struct ComplementBoundCache
    # internal usage
    upper::Dict{VI, VariableInfo}
    lower::Dict{VI, VariableInfo}
    ldual::Dict{CI, ConstraintInfo}
    # full map
    map::Dict{VI, VariableInfo}
    function ComplementBoundCache()
        return new(
            Dict{VI, VariableInfo}(),
            Dict{VI, VariableInfo}(),
            Dict{CI, ConstraintInfo}(),
            Dict{VI, VariableInfo}(),
        )
    end
end

abstract type AbstractBoundedMode{T} <: AbstractBilevelSolverMode{T} end

mutable struct FortunyAmatMcCarlMode{T} <: AbstractBoundedMode{T}
    with_slack::Bool
    primal_big_M::Float64
    dual_big_M::Float64
    cache::ComplementBoundCache
    function FortunyAmatMcCarlMode(;with_slack = false,
        primal_big_M = Inf, dual_big_M = Inf)
        return new{Float64}(
            with_slack,
            primal_big_M,
            dual_big_M,
            ComplementBoundCache()
        )
    end
end

mutable struct MixedMode{T} <: AbstractBoundedMode{T}
    default::AbstractBilevelSolverMode{T}
    constraint_mode_map_c::Dict{CI, AbstractBilevelSolverMode{T}}
    constraint_mode_map_v::Dict{VI, AbstractBilevelSolverMode{T}}
    cache::ComplementBoundCache
    function MixedMode(;default = SOS1Mode())
        return new{Float64}(
            default,
            Dict{CI, AbstractBilevelSolverMode{Float64}}(),
            Dict{VI, AbstractBilevelSolverMode{Float64}}(),
            ComplementBoundCache()
        )
    end
end

function reset!(::AbstractBilevelSolverMode)
    return nothing
end
function reset!(mode::AbstractBoundedMode)
    mode.cache = ComplementBoundCache()
    return nothing
end

function appush!(col, element::AbstractVector)
    append!(col, element)
    return nothing
end
function appush!(col, element)
    push!(col, element)
    return nothing
end

function build_full_map!(mode,
    upper_idxmap, lower_idxmap, lower_dual_idxmap, lower_primal_dual_map)
    return nothing
end
function build_full_map!(mode::AbstractBoundedMode,
        upper_idxmap, lower_idxmap, lower_dual_idxmap, lower_primal_dual_map)
    _build_bound_map!(mode.cache,
        upper_idxmap, lower_idxmap, lower_dual_idxmap, lower_primal_dual_map)
    return nothing
end
function _build_bound_map!(mode::ComplementBoundCache,
    upper_idxmap, lower_idxmap, lower_dual_idxmap, lower_primal_dual_map)
    empty!(mode.map)
    for (k,v) in mode.upper
        mode.map[upper_idxmap[k]] = v
    end
    for (k,v) in mode.lower
        mode.map[lower_idxmap[k]] = v
    end
    for (k,v) in mode.ldual
        vec = lower_primal_dual_map.primal_con_dual_var[k]#[1] # TODO check this scalar
        for var in vec
            mode.map[lower_dual_idxmap[var]] = VariableInfo(v)
        end
    end
    return nothing
end

mutable struct StrongDualityMode{T} <: AbstractBilevelSolverMode{T}
    inequality::Bool
    epsilon::T
    function StrongDualityMode(eps::T=zero(Float64); inequality = true) where T
        return new{T}(inequality, eps)
    end
end

ignore_dual_objective(::AbstractBilevelSolverMode{T}) where T = true
ignore_dual_objective(::StrongDualityMode{T}) where T = false

function accept_vector_set(mode::AbstractBilevelSolverMode{T}, con::Complement) where T
    if con.is_vec
        error("Set $(typeof(con.set_w_zero)) is not accepted when solution method is $(typeof(mode))")
    end
    return nothing
end
accept_vector_set(::ProductMode{T}, ::Complement) where T = nothing
accept_vector_set(::MixedMode{T}, ::Complement) where T = nothing

function get_canonical_complements(primal_model, primal_dual_map)
    map = primal_dual_map.primal_con_dual_var
    out = Complement[]
    for ci in keys(map)
        con = get_canonical_complement(primal_model, map, ci)
        push!(out, con)
    end
    return out
end
function get_canonical_complement(primal_model, map,
    ci::CI{F,S}) where {F, S<:VECTOR_SETS}
    T = Float64
    func = MOI.copy(MOI.get(primal_model, MOI.ConstraintFunction(), ci))::F
    set = MOI.copy(MOI.get(primal_model, MOI.ConstraintSet(), ci))::S
    dim = MOI.dimension(set)
    # vector sets have no constant
    # for i in 1:dim
    #     func.constant[i] = Dualization.set_dot(i, set, T) *
    #         Dualization.get_scalar_term(primal_model, i, ci)
    # end
    # todo - set dot on function
    con = Complement(true, ci, func, set_with_zero(set), map[ci])
    return con
end
function get_canonical_complement(primal_model, map,
    ci::CI{F,S}) where {F, S<:SCALAR_SETS}
    T = Float64
    func = MOI.copy(MOI.get(primal_model, MOI.ConstraintFunction(), ci))::F
    set = MOI.copy(MOI.get(primal_model, MOI.ConstraintSet(), ci))::S
    constant = Dualization.set_dot(1, set, T) *
        Dualization.get_scalar_term(primal_model, 1, ci)
    if F == MOI.SingleVariable
        func = MOIU.operate(+, T, func, constant)
    else
        func.constant = constant
    end
    # todo - set dot on function
    con = Complement(false, ci, func, set_with_zero(set), map[ci][1])
    return con
end

function set_with_zero(set::S) where {S<:SCALAR_SETS}
    return S(0.0)
end
function set_with_zero(set)
    return MOI.copy(set)
end

function build_bilevel(
    upper::MOI.ModelLike, lower::MOI.ModelLike,
    link::Dict{VI,VI}, upper_variables::Vector{VI},
    mode,
    upper_var_lower_ctr::Dict{VI,CI} = Dict{VI,CI}();
    copy_names::Bool = false,
    pass_start::Bool = false
    )

    # Start with an empty problem
    moi_mode = MOIU.AUTOMATIC
    m = MOIU.CachingOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()), moi_mode)

    #=
        Initialize Lower DUAL level model
    =#
    # dualize the second level
    dual_problem = Dualization.dualize(lower,
        dual_names = DualNames("dual_","dual_"),
        variable_parameters = upper_variables,
        ignore_objective = ignore_dual_objective(mode))
    lower_dual = dual_problem.dual_model
    lower_primal_dual_map = dual_problem.primal_dual_map

    #=
        Pass Upper level model
    =#

    # key are from src, value are from dest
    upper_idxmap = MOIU.default_copy_to(m, upper, copy_names)
    if copy_names
        pass_names(m, upper, upper_idxmap)
    end

    #=
        Pass Lower level model
    =#

    handle_lower_objective_sense(lower)

    # cache and delete lower objective
    if !ignore_dual_objective(mode)
        # get primal obj
        tp_primal_obj = MOI.get(lower, MOI.ObjectiveFunctionType())
        @assert tp_primal_obj !== nothing
        lower_primal_obj = MOI.get(lower, MOI.ObjectiveFunction{tp_primal_obj}())
        # deepcopy and delete dual obj
        # MOI.set(lower, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 0.0))
    end

    # initialize map to lower level model
    lower_idxmap = MOIU.IndexMap()
    for (upper_key, lower_val) in link
        lower_idxmap[lower_val] = upper_idxmap[upper_key]
    end

    # append the second level primal
    append_to(m, lower, lower_idxmap, copy_names, allow_single_bounds = true)
    if copy_names
        pass_names(m, lower, lower_idxmap)
    end

    #=
        Pass Dual of Lower level model
    =#

    # initialize map to lower level model
    if !ignore_dual_objective(mode)
        # get dual obj
        tp_dual_obj = MOI.get(lower_dual, MOI.ObjectiveFunctionType())
        @assert tp_dual_obj !== nothing
        lower_dual_obj = MOI.get(lower_dual, MOI.ObjectiveFunction{tp_dual_obj}())
        # delete dual obj
        # MOI.set(lower_dual, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 0.0))
    end

    # initialize map to lower level model
    lower_dual_idxmap = MOIU.IndexMap()
    # for QP's there are dual variable that are tied to:
    # primal variables
    for (lower_primal_var_key, lower_dual_quad_slack_val) in lower_primal_dual_map.primal_var_dual_quad_slack
        lower_dual_idxmap[lower_dual_quad_slack_val] = lower_idxmap[lower_primal_var_key]
    end
    # and to upper level variable which are lower level parameters
    for (lower_primal_param_key, lower_dual_param_val) in lower_primal_dual_map.primal_parameter
        lower_dual_idxmap[lower_dual_param_val] = lower_idxmap[lower_primal_param_key]
    end
    # Dual variables might appear in the upper level
    for (upper_var, lower_con) in upper_var_lower_ctr
        var = lower_primal_dual_map.primal_con_dual_var[lower_con][1] # TODO check this scalar
        lower_dual_idxmap[var] = upper_idxmap[upper_var]
    end

    # append the second level dual
    append_to(m, lower_dual, lower_dual_idxmap, copy_names)
    if copy_names
        pass_names(m, lower_dual, lower_dual_idxmap)
    end

    #=
        Additional Optimiality conditions (to complete the KKT)
    =#

    # build map bound map for FortunyAmatMcCarlMode
    build_full_map!(mode,
    upper_idxmap, lower_idxmap, lower_dual_idxmap, lower_primal_dual_map)

    if ignore_dual_objective(mode)
        # complementary slackness
        comps = get_canonical_complements(lower, lower_primal_dual_map)
        for comp in comps
            if !is_equality(comp.set_w_zero)
                accept_vector_set(mode, comp)
                add_complement(mode, m, comp,
                    lower_idxmap, lower_dual_idxmap, copy_names, pass_start)
            else
                # println("eq in complement")
            end
        end
    else
        # strong duality
        lower_dual_obj
        lower_primal_obj
        add_strong_duality(mode, m, lower_primal_obj, lower_dual_obj, lower_idxmap, lower_dual_idxmap)
    end

    return m, upper_idxmap, lower_idxmap, lower_primal_dual_map, lower_dual_idxmap
end

function add_strong_duality(mode::StrongDualityMode{T}, m, primal_obj, dual_obj,
    idxmap_primal, idxmap_dual) where T

    primal = MOIU.map_indices.(Ref(idxmap_primal), primal_obj)
    dual   = MOIU.map_indices.(Ref(idxmap_dual), dual_obj)

    func = MOIU.operate(-, T, primal, dual)

    if !mode.inequality
        c = MOIU.normalize_and_add_constraint(m, func, MOI.EqualTo(zero(T)))
        MOI.set(m, MOI.ConstraintName(), c, "lower_strong_duality")
        return CI[c]
    else
        func_up = MOIU.operate(-, T, func, mode.epsilon)
        c_up = MOIU.normalize_and_add_constraint(m, func_up, MOI.LessThan(zero(T)))
        MOI.set(m, MOI.ConstraintName(), c_up, "lower_strong_duality_up")

        func_lo = MOIU.operate(+, T, func, mode.epsilon)
        c_lo = MOIU.normalize_and_add_constraint(m, func_lo, MOI.GreaterThan(zero(T)))
        MOI.set(m, MOI.ConstraintName(), c_lo, "lower_strong_duality_lo")

        return CI[c_up, c_lo]
    end
end

function add_complement(mode::MixedMode{T}, m, comp::Complement,
        idxmap_primal, idxmap_dual, copy_names::Bool, pass_start::Bool) where T
    _mode = get_mode(mode, comp.constraint, idxmap_primal)
    accept_vector_set(_mode, comp)
    add_complement(_mode, m, comp,
        idxmap_primal, idxmap_dual, copy_names, pass_start)
end

function get_mode(mode::MixedMode{T}, ci::CI{F,S}, map) where {
    T,
    F<:MOI.SingleVariable,
    S<:Union{MOI.EqualTo{T}, MOI.LessThan{T}, MOI.GreaterThan{T}}
}
    # key = map[VI(ci.value)]
    key = VI(ci.value)
    if haskey(mode.constraint_mode_map_v, key)
        return mode.constraint_mode_map_v[key]
    else
        return mode.default
    end
end
function get_mode(mode::MixedMode{T}, ci::CI{F,S}, map) where T where {F,S}
    key = ci#map[ci]
    # @show ci, key, map
    # @show mode.constraint_mode_map_c
    if haskey(mode.constraint_mode_map_c, key)
        return mode.constraint_mode_map_c[key]
    else
        return mode.default
    end
end

function add_complement(mode::ComplementMode{T}, m, comp::Complement,
    idxmap_primal, idxmap_dual, copy_names::Bool, pass_start::Bool) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable
    out_var = VI[]
    out_ctr = CI[]

    with_slack = mode.with_slack

    f_dest = MOIU.map_indices.(Ref(idxmap_primal), f)

    dual = idxmap_dual[v]

    if with_slack
        slack, slack_in_set = MOI.add_constrained_variable(m, s)
        new_f = MOIU.operate(-, T, f_dest, MOI.SingleVariable(slack))
        equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.EqualTo(zero(T)))

        if pass_start
            val = MOIU.eval_variables(
                x-> nothing_to_nan(MOI.get(m, MOI.VariablePrimalStart(), x)), f_dest)
            if !isnan(val)
                MOI.set(m, MOI.VariablePrimalStart(), slack, val)
            end
        end

        c = MOI.add_constraint(m,
            MOI.VectorOfVariables([slack, dual]),
            MOI.Complements(1))
        if copy_names
            nm = MOI.get(m, MOI.VariableName(), dual)
            MOI.set(m, MOI.VariableName(), slack, "slk_($(nm))")
            MOI.set(m, MOI.ConstraintName(), slack_in_set, "bound_slk_($(nm))")
            MOI.set(m, MOI.ConstraintName(), equality, "eq_slk_($(nm))")
            MOI.set(m, MOI.ConstraintName(), c, "compl_complWslk_($(nm))")
        end

        appush!(out_var, slack)
        appush!(out_ctr, slack_in_set)
        appush!(out_ctr, equality)
        appush!(out_ctr, c)
    else
        new_f = MOIU.operate(vcat, T, f_dest, MOI.SingleVariable(dual))

        c = MOI.add_constraint(m,
            new_f,
            MOI.Complements(1))

        if copy_names
            nm = MOI.get(m, MOI.VariableName(), dual)
            MOI.set(m, MOI.ConstraintName(), c, "compl_compl_($(nm))")
        end

        appush!(out_ctr, c)
    end

    return out_var, out_ctr
end

function add_complement(mode::SOS1Mode{T}, m, comp::Complement,
    idxmap_primal, idxmap_dual, copy_names::Bool, pass_start::Bool) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    if comp.is_vec
        error("Vector constraint is not supported by SOS1 mode")
    end

    slack, slack_in_set = MOI.add_constrained_variable(m, s)
    f_dest = MOIU.map_indices.(Ref(idxmap_primal), f)
    new_f = MOIU.operate(-, T, f_dest, MOI.SingleVariable(slack))
    equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.EqualTo(zero(T)))

    dual = idxmap_dual[v]
    c1 = MOI.add_constraint(m,
        MOI.VectorOfVariables([slack, dual]),
        MOI.SOS1([1.0, 2.0]))

    if copy_names
        nm = MOI.get(m, MOI.VariableName(), dual)
        MOI.set(m, MOI.VariableName(), slack, "slk_($(nm))")
        MOI.set(m, MOI.ConstraintName(), slack_in_set, "bound_slk_($(nm))")
        MOI.set(m, MOI.ConstraintName(), equality, "eq_slk_($(nm))")
        MOI.set(m, MOI.ConstraintName(), c1, "compl_sos1_($(nm))")
    end

    return slack, slack_in_set, equality, c1
end

is_equality(set::S) where {S<:MOI.AbstractSet} = false
is_equality(set::MOI.EqualTo{T}) where T = true
is_equality(set::MOI.Zeros) = true

only_variable_functions(v::MOI.VariableIndex) = MOI.SingleVariable(v)
only_variable_functions(v::Vector{MOI.VariableIndex}) = MOI.VectorOfVariables(v)

nothing_to_nan(val) = ifelse(val === nothing, NaN, val)

function add_complement(mode::ProductMode{T}, m, comp::Complement,
    idxmap_primal, idxmap_dual, copy_names::Bool, pass_start::Bool) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    out_var = VI[]
    out_ctr = CI[]

    eps = mode.epsilon
    with_slack = mode.with_slack

    f_dest = MOIU.map_indices(x->idxmap_primal[x], f)

    dual = comp.is_vec ? map(x->idxmap_dual[x], v) : idxmap_dual[v]

    if with_slack
        slack, slack_in_set = if comp.is_vec
            MOI.add_constrained_variables(m, s)
        else
            MOI.add_constrained_variable(m, s)
        end
        new_f = MOIU.operate(-, T, f_dest, only_variable_functions(slack))
        if comp.is_vec
            equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.Zeros(length(slack)))
        else
            equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.EqualTo(zero(T)))
        end

        prod_f = MOIU.operate(dot, T, only_variable_functions(slack), only_variable_functions(dual))

        prod_f1 = MOIU.operate(-, T, prod_f, eps)
        c1 = MOIU.normalize_and_add_constraint(m,
            prod_f1,
            MOI.LessThan{Float64}(0.0))
        if comp.is_vec
            prod_f2 = MOIU.operate(+, T, prod_f, eps)
            c2 = MOIU.normalize_and_add_constraint(m,
                prod_f2,
                MOI.GreaterThan{Float64}(0.0))
        end

        appush!(out_var, slack)
        appush!(out_ctr, slack_in_set)
        appush!(out_ctr, equality)
        appush!(out_ctr, c1)
        if comp.is_vec
            appush!(out_ctr, c2)
        end

        if pass_start
            val = MOIU.eval_variables(
                x-> nothing_to_nan(MOI.get(m, MOI.VariablePrimalStart(), x)), f_dest)
            if comp.is_vec
                for i in eachindex(val)
                    if !isnan(val[i])
                        MOI.set(m, MOI.VariablePrimalStart(), slack[i], val[i])
                    end
                end
            else
                if !isnan(val)
                    MOI.set(m, MOI.VariablePrimalStart(), slack, val)
                end
            end
        end

        if copy_names
            nm = MOI.get(m, MOI.VariableName(), dual)
            MOI.set(m, MOI.VariableName(), slack, "slk_($(nm))")
            MOI.set(m, MOI.ConstraintName(), slack_in_set, "bound_slk_($(nm))")
            MOI.set(m, MOI.ConstraintName(), equality, "eq_slk_($(nm))")
            MOI.set(m, MOI.ConstraintName(), c1, "compl_prodWslk_($(nm))")
            if comp.is_vec
                MOI.set(m, MOI.ConstraintName(), c1, "compl_prodWslk2_($(nm))")
            end
        end
    else
        new_f = MOIU.operate(dot, T, f_dest, only_variable_functions(dual))
        new_f1 = MOIU.operate(-, T, new_f, eps)
        c1 = MOIU.normalize_and_add_constraint(m,
            new_f1,
            MOI.LessThan{T}(0.0))
        if comp.is_vec # conic
            new_f2 = MOIU.operate(+, T, new_f, eps)
            c2 = MOIU.normalize_and_add_constraint(m,
                new_f2,
                MOI.GreaterThan{T}(0.0))
        end

        # TODO(?): if eps == 0 then add equality

        appush!(out_ctr, c1)
        if comp.is_vec
            appush!(out_ctr, c2)
        end

        if copy_names
            nm = if comp.is_vec
                MOI.get.(m, MOI.VariableName(), dual)
            else
                MOI.get(m, MOI.VariableName(), dual)
            end
            MOI.set(m, MOI.ConstraintName(), c1, "compl_prod_($(nm))")
            if comp.is_vec
                MOI.set(m, MOI.ConstraintName(), c2, "compl_prod2_($(nm))")
            end
        end
    end
    return out_var, out_ctr
end

function add_complement(mode::IndicatorMode{T}, m, comp::Complement,
    idxmap_primal, idxmap_dual, copy_names::Bool, pass_start::Bool) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    method = mode.mode

    is_tight = false
    has_start = false

    f_dest = MOIU.map_indices(x->idxmap_primal[x], f)

    dual = idxmap_dual[v]

    if comp.is_vec
        error("Vector constraint is (currently) not supported by indicator mode")
    end

    if copy_names
        nm = if comp.is_vec
            MOI.get.(m, MOI.VariableName(), dual)
        else
            MOI.get(m, MOI.VariableName(), dual)
        end
    end

    vb1 = MOI.add_variable(m)
    if copy_names
        MOI.set(m, MOI.VariableName(), vb1, "compl_bin1_($(nm))")
    end

    cb1 = MOI.add_constraint(m, SVF(vb1), MOI.ZeroOne())
    if method == ONE_ONE || method == ZERO_ZERO
        # second binary
        vb2 = MOI.add_variable(m)
        cb2 = MOI.add_constraint(m, SVF(vb2), MOI.ZeroOne())
        if copy_names
            MOI.set(m, MOI.VariableName(), vb2, "compl_bin2_($(nm))")
        end

        # z1 + z2 == 1
        fb = MOIU.operate(+, T, SVF(vb1), SVF(vb2))
        cb = MOI.add_constraint(m, fb, MOI.EqualTo{T}(one(T)))
        if copy_names
            MOI.set(m, MOI.ConstraintName(), cb, "compl_sum_bin_($(nm))")
        end
    else
        vb2 = vb1
    end

    pre_f1, pre_s1 = MOIU.normalize_constant(f_dest, MOI.EqualTo(zero(T)))
    f1 = MOIU.operate(vcat, T, SVF(vb1), pre_f1)
    f2 = MOIU.operate(vcat, T, SVF(vb2), SVF(dual))

    if pass_start
        val = MOIU.eval_variables(
            x-> nothing_to_nan(MOI.get(m, MOI.VariablePrimalStart(), x)), f_dest)
        if !isnan(val)
            is_tight = abs(val) < 1e-8
            has_start = true
        end
    end

    if method == ONE_ONE
        s1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(pre_s1)
        s2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.EqualTo(zero(T)))
        if pass_start && has_start
            MOI.set(m, MOI.VariablePrimalStart(), vb1, ifelse(is_tight, 1.0, 0.0))
            MOI.set(m, MOI.VariablePrimalStart(), vb2, ifelse(is_tight, 0.0, 1.0))
        end
    elseif method == ZERO_ZERO
        s1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(pre_s1)
        s2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(MOI.EqualTo(zero(T)))
        if pass_start && has_start
            MOI.set(m, MOI.VariablePrimalStart(), vb1, ifelse(is_tight, 0.0, 1.0))
            MOI.set(m, MOI.VariablePrimalStart(), vb2, ifelse(is_tight, 1.0, 0.0))
        end
    else
        s1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(pre_s1)
        s2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(MOI.EqualTo(zero(T)))
        if pass_start && has_start
            MOI.set(m, MOI.VariablePrimalStart(), vb1, ifelse(is_tight, 1.0, 0.0))
        end
    end

    # do not MOIU.normalize_and_add_constraint because are vector functions
    c1 = MOI.add_constraint(m, to_vector_affine(f1), s1)
    c2 = MOI.add_constraint(m, to_vector_affine(f2), s2)

    if copy_names
        MOI.set(m, MOI.ConstraintName(), c1, "compl_ind1_($(nm))")
        MOI.set(m, MOI.ConstraintName(), c2, "compl_ind2_($(nm))")
    end
    return c1
end

function flip_set(set::MOI.LessThan{T}) where T
    return MOI.GreaterThan{T}(0.0)
end
function flip_set(set::MOI.GreaterThan{T}) where T
    return MOI.LessThan{T}(0.0)
end

function get_bounds(var, map, fallback_bound = Inf)
    if haskey(map, var)
        info = map[var]
        # TODO deal with precision and performance
        lower = ifelse(info.lower != -Inf, info.lower, -fallback_bound)
        upper = ifelse(info.upper != +Inf, info.upper, +fallback_bound)
        return IntervalArithmetic.interval(lower, upper)
    elseif 0.0 <= fallback_bound <= Inf
        return IntervalArithmetic.interval(-fallback_bound, fallback_bound)
    else
        error("variable $var has no finite bounds defined")
    end
end

function set_bound(inter::IntervalArithmetic.Interval, ::LT{T}) where T
    return inter.hi
end
function set_bound(inter::IntervalArithmetic.Interval, ::GT{T}) where T
    return inter.lo
end

function add_complement(mode::FortunyAmatMcCarlMode{T}, m, comp::Complement,
    idxmap_primal, idxmap_dual, copy_names::Bool, pass_start::Bool) where T

    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    is_tight = false
    has_start = false

    if mode.with_slack
        slack, slack_in_set = MOI.add_constrained_variable(m, s)
    end
    f_dest = MOIU.map_indices.(Ref(idxmap_primal), f)

    f_bounds = MOIU.eval_variables(vi -> get_bounds(vi, mode.cache.map, mode.primal_big_M), f_dest)

    if pass_start
        val = MOIU.eval_variables(
            x-> nothing_to_nan(MOI.get(m, MOI.VariablePrimalStart(), x)), f_dest)
        if !isnan(val)
            is_tight = abs(val) < 1e-8
            has_start = true
        end
    end

    if mode.with_slack
        new_f = MOIU.operate(-, T, f_dest, MOI.SingleVariable(slack))
        equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.EqualTo(zero(T)))
        if pass_start && has_start
            MOI.set(m, MOI.VariablePrimalStart(), slack, val)
        end
    end

    dual = idxmap_dual[v]
    v_bounds = get_bounds(dual, mode.cache.map, mode.dual_big_M)

    bin = MOI.add_variable(m)
    if pass_start && has_start && is_tight
        MOI.set(m, MOI.VariablePrimalStart(), bin, 1.0)
        MOI.set(m, MOI.VariablePrimalStart(), dual, 0.0)
    else
        MOI.set(m, MOI.VariablePrimalStart(), bin, 0.0)
    end

    s1 = flip_set(s)
    s2 = flip_set(s)

    Ms = set_bound(f_bounds, s1)
    Mv = set_bound(v_bounds, s2)

    if isnan(Ms) || abs(Ms) >= Inf || isnan(Mv) || abs(Mv) >= Inf
        error("It was not possible to automatically compute bounds"*
            " for a complementarity pair, please add the arguments"*
            " primal_big_M and dual_big_M to FortunyAmatMcCarlMode")
    end

    if mode.with_slack
        f1 = MOI.ScalarAffineFunction{T}(
            MOI.ScalarAffineTerm{T}.(
                [one(T), -Ms], [slack, bin]
            ),
            0.0
        )
    else
        push!(f_dest.terms, MOI.ScalarAffineTerm{T}(-Ms, bin))
        f1 = f_dest
    end

    f2 = MOI.ScalarAffineFunction{T}(
        MOI.ScalarAffineTerm{T}.(
            [one(T), Mv], [dual, bin]
        ),
        -Mv
    )

    c1 = MOIU.normalize_and_add_constraint(m, f1, s2)
    c2 = MOIU.normalize_and_add_constraint(m, f2, s2)
    c3 = MOI.add_constraint(m, MOI.SingleVariable(bin), MOI.ZeroOne())

    if copy_names
        nm = MOI.get(m, MOI.VariableName(), dual)
        if mode.with_slack
            MOI.set(m, MOI.VariableName(), slack, "slk_($(nm))")
            MOI.set(m, MOI.ConstraintName(), slack_in_set, "bound_slk_($(nm))")
            MOI.set(m, MOI.ConstraintName(), equality, "eq_slk_($(nm))")
        end
        MOI.set(m, MOI.VariableName(), bin, "bin_($(nm))")
        MOI.set(m, MOI.ConstraintName(), c1, "compl_fa_sl_($(nm))")
        MOI.set(m, MOI.ConstraintName(), c2, "compl_fa_dl_($(nm))")
        MOI.set(m, MOI.ConstraintName(), c2, "compl_fa_bn_($(nm))")
    end

    # if mode.with_slack
    #     return slack, slack_in_set, equality, c1
    # else
    # end
    return c1
end

function to_vector_affine(f::MOI.VectorAffineFunction{T}) where T
    return f
end
function to_vector_affine(f::MOI.VectorOfVariables)
    return MOI.VectorAffineFunction{Float64}(f)
end

function handle_lower_objective_sense(lower::MOI.ModelLike)
    lower_objective_sense = MOI.get(lower, MOI.ObjectiveSense())
    if lower_objective_sense == MOI.FEASIBILITY_SENSE
        throw(ErrorException("Lower level models with objective_sense: " *
                            lower_objective_sense *
                            " are not supported."))
    end
    return
end


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;