text:
"""
    SerendipityRefFE(::Type{T},p::Polytope,order::Int) where T

Returns an instance of `LagrangianRefFE`, whose underlying approximation space
is the serendipity space of order `order`. Implemented for order from 1 to 4.
The type of the polytope `p` has to implement all the queries detailed in the
constructor [`LagrangianRefFE(::Type{T},p::Polytope{D},orders) where {T,D}`](@ref)
plus this additional method:

- [`is_serendipity_compatible(p::Polytope)`](@ref)

# Examples

```jldoctest
using InplaceArrays.ReferenceFEs

order = 2
reffe = SerendipityRefFE(Float64,QUAD,order)

println( num_dofs(reffe) )

# output
8

```
"""
function SerendipityRefFE(::Type{T},p::Polytope,order::Int) where T
  @assert is_serendipity_compatible(p) "Polytope not compatible with serendipity elements"
  if order > 0
    sp = SerendipityPolytope(p)
  else
    sp = p
  end
  LagrangianRefFE(T,sp,order)
end

"""
    is_serendipity_compatible(p::Polytope) -> Bool

Returns `true` if the polytope `p` is compatible with the
serendipity space (i.e., it is a n-cube). This method is implemented
for `ExtrusionPolytope`, and should be implemented by new
polytope types if they are to be used to build serendipity spaces.
"""
function is_serendipity_compatible(p::Polytope)
  @abstractmethod
end

# Concrete implementation for ExtrusionPolytope

function is_serendipity_compatible(p::ExtrusionPolytope)
  all(p.extrusion.array .== HEX_AXIS)
end

# Helper private type
struct SerendipityPolytope{D,P} <: Polytope{D}
  hex::P
  SerendipityPolytope(p::Polytope{D}) where D = new{D,typeof(p)}(p)
end

# Implemented polytope interface

function get_faces(p::SerendipityPolytope)
  get_faces(p.hex)
end

function get_dimrange(p::SerendipityPolytope)
  get_dimrange(p.hex)
end

function Polytope{N}(p::SerendipityPolytope,Nfaceid::Integer) where N
  face_hex = Polytope{N}(p.hex, Nfaceid)
  SerendipityPolytope(face_hex)
end

function Polytope{D}(p::SerendipityPolytope{D},Dfaceid::Integer) where D
  @assert Dfaceid == 1 "Only one D-face"
  p
end

function get_vertex_coordinates(p::SerendipityPolytope)
  get_vertex_coordinates(p.hex)
end

function (==)(a::SerendipityPolytope{D},b::SerendipityPolytope{D}) where D
  a.hex == b.hex
end

function get_vertex_permutations(p::SerendipityPolytope)
  get_vertex_permutations(p.hex)
end

# Implemented polytope interface for LagrangianRefFEs

function _s_filter(e,order)
  sum( [ i for i in e if i>1 ] ) <= order
end

function compute_monomial_basis(::Type{T},p::SerendipityPolytope{D},orders) where {T,D}
  MonomialBasis{D}(T,orders,_s_filter)
end

function compute_own_nodes(p::SerendipityPolytope{0},orders)
  compute_own_nodes(p.hex,orders)
end

function compute_own_nodes(p::SerendipityPolytope{1},orders)
  compute_own_nodes(p.hex,orders)
end

function compute_own_nodes(p::SerendipityPolytope{2},orders)
  order, = orders
  if order == 4
    o = (2,2)
  elseif order in (0,1,2,3)
    o=(1,1)
  else
    @unreachable "Serendipity elements only up to order 4"
  end
  compute_own_nodes(p.hex,o)
end

function compute_own_nodes(p::SerendipityPolytope{3},orders)
  Point{3,Float64}[]
end

function compute_own_nodes(p::SerendipityPolytope,orders)
  @unreachable "Serendipity elements only up to 3d"
end

function compute_face_orders(
  p::SerendipityPolytope{D},face::SerendipityPolytope{N},iface::Int,orders) where {D,N}
  compute_face_orders(p.hex,face.hex,iface,orders)
end


text:
module SameReflect

using Test
using Configurations
using Configurations: has_same_reflect_field

@option struct OptionA
    type::Reflect
end

@option struct OptionB
    type::Reflect
end

@option struct OptionC
    type::Reflect
end

@option struct OptionD
    name::Reflect
end

@option struct OptionE end

@option struct Composite
    data::Union{OptionA,OptionB,OptionC}
end

@testset "has_same_reflect_field" begin
    @test has_same_reflect_field([OptionA, OptionB, OptionC]) == true
    @test has_same_reflect_field([OptionA, OptionB, OptionC, OptionD]) == false
    @test has_same_reflect_field([OptionA, OptionB, OptionC, OptionE]) == false
end

end

using Configurations
using Configurations: from_dict_specialize, from_dict_generated

d = Dict{String,Any}("data" => Dict{String,Any}("type" => "Main.SameReflect.OptionA"))
from_dict(SameReflect.Composite, d)

d = Dict{String,Any}("data" => Dict{String,Any}("type" => "Main.SameReflect.OptionB"))
from_dict(SameReflect.Composite, d)

@time from_dict(SameReflect.Composite, d)

@time from_dict_specialize(SameReflect.Composite, d)

Configurations.field_default(SameReflect.Composite, :data)

from_dict_generated(SameReflect.Composite, :d)
from_dict_generated(SameReflect.OptionA, :d)

@generated function goo(::Type{T}) where {T}
    quote
        println("hello")
    end
end

@generated function foo()
    quote
        goo()
    end
end

foo()

ext: jl


text:
@get_units @with_kw struct SurfaceFlow{T,R,L}
    β::T | "-"                              # constant in Manning's equation
    sl::Vector{T} | "m m-1"                 # Slope [m m⁻¹]
    n::Vector{T} | "s m-1/3"                # Manning's roughness [s m⁻⅓]
    dl::Vector{T} | "m"                     # Drain length [m]
    q::Vector{T} | "m3 s-1"                 # Discharge [m³ s⁻¹]
    qin::Vector{T} | "m3 s-1"               # Inflow from upstream cells [m³ s⁻¹]
    q_av::Vector{T} | "m3 s-1"              # Average discharge [m³ s⁻¹]
    qlat::Vector{T} | "m2 s-1"              # Lateral inflow per unit length [m² s⁻¹]
    inwater::Vector{T} | "m3 s-1"           # Lateral inflow [m³ s⁻¹]
    inflow::Vector{T} | "m3 s-1"            # External inflow (abstraction/supply/demand) [m³ s⁻¹]
    volume::Vector{T} | "m3"                # Kinematic wave volume [m³] (based on water level h)
    h::Vector{T} | "m"                      # Water level [m]
    h_av::Vector{T} | "m"                   # Average water level [m]
    bankfull_depth::Vector{T} | "m"         # Bankfull water level [m]
    Δt::T | "s"                             # Model time step [s]
    its::Int | "-"                          # Number of fixed iterations
    width::Vector{T} | "m"                  # Flow width [m]
    alpha_pow::T | "-"                      # Used in the power part of α
    alpha_term::Vector{T} | "-"             # Term used in computation of α
    α::Vector{T} | "s3/5 m1/5"              # Constant in momentum equation A = αQᵝ, based on Manning's equation
    cel::Vector{T} | "m s-1"                # Celerity of the kinematic wave
    to_river::Vector{T} | "m3 s-1"          # Part of overland flow [m³ s⁻¹] that flows to the river
    rivercells::Vector{Bool} | "-"          # Location of river cells (0 or 1)
    wb_pit::Vector{Bool} | "-"              # Boolean location (0 or 1) of a waterbody (wb, reservoir or lake).
    reservoir_index::Vector{Int} | "-"      # map cell to 0 (no reservoir) or i (pick reservoir i in reservoir field)
    lake_index::Vector{Int} | "-"           # map cell to 0 (no lake) or i (pick lake i in lake field)
    reservoir::R                            # Reservoir model struct of arrays
    lake::L                                 # Lake model struct of arrays
    kinwave_it::Bool                        # Boolean for iterations kinematic wave

    # TODO unclear why this causes a MethodError
    # function SurfaceFlow{T,R,L}(args...) where {T,R,L}
    #     equal_size_vectors(args)
    #     return new(args...)
    # end
end

function initialize_surfaceflow_land(
    nc,
    config,
    inds;
    sl,
    dl,
    width,
    wb_pit,
    iterate,
    tstep,
    Δt,
)
    @info "Kinematic wave approach is used for overland flow." iterate
    if tstep > 0
        @info "Using a fixed sub-timestep (seconds) $tstep for kinematic wave overland flow."
    end

    n_land = ncread(
        nc,
        config,
        "lateral.land.n";
        sel = inds,
        defaults = 0.072,
        type = Float,
    )
    n = length(inds)

    sf_land = SurfaceFlow(
        β = Float(0.6),
        sl = sl,
        n = n_land,
        dl = dl,
        q = zeros(Float, n),
        qin = zeros(Float, n),
        q_av = zeros(Float, n),
        qlat = zeros(Float, n),
        inwater = zeros(Float, n),
        inflow = zeros(Float, n),
        volume = zeros(Float, n),
        h = zeros(Float, n),
        h_av = zeros(Float, n),
        bankfull_depth = zeros(Float, n),
        Δt = Float(tosecond(Δt)),
        its = tstep > 0 ? Int(cld(tosecond(Δt), tstep)) : tstep,
        width = width,
        wb_pit = wb_pit,
        alpha_pow = Float((2.0 / 3.0) * 0.6),
        alpha_term = fill(mv, n),
        α = fill(mv, n),
        cel = zeros(Float, n),
        to_river = zeros(Float, n),
        rivercells = fill(false, n),
        reservoir_index = fill(0, n),
        lake_index = fill(0, n),
        reservoir = nothing,
        lake = nothing,
        kinwave_it = iterate,
    )

    return sf_land
end

function initialize_surfaceflow_river(
    nc,
    config,
    inds;
    dl,
    width,
    wb_pit,
    reservoir_index,
    reservoir,
    lake_index,
    lake,
    river,
    iterate,
    tstep,
    Δt,
)
    @info "Kinematic wave approach is used for river flow." iterate
    if tstep > 0
        @info "Using a fixed sub-timestep (seconds) $tstep for kinematic wave river flow."
    end

    n_river = ncread(
        nc,
        config,
        "lateral.river.n";
        sel = inds,
        defaults = 0.036,
        type = Float,
    )
    bankfull_depth = ncread(
        nc,
        config,
        "lateral.river.bankfull_depth";
        alias = "lateral.river.h_bankfull",
        sel = inds,
        defaults = 1.0,
        type = Float,
    )
    if haskey(config.input.lateral.river, "h_bankfull")
        @warn string(
            "The `h_bankfull` key in `[input.lateral.river]` is now called ",
            "`bankfull_depth`. Please update your TOML file.",
        )
    end
    sl = ncread(
        nc,
        config,
        "lateral.river.slope";
        optional = false,
        sel = inds,
        type = Float,
    )
    clamp!(sl, 0.00001, Inf)

    n = length(inds)

    sf_river = SurfaceFlow(
        β = Float(0.6),
        sl = sl,
        n = n_river,
        dl = dl,
        q = zeros(Float, n),
        qin = zeros(Float, n),
        q_av = zeros(Float, n),
        qlat = zeros(Float, n),
        inwater = zeros(Float, n),
        inflow = zeros(Float, n),
        volume = zeros(Float, n),
        h = zeros(Float, n),
        h_av = zeros(Float, n),
        bankfull_depth = bankfull_depth,
        Δt = Float(tosecond(Δt)),
        its = tstep > 0 ? Int(cld(tosecond(Δt), tstep)) : tstep,
        width = width,
        wb_pit = wb_pit,
        alpha_pow = Float((2.0 / 3.0) * 0.6),
        alpha_term = fill(mv, n),
        α = fill(mv, n),
        cel = zeros(Float, n),
        to_river = zeros(Float, n),
        reservoir_index = reservoir_index,
        lake_index = lake_index,
        reservoir = reservoir,
        lake = lake,
        rivercells = river,
        kinwave_it = iterate,
    )

    return sf_river
end

statevars(::SurfaceFlow) = (:q, :h, :h_av)

function update(
    sf::SurfaceFlow,
    network;
    frac_toriver = nothing,
    inflow_wb = nothing,
    doy = 0,
)
    @unpack graph,
    order,
    subdomain_order,
    topo_subdomain,
    indices_subdomain,
    upstream_nodes = network

    n = length(order)
    ns = length(subdomain_order)

    @. sf.alpha_term = pow(sf.n / sqrt(sf.sl), sf.β)
    # use fixed alpha value based on 0.5 * bankfull_depth
    @. sf.α = sf.alpha_term * pow(sf.width + sf.bankfull_depth, sf.alpha_pow)

    @. sf.qlat = sf.inwater / sf.dl

    # two options for iteration, fixed or based on courant number.
    if sf.kinwave_it
        if sf.its > 0
            its = sf.its
        else
            # calculate celerity
            courant = zeros(n)
            for k = 1:ns
                for m in subdomain_order[k]
                    for v in topo_subdomain[m]
                        if sf.q[v] > 0.0
                            sf.cel[v] = 1.0 / (sf.α[v] * sf.β * pow(sf.q[v], (sf.β - 1.0)))
                            courant[v] = (sf.Δt / sf.dl[v]) * sf.cel[v]
                        end
                    end
                end
            end
            filter!(x -> x ≠ 0.0, courant)
            its = isempty(courant) ? 1 : ceil(Int, (1.25 * quantile!(courant, 0.95)))
        end
    else
        its = 1
    end

    # sub time step
    adt = sf.Δt / its

    sf.q_av .= 0.0
    sf.h_av .= 0.0
    sf.to_river .= 0.0
    # because of possible iterations set reservoir and lake inflow and total outflow at
    # start to zero, the total sum of inflow and outflow at each sub time step is calculated
    if !isnothing(sf.reservoir)
        sf.reservoir.inflow .= 0.0
        sf.reservoir.totaloutflow .= 0.0
    end
    if !isnothing(sf.lake)
        sf.lake.inflow .= 0.0
        sf.lake.totaloutflow .= 0.0
    end

    for _ = 1:its
        sf.qin .= 0.0
        for k = 1:ns
            @threads for m in subdomain_order[k]
                for (n, v) in zip(indices_subdomain[m], topo_subdomain[m])

                    # for overland flow frac_toriver needs to be defined
                    if frac_toriver !== nothing # run kinematic wave for land domain
                        # for a river cell without a reservoir or lake (wb_pit is false) part of the
                        # upstream surface flow goes to the river (frac_toriver) and part goes to
                        # the surface flow reservoir (1.0 - frac_toriver), upstream nodes with a
                        # reservoir or lake are excluded
                        sf.to_river[v] += sum_at(
                            i -> sf.q[i] * frac_toriver[i],
                            upstream_nodes[n],
                            eltype(sf.to_river),
                        )
                        if sf.width[v] > 0.0
                            sf.qin[v] = sum_at(
                                i -> sf.q[i] * (1.0 - frac_toriver[i]),
                                upstream_nodes[n],
                                eltype(sf.q),
                            )
                        end
                    else # run kinematic wave for river domain (including reservoirs and lakes)
                        # sf.qin by outflow from upstream reservoir or lake location is added
                        sf.qin[v] += sum_at(sf.q, upstream_nodes[n])
                    end

                    # Inflow supply/abstraction is added to qlat (divide by flow length)
                    # If inflow < 0, abstraction is limited
                    if sf.inflow[v] < 0.0
                        max_abstract = min(sf.qin[v] + sf.qlat[v] * sf.dl[v], -sf.inflow[v])
                        inflow = -max_abstract / sf.dl[v]
                    else
                        inflow = sf.inflow[v] / sf.dl[v]
                    end

                    sf.q[v] = kinematic_wave(
                        sf.qin[v],
                        sf.q[v],
                        sf.qlat[v] + inflow,
                        sf.α[v],
                        sf.β,
                        adt,
                        sf.dl[v],
                    )

                    if !isnothing(sf.reservoir) && sf.reservoir_index[v] != 0
                        # run reservoir model and copy reservoir outflow to inflow (qin) of
                        # downstream river cell
                        i = sf.reservoir_index[v]
                        update(sf.reservoir, i, sf.q[v] + inflow_wb[v], adt)

                        downstream_nodes = outneighbors(graph, v)
                        n_downstream = length(downstream_nodes)
                        if n_downstream == 1
                            j = only(downstream_nodes)
                            sf.qin[j] = sf.reservoir.outflow[i]
                        elseif n_downstream == 0
                            error(
                                """A reservoir without a downstream river node is not supported.
                                Add a downstream river node or move the reservoir to an upstream node (model schematization).
                                """,
                            )
                        else
                            error("bifurcations not supported")
                        end

                    elseif !isnothing(sf.lake) && sf.lake_index[v] != 0
                        # run lake model and copy lake outflow to inflow (qin) of downstream river
                        # cell
                        i = sf.lake_index[v]
                        update(sf.lake, i, sf.q[v] + inflow_wb[v], doy, adt)

                        downstream_nodes = outneighbors(graph, v)
                        n_downstream = length(downstream_nodes)
                        if n_downstream == 1
                            j = only(downstream_nodes)
                            sf.qin[j] = sf.lake.outflow[i]
                        elseif n_downstream == 0
                            error(
                                """A lake without a downstream river node is not supported.
                                Add a downstream river node or move the lake to an upstream node (model schematization).
                                """,
                            )
                        else
                            error("bifurcations not supported")
                        end
                    end

                    # update h, only if surface width (overland flow) > 0.0
                    if sf.width[v] > 0.0
                        crossarea = sf.α[v] * pow(sf.q[v], sf.β)
                        sf.h[v] = crossarea / sf.width[v]
                    end

                    sf.q_av[v] += sf.q[v]
                    sf.h_av[v] += sf.h[v]
                end

            end
        end
    end
    sf.q_av ./= its
    sf.h_av ./= its
    sf.to_river ./= its
    sf.volume .= sf.dl .* sf.width .* sf.h
end

@get_units @with_kw struct LateralSSF{T}
    kh₀::Vector{T} | "m d-1"               # Horizontal hydraulic conductivity at soil surface [m d⁻¹]
    f::Vector{T} | "m-1"                   # A scaling parameter [m⁻¹] (controls exponential decline of kh₀)
    soilthickness::Vector{T} | "m"         # Soil thickness [m]
    θₛ::Vector{T} | "-"                     # Saturated water content (porosity) [-]
    θᵣ::Vector{T} | "-"                    # Residual water content [-]
    Δt::T | "d"                            # model time step [d]
    βₗ::Vector{T} | "m m-1"                 # Slope [m m⁻¹]
    dl::Vector{T} | "m"                    # Drain length [m]
    dw::Vector{T} | "m"                    # Flow width [m]
    zi::Vector{T} | "m"                    # Pseudo-water table depth [m] (top of the saturated zone)
    exfiltwater::Vector{T} | "m Δt-1"      # Exfiltration [m Δt⁻¹] (groundwater above surface level, saturated excess conditions)
    recharge::Vector{T} | "m Δt-1"         # Net recharge to saturated store [m Δt⁻¹]
    ssf::Vector{T} | "m3 d-1"              # Subsurface flow [m³ d⁻¹]
    ssfin::Vector{T} | "m3 d-1"            # Inflow from upstream cells [m³ d⁻¹]
    ssfmax::Vector{T} | "m2 d-1"           # Maximum subsurface flow [m² d⁻¹]
    to_river::Vector{T} | "m3 d-1"         # Part of subsurface flow [m³ d⁻¹] that flows to the river
    wb_pit::Vector{Bool} | "-"             # Boolean location (0 or 1) of a waterbody (wb, reservoir or lake).

    function LateralSSF{T}(args...) where {T}
        equal_size_vectors(args)
        return new(args...)
    end
end

statevars(::LateralSSF) = (:ssf,)

function update(ssf::LateralSSF, network, frac_toriver)
    @unpack subdomain_order, topo_subdomain, indices_subdomain, upstream_nodes = network

    ns = length(subdomain_order)
    for k = 1:ns
        @threads for m in subdomain_order[k]
            for (n, v) in zip(indices_subdomain[m], topo_subdomain[m])
                # for a river cell without a reservoir or lake (wb_pit is false) part of the
                # upstream subsurface flow goes to the river (frac_toriver) and part goes to the
                # subsurface flow reservoir (1.0 - frac_toriver) upstream nodes with a reservoir or
                # lake are excluded
                ssf.ssfin[v] = sum_at(
                    i -> ssf.ssf[i] * (1.0 - frac_toriver[i]),
                    upstream_nodes[n],
                    eltype(ssf.ssfin),
                )
                ssf.to_river[v] = sum_at(
                    i -> ssf.ssf[i] * frac_toriver[i],
                    upstream_nodes[n],
                    eltype(ssf.to_river),
                )

                ssf.ssf[v], ssf.zi[v], ssf.exfiltwater[v] = kinematic_wave_ssf(
                    ssf.ssfin[v],
                    ssf.ssf[v],
                    ssf.zi[v],
                    ssf.recharge[v],
                    ssf.kh₀[v],
                    ssf.βₗ[v],
                    ssf.θₛ[v] - ssf.θᵣ[v],
                    ssf.f[v],
                    ssf.soilthickness[v],
                    ssf.Δt,
                    ssf.dl[v],
                    ssf.dw[v],
                    ssf.ssfmax[v],
                )
            end
        end
    end
end

@get_units @with_kw struct GroundwaterExchange{T}
    Δt::T | "d"                         # model time step [d]
    exfiltwater::Vector{T} | "m Δt-1"   # Exfiltration [m Δt⁻¹]  (groundwater above surface level, saturated excess conditions)
    zi::Vector{T} | "m"                 # Pseudo-water table depth [m] (top of the saturated zone)
    to_river::Vector{T} | "m3 d-1"      # Part of subsurface flow [m³ d⁻¹] that flows to the river
    ssf::Vector{T} | "m3 d-1"           # Subsurface flow [m³ d⁻¹]
end

@get_units @with_kw struct ShallowWaterRiver{T,R,L}
    n::Int | "-"                        # number of cells
    ne::Int | "-"                       # number of edges/links
    g::T | "m s-2"                      # acceleration due to gravity
    α::T | "-"                          # stability coefficient (Bates et al., 2010)
    h_thresh::T | "m"                   # depth threshold for calculating flow
    Δt::T | "s"                         # model time step [s]
    q::Vector{T} | "m3 s-1"             # river discharge (subgrid channel)
    q0::Vector{T} | "m3 s-1"            # river discharge (subgrid channel) at previous time step
    q_av::Vector{T} | "m3 s-1"          # average river discharge [m³ s⁻¹]
    zb_max::Vector{T} | "m"             # maximum channel bed elevation
    mannings_n::Vector{T} | "s m-1/3"   # Manning's roughness at edge/link
    h::Vector{T} | "m"                  # water depth
    η_max::Vector{T} | "m"              # maximum water elevation
    hf::Vector{T} | "m"                 # water depth at edge/link
    h_av::Vector{T} | "m"               # average water depth
    dl::Vector{T} | "m"                 # river length
    dl_at_link::Vector{T} | "m"         # river length at edge/link
    width::Vector{T} | "m"              # river width
    width_at_link::Vector{T} | "m"      # river width at edge/link
    a::Vector{T} | "m2"                 # flow area at edge/link
    r::Vector{T} | "m"                  # wetted perimeter at edge/link
    volume::Vector{T} | "m3"            # river volume
    error::Vector{T} | "m3"             # error volume
    inwater::Vector{T} | "m3 s-1"       # lateral inflow [m³ s⁻¹]
    inflow::Vector{T} | "m3 s-1"        # External inflow (abstraction/supply/demand) [m³ s⁻¹]
    bankfull_volume::Vector{T} | "m3"   # bankfull volume
    bankfull_depth::Vector{T} | "m"     # bankfull depth
    zb::Vector{T} | "m"                 # river bed elevation
    froude_limit::Bool | "-"            # if true a check is performed if froude number > 1.0 (algorithm is modified)
    reservoir_index::Vector{Int} | "-"  # map cell to 0 (no reservoir) or i (pick reservoir i in reservoir field)
    lake_index::Vector{Int} | "-"       # map cell to 0 (no lake) or i (pick lake i in lake field)
    reservoir::R                        # Reservoir model struct of arrays
    lake::L                             # Lake model struct of arrays
end

function initialize_shallowwater_river(
    nc,
    config,
    inds;
    graph,
    ldd,
    dl,
    width,
    reservoir_index,
    reservoir,
    lake_index,
    lake,
    Δt,
)
    # The local inertial approach makes use of a staggered grid (Bates et al. (2010)),
    # with nodes and links. This information is extracted from the directed graph of the
    # river. Discharge q is calculated at links between nodes and mapped to the source
    # nodes for gridded output (index of link is equal to source node index, e.g.:
    # Edge 1 => 5
    # Edge 2 => 1
    # Edge 3 => 2
    # Edge 4 => 9
    # ⋮ )
    riverlength_bc = get(config.model, "riverlength_bc", 1.0e04)::Float64 # river length at boundary point (ghost point)
    alpha = get(config.model, "inertial_flow_alpha", 0.7)::Float64 # stability coefficient for model time step (0.2-0.7)
    h_thresh = get(config.model, "h_thresh", 1.0e-03)::Float64 # depth threshold for flow at link
    froude_limit = get(config.model, "froude_limit", true)::Bool # limit flow to subcritical according to Froude number

    @info "Local inertial approach is used for river flow." alpha h_thresh froude_limit riverlength_bc

    bankfull_elevation_2d = ncread(
        nc,
        config,
        "lateral.river.bankfull_elevation";
        optional = false,
        type = Float,
        fill = 0,
    )
    bankfull_depth_2d = ncread(
        nc,
        config,
        "lateral.river.bankfull_depth";
        optional = false,
        type = Float,
        fill = 0,
    )
    bankfull_depth = bankfull_depth_2d[inds]
    zb = bankfull_elevation_2d[inds] - bankfull_depth # river bed elevation

    bankfull_volume = bankfull_depth .* width .* dl

    n_river = ncread(
        nc,
        config,
        "lateral.river.n";
        sel = inds,
        defaults = 0.036,
        type = Float,
    )

    n = length(inds)
    # set ghost points for boundary condition (downstream river outlet): river width, bed
    # elevation, manning n is copied from the upstream cell, river depth h is set at 0.0
    # (fixed). river length at boundary point is by default 1.0e4 m (riverlength_bc).
    index_pit = findall(x -> x == 5, ldd)
    npits = length(index_pit)
    add_vertex_edge_graph!(graph, index_pit)
    append!(zb, zb[index_pit])
    append!(dl, fill(riverlength_bc, npits))
    append!(width, width[index_pit])
    append!(n_river, n_river[index_pit])

    # for each link the src and dst node is required
    nodes_at_link = adjacent_nodes_at_link(graph)
    _ne = ne(graph)

    # determine z, width, length and manning's n at links
    zb_max = fill(Float(0), _ne)
    width_at_link = fill(Float(0), _ne)
    length_at_link = fill(Float(0), _ne)
    mannings_n = fill(Float(0), _ne)
    for i = 1:_ne
        src_node = nodes_at_link.src[i]
        dst_node = nodes_at_link.dst[i]
        zb_max[i] = max(zb[src_node], zb[dst_node])
        width_at_link[i] = min(width[src_node], width[dst_node])
        length_at_link[i] = 0.5 * (dl[dst_node] + dl[src_node])
        mannings_n[i] =
            (n_river[dst_node] * dl[dst_node] + n_river[src_node] * dl[src_node]) /
            (dl[dst_node] + dl[src_node])
    end

    # set depth h of reservoir and lake location to bankfull depth
    h = fill(0.0, n + length(index_pit))
    if !isnothing(reservoir)
        inds_reservoir = findall(>(0), reservoir_index)
        h[inds_reservoir] = bankfull_depth[inds_reservoir]
    end
    if !isnothing(lake)
        inds_lake = findall(>(0), lake_index)
        h[inds_lake] = bankfull_depth[inds_lake]
    end

    sw_river = ShallowWaterRiver(
        n = n,
        ne = _ne,
        g = 9.80665,
        α = alpha,
        h_thresh = h_thresh,
        Δt = tosecond(Δt),
        q = zeros(_ne),
        q0 = zeros(_ne),
        q_av = zeros(_ne),
        zb_max = zb_max,
        mannings_n = mannings_n,
        h = fill(0.0, n + length(index_pit)),
        η_max = zeros(_ne),
        hf = zeros(_ne),
        h_av = zeros(n),
        width = width,
        width_at_link = width_at_link,
        a = zeros(_ne),
        r = zeros(_ne),
        volume = fill(0.0, n),
        error = zeros(n),
        inflow = zeros(n),
        inwater = zeros(n),
        dl = dl,
        dl_at_link = length_at_link,
        bankfull_volume = bankfull_volume,
        bankfull_depth = bankfull_depth,
        zb = zb,
        froude_limit = froude_limit,
        reservoir_index = reservoir_index,
        lake_index = lake_index,
        reservoir = reservoir,
        lake = lake,
    )
    return sw_river, nodes_at_link
end

function shallowwater_river_update(
    sw::ShallowWaterRiver,
    network,
    Δt,
    inflow_wb,
    doy,
    update_h,
)

    @unpack nodes_at_link, links_at_node = network

    sw.q0 .= sw.q
    @threads for i = 1:sw.ne
        # For reservoir and lake locations the local inertial solution is replaced by the
        # reservoir or lake model. These locations are handled as boundary conditions in the
        # local inertial model (fixed h).
        if !isnothing(sw.reservoir) && sw.reservoir_index[i] != 0
            v = sw.reservoir_index[i]
            update(sw.reservoir, v, sum_at(sw.q0, links_at_node.src[i]) + inflow_wb[i], Δt)
            sw.q[i] = sw.reservoir.outflow[v]
        elseif !isnothing(sw.lake) && sw.lake_index[i] != 0
            v = sw.lake_index[i]
            update(sw.lake, v, sum_at(sw.q0, links_at_node.src[i]) + inflow_wb[i], doy, Δt)
            sw.q[i] = sw.lake.outflow[v]
        else
            ηsrc = sw.zb[nodes_at_link.src[i]] + sw.h[nodes_at_link.src[i]]
            ηdst = sw.zb[nodes_at_link.dst[i]] + sw.h[nodes_at_link.dst[i]]

            sw.η_max[i] = max(ηsrc, ηdst)
            sw.hf[i] = (sw.η_max[i] - sw.zb_max[i])

            if sw.hf[i] > sw.h_thresh
                sw.a[i] = sw.width_at_link[i] * sw.hf[i] # cross area (rectangular channel)
                sw.r[i] = sw.a[i] / (sw.width_at_link[i] + 2.0 * sw.hf[i]) # wetted perimeter (rectangular channel)
                sw.q[i] = local_inertial_riverflow(
                    sw.q0[i],
                    ηsrc,
                    ηdst,
                    sw.hf[i],
                    sw.a[i],
                    sw.r[i],
                    sw.dl_at_link[i],
                    sw.mannings_n[i],
                    sw.g,
                    sw.froude_limit,
                    Δt,
                )
            else
                sw.q[i] = 0.0
            end
            sw.q_av[i] += sw.q[i] * Δt
        end
    end
    if update_h
        @threads for i = 1:sw.n
            if sw.reservoir_index[i] == 0 || sw.lake_index[i] == 0
                sw.volume[i] =
                    sw.volume[i] +
                    (
                        sum_at(sw.q, links_at_node.src[i]) -
                        sum_at(sw.q, links_at_node.dst[i]) + sw.inwater[i]
                    ) * Δt
                if sw.volume[i] < 0.0
                    sw.error[i] = sw.error[i] + abs(sw.volume[i])
                    sw.volume[i] = 0.0 # set volume to zero
                end
                sw.volume[i] = max(sw.volume[i] + sw.inflow[i] * Δt, 0.0) # add external inflow
                sw.h[i] = sw.volume[i] / (sw.dl[i] * sw.width[i])
                sw.h_av[i] += sw.h[i] * Δt
            end
        end
    end
end

function update(
    sw::ShallowWaterRiver{T},
    network;
    inflow_wb = nothing,
    doy = 0,
    update_h = true,
) where {T}
    @unpack nodes_at_link, links_at_node = network

    if !isnothing(sw.reservoir)
        sw.reservoir.inflow .= 0.0
        sw.reservoir.totaloutflow .= 0.0
    end
    if !isnothing(sw.lake)
        sw.lake.inflow .= 0.0
        sw.lake.totaloutflow .= 0.0
    end
    sw.q_av .= 0.0
    sw.h_av .= 0.0

    t = T(0.0)
    while t < sw.Δt
        Δt = stable_timestep(sw)
        if t + Δt > sw.Δt
            Δt = sw.Δt - t
        end
        shallowwater_river_update(sw, network, Δt, inflow_wb, doy, update_h)
        t = t + Δt
    end
    sw.q_av ./= sw.Δt
    sw.h_av ./= sw.Δt

    return nothing
end

# Stores links in x and y direction between cells of a Vector with CartesianIndex(x, y), for
# staggered grid calculations.
@with_kw struct Indices
    xu::Vector{Int}     # index of neighbor cell in the (+1, 0) direction
    xd::Vector{Int}     # index of neighbor cell in the (-1, 0) direction
    yu::Vector{Int}     # index of neighbor cell in the (0, +1) direction
    yd::Vector{Int}     # index of neighbor cell in the (0, -1) direction
end

# maps the fields of struct Indices to the defined Wflow cartesian indices of const
# neigbors.
const dirs = (:yd, :xd, :xu, :yu)

@get_units @with_kw struct ShallowWaterLand{T}
    n::Int | "-"                        # number of cells
    xl::Vector{T} | "m"                 # cell length x direction
    yl::Vector{T} | "m"                 # cell length y direction
    xwidth::Vector{T} | "m"             # effective flow width x direction (floodplain)
    ywidth::Vector{T} | "m"             # effective flow width y direction (floodplain)
    g::T | "m2 s-1"                     # acceleration due to gravity
    θ::T | "-"                          # weighting factor (de Almeida et al., 2012)
    α::T | "-"                          # stability coefficient (de Almeida et al., 2012)
    h_thresh::T | "m"                   # depth threshold for calculating flow
    Δt::T | "s"                         # model time step [s]
    qy0::Vector{T} | "m3 s-1"           # flow in y direction at previous time step
    qx0::Vector{T} | "m3 s-1"           # flow in x direction at previous time step
    qx::Vector{T} | "m3 s-1"            # flow in x direction
    qy::Vector{T} | "m3 s-1"            # flow in y direction
    zx_max::Vector{T} | "m"             # maximum cell elevation (x direction)
    zy_max::Vector{T} | "m"             # maximum cell elevation (y direction)
    mannings_n::Vector{T} | "s m-1/3"   # Manning's roughness
    volume::Vector{T} | "m3"            # total volume of cell (including river volume for river cells)
    error::Vector{T} | "m3"             # error volume
    runoff::Vector{T} | "m3 s-1"        # runoff from hydrological model
    h::Vector{T} | "m"                  # water depth of cell (for river cells the reference is the river bed elevation `zb`)
    z::Vector{T} | "m"                  # elevation of cell
    froude_limit::Bool | "-"            # if true a check is performed if froude number > 1.0 (algorithm is modified)
    rivercells::Vector{Bool} | "-"      # river cells
    h_av::Vector{T} | "m"               # average water depth (for river cells the reference is the river bed elevation `zb`)
end

function initialize_shallowwater_land(
    nc,
    config,
    inds;
    modelsize_2d,
    indices_reverse, # maps from the 2D external domain to the 1D internal domain (Int for linear indexing).
    xlength,
    ylength,
    riverwidth,
    graph_riv,
    ldd_riv,
    inds_riv,
    river,
    Δt,
)
    froude_limit = get(config.model, "froude_limit", true)::Bool # limit flow to subcritical according to Froude number
    alpha = get(config.model, "inertial_flow_alpha", 0.7)::Float64 # stability coefficient for model time step (0.2-0.7)
    theta = get(config.model, "inertial_flow_theta", 0.8)::Float64 # weighting factor
    h_thresh = get(config.model, "h_thresh", 1.0e-03)::Float64 # depth threshold for flow at link

    @info "Local inertial approach is used for overlandflow." alpha theta h_thresh froude_limit

    n_land = ncread(
        nc,
        config,
        "lateral.land.n";
        sel = inds,
        defaults = 0.072,
        type = Float,
    )
    elevation_2d = ncread(
        nc,
        config,
        "lateral.land.elevation";
        optional = false,
        type = Float,
        fill = 0,
    )
    elevation = elevation_2d[inds]
    n = length(inds)

    # initialize links between cells in x and y direction.
    indices = Indices(xu = zeros(n), xd = zeros(n), yu = zeros(n), yd = zeros(n))

    # links without neigbors are handled by an extra index (at n + 1, with n links), which
    # is set to a value of 0.0 m³ s⁻¹ for qx and qy fields at initialization.
    # links are defined as follows for the x and y direction, respectively:
    # node i => node xu (node i + CartesianIndex(1, 0))
    # node i => node yu (node i + CartesianIndex(0, 1))
    # where i is the index of inds
    nrow, ncol = modelsize_2d
    for (v, i) in enumerate(inds)
        for (m, neighbor) in enumerate(neighbors)
            j = i + neighbor
            dir = dirs[m]
            if (1 <= j[1] <= nrow) && (1 <= j[2] <= ncol) && (indices_reverse[j] != 0)
                getfield(indices, dir)[v] = indices_reverse[j]
            else
                getfield(indices, dir)[v] = n + 1
            end
        end
    end

    # determine z at links in x and y direction
    zx_max = fill(Float(0), n)
    zy_max = fill(Float(0), n)
    for i = 1:n
        xu = indices.xu[i]
        if xu <= n
            zx_max[i] = max(elevation[i], elevation[xu])
        end
        yu = indices.yu[i]
        if yu <= n
            zy_max[i] = max(elevation[i], elevation[yu])
        end
    end

    # set the effective flow width for river cells in the x and y direction at cell edges.
    we_x = copy(xlength)
    we_y = copy(ylength)
    set_effective_flowwidth!(
        we_x,
        we_y,
        indices,
        graph_riv,
        riverwidth,
        ldd_riv,
        indices_reverse[inds_riv],
    )

    sw_land = ShallowWaterLand{Float}(
        n = n,
        xl = xlength,
        yl = ylength,
        xwidth = we_x,
        ywidth = we_y,
        g = 9.80665,
        θ = theta,
        α = alpha,
        h_thresh = h_thresh,
        Δt = tosecond(Δt),
        qx0 = zeros(n + 1),
        qy0 = zeros(n + 1),
        qx = zeros(n + 1),
        qy = zeros(n + 1),
        zx_max = zx_max,
        zy_max = zy_max,
        mannings_n = n_land,
        volume = zeros(n),
        error = zeros(n),
        runoff = zeros(n),
        h = zeros(n),
        h_av = zeros(n),
        z = elevation,
        froude_limit = froude_limit,
        rivercells = river,
    )

    return sw_land, indices
end

"""
    stable_timestep(sw::ShallowWaterRiver)
    stable_timestep(sw::ShallowWaterLand)

Compute a stable timestep size for the local inertial approach, based on Bates et al. (2010).

Δt = α * (Δx / sqrt(g max(h))
"""
function stable_timestep(sw::ShallowWaterRiver{T})::T where {T}
    Δtₘᵢₙ = T(Inf)
    for i = 1:sw.n
        Δt = sw.α * sw.dl[i] / sqrt(sw.g * sw.h[i])
        Δtₘᵢₙ = Δt < Δtₘᵢₙ ? Δt : Δtₘᵢₙ
    end
    Δtₘᵢₙ = isinf(Δtₘᵢₙ) ? T(10.0) : Δtₘᵢₙ
    return Δtₘᵢₙ
end

function stable_timestep(sw::ShallowWaterLand{T})::T where {T}
    Δtₘᵢₙ = T(Inf)
    for i = 1:sw.n
        if !sw.rivercells[i]
            Δt = sw.α * min(sw.xl[i], sw.yl[i]) / sqrt(sw.g * sw.h[i])
            Δtₘᵢₙ = Δt < Δtₘᵢₙ ? Δt : Δtₘᵢₙ
        end
    end
    Δtₘᵢₙ = isinf(Δtₘᵢₙ) ? T(10.0) : Δtₘᵢₙ
    return Δtₘᵢₙ
end

function update(
    sw::ShallowWaterLand{T},
    swr::ShallowWaterRiver{T},
    network;
    inflow_wb = nothing,
    doy = 0,
    update_h = false,
) where {T}

    @unpack nodes_at_link, links_at_node = network.river

    if !isnothing(swr.reservoir)
        swr.reservoir.inflow .= 0.0
        swr.reservoir.totaloutflow .= 0.0
    end
    if !isnothing(swr.lake)
        swr.lake.inflow .= 0.0
        swr.lake.totaloutflow .= 0.0
    end
    swr.q_av .= 0.0
    swr.h_av .= 0.0
    sw.h_av .= 0.0

    t = T(0.0)
    while t < swr.Δt
        Δt_river = stable_timestep(swr)
        Δt_land = stable_timestep(sw)
        Δt = min(Δt_river, Δt_land)
        shallowwater_river_update(swr, network.river, Δt, inflow_wb, doy, update_h)
        update(sw, swr, network, Δt)

        if t + Δt > swr.Δt
            Δt = swr.Δt - t
        end
        t = t + Δt
    end
    swr.q_av ./= swr.Δt
    swr.h_av ./= swr.Δt
    sw.h_av ./= sw.Δt
end

function update(sw::ShallowWaterLand{T}, swr::ShallowWaterRiver{T}, network, Δt) where {T}

    indices = network.land.staggered_indices
    inds_riv = network.land.index_river

    @unpack nodes_at_link, links_at_node = network.river

    sw.qx0 .= sw.qx
    sw.qy0 .= sw.qy

    # update qx
    @threads for i = 1:sw.n
        yu = indices.yu[i]
        yd = indices.yd[i]
        xu = indices.xu[i]
        xd = indices.xd[i]

        # the effective flow width is zero when the river width exceeds the cell width (dy
        # for flow in x dir) and floodplain flow is not calculated.
        if xu <= sw.n && sw.ywidth[i] != 0.0

            η_x = sw.z[i] + sw.h[i]
            η_xu = sw.z[xu] + sw.h[xu]
            η_max = max(η_x, η_xu)
            hf = (η_max - sw.zx_max[i])

            if hf > sw.h_thresh
                length = 0.5 * (sw.xl[i] + sw.xl[xu]) # can be precalculated
                sw.qx[i] = local_inertial_flow(
                    sw.θ,
                    sw.qx0[i],
                    sw.qx0[xd],
                    sw.qx0[xu],
                    η_x,
                    η_xu,
                    hf,
                    sw.ywidth[i],
                    length,
                    sw.mannings_n[i],
                    sw.g,
                    sw.froude_limit,
                    Δt,
                )
            else
                sw.qx[i] = 0.0
            end
        end

        # update qy

        # the effective flow width is zero when the river width exceeds the cell width (dx
        # for flow in y dir) and floodplain flow is not calculated.
        if yu <= sw.n && sw.xwidth[i] != 0.0

            η_y = sw.z[i] + sw.h[i]
            η_yu = sw.z[yu] + sw.h[yu]
            η_max = max(η_y, η_yu)
            hf = (η_max - sw.zy_max[i])

            if hf > sw.h_thresh
                length = T(0.5) * (sw.yl[i] + sw.yl[yu]) # can be precalculated
                sw.qy[i] = local_inertial_flow(
                    sw.θ,
                    sw.qy0[i],
                    sw.qy0[yd],
                    sw.qy0[yu],
                    η_y,
                    η_yu,
                    hf,
                    sw.xwidth[i],
                    length,
                    sw.mannings_n[i],
                    sw.g,
                    sw.froude_limit,
                    Δt,
                )
            else
                sw.qy[i] = 0.0
            end
        end
    end

    # change in volume and water levels based on horizontal fluxes for river and land cells
    @threads for i = 1:sw.n
        yd = indices.yd[i]
        xd = indices.xd[i]

        # first add runoff (continuity equation)
        sw.volume[i] += sw.runoff[i] * Δt

        if sw.rivercells[i]
            sw.volume[i] +=
                (
                    sum_at(swr.q, links_at_node.src[inds_riv[i]]) -
                    sum_at(swr.q, links_at_node.dst[inds_riv[i]]) + sw.qx[xd] - sw.qx[i] +
                    sw.qy[yd] - sw.qy[i]
                ) * Δt
            if sw.volume[i] < 0.0
                sw.error[i] = sw.error[i] + abs(sw.volume[i])
                sw.volume[i] = 0.0 # set volume to zero
            end
            if sw.volume[i] >= swr.bankfull_volume[inds_riv[i]]
                swr.h[inds_riv[i]] =
                    swr.bankfull_depth[inds_riv[i]] +
                    (sw.volume[i] - swr.bankfull_volume[inds_riv[i]]) /
                    (sw.xl[i] * sw.yl[i])
                sw.h[i] = swr.h[inds_riv[i]] - swr.bankfull_depth[inds_riv[i]]
                swr.volume[inds_riv[i]] =
                    swr.h[inds_riv[i]] * swr.dl[inds_riv[i]] * swr.width[inds_riv[i]]
            else
                swr.h[inds_riv[i]] =
                    sw.volume[i] / (swr.dl[inds_riv[i]] * swr.width[inds_riv[i]])
                sw.h[i] = 0.0
                swr.volume[inds_riv[i]] = sw.volume[i]
            end
            swr.h_av[inds_riv[i]] += swr.h[inds_riv[i]] * Δt
        else
            sw.volume[i] += (sw.qx[xd] - sw.qx[i] + sw.qy[yd] - sw.qy[i]) * Δt
            if sw.volume[i] < 0.0
                sw.error[i] = sw.error[i] + abs(sw.volume[i])
                sw.volume[i] = 0.0 # set volume to zero
            end
            sw.h[i] = sw.volume[i] / (sw.xl[i] * sw.yl[i])
        end
        sw.h_av[i] += sw.h[i] * Δt
    end
end

"""
    set_river_inwater(model::Model{N,L,V,R,W,T}, ssf_toriver) where {N,L,V<:SBM,R,W,T}

Set `inwater` of the river component for a `Model` with vertical `SBM` concept.
`ssf_toriver` is the subsurface flow to the river.
"""
function set_river_inwater(model::Model{N,L,V,R,W,T}, ssf_toriver) where {N,L,V<:SBM,R,W,T}
    @unpack lateral, vertical, network = model
    inds = network.index_river

    @. lateral.river.inwater = (
        ssf_toriver[inds] +
        lateral.land.to_river[inds] +
        # net_runoff_river
        (
            (
                vertical.net_runoff_river[inds] *
                network.land.xl[inds] *
                network.land.yl[inds] *
                0.001
            ) / vertical.Δt
        )
    )
end

"""
    set_river_inwater(model, ssf_toriver)

Set `inwater` of the river component (based on overland flow).
"""
function set_river_inwater(model, ssf_toriver)
    @unpack lateral, network = model
    inds = network.index_river
    lateral.river.inwater .= lateral.land.to_river[inds]
end

"""
    set_land_inwater(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmGwfModel}

Set `inwater` of the land component for the `SbmGwgModel` type.
"""
function set_land_inwater(model::Model{N,L,V,R,W,T}) where {N,L,V,R,W,T<:SbmGwfModel}
    @unpack lateral, vertical, network, config = model

    do_drains = get(config.model, "drains", false)::Bool
    drainflux = zeros(vertical.n)
    if do_drains
        drainflux[lateral.subsurface.drain.index] =
            -lateral.subsurface.drain.flux ./ tosecond(basetimestep)
    end

    lateral.land.inwater .=
        (vertical.runoff .* network.land.xl .* network.land.yl .* 0.001) ./
        lateral.land.Δt .+ drainflux
end

"""
    set_land_inwater(model)

Set `inwater` of the land component, based on `runoff` of the `vertical` concept.
"""
function set_land_inwater(model)
    @unpack lateral, vertical, network = model
    lateral.land.inwater .=
        (vertical.runoff .* network.land.xl .* network.land.yl .* 0.001) ./ lateral.land.Δt
end

"""
    get_inflow_waterbody(model)

Get inflow to a water body (reservoir or lake) `inflow_wb` based on overland flow.
"""
function get_inflow_waterbody(model)
    @unpack lateral, network = model
    inds = network.index_river
    if !isnothing(lateral.river.reservoir) || !isnothing(lateral.river.lake)
        inflow_wb = lateral.land.q_av[inds]
    else
        inflow_wb = nothing
    end
    return inflow_wb
end

"""
    get_inflow_waterbody(
        model::Model{N,L,V,R,W,T},
    ) where {N,L<:NamedTuple{<:Any,<:Tuple{LateralSSF,SurfaceFlow,Any}},V,R,W,T}

Get inflow to a water body (reservoir or lake) `inflow_wb` from a model type that contains
the lateral components `LateralSSF` and `SurfaceFlow`.
"""
function get_inflow_waterbody(
    model::Model{N,L,V,R,W,T},
) where {N,L<:NamedTuple{<:Any,<:Tuple{LateralSSF,SurfaceFlow,Any}},V,R,W,T}
    @unpack lateral, network = model

    inds = network.index_river
    if !isnothing(lateral.river.reservoir) || !isnothing(lateral.river.lake)
        inflow_wb =
            lateral.subsurface.ssf[inds] ./ tosecond(basetimestep) .+
            lateral.land.q_av[inds]
    else
        inflow_wb = nothing
    end
    return inflow_wb
end

"""
    get_inflow_waterbody(
        model::Model{N,L,V,R,W,T},
    ) where {N,L<:NamedTuple{<:Any,<:Tuple{LateralSSF,ShallowWaterLand,Any}},V,R,W,T}

Get inflow to a water body (reservoir or lake) `inflow_wb` from a model type that contains
the lateral components `LateralSSF` and `ShallowWaterLand`.
"""
function get_inflow_waterbody(
    model::Model{N,L,V,R,W,T},
) where {N,L<:NamedTuple{<:Any,<:Tuple{LateralSSF,ShallowWaterLand,Any}},V,R,W,T}
    @unpack lateral, network = model

    inds = network.index_river
    if !isnothing(lateral.river.reservoir) || !isnothing(lateral.river.lake)
        inflow_wb = lateral.subsurface.ssf[inds] ./ tosecond(basetimestep)
    else
        inflow_wb = nothing
    end
    return inflow_wb
end

"""
    surface_routing(model; ssf_toriver = 0.0)

Run surface routing (land and river). Kinematic wave for overland flow and kinematic wave or
local inertial model for river flow.
"""
function surface_routing(model; ssf_toriver = 0.0)
    @unpack lateral, network, clock = model

    # run kinematic wave for overland flow
    set_land_inwater(model)
    update(lateral.land, network.land, frac_toriver = network.frac_toriver)

    # run river flow
    set_river_inwater(model, ssf_toriver)
    update(
        lateral.river,
        network.river,
        inflow_wb = get_inflow_waterbody(model),
        doy = dayofyear(clock.time),
    )
end

"""
    surface_routing(
        model::Model{N,L,V,R,W,T};
        ssf_toriver = 0.0,
    ) where {N,L<:NamedTuple{<:Any,<:Tuple{Any,ShallowWaterLand,ShallowWaterRiver}},V,R,W,T}

Run surface routing (land and river) for a model type that contains the lateral components
`ShallowWaterLand` and `ShallowWaterRiver`.
"""
function surface_routing(
    model::Model{N,L,V,R,W,T};
    ssf_toriver = 0.0,
) where {N,L<:NamedTuple{<:Any,<:Tuple{Any,ShallowWaterLand,ShallowWaterRiver}},V,R,W,T}

    @unpack lateral, vertical, network, clock = model

    @. lateral.land.runoff = (
        (vertical.runoff / 1000.0) * (network.land.xl * network.land.yl) / vertical.Δt +
        ssf_toriver +
        # net_runoff_river
        (
            (vertical.net_runoff_river * network.land.xl * network.land.yl * 0.001) /
            vertical.Δt
        )
    )

    update(
        lateral.land,
        lateral.river,
        network,
        inflow_wb = get_inflow_waterbody(model),
        doy = dayofyear(clock.time),
    )
end


text:
using OrdinaryDiffEq, Test

f(u,p,t) = zero(u)
prob = ODEProblem(f, 1.0, (0.0, 10.0))

# Shouldn't error
sol = solve(prob, Tsit5(), controller=IController())
@test sol.retcode === :Success

sol = solve(prob, Tsit5(), controller=PIController(7//50, 2//25))
@test sol.retcode === :Success

sol = solve(prob, Tsit5(), controller=PIDController(0.7, -0.4))
@test sol.retcode === :Success


text:
using SimJulia

function client(sim::Simulation, res::Resource, i::Int, priority::Int)
  println("$(now(sim)), client $i is waiting")
  yield(sim, request(sim, res, priority=priority))
  println("$(now(sim)), client $i is being served")
  yield(sim, timeout(sim, rand()))
  println("$(now(sim)), client $i has been served")
  yield(sim, release(sim, res))
end

function generate(sim::Simulation, res::Resource)
  for i = 1:10
    Process(sim, client, res, i, 10-i)
    yield(sim, timeout(sim, 0.5*rand()))
  end
end

sim = Simulation()
res = Resource(2, level=1)
Process(sim, generate, res)
run(sim)

function consumer(sim::Simulation, con::Container)
  for i = 1:10
    amount = 3*rand()
    println("$(now(sim)), consumer is demanding $amount")
    yield(sim, timeout(sim, rand()))
    yield(sim, get(sim, con, amount))
    println("$(now(sim)), consumer is being served, level is $(con.level)")
    yield(sim, timeout(sim, 5*rand()))
  end
end

function producer(sim::Simulation, con::Container)
  for i = 1:10
    amount = 2*rand()
    println("$(now(sim)), producer is offering $amount")
    yield(sim, timeout(sim, rand()))
    yield(sim, put(sim, con, amount))
    println("$(now(sim)), producer is being served, level is $(con.level)")
    yield(sim, timeout(sim, 5*rand()))
  end
end

sim = Simulation()
con = Container(10.0, level=5.0)
Process(sim, consumer, con)
Process(sim, producer, con)
run(sim)


text:
# Functions for various model reduction techniques

function MinimumRealization(S::StateSpace)::StateSpace
    # Todo
end

function BalancedReduction(S::StateSpace)::StateSpace
    # Todo
end

function HankelSVD(S::StateSpace)::Vector{Float64}
    # Todo
end

function EigenSystemRealization(S::StateSpace)::StateSpace
    # Todo
end

function MarkovParameters(S::StateSpace)::StateSpace
    # Todo
end

max_stars_repo_name: isrlab/RobustControl

max_issues_repo_name: isrlab/RobustControl

max_forks_repo_name: isrlab/RobustControl


text:
function get_below(subtable, arg)
    _tab = Iterators.filter(r -> r[arg.first] ≤ arg.second, eachrow(subtable))
    y = maximum(r -> r[arg.first], _tab)
    return filter(r -> r[arg.first] ≈ y, subtable)
end

function get_above(subtable, arg)
    _tab = Iterators.filter(r -> r[arg.first] > arg.second, eachrow(subtable))
    y = minimum(r -> r[arg.first], _tab)
    return filter(r -> r[arg.first] ≈ y, subtable)
end

function table_interpolate(table, arg1, arg2, arg3)
    sub = copy(table)
    try
        sub0 = get_below(sub, arg1)
        # isempty(sub0) == 0 && return NaN
        x0 = x1 = sub0[1, arg1.first]

        # check bounds
        # bounds2 = extrema(sub0[!, arg2.first])
        # if  arg2.second < bounds2[1] || arg2.second > bounds2[2]
        #     @debug "out of bounds" sub0 arg2
        #     return NaN
        # end

        tab0 = get_below(sub0, arg2)
        y0 = tab0[1, arg2.first]
        z0 = tab0[1, arg3]

        tab1 = get_above(sub0, arg2)
        y1 = tab1[1, arg2.first]
        z1 = tab1[1, arg3]

        sub1 = get_above(sub, arg1)
        # isempty(sub1) == 0 && return NaN
        x2 = x3 = sub1[1, arg1.first]
        # bounds2 = extrema(sub1[!, arg2.first])
        # if  arg2.second < bounds2[1] || arg2.second > bounds2[2]
        #     @debug "out of bounds" sub1 arg2
        #     return NaN
        # end

        tab2 = get_below(sub1, arg2)
        y2 = tab2[1, arg2.first]
        z2 = tab2[1, arg3]

        tab3 = get_above(sub1, arg2)
        y3 = tab3[1, arg2.first]
        z3 = tab3[1, arg3]

        xs = [x0, x1, x2, x3]
        ys = [y0, y1, y2, y3]
        zs = [z0, z1, z2, z3]
        spl = Spline2D(xs, ys, zs, kx=1, ky=1)
        return spl(arg1.second, arg2.second)
    catch
        return NaN
    end
end

#################################################

function read_ATMO2020_grid(dir)
    filenames = filter(f -> endswith(f, "vega.txt"), readdir(dir))

    header = ["Mass", "Age", "Teff", "Luminosity", "Radius", "Gravity", "MKO_Y", "MKO_J", "MKO_H",
              "MKO_K", "MKO_Lp", "MKO_Mp", "W1", "W2", "W3", "W4", "IRAC_CH1", "IRAC_CH2"]

    full_table = mapreduce(vcat, filenames) do filename
        mat = readdlm(joinpath(dir, filename), skipstart=2)
        # a bunch of magnitudes listed with 0, treat as NaN
        @. mat[iszero(mat)] = NaN
        return DataFrame(mat, header)
    end
    full_table.Mass_MJ = full_table.Mass .* 1047.57
    return full_table
end

const ATMO2020 = read_ATMO2020_grid(srcdir("mass_models", "ATMO2020_NEQ_weak"))

#################################################

function read_Sonora_grid(dir; metallicity=0, CO_ratio=1.0)
    # prepare file names
    if CO_ratio ≈ 1
        dex_token = @sprintf "%+02.1f" metallicity
        CO_token = ""
    else
        dex_token = @sprintf "_m%+02.1f" metallicity
        CO_token = @sprintf "_co%02.1f" CO_ratio
    end
    evo_path = joinpath(dir, "evolution_tables", @sprintf("evo_tables%+02.1f", metallicity), @sprintf("nc%+02.1f_co%02.1f_age", metallicity, CO_ratio))
    mag_path = joinpath(dir, "photometry_tables", @sprintf("mag_table%s%s", dex_token, CO_token))

    # read evolution_table
    evo_header = ["age_Gyr", "M_Msun", "logL", "Teff", "logg", "R_Rsun"]
    evo_mat = open(evo_path, "r") do fh
        # skip first (header) row
        readline(fh)
        # iterate over rows
        rows = map(readlines(fh)) do row
            tokens = split(replace(row, "\n" => ""))
            if length(tokens) == 1
                nothing
            else
                parse.(Float64, tokens)
            end
        end
        rows = filter(!isnothing, rows)
        return mapreduce(r -> reduce(hcat, r), vcat, rows)
    end
    evo_table = DataFrame(evo_mat, evo_header)
    evo_table.M_Mjup = evo_table.M_Msun .* 1047.57

    # read photometry_table
    mag_header = ["Teff", "logg", "M_Mjup", "R_Rsun", "Y", "log_Kzz", "MKO_Y", "MKO_Z", "MKO_J", "MKO_H", "MKO_K",
    "MKO_Lp", "MKO_Mp", "2MASS_J", "2MASS_H", "2MASS_Ks", "Keck_Ks", "Keck_Lp", "Keck_Ms", "SDSS_g",
    "SDSS_r", "SDSS_i", "SDSS_z", "IRAC_3.6", "IRAC_4.5", "IRAC_5.7", "IRAC_7.9", "WISE_W1", "WISE_W2",
    "WISE_W3", "WISE_W4"]
    mag_mat = readdlm(mag_path, skipstart=11)
    # remove values with R/Rsun has a *
    mag_mat = @. Float64(mag_mat[mag_mat[:, 4] isa Float64, :])
    # replace Kzz values
    @. mag_mat[mag_mat[:, 6] ≈ -99, 6] = -Inf
    mag_table = DataFrame(mag_mat, mag_header)

    # now we need to combine the magnitude and evolution tables
    # this is accomplished by taking (Teff, logg) pairs and interpolating age from the evo table
    for key in mag_header[7:end]
        evo_table[!, key] = map(eachrow(evo_table)) do row
            try
                table_interpolate(mag_table, :Teff => row.Teff, :logg => row.logg, Symbol(key))
            catch
                missing
            end
        end
    end
    dropmissing!(evo_table)
    return evo_table
end

const SonoraSolar = read_Sonora_grid(srcdir("mass_models", "Sonora_Bobcat"))
const SonoraMetalRich = read_Sonora_grid(srcdir("mass_models", "Sonora_Bobcat"); metallicity=0.5)

max_stars_repo_name: mileslucas/sirius-b

max_issues_repo_name: mileslucas/sirius-b

max_forks_repo_name: mileslucas/sirius-b


text:
using Yao
import Base:sparse
import Yao.Blocks: PrimitiveBlock
using Yao.Registers
using Yao.Intrinsics: log2i

################################################
#              Grover Operator Block           #
################################################
# psi and oracle are needed
struct Reflect{N, T} <: PrimitiveBlock{N, T}
    state :: Vector{T}
end
Reflect(state::Vector{T}) where T = Reflect{log2i(length(state)), T}(state)
Reflect(psi::AbstractRegister) = Reflect(statevec(psi))

# NOTE: this should not be matrix multiplication based
import Yao.Blocks: apply!
function apply!(r::AbstractRegister, g::Reflect)
    @views r.state[:,:] .= 2* (g.state'*r.state) .* reshape(g.state, :, 1) - r.state
    r
end
# since julia does not allow call overide on AbstractGate.
(rf::Reflect)(reg::AbstractRegister) = apply!(reg, rf)

import Base: show
function show(io::IO, g::Reflect{N, T}) where {N, T}
    print("Reflect(N = $N")
end

max_stars_repo_name: Roger-luo/QCBM

max_issues_repo_name: Roger-luo/QCBM

max_forks_repo_name: Roger-luo/QCBM


text:
@with_kw_noshow struct PatMat{S<:AbstractSurrogate, T1<:Real, T2<:Real} <: AbstractPatMat{S}
    τ::T1
    C::T2 = 1
    l1::S = Hinge(ϑ = 1.0)
    l2::S = Hinge(ϑ = 1.0)

    @assert 0 < τ < 1  "The vaule of `τ` must lay in the interval (0,1)."
end

PatMat(τ::Real; kwargs...) = PatMat(; τ = τ, kwargs...)

show(io::IO, model::PatMat) =
    print(io, "PatMat($(model.τ), $(model.C), $(model.l1), $(model.l2))")

@with_kw_noshow struct PatMatNP{S<:AbstractSurrogate, T1<:Real, T2<:Real} <: AbstractPatMat{S}
    τ::T1
    C::T2 = 1
    l1::S = Hinge(ϑ = 1.0)
    l2::S = Hinge(ϑ = 1.0)

    @assert 0 < τ < 1  "The vaule of `τ` must lay in the interval (0,1)."
end

PatMatNP(τ::Real; kwargs...) = PatMatNP(; τ = τ, kwargs...)

show(io::IO, model::PatMatNP) =
    print(io, "PatMatNP($(model.τ), $(model.C), $(model.l1), $(model.l2))")

# -------------------------------------------------------------------------------
# Primal problem - General solver
# -------------------------------------------------------------------------------
function optimize(solver::General, model::M, data::Primal) where {M <: AbstractPatMat}

    Xpos = @view data.X[data.ind_pos, :]
    w    = Convex.Variable(data.dim)
    t    = Convex.Variable()
    y    = Convex.Variable(data.npos)

    if M <: PatMat
        X = data.X
        z = Convex.Variable(data.n)
        n = data.n
    elseif M <: PatMatNP
        X = @view data.X[data.ind_neg, :]
        z = Convex.Variable(data.nneg)
        n = data.nneg
    end

    objective   = Convex.sumsquares(w)/2 + model.C * Convex.sum(model.l1.value_exact.(y))
    constraints = [Convex.sum(model.l2.value_exact.(z)) <= n * model.τ,
                   y == t - Xpos*w,
                   z == X*w - t]

    problem = Convex.minimize(objective, constraints)
    Convex.solve!(problem, solver.solver)

    return vec(w.value), t.value
end

# -------------------------------------------------------------------------------
# Primal problem - Gradient descent solver
# -------------------------------------------------------------------------------
function initialization(model::AbstractPatMat, data::Primal)
    w = zeros(eltype(data.X), data.dim)
    s = scores(model, data, w)
    Δ = zero(w)
    return w, s, Δ
end

function objective(model::AbstractPatMat, data::Primal, w, t, s = data.X*w)
    w'*w/2 + model.C * sum(model.l1.value.(t .- s[data.ind_pos]))
end

function threshold(model::PatMat, data::Primal, s)
   Roots.find_zero(t -> sum(model.l2.value.(s .- t)) - data.n*model.τ, (-Inf, Inf))
end

function gradient!(model::PatMat, data::Primal, w, s, Δ)
    t   = threshold(model, data, s)
    ∇l1 = model.l1.gradient.(t .- s[data.ind_pos])
    ∇l2 = model.l2.gradient.(s .- t)
    ∇t  = data.X'*∇l2/sum(∇l2)

    Δ .= w .+ model.C .* (sum(∇l1)*∇t .- data.X[data.ind_pos,:]'*∇l1)
    return t
end

function threshold(model::PatMatNP, data::Primal, s)
   Roots.find_zero(t -> sum(model.l2.value.(s[data.ind_neg] .- t)) - data.nneg*model.τ, (-Inf, Inf))
end

function gradient!(model::PatMatNP, data::Primal, w, s, Δ)
    t   = threshold(model, data, s)
    ∇l1 = model.l1.gradient.(t .- s[data.ind_pos])
    ∇l2 = model.l2.gradient.(s[data.ind_neg] .- t)
    ∇t  = data.X[data.ind_neg,:]'*∇l2/sum(∇l2)

    Δ .= w .+ model.C .* (sum(∇l1)*∇t .- data.X[data.ind_pos,:]'*∇l1)
    return t
end

# -------------------------------------------------------------------------------
# Dual problem - General solver
# -------------------------------------------------------------------------------
# Hinge loss
function optimize(solver::General, model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}; ε::Real = 1e-6)

    α = Convex.Variable(data.nα)
    β = Convex.Variable(data.nβ)
    δ = Convex.Variable()
    K = data.K + ε .* I

    objective   = - Convex.quadform(vcat(α, β), K)/2 + Convex.sum(α)/model.l1.ϑ +
                    Convex.sum(β)/model.l2.ϑ - δ*data.nβ*model.τ
    constraints = [Convex.sum(α) == Convex.sum(β),
                   α <= model.l1.ϑ*model.C,
                   β <= model.l2.ϑ*δ,
                   α >= 0,
                   β >= 0]

    problem = Convex.maximize(objective, constraints)
    Convex.solve!(problem, solver.solver)

    return vec(α.value), vec(β.value), δ.value
end

# Truncated quadratic loss
function optimize(solver::General, model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}; ε::Real = 1e-6)

    α = Convex.Variable(data.nα)
    β = Convex.Variable(data.nβ)
    δ = Convex.Variable()
    K = data.K + ε .* I

    objective   = - Convex.quadform(vcat(α, β), K)/2 +
                    Convex.sum(α)/model.l1.ϑ - Convex.sumsquares(α)/(4*model.C*model.l1.ϑ^2) +
                    Convex.sum(β)/model.l2.ϑ - Convex.quadoverlin(β, δ)/(4*model.l2.ϑ^2) -
                    δ*data.nβ*model.τ
    constraints = [Convex.sum(α) == Convex.sum(β),
                   α >= 0,
                   β >= 0,
                   δ >= 0]

    problem = Convex.maximize(objective, constraints)
    Convex.solve!(problem, solver.solver)

    return vec(α.value), vec(β.value), δ.value
end

# -------------------------------------------------------------------------------
# Dual problem - Graient descent solver
# -------------------------------------------------------------------------------
function initialization(model::AbstractPatMat{S}, data::Dual{<:DTrain}) where {S<:AbstractSurrogate}
    αβδ     = rand(eltype(data.K), data.nα + data.nβ + 1)
    α, β, δ = @views αβδ[data.ind_α], αβδ[data.ind_β], αβδ[[end]]
    if S <: Hinge
        δ .= maximum(β)/model.l1.ϑ
    else
        δ .= sqrt(max(sum(abs2, β)/(4*data.nβ*model.τ*model.l2.ϑ^2), 0))
    end
    projection!(model, data, α, β, δ)

    s = data.K * vcat(α, β)
    return α, β, δ, αβδ, s
end

# Hinge loss
function objective(model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}, α, β, δ, s = data.K * vcat(α, β))
    return - s'*vcat(α, β)/2 + sum(α)/model.l1.ϑ + sum(β)/model.l2.ϑ - δ*data.nβ*model.τ
end

function gradient!(model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}, α, β, δ, s, Δ)
    Δ[1:end-1]    .= .- s
    Δ[data.ind_α] .+= 1/model.l1.ϑ
    Δ[data.ind_β] .+= 1/model.l2.ϑ
    Δ[end]         = - data.nβ*model.τ
end

function projection!(model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}, α, β, δ)
     αs, βs, δs = projection(α, β, δ[1], model.l1.ϑ*model.C, model.l2.ϑ)
     α .= αs
     β .= βs
     δ .= δs
     return α, β, δ
end

# Truncated quadratic loss
function objective(model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}, α, β, δ, s = data.K * vcat(α, β))
    return - s'*vcat(α, β)/2 +
           sum(α)/model.l1.ϑ - sum(abs2, α)/(4*model.C*model.l1.ϑ^2) +
           sum(β)/model.l2.ϑ - sum(abs2, β)/(4*δ*model.l2.ϑ^2) -
           δ*data.nβ*model.τ
end

function gradient!(model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}, α, β, δ, s, Δ)
    Δ[1:end-1]    .= .- s
    Δ[data.ind_α] .+= 1/model.l1.ϑ .- α/(2*model.l1.ϑ^2*model.C)
    Δ[data.ind_β] .+= 1/model.l2.ϑ .- β/(2*model.l2.ϑ^2*δ[1])
    Δ[end]         = sum(abs2, β)/(4*model.l2.ϑ^2*δ[1]^2) - data.nβ*model.τ
end

function projection!(model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}, α, β, δ)
     αs, βs, δs = projection(α, β, δ[1])
     α .= αs
     β .= βs
     δ .= δs
     return α, β, δ
end

function threshold(model::PatMat, data::Dual{<:DTrain}, s)
   Roots.find_zero(t -> sum(model.l2.value.(.- s[data.type.inv_perm] .- t)) - data.n*model.τ, (-Inf, Inf))
end

function threshold(model::PatMatNP, data::Dual{<:DTrain}, s)
   Roots.find_zero(t -> sum(model.l2.value.(.- s[data.ind_β] .- t)) - data.nβ*model.τ, (-Inf, Inf))
end

function primal_objective(model::AbstractPatMat, data::Dual{<:DTrain}, α, β, δ, s)
    t = threshold(model, data, s)

    return s'*vcat(α, β)/2 + model.C * sum(model.l1.value.(t .- s[data.ind_α]))
end

# -------------------------------------------------------------------------------
# Dual problem - Coordinate descent solver
# -------------------------------------------------------------------------------
function select_k(model::AbstractPatMat, data::Dual{<:DTrain}, α, β, δ)
    rand(1:(data.nα + data.nβ))
end

# Hinge loss
function loss(model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}, a::Real, b::Real, δ::Real, Δ::Real)
    a*Δ^2/2 + b*Δ - δ*data.nβ*model.τ
end

function apply!(model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}, best::BestUpdate, α, β, δ, αβδ, s, βsort)
    βsorted!(data, best, β, βsort)
    αβδ[best.k] = best.vars[1]
    αβδ[best.l] = best.vars[2]
    αβδ[end]    = best.vars[3]
    scores!(data, best, s)
end

function rule_αα!(model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}, best::BestUpdate, k, l, α, β, δ, s, βsort)

    αk, αl   = α[k], α[l]
    n, C, ϑ1 = data.nβ, model.C, model.l1.ϑ

    a = - data.K[k,k] + 2*data.K[k,l] - data.K[l,l]
    b = - s[k] + s[l]
    Δ = solution(a, b, max(-αk, αl - ϑ1*C), min(ϑ1*C - αk, αl))

    vars = (αk = αk + Δ, αl = αl - Δ, δ = δ[1])
    L    = loss(model, data, a, b, vars.δ, Δ)
    update!(best, k, l, Δ, L, vars)
end

function rule_αβ!(model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}, best::BestUpdate, k, l, α, β, δ, s, βsort)

    αk, βl = α[k], β[l - data.nα]
    n, τ, C, ϑ1, ϑ2 = data.nβ, model.τ, model.C, model.l1.ϑ, model.l2.ϑ
    βmax   = find_βmax(βsort, β, l - data.nα)

    a = - data.K[k,k] - 2*data.K[k,l] - data.K[l,l]
    b = - s[k] - s[l] + 1/ϑ1 + 1/ϑ2

    # solution 1
    Δ = solution(a, b, max(-αk, -βl), ϑ1*C - αk)
    if βl + Δ <= βmax
        vars = (αk = αk + Δ, βl = βl + Δ, δ = βmax/ϑ2)
        L    = loss(model, data, a, b, vars.δ, Δ)
        update!(best, k, l, Δ, L, vars)
    end

    # solution 2
    Δ = solution(a, b - n*τ/ϑ2, max(-αk, -βl), ϑ1*C - αk)
    if βl + Δ >= βmax
        vars = (αk = αk + Δ, βl = βl + Δ, δ = (βl + Δ)/ϑ2)
        L    = loss(model, data, a, b - n*τ/ϑ2, vars.δ, Δ)
        update!(best, k, l, Δ, L, vars)
    end
end

function rule_ββ!(model::AbstractPatMat{<:Hinge}, data::Dual{<:DTrain}, best::BestUpdate, k, l, α, β, δ, s, βsort)

    βk, βl   = β[k - data.nα], β[l - data.nα]
    n, τ, ϑ2 = data.nβ, model.τ, model.l2.ϑ
    βmax     = find_βmax(βsort, β, k - data.nα, l - data.nα)

    a = - data.K[k,k] + 2*data.K[k,l] - data.K[l,l]
    b = - s[k] + s[l]

    # solution 1
    Δ = solution(a, b, -βk, βl)
    if βmax >= max(βk + Δ, βl - Δ)
        vars = (βk = βk + Δ, βl = βl - Δ, δ = βmax/ϑ2)
        L    = loss(model, data, a, b, vars.δ, Δ)
        update!(best, k, l, Δ, L, vars)
    end

    # solution 2
    Δ = solution(a, b - n*τ/ϑ2, -βk, βl)
    if βk + Δ >= max(βmax, βl - Δ)
        vars = (βk = βk + Δ, βl = βl - Δ, δ = (βk + Δ)/ϑ2)
        L    = loss(model, data, a, b - n*τ/ϑ2, vars.δ, Δ)
        update!(best, k, l, Δ, L, vars)
    end

    # solution 3
    Δ = solution(a, b + n*τ/ϑ2, -βk, βl)
    if βl - Δ >= max(βk + Δ, βmax)
        vars = (βk = βk + Δ, βl = βl - Δ, δ = (βl - Δ)/ϑ2)
        L    = loss(model, data, a, b + n*τ/ϑ2, vars.δ, Δ)
        update!(best, k, l, Δ, L, vars)
    end
end

# Truncated quadratic loss
function loss(model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}, a::Real, b::Real, δ::Real, Δ::Real, β2sum)
    a*Δ^2/2 + b*Δ - β2sum[1]/δ - δ*data.nβ*model.τ
end

function apply!(model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}, best::BestUpdate, α, β, δ, αβδ, s, β2sum)
    if best.k > data.nα || best.l > data.nα
        if best.k <= data.nα && best.l > data.nα
            β2sum .+= best.Δ*(best.Δ + 2*β[best.l - data.nα])/(4*model.l2.ϑ^2)
        else
            β2sum .+= best.Δ*(2*best.Δ + 2*(β[best.k - data.nα] - β[best.l - data.nα]))/(4*model.l2.ϑ^2)
        end
    end
    αβδ[best.k] = best.vars[1]
    αβδ[best.l] = best.vars[2]
    αβδ[end]    = best.vars[3]
    scores!(data, best, s)
end

function rule_αα!(model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}, best::BestUpdate, k, l, α, β, δ, s, β2sum)

    αk, αl   = α[k], α[l]
    n, C, ϑ1 = data.nβ, model.C, model.l1.ϑ

    a = - data.K[k,k] + 2*data.K[k,l] - data.K[l,l] - 1/(C*ϑ1^2)
    b = - s[k] + s[l] - (αk - αl)/(2*C*ϑ1^2)
    Δ = solution(a, b, - αk, αl)

    vars = (αk = αk + Δ, αl = αl - Δ, δ = δ[1])
    L    = loss(model, data, a, b, vars.δ, Δ, β2sum)
    update!(best, k, l, Δ, L, vars)
end

function rule_αβ!(model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}, best::BestUpdate, k, l, α, β, δ, s, β2sum)

    αk, βl = α[k], β[l - data.nα]
    n, τ, C, ϑ1, ϑ2 = data.nβ, model.τ, model.C, model.l1.ϑ, model.l2.ϑ

    a    = - data.K[k,k] - 2*data.K[k,l] - data.K[l,l] - 1/(2*C*ϑ1^2) - 1/(2*δ[1]*ϑ2^2)
    b    = - s[k] - s[l] + 1/ϑ1 - αk/(2*C*ϑ1^2) + 1/ϑ2 - βl/(2*δ[1]*ϑ2^2)
    Δ    = solution(a, b, max(-αk, -βl), Inf)
    δnew = sqrt(max(δ[1]^2 + (Δ^2 + 2*Δ*βl)/(4*ϑ2^2*n*τ), 0))

    a +=    (1/δ[1] - 1/δnew)/(2*ϑ2^2)
    b += βl*(1/δ[1] - 1/δnew)/(2*ϑ2^2)

    vars = (αk = αk + Δ, βl = βl + Δ, δ = δnew)
    L    = loss(model, data, a, b, vars.δ, Δ, β2sum)
    update!(best, k, l, Δ, L, vars)
end

function rule_ββ!(model::AbstractPatMat{<:Quadratic}, data::Dual{<:DTrain}, best::BestUpdate, k, l, α, β, δ, s, β2sum)

    βk, βl = β[k - data.nα], β[l - data.nα]
    n, τ, C, ϑ1, ϑ2 = data.nβ, model.τ, model.C, model.l1.ϑ, model.l2.ϑ

    a    = - data.K[k,k] + 2*data.K[k,l] - data.K[l,l] - 1/(δ[1]*ϑ2^2)
    b    = - s[k] + s[l] - (βk - βl)/(2*δ[1]*ϑ2^2)
    Δ    = solution(a, b, -βk, βl)
    δnew = sqrt(max(δ[1]^2 + (Δ^2 + Δ*(βk - βl))/(2*ϑ2^2*n*τ), 0))

    a += (1/δ[1] - 1/δnew)/(ϑ2^2)
    b += (1/δ[1] - 1/δnew)*(βk - βl)/(2*ϑ2^2)

    vars = (βk = βk + Δ, βl = βl - Δ, δ = δnew)
    L    = loss(model, data, a, b, vars.δ, Δ, β2sum)
    update!(best, k, l, Δ, L, vars)
end


text:
module UAVisualization
using GLVisualize, GeometryTypes, Colors
import GLVisualize: mm, Screen, labeled_slider, play_slider, center!, rotationmatrix_z
using Geodesy, MAT

export init, main

datafile = Pkg.dir("UAVisualization" * "/tools/data.mat")

function init()
	global window
	window = glscreen("UAVisualization")
	window
end

function rgbaxis(rot)
	scale = 10
	origin = Point3f0(0, 0, 0)
	ptx = origin + Point3f0(scale, 0, 0)
	pty = origin + Point3f0(0, scale, 0)
	ptz = origin + Point3f0(0, 0, scale)
	colors = RGBA{Float32}[
	RGBA{Float32}(1,0,0,1),
	RGBA{Float32}(1,0,0,1),

	RGBA{Float32}(0,1,0,1),
	RGBA{Float32}(0,1,0,1),

	RGBA{Float32}(0,0,1,1),
	RGBA{Float32}(0,0,1,1),
	]
	lines = [origin, ptx, origin, pty, origin, ptz]
	return visualize(lines, :linesegment, color=colors)
end

function loaddata()
	# TODO: if data file doesn't exist generate some fake data
	vars = matread(datafile)
	alt = vars["alt"]
	#@show length(alt)
	alt_nans = isnan.(alt)
	alt = alt[.!alt_nans]
	#@show length(alt)
	t = vars["t"][.!alt_nans]
	lat = rad2deg.(vars["lat"][.!alt_nans])
	lon = rad2deg.(vars["lon"][.!alt_nans])
	lla = LLA.(lat, lon, alt)
	trans = ENUfromLLA(lla[1], wgs84)
	hdg = vars["hdg"][.!alt_nans]
	enu = trans.(lla)
	return t - t[1], enu, hdg
end

function main()

	t, enu, hdg = loaddata()

	"Vehicle's position in ENU at a given index"
	function get_pos_enu(i)
		Point3f0(enu[i]...)
	end

	"Vehicle's rotatation at a given index"
	function get_rot(i)
		rotationmatrix_z(hdg[i])
	end

	iconsize = 8mm
	editarea, viewarea = x_partition_abs(window.area, round(Int, 8.2 * iconsize))
	edit_screen = Screen(window, area = editarea)
	view_screen = Screen(window, area = viewarea)
	#cubecamera(view_screen)

	time_viz, time_value = labeled_slider(1:1:length(t), edit_screen)

	# i_cur is the main Signal: Current index into data arrays using playback slider
	play_viz, i_cur = play_slider(edit_screen, iconsize, 1:1:length(t))

	cur_rot = map(get_rot, i_cur)
	curtime = map(i -> t[i], i_cur)
	curtime_s = map(t -> @sprintf("%.2f", t), curtime)
	pos_animation = map(i_cur) do i
		#pos(Int(i))
		#Sphere{Float32}(Point3f0(0.0), 1f0)
		[get_pos_enu(i)]
	end

	#mesh = loadasset(Pkg.dir("UAVisualization") * "/tools/Canopy_v2.STL")

	controls = Pair["playback" => play_viz, "time" => curtime_s]

	vehmarker = visualize((Sphere{Float32}(Point3f0(0.0), 1f0), pos_animation), color=RGBA(1f0,0f0,0f0,1f0))
	enutraj = visualize(Point3f0.(enu)[:], :lines, thickness=2f0, line=4.0)

	#_view(visualize(mesh), view_screen, camera = :perspective)
	_view(vehmarker, view_screen, camera = :perspective)
	_view(visualize(controls, text_scale = 4mm, width = 8*iconsize), edit_screen, camera = :fixed_pixel)
	_view(enutraj, view_screen, camera = :perspective)
	_view(rgbaxis(cur_rot), view_screen, camera = :perspective)

	push!(view_screen.cameras[:perspective].nearclip, 0.02f0)
	center!(view_screen, :perspective)
	renderloop(window)
end

end # module


text:
ENV["GKSwstype"] = "100"
using ADNLPModels
using Documenter
using Printf
using ARCTR

pages = [
    "Introduction" => "index.md",
    "Tutorial" => "benchmark.md",
    "Reference" => "reference.md",
]

makedocs(
    sitename = "ARCTR.jl",
    format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
    modules = [ARCTR],
    pages = pages,
)

deploydocs(repo = "github.com/tmigot/ARCTR.jl.git", push_preview = true, devbranch = "main")


text:
# This file was generated, do not modify it. # hide
using PyPlot
figure(figsize=(8,6))
plot(r.parameter_values, r.measurements)

xticks(1:5, fontsize=12)
yticks(fontsize=12)
xlabel("Maximum depth", fontsize=14)
ylabel("Misclassification rate", fontsize=14)
ylim([0, 1])

savefig("assets/literate/A-model-tuning-hpt.svg") # hide


text:
using ModelingToolkit, OrdinaryDiffEq, Unitful

using Unitful.DefaultSymbols
import Unitful: Voltage, mV, mS, cm, µF, mF
import Base: show, display

@derived_dimension SpecificConductance 𝐈^2*𝐋^-4*𝐌^-1*𝐓^3 # conductance per unit area
@derived_dimension SpecificCapacitance 𝐈^2*𝐋^-4*𝐌^-1*𝐓^4 # capacitance per unit area

@parameters t
const D = Differential(t)

# Ion species -- there could be more to do here (e.g. non-specific cation/anion etc)
abstract type Ion end
abstract type Calcium <: Ion end
abstract type Sodium <: Ion end
abstract type Potassium <: Ion end
abstract type Chloride <: Ion end
abstract type Leak <: Ion end # placeholder for non-specific ion conductance

# Convenience aliases
const Ca = Calcium
const Na = Sodium
const K = Potassium
const Cl = Chloride

abstract type AbstractGateModel end

struct SSTau <: AbstractGateModel
    g∞::Function
    τg::Function
    p::Real # exponent
end

struct AlphaBeta <: AbstractGateModel
    α::Function
    β::Function
    p::Real # exponent
end

AlphaBeta(α, β; exponent = 1) = AlphaBeta(α,β,exponent)
SSTau(g∞,τg; exponent = 1) = SSTau(g∞, τg, exponent)

# Steady-state channel gating
steady_state(gate::AlphaBeta, V) = gate.α(V)/(gate.α(V) + gate.β(V))
steady_state(gate::SSTau, V) = gate.g∞(V)

# Differential equations
function gate_equation(model::AlphaBeta, gate::Num, V::Num)
    α, β = model.α, model.β
    D(gate) ~ α(V) * (1 - gate) - β(V)*gate
end

function gate_equation(model::SSTau, gate::Num, V::Num)
    τ, g∞ = model.τg, model.g∞
    D(gate) ~ (1/τ(V))*(g∞(V) - gate)
end

# Conductance types
abstract type AbstractConductance end

# Eventually define more concrete types
abstract type ChemicalSynapse <: AbstractConductance end
abstract type ElectricalSynapse <: AbstractConductance end
abstract type ExchangePump <: AbstractConductance end

# This should be better (e.g. provide a callback function Iapp(t) or similar)
struct Stimulus <: AbstractConductance
    Iapp::Real
end

function Stimulus(; amps = 800pA)
    amp_val = ustrip(Float64, mA, amps)
    return Stimulus(amp_val)
end

struct PassiveChannel{I} <: AbstractConductance where {I <: Ion}
    gbar::Real
end

# FIXME: This is lazy--likely define this as a trivial ODESystem like
# "Pin" or "Ground" in the MTK tutorial
function PassiveChannel{I}(; max_g::SpecificConductance) where {I<:Ion}
    gbar_val = ustrip(Float64, mS/cm^2, max_g)
    return PassiveChannel{I}(gbar_val)
end

struct VoltageGatedChannel{I} <: AbstractConductance where {I <: Ion}
    sys::ODESystem
end

struct VoltageIonGatedChannel{I,C} <: AbstractConductance where {I <: Ion, C <: Ion}
    sys::ODESystem
end

# Return ODESystem pretty printing for our wrapper types
Base.show(io::IO, ::MIME"text/plain", x::VoltageGatedChannel) = Base.display(x.sys)

const ChordConductance{I} = Union{VoltageGatedChannel{I},PassiveChannel{I}}

ion_type(x::ChordConductance{I}) where {I<:Ion} = I

function VoltageGatedChannel{I}(gate_models::Vector{<:AbstractGateModel}; max_g::SpecificConductance,
                           Vinit::Voltage = -65mV, name) where {I<:Ion}

    # Strip off units
    gbar_val = ustrip(Float64, mS/cm^2, max_g)
    Vinit_val = ustrip(Float64, mV, Vinit)

    # Where gates are m, h etc and g(t) is channel conductance at time t
    states = @variables V(t), gates[1:length(gate_models)](t), g(t)
    params = @parameters gbar # ̄g = maximal conductance
    eqs = [gate_equation(x, gates[i], V) for (i,x) in enumerate(gate_models)]

    # g*x^a*y^b...
    total_g = g ~ gbar * *((gates[i]^(x.p) for (i,x) in enumerate(gate_models))...)
    eqs = [eqs..., total_g]

    defaultmap = [V => Vinit_val, gbar => gbar_val,
                  (gates[i] => steady_state(x, Vinit_val) for (i,x) in enumerate(gate_models))...]

    system = ODESystem(eqs, t, [states...], [params...]; name = name, defaults = defaultmap)

    return VoltageGatedChannel{I}(system)
end

function VoltageIonGatedChannel{I,C}(gate_models::Vector{<:AbstractGateModel}; max_g::SpecificConductance,
                                     Vinit::Voltage = -65mV, name) where {I<:Ion, C<:Ion}

abstract type AbstractNeuron end

struct Neuron <: AbstractNeuron
    sys::ODESystem
end

function Neuron(channels::Vector{<:AbstractConductance}, reversals; cm::SpecificCapacitance = 1µF/cm^2)

    #Strip off units
    area = 1.0 # for now assume surface area 1.0
    cm_val = ustrip(Float64, mF/cm^2, cm)
    rev_vals = IdDict(reversals)
    states = @variables V(t), I[1:length(channels)](t)
    params = @parameters Cm E[1:length(reversals)]

    ions = unique([ion_type(x) for x in channels])
    eq = [D(V) ~ -sum(chan.sys.g * area * (V - E[]))]

end

max_stars_repo_name: wsphillips/FutureNeuralSimulation

max_issues_repo_name: wsphillips/FutureNeuralSimulation

max_forks_repo_name: wsphillips/FutureNeuralSimulation


text:
using Revise
using DispatchOps
using DataFrames, Query, JuMP

#= SIM INPUT
    1. khazanah (dr readData)
    2. demand estim (dr readData)
    3. demand realization (bikin generator)
    4. initial stock (dr readData)
    5. trayek used (dr readData)
=#
khazanah, demand, init, trayek = readData("/home/kreiton/.julia/dev/DispatchOps/data");

# SIM INPUT (demand realization)
demand_realization = DataFrame(
    id = demand.id,
    periode = demand.periode,
    pecahan = demand.pecahan,
    value = demand.value .+ rand([-5,5])
)

# SIM ACCUMULATOR - INVENTORY
inventory = DataFrame(
    id = String[],
    periode = Int[],
    pecahan = String[],
    value = Float64[]
)
append!(inventory,
    DataFrame(
        id = init.id,
        periode = 0,
        pecahan = init.pecahan,
        value = init.value
    )
)

# SIM ACCUMULATOR - DISPATCH QUEUE

# SIM ACCUMULATOR - EXECUTED DISPATCH

# SIM PARAM
pH = 6
rH = 1
gap = 0.2

# SIM START
t = 0

stock = filter(p -> p.periode == t, inventory)
estim = filter(p -> p.periode >= t + 1 && p.periode <= t + pH, demand)

EG = buildGraph(khazanah, trayek, t, pH)
model = buildModel(EG, estim, stock)
optimizeModel!(model, gap = gap)

plannedDispatch = MetaDigraph{locper}()
for a in filter_arcs(EG, :type, "transport")
    v = value.(model[:flow][a,:]) |> sum
    if v > 0
        add_arc!(plannedDispatch, a)
        set_props!(plannedDispatch, a, EG[a])
        set_props!(plannedDispatch, a,
            Dict(
                :flow => Dict(
                    p => value(model[:flow][a,p])
                    for p in model[:flow][a,:].axes[1]
                ),
                :aggr => v,
                :trip => value(model[:trip][a]) |> round
            )
        )
    end
end

executedDispatch = MetaDigraph{locper}()
for a in filter(p -> src(p).per == t, arcs(plannedDispatch) |> collect)
    add_arc!(executedDispatch, a)
    set_props!(executedDispatch, a, plannedDispatch[a])
end

# TRANSPORT EXECUTION - BEGIN WITH EMPTY DF FOR NEXT PERIOD INV
new_stock = DataFrame(
    id = stock.id,
    periode = t + 1,
    pecahan = stock.pecahan,
    value = stock.value
)
new_stock

# TRANSPORT EXECUTION - EXECUTE ALL DISPATCH
for a in arcs(executedDispatch)
    for p in keys(executedDispatch[a][:flow])
        # stock reduction for source nodes
        new_stock[
            new_stock.id .== src(a).loc .&&
            new_stock.periode .== src(a).per + 1 .&&
            new_stock.pecahan .== p,:value
        ] .-= executedDispatch[a][:flow][p]
        # stock addition for target nodes
        new_stock[
            new_stock.id .== tgt(a).loc .&&
            new_stock.periode .== tgt(a).per .&&
            new_stock.pecahan .== p, :value
        ] .+= executedDispatch[a][:flow][p]
    end
end
new_stock

# DEMAND REALIZATION
realdemand = filter(p -> p.periode == t + 1, demand_realization)
for r in eachrow(realdemand)
    new_stock[
        new_stock.id .== r.id .&&
        new_stock.periode .== r.periode .&&
        new_stock.pecahan .== r.pecahan, :value
    ] .= max(0, new_stock[
        new_stock.id .== r.id .&&
        new_stock.periode .== r.periode .&&
        new_stock.pecahan .== r.pecahan, :value
    ][1] - r.value)
end

append!(inventory, new_stock)

# MOVE FORWARD TIMESTEP
t += 1


text:
module Network

import Microwave: AbstractParams, AbstractData
using ..Circuit
import ..Circuit: check_frequency_identical
import Base: +, -, *, /, ^, convert, promote_rule, show, getindex

export NetworkParams, Sparams, Yparams, Zparams, TwoPortParams, ABCDparams
export Port, NetworkData, impedances, swap_ports!, permute_ports!, check_frequency_identical
export connect_ports, innerconnect_ports, cascade
export reflection_coefficient, transmission_coefficient, impedance_step
export series_network, parallel_network, terminated_network, π_network, t_network

include("NetworkParams.jl")
include("Port.jl")
include("NetworkData.jl")
include("convert.jl")
include("operations.jl")
include("connect.jl")
include("twoport.jl")

end


text:
# import VimBindings: verb_part, text_object_part, well_formed, parse_value, command, matched_rule, parse_command
using VimBindings.Commands
using VimBindings.Parse
import VimBindings.Parse: command

@testset "parse commands into parts" begin
    cmd = "d10w"
    @test well_formed(cmd) == true
    r = matched_rule(cmd)
    @test command(match(r, cmd)) == OperatorCommand(1,
                                                    'd',
                                                    10,
                                                    'w')

end

@testset "parse line operator commands" begin
    @test parse_command("5dd") == LineOperatorCommand(5,
                                                      'd')
    @test parse_command("100yy") == LineOperatorCommand(100,
                                                        'y')

    @test parse_command("100yd") === nothing
    @test parse_command("yy") === LineOperatorCommand(1, 'y')
end

@testset "parse motion commands" begin
    @test parse_command("l") == MotionCommand(1, 'l')
    @test parse_command("B") == MotionCommand(1, 'B')
    @test parse_command("^") == MotionCommand(1, '^')
    @test parse_command("10w") == MotionCommand(10, 'w')
end


text:
# TODO: create test functions for the methods

# """
#     Left hand side of eqn(6) in https://arxiv.org/pdf/1703.00443.pdf
# """
# function create_LHS_matrix(z, λ, Q, G, h, A=nothing)
#     if A == nothing || size(A)[1] == 0
#         return [Q                G';
#                 Diagonal(λ) * G    Diagonal(G * z - h)]
#     else
#         @assert size(A)[2] == size(G)[2]
#         p, n = size(A)
#         m    = size(G)[1]
#         return [Q                  G'                    A';
#                 Diagonal(λ) * G    Diagonal(G * z - h)   zeros(m, p);
#                 A                  zeros(p, m)           zeros(p, p)]
#     end
# end

"""
    create_LHS_matrix(z, λ, Q, G, h, A=nothing)

Inverse matrix specified on RHS of eqn(7) in https://arxiv.org/pdf/1703.00443.pdf

Helper method while calling [`backward_quad`](@ref)
"""
function create_LHS_matrix(z, λ, Q, G, h, A=nothing)::AbstractMatrix{Float64}
    if A === nothing || size(A)[1] == 0
        return [Q         G' * Diagonal(λ);
                G         Diagonal(G * z - h)]
    else
        p, n = size(A)
        m    = size(G, 1)
        if n != size(G, 2)
            throw(DimensionError("Sizes of $A and $G do not match"))
        end
        return [Q         G' * Diagonal(λ)       A';
                G         Diagonal(G * z - h)    spzeros(m, p);
                A         spzeros(p, m)          spzeros(p, p)]
    end
end

# """
#     Right hand side of eqn(6) in https://arxiv.org/pdf/1703.00443.pdf
# """
# function create_RHS_matrix(z, dQ, dq, λ, dG, dh, ν=nothing, dA=nothing, db=nothing)
#     if dA == nothing || size(dA)[1] == 0
#         return -[dQ * z + dq + dG' * λ      ;
#                  Diagonal(λ) * (dG * z - dh)]
#     else
#         return -[dQ * z + dq + dG' * λ + dA' * ν;
#                  Diagonal(λ) * (dG * z - dh)    ;
#                  dA * z - db                    ]
#     end
# end

"""
    get_problem_data(model::MOI.AbstractOptimizer)

Return problem parameters as matrices along with other program info such as number of constraints, variables, etc
"""
function get_problem_data(model::MOI.AbstractOptimizer)
    var_list = MOI.get(model, MOI.ListOfVariableIndices())
    nz = length(var_list)

    # handle inequality constraints
    le_con_idx = MOI.get(
                        model,
                        MOI.ListOfConstraintIndices{
                            MOI.ScalarAffineFunction{Float64},
                            MOI.LessThan{Float64},
                        }())
    ge_con_idx = MOI.get(
                        model,
                        MOI.ListOfConstraintIndices{
                            MOI.ScalarAffineFunction{Float64},
                            MOI.GreaterThan{Float64},
                        }())
    nineq_le = length(le_con_idx)
    nineq_ge = length(ge_con_idx)
    le_con_sv_idx = MOI.get(
                        model,
                        MOI.ListOfConstraintIndices{
                            MOI.SingleVariable,
                            MOI.LessThan{Float64},
                        }())
    ge_con_sv_idx = MOI.get(
                        model,
                        MOI.ListOfConstraintIndices{
                            MOI.SingleVariable,
                            MOI.GreaterThan{Float64},
                        }())
    nineq_sv_le = length(le_con_sv_idx)
    nineq_sv_ge = length(ge_con_sv_idx)

    G = spzeros(nineq_le + nineq_ge + nineq_sv_le + nineq_sv_ge, nz)
    h = spzeros(nineq_le + nineq_ge + nineq_sv_le + nineq_sv_ge)

    for i in 1:nineq_le
        con = le_con_idx[i]

        func = MOI.get(model, MOI.ConstraintFunction(), con)
        set = MOI.get(model, MOI.ConstraintSet(), con)

        for (j, var_idx) in enumerate(var_list)
            for term in func.terms
                if term.variable_index == var_idx
                    G[i,j] = MOI.coefficient(term)
                end
            end
        end
        h[i] = set.upper - func.constant
    end
    for i in 1:nineq_ge # note: ax >= b needs to be converted in Gx <= h form
        con = ge_con_idx[i]

        func = MOI.get(model, MOI.ConstraintFunction(), con)
        set = MOI.get(model, MOI.ConstraintSet(), con)

        for (j, var_idx) in enumerate(var_list)
            for term in func.terms
                if term.variable_index == var_idx
                    G[i+nineq_le,j] = -MOI.coefficient(term)
                end
            end
        end
        h[i] = func.constant - set.lower
    end
    for i in eachindex(le_con_sv_idx)
        con = le_con_sv_idx[i]
        func = MOI.get(model, MOI.ConstraintFunction(), con)
        set = MOI.get(model, MOI.ConstraintSet(), con)
        vidx = findfirst(v -> v == func.variable, var_list)
        G[i+nineq_le+nineq_ge,vidx] = 1
        h[i+nineq_le+nineq_ge] = MOI.constant(set)
    end
    for i in eachindex(ge_con_sv_idx)  # note: x >= b needs to be converted in Gx <= h form
        con = ge_con_sv_idx[i]
        func = MOI.get(model, MOI.ConstraintFunction(), con)
        set = MOI.get(model, MOI.ConstraintSet(), con)
        vidx = findfirst(v -> v == func.variable, var_list)
        G[i+nineq_le+nineq_ge+nineq_sv_le,vidx] = -1
        h[i+nineq_le+nineq_ge+nineq_sv_le] = -MOI.constant(set)
    end

    # handle equality constraints
    eq_con_idx = MOI.get(
                        model,
                        MOI.ListOfConstraintIndices{
                            MOI.ScalarAffineFunction{Float64},
                            MOI.EqualTo{Float64}
                        }())
    neq = length(eq_con_idx)

    eq_con_sv_idx = MOI.get(
        model,
        MOI.ListOfConstraintIndices{
            MOI.SingleVariable,
            MOI.EqualTo{Float64}
        }())
    neq_sv = length(eq_con_sv_idx)

    A = spzeros(neq + neq_sv, nz)
    b = spzeros(neq + neq_sv)

    for i in 1:neq
        con = eq_con_idx[i]

        func = MOI.get(model, MOI.ConstraintFunction(), con)
        set = MOI.get(model, MOI.ConstraintSet(), con)

        for x in func.terms
            # never nothing, variable is present
            vidx = findfirst(v -> v == x.variable_index, var_list)
            A[i, vidx] = x.coefficient
        end
        b[i] = set.value - func.constant
    end
    for i in 1:neq_sv
        con = eq_con_sv_idx[i]
        func = MOI.get(model, MOI.ConstraintFunction(), con)
        set = MOI.get(model, MOI.ConstraintSet(), con)
        vidx = findfirst(v -> v == func.variable, var_list)
        A[i+neq,vidx] = 1
        b[i+neq] = set.value
    end

    # handle objective
    # works both for affine and quadratic objective functions
    objective_function = MOI.get(model, MOI.ObjectiveFunction{MOI.ScalarQuadraticFunction{Float64}}())
    Q = spzeros(nz, nz)
    q = spzeros(nz)

    if objective_function isa MathOptInterface.ScalarAffineFunction{Float64}
        for x in objective_function.terms
            vidx = findfirst(v -> v == x.variable_index, var_list)
            q[vidx] = x.coefficient
        end
    elseif objective_function isa MathOptInterface.ScalarQuadraticFunction{Float64}

        var_to_id = Dict(var_list .=> 1:nz)

        for quad in objective_function.quadratic_terms
            i = var_to_id[quad.variable_index_1]
            j = var_to_id[quad.variable_index_2]
            Q[i,j] = quad.coefficient
            Q[j,i] = quad.coefficient
        end

        q = MOI.coefficient.(objective_function.affine_terms)
    end

    return (
        Q, q, G, h, A, b,
        nz, var_list,
        nineq_le, le_con_idx,
        nineq_ge, ge_con_idx,
        nineq_sv_le, le_con_sv_idx,
        nineq_sv_ge, ge_con_sv_idx,
        neq, eq_con_idx,
        neq_sv, eq_con_sv_idx,
    )
end

# used for testing mostly
# computes the tuple `(s, y)`, vectorized forms of the slack and duals of inequality constraints
function _slack_dual_vectors(model)
    cone_types = unique!([S for (_, S) in MOI.get(model.optimizer, MOI.ListOfConstraints())])
    conic_form = MatOI.GeometricConicForm{Float64, MatOI.SparseMatrixCSRtoCSC{Float64, Int, MatOI.OneBasedIndexing}, Vector{Float64}}(cone_types)
    index_map = MOI.copy_to(conic_form, model)

    # fix optimization sense
    if MOI.get(model, MOI.ObjectiveSense()) == MOI.MAX_SENSE
        conic_form.sense = MOI.MIN_SENSE
        conic_form.c = -conic_form.c
    end
    s = map_rows(model.optimizer, conic_form, index_map, Flattened{Float64}()) do (ci, r)
        MOI.get(model, MOI.ConstraintPrimal(), ci)
    end
    y = map_rows(model.optimizer, conic_form, index_map, Flattened{Float64}()) do (ci, r)
        MOI.get(model, MOI.ConstraintDual(), ci)
    end
    return (s, y)
end


text:
struct UnionArray{T, N, P} <: Abstract.UnionArray{T, N}
    parent::P

    function UnionArray(parent::P) where {T, N, P <: AbstractArray{T, N}}
        # Make sure that `A.parent` is not a UnionArray/Vector
        P <: UnionArrayImpls && throw(MethodError(UnionArray, (A,)))
        return new{T, N, P}(parent)
    end
end

const UnionArrayImpls = Union{UnionVector, UnionArray}

default_reshape(A::UnionVector, dims::T) where T =
    invoke(reshape, Tuple{AbstractArray, T}, A, dims)

ua_reshape(A::UnionVector, dims) = UnionArray(default_reshape(A, dims))
ua_reshape(A::UnionArray, dims) = UnionArray(reshape(A.parent, dims))
ua_reshape(A::UnionVector, ::Tuple{Colon}) = A
function ua_reshape(A::UnionArray, ::Tuple{Colon})
    v = reshape(A.parent, :)
    if v isa UnionVector
        return v
    else
        return UnionArray(v)
    end
end

Base.reshape(A::UnionArrayImpls, dims::Dims) =
    ua_reshape(A, dims) :: UnionArrayImpls
Base.reshape(A::UnionArrayImpls, dims::Tuple{Vararg{Union{Int,Colon}}}) =
    ua_reshape(A, dims) :: UnionArrayImpls

Base.parent(A::UnionArray) = A.parent
Base.size(A::UnionArray) = size(A.parent)
Base.getindex(A::UnionArray, I...) = A.parent[I...]
Base.setindex!(A::UnionArray, v, I...) = setindex!(A.parent, v, I...)

# Base.showarg(io::IO, A::UnionArray, toplevel) =
#     Base.showarg(io, A.parent, toplevel)


text:
using JuMP

using DataStructures
using Parameters
using Printf
using SparseArrays, LinearAlgebra

import Xpress

root_path = dirname(@__DIR__)
push!(LOAD_PATH, root_path);
cd(root_path)
include(joinpath(root_path, "src", "PTDF.jl"));

## programme d'optimisation pour calculer theta
function compute_theta(B,P)
    nb_buses = length(P);
    model = Model(Xpress.Optimizer);
    set_optimizer_attribute(model, "OUTPUTLOG", 0);

    @variable(model, theta[1:nb_buses]);
    @objective(model,Min,0);

    for i in range(1,nb_buses)
        @constraint(model,sum(B[i,j] * theta[j] for j in range(1,nb_buses)) == P[i]);
    end

    # print(model);

    optimize!(model);

    return [value(theta[i]) for i in range(1,nb_buses)];
end;

function compute_ptdf_by_optim(network::PTDF.Network,ref_bus_num::Int, EPS_DIAG::Float64)
    nb_buses = length(network.buses);
    nb_branches = length(network.branches);

    # calcul de B
    B=PTDF.get_B(network,EPS_DIAG,false);

    # calcul de la colonne de la PTDF pour chaque noeud
    PTDF_matrix = zeros(nb_branches,nb_buses);
    # println(network.buses)
    for bus in network.buses
        bus_num = bus[1]
        # println("$bus, id is $bus_num")
        #définition de la distribution des injections
        P = zeros(nb_buses);
        if bus_num != ref_bus_num
            P[bus_num]=1;
            P[ref_bus_num]+=-1;
        end

        #calcul de theta tel que B*theta = P
        theta = compute_theta(B,P);

        #calcul de la PTDF
        for branch in collect(network.branches)
            branch_num = branch[1];
            x_inv = PTDF.get_b(branch[2]);

            ior = branch[2].from;
            iex = branch[2].to;

            PTDF_matrix[branch_num,bus_num] = x_inv*(theta[ior]-theta[iex]);
        end
    end
    return PTDF_matrix;
end;

function compare_ptdfs(network::PTDF.Network,ref_bus_num::Int, EPS_DIAG::Float64)

    ptdf_optim = compute_ptdf_by_optim(network,ref_bus_num, EPS_DIAG);
    ptdf_with_inv = PTDF.compute_ptdf(network,ref_bus_num, EPS_DIAG);

    return ptdf_optim-ptdf_with_inv;
end;

##### Example
EPS_DIAG = 0e-6;
ref_bus_num = 1;
data_path_sparse = joinpath(@__DIR__, "..", "data_matpower", "case1354pegase");
data_path = data_path_sparse;

network = PTDF.read_network(data_path);

# ptdf_optim = compute_ptdf_by_optim(network,ref_bus_num);
ptdf_difference = norm(compare_ptdfs(network, ref_bus_num, EPS_DIAG));
println("error_max = $ptdf_difference");

max_stars_repo_name: klorel/pscopf

max_issues_repo_name: klorel/pscopf

max_forks_repo_name: klorel/pscopf