text:
# binomial.jl
# Binomial tests
#
# Copyright (C) 2013   Simon Kornblith
#
# 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.

export BinomialTest, SignTest

## BINOMIAL TEST

immutable BinomialTest <: HypothesisTest
    p::Float64
    x::Int
    n::Int

    BinomialTest(x::Real, n::Real, p::Real=0.5) = new(p, x, n)
end

BinomialTest(x::AbstractVector{Bool}, p=0.5) =
    BinomialTest(sum(x), length(x), p)

"""
```julia
testname(::HypothesisTest)
```

Returns the string value. E.g. "Binomial test", "Sign Test"
"""
testname(::BinomialTest) = "Binomial test"
population_param_of_interest(x::BinomialTest) = ("Probability of success", x.p, x.x/x.n) # parameter of interest: name, value under h0, point estimate

function show_params(io::IO, x::BinomialTest, ident="")
    println(io, ident, "number of observations: $(x.n)")
    println(io, ident, "number of successes:    $(x.x)")
end

pvalue(x::BinomialTest; tail=:both) = pvalue(Binomial(x.n, x.p), x.x; tail=tail)

# Confidence interval

"""
```julia
function ci(x::HypothesisTest, alpha::Float64=0.05; tail=:both, method=:clopper_pearson)
```
Compute a confidence interval with coverage 1-alpha for multinomial proportions using one of the following methods. Possible values for method are:

Sison, Glaz intervals :sison_glaz (default):
Bootstrap intervals :bootstrap :
Quesenberry, Hurst intervals :quesenberry_hurst :
Gold intervals :gold (Asymptotic simultaneous intervals):
"""
function ci(x::BinomialTest, alpha::Float64=0.05; tail=:both, method=:clopper_pearson)
    check_alpha(alpha)

    if tail == :left
        (0.0, ci(x, alpha*2, method=method)[2])
    elseif tail == :right
        (ci(x, alpha*2, method=method)[1], 1.0)
    elseif tail == :both
        if method == :clopper_pearson
            ci_clopper_pearson(x, alpha)
        elseif method == :wald
            ci_wald(x, alpha)
        elseif method == :wilson
            ci_wilson(x, alpha)
        elseif method == :jeffrey
            ci_jeffrey(x, alpha)
        elseif method == :agresti_coull
            ci_agresti_coull(x, alpha)
        else
            throw(ArgumentError("method=$(method) is not implemented yet"))
        end
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end
end

# Clopper-Pearson interval (confidence interval by inversion)
function ci_clopper_pearson(x::BinomialTest, alpha::Float64=0.05)
    (quantile(Beta(x.x, x.n - x.x + 1), alpha/2), quantile(Beta(x.x + 1, x.n - x.x), 1-alpha/2))
end

# Wald interval / normal approximation interval
function ci_wald(x::BinomialTest, alpha::Float64=0.05)
    μ = x.x / x.n
    σ = sqrt(μ*(1-μ)/x.n)
    (quantile(Normal(μ, σ), alpha/2), quantile(Normal(μ, σ), 1-alpha/2))
end

# Jeffreys interval
function ci_jeffrey(x::BinomialTest, alpha::Float64=0.05)
    (quantile(Beta(x.x + 1/2, x.n - x.x + 1/2), alpha/2), quantile(Beta(x.x + 1/2, x.n - x.x + 1/2), 1-alpha/2))
end

# Agresti-Coull interval
function ci_agresti_coull(x::BinomialTest, alpha::Float64=0.05)
    q = quantile(Normal(), 1-alpha/2)
    n = x.n + q^2
    μ = (x.x + q^2/2)/n
    σ = sqrt(μ*(1-μ)/n)
    (μ-q*σ, μ+q*σ)
end

# Wilson score interval
function ci_wilson(x::BinomialTest, alpha::Float64=0.05)
    q = quantile(Normal(), 1-alpha/2)
    p = x.x / x.n
    denominator = 1 + q^2/x.n
    μ = p + q^2/(2*x.n)
    μ /= denominator
    σ = sqrt(p*(1-p)/x.n + q^2/(4x.n^2))
    σ /= denominator
    (μ-q*σ, μ+q*σ)
end

## SIGN TEST

immutable SignTest <: HypothesisTest
    median::Float64
    x::Int
    n::Int
    data
end

SignTest{T<:Real}(x::AbstractVector{T}, median::Real=0) =
    SignTest(median, sum(x .> median), sum(x .!= median), sort(x))
SignTest{T<:Real, S<:Real}(x::AbstractVector{T}, y::AbstractVector{S}) = SignTest(x - y, 0.0)

testname(::SignTest) = "Sign Test"
population_param_of_interest(x::SignTest) = ("Median", x.median, median(x.data)) # parameter of interest: name, value under h0, point estimate

function show_params(io::IO, x::SignTest, ident="")
    text1 = "number of observations:"
    text2 = "observations larger than $(x.median): "
    maxlen = length(text2)

    println(io, ident, text1, repeat(" ", maxlen-length(text1)), x.n)
    println(io, ident, text2, x.x)
end

"""
```julia
pvalue(x::HypothesisTest; tail=:both)
```

Compute the p-value for a given significance test.

If tail is :both (default), then the p-value for the two-sided test is returned. If tail is :left or :right, then a one-sided test is performed.
"""
pvalue(x::SignTest; tail=:both) = pvalue(Binomial(x.n, 0.5), x.x; tail=tail)

"""
```julia
function ci(x::HypothesisTest, alpha::Float64=0.05; tail=:both)
```

Compute a confidence interval C with coverage 1-alpha.

If tail is :both (default), then a two-sided confidence interval is returned. If tail is :left or :right, then a one-sided confidence interval is returned
"""
@compat function ci(x::SignTest, alpha::Float64=0.05; tail=:both)
    check_alpha(alpha)

    if tail == :left
        q = Int(quantile(Binomial(x.n, 0.5), alpha))
        (x.data[q+1], x.median)
    elseif tail == :right
        q = Int(quantile(Binomial(x.n, 0.5), alpha))
        (x.median, x.data[end-q])
    elseif tail == :both
        q = Int(quantile(Binomial(x.n, 0.5), alpha/2))
        (x.data[q+1], x.data[end-q])
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end
end


text:
"""
    Revise.WatchList

A struct for holding files that live inside a directory.
Some platforms (OSX) have trouble watching too many files. So we
watch parent directories, and keep track of which files in them
should be tracked.

Fields:
- `timestamp`: mtime of last update
- `trackedfiles`: Set of filenames
"""
mutable struct WatchList
    timestamp::Float64         # unix time of last revision
    trackedfiles::Set{String}
end

const DefMapValue = Union{RelocatableExpr,Tuple{Vector{Any},Int}}  # ([sigt1,...], lineoffset)

"""
    DefMap

Maps `def`=>nothing or `def=>([sigt1,...], lineoffset)`, where:

- `def` is an expression
- the value is `nothing` if `def` does not define a method
- if it does define a method, `sigt1...` are the signature-types and `lineoffset` is the
  difference between the line number when the method was compiled and the current state
  of the source file.

See the documentation page [How Revise works](@ref) for more information.
"""
const DefMap = OrderedDict{RelocatableExpr,Union{DefMapValue,Nothing}}

"""
    SigtMap

Maps `sigt=>def`, where `sigt` is the signature-type of a method and `def` the expression
defining the method.

See the documentation page [How Revise works](@ref) for more information.
"""
const SigtMap = IdDict{Any,RelocatableExpr}   # sigt=>def

"""
    FMMaps

`source=>sigtypes` and `sigtypes=>source` mappings for a particular file/module combination.
See the documentation page [How Revise works](@ref) for more information.
"""
struct FMMaps
    defmap::DefMap
    sigtmap::SigtMap
end
FMMaps() = FMMaps(DefMap(), SigtMap())

Base.isempty(fmm::FMMaps) = isempty(fmm.defmap)

function Base.show(io::IO, fmm::FMMaps)
    limit = get(io, :limit, true)
    if limit
        print(io, "FMMaps(<$(length(fmm.defmap)) expressions>, <$(length(fmm.sigtmap)) signatures>)")
    else
        println(io, "FMMaps with the following expressions:")
        for def in keys(fmm.defmap)
            print(io, "  ")
            Base.show_unquoted(io, def, 2)
            print(io, '\n')
        end
    end
end

"""
    FileModules

For a particular source file, the corresponding `FileModules` is an
`OrderedDict(mod1=>fmm1, mod2=>fmm2)`,
mapping the collection of modules "active" in the file (the parent module and any
submodules it defines) to their corresponding [`FMMaps`](@ref).

The first key is guaranteed to be the module into which this file was `include`d.

To create a `FileModules` from a source file, see [`parse_source`](@ref).
"""
const FileModules = OrderedDict{Module,FMMaps}

"""
    fm = FileModules(mod::Module)

Initialize an empty `FileModules` for a file that is `include`d into `mod`.
"""
FileModules(mod::Module) = FileModules(mod=>FMMaps())

Base.isempty(fm::FileModules) = length(fm) == 1 && isempty(first(values(fm)))

"""
    FileInfo(fm::FileModules, cachefile="")

Structure to hold the per-module expressions found when parsing a
single file.
`fm` holds the [`FileModules`](@ref) for the file.

Optionally, a `FileInfo` can also record the path to a cache file holding the original source code.
This is applicable only for precompiled modules and `Base`.
(This cache file is distinct from the original source file that might be edited by the
developer, and it will always hold the state
of the code when the package was precompiled or Julia's `Base` was built.)
When a cache is available, `fm` will be empty until the file gets edited:
the original source code gets parsed only when a revision needs to be made.

Source cache files greatly reduce the overhead of using Revise.
"""
struct FileInfo
    fm::FileModules
    cachefile::String
end
FileInfo(fm::FileModules) = FileInfo(fm, "")

"""
    FileInfo(mod::Module, cachefile="")

Initialze an empty FileInfo for a file that is `include`d into `mod`.
"""
FileInfo(mod::Module, cachefile::AbstractString="") = FileInfo(FileModules(mod), cachefile)

FileInfo(fm::FileModules, fi::FileInfo) = FileInfo(fm, fi.cachefile)

ext: jl


text:
using OrdinaryDiffEq, Flux, Random
using DiffEqFlux
using BSON: @save, bson

function model(ode_data::AbstractArray, batch_size::Int=4, maxiters::Int=150, lr::Float64=0.01; log_dir="")
  # u0, ode_data = prepare_data("basis_steps", bin_size)
  u0 = ode_data[:, 1]
  data_dim = length(u0)
  datasize = length(ode_data[1,:])
  tspan = (0.0, Float64(datasize) - 1)
  tsteps = range(tspan[1], tspan[2], length = datasize)

  train_t = tsteps
  train_y = ode_data

  function neural_ode(t, data_dim; saveat = t)
      # println(data_dim)
      f = FastChain(FastDense(data_dim, 64, swish),
                    FastDense(64, 32, swish),
                    FastDense(32, data_dim))

      node = NeuralODE(f, (minimum(t), maximum(t)), Tsit5(),
                       saveat = saveat, abstol = 1e-9,
                       reltol = 1e-9)
  end

  function train_one_round(node, θ, y, opt, maxiters,
                           y0 = y[:, 1]; kwargs...)

      predict(θ) = Array(node(y0, θ))
      loss(θ) = begin
          pred = predict(θ)
          Flux.mse(pred, y)
      end

      θ = θ == nothing ? node.p : θ
      res = DiffEqFlux.sciml_train(
          loss, θ, opt,
          maxiters = maxiters;
          kwargs...
      )
      return res.minimizer
  end

  function train(θ = nothing, maxiters = maxiters, lr = lr)
      log_results(θs, losses) =
          (θ, loss) -> begin
              push!(θs, copy(θ))
              push!(losses, loss)
              false
          end

      θs, losses = [], []
      num_obs = batch_size:batch_size:length(train_t)
      for k in num_obs
        println(k)
        node = neural_ode(train_t[1:k], data_dim)
	println(log_dir)
        if log_dir != ""
          bson(joinpath(log_dir, "model$(k).bson"), nn_model=node, thetas=θs)
        end
        θ = train_one_round(
            node, θ, train_y[:, 1:k],
            ADAMW(lr), maxiters;
            cb = log_results(θs, losses)
        )
      end
      # last iteration
      node = neural_ode(train_t, data_dim)
      θ = train_one_round(
            node, θ, train_y,
            ADAMW(lr), maxiters, train_y[:, 1];
            cb = log_results(θs, losses)
      )
      θs, losses, node
  end

  # Random.seed!(1)
  θs, losses, node = train();
  # node = neural_ode(train_t, input_data_dim, output_data_dim)
  idx = findmin(losses)[2]
  θ = θs[idx]
  y0 = train_y[:, 1]
  resol = Array(node(y0, θ))

  # v = "Steps"
  pl = plot(tsteps, ode_data[1,:], color = :red, label = "Data: CGM", xlabel = "t", title = "CGM, Steps, Bolus, Meals")
  plot!(tsteps, ode_data[2,:], color = :blue, label = "Data: Steps")
  plot!(tsteps, ode_data[3,:], color = :green, label = "Data: Bolus")
  plot!(tsteps, ode_data[4,:], color = :purple, label = "Data: Meals")
  plot!(tsteps,resol[1,:], alpha=0.5, color = :red, label = "CGM")
  plot!(tsteps,resol[2,:], alpha=0.5, color = :blue, label = "Steps")
  plot!(tsteps,resol[3,:], alpha=0.5, color = :red, label = "Bolus")
  plot!(tsteps,resol[4,:], alpha=0.5, color = :blue, label = "Meals")

  println(log_dir)
  println("idx")
  println(idx)
  println("theta")
  println(θ)
  println("y0")
  println(y0)
  println("resol")
  println(resol)
  if log_dir != ""
    println(joinpath(log_dir, "model.bson"))
    bson(joinpath(log_dir, "4_model.bson"), nn_model=node, theta=θ, loss=losses[idx], best_pred=resol, y0=y0)
  end

  θ, losses[idx], resol, pl

end

function model_exo(non_exo_data::AbstractArray, exo_data::AbstractArray, batch_size::Int64=4, maxiters::Int=150, lr=0.01; log_dir="")
  # u0, ode_data = prepare_data("basis_steps", bin_size)
  u0 = non_exo_data[:, 1]
  datasize = length(non_exo_data[1,:])
  tspan = (0.0, 1.0) # (0.0, Float64(datasize) - 1)
  tsteps = range(tspan[1], tspan[2], length = datasize)

  @show maxiters

  input_data_dim = size(non_exo_data, 1) + size(exo_data, 1)
  output_data_dim = size(non_exo_data, 1)

  println("input_data_dim")
  println(input_data_dim)

  println("output_data_dim")
  println(output_data_dim)

  train_t = tsteps
  train_y = non_exo_data
  println("datasize")
  println(datasize)
  println("train_y")
  println(size(train_y))
  println(train_y)
  println("train_t")
  println(train_t)

  function neural_ode(t, input_data_dim, output_data_dim; saveat = t)
      # println(data_dim)
      f = FastChain(FastDense(input_data_dim, 64, swish),
                    FastDense(64, 32, swish),
                    FastDense(32, output_data_dim))
      p_model = initial_params(f)
      function f_exo(u, p, t)
        f(vcat(u, exo_data[:, Int(round(t*(datasize - 1)/(tspan[2] - tspan[1]) + (1 - (datasize - 1)*tspan[1]/(tspan[2] - tspan[1]))))]), p)
      end
      prob = ODEProblem(f_exo, u0, tspan, nothing)
      prob, p_model
  end

  function train_one_round(prob, p_model, θ, y, opt, maxiters,
                            y0 = y[:, 1]; kwargs...)

      # println("SHOW HERE 2")
      # @show size(y)
      function predict(θ)
        _prob = remake(prob, p=θ)
        Array(solve(_prob, Tsit5(), saveat=tsteps, abstol = 1e-9, reltol = 1e-9))
      end
      # predict(θ) = Array(node(y0, θ))
      function loss(θ)
        pred = predict(θ)
        # println("SHOW HERE 3")
        # @show size(pred)
        Flux.mse(pred[:, 1:size(y)[2]], y)
      end

      @show maxiters
      θ = θ == nothing ? p_model : θ
      res = DiffEqFlux.sciml_train(
          loss, θ, opt,
          maxiters=maxiters;
          kwargs...
      )
      return res.minimizer
  end

  function train(θ = nothing, maxiters = maxiters, lr = lr)
      log_results(θs, losses) =
          (θ, loss) -> begin
              push!(θs, copy(θ))
              push!(losses, loss)
              false
          end

      θs, losses = [], []
      num_obs = batch_size:batch_size:length(train_t)
      for k in num_obs
        println(k)
        prob, p_model = neural_ode(train_t[1:k], input_data_dim, output_data_dim)
        # println("LOOK HERE")
        # @show size(train_y[:, 1:k])
        if log_dir != ""
          bson(joinpath(log_dir, "model$(k).bson"), nn_model=prob, thetas=θs)
        end
        @show train_y[:, 1:k]
        θ = train_one_round(
            prob, p_model, θ, train_y[:, 1:k],
            ADAMW(lr), maxiters, train_y[:, 1];
            cb = log_results(θs, losses)
        )
        @show length(losses)
      end
      # last iteration
      prob, p_model = neural_ode(train_t, input_data_dim, output_data_dim)
      θ = train_one_round(
            prob, p_model, θ, train_y,
            ADAMW(lr), maxiters, train_y[:, 1];
            cb = log_results(θs, losses)
      )
      @show length(losses)
      θs, losses, prob
  end

  # Random.seed!(1)
  θs, losses, prob = train();
  # node = neural_ode(train_t, input_data_dim, output_data_dim)
  # idx = findmin(losses)[2]
  # @show length(losses)
  # @show idx

  @show length(losses)
  @show findmin(losses)

  @show findmin(losses[(end - maxiters):end])
  idx = findmin(losses[(end - maxiters):end])[2] + length(losses) - maxiters - 1
  θ = θs[idx]
  y0 = train_y[:, 1]
  _prob = remake(prob, p=θ)
  resol = Array(solve(_prob, Tsit5(), saveat=tsteps, abstol = 1e-9, reltol = 1e-9))

  println(log_dir)
  println("idx")
  println(idx)
  println("theta")
  println(θ)
  println("y0")
  println(y0)
  println("resol")
  println(resol)

  pl = plot(tsteps, non_exo_data[1,:], color = :red, label = "Data: CGM", xlabel = "t", title = "CGM, Steps, Bolus, Meals")
  plot!(tsteps, non_exo_data[2,:], color = :blue, label = "Data: Steps")
  plot!(tsteps, exo_data[1,:], color = :green, label = "Data: Bolus")
  plot!(tsteps, exo_data[2,:], color = :purple, label = "Data: Meals")
  plot!(tsteps,resol[1,:], alpha=0.5, color = :red, label = "CGM")
  plot!(tsteps,resol[2,:], alpha=0.5, color = :blue, label = "Steps")
  pl

  num = string(now()) # rand(1:100)
  println(log_dir)
  if log_dir != ""
    println(joinpath(log_dir, "1_model$(num).bson"))
    bson(joinpath(log_dir, "1_model$(num).bson"), nn_model=prob, theta=θ, loss=losses[idx], best_pred=resol, y0=y0)
    savefig(pl, joinpath(log_dir, "1_image$(num).png"))
  end

  θ, losses[idx], resol
end

# for bin_size in [20, 15, 10, 5, 1]
#   for batch_size in [10, 8]
#     for maxiters in [150]
#       for lr in [0.01]
#         for i in 1:5
#           println("images/steps_bin_$(bin_size)_batch_$(batch_size)_maxIters_$(maxiters)_lr_$(lr)_i_$(i)")
#           pl = model(bin_size, batch_size, maxiters, lr)
#           savefig(pl, "images/bolus_and_steps_images/bolus_steps_bin_$(bin_size)_batch_$(batch_size)_maxIters_$(maxiters)_lr_$(lr)_i_$(i).png")
#         end
#       end
#     end
#   end
# end

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
abstract type AbstractDCLineFormulation <: AbstractBranchFormulation end

struct HVDCLossless <: AbstractDCLineFormulation end

struct HVDCDispatch <: AbstractDCLineFormulation end

struct VoltageSourceDC <: AbstractDCLineFormulation end

#################################### Branch Variables ##################################################

function flow_variables!(canonical::Canonical,
                        system_formulation::Type{S},
                        devices::IS.FlattenIteratorWrapper{B}) where {B<:PSY.DCBranch,
                                                                        S<:PM.AbstractPowerModel}

    return

end

function flow_variables!(canonical::Canonical,
                        system_formulation::Type{StandardPTDFModel},
                        devices::IS.FlattenIteratorWrapper{B}) where {B<:PSY.DCBranch}

    time_steps = model_time_steps(canonical)
    var_name = Symbol("Fp_$(B)")
    canonical.variables[var_name] = PSI._container_spec(canonical.JuMPmodel,
                                                  (PSY.get_name(d) for d in devices),
                                                   time_steps)

    for d in devices
        bus_fr = PSY.get_number(PSY.get_arc(d).from)
        bus_to = PSY.get_number(PSY.get_arc(d).to)
        for t in time_steps
            canonical.variables[var_name][PSY.get_name(d), t] = JuMP.@variable(canonical.JuMPmodel,
                                                                base_name="$(bus_fr), $(bus_to)_{$(PSY.get_name(d)), $(t)}")
            _add_to_expression!(canonical.expressions[:nodal_balance_active],
                                PSY.get_number(PSY.get_arc(d).from),
                                t,
                                canonical.variables[var_name][PSY.get_name(d), t],
                                -1.0)
            _add_to_expression!(canonical.expressions[:nodal_balance_active],
                                PSY.get_number(PSY.get_arc(d).to),
                                t,
                                canonical.variables[var_name][PSY.get_name(d), t],
                                1.0)
        end
    end

    return

end

#################################### Flow Variable Bounds ##################################################

#################################### Rate Limits Constraints ##################################################

function branch_rate_constraint!(canonical::Canonical,
                                devices::IS.FlattenIteratorWrapper{B},
                                device_formulation::Type{D},
                                system_formulation::Type{S}) where {B<:PSY.DCBranch,
                                                                    D<:AbstractDCLineFormulation,
                                                                    S<:PM.AbstractDCPModel}

    var_name = Symbol("Fp_$(B)")
    con_name = Symbol("RateLimit_$(B)")
    time_steps = model_time_steps(canonical)
    canonical.constraints[con_name] = JuMPConstraintArray(undef, (PSY.get_name(d) for d in devices), time_steps)

    for t in time_steps, d in devices
        min_rate = max(PSY.get_activepowerlimits_from(d).min, PSY.get_activepowerlimits_to(d).min)
        max_rate = min(PSY.get_activepowerlimits_from(d).max, PSY.get_activepowerlimits_to(d).max)
        canonical.constraints[con_name][PSY.get_name(d), t] = JuMP.@constraint(canonical.JuMPmodel, min_rate <= canonical.variables[var_name][PSY.get_name(d), t] <= max_rate)
    end

    return

end

function branch_rate_constraint!(canonical::Canonical,
                                devices::IS.FlattenIteratorWrapper{B},
                                device_formulation::Type{HVDCLossless},
                                system_formulation::Type{S}) where {B<:PSY.DCBranch,
                                                                    S<:PM.AbstractActivePowerModel}

    for dir in ("FT", "TF")
        var_name = Symbol("Fp$(dir)_$(B)")
        con_name = Symbol("RateLimit_$(dir)_$(B)")
        time_steps = model_time_steps(canonical)
        canonical.constraints[con_name] = JuMPConstraintArray(undef, (PSY.get_name(d) for d in devices), time_steps)

        for t in time_steps, d in devices
            min_rate = max(PSY.get_activepowerlimits_from(d).min, PSY.get_activepowerlimits_to(d).min)
            max_rate = min(PSY.get_activepowerlimits_from(d).max, PSY.get_activepowerlimits_to(d).max)
            name = PSY.get_name(d)
            canonical.constraints[con_name][name, t] = JuMP.@constraint(canonical.JuMPmodel, min_rate <= canonical.variables[var_name][name, t] <= max_rate)
        end
    end

    return

end

function branch_rate_constraint!(canonical::Canonical,
                                devices::IS.FlattenIteratorWrapper{B},
                                device_formulation::Type{HVDCLossless},
                                system_formulation::Type{S}) where {B<:PSY.DCBranch,
                                                                    S<:PM.AbstractPowerModel}

    for dir in ("FT", "TF")
        var_name = Symbol("Fp$(dir)_$(B)")
        con_name = Symbol("RateLimit_$(dir)_$(B)")
        time_steps = model_time_steps(canonical)
        canonical.constraints[con_name] = JuMPConstraintArray(undef, (PSY.get_name(d) for d in devices), time_steps)

        for t in time_steps, d in devices
            min_rate = max(PSY.get_activepowerlimits_from(d).min, PSY.get_activepowerlimits_to(d).min)
            max_rate = min(PSY.get_activepowerlimits_from(d).max, PSY.get_activepowerlimits_to(d).max)
            name = PSY.get_name(d)
            canonical.constraints[con_name][name, t] = JuMP.@constraint(canonical.JuMPmodel, min_rate <= canonical.variables[var_name][name, t] <= max_rate)
        end
    end

    return

end

function branch_rate_constraint!(canonical::Canonical,
                                devices::IS.FlattenIteratorWrapper{B},
                                device_formulation::Type{D},
                                system_formulation::Type{S}) where {B<:PSY.DCBranch,
                                                                    D<:AbstractDCLineFormulation,
                                                                    S<:PM.AbstractActivePowerModel}

    time_steps = model_time_steps(canonical)

    for dir in ("FT", "TF")
        var_name = Symbol("Fp$(dir)_$(B)")
        con_name = Symbol("RateLimit$(dir)_$(B)")
        canonical.constraints[con_name] = JuMPConstraintArray(undef, (PSY.get_name(d) for d in devices), time_steps)

        for t in time_steps, d in devices
            min_rate = max(PSY.get_activepowerlimits_from(d).min, PSY.get_activepowerlimits_to(d).min)
            max_rate = min(PSY.get_activepowerlimits_from(d).max, PSY.get_activepowerlimits_to(d).max)
            canonical.constraints[con_name][PSY.get_name(d), t] = JuMP.@constraint(canonical.JuMPmodel, min_rate <= canonical.variables[var_name][PSY.get_name(d), t] <= max_rate)
            _add_to_expression!(canonical.expressions[:nodal_balance_active],
                                PSY.get_number(PSY.get_arc(d).to),
                                t,
                                canonical.variables[var_name][PSY.get_name(d), t],
                                -PSY.get_loss(d).l1,
                                -PSY.get_loss(d).l0)
        end
    end

    return

end

function branch_rate_constraint!(canonical::Canonical,
                                devices::IS.FlattenIteratorWrapper{B},
                                device_formulation::Type{D},
                                system_formulation::Type{S}) where {B<:PSY.DCBranch,
                                                                    D<:AbstractDCLineFormulation,
                                                                    S<:PM.AbstractPowerModel}

    time_steps = model_time_steps(canonical)

    for dir in ("FT", "TF")
        var_name = Symbol("Fp$(dir)_$(B)")
        con_name = Symbol("RateLimit$(dir)_$(B)")
        canonical.constraints[con_name] = JuMPConstraintArray(undef, (PSY.get_name(d) for d in devices), time_steps)

        for t in time_steps, d in devices
            min_rate = max(PSY.get_activepowerlimits_from(d).min, PSY.get_activepowerlimits_to(d).min)
            max_rate = min(PSY.get_activepowerlimits_from(d).max, PSY.get_activepowerlimits_to(d).max)
            canonical.constraints[con_name][PSY.get_name(d), t] = JuMP.@constraint(canonical.JuMPmodel, min_rate <= canonical.variables[var_name][PSY.get_name(d), t] <= max_rate)
            _add_to_expression!(canonical.expressions[:nodal_balance_active],
                                PSY.get_number(PSY.get_arc(d).to),
                                t,
                                canonical.variables[var_name][PSY.get_name(d), t],
                                -PSY.get_loss(d).l1,
                                -PSY.get_loss(d).l0)
        end
    end

    return

end


text:
module TestExamplesParallel2D

using Test
using Trixi

include("test_trixi.jl")

# Start with a clean environment: remove Trixi output directory if it exists
outdir = "out"
Trixi.mpi_isroot() && isdir(outdir) && rm(outdir, recursive=true)

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
const EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "2d")

# Run basic tests
@testset "Examples 2D" begin
  @testset "parameters.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [9.144681765639205e-6],
            linf = [6.437440532547356e-5])
  end
  @testset "parameters.toml with polydeg=1" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [0.05264106093598111],
            linf = [0.08754218386076518],
            polydeg=1)
  end
  @testset "parameters_ec.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_ec.toml"),
            l2   = [0.06159341742582756, 0.05012484425381723, 0.05013298724507752, 0.22537740506116724],
            linf = [0.29912627861573327, 0.30886767304359375, 0.3088108573487326, 1.0657556075017878])
  end
  @testset "parameters_density_wave.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_density_wave.toml"),
            l2   = [0.001060077845747576, 0.00010600778457107525, 0.00021201556914875742, 2.6501946139091318e-5],
            linf = [0.0065356386867677085, 0.0006535638688170142, 0.0013071277374487877, 0.0001633909674296774],
            extra_analysis_quantities=["l2_error_primitive", "linf_error_primitive"], t_end=0.5)
  end
  @testset "parameters_vortex.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_vortex.toml"),
            l2   = [3.6343138447409784e-6, 0.0032111379843728876, 0.0032111482778261658, 0.004545715889714643],
            linf = [7.901869034399045e-5, 0.030511158864742205, 0.030451936462313256, 0.04361908901631395])
  end
  # MHD + MPI not yet implemented
  # @testset "parameters_ec_mhd.toml" begin
  #   test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_ec_mhd.toml"),
  #           l2   = [0.03607862694368351, 0.04281395008247395, 0.04280207686965749, 0.025746770192645763, 0.1611518499414067, 0.017455917249117023, 0.017456981264942977, 0.02688321120361229, 0.00015024027267648003],
  #           linf = [0.23502083666166018, 0.3156846367743936, 0.31227895161037256, 0.2118146956106238, 0.9743049414302711, 0.09050624115026618, 0.09131633488909774, 0.15693063355520998, 0.0038394720095667593])
  # end
  @testset "parameters_hyp_diff_harmonic_nonperiodic.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_hyp_diff_harmonic_nonperiodic.toml"),
            l2   = [8.618132353932638e-8, 5.619399844708813e-7, 5.619399845476024e-7],
            linf = [1.124861862326869e-6, 8.622436471483752e-6, 8.622436469707395e-6])
  end
  @testset "parameters_hyp_diff_llf.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_hyp_diff_llf.toml"),
            l2   = [0.00015687751088073104, 0.0010259867353397119, 0.0010259867353398994],
            linf = [0.001198695640053704, 0.006423873515701395, 0.006423873515686296])
  end
  @testset "parameters_hyp_diff_nonperiodic.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_hyp_diff_nonperiodic.toml"),
            l2   = [8.523077654037775e-6, 2.877932365308637e-5, 5.454942769137812e-5],
            linf = [5.484978959957587e-5, 0.00014544895979200218, 0.000324491268921534])
  end
  @testset "parameters_hyp_diff_upwind.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_hyp_diff_upwind.toml"),
            l2   = [5.868147556488962e-6, 3.8051792732628014e-5, 3.8051792732620214e-5],
            linf = [3.70196549871471e-5, 0.0002072058411455302, 0.00020720584114464202])
  end
  @testset "parameters_nonperiodic.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_nonperiodic.toml"),
            l2   = [2.3652137675654753e-6, 2.1386731303685556e-6, 2.138673130413185e-6, 6.009920290578574e-6],
            linf = [1.4080448659026246e-5, 1.7581818010814487e-5, 1.758181801525538e-5, 5.9568540361709665e-5])
  end
  @testset "parameters_source_terms.toml" begin
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters_source_terms.toml"),
            l2   = [8.517783186497567e-7, 1.2350199409361865e-6, 1.2350199409828616e-6, 4.277884398786315e-6],
            linf = [8.357934254688004e-6, 1.0326389653148027e-5, 1.0326389654924384e-5, 4.4961900057316484e-5])
  end
  @testset "parameters.toml with restart and t_end=2" begin
    Trixi.run(joinpath(EXAMPLES_DIR, "parameters.toml"))
    test_trixi_run(joinpath(EXAMPLES_DIR, "parameters.toml"),
            l2   = [1.2148032444677485e-5],
            linf = [6.495644794757283e-5],
            t_end = 2,
            restart = true, restart_filename = "out/restart_000040.h5")
  end
end

# Clean up afterwards: delete Trixi output directory
Trixi.mpi_isroot() && @test_nowarn rm(outdir, recursive=true)

end #module


text:
# in this file we use cs and cs_i variable names that mean "target columns specification"

# this constant defines which types of values returned by aggregation function
# in combine are considered to produce multiple columns in the resulting data frame
const MULTI_COLS_TYPE = Union{AbstractDataFrame, NamedTuple, DataFrameRow, AbstractMatrix}

# use a constant Vector{Int} as a sentinel to signal that idx_agg has not been computed yet
# we do not use nothing to avoid excessive specialization
const NOTHING_IDX_AGG = Int[]

function gen_groups(idx::Vector{Int})
    groups = zeros(Int, length(idx))
    groups[1] = 1
    j = 1
    last_idx = idx[1]
    @inbounds for i in 2:length(idx)
        cur_idx = idx[i]
        j += cur_idx != last_idx
        last_idx = cur_idx
        groups[i] = j
    end
    return groups
end

function _combine_prepare(gd::GroupedDataFrame,
                          (cs,)::Ref{Any};
                          keepkeys::Bool, ungroup::Bool, copycols::Bool,
                          keeprows::Bool, renamecols::Bool)
    for cei in cs
        @assert cei isa Union{Pair, Base.Callable, ColumnIndex, MultiColumnIndex,
                              AbstractVecOrMat{<:Pair}}
    end
    if !ungroup && !keepkeys
        throw(ArgumentError("keepkeys=false when ungroup=false is not allowed"))
    end

    cs_vec = []
    for p in cs
        if p === nrow
            push!(cs_vec, nrow => :nrow)
        elseif p isa AbstractVecOrMat{<:Pair}
            append!(cs_vec, p)
        else
            push!(cs_vec, p)
        end
    end
    return _combine_prepare_norm(gd, cs_vec, keepkeys, ungroup, copycols,
                                 keeprows, renamecols)
end

function _combine_prepare_norm(gd::GroupedDataFrame,
                               cs_vec::Vector{Any},
                               keepkeys::Bool, ungroup::Bool, copycols::Bool,
                               keeprows::Bool, renamecols::Bool)
    if any(x -> x isa Pair && first(x) isa Tuple, cs_vec)
        x = cs_vec[findfirst(x -> first(x) isa Tuple, cs_vec)]
        # an explicit error is thrown as this was allowed in the past
        throw(ArgumentError("passing a Tuple $(first(x)) as column selector is not " *
                            "supported, use a vector $(collect(first(x))) instead"))
    end

    cs_norm = []
    optional_transform = Bool[]
    for c in cs_vec
        arg = normalize_selection(index(parent(gd)), c, renamecols)
        if arg isa AbstractVector{Int}
            for col_idx in arg
                push!(cs_norm, col_idx => identity => _names(gd)[col_idx])
                push!(optional_transform, true)
            end
        else
            push!(cs_norm, arg)
            push!(optional_transform, false)
        end
    end

    # cs_norm holds now either src => fun => dst or just fun
    # if optional_transform[i] is true then the transformation will be skipped
    # if earlier column with a column with the same name was created

    idx, valscat = _combine(gd, cs_norm, optional_transform, copycols, keeprows, renamecols)

    !keepkeys && ungroup && return valscat

    gd_keys = groupcols(gd)
    for key in gd_keys
        if hasproperty(valscat, key)
            if (keeprows && !isequal(valscat[!, key], parent(gd)[!, key])) ||
                (!keeprows && !isequal(valscat[!, key], view(parent(gd)[!, key], idx)))
                throw(ArgumentError("column :$key in returned data frame " *
                                    "is not equal to grouping key :$key"))
            end
        end
    end
    if keeprows
        newparent = select(parent(gd), gd.cols, copycols=copycols)
    else
        newparent = length(gd) > 0 ? parent(gd)[idx, gd.cols] : parent(gd)[1:0, gd.cols]
    end
    added_cols = select(valscat, Not(intersect(gd_keys, _names(valscat))), copycols=false)
    hcat!(newparent, length(gd) > 0 ? added_cols : similar(added_cols, 0), copycols=false)
    ungroup && return newparent

    if length(idx) == 0 && !(keeprows && length(gd_keys) > 0)
        @assert nrow(newparent) == 0
        return GroupedDataFrame(newparent, copy(gd.cols), Int[],
                                Int[], Int[], Int[], 0, Dict{Any, Int}(),
                                Threads.ReentrantLock())
    elseif keeprows
        @assert length(gd_keys) > 0 || idx == gd.idx
        # in this case we are sure that the result GroupedDataFrame has the
        # same structure as the source except that grouping columns are at the start
        return Threads.lock(gd.lazy_lock) do
            return GroupedDataFrame(newparent, copy(gd.cols), gd.groups,
                                    getfield(gd, :idx), getfield(gd, :starts),
                                    getfield(gd, :ends), gd.ngroups,
                                    getfield(gd, :keymap), Threads.ReentrantLock())
        end
    else
        groups = gen_groups(idx)
        @assert groups[end] <= length(gd)
        return GroupedDataFrame(newparent, copy(gd.cols), groups,
                                nothing, nothing, nothing, groups[end], nothing,
                                Threads.ReentrantLock())
    end
end

# Find first value matching condition for each group
# Optimized for situations where a matching value is typically encountered
# among the first rows for each group
function fillfirst!(condf, outcol::AbstractVector, incol::AbstractVector,
                    gd::GroupedDataFrame; rev::Bool=false)
    ngroups = gd.ngroups
    # Use group indices if they have already been computed
    idx = getfield(gd, :idx)
    if idx !== nothing && condf === nothing
        v = rev ? gd.ends : gd.starts
        @inbounds for i in 1:ngroups
            outcol[i] = incol[idx[v[i]]]
        end
    elseif idx !== nothing
        nfilled = 0
        starts = gd.starts
        @inbounds for i in eachindex(outcol)
            s = starts[i]
            offsets = rev ? (nrow(gd[i])-1:-1:0) : (0:nrow(gd[i])-1)
            for j in offsets
                x = incol[idx[s+j]]
                if !condf === nothing || condf(x)
                    outcol[i] = x
                    nfilled += 1
                    break
                end
            end
        end
        if nfilled < length(outcol)
            throw(ArgumentError("some groups contain only missing values"))
        end
    else # Finding first row is faster than computing all group indices
        groups = gd.groups
        if rev
            r = length(groups):-1:1
        else
            r = 1:length(groups)
        end
        filled = fill(false, ngroups)
        nfilled = 0
        @inbounds for i in r
            gix = groups[i]
            x = incol[i]
            if gix > 0 && (condf === nothing || condf(x)) && !filled[gix]
                filled[gix] = true
                outcol[gix] = x
                nfilled += 1
                nfilled == ngroups && break
            end
        end
        if nfilled < length(outcol)
            throw(ArgumentError("some groups contain only missing values"))
        end
    end
    outcol
end

function _agg2idx_map_helper(idx::Vector{Int}, idx_agg::Vector{Int})
    agg2idx_map = fill(-1, length(idx))
    aggj = 1
    @inbounds for (j, idxj) in enumerate(idx)
        while idx_agg[aggj] != idxj
            aggj += 1
            @assert aggj <= length(idx_agg)
        end
        agg2idx_map[j] = aggj
    end
    return agg2idx_map
end

struct TransformationResult
    col_idx::Vector{Int} # index for a column
    col::AbstractVector # computed value of a column
    name::Symbol # name of a column
    optional::Bool # whether a column is allowed to be replaced in the future
end

# the transformation is an aggregation for which we have the fast path
function _combine_process_agg((cs_i,)::Ref{Any},
                              optional_i::Bool,
                              parentdf::AbstractDataFrame,
                              gd::GroupedDataFrame,
                              seen_cols::Dict{Symbol, Tuple{Bool, Int}},
                              trans_res::Vector{TransformationResult},
                              idx_agg::Vector{Int})
    @assert cs_i isa Pair{Int, <:Pair{<:Function, Symbol}}
    @assert isagg(cs_i, gd)
    @assert !optional_i
    out_col_name = last(last(cs_i))
    incol = parentdf[!, first(cs_i)]
    agg = check_aggregate(first(last(cs_i)), incol)
    outcol = agg(incol, gd)

    return function()
        if haskey(seen_cols, out_col_name)
            optional, loc = seen_cols[out_col_name]
            # we have seen this col but it is not allowed to replace it
            optional || throw(ArgumentError("duplicate output column name: :$out_col_name"))
            @assert trans_res[loc].optional && trans_res[loc].name == out_col_name
            trans_res[loc] = TransformationResult(idx_agg, outcol, out_col_name, optional_i)
            seen_cols[out_col_name] = (optional_i, loc)
        else
            push!(trans_res, TransformationResult(idx_agg, outcol, out_col_name, optional_i))
            seen_cols[out_col_name] = (optional_i, length(trans_res))
        end
    end
end

# move one column without transorming it
function _combine_process_noop(cs_i::Pair{<:Union{Int, AbstractVector{Int}}, Pair{typeof(identity), Symbol}},
                               optional_i::Bool,
                               parentdf::AbstractDataFrame,
                               seen_cols::Dict{Symbol, Tuple{Bool, Int}},
                               trans_res::Vector{TransformationResult},
                               idx_keeprows::AbstractVector{Int},
                               copycols::Bool)
    source_cols = first(cs_i)
    out_col_name = last(last(cs_i))
    if length(source_cols) != 1
        throw(ArgumentError("Exactly one column can be transformed to one output column " *
                            "when using identity transformation"))
    end
    outcol = parentdf[!, first(source_cols)]

    return function()
        if haskey(seen_cols, out_col_name)
            optional, loc = seen_cols[out_col_name]
            @assert trans_res[loc].name == out_col_name
            if optional
                if !optional_i
                    @assert trans_res[loc].optional
                    trans_res[loc] = TransformationResult(idx_keeprows, copycols ? copy(outcol) : outcol,
                                                        out_col_name, optional_i)
                    seen_cols[out_col_name] = (optional_i, loc)
                end
            else
                # if optional_i is true, then we ignore processing this column
                optional_i || throw(ArgumentError("duplicate output column name: :$out_col_name"))
            end
        else
            push!(trans_res, TransformationResult(idx_keeprows, copycols ? copy(outcol) : outcol,
                                                out_col_name, optional_i))
            seen_cols[out_col_name] = (optional_i, length(trans_res))
        end
    end
end

# perform a transformation taking SubDataFrame as an input
function _combine_process_callable(wcs_i::Ref{Any},
                                   optional_i::Bool,
                                   parentdf::AbstractDataFrame,
                                   gd::GroupedDataFrame,
                                   seen_cols::Dict{Symbol, Tuple{Bool, Int}},
                                   trans_res::Vector{TransformationResult},
                                   idx_agg::Ref{Vector{Int}})
    cs_i = only(wcs_i)
    @assert cs_i isa Base.Callable
    firstres = length(gd) > 0 ? cs_i(gd[1]) : cs_i(similar(parentdf, 0))
    idx, outcols, nms = _combine_multicol(Ref{Any}(firstres), wcs_i, gd, Ref{Any}(nothing))

    if !(firstres isa Union{AbstractVecOrMat, AbstractDataFrame,
                            NamedTuple{<:Any, <:Tuple{Vararg{AbstractVector}}}})
        lock(gd.lazy_lock) do
            # if idx_agg was not computed yet it is NOTHING_IDX_AGG
            # in this case if we are not passed a vector compute it.
            if idx_agg[] === NOTHING_IDX_AGG
                idx_agg[] = Vector{Int}(undef, length(gd))
                fillfirst!(nothing, idx_agg[], 1:length(gd.groups), gd)
            end
            @assert idx == idx_agg[]
            idx = idx_agg[]
        end
    end
    @assert length(outcols) == length(nms)
    return function()
        for j in eachindex(outcols)
            outcol = outcols[j]
            out_col_name = nms[j]
            if haskey(seen_cols, out_col_name)
                optional, loc = seen_cols[out_col_name]
                # if column was seen and it is optional now ignore it
                if !optional_i
                    optional, loc = seen_cols[out_col_name]
                    # we have seen this col but it is not allowed to replace it
                    optional || throw(ArgumentError("duplicate output column name: :$out_col_name"))
                    @assert trans_res[loc].optional && trans_res[loc].name == out_col_name
                    trans_res[loc] = TransformationResult(idx, outcol, out_col_name, optional_i)
                    seen_cols[out_col_name] = (optional_i, loc)
                end
            else
                push!(trans_res, TransformationResult(idx, outcol, out_col_name, optional_i))
                seen_cols[out_col_name] = (optional_i, length(trans_res))
            end
        end
    end
end

# perform a transformation specified using the Pair notation with a single output column
function _combine_process_pair_symbol(optional_i::Bool,
                                      gd::GroupedDataFrame,
                                      seen_cols::Dict{Symbol, Tuple{Bool, Int}},
                                      trans_res::Vector{TransformationResult},
                                      idx_agg::Ref{Vector{Int}},
                                      out_col_name::Symbol,
                                      firstmulticol::Bool,
                                      (firstres,)::Ref{Any},
                                      wfun::Ref{Any},
                                      wincols::Ref{Any})
    @assert only(wfun) isa Base.Callable
    @assert only(wincols) isa Union{Tuple, NamedTuple}

    if firstmulticol
        throw(ArgumentError("a single value or vector result is required (got $(typeof(firstres)))"))
    end
    # if idx_agg was not computed yet it is NOTHING_IDX_AGG
    # in this case if we are not passed a vector compute it.
    lock(gd.lazy_lock) do
        if !(firstres isa AbstractVector) && idx_agg[] === NOTHING_IDX_AGG
            idx_agg[] = Vector{Int}(undef, length(gd))
            fillfirst!(nothing, idx_agg[], 1:length(gd.groups), gd)
        end
    end
    # TODO: if firstres is a vector we recompute idx for every function
    # this could be avoided - it could be computed only the first time
    # and later we could just check if lengths of groups match this first idx

    # the last argument passed to _combine_with_first informs it about precomputed
    # idx. Currently we do it only for single-row return values otherwise we pass
    # NOTHING_IDX_AGG to signal that idx has to be computed in _combine_with_first
    idx, outcols, _ = _combine_with_first(Ref{Any}(wrap(firstres)), wfun, gd, wincols,
                                          firstmulticol,
                                          firstres isa AbstractVector ? NOTHING_IDX_AGG : idx_agg[])
    @assert length(outcols) == 1
    outcol = outcols[1]

    return function()
        if haskey(seen_cols, out_col_name)
            # if column was seen and it is optional now ignore it
            if !optional_i
                optional, loc = seen_cols[out_col_name]
                # we have seen this col but it is not allowed to replace it
                optional || throw(ArgumentError("duplicate output column name: :$out_col_name"))
                @assert trans_res[loc].optional && trans_res[loc].name == out_col_name
                trans_res[loc] = TransformationResult(idx, outcol, out_col_name, optional_i)
                seen_cols[out_col_name] = (optional_i, loc)
            end
        else
            push!(trans_res, TransformationResult(idx, outcol, out_col_name, optional_i))
            seen_cols[out_col_name] = (optional_i, length(trans_res))
        end
    end
end

# perform a transformation specified using the Pair notation with multiple output columns
function _combine_process_pair_astable(optional_i::Bool,
                                       gd::GroupedDataFrame,
                                       seen_cols::Dict{Symbol, Tuple{Bool, Int}},
                                       trans_res::Vector{TransformationResult},
                                       idx_agg::Ref{Vector{Int}},
                                       out_col_name::Union{Type{AsTable}, AbstractVector{Symbol}},
                                       firstmulticol::Bool,
                                       (firstres,)::Ref{Any},
                                       wfun::Ref{Any},
                                       wincols::Ref{Any})
    fun = only(wfun)
    @assert fun isa Base.Callable
    @assert only(wincols) isa Union{Tuple, NamedTuple}
    if firstres isa AbstractVector
        idx, outcol_vec, _ = _combine_with_first(Ref{Any}(wrap(firstres)), wfun, gd, wincols,
                                                 firstmulticol, NOTHING_IDX_AGG)
        @assert length(outcol_vec) == 1
        res = outcol_vec[1]
        @assert length(res) > 0

        kp1 = keys(res[1])
        prepend = all(x -> x isa Integer, kp1)
        if !(prepend || all(x -> x isa Symbol, kp1) || all(x -> x isa AbstractString, kp1))
            throw(ArgumentError("keys of the returned elements must be " *
                                "`Symbol`s, strings or integers"))
        end
        if any(x -> !isequal(keys(x), kp1), res)
            throw(ArgumentError("keys of the returned elements must be identical"))
        end
        outcols = [[x[n] for x in res] for n in kp1]
        nms = [prepend ? Symbol("x", n) : Symbol(n) for n in kp1]
    else
        if !firstmulticol
            firstres = Tables.columntable(firstres)
            oldfun = fun
            fun = (x...) -> Tables.columntable(oldfun(x...))
        end
        idx, outcols, nms = _combine_multicol(Ref{Any}(firstres), Ref{Any}(fun), gd, wincols)

        if !(firstres isa Union{AbstractVecOrMat, AbstractDataFrame,
            NamedTuple{<:Any, <:Tuple{Vararg{AbstractVector}}}})
            lock(gd.lazy_lock) do
                # if idx_agg was not computed yet it is nothing
                # in this case if we are not passed a vector compute it.
                if idx_agg[] === NOTHING_IDX_AGG
                    idx_agg[] = Vector{Int}(undef, length(gd))
                    fillfirst!(nothing, idx_agg[], 1:length(gd.groups), gd)
                end
                @assert idx == idx_agg[]
                idx = idx_agg[]
            end
        end
        @assert length(outcols) == length(nms)
    end
    if out_col_name isa AbstractVector{Symbol}
        if length(out_col_name) != length(nms)
            throw(ArgumentError("Number of returned columns does not " *
                                "match the length of requested output"))
        else
            nms = out_col_name
        end
    end
    return function()
        for j in eachindex(outcols)
            outcol = outcols[j]
            out_col_name = nms[j]
            if haskey(seen_cols, out_col_name)
                optional, loc = seen_cols[out_col_name]
                # if column was seen and it is optional now ignore it
                if !optional_i
                    optional, loc = seen_cols[out_col_name]
                    # we have seen this col but it is not allowed to replace it
                    optional || throw(ArgumentError("duplicate output column name: :$out_col_name"))
                    @assert trans_res[loc].optional && trans_res[loc].name == out_col_name
                    trans_res[loc] = TransformationResult(idx, outcol, out_col_name, optional_i)
                    seen_cols[out_col_name] = (optional_i, loc)
                end
            else
                push!(trans_res, TransformationResult(idx, outcol, out_col_name, optional_i))
                seen_cols[out_col_name] = (optional_i, length(trans_res))
            end
        end
    end
end

# perform a transformation specified using the Pair notation
# cs_i is a Pair that has many possible forms so this function is used to dispatch
# to an appropriate more specialized function
function _combine_process_pair((cs_i,)::Ref{Any},
                               optional_i::Bool,
                               parentdf::AbstractDataFrame,
                               gd::GroupedDataFrame,
                               seen_cols::Dict{Symbol, Tuple{Bool, Int}},
                               trans_res::Vector{TransformationResult},
                               idx_agg::Ref{Vector{Int}})
    @assert cs_i isa Pair

    source_cols, (fun, out_col_name) = cs_i

    if source_cols isa Int
        incols = (parentdf[!, source_cols],)
    elseif source_cols isa AsTable
        incols = Tables.columntable(select(parentdf,
                                           source_cols.cols,
                                           copycols=false))
    else
        @assert source_cols isa AbstractVector{Int}
        incols = ntuple(i -> parentdf[!, source_cols[i]], length(source_cols))
    end

    firstres = length(gd) > 0 ?
               do_call(fun, gd.idx, gd.starts, gd.ends, gd, incols, 1) :
               do_call(fun, Int[], 1:1, 0:0, gd, incols, 1)
    firstmulticol = firstres isa MULTI_COLS_TYPE

    if out_col_name isa Symbol
        return _combine_process_pair_symbol(optional_i, gd, seen_cols, trans_res, idx_agg,
                                            out_col_name, firstmulticol, Ref{Any}(firstres),
                                            Ref{Any}(fun), Ref{Any}(incols))
    end
    if out_col_name == AsTable || out_col_name isa AbstractVector{Symbol}
        return _combine_process_pair_astable(optional_i, gd, seen_cols, trans_res, idx_agg,
                                             out_col_name, firstmulticol, Ref{Any}(firstres),
                                             Ref{Any}(fun), Ref{Any}(incols))
    end
    throw(ArgumentError("unsupported target column name specifier $out_col_name"))
end

function prepare_idx_keeprows(idx::AbstractVector{<:Integer},
                              starts::AbstractVector{<:Integer},
                              ends::AbstractVector{<:Integer},
                              nrowparent::Integer)
    idx_keeprows = Vector{Int}(undef, nrowparent)
    i = 0
    for (s, e) in zip(starts, ends)
        v = idx[s]
        for k in s:e
            i += 1
            idx_keeprows[i] = v
        end
    end
    @assert i == nrowparent
    return idx_keeprows
end

function _combine(gd::GroupedDataFrame,
                  cs_norm::Vector{Any}, optional_transform::Vector{Bool},
                  copycols::Bool, keeprows::Bool, renamecols::Bool)
    if isempty(cs_norm)
        if keeprows && nrow(parent(gd)) > 0 && minimum(gd.groups) == 0
            throw(ArgumentError("select and transform do not support " *
                                "`GroupedDataFrame`s from which some groups have "*
                                "been dropped (including skipmissing=true)"))
        end
        return Int[], DataFrame()
    end

    if keeprows
        if nrow(parent(gd)) > 0 && minimum(gd.groups) == 0
            throw(ArgumentError("select and transform do not support " *
                                "`GroupedDataFrame`s from which some groups have "*
                                "been dropped (including skipmissing=true)"))
        end
        idx_keeprows = prepare_idx_keeprows(gd.idx, gd.starts, gd.ends, nrow(parent(gd)))
    else
        idx_keeprows = Int[]
    end

    idx_agg = Ref(NOTHING_IDX_AGG)
    if length(gd) > 0 && any(x -> isagg(x, gd), cs_norm)
        # Compute indices of representative rows only once for all AbstractAggregates
        idx_agg[] = Vector{Int}(undef, length(gd))
        fillfirst!(nothing, idx_agg[], 1:length(gd.groups), gd)
    elseif length(gd) == 0 || !all(x -> isagg(x, gd), cs_norm)
        # Trigger computation of indices
        # This can speed up some aggregates that would not trigger this on their own
        @assert gd.idx !== nothing
    end

    trans_res = Vector{TransformationResult}()

    # seen_cols keeps an information about location of columns already processed
    # and if a given column can be replaced in the future
    seen_cols = Dict{Symbol, Tuple{Bool, Int}}()

    tasks = similar(cs_norm, Task)

    parentdf = parent(gd)
    # Operations are run in separate tasks, except two parts:
    # - the first task that needs idx_agg computes it;
    #   a lock ensures others wait for it to complete
    # - once all tasks are done, they need sequential postprocessing
    #   since their order affects that of columns
    for i in eachindex(cs_norm, optional_transform, tasks)
        cs_i = cs_norm[i]
        optional_i = optional_transform[i]

        tasks[i] = @spawn if length(gd) > 0 && isagg(cs_i, gd)
            _combine_process_agg(Ref{Any}(cs_i), optional_i, parentdf, gd, seen_cols, trans_res, idx_agg[])
        elseif keeprows && cs_i isa Pair && first(last(cs_i)) === identity &&
               !(first(cs_i) isa AsTable) && (last(last(cs_i)) isa Symbol)
            # this is a fast path used when we pass a column or rename a column in select or transform
            _combine_process_noop(cs_i, optional_i, parentdf, seen_cols, trans_res, idx_keeprows, copycols)
        elseif cs_i isa Base.Callable
            _combine_process_callable(Ref{Any}(cs_i), optional_i, parentdf, gd, seen_cols, trans_res, idx_agg)
        else
            @assert cs_i isa Pair
            _combine_process_pair(Ref{Any}(cs_i), optional_i, parentdf, gd, seen_cols, trans_res, idx_agg)
        end
    end
    # Workaround JuliaLang/julia#38931:
    # we want to preserve the exception type thrown in user code,
    # and print the backtrace corresponding to it
    for t in tasks
        try
            wait(t)
        catch e
            @static if VERSION > v"1.3"
                if e isa TaskFailedException
                    throw(t.exception)
                else
                    rethrow(e)
                end
            else
                if e isa ErrorException
                    throw(t.exception)
                else
                    rethrow(e)
                end
            end
        end
    end
    # Post-processing has to be run sequentially
    # since the order of operations determines that of columns
    for t in tasks
        postprocessf = fetch(t)
        postprocessf()
    end

    isempty(trans_res) && return Int[], DataFrame()
    # idx_agg[] === NOTHING_IDX_AGG then we have only functions that
    # returned multiple rows and idx_loc = 1
    idx_loc = findfirst(x -> x.col_idx !== idx_agg[], trans_res)
    if !keeprows && isnothing(idx_loc)
        @assert idx_agg[] !== NOTHING_IDX_AGG
        idx = idx_agg[]
    else
        idx = keeprows ? idx_keeprows : trans_res[idx_loc].col_idx
        agg2idx_map = nothing
        for i in 1:length(trans_res)
            if trans_res[i].col_idx !== idx
                if trans_res[i].col_idx === idx_agg[]
                    # we perform pseudo broadcasting here
                    # keep -1 as a sentinel for errors
                    if isnothing(agg2idx_map)
                        agg2idx_map = _agg2idx_map_helper(idx, idx_agg[])
                    end
                    trans_res[i] = TransformationResult(idx_agg[], trans_res[i].col[agg2idx_map],
                                                        trans_res[i].name, trans_res[i].optional)
                elseif idx != trans_res[i].col_idx
                    if keeprows
                        throw(ArgumentError("all functions must return vectors with " *
                                            "as many values as rows in each group"))
                    else
                        throw(ArgumentError("all functions must return vectors of the same length"))
                    end
                end
            end
        end
    end

    # here first field in trans_res[i] is used to keep track how the column was generated
    # a correct index is stored in idx variable

    @sync for i in eachindex(trans_res)
        let i=i
            @spawn reorder_cols!(trans_res, i, trans_res[i].col, trans_res[i].col_idx,
                                 keeprows, idx_keeprows, gd)
        end
    end

    outcols = AbstractVector[x.col for x in trans_res]
    nms = Symbol[x.name for x in trans_res]
    # this check is redundant given we check idx above
    # but it is safer to double check and it is cheap
    @assert all(x -> length(x) == length(outcols[1]), outcols)
    return idx, DataFrame(outcols, nms, copycols=false)
end

function reorder_cols!(trans_res::Vector{TransformationResult}, i::Integer,
                       col::AbstractVector, col_idx::Vector{Int}, keeprows::Bool,
                       idx_keeprows::Vector{Int}, gd::GroupedDataFrame)
    if keeprows && col_idx !== idx_keeprows # we need to reorder the column
        newcol = similar(col)
        # we can probably make it more efficient, but I leave it as an optimization for the future
        gd_idx = gd.idx
        k = 0
        for (s, e) in zip(gd.starts, gd.ends)
            for j in s:e
                k += 1
                @inbounds newcol[gd_idx[j]] = col[k]
            end
        end
        @assert k == length(gd_idx)
        trans_res[i] = TransformationResult(col_idx, newcol, trans_res[i].name, trans_res[i].optional)
    end
end

function combine(@nospecialize(f::Base.Callable), gd::GroupedDataFrame;
                 keepkeys::Bool=true, ungroup::Bool=true, renamecols::Bool=true)
    if f isa Colon
        throw(ArgumentError("First argument must be a transformation if the second argument is a GroupedDataFrame"))
    end
    return combine(gd, f, keepkeys=keepkeys, ungroup=ungroup, renamecols=renamecols)
end

combine(@nospecialize(f::Pair), gd::GroupedDataFrame;
        keepkeys::Bool=true, ungroup::Bool=true, renamecols::Bool=true) =
    throw(ArgumentError("First argument must be a transformation if the second argument is a GroupedDataFrame. " *
                        "You can pass a `Pair` as the second argument of the transformation. If you want the return " *
                        "value to be processed as having multiple columns add `=> AsTable` suffix to the pair."))

combine(gd::GroupedDataFrame,
        @nospecialize(cs::Union{Pair, Base.Callable, ColumnIndex, MultiColumnIndex,
                                AbstractVecOrMat{<:Pair}}...);
        keepkeys::Bool=true, ungroup::Bool=true, renamecols::Bool=true) =
    _combine_prepare(gd, Ref{Any}(cs), keepkeys=keepkeys, ungroup=ungroup,
                     copycols=true, keeprows=false, renamecols=renamecols)

function select(@nospecialize(f::Base.Callable), gd::GroupedDataFrame; copycols::Bool=true,
                keepkeys::Bool=true, ungroup::Bool=true, renamecols::Bool=true)
    if f isa Colon
        throw(ArgumentError("First argument must be a transformation if the second argument is a grouped data frame"))
    end
    return select(gd, f, copycols=copycols, keepkeys=keepkeys, ungroup=ungroup)
end

select(gd::GroupedDataFrame, @nospecialize(args::Union{Pair, Base.Callable, ColumnIndex, MultiColumnIndex,
                                                       AbstractVecOrMat{<:Pair}}...);
       copycols::Bool=true, keepkeys::Bool=true, ungroup::Bool=true, renamecols::Bool=true) =
    _combine_prepare(gd, Ref{Any}(args), copycols=copycols, keepkeys=keepkeys,
                     ungroup=ungroup, keeprows=true, renamecols=renamecols)

function transform(@nospecialize(f::Base.Callable), gd::GroupedDataFrame; copycols::Bool=true,
                keepkeys::Bool=true, ungroup::Bool=true, renamecols::Bool=true)
    if f isa Colon
        throw(ArgumentError("First argument must be a transformation if the second argument is a grouped data frame"))
    end
    return transform(gd, f, copycols=copycols, keepkeys=keepkeys, ungroup=ungroup)
end

function transform(gd::GroupedDataFrame, @nospecialize(args::Union{Pair, Base.Callable, ColumnIndex, MultiColumnIndex,
                                                                   AbstractVecOrMat{<:Pair}}...);
                   copycols::Bool=true, keepkeys::Bool=true, ungroup::Bool=true, renamecols::Bool=true)
    res = select(gd, :, args..., copycols=copycols, keepkeys=keepkeys,
                 ungroup=ungroup, renamecols=renamecols)
    # res can be a GroupedDataFrame based on DataFrame or a DataFrame,
    # so parent always gives a data frame
    select!(parent(res), propertynames(parent(gd)), :)
    return res
end

function select!(@nospecialize(f::Base.Callable), gd::GroupedDataFrame; ungroup::Bool=true, renamecols::Bool=true)
    if f isa Colon
        throw(ArgumentError("First argument must be a transformation if the second argument is a grouped data frame"))
    end
    return select!(gd, f, ungroup=ungroup)
end

function select!(gd::GroupedDataFrame,
                 @nospecialize(args::Union{Pair, Base.Callable, ColumnIndex, MultiColumnIndex,
                                           AbstractVecOrMat{<:Pair}}...);
                 ungroup::Bool=true, renamecols::Bool=true)
    df = parent(gd)
    if df isa DataFrame
        newdf = select(gd, args..., copycols=false, renamecols=renamecols)
    else
        @assert df isa SubDataFrame
        newdf = select(gd, args..., copycols=true, renamecols=renamecols)
    end
    _replace_columns!(df, newdf)
    return ungroup ? df : gd
end

function transform!(@nospecialize(f::Base.Callable), gd::GroupedDataFrame; ungroup::Bool=true, renamecols::Bool=true)
    if f isa Colon
        throw(ArgumentError("First argument must be a transformation if the second argument is a grouped data frame"))
    end
    return transform!(gd, f, ungroup=ungroup)
end

function transform!(gd::GroupedDataFrame,
                    @nospecialize(args::Union{Pair, Base.Callable, ColumnIndex, MultiColumnIndex,
                                              AbstractVecOrMat{<:Pair}}...);
                    ungroup::Bool=true, renamecols::Bool=true)
    df = parent(gd)
    if df isa DataFrame
        newdf = select(gd, :, args..., copycols=false, renamecols=renamecols)
    else
        @assert df isa SubDataFrame
        newdf = select(gd, :, args..., copycols=true, renamecols=renamecols)
    end
    select!(newdf, propertynames(df), :)
    _replace_columns!(df, newdf)
    return ungroup ? df : gd
end


text:
#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    ReverseGraphView, reversegraph

struct ReverseGraphView{T<:DirectedGraph} <: DiGraphView
    graph::T
end

reversegraph(graph::DirectedGraph) = ReverseGraphView(graph)

function getedge(view::ReverseGraphView, idx)
    edge = getedge(view.graph, idx)
    return reverseedge(edge)
end

getedge(
    view::ReverseGraphView, s, t
) = getedge(view, outneighbors(view.graph, s)[t])

edgesiter(
    view::ReverseGraphView
) = [i => reverseedge(e) for (i, e) in edgesiter(view.graph)]

function reverseedge(edge)
    return setnodes(edge, edge.target, edge.source)
end

outneighbors(view::ReverseGraphView, idx) = view.graph.inneighbormap[idx]
inneighbors(view::ReverseGraphView, idx) = view.graph.outneighbormap[idx]


text:
# Gibbs Sampling Bayesian Probabilistic Matrix Factorization
#
# Reference
# ----------
#   Salakhutdinov, Ruslan, and Andriy Mnih. "Bayesian probabilistic matrix factorization using Markov chain Monte Carlo."
#   Proceedings of the 25th international conference on Machine learning. ACM, 2008.
#
mutable struct GBPMFModel
    R::Array
    N::Int64
    M::Int64
    D::Int64
    T::Int64
    U::Array
    V::Array
    α::Float64
    β₀::Float64
    μ₀::Float64
    ν₀::Float64
    W₀::Array
end

function fit(model::GBPMFModel)
    #set initial value
    model.U = zeros((model.T, model.N, model.D)) #user latent matrix
    model.V = zeros((model.T, model.M, model.D)) #item latent matrix
    model.U[1, :, :] = rand(model.N, model.D)
    model.V[1, :, :] = rand(model.M, model.D)

    T = model.T
    N = model.N
    M = model.M

    for t in 1:(T-1) #1 is initial value
        #u matrix
        Θᵤ = hyper_prameter_posterior(model, model.U[t, :, :])
        μᵤ, Λᵤ = hyper_parameter_sampling(Θᵤ[1], Θᵤ[2], Θᵤ[3], Θᵤ[4])

        #v matrix
        Θᵥ = hyper_prameter_posterior(model, model.V[t, :, :])
        μᵥ, Λᵥ = hyper_parameter_sampling(Θᵥ[1], Θᵥ[2], Θᵥ[3], Θᵥ[4])

        for n in 1:N
            μᵢ_asterisk, Λᵢ_asterisk_inv = u_matrix_posterior(model, model.V[t, :, :], μᵤ, Λᵤ, n)
            model.U[t + 1, n, :] = latent_matrix_sampling(μᵢ_asterisk, Λᵢ_asterisk_inv)
        end

        for m in 1:M
            μᵢ_asterisk, Λᵢ_asterisk_inv = v_matrix_posterior(model, model.U[t + 1, :, :], μᵥ, Λᵥ, m)
            model.V[t + 1, m, :] = latent_matrix_sampling(μᵢ_asterisk, Λᵢ_asterisk_inv)
        end

    end
end

function hyper_prameter_posterior(model::GBPMFModel, latent_matrix)
    N = size(latent_matrix)[2]
    β₀ = model.β₀
    μ₀ = model.μ₀
    ν₀ = model.ν₀
    W₀ = model.W₀

    latent_matrix_bar = (1 / N) * sum(latent_matrix[:, :], dims = 1)
    S̄ = (1 / N) .* sum([latent_matrix[i, :]' .* latent_matrix[i, :] for i = 1:N], dims = 1)[1]

    μ₀_asterisk = ((β₀ * μ₀ .+ N * latent_matrix_bar) / (β₀ + N))[1, :]
    β₀_asterisk = β₀ + N
    ν₀_asterisk = ν₀ + N

    W₀_asterisk = inv(inv(W₀) .+ N .* S̄ .+ ((β₀ * N) / (β₀ + N)) .* (μ₀ .- latent_matrix_bar)' * (μ₀ .- latent_matrix_bar))
    W₀_asterisk = Array(Symmetric(W₀_asterisk))

    return μ₀_asterisk, β₀_asterisk, ν₀_asterisk, W₀_asterisk
end

function hyper_parameter_sampling(μ₀_asterisk, β₀_asterisk, ν₀_asterisk, W₀_asterisk)
    Λ = rand(Wishart(ν₀_asterisk, W₀_asterisk), 1)[1]
    μ = rand(MvNormal(μ₀_asterisk, Array(Symmetric(inv(β₀_asterisk * Λ)))), 1)

    return μ, Λ
end

function u_matrix_posterior(model::GBPMFModel, latent_matrix, μᵤ, Λᵤ, n)
    R = model.R
    α = model.α

    r = R[R[:, 1] .==  n - 1, :] #user's rate item
    nonzero_ix = (R[R[:, 1] .==  n - 1, :][:, 2] .+ 1) #user's rate item index

    if nonzero_ix == [] #when no rate then return random value
        μᵢ_asterisk = zeros(size(latent_matrix)[2])
        Λᵢ_asterisk = one(zeros(size(latent_matrix)[2], size(latent_matrix)[2]))
        return μᵢ_asterisk, Λᵢ_asterisk
    end

    Λᵢ_asterisk = Λᵤ .+ α * sum([latent_matrix[i, :]' .*  latent_matrix[i, :] for i in nonzero_ix], dims = 1)[1]
    Λᵢ_asterisk_inv = Array(Symmetric(inv(Λᵢ_asterisk)))

    μᵢ_asterisk = (inv(Λᵢ_asterisk) * ((α * sum([latent_matrix[i, :] .* r[r[:, 2] .== i - 1, :][:, 3] for i in nonzero_ix], dims = 1)[1]') + (Λᵤ * μᵤ)')')[:, 1]

    return μᵢ_asterisk, Λᵢ_asterisk_inv
end

function v_matrix_posterior(model::GBPMFModel, latent_matrix, μᵥ, Λᵥ, m)
    R = model.R
    α = model.α

    r = R[R[:, 2] .==  m - 1, :] #item rate user
    nonzero_ix = (R[R[:, 2] .==  m - 1, :][:, 1] .+ 1) #item rate user index

    if nonzero_ix == [] #when no rate then return random value
        μᵢ_asterisk = zeros(size(latent_matrix)[2])
        Λᵢ_asterisk = one(zeros(size(latent_matrix)[2], size(latent_matrix)[2]))
        return μᵢ_asterisk, Λᵢ_asterisk
    end

    Λᵢ_asterisk = Λᵥ .+ α * sum([latent_matrix[i, :]' .*  latent_matrix[i, :] for i in nonzero_ix], dims = 1)[1]
    Λᵢ_asterisk_inv = Array(Symmetric(inv(Λᵢ_asterisk)))

    μᵢ_asterisk = (inv(Λᵢ_asterisk) * ((α * sum([latent_matrix[i, :] .* r[r[:, 1] .== i - 1, :][:, 3] for i in nonzero_ix], dims = 1)[1]') + (Λᵥ * μᵥ)')')[:, 1]

    return μᵢ_asterisk, Λᵢ_asterisk_inv
end

function latent_matrix_sampling(μᵢ_asterisk, Λᵢ_asterisk_inv)
    return  rand(MvNormal(μᵢ_asterisk, Λᵢ_asterisk_inv), 1)
end

function predict_all(model::GBPMFModel, burn_in)
    T = model.T
    sum([model.U[i, :, :] *  model.V[i,:,  :]' for i = burn_in:T], dims = 1)[1]  / (T - burn_in )
end

function predict(model::GBPMFModel, new_data, burn_in)
    R_p = predict_all(model, burn_in)
    r = zeros(size(new_data)[1])
    for i in 1:size(new_data)[1]
        r[i] = R_p[new_data[i, 1] + 1, new_data[i, 2] + 1]
    end
    return r
end


text:
# helper functions to check that the given system can be applied to the states/inputs
@inline _is_conformable_state(system, x) = statedim(system) == length(x)
@inline _is_conformable_input(system, u) = inputdim(system) == length(u)
@inline _is_conformable_noise(system, w) = noisedim(system) == length(w)
@inline _in_stateset(system, x) = x ∈ stateset(system)
@inline _in_inputset(system, u) = u ∈ inputset(system)
@inline _in_noiseset(system, w) = w ∈ noiseset(system)

@inline function _argument_error(sym, set=:none)
    if sym == :x
        vector = "state"
    elseif sym == :u
        vector = "input"
    elseif sym == :w
        vector = "noise"
    end
    if set == :none
        throw(ArgumentError("the $vector vector has the wrong dimensions"))
    else
        throw(ArgumentError("the $vector vector is not contained in the $(vector) set"))
    end
end

"""
    successor(system::DiscreteIdentitySystem, x::AbstractVector)

Return the successor state of a `DiscreteIdentitySystem`.

### Input

- `system` -- `DiscreteIdentitySystem`
- `x`      -- state (it should be any vector type)

### Output

The same state `x`.
"""
function successor(system::DiscreteIdentitySystem, x::AbstractVector)
    !_is_conformable_state(system, x) && _argument_error(:x)
    return x
end

"""
    successor(system::ConstrainedDiscreteIdentitySystem, x::AbstractVector;
              [check_constraints]=true)

Return the successor state of a `ConstrainedDiscreteIdentitySystem`.

### Input

- `system`            -- `ConstrainedDiscreteIdentitySystem`
- `x`                 -- state (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state belongs to
                         the state set

### Output

The same state `x`.
"""
function successor(system::ConstrainedDiscreteIdentitySystem, x::AbstractVector;
                   check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
    end
    return x
end

"""
    successor(system::LinearDiscreteSystem, x::AbstractVector)

Return the successor state of a `LinearDiscreteSystem`.

### Input

- `system` -- `LinearDiscreteSystem`
- `x`      -- state (it should be any vector type)

### Output

The result of applying the system to `x`.
"""
function successor(system::LinearDiscreteSystem, x::AbstractVector)
    !_is_conformable_state(system, x) && _argument_error(:x)
    return system.A * x
end

"""
    successor(system::AffineDiscreteSystem, x::AbstractVector)

Return the successor state of a `AffineDiscreteSystem`.

### Input

- `system` -- `AffineDiscreteSystem`
- `x`      -- state (it should be any vector type)

### Output

The result of applying the system to `x`.
"""
function successor(system::AffineDiscreteSystem, x::AbstractVector)
    !_is_conformable_state(system, x) && _argument_error(:x)
    return system.A * x + system.c
end

"""
    successor(system::LinearControlDiscreteSystem, x::AbstractVector, u::AbstractVector)

Return the successor state of a `LinearControlDiscreteSystem`.

### Input

- `system` -- `LinearControlDiscreteSystem`
- `x`      -- state (it should be any vector type)
- `u`      -- input (it should be any vector type)

### Output

The result of applying the system to `x`, with input `u`.
"""
function successor(system::LinearControlDiscreteSystem, x::AbstractVector, u::AbstractVector)
    !_is_conformable_state(system, x) && _argument_error(:x)
    !_is_conformable_input(system, u) && _argument_error(:u)
    return system.A * x + system.B * u
end

"""
    successor(system::ConstrainedLinearDiscreteSystem, x::AbstractVector;
              [check_constraints]=true)

Return the successor state of a `ConstrainedLinearDiscreteSystem`.

### Input

- `system`            -- `ConstrainedLinearDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state belongs to
                         the state set

### Output

The result of applying the system to `x`.
"""
function successor(system::ConstrainedLinearDiscreteSystem, x::AbstractVector;
                   check_constraints::Bool=true)
   !_is_conformable_state(system, x) && _argument_error(:x)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
    end
    return system.A * x
end

"""
    successor(system::ConstrainedAffineDiscreteSystem, x::AbstractVector;
              [check_constraints])

Return the successor state of a `ConstrainedAffineDiscreteSystem`.

### Input

- `system`            -- `ConstrainedAffineDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state belongs to
                         the state set

### Output

The result of applying the system to `x`.
"""
function successor(system::ConstrainedAffineDiscreteSystem, x::AbstractVector;
                   check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
    end
    return system.A * x + system.c
end

"""
    successor(system::ConstrainedLinearControlDiscreteSystem, x::AbstractVector,
              u::AbstractVector; [check_constraints]=true)

Return the successor state of a `ConstrainedLinearControlDiscreteSystem`.

### Input

- `system`            -- `ConstrainedLinearControlDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `u`                 -- input (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state (resp. input)
                         belongs to the state set (resp. input set)

### Output

The result of applying the system to `x`, with input `u`.
"""
function successor(system::ConstrainedLinearControlDiscreteSystem, x::AbstractVector,
                   u::AbstractVector; check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    !_is_conformable_input(system, u) && _argument_error(:u)
    if check_constraints
        !_in_stateset(system, x) &&_argument_error(:x,:X)
        !_in_inputset(system, u) && _argument_error(:u,:U)
    end
    return system.A * x + system.B * u
end

"""
    successor(system::ConstrainedAffineControlDiscreteSystem, x::AbstractVector,
              u::AbstractVector; [check_constraints]=true)

Return the successor state of a `ConstrainedAffineControlDiscreteSystem`.

### Input

- `system`            -- `ConstrainedAffineControlDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `u`                 -- input (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state (resp. input)
                         belongs to the state set (resp. input set)

### Output

The result of applying the system to `x`, with input `u`.
"""
function successor(system::ConstrainedAffineControlDiscreteSystem,
                   x::AbstractVector, u::AbstractVector; check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    !_is_conformable_input(system, u) && _argument_error(:u)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
        !_in_inputset(system, u) && _argument_error(:u,:U)
    end
    return system.A * x + system.B * u + system.c
end

"""
    successor(system::BlackBoxDiscreteSystem, x::AbstractVector)

Return the successor state of a `BlackBoxDiscreteSystem`.

### Input

- `system` -- `BlackBoxDiscreteSystem`
- `x`      -- state (it should be any vector type)

### Output

The result of applying the system to `x`.
"""
function successor(system::BlackBoxDiscreteSystem, x::AbstractVector)
    !_is_conformable_state(system, x) && _argument_error(:x)
    return system.f(x)
end

"""
    successor(system::ConstrainedBlackBoxDiscreteSystem, x::AbstractVector;
              [check_constraints]=true)

Return the successor state of a `ConstrainedBlackBoxDiscreteSystem`.

### Input

- `system`            -- `ConstrainedBlackBoxDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state belongs to
                         the state set

### Output

The result of applying the system to `x`.
"""
function successor(system::ConstrainedBlackBoxDiscreteSystem, x::AbstractVector;
                   check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
    end
    return system.f(x)
end

"""
    successor(system::ConstrainedBlackBoxControlDiscreteSystem, x::AbstractVector,
              u::AbstractVector; [check_constraints]=true)

Return the successor state of a `ConstrainedBlackBoxControlDiscreteSystem`.

### Input

- `system`            -- `ConstrainedBlackBoxControlDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `u`                 -- input (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state belongs to
                         the state set

### Output

The result of applying the system to `x`, with input `u`.
"""
function successor(system::ConstrainedBlackBoxControlDiscreteSystem, x::AbstractVector,
                   u::AbstractVector; check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    !_is_conformable_input(system, u) && _argument_error(:u)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
        !_in_inputset(system, u) && _argument_error(:u,:U)
    end
    return system.f(x, u)
end

"""
    successor(system::NoisyConstrainedLinearDiscreteSystem, x::AbstractVector,
              w::AbstractVector; [check_constraints]=true)

Return the successor state of a `NoisyConstrainedLinearDiscreteSystem`.

### Input

- `system`            -- `NoisyConstrainedLinearDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `w`                 -- noise (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state belongs to
                         the state set

### Output

The result of applying the system to `x`, with noise `w`.
"""
function successor(system::NoisyConstrainedLinearDiscreteSystem, x::AbstractVector,
                   w::AbstractVector; check_constraints::Bool=true)
   !_is_conformable_state(system, x) && _argument_error(:x)
   !_is_conformable_noise(system, w) && _argument_error(:w)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
        !_in_noiseset(system, w) && _argument_error(:w,:W)
    end
    return system.A * x + system.D * w
end

"""
    successor(system::NoisyConstrainedLinearControlDiscreteSystem,
              x::AbstractVector, u::AbstractVector, w::AbstractVector; [check_constraints]=true)

Return the successor state of a `NoisyConstrainedLinearControlDiscreteSystem`.

### Input

- `system`            -- `NoisyConstrainedLinearControlDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `u`                 -- input (it should be any vector type)
- `w`                 -- noise (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state (resp. input)
                         belongs to the state set (resp. input set)

### Output

The result of applying the system to `x`, with input `u` and noise `w`.
"""
function successor(system::NoisyConstrainedLinearControlDiscreteSystem,
                   x::AbstractVector, u::AbstractVector, w::AbstractVector; check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    !_is_conformable_input(system, u) && _argument_error(:u)
    !_is_conformable_noise(system, w) && _argument_error(:w)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
        !_in_inputset(system, u) && _argument_error(:u,:U)
        !_in_noiseset(system, w) && _argument_error(:w,:W)
    end
    return system.A * x + system.B * u + system.D * w
end

"""
    successor(system::NoisyConstrainedAffineControlDiscreteSystem, x::AbstractVector,
              u::AbstractVector; [check_constraints]=true)

Return the successor state of a `NoisyConstrainedAffineControlDiscreteSystem`.

### Input

- `system`            -- `NoisyConstrainedAffineControlDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `u`                 -- input (it should be any vector type)
- `w`                 -- noise (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state (resp. input)
                         belongs to the state set (resp. input set)

### Output

The result of applying the system to `x`, with input `u` and noise `w`.
"""
function successor(system::NoisyConstrainedAffineControlDiscreteSystem,
                   x::AbstractVector, u::AbstractVector, w::AbstractVector;
                   check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    !_is_conformable_input(system, u) && _argument_error(:u)
    !_is_conformable_noise(system, w) && _argument_error(:w)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
        !_in_inputset(system, u) && _argument_error(:u,:U)
        !_in_noiseset(system, w) && _argument_error(:w,:W)
    end
    return system.A * x + system.B * u + system.c + system.D * w
end

"""
    successor(system::NoisyConstrainedBlackBoxControlDiscreteSystem, x::AbstractVector,
              u::AbstractVector, w::AbstractVector; [check_constraints]=true)

Return the successor state of a `NoisyConstrainedBlackBoxControlDiscreteSystem`.

### Input

- `system`            -- `NoisyConstrainedBlackBoxControlDiscreteSystem`
- `x`                 -- state (it should be any vector type)
- `u`                 -- input (it should be any vector type)
- `w`                 -- noise (it should be any vector type)
- `check_constraints` -- (optional, default: `true`) check if the state belongs to
                         the state set

### Output

The result of applying the system to `x`, with input `u` and noise `w`.
"""
function successor(system::NoisyConstrainedBlackBoxControlDiscreteSystem,
                   x::AbstractVector, u::AbstractVector, w::AbstractVector; check_constraints::Bool=true)
    !_is_conformable_state(system, x) && _argument_error(:x)
    !_is_conformable_input(system, u) && _argument_error(:u)
    !_is_conformable_noise(system, w) && _argument_error(:w)
    if check_constraints
        !_in_stateset(system, x) && _argument_error(:x,:X)
        !_in_inputset(system, u) && _argument_error(:u,:U)
        !_in_noiseset(system, w) && _argument_error(:w,:W)
    end
    return system.f(x, u, w)
end


text:
using LazyArrays, Test
using LazyArrays.SetOperations
import LazyArrays.SetOperations: ∪, ∩, unioned, intersected, setdiffed

@testset "Set operations" begin
    @test SetOperations.Style(Int) == SetOperations.VectorStyle()
    @test SetOperations.Style(Set) == SetOperations.SetStyle()
    @test 1 ∈ unioned(1,2.0)
    @test 2 ∈ unioned(1,2.0)
    @test 3 ∉ unioned(1,2.0)
    @test issubset([1,2],  unioned(1,2.0))

    @test 1 ∪ 2.0 == [1.0,2.0]
    @test [1,2] ∪ Set([2,4]) ==  Set([2,4]) ∪ [1,2] == SetOperations.union(Set([2,4]) , [1,2]) ==
        SetOperations.union([1,2], Set([2,4])) == Set([1,2,4])
    @test Base.union([1,2], Set([2,4])) == Base.:∪([1,2], Set([2,4])) ==
                [1,2,4]

    @test 1 ∈ intersected([1,2], Set([1.0]))
    @test 2 ∉ intersected([1,2], Set([1.0]))
    @test 1 ∉ setdiffed([1,2], Set([1.0]))
    @test 2 ∈ setdiffed([1,2], Set([1.0]))

    @test 1 ∈ intersected([1,2], Set([1.0]), 1f0)
end


text:
function bar1()::Int64
    return 0
end