text:
#
# Combinator algebra of vectorized data transformations.
#

import Base:
    OneTo,
    show

using Base.Cartesian

#
# Pipeline interface.
#

"""
    Runtime()

Runtime state for pipeline evaluation.
"""
mutable struct Runtime
end

"""
    Pipeline(op, args...)

A pipeline object represents a vectorized data transformation.

Parameter `op` is a function that performs the transformation; `args` are extra
arguments to be passed to the function.

The pipeline transforms any input vector by invoking `op` with the following
arguments:

    op(rt::Runtime, input::AbstractVector, args...)

The result of `op` must be the output vector, which should be of the same
length as the input vector.
"""
struct Pipeline
    op
    args::Vector{Any}
    sig::Signature

    Pipeline(op; args::Vector{Any}=Any[], sig::Signature=Signature()) =
        new(op, args, sig)
end

Pipeline(op, args...) =
    Pipeline(op, args=Any[args...])

"""
    designate(::Pipeline, ::Signature) :: Pipeline
    designate(::Pipeline, ::AbstractShape, ::AbstractShape) :: Pipeline
    p::Pipeline |> designate(::Signature) :: Pipeline
    p::Pipeline |> designate(::AbstractShape, ::AbstractShape) :: Pipeline

Sets the pipeline signature.
"""
function designate end

designate(p::Pipeline, sig::Signature) =
    Pipeline(p.op, args=p.args, sig=sig)

designate(p::Pipeline, src::Union{AbstractShape,Type}, tgt::Union{AbstractShape,Type}) =
    Pipeline(p.op, args=p.args, sig=Signature(src, tgt))

designate(sig::Signature) =
    p::Pipeline -> designate(p, sig)

designate(src::Union{AbstractShape,Type}, tgt::Union{AbstractShape,Type}) =
    p::Pipeline -> designate(p, Signature(src, tgt))

"""
    signature(::Pipeline) :: Signature

Returns the pipeline signature.
"""
signature(p::Pipeline) = p.sig

source(p::Pipeline) = source(p.sig)

target(p::Pipeline) = target(p.sig)

function (p::Pipeline)(input::DataKnot)
    @assert fits(shape(input), source(p))
    DataKnot(target(p), p(cell(input)))
end

function (p::Pipeline)(@nospecialize input::AbstractVector)::AbstractVector
    rt = Runtime()
    output = p(rt, input)
end

function (p::Pipeline)(rt::Runtime, @nospecialize input::AbstractVector)::AbstractVector
    p.op(rt, input, p.args...)
end

quoteof(p::Pipeline) =
    quoteof(p.op, p.args)

show(io::IO, p::Pipeline) =
    print_expr(io, quoteof(p))

#
# Vectorizing scalar functions.
#

"""
    lift(f) :: Pipeline

`f` is any scalar unary function.

The pipeline applies `f` to each element of the input vector.
"""
lift(f) = Pipeline(lift, f)

lift(f, args...) = Pipeline(lift, f, args)

lift(rt::Runtime, input::AbstractVector, f) =
    f.(input)

lift(rt::Runtime, input::AbstractVector, f, args) =
    f.(input, Ref.(args)...)

"""
    tuple_lift(f) :: Pipeline

`f` is an n-ary function.

The pipeline applies `f` to each row of an n-tuple vector.
"""
tuple_lift(f) = Pipeline(tuple_lift, f)

function tuple_lift(rt::Runtime, input::AbstractVector, f)
    @assert input isa TupleVector
    _tuple_lift(f, length(input), columns(input)...)
end

@generated function _tuple_lift(f, len::Int, cols::AbstractVector...)
    D = length(cols)
    return quote
        I = Tuple{Any[eltype(col) for col in cols]...}
        O = Core.Compiler.return_type(f, I)
        output = Vector{O}(undef, len)
        @inbounds for k = 1:len
            output[k] = @ncall $D f (d -> cols[d][k])
        end
        output
    end
end

"""
    block_lift(f) :: Pipeline
    block_lift(f, default) :: Pipeline

`f` is a function that expects a vector argument.

The pipeline applies `f` to each block of the input block vector.  When a block
is empty, `default` (if specified) is used as the output value.
"""
function block_lift end

block_lift(f) = Pipeline(block_lift, f)

function block_lift(rt::Runtime, input::AbstractVector, f)
    @assert input isa BlockVector
    _block_lift(f, input)
end

block_lift(f, default) = Pipeline(block_lift, f, default)

function block_lift(rt::Runtime, input::AbstractVector, f, default)
    @assert input isa BlockVector
    _block_lift(f, default, input)
end

function _block_lift(f, input)
    cr = cursor(input)
    I = Tuple{typeof(cr)}
    O = Core.Compiler.return_type(f, I)
    output = Vector{O}(undef, length(input))
    @inbounds while next!(cr)
        output[cr.pos] = f(cr)
    end
    output
end

function _block_lift(f, default, input)
    cr = cursor(input)
    I = Tuple{typeof(cr)}
    O = Union{Core.Compiler.return_type(f, I), typeof(default)}
    output = Vector{O}(undef, length(input))
    @inbounds while next!(cr)
        output[cr.pos] = !isempty(cr) ? f(cr) : default
    end
    output
end

"""
    filler(val) :: Pipeline

This pipeline produces a vector filled with the given value.
"""
filler(val) = Pipeline(filler, val)

filler(rt::Runtime, input::AbstractVector, val) =
    fill(val, length(input))

"""
    null_filler() :: Pipeline

This pipeline produces a block vector with empty blocks.
"""
null_filler() = Pipeline(null_filler)

null_filler(rt::Runtime, input::AbstractVector) =
    BlockVector(fill(1, length(input)+1), Union{}[], x0to1)

"""
    block_filler(block::AbstractVector, card::Cardinality) :: Pipeline

This pipeline produces a block vector filled with the given block.
"""
block_filler(block, card::Cardinality=x0toN) = Pipeline(block_filler, block, card)

function block_filler(rt::Runtime, input::AbstractVector, block::AbstractVector, card::Cardinality)
    if isempty(input)
        return BlockVector(:, block[[]], card)
    elseif length(input) == 1
        return BlockVector([1, length(block)+1], block, card)
    else
        len = length(input)
        sz = length(block)
        perm = Vector{Int}(undef, len*sz)
        for k in eachindex(input)
            copyto!(perm, 1 + sz * (k - 1), 1:sz)
        end
        return BlockVector(1:sz:(len*sz+1), block[perm], card)
    end
end

#
# Converting regular vectors to columnar vectors.
#

"""
    adapt_missing() :: Pipeline

This pipeline transforms a vector that contains `missing` elements to a block
vector with `missing` elements replaced by empty blocks.
"""
adapt_missing() = Pipeline(adapt_missing)

function adapt_missing(rt::Runtime, input::AbstractVector)
    if !(Missing <: eltype(input))
        return BlockVector(:, input, x0to1)
    end
    sz = 0
    for elt in input
        if elt !== missing
            sz += 1
        end
    end
    O = Base.nonmissingtype(eltype(input))
    if sz == length(input)
        return BlockVector(:, collect(O, input), x0to1)
    end
    offs = Vector{Int}(undef, length(input)+1)
    elts = Vector{O}(undef, sz)
    @inbounds offs[1] = top = 1
    @inbounds for k in eachindex(input)
        elt = input[k]
        if elt !== missing
            elts[top] = elt
            top += 1
        end
        offs[k+1] = top
    end
    return BlockVector(offs, elts, x0to1)
end

"""
    adapt_vector() :: Pipeline

This pipeline transforms a vector with vector elements to a block vector.
"""
adapt_vector() = Pipeline(adapt_vector)

function adapt_vector(rt::Runtime, input::AbstractVector)
    @assert eltype(input) <: AbstractVector
    if eltype(input) === Union{}
        return BlockVector(fill(1, length(input)+1), Union{}[], x0toN)
    end
    sz = 0
    for v in input
        sz += length(v)
    end
    O = eltype(eltype(input))
    offs = Vector{Int}(undef, length(input)+1)
    elts = Vector{O}(undef, sz)
    @inbounds offs[1] = top = 1
    @inbounds for k in eachindex(input)
        v = input[k]
        copyto!(elts, top, v)
        top += length(v)
        offs[k+1] = top
    end
    return BlockVector(offs, elts, x0toN)
end

"""
    adapt_tuple() :: Pipeline

This pipeline transforms a vector of tuples to a tuple vector.
"""
adapt_tuple() = Pipeline(adapt_tuple)

function adapt_tuple(rt::Runtime, input::AbstractVector)
    @assert eltype(input) <: Union{Tuple,NamedTuple}
    lbls = Symbol[]
    I = eltype(input)
    if typeof(I) == DataType && I <: NamedTuple
        lbls = collect(Symbol, I.parameters[1])
        I = I.parameters[2]
    end
    cols = _adapt_tuple(input, Val(Tuple{I.parameters...}))
    TupleVector(lbls, length(input), cols)
end

@generated function _adapt_tuple(input, vty)
    Is = (vty.parameters[1].parameters...,)
    D = length(Is)
    return quote
        len = length(input)
        @nexprs $D j -> col_j = Vector{$Is[j]}(undef, len)
        @inbounds for k in eachindex(input)
            t = input[k]
            @nexprs $D j -> col_j[k] = t[j]
        end
        @nref $D AbstractVector j -> col_j
    end
end

"""
    assert_type(T::Type, lbl) :: Pipeline

This pipeline asserts the types of the elements of the input vector.
"""
assert_type(T::Type, lbl::Union{Symbol,Nothing}) =
    Pipeline(assert_type, T, lbl)

assert_type(T::Type) =
    Pipeline(assert_type, T)

function assert_type(rt::Runtime, input::AbstractVector, T, lbl=nothing)
    if eltype(input) <: T
        return input
    end
    if isempty(input)
        return T[]
    end
    output = Vector{T}(undef, length(input))
    @inbounds for i in eachindex(input)
        val = input[i]
        val isa T || error(lbl !== nothing ? "\"$lbl\"" : "", "[$i]: expected a value of type $T; got $(typeof(val))")
        output[i] = val
    end
    output
end

#
# Identity and composition.
#

"""
    pass() :: Pipeline

This pipeline returns its input unchanged.
"""
pass() = Pipeline(pass)

pass(rt::Runtime, @nospecialize input::AbstractVector) =
    input

"""
    chain_of(p₁::Pipeline, p₂::Pipeline … pₙ::Pipeline) :: Pipeline

This pipeline sequentially applies `p₁`, `p₂` … `pₙ`.
"""
function chain_of end

chain_of() = pass()

chain_of(p) = p

function chain_of(ps...)
    ps′ = filter(p -> !(p isa Pipeline && p.op == pass), collect(ps))
    isempty(ps′) ? pass() : length(ps′) == 1 ? ps′[1] : chain_of(ps′)
end

chain_of(ps::Vector) =
    Pipeline(chain_of, ps)

quoteof(::typeof(chain_of), args::Vector{Any}) =
    if length(args) == 1 && args[1] isa Vector
        Expr(:call, chain_of, Any[quoteof(arg) for arg in args[1]]...)
    else
        Expr(:call, chain_of, Any[quoteof(arg) for arg in args]...)
    end

function chain_of(rt::Runtime, @nospecialize(input::AbstractVector), ps)
    output = input
    for p in ps
        output = p(rt, output)
    end
    output
end

#
# Operations on tuple vectors.
#

"""
    tuple_of(p₁::Pipeline, p₂::Pipeline … pₙ::Pipeline) :: Pipeline

This pipeline produces an n-tuple vector, whose columns are generated by
applying `p₁`, `p₂` … `pₙ` to the input vector.
"""
tuple_of(ps...) =
    tuple_of(Symbol[], collect(ps))

tuple_of(lps::Pair{Symbol}...) =
    tuple_of(Symbol[first(lp) for lp in lps], Pipeline[last(lp) for lp in lps])

tuple_of(lbls::Vector{Symbol}, ps::Vector) = Pipeline(tuple_of, lbls, ps)

tuple_of(lbls::Vector{Symbol}, width::Int) = Pipeline(tuple_of, lbls, width)

tuple_of(width::Int) = tuple_of(Symbol[], width)

quoteof(::typeof(tuple_of), args::Vector{Any}) =
    if length(args) == 2 && args[1] isa Vector{Symbol} && args[2] isa Vector
        if isempty(args[1])
            Expr(:call, tuple_of, Any[quoteof(p) for p in args[2]]...)
        else
            Expr(:call, tuple_of, Any[quoteof(l => p) for (l, p) in zip(args[1], args[2])]...)
        end
    else
        Expr(:call, tuple_of, Any[quoteof(arg) for arg in args]...)
    end

function tuple_of(rt::Runtime, @nospecialize(input::AbstractVector), lbls::Vector{Symbol}, ps::Vector{Pipeline})
    len = length(input)
    cols = AbstractVector[p(rt, input) for p in ps]
    TupleVector(lbls, len, cols)
end

function tuple_of(rt::Runtime, @nospecialize(input::AbstractVector), lbls::Vector{Symbol}, width::Int)
    len = length(input)
    cols = AbstractVector[input for k in 1:width]
    TupleVector(lbls, len, cols)
end

"""
    column(lbl::Union{Int,Symbol}) :: Pipeline

This pipeline extracts the specified column of a tuple vector.
"""
column(lbl::Union{Int,Symbol}) = Pipeline(column, lbl)

function column(rt::Runtime, input::AbstractVector, lbl)
    @assert input isa TupleVector
    j = locate(input, lbl)
    column(input, j)
end

"""
    with_column(lbl::Union{Int,Symbol}, p::Pipeline) :: Pipeline

This pipeline transforms a tuple vector by applying `p` to the specified
column.
"""
with_column(lbl::Union{Int,Symbol}, p) = Pipeline(with_column, lbl, p)

function with_column(rt::Runtime, input::AbstractVector, lbl, p)
    @assert input isa TupleVector
    j = locate(input, lbl)
    cols′ = copy(columns(input))
    cols′[j] = p(rt, cols′[j])
    TupleVector(labels(input), length(input), cols′)
end

#
# Operations on block vectors.
#

"""
    wrap() :: Pipeline

This pipeline produces a block vector with one-element blocks wrapping the
values of the input vector.
"""
wrap() = Pipeline(wrap)

wrap(rt::Runtime, input::AbstractVector) =
    BlockVector(:, input, x1to1)

"""
    with_elements(p::Pipeline) :: Pipeline

This pipeline transforms a block vector by applying `p` to its vector of
elements.
"""
with_elements(p) = Pipeline(with_elements, p)

function with_elements(rt::Runtime, input::AbstractVector, p)
    @assert input isa BlockVector
    BlockVector(offsets(input), p(rt, elements(input)), cardinality(input))
end

"""
    flatten() :: Pipeline

This pipeline flattens a nested block vector.
"""
flatten() = Pipeline(flatten)

function flatten(rt::Runtime, input::AbstractVector)
    @assert input isa BlockVector && elements(input) isa BlockVector
    offs = offsets(input)
    nested = elements(input)
    nested_offs = offsets(nested)
    elts = elements(nested)
    card = cardinality(input)|cardinality(nested)
    BlockVector(_flatten(offs, nested_offs), elts, card)
end

_flatten(offs1::AbstractVector{Int}, offs2::AbstractVector{Int}) =
    Int[offs2[off] for off in offs1]

_flatten(offs1::OneTo{Int}, offs2::OneTo{Int}) = offs1

_flatten(offs1::OneTo{Int}, offs2::AbstractVector{Int}) = offs2

_flatten(offs1::AbstractVector{Int}, offs2::OneTo{Int}) = offs1

"""
    distribute(lbl::Union{Int,Symbol}) :: Pipeline

This pipeline transforms a tuple vector with a column of blocks to a block
vector with tuple elements.
"""
distribute(lbl) = Pipeline(distribute, lbl)

function distribute(rt::Runtime, input::AbstractVector, lbl)
    @assert input isa TupleVector && column(input, lbl) isa BlockVector
    j = locate(input, lbl)
    _distribute(column(input, j), input, j)
end

function _distribute(col::BlockVector, tv::TupleVector, j)
    lbls = labels(tv)
    cols′ = copy(columns(tv))
    len = length(col)
    card = cardinality(col)
    offs = offsets(col)
    col′ = elements(col)
    if offs isa OneTo{Int}
        cols′[j] = col′
        return BlockVector{card}(offs, TupleVector(lbls, len, cols′))
    end
    len′ = length(col′)
    perm = Vector{Int}(undef, len′)
    l = r = 1
    @inbounds for k = 1:len
        l = r
        r = offs[k+1]
        for n = l:r-1
            perm[n] = k
        end
    end
    for i in eachindex(cols′)
        cols′[i] =
            if i == j
                col′
            else
                cols′[i][perm]
            end
    end
    return BlockVector{card}(offs, TupleVector(lbls, len′, cols′))
end

"""
    distribute_all() :: Pipeline

This pipeline transforms a tuple vector with block columns to a block vector
with tuple elements.
"""
distribute_all() = Pipeline(distribute_all)

function distribute_all(rt::Runtime, input::AbstractVector)
    @assert input isa TupleVector && all(Bool[col isa BlockVector for col in columns(input)])
    cols = columns(input)
    _distribute_all(labels(input), length(input), cols...)
end

@generated function _distribute_all(lbls::Vector{Symbol}, len::Int, cols::BlockVector...)
    D = length(cols)
    CARD = x1to1
    for col in cols
        CARD |= cardinality(col)
    end
    return quote
        @nextract $D offs (d -> offsets(cols[d]))
        @nextract $D elts (d -> elements(cols[d]))
        if @nall $D (d -> offs_d isa OneTo{Int})
            return BlockVector{$CARD}(:, TupleVector(lbls, len, AbstractVector[(@ntuple $D elts)...]))
        end
        len′ = 0
        regular = true
        @inbounds for k = 1:len
            sz = @ncall $D (*) (d -> (offs_d[k+1] - offs_d[k]))
            len′ += sz
            regular = regular && sz == 1
        end
        if regular
            return BlockVector{$CARD}(:, TupleVector(lbls, len, AbstractVector[(@ntuple $D elts)...]))
        end
        offs′ = Vector{Int}(undef, len+1)
        @nextract $D perm (d -> Vector{Int}(undef, len′))
        @inbounds offs′[1] = top = 1
        @inbounds for k = 1:len
            @nloops $D n (d -> offs_{$D-d+1}[k]:offs_{$D-d+1}[k+1]-1) begin
                @nexprs $D (d -> perm_{$D-d+1}[top] = n_d)
                top += 1
            end
            offs′[k+1] = top
        end
        cols′ = @nref $D AbstractVector (d -> elts_d[perm_d])
        return BlockVector{$CARD}(offs′, TupleVector(lbls, len′, cols′))
    end
end

"""
    block_cardinality(card::Cardinality, src_lbl, tgt_lbl) :: Pipeline

This pipeline asserts the cardinality of the input block vector.
"""
block_cardinality(card::Cardinality, src_lbl::Union{Symbol,Nothing}, tgt_lbl::Union{Symbol,Nothing}) =
    Pipeline(block_cardinality, card, src_lbl, tgt_lbl)

block_cardinality(card::Cardinality) =
    Pipeline(block_cardinality, card)

function block_cardinality(rt::Runtime, input::AbstractVector, card, src_lbl=nothing, tgt_lbl=nothing)
    @assert input isa BlockVector
    cardinality(input) != card || return input
    if !fits(cardinality(input), card) && !(offsets(input) isa OneTo)
        offs = offsets(input)
        @inbounds off = offs[1]
        for k = 2:lastindex(offs)
            @inbounds off′ = offs[k]
            !issingular(card) || off′ <= off+1 || _cardinality_error(src_lbl, tgt_lbl, "singular")
            !ismandatory(card) || off′ >= off+1 || _cardinality_error(src_lbl, tgt_lbl, "mandatory")
            off = off′
        end
    end
    card != x1to1 || return BlockVector(:, elements(input))
    @inbounds output = BlockVector(offsets(input), elements(input), card)
    return output
end

_cardinality_error(src_lbl, tgt_lbl, kind) =
    error(tgt_lbl !== nothing ? "\"$tgt_lbl\": " : "",
          "expected a $kind value",
          src_lbl !== nothing ? ", relative to \"$src_lbl\"" : "")

"""
    block_length() :: Pipeline

This pipeline converts a block vector to a vector of block lengths.
"""
block_length() = Pipeline(block_length)

function block_length(rt::Runtime, input::AbstractVector)
    @assert input isa BlockVector
    _block_length(offsets(input))
end

_block_length(offs::OneTo{Int}) =
    fill(1, length(offs)-1)

function _block_length(offs::AbstractVector{Int})
    len = length(offs) - 1
    output = Vector{Int}(undef, len)
    @inbounds for k = 1:len
        output[k] = offs[k+1] - offs[k]
    end
    output
end

"""
    block_not_empty() :: Pipeline

This pipeline converts a block vector to a vector of Boolean values, where each
value indicates whether the corresponding block is empty or not.
"""
block_not_empty() = Pipeline(block_not_empty)

function block_not_empty(rt::Runtime, input::AbstractVector)
    @assert input isa BlockVector
    _block_not_empty(offsets(input))
end

_block_not_empty(offs::OneTo{Int}) =
    fill(true, length(offs)-1)

function _block_not_empty(offs::AbstractVector{Int})
    len = length(offs) - 1
    output = Vector{Bool}(undef, len)
    @inbounds for k = 1:len
        output[k] = offs[k+1] > offs[k]
    end
    output
end

"""
    block_any() :: Pipeline

This pipeline applies `any` to a block vector with `Bool` elements.
"""
block_any() = Pipeline(block_any)

function block_any(rt::Runtime, input::AbstractVector)
    @assert input isa BlockVector && eltype(elements(input)) <: Bool
    len = length(input)
    offs = offsets(input)
    elts = elements(input)
    if offs isa OneTo
        return elts
    end
    output = Vector{Bool}(undef, len)
    l = r = 1
    @inbounds for k = 1:len
        val = false
        l = r
        r = offs[k+1]
        for i = l:r-1
            if elts[i]
                val = true
                break
            end
        end
        output[k] = val
    end
    return output
end

#
# Filtering.
#

"""
    sieve_by() :: Pipeline

This pipeline filters a vector of pairs by the second column.  It expects a
pair vector, whose second column is a `Bool` vector, and produces a block
vector with 0- or 1-element blocks containing the elements of the first column.
"""
sieve_by() = Pipeline(sieve_by)

function sieve_by(rt::Runtime, input::AbstractVector)
    @assert input isa TupleVector && width(input) == 2 && eltype(column(input, 2)) <: Bool
    val_col, pred_col = columns(input)
    _sieve_by(val_col, pred_col)
end

function _sieve_by(@nospecialize(v), bv)
    len = length(bv)
    sz = count(bv)
    if sz == len
        return BlockVector(:, v, x0to1)
    elseif sz == 0
        return BlockVector(fill(1, len+1), v[[]], x0to1)
    end
    offs = Vector{Int}(undef, len+1)
    perm = Vector{Int}(undef, sz)
    @inbounds offs[1] = top = 1
    for k = 1:len
        if bv[k]
            perm[top] = k
            top += 1
        end
        offs[k+1] = top
    end
    return BlockVector(offs, v[perm], x0to1)
end

#
# Extracting and slicing.
#

"""
    get_by(N::Int) :: Pipeline

This pipeline extracts the `N`-th element from the given block vector.
"""
get_by(N::Int) =
    Pipeline(get_by, N)

get_by(rt::Runtime, input::AbstractVector, N::Int) =
    get_by(rt, input, N, x0to1)

get_by(N::Int, card::Cardinality) =
    Pipeline(get_by, N, card)

function get_by(::Runtime, input::AbstractVector, N::Int, card::Cardinality)
    @assert input isa BlockVector
    len = length(input)
    offs = offsets(input)
    elts = elements(input)
    sz = 0
    R = 1
    for k = 1:len
        L = R
        @inbounds R = offs[k+1]
        sz += checkindex(Bool, L:R-1, _get_index(L, R-1, N))
    end
    if sz == len
        if N == 1
            elts′ = elts[view(offs, 1:len)]
            return BlockVector(:, elts′, card)
        end
        perm = Vector{Int}(undef, sz)
        R = 1
        for k = 1:len
            L = R
            @inbounds R = offs[k+1]
            @inbounds perm[k] = _get_index(L, R-1, N)
        end
        elts′ = elts[perm]
        return BlockVector(:, elts′, card)
    end
    offs′ = Vector{Int}(undef, len+1)
    perm = Vector{Int}(undef, sz)
    @inbounds offs′[1] = top = 1
    R = 1
    for k = 1:len
        L = R
        @inbounds R = offs[k+1]
        i = _get_index(L, R-1, N)
        if checkindex(Bool, L:R-1, i)
            perm[top] = i
            top += 1
        end
        offs′[k+1] = top
    end
    elts′ = elts[perm]
    return BlockVector(offs′, elts′, card)
end

"""
    get_by() :: Pipeline

This pipeline takes a pair vector of blocks and integers, and returns the first
column indexed by the second column.
"""
get_by() =
    Pipeline(get_by)

function get_by(::Runtime, input::AbstractVector)
    @assert input isa TupleVector
    cols = columns(input)
    @assert length(cols) == 2
    vals, Ns = cols
    @assert vals isa BlockVector
    @assert eltype(Ns) <: Int
    _get_by(elements(vals), offsets(vals), Ns)
end

function _get_by(@nospecialize(elts), offs, Ns)
    len = length(Ns)
    R = 1
    sz = 0
    for k = 1:len
        L = R
        @inbounds N = Ns[k]
        @inbounds R = offs[k+1]
        sz += checkindex(Bool, L:R-1, _get_index(L, R-1, N))
    end
    if sz == len
        perm = Vector{Int}(undef, sz)
        R = 1
        for k = 1:len
            L = R
            @inbounds N = Ns[k]
            @inbounds R = offs[k+1]
            @inbounds perm[k] = _get_index(L, R-1, N)
        end
        elts′ = elts[perm]
        return BlockVector(:, elts′, x0to1)
    end
    offs′ = Vector{Int}(undef, len+1)
    perm = Vector{Int}(undef, sz)
    @inbounds offs′[1] = top = 1
    R = 1
    for k = 1:len
        L = R
        @inbounds N = Ns[k]
        @inbounds R = offs[k+1]
        i = _get_index(L, R-1, N)
        if checkindex(Bool, L:R-1, i)
            perm[top] = i
            top += 1
        end
        offs′[k+1] = top
    end
    elts′ = elts[perm]
    return BlockVector(offs′, elts′, x0to1)
end

@inline _get_index(l, r, n) =
    n >= 0 ? l + n - 1 : r + n + 1

"""
    slice_by(N::Int, inv::Bool=false) :: Pipeline

This pipeline transforms a block vector by keeping the first `N` elements of
each block.  If `inv` is true, the pipeline drops the first `N` elements of
each block.
"""
slice_by(N::Union{Int,Missing}, inv::Bool=false) =
    Pipeline(slice_by, N, inv)

slice_by(N::Union{Int,Missing}, card::Cardinality, inv::Bool=false) =
    Pipeline(splice_by, N, card, inv)

slice_by(rt::Runtime, input::AbstractVector, N::Union{Int,Missing}, inv::Bool) =
    slice_by(rt, input, N, cardinality(input)|x0to1, inv)

function slice_by(rt::Runtime, input::AbstractVector, ::Missing, card::Cardinality, inv::Bool)
    @assert input isa BlockVector
    offs′ = !inv ? offsets(input) : fill(1, length(input)+1)
    elts′ = !inv ? elements(input) : elements(input)[Int[]]
    return BlockVector(offs′, elts′, card)
end

function slice_by(rt::Runtime, input::AbstractVector, N::Int, card::Cardinality, inv::Bool)
    @assert input isa BlockVector
    len = length(input)
    offs = offsets(input)
    elts = elements(input)
    sz = 0
    R = 1
    for k = 1:len
        L = R
        @inbounds R = offs[k+1]
        (l, r) = _slice_range(N, R-L, inv)
        sz += r - l + 1
    end
    if sz == length(elts)
        return BlockVector(offs, elts, card)
    end
    offs′ = Vector{Int}(undef, len+1)
    perm = Vector{Int}(undef, sz)
    @inbounds offs′[1] = top = 1
    R = 1
    for k = 1:len
        L = R
        @inbounds R = offs[k+1]
        (l, r) = _slice_range(N, R-L, inv)
        for j = (L + l - 1):(L + r - 1)
            perm[top] = j
            top += 1
        end
        offs′[k+1] = top
    end
    elts′ = elts[perm]
    return BlockVector(offs′, elts′, card)
end

"""
    slice_by(inv::Bool=false) :: Pipeline

This pipeline takes a pair vector of blocks and integers, and returns the first
column sliced by the second column.
"""
slice_by(inv::Bool=false) =
    Pipeline(slice_by, inv)

function slice_by(rt::Runtime, input::AbstractVector, inv::Bool)
    @assert input isa TupleVector
    cols = columns(input)
    @assert length(cols) == 2
    vals, Ns = cols
    @assert vals isa BlockVector
    @assert eltype(Ns) <: Union{Missing,Int}
    _slice_by(elements(vals), offsets(vals), cardinality(vals), Ns, inv)
end

function _slice_by(@nospecialize(elts), offs, card, Ns, inv)
    card′ = card|x0to1
    len = length(Ns)
    R = 1
    sz = 0
    for k = 1:len
        L = R
        @inbounds N = Ns[k]
        @inbounds R = offs[k+1]
        (l, r) = _slice_range(N, R-L, inv)
        sz += r - l + 1
    end
    if sz == length(elts)
        return BlockVector(offs, elts, card′)
    end
    offs′ = Vector{Int}(undef, len+1)
    perm = Vector{Int}(undef, sz)
    @inbounds offs′[1] = top = 1
    R = 1
    for k = 1:len
        L = R
        @inbounds N = Ns[k]
        @inbounds R = offs[k+1]
        (l, r) = _slice_range(N, R-L, inv)
        for j = (L + l - 1):(L + r - 1)
            perm[top] = j
            top += 1
        end
        offs′[k+1] = top
    end
    elts′ = elts[perm]
    return BlockVector(offs′, elts′, card′)
end

@inline _slice_range(n::Int, l::Int, inv::Bool) =
    if !inv
        (1, n >= 0 ? min(l, n) : max(0, l + n))
    else
        (n >= 0 ? min(l + 1, n + 1) : max(1, l + n + 1), l)
    end

@inline _slice_range(::Missing, l::Int, inv::Bool) =
    !inv ? (1, l) : (1, 0)

#
# Ordering extensions.
#

struct LexOrdering{O<:Base.Ordering} <: Base.Ordering
    order::O
end

Base.@propagate_inbounds function Base.lt(o::LexOrdering, a, b)
    b > zero(b) && (a == zero(a) || Base.lt(o.order, a, b))
end

Base.ordtype(o::LexOrdering, vs::AbstractArray) = ordtype(o.order, vs)

ord_eq(o::Base.ForwardOrdering, a, b) = isequal(a, b)
ord_eq(o::Base.ReverseOrdering, a, b) = ord_eq(o.fwd, b, a)
ord_eq(o::Base.By, a, b) = isequal(o.by(a), o.by(b))
ord_eq(o::Base.Lt, a, b) = !o.lt(a, b) && !o.lt(b, a)
Base.@propagate_inbounds ord_eq(o::Base.Perm, a, b) =
    ord_eq(o.order, o.data[a], o.data[b])   # intentionally does not match `lt(::Perm)`
Base.@propagate_inbounds ord_eq(o::LexOrdering, a, b) =
    isequal(a, b) || (a > zero(a) && b > zero(b) && ord_eq(o.order, a, b))

#
# Grouping.
#

unique_by() = Pipeline(unique_by)

function unique_by(::Runtime, input::AbstractVector)
    @assert input isa BlockVector
    card = cardinality(input)
    offs = offsets(input)
    elts = elements(input)
    len = length(elts)
    perm = collect(1:len)
    sep = falses(len+1)
    for off in offs
        sep[off] = true
    end
    if len > 1
        _group_by!(elts, sep, perm, OneTo(len))
    end
    offs_outer, offs_inner = _partition(offs, sep)
    elts′ = elts[perm[view(offs_inner, 1:length(offs_inner)-1)]]
    BlockVector{card}(offs_outer, elts′)
end

sort_by() = Pipeline(sort_by)

function sort_by(::Runtime, input::AbstractVector)
    @assert input isa BlockVector
    card = cardinality(input)
    offs = offsets(input)
    elts = elements(input)
    @assert elts isa TupleVector
    cols = columns(elts)
    @assert length(cols) == 2
    vals, keys = cols
    len = length(elts)
    perm = collect(1:len)
    sep = falses(len+1)
    for off in offs
        sep[off] = true
    end
    if len > 1
        _group_by!(keys, sep, perm, OneTo(len))
    end
    output = BlockVector{card}(offs, vals[perm])
    output
end

group_by() = Pipeline(group_by)

function group_by(::Runtime, input::AbstractVector)
    @assert input isa BlockVector
    card = cardinality(input)
    offs = offsets(input)
    elts = elements(input)
    @assert elts isa TupleVector
    cols = columns(elts)
    @assert length(cols) == 2
    vals, keys = cols
    len = length(elts)
    perm = collect(1:len)
    sep = falses(len+1)
    for off in offs
        sep[off] = true
    end
    if len > 1
        _group_by!(keys, sep, perm, OneTo(len))
    end
    offs_outer, offs_inner = _partition(offs, sep)
    card′ = card & x1toN
    vals′ = BlockVector{card′}(offs_inner, vals[perm])
    sz = length(vals′)
    keys′ = keys[perm[view(offs_inner, 1:sz)]]
    output = BlockVector{card}(offs_outer, TupleVector(sz, AbstractVector[vals′, keys′]))
    output
end

function _group_by!(keys::TupleVector, sep, perm, idxs)
    for col in columns(keys)
        _group_by!(col, sep, perm, idxs)
    end
end

_group_by!(keys::BlockVector{CARD,OneTo{Int}}, sep, perm, idxs) where {CARD} =
    _group_by!(elements(keys), sep, perm, idxs)

function _group_by!(keys::BlockVector{CARD}, sep, perm, idxs) where {CARD}
    if CARD == x1to1
        return _group_by!(elements(keys), sep, perm, idxs)
    end
    offs = offsets(keys)
    elts = elements(keys)
    pos = 0
    idxs′ = Vector{Int}(undef, length(perm))
    done = false
    while !done
        for k = 1:length(perm)
            idx = idxs[k]
            if idx > 0
                l = offs[idx]
                r = offs[idx+1]
                if l + pos < r
                    idxs′[k] = l + pos
                else
                    idxs′[k] = 0
                end
            else
                idxs′[k] = 0
            end
        end
        done = true
        p = idxs′[perm[1]]
        for k = 2:length(perm)
            q = idxs′[perm[k]]
            if !sep[k] && p != q
                done = false
                break
            end
            p = q
        end
        if !done
            _group_by!(elts, sep, perm, idxs′)
        end
        if CARD == x0to1
            done = true
        end
        pos += 1
    end
end

_group_order(keys, ::OneTo) =
    Base.Perm(Base.Forward, keys)

_group_order(keys, idxs) =
    Base.Perm(LexOrdering(Base.Perm(Base.Forward, keys)), idxs)

function _group_by!(keys, sep, perm, idxs)
    o = _group_order(keys, idxs)
    _group_by!(o, sep, perm)
end

function _group_by!(o::Base.Perm, sep, perm)
    l = 1
    for r = 2:length(sep)
        if sep[r]
            if l < r-1
                sort!(perm, l, r-1, QuickSort, o)
            end
            da = o.data[perm[l]]
            for n = l+1:r-1
                db = o.data[perm[n]]
                if !ord_eq(o.order, da, db)
                    sep[n] = true
                end
                da = db
            end
            l = r
        end
    end
end

function _partition(offs, sep)
    sz = 0
    l = 1
    for k = 1:length(offs)-1
        r = offs[k+1]
        for n = l+1:r
            sz += sep[n]
        end
        l = r
    end
    offs_inner = Vector{Int}(undef, sz+1)
    offs_inner[1] = 1
    offs_outer = Vector{Int}(undef, length(offs))
    offs_outer[1] = top = 1
    l = 1
    for k = 1:length(offs)-1
        r = offs[k+1]
        for n = l+1:r
            if sep[n]
                top += 1
                offs_inner[top] = n
            end
        end
        offs_outer[k+1] = top
        l = r
    end
    return (offs_outer, offs_inner)
end

#
# Pipeline matching.
#

macro match_pipeline(ex)
    return esc(_match_pipeline(ex))
end

function _match_pipeline(@nospecialize ex)
    if Meta.isexpr(ex, :call, 3) && ex.args[1] == :~
        val = gensym()
        p = ex.args[2]
        pat = ex.args[3]
        cs = Expr[]
        as = Expr[]
        _match_pipeline!(val, pat, cs, as)
        c = foldl((l, r) -> :($l && $r), cs)
        quote
            local $val = $p
            if $c
                $(as...)
                true
            else
                false
            end
        end
    elseif ex isa Expr
        Expr(ex.head, Any[_match_pipeline(arg) for arg in ex.args]...)
    else
        ex
    end
end

function _match_pipeline!(val, pat, cs, as)
    if pat === :_
        return
    elseif pat isa Symbol
        push!(as, :(local $pat = $val))
        return
    elseif Meta.isexpr(pat, :(::), 2)
        ty = pat.args[2]
        pat = pat.args[1]
        push!(cs, :($val isa $ty))
        _match_pipeline!(val, pat, cs, as)
        return
    elseif Meta.isexpr(pat, :call) && length(pat.args) >= 1 && pat.args[1] isa Symbol
        fn = pat.args[1]
        args = pat.args[2:end]
        push!(cs, :($val isa $DataKnots.Pipeline))
        push!(cs, :($val.op === $DataKnots.$fn))
        _match_pipeline!(:($val.args), args, cs, as)
        return
    elseif Meta.isexpr(pat, :vect)
        push!(cs, :($val isa Vector{$DataKnots.Pipeline}))
        _match_pipeline!(val, pat.args, cs, as)
        return
    end
    error("expected a pipeline pattern; got $(repr(pat))")
end

function _match_pipeline!(val, pats::Vector{Any}, cs, as)
    minlen = 0
    varlen = false
    for pat in pats
        if Meta.isexpr(pat, :..., 1)
            !varlen || error("duplicate vararg pattern in $(repr(:([$pat])))")
            varlen = true
        else
            minlen += 1
        end
    end
    push!(cs, !varlen ? :(length($val) == $minlen) : :(length($val) >= $minlen))
    nearend = false
    for (k, pat) in enumerate(pats)
        if Meta.isexpr(pat, :..., 1)
            pat = pat.args[1]
            k = Expr(:call, :(:), k, Expr(:call, :-, :end, minlen-k+1))
            nearend = true
        elseif nearend
            k = Expr(:call, :-, :end, minlen-k+1)
        end
        _match_pipeline!(:($val[$k]), pat, cs, as)
    end
end


text:
using GNSSTools
using Statistics
using PyPlot
pygui(true)

N = 10
# idx = (rand(1:N), rand(1:N))
idx = (4, 4)
background_val = 1
peak_val = 5
corr = fill(background_val, N, N)
corr[idx[1],idx[2]] = peak_val

# Get codes for plot
unknown_code = circshift(deepcopy(l1ca_codes[1]), 500)
prn_1_code = deepcopy(l1ca_codes[1])
prn_2_code = deepcopy(l1ca_codes[2])
# Make zeros -1
unknown_code[unknown_code .== 0] .= -1
prn_1_code[prn_1_code .== 0] .= -1
prn_2_code[prn_2_code .== 0] .= -1
# Perform auto and cross correlations
auto_correlation = real.(fft_correlate(unknown_code, prn_1_code))
cross_correlation = real.(fft_correlate(unknown_code, prn_2_code))

xlims = [-50, length(auto_correlation)+50]
ylims = [mean(cross_correlation)-120, maximum(auto_correlation)+53]

# Simulate L1 C/A signal and perform course acquisition
signal = definesignal(define_l1ca_code_type(), 5e6, 1; prn=1, f_d=1000,
                      code_start_idx=1000, include_phase_noise=true)
generatesignal!(signal)
replica = definesignal(define_l1ca_code_type(), 5e6, 1e-3)
fd_est, n0_est, SNR_est, corr_result = courseacquisition(signal, replica, 1;
                                                         return_corrresult=true,
                                                         Δfd=1000)
corr_crop = corr_result[:,n0_est-5:n0_est+5]
idx_crop = argmax(corr_crop)
box_x = [idx_crop[2]-0.5, idx_crop[2]-0.5, idx_crop[2]+0.5, idx_crop[2]+0.5, idx_crop[2]-0.5] .- 1
box_y = [idx_crop[1]+0.5, idx_crop[1]-0.5, idx_crop[1]-0.5, idx_crop[1]+0.5, idx_crop[1]+0.5] .- 1

title_font_size = 11
axis_label_font_size = 10

fig = figure(figsize=(6.5,2))
ax1 = subplot2grid((2,4), (0,0), rowspan=2, colspan=2)
ax1.imshow(corr_crop, aspect="equal", cmap="gray", vmin=-9*mean(corr_crop),
           vmax=1.5*maximum(corr_crop))
ax1.plot(box_x, box_y, "r")
ax1.set_xticks(Array(1:N).-0.5)
ax1.set_yticks(Array(1:N).-0.5)
rc("grid", linestyle=":", color="gray", linewidth=1)
grid(true)
xlabel("Code Phase", fontsize=axis_label_font_size)
ylabel("Doppler Frequency", fontsize=axis_label_font_size)
tick_params(left=false, bottom=false, labelleft=false, labelbottom=false)
title("(a)")
ax2 = subplot2grid((2,4), (0,2), rowspan=1, colspan=2)
ax2.plot(auto_correlation, "k")
xlim(xlims)
ylim(ylims)
title("(b)")
# xlabel("Time", fontsize=axis_label_font_size)
tick_params(left=false, bottom=false, labelleft=false, labelbottom=false)
ax3 = subplot2grid((2,4), (1,2), rowspan=1, colspan=2)
ax3.plot(cross_correlation, "k")
xlim(xlims)
ylim(ylims)
xlabel("Time", fontsize=axis_label_font_size)
title("(c)")
tick_params(left=false, bottom=false, labelleft=false, labelbottom=false)
subplots_adjust(hspace=0.5, wspace=0.4, top=0.85, left=0.08, right=0.93, bottom=0.1)
savefig("figures/ch3_course_acquisition.pdf", dpi=300)

ext: jl


text:
"""
Exact inference using factors and variable eliminations
"""
struct ExactInference <: InferenceMethod
end
function infer(im::ExactInference, inf::InferenceState{BN}) where {BN<:DiscreteBayesNet}

    bn = inf.pgm
    nodes = names(bn)
    query = inf.query
    evidence = inf.evidence
    hidden = setdiff(nodes, vcat(query, names(evidence)))

    factors = map(n -> Factor(bn, n, evidence), nodes)

    # successively remove the hidden nodes
    # order impacts performance, but we currently have no ordering heuristics
    for h in hidden
        # find the facts that contain the hidden variable
        contain_h = filter(ϕ -> h in ϕ, factors)
        # add the product of those factors to the set
        if !isempty(contain_h)
            # remove those factors
            factors = setdiff(factors, contain_h)
            push!(factors, sum(reduce((*), contain_h), h))
        end
    end

    ϕ = normalize!(reduce((*), factors))
    return ϕ
end
infer(inf::InferenceState{BN}) where {BN<:DiscreteBayesNet} = infer(ExactInference(), inf)
infer(bn::BN, query::NodeNameUnion; evidence::Assignment=Assignment()) where {BN<:DiscreteBayesNet} = infer(ExactInference(), InferenceState(bn, query, evidence))


text:
@defstruct ArgmaxLayer Layer (
  name :: AbstractString = "argmax",
  (dim :: Int = -2, dim != 0),
  (tops :: Vector{Symbol} = Symbol[], length(tops) > 0),
  (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) == length(tops)),
)

type ArgmaxLayerState <: LayerState
  layer :: ArgmaxLayer
  blobs :: Vector{Blob}

  dims  :: Vector{Int}
end

function setup(backend::Backend, layer::ArgmaxLayer, inputs::Vector{Blob}, diffs::Vector{Blob})
  dims = Array{Int}(length(inputs))
  blobs = Array{Blob}(length(inputs))
  for i = 1:length(inputs)
    total_dim = ndims(inputs[i])
    dim = layer.dim < 0 ? layer.dim + total_dim + 1 : layer.dim
    @assert 1 <= dim <= total_dim
    @assert dim != total_dim
    dims[i] = dim
    shape = [size(inputs[i])...]
    shape[dim] = 1
    blobs[i] = make_blob(backend, eltype(inputs[i]), shape...)
  end

  return ArgmaxLayerState(layer, blobs, dims)
end

function forward(backend::CPUBackend, state::ArgmaxLayerState, inputs::Vector{Blob})
  for i = 1:length(inputs)
    input = inputs[i].data
    output = state.blobs[i].data
    pre_dim, mid_dim, post_dim = split_dims(input, state.dims[i])
    for x = 0:pre_dim-1
      for z = 0:post_dim-1
        idx = Int[x + pre_dim*(y + mid_dim*z) for y=0:mid_dim-1] + 1
        maxc = 1
        @inbounds maxval = input[idx[1]]
        for y = 2:length(idx)
          @inbounds val = input[idx[y]]
          if val > maxval
            maxval = val
            maxc = y
          end
        end
        @inbounds output[x + pre_dim*z + 1] = maxc-1
      end
    end
  end
end

function backward(backend::Backend, state::ArgmaxLayerState, inputs::Vector{Blob}, diffs::Vector{Blob})
  # no backward for argmax layer
end

function shutdown(backend::Backend, state::ArgmaxLayerState)
end


text:
using Ipopt, Juniper, JuMP, Cbc, Gurobi, Cbc, XLSX, DataFrames, Dates, CSV
import cordoba; const _CBD = cordoba
import PowerModelsACDC; const _PMACDC = PowerModelsACDC
import PowerModels; const _PM = PowerModels
import FlexPlan; const _FP = FlexPlan
import InfrastructureModels; const _IM = InfrastructureModels
include("../test/data/conv_spec.jl")

#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
#NOTE EENS is set to zero for CORDOBA/FlexPlan!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

ic_mva=500;#MVA - interconnector
owpp_mva=1000#mva
ic_length=600;#km
owpp_km=300#km
candidates=[[0.5,0.4],[0.5,0.4]]#[[ic],[owpp]]=#
d1="BEdata";d2="UKdata"
#800MWh with ic_mva=500;owpp_mva=1000#mva,ic_length=600;#km,owpp_km=300#km,candidates=[[0.5,0.4],[0.5,0.4]],x=3 in storage_costs
##############################
function cordoba_go!(d1,d2,ic_mva,owpp_mva,ic_length,owpp_km,candidates::Vector{Array{Float64,1}}=[[1.0],[1.0]])
    casename = "ic"
    file = "./test/data/input/$casename.m"
        data = _PM.parse_file(file)#load data in PM format
        data=_CBD.additional_candidates(data,candidates)
        _PMACDC.process_additional_data!(data)#add extra DC model data
        ##################### Network cost/elec PARAMETERS ######################
        hoa=_CBD.hoa_datastruct_candidateICs(ic_mva,owpp_mva,ic_length,owpp_km,candidates)
        data=_CBD.pm_dict_candidateICs(data,hoa,candidates)
        converter_parameters_rxb(data)#adjust converter parameters
        ##########################################################################

        #################### Multi-period INPUT PARAMETERS #######################
        n=number_of_hours = 1000 # Number of time points
        scenario = Dict{String, Any}("hours" => number_of_hours, "sc_years" => Dict{String, Any}())
        scenario["sc_years"]["1"] = Dict{String, Any}()
        scenario["sc_years"]["1"]["year"] = 2020#year of data
        scenario["sc_years"]["1"]["start"] = Dates.datetime2epochms(DateTime(scenario["sc_years"]["1"]["year"]))   #seconds from year 0 (2020 = 63745056000000)
        scenario["sc_years"]["1"]["probability"] = 1
        scenario["planning_horizon"] = 30 # in years, to scale generation cost
        ###########################################################################
         # get wind, generation and load time series
        #data, windgenprofile_beuk, gencostid_beuk, gencost_beuk, nFlow_losses = get_profile_data_UKBE(data, scenario)
        data, z0,z1 = _CBD.get_profile_data_sets(d1,d2,data, n, scenario)
        n=number_of_hours = length(z0.time_stamp)
        #set problem dimension
        dim = scenario["hours"] * length(data["scenario"])

         # create a dictionary to pass time series data to data dictionary
        extradata = _CBD.create_profile_sets(dim, data, z0,z1,ic_mva,owpp_mva)

        # Scale cost data
        _FP.scale_cost_data_cordoba!(data, scenario)
        _FP.scale_cost_data_cordoba!(extradata, scenario)

        # Create data dictionary where time series data is included at the right place
        mn_data = _PMACDC.multinetwork_data(data, extradata, Set{String}(["source_type","scenario","name", "source_version", "per_unit"]))

        #select solver
        gurobi = JuMP.optimizer_with_attributes(Gurobi.Optimizer,"OutputFlag" => 1)

        #settings
        s = Dict("output" => Dict("branch_flows" => true), "conv_losses_mp" => false, "process_data_internally" => false)

        #run optimization
        resultACDC = _PMACDC.run_mp_acdctnepopf(mn_data, _PM.DCPPowerModel, gurobi, multinetwork=true; setting = s)

    return resultACDC, mn_data
end


text:
export iauPvm
"""
Modulus of pv-vector.

This function is part of the International Astronomical Union's
SOFA (Standards Of Fundamental Astronomy) software collection.

Status:  vector/matrix support function.

Given:
    pv     double[2][3]   pv-vector

Returned:
    r      double         modulus of position component
    s      double         modulus of velocity component

Called:
    iauPm        modulus of p-vector

This revision:  2013 June 18

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauPvm(double pv[2][3], double *r, double *s)
function iauPvm(pv::AbstractMatrix{<:Real})
    # Preallocate return values
    ref_r = Ref{Float64}(0.0)
    ref_s = Ref{Float64}(0.0)

    ccall((:iauPvm, libsofa_c), Cvoid,
          (Ptr{Cdouble}, Ref{Cdouble}, Ref{Cdouble}),
           convert(Array{Float64, 2}, pv'),
           ref_r, ref_s)

    return ref_r[], ref_s[]
 end


text:
@everywhere using Statistics, Distributions, StatsBase, Random, LinearAlgebra, DelimitedFiles, Distributed

@everywhere using Revise
@everywhere using Pkg
@everywhere Pkg.activate(homedir()*"/GIT/SuSiEGLMM.jl")
@everywhere using SuSiEGLMM

### causal5:pop
@time info=readdlm(homedir()*"/GIT/SuSiEGLMM.jl/testdata/causal5/pop_folder/ascertained_pop_12_10.bim");
@time geno=readdlm(homedir()*"/GIT/SuSiEGLMM.jl/testdata/causal5/pop_folder/ascertained_pop_genotype_12_10.txt";header=true);
@time pheno=readdlm(homedir()*"/GIT/SuSiEGLMM.jl/testdata/causal5/pop_folder/ascertained_pop_phenotype_12_10.txt";header=true); #518 x 4000 snps
snps =[1842,1976,2007,2119,2129]

# # covariate: sex
# C = pheno[1][:,end-1]
# C[C.==1].=0.0
# C[C.==2].=1.0

#last col: trait
# y=convert(Vector{Float64},data[1][:,end])
y=pheno[1][:,end]

# kinship
K=readdlm(homedir()*"/GIT/SuSiEGLMM.jl/testdata/causal5/pop_folder/pop_grm_ped.txt") #518
K0=readdlm(homedir()*"/GIT/SuSiEGLMM.jl/testdata/causal5/pop_folder/pop_grm.txt");
K0=Symmetric(K0);K0=convert(Matrix{Float64}, K0);
 n=size(K,1)
K1=zeros(n,n,2);
K1[:,:,1]=K; K1[:,:,2]=K0;

G= GenoInfo(info[:,2],info[:,1],info[:,3])

# fill out "NA"
X = geno[1][:,6:end]

for j =axes(X,2)
    idx = findall(X[:,j].=="NA")
    X[idx,j].= missing
    X[idx,j] .= mean(skipmissing(X[:,j]))
end

# for j =axes(X,2)
#    println(sum(ismissing.(X[:,j])))
# end

X = convert(Matrix{Float64},X)
n,p = size(X)
# L=3; Π = ones(p)/p
#score test
@time Tstat, pval= SuSiEGLMM.scoreTest(K0,G,y,X;LOCO=false);

    T, S = svdK(K0;LOCO=false)
    Xt, Xt₀, yt,init00= initialization(y,X,ones(n,1),T,S;tol=1e-4)
    T0= computeT(init00,yt,Xt₀,Xt)

@time Tstat1, pval1= SuSiEGLMM.scoreTest(K1,G,y,X)

#susie-glm
@time est0= fineQTL_glm(G,y,X;tol=1e-4,L=5)
# @time est1= fineQTL_glm(G,y,X,C;tol=1e-4)

a1=sum(est0[1].α,dims=2)
a2=sum(est0[2].α,dims=2)

#susie-glmm: verion 1
# T, S= svdK(K;LOCO=false)
# @time est2 = fineQTL_glmm(G,y,X,ones(n,1),T,S;LOCO=false,L=3)
# T1, S1= svdK(K1)
# @time est3 = fineQTL_glmm(G,y,X,ones(n,1),T1,S1;LOCO=true,L=3)

#susie-glmm: version 2
@time est4 =fineQTL_glmm(K,G,y,X;LOCO=false,tol=1e-4,L=5)
a3=sum(est4[1].α,dims=2)
a4=sum(est4[2].α,dims=2)
@time est5 =fineQTL_glmm(K1,G,y,X;tol=1e-4,L=5)

# # K=I
# @time Xt, Ct, yt, init0= initialization(y,X1,ones(n,1),Matrix(1.0I,n,n),ones(n);tol=1e-5)
# @time est1= fineMapping_GLMM(G1,y,X1,ones(n,1),Matrix(1.0I,n,n),ones(n);LOCO=false, tol=1e-5)
# @time res=susieGLMM(10, ones(p)/p,y,X1,ones(n,1),ones(n);tol=1e-5)

# @time glmr=susieGLM(10, ones(p)/p,y,X1,ones(n,1);tol=1e-4)
# #pip
# p=size(X,2)
# [[1.0.-prod(1.0.-est1.α[j,:]) for j =1:p] [1.0.-prod(1.0.-res.α[j,:]) for j =1:p]]
# [1.0.-prod(1.0.-res.α[j,:]) for j =1:p]

######## the same simulation in R-version
Seed(124)

 L=5; B=100;τ2=0.2;
K0=K0+ (abs(eigmin(K0))+0.001)*I
#GLM
b_true=zeros(p);
b_1s=zeros(B);b_2s=zeros(B); b_3s=zeros(B);
b_4s=zeros(B);b_5s=zeros(B);
res=[];Tm=zeros(B);

for j = 1:B
    b_true[1]= randn(1)[1]
    b_true[2]=randn(1)[1]
    b_true[3]=randn(1)[1]
    b_true[4]=randn(1)[1]
    b_true[5]=randn(1)[1]
    b_1s[j] = b_true[1]
    b_2s[j] = b_true[2]
    b_3s[j] = b_true[3]
    b_4s[j] = b_true[4]
    b_5s[j] = b_true[5]
    # X=randn(n,p)
    g=rand(MvNormal(τ2*K0))
    Y= logistic.(X*b_true+g) .>rand(n) #generating binary outcome
    Y=convert(Vector{Float64},Y)

  t0=@elapsed  res0= fineQTL_glm(G,Y,X;tol=1e-4)
    res=[res;res0]; Tm[j]=t0
end

# writedlm("./test/glm-cau1.txt",[b̂ α̂ b_1s])
println("min, median, max times for glm are $(minimum(Tm)), $(median(Tm)),$(maximum(Tm)).")
#min, median, max times for glm are 5.666828984, 7.7632395855,10.995734686. for K0

b̂1=zeros(B);b̂2=zeros(B);b̂3=zeros(B);b̂4=zeros(B);b̂5=zeros(B);
for j=1:B
A = sum(res[2j-1].α.*res[2j-1].ν,dims=2)[:,1]
b̂1[j]=A[1]
b̂2[j]=A[2]
b̂3[j]=A[3];b̂4[j]=A[4];b̂5[j]=A[5]
end
α̂1 = [maximum(res[2j-1].α[1,:]) for j=1:B]
α̂2=[maximum(res[2j-1].α[2,:]) for j=1:B]
α̂3=[maximum(res[2j-1].α[3,:]) for j=1:B]
α̂4=[maximum(res[2j-1].α[4,:]) for j=1:B]
α̂5=[maximum(res[2j-1].α[5,:]) for j=1:B]

# writedlm("./test/susie-glm-3causal-fam.txt",[b_1s b̂1 α̂1 b_2s b̂2 α̂2 b_3s b̂3 α̂3 ])

using UnicodePlots
scatterplot(b_1s,b̂1,title="susie-glm",xlabel= "True effects", ylabel="Posterior estimate1")
scatterplot(b_1s,α̂1, title="susie-glm", xlabel="True effects",ylabel="pip1")
scatterplot(b_2s,b̂2,title="susie-glm",xlabel= "True effects", ylabel="Posterior estimate2")
scatterplot(b_2s,α̂2,title="susie-glm", xlabel="True effects",ylabel="pip2")
scatterplot(b_3s,b̂3,title="susie-glm",xlabel= "True effects", ylabel="Posterior estimate3")
scatterplot(b_3s,α̂3,title="susie-glm", xlabel="True effects",ylabel="pip3")
scatterplot(b_4s,b̂4,title="susie-glm",xlabel= "True effects", ylabel="Posterior estimate3")
scatterplot(b_4s,α̂4,title="susie-glm", xlabel="True effects",ylabel="pip4")
scatterplot(b_5s,b̂5,title="susie-glm",xlabel= "True effects", ylabel="Posterior estimate3")
scatterplot(b_5s,α̂5,title="susie-glm", xlabel="True effects",ylabel="pip5")

#GLMM :scroe test

# n=100; p=10;

# K2=Symmetric(K0)
b_true=zeros(p);
# b_1s=zeros(B); init0=[];
Ps=zeros(p,B); Tscore=zeros(B);tt=zeros(p,B);

# F =cholesky(K2)
# f = svd(F.U)
# T,S = f.Vt, f.S.^2
# T,S = svdK(K0;LOCO=false)
# H=svd(K2);
for j = 1:B

    b_true[1]= b_1s[j]
    b_true[2]= b_2s[j]
    b_true[3]= b_3s[j]
    b_true[4]= b_4s[j]
    b_true[5]= b_5s[j]
    # X=randn(n,p)
    g=rand(MvNormal(τ2*K0))
    # writedlm("./testdata/dataX-julia.csv",X)
    Y= logistic.(X*b_true+g) .>rand(n) #generating binary outcome
    Y=convert(Vector{Float64},Y)
    # writedlm("./testdata/dataY-julia.csv",Y)

    # Xt, Xt₀, yt,init00= initialization(Y,X1,ones(n,1),T,S;tol=1e-4)
    # T0= computeT(init00,yt,Xt₀,Xt)
    # init0=[init0;init00]
    # Ts[:,j]=T0
   t0= @elapsed Ts,P0= scoreTest(K0,G,Y,X;LOCO=false);
   Ps[:,j]=P0; Tscore[j]=t0; tt[:,j]=Ts
end

# writedlm("./test/glmm-scoretest.txt",Ps)
println("min, median, max times for score test are $(minimum(Tscore)),$(median(Tscore)),$(maximum(Tscore)).")
#for thoeretical K
histogram(Ps[1,:],bins=20) #80% at α=0.05
histogram(Ps[2,:],bins=20) # 82%
histogram(Ps[3,:],bins=20) #65%
histogram(Ps[4,:],bins=20)
histogram(Ps[5,:],bins=20)

#susie-GLMM
b_true=zeros(p);
# b_1s=zeros(B);
res1=[]; Tmm=zeros(B)
#  K=Matrix(1.0I,n,n);

for j = 1:B

    # b_true[1]= randn(1)[1]
    b_true[1]= b_1s[j]
    b_true[2]= b_2s[j]
    b_true[3]= b_3s[j]
    b_true[4]= b_4s[j]
    b_true[5]= b_5s[j]
    # X=randn(n,p)

    g=rand(MvNormal(τ2*K0))
    # writedlm("./testdata/dataX-julia.csv",X)
    Y= logistic.(X*b_true+g) .>rand(n) #generating binary outcome
    Y=convert(Vector{Float64},Y)
    # writedlm("./testdata/dataY-julia.csv",Y)
    # T, S = svdK(K;LOCO=false)
    # Xt, Xt₀, yt, init0 = initialization(Y,X,ones(n,1),T,S;tol=1e-4)
    # res10 = susieGLMM(L,ones(p)/p,yt,Xt,Xt₀,S,init0;tol=1e-4)
    # @time res10 =susieGLMM(1,ones(p)/p,Y,X1,ones(n,1),T,S)
   t0= @elapsed res10=fineQTL_glmm(K0,G,Y,X;LOCO=false)
    res1=[res1;res10]; Tmm[j]=t0
end

# writedlm("./test/glmm-score-susie.txt",[b̂ α̂ b_1s])
println("min, median, max times for score test are $(minimum(Tmm)), $(median(Tmm)),$(maximum(Tmm)).")
#min, median, max times for score test are 3.434145138, 7.687690389,34.555872684. for K0
b̂1=zeros(B);b̂2=zeros(B);b̂3=zeros(B);b̂4=zeros(B);b̂5=zeros(B);
for j=1:B
A = sum(res1[2j-1].α.*res1[2j-1].ν,dims=2)[:,1]
b̂1[j]=A[1]
b̂2[j]=A[2]
b̂3[j]=A[3];b̂4[j]=A[4];b̂5[j]=A[5];
end
α̂1 = [maximum(res1[2j-1].α[1,:]) for j=1:B]
α̂2=[maximum(res1[2j-1].α[2,:]) for j=1:B]
α̂3=[maximum(res1[2j-1].α[3,:]) for j=1:B]
α̂4=[maximum(res[2j-1].α[4,:]) for j=1:B]
α̂5=[maximum(res[2j-1].α[5,:]) for j=1:B]

scatterplot(b_1s,b̂1,title="susie-glmm",xlabel= "True effects", ylabel="Posterior estimate1")
scatterplot(b_1s,α̂1,title="susie-glmm", xlabel="True effects",ylabel="pip1")
scatterplot(b_2s,b̂2,title="susie-glmm",xlabel= "True effects", ylabel="Posterior estimate2")
scatterplot(b_2s,α̂2,title="susie-glmm", xlabel="True effects",ylabel="pip2")
scatterplot(b_3s,b̂3,title="susie-glmm",xlabel= "True effects", ylabel="Posterior estimate3")
scatterplot(b_3s,α̂3,title="susie-glmm", xlabel="True effects",ylabel="pip3")
scatterplot(b_4s,b̂4,title="susie-glmm",xlabel= "True effects", ylabel="Posterior estimate3")
scatterplot(b_4s,α̂4,title="susie-glmm", xlabel="True effects",ylabel="pip4")
scatterplot(b_5s,b̂5,title="susie-glmm",xlabel= "True effects", ylabel="Posterior estimate3")
scatterplot(b_5s,α̂5,title="susie-glmm", xlabel="True effects",ylabel="pip5")


text:
module SymbolicTensors

using SymPy
using PyCall

import SymPy: SymbolicObject,Sym,diff
import SymPy: jprint,factor,simplify
import SymPy: _convert

import Base: show
import Base: convert, promote_rule
import Base: getproperty
import Base: hash, ==
import Base: length, size
import Base.iterate
import Base: +, -, *, /, //, \, ^, log

export TensorIndexType, tensor_indices, TensorIndex, TensorHead
export IndexedTensor, TensAdd, TensMul, TensScalar
export TensorSymmetry
#export @heads,@indices
export @indices, @heads
export tensor
export diff
export canon_bp, contract_metric, factor, simplify
export get_tsymmetry
export scalarIsEqual
export replace_with_arrays
export Quote
export get_christoffel, get_riemann, get_ricci, get_ricci_scalar, get_einstein
export contract_delta

include("types.jl")
include("tensor.jl")
include("indices.jl")
include("derivatives.jl")
include("quoting.jl")
include("gr.jl")
include("replace.jl")

const tensor = PyCall.PyNULL()
const toperations = PyCall.PyNULL()
const get_tsymmetry = convert(TensorSymmetry,PyCall.PyNULL())

function __init__()
    copy!(tensor,PyCall.pyimport_conda("sympy.tensor.tensor","sympy"))
    copy!(get_tsymmetry.__pyobject__,tensor.TensorSymmetry)
    #pytype_mapping(tensor.TensorSymmetry,TensorSymmetry)
#    copy!(toperations,PyCall.pyimport_conda("sympy.tensor.toperations","sympy"))
end

end


text:
# The file calculate the uniform induced velocity
# 给出均匀入流模型

function uniforminflow(ψ,θcp,θlat,θlon,betanow,betaxnow,Mbeta_aero=0.0)
  vind_s = [0.0,0.0,-T/(2*ρ*A*norm(v_air))] # 固定坐标系诱导速度初值
  vall_s = v_air + vind_s # 固定坐标系总的滑流速度初值
  # lmdaui = -T/(2*ρ*A*norm(vall_s))/(Ω*R) # 前飞均匀诱导速度系数
  vindui_s = [0.0,0.0,-T/(2*ρ*A*norm(vall_s))]  # 第一步迭代的固定坐标系均匀入流值

  while true  # 迭代求解均匀入流
    if norm(vindui_s-vind_s)<=1e-3
      break
    else
      vind_s = vindui_s
      vall_s = v_air + vind_s
      vindui_s = [0.0,0.0,-T/(2*ρ*A*norm(vall_s))]
    end
  end
  lmdaui = vindui_s[3]/(Ω*R) # 收敛前飞均匀诱导速度系数

  vinduitmp_r = Array{Vector}(NR,Nbe)
  valluitmp_r = Array{Vector}(NR,Nbe)

  # 准静态挥舞解方法
  β = sflapre(ψ,vall_s[3]/(Ω*R),θcp,θlat,θlon)[1]
  dβ = sflapre(ψ,vall_s[3]/(Ω*R),θcp,θlat,θlon)[2]
  vβ_β = [0.0,0.0,-dβ] # 挥舞坐标系下挥舞引起的气流速度分量
  vβ_r = betatoro(vβ_β,β)

  # if Mbeta_aero==0.0 # 入流场无挥舞初始化
  #   vβ_r = [0.0,0.0,0.0]
  #   β = 0.0
  #   dβ = 0.0
  # else
  #   # 时间步进挥舞解方法
  #   betatmp = dflapre(betanow,betaxnow,Mbeta_aero)
  #   β = betatmp[1]
  #   dβ = betatmp[2]
  #   vβ_β = [0.0,0.0,-dβ] # 挥舞坐标系下挥舞引起的气流速度分量
  #   vβ_r = betatoro(vβ_β,β)
  # end

  for k in 1:NR
    ψk = ψ+(k-1)*2*π/NR
    for i in 1:Nbe
      # 将速度从固定坐标系转为旋转坐标系
      vinduitmp_r[k,i] = systoro(vindui_s,ψk)
      vOmega_r = [0,-Ω*rb[k,i],0] # 旋转坐标系下旋转导致的气流分量
      valluitmp_r[k,i] = systoro(v_air,ψk)+vinduitmp_r[k,i]+vOmega_r+vβ_r*(rb[k,i]-ecut*R)
    end
  end

  return vinduitmp_r,valluitmp_r,lmdaui,β,dβ
end

max_stars_repo_name: dingzhw/RotorTrim

max_issues_repo_name: helicopterliu/RotorTrim

max_forks_repo_name: helicopterliu/RotorTrim


text:
module problem117

description = """
Red, green, and blue tiles
Using a combination of black square tiles and oblong tiles chosen from:
 red tiles measuring two units,
 green tiles measuring three units,
 and blue tiles measuring four units,
 it is possible to tile a row measuring five units in length in exactly fifteen different ways.

How many ways can a row measuring fifty units in length be tiled?"""

# Fill function, based on problems 114, 115, 116
function F(n :: Integer, memo = zeros(Int, n))
  # Base case: no space remaining in the row
  if n < 2 return 0 end

  # Recursive case, count 1 for each placement of block at positions 1:n-m+1, and all ways of
  # filling the remaining space at the end of the row after each placement, memoizing results.
  if memo[n] == 0
    total = 0
    for m=2:4, blockpos = 1:(n-m+1)
      total += 1 + F(n - blockpos - m + 1, memo)
    end
    memo[n] = total
  end
  return memo[n]
end

function mixed_fill_count(n :: Integer)
  return 1 + F(n) # +1 for empty row being a valid solution
end

using Base.Test

@test mixed_fill_count(5) == 15

end

max_stars_repo_name: mbuhot/mbuhot-euler-solutions

max_issues_repo_name: mbuhot/mbuhot-euler-solutions

max_forks_repo_name: mbuhot/mbuhot-euler-solutions


text:
## sin contractor: alters x and y
"""
Contractor for "main branch" of sin, from x = -π/2 to π/2.
"""
function sin_main(X::IntervalBox)

    x, y = X

    x_range = Interval(-half_pi.hi, half_pi.hi)
    y_range = -1..1

    x = x ∩ x_range
    y = y ∩ y_range

    isempty(IntervalBox(x, y)) && return IntervalBox(x, y)

    y = y ∩ sin(x)
    x = x ∩ asin(y)

    return IntervalBox(x, y)

end

# TODO: Be careful with the pi constants if using e.g. BigFloats
sin_reverse = symmetrise(sin_main, reflect_x(half_pi))

"""
    sin!(X::IntervalBox)

Contractor for `sin`.
Takes an `IntervalBox` containing the `x` and `y` component intervals.
Returns an `IntervalBox` contracted down to the set ``y = \\sin(x)``.
"""
sin!(X::IntervalBox) = periodise(sin_main, two_pi)(X) ∪ periodise(sin_reverse, two_pi)(X)

# Reverse function for sin; does not alter y
"""
    sin_rev(y::Interval, x::Interval)

Reverse function for `sin`:
- find the subset of `x` such that ``y = \\sin(x)`` for the given `y`.
"""
function sin_rev(y::Interval, x::Interval)

    X = IntervalBox(x, y)

    X_new = sin!(X)

    return X_new[2], X_new[1]   # return in order y, x
end

## cos contractor: alters x and y
"""
Contractor for main branch of cos, from x = 0 to π.
"""
function cos_main(X::IntervalBox)

    x, y = X

    x_range = Interval(0, Interval{Float64}(π).lo)
    y_range = -1..1

    x = x ∩ x_range
    y = y ∩ y_range

    isempty(IntervalBox(x, y)) && return IntervalBox(x, y)

    y = y ∩ cos(x)
    x = x ∩ acos(y)

    return IntervalBox(x, y)

end

# TODO: Be careful with the pi constants if using e.g. BigFloats
cos_reverse = symmetrise(cos_main, reflect_x(0.0))

"""
    cos!(X::IntervalBox)

Contractor for `cos`.
Takes an `IntervalBox` containing the `x` and `y` component intervals.
Returns an `IntervalBox` contracted down to the set ``y = \\cos(x)``.
"""
cos!(X::IntervalBox) = periodise(cos_main, two_pi)(X) ∪ periodise(cos_reverse, two_pi)(X)

# Reverse function for cos; does not alter y
"""
    cos_rev(y::Interval, x::Interval)

Reverse function for `cos`:
- find the subset of `x` such that ``y = \\cos(x)`` for the given `y`.
"""
function cos_rev(y::Interval, x::Interval)

    X = IntervalBox(x, y)

    X_new = cos!(X)

    return X_new[2], X_new[1]   # return in order y, x
end

"""
Contractor for "main branch" of tan, from x = -π/2 to π/2.
"""
function tan_main(X::IntervalBox)

    x, y = X

    x_range = Interval(-half_pi.hi, half_pi.hi)

    x = x ∩ x_range

    isempty(x) && return IntervalBox(x, y)

    y = y ∩ tan(x)
    x = x ∩ atan(y)

    return IntervalBox(x, y)

end

tan!(X::IntervalBox) = periodise(tan_main, Interval{Float64}(π))(X)

"""
    tan_rev(y::Interval, x::Interval)

Reverse function for `tan`:
- find the subset of `x` such that ``y = \\tan(x)`` for the given `y`.
"""
function tan_rev(y::Interval, x::Interval)

    X = IntervalBox(x, y)

    X_new = tan!(X)

    return X_new[2], X_new[1]   # return in order y, x
end


text:
module Mean

require("Corpora.jl")
require("CommonTools.jl")
require("CVMAbstract.jl")
require("DataTypes.jl")

## Imports
using DataTypes
using Base.LinAlg.BLAS.axpy!
using Corpora.Token
using Corpora.text
using CommonTools: tanh, square!, l2sq, l2sq_distance, mult_by_0!, add!
using CommonTools.CVMIntermediates

using CVMAbstractModel
import CVMAbstractModel: compose!, store_gradients!

## Exports
export MeanModel
export compose, step!, compose_w_l2
export compose!
export store_gradients!
export MeanParameters

immutable MeanModel <: AbstractModel
    wordvecdim::Int64;
    margin::Float;
    regularizer::Float;
    lr_init::Float;
    lr_min::Float;
    niter::Int64;
    gamma::Float;
    nnoise::Int64;
    batchsize::Int64;
end

type MeanParameters <: AbstractParameters
    wordvecs::Matrix{Float};
    gradients::Matrix{Float};
end

MeanParameters() = MeanParameters(Array(Float,0,0), Array(Float,0,0))

function compose_w_l2(model::MeanModel, s::Vector{Token},
                            meanvec::Vector{Float})
    """
        Compose a vector representation for a sentence while calculating
        the l2 norm of the sentence's word vectors at the same time.
        The sentence is given as a list of tokens.
    """

    Base.LinAlg.BLAS.scal!(model.wordvecdim,0.0,meanvec,1);
    l2reg = 0.0
    for iword=1:size(s,1)
        v1 = s[iword].vector
        axpy!(1.0, v1, meanvec)
        l2reg += l2sq(v1)*0.5
    end
    Base.LinAlg.BLAS.scal!(model.wordvecdim,1.0/size(s,1),meanvec,1);
    return meanvec, l2reg
end

function compose_w_l2(model::MeanModel, s::Vector{Token})
    return compose_w_l2(model, s, zeros(Float, model.wordvecdim));
end

function compose!(model::MeanModel, s::Vector{Token}, o::Vector{Float})
    """
        Compose a sentence vector from a list of tokens while
        storing intermediates results.
        In-place
    """
    mult_by_0!(o);
    for iword=1:size(s,1)
        add!(s[iword].vector, o)
    end
    Base.LinAlg.BLAS.scal!(model.wordvecdim,1.0/size(s,1),o,1);
    return o
end

function compose!(model::MeanModel, s::Vector{Token},
                                intermediate::CVMIntermediates)
    return compose!(model, s, intermediate.result)
end

function compose(model::MeanModel, s::Vector{Token})
    """
        Compose a sentence vector from a list of tokens while
        storing intermediates results.
    """
    o = zeros(Float, model.wordvecdim)
    for iword=1:size(s,1)
        o += s[iword].vector
    end
    Base.LinAlg.BLAS.scal!(model.wordvecdim,1.0/size(s,1),o,1);
    return o
end

function store_gradients!(model::MeanModel,
                    pre_grad::Vector{Float},
                    s::Vector{Token}, intermediate::CVMIntermediates;
                    intermediates_precomputed::Bool=false)
    """
        Given a gradient from a higher level in the neural network,
        back-propagate the gradient through to the word representations.
    """
    scaling = 1.0/size(s,1)
    for token in s
        # Contrastive
        #add!(pre_grad, token.gradient)
        axpy!(scaling, pre_grad, token.gradient)
    end
end

end

max_stars_repo_name: ogh/binclusion

max_issues_repo_name: ogh/binclusion

max_forks_repo_name: ogh/binclusion


text:
#!/usr/bin/env julia
using ArgParse

s = ArgParseSettings()

@add_arg_table s begin
    "--geometry"
        required = true
    "--N"
        required = true
        arg_type = Int
    "--a"
        required = true
        arg_type = Float64
    "--b"
        arg_type = Float64
    "--o"
        required = true
end

parameters = parse_args(s)

# Output
const odir = parameters["o"]
@assert isdir(odir)

# System geometry
using QuantumOptics, CollectiveSpins
const N = parameters["N"]
const a = parameters["a"]
const geomstring = parameters["geometry"]

keyparameters = Dict(
        "a"=>parameters["a"],
        "N"=>parameters["N"],
    )

omega_eff = NaN
gamma_eff = NaN

tic()
if geomstring=="chain_orthogonal"
    omega_eff, gamma_eff = CollectiveSpins.effective_interaction.chain_orthogonal(a, N)
elseif geomstring=="squarelattice_orthogonal"
    omega_eff, gamma_eff = CollectiveSpins.effective_interaction.squarelattice_orthogonal(a, N)
elseif geomstring=="hexagonallattice_orthogonal"
    omega_eff, gamma_eff = CollectiveSpins.effective_interaction.hexagonallattice_orthogonal(a, N)
elseif geomstring=="cubiclattice_orthogonal"
    omega_eff, gamma_eff = CollectiveSpins.effective_interaction.cubiclattice_orthogonal(a, N)
elseif geomstring=="tetragonallattice_orthogonal"
    const b = parameters["b"]
    omega_eff, gamma_eff = CollectiveSpins.effective_interaction.tetragonallattice_orthogonal(a, b, N)
    keyparameters["b"]  = parameters["b"]
elseif geomstring=="hexagonallattice3d_orthogonal"
    const b = parameters["b"]
    omega_eff, gamma_eff = CollectiveSpins.effective_interaction.hexagonallattice3d_orthogonal(a, b, N)
    keyparameters["b"]  = parameters["b"]
else
    error("Unknown geometry.")
end
t = toc()

name = QuantumOptics.io.dict2filename(keyparameters)
f = open(joinpath(odir, name), "w")
write(f, "$omega_eff;$gamma_eff\n")
write(f, "\n#calctime=$t\n")
close(f)


text:
"""
`module  JAC.PairProduction`
    ... a submodel of JAC that contains all methods for computing positron-bound-electron pair production (PEPP) by single-photon
        absorption cross sections and rates; photon + |i(N)> --> |f(N+1)> + e^+.
"""
module PairProduction

    using JAC, ..ManyElectron

    """
    `struct  PairProduction.Settings`
        ... defines a type for the details and parameters in computing positron-bound-electron pair production (PEPP) with single-photon
            absorption (positron) lines; photon + |i(N)> --> |f(N+1)> + e^+.

        + multipoles              ::Array{EmMultipole}           ... Specifies the multipoles of the radiation field that are to be included.
        + gauges                  ::Array{UseGauge}              ... Specifies the gauges to be included into the computations.
        + photonEnergies          ::Array{Float64,1}             ... List of photon energies.
        + printBefore             ::Bool                         ... True, if all energies and lines are printed before their evaluation.
        + selectLines             ::Bool                         ... True, if lines are selected individually for the computations.
        + selectedLines           ::Array{Tuple{Int64,Int64},1}  ... List of lines, given by tupels (inital-level, final-level).
    """
    struct Settings
        multipoles                ::Array{EmMultipole}
        gauges                    ::Array{UseGauge}
        photonEnergies            ::Array{Float64,1}
        printBefore               ::Bool
        selectLines               ::Bool
        selectedLines             ::Array{Tuple{Int64,Int64},1}
    end

    """
    `PairProduction.Settings()`  ... constructor for the default values of pair-production positron line computations.
    """
    function Settings()
        Settings(EmMultipole[], UseGauge[], Float64[], false, false, Tuple{Int64,Int64}[])
    end

    # `Base.show(io::IO, settings::PairProduction.Settings)`  ... prepares a proper printout of the variable settings::PairProduction.Settings.
    function Base.show(io::IO, settings::PairProduction.Settings)
        println(io, "multipoles:               $(settings.multipoles)  ")
        println(io, "gauges:                   $(settings.gauges)  ")
        println(io, "photonEnergies:           $(settings.photonEnergies)  ")
        println(io, "printBefore:              $(settings.printBefore)  ")
        println(io, "selectLines:              $(settings.selectLines)  ")
        println(io, "selectedLines:            $(settings.selectedLines)  ")
    end

    """
    `struct  PairProduction.Channel`  ... defines a type for a positron-bound-electron pair production (PEPP) with single-photon absorption
                                          (positron) channel that specifies all quantum numbers, phases and amplitudes.

        + multipole      ::EmMultipole          ... Multipole of the photon absorption.
        + gauge          ::EmGauge              ... Gauge for dealing with the (coupled) radiation field.
        + kappa          ::Int64                ... partial-wave of the free electron
        + symmetry       ::LevelSymmetry        ... total angular momentum and parity of the scattering state
        + phase          ::Float64              ... phase of the partial wave
        + amplitude      ::Complex{Float64}     ... Photoionization amplitude associated with the given channel.
   """
    struct  Channel
        multipole        ::EmMultipole
        gauge            ::EmGauge
        kappa            ::Int64
        symmetry         ::LevelSymmetry
        phase            ::Float64
        amplitude        ::Complex{Float64}
    end

    """
    `struct  PairProduction.Line`  ... defines a type for a pair-production (positron) line that may include the definition of channels.

        + initialLevel   ::Level                  ... initial-(state) level
        + finalLevel     ::Level                  ... final-(state) level
        + positronEnergy ::Float64                ... Energy of the (outgoing free) positron.
        + photonEnergy   ::Float64                ... Energy of the absorbed photon.
        + crossSection   ::EmProperty             ... Cross section for this pair-production (positron) line.
        + hasChannels    ::Bool                   ... Determines whether the individual (sub-) channels are defined in terms of their
                                                      free-electron energy, kappa, multipole, etc., or not.
        + channels       ::Array{PairProduction.Channel,1}  ... List of PairProduction.Channels of this line.
    """
    struct  Line
        initialLevel     ::Level
        finalLevel       ::Level
        positronEnergy   ::Float64
        photonEnergy     ::Float64
        crossSection     ::EmProperty
        hasChannels      ::Bool
        channels         ::Array{PairProduction.Channel,1}
    end

    """
    `PairProduction.Line()`  ... 'empty' constructor for a pair-production (positron) line between a specified initial and final level.
    """
    function Line()
        Line(Level(), Level(), 0., 0., EmProperty(0., 0.), false, PairProduction[] )
    end

    # `Base.show(io::IO, line::PairProduction.Line)`  ... prepares a proper printout of the variable line::PairProduction.Line.
    function Base.show(io::IO, line::PairProduction.Line)
        println(io, "initialLevel:      $(line.initialLevel)  ")
        println(io, "finalLevel:        $(line.finalLevel)  ")
        println(io, "positronEnergy:    $(line.positronEnergy)  ")
        println(io, "photonEnergy:      $(line.photonEnergy)  ")
        println(io, "crossSection:      $(line.crossSection)  ")
        println(io, "hasChannels:       $(line.hasChannels)  ")
        println(io, "channels:          $(line.channels)  ")
    end

end # module


text:
################################################################################
#
#   HYPERKAGOME LATTICE
#
################################################################################

# Implementation
# - implementation 1 (only NN)
# - labels <: Any
function getUnitcellHyperkagome(
			unitcell_type  :: Type{U},
			implementation :: Val{1}
		) :: U where {LS,LB,S<:AbstractSite{LS,3},B<:AbstractBond{LB,3},U<:AbstractUnitcell{S,B}}

	# return a new Unitcell
	return newUnitcell(
		# Type of the unitcell
		U,
		# lattice vectors
		Vector{Float64}[
			Float64[2*sqrt(2), 0, 0],
			Float64[0, 2*sqrt(2), 0],
			Float64[0, 0, 2*sqrt(2)]
		],
		# sites
		S[
            newSite(S, Float64[7/2/sqrt(2), 1/2/sqrt(2), 1/2/sqrt(2)], getDefaultLabelN(LS,1)),
            newSite(S, Float64[5/2/sqrt(2), 3/2/sqrt(2), 1/2/sqrt(2)], getDefaultLabelN(LS,2)),
            newSite(S, Float64[7/2/sqrt(2), 3/2/sqrt(2), 3/2/sqrt(2)], getDefaultLabelN(LS,3)),
            newSite(S, Float64[7/2/sqrt(2), 5/2/sqrt(2), 5/2/sqrt(2)], getDefaultLabelN(LS,4)),
            newSite(S, Float64[5/2/sqrt(2), 5/2/sqrt(2), 7/2/sqrt(2)], getDefaultLabelN(LS,5)),
            newSite(S, Float64[3/2/sqrt(2), 3/2/sqrt(2), 7/2/sqrt(2)], getDefaultLabelN(LS,6)),
            newSite(S, Float64[3/2/sqrt(2), 1/2/sqrt(2), 5/2/sqrt(2)], getDefaultLabelN(LS,7)),
            newSite(S, Float64[1/2/sqrt(2), 1/2/sqrt(2), 7/2/sqrt(2)], getDefaultLabelN(LS,8)),
            newSite(S, Float64[5/2/sqrt(2), 7/2/sqrt(2), 5/2/sqrt(2)], getDefaultLabelN(LS,9)),
            newSite(S, Float64[3/2/sqrt(2), 7/2/sqrt(2), 3/2/sqrt(2)], getDefaultLabelN(LS,10)),
            newSite(S, Float64[1/2/sqrt(2), 7/2/sqrt(2), 1/2/sqrt(2)], getDefaultLabelN(LS,11)),
            newSite(S, Float64[1/2/sqrt(2), 5/2/sqrt(2), 3/2/sqrt(2)], getDefaultLabelN(LS,12)),
		],
		# bonds
		B[
            newBond(B, 1, 2, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 2, 3, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 1, 3, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 3, 4, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 4, 5, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 4, 9, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 5, 9, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 5, 6, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 6, 7, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 6, 8, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 7, 8, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 7, 9, getDefaultLabel(LB), (0, -1, 0)),
            newBond(B, 7, 10, getDefaultLabel(LB), (0, -1, 0)),
            newBond(B, 9, 10, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 10, 11, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 10, 12, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 11, 12, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 1, 8, getDefaultLabel(LB), (1, 0, -1)),
            newBond(B, 1, 11, getDefaultLabel(LB), (1, -1, 0)),
            newBond(B, 3, 12, getDefaultLabel(LB), (1, 0, 0)),
            newBond(B, 4, 12, getDefaultLabel(LB), (1, 0, 0)),
            newBond(B, 5, 2, getDefaultLabel(LB), (0, 0, 1)),
            newBond(B, 6, 2, getDefaultLabel(LB), (0, 0, 1)),
            newBond(B, 8, 11, getDefaultLabel(LB), (0, -1, 1)),

            newBond(B, 2, 1, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 3, 2, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 3, 1, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 4, 3, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 5, 4, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 9, 4, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 9, 5, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 6, 5, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 7, 6, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 8, 6, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 8, 7, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 9, 7, getDefaultLabel(LB), (0, 1, 0)),
            newBond(B, 10, 7, getDefaultLabel(LB), (0, 1, 0)),
            newBond(B, 10, 9, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 11, 10, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 12, 10, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 12, 11, getDefaultLabel(LB), (0, 0, 0)),
            newBond(B, 8, 1, getDefaultLabel(LB), (-1, 0, 1)),
            newBond(B, 11, 1, getDefaultLabel(LB), (-1, 1, 0)),
            newBond(B, 12, 3, getDefaultLabel(LB), (-1, 0, 0)),
            newBond(B, 12, 4, getDefaultLabel(LB), (-1, 0, 0)),
            newBond(B, 2, 5, getDefaultLabel(LB), (0, 0, -1)),
            newBond(B, 2, 6, getDefaultLabel(LB), (0, 0, -1)),
            newBond(B, 11, 8, getDefaultLabel(LB), (0, 1, -1)),
		]
	)
end


text:
function stonemask(
    x::Vector{Float64},
    fs::Int,
    temporal_positions::Vector{Float64},
    f0::Vector{Float64})

    refined_f0 = similar(f0)

    @ccall libworld.StoneMask(
        x::Ptr{Float64},
        length(x)::Cint,
        fs::Cint,
        temporal_positions::Ptr{Float64},
        f0::Ptr{Float64},
        length(f0)::Cint,
        refined_f0::Ptr{Float64})::Cvoid

    refined_f0
end


text:
Base.:(==)(x::nodeTypeFE, y::Real) = ==(Int(x), y)
Base.:(==)(x::Real, y::nodeTypeFE) = ==(x, Int(y))
Base.convert(::Type{nodeTypeFE}, x::Real) = nodeTypeFE(Int(x))
Base.zero(::Type{nodeTypeFE}) = nodeTypeFE(0)
function findSurfaceElemCoord(mesh::AbstractMesh, idx = nothing)
    isnothing(idx) ? idx = (1:size(mesh.surfElemNode, 1)) : nothing
    return reshape(mesh.coord[:, vec(mesh.surfElemNode[idx, :])], :, 4)
end

max_stars_repo_name: dcelisgarza/DDD

max_forks_repo_name: dcelisgarza/DDD


text:
using Pkg
# change to course environment
Pkg.activate(dirname(@__FILE__))
# also change pwd for git functionality!
cd(dirname(@__FILE__))

using GitCommand
# pull latest updates
git() do git
    run(`$git config pull.rebase false`)
    # if on CDN
    #run(`$git config --global http.proxy http://CDNusername:CDNpassword@dmzproxy005.idcn.mil.intra:8080`)
    run(`$git pull`)
end

include("config.jl")


text: include(joinpath("abstract","all.jl"))


text:
using Test
using NeXLMatrixCorrection
using Random

function testIterate(unk, stds, e0, θ)
    toa = deg2rad(θ)
    props = Dict{Symbol,Any}(:BeamEnergy => e0, :TakeOffAngle => toa)
    krs = KRatio[]
    for (xrays, std) in stds
        elm = element(xrays[1])
        zu = zafcorrection(XPP, ReedFluorescence, Coating, unk, xrays, e0)
        zs = zafcorrection(XPP, ReedFluorescence, Coating, std, xrays, e0)
        k = gZAFc(zu, zs, toa, toa) * unk[elm] / std[elm]
        push!(krs, KRatio(xrays, props, props, std, k))
    end
    up = RecordingUpdateRule(NeXLMatrixCorrection.WegsteinUpdateRule())
    iter = Iteration(XPP, ReedFluorescence, Coating, updater = up)
    return quantify(iter, "Result", krs)
end

randomize(mat::Material, qty::Float64)::Material =
    material(mat.name, (elm => mat[elm] * (1.0 + qty * randn()) for elm in keys(mat))...)

function evaluate(mat, stds, e0, θ, tol = 0.0001)
    tr(x) = x>zero(typeof(x)) ? x : zero(typeof(x))
    res = testIterate(mat, stds, e0, θ)
    res.converged || println("$mat: Not converged.")
    for elm in filter(elm->abs(value(res.comp[elm]) - tr(value(mat[elm]))) >= tol, keys(mat))
        println("$elm in $mat: $(value(res.comp[elm])) - $(tr(value(mat[elm]))) > $tol")
    end
    return res.converged && #
        all(abs(value(res.comp[elm]) - tr(value(mat[elm]))) < tol for elm in keys(mat)) && #
        all(value(res.comp[elm]) >= 0 for elm in keys(mat))
end

Base.eps(::Type{UncertainValue}) = eps(Float64)
Base.isapprox(uv1::UncertainValue, uv2::UncertainValue; atol::Real=0.0) =
    isapprox(value(uv1),value(uv2),atol=atol)
Base.isapprox(uv1::UncertainValue, v2::Real; atol::Real=0.0) =
    isapprox(value(uv1), v2, atol=atol)
Base.isapprox(v1::Real, uv2::UncertainValue; atol::Real=0.0) =
    isapprox(v1, value(uv2), atol=atol)

@testset "Iteration tests" begin
    @testset "FeCr series at 15 keV (normalized inputs)" begin
        mat = material("0.6Fe+0.4Cr", n"Fe" => 0.6, n"Cr" => 0.4)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)

        mat = material("0.2Fe+0.8Cr", n"Fe" => 0.2, n"Cr" => 0.8)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)

        mat = material("0.8Fe+0.2Cr", n"Fe" => 0.8, n"Cr" => 0.2)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)

        mat = material("0.99Fe+0.01Cr", n"Fe" => 0.99, n"Cr" => 0.01)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)

        mat = material("0.01Fe+0.99Cr", n"Fe" => 0.01, n"Cr" => 0.99)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00055)
    end

    @testset "FeCr series at 15 keV (non-normalized inputs)" begin
        mat = material("0.62Fe+0.4Cr", n"Fe" => 0.62, n"Cr" => 0.4)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)

        mat = material("0.18Fe+0.84Cr", n"Fe" => 0.18, n"Cr" => 0.84)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)

        mat = material("0.85Fe+0.22Cr", n"Fe" => 0.85, n"Cr" => 0.22)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)

        mat = material("0.95Fe+0.01Cr", n"Fe" => 0.95, n"Cr" => 0.01)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)

        mat = material("0.03Fe+0.99Cr", n"Fe" => 0.03, n"Cr" => 0.99)
        res = testIterate(mat, Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr")), 15.0e3, 40.0)

        @test res.converged
        @test isapprox(res.comp[n"Fe"], mat[n"Fe"], atol = 0.00005)
        @test isapprox(res.comp[n"Cr"], mat[n"Cr"], atol = 0.00005)
    end

    @testset "K240 tests - Pure standards" begin
        mat = material(
            "K240",
            n"O" => 0.340023,
            n"Mg" => 0.030154,
            n"Si" => 0.186986,
            n"Ti" => 0.059950,
            n"Zn" => 0.040168,
            n"Zr" => 0.074030,
            n"Ba" => 0.268689,
        )
        stds = Dict(
            [n"O K-L3"] => pure(n"O"),
            [n"Si K-L3"] => pure(n"Si"),
            [n"Mg K-L3"] => pure(n"Mg"),
            [n"Ba L3-M5"] => pure(n"Ba"),
            [n"Ti K-L3"] => pure(n"Ti"),
            [n"Zn K-L3"] => pure(n"Zn"),
            [n"Zr L3-M5"] => pure(n"Zr"),
        )
        res = testIterate(mat, stds, 20.0e3, 40)
        @test res.converged
        @test isapprox(res.comp[n"O"], mat[n"O"], atol = 0.00005)
        @test isapprox(res.comp[n"Si"], mat[n"Si"], atol = 0.00005)
        @test isapprox(res.comp[n"Mg"], mat[n"Mg"], atol = 0.00005)
        @test isapprox(res.comp[n"Ba"], mat[n"Ba"], atol = 0.00005)
        @test isapprox(res.comp[n"Ti"], mat[n"Ti"], atol = 0.00005)
        @test isapprox(res.comp[n"Zn"], mat[n"Zn"], atol = 0.00005)
        @test isapprox(res.comp[n"Zr"], mat[n"Zr"], atol = 0.00005)
    end

    @testset "K240 tests - Simple standards" begin
        mat = material(
            "K240",
            n"O" => 0.340023,
            n"Mg" => 0.030154,
            n"Si" => 0.186986,
            n"Ti" => 0.059950,
            n"Zn" => 0.040168,
            n"Zr" => 0.074030,
            n"Ba" => 0.268689,
        )
        stds = Dict(
            [n"O K-L3"] => parse(Material, "SiO2"),
            [n"Si K-L3"] => parse(Material, "SiO2"),
            [n"Mg K-L3"] => parse(Material, "MgO"),
            [n"Ba L3-M5"] => parse(Material, "BaF2"),
            [n"Ti K-L3"] => pure(n"Ti"),
            [n"Zn K-L3"] => pure(n"Zn"),
            [n"Zr L3-M5"] => pure(n"Zr"),
        )
        res = testIterate(mat, stds, 20.0e3, 40)
        @test res.converged
        @test isapprox(res.comp[n"O"], mat[n"O"], atol = 0.00005)
        @test isapprox(res.comp[n"Si"], mat[n"Si"], atol = 0.00005)
        @test isapprox(res.comp[n"Mg"], mat[n"Mg"], atol = 0.00005)
        @test isapprox(res.comp[n"Ba"], mat[n"Ba"], atol = 0.00005)
        @test isapprox(res.comp[n"Ti"], mat[n"Ti"], atol = 0.00005)
        @test isapprox(res.comp[n"Zn"], mat[n"Zn"], atol = 0.00005)
        @test isapprox(res.comp[n"Zr"], mat[n"Zr"], atol = 0.00005)
    end
    @testset "K240 tests - Simple standards - EDS mode" begin
        mat = randomize(
            material(
                "K240",
                n"O" => 0.340023,
                n"Mg" => 0.030154,
                n"Si" => 0.186986,
                n"Ti" => 0.059950,
                n"Zn" => 0.040168,
                n"Zr" => 0.074030,
                n"Ba" => 0.268689,
            ),
            0.1,
        )
        stds = Dict(
            characteristic(n"O", ktransitions) => parse(Material, "SiO2"),
            characteristic(n"Si", ktransitions) => parse(Material, "SiO2"),
            characteristic(n"Mg", ktransitions) => parse(Material, "MgO"),
            characteristic(n"Ba", ltransitions) => parse(Material, "BaF2"),
            characteristic(n"Ti", ktransitions) => pure(n"Ti"),
            characteristic(n"Zn", ktransitions) => pure(n"Zn"),
            characteristic(n"Zr", kalpha) => pure(n"Zr"),
        )
        res = testIterate(mat, stds, 20.0e3, 40)
        @test res.converged
        @test isapprox(res.comp[n"O"], mat[n"O"], atol = 0.00005)
        @test isapprox(res.comp[n"Si"], mat[n"Si"], atol = 0.00005)
        @test isapprox(res.comp[n"Mg"], mat[n"Mg"], atol = 0.00005)
        @test isapprox(res.comp[n"Ba"], mat[n"Ba"], atol = 0.00005)
        @test isapprox(res.comp[n"Ti"], mat[n"Ti"], atol = 0.00005)
        @test isapprox(res.comp[n"Zn"], mat[n"Zn"], atol = 0.00005)
        @test isapprox(res.comp[n"Zr"], mat[n"Zr"], atol = 0.00005)
    end
    @testset "K240 tests - Unnorm" begin
        mat = material(
            "K240",
            n"O" => 0.340023,
            n"Mg" => 0.030154,
            n"Si" => 0.186986,
            n"Ti" => 0.059950,
            n"Zn" => 0.040168,
            n"Zr" => 0.074030,
            n"Ba" => 0.268689,
        )
        stds = Dict(
            [n"O K-L3"] => parse(Material, "SiO2"),
            [n"Si K-L3"] => parse(Material, "SiO2"),
            [n"Mg K-L3"] => parse(Material, "MgO"),
            [n"Ba L3-M5"] => parse(Material, "BaF2"),
            [n"Ti K-L3"] => pure(n"Ti"),
            [n"Zn K-L3"] => pure(n"Zn"),
            [n"Zr L3-M5"] => pure(n"Zr"),
        )
        Random.seed!(0xEA7BADF00D0)
        @test evaluate(randomize(mat, 0.01), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.1), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.2), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.3), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.4), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
    end
    @testset "FeCr tests - Unnorm" begin
        mat = material("FeCr", n"Fe" => 0.6, n"Cr" => 0.4)
        stds = Dict([n"Fe K-L3"] => pure(n"Fe"), [n"Cr K-L3"] => pure(n"Cr"))
        Random.seed!(0xEA7BADF00D)
        @test evaluate(randomize(mat, 0.01), stds, 15.0e3, 40)
        @test evaluate(randomize(mat, 0.1), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.2), stds, 17.0e3, 40)
        @test evaluate(randomize(mat, 0.3), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.4), stds, 14.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 16.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 22.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 19.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 17.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 20.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 14.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 26.0e3, 40)
        @test evaluate(randomize(mat, 0.5), stds, 30.0e3, 40)
    end
    @testset "K309 test" begin
        # From a measurement
        props = Dict(:BeamEnergy=>15.0e3, :TakeOffAngle=>deg2rad(40.0))
        fe = KRatio(characteristic(n"Fe",kalpha),props,props,mat"Fe",UncertainValue(0.08884060796139037,0.000363897))
        al = KRatio(characteristic(n"Al",ktransitions),props,props,mat"Al2O3",UncertainValue(0.12212761167159845,0.000257044))
        si = KRatio(characteristic(n"Si",ktransitions),props,props,mat"Si",UncertainValue(0.13931565459523768,0.000175283))
        ba = KRatio(characteristic(n"Ba",ltransitions),props,props,mat"BaF2",UncertainValue(0.14026640286207112,0.000494559))
        o = KRatio(characteristic(n"O",ktransitions),props,props,mat"Al2O3",UncertainValue(0.636480599484999,0.000823425))
        ca = KRatio(characteristic(n"Ca",ktransitions),props,props,mat"CaF2",UncertainValue(0.20597043348028954,0.000417253))
        res = quantify(nl"K309", [al,ba,ca,fe,o,si])
        @test isapprox(res.comp[n"Fe"],0.1038,atol=0.0001) # 0.1036 in DTSA-II
        @test isapprox(res.comp[n"Al"],0.0756,atol=0.0001) # 0.0756
        @test isapprox(res.comp[n"Si"],0.1829,atol=0.0001) # 0.1823
        @test isapprox(res.comp[n"Ba"],0.1404,atol=0.0001) # 0.1402
        @test isapprox(res.comp[n"O"],0.3899,atol=0.0001)  # 0.3899
        @test isapprox(res.comp[n"Ca"],0.1074,atol=0.0001) # 0.1036
        @test isapprox(analyticaltotal(res.comp), 1.0, atol=0.00002) # 0.9986
    end
    @testset "20 nm C on SiO2 at 5 keV" begin
        props = Dict(:BeamEnergy=>5.0e3, :TakeOffAngle=>deg2rad(40.0))
        si = KRatio([n"Si K-L3"],props,props,mat"SiO2",0.9219)
        o = KRatio([n"O K-L3"],props,props,mat"SiO2",0.8950)
        c = KRatio([n"C K-L2"],props,props,mat"C",0.0548)
        res = quantify(nl"20 nm C on SiO2 at 5 keV", [si, o, c], coating=n"C K-L2"=>parse(Material,"C",density=1.9))
        @test isapprox(thickness(si.unkProps[:Coating]), 2.41e-6, atol=0.01e-6)
    end
    @testset "4 nm of Al2O3 on Al at 3 keV" begin
        props = Dict(:BeamEnergy=>3.0e3, :TakeOffAngle=>deg2rad(40.0))
        al = KRatio([n"Al K-L3"], props, props, mat"Al", 4402. / 4657.)
        o = KRatio([n"O K-L3"], props, props, mat"SiO2", 629.7 / 10061.)
        res = quantify(nl"4 nm Al2O3 on Al at 3 keV", [al, o], coating=n"O K-L3"=>parse(Material,"Al2O3",density=3.987))
        @test isapprox(thickness(al.unkProps[:Coating]), 4.85e-7, atol=0.01e-6)
    end
end

max_stars_repo_name: NicholasWMRitchie/NeXLMatrixCorrection

max_issues_repo_name: NicholasWMRitchie/NeXLMatrixCorrection

max_forks_repo_name: NicholasWMRitchie/NeXLMatrixCorrection


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    variation(table)

Return the variation matrix `Τ` of the `table` such that:

- `Τ[i,j] = Var(log(x[i]/x[j]))` for `i, j = 1, ..., D`
"""
function variation(table)
  X = Tables.matrix(table)
  n = Tables.columnnames(table) |> collect
  D = size(X, 2)
  L = log.(X .+ eps())

  T = Matrix{Float64}(undef, D, D)
  for j in 1:D
    for i in j+1:D
      lr = L[:,i] .- L[:,j]
      T[i,j] = var(lr)
    end
    T[j,j] = 0.0
    for i in 1:j-1
      T[i,j] = T[j,i]
    end
  end

  AxisArray(T, row=n, col=n)
end

"""
    alrcov(table)

Return the log-ratio covariance matrix `Σ` of the `table` such that:

- `Σ[i,j] = cov(log(x[i]/x[D]), log(x[j]/x[D]))` for `i, j = 1, ..., d`
"""
function alrcov(table)
  alrtable = table |> ALR()

  Σ = cov(Tables.matrix(alrtable), dims=1)

  vars = Tables.columnnames(alrtable) |> collect
  AxisArray(Σ, row=vars, col=vars)
end

"""
    clrcov(table)

Return the centered log-ratio covariance matrix `Γ` of the `table` such that:

- `Γ[i,j] = cov(log(x[i]/g(x)), log(x[j]/g(x)))` for `i, j = 1, ..., D`,
where `g(x)` is the geometric mean.
"""
function clrcov(table)
  clrtable = table |> CLR()

  Γ = cov(Tables.matrix(clrtable), dims=1)

  vars = Tables.columnnames(clrtable) |> collect
  AxisArray(Γ, row=vars, col=vars)
end

"""
    lrarray(table)

Return the variation array `A` of the `table` such that:

- `A[i,j] = E[log(x[i]/x[j])]` for `i > j`
- `A[i,j] = Var(log(x[i]/x[j]))` for `i < j`
- `A[i,j] = 0` for `i = j`
"""
function lrarray(table)
  X = Tables.matrix(table)
  D = size(X, 2)
  L = log.(X .+ eps())

  A = Matrix{Float64}(undef, D, D)
  for i in 1:D
    for j in i+1:D
      lr = L[:,i] .- L[:,j]
      A[j,i] = mean(lr)
      A[i,j] = var(lr)
    end
    A[i,i] = 0.0
  end

  vars = Tables.columnnames(table) |> collect
  AxisArray(A, row=vars, col=vars)
end


text:
# Alternative implementation

function to_int(x::Vector{UInt8})
    if isempty(x)
        return 0
    end
    if length(x) > div(Sys.WORD_SIZE, 8)
        T = BigInt
    else
        T = Int
    end
    result = zero(T)
    if ENDIAN_BOM == 0x01020304
        reverse!(x)
    end
    for c in x
        result <<= 8
        result += c
    end
    return result
end

function faster_big(x::Vector{UInt8})
    xsize = cld(length(x), Base.GMP.BITS_PER_LIMB / 8)
    if ENDIAN_BOM == 0x04030201
        reverse!(x)
    end
    result = Base.GMP.MPZ.realloc2(xsize * Base.GMP.BITS_PER_LIMB)
    result.size = xsize
    unsafe_copyto!(result.d, convert(Ptr{Base.GMP.Limb}, pointer(x)), xsize)
    return result
end


text:
export DeltaRCWAProblem, solve

"""
    DeltaRCWAProblem(
        stack::T,
        pw::PlaneWaves{N},
        I₁::AbstractArray{<:Number, N},
        I₂::AbstractArray{<:Number, N},
    ) where {N, T<:SheetStack}

This type of problem assumes that the only type of scatterer are sheets and that
the solution can be expressed entirely by one set of pw
"""
struct DeltaRCWAProblem{N, T<:SheetStack}
    stack::T
    pw::PlaneWaves{N}
    I₁::Array{ComplexF64, N}
    I₂::Array{ComplexF64, N}
    function DeltaRCWAProblem(
        stack::T,
        pw::PlaneWaves{N},
        I₁::AbstractArray{<:Number, N},
        I₂::AbstractArray{<:Number, N},
    ) where {N, T<:SheetStack}
        check_mode_size(N, pw.dims, I₁, I₂)
        new{N, T}(stack, pw, I₁, I₂)
    end
end

function check_mode_size(N, dims, weights...)
    for w in weights
        @assert size(w) == Tuple(i == length(dims) ? N * dims[i][1] : dims[i][1] for i in eachindex(dims)) ""
    end
end

"""
    DeltaRCWAProblem(structure, dims, ω, I₁, I₂)

Outer constructor for convenience
"""
function DeltaRCWAProblem(structure, dims, ω, I₁, I₂)
    DeltaRCWAProblem(convert(SheetStack, structure), PlaneWaves(ω, dims), I₁, I₂)
end

FieldStyle(prob::DeltaRCWAProblem) = FieldStyle(typeof(prob))
FieldStyle(::Type{<:DeltaRCWAProblem}) = HField()

struct DeltaRCWASolution{N, T<:SheetStack}
    stack::T
    pw::PlaneWaves{N}
    I₁::Array{ComplexF64, N}
    I₂::Array{ComplexF64, N}
    O₁::Array{ComplexF64, N}
    O₂::Array{ComplexF64, N}
    function DeltaRCWASolution(
        stack::T,
        pw::PlaneWaves{N},
        I₁::AbstractArray{<:Number, N},
        I₂::AbstractArray{<:Number, N},
        O₁::AbstractArray{<:Number, N},
        O₂::AbstractArray{<:Number, N},
    ) where {N, T<:SheetStack}
        check_mode_size(N, pw.dims, I₁, I₂, O₁, O₂)
        new{N, T}(stack, pw, I₁, I₂, O₁, O₂)
    end
end

"""
    solve(::DeltaRCWAProblem; T=Matrix)::DeltaRCWASolution

Evaluates the solution to the problem, with the solver method to be selected by
the user. `T` can be `Matrix`, `BlockMatrix`, `LinearMap`, or `BlockMap` and
selects the underlying representation of the scattering matrices.
"""
function solve(prob::DeltaRCWAProblem{N}; T=Matrix)::DeltaRCWASolution{N} where N
    stack, pw, I₁, I₂ = prob.stack, prob.pw, prob.I₁, prob.I₂
    sol = smatrix(T, FieldStyle(prob), pw, stack) * vcat(reshape.([I₁, I₂], :)...)
    O₁ = reshape(sol[1:length(I₁)], size(I₁))
    O₂ = reshape(sol[(length(I₁)+1):sum(length.([I₁, I₂]))], size(I₂))
    DeltaRCWASolution(stack, pw, I₁, I₂, O₁, O₂)
end


text:
using Test
using NQCBase
using StaticArrays

atoms = Atoms{Float64}([:C, :H])
@test atoms isa Atoms{2, Float64}
@test length(atoms) == 2
@test atoms.numbers == [6, 1]
@test atoms.types == [:C, :H]
@test atoms.masses isa SVector

@test range(atoms) == 1:2

@test atoms[1] isa Atoms{1,Float64}
@test atoms[1:2] isa Atoms{2,Float64}

@test Atoms(2000) isa Atoms{1, Float64}


text:
using Surrogates, XGBoost

#1D
obj_1D = x -> 3*x+1
x = [1.0,2.0,3.0,4.0,5.0]
y = obj_1D.(x)
a = 0.0
b = 10.0
num_round = 2
my_forest_1D = RandomForestSurrogate(x,y,a,b,num_round)
val = my_forest_1D(3.5)
add_point!(my_forest_1D,6.0,19.0)
add_point!(my_forest_1D,[7.0,8.0],obj_1D.([7.0,8.0]))

#ND
lb = [0.0,0.0,0.0]
ub = [10.0,10.0,10.0]
x = sample(5,lb,ub,SobolSample())
obj_ND = x -> x[1] * x[2]^2 * x[3]
y = obj_ND.(x)
my_forest_ND = RandomForestSurrogate(x,y,lb,ub,num_round)
val = my_forest_ND((1.0,1.0,1.0))
add_point!(my_forest_ND,(1.0,1.0,1.0),1.0)
add_point!(my_forest_ND,[(1.2,1.2,1.0),(1.5,1.5,1.0)],[1.728,3.375])