text:
_add_sum(x, y) = Base.add_sum(x, y)
_add_sum(x, ::Missing) = x
_add_sum(::Missing, x) = x
_add_sum(::Missing, ::Missing) = missing
_mul_prod(x, y) = Base.mul_prod(x, y)
_mul_prod(x, ::Missing) = x
_mul_prod(::Missing, x) = x
_mul_prod(::Missing, ::Missing) = missing
_min_fun(x, y) = min(x, y)
_min_fun(x, ::Missing) = x
_min_fun(::Missing, y) = y
_min_fun(::Missing, ::Missing) = missing
_max_fun(x, y) = max(x, y)
_max_fun(x, ::Missing) = x
_max_fun(::Missing, y) = y
_max_fun(::Missing, ::Missing) = missing
_bool(f) = x->f(x)::Bool

const __NCORES = Threads.nthreads()

struct _Prehashed
    hash::UInt64
end
Base.hash(x::_Prehashed) = x.hash

Base.@propagate_inbounds function _op_for_sum!(x, y, f, lo, hi)
    @simd for i in lo:hi
        x[i] = _add_sum(x[i], f(y[i]))
    end
    x
end

function row_sum(ds::AbstractDataset, f::Function,  cols = names(ds, Union{Missing, Number}); threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    CT = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    T = Core.Compiler.return_type(f, (CT,))
    init0 = _missings(T, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_sum!(x,y, f, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_sum!(x,y, f, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end

row_sum(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true) = row_sum(ds, identity, cols, threads = threads)

Base.@propagate_inbounds function _op_for_prod!(x, y, f, lo, hi)
    @simd for i in lo:hi
        x[i] = _mul_prod(x[i], f(y[i]))
    end
    x
end

function row_prod(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    CT = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    T = Core.Compiler.return_type(f, (CT,))
    init0 = _missings(T, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_prod!(x,y, f, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_prod!(x,y, f, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end

row_prod(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true) = row_prod(ds, identity, cols; threads = threads)

Base.@propagate_inbounds function _op_for_count!(x, y, f, lo, hi)
    @simd for i in lo:hi
        x[i] += f(y[i])
    end
    x
end

function row_count(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    init0 = zeros(Int32, size(ds,1))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_count!(x, y, _bool(f), lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_count!(x, y, _bool(f), 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end

row_count(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true) = row_count(ds, x->true, cols; threads = threads)

_op_bool_add(x::Bool,y::Bool) = x | y ? true : false
Base.@propagate_inbounds function op_for_any!(x, y, f, lo, hi)
    @simd for i in lo:hi
        x[i] = _op_bool_add(x[i], f(y[i]))
    end
    x
end

function row_any(ds::AbstractDataset, f::Union{AbstractVector{<:Function}, Function}, cols = :; threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    init0 = zeros(Bool, size(ds,1))

    multi_f = false
    if f isa AbstractVector
        @assert length(f) == length(colsidx) "number of provided functions must match the number of selected columns"
        multi_f = true
    end

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if multi_f
                mapreduce_index(f, (x, y, func) -> op_for_any!(x, y, _bool(func), lo, hi), view(_columns(ds),colsidx), init0)
            else
                mapreduce(identity, (x,y) -> op_for_any!(x,y, _bool(f), lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        if multi_f
            mapreduce_index(f, (x, y, func) -> op_for_any!(x, y, _bool(func), 1, length(x)), view(_columns(ds),colsidx), init0)
        else
            mapreduce(identity, (x,y) -> op_for_any!(x,y, _bool(f), 1, length(x)), view(_columns(ds),colsidx), init = init0)
        end
    end
    init0
end

row_any(ds::AbstractDataset, cols = :; threads = true) = row_any(ds, isequal(true), cols; threads = threads)

_op_bool_mult(x::Bool,y::Bool) = x & y ? true : false

Base.@propagate_inbounds function op_for_all!(x, y, f, lo, hi)
    @simd for i in lo:hi
        x[i] = _op_bool_mult(x[i], f(y[i]))
    end
    x
end

function row_all(ds::AbstractDataset, f::Union{AbstractVector{<:Function}, Function}, cols = :; threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    init0 = ones(Bool, size(ds,1))

    multi_f = false
    if f isa AbstractVector
        @assert length(f) == length(colsidx) "number of provided functions must match the number of selected columns"
        multi_f = true
    end

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if multi_f
                mapreduce_index(f, (x, y, func) -> op_for_all!(x, y, _bool(func), lo, hi), view(_columns(ds),colsidx), init0)
            else
                mapreduce(identity, (x,y) -> op_for_all!(x,y, _bool(f), lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        if multi_f
            mapreduce_index(f, (x, y, func) -> op_for_all!(x, y, _bool(func), 1, length(x)), view(_columns(ds),colsidx), init0)
        else
            mapreduce(identity, (x,y) -> op_for_all!(x,y, _bool(f), 1, length(x)), view(_columns(ds),colsidx), init = init0)
        end
    end
    init0
end
row_all(ds::AbstractDataset, cols = :; threads = true) = row_all(ds, isequal(true), cols; threads = threads)

# this is a general rule for order of arguments in isequal, isless, findfirst, ...
# if the keyword argument is `with` then eq(y, with)
# if the keyword argument is `item` then eq(item, y)

Base.@propagate_inbounds function _op_for_isequal!(x, y, x1, lo, hi)
    @simd for i in lo:hi
        x[i] &= isequal(y[i], x1[i])
    end
    x
end

function row_isequal(ds::AbstractDataset, cols = :; by::Union{AbstractVector, DatasetColumn, SubDatasetColumn, ColumnIndex, Nothing} = nothing, threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    if !(by isa ColumnIndex) && by !== nothing
        @assert length(by) == nrow(ds) "to compare values of selected columns in each row, the length of the passed vector and the number of rows must match"
    end
    if by isa SubDatasetColumn || by isa DatasetColumn
        x1 = __!(by)
    elseif by isa ColumnIndex
        x1 = _columns(ds)[index(ds)[by]]
    elseif by === nothing
        x1 = _columns(ds)[colsidx[1]]
    else
        x1 = by
    end
    init0 = ones(Bool, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_isequal!(x,y, x1, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_isequal!(x,y, x1, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end

Base.@propagate_inbounds function _op_for_isless!(x, y, vals, rev,lt, lo, hi)
    if !rev
        @simd for i in lo:hi
            x[i] &= lt(y[i], vals[i])
        end
    else
        @simd for i in lo:hi
            x[i] &= lt(vals[i], y[i])
        end
    end
    x
end

function row_isless(ds::AbstractDataset, cols, colselector::Union{AbstractVector, DatasetColumn, SubDatasetColumn, ColumnIndex}; threads = true, rev = false, lt = isless)
    if !(colselector isa ColumnIndex)
        @assert length(colselector) == nrow(ds) "to compare values of selected columns in each row, the length of the passed vector and the number of rows must match"
    end
    colsidx = multiple_getindex(index(ds), cols)
    if colselector isa SubDatasetColumn || colselector isa DatasetColumn
        colselector = __!(colselector)
    elseif colselector isa ColumnIndex
        colselector = _columns(ds)[index(ds)[colselector]]
    end
    init0 = ones(Bool, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_isless!(x, y, colselector, rev, lt, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_isless!(x, y, colselector, rev, lt, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0

end

# TODO probably we should use this approach instead of mapreduce_indexed
Base.@propagate_inbounds function _op_for_findfirst!(x, y, f, idx, missref, lo, hi)
    idx[] += 1
    @simd for i in lo:hi
        x[i] = ifelse(isequal(missref, x[i]) && isequal(true, f(y[i])), idx[], x[i])
    end
    x
end

Base.@propagate_inbounds function _op_for_findfirst!(x, y, item, idx, missref, eq, lo, hi)
    idx[] += 1
    @simd for i in lo:hi
        x[i] = ifelse(isequal(missref, x[i]) && eq(item[i], y[i]), idx[], x[i])
    end
    x
end

Base.@propagate_inbounds function _op_for_findlast!(x, y, f, idx, missref, lo, hi)
    idx[] += 1
    @simd for i in lo:hi
        x[i] = ifelse(isequal(true, f(y[i])), idx[], x[i])
    end
    x
end

Base.@propagate_inbounds function _op_for_findlast!(x, y, item, idx, missref, eq, lo, hi)
    idx[] += 1
    @simd for i in lo:hi
        x[i] = ifelse(eq(item[i], y[i]), idx[], x[i])
    end
    x
end

function row_findfirst(ds::AbstractDataset, f, cols = names(ds, Union{Missing, Number});item::Union{AbstractVector, DatasetColumn, SubDatasetColumn, ColumnIndex, Nothing} = nothing, threads = true, eq = isequal)
    if !(item isa ColumnIndex) && item !== nothing
        @assert length(item) == nrow(ds) "length of item values must be the same as the number of rows"
    elseif item isa SubDatasetColumn || item isa DatasetColumn
        item = __!(item)
    elseif item isa ColumnIndex
        item = _columns(ds)[index(ds)[item]]
    end
    colsidx = multiple_getindex(index(ds), cols)

    colnames_pa = allowmissing(PooledArray(_names(ds)[colsidx]))
    push!(colnames_pa, missing)
    missref = get(colnames_pa.invpool, missing, 0)
    init0 = fill(missref, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        idx = [Ref{Int}(0) for _ in 1:__NCORES]
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if item === nothing
                mapreduce(identity, (x,y) -> _op_for_findfirst!(x, y, f, idx[i], missref, lo, hi), view(_columns(ds),colsidx), init = init0)
            else
                mapreduce(identity, (x,y) -> _op_for_findfirst!(x, y, item, idx[i], missref, eq, lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        idx = Ref{Int64}(0)
        if item === nothing
            mapreduce(identity, (x,y) -> _op_for_findfirst!(x, y, f, idx, missref, 1, length(x)), view(_columns(ds),colsidx), init = init0)
        else
            mapreduce(identity, (x,y) -> _op_for_findfirst!(x, y, item, idx, missref, eq, 1, length(x)), view(_columns(ds),colsidx), init = init0)
        end
    end
    colnames_pa.refs = init0
    colnames_pa
end

function row_findlast(ds::AbstractDataset, f, cols = names(ds, Union{Missing, Number}); item::Union{AbstractVector, DatasetColumn, SubDatasetColumn, ColumnIndex, Nothing} = nothing, threads = true, eq = isequal)
    if !(item isa ColumnIndex) && item !== nothing
        @assert length(item) == nrow(ds) "length of item values must be the same as the number of rows"
    elseif item isa SubDatasetColumn || item isa DatasetColumn
        item = __!(item)
    elseif item isa ColumnIndex
        item = _columns(ds)[index(ds)[item]]
    end
    colsidx = multiple_getindex(index(ds), cols)
    colnames_pa = allowmissing(PooledArray(_names(ds)[colsidx]))
    push!(colnames_pa, missing)
    missref = get(colnames_pa.invpool, missing, 0)
    init0 = fill(missref, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        idx = [Ref{Int}(0) for _ in 1:__NCORES]
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if item === nothing
                mapreduce(identity, (x,y) -> _op_for_findlast!(x, y, f, idx[i], missref, lo, hi), view(_columns(ds),colsidx), init = init0)
            else
                mapreduce(identity, (x,y) -> _op_for_findlast!(x, y, item, idx[i], missref, eq, lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        idx = Ref{Int}(0)
        if item === nothing
            mapreduce(identity, (x,y) -> _op_for_findlast!(x, y, f, idx, missref, 1, length(x)), view(_columns(ds),colsidx), init = init0)
        else
            mapreduce(identity, (x,y) -> _op_for_findlast!(x, y, item, idx, missref, eq, 1, length(x)), view(_columns(ds),colsidx), init = init0)
        end
    end
    colnames_pa.refs = init0
    colnames_pa
end

Base.@propagate_inbounds function _op_for_in!(x, y, x1, eq, lo, hi)
    @simd for i in lo:hi
        !x[i] ? x[i] = eq(x1[i], y[i]) : nothing
    end
    x
end

function row_in(ds::AbstractDataset, collections, items::Union{AbstractVector, DatasetColumn, SubDatasetColumn, ColumnIndex}; threads = true, eq = isequal)
    if !(items isa ColumnIndex)
        @assert length(items) == nrow(ds) "length of item values must be the same as the number of rows"
    end
    colsidx = index(ds)[collections]
    if items isa SubDatasetColumn || items isa DatasetColumn
        items = __!(items)
    elseif items isa ColumnIndex
        items = _columns(ds)[index(ds)[items]]
    end
    init0 = zeros(Bool, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_in!(x, y, items, eq, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_in!(x, y, items, eq, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end

Base.@propagate_inbounds function _op_for_select!(x, y, colselector, dsnames, idx, lo, hi)
    idx[] += 1
    @simd for i in lo:hi
        if isequal(colselector[i], dsnames[idx[]])
            x[i] = y[i]
        end
    end
    x
end

function row_select(ds::AbstractDataset, cols, colselector::Union{AbstractVector, DatasetColumn, SubDatasetColumn, ColumnIndex}; threads = true)
    if !(colselector isa ColumnIndex)
        @assert length(colselector) == nrow(ds) "to pick values of selected columns in each row, the length of the column names and the number of rows must match, i.e. the length of the vector passed as `by` must be $(nrow(ds))."
    end
    colsidx = multiple_getindex(index(ds), cols)
    CT = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    if colselector isa SubDatasetColumn || colselector isa DatasetColumn
        colselector = __!(colselector)
    end
    if colselector isa ColumnIndex
        colselector = _columns(ds)[index(ds)[colselector]]
    end
    if eltype(colselector) <: Union{Missing, Symbol}
        nnames = Symbol.(names(ds, colsidx))
    elseif eltype(colselector) <: Union{Missing, AbstractString}
        nnames = names(ds, colsidx)
    elseif eltype(colselector) <: Union{Missing, Integer}
        nnames = 1:length(colsidx)
    end
    init0 = _missings(CT, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        idx = [Ref{Int}(0) for _ in 1:__NCORES]
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_select!(x, y, colselector, nnames, idx[i], lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        idx = Ref{Int}(0)
        mapreduce(identity, (x,y) -> _op_for_select!(x, y, colselector, nnames, idx, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end

Base.@propagate_inbounds function _op_for_fill!(x, y, f, lo, hi)
    @simd for i in lo:hi
       y[i] = ifelse(f(y[i]), x[i], y[i])
    end
    x
end

Base.@propagate_inbounds function _op_for_fill_roll!(x, y, f, lo, hi)
    @simd for i in lo:hi
       y[i] = ifelse(f(y[i]), x[i], y[i])
    end
    y
end

# rolling = true, means fill a column an use its value for next filling
function row_fill!(ds::AbstractDataset, cols, val::Union{AbstractVector, DatasetColumn, SubDatasetColumn, ColumnIndex}; f = ismissing, threads = true, rolling = false)
    if !(val isa ColumnIndex)
        @assert length(val) == nrow(ds) "to fill values in each row, the length of passed values and the number of rows must match."
    end
    colsidx = multiple_getindex(index(ds), cols)
    if val isa SubDatasetColumn || val isa DatasetColumn
        val = __!(val)
    end
    if val isa ColumnIndex
        val = _columns(ds)[index(ds)[val]]
    end
    init0 = val

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if rolling
                mapreduce(identity, (x,y) -> _op_for_fill_roll!(x, y, f, lo, hi), view(_columns(ds),colsidx), init = init0)
            else
                mapreduce(identity, (x,y) -> _op_for_fill!(x, y, f, lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        if rolling
            mapreduce(identity, (x,y) -> _op_for_fill_roll!(x, y, f, 1, length(x)), view(_columns(ds),colsidx), init = init0)
        else
            mapreduce(identity, (x,y) -> _op_for_fill!(x, y, f, 1, length(x)), view(_columns(ds),colsidx), init = init0)
        end
    end
    any(index(parent(ds)).sortedcols .∈ Ref(IMD.parentcols.(Ref(IMD.index(ds)), cols))) && _reset_grouping_info!(parent(ds))
    _modified(_attributes(parent(ds)))
    ds
end

Base.@propagate_inbounds function _op_for_coalesce!(x, y, lo, hi)
    if all(!ismissing, view(x, lo:hi))
        x
    else
        @simd for i in lo:hi
            x[i] = ifelse(ismissing(x[i]), y[i], x[i])
        end
    end
    x
end

function row_coalesce(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    CT = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))

    init0 = _missings(CT, size(ds,1))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_coalesce!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_coalesce!(x, y, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end

function row_mean(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); threads = true)
    row_sum(ds, f, cols; threads = threads) ./ row_count(ds, x -> !ismissing(x), cols; threads = threads)
end
row_mean(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true) = row_mean(ds, identity, cols; threads = threads)

Base.@propagate_inbounds function _op_for_min!(x, y, f, lo, hi)
    @simd for i in lo:hi
        x[i] = _min_fun(x[i], f(y[i]))
    end
    x
end

Base.@propagate_inbounds function _op_for_max!(x, y, f, lo, hi)
    @simd for i in lo:hi
        x[i] = _max_fun(x[i], f(y[i]))
    end
    x
end

function row_minimum(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    CT = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    T = Core.Compiler.return_type(f, (CT,))
    init0 = _missings(T, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_min!(x, y, f, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_min!(x, y, f, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end
row_minimum(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true) = row_minimum(ds, identity, cols; threads = threads)

function row_maximum(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    CT = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    T = Core.Compiler.return_type(f, (CT,))
    init0 = _missings(T, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_max!(x, y, f, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_max!(x, y, f, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end
row_maximum(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true) = row_maximum(ds, identity, cols; threads = threads)

Base.@propagate_inbounds function _op_for_argminmax!(x, y, f, vals, idx, missref, lo, hi)
    idx[] += 1
    @simd for i in lo:hi
        if !ismissing(vals[i]) && isequal(vals[i], f(y[i])) && isequal(x[i], missref)
            x[i] = idx[]
        end
    end
    x
end

function row_argmin(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    minvals = row_minimum(ds, f, cols)
    colnames_pa = allowmissing(PooledArray(_names(ds)[colsidx]))
    push!(colnames_pa, missing)
    missref = get(colnames_pa.invpool, missing, missing)
    init0 = fill(missref, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        idx = [Ref{Int}(0) for _ in 1:__NCORES]
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_argminmax!(x, y, f, minvals, idx[i], missref, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        idx = Ref{Int}(0)
        mapreduce(identity, (x,y) -> _op_for_argminmax!(x, y, f, minvals, idx, missref, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end

    colnames_pa.refs = init0
    colnames_pa
end
row_argmin(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true) = row_argmin(ds, identity, cols, threads = threads)

function row_argmax(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    maxvals = row_maximum(ds, f, cols)
    colnames_pa = allowmissing(PooledArray(_names(ds)[colsidx]))
    push!(colnames_pa, missing)
    missref = get(colnames_pa.invpool, missing, missing)
    init0 = fill(missref, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        idx = [Ref{Int}(0) for _ in 1:__NCORES]
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_argminmax!(x, y, f, maxvals, idx[i], missref, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        idx = Ref{Int}(0)
        mapreduce(identity, (x,y) -> _op_for_argminmax!(x, y, f, maxvals, idx, missref, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end

    colnames_pa.refs = init0
    colnames_pa
end
row_argmax(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); threads = true) = row_argmax(ds, identity, cols, threads = threads)

# TODO better function for the first component of operator
function _row_wise_var(ss, sval, n, dof, T)
    res = Vector{T}(undef, length(ss))
    for i in 1:length(ss)
        if n[i] == 0
            res[i] = missing
        elseif n[i] == 1
            res[i] = zero(T)
        else
            res[i] = ss[i]/n[i] - (sval[i]/n[i])*(sval[i]/n[i])
            if dof
                res[i] = (n[i] * res[i])/(n[i]-1)
            end
        end
    end
    res
end

# TODO needs type stability
# TODO need abs2 for calculations
function row_var(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); dof = true, threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    CT = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    T = Core.Compiler.return_type(f, (CT,))
    _sq_(x) = x^2
    ss = row_sum(ds, _sq_ ∘ f, cols; threads = threads)
    sval = row_sum(ds, f, cols; threads = threads)
    n = row_count(ds, x -> !ismissing(x), cols; threads = threads)
    T2 = Core.Compiler.return_type(/, (eltype(ss), eltype(n)))
    res = Vector{Union{Missing, T2}}(undef, length(ss))
    res .= ss ./ n .- (sval ./ n) .^ 2
    if dof
        res .= (n .* res) ./ (n .- 1)
        res .= ifelse.(n .== 1, missing, res)
    end
    res
    # _row_wise_var(ss, sval, n, dof, T)
end
row_var(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); dof = true, threads = true) = row_var(ds, identity, cols, dof = dof, threads = threads)

function row_std(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); dof = true, threads = true)
    sqrt.(row_var(ds, f, cols, dof = dof, threads = threads))
end
row_std(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); dof = true, threads = true) = row_std(ds, identity, cols, dof = dof, threads = threads)

Base.@propagate_inbounds function _op_for_cumsum_skip!(x, y, lo, hi)
    @simd for i in lo:hi
        x[i] = _add_sum(x[i], y[i])
    end
    @simd for i in lo:hi
        y[i] = ifelse(ismissing(y[i]), missing, x[i])
    end
    x
end

Base.@propagate_inbounds function _op_for_cumsum_ignore!(x, y, lo, hi)
    @simd for i in lo:hi
        y[i] = _add_sum(x[i], y[i])
    end
    y
end

function row_cumsum!(ds::Dataset, cols = names(ds, Union{Missing, Number}); missings = :ignore, threads = true)
    !(missings in (:ignore, :skip)) && throw(ArgumentError("`missings` can be either `:ignore` or `:skip`"))
    colsidx = index(ds)[cols]
    T = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    if T <: Union{Missing, INTEGERS}
        T <: Union{Missing, Base.SmallSigned}
        T = T <: Union{Missing, Base.SmallSigned, Bool} ? Union{Int, Missing} : T
        T = T <: Union{Missing, Base.SmallUnsigned} ?  Union{Missing, UInt} : T
    end
    for i in colsidx
        if eltype(ds[!, i]) >: Missing
            _columns(ds)[i] = convert(Vector{Union{Missing, T}}, _columns(ds)[i])
        else
            _columns(ds)[i] = convert(Vector{T}, _columns(ds)[i])
        end
    end
    init0 = _missings(T, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if missings == :ignore
                mapreduce(identity, (x,y) -> _op_for_cumsum_ignore!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
            else
                mapreduce(identity, (x,y) -> _op_for_cumsum_skip!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        if missings == :ignore
            mapreduce(identity, (x,y) -> _op_for_cumsum_ignore!(x, y, 1, length(init0)), view(_columns(ds),colsidx), init = init0)
        else
            mapreduce(identity, (x,y) -> _op_for_cumsum_skip!(x, y, 1, length(init0)), view(_columns(ds),colsidx), init = init0)
        end
    end

    removeformat!(ds, cols)
    any(index(ds).sortedcols .∈ Ref(colsidx)) && _reset_grouping_info!(ds)
    _modified(_attributes(ds))
    ds
end

function row_cumsum(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); missings = :ignore, threads = true)
    dscopy = copy(ds)
    row_cumsum!(dscopy, cols, missings = missings, threads = threads)
    dscopy
end

Base.@propagate_inbounds function _op_for_cumprod_skip!(x, y, lo, hi)
    @simd for i in lo:hi
        x[i] = _mul_prod(x[i], y[i])
    end
    @simd for i in lo:hi
        y[i] = ifelse(ismissing(y[i]), missing, x[i])
    end
    x
end

Base.@propagate_inbounds function _op_for_cumprod_ignore!(x, y, lo, hi)
    @simd for i in lo:hi
        y[i] = _mul_prod(x[i], y[i])
    end
    y
end

function row_cumprod!(ds::Dataset, cols = names(ds, Union{Missing, Number}); missings = :ignore, threads = true)
    !(missings in (:ignore, :skip)) && throw(ArgumentError("`missings` can be either `:ignore` or `:skip`"))
    colsidx = index(ds)[cols]
    T = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    for i in colsidx
        if eltype(ds[!, i]) >: Missing
            _columns(ds)[i] = convert(Vector{Union{Missing, T}}, _columns(ds)[i])
        else
            _columns(ds)[i] = convert(Vector{T}, _columns(ds)[i])
        end
    end
    init0 = _missings(T, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if missings == :ignore
                mapreduce(identity, (x,y) -> _op_for_cumprod_ignore!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
            else
                mapreduce(identity, (x,y) -> _op_for_cumprod_skip!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        if missings == :ignore
            mapreduce(identity, (x,y) -> _op_for_cumprod_ignore!(x, y, 1, length(init0)), view(_columns(ds),colsidx), init = init0)
        else
            mapreduce(identity, (x,y) -> _op_for_cumprod_skip!(x, y, 1, length(init0)), view(_columns(ds),colsidx), init = init0)
        end
    end

    removeformat!(ds, cols)
    any(index(ds).sortedcols .∈ Ref(colsidx)) && _reset_grouping_info!(ds)
    _modified(_attributes(ds))
    ds
end

function row_cumprod(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); missings = :ignore, threads = true)
    dscopy = copy(ds)
    row_cumprod!(dscopy, cols; missings = missings, threads = threads)
    dscopy
end

Base.@propagate_inbounds function _op_for_cummin_skip!(x, y, lo, hi)
    @simd for i in lo:hi
        x[i] = _min_fun(x[i], y[i])
    end
    @simd for i in lo:hi
        y[i] = ifelse(ismissing(y[i]), missing, x[i])
    end
    x
end

Base.@propagate_inbounds function _op_for_cummin_ignore!(x, y, lo, hi)
    @simd for i in lo:hi
        y[i] = _min_fun(x[i], y[i])
    end
    y
end

function row_cummin!(ds::Dataset, cols = names(ds, Union{Missing, Number}); missings = :ignore, threads = true)
    !(missings in (:ignore, :skip)) && throw(ArgumentError("`missings` can be either `:ignore` or `:skip`"))
    colsidx = index(ds)[cols]
    T = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    for i in colsidx
        if eltype(ds[!, i]) >: Missing
            _columns(ds)[i] = convert(Vector{Union{Missing, T}}, _columns(ds)[i])
        else
            _columns(ds)[i] = convert(Vector{T}, _columns(ds)[i])
        end
    end
    init0 = _missings(T, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if missings == :ignore
                mapreduce(identity, (x,y) -> _op_for_cummin_ignore!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
            else
                mapreduce(identity, (x,y) -> _op_for_cummin_skip!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        if missings == :ignore
            mapreduce(identity, (x,y) -> _op_for_cummin_ignore!(x, y, 1, length(init0)), view(_columns(ds),colsidx), init = init0)
        else
            mapreduce(identity, (x,y) -> _op_for_cummin_skip!(x, y, 1, length(init0)), view(_columns(ds),colsidx), init = init0)
        end
    end

    removeformat!(ds, cols)
    any(index(ds).sortedcols .∈ Ref(colsidx)) && _reset_grouping_info!(ds)
    _modified(_attributes(ds))
    ds
end
# row_cumsum!(ds::AbstractDataset, cols = names(ds, Union{Missing, Number})) = row_cumsum!(identity, ds, cols)

function row_cummin(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); missings = :ignore, threads = true)
    dscopy = copy(ds)
    row_cummin!(dscopy, cols, missings = missings, threads = threads)
    dscopy
end

Base.@propagate_inbounds function _op_for_cummax_skip!(x, y, lo, hi)
    @simd for i in lo:hi
        x[i] = _max_fun(x[i], y[i])
    end
    @simd for i in lo:hi
        y[i] = ifelse(ismissing(y[i]), missing, x[i])
    end
    x
end

Base.@propagate_inbounds function _op_for_cummax_ignore!(x, y, lo, hi)
    @simd for i in lo:hi
        y[i] = _max_fun(x[i], y[i])
    end
    y
end

function row_cummax!(ds::Dataset, cols = names(ds, Union{Missing, Number}); missings = :ignore, threads = true)
    !(missings in (:ignore, :skip)) && throw(ArgumentError("`missings` can be either `:ignore` or `:skip`"))
    colsidx = index(ds)[cols]
    T = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    for i in colsidx
        if eltype(ds[!, i]) >: Missing
            _columns(ds)[i] = convert(Vector{Union{Missing, T}}, _columns(ds)[i])
        else
            _columns(ds)[i] = convert(Vector{T}, _columns(ds)[i])
        end
    end
    init0 = _missings(T, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if missings == :ignore
                mapreduce(identity, (x,y) -> _op_for_cummax_ignore!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
            else
                mapreduce(identity, (x,y) -> _op_for_cummax_skip!(x, y, lo, hi), view(_columns(ds),colsidx), init = init0)
            end
        end
    else
        if missings == :ignore
            mapreduce(identity, (x,y) -> _op_for_cummax_ignore!(x, y, 1, length(init0)), view(_columns(ds),colsidx), init = init0)
        else
            mapreduce(identity, (x,y) -> _op_for_cummax_skip!(x, y, 1, length(init0)), view(_columns(ds),colsidx), init = init0)
        end
    end

    removeformat!(ds, cols)
    any(index(ds).sortedcols .∈ Ref(colsidx)) && _reset_grouping_info!(ds)
    _modified(_attributes(ds))
    ds
end
# row_cumsum!(ds::AbstractDataset, cols = names(ds, Union{Missing, Number})) = row_cumsum!(identity, ds, cols)

function row_cummax(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); missings = :ignore, threads = true)
    dscopy = copy(ds)
    row_cummax!(dscopy, cols, missings = missings, threads = threads)
    dscopy
end

function row_stdze!(ds::Dataset , cols = names(ds, Union{Missing, Number}))
    meandata = row_mean(ds, cols)
    stddata = row_std(ds, cols)
    _stdze_fun(x) = ifelse.(isequal.(stddata, 0), missing, (x .- meandata) ./ stddata)
    colsidx = index(ds)[cols]

    for i in 1:length(colsidx)
        _columns(ds)[colsidx[i]] = _stdze_fun(_columns(ds)[colsidx[i]])
    end
    removeformat!(ds, cols)
    any(index(ds).sortedcols .∈ Ref(colsidx)) && _reset_grouping_info!(ds)
    _modified(_attributes(ds))
end

function row_stdze(ds::AbstractDataset , cols = names(ds, Union{Missing, Number}))
    dscopy = copy(ds)
    row_stdze!(dscopy, cols)
    dscopy
end

function row_sort!(ds::Dataset, cols = names(ds, Union{Missing, Number}); kwargs...)
    colsidx = index(ds)[cols]
    T = mapreduce(eltype, promote_type, eachcol(ds)[colsidx])
    m = Matrix{T}(ds[!, colsidx])
    sort!(m; dims = 2, kwargs...)
    for i in 1:length(colsidx)
        _columns(ds)[colsidx[i]] = m[:, i]
    end
    removeformat!(ds, cols)
    any(index(ds).sortedcols .∈ Ref(colsidx)) && _reset_grouping_info!(ds)
    _modified(_attributes(ds))
    ds
end

function row_sort(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); kwargs...)
    dscopy = copy(ds)
    row_sort!(dscopy, cols; kwargs...)
    dscopy
end

Base.@propagate_inbounds function _op_for_issorted!(x, y, res, lt, lo, hi)
    @simd for i in lo:hi
        res[i] &= !lt(y[i], x[i])
    end
    y
end
Base.@propagate_inbounds function _op_for_issorted_rev!(x, y, res, lt, lo, hi)
    @simd for i in lo:hi
        res[i] &= !lt(x[i], y[i])
    end
    y
end

function row_issorted(ds::AbstractDataset, cols; rev = false, lt = isless, threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    init0 = ones(Bool, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            if rev
                mapreduce(identity, (x,y) -> _op_for_issorted_rev!(x, y, init0, lt, lo, hi), view(_columns(ds),colsidx))
            else
                mapreduce(identity, (x,y) -> _op_for_issorted!(x, y, init0, lt, lo, hi), view(_columns(ds),colsidx))
            end
        end
    else
        if rev
            mapreduce(identity, (x,y) -> _op_for_issorted_rev!(x, y, init0, lt, 1, length(init0)), view(_columns(ds),colsidx))
        else
            mapreduce(identity, (x,y) -> _op_for_issorted!(x, y, init0, lt, 1, length(init0)), view(_columns(ds),colsidx))
        end
    end
    init0
end

# TODO is it possible to have a faster row_count_unique??
function _fill_prehashed!(prehashed, y, f, n, j)
    @views copy!(prehashed[:, j] , _Prehashed.(hash.(f.(y))))
end

function _fill_dict_and_add!(init0, dict, prehashed, n, p)
    for i in 1:n
        for j in 1:p
            if !haskey(dict, prehashed[i, j])
                get!(dict, prehashed[i, j], nothing)
                init0[i] += 1
            end
        end
        empty!(dict)
    end
end

function row_nunique(ds::AbstractDataset, f::Function, cols = names(ds, Union{Missing, Number}); count_missing = true)
    colsidx = multiple_getindex(index(ds), cols)
    prehashed = Matrix{_Prehashed}(undef, size(ds,1), length(colsidx))
    allcols = view(_columns(ds),colsidx)

    for j in 1:size(prehashed,2)
        _fill_prehashed!(prehashed, allcols[j], f, size(ds,1), j)
    end

    init0 = zeros(Int32, size(ds,1))
    dict = Dict{_Prehashed, Nothing}()
    _fill_dict_and_add!(init0, dict, prehashed, size(ds,1), length(colsidx))
    if count_missing
        return init0
    else
        return init0 .- row_any(ds, ismissing, cols)
    end
end
row_nunique(ds::AbstractDataset, cols = names(ds, Union{Missing, Number}); count_missing = true) = row_nunique(ds, identity, cols; count_missing = count_missing)

Base.@propagate_inbounds function _op_for_hash!(x, y, f, lo, hi)
    @simd for i in lo:hi
        x[i] = hash(f(y[i]), x[i])
    end
    x
end

function row_hash(ds::AbstractDataset, f::Function, cols = :; threads = true)
    colsidx = multiple_getindex(index(ds), cols)
    init0 = zeros(UInt64, nrow(ds))

    if threads
        cz = div(length(init0), __NCORES)
        Threads.@threads for i in 1:__NCORES
            lo = (i-1)*cz+1
            i == __NCORES ? hi = length(init0) : hi = i*cz
            mapreduce(identity, (x,y) -> _op_for_hash!(x, y, f, lo, hi), view(_columns(ds),colsidx), init = init0)
        end
    else
        mapreduce(identity, (x,y) -> _op_for_hash!(x, y, f, 1, length(x)), view(_columns(ds),colsidx), init = init0)
    end
    init0
end
row_hash(ds::AbstractDataset, cols = :; threads = true) = row_hash(ds, identity, cols; threads = threads)

function _convert_uint8_to_string!(res, init0, curr_pos, ds, threads)
    if threads
        Threads.@threads for i in 1:nrow(ds)
            res[i] = String(view(init0, 1:curr_pos[i]-2, i))
        end
    else
        for i in 1:nrow(ds)
            res[i] = String(view(init0, 1:curr_pos[i]-2, i))
        end
    end
end
function _add_last_for_join!(init0, curr_pos, ds, last_uint, last_len, threads)
    if threads
        Threads.@threads for i in 1:nrow(ds)
            init0[curr_pos[i]-1:curr_pos[i]+last_len-2, i] = last_uint
            curr_pos[i] += last_len-1
        end
    else
        for i in 1:nrow(ds)
            init0[curr_pos[i]-1:curr_pos[i]+last_len-2, i] = last_uint
            curr_pos[i] += last_len-1
        end
    end
end

function row_join(ds::AbstractDataset, cols = :; threads = true, delim::AbstractString = ",", last::AbstractString = "")
    colsidx = multiple_getindex(index(ds), cols)

    max_line_size = maximum(byrow(ds, sum, colsidx, by = y->length(__STRING(y)), threads = threads))
    max_line_size += length(delim)*(length(colsidx)) + length(last)+1
    init0 = Matrix{UInt8}(undef, max_line_size, nrow(ds))
    curr_pos = ones(Int, nrow(ds))

    delimiter = Base.CodeUnits(delim)
    row_join!(init0, curr_pos, ds, repeat([identity], length(colsidx)-1), view(colsidx, 1:length(colsidx)-1); delim = delimiter, quotechar = nothing, threads = threads)
    if length(last)>0
        delimiter = Base.CodeUnits(last)
    end
    if length(colsidx) > 1
        _add_last_for_join!(init0, curr_pos, ds, delimiter, length(delimiter), threads)
    end
    row_join!(init0, curr_pos, ds, [identity], colsidx[length(colsidx)]; delim = delimiter, quotechar = nothing, threads = threads)

    res = Vector{Union{String, Missing}}(undef, nrow(ds))
    _convert_uint8_to_string!(res, init0, curr_pos, ds, threads)
    res

end

function _fill_col!(inmat, column, rows, j)
    for i in 1:length(rows)
        inmat[j, i] = column[rows[i]]
    end
end
function _fill_matrix!(inmat, all_data, rows, cols)
    for j in 1:length(cols)
        _fill_col!(inmat, all_data[j], rows, j)
    end
end

function row_generic(ds::AbstractDataset, f::Function, cols::MultiColumnIndex)
    colsidx = multiple_getindex(index(ds), cols)
    if length(colsidx) == 2
        try
            allowmissing(f.(_columns(ds)[colsidx[1]], _columns(ds)[colsidx[2]]))
        catch e
            if e isa MethodError
                _row_generic(ds, f, colsidx)
            else
                rethrow(e)
            end
        end
    else
        _row_generic(ds, f, colsidx)
    end
end
function _row_generic(ds::AbstractDataset, f::Function, colsidx)
    T = mapreduce(eltype, promote_type, view(_columns(ds),colsidx))
    inmat = Matrix{T}(undef, length(colsidx), min(1000, nrow(ds)))

    all_data = view(_columns(ds), colsidx)
    _fill_matrix!(inmat, all_data, 1:min(1000, nrow(ds)), colsidx)
    res_temp = allowmissing(f.(eachcol(inmat)))
    if !(typeof(res_temp) <:  AbstractVector)
        throw(ArgumentError("output of `f` must be a vector"))
    end

    # if length(res_temp[1]) > 1
    #     throw(ArgumentError("The matrix output is not supported"))
    #     res = similar(res_temp, nrow(ds), size(res_temp,2))
    # elseif length(res_temp[1]) == 1
    res = similar(res_temp, nrow(ds))
    # else
        # throw(ArgumentError("the result cannot be with zero dimension"))

    if nrow(ds)>1000
        if size(res, 2) == 1
            view(res, 1:1000) .= res_temp
            _row_generic_vec!(res, ds, f, colsidx, Val(T))
        else
            view(res, 1:1000, :) .= res_temp
            _row_generic_mat!(res, ds, f, colsidx)
        end
    else
        return res_temp
    end
    return res
end

function _row_generic_vec!_barrier(res, f, inmat, all_data, chunck, colsidx)

    for i in 1:chunck
        t_st = i*1000 + 1
        i == chunck ? t_en = length(res) : t_en = (i+1)*1000
        _fill_matrix!(inmat, all_data, t_st:t_en, colsidx)
        for k in t_st:t_en
            res[k] = f(view(inmat, :, k - t_st + 1))
        end
    end
end

function _row_generic_vec!(res, ds, f, colsidx, ::Val{T}) where T
    all_data = view(_columns(ds), colsidx)
    chunck = div(length(res) - 1000, 1000)
    max_cz = length(res) - 1000 - (chunck - 1)* 1000
    inmat = Matrix{T}(undef, length(colsidx), max_cz)
    # make sure that the variable inside the loop are not the same as the out of scope one
    # for i in 1:chunck
    #     t_st = i*1000 + 1
    #     i == chunck ? t_en = length(res) : t_en = (i+1)*1000
    #     _fill_matrix!(inmat, all_data, t_st:t_en, colsidx)
    #     for k in t_st:t_en
    #         res[k] = f(view(inmat, k - t_st + 1, :))
    #     end
    # end
    _row_generic_vec!_barrier(res, f, inmat, all_data, chunck, colsidx)
end


text:
#=
Conic reformulation approach for the shortest path problem.
=#

using Parameters, Gurobi, JuMP, Clp, Gadfly, SparseArrays, LinearAlgebra

## Type for reformulation algorithm parameters with default values
@with_kw struct reformulation_path_parms
    lambda_max::Union{Missing, Float64} = missing
    lambda_min_ratio::Float64 = 0.0001
    num_lambda::Int = 100
    solver::Symbol = :Gurobi
    gurobiEnv::Gurobi.Env = Gurobi.Env()
    regularization::Symbol = :ridge
    regularize_first_column_B::Bool = false
    upper_bound_B_present::Bool = false
    upper_bound_B::Float64 = 10.0^6
    po_loss_function::Symbol = :leastSquares
    huber_delta::Float64 = 0.0
    verbose::Bool = false
    algorithm_type::Symbol = :fake_algorithm
end

## Type for specifying a shortest path graph instance
@with_kw struct shortest_path_graph
    sources::Array{Int,1}
    destinations::Array{Int,1}
    start_node::Int
    end_node::Int
    acyclic::Bool = false
end

"""
    sp_reformulation_path_jump(X, c, sp_graph::shortest_path_graph;
        sp_oracle = missing, reform_alg_parms::reformulation_path_parms = reformulation_path_parms())
Solves the empirical risk SPO+ problem for the specific case of the shortest path problem
with the reformulation approach using JuMP. Solves over a grid of `num_lamba` regularization
parameter values that are evenly spaced on a log scale between `lambda_min_ratio*lambda_max` and `lambda_max`.
Both Ridge and Lasso are implemented. Returns `B_soln_list, lambdas` where `B_soln_list` is a list of
lienar models such that `B_soln_list[i]` is the solution with parameter `lambdas[i]`.
The default values for the algorithm parameters `reform_alg_parms` are (explanations are given below):
    lambda_max::Number = -1
    lambda_min_ratio::Float64 = 0.0001
    num_lambda::Int = 100
    solver::Symbol = :Gurobi
    regularization::Symbol = :ridge
# Arguments
- `sources`:  list of sources for each edge, length of this vector is number of edges
- `destinations`:  list of destinations of each edge, length of this vector is number of edges
- `start_node`:  start node for the shortest path problem
- `end_node`:  end node for the shortest path problem
- `X`: p x n training set feature matrix
- `c`: d x n training set matrix of cost vectors
- `lambda_max`:  largest value of lambda, default is `(d/n)*(norm(X)^2)`
- `lambda_min_ratio`:  lambda_min is `lambda_max*lambda_min_ratio`
- `num_lambda`:  number of lambda values, evenly spaced on log scale between [lambda_min, lambda_max]
- `solver`:  either `:Gurobi` or `:Clp`
- `sp_oracle`:  optionally provide the optimization oracle corresponding to `sp_graph`
- `gurobiEnv` is a Gurobi enviornment which can be passed in (helpful to minimize the number of license creations)
- `regularization`:  either `:ridge` or `:lasso`
"""
function sp_reformulation_path_jump(X, c, sp_graph::shortest_path_graph;
    sp_oracle = missing, reform_alg_parms::reformulation_path_parms = reformulation_path_parms())

    # Unpack parms
    @unpack lambda_max,
        lambda_min_ratio,
        num_lambda,
        solver,
        gurobiEnv,
        regularization,
        regularize_first_column_B,
        upper_bound_B_present,
        upper_bound_B,
        verbose,
        algorithm_type = reform_alg_parms

    @unpack sources, destinations, start_node, end_node, acyclic = sp_graph

    # Dimension check
    (p, n) = size(X)
    (d, n2) = size(c)
    if n != n2
        error("Dimensions of the input are mismatched.")
    end

    # Process graph
    nodes = unique(union(sources, destinations))
    n_nodes = length(nodes)
    n_edges = length(sources)
    if n_edges != d
        error("Dimensions of the input are mismatched.")
    end

    # Hard code sparse incidence matrix!
    I_vec = [sources; destinations]
    J_vec = [collect(1:n_edges); collect(1:n_edges)]
    V_vec = [-ones(n_edges); ones(n_edges)]
    A_mat = sparse(I_vec, J_vec, V_vec)
    A_mat_trans = transpose(A_mat)

    # Set up RHS
    b_vec = sparse(zeros(n_nodes))
    b_vec[start_node] = -1
    b_vec[end_node] = 1

    # Remove sparse
    A_mat = round.(Int, Matrix(A_mat))
    b_vec = round.(Int, Vector(b_vec))

    # Get oracle and w^*
    if ismissing(sp_oracle)
        sp_oracle = sp_flow_jump_setup(sources, destinations, start_node, end_node; solver = solver)
    end
    (z_star_data, w_star_data) = oracle_dataset(c, sp_oracle)

    # Start creating dual JuMP model
    if solver == :Gurobi
        mod = Model(with_optimizer(Gurobi.Optimizer, gurobiEnv))
    elseif solver == :Clp
        mod = Model(with_optimizer(Clp.Optimizer))
    else
        error("Not a valid solver: either :Clp or :Gurobi")
    end

    # Add variables
    @variable(mod, p_var[1:n_nodes, 1:n])
    if !acyclic
        @variable(mod, s_var[1:d, 1:n] >= 0)
    end
    if upper_bound_B_present
        @variable(mod, -upper_bound_B <= B_var[1:d, 1:p] <= upper_bound_B)
    else
        @variable(mod, B_var[1:d, 1:p])
    end

    # Add constraints
    for i = 1:n
        if acyclic
            @constraint(mod, -A_mat_trans*p_var[:,i] .>= c[:,i] - 2*B_var*X[:,i])
        else
            @constraint(mod, s_var[:,i] - A_mat_trans*p_var[:,i] .>= c[:,i] - 2*B_var*X[:,i])
        end
    end

    # Add objective part without regularization
    if regularization == :ridge
        obj_expr_noreg = QuadExpr()
    elseif regularization == :lasso
        obj_expr_noreg = AffExpr()
    else
        error("enter valid regularization: :ridge or :lasso")
    end
    for i = 1:n
        if acyclic
            cur_expr = -dot(b_vec, p_var[:,i]) + 2*(w_star_data[:,i]')*B_var*X[:,i] - z_star_data[i]
        else
            cur_expr = -dot(b_vec, p_var[:,i]) + dot(ones(d), s_var[:,i]) + 2*(w_star_data[:,i]')*B_var*X[:,i] - z_star_data[i]
        end
        add_to_expression!(obj_expr_noreg, cur_expr)
        #obj_expr_noreg = obj_expr_noreg + cur_expr
    end

    ### Add regularization part and solve for path
    if ismissing(lambda_max)
        lambda_max = (d/n)*(norm(X)^2)
    end

    # Construct lambda sequence
    if num_lambda == 1 && lambda_max == 0
        lambdas = [0.0]
    else
        lambda_min = lambda_max*lambda_min_ratio
        #println("Lambda min is $lambda_min")
        log_lambdas = collect(range(log(lambda_min), stop = log(lambda_max), length = num_lambda))
        lambdas = exp.(log_lambdas)
    end

    # Add theta variables for Lasso
    if regularization == :lasso
        @variable(mod, theta_var[1:d, 1:p])
        @constraint(mod, theta_var .>= B_var)
        @constraint(mod, theta_var .>= -B_var)
    end

    # Make list of B models and solve path
    B_soln_list = Matrix{Float64}[]
    for i = 1:num_lambda
        cur_lambda = lambdas[i]
        if verbose
            println("Trying lambda = $cur_lambda")
        end

        if regularization == :ridge && regularize_first_column_B
            obj_expr_full = obj_expr_noreg + n*(cur_lambda/2)*dot(B_var, B_var)
        elseif regularization == :ridge && !regularize_first_column_B
            obj_expr_full = obj_expr_noreg + n*(cur_lambda/2)*dot(B_var[:, 2:p], B_var[:, 2:p])
        elseif regularization == :lasso && regularize_first_column_B
            obj_expr_full = obj_expr_noreg + n*cur_lambda*sum(theta_var)
        elseif regularization == :lasso && !regularize_first_column_B
            obj_expr_full = obj_expr_noreg + n*cur_lambda*sum(theta_var[:, 2:p])
        else
            error("enter valid regularization: :ridge or :lasso")
        end
        @objective(mod, Min, obj_expr_full)

        # solve
        optimize!(mod)
        #z_ast = objective_value(mod)
        mod_status = termination_status(mod)
        B_ast = value.(B_var)

        if any(isnan, B_ast)
            # set B_ast to zero
            B_ast = zeros(d, p)

            # print reason
            println("We got NaNs in B_ast and the reason is:  $mod_status. The algoritim is $algorithm_type.")
        elseif Int(mod_status) != 1
            println("There are no NaNs in B_ast, but the model status is:  $mod_status. The algoritim is $algorithm_type.")
        end

        # if Int(mod_status) != 1
        #     println("Status of the model is: $mod_status")
        #     println("Current lambda is:  $cur_lambda")
        #
        #     B_val = value.(B_var)
        #     p_val = value.(p_var)
        #     s_val = value.(s_var)
        #     theta_val = value.(theta_var)
        #     println("Value of B_var is:  $B_val")
        #     println("Value of p_var is:  $p_val")
        #     println("Value of s_var is:  $s_val")
        #     println("Value of theta_var is:  $theta_val")
        # end

        push!(B_soln_list, B_ast)
    end

    return B_soln_list, lambdas
end

"""
    leastSquares_path_jump(X, c;
        reform_alg_parms = reformulation_path_parms())
Solves the empirical risk least squares problem using JuMP.
Solves over a grid of `num_lamba` regularization parameter values that are evenly spaced
on a log scale between `lambda_min_ratio*lambda_max` and `lambda_max`.
Both Ridge and Lasso are implemented. Returns `B_soln_list, lambdas` where `B_soln_list` is a list of
lienar models such that `B_soln_list[i]` is the solution with parameter `lambdas[i]`.
The default values for the algorithm parameters `reform_alg_parms` are (explanations are given below):
    lambda_max::Number = -1 (actual default is `(d/n)*(norm(X)^2)`)
    lambda_min_ratio::Float64 = 0.0001
    num_lambda::Int = 100
    solver::Symbol = :Gurobi
    regularization::Symbol = :ridge
    po_loss_function::Symbol = :leastSquares
    huber_delta::Float64 = 0.0
# Arguments
- `X`: p x n training set feature matrix
- `c`: d x n training set matrix of cost vectors
- `lambda_max`:  largest value of lambda, default is `(d/n)*(norm(X)^2)`
- `lambda_min_ratio`:  lambda_min is `lambda_max*lambda_min_ratio`
- `num_lambda`:  number of lambda values, evenly spaced on log scale between [lambda_min, lambda_max]
- `solver`:  either `:Gurobi` or `:Clp`
- `gurobiEnv` is a Gurobi enviornment which can be passed in (helpful to minimize the number of license creations)
- `regularization`:  either `:ridge` or `:lasso`
- `regularize_first_column_B`:  either true or false
- `po_loss_function`:  either `:leastSquares` or `:huber` or `:absolute`
- `huber_delta`:  the delta parameter for Huber loss
"""
function leastSquares_path_jump(X, c;
    reform_alg_parms::reformulation_path_parms = reformulation_path_parms())

    # Unpack parms
    @unpack lambda_max,
        lambda_min_ratio,
        num_lambda,
        solver,
        gurobiEnv,
        regularization,
        regularize_first_column_B,
        upper_bound_B,
        upper_bound_B_present,
        po_loss_function,
        huber_delta,
        verbose,
        algorithm_type = reform_alg_parms

    # Dimension check
    (p, n) = size(X)
    (d, n2) = size(c)
    if n != n2
        error("Dimensions of the input are mismatched.")
    end

    # Start creating JuMP model
    if solver == :Gurobi
        mod = Model(with_optimizer(Gurobi.Optimizer, gurobiEnv))
    elseif solver == :Clp
        mod = Model(with_optimizer(Clp.Optimizer))
    else
        error("Not a valid solver: either :Clp or :Gurobi")
    end

    # Add variables
    if upper_bound_B_present
        @variable(mod, -upper_bound_B <= B_var[1:d, 1:p] <= upper_bound_B)
    else
        @variable(mod, B_var[1:d, 1:p])
    end

    # Add least squares or Huber objective
    if po_loss_function == :leastSquares
        obj_expr_noreg = dot(c - B_var*X, c - B_var*X)
    elseif po_loss_function == :huber
        @variable(mod, w_var[1:d, 1:n] >= 0)
        @variable(mod, v_var[1:d, 1:n])
        @constraint(mod, c - B_var*X .<= v_var + w_var)
        @constraint(mod, -(c - B_var*X) .<= v_var + w_var)
        @constraint(mod, w_var .<= huber_delta*ones(d, n))

        obj_expr_noreg = dot(w_var, w_var) + 2*huber_delta*dot(v_var, ones(d, n))
    elseif po_loss_function == :absolute
        @variable(mod, w_var[1:d, 1:n] >= 0)
        @constraint(mod, c - B_var*X .<= w_var)
        @constraint(mod, -(c - B_var*X) .<= w_var)

        obj_expr_noreg = 2*dot(w_var, ones(d, n))
    else
        error("Enter a valid loss function: either :leastSquares or :huber or :absolute")
    end

    ### Add regularization part and solve for path
    if ismissing(lambda_max)
        lambda_max = (d/n)*(norm(X)^2)
    end

    # Construct lambda sequence
    if num_lambda == 1 && lambda_max == 0
        lambdas = [0.0]
    else
        lambda_min = lambda_max*lambda_min_ratio
        #println("Lambda min is $lambda_min")
        log_lambdas = collect(range(log(lambda_min), stop = log(lambda_max), length = num_lambda))
        lambdas = exp.(log_lambdas)
    end

    # Add theta variables for Lasso
    if regularization == :lasso
        @variable(mod, theta_var[1:d, 1:p])
        @constraint(mod, theta_var .>= B_var)
        @constraint(mod, theta_var .>= -B_var)
    end

    # Make list of B models and solve path
    B_soln_list = Matrix{Float64}[]
    for i = 1:num_lambda
        cur_lambda = lambdas[i]
        if verbose
            println("Trying lambda = $cur_lambda")
        end

        if regularization == :ridge && regularize_first_column_B
            obj_expr_full = obj_expr_noreg + n*cur_lambda*dot(B_var, B_var)
        elseif regularization == :ridge && !regularize_first_column_B
            obj_expr_full = obj_expr_noreg + n*cur_lambda*dot(B_var[:, 2:p], B_var[:, 2:p])
        elseif regularization == :lasso && regularize_first_column_B
            obj_expr_full = obj_expr_noreg + 2*n*cur_lambda*sum(theta_var)
        elseif regularization == :lasso && !regularize_first_column_B
            obj_expr_full = obj_expr_noreg + 2*n*cur_lambda*sum(theta_var[:, 2:p])
        else
            error("enter valid regularization: :ridge or :lasso")
        end
        @objective(mod, Min, obj_expr_full)

        # solve
        optimize!(mod)
        #z_ast = objective_value(mod)
        mod_status = termination_status(mod)
        B_ast = value.(B_var)

        if any(isnan, B_ast)
            # set B_ast to zero
            B_ast = zeros(d, p)

            # print reason
            println("We got NaNs in B_ast and the reason is:  $mod_status. The algoritim is $algorithm_type.")
        elseif Int(mod_status) != 1
            println("There are no NaNs in B_ast, but the model status is:  $mod_status. The algoritim is $algorithm_type.")
        end

        push!(B_soln_list, B_ast)
    end

    return B_soln_list, lambdas
end

ext: jl

max_stars_repo_name: katerinagiannoutsou/dso-coding-project

max_issues_repo_name: katerinagiannoutsou/dso-coding-project

max_forks_repo_name: katerinagiannoutsou/dso-coding-project


text:
const sample_stats_eltypes = Dict(
    "lp" => Float64,
    "step_size" => Float64,
    "step_size_nom" => Float64,
    "tree_depth" => Int,
    "n_steps" => Int,
    "diverging" => Bool,
    "energy" => Float64,
    "energy_error" => Float64,
    "max_energy_error" => Float64,
    "int_time" => Float64,
)

@doc doc"""
    with_interactive_backend(f; backend::Symbol = nothing)

Execute the thunk `f` in a temporary interactive context with the chosen `backend`, or
provide no arguments to use a default.

# Examples

```julia
idata = load_arviz_data("centered_eight")
plot_posterior(idata) # inline
with_interactive_backend() do
    plot_density(idata) # interactive
end
plot_trace(idata) # inline
```
"""
with_interactive_backend

function with_interactive_backend(f; backend=nothing)
    oldisint = PyPlot.isinteractive()
    oldgui = pygui()
    backend === nothing || pygui(Symbol(backend))
    pygui(true)
    ret = f()
    pygui(oldisint)
    pygui(oldgui)
    return ret
end

"""
    use_style(style::String)
    use_style(style::Vector{String})

Use matplotlib style settings from a style specification `style`.

The style name of "default" is reserved for reverting back to the default style settings.

ArviZ-specific styles are
`["arviz-whitegrid", "arviz-darkgrid", "arviz-colors", "arviz-white"]`.
To see all available style specifications, use [`styles()`](@ref).

If a `Vector` of styles is provided, they are applied from first to last.
"""
use_style(style) = plt.style.use(style)

"""
    styles() -> Vector{String}

Get all available matplotlib styles for use with [`use_style`](@ref)
"""
styles() = plt.style.available

"""
    convert_arguments(f, args...; kwargs...) -> NTuple{2}

Convert arguments to the function `f` before calling.

This function is used primarily for pre-processing arguments within macros before sending
to arviz.
"""
convert_arguments(::Any, args...; kwargs...) = args, kwargs

"""
    convert_result(f, result, args...)

Convert result of the function `f` before returning.

This function is used primarily for post-processing outputs of arviz before returning.
The `args` are primarily used for dispatch.
"""
convert_result(f, result, args...) = result

load_backend(backend) = nothing

function forwarddoc(f::Symbol)
    return "See documentation for [`arviz.$(f)`](https://python.arviz.org/en/latest/api/generated/arviz.$(f).html)."
end

forwardgetdoc(f::Symbol) = Docs.getdoc(getproperty(arviz, f))

"""
    @forwardfun f
    @forwardfun(f)

Wrap a function `arviz.f` in `f`, forwarding its docstrings.

Use [`convert_arguments`](@ref) and [`convert_result`](@ref) to customize what is passed to
and returned from `f`.
"""
macro forwardfun(f)
    fdoc = forwarddoc(f)
    return esc(
        quote
            @doc $fdoc $f

            function $(f)(args...; kwargs...)
                args, kwargs = convert_arguments($(f), args...; kwargs...)
                result = arviz.$(f)(args...; kwargs...)
                return convert_result($(f), result)
            end

            Docs.getdoc(::typeof($(f))) = forwardgetdoc(Symbol($(f)))
        end,
    )
end

"""
    @forwardplotfun f
    @forwardplotfun(f)

Wrap a plotting function `arviz.f` in `f`, forwarding its docstrings.

This macro also ensures that outputs for the different backends are correctly handled.
Use [`convert_arguments`](@ref) and [`convert_result`](@ref) to customize what is passed to
and returned from `f`.
"""
macro forwardplotfun(f)
    fdoc = forwarddoc(f)
    return esc(
        quote
            @doc $fdoc $f

            function $(f)(args...; backend=nothing, kwargs...)
                if backend === nothing
                    backend = get(rcParams, "plot.backend", nothing)
                end
                backend = Symbol(backend)
                backend_val = Val(backend)
                load_backend(backend_val)
                args, kwargs = convert_arguments($(f), args...; kwargs...)
                result = arviz.$(f)(args...; kwargs..., backend)
                return convert_result($(f), result, backend_val)
            end

            Docs.getdoc(::typeof($(f))) = forwardgetdoc(Symbol($(f)))
        end,
    )
end

# Replace `missing` values with `NaN` and do type inference on the result
replacemissing(x) = map(identity, replace(x, missing => NaN))
replacemissing(x::AbstractArray{<:AbstractArray}) = map(replacemissing, x)
@inline replacemissing(x::AbstractArray{<:Real}) = x
@inline replacemissing(::Missing) = NaN
@inline replacemissing(x::Number) = x

"""
    flatten(x)

If `x` is an array of arrays, flatten into a single array whose dimensions are ordered with
dimensions of the outermost container first and innermost container last.
"""
flatten(x) = x
flatten(x::AbstractArray{<:Number}) = convert(Array, x)
function flatten(x::AbstractArray{S}) where {T<:Number,N,S<:AbstractArray{T,N}}
    ret = Array{T}(undef, (size(x)..., size(x[1])...))
    for k in keys(x)
        setindex!(ret, x[k], k, (Colon() for _ in 1:N)...)
    end
    return ret
end

# Convert python types to Julia types if possible
@inline frompytype(x) = x
@inline frompytype(x::PyObject) = PyAny(x)
frompytype(x::AbstractArray{PyObject}) = map(frompytype, x)
frompytype(x::AbstractArray{Any}) = map(frompytype, x)
frompytype(x::AbstractArray{<:AbstractArray}) = map(frompytype, x)

function rekey(d, keymap)
    dnew = empty(d)
    for (k, v) in d
        knew = get(keymap, k, k)
        haskey(dnew, knew) && throw(ArgumentError("$knew in `keymap` is already in `d`."))
        dnew[knew] = d[k]
    end
    return dnew
end

removekeys!(dict, keys) = map(k -> delete!(dict, k), keys)

function popsubdict!(dict, names)
    (dict === nothing || names === nothing) && return nothing
    subdict = empty(dict)
    for k in names
        subdict[k] = pop!(dict, k)
    end
    return subdict
end
popsubdict!(dict, key::String) = popsubdict!(dict, [key])

snakecase(s) = replace(lowercase(s), " " => "_")

@inline _asarray(x) = [x]
@inline _asarray(x::AbstractArray) = x

_asstringkeydict(x) = Dict(String(k) => v for (k, v) in pairs(x))
_asstringkeydict(x::Dict{String}) = x

enforce_stat_eltypes(stats) = convert_to_eltypes(stats, sample_stats_eltypes)

function convert_to_eltypes(data::Dict, data_eltypes)
    return Dict(k => convert(Array{get(data_eltypes, k, eltype(v))}, v) for (k, v) in data)
end
function convert_to_eltypes(data::NamedTuple, data_eltypes)
    return NamedTuple(
        k => convert(Array{get(data_eltypes, k, eltype(v))}, v) for (k, v) in pairs(data)
    )
end

"""
    todataframes(df; index_name = nothing) -> DataFrames.DataFrame

Convert a Python `pandas.DataFrame` or `pandas.Series` into a `DataFrames.DataFrame`.

If `index_name` is not `nothing`, the index is converted into a column with `index_name`.
Otherwise, it is discarded.
"""
function todataframes(::Val{:DataFrame}, df::PyObject; index_name=nothing)
    initialize_pandas()
    col_vals = map(df.columns) do name
        series = py"$(df)[$(name)]"
        vals = series.values
        return Symbol(name) => frompytype(vals)
    end
    if index_name !== nothing
        index_vals = frompytype(df.index.values)
        col_vals = [Symbol(index_name) => index_vals; col_vals]
    end
    return DataFrames.DataFrame(col_vals)
end
function todataframes(::Val{:Series}, series::PyObject; kwargs...)
    initialize_pandas()
    colnames = map(i -> Symbol(frompytype(i)), series.index)
    colvals = map(x -> [frompytype(x)], series.values)
    return DataFrames.DataFrame(colvals, colnames)
end
function todataframes(df::PyObject; kwargs...)
    initialize_pandas()
    if pyisinstance(df, pandas.Series)
        return todataframes(Val(:Series), df; kwargs...)
    end
    return todataframes(Val(:DataFrame), df; kwargs...)
end

"""
    topandas(::Type{:DataFrame}, df; index_name = nothing) -> PyObject
    topandas(::Type{:Series}, df) -> PyObject
    topandas(::Val{:ELPDData}, df) -> PyObject

Convert a `DataFrames.DataFrame` to the specified pandas type.

If `index_name` is not `nothing`, the corresponding column is made the index of the
returned dataframe.
"""
function topandas(::Val{:DataFrame}, df; index_name=nothing)
    initialize_pandas()
    df = DataFrames.DataFrame(df)
    colnames = names(df)
    rowvals = map(Array, eachrow(df))
    pdf = pandas.DataFrame(rowvals; columns=colnames)
    index_name !== nothing && pdf.set_index(index_name; inplace=true)
    return pdf
end
function topandas(::Val{:Series}, df)
    initialize_pandas()
    df = DataFrames.DataFrame(df)
    rownames = names(df)
    colvals = Array(only(eachrow(df)))
    return pandas.Series(colvals, rownames)
end
function topandas(::Val{:ELPDData}, df)
    initialize_pandas()
    df = DataFrames.DataFrame(df)
    rownames = names(df)
    colvals = Array(only(eachrow(df)))
    return ArviZ.arviz.stats.ELPDData(colvals, rownames)
end


text:
using DiffEqEnvironments
using DifferentialEquations
using ReinforcementLearning
using LinearAlgebra
using ControlSystems
using Test

"""
This test defines an SISO system of form
    ṡ = A*s + B*a
    o = C*s
for the rocket car problem, where
    s = [p, ṗ]  ,
p being the position of the car.
The action used as an input to the system is
the car's acceleration p̈ ∈ [-1, 1]. Therefore
    ṡ = [ṗ, p̈] = [0 1; 0 0] * s + [0; 1] * u .
Full observation of the state is given, as
    o = s   .
"""
T = Float32

# System matrices
A = [0 1; 0 0]
B = [0; 1]

# weights for quadratic cost function
Q = [1 0; 0 1] * 10
R = [1][:,:]

# Set range of control input
a_lb = -1.0; a_ub = 1.0

# Set IC & tspan for integration
s0 = [-0.75, 0.0]
tspan = (0f0, 5f0)

# Define ODEProblem of SISO LTI system
function lti(s, a, t)
    ṡ = A * s + B * a
end
prob_lti = ODEProblem(lti, s0, tspan)

# Get LQR from ControlSystems.jl
K = lqr(A, B, Q, R)
π(s) = clamp.(-dot(K, s), a_lb, a_ub) #  LQR controller w/ saturation

# Define DiffEqEnv
r = QuadraticReward(Q, R)
n_actions = 1
dt = 0.1
env = DiffEqEnv(prob_lti, r, n_actions, dt, a_lb=a_lb, a_ub=a_ub)

# Manually step through environment
for _ ∈ 1:200
    s = RLBase.get_state(env)
    a = Float32.(π(s))[]
    env(a) # perform step
end

# Check if state [0,0] was reached
@test isapprox(env.state[1], 0, atol=1e-8)
@test isapprox(env.state[2], 0, atol=1e-8)

# Run random policy using ReinforcementLearning.jl
lqr_policy = LQRPolicy(A, B, Q, R, dt, a_lb, a_ub)

RLBase.reset!(env)
hook = TotalRewardPerEpisode()
run(
    Agent(
        ;policy=lqr_policy,
        trajectory=VectCompactSARTSATrajectory(
            state_type=Any,
            action_type=Any,
            reward_type=Real,
            terminal_type=Bool,
        ),
    ),
    env,
    StopAfterEpisode(10),
    hook
)


text:
module EHTIM
using PyCall, StructArrays, Dates

export structarray, bldict, get_obstime

new() = pyimport("ehtim")

function get_obstime(x::PyObject)
    if :unpack in keys(x)
        julian2datetime.(get(x.data, :time)/24 .+ 2400000.5 .+ x.mjd)
    else
        error("Use get_obstime only on Obsdata object.")
    end
end

function structarray(x::PyCall.PyObject; datetime = true)
    ref = :unpack in keys(x) ? x.data : x
    labels = ref.dtype.names |> collect
    data = []
    for (i,n) in enumerate(labels)
        # Handle for Unicode arrays that are not supported by PyArray
        # https://www.oreilly.com/library/view/python-for-finance/9781491945360/ch04.html
        push!(data, get(ref.dtype, i-1).str[2] == 'U' ? get.(ref, i-1) : PyArray(py"$ref[$n]"o))
    end
    if :unpack in keys(x) && datetime && any(ref.dtype.names .== "time")
        push!(labels, "datetime")
        push!(data, get_obstime(x))
    end

    (;zip(Symbol.(labels), data)...) |> StructArray
end

function structarray(x::Array{PyCall.PyObject,1})
    array = StructArray[]
    for element in x
        push!(array, structarray(element))
    end
    return array
end

function bldict(x::Array{StructArrays.StructArray,1})
    dict = Dict{String,StructArray}()
    for element in x
        if hasproperty(element, :t1) && hasproperty(element, :t2)
            dict[element.t1[1] > element.t2[1] ? element.t2[1]*"-"*element.t1[1] :
                 element.t1[1]*"-"*element.t2[1]] = element
        else
            error("No baseline information in data.")
        end
    end
    return dict
end

end # EHTIM


text:
@noinline throwboundserror() = throw(BoundsError())

struct ResizableMatrix{T,NC} <: AbstractMatrix{T}
    data::Vector{T}
    nrows::Base.RefValue{Int}
end
@inline Base.size(rm::ResizableMatrix{T,NC}) where {T,NC} = (rm.nrows[], NC)
@inline Base.length(rm::ResizableMatrix) = length(rm.data)
@inline Base.getindex(rm::ResizableMatrix, i::Integer) = rm.data[i]
@inline Base.getindex(rm::ResizableMatrix, I::CartesianIndex{2}) = rm.data[I[1] + (I[2]-1)*rm.nrows[]]
@inline Base.getindex(rm::ResizableMatrix, i::Integer, j::Integer) = rm.data[ i + (j-1)*rm.nrows[] ]
@inline Base.setindex!(rm::ResizableMatrix, v, i) = rm.data[i] = v
@inline Base.setindex!(rm::ResizableMatrix, v, i::Integer, j::Integer) = rm.data[i + (j-1)*rm.nrows[]] = v
@inline Base.pointer(rm::ResizableMatrix) = pointer(rm.data)
@inline Base.pointer(rm::ResizableMatrix, i::Integer) = pointer(rm.data, i)
Base.resize!(rm::ResizableMatrix{T,NC}, N) where {T,NC} = (rm.nrows[] = N; resize!(rm.data, NC*N))
Base.IndexStyle(::Type{<:ResizableMatrix}) = IndexLinear()
function ResizableMatrix{T,NC}(::UndefInitializer, N) where {T,NC}
    data = Vector{T}(undef, N * NC)
    ResizableMatrix{T,NC}(data, Ref(N))
end

function Base.getindex(rm::ResizableMatrix{T,NC}, I::Union{<:AbstractArray{Bool},<:BitArray}, ::Colon) where {T,NC}
    @boundscheck rm.nrows[] < length(I) && throwboundserror()
    out = Matrix{T}(undef, sum(I), NC)
    nrows = rm.nrows[]
    @inbounds for c ∈ 1:NC
        out[:,c] .= @view(rm.data[1+(c-1)*nrows:c*nrows])[I]
    end
    out
end
@inline function Base.view(rm::ResizableMatrix{T,NC}, ::Colon, i::Integer) where {T,NC}
    @boundscheck i > NC && throwboundserror()
    nrows = rm.nrows[]
    @view rm.data[ (1:nrows) .+ (i-1)*nrows ]
end

@generated function findall_zerorows(X::AbstractMatrix{T}, ::Val{ncol}) where {T,ncol}
    @assert ncol > 1
    ex = :((X[i,$(ncol-1)] != zero($T)) | (X[i,$ncol] != zero($T)))
    for j ∈ ncol-2:-1:1
        ex = :( (X[i,$j] != zero($T)) | $ex )
    end
    q = quote
        tb_1 = $(Expr(:tuple, [:(X[iu+$(j-1),1] != zero($T)) for j ∈ 1:16]...))
    end
    for c ∈ 2:ncol
        push!(q.args, quote
            $(Symbol(:tb_, c)) = ($(Symbol(:tb_, c-1)) .| $(Expr(:tuple, [:(X[iu+$(j-1),$c] != zero($T)) for j ∈ 1:16]...)))
        end)
    end
    quote
        N = size(X,1)
        N16, r = divrem(N,16)
        z0 = Vector{Bool}(undef, N)
        @inbounds @fastmath for iu ∈ 1:16:N
            $q
            z0[iu:iu+15] .= $(Symbol(:tb_, ncol))
        end
        @inbounds @fastmath for i ∈ N-r+1:N
            z0[i] = $ex
        end
        z0
    end
end


text:
#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

import ..DomainIR

function relabel(exprs::Array{Any}, irState)
  labelDict = Dict{Int, Int}()
  for i = 1:length(exprs)
    expr = exprs[i]
    if isa(expr, LabelNode)
      new_label = next_label(irState)
      labelDict[expr.label] = new_label
    end
  end
  for i = 1:length(exprs)
    expr = exprs[i]
    if isa(expr, LabelNode)
      exprs[i] = LabelNode(labelDict[expr.label])
    elseif isa(expr, GotoNode)
      exprs[i] = GotoNode(labelDict[expr.label])
    elseif isa(expr, Expr) && (expr.head === :gotoifnot)
      exprs[i] = TypedExpr(expr.typ, expr.head, expr.args[1], labelDict[expr.args[2]])
    end
  end
  return exprs
end

function simplifyBodyExpr(kernelF, state)
    LambdaVarInfo, body = CompilerTools.LambdaHandling.lambdaToLambdaVarInfo(kernelF)
    cfg = CompilerTools.CFGs.from_lambda(body)
    body = CompilerTools.LambdaHandling.getBody(CompilerTools.CFGs.createFunctionBody(cfg), CompilerTools.LambdaHandling.getReturnType(LambdaVarInfo))
    non_array_params = Set{LHSVar}()
    changed = true
    lives = computeLiveness(body, LambdaVarInfo)
    while changed
        @dprintln(1,"Removing statement with no dependencies from the AST with parameters")
        rnd_state = RemoveNoDepsState(lives, non_array_params)
        body = AstWalk(body, remove_no_deps, rnd_state)
        @dprintln(3,"body after no dep stmts removed = ", body)

        @dprintln(3,"top_level_no_deps = ", rnd_state.top_level_no_deps)

        @dprintln(1,"Adding statements with no dependencies to the start of the AST.")
        body = CompilerTools.LambdaHandling.prependStatements(body, rnd_state.top_level_no_deps)
        @dprintln(3,"body after no dep stmts re-inserted = ", body)

        @dprintln(1,"Re-starting liveness analysis.")
        lives = computeLiveness(body, LambdaVarInfo)
        @dprintln(1,"Finished liveness analysis.")

        changed = rnd_state.change
    end
    body = AstWalk(body, remove_dead, RemoveDeadState(lives,LambdaVarInfo))
    return LambdaVarInfo, body
end

function mk_parfor_args_from_stencil(typ, head, args, irState)
  assert(length(args) == 4)
  local stat = args[1]
  local border = stat.borderSty
  local border_inloop = false
  local oob_skip = false
  local oob_dst_zero = false
  local oob_src_zero = false
  local oob_wraparound = false
  assert(isa(border, Symbol))
  oob_skip = (border === :oob_skip)
  oob_dst_zero = (border === :oob_dst_zero)
  oob_src_zero = (border === :oob_src_zero)
  oob_wraparound = (border === :oob_wraparound)
  local iterations = args[2]
  # convert all TypedVar in bufs to just Symbol
  local bufs = args[3]
  local kernelF :: DomainLambda = args[4]
  local linfo = irState.LambdaVarInfo
  # println(stat)
  local buf = bufs[1] # assumes that first buffer has the same dimension as all other buffers
  main_length_correlation = getOrAddArrayCorrelation(toLHSVar(bufs[1]), irState)
  local n = stat.dimension
  local sizeNodes = [ addTempVariable(Int, linfo) for i = 1:n ]
  local sizeInitExpr = [ TypedExpr(Int, :(=), sizeNodes[i],
                                   TypedExpr(Int, :call, GlobalRef(Base, :arraysize), buf, i))
                         for i in 1:n ]
  local strideNodes = [ addTempVariable(Int, linfo) for s in stat.strideSym ]
  local strideInitExpr = Array{Any}(n)
  strideInitExpr[1] = TypedExpr(Int, :(=), strideNodes[1], 1)
  # the following assumes contiguous column major layout for multi-dimensional arrays
  for i = 2:n
    strideInitExpr[i] = TypedExpr(Int, :(=), strideNodes[i],
                                  DomainIR.mul_expr(strideNodes[i-1], sizeNodes[i-1]))
  end
  local idxNodes = Any[ DomainIR.addFreshLocalVariable(string("i",s.id), Int, ISASSIGNED | ISASSIGNEDONCE, linfo) for s in stat.idxSym ]
  local loopNest = [ PIRLoopNest(idxNodes[i],
                                 border_inloop ? 1 : 1-stat.shapeMin[i],
                                 border_inloop ? sizeNodes[i] : DomainIR.sub_expr(sizeNodes[i], stat.shapeMax[i]),
                                 1)
                     for i in n:-1:1 ]
  local nbufs = length(bufs)
  tmpBufs = Array{TypedVar}(nbufs)
  for i = 1:length(bufs)
    elmTyp = DomainIR.elmTypOf(getType(bufs[i], linfo))
    arrTyp = Array{elmTyp, n}
    # TODO: enforce buffer type
    # bufs[i].typ = arrTyp
    tmpBufs[i] = DomainIR.addFreshLocalVariable("tmp", arrTyp, ISASSIGNED, linfo)

    this_correlation = getOrAddArrayCorrelation(toLHSVar(bufs[i]), irState)
    if this_correlation != main_length_correlation
      merge_correlations(irState, main_length_correlation, this_correlation)
    end
  end
  @dprintln(3, "before simplifyBodyExpr body=", kernelF.body)
  bodyLinfo, bodyExpr = simplifyBodyExpr(kernelF, irState)
  @dprintln(3, "before simplifyBodyExpr body=", bodyExpr)
  bodyExpr = relabel(DomainIR.stencilGenBody(stat, bodyLinfo, bodyExpr, idxNodes, strideNodes, bufs, linfo, getLoopPrivateFlags()), irState)
  # rotate
  assert((bodyExpr[end].head === :tuple))
  @dprintln(3,"bodyExpr = ")
  printBody(3, bodyExpr)
  local rotateExpr = Array{Any}(0)
  local revertExpr = Array{Any}(0)
  # warn(string("last return=", bodyExpr[end]))
  if bodyExpr[end].args[1] != nothing
    rets = bodyExpr[end].args
    @dprintln(3,"bodyExpr[end].args = ", rets)
    assert(length(rets) == nbufs)
    for i = 1:nbufs
      push!(rotateExpr, TypedExpr(CompilerTools.LambdaHandling.getType(rets[i], linfo), :(=), toLHSVar(tmpBufs[i]), rets[i]))
      push!(revertExpr, TypedExpr(CompilerTools.LambdaHandling.getType(rets[i], linfo), :(=), toLHSVar(tmpBufs[i]), bufs[i]))
    end
    for i = 1:nbufs
      push!(rotateExpr, TypedExpr(CompilerTools.LambdaHandling.getType(tmpBufs[i], linfo), :(=), toLHSVar(bufs[i]), tmpBufs[i]))
      push!(revertExpr, TypedExpr(CompilerTools.LambdaHandling.getType(tmpBufs[i], linfo), :(=), toLHSVar(rets[i]), tmpBufs[i]))
    end
  end
  bodyExpr = bodyExpr[1:end-1]
  # border handling
  borderLabel = next_label(irState)
  afterBorderLabel = next_label(irState)
  lowerExprs = [ DomainIR.box_ty(Bool, Expr(:call, GlobalRef(Base, :sle_int), 1-stat.shapeMin[i], idxNodes[i]))
                 for i in 1:n ]
  upperExprs = [ DomainIR.box_ty(Bool, Expr(:call, GlobalRef(Base, :sle_int), idxNodes[i],
                      DomainIR.sub_expr(sizeNodes[i], stat.shapeMax[i])))
                 for i in 1:n ]
  lowerGotos = [ Expr(:gotoifnot, e, borderLabel) for e in lowerExprs ]
  upperGotos = [ Expr(:gotoifnot, e, borderLabel) for e in upperExprs ]
  borderHead = Any[ TypedExpr(Int, :(=), toLHSVar(idxNodes[1]),
                      DomainIR.sub_expr(sizeNodes[1], stat.shapeMax[1])),
                    GotoNode(afterBorderLabel),
                    LabelNode(borderLabel),
                  ]
  borderExpr = Any[]
  if oob_dst_zero
    for expr in bodyExpr
      expr = deepcopy(expr)
      if (expr.head === :call) && isBaseFunc(expr.args[1], :unsafe_arrayset)
        zero = Base.convert(DomainIR.elmTypOf(getType(expr.args[2], linfo)), 0)
        push!(borderExpr, TypedExpr(expr.typ, :call, expr.args[1], expr.args[2], zero, expr.args[4:end]...))
      end
    end
  elseif oob_src_zero
    for expr in bodyExpr
      expr = deepcopy(expr)
      if isa(expr, Expr) && (expr.head === :(=))
        lhs = expr.args[1]
        rhs = expr.args[2]
        if isa(rhs, Expr) && (rhs.head === :call) && isBaseFunc(rhs.args[1], :unsafe_arrayref)
          zero = Base.convert(rhs.typ, 0)
          expr = TypedExpr(expr.typ, :(=), lhs,
                  TypedExpr(rhs.typ, :call, GlobalRef(Base, :safe_arrayref),
                      rhs.args[2], zero, rhs.args[3:end]...))
        end
      end
      push!(borderExpr, expr)
    end
  elseif oob_wraparound
    for expr in bodyExpr
      expr = deepcopy(expr)
      if isa(expr, Expr) && (expr.head === :(=))
        lhs = expr.args[1]
        rhs = expr.args[2]
        if isa(rhs, Expr) && (rhs.head === :call) && isBaseFunc(rhs.args[1], :unsafe_arrayref)
          indices = Expr[ TypedExpr(Int, :call, GlobalRef(Base, :select_value),
                            DomainIR.box_ty(Bool, Expr(:call, GlobalRef(Base, :sle_int), deepcopy(rhs.args[2+i]), 0)),
                            DomainIR.add_expr(deepcopy(rhs.args[2+i]), sizeNodes[i]),
                            TypedExpr(Int, :call, GlobalRef(Base, :select_value),
                              DomainIR.box_ty(Bool, Expr(:call, GlobalRef(Base, :sle_int),
                                DomainIR.add_expr(sizeNodes[i], 1), deepcopy(rhs.args[2+i]))),
                              DomainIR.sub_expr(deepcopy(rhs.args[2+i]), sizeNodes[i]),
                              deepcopy(rhs.args[2+i]))) for i = 1:n ]
          expr = TypedExpr(expr.typ, :(=), lhs,
                  TypedExpr(rhs.typ, :call, GlobalRef(Base, :unsafe_arrayref),
                      rhs.args[2], indices...))
        end
      end
      push!(borderExpr, expr)
    end
  else #FIXME: the following is to make a dummy node to avoid an IR bug
#    push!(borderExpr, TypedExpr(Int, :(=), idxNodes[1].name, idxNodes[1]))
  end
  borderExpr = vcat(borderHead, relabel(borderExpr, irState), LabelNode(afterBorderLabel))
  # borderCond = [ borderHead, lowerGotos, upperGotos, borderExpr, borderTail ]
  borderCond = oob_skip ? Any[] : Any[TypedExpr(Void, head,
        PIRParForAst(InputInfo(toLHSVar(buf)),
            vcat(lowerGotos, upperGotos, borderExpr),
            [],
            [],
            [ PIRLoopNest(idxNodes[i], 1, sizeNodes[i], 1) for i = n:-1:1 ],
            PIRReduction[],
            [], [], irState.top_level_number, get_unique_num(), Set{LHSVar}(), Set{LHSVar}([toLHSVar(x) for x in bufs])))]

  stepNode = DomainIR.addFreshLocalVariable("step", Int, ISASSIGNED | ISASSIGNEDONCE, linfo)
  # Sequential loop for multi-iterations
  iterPre  = (isa(iterations, Number) && iterations == 1) ?
            Any[] : Any[ Expr(:loophead, stepNode, 1, iterations) ]
  iterPost = (isa(iterations, Number) && iterations == 1) ?
            Any[] : vcat(rotateExpr, Expr(:loopend, stepNode))
  preExpr = vcat(sizeInitExpr, strideInitExpr, iterPre, borderCond)
  postExpr = vcat(iterPost)
  expr = PIRParForAst(
    InputInfo(buf),
    bodyExpr,
    [],
    [],
    loopNest,
    PIRReduction[],
    Any[ length(bufs) > 2 ? tuple(bufs...) : bufs[1] ],
    DomainOperation[ DomainOperation(:stencil!, args) ],
    irState.top_level_number,
    get_unique_num(),
    Set{LHSVar}(),
    Set{LHSVar}([toLHSVar(x) for x in bufs]))
  return vcat(preExpr, TypedExpr(typ, head, expr), postExpr)
end


text:
using InvariantMeasures
using ValidatedNumerics
using LinearAlgebra

@testset "Hat assembler" begin

D = Mod1Dynamic(x->2*x)
B = Hat(8)
P = assemble(B, D)

Ptrue = [
        0.5 0.25 0    0    0    0    0   0.25;
        0   0.25 0.5  0.25 0    0    0   0   ;
        0   0    0    0.25 0.5  0.25 0   0   ;
        0   0    0    0    0    0.25 0.5 0.25;
        0.5 0.25 0    0    0    0    0   0.25;
        0   0.25 0.5  0.25 0    0    0   0   ;
        0   0    0    0.25 0.5  0.25 0   0   ;
        0   0    0    0    0    0.25 0.5 0.25;
		]
Ptrue = Ptrue'

@test all(contains_zero.(P-Ptrue))

@test opnormbound(L1, DiscretizedOperator(B, D)) == 1
@test opnormbound(Linf, DiscretizedOperator(B, D)) == 1

Q = DiscretizedOperator(B, D)
@test size(Q) == (8,8)

n = size(Q)[1]
e = randn(n)
x = similar(e)
x[:] .= 2.

mQ = mid(Q)
@test mul!(x, mQ, e, 1, 0) == mQ*e

fQ = mQ.L + mQ.e * mQ.w
@test fQ * e ≈ mQ * e

end


text:
# ----------------------------------------------------
# --- Import modules
# ----------------------------------------------------
using DigitalComm
using Test
# ----------------------------------------------------
# --- Tests
# ----------------------------------------------------
# Note --> The mapping system is described in bitMapping.jl
println("Tests for symbol mapper with QAM sequences");

@testset  "bpsk" begin
	# Create a bit squence (already tested)
	nbBits	= 2 * 2048;
	bitSeq  = genBitSequence(nbBits);
	# Pass trough the function
	buff	= zeros(Complex{Float64},nbBits)
	# Call
	bitMappingQAM!(buff,2,bitSeq);
	buff2	= bitMappingQAM(2,bitSeq);
	@test all( buff .== buff2)
	# --- Good const point
	# For BPSK, we should have |r|^2 = 1
	@test all(abs.(real(buff)) .== 1);
	# Some manual check due to gray coding
	buff  = [0x01 0x00 0x00 0x00 0x00 0x01 0x01 0x01 ];
	@show bitMappingQAM(2,buff)
	@test all(bitMappingQAM(2,buff) .≈ [1; -1; -1; -1; -1; 1; 1; 1]);
end

@testset  "qpsk" begin
	# Create a bit squence (already tested)
	nbBits	= 2 * 2048;
	bitSeq  = genBitSequence(nbBits);
	# Pass trough the function
	buff	= zeros(Complex{Float64},nbBits÷2)
	# Call
	bitMappingQAM!(buff,4,bitSeq);
	buff2	= bitMappingQAM(4,bitSeq);
	@test all( buff .== buff2)
	# --- Good const point
	# For QPSK, we should have |r|^2 = 1
	@test all(abs.(real(buff)) .== 1/sqrt(2));
	# Some manual check due to gray coding
	buff  = [0x01 0x00 0x00 0x00 0x00 0x01 0x01 0x01 ];
	@test all(bitMappingQAM(4,buff) .≈ 1/sqrt(2)*[1-1im; -1-1im; -1+1im; 1+1im]);
end

@testset "16-QAM" begin
	# Create a bit squence (already tested)
	mcs		= 16;
	n		= Int(log2(mcs));
	nbBits	= n * 2048;
	bitSeq  = genBitSequence(nbBits);
	# Pass trough the function
	buff	= zeros(Complex{Float64},nbBits÷n)
	# Call
	bitMappingQAM!(buff,mcs,bitSeq);
	buff2	= bitMappingQAM(mcs,bitSeq);
	@test all( buff .== buff2)
	# Some manual check due to gray coding
	buff  = [0x01 0x00 0x00 0x00 0x00 0x01 0x01 0x01 ];
	@test all(bitMappingQAM(mcs,buff) .≈ 1/sqrt(10)*[-3+3im;1-1im]);
end

@testset "64-QAM" begin
	# Create a bit squence (already tested)
	mcs		= 64;
	n		= Int(log2(mcs));
	nbBits	= n * 2048;
	bitSeq  = genBitSequence(nbBits);
	# Pass trough the function
	buff	= zeros(Complex{Float64},nbBits÷n)
	# Call
	bitMappingQAM!(buff,mcs,bitSeq);
	buff2	= bitMappingQAM(mcs,bitSeq);
	@test all( buff .== buff2)
	# Some manual check due to gray coding
	buff  = [0x01 0x00 0x00 0x00 0x00 0x01 0x01 0x01 0x01 0x00 0x01 0x00];
	@test all(bitMappingQAM(mcs,buff) .≈ 1/sqrt(42)*[-7+5im;-3-7im]);
end

@testset "256-QAM" begin
	# Create a bit squence (already tested)
	mcs		= 256;
	n		= Int(log2(mcs));
	nbBits	= n * 2048;
	bitSeq  = genBitSequence(nbBits);
	# Pass trough the function
	buff	= zeros(Complex{Float64},nbBits÷n)
	# Call
	bitMappingQAM!(buff,mcs,bitSeq);
	buff2	= bitMappingQAM(mcs,bitSeq);
	@test all( buff .== buff2)
	# Some manual check due to gray coding
	buff  = [0x01 0x00 0x00 0x00 0x00 0x01 0x01 0x01 0x01 0x00 0x01 0x00 0x01 0x01 0x00 0x00];
	@test all(bitMappingQAM(mcs,buff) .≈ 1/sqrt(170)*[3-5im;9-7im]);
end