### Soil erosion ###
@get_units @with_kw struct LandSediment{T}
    # number of cells
    n::Int | "-"
    ### Soil erosion part ###
    # length of cells in y direction [m]
    yl::Vector{T} | "m"
    # length of cells in x direction [m]
    xl::Vector{T} | "m"
    # Fraction of river [-]
    riverfrac::Vector{T} | "-"
    # Waterbodies areas [-]
    wbcover::Vector{T} | "-"
    # Depth of overland flow [m]
    h_land::Vector{T} | "m"
    # Canopy interception [mm]
    interception::Vector{T} | "mm"
    # Precipitation [mm]
    precipitation::Vector{T} | "mm"
    # Overland flow [m3/s]
    q_land::Vector{T} | "m3 s-1"
    # Canopy height [m]
    canopyheight::Vector{T} | "m"
    # Canopy gap fraction [mm]
    canopygapfraction::Vector{T} | "mm"
    # Coefficient for EUROSEM rainfall erosion [-]
    erosk::Vector{T} | "-"
    # Exponent for EUROSEM rainfall erosion [-]
    erosspl::Vector{T} | "-"
    # Coefficient for ANSWERS overland flow erosion [-]
    erosov::Vector{T} | "-"
    # Fraction of impervious area per grid cell [-]
    pathfrac::Vector{T} | "-"
    # Land slope [-]
    slope::Vector{T} | "-"
    # USLE crop management factor [-]
    usleC::Vector{T} | "-"
    # USLE soil erodibility factor [-]
    usleK::Vector{T} | "-"
    # Sediment eroded by rainfall [ton]
    sedspl::Vector{T} | "t"
    # Sediment eroded by overland flow [ton]
    sedov::Vector{T} | "t"
    # Total eroded soil [ton]
    soilloss::Vector{T} | "t"
    # Eroded soil per particle class [ton]
    erosclay::Vector{T} | "t"  # clay
    erossilt::Vector{T} | "t"  # silt
    erossand::Vector{T} | "t"  # sand
    erossagg::Vector{T} | "t"  # small aggregates
    eroslagg::Vector{T} | "t"  # large aggregates
    ## Interception related to leaf_area_index climatology ###
    # Specific leaf storage [mm]
    sl::Vector{T} | "mm"
    # Storage woody part of vegetation [mm]
    swood::Vector{T} | "mm"
    # Extinction coefficient [-] (to calculate canopy gap fraction)
    kext::Vector{T} | "-"
    # Leaf area index [m² m⁻²]
    leaf_area_index::Vector{T} | "m2 m-2"
    ### Transport capacity part ###
    # Drain length [m]
    dl::Vector{T} | "m"
    # Flow width [m]
    dw::Vector{T} | "m"
    # Govers transport capacity coefficients [-]
    cGovers::Vector{T} | "-"
    nGovers::Vector{T} | "-"
    # Median particle diameter of the topsoil [mm]
    D50::Vector{T} | "mm"
    # Median diameter per particle class [um]
    dmclay::Vector{T} | "µm"
    dmsilt::Vector{T} | "µm"
    dmsand::Vector{T} | "µm"
    dmsagg::Vector{T} | "µm"
    dmlagg::Vector{T} | "µm"
    # Fraction of the different particle class [-]
    fclay::Vector{T} | "-"
    fsilt::Vector{T} | "-"
    fsand::Vector{T} | "-"
    fsagg::Vector{T} | "-"
    flagg::Vector{T} | "-"
    # Density of sediment [kg/m3]
    rhos::Vector{T} | "kg m-3"
    # Filter with river cells
    rivcell::Vector{T} | "-"
    # Total transport capacity of overland flow [ton]
    TCsed::Vector{T} | "t"
    # Transport capacity of overland flow per particle class [ton]
    TCclay::Vector{T} | "t"
    TCsilt::Vector{T} | "t"
    TCsand::Vector{T} | "t"
    TCsagg::Vector{T} | "t"
    TClagg::Vector{T} | "t"

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

statevars(::LandSediment) = ()

function initialize_landsed(nc, config, river, riverfrac, xl, yl, inds)
    # Initialize parameters for the soil loss part
    n = length(inds)
    do_river = get(config.model, "runrivermodel", false)::Bool
    # Reservoir / lake
    do_reservoirs = get(config.model, "doreservoir", false)::Bool
    do_lakes = get(config.model, "dolake", false)::Bool
    # Rainfall erosion equation: ["answers", "eurosem"]
    rainerosmethod = get(config.model, "rainerosmethod", "answers")::String
    # Overland flow transport capacity method: ["yalinpart", "govers", "yalin"]
    landtransportmethod = get(config.model, "landtransportmethod", "yalinpart")::String

    altitude =
        ncread(nc, param(config, "input.vertical.altitude"); sel = inds, type = Float)
    canopyheight = ncread(
        nc,
        param(config, "input.vertical.canopyheight", nothing);
        sel = inds,
        defaults = 3.0,
        type = Float,
    )
    erosk = ncread(
        nc,
        param(config, "input.vertical.erosk", nothing);
        sel = inds,
        defaults = 0.6,
        type = Float,
    )
    erosspl = ncread(
        nc,
        param(config, "input.vertical.erosspl", nothing);
        sel = inds,
        defaults = 2.0,
        type = Float,
    )
    erosov = ncread(
        nc,
        param(config, "input.vertical.erosov", nothing);
        sel = inds,
        defaults = 0.9,
        type = Float,
    )
    pathfrac = ncread(
        nc,
        param(config, "input.vertical.pathfrac", nothing);
        sel = inds,
        defaults = 0.01,
        type = Float,
    )
    slope = ncread(
        nc,
        param(config, "input.vertical.slope", nothing);
        sel = inds,
        defaults = 0.01,
        type = Float,
    )
    usleC = ncread(
        nc,
        param(config, "input.vertical.usleC", nothing);
        sel = inds,
        defaults = 0.01,
        type = Float,
    )
    usleK = ncread(
        nc,
        param(config, "input.vertical.usleK", nothing);
        sel = inds,
        defaults = 0.1,
        type = Float,
    )

    cmax, _, canopygapfraction, sl, swood, kext = initialize_canopy(nc, config, inds)

    # Initialise parameters for the transport capacity part
    βₗ = slope
    clamp!(βₗ, 0.00001, Inf)
    ldd_2d = ncread(nc, param(config, "input.ldd"); allow_missing = true)
    ldd = ldd_2d[inds]
    dl = fill(mv, n)
    dw = fill(mv, n)

    for i = 1:n
        dl[i] = detdrainlength(ldd[i], xl[i], yl[i])
        dw[i] = detdrainwidth(ldd[i], xl[i], yl[i])
    end
    dmclay = ncread(
        nc,
        param(config, "input.vertical.dmclay", nothing);
        sel = inds,
        defaults = 2.0,
        type = Float,
    )
    dmsilt = ncread(
        nc,
        param(config, "input.vertical.dmsilt", nothing);
        sel = inds,
        defaults = 10.0,
        type = Float,
    )
    dmsand = ncread(
        nc,
        param(config, "input.vertical.dmsand", nothing);
        sel = inds,
        defaults = 200.0,
        type = Float,
    )
    dmsagg = ncread(
        nc,
        param(config, "input.vertical.dmsagg", nothing);
        sel = inds,
        defaults = 30.0,
        type = Float,
    )
    dmlagg = ncread(
        nc,
        param(config, "input.vertical.dmlagg", nothing);
        sel = inds,
        defaults = 500.0,
        type = Float,
    )
    pclay = ncread(
        nc,
        param(config, "input.vertical.pclay", nothing);
        sel = inds,
        defaults = 0.1,
        type = Float,
    )
    psilt = ncread(
        nc,
        param(config, "input.vertical.psilt", nothing);
        sel = inds,
        defaults = 0.1,
        type = Float,
    )
    rhos = ncread(
        nc,
        param(config, "input.vertical.rhosed", nothing);
        sel = inds,
        defaults = 2650.0,
        type = Float,
    )

    ### Initialize transport capacity variables ###
    rivcell = float(river)
    # Percent Sand
    psand = 100 .- pclay .- psilt
    # Govers coefficient for transport capacity
    if landtransportmethod != "yalinpart"
        # Calculation of D50 and fraction of fine and very fine sand (fvfs) from Fooladmand et al, 2006
        psand999 = psand .* ((999 - 25) / (1000 - 25))
        vd50 = log.((1 ./ (0.01 .* (pclay .+ psilt)) .- 1) ./ (1 ./ (0.01 .* pclay) .- 1))
        wd50 =
            log.(
                (1 ./ (0.01 .* (pclay .+ psilt .+ psand999)) .- 1) ./
                (1 ./ (0.01 .* pclay) .- 1),
            )
        ad50 = 1 / log((25 - 1) / (999 - 1))
        bd50 = ad50 ./ log((25 - 1) / 1)
        cd50 = ad50 .* log.(vd50 ./ wd50)
        ud50 = (.-vd50) .^ (1 .- bd50) ./ ((.-wd50) .^ (.-bd50))
        D50 = 1 .+ (-1 ./ ud50 .* log.(1 ./ (1 ./ (0.01 .* pclay) .- 1))) .^ (1 ./ cd50) #[um]
        D50 = D50 ./ 1000 # [mm]
    else
        D50 = fill(mv, n)
    end
    if landtransportmethod == "govers"
        cGovers = ((D50 .* 1000 .+ 5) ./ 0.32) .^ (-0.6)
        nGovers = ((D50 .* 1000 .+ 5) ./ 300) .^ (0.25)
    else
        cGovers = fill(mv, n)
        nGovers = fill(mv, n)
    end
    if do_river || landtransportmethod == "yalinpart"
        # Determine sediment size distribution, estimated from primary particle size distribution (Foster et al., 1980)
        fclay = 0.20 .* pclay ./ 100
        fsilt = 0.13 .* psilt ./ 100
        fsand = 0.01 .* psand .* (1 .- 0.01 .* pclay) .^ (2.4)
        fsagg = 0.28 .* (0.01 .* pclay .- 0.25) .+ 0.5
        for i = 1:n
            if pclay[i] > 50.0
                fsagg[i] = 0.57
            elseif pclay[i] < 25
                fsagg[i] = 2.0 * 0.01 * pclay[i]
            end
        end
        flagg = 1.0 .- fclay .- fsilt .- fsand .- fsagg
    else
        fclay = fill(mv, n)
        fsilt = fill(mv, n)
        fsand = fill(mv, n)
        fsagg = fill(mv, n)
        flagg = fill(mv, n)
    end

    # Reservoir and lakes
    wbcover = zeros(Float, n)
    if do_reservoirs
        rescoverage_2d = ncread(
            nc,
            param(config, "input.vertical.resareas");
            sel = inds,
            type = Float,
            fill = 0.0,
        )
        wbcover = wbcover .+ rescoverage_2d
    end
    if do_lakes
        lakecoverage_2d = ncread(
            nc,
            param(config, "input.vertical.lakeareas");
            sel = inds,
            type = Float,
            fill = 0.0,
        )
        wbcover = wbcover .+ lakecoverage_2d
    end

    eros = LandSediment{Float}(
        n = n,
        yl = yl,
        xl = xl,
        riverfrac = riverfrac,
        wbcover = wbcover,
        ### Soil erosion part ###
        # Forcing
        interception = fill(mv, n),
        h_land = fill(mv, n),
        precipitation = fill(mv, n),
        q_land = fill(mv, n),
        # Parameters
        canopyheight = canopyheight,
        canopygapfraction = canopygapfraction,
        erosk = erosk,
        erosspl = erosspl,
        erosov = erosov,
        pathfrac = pathfrac,
        slope = slope,
        usleC = usleC,
        usleK = usleK,
        # Interception related to climatology (leaf_area_index)
        sl = sl,
        swood = swood,
        kext = kext,
        leaf_area_index = fill(mv, n),
        # Outputs
        sedspl = fill(mv, n),
        sedov = fill(mv, n),
        soilloss = fill(mv, n),
        erosclay = fill(mv, n),
        erossilt = fill(mv, n),
        erossand = fill(mv, n),
        erossagg = fill(mv, n),
        eroslagg = fill(mv, n),
        ### Transport capacity part ###
        # Parameters
        dl = dl,
        dw = dw,
        cGovers = cGovers,
        D50 = D50,
        dmclay = dmclay,
        dmsilt = dmsilt,
        dmsand = dmsand,
        dmsagg = dmsagg,
        dmlagg = dmlagg,
        fclay = fclay,
        fsilt = fsilt,
        fsand = fsand,
        fsagg = fsagg,
        flagg = flagg,
        nGovers = nGovers,
        rhos = rhos,
        rivcell = rivcell,
        # Outputs
        TCsed = fill(mv, n),
        TCclay = fill(mv, n),
        TCsilt = fill(mv, n),
        TCsand = fill(mv, n),
        TCsagg = fill(mv, n),
        TClagg = fill(mv, n),
    )

    return eros
end

# Soil erosion
function update_until_ols(eros::LandSediment, config)
    # Options from config
    do_lai = haskey(config.input.vertical, "leaf_area_index")
    rainerosmethod = get(config.model, "rainerosmethod", "answers")::String
    Δt = Second(config.timestepsecs)
    ts = Float(Δt.value)

    for i = 1:eros.n

        ### Splash / Rainfall erosion ###
        # ANSWERS method
        if rainerosmethod == "answers"
            # calculate rainfall intensity [mm/min]
            rintnsty = eros.precipitation[i] / (ts / 60)
            # splash erosion [kg/min]
            sedspl =
                0.108 * eros.usleC[i] * eros.usleK[i] * eros.xl[i] * eros.yl[i] * rintnsty^2
            # [ton/timestep]
            sedspl = sedspl * (ts / 60) * 1e-3
        end
        # TODO check eurosem method output values (too high!)
        if rainerosmethod == "eurosem"
            if do_lai
                cmax = eros.sl[i] * eros.leaf_area_index[i] + eros.swood[i]
                canopygapfraction = exp(-eros.kext[i] * eros.leaf_area_index[i])
            end
            # calculate rainfall intensity [mm/h]
            rintnsty = eros.precipitation[i] / (ts / 3600)
            # Kinetic energy of direct throughfall [J/m2/mm]
            # kedir = max(11.87 + 8.73 * log10(max(0.0001, rintnsty)),0.0) #basis used in USLE
            kedir = max(8.95 + 8.44 * log10(max(0.0001, rintnsty)), 0.0) #variant used in most distributed mdoels
            # Kinetic energy of leaf drainage [J/m2/mm]
            pheff = 0.5 * eros.canopyheight[i]
            keleaf = max((15.8 * pheff^0.5) - 5.87, 0.0)

            #Depths of rainfall (total, leaf drianage, direct) [mm]
            rdtot = eros.precipitation[i]
            rdleaf = rdtot * 0.1 * canopygapfraction #stemflow
            rddir = max(rdtot - rdleaf - eros.interception[i], 0.0) #throughfall

            #Total kinetic energy by rainfall [J/m2]
            ketot = (rddir * kedir + rdleaf * keleaf) * 0.001
            # Rainfall / splash erosion [g/m2]
            sedspl = eros.erosk[i] * ketot * exp(-eros.erosspl[i] * eros.h_land[i])
            sedspl = sedspl * eros.xl[i] * eros.yl[i] * 1e-6 # ton/cell
        end

        # Remove the impervious area
        sedspl = sedspl * (1.0 - eros.pathfrac[i])

        ### Overland flow erosion ###
        # ANWERS method
        # Overland flow rate [m2/min]
        qr_land = eros.q_land[i] * 60 / ((eros.xl[i] + eros.yl[i]) / 2)
        # Sine of the slope
        sinslope = sin(atan(eros.slope[i]))

        # Overland flow erosion [kg/min]
        # For a wide range of slope, it is better to use the sine of slope rather than tangeant
        sedov =
            eros.erosov[i] *
            eros.usleC[i] *
            eros.usleK[i] *
            eros.xl[i] *
            eros.yl[i] *
            sinslope *
            qr_land
        # [ton/timestep]
        sedov = sedov * (ts / 60) * 1e-3
        # Remove the impervious area
        sedov = sedov * (1.0 - eros.pathfrac[i])

        # Assume no erosion in reservoir/lake areas
        if eros.wbcover[i] > 0
            sedspl = 0.0
            sedov = 0.0
        end

        # Total soil loss [ton/cell/timestep]
        soilloss = sedspl + sedov

        # update the outputs and states
        eros.sedspl[i] = sedspl
        eros.sedov[i] = sedov
        eros.soilloss[i] = soilloss
        # Eroded amount per particle class
        eros.erosclay[i] = soilloss * eros.fclay[i]
        eros.erossilt[i] = soilloss * eros.fsilt[i]
        eros.erossand[i] = soilloss * eros.fsand[i]
        eros.erossagg[i] = soilloss * eros.fsagg[i]
        eros.eroslagg[i] = soilloss * eros.flagg[i]
    end

end

### Sediment transport capacity in overland flow ###
function update_until_oltransport(ols::LandSediment, config::Config)

    do_river = get(config.model, "runrivermodel", false)::Bool
    tcmethod = get(config.model, "landtransportmethod", "yalinpart")::String
    Δt = Second(config.timestepsecs)
    ts = Float(Δt.value)

    for i = 1:ols.n
        sinslope = sin(atan(ols.slope[i]))

        if !do_river
            # Total transport capacity without particle differentiation
            if tcmethod == "govers"
                TC = tc_govers(ols, i, sinslope, ts)
            end

            if tcmethod == "yalin"
                TC = tc_yalin(ols, i, sinslope, ts)
            end

            # Filter TC land for river cells (0 in order for sediment from land to stop when entering the river)
            if ols.rivcell[i] == 1.0
                TC = 0.0
            end

            # Set particle TC to 0
            TCclay = 0.0
            TCsilt = 0.0
            TCsand = 0.0
            TCsagg = 0.0
            TClagg = 0.0
        end

        if do_river || tcmethod == "yalinpart"
            TC, TCclay, TCsilt, TCsand, TCsagg, TClagg = tc_yalinpart(ols, i, sinslope, ts)
        end

        # update the outputs and states
        ols.TCsed[i] = TC
        ols.TCclay[i] = TCclay
        ols.TCsilt[i] = TCsilt
        ols.TCsand[i] = TCsand
        ols.TCsagg[i] = TCsagg
        ols.TClagg[i] = TClagg
    end

end

function tc_govers(ols::LandSediment, i::Int, sinslope::Float, ts::Float)::Float
    # Transport capacity from govers 1990
    # Unit stream power
    if ols.h_land[i] > 0.0
        velocity = ols.q_land[i] / (ols.dw[i] * ols.h_land[i])
    else
        velocity = 0.0
    end
    omega = 10 * sinslope * 100 * velocity #cm/s
    if omega > 0.4
        TCf = ols.cGovers[i] * (omega - 0.4)^(ols.nGovers[i]) * 2650 #kg/m3
    else
        TCf = 0.0
    end
    TC = TCf * ols.q_land[i] * ts * 1e-3 #[ton/cell]

    return TC
end

function tc_yalin(ols::LandSediment, i::Int, sinslope::Float, ts::Float)::Float
    # Transport capacity from Yalin without particle differentiation
    delta = max(
        (
            ols.h_land[i] * sinslope / (ols.D50[i] * 0.001 * (ols.rhos[i] / 1000 - 1)) /
            0.06 - 1
        ),
        0.0,
    )
    alphay = delta * 2.45 / (0.001 * ols.rhos[i])^0.4 * 0.06^(0.5)
    if ols.q_land[i] > 0.0 && alphay != 0.0
        TC = (
            ols.dw[i] / ols.q_land[i] *
            (ols.rhos[i] - 1000) *
            ols.D50[i] *
            0.001 *
            (9.81 * ols.h_land[i] * sinslope) *
            0.635 *
            delta *
            (1 - log(1 + alphay) / (alphay))
        ) # [kg/m3]
        TC = TC * ols.q_land[i] * ts * 1e-3 #[ton]
    else
        TC = 0.0
    end
    return TC
end

function tc_yalinpart(ols::LandSediment, i::Int, sinslope::Float, ts::Float)
    # Transport capacity from Yalin with particle differentiation
    # Delta parameter of Yalin for each particle class
    delta = ols.h_land[i] * sinslope / (1e-6 * (ols.rhos[i] / 1000 - 1)) / 0.06
    dclay = max(1 / ols.dmclay[i] * delta - 1, 0.0)
    dsilt = max(1 / ols.dmsilt[i] * delta - 1, 0.0)
    dsand = max(1 / ols.dmsand[i] * delta - 1, 0.0)
    dsagg = max(1 / ols.dmsagg[i] * delta - 1, 0.0)
    dlagg = max(1 / ols.dmlagg[i] * delta - 1, 0.0)
    # Total transportability
    dtot = dclay + dsilt + dsand + dsagg + dlagg

    # Yalin transport capacity of overland flow for each particle class
    if ols.q_land[i] > 0.0
        TCa =
            ols.dw[i] / ols.q_land[i] *
            (ols.rhos[i] - 1000) *
            1e-6 *
            (9.81 * ols.h_land[i] * sinslope)
    else
        TCa = 0.0
    end
    TCb = 2.45 / (ols.rhos[i] / 1000)^0.4 * 0.06^0.5
    if dtot != 0.0 && dclay != 0.0
        TCclay =
            TCa * ols.dmclay[i] * dclay / dtot *
            0.635 *
            dclay *
            (1 - log(1 + dclay * TCb) / dclay * TCb) # [kg/m3]
        TCclay = TCclay * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TCclay = 0.0
    end
    if dtot != 0.0 && dsilt != 0.0
        TCsilt =
            TCa * ols.dmsilt[i] * dsilt / dtot *
            0.635 *
            dsilt *
            (1 - log(1 + dsilt * TCb) / dsilt * TCb) # [kg/m3]
        TCsilt = TCsilt * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TCsilt = 0.0
    end
    if dtot != 0.0 && dsand != 0.0
        TCsand =
            TCa * ols.dmsand[i] * dsand / dtot *
            0.635 *
            dsand *
            (1 - log(1 + dsand * TCb) / dsand * TCb) # [kg/m3]
        TCsand = TCsand * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TCsand = 0.0
    end
    if dtot != 0.0 && dsagg != 0.0
        TCsagg =
            TCa * ols.dmsagg[i] * dsagg / dtot *
            0.635 *
            dsagg *
            (1 - log(1 + dsagg * TCb) / dsagg * TCb) # [kg/m3]
        TCsagg = TCsagg * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TCsagg = 0.0
    end
    if dtot != 0.0 && dlagg != 0.0
        TClagg =
            TCa * ols.dmlagg[i] * dlagg / dtot *
            0.635 *
            dlagg *
            (1 - log(1 + dlagg * TCb) / dlagg * TCb) # [kg/m3]
        TClagg = TClagg * ols.q_land[i] * ts * 1e-3 # [ton]
    else
        TClagg = 0.0
    end

    # Assume that ols all reach the river in reservoir/lake areas (very high TC)
    if ols.wbcover[i] > 0
        TC = 1e9
        TCclay = 1e9
        TCsilt = 1e9
        TCsand = 1e9
        TCsagg = 1e9
        TClagg = 1e9
    end

    # Filter TC land for river cells (0 in order for sediment from land to stop when entering the river)
    if ols.rivcell[i] == 1.0
        TCclay = 0.0
        TCsilt = 0.0
        TCsand = 0.0
        TCsagg = 0.0
        TClagg = 0.0
    end

    # Set total TC to 0
    TC = 0.0

    return TC, TCclay, TCsilt, TCsand, TCsagg, TClagg
end

### Sediment transport in overland flow ###
@get_units @with_kw struct OverlandFlowSediment{T}
    # number of cells
    n::Int | "-"
    # Filter with river cells
    rivcell::Vector{T} | "-"
    # Total eroded soil [ton]
    soilloss::Vector{T} | "t"
    # Eroded soil per particle class [ton]
    erosclay::Vector{T} | "t"
    erossilt::Vector{T} | "t"
    erossand::Vector{T} | "t"
    erossagg::Vector{T} | "t"
    eroslagg::Vector{T} | "t"
    # Total transport capacity of overland flow [ton]
    TCsed::Vector{T} | "t"
    # Transport capacity of overland flow per particle class [ton]
    TCclay::Vector{T} | "t"
    TCsilt::Vector{T} | "t"
    TCsand::Vector{T} | "t"
    TCsagg::Vector{T} | "t"
    TClagg::Vector{T} | "t"
    # Sediment flux in overland flow [ton]
    olsed::Vector{T} | "t"
    olclay::Vector{T} | "t"
    olsilt::Vector{T} | "t"
    olsand::Vector{T} | "t"
    olsagg::Vector{T} | "t"
    ollagg::Vector{T} | "t"
    # Sediment reaching the river with overland flow [ton]
    inlandsed::Vector{T} | "t"
    inlandclay::Vector{T} | "t"
    inlandsilt::Vector{T} | "t"
    inlandsand::Vector{T} | "t"
    inlandsagg::Vector{T} | "t"
    inlandlagg::Vector{T} | "t"

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

statevars(::OverlandFlowSediment) = ()

function partial_update!(inland, rivcell, eroded)
    no_erosion = zero(eltype(eroded))
    for i in eachindex(inland)
        inland[i] = rivcell[i] == 1 ? eroded[i] : no_erosion
    end
    return inland
end

function update(ols::OverlandFlowSediment, network, config)
    do_river = get(config.model, "runrivermodel", false)::Bool
    tcmethod = get(config.model, "landtransportmethod", "yalinpart")::String
    zeroarr = fill(0.0, ols.n)

    # transport sediment down the network
    if do_river || tcmethod == "yalinpart"
        # clay
        accucapacityflux!(ols.olclay, ols.erosclay, network, ols.TCclay)
        partial_update!(ols.inlandclay, ols.rivcell, ols.erosclay)
        # silt
        accucapacityflux!(ols.olsilt, ols.erossilt, network, ols.TCsilt)
        partial_update!(ols.inlandsilt, ols.rivcell, ols.erossilt)
        # sand
        accucapacityflux!(ols.olsand, ols.erossand, network, ols.TCsand)
        partial_update!(ols.inlandsand, ols.rivcell, ols.erossand)
        # small aggregates
        accucapacityflux!(ols.olsagg, ols.erossagg, network, ols.TCsagg)
        partial_update!(ols.inlandsagg, ols.rivcell, ols.erossagg)
        # large aggregates
        accucapacityflux!(ols.ollagg, ols.eroslagg, network, ols.TClagg)
        partial_update!(ols.inlandlagg, ols.rivcell, ols.eroslagg)

        # total sediment, all particle classes
        ols.olsed .= ols.olclay .+ ols.olsilt .+ ols.olsand .+ ols.olsagg .+ ols.ollagg
        ols.inlandsed .=
            ols.inlandclay .+ ols.inlandsilt .+ ols.inlandsand .+ ols.inlandsagg .+
            ols.inlandlagg
    else
        accucapacityflux!(ols.olsed, ols.soilloss, network, ols.TCsed)
        ols.inlandsed .= ifelse.(ols.rivcell .== 1, ols.soilloss, zeroarr)
    end
end

### River transport and processes ###
@get_units @with_kw struct RiverSediment{T}
    # number of cells
    n::Int | "-"
    # Timestep [s]
    Δt::T | "s"
    # River geometry (slope [-], length [m], width [m])
    sl::Vector{T} | "m"
    dl::Vector{T} | "m"
    width::Vector{T} | "m"
    # Sediment mean diameter in the river bed [mm]
    d50::Vector{T}
    # Particle mean diameter [mm]
    dmclay::Vector{T}
    dmsilt::Vector{T}
    dmsand::Vector{T}
    dmsagg::Vector{T}
    dmlagg::Vector{T}
    dmgrav::Vector{T}
    # River bed and bank particle fraction composition [-]
    fclayriv::Vector{T} | "-"
    fsiltriv::Vector{T} | "-"
    fsandriv::Vector{T} | "-"
    fgravriv::Vector{T} | "-"
    # Sediment mean diameter for Engelund and Hansen transport equation [mm]
    d50engelund::Vector{T}
    # Parameters for Bagnold transport equation
    cbagnold::Vector{T} | "-"
    ebagnold::Vector{T} | "-"
    # Parameters for Kodatie transport equation
    ak::Vector{T} | "-"
    bk::Vector{T} | "-"
    ck::Vector{T} | "-"
    dk::Vector{T} | "-"
    # Critical bed and bank shear stress [N/m2]
    TCrbank::Vector{T} | "N m-2"
    TCrbed::Vector{T} | "N m-2"
    # Bed and bank erodibilities [m3/N.s]
    kdbank::Vector{T} | "m3 N-1 s-1"
    kdbed::Vector{T} | "m3 N-1 s-1"
    # Sediment density [kg/m3]
    rhos::Vector{T} | "kg m-3"
    # River water level [m]
    h_riv::Vector{T} | "m"
    # River discharge [m3/s]
    q_riv::Vector{T} | "m3 s-1"
    # Sediment input from land erosion [ton]
    inlandclay::Vector{T} | "t"
    inlandsilt::Vector{T} | "t"
    inlandsand::Vector{T} | "t"
    inlandsagg::Vector{T} | "t"
    inlandlagg::Vector{T} | "t"
    inlandsed::Vector{T} | "t"
    # Sediment / particle left in the cell [ton]
    sedload::Vector{T} | "t"
    clayload::Vector{T} | "t"
    siltload::Vector{T} | "t"
    sandload::Vector{T} | "t"
    saggload::Vector{T} | "t"
    laggload::Vector{T} | "t"
    gravload::Vector{T} | "t"
    # Sediment / particle stored on the river bed after deposition [ton]
    sedstore::Vector{T} | "t"
    claystore::Vector{T} | "t"
    siltstore::Vector{T} | "t"
    sandstore::Vector{T} | "t"
    saggstore::Vector{T} | "t"
    laggstore::Vector{T} | "t"
    gravstore::Vector{T} | "t"
    # Sediment / particle flux [ton]
    outsed::Vector{T} | "t"
    outclay::Vector{T} | "t"
    outsilt::Vector{T} | "t"
    outsand::Vector{T} | "t"
    outsagg::Vector{T} | "t"
    outlagg::Vector{T} | "t"
    outgrav::Vector{T} | "t"
    # Sediment concentrations [kg/m3]
    Sedconc::Vector{T} | "kg m-3"
    SSconc::Vector{T} | "kg m-3"
    Bedconc::Vector{T} | "kg m-3"
    # River transport capacity
    maxsed::Vector{T} | "t"
    # Eroded sediment (total, bank and bed)
    erodsed::Vector{T} | "t"
    erodsedbank::Vector{T} | "t"
    erodsedbed::Vector{T} | "t"
    # Deposited sediment
    depsed::Vector{T} | "t"
    # Sediment in
    insed::Vector{T} | "t"
    # Reservoir and lakes
    wbcover::Vector{T} | "-"
    wblocs::Vector{T} | "-"
    wbarea::Vector{T} | "m2"
    wbtrap::Vector{T} | "-"

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

statevars(::RiverSediment) = (
    :clayload,
    :siltload,
    :sandload,
    :saggload,
    :laggload,
    :gravload,
    :claystore,
    :siltstore,
    :sandstore,
    :saggstore,
    :laggstore,
    :gravstore,
    :outclay,
    :outsilt,
    :outsand,
    :outsagg,
    :outlagg,
    :outgrav,
)

function initialize_riversed(nc, config, riverwidth, riverlength, inds_riv)
    # Initialize river parameters
    nriv = length(inds_riv)
    # River flow transport capacity method: ["bagnold", "engelund", "yang", "kodatie", "molinas"]
    tcmethodriv = get(config.model, "rivtransportmethod", "bagnold")::String
    Δt = Second(config.timestepsecs)
    # Reservoir / lakes
    do_reservoirs = get(config.model, "doreservoir", false)::Bool
    do_lakes = get(config.model, "dolake", false)::Bool
    wbcover = zeros(Float, nriv)
    wblocs = zeros(Float, nriv)
    wbarea = zeros(Float, nriv)
    wbtrap = zeros(Float, nriv)

    if do_reservoirs
        reslocs = ncread(
            nc,
            param(config, "input.lateral.river.reslocs");
            sel = inds_riv,
            type = Float,
            fill = 0,
        )
        rescoverage_2d = ncread(
            nc,
            param(config, "input.lateral.river.resareas");
            sel = inds_riv,
            type = Float,
            fill = 0,
        )
        resarea = ncread(
            nc,
            param(config, "input.lateral.river.resarea");
            sel = inds_riv,
            type = Float,
            fill = 0.0,
        )
        restrapefficiency = ncread(
            nc,
            param(config, "input.lateral.river.restrapeff");
            sel = inds_riv,
            type = Float,
            defaults = 1.0,
            fill = 0.0,
        )

        wbcover = wbcover .+ rescoverage_2d
        wblocs = wblocs .+ reslocs
        wbarea = wbarea .+ resarea
        wbtrap = wbtrap .+ restrapefficiency
    end

    if do_lakes
        lakelocs = ncread(
            nc,
            param(config, "input.lateral.river.lakelocs");
            sel = inds_riv,
            type = Float,
            fill = 0,
        )
        lakecoverage_2d = ncread(
            nc,
            param(config, "input.lateral.river.lakeareas");
            sel = inds_riv,
            type = Float,
            fill = 0,
        )
        lakearea = ncread(
            nc,
            param(config, "input.lateral.river.lakearea");
            sel = inds_riv,
            type = Float,
            fill = 0.0,
        )

        wbcover = wbcover .+ lakecoverage_2d
        wblocs = wblocs .+ lakelocs
        wbarea = wbarea .+ lakearea
    end

    riverslope =
        ncread(nc, param(config, "input.lateral.river.slope"); sel = inds_riv, type = Float)
    clamp!(riverslope, 0.00001, Inf)
    rhos = ncread(
        nc,
        param(config, "input.lateral.river.rhosed", nothing);
        sel = inds_riv,
        defaults = 2650.0,
        type = Float,
    )
    dmclay = ncread(
        nc,
        param(config, "input.lateral.river.dmclay", nothing);
        sel = inds_riv,
        defaults = 2.0,
        type = Float,
    )
    dmsilt = ncread(
        nc,
        param(config, "input.lateral.river.dmsilt", nothing);
        sel = inds_riv,
        defaults = 10.0,
        type = Float,
    )
    dmsand = ncread(
        nc,
        param(config, "input.lateral.river.dmsand", nothing);
        sel = inds_riv,
        defaults = 200.0,
        type = Float,
    )
    dmsagg = ncread(
        nc,
        param(config, "input.lateral.river.dmsagg", nothing);
        sel = inds_riv,
        defaults = 30.0,
        type = Float,
    )
    dmlagg = ncread(
        nc,
        param(config, "input.lateral.river.dmlagg", nothing);
        sel = inds_riv,
        defaults = 500.0,
        type = Float,
    )
    dmgrav = ncread(
        nc,
        param(config, "input.lateral.river.dmgrav", nothing);
        sel = inds_riv,
        defaults = 2000.0,
        type = Float,
    )
    fclayriv = ncread(
        nc,
        param(config, "input.lateral.river.fclayriv");
        sel = inds_riv,
        type = Float,
    )
    fsiltriv = ncread(
        nc,
        param(config, "input.lateral.river.fsiltriv");
        sel = inds_riv,
        type = Float,
    )
    fsandriv = ncread(
        nc,
        param(config, "input.lateral.river.fsandriv");
        sel = inds_riv,
        type = Float,
    )
    fgravriv = ncread(
        nc,
        param(config, "input.lateral.river.fgravriv");
        sel = inds_riv,
        type = Float,
    )
    d50riv =
        ncread(nc, param(config, "input.lateral.river.d50"); sel = inds_riv, type = Float)
    d50engelund = ncread(
        nc,
        param(config, "input.lateral.river.d50engelund");
        sel = inds_riv,
        type = Float,
    )
    cbagnold = ncread(
        nc,
        param(config, "input.lateral.river.cbagnold");
        sel = inds_riv,
        type = Float,
    )
    ebagnold = ncread(
        nc,
        param(config, "input.lateral.river.ebagnold");
        sel = inds_riv,
        type = Float,
    )

    # Initialisation of parameters for Kodatie transport capacity
    ak = zeros(Float, nriv)
    bk = zeros(Float, nriv)
    ck = zeros(Float, nriv)
    dk = zeros(Float, nriv)
    if tcmethodriv == "kodatie"
        for i = 1:nriv
            if d50riv[i] <= 0.05
                ak[i] = 281.4
                bk[i] = 2.622
                ck[i] = 0.182
                dk[i] = 0.0
            elseif d50riv[i] <= 0.25
                ak[i] = 2829.6
                bk[i] = 3.646
                ck[i] = 0.406
                dk[i] = 0.412
            elseif d50riv[i] <= 2.0
                ak[i] = 2123.4
                bk[i] = 3.3
                ck[i] = 0.468
                dk[i] = 0.613
            else
                ak[i] = 431884.8
                bk[i] = 1.0
                ck[i] = 1.0
                dk[i] = 2.0
            end
        end
    end
    # Initialisation of parameters for river erosion
    # Bed and Bank from Shields diagram, Da Silva & Yalin (2017)
    E_ = (2.65 - 1) * 9.81
    E = (E_ .* (d50riv .* 1e-3) .^ 3 ./ 1e-12) .^ 0.33
    TCrbed = @. Float(
        E_ *
        d50riv *
        (0.13 * E^(-0.392) * exp(-0.015 * E^2) + 0.045 * (1 - exp(-0.068 * E))),
    )
    TCrbank = TCrbed
    # kd from Hanson & Simon 2001
    kdbank = @. Float(0.2 * TCrbank^(-0.5) * 1e-6)
    kdbed = @. Float(0.2 * TCrbed^(-0.5) * 1e-6)

    rs = RiverSediment(
        n = nriv,
        Δt = Float(Δt.value),
        # Parameters
        sl = riverslope,
        dl = riverlength,
        width = riverwidth,
        dmclay = dmclay,
        dmsilt = dmsilt,
        dmsand = dmsand,
        dmsagg = dmsagg,
        dmlagg = dmlagg,
        dmgrav = dmgrav,
        fclayriv = fclayriv,
        fsiltriv = fsiltriv,
        fsandriv = fsandriv,
        fgravriv = fgravriv,
        d50 = d50riv,
        d50engelund = d50engelund,
        cbagnold = cbagnold,
        ebagnold = ebagnold,
        ak = ak,
        bk = bk,
        ck = ck,
        dk = dk,
        kdbank = kdbank,
        kdbed = kdbed,
        TCrbank = TCrbank,
        TCrbed = TCrbed,
        rhos = rhos,
        # Forcing
        h_riv = fill(mv, nriv),
        q_riv = fill(mv, nriv),
        # Input from land
        inlandclay = zeros(Float, nriv),
        inlandsilt = zeros(Float, nriv),
        inlandsand = zeros(Float, nriv),
        inlandsagg = zeros(Float, nriv),
        inlandlagg = zeros(Float, nriv),
        inlandsed = zeros(Float, nriv),
        # States
        sedload = zeros(Float, nriv),
        clayload = zeros(Float, nriv),
        siltload = zeros(Float, nriv),
        sandload = zeros(Float, nriv),
        saggload = zeros(Float, nriv),
        laggload = zeros(Float, nriv),
        gravload = zeros(Float, nriv),
        sedstore = zeros(Float, nriv),
        claystore = zeros(Float, nriv),
        siltstore = zeros(Float, nriv),
        sandstore = zeros(Float, nriv),
        saggstore = zeros(Float, nriv),
        laggstore = zeros(Float, nriv),
        gravstore = zeros(Float, nriv),
        outsed = zeros(Float, nriv),
        outclay = zeros(Float, nriv),
        outsilt = zeros(Float, nriv),
        outsand = zeros(Float, nriv),
        outsagg = zeros(Float, nriv),
        outlagg = zeros(Float, nriv),
        outgrav = zeros(Float, nriv),
        # Outputs
        Sedconc = zeros(Float, nriv),
        SSconc = zeros(Float, nriv),
        Bedconc = zeros(Float, nriv),
        maxsed = zeros(Float, nriv),
        erodsed = zeros(Float, nriv),
        erodsedbank = zeros(Float, nriv),
        erodsedbed = zeros(Float, nriv),
        depsed = zeros(Float, nriv),
        insed = zeros(Float, nriv),
        # Reservoir / lake
        wbcover = wbcover,
        wblocs = wblocs,
        wbarea = wbarea,
        wbtrap = wbtrap,
    )

    return rs
end

function update(rs::RiverSediment, network, config)
    @unpack graph, order = network
    tcmethod = get(config.model, "rivtransportmethod", "bagnold")::String

    # River sediment loads are separated into different particle class.
    # Clay, silt and sand can both come from land, resuspension or river channel erosion.
    # Small and large aggregates only come from land erosion or resuspension.
    # Gravel only comes from resuspension or river channel erosion.

    for v in order
        ### Sediment input in the cell (left from previous timestep + from land + from upstream outflux) ###
        upstream_nodes = inneighbors(graph, v)

        inrivclay = 0.0
        inrivsilt = 0.0
        inrivsand = 0.0
        inrivsagg = 0.0
        inrivlagg = 0.0
        inrivgrav = 0.0
        if !isempty(upstream_nodes)
            for i in upstream_nodes
                if rs.outclay[i] >= 0.0 # avoid NaN from upstream non-river cells
                    inrivclay += rs.outclay[i]
                    inrivsilt += rs.outsilt[i]
                    inrivsand += rs.outsand[i]
                    inrivsagg += rs.outsagg[i]
                    inrivlagg += rs.outlagg[i]
                    inrivgrav += rs.outgrav[i]
                end
            end
        end

        inclay = rs.clayload[v] + rs.inlandclay[v] + inrivclay
        insilt = rs.siltload[v] + rs.inlandsilt[v] + inrivsilt
        insand = rs.sandload[v] + rs.inlandsand[v] + inrivsand
        insagg = rs.saggload[v] + rs.inlandsagg[v] + inrivsagg
        inlagg = rs.laggload[v] + rs.inlandlagg[v] + inrivlagg
        ingrav = rs.gravload[v] + inrivgrav

        insed = inclay + insilt + insand + insagg + inlagg + ingrav
        rs.insed[v] = insed
        rs.inlandsed[v] =
            rs.inlandclay[v] +
            rs.inlandsilt[v] +
            rs.inlandsand[v] +
            rs.inlandsagg[v] +
            rs.inlandlagg[v]

        ### Transport capacity of the flow ###
        # Hydraulic radius of the river [m] (rectangular channel)
        hydrad = rs.h_riv[v] * rs.width[v] / (rs.width[v] + 2 * rs.h_riv[v])

        # Engelund and Hansen transport formula
        if tcmethod == "engelund"
            vmean =
                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
            vshear = (9.81 * hydrad * rs.sl[v])^0.5
            # Concentration by weight
            cw = ifelse(
                hydrad > 0.0,
                (
                    rs.rhos[v] / 1000 * 0.05 * vmean * vshear^3 /
                    ((rs.rhos[v] / 1000 - 1)^2 * 9.81^2 * rs.d50engelund[v] * hydrad)
                ),
                0.0,
            )
            cw = min(1.0, cw)
            # Transport capacity [tons/m3]
            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
        elseif tcmethod == "bagnold"
            maxsed =
                rs.cbagnold[v] * (rs.q_riv[v] / (rs.h_riv[v] * rs.width[v]))^rs.ebagnold[v]
        elseif tcmethod == "kodatie"
            vmean =
                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
            maxsed = rs.ak[v] * vmean^rs.bk[v] * rs.h_riv[v]^rs.ck[v] * rs.sl[v]^rs.dk[v]
            # Transport capacity [tons/m3]
            maxsed =
                ifelse(rs.q_riv[v] > 0.0, maxsed * rs.width[v] / (rs.q_riv[v] * rs.Δt), 0.0)
        elseif tcmethod == "yang"
            ws = 411 * rs.d50[v]^2 / 3600
            vshear = (9.81 * hydrad * rs.sl[v])^0.5
            var1 = vshear * rs.d50[v] / 1000 / (1.16 * 1e-6)
            var2 = ws * rs.d50[v] / 1000 / (1.16 * 1e-6)
            vcr = min(
                0.0,
                ifelse(var1 >= 70.0, 2.05 * ws, ws * (2.5 / (log10(var1) - 0.06) + 0.66)),
            )
            # Sand equation
            if (rs.width[v] * rs.h_riv[v]) >= vcr && rs.d50[v] < 2.0
                logcppm = (
                    5.435 - 0.286 * log10(var2) - 0.457 * log10(vshear / ws) + 1.799 -
                    0.409 * log10(var2) -
                    0.314 *
                    log10(vshear / ws) *
                    log10(
                        (rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]) - vcr) * rs.sl[v] / ws,
                    )
                )
                # Gravel equation
            elseif (rs.width[v] * rs.h_riv[v]) >= vcr && rs.d50[v] < 2.0
                logcppm = (
                    6.681 - 0.633 * log10(var2) - 4.816 * log10(vshear / ws) + 2.784 -
                    0.305 * log10(var2) -
                    0.282 *
                    log10(vshear / ws) *
                    log10(
                        (rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]) - vcr) * rs.sl[v] / ws,
                    )
                )
            else
                logcppm = 0.0
            end
            # Sediment concentration by weight
            cw = 10^logcppm * 1e-6
            # Transport capacity [ton/m3]
            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
        elseif tcmethod == "molinas"
            ws = 411 * rs.d50[v]^2 / 3600
            vmean =
                ifelse(rs.h_riv[v] > 0.0, rs.q_riv[v] / (rs.width[v] * rs.h_riv[v]), 0.0)
            if rs.h_riv[v] > 0.0
                psi = (
                    vmean^3 / (
                        (rs.rhos[v] / 1000 - 1) *
                        9.81 *
                        rs.h_riv[v] *
                        ws *
                        log10(1000 * rs.h_riv[v] / rs.d50[v])^2
                    )
                )
            else
                psi = 0.0
            end
            # Concentration by weight
            cw = 1430 * (0.86 + psi^0.5) * psi^1.5 / (0.016 + psi) * 1e-6
            # Transport capacity [ton/m3]
            maxsed = max(cw / (cw + (1 - cw) * rs.rhos[v] / 1000) * rs.rhos[v] / 1000, 0.0)
        end

        # 1285 g/L: boundary between streamflow and debris flow (Costa, 1988)
        maxsed = min(maxsed, 1.285)
        # Transport capacity [ton]
        maxsed = maxsed * (rs.h_riv[v] * rs.width[v] * rs.dl[v] + rs.q_riv[v] * rs.Δt)
        rs.maxsed[v] = maxsed

        ### River erosion ###
        # Erosion only if the load is below the transport capacity of the flow.
        sedex = max(maxsed - insed, 0.0)
        # No erosion in lake and reservoir cells
        if rs.wbcover[v] > 0.0
            sedex = 0.0
        end
        # Bed and bank are eroded only if the previously deposited material is not enough
        rs.sedstore[v] =
            rs.claystore[v] +
            rs.siltstore[v] +
            rs.sandstore[v] +
            rs.saggstore[v] +
            rs.laggstore[v] +
            rs.gravstore[v]
        if sedex > 0.0 && sedex > rs.sedstore[v]
            # Effective sediment needed fom river bed and bank erosion [ton]
            effsedex = sedex - rs.sedstore[v]

            # Repartition of the effective shear stress between the bank and the bed from Knight et al. 1984 [%]
            SFbank = ifelse(
                rs.h_riv[v] > 0.0,
                exp(-3.23 * log10(rs.width[v] / rs.h_riv[v] + 3) + 6.146),
                0.0,
            )
            # Effective shear stress on river bed and banks [N/m2]
            TEffbank = ifelse(
                rs.h_riv[v] > 0.0,
                1000 * 9.81 * hydrad * rs.sl[v] * SFbank / 100 *
                (1 + rs.width[v] / (2 * rs.h_riv[v])),
                0.0,
            )
            TEffbed =
                1000 *
                9.81 *
                hydrad *
                rs.sl[v] *
                (1 - SFbank / 100) *
                (1 + 2 * rs.h_riv[v] / rs.width[v])
            # Potential erosion rates of the bed and bank [t/cell/timestep]
            #(assuming only one bank is eroding)
            Tex = max(TEffbank - rs.TCrbank[v], 0.0)
            # 1.4 is bank default bulk density
            ERbank = max(0.0, rs.kdbank[v] * Tex * rs.dl[v] * rs.h_riv[v] * 1.4 * rs.Δt)
            # 1.5 is bed default bulk density
            ERbed = max(
                0.0,
                rs.kdbed[v] *
                (TEffbed - rs.TCrbed[v]) *
                rs.dl[v] *
                rs.width[v] *
                1.5 *
                rs.Δt,
            )
            # Relative potential erosion rates of the bed and the bank [-]
            RTEbank = ifelse(ERbank + ERbed > 0.0, ERbank / (ERbank + ERbed), 0.0)
            RTEbed = 1.0 - RTEbank

            # Bank erosion (difference between effective and potential erosion) [ton]
            sedbank = ifelse(effsedex * RTEbank <= ERbank, effsedex * RTEbank, ERbank)
            claybank = rs.fclayriv[v] * sedbank
            siltbank = rs.fsiltriv[v] * sedbank
            sandbank = rs.fsandriv[v] * sedbank
            gravbank = rs.fgravriv[v] * sedbank

            # Bed erosion [ton]
            sedbed = ifelse(effsedex * RTEbed <= ERbed, effsedex * RTEbed, ERbed)
            claybed = rs.fclayriv[v] * sedbed
            siltbed = rs.fsiltriv[v] * sedbed
            sandbed = rs.fsandriv[v] * sedbed
            gravbed = rs.fgravriv[v] * sedbed
        else
            sedbank = 0.0
            claybank = 0.0
            siltbank = 0.0
            sandbank = 0.0
            gravbank = 0.0

            sedbed = 0.0
            claybed = 0.0
            siltbed = 0.0
            sandbed = 0.0
            gravbed = 0.0
        end

        # Erosion/degradation of the previously deposited sediment (from clay to gravel) [ton]
        if sedex > 0.0
            degstoreclay = ifelse(rs.claystore[v] >= sedex, sedex, rs.claystore[v])
            rs.claystore[v] = rs.claystore[v] - degstoreclay
            #Update amount of sediment that need to be degraded
            sedex = sedex - degstoreclay
            degstoresilt = ifelse(rs.siltstore[v] >= sedex, sedex, rs.siltstore[v])
            rs.siltstore[v] = rs.siltstore[v] - degstoresilt
            sedex = max(0.0, sedex - degstoresilt)
            degstoresagg = ifelse(rs.saggstore[v] >= sedex, sedex, rs.saggstore[v])
            rs.saggstore[v] = rs.saggstore[v] - degstoresagg
            sedex = max(0.0, sedex - degstoresagg)
            degstoresand = ifelse(rs.sandstore[v] >= sedex, sedex, rs.sandstore[v])
            rs.sandstore[v] = rs.sandstore[v] - degstoresand
            sedex = max(0.0, sedex - degstoresand)
            degstorelagg = ifelse(rs.laggstore[v] >= sedex, sedex, rs.laggstore[v])
            rs.laggstore[v] = rs.laggstore[v] - degstorelagg
            sedex = max(0.0, sedex - degstorelagg)
            degstoregrav = ifelse(rs.gravstore[v] >= sedex, sedex, rs.gravstore[v])
            rs.gravstore[v] = rs.gravstore[v] - degstoregrav
            sedex = max(0.0, sedex - degstoregrav)
            degstoresed =
                degstoreclay +
                degstoresilt +
                degstoresagg +
                degstoresand +
                degstorelagg +
                degstoregrav
        else
            degstoreclay = 0.0
            degstoresilt = 0.0
            degstoresagg = 0.0
            degstoresand = 0.0
            degstorelagg = 0.0
            degstoregrav = 0.0
            degstoresed = 0.0
        end

        # Sum all erosion sources per particle class
        erodsed = sedbank + sedbed + degstoresed
        erodclay = claybank + claybed + degstoreclay
        erodsilt = siltbank + siltbed + degstoresilt
        erodsand = sandbank + sandbed + degstoresand
        erodsagg = degstoresagg
        erodlagg = degstorelagg
        erodgrav = gravbank + gravbed + degstoregrav

        rs.erodsed[v] = erodsed
        rs.erodsedbank[v] = sedbank
        rs.erodsedbed[v] = sedbed

        ### Deposition / settling ###
        # Fractions of deposited particles in river cells from the Einstein formula [-]
        # Particle fall velocity [m/s] from Stokes
        xs = ifelse(rs.q_riv[v] > 0.0, 1.055 * rs.dl[v] / (rs.q_riv[v] / rs.width[v]), 0.0)
        xclay = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmclay[v] / 1000)^2 / 3600)))
        xsilt = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsilt[v] / 1000)^2 / 3600)))
        xsand = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsand[v] / 1000)^2 / 3600)))
        xsagg = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmsagg[v] / 1000)^2 / 3600)))
        xlagg = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmlagg[v] / 1000)^2 / 3600)))
        xgrav = min(1.0, 1.0 - 1.0 / exp(xs * (411.0 * (rs.dmgrav[v] / 1000)^2 / 3600)))

        # Sediment deposited in the channel [ton]
        # From natural settling with Einstein formula (density controlled).
        settclay = xclay * (inclay + erodclay)
        settsilt = xsilt * (insilt + erodsilt)
        settsand = xsand * (insand + erodsand)
        settsagg = xsagg * (insagg + erodsagg)
        settlagg = xlagg * (inlagg + erodlagg)
        settgrav = xgrav * (ingrav + erodgrav)

        # Sediment deposited in the channel (from gravel to clay) [ton]
        # From transport capacity exceedance (insed > maxsed)
        insedex = max(insed - maxsed, 0.0)
        if insedex > 0.0
            depgrav = ifelse(ingrav >= insedex, insedex, ingrav)
            insedex = max(insedex - depgrav, 0.0)
            deplagg = ifelse(inlagg >= insedex, insedex, inlagg)
            insedex = max(insedex - deplagg, 0.0)
            depsand = ifelse(insand >= insedex, insedex, insand)
            insedex = max(insedex - depsand, 0.0)
            depsagg = ifelse(insagg >= insedex, insedex, insagg)
            insedex = max(insedex - depsagg, 0.0)
            depsilt = ifelse(insilt >= insedex, insedex, insilt)
            insedex = max(insedex - depsilt, 0.0)
            depclay = ifelse(inclay >= insedex, insedex, inclay)
            insedex = max(insedex - depclay, 0.0)
        else
            depclay = settclay
            depsilt = settsilt
            depsand = settsand
            depsagg = settsagg
            deplagg = settlagg
            depgrav = settgrav
        end

        # No deposition in regular lake and reservoir cells, only at the outlet
        # Deposition in lake/reservoir from Camp 1945
        # Extra trapping of large particles for dams
        if rs.wbcover[v] > 0.0 && rs.wblocs[v] > 0.0
            # Compute deposition
            vcres = rs.q_riv[v] / rs.wbarea[v]
            DCres = 411 / 3600 / vcres
            depclay = (inclay + erodclay) * min(1.0, (DCres * (rs.dmclay[v] / 1000)^2))
            depsilt = (insilt + erodsilt) * min(1.0, (DCres * (rs.dmsilt[v] / 1000)^2))
            depsand = (insand + erodsand) * min(1.0, (DCres * (rs.dmsand[v] / 1000)^2))
            depsagg = (insagg + erodsagg) * min(1.0, (DCres * (rs.dmsagg[v] / 1000)^2))
            deplagg = (inlagg + erodlagg) * min(1.0, (DCres * (rs.dmlagg[v] / 1000)^2))
            depgrav = (ingrav + erodgrav) * min(1.0, (DCres * (rs.dmgrav[v] / 1000)^2))
            # Trapping of large particles
            # Use the rouse number for sagg (suspension or bedload)
            depsand = max(depsand, rs.wbtrap[v] * (insand + erodsand))
            deplagg = max(deplagg, rs.wbtrap[v] * (inlagg + erodlagg))
            depgrav = max(depgrav, rs.wbtrap[v] * (ingrav + erodgrav))
            # threshold diameter between suspended load and mixed load using Rouse number
            dsuspf =
                1e3 * (1.2 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
            depsagg = ifelse(
                rs.dmsagg[v] > dsuspf,
                depsagg,
                max(depsagg, rs.wbtrap[v] * (insagg + erodsagg)),
            )
        elseif rs.wbcover[v] > 0.0
            depsed = 0.0
            depclay = 0.0
            depsilt = 0.0
            depsand = 0.0
            depsagg = 0.0
            deplagg = 0.0
            depgrav = 0.0
        end

        depsed = depclay + depsilt + depsand + depsagg + deplagg + depgrav
        rs.depsed[v] = depsed

        # Update the river deposited sediment storage
        rs.sedstore[v] = rs.sedstore[v] + depsed
        rs.claystore[v] = rs.claystore[v] + depclay
        rs.siltstore[v] = rs.siltstore[v] + depsilt
        rs.sandstore[v] = rs.sandstore[v] + depsand
        rs.saggstore[v] = rs.saggstore[v] + depsagg
        rs.laggstore[v] = rs.laggstore[v] + deplagg
        rs.gravstore[v] = rs.gravstore[v] + depgrav

        ### Ouput loads ###
        # Sediment transported out of the cell during the timestep [ton]
        # 0 in case all sediment are deposited in the cell
        # Reduce the fraction so that there is still some sediment staying in the river cell
        fwaterout = min(rs.q_riv[v] * rs.Δt / (rs.h_riv[v] * rs.width[v] * rs.dl[v]), 1.0)
        rs.outsed[v] = fwaterout * (insed + erodsed - depsed)
        rs.outclay[v] = fwaterout * (inclay + erodclay - depclay)
        rs.outsilt[v] = fwaterout * (insilt + erodsilt - depsilt)
        rs.outsand[v] = fwaterout * (insand + erodsand - depsand)
        rs.outsagg[v] = fwaterout * (insagg + erodsagg - depsagg)
        rs.outlagg[v] = fwaterout * (inlagg + erodlagg - deplagg)
        rs.outgrav[v] = fwaterout * (ingrav + erodgrav - depgrav)

        ### Mass balance ###
        # Sediment left in the cell [ton]
        rs.sedload[v] = insed + erodsed - depsed - rs.outsed[v]
        rs.clayload[v] = inclay + erodclay - depclay - rs.outclay[v]
        rs.siltload[v] = insilt + erodsilt - depsilt - rs.outsilt[v]
        rs.sandload[v] = insand + erodsand - depsand - rs.outsand[v]
        rs.saggload[v] = insagg + erodsagg - depsagg - rs.outsagg[v]
        rs.gravload[v] = ingrav + erodgrav - depgrav - rs.outgrav[v]

        ### Concentrations and suspended sediments ###
        # Conversion from load [ton] to concentration for rivers [mg/L]
        toconc = ifelse(rs.q_riv[v] > 0.0, 1e6 / (rs.q_riv[v] * rs.Δt), 0.0)
        rs.Sedconc[v] = rs.outsed[v] * toconc

        # Differentiation of bed and suspended load using Rouse number for suspension
        # threshold diameter between bed load and mixed load using Rouse number
        dbedf = 1e3 * (2.5 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
        # threshold diameter between suspended load and mixed load using Rouse number
        dsuspf = 1e3 * (1.2 * 3600 * 0.41 / 411 * (9.81 * rs.h_riv[v] * rs.sl[v])^0.5)^0.5
        # Rouse with diameter
        SSclay = ifelse(
            rs.dmclay[v] <= dsuspf,
            rs.outclay[v],
            ifelse(rs.dmclay[v] <= dbedf, rs.outclay[v] / 2, 0.0),
        )
        SSsilt = ifelse(
            rs.dmsilt[v] <= dsuspf,
            rs.outsilt[v],
            ifelse(rs.dmsilt[v] <= dbedf, rs.outsilt[v] / 2, 0.0),
        )
        SSsagg = ifelse(
            rs.dmsagg[v] <= dsuspf,
            rs.outsagg[v],
            ifelse(rs.dmsagg[v] <= dbedf, rs.outsagg[v] / 2, 0.0),
        )
        SSsand = ifelse(
            rs.dmsand[v] <= dsuspf,
            rs.outsand[v],
            ifelse(rs.dmsand[v] <= dbedf, rs.outsand[v] / 2, 0.0),
        )
        SSlagg = ifelse(
            rs.dmlagg[v] <= dsuspf,
            rs.outlagg[v],
            ifelse(rs.dmlagg[v] <= dbedf, rs.outlagg[v] / 2, 0.0),
        )
        SSgrav = ifelse(
            rs.dmgrav[v] <= dsuspf,
            rs.outgrav[v],
            ifelse(rs.dmgrav[v] <= dbedf, rs.outgrav[v] / 2, 0.0),
        )

        SS = SSclay + SSsilt + SSsagg + SSsand + SSlagg + SSgrav
        Bed = rs.outsed[v] - SS

        rs.SSconc[v] = SS * toconc
        rs.Bedconc[v] = Bed * toconc

    end

end


module TestEnsembleMethods

include(joinpath("..", "fixture_learners.jl"))
using .FixtureLearners
nfcp = NumericFeatureClassification()

using FactCheck

importall Orchestra.Transformers.EnsembleMethods
import Orchestra.Transformers.DecisionTreeWrapper: fit!, transform!
import Orchestra.Transformers.DecisionTreeWrapper: PrunedTree
import Orchestra.Transformers.DecisionTreeWrapper: RandomForest
import Orchestra.Transformers.DecisionTreeWrapper: DecisionStumpAdaboost

facts("Ensemble learners") do
  context("VoteEnsemble predicts according to majority") do
    always_a_options = { :label => "a" }
    always_b_options = { :label => "b" }
    learner = VoteEnsemble({
      :learners => [
        AlwaysSameLabelLearner(always_a_options),
        AlwaysSameLabelLearner(always_a_options),
        AlwaysSameLabelLearner(always_b_options)
      ]
    })
    fit!(learner, nfcp.train_instances, nfcp.train_labels)
    predictions = transform!(learner, nfcp.test_instances)
    expected_predictions = fill("a", size(nfcp.test_instances, 1))

    @fact predictions => expected_predictions
  end

  context("StackEnsemble predicts with combined learners") do
    # Fix random seed, due to stochasticity in stacker.
    srand(2)

    always_a_options = { :label => "a" }
    learner = StackEnsemble({
      :learners => [
        AlwaysSameLabelLearner(always_a_options),
        AlwaysSameLabelLearner(always_a_options),
        PerfectScoreLearner()
      ],
      :keep_original_features => true
    })
    fit!(learner, nfcp.train_instances, nfcp.train_labels)
    predictions = transform!(learner, nfcp.test_instances)
    unexpected_predictions = fill("a", size(nfcp.test_instances, 1))

    @fact predictions => not(unexpected_predictions)
  end

  context("BestLearner picks the best learner") do
    always_a_options = { :label => "a" }
    always_b_options = { :label => "b" }
    learner = BestLearner({
      :learners => [
        AlwaysSameLabelLearner(always_a_options),
        PerfectScoreLearner(),
        AlwaysSameLabelLearner(always_b_options)
      ]
    })
    fit!(learner, nfcp.train_instances, nfcp.train_labels)

    @fact learner.model[:best_learner_index] => 2
  end

  context("BestLearner conducts grid search") do
    learner = BestLearner({
      :learners => [PrunedTree(), DecisionStumpAdaboost(), RandomForest()],
      :learner_options_grid => [
        {
          :impl_options => {
            :purity_threshold => [0.5, 1.0]
          }
        },
        Dict(),
        {
          :impl_options => {
            :num_trees => [5, 10, 20],
            :partial_sampling => [0.5, 0.7]
          }
        }
      ]
    })
    fit!(learner, nfcp.train_instances, nfcp.train_labels)

    @fact length(learner.model[:learners]) => 8
  end
end

end # module


struct JacobianEntry{E} <: Entry
    index::Int
    e::E
end
# struct Jacobian{G <: Gradient, I}
#     inner::I
#     ds::Vector{Tuple{Int,G}}
# end
# @inline function (J::Jacobian{G,I})(y,x,p=nothing) where {G,I}
#     inner(J)(y,x)
#     @simd for i in eachindex(J.ds)
#         (j,d) = J.ds[i]
#         d((j,y),x,p)
#     end
# end
# @inline function (J::Jacobian{G,Nothing})(y,x,p=nothing) where G
#     @simd for i in eachindex(J.ds)
#         (j,d) = J.ds[i]
#         d((j,y),x,p)
#     end
# end
@inline (e::JacobianEntry{E})(y,x,p=nothing) where {E} = e.e((index(e),y),x,p)
# Jacobian(s::Sink{Field}) = Jacobian(inner(s))
Jacobian(f::Field1{G,I},indexer = nothing) where {G,I} = Field1(Jacobian(inner(f),indexer),[JacobianEntry(index(ie),Gradient(ie.e,(index(ie),indexer))) for ie in f.es])
Jacobian(f::Field1{G,Nothing},indexer = nothing) where G = Field1(nothing,[JacobianEntry(index(ie),Gradient(ie.e,(index(ie),indexer))) for ie in f.es])


__precompile__()

module Utils

    # export @define, @reexport
    #
    # include("utils/macro_utils.jl")

    export istriustrict, istrilstrict
    export simd_scale!, simd_copy!, simd_copy_scale!,
           simd_copy_xy_first!, simd_copy_yx_first!,
           simd_copy_yx_first_last!,
           simd_xpy!, simd_axpy!, simd_wxpy!, simd_waxpy!, simd_mult!

    include("utils/matrix_utils.jl")

end


include("../src/qspice.jl")
include("../src/netlist.jl")

include("gates.jl")


@testset "Escape.run(::HLAModel, data::AbstractHLAData)" begin
    ds = Escape.HLADataset("Test")
    m = Escape.HLAModel{4}()
    data = ds.data[1]

    @test @suppress Escape.run(m, data, iter = 5, warmup = 5, chains = 1) isa
        Escape.HLAModelResult
end

@testset "Escape.reduce_size!(::StanInterface.Stanfit)" begin
    ds = Escape.HLADataset("Test")
    m = Escape.HLAModel{4}()
    data = ds.data[1]

    res = @suppress Escape.run(m, data, iter = 5, warmup = 5, chains = 1,
        keep_all_parameters = true)
    previous_size = Base.summarysize(res.sf)

    @test "aux1_tau.1" in keys(res.sf.result[1])
    Escape.reduce_size!(res.sf)

    new_size = Base.summarysize(res.sf)
    @test !("aux1_tau.1" in keys(res.sf.result[1]))
    @test new_size < previous_size
end

@testset "Escape.run_model(::HLAModel, ::AbstractHLAData)" begin
    alignment_file_2_digits = joinpath(@__DIR__, "data",
        "test_large_annotated_2_digits.fasta")
    tree_file = joinpath(@__DIR__, "data", "phylogeny.tree")

    data = @suppress HLAData(
        alignment_file = alignment_file_2_digits,
        tree_file = tree_file
    )

    @test @suppress Escape.run_model(
        Escape.HLAModel{3}(),
        data, iter = 5, warmup = 5, chains = 1
    ) isa Escape.HLAModelResult

    save_file = tempname() * ".jld2"
    @test !isfile(save_file)
    @suppress Escape.run_model(
        Escape.HLAModel{3}(),
        data, iter = 5, warmup = 5, chains = 1, save_file = save_file
    )
    @test isfile(save_file)
    rm(save_file)
end

@testset "Escape.run_model(::AbstractHLAData)" begin
    alignment_file_2_digits = joinpath(@__DIR__, "data",
    "test_large_annotated_2_digits.fasta")
    tree_file = joinpath(@__DIR__, "data", "phylogeny.tree")

    data = @suppress HLAData(
        alignment_file = alignment_file_2_digits,
        tree_file = tree_file
    )

    @test @suppress Escape.run_model(
        data, iter = 5, warmup = 5, chains = 1
    ) isa Escape.HLAModelResult

    save_file = tempname() * ".jld2"
    @test !isfile(save_file)
    @suppress Escape.run_model(
        data, iter = 5, warmup = 5, chains = 1, save_file = save_file
    )
    @test isfile(save_file)
    rm(save_file)
end


function backward{N}(backend::GPUBackend, state::SplitLayerState{N}, inputs::Vector{Blob}, diffs::Vector{Blob})
  if !isa(diffs[1], NullBlob)
    diff = diffs[1]
    len = length(diff)
    one = convert(eltype(diff), 1)
    for i = 2:N
      CuBLAS.axpy(backend.cublas_ctx, len, one, state.blobs_diff[i].ptr, 1, diff.ptr, 1)
    end
  end
end


function init_constraint!(
    com::CS.CoM,
    constraint::XNorConstraint,
    fct,
    set::XNorSet;
)
    !init_lhs_and_rhs!(com, constraint, fct, set) && return false

    if constraint.complement_lhs !== nothing
        init_constraint!(com, constraint.complement_lhs, constraint.complement_lhs.fct, constraint.complement_lhs.set)
    end
    if constraint.complement_rhs !== nothing
        init_constraint!(com, constraint.complement_rhs, constraint.complement_rhs.fct, constraint.complement_rhs.set)
    end
    return true
end

"""
    function is_constraint_violated(
        com::CoM,
        constraint::BoolConstraint,
        fct,
        set::XNorSet,
    )

Check whether one side is solved and other is violated
"""
function is_constraint_violated(
    com::CoM,
    constraint::BoolConstraint,
    fct,
    set::XNorSet,
)
    lhs_solved = is_constraint_solved(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    rhs_solved = is_constraint_solved(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    # neither of them is solved => it's not violated yet
    if !lhs_solved && !rhs_solved
        return false
    end
    if lhs_solved && rhs_solved
        return false
    end

    lhs_violated = is_lhs_constraint_violated(com, constraint)
    rhs_violated = is_rhs_constraint_violated(com, constraint)
    if lhs_solved && rhs_violated
        return true
    end
    if rhs_solved && lhs_violated
        return true
    end

    return false
end

"""
    still_feasible(com::CoM, constraint::XNorConstraint, fct, set::XNorSet, vidx::Int, value::Int)

Return whether the constraint can be still fulfilled when setting a variable with index `vidx` to `value`.
**Attention:** This assumes that it isn't violated before.
"""
function still_feasible(
    com::CoM,
    constraint::XNorConstraint,
    fct,
    set::XNorSet,
    vidx::Int,
    value::Int,
)
    lhs_violated = is_constraint_violated(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    lhs_feasible = !lhs_violated
    lhs_solved = false
    if lhs_feasible
        lhs_solved = is_constraint_solved_when_fixed(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set, vidx, value)
    end
    if lhs_feasible && !lhs_solved
        lhs_indices = constraint.lhs.indices
        for i in 1:length(lhs_indices)
            if lhs_indices[i] == vidx
                lhs_feasible = still_feasible(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set, vidx, value)
                break
            end
        end
    end
    rhs_violated = is_constraint_violated(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    rhs_feasible = !rhs_violated
    rhs_solved = false
    if rhs_feasible
        rhs_solved = is_constraint_solved_when_fixed(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set, vidx, value)
    end
    # both must be violated or both must be solved in the end
    if !lhs_feasible && rhs_solved
        return false
    end
    if rhs_feasible && !rhs_solved
        rhs_indices = constraint.rhs.indices
        for i in 1:length(rhs_indices)
            if rhs_indices[i] == vidx
                rhs_feasible = still_feasible(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set, vidx, value)
                break
            end
        end
    end

    if !rhs_feasible && lhs_solved
        return false
    end
    return true
end

"""
    prune_constraint!(com::CS.CoM, constraint::XNorConstraint, fct, set::XNorSet; logs = true)

Reduce the number of possibilities given the `XNorConstraint` by pruning both parts
Return whether still feasible
"""
function prune_constraint!(
    com::CS.CoM,
    constraint::XNorConstraint,
    fct,
    set::XNorSet;
    logs = true,
)
    lhs_violated = is_constraint_violated(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    rhs_violated = is_constraint_violated(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    if lhs_violated && rhs_violated
        return true
    end
    # check if one is already solved
    lhs_solved = is_constraint_solved(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set)
    rhs_solved = is_constraint_solved(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set)
    if lhs_solved && rhs_solved
        return true
    end

    # if one is solved => prune the other
    if lhs_solved
        activate_rhs!(com, constraint)
        return prune_constraint!(com, constraint.rhs, constraint.rhs.fct, constraint.rhs.set; logs=logs)
    end
    if rhs_solved
        activate_lhs!(com, constraint)
        return prune_constraint!(com, constraint.lhs, constraint.lhs.fct, constraint.lhs.set; logs=logs)
    end

    # if one is violated complement prune the other
    # Todo implement for activated complement constraints
    if lhs_violated && constraint.complement_rhs !== nothing && !implements_activate(typeof(constraint.complement_rhs), typeof(constraint.complement_rhs.fct), typeof(constraint.complement_rhs.set))
        return prune_constraint!(com, constraint.complement_rhs, constraint.complement_rhs.fct, constraint.complement_rhs.set; logs=logs)
    end
    if rhs_violated && constraint.complement_lhs !== nothing && !implements_activate(typeof(constraint.complement_lhs), typeof(constraint.complement_lhs.fct), typeof(constraint.complement_lhs.set))
        return prune_constraint!(com, constraint.complement_lhs, constraint.complement_lhs.fct, constraint.complement_lhs.set; logs=logs)
    end

    return true
end


__precompile__(true)

module MLSuite

using Statistics, Random, Distributed, Printf, LinearAlgebra, DelimitedFiles, Mmap
using DataStructures, Parameters, ProgressMeter, PyCall, BSON, Requires
using PyCallUtils, PandasLite, HDF5Utils, SVMLightWriter, Iconv
using ScikitLearnBase: BaseEstimator, BaseClassifier, BaseRegressor
import ScikitLearnBase: fit!, predict, predict_proba, is_classifier
import MLSuiteBase: gridparams, reset!, istree, isrnn, is_ranker, support_multiclass, modelhash, signone

const JULIA = joinpath(Sys.BINDIR, Base.julia_exename())
const VW = joinpath(DEPOT_PATH[1], "mlsuite/bin/vw")
const VW_VARINFO = joinpath(DEPOT_PATH[1], "mlsuite/bin/vw-varinfo")
const SVM_RANK_LEARN = joinpath(DEPOT_PATH[1], "mlsuite/bin/svm_rank_learn")
const SVM_RANK_CLASSIFY = joinpath(DEPOT_PATH[1], "mlsuite/bin/svm_rank_classify")
const RANKLIB = joinpath(DEPOT_PATH[1], "mlsuite/bin/ranklib")
const LIGHTGBM = joinpath(DEPOT_PATH[1], "mlsuite/bin/lightgbm")

export accuracy_score, r2_score
accuracy_score(a...; ka...) = pyimport("sklearn.metrics").accuracy_score(a...; ka...)
r2_score(a...; ka...) = pyimport("sklearn.metrics").r2_score(a...; ka...)

include("util.jl")

include("base/vw.jl")
include("base/h2o.jl")
include("base/dai.jl")
include("base/gbm.jl")
include("base/lgbm.jl")

include("class/scikit.jl")

include("regress/grid.jl")
include("regress/scikit.jl")

include("rank/svm.jl")
include("rank/ranklib.jl")
include("rank/tf.jl")

function __init__()
    boostlib = joinpath(DEPOT_PATH[1], "mlsuite/boost/lib")
    ENV["LD_LIBRARY_PATH"] = boostlib * ":" * get(ENV, "LD_LIBRARY_PATH", "")
    @require Elemental="902c3f28-d1ec-5e7e-8399-a24c3845ee38" include("regress/el.jl")
end

end


using Turing, MCMCChains, Test

@testset "DataFrame constructor tests" begin
    @model gdemo(x) = begin
        m ~ Normal(1, 0.01)
        s ~ Normal(5, 0.01)
    end

    model = gdemo([1.5, 2.0])
    sampler = HMC(1000, 0.01, 5)

    chns = [sample(model, sampler, save_state=true) for i in 1:4]
    chn = chainscat(chns...)

    df = DataFrame(chn, showall=true)
    @test size(df) == (4000, 8)

    df1 = DataFrame(chn, [:parameters])
    @test size(df1) == (4000, 2)

    df2 = DataFrame(chn, [:internals, :parameters])
    @test size(df2) == (4000, 8)

    df3 = DataFrame(chn[:s])
    @test size(df3) == (4000, 1)

    df4 = DataFrame(chn[:lp])
    @test size(df4) == (4000, 1)

    df5 = DataFrame(chn, [:parameters], append_chains=false)
    @test size(df5) == (4, )
    @test size(df5[1]) == (1000, 2)

    df6 = DataFrame(chn, [:parameters, :internals], append_chains=false)
    @test size(df6) == (4, )
    @test size(df6[1]) == (1000, 8)

    df7 = DataFrame(chn, [:parameters, :internals],
        remove_missing_union=false)
    @test size(df7) == (4000, 8)

    df8 = DataFrame(chn, [:parameters, :internals],
        remove_missing_union=false, append_chains=false)
    @test size(df8) == (4, )
    @test size(df8[1]) == (1000, 8)
end


using Plots
pyplot()
using Statistics
using ProgressMeter

function visualise_ballbott_compare(X,X2,u,u2,t, filename::String, des_traj; options::Dict = Dict())
    # Find plotting limits
    xmin = minimum(X[1, :]; dims = 1)[1] .- .2
    xmax = maximum(X[1, :]; dims = 1)[1] .+ .2
    xlim = (xmin, xmax)
    x_mean = sum(xlim)/2
    x_range = abs.(x_mean-xmin)*2

    ymin = minimum(X[2, :])#-2
    ymax = maximum(X[2, :])#+2
    ylim = (ymin, ymax)
    y_mean = sum(ylim) / 2
    y_range = abs.(y_mean-ymin)*2

    #Idea scale equally by having same range:
    range_max = max.(x_range, y_range)
    dx = 0.5*(range_max / x_range)*x_range
    dy = 0.5*(range_max / y_range)*y_range

    xlim = (x_mean - dx, x_mean + dx)
    ylim = (y_mean - dy, y_mean + dy)

    fps::Int = get(options, :fps, 25)

    # Create animation
    plotsize::Tuple{Int, Int} = get(options, :size, (1280, 720))
    #plotsize::Tuple{Int, Int} = get(options, :size, (640, 360))
    anim = Animation()

    # Create timestamps of the visualisation and interpolate states
    t_plot = t[1]:(1/fps):t[end]
    #X_plot = Array{Float64}(size(X, 1), length(t_plot))

    prog = Progress(length(t))
    n_step = 12
    for i = 1:n_step:length(t)

        fig = plot(;size = plotsize, layout = grid(3, 1, heights=[0.8,0.1,0.1]))
        plot_ballbot_3d_compare!(fig[1], X[:,i], X2[:,i], xlim, ylim, des_traj, i, X, X2)
        plot!(fig[2], t[1:i], u[1,1:i], xlim = (0, t[end]), ylim = (-4, 4), lab = "u1_hard" )
        plot!(fig[2], t[1:i], u[2,1:i], lab = "u2_hard")
        plot!(fig[2], t[1:i], u[3,1:i], lab = "u3_hard")
        plot!(fig[3], t[1:i], u2[1,1:i], xlim = (0, t[end]), ylim = (-4, 4), lab = "u1_soft")
        plot!(fig[3], t[1:i], u2[2,1:i], lab = "u2_soft")
        plot!(fig[3], t[1:i], u2[2,1:i], lab = "u3_soft")
        frame(anim)

        ProgressMeter.update!(prog,i)
    end

    run(`ffmpeg -v 0 -framerate $(fps) -loop 0 -i $(anim.dir)/%06d.png -y $(filename)`)
    return nothing
end

function plot_ballbot_3d_compare!(fig, xt, xt2, x_lim, y_lim, des_traj, i, X1, X2)
    #Trajectory to do:
    plot!(fig,des_traj[1:i,1],des_traj[1:i,2],des_traj[1:i,1] .*0, lab ="desired position", c=:black)

    #Trajectory done:
    plot!(fig,X1[1,1:i],X1[2,1:i],X1[1,1:i].*0, lab="position PPO", c=:blue)

    #Trajectory done:
    plot!(fig,X2[1,1:i],X2[2,1:i],X2[1,1:i].*0, lab="position PPO-soft", c=:green)

    p_close = [2.0, 0.0]

    d_sim  = (p_close[1] - xt[1] )^2 + (p_close[2] - xt[2] )^2
    d_real = (p_close[1] - xt2[1])^2 + (p_close[2] - xt2[2])^2

    scale_it = 0.25
    range = x_lim[2] - x_lim[1]
    z_lim = (0, max(3*scale_it,range))

    dr1 = 0.4
    dr2 = 0.4

    if d_sim > d_real
        """
        1st Ball
        """
        plot!(fig, xlim = x_lim, ylim = y_lim, zlim = z_lim, ratio = :equal) # TODO zlim problematesch

        dx = xt[1]
        dy = xt[2]

        ϕ_y = xt[4]#*(pi/180)
        ϕ_x = -xt[5]#*(pi/180)
        ϕ_z = xt[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = 0.0:dr1:2pi+ dr1
        v = 0.0:dr2:pi+ dr2

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        surface!(x,y,z, c=:blue);
        #surf(x,y,z)

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo .+= sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo .+= dx + sin(ϕ_y)*2*radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo .+= sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo .+= dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo .-= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo .-= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo .+= dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, c=:blue)
        end

        """
        2nd ball
        """

        dx = xt2[1]
        dy = xt2[2]

        ϕ_y = xt2[4]#*(pi/180)
        ϕ_x = -xt2[5]#*(pi/180)
        ϕ_z = xt2[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = 0.0:dr1:2pi+ dr1
        v = 0.0:dr2:pi+ dr2

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        surface!(x,y,z, c=:green);
        #surf(x,y,z)

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo .+= sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo .+= dx + sin(ϕ_y)*2*radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo .+= sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo .+= dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo .-= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo .-= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo .+= dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, α = 0.5, c=:green)
        end

    else

        """
        2nd ball
        """

        plot!(fig, xlim = x_lim, ylim = y_lim, zlim = z_lim, ratio = :equal) # TODO zlim problematesch

        dx = xt2[1]
        dy = xt2[2]

        ϕ_y = xt2[4]#*(pi/180)
        ϕ_x = -xt2[5]#*(pi/180)
        ϕ_z = xt2[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = 0.0:dr1:2pi+ dr1
        v = 0.0:dr2:pi+ dr2

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        #plot!(x,y,z);
        surface!(x,y,z, color = "blue");

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo .+= sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo .+= dx + sin(ϕ_y) * 2 * radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo .+= sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo .+= dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo .-= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo .-= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo .+= dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, α = 0.5, c=:green)
        end

        """
        1st ball
        """

        dx = xt[1]
        dy = xt[2]

        ϕ_y = xt[4]#*(pi/180)
        ϕ_x = -xt[5]#*(pi/180)
        ϕ_z = xt[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = 0.0:dr1:2pi+ dr1
        v = 0.0:dr2:pi+ dr2

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        surface!(x,y,z, c=:blue);
        #surf(x,y,z)

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo .+= sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo .+= dx + sin(ϕ_y)*2*radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo .+= sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo .+= dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo .-= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo .-= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo .+= dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, c=:blue, α = 0.5)
        end

    end
    gui()

end

using JLD2
using FileIO

filename = "results_ballvsball.jld2"
a = load(filename)
data = a["(X, X2, u, u2, t, des_traj_real)"]
X  = data[1].parent'
X2 = data[2].parent'
u = data[3].parent'
u2 = data[4].parent'
t = data[5]
des_traj_real = data[6]

newfile = ("ballvsball.mp4")
visualise_ballbott_compare(X,X2,u,u2,t, newfile, des_traj_real; options= Dict())

using Plots

fig = plot(0,0)
pyplot()

scale_it= 0.25
xt2 = randn(12)
xt = randn(12)
dr1 = 0.2
dr2 = 0.2

dx = xt2[1]
dy = xt2[2]

ϕ_y = xt2[4]#*(pi/180)
ϕ_x = -xt2[5]#*(pi/180)
ϕ_z = xt2[6]#*(pi/180)

radi = scale_it * 0.5

fac_x = scale_it * 0.4
fac_y = scale_it * 0.4
fac_z = scale_it * 1

dz   = 2 * radi + fac_z

#Draw the Circle:
u = 0.0:dr1:2pi+ dr1
v = 0.0:dr2:pi+ dr2

lu = length(u);
lv = length(v);

x = zeros(lu,lv);
y = zeros(lu,lv);
z = zeros(lu,lv);

colors = zeros(lu,lv,3) ;

colortube = 0.1:0.1:1.0;

for uu=1:lu
    for vv=1:lv
        x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
        y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
        z[uu,vv]= radi*cos(v[vv]) + radi
    end
end

#plot!(x,y,z);
surface!(x,y,z, surfacecolor = z .*1 .+ 3, α= 0.5);
gui()

#Draw the Quader on Top
x = -1:1:1
y = -1:1:1

lx = length(x)
ly = length(y)

X = zeros(lx,ly)
Y = zeros(lx,ly)
Z = zeros(lx,ly)

for i_x = 1:lx
    for i_y = 1:ly
        X[i_x,i_y] = x[i_x]
        Y[i_x,i_y] = y[i_y]
        Z[i_x,i_y] = 1
    end
end

arr_X = [X,Y,Z,-Y,-Z,-Y]
arr_Y = [Y,Z,X,-Z,-X,-Z]
arr_Z = [Z,X,Y,-X,-Y,-Z]

for i_side = 1:6
    #x-Koordinate Calculation
    x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
    x_koo .+= sin(ϕ_y)* arr_Z[i_side] * fac_z
    x_koo .+= dx + sin(ϕ_y) * 2 * radi
    #y-Koordinate Calculation
    y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
    y_koo .+= sin(ϕ_x)*arr_Z[i_side]* fac_z
    y_koo .+= dy + sin(ϕ_x)*2*radi
    #z-Koordinate Calculation
    z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
    z_koo .-= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
    z_koo .-= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
    z_koo .+= dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

    surface!(fig,x_koo,y_koo,z_koo, surfacecolor = z .* 1 .+3, α= 0.5)
end

gui()

        dx = xt[1]
        dy = xt[2]

        ϕ_y = xt[4]#*(pi/180)
        ϕ_x = -xt[5]#*(pi/180)
        ϕ_z = xt[6]#*(pi/180)

        radi = scale_it * 0.5

        fac_x = scale_it * 0.4
        fac_y = scale_it * 0.4
        fac_z = scale_it * 1

        dz   = 2 * radi + fac_z

        #Draw the Circle:
        u = 0.0:dr1:2pi+ dr1
        v = 0.0:dr2:pi+ dr2

        lu = length(u);
        lv = length(v);

        x = zeros(lu,lv);
        y = zeros(lu,lv);
        z = zeros(lu,lv);

        colors = zeros(lu,lv,3) ;

        colortube = 0.1:0.1:1.0;

        for uu=1:lu
            for vv=1:lv
                x[uu,vv]= radi*cos(u[uu])*sin(v[vv]) + dx
                y[uu,vv]= radi*sin(u[uu])*sin(v[vv]) + dy
                z[uu,vv]= radi*cos(v[vv]) + radi
            end
        end

        surface!(x,y,z, surfacecolor = z .* 2 .+1, α= 0.5)
        #surf(x,y,z)

        #Draw the Quader on Top
        x = -1:1:1
        y = -1:1:1

        lx = length(x)
        ly = length(y)

        X = zeros(lx,ly)
        Y = zeros(lx,ly)
        Z = zeros(lx,ly)

        for i_x = 1:lx
            for i_y = 1:ly
                X[i_x,i_y] = x[i_x]
                Y[i_x,i_y] = y[i_y]
                Z[i_x,i_y] = 1
            end
        end

        arr_X = [X,Y,Z,-Y,-Z,-Y]
        arr_Y = [Y,Z,X,-Z,-X,-Z]
        arr_Z = [Z,X,Y,-X,-Y,-Z]

        for i_side = 1:6
            #x-Koordinate Calculation
            x_koo  = cos(ϕ_y)*(cos(ϕ_z)*arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            x_koo .+= sin(ϕ_y)* arr_Z[i_side] * fac_z
            x_koo .+= dx + sin(ϕ_y)*2*radi
            #y-Koordinate Calculation
            y_koo  = cos(ϕ_x)*(cos(ϕ_z)*arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            y_koo .+= sin(ϕ_x)*arr_Z[i_side]* fac_z
            y_koo .+= dy + sin(ϕ_x)*2*radi
            #z-Koordinate Calculation
            z_koo  =                            (cos(ϕ_y)*cos(ϕ_x)*arr_Z[i_side]) * fac_z
            z_koo .-= sin(ϕ_y)*(cos(ϕ_z) * arr_X[i_side] + sin(ϕ_z)*arr_Y[i_side]) * fac_x
            z_koo .-= sin(ϕ_x)*(cos(ϕ_z) * arr_Y[i_side] - sin(ϕ_z)*arr_X[i_side]) * fac_y
            z_koo .+= dz - (2-cos(ϕ_x)-cos(ϕ_y))*2*radi

            surface!(fig,x_koo,y_koo,z_koo, surfacecolor = z .* 2 .+1, α= 0.5)
        end

    gui()


###############################################################################
# Stochiometry for a given reaction is a vector of pairs mapping species id to
# stochiometric coefficient.
###############################################################################

@inline @fastmath function evalrxrate(speciesvec::AbstractVector{T}, rxidx::S,
                              majump::MassActionJump{U,V,W})::R where {T,S,R,U <: AbstractVector{R},V,W}
    val = one(T)
    @inbounds for specstoch in majump.reactant_stoch[rxidx]
        specpop = speciesvec[specstoch[1]]
        val    *= specpop
        @inbounds for k = 2:specstoch[2]
            specpop -= one(specpop)
            val     *= specpop
        end
    end

    @inbounds return val * majump.scaled_rates[rxidx]
end

@inline @fastmath function executerx!(speciesvec::AbstractVector{T}, rxidx::S,
                                      majump::MassActionJump{U,V,W}) where {T,S,U,V,W}
    @inbounds net_stoch = majump.net_stoch[rxidx]
    @inbounds for specstoch in net_stoch
        speciesvec[specstoch[1]] += specstoch[2]
    end
    nothing
end

function scalerates!(unscaled_rates::AbstractVector{U}, stochmat::AbstractVector{V}) where {U,S,T,W <: Pair{S,T}, V <: AbstractVector{W}}
    @inbounds for i in eachindex(unscaled_rates)
        coef = one(T)
        @inbounds for specstoch in stochmat[i]
            coef *= factorial(specstoch[2])
        end
        unscaled_rates[i] /= coef
    end
    nothing
end

function scalerate(unscaled_rate::U, stochmat::AbstractVector{Pair{S,T}}) where {U <: Number, S, T}
    coef = one(T)
    @inbounds for specstoch in stochmat
        coef *= factorial(specstoch[2])
    end
    unscaled_rate /= coef
end

###############################################################################
# dependency graph when MassActionJump uses pairs to represent (species,stoich)
###############################################################################

# map from species to Set of reactions depending on that species
function spec_to_dep_rxs_map(numspec, ma_jumps::MassActionJump)

    numrxs = get_num_majumps(ma_jumps)

    # map from a species to reactions that depend on it
    spec_to_dep_rxs = [Set{Int}() for n = 1:numspec]
    for rx in 1:numrxs
        for (spec,stoch) in ma_jumps.reactant_stoch[rx]
            push!(spec_to_dep_rxs[spec], rx)
        end
    end

    spec_to_dep_rxs
end

# dependency graph is a map from a reaction to a vector of reactions
# that should depend on species it changes
function make_dependency_graph(numspec, ma_jumps::MassActionJump)

    numrxs          = get_num_majumps(ma_jumps)
    spec_to_dep_rxs = spec_to_dep_rxs_map(numspec, ma_jumps)

    # create map from rx to reactions depending on it
    dep_sets = [SortedSet{Int}() for n = 1:numrxs]
    for rx in 1:numrxs

        # rx changes spec, hence rxs depending on spec depend on rx
        for (spec,stoch) in ma_jumps.net_stoch[rx]
            for dependent_rx in spec_to_dep_rxs[spec]
                push!(dep_sets[rx], dependent_rx)
            end
        end
    end

    # convert to Vectors of Vectors
    dep_graph = Vector{Vector{Int}}(undef, numrxs)
    for rx = 1:numrxs
        dep_graph[rx] = [dep for dep in dep_sets[rx]]
    end

    dep_graph
end

# update dependency graph to make sure jumps depend on themselves
function add_self_dependencies!(dg)
    for (i,jump_deps) in enumerate(dg)
        if !any(y->isequal(y,i), jump_deps)
            push!(jump_deps, i)
            sort!(jump_deps)
        end
    end
end


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

# Quadratic decision function #
# ========================== #
struct QuadraticDecisionObjectiveBridge{T} <: MOIB.Objective.AbstractBridge
    decision_function::QuadraticDecisionFunction{T}
end

function MOIB.Objective.bridge_objective(::Type{QuadraticDecisionObjectiveBridge{T}}, model::MOI.ModelLike,
                                         f::QuadraticDecisionFunction{T}) where T
    # All decisions have been mapped to the variable part terms
    # at this point.
    F = MOI.ScalarQuadraticFunction{T}
    # Set the bridged objective
    MOI.set(model, MOI.ObjectiveFunction{F}(),
            MOI.ScalarQuadraticFunction(
                f.variable_part.affine_terms,
                f.variable_part.quadratic_terms,
                zero(T)))
    # Save decision function to allow modifications
    return QuadraticDecisionObjectiveBridge{T}(f)
end

function MOIB.Objective.supports_objective_function(
    ::Type{<:QuadraticDecisionObjectiveBridge}, ::Type{<:QuadraticDecisionFunction})
    return true
end
MOIB.added_constrained_variable_types(::Type{<:QuadraticDecisionObjectiveBridge}) = Tuple{DataType}[]
function MOIB.added_constraint_types(::Type{<:QuadraticDecisionObjectiveBridge})
    return Tuple{DataType, DataType}[]
end
function MOIB.Objective.concrete_bridge_type(::Type{<:QuadraticDecisionObjectiveBridge{T}},
                                             ::Type{QuadraticDecisionFunction{T}}) where T
    return QuadraticDecisionObjectiveBridge{T}
end
function MOIB.set_objective_function_type(::Type{QuadraticDecisionObjectiveBridge{T}}) where T
    return MOI.ScalarQuadraticFunction{T}
end

function MOI.get(::QuadraticDecisionObjectiveBridge, ::MOI.NumberOfVariables)
    return 0
end

function MOI.get(::QuadraticDecisionObjectiveBridge, ::MOI.ListOfVariableIndices)
    return MOI.VariableIndex[]
end

function MOI.delete(::MOI.ModelLike, ::QuadraticDecisionObjectiveBridge)
    # Nothing to delete
    return nothing
end

function MOI.set(::MOI.ModelLike, ::MOI.ObjectiveSense,
                 ::QuadraticDecisionObjectiveBridge, ::MOI.OptimizationSense)
    # Nothing to handle if sense changes
    return nothing
end

function MOI.get(model::MOI.ModelLike,
                 attr::MOIB.ObjectiveFunctionValue{QuadraticDecisionFunction{T}},
                 bridge::QuadraticDecisionObjectiveBridge{T}) where T
    f = bridge.decision_function
    G = MOI.ScalarQuadraticFunction{T}
    obj_val = MOI.get(model, MOIB.ObjectiveFunctionValue{G}(attr.result_index))
    # Calculate and add constant
    constant = f.variable_part.constant +
        f.decision_part.constant
    return obj_val + f.variable_part.constant
end

function MOI.get(model::MOI.ModelLike,
                 attr::MOI.ObjectiveFunction{QuadraticDecisionFunction{T}},
                 bridge::QuadraticDecisionObjectiveBridge{T}) where T
    return bridge.decision_function
end

function MOI.modify(model::MOI.ModelLike, bridge::QuadraticDecisionObjectiveBridge{T}, change::MOI.ScalarConstantChange) where T
    f = bridge.decision_function
    f = f.linear_quadratic_terms
    # Modify constant of variable part
    lq.variable_part.constant = change.new_constant
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::QuadraticDecisionObjectiveBridge{T}, change::MOI.ScalarCoefficientChange) where T
    f = bridge.decision_function
    # Update variable part
    modify_coefficient!(f.variable_part.affine_terms, change.variable, change.new_coefficient)
    # Modify variable part of objective as usual
    F = MOI.ScalarQuadraticFunction{T}
    MOI.modify(model, MOI.ObjectiveFunction{F}(), change)
    return nothing
end