text:
#
# methods.jl -
#
# Implement non-specific methods for mappings.
#
#-------------------------------------------------------------------------------
#
# This file is part of LazyAlgebra (https://github.com/emmt/LazyAlgebra.jl)
# released under the MIT "Expat" license.
#
# Copyright (c) 2017-2021 Éric Thiébaut.
#

@noinline function unimplemented(::Type{P},
                                 ::Type{T}) where {P<:Operations, T<:Mapping}
    throw(UnimplementedOperation("unimplemented operation `$P` for mapping $T"))
end

@noinline function unimplemented(func::Union{AbstractString,Symbol},
                                 ::Type{T}) where {T<:Mapping}
    throw(UnimplementedMethod("unimplemented method `$func` for mapping $T"))
end

"""
    @callable T

makes instances of concrete type `T` callable as a regular `LazyAlgebra`
mapping, that is `A(x)` yields `apply(A,x)` for any `A` of type `T`.

"""
macro callable(T)
    quote
	(A::$(esc(T)))(x) = apply(A, x)
    end
end
@callable Adjoint
@callable Inverse
@callable InverseAdjoint
@callable Jacobian
@callable Gram
@callable Scaled
@callable Sum
@callable Composition

show(io::IO, ::MIME"text/plain", A::Mapping) = show(io, A)

show(io::IO, A::Identity) = print(io, "Id")

function show(io::IO, A::Scaled)
    λ, M = multiplier(A), unscaled(A)
    if λ == -1
        print(io, "-")
    elseif λ != 1
        print(io, λ, "⋅")
    end
    show(io, M)
end

function show(io::IO, A::Scaled{<:Sum})
    λ, M = multiplier(A), unscaled(A)
    if λ == -1
        print(io, "-(")
    elseif λ != 1
        print(io, λ, "⋅(")
    end
    show(io, M)
    if λ != 1
        print(io, ")")
    end
end

function show(io::IO, A::Adjoint{<:Mapping})
    show(io, unveil(A))
    print(io, "'")
end

function show(io::IO, A::Adjoint{T}) where {T<:Union{Scaled,Composition,Sum}}
    print(io, "(")
    show(io, unveil(A))
    print(io, ")'")
end

function show(io::IO, A::Inverse{<:Mapping})
    print(io, "inv(")
    show(io, unveil(A))
    print(io, ")")
end

function show(io::IO, A::InverseAdjoint{<:Mapping})
    print(io, "inv(")
    show(io, unveil(A))
    print(io, ")'")
end

function show(io::IO, A::Jacobian{<:Mapping})
    print(io, "∇(")
    show(io, primitive(A))
    print(io, ",x)")
end

function show(io::IO, A::Sum{N}) where {N}
    function show_term(io::IO, A::Sum)
        print(io, "(")
        show(io, A)
        print(io, ")")
    end
    show_term(io::IO, A::Mapping) = show(io, A)

    for i in 1:N
        let B = A[i]
            if isa(B, Scaled)
                λ, M = multiplier(B), unscaled(B)
                if λ < 0
                    print(io, (i == 1 ? "-" : " - "))
                    λ = -λ
                elseif i > 1
                    print(io, " + ")
                end
                if λ != 1
                    print(io, λ, "⋅")
                end
                show_term(io, M)
            else
                if i > 1
                    print(io, " + ")
                end
                show_term(io, B)
            end
        end
    end
end

function show(io::IO, A::Composition{N}) where {N}
    for i in 1:N
        let B = A[i]
            if i > 1
                print(io, "⋅")
            end
            if isa(B, Sum) || isa(B, Scaled)
                print(io, "(")
                show(io, B)
                print(io, ")")
            else
                show(io, B)
            end
        end
    end
end

"""
    terms(A)

yields the list (as a tuple) of terms that compose mapping `A`.  If `A` is a
sum or a composition of mappings, the list of terms is returned; otherwise, the
1-tuple `(A,)` is returned.

If `A` is sum or a composition of mappings, `Tuple(A)` yields the same result
as `terms(A)`.

"""
terms(A::Union{Sum,Composition}) = getfield(A, :ops)
terms(A::Mapping) = (A,)

"""
    unveil(A)

unveils the mapping embedded in mapping `A` if it is a *decorated* mapping (see
[`LazyAlgebra.DecoratedMapping`](@ref)); otherwise, just returns `A` if it is
not a *decorated* mapping.

As a special case, `A` may be an instance of `LinearAlgebra.UniformScaling` and
the result is the LazyAlgebra mapping corresponding to `A`.

"""
unveil(A::DecoratedMapping) = getfield(A, :op)
unveil(A::Mapping) = A
unveil(A::UniformScaling) = multiplier(A)*Id
Mapping(A::UniformScaling) = unveil(A)

"""
    unscaled(A)

and

    multiplier(A)

respectively yield the mapping `M` and the multiplier `λ` if `A = λ*M` is a
scaled mapping (see [`Scaled`](@ref)); `A` and `1` otherwise.  Note that these
methods also work for intances of `LinearAlgebra.UniformScaling`.

"""
unscaled(A::Mapping) = A
unscaled(A::Scaled) = getfield(A, :M)
unscaled(A::UniformScaling) = Id
multiplier(A::Scaled) = getfield(A, :λ)
multiplier(A::Mapping) = 1
multiplier(A::UniformScaling) = getfield(A, :λ)

@doc @doc(unscaled) multiplier

"""
    primitive(J)

and

    variables(J)

respectively yield the mapping `A` and the variables `x` embedded in Jacobian
`J = ∇(A,x)`.

"""
primitive(J::Jacobian) = getfield(J, :A)
variables(J::Jacobian) = getfield(J, :x)
@doc @doc(primitive) variables

"""
    identifier(A)

yields a hash value identifying almost uniquely the unscaled mapping `A`.  This
identifier is used for sorting terms in a sum of mappings.

!!! warning
    For now, the identifier is computed as `objectid(unscaled(A))` and is
    unique with a very high probability.

"""
identifier(A::Mapping) = objectid(unscaled(A))

Base.isless(A::Mapping, B::Mapping) = isless(identifier(A), identifier(B))

# Extend base methods to simplify the code for reducing expressions.
first(A::Mapping) = A
last(A::Mapping) = A
first(A::Union{Sum,Composition}) = @inbounds A[1]
last(A::Union{Sum{N},Composition{N}}) where {N} = @inbounds A[N]
firstindex(A::Union{Sum,Composition}) = 1
lastindex(A::Union{Sum{N},Composition{N}}) where {N} = N
length(A::Union{Sum{N},Composition{N}}) where {N} = N
eltype(::Type{<:Sum{N,T}}) where {N,T} = eltype(T)
eltype(::Type{<:Composition{N,T}}) where {N,T} = eltype(T)
Tuple(A::Union{Sum,Composition}) = terms(A)

@inline @propagate_inbounds getindex(A::Union{Sum,Composition}, i) =
    getindex(terms(A), i)

"""
    input_type([P=Direct,] A)
    output_type([P=Direct,] A)

yield the (preferred) types of the input and output arguments of the operation
`P` with mapping `A`.  If `A` operates on Julia arrays, the element type, list
of dimensions, `i`-th dimension and number of dimensions for the input and
output are given by:

    input_eltype([P=Direct,] A)          output_eltype([P=Direct,] A)
    input_size([P=Direct,] A)            output_size([P=Direct,] A)
    input_size([P=Direct,] A, i)         output_size([P=Direct,] A, i)
    input_ndims([P=Direct,] A)           output_ndims([P=Direct,] A)

For mappings operating on Julia arrays, only `input_size(A)` and
`output_size(A)` have to be implemented.

Also see: [`vcreate`](@ref), [`apply!`](@ref), [`LinearMapping`](@ref),
[`Operations`](@ref).

"""
function input_type end

for sfx in (:size, :eltype, :ndims, :type),
    pfx in (:output, :input)

    fn1 = Symbol(pfx, "_", sfx)

    for P in (Direct, Adjoint, Inverse, InverseAdjoint)

        fn2 = Symbol(P === Adjoint || P === Inverse ?
                     (pfx === :output ? :input : :output) : pfx, "_", sfx)

        T = (P === Adjoint || P === InverseAdjoint ? LinearMapping : Mapping)

        # Provide basic methods for the different operations and for tagged
        # mappings.
        @eval $fn1(::Type{$P}, A::$T) = $fn2(A)
        if P !== Direct
            @eval $fn1(A::$P{<:$T}) = $fn2(unveil(A))
        end
        if sfx === :size
            if P !== Direct
                @eval $fn1(A::$P{<:$T}, dim...) = $fn2(unveil(A), dim...)
            end
            @eval $fn1(::Type{$P}, A::$T, dim...) = $fn2(A, dim...)
        end
    end

    # Link documentation for the basic methods.
    if fn1 !== :input_type
        @eval @doc @doc(:input_type) $fn1
    end

end

# Provide default methods for `$(sfx)_size(A, dim...)` and `$(sfx)_ndims(A)`.
for pfx in (:input, :output)
    get_size = Symbol(pfx, "_size")
    get_ndims = Symbol(pfx, "_ndims")
    @eval begin
        $get_ndims(A::Mapping) = length($get_size(A))
        $get_size(A::Mapping, dim) = $get_size(A)[dim]
        function $get_size(A::Mapping, dim...)
            dims = $get_size(A)
            ntuple(i -> dims[dim[i]], length(dim))
        end
    end
end

for f in (:input_eltype, :output_eltype, :input_size, :output_size)
    @eval $f(::T) where {T<:Mapping} = unimplemented($(string(f)), T)
end

"""
    nrows(A)

yields the *equivalent* number of rows of the linear operator `A`.  Not all
operators extend this method.

In the implemented generalization of linear operators, the equivalent number of
rows is the number of element of the result of applying the operator be it
single- or multi-dimensional.

"""
nrows(A::LinearMapping) = prod(row_size(A))
@noinline nrows(A::Mapping) =
    throw(ArgumentError("`nrows` is only implemented for linear mappings"))

"""
    ncols(A)

yields the *equivalent* number of columns of the linear operator `A`.  Not all
operators extend this method.

In the implemented generalization of linear operators, the equivalent number of
columns is the number of element of an argument of the operator be it single-
or multi-dimensional.

"""
ncols(A::LinearMapping) = prod(col_size(A))
@noinline ncols(A::Mapping) =
    throw(ArgumentError("`ncols` is only implemented for linear mappings"))

"""
    row_size(A)

yields the dimensions of the result of applying the linear operator `A`, this
is equivalent to `output_size(A)`.  Not all operators extend this method.

"""
row_size(A::LinearMapping) = output_size(A)
@noinline row_size(A::Mapping) =
    throw(ArgumentError("`row_size` is only implemented for linear mappings"))

"""
    col_size(A)

yields the dimensions of the argument of the linear operator `A`, this
is equivalent to `input_size(A)`.  Not all
operators extend this method.

"""
col_size(A::LinearMapping) = input_size(A)
@noinline col_size(A::Mapping) =
    throw(ArgumentError("`col_size` is only implemented for linear mappings"))

"""
    coefficients(A)

yields the object backing the storage of the coefficients of the linear mapping
`A`.  Not all linear mappings extend this method.

""" coefficients

"""
    check(A) -> A

checks integrity of mapping `A` and returns it.

"""
check(A::Mapping) = A

"""
    checkmapping(y, A, x) -> (v1, v2, v1 - v2)

yields `v1 = vdot(y, A*x)`, `v2 = vdot(A'*y, x)` and their difference for `A` a
linear mapping, `y` a *vector* of the output space of `A` and `x` a *vector* of
the input space of `A`.  In principle, the two inner products should be equal
whatever `x` and `y`; otherwise the mapping has a bug.

Simple linear mappings operating on Julia arrays can be tested on random
*vectors* with:

    checkmapping([T=Float64,] outdims, A, inpdims) -> (v1, v2, v1 - v2)

with `outdims` and `outdims` the dimensions of the output and input *vectors*
for `A`.  Optional argument `T` is the element type.

If `A` operates on Julia arrays and methods `input_eltype`, `input_size`,
`output_eltype` and `output_size` have been specialized for `A`, then:

    checkmapping(A) -> (v1, v2, v1 - v2)

is sufficient to check `A` against automatically generated random arrays.

See also: [`vdot`](@ref), [`vcreate`](@ref), [`apply!`](@ref),
[`input_type`](@ref).

"""
function checkmapping(y::Ty, A::Mapping, x::Tx) where {Tx, Ty}
    is_linear(A) || bad_argument("expecting a linear mapping")
    v1 = vdot(y, A*x)
    v2 = vdot(A'*y, x)
    (v1, v2, v1 - v2)
end

function checkmapping(::Type{T},
                      outdims::Tuple{Vararg{Int}},
                      A::Mapping,
                      inpdims::Tuple{Vararg{Int}}) where {T<:AbstractFloat}
    checkmapping(randn(T, outdims), A, randn(T, inpdims))
end

function checkmapping(outdims::Tuple{Vararg{Int}},
                      A::Mapping,
                      inpdims::Tuple{Vararg{Int}})
    checkmapping(Float64, outdims, A, inpdims)
end

checkmapping(A::LinearMapping) =
    checkmapping(randn(output_eltype(A), output_size(A)), A,
                 randn(input_eltype(A), input_size(A)))

"""
    identical(A, B)

yields whether `A` is the same mapping as `B` in the sense that their effects
will always be the same.  This method is used to perform some simplifications
and optimizations and may have to be specialized for specific mapping types.
The default implementation is to return `A === B`.

!!! note
    The returned result may be true although `A` and `B` are not necessarily
    the same objects.  For instance, if `A` and `B` are two sparse matrices
    whose coefficients and indices are stored in the same arrays (as can be
    tested with the `===` or `≡` operators, `identical(A,B)` should return
    `true` because the two operators will always behave identically (any
    changes in the coefficients or indices of `A` will be reflected in `B`).
    If any of the arrays storing the coefficients or the indices are not the
    same objects, then `identical(A,B)` must return `false` even though the
    stored values may be the same because it is possible, later, to change one
    operator without affecting identically the other.

"""
@inline identical(::Mapping, ::Mapping) = false # false if not same types
@inline identical(A::T, B::T) where {T<:Mapping} = (A === B)

"""
    gram(A) -> A'*A

yields the Gram operator built out of the linear mapping `A`.  The result is
equivalent to `A'*A` but its type depends on simplifications that may occur.

See also [`Gram`](@ref).

"""
gram(A::LinearMapping) = A'*A
gram(A::Mapping) =
    is_linear(A) ? A'*A : throw_forbidden_Gram_of_non_linear_mapping()

@noinline throw_forbidden_Gram_of_non_linear_mapping() =
    bad_argument("making a Gram operator out of a non-linear mapping is not allowed")

# Inlined functions called to perform `α*x + β*y` for specific values of the
# multipliers `α` and `β`.  Passing these (simple) functions to another method
# is to simplify the coding of vectorized methods and of the the `apply!`
# method by mappings.  NOTE: Forcing inlining may not be necessary but it does
# not hurt.
@inline axpby_yields_zero( α, x, β, y) = zero(typeof(y)) # α = 0, β = 0
@inline axpby_yields_y(    α, x, β, y) = y               # α = 0, β = 1
@inline axpby_yields_my(   α, x, β, y) = -y              # α = 0, β = -1
@inline axpby_yields_by(   α, x, β, y) = β*y             # α = 0, any β
@inline axpby_yields_x(    α, x, β, y) = x               # α = 1, β = 0
@inline axpby_yields_xpy(  α, x, β, y) = x + y           # α = 1, β = 1
@inline axpby_yields_xmy(  α, x, β, y) = x - y           # α = 1, β = -1
@inline axpby_yields_xpby( α, x, β, y) = x + β*y         # α = 1, any β
@inline axpby_yields_mx(   α, x, β, y) = -x              # α = -1, β = 0
@inline axpby_yields_ymx(  α, x, β, y) = y - x           # α = -1, β = 1
@inline axpby_yields_mxmy( α, x, β, y) = -x - y          # α = -1, β = -1
@inline axpby_yields_bymx( α, x, β, y) = β*y - x         # α = -1, any β
@inline axpby_yields_ax(   α, x, β, y) = α*x             # any α, β = 0
@inline axpby_yields_axpy( α, x, β, y) = α*x + y         # any α, β = 1
@inline axpby_yields_axmy( α, x, β, y) = α*x - y         # any α, β = -1
@inline axpby_yields_axpby(α, x, β, y) = α*x + β*y       # any α, any β

#------------------------------------------------------------------------------
# VCREATE, APPLY AND APPLY!

"""
```julia
vcreate([P,] A, x, scratch=false) -> y
```

yields a new instance `y` suitable for storing the result of applying mapping
`A` to the argument `x`.  Optional parameter `P ∈ Operations` is one of
`Direct` (the default), `Adjoint`, `Inverse` and/or `InverseAdjoint` and can be
used to specify how `A` is to be applied as explained in the documentation of
the [`apply`](@ref) method.

Optional argument `scratch` indicates whether input argument `x` can be
overwritten by the operation and thus used to store the result.  This may be
exploited by some mappings (which are able to operate *in-place*) to avoid
allocating a new object for the result `y`.

The caller should set `scratch = true` if `x` is not needed after calling
`apply`.  If `scratch = true`, then it is possible that `y` be the same object
as `x`; otherwise, `y` is a new object unless applying the operation yields the
same contents as `y` for the result `x` (this is always true for the identity
for instance).  Thus, in general, it should not be assumed that the returned
`y` is different from the input `x`.

The method `vcreate(::Type{P}, A, x)` should be implemented by linear mappings
for any supported operations `P` and argument type for `x`.  The result
returned by `vcreate` should be of predictible type to ensure *type-stability*.
Checking the validity (*e.g.* the size) of argument `x` in `vcreate` may be
skipped because this argument will be eventually checked by the `apply!`
method.

See also: [`Mapping`](@ref), [`apply`](@ref).

"""
vcreate(A::Mapping, x, scratch::Bool=false) = vcreate(Direct, A, x, scratch)
vcreate(::Type{P}, A::Mapping, x) where {P<:Operations} =
    vcreate(P, A, x, false)

"""
    vmul(A, x) -> y

yields `y = A*x`.  The default behavior is to call `apply(Direct,A,x,false)`.
Method [`vmul!`](@ref) is the in-place version.

"""
vmul(A, x) = apply(Direct, A, x, false)

"""
    vmul!(y, A, x) -> y

overwrites `y` with the result of `A*x` and returns `y`.  The default behavior
is to call `apply!(1,Direct,A,x,false,0,y)`.

!!! note

    This method is intended to be used by algorithms such as the conjugate
    gradient to apply operators.  It may be specialized by the caller for its
    needs which is much easier than specializing [`apply!`](@ref) which
    requires to consider the specific values of the multipliers `α` and `β`.

"""
vmul!(y, A, x) = apply!(1, Direct, A, x, false, 0, y)

"""
```julia
apply([P,] A, x, scratch=false) -> y
```

yields the result `y` of applying mapping `A` to the argument `x`.
Optional parameter `P` can be used to specify how `A` is to be applied:

* `Direct` (the default) to apply `A` and yield `y = A⋅x`;
* `Adjoint` to apply the adjoint of `A` and yield `y = A'⋅x`;
* `Inverse` to apply the inverse of `A` and yield `y = A\\x`;
* `InverseAdjoint` or `AdjointInverse` to apply the inverse of `A'` and
  yield `y = A'\\x`.

Not all operations may be implemented by the different types of mappings and
`Adjoint` and `InverseAdjoint` may only be applicable for linear mappings.

Optional argument `scratch` indicates whether input argument `x` can be
overwritten by the operation.  This may be exploited to avoid allocating
temporary workspace(s).  The caller should set `scratch = true` if `x` is not
needed after calling `apply`.  If `scratch = true`, then it is possible that
`y` be the same object as `x`; otherwise, `y` is a new object unless applying
the operation yields the same contents as `y` for the result `x` (this is
always true for the identity for instance).  Thus, in general, it should not be
assumed that the result of applying a mapping is different from the input.

Julia methods are provided so that `apply(A', x)` automatically calls
`apply(Adjoint, A, x)` so the shorter syntax may be used without impacting
performances.

See also: [`Mapping`](@ref), [`apply!`](@ref), [`vcreate`](@ref).

"""
apply(A::Mapping, x, scratch::Bool=false) = apply(Direct, A, x, scratch)
apply(P::Type{<:Operations}, A::Mapping, x, scratch::Bool=false) =
    apply!(1, P, A, x, scratch, 0, vcreate(P, A, x, scratch))

*(A::Mapping, x) = apply(Direct, A, x, false)
\(A::Mapping, x) = apply(Inverse, A, x, false)

"""
    apply!([α=1,] [P=Direct,] A::Mapping, x, [scratch=false,] [β=0,] y) -> y

overwrites `y` with `α*P(A)⋅x + β*y` where `P ∈ Operations` can be `Direct`,
`Adjoint`, `Inverse` and/or `InverseAdjoint` to indicate which variant of the
mapping `A` to apply.  The convention is that the prior contents of `y` is not
used at all if `β = 0` so `y` can be directly used to store the result even
though it is not initialized.  The `scratch` optional argument indicates
whether the input `x` is no longer needed by the caller and can thus be used as
a scratch array.  Having `scratch = true` or `β = 0` may be exploited by the
specific implementation of the `apply!` method for the mapping type to avoid
allocating temporary workspace(s).

The `apply!` method can be seen as a generalization of the `LinearAlgebra.mul!`
method.

The order of arguments can be changed and the same result as above is obtained
with:

    apply!([β=0,] y, [α=1,] [P=Direct,] A::Mapping, x, scratch=false) -> y

The result `y` may have been allocated by:

    y = vcreate([P=Direct,] A, x, scratch=false)

Mapping sub-types only need to extend `vcreate` and `apply!` with the specific
signatures:

    vcreate(::Type{P}, A::M, x, scratch::Bool=false) -> y
    apply!(α::Number, ::Type{P}, A::M, x, scratch::Bool, β::Number, y) -> y

for any supported operation `P` and where `M` is the type of the mapping.  Of
course, the types of arguments `x` and `y` may be specified as well.

Optionally, the method with signature:

    apply(::Type{P}, A::M, x, scratch::Bool=false) -> y

may also be extended to improve the default implementation which is:

    apply(P::Type{<:Operations}, A::Mapping, x, scratch::Bool=false) =
        apply!(1, P, A, x, scratch, 0, vcreate(P, A, x, scratch))

See also: [`Mapping`](@ref), [`apply`](@ref), [`vcreate`](@ref).

""" apply!

# Provide fallbacks so that `Direct` is the default operation and only the
# method with signature:
#
#     apply!(α::Number, ::Type{P}, A::MappingType, x::X, scratch::Bool,
#            β::Number, y::Y) where {P<:Operations,X,Y}
#
# has to be implemented (possibly with restrictions on X and Y) by subtypes of
# Mapping so we provide the necessary mechanism to dispatch derived methods.
apply!(A::Mapping, x, y) =
    apply!(1, Direct, A, x, false, 0, y)
apply!(α::Number, A::Mapping, x, y) =
    apply!(α, Direct, A, x, false, 0, y)
apply!(A::Mapping, x, β::Number, y) =
    apply!(1, Direct, A, x, false, β, y)
apply!(α::Number, A::Mapping, x, β::Number, y) =
    apply!(α, Direct, A, x, false, β, y)

apply!(P::Type{<:Operations}, A::Mapping, x, y) =
    apply!(1, P, A, x, false, 0, y)
apply!(α::Number, P::Type{<:Operations}, A::Mapping, x, y) =
    apply!(α, P, A, x, false, 0, y)
apply!(P::Type{<:Operations}, A::Mapping, x, β::Number, y) =
    apply!(1, P, A, x, false, β, y)
apply!(α::Number, P::Type{<:Operations}, A::Mapping, x, β::Number, y) =
    apply!(α, P, A, x, false, β, y)

apply!(A::Mapping, x, scratch::Bool, y) =
    apply!(1, Direct, A, x, scratch, 0, y)
apply!(α::Number, A::Mapping, x, scratch::Bool, y) =
    apply!(α, Direct, A, x, scratch, 0, y)
apply!(A::Mapping, x, scratch::Bool, β::Number, y) =
    apply!(1, Direct, A, x, scratch, β, y)

apply!(P::Type{<:Operations}, A::Mapping, x, scratch::Bool, y) =
    apply!(1, P, A, x, scratch, 0, y)
apply!(α::Number, P::Type{<:Operations}, A::Mapping, x, scratch::Bool, y) =
    apply!(α, P, A, x, scratch, 0, y)
apply!(P::Type{<:Operations}, A::Mapping, x, scratch::Bool, β::Number, y) =
    apply!(1, P, A, x, scratch, β, y)

# Change order of arguments.
apply!(y, A::Mapping, x, scratch::Bool=false) =
    apply!(1, Direct, A, x, scratch, 0, y)
apply!(y, P::Type{<:Operations}, A::Mapping, x, scratch::Bool=false) =
    apply!(1, P, A, x, scratch, 0, y)
apply!(y, α::Number, A::Mapping, x, scratch::Bool=false) =
    apply!(α, Direct, A, x, scratch, 0, y)
apply!(y, α::Number, P::Type{<:Operations}, A::Mapping, x, scratch::Bool=false) =
    apply!(α, P, A, x, scratch, 0, y)
apply!(β::Number, y, A::Mapping, x, scratch::Bool=false) =
    apply!(1, Direct, A, x, scratch, β, y)
apply!(β::Number, y, P::Type{<:Operations}, A::Mapping, x, scratch::Bool=false) =
    apply!(1, P, A, x, scratch, β, y)
apply!(β::Number, y, α::Number, A::Mapping, x, scratch::Bool=false) =
    apply!(α, Direct, A, x, scratch, β, y)
apply!(β::Number, y, α::Number, P::Type{<:Operations}, A::Mapping, x, scratch::Bool=false) =
    apply!(α, P, A, x, scratch, β, y)

# Extend `LinearAlgebra.mul!` so that `A'*x`, `A*B*C*x`, etc. yield the
# expected result.  FIXME: This should be restricted to linear mappings but
# this is not possible without overheads.
mul!(y, A::Mapping, x) = apply!(1, Direct, A, x, false, 0, y)
mul!(y, A::Mapping, x, α::Number, β::Number) =
    apply!(α, Direct, A, x, false, β, y)

# Implemention of the `apply!(α,P,A,x,scratch,β,y)` and
# `vcreate(P,A,x,scratch)` methods for a scaled mapping.
for (P, expr) in ((:Direct, :(α*multiplier(A))),
                  (:Adjoint, :(α*conj(multiplier(A)))),
                  (:Inverse, :(α/multiplier(A))),
                  (:InverseAdjoint, :(α/conj(multiplier(A)))))
    @eval begin

        apply!(α::Number, ::Type{$P}, A::Scaled, x, scratch::Bool, β::Number, y) =
            apply!($expr, $P, unscaled(A), x, scratch, β, y)

    end
end

"""
    overwritable(scratch, x, y) -> bool

yields whether the result `y` of applying a mapping to `x` with scratch flag
`scratch` can overwritten.  Arguments `x` and `y` can be reversed.

"""
overwritable(scratch::Bool, x, y) = (scratch || x !== y)

# Implement `apply` for scaled operators to avoid the needs of explicitly
# calling `vcreate` as done by the default implementation of `apply`.  This is
# needed for scaled compositions among others.
function apply(::Type{Direct}, A::Scaled, x, scratch::Bool)
    y = apply(Direct, unscaled(A), x, scratch)
    vscale!((overwritable(scratch, x, y) ? y : vcopy(y)), multiplier(A))
end

function apply(::Type{Adjoint}, A::Scaled, x, scratch::Bool)
    y = apply(Direct, unscaled(A), x, scratch)
    vscale!((overwritable(scratch, x, y) ? y : vcopy(y)), conj(multiplier(A)))
end

function apply(::Type{Inverse}, A::Scaled, x, scratch::Bool)
    y = apply(Direct, unscaled(A), x, scratch)
    vscale!((overwritable(scratch, x, y) ? y : vcopy(y)), 1/multiplier(A))
end

function apply(::Type{InverseAdjoint}, A::Scaled, x, scratch::Bool)
    y = apply(Direct, unscaled(A), x, scratch)
    vscale!((overwritable(scratch, x, y) ? y : vcopy(y)), 1/conj(multiplier(A)))
end

vcreate(P::Type{<:Operations}, A::Scaled, x, scratch::Bool) =
    vcreate(P, unscaled(A), x, scratch)

# Implemention of the `vcreate(P,A,x,scratch)` and
# `apply!(α,P,A,x,scratch,β,y)` methods for the various decorations of a
# mapping so as to automatically unveil the embedded mapping.
for (T1, T2, T3) in ((:Direct,         :Adjoint,        :Adjoint),
                     (:Adjoint,        :Adjoint,        :Direct),
                     (:Inverse,        :Adjoint,        :InverseAdjoint),
                     (:InverseAdjoint, :Adjoint,        :Inverse),
                     (:Direct,         :Inverse,        :Inverse),
                     (:Adjoint,        :Inverse,        :InverseAdjoint),
                     (:Inverse,        :Inverse,        :Direct),
                     (:InverseAdjoint, :Inverse,        :Adjoint),
                     (:Direct,         :InverseAdjoint, :InverseAdjoint),
                     (:Adjoint,        :InverseAdjoint, :Inverse),
                     (:Inverse,        :InverseAdjoint, :Adjoint),
                     (:InverseAdjoint, :InverseAdjoint, :Direct))
    @eval begin

        vcreate(::Type{$T1}, A::$T2, x, scratch::Bool) =
            vcreate($T3, unveil(A), x, scratch)

        apply!(α::Number, ::Type{$T1}, A::$T2, x, scratch::Bool, β::Number, y) =
            apply!(α, $T3, unveil(A), x, scratch, β, y)

    end
end

# Implementation of the `vcreate(P,A,x,scratch)` and
# `apply!(α,P,A,x,scratch,β,y)` and methods for a sum of mappings.  Note that
# `Sum` instances are warranted to have at least 2 components.

function vcreate(::Type{P}, A::Sum, x,
                 scratch::Bool) where {P<:Union{Direct,Adjoint}}
    # The sum only makes sense if all mappings yields the same kind of result.
    # Hence we just call the vcreate method for the first mapping of the sum.
    vcreate(P, A[1], x, scratch)
end

function apply!(α::Number, P::Type{<:Union{Direct,Adjoint}}, A::Sum{N},
                x, scratch::Bool, β::Number, y) where {N}
    if α == 0
        # Just scale the destination.
        vscale!(y, β)
    else
        # Apply first mapping with β and then other with β=1.  Scratch flag is
        # always false until last mapping because we must preserve x as there
        # is more than one term.
        apply!(α, P, A[1], x, false, β, y)
        for i in 2:N
            apply!(α, P, A[i], x, (scratch && i == N), 1, y)
        end
    end
    return y
end

vcreate(::Type{<:Union{Inverse,InverseAdjoint}}, A::Sum, x, scratch::Bool) =
    throw_unsupported_inverse_of_sum()

apply(::Type{<:Union{Inverse,InverseAdjoint}}, A::Sum, x, scratch::Bool) =
    throw_unsupported_inverse_of_sum()

function apply!(α::Number, ::Type{<:Union{Inverse,InverseAdjoint}}, A::Sum,
                x, scratch::Bool, β::Number, y)
    throw_unsupported_inverse_of_sum()
end

throw_unsupported_inverse_of_sum() =
    error("automatic dispatching of the inverse of a sum of mappings is not supported")

# Implementation of the `apply!(α,P,A,x,scratch,β,y)` method for a composition
# of mappings.  There is no possible `vcreate(P,A,x,scratch)` method for a
# composition so we directly extend the `apply(P,A,x,scratch)` method.  Note
# that `Composition` instances are warranted to have at least 2 components.
#
# The unrolled code (taking care of allowing as few temporaries as possible and
# for the Direct or InverseAdjoint operation) writes:
#
#     w1 = apply(P, A[N], x, scratch)
#     scratch = overwritable(scratch, x, w1)
#     w2 = apply!(1, P, A[N-1], w1, scratch)
#     scratch = overwritable(scratch, w1, w2)
#     w3 = apply!(1, P, A[N-2], w2, scratch)
#     scratch = overwritable(scratch, w2, w3)
#     ...
#     return apply!(α, P, A[1], wNm1, scratch, β, y)
#
# To break the type barrier, this is done by a recursion.  The recursion is
# just done in the other direction for the Adjoint or Inverse operation.

function vcreate(::Type{<:Operations},
                 A::Composition{N}, x, scratch::Bool) where {N}
    error("it is not possible to create the output of a composition of mappings")
end

function apply!(α::Number, ::Type{P}, A::Composition{N}, x, scratch::Bool,
                β::Number, y) where {N,P<:Union{Direct,InverseAdjoint}}
    if α == 0
        # Just scale the destination.
        vscale!(y, β)
    else
        ops = terms(A)
        w = apply(P, *, ops[2:N], x, scratch)
        scratch = overwritable(scratch, w, x)
        apply!(α, P, ops[1], w, scratch, β, y)
    end
    return y
end

function apply(::Type{P}, A::Composition{N}, x,
               scratch::Bool) where {N,P<:Union{Direct,InverseAdjoint}}
    apply(P, *, terms(A), x, scratch)
end

function apply(::Type{P}, ::typeof(*), ops::NTuple{N,Mapping}, x,
               scratch::Bool) where {N,P<:Union{Direct,InverseAdjoint}}
    w = apply(P, ops[N], x, scratch)
    N == 1 && return w
    scratch = overwritable(scratch, w, x)
    apply(P, *, ops[1:N-1], w, scratch)
end

function apply!(α::Number, ::Type{P}, A::Composition{N}, x, scratch::Bool,
                β::Number, y) where {N,P<:Union{Adjoint,Inverse}}
    if α == 0
        # Just scale the destination.
        vscale!(y, β)
    else
        ops = terms(A)
        w = apply(P, *, ops[1:N-1], x, scratch)
        scratch = overwritable(scratch, w, x)
        apply!(α, P, ops[N], w, scratch, β, y)
    end
    return y
end

function apply(::Type{P}, A::Composition{N}, x,
               scratch::Bool) where {N,P<:Union{Adjoint,Inverse}}
    apply(P, *, terms(A), x, scratch)
end

function apply(::Type{P}, ::typeof(*), ops::NTuple{N,Mapping}, x,
               scratch::Bool) where {N,P<:Union{Adjoint,Inverse}}
    w = apply(P, ops[1], x, scratch)
    N == 1 && return w
    scratch = overwritable(scratch, w, x)
    apply(P, *, ops[2:N], w, scratch)
end

# Default rules to apply a Gram operator.  Gram matrices are Hermitian by
# construction which left only 2 cases to deal with.

apply!(α::Number, ::Type{Adjoint}, A::Gram, x, scratch::Bool, β::Number, y) =
    apply!(α, Direct, A, x, scratch, β, y)

apply!(α::Number, ::Type{InverseAdjoint}, A::Gram, x, scratch::Bool, β::Number, y) =
    apply!(α, Inverse, A, x, scratch, β, y)

function apply!(α::Number, ::Type{Direct}, A::Gram, x, scratch::Bool, β::Number, y)
    if α == 0
        vscale!(y, β)
    else
        B = unveil(A) # A ≡ B'*B
        z = apply(Direct, B, x, scratch) # z <- B⋅x
        apply!(α, Adjoint, B, z, (z !== x), β, y) # y <- α⋅B'⋅z + β⋅y
    end
    return y
end

function apply!(α::Number, ::Type{Inverse}, A::Gram, x, scratch::Bool, β::Number, y)
    if α == 0
        vscale!(y, β)
    else
        B = unveil(A) # A ≡ B'⋅B
        # Compute α⋅inv(A)⋅x + β⋅y = α⋅inv(B'⋅B)⋅x + β⋅y
        #                          = α⋅inv(B)⋅inv(B')⋅x + β⋅y
        z = apply(InverseAdjoint, B, x, scratch) # z <- inv(B')⋅x
        apply!(α, Inverse, B, z, (z !== x), β, y) # y <- α⋅inv(B)⋅z + β⋅y
    end
    return y
end

# A Gram operator is self-adjoint by construction and yields result of same
# kind as input.
vcreate(::Type{<:Operations}, ::Gram, x, scratch::Bool) = vcreate(x)


text:
module TimeSteppers

export
    AdamsBashforthTimeStepper,
    time_step!,
    compute_w_from_continuity!

using GPUifyLoops: @launch, @loop, @unroll

import Oceananigans: TimeStepper

using Oceananigans.Architectures: @hascuda
@hascuda using CUDAnative, CuArrays

using Oceananigans.Architectures
using Oceananigans.Grids
using Oceananigans.Fields
using Oceananigans.Operators
using Oceananigans.Coriolis
using Oceananigans.Buoyancy
using Oceananigans.SurfaceWaves
using Oceananigans.BoundaryConditions
using Oceananigans.Solvers
using Oceananigans.Models
using Oceananigans.Utils

using Oceananigans.TurbulenceClosures:
    calculate_diffusivities!, ∂ⱼ_2ν_Σ₁ⱼ, ∂ⱼ_2ν_Σ₂ⱼ, ∂ⱼ_2ν_Σ₃ⱼ, ∇_κ_∇c

"""
    AbstractTimeStepper

Abstract supertype for time steppers.
"""
abstract type AbstractTimeStepper end

"""
    TimeStepper(name, args...)

Returns a timestepper with name `name`, instantiated with `args...`.

Example
=======

julia> stepper = TimeStepper(:AdamsBashforth, Float64, CPU(), grid, tracernames)
"""
function TimeStepper(name::Symbol, args...)
    fullname = Symbol(name, :TimeStepper)
    return eval(Expr(:call, fullname, args...))
end

# Fallback
TimeStepper(stepper, args...) = stepper

"""Returns the arguments passed to boundary conditions functions."""
boundary_condition_function_arguments(model) =
    (model.clock.time, model.clock.iteration, datatuple(model.velocities),
     datatuple(model.tracers), model.parameters)

#####
##### Time-stepping kernels/function that are independent of the TimeStepper
#####

"""
    calculate_explicit_substep!(tendencies, velocities, tracers, pressures, diffusivities, model)

Calculate the initial and explicit substep of the two-step fractional step method with pressure correction.
"""
function calculate_explicit_substep!(tendencies, velocities, tracers, pressures, diffusivities, model)
    time_step_precomputations!(diffusivities, pressures, velocities, tracers, model)
    calculate_tendencies!(tendencies, velocities, tracers, pressures, diffusivities, model)
    return nothing
end

"""
    time_step_precomputations!(diffusivities, pressures, velocities, tracers, model)

Perform precomputations necessary for an explicit timestep or substep.
"""
function time_step_precomputations!(diffusivities, pressures, velocities, tracers, model)

    fill_halo_regions!(merge(model.velocities, model.tracers), model.architecture,
                       boundary_condition_function_arguments(model)...)

    calculate_diffusivities!(diffusivities, model.architecture, model.grid, model.closure,
                             model.buoyancy, velocities, tracers)

    fill_halo_regions!(model.diffusivities, model.architecture)

    @launch(device(model.architecture), config=launch_config(model.grid, :xy),
            update_hydrostatic_pressure!(pressures.pHY′, model.grid, model.buoyancy, tracers))

    fill_halo_regions!(model.pressures.pHY′, model.architecture)

    return nothing
end

"""
    calculate_tendencies!(diffusivities, pressures, velocities, tracers, model)

Calculate the interior and boundary contributions to tendency terms without the
contribution from non-hydrostatic pressure.
"""
function calculate_tendencies!(tendencies, velocities, tracers, pressures, diffusivities, model)

    calculate_interior_source_terms!(
        tendencies, model.architecture, model.grid, model.coriolis, model.buoyancy,
        model.surface_waves, model.closure, velocities, tracers, pressures.pHY′,
        diffusivities, model.forcing, model.parameters, model.clock.time
    )

    calculate_boundary_source_terms!(
        model.timestepper.Gⁿ, model.architecture, model.velocities,
        model.tracers, boundary_condition_function_arguments(model)...
    )

    return nothing
end

"""
    calculate_pressure_correction!(nonhydrostatic_pressure, Δt, tendencies, velocities, model)

Calculate the (nonhydrostatic) pressure correction associated `tendencies`, `velocities`, and step size `Δt`.
"""
function calculate_pressure_correction!(nonhydrostatic_pressure, Δt, tendencies, velocities, model)
    velocity_tendencies = (u=model.timestepper.Gⁿ.u, v=model.timestepper.Gⁿ.v, w=model.timestepper.Gⁿ.w)

    fill_halo_regions!(velocity_tendencies, model.architecture)

    solve_for_pressure!(nonhydrostatic_pressure, model.pressure_solver,
                        model.architecture, model.grid, velocities, tendencies, Δt)

    fill_halo_regions!(model.pressures.pNHS, model.architecture)

    return nothing
end

calculate_pressure_correction!(::Nothing, args...) = nothing

"""
    complete_pressure_correction_step!(velocities, Δt, tracers, pressures, tendencies, model)

After calculating the pressure correction, complete the pressure correction step by updating
the velocity and tracer fields.
"""
function complete_pressure_correction_step!(velocities, Δt, tracers, pressures, tendencies, model)

    update_solution!(velocities, tracers, model.architecture,
                     model.grid, Δt, tendencies, pressures.pNHS)

    fill_halo_regions!(model.velocities, model.architecture,
                       boundary_condition_function_arguments(model)...)

    compute_w_from_continuity!(model)

    return nothing
end

include("kernels.jl")
include("adams_bashforth.jl")

end # module

ext: jl


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 = "CryptoMiniSat"
version = v"5.8.0"

# Collection of sources required to complete build
sources = [
    GitSource("https://github.com/msoos/cryptominisat.git", "e7079937ed2bfe9160a104378e5a344028e4ab78"),
    DirectorySource("./bundled")
]

# Bash recipe for building across all platforms
script = raw"""

cd $WORKSPACE/srcdir/cryptominisat
atomic_patch -p1 ../Yalsatpatch.patch
atomic_patch -p1 ../feenablepatch.patch
mkdir build && cd build
cmake -DCMAKE_INSTALL_PREFIX=$prefix \
    -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
    -DCMAKE_BUILD_TYPE=Release \
    -DENABLE_PYTHON_INTERFACE=OFF \
    -DIPASIR=ON \
    -DNOM4RI=ON \
    -DBoost_USE_STATIC_LIBS=OFF \
    ..
make -j${nproc}
make install
install_license ${WORKSPACE}/srcdir/cryptominisat/LICENSE.txt
"""

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

# The products that we will ensure are always built
products = Product[
    ExecutableProduct(["cryptominisat5", "cryptominisat5win"], :cryptominisat5),
    ExecutableProduct("cryptominisat5_simple", :cryptominisat5_simple),
    LibraryProduct(["libcryptominisat5", "libcryptominisat5win"], :libcryptominisat5),
    LibraryProduct("libipasircryptominisat5", :libipasircryptominisat5)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency("boost_jll"),
    Dependency("Zlib_jll"),
    Dependency("SQLite_jll"),
    Dependency("MPICH_jll")
]

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

max_stars_repo_name: waralex/Yggdrasil

max_issues_repo_name: waralex/Yggdrasil

max_forks_repo_name: waralex/Yggdrasil


text:
# output from JPL solution #197 SPK file for Apophis via SPICE.jl
const x0_JPL_s197 = [-0.9633018953468989, 0.5100289806011301, 0.16528024397505386,
    -0.0071188720570829036, -0.012061235689040535, -0.0046695143453363164]
# output from JPL solution #199 SPK file for Apophis via SPICE.jl
const x0_JPL_s199 = [-0.96330181481543, 0.5100291398744977, 0.16528030015848072,
    -0.007118874645519018, -0.0120612341605072, -0.004669513801491289]

function initialcond(dq::Vector=zeros(7); quadmath::Bool=false)
    if quadmath
        q0 = Vector{Float128}(undef, 7) #initial condition array
    else
        q0 = Vector{Float64}(undef, 7) #initial condition array
    end
    # q0 .= x0_JPL_s199
    q0[1:6] .= x0_JPL_s197
    q0[7] = 0.0 # A2 Yarkovsky coefficient
    S = typeof(q0[1]+dq[1])
    q = Vector{S}(undef, 7)
    for i in 1:6
        q[i] = q0[i] + dq[i]
    end
    if length(dq) == 7
        q[7] = q0[7] + dq[7]
    elseif length(dq) == 6
        q[7] = q0[7] + 0dq[1]
    end
    return q
end


text:
macro varinfo()
    :(throw(_error_msg()))
end
macro logpdf()
    :(throw(_error_msg()))
end
macro sampler()
    :(throw(_error_msg()))
end
function _error_msg()
    return "This macro is only for use in the `@model` macro and not for external use."
end

# Check if the right-hand side is a distribution.
function assert_dist(dist; msg)
    isa(dist, Distribution) || throw(ArgumentError(msg))
end
function assert_dist(dist::AbstractVector; msg)
    all(d -> isa(d, Distribution), dist) || throw(ArgumentError(msg))
end

function wrong_dist_errormsg(l)
    return "Right-hand side of a ~ must be subtype of Distribution or a vector of " *
        "Distributions on line $(l)."
end

"""
    @isassumption(model, expr)

Let `expr` be `x[1]`. `vn` is an assumption in the following cases:
    1. `x` was not among the input data to the model,
    2. `x` was among the input data to the model but with a value `missing`, or
    3. `x` was among the input data to the model with a value other than missing,
       but `x[1] === missing`.
When `expr` is not an expression or symbol (i.e., a literal), this expands to `false`.
"""
macro isassumption(model, expr::Union{Symbol, Expr})
    # Note: never put a return in this... don't forget it's a macro!
    vn = gensym(:vn)

    return quote
        $vn = @varname($expr)

        # This branch should compile nicely in all cases except for partial missing data
        # For example, when `expr` is `x[i]` and `x isa Vector{Union{Missing, Float64}}`
        if !DynamicPPL.inargnames($vn, $model) || DynamicPPL.inmissings($vn, $model)
            true
        else
            if DynamicPPL.inargnames($vn, $model)
                # Evaluate the lhs
                $expr === missing
            else
                throw("This point should not be reached. Please report this error.")
            end
        end
    end |> esc
end

macro isassumption(model, expr)
    # failsafe: a literal is never an assumption
    false
end

#################
# Main Compiler #
#################

"""
    @model(body)

Macro to specify a probabilistic model.

Example:

Model definition:

```julia
@model model_generator(x = default_x, y) = begin
    ...
end
```

To generate a `Model`, call `model_generator(x_value)`.
"""
macro model(input_expr)
    build_model_info(input_expr) |> replace_tilde! |> replace_vi! |>
        replace_logpdf! |> replace_sampler! |> build_output
end

"""
    build_model_info(input_expr)

Builds the `model_info` dictionary from the model's expression.
"""
function build_model_info(input_expr)
    # Extract model name (:name), arguments (:args), (:kwargs) and definition (:body)
    modeldef = MacroTools.splitdef(input_expr)
    # Function body of the model is empty
    warn_empty(modeldef[:body])
    # Construct model_info dictionary

    # Extracting the argument symbols from the model definition
    arg_syms = map(modeldef[:args]) do arg
        # @model demo(x)
        if (arg isa Symbol)
            arg
        # @model demo(::Type{T}) where {T}
        elseif MacroTools.@capture(arg, ::Type{T_} = Tval_)
            T
        # @model demo(x::T = 1)
        elseif MacroTools.@capture(arg, x_::T_ = val_)
            x
        # @model demo(x = 1)
        elseif MacroTools.@capture(arg, x_ = val_)
            x
        else
            throw(ArgumentError("Unsupported argument $arg to the `@model` macro."))
        end
    end
    if length(arg_syms) == 0
        args_nt = :(NamedTuple())
    else
        nt_type = Expr(:curly, :NamedTuple,
            Expr(:tuple, QuoteNode.(arg_syms)...),
            Expr(:curly, :Tuple, [:(Core.Typeof($x)) for x in arg_syms]...)
        )
        args_nt = Expr(:call, :(DynamicPPL.namedtuple), nt_type, Expr(:tuple, arg_syms...))
    end
    args = map(modeldef[:args]) do arg
        if (arg isa Symbol)
            arg
        elseif MacroTools.@capture(arg, ::Type{T_} = Tval_)
            if in(T, modeldef[:whereparams])
                S = :Any
            else
                ind = findfirst(modeldef[:whereparams]) do x
                    MacroTools.@capture(x, T1_ <: S_) && T1 == T
                end
                ind !== nothing || throw(ArgumentError("Please make sure type parameters are properly used. Every `Type{T}` argument need to have `T` in the a `where` clause"))
            end
            Expr(:kw, :($T::Type{<:$S}), Tval)
        else
            arg
        end
    end
    args_nt = to_namedtuple_expr(arg_syms)

    default_syms = []
    default_vals = []
    foreach(modeldef[:args]) do arg
        # @model demo(::Type{T}) where {T}
        if MacroTools.@capture(arg, ::Type{T_} = Tval_)
            push!(default_syms, T)
            push!(default_vals, Tval)
        # @model demo(x::T = 1)
        elseif MacroTools.@capture(arg, x_::T_ = val_)
            push!(default_syms, x)
            push!(default_vals, val)
        # @model demo(x = 1)
        elseif MacroTools.@capture(arg, x_ = val_)
            push!(default_syms, x)
            push!(default_vals, val)
        end
    end
    defaults_nt = to_namedtuple_expr(default_syms, default_vals)

    model_info = Dict(
        :name => modeldef[:name],
        :main_body => modeldef[:body],
        :arg_syms => arg_syms,
        :args_nt => args_nt,
        :defaults_nt => defaults_nt,
        :args => args,
        :whereparams => modeldef[:whereparams],
        :main_body_names => Dict(
            :ctx => gensym(:ctx),
            :vi => gensym(:vi),
            :sampler => gensym(:sampler),
            :model => gensym(:model)
        )
    )

    return model_info
end

"""
    replace_vi!(model_info)

Replaces `@varinfo()` expressions with a handle to the `VarInfo` struct.
"""
function replace_vi!(model_info)
    ex = model_info[:main_body]
    vi = model_info[:main_body_names][:vi]
    ex = MacroTools.postwalk(ex) do x
        if @capture(x, @varinfo())
            vi
        else
            x
        end
    end
    model_info[:main_body] = ex
    return model_info
end

"""
    replace_logpdf!(model_info)

Replaces `@logpdf()` expressions with the value of the accumulated `logpdf` in the `VarInfo` struct.
"""
function replace_logpdf!(model_info)
    ex = model_info[:main_body]
    vi = model_info[:main_body_names][:vi]
    ex = MacroTools.postwalk(ex) do x
        if @capture(x, @logpdf())
            :($vi.logp[])
        else
            x
        end
    end
    model_info[:main_body] = ex
    return model_info
end

"""
    replace_sampler!(model_info)

Replaces `@sampler()` expressions with a handle to the sampler struct.
"""
function replace_sampler!(model_info)
    ex = model_info[:main_body]
    spl = model_info[:main_body_names][:sampler]
    ex = MacroTools.postwalk(ex) do x
        if @capture(x, @sampler())
            spl
        else
            x
        end
    end
    model_info[:main_body] = ex
    return model_info
end

"""
    replace_tilde!(model_info)

Replace `~` and `.~` expressions with observation or assumption expressions, updating `model_info`.
"""
function replace_tilde!(model_info)
    # Apply the `@.` macro first.
    expr = model_info[:main_body]
    dottedexpr = MacroTools.postwalk(apply_dotted, expr)

    # Check for tilde operators.
    tildeexpr = MacroTools.postwalk(dottedexpr) do x
        # Check dot tilde first.
        dotargs = getargs_dottilde(x)
        if dotargs !== nothing
            L, R = dotargs
            return generate_dot_tilde(L, R, model_info)
        end

        # Check tilde.
        args = getargs_tilde(x)
        if args !== nothing
            L, R = args
            return generate_tilde(L, R, model_info)
        end

        return x
    end

    # Update the function body.
    model_info[:main_body] = tildeexpr

    return model_info
end

# """ Unbreak code highlighting in Emacs julia-mode

"""
    generate_tilde(left, right, model_info)

The `tilde` function generates `observe` expression for data variables and `assume`
expressions for parameter variables, updating `model_info` in the process.
"""
function generate_tilde(left, right, model_info)
    model = model_info[:main_body_names][:model]
    vi = model_info[:main_body_names][:vi]
    ctx = model_info[:main_body_names][:ctx]
    sampler = model_info[:main_body_names][:sampler]
    temp_right = gensym(:temp_right)
    out = gensym(:out)
    lp = gensym(:lp)
    vn = gensym(:vn)
    inds = gensym(:inds)
    isassumption = gensym(:isassumption)
    assert_ex = :(DynamicPPL.assert_dist($temp_right, msg = $(wrong_dist_errormsg(@__LINE__))))

    if left isa Symbol || left isa Expr
        ex = quote
            $temp_right = $right
            $assert_ex

            $vn, $inds = $(varname(left)), $(vinds(left))
            $isassumption = DynamicPPL.@isassumption($model, $left)
            if $isassumption
                $out = DynamicPPL.tilde_assume($ctx, $sampler, $temp_right, $vn, $inds, $vi)
                $left = $out[1]
                DynamicPPL.acclogp!($vi, $out[2])
            else
                DynamicPPL.acclogp!(
                    $vi,
                    DynamicPPL.tilde_observe($ctx, $sampler, $temp_right, $left, $vn, $inds, $vi),
                )
            end
        end
    else
        # we have a literal, which is automatically an observation
        ex = quote
            $temp_right = $right
            $assert_ex

            DynamicPPL.acclogp!(
                $vi,
                DynamicPPL.tilde_observe($ctx, $sampler, $temp_right, $left, $vi),
            )
        end
    end
    return ex
end

"""
    generate_dot_tilde(left, right, model_info)

This function returns the expression that replaces `left .~ right` in the model body. If
`preprocessed isa VarName`, then a `dot_assume` block will be run. Otherwise, a `dot_observe` block
will be run.
"""
function generate_dot_tilde(left, right, model_info)
    model = model_info[:main_body_names][:model]
    vi = model_info[:main_body_names][:vi]
    ctx = model_info[:main_body_names][:ctx]
    sampler = model_info[:main_body_names][:sampler]
    out = gensym(:out)
    temp_right = gensym(:temp_right)
    isassumption = gensym(:isassumption)
    lp = gensym(:lp)
    vn = gensym(:vn)
    inds = gensym(:inds)
    assert_ex = :(DynamicPPL.assert_dist($temp_right, msg = $(wrong_dist_errormsg(@__LINE__))))

    if left isa Symbol || left isa Expr
        ex = quote
            $temp_right = $right
            $assert_ex

            $vn, $inds = $(varname(left)), $(vinds(left))
            $isassumption = DynamicPPL.@isassumption($model, $left)

            if $isassumption
                $out = DynamicPPL.dot_tilde_assume($ctx, $sampler, $temp_right, $left, $vn, $inds, $vi)
                $left .= $out[1]
                DynamicPPL.acclogp!($vi, $out[2])
            else
                DynamicPPL.acclogp!(
                    $vi,
                    DynamicPPL.dot_tilde_observe($ctx, $sampler, $temp_right, $left, $vn, $inds, $vi),
                )
            end
        end
    else
        # we have a literal, which is automatically an observation
        ex = quote
            $temp_right = $right
            $assert_ex

            DynamicPPL.acclogp!(
                $vi,
                DynamicPPL.dot_tilde_observe($ctx, $sampler, $temp_right, $left, $vi),
            )
        end
    end
    return ex
end

const FloatOrArrayType = Type{<:Union{AbstractFloat, AbstractArray}}
hasmissing(T::Type{<:AbstractArray{TA}}) where {TA <: AbstractArray} = hasmissing(TA)
hasmissing(T::Type{<:AbstractArray{>:Missing}}) = true
hasmissing(T::Type) = false

"""
    build_output(model_info)

Builds the output expression.
"""
function build_output(model_info)
    # Construct user-facing function
    main_body_names = model_info[:main_body_names]
    ctx = main_body_names[:ctx]
    vi = main_body_names[:vi]
    model = main_body_names[:model]
    sampler = main_body_names[:sampler]

    # Arguments with default values
    args = model_info[:args]
    # Argument symbols without default values
    arg_syms = model_info[:arg_syms]
    # Arguments namedtuple
    args_nt = model_info[:args_nt]
    # Default values of the arguments
    # Arguments namedtuple
    defaults_nt = model_info[:defaults_nt]
    # Where parameters
    whereparams = model_info[:whereparams]
    # Model generator name
    model_gen = model_info[:name]
    # Main body of the model
    main_body = model_info[:main_body]

    unwrap_data_expr = Expr(:block)
    for var in arg_syms
        temp_var = gensym(:temp_var)
        varT = gensym(:varT)
        push!(unwrap_data_expr.args, quote
            local $var
            $temp_var = $model.args.$var
            $varT = typeof($temp_var)
            if $temp_var isa DynamicPPL.FloatOrArrayType
                $var = DynamicPPL.get_matching_type($sampler, $vi, $temp_var)
            elseif DynamicPPL.hasmissing($varT)
                $var = DynamicPPL.get_matching_type($sampler, $vi, $varT)($temp_var)
            else
                $var = $temp_var
            end
        end)
    end

    @gensym(evaluator, generator)
    generator_kw_form = isempty(args) ? () : (:($generator(;$(args...)) = $generator($(arg_syms...))),)
    model_gen_constructor = :(DynamicPPL.ModelGen{$(Tuple(arg_syms))}($generator, $defaults_nt))

    ex = quote
        function $evaluator(
            $model::Model,
            $vi::DynamicPPL.VarInfo,
            $sampler::DynamicPPL.AbstractSampler,
            $ctx::DynamicPPL.AbstractContext,
        )
            $unwrap_data_expr
            DynamicPPL.resetlogp!($vi)
            $main_body
        end

        $generator($(args...)) = DynamicPPL.Model($evaluator, $args_nt, $model_gen_constructor)
        $(generator_kw_form...)

        $model_gen = $model_gen_constructor
    end

    return esc(ex)
end

function warn_empty(body)
    if all(l -> isa(l, LineNumberNode), body.args)
        @warn("Model definition seems empty, still continue.")
    end
    return
end

"""
    get_matching_type(spl, vi, ::Type{T}) where {T}
Get the specialized version of type `T` for sampler `spl`. For example,
if `T === Float64` and `spl::Hamiltonian`, the matching type is `eltype(vi[spl])`.
"""
function get_matching_type end


text:
include("Graph.jl")

function AddLine(X, Y, PATH, edges, L, H)
    n = size(X, 1)
    m = size(edges, 1)
    leftX = ones(Int, L)
    rightX = ones(Int, L)
    upY = ones(Int, H)
    downY = ones(Int, H)

    deg = zeros(Int, n)
    for (u, v) in edges
        deg[u] += 1
        deg[v] += 1
    end
    for i = 1 : n
        x, y = X[i], Y[i]
        leftX[x] = max(leftX[x], div(deg[i]+3, 4)+1)
        rightX[x] = max(rightX[x], div(deg[i]+3, 4)+1)
        upY[y] = max(upY[y], div(deg[i]+3, 4)+1)
        downY[y] = max(downY[y], div(deg[i]+3, 4)+1)
    end

    dicX = Dict{Int, Bool}()
	dicY = Dict{Int, Bool}()
    nX = zeros(Int, L)
    nY = zeros(Int, H)
    nL, nH = 0, 0
    for i = 1 : L
        nL += (i == 1) ? (leftX[i]+1) : (leftX[i]+rightX[i-1]+1)
        nX[i] = nL
		dicX[nL] = true
    end
    nL += rightX[L]
    for i = 1 : H
        nH += (i == 1) ? (downY[i]+1) : (downY[i]+upY[i-1]+1)
        nY[i] = nH
		dicY[nL] = true
    end
    nH += upY[H]

    MAP = zeros(Bool, nL, nH)
    used = zeros(Bool, nL, nH)
    for i = 1 : n
        MAP[nX[X[i]], nY[Y[i]]] = true
        rr = div(deg[i]+3, 4)
        if rr > 0
            xx, yy = nX[X[i]], nY[Y[i]]
            di, dj = rr, 0
            while dj < rr
                MAP[xx+di, yy+dj] = true
                MAP[xx-dj, yy+di] = true
                MAP[xx+dj, yy-di] = true
                MAP[xx-di, yy-dj] = true
                di -= 1
                dj += 1
            end
        end
    end
    dx = [0, 0, 1, -1]
    dy = [1, -1, 0, 0]
    isInQ = zeros(Int32, nL, nH)

	cross(x1, y1, x2, y2) = (x1*y2 - x2*y1)
	isBend(x1, y1, x2, y2, x3, y3) = (cross(x1-x2, y1-y2, x3-x2, y3-y2) != 0)

    # find a shortest path from node u to node v
    SearchWay(u, v) = begin
        target = Dict{Tuple{Int, Int}, Bool}()
        Q = Array{Tuple{Int, Int}, 1}()
        fill!(isInQ, 0)
        prior = Array{Int, 1}()
		bend = Array{Int, 1}()
        front, rear = 1, 0
        pushQ(x, y) = begin
            rear += 1
            push!(Q, (x, y))
            isInQ[x, y] = rear
            push!(prior, 0)
			push!(bend, 0)
        end
        popQ() = begin
            tp = Q[front]
            front += 1
            return tp
        end

        ## push start points into the queue
        rr = div(deg[u]+3, 4)
        xx, yy = nX[X[u]], nY[Y[u]]
        di, dj = rr, 0
        while dj < rr
            (used[xx+di, yy+dj] == false) ? pushQ(xx+di, yy+dj) : nothing
            (used[xx-dj, yy+di] == false) ? pushQ(xx-dj, yy+di) : nothing
            (used[xx+dj, yy-di] == false) ? pushQ(xx+dj, yy-di) : nothing
            (used[xx-di, yy-dj] == false) ? pushQ(xx-di, yy-dj) : nothing
            di -= 1
            dj += 1
        end

        ## mark target points
        rr = div(deg[v]+3, 4)
        xx, yy = nX[X[v]], nY[Y[v]]
        di, dj = rr, 0
        while dj < rr
            (used[xx+di, yy+dj] == false) ? (target[(xx+di, yy+dj)] = true) : nothing
            (used[xx-dj, yy+di] == false) ? (target[(xx-dj, yy+di)] = true) : nothing
            (used[xx+dj, yy-di] == false) ? (target[(xx+dj, yy-di)] = true) : nothing
            (used[xx-di, yy-dj] == false) ? (target[(xx-di, yy-dj)] = true) : nothing
            di -= 1
            dj += 1
        end

        checkRange(x, y) = ((x > 0) && (y > 0) && (x <= nL) && (y <= nH))

        ans = Array{Tuple{Int, Int}, 1}()
		minBend = 1048576
        ## find the shortest path
        while front <= rear
            cntx, cnty = popQ()
			prex, prey = (prior[front-1] == 0) ? (nX[X[u]], nY[Y[u]]) : Q[prior[front-1]]
            for i = 1 : 4
                tx, ty = cntx+dx[i], cnty+dy[i]
				if (checkRange(tx, ty) == false) || (used[tx, ty] == true)
					continue
				end
				cntbend = isBend(prex, prey, cntx, cnty, tx, ty) ? (bend[front-1]+1) : bend[front-1]
				#cntbend += (!haskey(dicX, cntx)) + (!haskey(dicY, cnty))
                if haskey(target, (tx, ty)) && (cntbend+isBend(cntx, cnty, tx, ty, nX[X[v]], nY[Y[v]]) < minBend)
					minBend = cntbend + isBend(cntx, cnty, tx, ty, nX[X[v]], nY[Y[v]])
					tans = Array{Tuple{Int, Int}, 1}()
                    push!(tans, (nX[X[v]], nY[Y[v]]))
                    push!(tans, (tx, ty))
                    pt = front-1
                    while pt != 0
                        push!(tans, Q[pt])
                        pt = prior[pt]
                    end
                    push!(tans, (nX[X[u]], nY[Y[u]]))
                    ans = copy(tans)
                end
                if (MAP[tx, ty] == false) && (isInQ[tx, ty] == 0)
                    pushQ(tx, ty)
                    prior[rear] = front-1
					bend[rear] = cntbend
                elseif (MAP[tx, ty] == false) && (isInQ[tx, ty] >= front)
					if bend[isInQ[tx, ty]] > cntbend
						bend[isInQ[tx, ty]] = cntbend
						prior[isInQ[tx, ty]] = front-1
					end
				end
            end
        end
        return (minBend == 1048576) ? nothing : ans
    end

    PATH2 = Array{Array{Tuple{Int, Int}, 1}, 1}(undef, m)

	addToPATH2(eid, x, y) = begin
		tn = size(PATH2[eid], 1)
		if tn == 0
			push!(PATH2[eid], (x, y))
		elseif (x != PATH2[eid][tn][1]) || (y != PATH2[eid][tn][2])
			if tn > 1
				if cross(PATH2[eid][tn-1][1] - PATH2[eid][tn][1], PATH2[eid][tn-1][2] - PATH2[eid][tn][2], x - PATH2[eid][tn][1], y - PATH2[eid][tn][2]) == 0
					PATH2[eid][tn] = (x, y)
				else
					push!(PATH2[eid], (x, y))
				end
			else
				push!(PATH2[eid], (x, y))
			end
		end
	end

    dropedEdges = 0
    for i = 1 : m
        tmp = SearchWay(edges[i][1], edges[i][2])
        PATH2[i] = []
		if tmp == nothing
			dropedEdges += 1
			continue
		end
		for (xx, yy) in tmp
			addToPATH2(i, xx, yy)
			used[xx, yy] = true
		end
    end

	Xs = Array{Int, 1}()
	Ys = Array{Int, 1}()
	NewX = Dict{Int, Int}()
	NewY = Dict{Int, Int}()
	totalX, totalY = 0, 0
	for i = 1 : n
		push!(Xs, nX[X[i]])
		push!(Ys, nY[Y[i]])
	end
	for i = 1 : size(PATH, 1)
		for (xx, yy) in PATH[i]
			push!(Xs, nX[xx])
			push!(Ys, nY[yy])
		end
	end
	for i = 1 : size(PATH2, 1)
		for (xx, yy) in PATH2[i]
			push!(Xs, xx)
			push!(Ys, yy)
		end
	end
	sort!(Xs)
	sort!(Ys)
	for xx in Xs
		if !haskey(NewX, xx)
			totalX += 1
			NewX[xx] = totalX
		end
	end
	for yy in Ys
		if !haskey(NewY, yy)
			totalY += 1
			NewY[yy] = totalY
		end
	end

	realX = zeros(Int, n)
	realY = zeros(Int, n)
	realPATH = Array{Array{Tuple{Int, Int}, 1}, 1}(undef, size(PATH, 1))
	realPATH2 = Array{Array{Tuple{Int, Int}, 1}, 1}(undef, m)
	for i = 1 : n
		realX[i] = NewX[nX[X[i]]]
		realY[i] = NewY[nY[Y[i]]]
	end
	for i = 1 : size(PATH, 1)
		realPATH[i] = []
		for (xx, yy) in PATH[i]
			push!(realPATH[i], (NewX[nX[xx]], NewY[nY[yy]]))
		end
	end
	for i = 1 : size(PATH2, 1)
		realPATH2[i] = []
		for (xx, yy) in PATH2[i]
			push!(realPATH2[i], (NewX[xx], NewY[yy]))
		end
	end

	# reuse rows and columns
	statusX = zeros(Int8, totalX)
	statusY = zeros(Int8, totalY)
	for i = 1 : n
		statusX[realX[i]] |= 3
		statusY[realY[i]] |= 3
	end
	for i = 1 : size(realPATH, 1)
		for (xx, yy) in realPATH[i]
			statusX[xx] |= 1
			statusY[yy] |= 1
		end
	end
	for i = 1 : size(realPATH2, 1)
		for (xx, yy) in realPATH2[i]
			statusX[xx] |= 2
			statusY[yy] |= 2
		end
	end

	X2 = zeros(Int, totalX)
	Y2 = zeros(Int, totalY)
	ansX1 = 0
	ansX2 = 0
	ansY1 = 0
	ansY2 = 0

	for i = 1 : totalX
		if statusX[i] == 3
			tmp = max(ansX1, ansX2)
			ansX1 = tmp+1
			ansX2 = tmp+1
			X2[i] = tmp+1
		elseif statusX[i] == 2
			ansX2 += 1
			X2[i] = ansX2
		else
			ansX1 += 1
			X2[i] = ansX1
		end
	end

	for i = 1 : totalY
		if statusY[i] == 3
			tmp = max(ansY1, ansY2)
			ansY1 = tmp+1
			ansY2 = tmp+1
			Y2[i] = tmp+1
		elseif statusY[i] == 2
			ansY2 += 1
			Y2[i] = ansY2
		else
			ansY1 += 1
			Y2[i] = ansY1
		end
	end
	map!(x -> X2[x], realX, realX)
	map!(y -> Y2[y], realY, realY)
	foreach(i -> map!(p -> (X2[p[1]], Y2[p[2]]), realPATH[i], realPATH[i]), 1 : size(realPATH, 1))
	foreach(i -> map!(p -> (X2[p[1]], Y2[p[2]]), realPATH2[i], realPATH2[i]), 1 : size(realPATH2, 1))
	totalX = max(ansX1, ansX2)
	totalY = max(ansY1, ansY2)

	return realX, realY, realPATH, realPATH2, totalX, totalY, dropedEdges
end

function CalculatePreservedCentralityInformation(G, NodeC, maxNodeC, preservedEdgeC, deletedEdgeC, PATH2)
	nodeInfo = 0.0
	edgeInfo = preservedEdgeC
	for i = 1 : G.n
		nodeInfo += (NodeC[G.V[i]]*maxNodeC)
	end
	for i = 1 : size(PATH2, 1)
		if size(PATH2[i], 1) > 0
			edgeInfo += deletedEdgeC[i]
		end
	end

	return nodeInfo, edgeInfo, (nodeInfo+edgeInfo)/2
end


text:
__precompile__()

module CausalForest

import Base: length, show, convert, promote_rule, zero
using DecisionTree
using DelimitedFiles
using Random
using Statistics
using NonNegLeastSquares
using LinearAlgebra
import StatsBase

export TreeOOB, EnsembleOOB, apply_tree_oob, build_forest_oob, apply_forest_oob, load_data,
    build_tree, build_forest, apply_forest, shaff,
    get_occurence_frequencies, sample_U, apply_forest_oob # pour test

#####Includes#####

include("util.jl")
include("load_data.jl")
include("oob_random_forest.jl")
include("classification/tree.jl")
include("classification/main.jl")
include("regression/tree.jl")
include("regression/main.jl")

########## Types ##########

mutable struct LeafCausalH # honest
    inds_build   :: Vector{Int} # indices in the leaf during construction
    inds_pred    :: Vector{Int} # indices to predict these are the same if no honesty
end

mutable struct LeafCausalNH # nothonest
    inds_build   :: Vector{Int} # indices in the leaf
end

struct NodeCausalH{S}
    featid  :: Int
    featval :: S
    left    :: Union{LeafCausalH, NodeCausalH{S}}
    right   :: Union{LeafCausalH, NodeCausalH{S}}
end

struct NodeCausalNH{S}
    featid  :: Int
    featval :: S
    left    :: Union{LeafCausalNH, NodeCausalNH{S}}
    right   :: Union{LeafCausalNH, NodeCausalNH{S}}
end

LeafOrNodeCausalH{S} = Union{LeafCausalH, NodeCausalH{S}}

LeafOrNodeCausalNH{S} = Union{LeafCausalNH, NodeCausalNH{S}}

struct TreeCausalH{S} # honest causal tree
    tree       :: LeafOrNodeCausalH{S}
    inds_build :: Vector{Int}
    inds_pred  :: Vector{Int}
    oob        :: Vector{Int}
end

struct TreeCausalNH{S} # honest causal tree
    tree :: LeafOrNodeCausalNH{S}
    inds :: Vector{Int}
    oob  :: Vector{Int}
end

struct EnsembleCausal{S}
    trees     :: Union{Vector{TreeCausalH{S}}, Vector{TreeCausalNH{S}}}
    centering :: Bool
    bootstrap :: Bool
    honest    :: Bool
    X         :: AbstractMatrix{S}
    Y         :: AbstractVector{Float64}
    T         :: AbstractVector{Int}
    model_Y   :: Union{Nothing, EnsembleOOB{S, Float64}}
    model_T   :: Union{Nothing, EnsembleOOB{S, Int}}
    Y_center  :: Union{Nothing, AbstractVector{Float64}}
    T_center  :: Union{Nothing, AbstractVector{Int}}
end

is_leaf(l::LeafCausalH) = true
is_leaf(l::LeafCausalNH) = true
is_leaf(n::NodeCausalH) = false
is_leaf(n::NodeCausalNH) = false

length(forest::EnsembleCausal) = length(forest.trees)

#####Includes#####

include("causal/tree.jl")
include("causal/main.jl")
include("shaff/sampling.jl")
include("shaff/prf.jl")
include("shaff/shaff.jl")

end #module


text:
using TiledViews
using BenchmarkTools

function tv_test(s)
    @btime
    return
end


text:
using MedicalImagingUtils
using Documenter

DocMeta.setdocmeta!(MedicalImagingUtils, :DocTestSetup, :(using MedicalImagingUtils); recursive=true)

makedocs(;
    modules=[MedicalImagingUtils],
    authors="Dale <djblack@uci.edu> and contributors",
    repo="https://github.com/Dale-Black/MedicalImagingUtils.jl/blob/{commit}{path}#{line}",
    sitename="MedicalImagingUtils.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://Dale-Black.github.io/MedicalImagingUtils.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/Dale-Black/MedicalImagingUtils.jl",
)


text:
time_average_outputs(schedule, outputs, model, field_slicer) = schedule, outputs # fallback

"""
    time_average_outputs(schedule::AveragedTimeInterval, outputs, model, field_slicer)

Wrap each `output` in a `WindowedTimeAverage` on the time-averaged `schedule` and with `field_slicer`.

Returns the `TimeInterval` associated with `schedule` and a `NamedTuple` or `Dict` of the wrapped
outputs.
"""
function time_average_outputs(schedule::AveragedTimeInterval, outputs::Dict, model, field_slicer)
    averaged_outputs = Dict(name => WindowedTimeAverage(output, model; schedule=schedule, field_slicer=field_slicer)
                            for (name, output) in outputs)

    return TimeInterval(schedule), averaged_outputs
end

function time_average_outputs(schedule::AveragedTimeInterval, outputs::NamedTuple, model, field_slicer)
    output_names = Tuple(keys(outputs))

    output_values = Tuple(WindowedTimeAverage(outputs[name], model; schedule=schedule, field_slicer=field_slicer)
                          for name in output_names)

    averaged_outputs = NamedTuple{output_names}(output_values)

    return TimeInterval(schedule), averaged_outputs
end


text:
using Base.Test

using DataFrames
import FITSIO

import Celeste: AccuracyBenchmark
import Celeste: DeterministicVI
import Celeste: Model

@testset "accuracy benchmarks" begin

@testset "whole accuracy benchmark pipeline runs" begin
    images = SampleData.get_sdss_images(4263, 5, 119)

    rcf = RunCamcolField(4263, 5, 119)
    primary_df = AccuracyBenchmark.load_primary(rcf, SampleData.DATADIR)
    output_path = joinpath(SampleData.DATADIR, "test_primary.csv")
    new_csv = AccuracyBenchmark.write_catalog(output_path, primary_df)
    primary_df2 = AccuracyBenchmark.read_catalog(new_csv)
    catalog_entries = AccuracyBenchmark.make_initialization_catalog(primary_df2, true)
    # entry 8 is a star near [0.513037, 0.535631],
    # see http://legacysurvey.org/viewer/jpeg-cutout/?ra=0.5130&dec=0.5358&zoom=16&layer=sdss2
    target_sources = [8]

    patches = Model.get_sky_patches(images, catalog_entries)
    neighbor_map = Dict(i=>Model.find_neighbors(patches, i)
                        for i in target_sources)
    results = ParallelRun.one_node_single_infer(catalog_entries,
                                                patches,
                                                target_sources,
                                                neighbor_map, images,
                                                config=Celeste.Config())
    results_df = AccuracyBenchmark.celeste_to_df(results)

    coadd_path = joinpath(SampleData.DATADIR, "coadd_for_4263_5_119.fit")
    coadd_df = AccuracyBenchmark.load_coadd_catalog(coadd_path)
    coadd_path2 = joinpath(SampleData.DATADIR, "test_coadd.csv")
    coadd_csv = AccuracyBenchmark.write_catalog(coadd_path2, coadd_df)
    coadd_df2 = AccuracyBenchmark.read_catalog(coadd_csv)

    scores = AccuracyBenchmark.score_predictions(coadd_df2, [results_df])

    uncertainty_df = AccuracyBenchmark.get_uncertainty_df(coadd_df2, results_df)
    uq_scores = AccuracyBenchmark.score_uncertainty(uncertainty_df)
end

@testset "flux <-> mags conversion" begin
    # reference values based on `nmgy2lups()` from
    # https://github.com/esheldon/sdsspy/blob/683d6e0f16a998240a129942f80ad3ce6e7d5dfe/sdsspy/util.py
    @test isapprox(AccuracyBenchmark.flux_to_mag(15.0, 1), 19.559677, atol=1e-5)
    @test isapprox(AccuracyBenchmark.flux_to_mag(15.0, 3),19.559702, atol=1e-5)
    @test isapprox(AccuracyBenchmark.mag_to_flux(19.559677, 1), 15.0, atol=1e-5)
    @test isapprox(AccuracyBenchmark.mag_to_flux(19.559702, 3), 15.0, atol=1e-5)
end

@testset "color calculations" begin
    @test isapprox(AccuracyBenchmark.color_from_fluxes(15.0, 20.0), log(20 / 15))
    @test ismissing(AccuracyBenchmark.color_from_fluxes(15.0, 0.0))
    fluxes = AccuracyBenchmark.fluxes_from_colors(10.0, [-1.0, 0.0, 1.0, 2.0])
    @test isapprox(fluxes[1], exp(1.0) * 10.0)
    @test isapprox(fluxes[2], 10.0)
    @test isapprox(fluxes[3], 10.0)
    @test isapprox(fluxes[4], exp(1.0) * 10.0)
    @test isapprox(fluxes[5], exp(3.0) * 10.0)
end

@testset "angle canonicalization" begin
    @test AccuracyBenchmark.canonical_angle(95.0) == 95.0
    @test AccuracyBenchmark.canonical_angle(195.0) == 15.0
    @test AccuracyBenchmark.canonical_angle(-20.0) == 160.0

    @test isapprox(AccuracyBenchmark.degrees_to_diff(20., -30.), 50.)
    @test isapprox(AccuracyBenchmark.degrees_to_diff(-10., 190.), 20.)
end

@testset "variational params -> data frame -> catalog entry" begin
    variational_params = DeterministicVI.generic_init_source([1.0, 2.0])
    variational_params[Model.ids.gal_axis_ratio] = 0.5
    variational_params[Model.ids.gal_radius_px] = 10.0
    variational_params[Model.ids.gal_angle] = -pi / 4
    variational_params[Model.ids.flux_loc[2]] = log(20.0)
    variational_params[Model.ids.is_star[1]] = 0.01
    variational_params[Model.ids.is_star[2]] = 0.99

    data = AccuracyBenchmark.variational_parameters_to_data_frame_row(
        variational_params)
    # we'll just check a few particularly troublesome parameters :)
    @test isapprox(data[1, :gal_radius_px], 10 * sqrt(0.5))
    @test isapprox(data[1, :gal_angle_deg], 135.0)
    @test isapprox(data[1, :flux_r_nmgy], 20.0)

    data[1, :is_star] = false
    catalog_entry = AccuracyBenchmark.make_catalog_entry(first(eachrow(data)))
    @test isapprox(catalog_entry.gal_radius_px, 10.0)
    @test isapprox(catalog_entry.gal_angle, 3 * pi / 4)
    @test isapprox(catalog_entry.gal_fluxes[3], 20.0)
end

@testset "PSF serialize/deserialize" begin
    psf = PsfComponent[
        PsfComponent(0.3, StaticArrays.@SVector([1., 2.]), StaticArrays.@SMatrix([3. 1.; 1. 4.])),
        PsfComponent(0.7, StaticArrays.@SVector([-2., -3.]), StaticArrays.@SMatrix([2. 0.; 0. 1.]))
    ]
    header = FITSIO.FITSHeader(String[], [], String[])

    AccuracyBenchmark.serialize_psf_to_header(psf, header)
    new_psf = AccuracyBenchmark.make_psf_from_header(header)

    @test new_psf == psf
end

const FITS_HEADER_STRING = "SIMPLE  =                    T /                                                BITPIX  =                  -32 / 32 bit floating point                          NAXIS   =                    2                                                  NAXIS1  =                 2048                                                  NAXIS2  =                 1489                                                  EXTEND  =                    T /Extensions may be present                       BZERO   =              0.00000 /Set by MRD_SCALE                                BSCALE  =              1.00000 /Set by MRD_SCALE                                TAI     =        4537928038.58 / 1st row - Number of seconds since Nov 17 1858  RA      =            359.76830 / 1st row - Right ascension of telescope boresighDEC     =            0.000000  / 1st row - Declination of telescope boresight (dSPA     =              90.000  / 1st row - Camera col position angle wrt north (IPA     =             102.229  / 1st row - Instrument rotator position angle (deIPARATE =              0.0000  / 1st row - Instrument rotator angular velocity (AZ      =            14.483496 / 1st row - Azimuth  (encoder) of tele (0=N?) (deALT     =            56.247810 / 1st row - Altitude (encoder) of tele        (deFOCUS   =           -436.80000 / 1st row - Focus piston (microns?)              DATE-OBS= '2002-09-05'         / 1st row - TAI date                             TAIHMS  = '07:33:58.57'        / 1st row - TAI time (HH:MM:SS.SS) (TAI-UT = apprCOMMENT  TAI,RA,DEC,SPA,IPA,IPARATE,AZ,ALT,FOCUS at reading of col 0, row 0     ORIGIN  = 'SDSS    '                                                            "

@testset "parse FITS header" begin
    header = AccuracyBenchmark.parse_fits_header_from_string(FITS_HEADER_STRING)
    @test header["SIMPLE"] == true
    @test header["NAXIS1"] == 2048
    @test isapprox(header["RA"], 359.7683)
    @test strip(header["ORIGIN"]) == "SDSS"
end

@testset "filtering rows for scoring errors" begin
    function make_data()
        (
            DataFrame(
                gal_radius_px=Union{Float64, Missing}[10.0],
                gal_frac_dev=Union{Float64, Missing}[0.99],
                gal_axis_ratio=Union{Float64, Missing}[0.8],
            ),
            DataFrame(
                gal_axis_ratio=Union{Float64, Missing}[0.5],
                gal_angle_deg=Union{Float64, Missing}[10.0],
                dec=Union{Float64, Missing}[0.0],
            )
        )
    end

    truth, error = make_data()

    function check_row(column_name=:gal_axis_ratio)
        AccuracyBenchmark.is_good_row(first(eachrow(truth)), first(eachrow(error)), column_name)
    end

    @test check_row()

    truth, error = make_data()
    error[1, :gal_axis_ratio] = NaN
    @test !check_row()

    truth, error = make_data()
    error[1, :gal_axis_ratio] = missing
    @test !check_row()

    truth, error = make_data()
    truth[1, :gal_radius_px] = missing
    @test check_row()

    truth, error = make_data()
    truth[1, :gal_radius_px] = 25.0
    @test !check_row()

    truth, error = make_data()
    truth[1, :gal_frac_dev] = 0.5
    @test !check_row()
    @test check_row(:dec)

    truth, error = make_data()
    @test !check_row(:gal_angle_deg)
    truth[1, :gal_axis_ratio] = 0.2
    @test check_row(:gal_angle_deg)
end

@testset "match catalogs" begin
    ra = [0.0, 1.0, 2.0, 3.0]
    dec = [50.0, 51.0, 52.0, 53.0]
    truth = DataFrame(Any[ra, dec], [:ra, :dec])

    off = 0.2 / 3600.0
    ra1 = [0.0 + off, 1.0 - off, 5.0, 3.0 + off, 4.0]
    dec1 = [50.0 - off, 51.0 + off, 60.0, 53.0 - off, 50.0]
    pred1 = DataFrame(Any[ra1, dec1], [:ra, :dec])

    ra2 = [3.0 - off, 0.0 - off, 7.0, 5.0, 4.0]
    dec2 = [53.0 + off, 50.0 + off, 51.0, 60.0, 50.0]
    pred2 = DataFrame(Any[ra2, dec2], [:ra, :dec])

    truth_matched, predictions_matched =
        AccuracyBenchmark.match_catalogs(truth, [pred1, pred2])

    @test nrow(truth_matched) == 2
    @test length(predictions_matched) == 2
    for pred in predictions_matched
        @test nrow(pred) == 2
    end
    @test truth_matched[:ra] == [0.0, 3.0]
    @test truth_matched[:dec] == [50.0, 53.0]
    @test predictions_matched[1][:ra] == [0.0 + off, 3.0 + off]
    @test predictions_matched[1][:dec] == [50.0 - off, 53.0 - off]
    @test predictions_matched[2][:ra] == [0.0 - off, 3.0 - off]
    @test predictions_matched[2][:dec] == [50.0 + off, 53.0 + off]
end

end


text:
@testset "merge1" begin
    @test merge1([Dict{Any,Any}(), Dict{Int,Int}()]) == Dict{Any,Any}()
    @test merge1([Dict{Int,Int}( 1 => 2 ), Dict{Int,Int}()]) == Dict{Int,Int}(1=>2)
    @test merge1([Dict{Any,Any}( "cat" => 2 ), Dict{Int,Int}(1=>3)]) == Dict("cat"=>2, 1=>3)
end

@testset "foreachdict" begin
    a = Any[]
    d = Dict(:a=>1, :b=>2)
    foreachdict((k,v) -> append!(a,(k,v)),d)
    @test a == [:a, 1, :b, 2]
end

@testset "rgbtohex" begin
    @test rgbtohex(1,1,1) == "#010101"
end

@testset "arrcountmap" begin
    d = Dict( ["a","b"] => 1, ["b", "c"] => 2 )
    @test arrcountmap(keys(d),values(d)).map == Dict("c"=>2,"b"=>3,"a"=>1)
    @test arrcountmap(d).map == Dict("c"=>2,"b"=>3,"a"=>1)
end


text:
using Plots
pyplot(size=(1800, 1000))
using FileIO
using DataFrames
using Images
using Dates
using ScatteringLayers

echo_df = DataFrame(load("test\\data\\deimos-2019-04.csv"))
fmt = DateFormat("yyyymmddHH:MM:SS.s")
echo_df[:datetime] = DateTime.(string.(echo_df[:Date_M]) .* echo_df[:Time_M], fmt)

echo = unstack(echo_df, :Layer_depth_max, :datetime, :Sv_mean)[3:end, 2:end]
echo = disallowmissing(Matrix(echo))
echo = echo[end:-1:1, :]
echo[echo .< -90] .= -90
echo = 10 .^(echo/10)
m, n = size(echo)
dz = 1.0
order = 1

heatmap(10log10.(echo))
heatmap(imfilter(echo, Kernel.gaussian(0.75)), clim=(0, 1e-7))
heatmap(10log10.(imfilter(echo, Kernel.gaussian((2, 1)))))
# echo = imfilter(echo, Kernel.gaussian((1, 1)))

ping = imfilter(echo, Kernel.gaussian((2, 1)))[:, 10]
zz = Float64.(880:-1:8)
deimos_x = SpacePoint([36.712110, -122.187028])
ping = Ping(ping, zz, [deimos_x])

THRESH = 10^(-90/10)
μ, p, c = getpeaks(ping, thresh=THRESH)
mix = guesslayers(ping, thresh=THRESH)
plot(zz, ping, label="Ping")
scatter!(μ, p, label="Peaks")
# plot!(zz, mixpdf(mix, zz), label="Naive")

basis = getbasis(mix, zz)
N1 = refine_heights!(mix, basis, ping)
ping_nnls = mixpdf(mix, zz)
plot!(zz, ping_nnls, label="NNLS")

fit1 = refine_μσ!(mix, ping)
ping_refined1 = mixpdf(mix, zz)
plot!(zz, ping_refined1, label="Refined  1")
mix = guesslayers(ping, thresh=THRESH)
fit!(mix, ping)
ping_refined2 = mixpdf(mix, zz)
plot!(zz, ping_refined2, label="Refined 2")
scatter!(mix.μ, mixpdf(mix, mix.μ))

plot(cumsum(abs2.(ping .- ping_nnls)))
plot!(cumsum(abs2.(ping .- ping_refined1)))
plot!(cumsum(abs2.(ping .- ping_refined2)))

plot(zz, 10log10.(ping), label="Ping", ylim=(-90, -60))
plot!(zz, 10log10.(basis * N1), label="NNLS")
plot!(zz, 10log10.(mixpdf(mix, zz)), label="Refined σ")

plot(zz, ping, legend=false)
mix1 = deepcopy(mix)
mix1.μ .+= randn()
mix1.σ .+= max(0, randn())
mix1.N .+= 1e-9randn()
fit!(mix1, ping)
plot!(zz, mixpdf(mix1, zz), color=:black, alpha=0.2)

echogram = Echogram(imfilter(echo, Kernel.gaussian((2, 1))), zz, repeat([deimos_x], n))
THRESH = -90
layers = fitlayers(echogram, thresh=THRESH, refine=:none)
[l.k for l in layers]

p = heatmap(1:length(layers), -zz, 10log10.(echo), color=:grayscale, clim=(-90, NaN))
# p = heatmap(1:length(layers), -zz, echo, clim=(1e-8, 4e-7))
for i in 1:length(layers)
    scatter!(p, i*ones(layers[i].k), -layers[i].μ,
        color=:red, markeralpha=1, markerstrokecolor=:red, markersize=1e6*layers[i].N, legend=false)
end
p
savefig(p, "layerdots.png")

using LightGraphs
using MetaGraphs

function pushlayers!(g, mix, time)
    for i in 1:mix.k
        dict = Dict(:μ=>mix.μ[i], :σ=>mix.σ[i], :N=>mix.N[i], :k=>mix.k, :time=>time, :v=>0)
        add_vertex!(g, dict)
    end
end

function getnearest(x1, x2)
    n1, n2 = length(x1), length(x2)
    nearest = zeros(Int, n1)
    distances = similar(x1)
    for i in 1:n1
        dist = Inf
        for j in 1:n2
            dist1 = abs(x1[i] - x2[j])
            if dist1 < dist
                dist = dist1
                nearest[i] = j
                distances[i] = dist
            end
        end
    end
    return nearest, distances
end

@time nearest1, dists = getnearest(layers[1].μ, layers[2].μ)

using ProgressMeter
function link(layers, times, max_dist=Inf, max_rel_change=Inf, thresh=0)
    g = MetaDiGraph()
    # Add all layers to graph
    for (i, mix) in enumerate(layers)
        pushlayers!(g, mix, times[i])
    end
    t1 = times[1]
    printstyled("Linking...\n", color=:green)
    @showprogress for i in 2:length(times)
        t2 = times[i]
        v1 = collect(filter_vertices(g, :time, t1))
        v2 = collect(filter_vertices(g, :time, t2))
        μ1 = [get_prop(g, v, :μ) for v in v1]
        μ2 = [get_prop(g, v, :μ) for v in v2]
        N1 = [get_prop(g, v, :N) for v in v1]
        N2 = [get_prop(g, v, :N) for v in v2]
        nearest_μ, dists_μ = getnearest(μ1, μ2)
        nearest_N, dists_N = getnearest(N1, N2)

        N_order = sortperm(N1, rev=true)
        for j in N_order
            d = dists_μ[j]
            if (d <= max_dist) &
                (dists_N[j] / N1[j] < max_rel_change) &
                (length(inneighbors(g, v2[nearest_μ[j]])) == 0)
                add_edge!(g, v1[j], v2[nearest_μ[j]])
            end
        end
        t1 = t2
    end
    return g
end

g = link(layers, 1:420, 30, 2)

gr(size=(1800, 1000))
p = heatmap(1:length(layers), -zz, 10log10.(echo), color=:grayscale,
    clim=(THRESH, NaN))
for e in edges(g)
    v1 = src(e)
    v2 = dst(e)
    x1 = get_prop(g, v1, :μ)
    t1 = get_prop(g, v1, :time)
    x2 = get_prop(g, v2, :μ)
    t2 = get_prop(g, v2, :time)
    if get_prop(g, v1, :N) > 10^(-75/10)
        plot!(p, [t1, t2], -[x1, x2], color=:red, legend=false)
    end
end
p

μ = [get_prop(g, v, :μ) for v in vertices(g)]
σ = [get_prop(g, v, :σ) for v in vertices(g)]
N = [get_prop(g, v, :N) for v in vertices(g)]
using StatsPlots
histogram2d(μ, N, nbins=200)
histogram2d(log10.(N), log10.(σ), nbins=200, xlabel="N", ylabel="sigma")

max_stars_repo_name: ElOceanografo/ScatteringLayers

max_issues_repo_name: ElOceanografo/ScatteringLayers

max_forks_repo_name: ElOceanografo/ScatteringLayers


text:
using Example200530d
using Test

@testset "Example200530d.jl" begin
    @test my_f(2,1) == 7
    @test my_f(2,3) == 13
    @test my_f(1,3) == 11 # non-failing test
end

@testset "Derivative Tests" begin
    @test derivative_of_my_f(2,1) == 2
end


text:
module Day18

export day18

using CombinedParsers
using TextParse

function evaluate((start, operation_values))
    aggregated_value::Int = start
    for (op,val) in operation_values
        aggregated_value = eval(Expr(:call, Symbol(op),
       			             aggregated_value, val))
    end
    return aggregated_value
end

# Parser of '+' surrounded by whitepsace
plus = Sequence(2, Repeat(CombinedParsers.Regexp.whitespace_maybe), '+', Repeat(CombinedParsers.Regexp.whitespace_maybe))
# Parser of '*' surrounded by whitepsace
times = Sequence(2, Repeat(CombinedParsers.Regexp.whitespace_maybe), '*', Repeat(CombinedParsers.Regexp.whitespace_maybe))
# Parser of '+' or '*' surrounded by whitepsace
plustimes = Sequence(2, Repeat(CombinedParsers.Regexp.whitespace_maybe), CharIn("+*"), Repeat(CombinedParsers.Regexp.whitespace_maybe))

@syntax subterma = Either{Int}(Any[TextParse.Numeric(Int)])

@syntax for parenthesis in subterma
    term = evaluate |> join(subterma, plustimes, infix=:prefix )
    Sequence(2,'(',term,')')
end

@syntax expra = term

@syntax subtermb = Either{Int}(Any[TextParse.Numeric(Int)])

@syntax for parenthesis in subtermb
    adds = evaluate |> join(subtermb, plus, infix=:prefix )
    mult = evaluate |> join(adds, times, infix=:prefix )
    Sequence(2,'(',mult,')')
end

@syntax exprb = mult

day18(::Val{:a}, filename) = sum(map(expra, readlines("input_day18.txt")))

day18(::Val{:b}, filename) = sum(map(exprb, readlines("input_day18.txt")))

end


text:
import JSON

function main()
  text = open(readall, "1.json")
  jobj = JSON.parse(text)
  coordinates = jobj["coordinates"]
  len = length(coordinates)
  x = y = z = 0

  for coord in coordinates
    x += coord["x"]
    y += coord["y"]
    z += coord["z"]
  end

  println(x / len)
  println(y / len)
  println(z / len)
end

main()

max_stars_repo_name: japgolly/kostya-benchmarks

max_issues_repo_name: japgolly/kostya-benchmarks

max_forks_repo_name: japgolly/kostya-benchmarks


text:
import AlphaZero.GI
import AlphaZero.MCTS

include("../game.jl")
using .Mancala

# Late game (for debug purposes)
game = Game(Board([11, 11], [[1 2 4 0 0 0]; [1 2 4 0 0 0]]), true)
env = MCTS.Env{Game}(MCTS.RolloutOracle{Game}())
GI.interactive!(game, MCTS.AI(env, timeout=1.), GI.Human())


text:
# Initialization algorithms
#
#   Each algorithm is represented by a subtype of SeedingAlgorithm
#
#   Let alg be an instance of such an algorithm, then it should
#   support the following usage:
#
#       initseeds!(iseeds, alg, X)
#       initseeds_by_costs!(iseeds, alg, costs)
#
#   Here:
#       - iseeds:   a vector of resultant indexes of the chosen seeds
#       - alg:      the seeding algorithm instance
#       - X:        the data matrix, each column being a data point
#       - costs:    pre-computed pairwise cost matrix.
#
#   This function returns iseeds
#

"""
Base type for all seeding algorithms.

Each seeding algorithm should implement the two functions: [`initseeds!`](@ref)
and [`initseeds_by_costs!`](@ref).
"""
abstract type SeedingAlgorithm end

"""
    initseeds(alg::Union{SeedingAlgorithm, Symbol},
              X::AbstractMatrix, k::Integer)

Select `k` seeds from a ``d×n`` data matrix `X` using the `alg`
algorithm.

`alg` could be either an instance of [`SeedingAlgorithm`](@ref) or a symbolic
name of the algorithm.

Returns an integer vector of length `k` that contains the indices of
chosen seeds.
"""
initseeds(alg::SeedingAlgorithm, X::AbstractMatrix{<:Real}, k::Integer) =
    initseeds!(Vector{Int}(undef, k), alg, X)

"""
    initseeds_by_costs(alg::Union{SeedingAlgorithm, Symbol},
                       costs::AbstractMatrix, k::Integer)

Select `k` seeds from the ``n×n`` `costs` matrix using algorithm `alg`.

Here, ``\\mathrm{costs}_{ij}`` is the cost of assigning points ``i`` and ``j``
to the same cluster. One may, for example, use the squared Euclidean distance
between the points as the cost.

Returns an integer vector of length `k` that contains the indices of
chosen seeds.
"""
initseeds_by_costs(alg::SeedingAlgorithm, costs::AbstractMatrix{<:Real}, k::Integer) =
    initseeds_by_costs!(Vector{Int}(undef, k), alg, costs)

seeding_algorithm(s::Symbol) =
    s == :rand ? RandSeedAlg() :
    s == :kmpp ? KmppAlg() :
    s == :kmcen ? KmCentralityAlg() :
    error("Unknown seeding algorithm $s")

initseeds(algname::Symbol, X::AbstractMatrix{<:Real}, k::Integer) =
    initseeds(seeding_algorithm(algname), X, k)::Vector{Int}

initseeds_by_costs(algname::Symbol, costs::AbstractMatrix{<:Real}, k::Integer) =
    initseeds_by_costs(seeding_algorithm(algname), costs, k)

initseeds(iseeds::Vector{Int}, X::AbstractMatrix{<:Real}, k::Integer) = iseeds
initseeds_by_costs(iseeds::Vector{Int}, costs::AbstractMatrix{<:Real}, k::Integer) = iseeds

function copyseeds!(S::Matrix{<:AbstractFloat}, X::AbstractMatrix{<:Real},
                    iseeds::AbstractVector)
    d, n = size(X)
    k = length(iseeds)
    size(S) == (d, k) || throw(DimensionMismatch("Inconsistent array dimensions."))
    for j = 1:k
        copyto!(view(S, :, j), view(X, :, iseeds[j]))
    end
    return S
end

# NOTE: this should eventually be removed as only `copyseeds!` is used in `kmeans`.
copyseeds(X::AbstractMatrix{<:Real}, iseeds::AbstractVector) =
    copyseeds!(Matrix{eltype(X)}(undef, size(X, 1), length(iseeds)), X, iseeds)

function check_seeding_args(n::Integer, k::Integer)
    k >= 1 || error("The number of seeds must be positive.")
    k <= n || error("Attempted to select more seeds than data points.")
end

"""
Random seeding (`:rand`).

Chooses an arbitrary subset of ``k`` data points as cluster seeds.
"""
struct RandSeedAlg <: SeedingAlgorithm end

"""
    initseeds!(iseeds::AbstractVector{Int}, alg::SeedingAlgorithm,
               X::AbstractMatrix)

Initialize `iseeds` with the indices of cluster seeds for the `X` data matrix
using the `alg` seeding algorithm.

Returns `iseeds`.
"""
initseeds!(iseeds::IntegerVector, alg::RandSeedAlg, X::AbstractMatrix{<:Real}) =
    sample!(1:size(X, 2), iseeds; replace=false)

"""
    initseeds_by_costs!(iseeds::AbstractVector{Int}, alg::SeedingAlgorithm,
                        costs::AbstractMatrix)

Initialize `iseeds` with the indices of cluster seeds for the `costs` matrix
using the `alg` seeding algorithm.

Here, ``\\mathrm{costs}_{ij}`` is the cost of assigning points ``i`` and ``j``
to the same cluster. One may, for example, use the squared Euclidean distance
between the points as the cost.

Returns `iseeds`.
"""
initseeds_by_costs!(iseeds::IntegerVector, alg::RandSeedAlg, X::AbstractMatrix{<:Real}) =
    sample!(1:size(X,2), iseeds; replace=false)

"""
Kmeans++ seeding (`:kmpp`).

Chooses the seeds sequentially. The probability of a point to be chosen is
proportional to the minimum cost of assigning it to the existing seeds.

> D. Arthur and S. Vassilvitskii (2007).
> *k-means++: the advantages of careful seeding.*
> 18th Annual ACM-SIAM symposium on Discrete algorithms, 2007.
"""
struct KmppAlg <: SeedingAlgorithm end

function initseeds!(iseeds::IntegerVector, alg::KmppAlg,
                    X::AbstractMatrix{<:Real}, metric::PreMetric)
    n = size(X, 2)
    k = length(iseeds)
    check_seeding_args(n, k)

    # randomly pick the first center
    p = rand(1:n)
    iseeds[1] = p

    if k > 1
        mincosts = colwise(metric, X, view(X, :, p))
        mincosts[p] = 0

        # pick remaining (with a chance proportional to mincosts)
        tmpcosts = zeros(n)
        for j = 2:k
            p = wsample(1:n, mincosts)
            iseeds[j] = p

            # update mincosts
            c = view(X, :, p)
            colwise!(tmpcosts, metric, X, view(X, :, p))
            updatemin!(mincosts, tmpcosts)
            mincosts[p] = 0
        end
    end

    return iseeds
end

initseeds!(iseeds::IntegerVector, alg::KmppAlg, X::AbstractMatrix{<:Real}) =
    initseeds!(iseeds, alg, X, SqEuclidean())

function initseeds_by_costs!(iseeds::IntegerVector, alg::KmppAlg,
                             costs::AbstractMatrix{<:Real})
    n = size(costs, 1)
    k = length(iseeds)
    check_seeding_args(n, k)

    # randomly pick the first center
    p = rand(1:n)
    iseeds[1] = p

    if k > 1
        mincosts = costs[:, p]
        mincosts[p] = 0

        # pick remaining (with a chance proportional to mincosts)
        for j = 2:k
            p = wsample(1:n, mincosts)
            iseeds[j] = p

            # update mincosts
            updatemin!(mincosts, view(costs, :, p))
            mincosts[p] = 0
        end
    end

    return iseeds
end

"""
    kmpp(X, k)

Use *Kmeans++* to choose `k` seeds from the ``d×n`` data matrix `X`.
"""
kmpp(X::AbstractMatrix{<:Real}, k::Int) = initseeds(KmppAlg(), X, k)

"""
    kmpp_by_costs(C, k)

Use *Kmeans++* to choose `k` seeds based on the ``n×n`` cost matrix `C`.
"""
kmpp_by_costs(costs::AbstractMatrix{<:Real}, k::Int) = initseeds(KmppAlg(), costs, k)

"""
K-medoids initialization based on centrality (`:kmcen`).

Choose the ``k`` points with the highest *centrality* as seeds.

> Hae-Sang Park and Chi-Hyuck Jun.
> *A simple and fast algorithm for K-medoids clustering.*
> doi:10.1016/j.eswa.2008.01.039
"""
struct KmCentralityAlg <: SeedingAlgorithm end

function initseeds_by_costs!(iseeds::IntegerVector, alg::KmCentralityAlg,
                             costs::AbstractMatrix{<:Real})

    n = size(costs, 1)
    k = length(iseeds)
    k <= n || error("Attempted to select more seeds than points.")

    # compute score for each item
    coefs = vec(sum(costs, dims=2))
    for i = 1:n
        @inbounds coefs[i] = inv(coefs[i])
    end

    # scores[j] = \sum_j costs[i,j] / (\sum_{j'} costs[i,j'])
    #           = costs[i,j] * coefs[i]
    #
    # So this is matrix-vector multiplication
    scores = costs'coefs

    # lower score indicates better seeds
    sp = sortperm(scores)
    for i = 1:k
        @inbounds iseeds[i] = sp[i]
    end
    return iseeds
end

initseeds!(iseeds::IntegerVector, alg::KmCentralityAlg, X::AbstractMatrix{<:Real}, metric::PreMetric) =
    initseeds_by_costs!(iseeds, alg, pairwise(metric, X, dims=2))

initseeds!(iseeds::IntegerVector, alg::KmCentralityAlg, X::AbstractMatrix{<:Real}) =
    initseeds!(iseeds, alg, X, SqEuclidean())