text:
import Base.isempty

export AbstractPolytope,
       vertices_list,
       singleton_list,
       isempty,
       minkowski_sum

"""
    AbstractPolytope{N<:Real} <: AbstractPolyhedron{N}

Abstract type for compact convex polytopic sets.

### Notes

Every concrete `AbstractPolytope` must define the following functions:
- `vertices_list(::AbstractPolytope{N})::Vector{Vector{N}}` -- return a list of
    all vertices

```jldoctest; setup = :(using LazySets: subtypes)
julia> subtypes(AbstractPolytope)
4-element Array{Any,1}:
 AbstractCentrallySymmetricPolytope
 AbstractPolygon
 HPolytope
 VPolytope
```

A polytope is a bounded polyhedron (see [`AbstractPolyhedron`](@ref)).
Polytopes are compact convex sets with either of the following equivalent
properties:
1. They are the intersection of a finite number of closed half-spaces.
2. They are the convex hull of finitely many vertices.
"""
abstract type AbstractPolytope{N<:Real} <: AbstractPolyhedron{N} end

# =============================================
# Common AbstractPolytope functions
# =============================================

"""
    isbounded(P::AbstractPolytope)::Bool

Determine whether a polytopic set is bounded.

### Input

- `P` -- polytopic set

### Output

`true` (since a polytope must be bounded).
"""
function isbounded(::AbstractPolytope)::Bool
    return true
end

"""
    singleton_list(P::AbstractPolytope{N})::Vector{Singleton{N}} where {N<:Real}

Return the vertices of a polytopic set as a list of singletons.

### Input

- `P` -- polytopic set

### Output

List containing a singleton for each vertex.
"""
function singleton_list(P::AbstractPolytope{N}
                       )::Vector{Singleton{N}} where {N<:Real}
    return [Singleton(vi) for vi in vertices_list(P)]
end

"""
    isempty(P::AbstractPolytope)::Bool

Determine whether a polytope is empty.

### Input

- `P` -- abstract polytope

### Output

`true` if the given polytope contains no vertices, and `false` otherwise.

### Algorithm

This algorithm checks whether the `vertices_list` of the given polytope is empty
or not.
"""
function isempty(P::AbstractPolytope)::Bool
    return isempty(vertices_list(P))
end

# given a polytope P, apply the linear map P to each vertex of P
# it is assumed that the interface function `vertices_list(P)` is available
@inline function _linear_map_vrep(M::AbstractMatrix{N}, P::AbstractPolytope{N}) where {N<:Real}
    vertices = broadcast(v -> M * v, vertices_list(P))
    m = size(M, 1) # output dimension
    if m == 1
        # TODO: substitute with Interval(convex_hull(vertices)...): convex hull 1D
        return Interval(minimum(vertices)[1], maximum(vertices)[1])
    elseif m == 2
        return VPolygon(vertices)
    else
        return VPolytope(vertices)
    end
end

@inline function _linear_map_hrep(M::AbstractMatrix{N}, P::AbstractPolytope{N},
                                  use_inv::Bool;
                                  inverse::Union{Nothing, AbstractMatrix{N}}=nothing
                                 ) where {N<:Real}
    constraints = _linear_map_hrep_helper(M, P, use_inv; inverse=inverse)
    m = size(M, 1) # output dimension
    if m == 1
        return convert(Interval, HPolygon(constraints))
    elseif m == 2
        return HPolygon(constraints)
    else
        return HPolytope(constraints)
    end
end

# the "backend" argument is ignored, used for dispatch
function _vertices_list(P::AbstractPolytope, backend)
    return vertices_list(P)
end


text:
## MACHINE TYPE

mutable struct Machine{M<:Model} <: MLJType

    model::M
    old_model::M # for remembering the model used in last call to `fit!`
    fitresult
    cache
    args::Tuple{Vararg{AbstractNode}}
    report
    frozen::Bool
    old_rows
    state::Int
    old_upstream_state

    # cleared by fit!(::Node) calls; put! by `fit_only!(machine, true)` calls:
    fit_okay::Channel{Bool}

    function Machine(model::M, args::AbstractNode...) where M<:Model
        mach = new{M}(model)
        mach.frozen = false
        mach.state = 0
        mach.args = args
        mach.old_upstream_state = upstream(mach)
        mach.fit_okay = Channel{Bool}(1)
        return mach
    end

end

upstream(mach::Machine) = Tuple(m.state for m in ancestors(mach))

"""
    ancestors(mach::Machine; self=false)

All ancestors of `mach`, including `mach` if `self=true`.

"""
function ancestors(mach::Machine; self=false)
    ret = Machine[]
    self && push!(ret, mach)
    return vcat(ret, (machines(N) for N in mach.args)...) |> unique
end

## CONSTRUCTORS

# In these checks the args are abstract nodes but `full=true` only
# makes sense if they are actually source nodes.

function check(model::Supervised, args... ; full=false)

    nowarns = true

    nargs = length(args)
    if nargs == 2
        X, y = args
    elseif nargs > 2
        supports_weights(model) || elscitype(args[3]) <: Unknown ||
            @info("$(typeof(model)) does not support sample weights and "*
                  "the supplied weights will be ignored in training.\n"*
                  "However, supplied weights will be passed to "*
                  "weight-supporting measures on calls to `evaluate!` "*
                  "and in tuning. ")
        X, y, w = args
    else
        throw(ArgumentError("Use `machine(model, X, y)` or "*
                            "`machine(model, X, y, w)` for a supervised "*
                            "model."))
    end

    if full
        # checks on input type:
        input_scitype(model) <: Unknown ||
            elscitype(X) <: input_scitype(model) ||
            (@warn("The scitype of `X`, in `machine(model, X, ...)` "*
             "is incompatible with "*
             "`model=$model`:\nscitype(X) = $(elscitype(X))\n"*
             "input_scitype(model) = $(input_scitype(model))."); nowarns=false)

        # checks on target type:
        target_scitype(model) <: Unknown ||
            elscitype(y) <: target_scitype(model) ||
            (@warn("The scitype of `y`, in `machine(model, X, y, ...)` "*
            "is incompatible with "*
             "`model=$model`:\nscitype(y) = $(elscitype(y))\ntarget_scitype(model) "*
             "= $(target_scitype(model))."); nowarns=false)

        # checks on dimension matching:
        X() === nothing || # model fits a distribution to y
            nrows(X()) == nrows(y()) ||
            throw(DimensionMismatch("Differing number of observations "*
                                    "in input and target. "))
        if nargs > 2 && !(w.data isa Nothing)
            w.data isa AbstractVector{<:Real} ||
                throw(ArgumentError("Weights must be real."))
            nrows(w()) == nrows(y()) ||
                throw(DimensionMismatch("Weights and target "*
                                        "differ in length."))
        end
    end
    return nowarns
end

function check(model::Unsupervised, args...; full=false)
    nowarns = true

    nargs = length(args)
    nargs <= 1 ||
        throw(ArgumentError("Wrong number of arguments. Use "*
                            "`machine(model, X)` for an unsupervised model, "*
                            "or `machine(model)` if there are no training "*
                            "arguments (`Static` transformers).) "))
    if full && nargs == 1
        X = args[1]
        # check input scitype
        input_scitype(model) <: Unknown ||
            elscitype(X) <: input_scitype(model) ||
            (@warn("The scitype of `X`, in `machine(model, X)` is "*
        "incompatible with `model=$model`:\nscitype(X) = $(elscitype(X))\n"*
             "input_scitype(model) = $(input_scitype(model))."); nowarns=false)
    end
    return nowarns
end

"""
    machine(model, args...)

Construct a `Machine` object binding a `model`, storing
hyper-parameters of some machine learning algorithm, to some data,
`args`. When building a learning network, `Node` objects can be
substituted for concrete data.

    machine(Xs; oper1=node1, oper2=node2, ...)
    machine(Xs, ys; oper1=node1, oper2=node2, ...)
    machine(Xs, ys, ws; oper1=node1, oper2=node2, ...)

Construct a special machine called a *learning network machine*, that
"wraps" a learning network, usually in preparation to export the
network as a stand-alone composite model type. The keyword arguments
declare what nodes are called when operations, such as `predict` and
`transform`, are called on the machine.

In addition to the operations named in the constructor, the methods
`fit!`, `report`, and `fitted_params` can be applied as usual to the
machine constructed.

    machine(Probablistic(), args...; kwargs...)
    machine(Deterministic(), args...; kwargs...)
    machine(Unsupervised(), args...; kwargs...)
    machine(Static(), args...; kwargs...)

Same as above, but specifying explicitly the kind of model the
learning network is to meant to represent.

Learning network machines are not to be confused with an ordinary
machine that happens to be bound to a stand-alone composite model
(i.e., an *exported* learning network).

### Examples

Supposing a supervised learning network's final predictions are
obtained by calling a node `yhat`, then the code

```julia
mach = machine(Deterministic(), Xs, ys; predict=yhat)
fit!(mach; rows=train)
predictions = predict(mach, Xnew) # `Xnew` concrete data
```

is  equivalent to

```julia
fit!(yhat, rows=train)
predictions = yhat(Xnew)
```

Here `Xs` and `ys` are the source nodes receiving, respectively, the
input and target data.

In a unsupervised learning network for clustering, with single source
node `Xs` for inputs, and in which the node `Xout` delivers the output
of dimension reduction, and `yhat` the class labels, one can write

```julia
mach = machine(Unsupervised(), Xs; transform=Xout, predict=yhat)
fit!(mach)
transformed = transform(mach, Xnew) # `Xnew` concrete data
predictions = predict(mach, Xnew)
```

which is equivalent to

```julia
fit!(Xout)
fit!(yhat)
transformed = Xout(Xnew)
predictions = yhat(Xnew)
```

"""
function machine end

machine(T::Type{<:Model}, args...) =
    throw(ArgumentError("Model *type* provided where "*
                        "model *instance* expected. "))

static_error() =
    throw(ArgumentError("A `Static` transformer "*
                        "has no training arguments. "*
                        "Use `machine(model)`. "))

function machine(model::Static, args...)
    isempty(args) || static_error()
    return Machine(model)
end

function machine(model::Static, args::AbstractNode...)
    isempty(args) || static_error()
    return Machine(model)
end

machine(model::Model, raw_arg1, arg2::AbstractNode, args::AbstractNode...) =
    error("Mixing concrete data with `Node` training arguments "*
          "is not allowed. ")

machine(model::Model, arg1::AbstractNode, arg2, args...) =
    error("Mixing concrete data with `Node` training arguments "*
          "is not allowed. ")

function machine(model::Model, raw_arg1, raw_args...)
    args = source.((raw_arg1, raw_args...))
    check(model, args...; full=true)
    return Machine(model, args...)
end

function machine(model::Model, arg1::AbstractNode, args::AbstractNode...)
    check(model, arg1, args...)
    return Machine(model, arg1, args...)
end

## INSPECTION AND MINOR MANIPULATION OF FIELDS

# Note: freeze! and thaw! are possibly not used within MLJ itself.

"""
    freeze!(mach)

Freeze the machine `mach` so that it will never be retrained (unless
thawed).

See also [`thaw!`](@ref).
"""
function freeze!(machine::Machine)
    machine.frozen = true
end

"""
    thaw!(mach)

Unfreeze the machine `mach` so that it can be retrained.

See also [`freeze!`](@ref).
"""
function thaw!(machine::Machine)
    machine.frozen = false
end

params(mach::Machine) = params(mach.model)

machines(::Source) = Machine[]

## DISPLAY

function Base.show(io::IO, ::MIME"text/plain", mach::Machine)
    show(io, mach)
    print(io, " trained $(mach.state) time")
    if mach.state == 1
        println(io, ".")
    else
        println(io, "s.")
    end
    println(io, "  args: ")
    for i in eachindex(mach.args)
        arg = mach.args[i]
        print(io, "    $i:\t$arg")
        if arg isa Source
            println(io, " \u23CE `$(elscitype(arg))`")
        else
            println(io)
        end
    end
end

## FITTING

# Not one, but *two*, fit methods are defined for machines here,
# `fit!` and `fit_only!`.

# - `fit_only!`: trains a machine without touching the learned
#   parameters (`fitresult`) of any other machine. It may error if
#   another machine on which it depends (through its node training
#   arguments `N1, N2, ...`) has not been trained.

# - `fit!`: trains a machine after first progressively training all
#   machines on which the machine depends. Implicitly this involves
#   making `fit_only!` calls on those machines, scheduled by the node
#   `@tuple N1 N2 ... `.)

function fitlog(mach, action::Symbol, verbosity)
    if verbosity < -1000
            put!(MACHINE_CHANNEL, (action, mach))
    elseif verbosity > -1 && action == :frozen
        @warn "$mach not trained as it is frozen."
    elseif verbosity > 0
        action == :train && (@info "Training $mach."; return)
        action == :update && (@info "Updating $mach."; return)
        action == :skip && begin
            @info "Not retraining $mach. Use `force=true` to force."
            return
        end
    end
end

"""
    MLJBase.fit_only!(mach::Machine; rows=nothing, verbosity=1, force=false)

Without mutating any other machine on which it may depend, perform one of
the following actions to the machine `mach`, using the data and model
bound to it, and restricting the data to `rows` if specified:

- *Ab initio training.* Ignoring any previous learned parameters and
  cache, compute and store new learned parameters. Increment `mach.state`.

- *Training update.* Making use of previous learned parameters and/or
   cache, replace or mutate existing learned parameters. The effect is
   the same (or nearly the same) as in ab initio training, but may be
   faster or use less memory, assuming the model supports an update
   option (implements `MLJBase.update`). Increment `mach.state`.

- *No-operation.* Leave existing learned parameters untouched. Do not
   increment `mach.state`.

### Training action logic

For the action to be a no-operation, either `mach.frozen == true` or
none of the following apply:

- (i) `mach` has never been trained (`mach.state == 0`).

- (ii) `force == true`

- (iii) The `state` of some other machine on which `mach` depends has
  changed since the last time `mach` was trained (ie, the last time
  `mach.state` was last incremented)

- (iv) The specified `rows` have changed since the last retraining.

- (v) `mach.model` has changed since the last retraining.

In cases (i) - (iv), `mach` is trained ab initio. In case (v) a
training update is applied.

To freeze or unfreeze `mach`, use `freeze!(mach)` or `thaw!(mach)`.

### Implementation detail

The data to which a machine is bound is stored in `mach.args`. Each
element of `args` is either a `Node` object, or, in the case that
concrete data was bound to the machine, it is concrete data wrapped in
a `Source` node. In all cases, to obtain concrete data for actual
training, each argument `N` is called, as in `N()` or `N(rows=rows)`,
and either `MLJBase.fit` (ab initio training) or `MLJBase.update`
(training update) is dispatched on `mach.model` and this data. See the
"Adding models for general use" section of the MLJ documentation for
more on these lower-level training methods.

"""
function fit_only!(mach::Machine; rows=nothing, verbosity=1, force=false)

    if mach.frozen
        # no-op; do not increment `state`.
        fitlog(mach, :frozen, verbosity)
        return mach
    end

    # catch deserialized machines not bound to data:
    !(mach.model isa Static) && isempty(mach.args) &&
        error("This machine is not bound to any data and so "*
              "cannot be trained. ")

    # take action if model has been mutated illegally:
    warning = clean!(mach.model)
    isempty(warning) || verbosity < 0 || @warn warning

    rows === nothing && (rows = (:))

    rows_have_changed = !isdefined(mach, :old_rows) ||
                          rows != mach.old_rows

    upstream_state = upstream(mach)

    data_has_changed =
        rows_have_changed ||upstream_state != mach.old_upstream_state
    previously_fit = (mach.state > 0)

    raw_args = [N(rows=rows) for N in mach.args]

    # we `fit`, `update`, or return untouched:
    if !previously_fit || data_has_changed || force
        # fit the model:
        fitlog(mach, :train, verbosity)
        mach.fitresult, mach.cache, mach.report =
            try
                fit(mach.model, verbosity, raw_args...)
            catch exception
                @error "Problem fitting the machine $mach, "*
                "possibly because an upstream node in a learning "*
                "network is providing data of incompatible scitype. "
                _sources = sources(glb(mach.args...))
                length(_sources) > 2 ||
                    mach.model isa Composite ||
                    all((!isempty).(_sources)) ||
                    @warn "Some learning network source nodes are empty. "
                @info "Running type checks... "
                check(mach.model, source.(raw_args)... ; full=true) &&
                    @info "Type checks okay. "
                throw(exception)
            end
    elseif mach.model != mach.old_model
        # update the model:
        fitlog(mach, :update, verbosity)
        mach.fitresult, mach.cache, mach.report =
            update(mach.model,
                   verbosity,
                   mach.fitresult,
                   mach.cache,
                   raw_args...)
    else
        # don't fit the model and return without incrementing `state`:
        fitlog(mach, :skip, verbosity)
        return mach
    end

    # If we get to here it's because we have run `fit` or `update`!

    if rows_have_changed
        mach.old_rows = deepcopy(rows)
    end

    mach.old_model = deepcopy(mach.model)
    mach.old_upstream_state = upstream_state
    mach.state = mach.state + 1

    return mach
end

"""

    fit!(mach::Machine, rows=nothing, verbosity=1, force=false)

Fit the machine `mach`. In the case that `mach` has `Node` arguments,
first train all other machines on which `mach` depends.

To attempt to fit a machine without touching any other machine, use
`fit_only!`. For more on the internal logic of fitting see
[`fit_only!`](@ref)

"""
function fit!(mach::Machine; kwargs...)
    glb_node = glb(mach.args...) # greatest lower bound node of arguments
    fit!(glb_node; kwargs...)
    fit_only!(mach; kwargs...)
end

# version of fit_only! for calling by scheduler (a node), which waits
# on the specified `machines` to fit:
function fit_only!(mach::Machine, wait_on_downstream::Bool; kwargs...)

    wait_on_downstream || fit_only!(mach; kwargs...)

    upstream_machines = machines(glb(mach.args...))

    # waiting on upstream machines to fit:
    for m in upstream_machines
        fit_okay = fetch(m.fit_okay)
        if !fit_okay
            put!(mach.fit_okay, false)
            return mach
        end
    end

    # try to fit this machine:
    try
        fit_only!(mach; kwargs...)
    catch e
        put!(mach.fit_okay, false)
        @error "Problem fitting $mach"
        throw(e)
    end
    put!(mach.fit_okay, true)
    return mach

end

## INSPECTION OF TRAINING OUTCOMES

"""
    fitted_params(mach)

Return the learned parameters for a machine `mach` that has been
`fit!`, for example the coefficients in a linear model.

This is a named tuple and human-readable if possible.

If `mach` is a machine for a composite model, such as a model
constructed using `@pipeline`, then the returned named tuple has the
composite type's field names as keys. The corresponding value is the
fitted parameters for the machine in the underlying learning network
bound to that model. (If multiple machines share the same model, then the
value is a vector.)

```julia
using MLJ
@load LogisticClassifier pkg=MLJLinearModels
X, y = @load_crabs;
pipe = @pipeline Standardizer LogisticClassifier
mach = machine(pipe, X, y) |> fit!

julia> fitted_params(mach).logistic_classifier
(classes = CategoricalArrays.CategoricalValue{String,UInt32}["B", "O"],
 coefs = Pair{Symbol,Float64}[:FL => 3.7095037897680405, :RW => 0.1135739140854546, :CL => -1.6036892745322038, :CW => -4.415667573486482, :BD => 3.238476051092471],
 intercept = 0.0883301599726305,)
```

Additional keys, `machines` and `fitted_params_given_machine`, give a
list of *all* machines in the underlying network, and a dictionary of
fitted parameters keyed on those machines.

```

"""
function fitted_params(machine::Machine)
    if isdefined(machine, :fitresult)
        return fitted_params(machine.model, machine.fitresult)
    else
        throw(error("$machine has not been trained."))
    end
end

"""
    report(mach)

Return the report for a machine `mach` that has been
`fit!`, for example the coefficients in a linear model.

This is a named tuple and human-readable if possible.

If `mach` is a machine for a composite model, such as a model
constructed using `@pipeline`, then the returned named tuple has the
composite type's field names as keys. The corresponding value is the
report for the machine in the underlying learning network
bound to that model. (If multiple machines share the same model, then the
value is a vector.)

```julia
using MLJ
@load LinearBinaryClassifier pkg=GLM
X, y = @load_crabs;
pipe = @pipeline Standardizer LinearBinaryClassifier
mach = machine(pipe, X, y) |> fit!

julia> report(mach).linear_binary_classifier
(deviance = 3.8893386087844543e-7,
 dof_residual = 195.0,
 stderror = [18954.83496713119, 6502.845740757159, 48484.240246060406, 34971.131004997274, 20654.82322484894, 2111.1294584763386],
 vcov = [3.592857686311793e8 9.122732393971942e6 … -8.454645589364915e7 5.38856837634321e6; 9.122732393971942e6 4.228700272808351e7 … -4.978433790526467e7 -8.442545425533723e6; … ; -8.454645589364915e7 -4.978433790526467e7 … 4.2662172244975924e8 2.1799125705781363e7; 5.38856837634321e6 -8.442545425533723e6 … 2.1799125705781363e7 4.456867590446599e6],)

```

Additional keys, `machines` and `report_given_machine`, give a
list of *all* machines in the underlying network, and a dictionary of
reports keyed on those machines.

```

"""
report(mach::Machine) = mach.report

## SERIALIZATION

# saving:
"""
    MLJ.save(filename, mach::Machine; kwargs...)
    MLJ.save(io, mach::Machine; kwargs...)

    MLJBase.save(filename, mach::Machine; kwargs...)
    MLJBase.save(io, mach::Machine; kwargs...)

Serialize the machine `mach` to a file with path `filename`, or to an
input/output stream `io` (at least `IOBuffer` instances are
supported).

The format is JLSO (a wrapper for julia native or BSON serialization)
unless a custom format has been implemented for the model type of
`mach.model`. The keyword arguments `kwargs` are passed to
the format-specific serializer, which in the JSLO case include these:

keyword        | values                        | default
---------------|-------------------------------|-------------------------
`format`       | `:julia_serialize`, `:BSON`   | `:julia_serialize`
`compression`  | `:gzip`, `:none`              | `:none`

See (see
[https://github.com/invenia/JLSO.jl](https://github.com/invenia/JLSO.jl)
for details.

Machines are de-serialized using the `machine` constructor as shown in
the example below. Data (or nodes) may be optionally passed to the
constructor for retraining on new data using the saved model.

### Example

    using MLJ
    tree = @load DecisionTreeClassifier
    X, y = @load_iris
    mach = fit!(machine(tree, X, y))

    MLJ.save("tree.jlso", mach, compression=:none)
    mach_predict_only = machine("tree.jlso")
    predict(mach_predict_only, X)

    mach2 = machine("tree.jlso", selectrows(X, 1:100), y[1:100])
    predict(mach2, X) # same as above

    fit!(mach2) # saved learned parameters are over-written
    predict(mach2, X) # not same as above

    # using a buffer:
    io = IOBuffer()
    MLJ.save(io, mach)
    seekstart(io)
    predict_only_mach = machine(io)
    predict(predict_only_mach, X)

!!! warning "Only load files from trusted sources"
    Maliciously constructed JLSO files, like pickles, and most other
    general purpose serialization formats, can allow for arbitrary code
    execution during loading. This means it is possible for someone
    to use a JLSO file that looks like a serialized MLJ machine as a
    [Trojan
    horse](https://en.wikipedia.org/wiki/Trojan_horse_(computing)).

"""
function MMI.save(file, mach::Machine; verbosity=1, kwargs...)
    isdefined(mach, :fitresult)  ||
        error("Cannot save an untrained machine. ")
    MMI.save(file, mach.model, mach.fitresult, mach.report; kwargs...)
end

# restoring:
function machine(file::Union{String,IO}, args...; kwargs...)
    model, fitresult, report = MMI.restore(file; kwargs...)
    if isempty(args)
        mach = Machine(model)
    else
        mach = machine(model, args...)
    end
    mach.fitresult = fitresult
    mach.state = 1
    mach.report = report
    return mach
end

ext: jl


text:
int_rules_7_3_2 = @theory begin

    #= ::Subsection::Closed:: =#
    #= 7.3.2*(d*x)^m*(a+b*arctanh(c*x^n))^p =#
    @apply_utils Antiderivative((~(a') + ~(b') * atanh(~(c') * ~x)) / ~x, ~x) => (~a * log(~x) - (~b / 2) * PolyLog(2, -(~c) * ~x)) + (~b / 2) * PolyLog(2, ~c * ~x) <-- FreeQ([~a, ~b, ~c], ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * acoth(~(c') * ~x)) / ~x, ~x) => (~a * log(~x) + (~b / 2) * PolyLog(2, -1 / (~c * ~x))) - (~b / 2) * PolyLog(2, 1 / (~c * ~x)) <-- FreeQ([~a, ~b, ~c], ~x)
    @apply_utils Antiderivative((~(a') + ~(b') * atanh(~(c') * ~x)) ^ ~p / ~x, ~x) => 2 * (~a + ~b * atanh(~c * ~x)) ^ ~p * atanh(1 - 2 / (1 - ~c * ~x)) - 2 * ~b * ~c * ~p * Antiderivative(((~a + ~b * atanh(~c * ~x)) ^ (~p - 1) * atanh(1 - 2 / (1 - ~c * ~x))) / (1 - (~c) ^ 2 * (~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~c], ~x) && IGtQ(~p, 1)
    @apply_utils Antiderivative((~(a') + ~(b') * acoth(~(c') * ~x)) ^ ~p / ~x, ~x) => 2 * (~a + ~b * acoth(~c * ~x)) ^ ~p * acoth(1 - 2 / (1 - ~c * ~x)) - 2 * ~b * ~c * ~p * Antiderivative(((~a + ~b * acoth(~c * ~x)) ^ (~p - 1) * acoth(1 - 2 / (1 - ~c * ~x))) / (1 - (~c) ^ 2 * (~x) ^ 2), ~x) <-- FreeQ([~a, ~b, ~c], ~x) && IGtQ(~p, 1)
    @apply_utils Antiderivative((~(a') + ~(b') * atanh(~(c') * (~x) ^ ~n)) ^ ~(p') / ~x, ~x) => (1 / ~n) * Subst(Antiderivative((~a + ~b * atanh(~c * ~x)) ^ ~p / ~x, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~n], ~x) && IGtQ(~p, 0)
    @apply_utils Antiderivative((~(a') + ~(b') * acoth(~(c') * (~x) ^ ~n)) ^ ~(p') / ~x, ~x) => (1 / ~n) * Subst(Antiderivative((~a + ~b * acoth(~c * ~x)) ^ ~p / ~x, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~n], ~x) && IGtQ(~p, 0)
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~(n'))) ^ ~(p'), ~x) => ((~x) ^ (~m + 1) * (~a + ~b * atanh(~c * (~x) ^ ~n)) ^ ~p) / (~m + 1) - ((~b * ~c * ~n * ~p) / (~m + 1)) * Antiderivative(((~x) ^ (~m + ~n) * (~a + ~b * atanh(~c * (~x) ^ ~n)) ^ (~p - 1)) / (1 - (~c) ^ 2 * (~x) ^ (2 * ~n)), ~x) <-- FreeQ([~a, ~b, ~c, ~m, ~n], ~x) && (IGtQ(~p, 0) && ((EqQ(~p, 1) || EqQ(~n, 1) && IntegerQ(~m)) && NeQ(~m, -1)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~(n'))) ^ ~(p'), ~x) => ((~x) ^ (~m + 1) * (~a + ~b * acoth(~c * (~x) ^ ~n)) ^ ~p) / (~m + 1) - ((~b * ~c * ~n * ~p) / (~m + 1)) * Antiderivative(((~x) ^ (~m + ~n) * (~a + ~b * acoth(~c * (~x) ^ ~n)) ^ (~p - 1)) / (1 - (~c) ^ 2 * (~x) ^ (2 * ~n)), ~x) <-- FreeQ([~a, ~b, ~c, ~m, ~n], ~x) && (IGtQ(~p, 0) && ((EqQ(~p, 1) || EqQ(~n, 1) && IntegerQ(~m)) && NeQ(~m, -1)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (Simplify((~m + 1) / ~n) - 1) * (~a + ~b * atanh(~c * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~m, ~n], ~x) && (IGtQ(~p, 1) && IntegerQ(Simplify((~m + 1) / ~n)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~n)) ^ ~(p'), ~x) => (1 / ~n) * Subst(Antiderivative((~x) ^ (Simplify((~m + 1) / ~n) - 1) * (~a + ~b * acoth(~c * ~x)) ^ ~p, ~x), ~x, (~x) ^ ~n) <-- FreeQ([~a, ~b, ~c, ~m, ~n], ~x) && (IGtQ(~p, 1) && IntegerQ(Simplify((~m + 1) / ~n)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~n)) ^ ~p, ~x) => Antiderivative(ExpandIntegrand((~x) ^ ~m * ((~a + (~b * log(1 + ~c * (~x) ^ ~n)) / 2) - (~b * log(1 - ~c * (~x) ^ ~n)) / 2) ^ ~p, ~x), ~x) <-- FreeQ([~a, ~b, ~c], ~x) && (IGtQ(~p, 1) && (IGtQ(~n, 0) && IntegerQ(~m)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~n)) ^ ~p, ~x) => Antiderivative(ExpandIntegrand((~x) ^ ~m * ((~a + (~b * log(1 + (~x) ^ -(~n) / ~c)) / 2) - (~b * log(1 - (~x) ^ -(~n) / ~c)) / 2) ^ ~p, ~x), ~x) <-- FreeQ([~a, ~b, ~c], ~x) && (IGtQ(~p, 1) && (IGtQ(~n, 0) && IntegerQ(~m)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~n)) ^ ~p, ~x) => With([k = Denominator(~m)], k * Subst(Antiderivative((~x) ^ (k * (~m + 1) - 1) * (~a + ~b * atanh(~c * (~x) ^ (k * ~n))) ^ ~p, ~x), ~x, (~x) ^ (1 / k))) <-- FreeQ([~a, ~b, ~c], ~x) && (IGtQ(~p, 1) && (IGtQ(~n, 0) && FractionQ(~m)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~n)) ^ ~p, ~x) => With([k = Denominator(~m)], k * Subst(Antiderivative((~x) ^ (k * (~m + 1) - 1) * (~a + ~b * acoth(~c * (~x) ^ (k * ~n))) ^ ~p, ~x), ~x, (~x) ^ (1 / k))) <-- FreeQ([~a, ~b, ~c], ~x) && (IGtQ(~p, 1) && (IGtQ(~n, 0) && FractionQ(~m)))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~n)) ^ ~p, ~x) => Antiderivative((~x) ^ ~m * (~a + ~b * acoth((~x) ^ -(~n) / ~c)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~m], ~x) && (IGtQ(~p, 1) && ILtQ(~n, 0))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~n)) ^ ~p, ~x) => Antiderivative((~x) ^ ~m * (~a + ~b * atanh((~x) ^ -(~n) / ~c)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~m], ~x) && (IGtQ(~p, 1) && ILtQ(~n, 0))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~n)) ^ ~p, ~x) => With([k = Denominator(~n)], k * Subst(Antiderivative((~x) ^ (k * (~m + 1) - 1) * (~a + ~b * atanh(~c * (~x) ^ (k * ~n))) ^ ~p, ~x), ~x, (~x) ^ (1 / k))) <-- FreeQ([~a, ~b, ~c, ~m], ~x) && (IGtQ(~p, 1) && FractionQ(~n))
    @apply_utils Antiderivative((~x) ^ ~(m') * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~n)) ^ ~p, ~x) => With([k = Denominator(~n)], k * Subst(Antiderivative((~x) ^ (k * (~m + 1) - 1) * (~a + ~b * acoth(~c * (~x) ^ (k * ~n))) ^ ~p, ~x), ~x, (~x) ^ (1 / k))) <-- FreeQ([~a, ~b, ~c, ~m], ~x) && (IGtQ(~p, 1) && FractionQ(~n))
    @apply_utils Antiderivative((~d * ~x) ^ ~m * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~(n'))), ~x) => ((~d * ~x) ^ (~m + 1) * (~a + ~b * atanh(~c * (~x) ^ ~n))) / (~d * (~m + 1)) - ((~b * ~c * ~n) / ((~d) ^ ~n * (~m + 1))) * Antiderivative((~d * ~x) ^ (~m + ~n) / (1 - (~c) ^ 2 * (~x) ^ (2 * ~n)), ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && (IntegerQ(~n) && NeQ(~m, -1))
    @apply_utils Antiderivative((~d * ~x) ^ ~m * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~(n'))), ~x) => ((~d * ~x) ^ (~m + 1) * (~a + ~b * acoth(~c * (~x) ^ ~n))) / (~d * (~m + 1)) - ((~b * ~c * ~n) / ((~d) ^ ~n * (~m + 1))) * Antiderivative((~d * ~x) ^ (~m + ~n) / (1 - (~c) ^ 2 * (~x) ^ (2 * ~n)), ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && (IntegerQ(~n) && NeQ(~m, -1))
    @apply_utils Antiderivative((~d * ~x) ^ ~m * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~(n'))) ^ ~(p'), ~x) => (((~d) ^ IntPart(~m) * (~d * ~x) ^ FracPart(~m)) / (~x) ^ FracPart(~m)) * Antiderivative((~x) ^ ~m * (~a + ~b * atanh(~c * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && (IGtQ(~p, 0) && (EqQ(~p, 1) || RationalQ(~m, ~n)))
    @apply_utils Antiderivative((~d * ~x) ^ ~m * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~(n'))) ^ ~(p'), ~x) => (((~d) ^ IntPart(~m) * (~d * ~x) ^ FracPart(~m)) / (~x) ^ FracPart(~m)) * Antiderivative((~x) ^ ~m * (~a + ~b * acoth(~c * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n], ~x) && (IGtQ(~p, 0) && (EqQ(~p, 1) || RationalQ(~m, ~n)))
    @apply_utils Antiderivative((~(d') * ~x) ^ ~(m') * (~(a') + ~(b') * atanh(~(c') * (~x) ^ ~(n'))) ^ ~(p'), ~x) => Unintegrable((~d * ~x) ^ ~m * (~a + ~b * atanh(~c * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x)
    @apply_utils Antiderivative((~(d') * ~x) ^ ~(m') * (~(a') + ~(b') * acoth(~(c') * (~x) ^ ~(n'))) ^ ~(p'), ~x) => Unintegrable((~d * ~x) ^ ~m * (~a + ~b * acoth(~c * (~x) ^ ~n)) ^ ~p, ~x) <-- FreeQ([~a, ~b, ~c, ~d, ~m, ~n, ~p], ~x)
end

max_stars_repo_path: src/rules/7 Inverse hyperbolic functions/7.3 Inverse hyperbolic tangent/7.3.2 (d x)^m (a+b arctanh(c x^n))^p.jl


text:
function gw_ping(gw::GitWorker = gw_curr();
        tout = 30.0,
        verbose = false,
        onping::Function = () -> nothing
    )

    GitLinks.ping(gitlink(gw); tout, verbose, onping)
end

max_stars_repo_name: josePereiro/GitWorkers

max_issues_repo_name: josePereiro/GitWorkers

max_forks_repo_name: josePereiro/GitWorkers


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "SuperLU"
version = v"5.2.2"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://github.com/xiaoyeli/superlu/archive/v5.2.2.tar.gz", "470334a72ba637578e34057f46948495e601a5988a602604f5576367e606a28c")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/superlu*
mkdir build && cd build
cmake -DCMAKE_INSTALL_PREFIX=${prefix} \
    -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
    -DCMAKE_BUILD_TYPE=Release \
    -DBUILD_SHARED_LIBS=ON \
    -Denable_internal_blaslib=OFF \
    -Denable_matlabmex=OFF \
    -Denable_tests=OFF \
    -Denable_doc=OFF \
    -Denable_single=ON \
    -Denable_double=ON \
    -Denable_complex=ON \
    -Denable_complex16=ON \
    -DTPL_BLAS_LIBRARIES="${libdir}/libopenblas.${dlext}" \
    ..
make -j${nproc}
make install
if [[ "${target}" == *-mingw* ]]; then
    # Manually install the library
    cp "SRC/libsuperlu.${dlext}" "${libdir}/libsuperlu.${dlext}"
fi
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()

# The products that we will ensure are always built
products = [
    LibraryProduct("libsuperlu", :libsuperlu)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="OpenBLAS32_jll", uuid="656ef2d0-ae68-5445-9ca0-591084a874a2"))
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
# 2D Lattice Boltzmann (BGK) model of a fluid.
#  c4  c3   c2  D2Q9 model. At each timestep, particle densities propagate
#    \  |  /    outwards in the directions indicated in the figure. An
#  c5 -c9 - c1  equivalent 'equilibrium' density is found, and the densities
#    /  |  \    relax towards that state, in a proportion governed by omega.
#  c6  c7   c8
#
# The first Julia version has many traces of Matlab, and it only works before
# version 0.6. The second version looks more like Julia.
#
# There is a lazy version of shifted array package. You can try and replace the
# circshift function with that.
#
# Orignal version by Iain Haslam, March 2006
# First Julia version by Jeroen Wouters, Jun 2015
# Second Julia version by Hongyang Zhou, Mar 2019

using LinearAlgebra

##################
# Configuration
##################
Ω        = 1.0    # Relaxation factor
ρ0       = 1.0    # Initial density
fraction = 0.1    # Average fraction of occupied nodes in the domain
rbody    = 5.0    # Radius of the solid cylinder

const nx,ny = 64,32
const nStep = 4000

ΔU = 1e-7

####################
# Setup of variables
####################
avu, avuOld = 1, 0
ts = 0
#avus = zeros(Float64, 10nStep)

# Single-particle distribution function
# Particle densities conditioned on one of the 9 possible velocities
F = fill(ρ0/9,nx,ny,9)
FEQ = F
nNode = nx*ny
# Offsets of different directions e_a in the F matrix
CI = [0:nNode:nNode*7...]

function setboundary()
   #BOUND = rand(nx,ny) .> 1 - fraction; # random domain
   center = [nx/2, ny/2]
   global BOUND = [norm([i, j]-center) < rbody for i=1:nx, j=1:ny]

   # Matrix offset of each Occupied Node (not physical domain)
   global ON = findall(x->x==true, vec(BOUND))

   global nActiveNodes = sum(1 .- BOUND)

   # Linear indices in F of occupied nodes
   global TO_REFLECT = [ON.+CI[1] ON.+CI[2] ON.+CI[3] ON.+CI[4] ON.+CI[5] ON.+CI[6] ON.+CI[7] ON.+CI[8]]

   # Right <-> Left: 1 <-> 5; Up <-> Down: 3 <-> 7
   #(1,1) <-> (-1,-1): 2 <-> 6; (1,-1) <-> (-1,1): 4 <-> 8
   global REFLECTED = [ON.+CI[5] ON.+CI[6] ON.+CI[7] ON.+CI[8] ON.+CI[1] ON.+CI[2] ON.+CI[3] ON.+CI[4]]
end

setboundary()

###############
# Constants
###############
const w1,w2,w3 = 4/9, 1/9, 1/36 # Weights
const c_squ = 1/3

Ux, Uy = Array{Float64,3}(undef,nx,ny,1), Array{Float64,3}(undef,nx,ny,1)

###############
# main loop
###############
while (ts < nStep) & (abs((avu-avuOld)/avu) > 1e-10)
   global F,avu,avuOld,ts,Ux,Uy # This is only needed outside function call
   ##############
   # Streaming
   ##############
   # particles at (x,y)=[2,1] were at [1,1] before: 1 points right (1,0)
   F[:,:,1] = F[circshift([1:nx...],1),:,1]
   # particles at [1,2] were at [1,1] before: 3 points up (0,1)
   F[:,:,3] = F[:,circshift([1:ny...],1),3]
   # particles at [1,1] were at [2,1] before: 5 points left (-1,0)
   F[:,:,5] = F[circshift([1:nx...],-1),:,5]
   # particles at [1,1] were at [1,2] before: 7 points down (0,-1)
   F[:,:,7] = F[:,circshift([1:ny...],-1),7]

   # particles at [2,2] were at [1, 1] before: 2 points to (1,1)
   F[:,:,2] = F[circshift([1:nx...],1), circshift([1:ny...],1), 2]
   # particles at [1,2] were at [2, 1] before: 4 points to (-1, 1)
   F[:,:,4] = F[circshift([1:nx...],-1),circshift([1:ny...],1), 4]
   # particles at [1,1] were at [2, 2] before: 6 points to (-1, -1)
   F[:,:,6] = F[circshift([1:nx...],-1),circshift([1:ny...],-1),6]
   # particles at [2,1] were at [1, 2] before: 8 points down (1,-1)
   F[:,:,8] = F[circshift([1:nx...],1), circshift([1:ny...],-1),8]

   ρ = sum(F,dims=3)

   # 1,2,8 are moving to the right, 4,5,6 to the left
   # 3, 7 and 9 don't move in the x direction
   Ux = (sum(F[:,:,[1, 2, 8]],dims=3) - sum(F[:,:,[4, 5, 6]],dims=3))./ρ
   # 2,3,4 are moving up, 6,7,8 down
   # 1, 5 and 9 don't move in the y direction
   Uy = (sum(F[:,:,[2, 3, 4]],dims=3) - sum(F[:,:,[6, 7, 8]],dims=3))./ρ

   Ux[1,1:ny] .+= ΔU # Increase inlet pressure

   Ux[ON] .= 0; Uy[ON] .= 0; ρ[ON] .= 0 # Set false cells to 0

   U2 = Ux.^2 + Uy.^2
   U_C2 = Ux + Uy
   U_C4 = -Ux + Uy

   # Calculate equilibrium distribution: stationary (a = 0)
   FEQ[:,:,9] = w1*ρ.*(1 .- U2/(2*c_squ))

   # nearest-neighbours
   FEQ[:,:,1] = w2*ρ.*(1 .+ Ux/c_squ + 0.5*(Ux/c_squ).^2 - 0.5*U2/c_squ)
   FEQ[:,:,3] = w2*ρ.*(1 .+ Uy/c_squ + 0.5*(Uy/c_squ).^2 - 0.5*U2/c_squ)
   FEQ[:,:,5] = w2*ρ.*(1 .- Ux/c_squ + 0.5*(Ux/c_squ).^2 - 0.5*U2/c_squ)
   FEQ[:,:,7] = w2*ρ.*(1 .- Uy/c_squ + 0.5*(Uy/c_squ).^2 - 0.5*U2/c_squ)

   # next-nearest neighbours
   FEQ[:,:,2] = w3*ρ.*(1 .+ U_C2/c_squ + 0.5*(U_C2/c_squ).^2 - 0.5*U2/c_squ)
   FEQ[:,:,4] = w3*ρ.*(1 .+ U_C4/c_squ + 0.5*(U_C4/c_squ).^2 - 0.5*U2/c_squ)
   FEQ[:,:,6] = w3*ρ.*(1 .- U_C2/c_squ + 0.5*(U_C2/c_squ).^2 - 0.5*U2/c_squ)
   FEQ[:,:,8] = w3*ρ.*(1 .- U_C4/c_squ + 0.5*(U_C4/c_squ).^2 - 0.5*U2/c_squ)

   BOUNCEDBACK = F[TO_REFLECT] # Densities bouncing back at next timestep

   F = Ω*FEQ + (1-Ω)*F

   F[REFLECTED] = BOUNCEDBACK

   avuOld = avu
   avu = sum(Ux)/nActiveNodes
   #avus[ts+1] = avu
   ts += 1
end

###############
# Visualization
###############
using PyPlot

figure()
imshow(1 .- BOUND', cmap="hot", interpolation="None", vmin=0., vmax=1., origin="lower");
quiver(1:nx-1, 0:ny-1, Ux[2:nx,:]', Uy[2:nx,:]');
title(string("Flow field after \$ ", string(ts), " \\delta t\$"));
xlabel("x")
ylabel("y")


text:
using DistStat, Pkg, Test

if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
    using CUDA
    A = CuArray
else
    A = Array
end

type=[Float64,Float32]

for T in type
    N=9; M=7
    X=A{T}(reshape(collect(1:63),N,M))
    X_dist=distribute(X)
    X1=A{T}(undef,1, M); X1=distribute(X1)
    X2=A{T}(undef,N,1)
    cols=X_dist.partitioning[DistStat.Rank()+1][2]

    @test sum(X)==sum(X_dist)

    @test sum!(X1,X_dist).localarray==sum(X;dims=1)[:,cols]

    @test sum!(X2,X_dist)==sum(X;dims=2)

    @test sum(X_dist;dims=1).localarray==sum(X;dims=1)[:,cols]

    @test sum(X_dist;dims=2)==sum(X;dims=2)

end


text:
using HTTP

include("../src/api.jl")

@testset "api" begin

    @test works(doi="10.1038/nature12373") isa Dict

    @test works(query="Heisenberg", limit=1) isa Dict

    @test members(member_id=98, works=true, limit=1) isa Dict

    @test funders(funder_id="10.13039/100000001", works=true, limit=1) isa Dict

    @test types(type_id="reference-book", works=true, limit=1) isa Dict

    @test prefixes(owner_prefix="10.1016", works=true, limit=1) isa Dict

    @test licenses(query="creative", limit=1) isa Dict

    @test agency("10.1038/nature12373") isa Dict

end

max_stars_repo_name: naustica/Pitaya

max_issues_repo_name: naustica/Pitaya

max_forks_repo_name: naustica/Pitaya


text:
## header to provide surface and context
using Cairo
using Graphics

c = CairoRGBSurface(256,256);
cr = CairoContext(c);

save(cr);
set_source_rgb(cr,0.8,0.8,0.8);    # light gray
rectangle(cr,0.0,0.0,256.0,256.0); # background
fill(cr);
restore(cr);

save(cr);

# put image top left
s = read_from_png("data/mulberry.png");
Cairo.image(cr,s,0,0,128,128)

# copy inplace, so to apply copy as source, need coordinate translate
c1 = copy(cr)

save(cr)
translate(cr,128,0)
rectangle(cr,0,0,256,256)
set_source_surface(cr,c1.surface)
paint(cr)
restore(cr)

# copy centered part
c2 = copy(cr,Graphics.BoundingBox(64,192,0,128))

save(cr)

rectangle(cr,0,0,256,256)
set_source_surface(cr,c2.surface,64,128)
paint(cr)
restore(cr)

## mark picture with current date
restore(cr);
move_to(cr,0.0,12.0);
set_source_rgb(cr, 0,0,0);
show_text(cr,Libc.strftime(time()));
write_to_png(c,"sample_image_copy.png");


text:
#=
    # The four inputs t, b, c, d are defined as follows:

    * t = current time (in any unit measure, but same unit as duration)
    * b = starting value to interpolate
    * c = the total change in value of b that needs to occur
    * d = total time it should take to complete (duration)
=#

@testset "easing" begin
    currentTime = 0
    duration = 100
    startPositionX = 0
    finalPositionX = 30
    currentPositionX = startPositionX

    while (currentPositionX < finalPositionX)
        currentPositionX = Raylib.EaseSineIn(currentTime, startPositionX, finalPositionX - startPositionX, duration)
        currentTime += 1
        @show (currentTime, currentPositionX)
    end
end


text:
prodir = realpath(joinpath(dirname(dirname(@__FILE__))))
srcdir = joinpath(prodir, "src")
srcdir ∉ LOAD_PATH && push!(LOAD_PATH, srcdir)

using Documenter, IntegerSequences

makedocs(
   modules = [IntegerSequences],
   format = Documenter.HTML(prettyurls = get(ENV, "CI", nothing) == "true"),
   clean = true,
   doctest = false,
   sitename = "IntegerSequences.jl",
   pages = [
      "About" => "about.md",
      "Sequences" => "index.md",
      "Modules" => "modules.md",
      "Notebooks and demos" => "notebooks.md",
      "User guide" => "userguide.md",
      "Developer guide" => "developerguide.md",
      "Use of the OEIS" => "useofoeis.md",
      "License" => "license.md",
   ]
)

# repo = "github.com/OpenLibMathSeq/IntegerSequences.jl.git",

deploydocs(
    repo = "github.com/PeterLuschny/IntegerSequences.jl.git",
    push_preview = true
    )

max_stars_repo_name: OpenLibMathSeq/Sequences

max_issues_repo_name: OpenLibMathSeq/Sequences


text:
using Random, Statistics, LinearAlgebra, Test, FillArrays, Printf, TotalLeastSquares, StatsBase
using RandomizedLinAlg
Random.seed!(0)

@testset "TotalLeastSquares" begin

    @testset "TLS" begin
        @info "Testing TLS"
        x   = randn(3)
        A   = randn(50,3)
        σa  = 1
        σy  = 0.01
        An  = A + σa*randn(size(A))
        y   = A*x
        yn  = y + σy*randn(size(y))
        Qaa = σa^2*Eye(prod(size(A)))
        Qay = 0Eye(prod(size(A)),length(y))
        Qyy = σy^2*Eye(prod(size(y)))

        x̂ = An\yn
        @printf "Least squares error: %25.3e %10.3e %10.3e, Norm: %10.3e\n" (x-x̂)... norm(x-x̂)
        @test norm(x-x̂) < 1

        x̂ = wls(An,yn,Qyy)
        @printf "Weigthed Least squares error: %16.3e %10.3e %10.3e, Norm: %10.3e\n" (x-x̂)... norm(x-x̂)
        @test norm(x-x̂) < 1

        x̂ = tls(An,yn)
        @printf "Total Least squares error: %19.3e %10.3e %10.3e, Norm: %10.3e\n" (x-x̂)... norm(x-x̂)
        @test norm(x-x̂) < 1

        x̂ = wtls(An,yn,Qaa,Qay,Qyy,iters=10)
        @printf "Weighted Total Least squares error: %10.3e %10.3e %10.3e, Norm: %10.3e\n" (x-x̂)... norm(x-x̂)
        println("----------------------------")
        @test norm(x-x̂) < 1

        @test tls(An,yn) ≈ tls!([An yn], size(An,2))

        rowC = rowcovariance([[σa^2*Eye(3) zeros(3); zeros(1,3) σy^2] for _ in 1:50])
        @test rowC[1] ≈ Qaa
        @test rowC[2] ≈ Qay
        @test rowC[3] ≈ Qyy

    end

    @testset "Robust PCA" begin
        @info "Testing Robust PCA"
        D = [0.462911    0.365901  0.00204357    0.692873    0.935861;
        0.0446199    0.108606   0.0664309   0.0736707    0.264429;
        0.320581    0.287788    0.073133    0.188872    0.526404;
        0.356266    0.197536 0.000718338    0.513795    0.370094;
        0.677814    0.011651    0.818047   0.0457694    0.471477]

        A = [0.462911   0.365901  0.00204356   0.345428   0.623104;
        0.0446199  0.108606  0.0429271    0.0736707  0.183814;
        0.320581   0.203777  0.073133     0.188872   0.472217;
        0.30725    0.197536  0.000717701  0.201626   0.370094;
        0.234245   0.011651  0.103622     0.0457694  0.279032]

        E = [0.0        0.0        0.0        0.347445  0.312757 ;
        0.0        0.0        0.0235038  0.0       0.0806151;
        0.0        0.0840109  0.0        0.0       0.0541868;
        0.0490157  0.0        6.5061e-7  0.312169  0.0      ;
        0.443569   0.0        0.714425   0.0       0.192445]

        Â, Ê,_,_ = rpca(D, nonnegE=true, nonnegA=true, verbose=false)

        @test Â ≈ A atol=1.0e-6
        @test Ê ≈ E atol=1.0e-6
        @test norm(D - (Â + Ê))/norm(D) < sqrt(eps())

        Â, Ê,_,_ = rpca(D, nonnegE=false, nonnegA=false, verbose=false)
        @test norm(D - (Â + Ê))/norm(D) < sqrt(eps())

        @testset "Missing values" begin
            @info "Testing Missing values"

            res = map(1:20) do _
                N = 500
                y = sin.(0.1 .* (1:N)) .+ 0.1*randn(N)
                miss = rand(N) .< 0.1
                yn = y .+ miss .* 1e2
                yf = lowrankfilter(yn,40)
                mean(abs2,y-yf)/mean(abs2,y)
            end
            @test mean(res) < 0.025

        end

        @testset "Complex" begin
            @info "Testing Complex"
            u = randn(ComplexF64, 100)
            v = randn(ComplexF64, 20)
            E = randn(ComplexF64,100,20) .* 10 .* (rand.() .< 0.01)
            A = u*v'
            D = A .+ E
            Ah, Eh, s, sv = rpca(D)
            @test sum(abs2, Eh-E)/sum(abs2, E) < 1e-5
            @test sum(abs2, Ah-A)/sum(abs2, A) < 1e-5

        end

    end

@testset "rtls" begin
    @info "Testing rtls"
    passes = map(1:1000) do _
        x   = randn(3)
        A   = randn(50,3)
        σ   = 5
        An  = A + σ*randn(size(A)) .* (rand(size(A)...) .< 0.1)
        y   = A*x
        yn  = y + σ*randn(size(y)) .* (rand(size(y)...) .< 0.1)

        AA  = [An yn]
        Ah,Eh,_,_ = rpca(AA, verbose=false)

        sum(abs2, Ah - [A y])/sum(abs2,[A y]) < sum(abs2, AA - [A y])/sum(abs2,[A y])
    end
    @show mean(passes)
    @test mean(passes) > 0.9

    passes = map(1:1000) do _
        x   = randn(3)
        A   = randn(50,3)
        σ   = 50
        An  = A + σ*randn(size(A)) .* (rand(size(A)...) .< 0.1)
        y   = A*x
        yn  = y + σ*randn(size(y)) .* (rand(size(y)...) .< 0.1)

        x̂t = tls(An,yn)
        x̂r = rtls(An,yn)

        norm(x-x̂r) < norm(x-x̂t)
    end
    @show mean(passes)
    @test mean(passes) > 0.9

    # @testset "Elastic Net rtls" begin
    #     @info "Testing Elastic Net rtls"
    #     passes = map(1:500) do _
    #         x     = randn(10)
    #         A     = randn(100,3) * randn(3,10)
    #         σs    = 50
    #         σd    = 0.5
    #         S     = σs*randn(size(A)) .* (rand(size(A)...) .< 0.1)
    #         D     = σd*randn(size(A))
    #         N     = S + D
    #         An    = A + N
    #
    #         Ah,Eh,_,_ = rpca(An, verbose=false)
    #
    #         λ = 1/sqrt(100)
    #         μ = 0.001λ
    #         Ah2,Eh2,_,_ = rpca(An, verbose=false, proxE = ElasticNet(λ, μ))
    #
    #         sum(abs2, Ah - A)/sum(abs2,A) < sum(abs2, An - A)/sum(abs2,A),
    #         sum(abs2, Eh2 - N)/sum(abs2,N) < sum(abs2, Eh - N)/sum(abs2,N)
    #     end
    #     @show mean(getindex.(passes, 1))
    #     @test mean(getindex.(passes, 1)) > 0.9
    #     @show mean(getindex.(passes, 2))
    #     @test mean(getindex.(passes, 2)) > 0.9
    #
    #
    # end

end

# @testset "rtls ga" begin
#     @info "Testing rtls ga"
#
#     passes = map(1:1000) do _
#         x   = randn(5)
#         A   = randn(500,5)
#         σ   = 50
#         An  = A + σ*randn(size(A)) .* (rand(size(A)...) .< 0.01)
#         y   = A*x
#         yn  = y + σ*randn(size(y)) .* (rand(size(y)...) .< 0.01)
#
#         x̂t = tls(An,yn)
#         x̂r = rtls_ga(An,yn, μ=entrywise_median)
#
#         norm(x-x̂r) < norm(x-x̂t)
#     end
#     @show mean(passes)
#     @test mean(passes) > 0.7
#
# end

@testset "soft hankel" begin
    @info "Testing soft hankel"

    @test hankel(1:20,2) == [1:19 2:20]
    @test hankel(1:20,3,2) == [1:2:17 2:2:18 3:2:19]

    A = hankel(1:8,4)
    @test ishankel(A)
    An = A + 0.1randn(size(A))
    @test !ishankel(An)
    Anc = copy(An)
    TotalLeastSquares.soft_hankel!(An, 0.1)
    @test sum(abs2,An-A) < sum(abs2,Anc-A)

    An = -A + 0.1randn(size(A))
    Anc = copy(An)
    TotalLeastSquares.soft_hankel!(An, 0.1)
    @test sum(abs2,An+A) < sum(abs2,Anc+A)

    passes = map(1:100) do _
        y = randn(20)
        yn = y .+ 0.1 .* randn.()
        A = hankel(y,3)
        An = hankel(yn,3)
        Anc = copy(An)
        A1,E1,_,_ = rpca(An, verbose=false, nukeA=false)
        A2,E2,_,_ = rpca(An, verbose=false, nukeA=false, hankel=true)
        sum(abs2,A2-A) < sum(abs2,A1-A)
    end
    @show mean(passes)
    @test mean(passes) > 0.7

    @info "Small random Hankel"
    passes = map(1:500) do _
        y = randn(100)
        A = hankel(y,5)
        @test ishankel(A)
        A1,E1,_,_ = rpca(A, verbose=false, nukeA=false)
        A2,E2,_,_ = rpca(A, verbose=false, nukeA=false, hankel=true)
        @test ishankel(A2)
        @test ishankel(E2)
        mean(abs2,A2-A) < mean(abs2,A1-A)
    end
    @show mean(passes)
    @test mean(passes) > 0.8

    @info "Big random Hankel"
    passes = map(1:10) do _
        y = randn(1000)
        A = hankel(y,50)
        @test ishankel(A)
        A1,E1,_,_ = rpca(A, verbose=false, nukeA=false)
        A2,E2,_,_ = rpca(A, verbose=false, nukeA=false, hankel=true)
        @test ishankel(A2)
        @test ishankel(E2)
        mean(abs2,A2-A) < mean(abs2,A1-A)
    end
    @show mean(passes)
    @test mean(passes) >= 0.8

end

@testset "lowrankfilter" begin
    @info "Testing lowrankfilter"
    qn(x) = x ./ quantile(abs.(x), 0.9)
    T = 1000
    t = 1:T
    y = sin.(0.1 .* t) |> qn

    H = hankel(y, 2)
    @test ishankel(H)
    @test unhankel(H) == y

    H = hankel(y,2,2)
    yh = unhankel(H,2,T)
    @test yh == y

    H = hankel(y,5,2)
    yh = unhankel(H,2,T)
    @test yh[1:end-1] ≈ y[1:end-1]

    y2 = randn(T)
    H = hankel([y y2],5,2)
    yh = unhankel(H,2,T,2)
    @test yh[1:end-1,:] ≈ [y y2][1:end-1,:]

    n = 20randn(T) .* (rand(T) .< 0.01) + 0.1randn(T)
    yf = lowrankfilter(y+n) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.001
    @show mean(abs2, y-yf)/mean(abs2, n)

    # Randomized lin alg
    yf = lowrankfilter(y+n, opnorm=x->rnorm(x,10)) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.001
    @show mean(abs2, y-yf)/mean(abs2, n)

    yf = lowrankfilter(y+n, svd=rsvd_fnkz) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.05
    @show mean(abs2, y-yf)/mean(abs2, n)

    yf = lowrankfilter(y+n, opnorm=x->rnorm(x,10), svd=rsvd_fnkz) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.05
    @show mean(abs2, y-yf)/mean(abs2, n)

    yf = lowrankfilter(y+n, opnorm=x->rnorm(x,10), svd=rsvd_fnkz, maxrank=5) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.05
    @show mean(abs2, y-yf)/mean(abs2, n)

    # Only SSA
    n = randn(T)
    yf = lowrankfilter(y+n, sv=2) |> qn
    @test mean(abs2, y-yf)/mean(abs2, n) < 0.05
    @show mean(abs2, y-yf)/mean(abs2, n)

    @testset "Multi channel" begin
        @info "Testing Multi channel"

        T = 1000
        t = 1:T
        y1 = sin.(0.1 .* t)
        y2 = sin.(0.3 .* t)
        y = [y1 y2+0.5y1]
        yn = y .+ 0.1 .* randn.() .+ (rand.() .< 0.01) .* 5 .* randn.()
        yf = lowrankfilter(yn, 20)
        @show mean(abs2,y-yf)/mean(abs2,y-yn)
        @test mean(abs2,y-yf)/mean(abs2,y-yn) < 0.05

        # There shoud be benefit to joint filtering since the two signals are correlated
        yf1 = lowrankfilter(yn[:,1], 20)
        @test mean(abs2,y1-yf1)/mean(abs2,y1-yf[:,1]) > 1.02

        # Same freq, different phase
        y1 = sin.(0.1 .* t)
        y2 = sin.(0.1 .* t .+ 0.01)
        y = [y1 y2]
        yn = y .+ 0.1 .* randn.() .+ (rand.() .< 0.01) .* 5 .* randn.()
        yf = lowrankfilter(yn, 20)
        @show mean(abs2,y-yf)/mean(abs2,y-yn)
        @test mean(abs2,y-yf)/mean(abs2,y-yn) < 0.02

        yf1 = lowrankfilter(yn[:,1], 20)
        @show mean(abs2,y1-yf1)/mean(abs2,y1-yf[:,1])
        @test mean(abs2,y1-yf1)/mean(abs2,y1-yf[:,1]) > 1.2

    end

end

@testset "rpca_ga" begin
    @info "Testing rpca_ga"

    Random.seed!(1)
    for r = 1:10, ϵ = exp10.(LinRange(-8, 0, 20))
        s = svd(randn(10,40))
        u,v = s.U[:,1:r], s.V[:,1:r]
        A = u*Diagonal(10*(1:r))*v' .+ ϵ .* randn.()
        Q = rpca_ga(A, r, verbose=false)
        @test rank([Q u], atol=ϵ) == r # Q and u should span the same space, but they may not share any other properties
        @test norm(Q'Q - I) < sqrt(eps())
    end

    # Below is the reverse dimensions of above
    for r = 1:10, ϵ = exp10.(LinRange(-8, 0, 20))
        s = svd(randn(40,10))
        u,v = s.U[:,1:r], s.V[:,1:r]
        A = u*Diagonal(10*(1:r))*v' .+ ϵ .* randn.()
        Q = rpca_ga(A, r, verbose=false)
        @test rank([Q u], atol=ϵ) == r # Q and u should span the same space, but they may not share any other properties
        @test norm(Q'Q - I) < sqrt(eps())
    end

    @testset "entrywise trimmed mean" begin
        @info "Testing entrywise trimmed mean"

        U = randn(10,10)
        s = zeros(10)
        w = ones(10)
        m1 = TotalLeastSquares.μ!(s,w,U)
        @test m1 ≈ mean(U, dims=2)

        m2 = TotalLeastSquares.entrywise_trimmed_mean(s,w,U,0)
        @test m2 ≈ mean(U, dims=2)

        w = randn(10)
        m1 = TotalLeastSquares.μ!(s,w,U)
        @test m1 ≈ sum(U .* w', dims=2)./ sum(w)

        m2 = TotalLeastSquares.entrywise_trimmed_mean(s,w,U,0)
        @test m2 ≈ sum(U .* w', dims=2)./ sum(w)

        w = ones(10)
        m2 = TotalLeastSquares.entrywise_trimmed_mean(s,w,U,0.1)
        for i in eachindex(m2)
            @test m2[i] ≈ mean(StatsBase.trim(U[i,:], prop=0.1))
        end

        r = 3; ϵ = 1e-8
        Random.seed!(1)
        passes = map(Iterators.product(1:4, exp10.(LinRange(-8, -1, 3)))) do (r,ϵ)
            # @show r, ϵ
            s = svd(randn(10,1000))
            u,v = s.U[:,1:r], s.V[:,1:r]
            # A = u*Diagonal(10*(1:r))*v' .+ ϵ .* randn.()
            A = u*Diagonal(s.S[1:r])*v' .+ ϵ .* randn.()
            A .+= 1000*randn.() .* (rand.() .< 0.01)

            Qtm = rpca_ga(A, r, verbose=false, μ = TotalLeastSquares.entrywise_trimmed_mean, iters=120)
            Qm = rpca_ga(A, r, verbose=false, iters=120)
            sum(svdvals([Qtm u])[r+1:2r]) < sum(svdvals([Qm u])[r+1:2r])
            # @test norm(Q'Q - I) < r*sqrt(eps()) # This test should pass but doesn't
        end
        @show mean(passes)
        @test mean(passes) > 0.8

        passes = map(Iterators.product(1:4, exp10.(LinRange(-8, -1, 3)))) do (r,ϵ)
            # @show r, ϵ
            s = svd(randn(10,1000))
            u,v = s.U[:,1:r], s.V[:,1:r]
            # A = u*Diagonal(10*(1:r))*v' .+ ϵ .* randn.()
            A = u*Diagonal(s.S[1:r])*v' .+ ϵ .* randn.()
            A .+= 1000*randn.() .* (rand.() .< 0.01)

            Qmed = rpca_ga(A, r, verbose=false, μ = TotalLeastSquares.entrywise_median, iters=120)
            Qm = rpca_ga(A, r, verbose=false, iters=120)
            sum(svdvals([Qmed u])[r+1:2r]) < sum(svdvals([Qm u])[r+1:2r])
            # @test norm(Q'Q - I) < r*sqrt(eps())  # This test should pass but doesn't
        end
        @show mean(passes)
        @test mean(passes) > 0.9
    end
end

@testset "irls" begin
    @info "Testing irls"

    N = 10000
    n = 20
    A = randn(N, n)
    x = randn(n)
    y = A*x
    e = 100*(rand(N) .< 0.01) + 0.1randn(N)
    @time xh = irls(A,y+e, verbose=true,iters=100)
    @test norm(x-xh)/norm(x) < 0.01
    xhls = A\(y+e)
    @test norm(x-xh) < norm(x-xhls)
end

@testset "sls" begin
    @info "Testing sls"

    for i = 1:100
        x = TotalLeastSquares.proj_simplex!(randn(i))
        @test sum(x) ≈ 1
        @test all(x .>= 0)
    end

    s1(x) = x ./= sum(x)
    ##
    n = 10
    N = 2
    results = map(1:100) do _
        a = [(rand(n)) for _ in 1:N]
        A = reduce(hcat,a)
        x = s1(rand(N))
        y = A*x
        yn = y .+ 0.01.*randn.()
        xh1 = TotalLeastSquares.proj_simplex!(A\yn)

        xh2 = sls(A,yn,verbose=false)
        @test sum(xh2) ≈ 1

        norm(x-xh2) < norm(x-xh1) || norm(yn-A*xh2) < norm(yn-A*xh1)
    end
    @test mean(results) >= 0.98
end

@testset "flts" begin
    @info "Testing flts"
    # Example from paper
    x = 10randn(1000)
    a, b = 1, 2
    y = a*x .+ b
    y[801:end] = 5randn(200)
    x[801:end] = 5randn(200) .+ 50
    xb = [x ones(1000)]
    # default
    res = flts(xb,y, verbose = true)
    @test a ≈ res[1]
    @test b ≈ res[2]
    # define % of outliers
    res1 = flts(xb,y, N = 10, outliers = 0.20, verbose = true)
    @test a ≈ res1[1]
    @test b ≈ res1[2]
    # define length of set H and check the subset
    (H, res2, Q) = flts(xb,y, N = 10, h = 800, return_set = true, verbose = true)
    @test a ≈ res2[1]
    @test b ≈ res2[2]
    @test all(H .< 801)
    @test Q + 1 ≈ 1
end
end


text:
function areaBC(r)
    return pi.*r.^2
end


text:
# Creation macros #
# ========================== #
"""
    @define_scenario(def)

Define a scenario type compatible with StochasticPrograms using the syntax
```julia
@define_scenario name = begin
    ...structdef...

    [@zero begin
        ...
        return zero(scenario)
    end]

    [@expectation begin
        ...
        return expected(scenarios)
     end]
end
```
The generated type is referenced through `name` and a default constructor is always generated. This constructor accepts the keyword `probability` to set the probability of the scenario occuring. Otherwise, any internal variables and specialized constructors are defined in the @define_scenario block as they would be in any Julia struct.

If possible, a `zero` method and an `expected` method will be generated for the defined type. Otherwise, or if the default implementation is not desired, these can be user provided through [`@zero`](@ref) and [`@expectation`](@ref).

The defined scenario type will be available on all Julia processes.

## Examples

The following defines a simple scenario ``ξ`` with a single value.

```jldoctest
@define_scenario ExampleScenario = begin
    ξ::Float64
end

ExampleScenario(1.0, probability = 0.5)

# output

ExampleScenario with probability 0.5
  ξ: 1.0

```

See also: [`@zero`](@ref), [`@expectation`](@ref), [`@sampler`](@ref)
"""
macro define_scenario(arg)
    @capture(arg, scenarioname_Symbol = scenariodef_) || error("Invalid syntax. Expected: scenarioname = begin scenariodef end")
    vars = Vector{Symbol}()
    vartypes = Vector{Union{Expr,Symbol}}()
    vardefs = Vector{Expr}()
    zerodefs = Vector{Expr}()
    expectdefs = Vector{Expr}()
    def = postwalk(prettify(scenariodef)) do x
        if @capture(x, var_Symbol::vartype_)
            push!(vars, var)
            push!(vartypes, vartype)
            push!(vardefs, x)
            return x
        end
        if @capture(x, @zero zerodef_)
            push!(zerodefs, zerodef)
            return @q begin end
        end
        if @capture(x, @expectation expectdef_)
            push!(expectdefs, expectdef)
            return @q begin end
        end
        return x
    end
    # Handle zero definition
    if length(zerodefs) > 1
        error("Only provide one zero implementation")
    end
    provided_zerodef, zerodef = if length(zerodefs) == 1
        zerodef = zerodefs[1]
        true, zerodef
    else
        false, Expr(:block)
    end
    # Handle expectation definition
    if length(expectdefs) > 1
        error("Only provide one expectation implementation")
    end
    provided_expectdef, expectdef = if length(expectdefs) == 1
        expectdef = expectdefs[1]
        expectdef = postwalk(prettify(expectdef)) do x
            @capture(x, constructor_Symbol(args__)) && constructor == scenarioname && return :(StochasticPrograms.ExpectedScenario($x))
            return x
        end
        true, expectdef
    else
        false, Expr(:block)
    end
    # Prepare automatic expectation definition
    combine = Expr(:call)
    push!(combine.args, scenarioname)
    push!(combine.args, :($(Expr(:parameters, :(probability = 1.0)))))
    for var in vars
        push!(combine.args, :(probability(s1)*s1.$var + probability(s2)*s2.$var))
    end
    # Prepare automatic scenariotext definition
    textdef = Expr(:block)
    for var in vars
        key = Meta.quot(var)
        push!(textdef.args, :(print(io, "\n  $($key): $(scenario.$var)")))
    end
    # Define scenario type
    code = @q begin
        if StochasticPrograms.supports_expected([$(vartypes...)], $provided_expectdef) || $provided_expectdef
            if StochasticPrograms.supports_zero([$(vartypes...)], $provided_zerodef) || $provided_zerodef
                struct $scenarioname <: AbstractScenario
                    probability::Probability
                    $def
                    function (::Type{$scenarioname})($(vardefs...); probability::AbstractFloat = 1.0)
                        return new(Probability(probability), $(vars...))
                    end
                end
                if $provided_zerodef
                    function Base.zero(::Type{$scenarioname})
                        $zerodef
                    end
                else
                    function Base.zero(::Type{$scenarioname})
                        return $scenarioname(zero.([$(vartypes...)])...; probability = 1.0)
                    end
                end
                if $provided_expectdef
                    function StochasticPrograms.expected(scenarios::Vector{$scenarioname})
                        isempty(scenarios) && return StochasticPrograms.ExpectedScenario(zero($scenarioname))
                        $expectdef
                    end
                else
                    function StochasticPrograms.expected(scenarios::Vector{$scenarioname})
                        isempty(scenarios) && return StochasticPrograms.ExpectedScenario(zero($scenarioname))
                        return StochasticPrograms.ExpectedScenario(reduce(scenarios) do s1, s2
                                                                   $combine
                                                                   end)
                    end
                end
                function StochasticPrograms.scenariotext(io::IO, scenario::$scenarioname)
                    $textdef
                    return io
                end
            else
                @warn "The scenario type $(string($(Meta.quot(scenarioname)))) was not defined. A user-provided implementation \n\n    function zero(::Type{{$(string($(Meta.quot(scenarioname))))})\n        ...\n    end\n\nis required."
            end
        else
            @warn "The scenario type $(string($(Meta.quot(scenarioname)))) was not defined. A user-provided implementation \n\n    function expected(scenarios::Vector{$(string($(Meta.quot(scenarioname))))})\n        ...\n    end\n\nis required."
        end
    end
    code = prettify(code)
    code = @q begin
        if (@__MODULE__) == $(esc(:Main))
            @everywhere begin
                $code
            end
        else
            $(esc(code))
        end
    end
    return prettify(code)
end
"""
    @zero(def)

Define the additive zero scenario inside a @scenario block using the syntax:
```julia
@zero begin
    ...
    return zero_scenario
end
```

## Examples

The following defines a zero scenario for the example scenario defined in [`@define_scenario`](@ref)

```julia
@zero begin
    return ExampleScenario(0.0)
end
```

See also [`@define_scenario`](@ref)
"""
macro zero(def) @warn "@zero should be used inside a @define_scenario block." end
"""
    @expectation(def)

Define how to form the expected scenario inside a [`@define_scenario`](@ref) block. The scenario collection is accessed through the reserved keyword `scenarios`.

```julia
@zero begin
    ...
    return zero_scenario
end
```

## Examples

The following defines expectation for the example scenario defined in [`@scenario`](@ref)

```julia
@expectation begin
    return ExampleScenario(sum([probability(s)*s.ξ for s in scenarios]))
end
```

See also [`@define_scenario`](@ref)
"""
macro expectation(def) @warn "@expectation should be used inside a @scenario block." end
"""
    @scenario(args..., probability = )

Create [`Scenario`](@ref) matching some [`@uncertain`](@ref) declaration with a supplied probability.

    @scenario(var1 = val1, var2 = val2, ..., probability = 1.0)

Create [`Scenario`](@ref) matching [`@uncertain`](@ref) annotation of the form `@uncertain var1, var2, ...`

    @scenario(ξ[i=..., j=..., ...] = values, probability = 1.0)

Create [`Scenario`](@ref) matching [`@uncertain`](@ref) annotation of the form `@uncertain ξ[i=..., j=..., ...]`. `values` must have the same dimension as the specified index sets.

    @scenario(ξ[i=..., j=..., ...], expr, probability = 1.0, requested_container = :Auto)

Create [`Scenario`](@ref) matching [`@uncertain`](@ref) annotation of the form `@uncertain ξ[i=..., j=..., ...]`. Wraps JuMP's `@container` macro to create `DenseAxisArray` or `SparseAxisArray` as underlying data. See `@container` for further syntax information.

## Examples

The following are equivalent ways of creating an instance of the random vector ``[q₁(ξ) q₂(ξ) d₁(ξ) d₂(ξ)]``  of probability ``0.4`` and values ``[24.0 28.0 500.0 100.0]``.

```julia
@scenario q₁ = 24.0 q₂ = 28.0 d₁ = 500.0 d₂ = 100.0 probability = 0.4

@scenario ξ[i in 1:4] = [24.0, 28.0, 500.0, 100.0] probability = 0.4
```

"""
macro scenario(args...)
    args, kw_args, probability, requested_container = _extract_kw_args(args)
    if length(args) == 0
        if length(kw_args) > 1 || @capture(kw_args[1], var_Symbol = val_)
            return :(Scenario(; $(esc.(kw_args)...), probability = $(esc(probability))))
        else
            if @capture(kw_args[1], idx_ = val_)
                idxvars, indices = Containers._build_ref_sets(error, idx)
                values = @q begin
                    if $indices isa VectorizedProductIterator && all(idx -> idx isa Base.OneTo, $indices.prod.iterators)
                        $(esc(val))[$(esc.(idxvars)...)]
                    else
                        if $(esc(val)) isa Dict
                            $(esc(val))[($(esc.(idxvars)...),)]
                        elseif $(esc(val)) isa Array
                            zero(eltype($(esc(val))))
                        else
                            error("Unsupported right hand side in @scenario assignment.")
                        end
                    end
                end
                code = Containers.container_code(idxvars, indices, :($values), requested_container)
                return @q begin
                    ξ = Scenario($code; probability = $(esc(probability)))
                    if ξ.data isa Array
                        size(ξ.data) == size($val) || error("Dimensions of right hand side in @scenario assignment differ from left hand side.")
                    elseif ξ.data isa DenseAxisArray
                        ξ.data.data .= $val
                    elseif ξ.data isa SparseAxisArray
                        length(ξ.data.data) == length($val) || error("Dimensions of right hand side in @scenario assignment differ from left hand side.")
                    end
                    ξ
                end
            end
        end
    end
    @assert length(args) == 2
    @assert isempty(kw_args)
    idx, values = args
    code = Containers.parse_container(error, idx, esc(values), requested_container)
    return :(Scenario($code; probability = $(esc(probability))))
end
function _extract_kw_args(args)
    kw_args = filter(x -> isexpr(x, :(=)) && x.args[1] != :container && x.args[1] != :probability, collect(args))
    flat_args = filter(x->!isexpr(x, :(=)), collect(args))
    requested_container = :Auto
    probability = 1.0
    for kw in args
        if isexpr(kw, :(=)) && kw.args[1] == :container
            requested_container = kw.args[2]
        end
        if isexpr(kw, :(=)) && kw.args[1] == :probability
            probability = kw.args[2]
        end
    end
    return flat_args, kw_args, probability, requested_container
end
"""
    @sampler(def)

Define a sampler type for some `scenariotype` compatible with StochasticPrograms using the syntax
```julia
@sampler samplername = begin
    ...internals...

    @sample scenariotype begin
        ...
        return scenario
    end
end
```
Any internal state required by the sampler, as well as any specialized constructor, are defined in the @sampler block as they would be in any Julia struct. Define the sample operation inside the [`@sample`](@ref) block and specify the `scenariotype` that the sampler returns. The defined sampler will be available on all Julia processes.

## Examples

The following defines a simple dummy sampler, with some internal weight value, for the scenario defined in [`@scenario`](@ref), and samples one scenario.

```jldoctest; setup = :(@scenario ExampleScenario = ξ::Float64), filter = r".*"
@sampler ExampleSampler = begin
    w::Float64

    ExampleSampler(w::AbstractFloat) = new(w)

    @sample ExampleScenario begin
        @parameters w
        return ExampleScenario(w*randn(), probability = rand())
    end
end
s = ExampleSampler(2.0)
s()

# output

ExampleScenario with probability 0.29
  ξ: 1.48

```

See also: [`@sample`](@ref), [`@scenario`](@ref)
"""
macro sampler(arg)
    @capture(prettify(arg), samplername_Symbol = samplerdef_) || error("Invalid syntax. Expected: sampler = begin samplerdef end")
    scenarioname = :undefined
    sampledefs = Vector{Expr}()
    def = postwalk(prettify(samplerdef)) do x
        if @capture(x, @parameters args__)
            code = Expr(:block)
            for param in args
                push!(code.args, :($param = sampler.$param))
            end
            return code
        elseif @capture(x, @sample sname_ sampledef_)
            scenarioname = sname
            push!(sampledefs, sampledef)
            return @q begin end
        else
            return x
        end
    end
    # Handle sample definition
    length(sampledefs) == 1 || error("Must provide exactly one @sample operation.")
    sampledef = sampledefs[1]
    # Define sampler type
    code = @q begin
        struct $samplername <: AbstractSampler{$scenarioname}
            $def
        end
        function (sampler::$samplername)()
            $sampledef
        end
    end
    code = prettify(code)
    code = @q begin
        if (@__MODULE__) == $(esc(:Main))
            @everywhere begin
                $code
            end
        else
            $(esc(code))
        end
    end
    return prettify(code)
end
"""
    @sample(def)

Define the sample operaton inside a @sampler block, using the syntax
```julia
@sample begin
    ...
    return sampled_scenario
end
```
The sampler object is referenced through the reserved keyword `sampler`, from which any internals can be accessed.
"""
macro sample(def) @warn "@sample should be used inside a @sampler block." end
"""
    @first_stage(def)

Add a first stage model generation recipe to `stochasticprogram` using the syntax
```julia
@first_stage stochasticprogram::StochasticProgram = begin
    ...
end [defer]
```
where JuMP syntax is used inside the block to define the first stage model. During definition, the first stage model is referenced through the reserved keyword `model`.

## Examples

The following defines the first stage model given by:

```math
  minimize 100x₁ + 150x₂
    s.t  x₁ + x₂ ≤ 120
         x₁ ≥ 40
         x₂ ≥ 20
```

```julia
@first_stage sp = begin
    @decision(model, x₁ >= 40)
    @decision(model, x₂ >= 20)
    @objective(model, Min, 100*x₁ + 150*x₂)
    @constraint(model, x₁ + x₂ <= 120)
end
```

See also: [`@second_stage`](@ref)
"""
macro first_stage(arg)
    @capture(arg, sp_Symbol = def_) || error("Invalid syntax. Expected @first_stage sp = begin ... end")
    return esc(:(@stage 1 $arg))
end
"""
    @second_stage(def)

Add a second stage model generation recipe to `stochasticprogram` using the syntax
```julia
@second_stage stochasticprogram::StochasticProgram = begin
    @known var1 var2 ...
    ...
end
```
where JuMP syntax is used inside the block to define the second stage model. During definition, the second stage model is referenced through the reserved keyword `model`.

## Examples

The following defines the second stage model given by:
```math
  minimize q₁(ξ)y₁ + q₂(ξ)y₂
    s.t  6y₁ + 10y₂ ≤ 60x₁
         8y₁ + 5y₂ ≤ 60x₂
         0 ≤ y₁ ≤ d₁(ξ)
         0 ≤ y₂ ≤ d₂(ξ)
```
where ``q₁(ξ), q₂(ξ), d₁(ξ), d₂(ξ)`` depend on the scenario ``ξ`` and ``x₁, x₂`` are first stage variables. Two scenarios are added so that two second stage models are generated.

```julia
@second_stage sp = begin
    @known x₁ x₂
    @uncertain q₁ q₂ d₁ d₂
    @variable(model, 0 <= y₁ <= d₁)
    @variable(model, 0 <= y₂ <= d₂)
    @objective(model, Min, q₁*y₁ + q₂*y₂)
    @constraint(model, 6*y₁ + 10*y₂ <= 60*x₁)
    @constraint(model, 8*y₁ + 5*y₂ <= 80*x₂)
end
```

See also: [`@first_stage`](@ref)
"""
macro second_stage(arg)
    @capture(arg, sp_Symbol = def_) || error("Invalid syntax. Expected @second_stage sp = begin ... end")
    return esc(@q begin
        @stage 2 $arg
        generate!($sp)
    end)
end
"""
    @parameters(def)

Define the problem parameters in a @stage block
```julia
@parameters param1, param2, ...
```
possibly with default values. Any defined parameter without a default value must be supplied as a keyword argument to [`instantiate`](@ref) when creating models.

## Examples

```julia
@parameters d

@parameters begin
    Crops = [:wheat, :corn, :beets]
    Cost = Dict(:wheat=>150, :corn=>230, :beets=>260)
    Budget = 500
end
```

See also [`@decision`](@ref), [`@uncertain`](@ref), [`@stage`](@ref)
"""
macro parameters(def) @warn "@parameters should be used inside a @stage block." end
"""
    @decision(model, expr, args..., kw_args...)

Add a decision variable to `model` described by the expression `expr`. If used inside a [`@stage`](@ref) block, the created variable can be used in subsequent stage blocks. [`@recourse`](@ref) should be used to mark decisions in the final stage. See `@variable` for syntax details.

## Examples

```julia
@decision(model, x >= 40)
```

See also [`@recourse`](@ref), [`@parameters`](@ref), [`@uncertain`](@ref), [`@stage`](@ref)
"""
macro decision(def...) @warn "@decision should be used inside a @stage block." end
"""
    @recourse(model, expr, args..., kw_args...)

Add a recourse decision variable to `model` described by the expression `expr`. Replaces [`@decision`](@ref) in the [`@stage`](@ref) block of the final stage, and can only be used there. See `@variable` for syntax details.

## Examples

```julia
@recourse(model, 0 <= y <= 1)
```

See also [`@decision`](@ref), [`@parameters`](@ref), [`@uncertain`](@ref), [`@stage`](@ref)
"""
macro recourse(def...) @warn "@recourse should be used inside the @stage block of the final stage." end
macro _decision(args...)
    args = [args...]
    known = false
    recourse = false
    # Check context of decision definition
    if args[1] isa Symbol && args[1] == :known
        known = true
    elseif args[1] isa Symbol && args[1] == :recourse
        recourse = true
    else
        args[1] == :decision || error("Incorrect usage of @_decision.")
    end
    # Remove context from args
    deleteat!(args, 1)
    # Cache stage and remove from args
    stage = args[1]
    deleteat!(args, 1)
    # Move any set definitions
    set = NoSpecifiedConstraint()
    for (i, arg) in enumerate(args)
        set, arg, found = extract_set(arg)
        if found
            if arg == :()
                deleteat!(args, i)
            else
                args[i] = arg
            end
            break
        end
    end
    # Return expression based on context
    if known
        return esc(:(@variable $((args)...) set = StochasticPrograms.KnownSet($stage)))
    else
        return esc(:(@variable $((args)...) set = StochasticPrograms.DecisionSet($stage, constraint = $set, is_recourse = $recourse)))
    end
end

"""
    @known(def)

Annotate each decision taken in the previous stage. Any [`@decision`](@ref) included in a [`@stochastic_model`](@ref) definition will implicitly add `@known` annotations to subsequent stages.

## Examples

```julia
@known x₁, x₂
```

See also [`@decision`](@ref), [`@parameters`](@ref), [`@uncertain`](@ref), [`@stage`](@ref)
"""
macro known(args...)
    code = Expr(:block)
    for var in args
        varkey = Meta.quot(var)
        push!(code.args, :($(esc(var)) = $(esc(:model))[$varkey]))
    end
    return code
end
"""
    @uncertain(def)

In a @stage block, annotate each uncertain variable using the syntax
```julia
@uncertain var1, var2, ...
```
or using JuMP's container syntax
```julia
@uncertain ξ[i=..., j=..., ...]
```
This assumes that the [`Scenario`] type is used. Matching scenario data is then conveniently created using [`@scenario`](@ref).

Alternatively, user-defined scenarios can be specified by annotating the type. Also, inside a @stochastic_model block, user-defined scenarios can be created during the @uncertain annotation, using [`@define_scenario`](@ref) syntax.

## Examples

The following are equivalent ways of declaring a random vector ``[q₁(ξ) q₂(ξ) d₁(ξ) d₂(ξ)]`` in a `@stage` block, and creating a matching scenario instance of probability ``0.4`` and values ``[24.0 28.0 500.0 100.0]``.ö

```julia
@uncertain q₁ q₂ d₁ d₂
ξ₁ = @scenario q₁ = 24.0 q₂ = 28.0 d₁ = 500.0 d₂ = 100.0 probability = 0.4

@uncertain ξ[i in 1:4]
ξ₁ = @scenario ξ[i in 1:4] = [24.0, 28.0, 500.0, 100.0] probability = 0.4

@define_scenario SimpleScenario = begin
    q₁::Float64
    q₂::Float64
    d₁::Float64
    d₂::Float64
end
@uncertain ξ::SimpleScenario
ξ₁ = SimpleScenario(24.0, 28.0, 500.0, 100.0, probability = 0.4)

@stochastic_model begin
    ...
    @uncertain ξ::SimpleScenario = begin
        q₁::Float64
        q₂::Float64
        d₁::Float64
        d₂::Float64
    end
    ...
end
ξ₁ = SimpleScenario(24.0, 28.0, 500.0, 100.0 probability = 0.4)
```

See also [`@scenario`](@ref), [`@define_scenario`](@ref), [`@parameters`](@ref), [`@decision`](@ref), [`@stage`](@ref)
"""
macro uncertain(def) @warn "@uncertain should be used inside a @stage block." end
"""
    @stage(def)

Add a stage model generation recipe to `stochasticprogram` using the syntax
```julia
@stage stage stochasticprogram::StochasticProgram = begin
    @parameters param1 param2 ...
    @decision(model, var) ...
    @uncertain ξ
    ... JuMPdef ...
    ...
end
```
where JuMP syntax is used inside the block to define the stage model. During definition, the second stage model is referenced through the reserved keyword `model`.

## Examples

The following defines the first stage model given by:
```math
  minimize 100x₁ + 150x₂
    s.t  x₁ + x₂ ≤ 120
         x₁ ≥ 40
         x₂ ≥ 20
```
and the second-stage model given by:
```math
  maximize q₁(ξ)y₁ + q₂(ξ)y₂
    s.t  6y₁ + 10y₂ ≤ 60x₁
         8y₁ + 5y₂ ≤ 60x₂
         0 ≤ y₁ ≤ d₁(ξ)
         0 ≤ y₂ ≤ d₂(ξ)
```
where ``q₁(ξ), q₂(ξ), d₁(ξ), d₂(ξ)`` depend on the scenario ``ξ`` and ``x₁, x₂`` are first stage variables. Two scenarios are added so that two second stage models are generated.

```jldoctest
ξ₁ = @scenario q₁ = 24.0 q₂ = 28.0 d₁ = 500.0 d₂ = 100.0 probability = 0.4
ξ₂ = @scenario q₁ = 28.0 q₂ = 32.0 d₁ = 300.0 d₂ = 300.0 probability = 0.6

sp = StochasticProgram([ξ₁, ξ₂])

@stage 1 sp = begin
    @decision(model, x₁ >= 40)
    @decision(model, x₂ >= 20)
    @objective(model, Min, 100*x₁ + 150*x₂)
    @constraint(model, x₁ + x₂ <= 120)
end

@stage 2 sp = begin
    @uncertain q₁ q₂ d₁ d₂
    @variable(model, 0 <= y₁ <= d₁)
    @variable(model, 0 <= y₂ <= d₂)
    @objective(model, Max, q₁*y₁ + q₂*y₂)
    @constraint(model, 6*y₁ + 10*y₂ <= 60*x₁)
    @constraint(model, 8*y₁ + 5*y₂ <= 80*x₂)
end

# output

Stochastic program with:
 * 2 decision variables
 * 2 scenarios of type Scenario
Solver is default solver

```

See also: [`@parameters`](@ref), [`@decision`](@ref), [`@uncertain`](@ref)
"""
macro stage(stage, args)
    @capture(args, sp_Symbol = def_) || error("Invalid syntax. Expected @stage stage sp = begin ... end")
    # Flags for error checking
    seen_decision = :(false)
    seen_recourse = :(false)
    # Decision definitions might require parameter calculations,
    # so we first need to extract and save any such lines
    vardefs = Expr(:block)
    decisiondefs = postwalk(def) do x
        if  @capture(x, @variable(m_Symbol, variabledef__)) ||
            @capture(x, @known(m_Symbol, variabledef__)) ||
            @capture(x, @constraint(m_Symbol, constdef__)) ||
            @capture(x, @expression(m_Symbol, expressiondef__)) ||
            @capture(x, @objective(m_Symbol, objdef__)) ||
            @capture(x, @parameters args__) ||
            @capture(x, @uncertain args__)
            # Skip any line related to the JuMP model, stochastics, or unhandled @parameter lines
            return Expr(:block)
        elseif @capture(x, @decision args__) || @capture(x, @recourse args__)
            # Handle @decision
            return @q begin
                StochasticPrograms.@_decision known $stage $((args)...)
            end
        else
            # Anything else could be required for decision variable construction, and is therefore saved
            return x
        end
    end
    # Next, handle @decision annotations in main definition
    def = postwalk(def) do x
        if @capture(x, @decision args__)
            # Bookkeep for error checking
            seen_decision = :(true)
            return @q begin
                StochasticPrograms.@_decision decision $stage $((args)...)
            end
        end
        if @capture(x, @recourse args__)
            # Bookkeep for error checking
            seen_recourse = :(true)
            return @q begin
                StochasticPrograms.@_decision recourse $stage $((args)...)
            end
        end
        return x
    end
    # Handle parameters
    def = postwalk(def) do x
        if @capture(x, @parameters arg_)
            code = Expr(:block)
            for paramdef in block(prettify(arg)).args
                if @capture(paramdef, key_Symbol = val_) || @capture(paramdef, key_Symbol)
                    push!(code.args, :($key = stage.$key))
                else
                    error("Incorrect @parameters syntax. Specify parameter names, possibly with a default value.")
                end
            end
            # Extracted parameters might be required for decision variable construction
            pushfirst!(decisiondefs.args, code)
            return code
        elseif @capture(x, @parameters args__)
            code = Expr(:block)
            for param in args
                push!(code.args, :($param = stage.$param))
            end
            # Extracted parameters might be required for decision variable construction
            pushfirst!(decisiondefs.args, code)
            return code
         elseif @capture(x, @uncertain var_Symbol::t_Symbol)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            return :($var::$t = scenario)
        elseif @capture(x, @uncertain vars__ from scenvar_Symbol::t_Symbol)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                $scenvar::$t = scenario
            end
            for var in vars
                varkey = Meta.quot(var)
                push!(code.args, :($var = getproperty($scenvar, $varkey)))
            end
            return code
        elseif @capture(x, @uncertain vars__ from t_Symbol)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                scenvar::$t = scenario
            end
            for var in vars
                varkey = Meta.quot(var)
                push!(code.args, :($var = getproperty(scenvar, $varkey)))
            end
            return code
        elseif @capture(x, @uncertain var_Symbol::t_Symbol = scenariodef_)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            return @q begin
                $var::$t = scenario
            end
        elseif @capture(x, @uncertain vars__ from scenvar_Symbol::t_Symbol = scenariodef_)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                $scenvar::$t = scenario
            end
            for var in vars
                varkey = Meta.quot(var)
                push!(code.args, :($var = getproperty($scenvar, $varkey)))
            end
            return code
        elseif @capture(x, @uncertain vars__ from t_Symbol = scenariodef_)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                scenvar::$t = scenario
            end
            for var in vars
                varkey = Meta.quot(var)
                push!(code.args, :($var = getproperty(scenvar, $varkey)))
            end
            return code
        elseif @capture(x, @uncertain args_) && typeof(args) == Expr && args.head == :ref
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            scenvar = first(args.args)
            reference = :(Containers.@container $args "uncertain value")
            code = @q begin
                typeof(scenario) <: Scenario{D} where D <: Union{Array, StochasticPrograms.DenseAxisArray} || error("@uncertain declarations of type `@uncertain ξ[i = ..., j = ..., ...]` only support scenarios of type `Scenario` with Array or DenseAxisArray as underlying data. Consider declaring a custom scenario type.")
                size(scenario.data) == size($reference) || error("Given scenario \n\n$scenario \n\ndoes not match @uncertain declaration \n\n$($reference).")
                $scenvar = scenario.data
            end
            return code
        elseif @capture(x, @uncertain args_) && typeof(args) == Expr && args.head == :typed_vcat
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            scenvar = first(args.args)
            reference = :(Containers.@container $args "uncertain value")
            code = @q begin
                typeof(scenario) <: Scenario{D} where D <: StochasticPrograms.SparseAxisArray || error("@uncertain declarations of type `@uncertain ξ[i = ..., j = ..., ...; ...]` only support scenarios of type `Scenario` with SparseAxisArray as underlying data. Consider declaring a custom scenario type.")
                keys(scenario.data.data) == keys($reference.data) || error("Given scenario \n\n$scenario \n\ndoes not match @uncertain declaration \n\n$($reference).")
                $scenvar = scenario.data
            end
            return code
        elseif @capture(x, @uncertain args__)
            stage == 1 && error("@uncertain declarations cannot be used in the first stage.")
            code = @q begin
                typeof(scenario) <: Scenario || error("@uncertain declarations of type `@uncertain var1, var2, ...` only support scenarios of type `Scenario`. Consider declaring a custom scenario type.")
            end
            for var in args
                varkey = Meta.quot(var)
                code = @q begin
                    $code
                    $var = try
                        $var = scenario.data[$varkey]
                    catch err
                        error("Given scenario $scenario does not match @uncertain declaration.")
                    end
                end
            end
            return code
        else
            return x
        end
    end
    # The first stage is treated different from the subsequent stages
    generatordefs = if stage == 1
        @q begin
            $(esc(sp)).generator[:stage_1] = ($(esc(:model))::JuMP.Model, $(esc(:stage))) -> begin
                $(esc(def))
	            return $(esc(:model))
            end
        end
    else
        @q begin
            $(esc(sp)).generator[Symbol(:stage_,$stage)] = ($(esc(:model))::JuMP.Model, $(esc(:stage)), $(esc(:scenario))::AbstractScenario) -> begin
                $(esc(def))
	            return $(esc(:model))
            end
        end
    end
    # Create definition code
    code = @q begin
        isa($(esc(sp)), StochasticProgram) || error("Given object is not a stochastic program.")
        n = num_stages($(esc(sp)))
        $stage > n && error("Cannot specify stage $($stage) for stochastic program with $n stages.")
        $stage == n && $seen_decision && error("@decision declarations cannot be used in the final stage. Consider @recourse instead.")
        if $stage < n
            $seen_recourse && error("@recourse declarations can only be used in the final stage.")
            $(esc(sp)).generator[Symbol(:stage_,$stage,:_decisions)] = ($(esc(:model))::JuMP.Model, $(esc(:stage))) -> begin
                $(esc(decisiondefs))
	            return $(esc(:model))
            end
        end
        # Stage model generation code
        $generatordefs
        $(esc(sp))
    end
    # Return code
    return code
end
# ========================== #


text:
using Distributed
@everywhere using MCPhylo

## Data
surgical = Dict{Symbol, Any}(
  :r => [0, 18, 8, 46, 8, 13, 9, 31, 14, 8, 29, 24],
  :n => [47, 148, 119, 810, 211, 196, 148, 215, 207, 97, 256, 360]
)
surgical[:N] = length(surgical[:r])

## Model Specification
model = Model(

  r = Stochastic(1,
    (n, p, N) ->
      UnivariateDistribution[Binomial(n[i], p[i]) for i in 1:N],
    false
  ),

  p = Logical(1,
    b -> invlogit.(b)
  ),

  b = Stochastic(1,
    (mu, s2) -> Normal(mu, sqrt(s2)),
    false
  ),

  mu = Stochastic(
    () -> Normal(0, 1000)
  ),

  pop_mean = Logical(
    mu -> invlogit(mu)
  ),

  s2 = Stochastic(
    () -> InverseGamma(0.001, 0.001)
  )

)

## Initial Values
inits = [
  Dict(:r => surgical[:r], :b => fill(0.1, surgical[:N]), :s2 => 1, :mu => 0),
  Dict(:r => surgical[:r], :b => fill(0.5, surgical[:N]), :s2 => 10, :mu => 1)
]

## Sampling Scheme
scheme = [NUTS(:b),
          Slice([:mu, :s2], 1.0)]
setsamplers!(model, scheme)

## MCMC Simulations
sim = mcmc(model, surgical, inits, 10000, burnin=2500, thin=2, chains=2)
describe(sim)


text:
"""
    make_selectors_from_cliques(cliques, n)

Builds selector matrices `Tℓ` from `cliques`, a list of the cliques in a graph
with `n` vertices.

`Tℓ*X*Tℓ'` is the submatrix of `X` corresponding to clique `ℓ`.
"""
function make_selectors_from_cliques(cliques, n)
    selector_mats = [spzeros(length(cliques[i]), n) for i in 1:length(cliques)]
    for i in 1:length(cliques)
        idx = 1
        for node in cliques[i]
            selector_mats[i][idx, node] = 1.0
            idx += 1
        end
    end
    return selector_mats
end

# TODO: make this multithreaded
# TODO: remove n from args
"""
    make_selectors_from_clique_graph(cg::CliqueGraph, n)

Builds selector matrices `Tℓ` from `cg`, a CliqueGraph with `n` vertices.

`Tℓ*X*Tℓ'` is the submatrix of `X` corresponding to clique `ℓ`.
"""
function make_selectors_from_clique_graph(cg::CliqueGraph, n)
    # n = size(cg.membership_mat, 1)
    m = length(cg.active_cliques)
    selector_mats = Vector{SparseMatrixCSC}(undef, m)
    for (i, cnum) in enumerate(cg.active_cliques)
        clique = cg.membership_mat[:,cnum]
        selector_mat = spzeros(sum(clique), n)

        idx = 1
        for node in findnz(clique)[1]
            selector_mat[idx, node] = 1.0
            idx += 1
        end
        selector_mats[i] = selector_mat
    end
    return selector_mats
end

"""
    get_selectors(input_mat::SparseMatrixCSC; verbose=true, ret_cliques=true)

Returns the (merged) cliques of the graph corresponding to the sparsity pattern
of `input_mat` (after a permutation to reduce fill-in) and optionally the cliques.
Also returns the fill-reducing permutation and and inverse permutation used.
"""
function get_selectors(input_mat::SparseMatrixCSC; verbose=true, ret_cliques=true, perm="amd")
    n = size(input_mat)[1]
    preprocess!(input_mat)

    sp = sparsity_pattern(input_mat)
    # TODO: make this work for matrices that are block diagonal (under reordering)
    if is_separable(sp)
        error(ArgumentError("Input matrix should not be block diagonal; decompose the blocks separately."))
    end

    # L is the chordal extension of sp under reordering perm
    perm, iperm, L = get_chordal_extension(sp; verbose=verbose, perm=perm)

    cliques = get_cliques(L)
    cg = generate_clique_graph(cliques, n)
    merge_cliques!(cg; verbose=verbose)

    if ret_cliques
        Cls = get_cliques(cg)
        Tls = make_selectors_from_cliques(Cls, n)
        return perm, iperm, Tls, Cls
    end

    Tls = make_selectors_from_clique_graph(cg, n)
    return perm, iperm, Tls
end


text:
using StarAlgebras
using Test
using LinearAlgebra
using SparseArrays

include("test_example_words.jl")

@testset "StarAlgebras" begin
    include("mtables.jl")
    include("constructors.jl")
    include("arithmetic.jl")

    if VERSION >= v"1.3.0"
        using Pkg
        Pkg.add("Groups")
        include("sum_of_squares.jl")
    end
end


text:
# CuArray{ComplexF32}
function gpu_downconvert!(
    downconverted_signal::CuVector{ComplexF32},
    carrier::CuVector{ComplexF32},
    signal::CuVector{ComplexF32},
    start_sample::Integer,
    num_samples_left::Integer
)
    @. @views downconverted_signal[start_sample:num_samples_left + start_sample - 1] =
            signal[start_sample:num_samples_left + start_sample - 1] *
            conj(carrier[start_sample:num_samples_left + start_sample - 1])
end

# CuArray{ComplexF32} Matrix
function gpu_downconvert!(
    downconverted_signal::CuMatrix{ComplexF32},
    carrier::CuVector{ComplexF32},
    signal::CuMatrix{ComplexF32},
    start_sample::Integer,
    num_samples_left::Integer
)
    @. downconverted_signal = signal * conj(carrier)
end

# StructArray{CuArray} Vector
function gpu_downconvert!(
    downconverted_signal_re::CuVector{Float32},
    downconverted_signal_im::CuVector{Float32},
    carrier_re::CuVector{Float32},
    carrier_im::CuVector{Float32},
    signal_re::CuVector{Float32},
    signal_im::CuVector{Float32},
    start_sample::Integer,
    num_samples_left::Integer
    )
    @. downconverted_signal_re = signal_re * carrier_re + signal_im * carrier_im
    @. downconverted_signal_im = signal_im * carrier_re - signal_re * carrier_im
end

# StructArray{CuArray} Matrix
function gpu_downconvert!(
    downconverted_signal_re::CuMatrix{Float32},
    downconverted_signal_im::CuMatrix{Float32},
    carrier_re::CuVector{Float32},
    carrier_im::CuVector{Float32},
    signal_re::CuMatrix{Float32},
    signal_im::CuMatrix{Float32},
    start_sample::Integer,
    num_samples_left::Integer
    )
    @. downconverted_signal_re = signal_re * carrier_re + signal_im * carrier_im
    @. downconverted_signal_im = signal_im * carrier_re - signal_re * carrier_im
end

# Float32 implementation of the orig. CPU function
function cpu_downconvert!(
    downconverted_signal_re::Vector{Float32},
    downconverted_signal_im::Vector{Float32},
    carrier_re::Vector{Float32},
    carrier_im::Vector{Float32},
    signal_re::Vector{Float32},
    signal_im::Vector{Float32},
    start_sample::Integer,
    num_samples_left::Integer
)
    @avx unroll = 3 for i = start_sample:num_samples_left + start_sample - 1
        downconverted_signal_re[i] = signal_re[i] * carrier_re[i] +
            signal_im[i] * carrier_im[i]
        downconverted_signal_im[i] = signal_im[i] * carrier_re[i] -
            signal_re[i] * carrier_im[i]
    end
end

# Float32 implementation of the orig. CPU function for matrices
function cpu_downconvert!(
    downconverted_signal_re::Matrix{Float32},
    downconverted_signal_im::Matrix{Float32},
    carrier_re::Vector{Float32},
    carrier_im::Vector{Float32},
    signal_re::Matrix{Float32},
    signal_im::Matrix{Float32},
    start_sample::Integer,
    num_samples_left::Integer
)
    @avx unroll = 3 for i = start_sample:num_samples_left + start_sample - 1, j = 1:size(signal_re, 2)
        # Calculate signal * carrier'
        downconverted_signal_re[i, j] = signal_re[i, j] * carrier_re[i] +
            signal_im[i, j] * carrier_im[i]
        downconverted_signal_im[i, j] = signal_im[i, j] * carrier_re[i] -
            signal_re[i, j] * carrier_im[i]
    end
end