# This file is a part of Julia. License is MIT: https://julialang.org/license

## Functions to compute the reduced shape

# for reductions that expand 0 dims to 1
reduced_indices(a::AbstractArray, region) = reduced_indices(axes(a), region)

# for reductions that keep 0 dims as 0
reduced_indices0(a::AbstractArray, region) = reduced_indices0(axes(a), region)

function reduced_indices(inds::Indices{N}, d::Int, rd::AbstractUnitRange) where N
    d < 1 && throw(ArgumentError("dimension must be ≥ 1, got $d"))
    if d == 1
        return (oftype(inds[1], rd), tail(inds)...)
    elseif 1 < d <= N
        return tuple(inds[1:d-1]..., oftype(inds[d], rd), inds[d+1:N]...)::typeof(inds)
    else
        return inds
    end
end
reduced_indices(inds::Indices, d::Int) = reduced_indices(inds, d, OneTo(1))

function reduced_indices0(inds::Indices{N}, d::Int) where N
    d < 1 && throw(ArgumentError("dimension must be ≥ 1, got $d"))
    if d <= N
        ind = inds[d]
        return reduced_indices(inds, d, (isempty(ind) ? ind : OneTo(1)))
    else
        return inds
    end
end

function reduced_indices(inds::Indices{N}, region) where N
    rinds = [inds...]
    for i in region
        isa(i, Integer) || throw(ArgumentError("reduced dimension(s) must be integers"))
        d = Int(i)
        if d < 1
            throw(ArgumentError("region dimension(s) must be ≥ 1, got $d"))
        elseif d <= N
            rinds[d] = oftype(rinds[d], OneTo(1))
        end
    end
    tuple(rinds...)::typeof(inds)
end

function reduced_indices0(inds::Indices{N}, region) where N
    rinds = [inds...]
    for i in region
        isa(i, Integer) || throw(ArgumentError("reduced dimension(s) must be integers"))
        d = Int(i)
        if d < 1
            throw(ArgumentError("region dimension(s) must be ≥ 1, got $d"))
        elseif d <= N
            rind = rinds[d]
            rinds[d] = oftype(rind, (isempty(rind) ? rind : OneTo(1)))
        end
    end
    tuple(rinds...)::typeof(inds)
end

###### Generic reduction functions #####

## initialization
# initarray! is only called by sum!, prod!, etc.
for (Op, initfun) in ((:(typeof(add_sum)), :zero), (:(typeof(mul_prod)), :one))
    @eval initarray!(a::AbstractArray{T}, ::$(Op), init::Bool, src::AbstractArray) where {T} = (init && fill!(a, $(initfun)(T)); a)
end

for Op in (:(typeof(max)), :(typeof(min)))
    @eval initarray!(a::AbstractArray{T}, ::$(Op), init::Bool, src::AbstractArray) where {T} = (init && copyfirst!(a, src); a)
end

for (Op, initval) in ((:(typeof(&)), true), (:(typeof(|)), false))
    @eval initarray!(a::AbstractArray, ::$(Op), init::Bool, src::AbstractArray) = (init && fill!(a, $initval); a)
end

# reducedim_initarray is called by
reducedim_initarray(A::AbstractArray, region, v0, ::Type{R}) where {R} = fill!(similar(A,R,reduced_indices(A,region)), v0)
reducedim_initarray(A::AbstractArray, region, v0::T) where {T} = reducedim_initarray(A, region, v0, T)

function reducedim_initarray0(A::AbstractArray{T}, region, f, ops) where T
    ri = reduced_indices0(A, region)
    if isempty(A)
        if prod(length, reduced_indices(A, region)) != 0
            reducedim_initarray0_empty(A, region, f, ops) # ops over empty slice of A
        else
            R = f == identity ? T : Core.Compiler.return_type(f, (T,))
            similar(A, R, ri)
        end
    else
        R = f == identity ? T : typeof(f(first(A)))
        si = similar(A, R, ri)
        mapfirst!(f, si, A)
    end
end

reducedim_initarray0_empty(A::AbstractArray, region, f, ops) = mapslices(x->ops(f.(x)), A, region)
reducedim_initarray0_empty(A::AbstractArray, region,::typeof(identity), ops) = mapslices(ops, A, region)

# TODO: better way to handle reducedim initialization
#
# The current scheme is basically following Steven G. Johnson's original implementation
#
promote_union(T::Union) = promote_type(promote_union(T.a), promote_union(T.b))
promote_union(T) = T

function reducedim_init(f, op::Union{typeof(+),typeof(add_sum)}, A::AbstractArray, region)
    _reducedim_init(f, op, zero, sum, A, region)
end
function reducedim_init(f, op::Union{typeof(*),typeof(mul_prod)}, A::AbstractArray, region)
    _reducedim_init(f, op, one, prod, A, region)
end
function _reducedim_init(f, op, fv, fop, A, region)
    T = promote_union(eltype(A))
    if T !== Any && applicable(zero, T)
        x = f(zero(T))
        z = op(fv(x), fv(x))
        Tr = typeof(z) == typeof(x) && !isbits(T) ? T : typeof(z)
    else
        z = fv(fop(f, A))
        Tr = typeof(z)
    end
    return reducedim_initarray(A, region, z, Tr)
end

reducedim_init(f, op::typeof(max), A::AbstractArray{T}, region) where {T} = reducedim_initarray0(A, region, f, maximum)
reducedim_init(f, op::typeof(min), A::AbstractArray{T}, region) where {T} = reducedim_initarray0(A, region, f, minimum)
reducedim_init(f::Union{typeof(abs),typeof(abs2)}, op::typeof(max), A::AbstractArray{T}, region) where {T} =
    reducedim_initarray(A, region, zero(f(zero(T))))

reducedim_init(f, op::typeof(&), A::AbstractArray, region) = reducedim_initarray(A, region, true)
reducedim_init(f, op::typeof(|), A::AbstractArray, region) = reducedim_initarray(A, region, false)

# specialize to make initialization more efficient for common cases

let
    BitIntFloat = Union{BitInteger, Math.IEEEFloat}
    T = Union{
        [AbstractArray{t} for t in uniontypes(BitIntFloat)]...,
        [AbstractArray{Complex{t}} for t in uniontypes(BitIntFloat)]...}

    global function reducedim_init(f, op::Union{typeof(+),typeof(add_sum)}, A::T, region)
        z = zero(f(zero(eltype(A))))
        reducedim_initarray(A, region, op(z, z))
    end
    global function reducedim_init(f, op::Union{typeof(*),typeof(mul_prod)}, A::T, region)
        u = one(f(one(eltype(A))))
        reducedim_initarray(A, region, op(u, u))
    end
end

## generic (map)reduction

has_fast_linear_indexing(a::AbstractArray) = false
has_fast_linear_indexing(a::Array) = true

function check_reducedims(R, A)
    # Check whether R has compatible dimensions w.r.t. A for reduction
    #
    # It returns an integer value (useful for choosing implementation)
    # - If it reduces only along leading dimensions, e.g. sum(A, dims=1) or sum(A, dims=(1,2)),
    #   it returns the length of the leading slice. For the two examples above,
    #   it will be size(A, 1) or size(A, 1) * size(A, 2).
    # - Otherwise, e.g. sum(A, dims=2) or sum(A, dims=(1,3)), it returns 0.
    #
    ndims(R) <= ndims(A) || throw(DimensionMismatch("cannot reduce $(ndims(A))-dimensional array to $(ndims(R)) dimensions"))
    lsiz = 1
    had_nonreduc = false
    for i = 1:ndims(A)
        Ri, Ai = axes(R, i), axes(A, i)
        sRi, sAi = length(Ri), length(Ai)
        if sRi == 1
            if sAi > 1
                if had_nonreduc
                    lsiz = 0  # to reduce along i, but some previous dimensions were non-reducing
                else
                    lsiz *= sAi  # if lsiz was set to zero, it will stay to be zero
                end
            end
        else
            Ri == Ai || throw(DimensionMismatch("reduction on array with indices $(axes(A)) with output with indices $(axes(R))"))
            had_nonreduc = true
        end
    end
    return lsiz
end

"""
Extract first entry of slices of array A into existing array R.
"""
copyfirst!(R::AbstractArray, A::AbstractArray) = mapfirst!(identity, R, A)

function mapfirst!(f, R::AbstractArray, A::AbstractArray)
    lsiz = check_reducedims(R, A)
    iA = axes(A)
    iR = axes(R)
    t = []
    for i in 1:length(iR)
        iAi = iA[i]
        push!(t, iAi == iR[i] ? iAi : first(iAi))
    end
    map!(f, R, view(A, t...))
end

function _mapreducedim!(f, op, R::AbstractArray, A::AbstractArray)
    lsiz = check_reducedims(R,A)
    isempty(A) && return R

    if has_fast_linear_indexing(A) && lsiz > 16
        # use mapreduce_impl, which is probably better tuned to achieve higher performance
        nslices = div(_length(A), lsiz)
        ibase = first(linearindices(A))-1
        for i = 1:nslices
            @inbounds R[i] = op(R[i], mapreduce_impl(f, op, A, ibase+1, ibase+lsiz))
            ibase += lsiz
        end
        return R
    end
    indsAt, indsRt = safe_tail(axes(A)), safe_tail(axes(R)) # handle d=1 manually
    keep, Idefault = Broadcast.shapeindexer(indsAt, indsRt)
    if reducedim1(R, A)
        # keep the accumulator as a local variable when reducing along the first dimension
        i1 = first(indices1(R))
        @inbounds for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            r = R[i1,IR]
            @simd for i in axes(A, 1)
                r = op(r, f(A[i, IA]))
            end
            R[i1,IR] = r
        end
    else
        @inbounds for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            @simd for i in axes(A, 1)
                R[i,IR] = op(R[i,IR], f(A[i,IA]))
            end
        end
    end
    return R
end

mapreducedim!(f, op, R::AbstractArray, A::AbstractArray) =
    (_mapreducedim!(f, op, R, A); R)

reducedim!(op, R::AbstractArray{RT}, A::AbstractArray) where {RT} =
    mapreducedim!(identity, op, R, A)

"""
    mapreduce(f, op[, v0], A::AbstractArray; dims)

Evaluates to the same as `reduce(op, f(v0), map(f, A); dims)`, but is generally
faster because the intermediate array is avoided.

# Examples
```jldoctest
julia> a = reshape(Vector(1:16), (4,4))
4×4 Array{Int64,2}:
 1  5   9  13
 2  6  10  14
 3  7  11  15
 4  8  12  16

julia> mapreduce(isodd, *, a, dims=1)
1×4 Array{Bool,2}:
 false  false  false  false

julia> mapreduce(isodd, |, true, a, dims=1)
1×4 Array{Bool,2}:
 true  true  true  true
```
"""
mapreduce(f, op, v0, A::AbstractArray; dims=:) = _mapreduce_dim(f, op, v0, A, dims)

mapreduce(f, op, A::AbstractArray; dims=:) = _mapreduce_dim(f, op, A, dims)

_mapreduce_dim(f, op, v0, A::AbstractArray, ::Colon) = mapfoldl(f, op, v0, A)

_mapreduce_dim(f, op, A::AbstractArray, ::Colon) = _mapreduce(f, op, IndexStyle(A), A)

_mapreduce_dim(f, op, v0, A::AbstractArray, dims) =
    mapreducedim!(f, op, reducedim_initarray(A, dims, v0), A)

_mapreduce_dim(f, op, A::AbstractArray, dims) =
    mapreducedim!(f, op, reducedim_init(f, op, A, dims), A)

"""
    reduce(f[, v0], A; dims)

Reduce 2-argument function `f` along dimensions of `A`. `dims` is a vector specifying the
dimensions to reduce, and `v0` is the initial value to use in the reductions. For `+`, `*`,
`max` and `min` the `v0` argument is optional.

The associativity of the reduction is implementation-dependent; if you need a particular
associativity, e.g. left-to-right, you should write your own loop. See documentation for
[`reduce`](@ref).

# Examples
```jldoctest
julia> a = reshape(Vector(1:16), (4,4))
4×4 Array{Int64,2}:
 1  5   9  13
 2  6  10  14
 3  7  11  15
 4  8  12  16

julia> reduce(max, a, dims=2)
4×1 Array{Int64,2}:
 13
 14
 15
 16

julia> reduce(max, a, dims=1)
1×4 Array{Int64,2}:
 4  8  12  16
```
"""
reduce(op, v0, A::AbstractArray; dims=:) = _reduce_dim(op, v0, A, dims)

_reduce_dim(op, v0, A, dims) = mapreduce(identity, op, v0, A, dims=dims)
_reduce_dim(op, v0, A, ::Colon) = mapreduce(identity, op, v0, A)

reduce(op, A::AbstractArray; dims=:) = _reduce_dim(op, A, dims)

_reduce_dim(op, A, dims) = mapreduce(identity, op, A, dims=dims)
_reduce_dim(op, A, ::Colon) = mapreduce(identity, op, A)

##### Specific reduction functions #####
"""
    sum(A::AbstractArray; dims)

Sum elements of an array over the given dimensions.

# Examples
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> sum(A, dims=1)
1×2 Array{Int64,2}:
 4  6

julia> sum(A, dims=2)
2×1 Array{Int64,2}:
 3
 7
```
"""
sum(A::AbstractArray; dims)

"""
    sum!(r, A)

Sum elements of `A` over the singleton dimensions of `r`, and write results to `r`.

# Examples
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> sum!([1; 1], A)
2-element Array{Int64,1}:
 3
 7

julia> sum!([1 1], A)
1×2 Array{Int64,2}:
 4  6
```
"""
sum!(r, A)

"""
    prod(A::AbstractArray; dims)

Multiply elements of an array over the given dimensions.

# Examples
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> prod(A, dims=1)
1×2 Array{Int64,2}:
 3  8

julia> prod(A, dims=2)
2×1 Array{Int64,2}:
  2
 12
```
"""
prod(A::AbstractArray; dims)

"""
    prod!(r, A)

Multiply elements of `A` over the singleton dimensions of `r`, and write results to `r`.

# Examples
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> prod!([1; 1], A)
2-element Array{Int64,1}:
  2
 12

julia> prod!([1 1], A)
1×2 Array{Int64,2}:
 3  8
```
"""
prod!(r, A)

"""
    maximum(A::AbstractArray; dims)

Compute the maximum value of an array over the given dimensions. See also the
[`max(a,b)`](@ref) function to take the maximum of two or more arguments,
which can be applied elementwise to arrays via `max.(a,b)`.

```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> maximum(A, dims=1)
1×2 Array{Int64,2}:
 3  4

julia> maximum(A, dims=2)
2×1 Array{Int64,2}:
 2
 4
```
"""
maximum(A::AbstractArray; dims)

"""
    maximum!(r, A)

Compute the maximum value of `A` over the singleton dimensions of `r`, and write results to `r`.

# Examples
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> maximum!([1; 1], A)
2-element Array{Int64,1}:
 2
 4

julia> maximum!([1 1], A)
1×2 Array{Int64,2}:
 3  4
```
"""
maximum!(r, A)

"""
    minimum(A::AbstractArray; dims)

Compute the minimum value of an array over the given dimensions. See also the
[`min(a,b)`](@ref) function to take the minimum of two or more arguments,
which can be applied elementwise to arrays via `min.(a,b)`.

# Examples
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> minimum(A, dims=1)
1×2 Array{Int64,2}:
 1  2

julia> minimum(A, dims=2)
2×1 Array{Int64,2}:
 1
 3
```
"""
minimum(A::AbstractArray; dims)

"""
    minimum!(r, A)

Compute the minimum value of `A` over the singleton dimensions of `r`, and write results to `r`.

# Examples
```jldoctest
julia> A = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> minimum!([1; 1], A)
2-element Array{Int64,1}:
 1
 3

julia> minimum!([1 1], A)
1×2 Array{Int64,2}:
 1  2
```
"""
minimum!(r, A)

"""
    all(A; dims)

Test whether all values along the given dimensions of an array are `true`.

# Examples
```jldoctest
julia> A = [true false; true true]
2×2 Array{Bool,2}:
 true  false
 true   true

julia> all(A, dims=1)
1×2 Array{Bool,2}:
 true  false

julia> all(A, dims=2)
2×1 Array{Bool,2}:
 false
  true
```
"""
all(A::AbstractArray; dims)

"""
    all!(r, A)

Test whether all values in `A` along the singleton dimensions of `r` are `true`, and write results to `r`.

# Examples
```jldoctest
julia> A = [true false; true false]
2×2 Array{Bool,2}:
 true  false
 true  false

julia> all!([1; 1], A)
2-element Array{Int64,1}:
 0
 0

julia> all!([1 1], A)
1×2 Array{Int64,2}:
 1  0
```
"""
all!(r, A)

"""
    any(A; dims)

Test whether any values along the given dimensions of an array are `true`.

# Examples
```jldoctest
julia> A = [true false; true false]
2×2 Array{Bool,2}:
 true  false
 true  false

julia> any(A, dims=1)
1×2 Array{Bool,2}:
 true  false

julia> any(A, dims=2)
2×1 Array{Bool,2}:
 true
 true
```
"""
any(::AbstractArray; dims)

"""
    any!(r, A)

Test whether any values in `A` along the singleton dimensions of `r` are `true`, and write
results to `r`.

# Examples
```jldoctest
julia> A = [true false; true false]
2×2 Array{Bool,2}:
 true  false
 true  false

julia> any!([1; 1], A)
2-element Array{Int64,1}:
 1
 1

julia> any!([1 1], A)
1×2 Array{Int64,2}:
 1  0
```
"""
any!(r, A)

for (fname, _fname, op) in [(:sum,     :_sum,     :add_sum), (:prod,    :_prod,    :mul_prod),
                            (:maximum, :_maximum, :max),     (:minimum, :_minimum, :min)]
    @eval begin
        # User-facing methods with keyword arguments
        @inline ($fname)(a::AbstractArray; dims=:) = ($_fname)(a, dims)
        @inline ($fname)(f::Callable, a::AbstractArray; dims=:) = ($_fname)(f, a, dims)

        # Underlying implementations using dispatch
        ($_fname)(a, ::Colon) = ($_fname)(identity, a, :)
        ($_fname)(f, a, ::Colon) = mapreduce(f, $op, a)
    end
end

any(a::AbstractArray; dims=:)              = _any(a, dims)
any(f::Function, a::AbstractArray; dims=:) = _any(f, a, dims)
_any(a, ::Colon)                           = _any(identity, a, :)
all(a::AbstractArray; dims=:)              = _all(a, dims)
all(f::Function, a::AbstractArray; dims=:) = _all(f, a, dims)
_all(a, ::Colon)                           = _all(identity, a, :)

for (fname, op) in [(:sum, :add_sum), (:prod, :mul_prod),
                    (:maximum, :max), (:minimum, :min),
                    (:all, :&),       (:any, :|)]
    fname! = Symbol(fname, '!')
    _fname = Symbol('_', fname)
    @eval begin
        $(fname!)(f::Function, r::AbstractArray, A::AbstractArray; init::Bool=true) =
            mapreducedim!(f, $(op), initarray!(r, $(op), init, A), A)
        $(fname!)(r::AbstractArray, A::AbstractArray; init::Bool=true) = $(fname!)(identity, r, A; init=init)

        $(_fname)(A, dims)    = $(_fname)(identity, A, dims)
        $(_fname)(f, A, dims) = mapreduce(f, $(op), A, dims=dims)
    end
end

##### findmin & findmax #####
# The initial values of Rval are not used if the correponding indices in Rind are 0.
#
function findminmax!(f, Rval, Rind, A::AbstractArray{T,N}) where {T,N}
    (isempty(Rval) || isempty(A)) && return Rval, Rind
    lsiz = check_reducedims(Rval, A)
    for i = 1:N
        axes(Rval, i) == axes(Rind, i) || throw(DimensionMismatch("Find-reduction: outputs must have the same indices"))
    end
    # If we're reducing along dimension 1, for efficiency we can make use of a temporary.
    # Otherwise, keep the result in Rval/Rind so that we traverse A in storage order.
    indsAt, indsRt = safe_tail(axes(A)), safe_tail(axes(Rval))
    keep, Idefault = Broadcast.shapeindexer(indsAt, indsRt)
    ks = keys(A)
    k, kss = next(ks, start(ks))
    zi = zero(eltype(ks))
    if reducedim1(Rval, A)
        i1 = first(indices1(Rval))
        @inbounds for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            tmpRv = Rval[i1,IR]
            tmpRi = Rind[i1,IR]
            for i in axes(A,1)
                tmpAv = A[i,IA]
                if tmpRi == zi || (tmpRv == tmpRv && (tmpAv != tmpAv || f(tmpAv, tmpRv)))
                    tmpRv = tmpAv
                    tmpRi = k
                end
                k, kss = next(ks, kss)
            end
            Rval[i1,IR] = tmpRv
            Rind[i1,IR] = tmpRi
        end
    else
        @inbounds for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            for i in axes(A, 1)
                tmpAv = A[i,IA]
                tmpRv = Rval[i,IR]
                tmpRi = Rind[i,IR]
                if tmpRi == zi || (tmpRv == tmpRv && (tmpAv != tmpAv || f(tmpAv, tmpRv)))
                    Rval[i,IR] = tmpAv
                    Rind[i,IR] = k
                end
                k, kss = next(ks, kss)
            end
        end
    end
    Rval, Rind
end

"""
    findmin!(rval, rind, A) -> (minval, index)

Find the minimum of `A` and the corresponding linear index along singleton
dimensions of `rval` and `rind`, and store the results in `rval` and `rind`.
`NaN` is treated as less than all other values.
"""
function findmin!(rval::AbstractArray, rind::AbstractArray, A::AbstractArray;
                  init::Bool=true)
    findminmax!(isless, init && !isempty(A) ? fill!(rval, first(A)) : rval, fill!(rind,zero(eltype(keys(A)))), A)
end

"""
    findmin(A; dims) -> (minval, index)

For an array input, returns the value and index of the minimum over the given dimensions.
`NaN` is treated as less than all other values.

# Examples
```jldoctest
julia> A = [1.0 2; 3 4]
2×2 Array{Float64,2}:
 1.0  2.0
 3.0  4.0

julia> findmin(A, dims=1)
([1.0 2.0], CartesianIndex{2}[CartesianIndex(1, 1) CartesianIndex(1, 2)])

julia> findmin(A, dims=2)
([1.0; 3.0], CartesianIndex{2}[CartesianIndex(1, 1); CartesianIndex(2, 1)])
```
"""
findmin(A::AbstractArray; dims=:) = _findmin(A, dims)

function _findmin(A, region)
    ri = reduced_indices0(A, region)
    if isempty(A)
        if prod(map(length, reduced_indices(A, region))) != 0
            throw(ArgumentError("collection slices must be non-empty"))
        end
        (similar(A, ri), similar(dims->zeros(eltype(keys(A)), dims), ri))
    else
        findminmax!(isless, fill!(similar(A, ri), first(A)),
                    similar(dims->zeros(eltype(keys(A)), dims), ri), A)
    end
end

isgreater(a, b) = isless(b,a)

"""
    findmax!(rval, rind, A) -> (maxval, index)

Find the maximum of `A` and the corresponding linear index along singleton
dimensions of `rval` and `rind`, and store the results in `rval` and `rind`.
`NaN` is treated as greater than all other values.
"""
function findmax!(rval::AbstractArray, rind::AbstractArray, A::AbstractArray;
                  init::Bool=true)
    findminmax!(isgreater, init && !isempty(A) ? fill!(rval, first(A)) : rval, fill!(rind,zero(eltype(keys(A)))), A)
end

"""
    findmax(A; dims) -> (maxval, index)

For an array input, returns the value and index of the maximum over the given dimensions.
`NaN` is treated as greater than all other values.

# Examples
```jldoctest
julia> A = [1.0 2; 3 4]
2×2 Array{Float64,2}:
 1.0  2.0
 3.0  4.0

julia> findmax(A, dims=1)
([3.0 4.0], CartesianIndex{2}[CartesianIndex(2, 1) CartesianIndex(2, 2)])

julia> findmax(A, dims=2)
([2.0; 4.0], CartesianIndex{2}[CartesianIndex(1, 2); CartesianIndex(2, 2)])
```
"""
findmax(A::AbstractArray; dims=:) = _findmax(A, dims)

function _findmax(A, region)
    ri = reduced_indices0(A, region)
    if isempty(A)
        if prod(map(length, reduced_indices(A, region))) != 0
            throw(ArgumentError("collection slices must be non-empty"))
        end
        similar(A, ri), similar(dims->zeros(eltype(keys(A)), dims), ri)
    else
        findminmax!(isgreater, fill!(similar(A, ri), first(A)),
                    similar(dims->zeros(eltype(keys(A)), dims), ri), A)
    end
end

reducedim1(R, A) = length(indices1(R)) == 1


"""
    load_timeseriesdata(data_path::String; T-#Segments,
years::Array{Int64,1}=# years to be selected for the time series, att::Array{String,1}=# attributes to be loaded)
- Loading all `*.csv` files in the folder or the file `data_path`
The `*.csv` files shall have the following structure and must have the same length:
|Timestamp |[column names...]|
|[iterator]|[values]         |
The first column should be called `Timestamp` if it contains a time iterator
The other columns can specify the single timeseries like specific geolocation.
Each column in `[file name].csv` file will be added to the ClustData.data called `"[file name]-[column name]"`
Loads all attributes if the `att`-Array is empty or only the ones specified in it
"""
function load_timeseries_data(data_path;
                              region::String="none",
                              T::Int64=24,
                              years::Array{Int64,1}=[2016],
                              att::Array{String,1}=Array{String,1}())
  dt = Dict{String,Array}()
  num=0
  K=0
  #Check if data_path is directory or file
  if isdir(data_path)
      for full_data_name in readdir(data_path)
          if split(full_data_name,".")[end]=="csv"
              data_name=split(full_data_name,".")[1]
              if isempty(att) || data_name in att
                  # Add the
                  K=add_timeseries_data!(dt, data_name, data_path; K=K, T=T, years=years)
              end
          end
      end
  elseif isfile(data_path)
      full_data_name=splitdir(data_path)[end]
      data_name=split(full_data_name,".")[1]
      K=add_timeseries_data!(dt, data_name, dirname(data_path); K=K, T=T, years=years)
  else
      throw(@error("The path $data_path is neither recognized as a directory nor as a file"))
  end
  # Store the data
  ts_input_data =  ClustData(FullInputData(region, years, num, dt),K,T)
  return ts_input_data
end #load_timeseries_data

"""
    add_timeseries_data!(dt::Dict{String,Array}, data::DataFrame; K::Int64=0, T::Int64=24, years::Array{Int64,1}=[2016])
selects first the years and second the data_points so that their number is a multiple of T and same with the other timeseries
"""
function add_timeseries_data!(dt::Dict{String,Array},
                            data_name::SubString,
                            data_path::String;
                            K::Int64=0,
                            T::Int64=24,
                            years::Array{Int64,1}=[2016])
    #Load the data
    data_df=CSV.read(joinpath(data_path,data_name*".csv");allowmissing=:none)
    # Add it to the dictionary
    return add_timeseries_data!(dt,data_name, data_df; K=K, T=T, years=years)
end

"""
    add_timeseries_data!(dt::Dict{String,Array}, data::DataFrame; K::Int64=0, T::Int64=24, years::Array{Int64,1}=[2016])
selects first the years and second the data_points so that their number is a multiple of T and same with the other timeseries
"""
function add_timeseries_data!(dt::Dict{String,Array},
                            data_name::SubString,
                            data::DataFrame;
                            K::Int64=0,
                            T::Int64=24,
                            years::Array{Int64,1}=[2016])
    # find the right years to select
    data_selected=data[in.(data[:year],[years]),:]
    for column in eachcol(data_selected, true)
        # check that this column isn't time or year
        if !(column[1] in [:Timestamp,:time,:Time,:Zeit,:year])
            K_calc=Int(floor(length(column[2])/T))
            if K_calc!=K && K!=0
                @error("The time_series $(column[1]) has K=$K_calc != K=$K of the previous")
            else
                K=K_calc
            end
            dt[data_name*"-"*string(column[1])]=Float64.(column[2][1:(Int(T*K))])
        end
    end
    return K
end

"""
        combine_timeseries_weather_data(ts::ClustData,ts_weather::ClustData)
-`ts` is the shorter timeseries with e.g. the demand
-`ts_weather` is the longer timeseries with the weather information
The `ts`-timeseries is repeated to match the number of periods of the longer `ts_weather`-timeseries.
If the number of periods of the `ts_weather` data isn't a multiple of the `ts`-timeseries, the necessary number of the `ts`-timeseries periods 1 to x are attached to the end of the new combined timeseries.
"""
function combine_timeseries_weather_data(ts::ClustData,
                                        ts_weather::ClustData)
    ts.T==ts_weather.T || throw(@error "The number of timesteps per period is not the same: `ts.T=$(ts.T)≢$(ts_weather.T)=ts_weather.T`")
    ts.K<=ts_weather.K || throw(@error "The number of timesteps in the `ts`-timeseries isn't shorter or equal to the ones in the `ts_weather`-timeseries.")
    ts_weather.K%ts.K==0 || @warn "The number of periods of the `ts_weather` data isn't a multiple of the other `ts`-timeseries: periods 1 to $(ts_weather.K%ts.K) are attached to the end of the new combined timeseries."
    ts_data=deepcopy(ts_weather.data)
    ts_mean=deepcopy(ts_weather.mean)
    ts_sdv=deepcopy(ts_weather.sdv)
    for (k,v) in ts.data
        ts_data[k]=repeat(v, 1, ceil(Int,ts_weather.K/ts.K))[:,1:ts_weather.K]
    end
    for (k,v) in ts.mean
        ts_mean[k]=v
    end
    for (k,v) in ts.sdv
        ts_sdv[k]=v
    end

    return ClustData(ts.region, ts_weather.years, ts_weather.K, ts_weather.T, ts_data, ts_weather.weights, ts_mean, ts_sdv, ts_weather.deltas, ts_weather.k_ids)
end


import Pkg;

Pkg.add("MLJ")

## 通过 clone 的方式安装引入
Pkg.add("mocha")

Pkg.add("DataFrames")

Pkg.add("Plots")

Pkg.add("Flux")

Pkg.add("ECharts")


# This file is a part of BayesianBlocks.jl, licensed under the MIT License (MIT).

__precompile__(true)

"""
    BayesianBlocks

Bayesian Blocks algorithm described in [Scargle 2021](https://doi.org/10.1088/0004-637X/764/2/167).
"""
module BayesianBlocks

include("core.jl")

end # module


import Base: iterate, indexed_iterate

### ITERATION: We define iterate in terms of getindex which does the necessary tracking.

# Iterate is used in `for x in a` loops
iterate(a::Value,i=nothing)=throw(MethodError(iterate,(a,i)))
iterate(a::Value{T},i=1) where {T<:AbstractArray} = i > length(a) ? nothing : (a[i],i+1)
iterate(a::Value{T},i=1) where {T<:Tuple}         = i > length(a) ? nothing : (a[i],i+1)
iterate(a::Value{T})   where {T<:AbstractDict} = (v=iterate(value(a));   v==nothing ? v : (((k,_),j)=v;(k=>a[k],j)))
iterate(a::Value{T},i) where {T<:AbstractDict} = (v=iterate(value(a),i); v==nothing ? v : (((k,_),j)=v;(k=>a[k],j)))
iterate(a::Value{T})   where {T<:Number} = (a,nothing)
iterate(a::Value{T},i) where {T<:Number} = nothing

# indexed_iterate for `(x,y)=a` multiple assignments.
indexed_iterate(a::Value{T},i::Int,state=1) where {T<:Tuple} = (a[i],i+1)
indexed_iterate(a::Value{T},i::Int,state=1) where {T<:Array} = (a[i],i+1)


abstract type AbstractVectorBounds end

function set_from_constants(
    b::AbstractVectorBounds,
    ::Type{<:MOI.EqualTo},
    index,
)
    return MOI.EqualTo(b.lower[index])
end

function set_from_constants(
    b::AbstractVectorBounds,
    S::Type{<:Union{MOI.GreaterThan,MOI.EqualTo}},
    index,
)
    # Lower and upper bounds are equal for `EqualTo`, we can take either of them.
    return S(b.lower[index])
end

function set_from_constants(
    b::AbstractVectorBounds,
    S::Type{<:MOI.LessThan},
    index,
)
    return S(b.upper[index])
end

function set_from_constants(
    b::AbstractVectorBounds,
    S::Type{<:Union{MOI.Interval,MOI.Semicontinuous,MOI.Semiinteger}},
    index,
)
    return S(b.lower[index], b.upper[index])
end

function set_from_constants(
    ::AbstractVectorBounds,
    S::Type{<:Union{MOI.Integer,MOI.ZeroOne}},
    index,
)
    return S()
end

"""
    SUPPORTED_VARIABLE_SCALAR_SETS{T}

The union of scalar sets for `VariableIndex` constraints supported by
`Utilities.Hyperrectangle` and `Utilities.VariablesContainer`.
"""
const SUPPORTED_VARIABLE_SCALAR_SETS{T} = Union{
    MOI.EqualTo{T},
    MOI.GreaterThan{T},
    MOI.LessThan{T},
    MOI.Interval{T},
    MOI.Integer,
    MOI.ZeroOne,
    MOI.Semicontinuous{T},
    MOI.Semiinteger{T},
}

# 0xcb = 0x0080 | 0x0040 | 0x0008 | 0x0002 | 0x0001
const _LOWER_BOUND_MASK = 0x00cb
# 0xcd = 0x0080 | 0x0040 | 0x0008 | 0x0004 | 0x0001
const _UPPER_BOUND_MASK = 0x00cd

const _DELETED_VARIABLE = 0x8000

_single_variable_flag(::Type{<:MOI.EqualTo}) = 0x0001
_single_variable_flag(::Type{<:MOI.GreaterThan}) = 0x0002
_single_variable_flag(::Type{<:MOI.LessThan}) = 0x0004
_single_variable_flag(::Type{<:MOI.Interval}) = 0x0008
_single_variable_flag(::Type{MOI.Integer}) = 0x0010
_single_variable_flag(::Type{MOI.ZeroOne}) = 0x0020
_single_variable_flag(::Type{<:MOI.Semicontinuous}) = 0x0040
_single_variable_flag(::Type{<:MOI.Semiinteger}) = 0x0080

function _flag_to_set_type(flag::UInt16, ::Type{T}) where {T}
    if flag == 0x0001
        return MOI.EqualTo{T}
    elseif flag == 0x0002
        return MOI.GreaterThan{T}
    elseif flag == 0x0004
        return MOI.LessThan{T}
    elseif flag == 0x0008
        return MOI.Interval{T}
    elseif flag == 0x0010
        return MOI.Integer
    elseif flag == 0x0020
        return MOI.ZeroOne
    elseif flag == 0x0040
        return MOI.Semicontinuous{T}
    else
        @assert flag == 0x0080
        return MOI.Semiinteger{T}
    end
end

# Julia doesn't infer `S1` correctly, so we use a function barrier to improve
# inference.
function _throw_if_lower_bound_set_inner(variable, S2, mask, T)
    S1 = _flag_to_set_type(mask, T)
    throw(MOI.LowerBoundAlreadySet{S1,S2}(variable))
    return
end

function _throw_if_lower_bound_set(variable, S2, mask, T)
    lower_mask = mask & _LOWER_BOUND_MASK
    if iszero(lower_mask)
        return  # No lower bound set.
    elseif iszero(_single_variable_flag(S2) & _LOWER_BOUND_MASK)
        return  # S2 isn't related to the lower bound.
    end
    return _throw_if_lower_bound_set_inner(variable, S2, lower_mask, T)
end

# Julia doesn't infer `S1` correctly, so we use a function barrier to improve
# inference.
function _throw_if_upper_bound_set_inner(variable, S2, mask, T)
    S1 = _flag_to_set_type(mask, T)
    throw(MOI.UpperBoundAlreadySet{S1,S2}(variable))
    return
end

function _throw_if_upper_bound_set(variable, S2, mask, T)
    upper_mask = mask & _UPPER_BOUND_MASK
    if iszero(upper_mask)
        return  # No upper bound set.
    elseif iszero(_single_variable_flag(S2) & _UPPER_BOUND_MASK)
        return  # S2 isn't related to the upper bound.
    end
    return _throw_if_upper_bound_set_inner(variable, S2, upper_mask, T)
end

function _lower_bound(
    set::Union{MOI.GreaterThan,MOI.Interval,MOI.Semicontinuous,MOI.Semiinteger},
)
    return set.lower
end

_lower_bound(set::MOI.EqualTo) = set.value

function _upper_bound(
    set::Union{MOI.LessThan,MOI.Interval,MOI.Semicontinuous,MOI.Semiinteger},
)
    return set.upper
end

_upper_bound(set::MOI.EqualTo) = set.value

# Use `-Inf` and `Inf` for `AbstractFloat` subtypes.
_no_lower_bound(::Type{T}) where {T} = zero(T)
_no_lower_bound(::Type{T}) where {T<:AbstractFloat} = typemin(T)
_no_upper_bound(::Type{T}) where {T} = zero(T)
_no_upper_bound(::Type{T}) where {T<:AbstractFloat} = typemax(T)

###
### VariablesContainer
###
### For use in MOI.Utilities.Model
###

"""
    struct VariablesContainer{T} <: AbstractVectorBounds
        set_mask::Vector{UInt16}
        lower::Vector{T}
        upper::Vector{T}
    end

A struct for storing variables and VariableIndex-related constraints. Used in
`MOI.Utilities.Model` by default.
"""
mutable struct VariablesContainer{T} <: AbstractVectorBounds
    set_mask::Vector{UInt16}
    lower::Vector{T}
    upper::Vector{T}
end

function VariablesContainer{T}() where {T}
    return VariablesContainer{T}(UInt16[], T[], T[])
end

function Base.:(==)(a::VariablesContainer, b::VariablesContainer)
    return a.set_mask == b.set_mask && a.lower == b.lower && a.upper == b.upper
end

function MOI.empty!(b::VariablesContainer)
    empty!(b.set_mask)
    empty!(b.lower)
    empty!(b.upper)
    return b
end

function MOI.is_empty(b::VariablesContainer)
    if length(b.set_mask) == 0
        return true
    end
    return all(isequal(_DELETED_VARIABLE), b.set_mask)
end

function Base.resize!(b::VariablesContainer, n)
    resize!(b.set_mask, n)
    resize!(b.lower, n)
    resize!(b.upper, n)
    return
end

function MOI.add_variable(b::VariablesContainer{T}) where {T}
    push!(b.set_mask, 0x0000)
    push!(b.lower, _no_lower_bound(T))
    push!(b.upper, _no_upper_bound(T))
    x = MOI.VariableIndex(length(b.set_mask))
    return x
end

function MOI.get(b::VariablesContainer, ::MOI.ListOfVariableIndices)
    return MOI.VariableIndex[
        MOI.VariableIndex(i) for
        i in 1:length(b.set_mask) if b.set_mask[i] != _DELETED_VARIABLE
    ]
end

function MOI.is_valid(b::VariablesContainer, x::MOI.VariableIndex)
    mask = get(b.set_mask, x.value, _DELETED_VARIABLE)
    return mask != _DELETED_VARIABLE
end

function MOI.get(b::VariablesContainer, ::MOI.NumberOfVariables)::Int64
    if length(b.set_mask) == 0
        return 0
    end
    return sum(x != _DELETED_VARIABLE for x in b.set_mask)
end

function MOI.supports_constraint(
    ::VariablesContainer{T},
    ::Type{MOI.VariableIndex},
    ::Type{<:SUPPORTED_VARIABLE_SCALAR_SETS{T}},
) where {T}
    return true
end

function MOI.add_constraint(
    b::VariablesContainer{T},
    f::MOI.VariableIndex,
    set::S,
) where {T,S<:SUPPORTED_VARIABLE_SCALAR_SETS}
    flag = _single_variable_flag(S)
    mask = b.set_mask[f.value]
    _throw_if_lower_bound_set(f, S, mask, T)
    _throw_if_upper_bound_set(f, S, mask, T)
    if !iszero(flag & _LOWER_BOUND_MASK)
        b.lower[f.value] = _lower_bound(set)
    end
    if !iszero(flag & _UPPER_BOUND_MASK)
        b.upper[f.value] = _upper_bound(set)
    end
    b.set_mask[f.value] = mask | flag
    return MOI.ConstraintIndex{MOI.VariableIndex,S}(f.value)
end

function MOI.delete(
    b::VariablesContainer{T},
    ci::MOI.ConstraintIndex{MOI.VariableIndex,S},
) where {T,S}
    MOI.throw_if_not_valid(b, ci)
    flag = _single_variable_flag(S)
    b.set_mask[ci.value] &= ~flag
    if !iszero(flag & _LOWER_BOUND_MASK)
        b.lower[ci.value] = _no_lower_bound(T)
    end
    if !iszero(flag & _UPPER_BOUND_MASK)
        b.upper[ci.value] = _no_upper_bound(T)
    end
    return
end

function MOI.delete(b::VariablesContainer, x::MOI.VariableIndex)
    MOI.throw_if_not_valid(b, x)
    b.set_mask[x.value] = _DELETED_VARIABLE
    return
end

function MOI.is_valid(
    b::VariablesContainer,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,S},
) where {S}
    if !(1 <= ci.value <= length(b.set_mask))
        return false
    end
    return !iszero(b.set_mask[ci.value] & _single_variable_flag(S))
end

function MOI.get(
    model::VariablesContainer,
    ::MOI.ConstraintFunction,
    ci::CI{MOI.VariableIndex},
)
    MOI.throw_if_not_valid(model, ci)
    return MOI.VariableIndex(ci.value)
end

function MOI.set(
    ::VariablesContainer,
    ::MOI.ConstraintFunction,
    ::MOI.ConstraintIndex{MOI.VariableIndex},
    ::MOI.VariableIndex,
)
    return throw(MOI.SettingVariableIndexNotAllowed())
end

function MOI.get(
    model::VariablesContainer,
    ::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,S},
) where {S}
    MOI.throw_if_not_valid(model, ci)
    return set_from_constants(model, S, ci.value)
end

function MOI.set(
    model::VariablesContainer{T},
    ::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex{MOI.VariableIndex,S},
    set::S,
) where {T,S<:SUPPORTED_VARIABLE_SCALAR_SETS{T}}
    MOI.throw_if_not_valid(model, ci)
    flag = _single_variable_flag(S)
    if !iszero(flag & _LOWER_BOUND_MASK)
        model.lower[ci.value] = _lower_bound(set)
    end
    if !iszero(flag & _UPPER_BOUND_MASK)
        model.upper[ci.value] = _upper_bound(set)
    end
    return
end

function MOI.get(
    b::VariablesContainer,
    ::MOI.NumberOfConstraints{MOI.VariableIndex,S},
)::Int64 where {S}
    flag = _single_variable_flag(S)
    return count(mask -> !iszero(flag & mask), b.set_mask)
end

function _add_constraint_type(
    list,
    b::VariablesContainer,
    S::Type{<:MOI.AbstractScalarSet},
)
    flag = _single_variable_flag(S)::UInt16
    if any(mask -> !iszero(flag & mask), b.set_mask)
        push!(list, (MOI.VariableIndex, S))
    end
    return
end

function MOI.get(
    b::VariablesContainer{T},
    ::MOI.ListOfConstraintTypesPresent,
) where {T}
    list = Tuple{Type,Type}[]
    _add_constraint_type(list, b, MOI.EqualTo{T})
    _add_constraint_type(list, b, MOI.GreaterThan{T})
    _add_constraint_type(list, b, MOI.LessThan{T})
    _add_constraint_type(list, b, MOI.Interval{T})
    _add_constraint_type(list, b, MOI.Semicontinuous{T})
    _add_constraint_type(list, b, MOI.Semiinteger{T})
    _add_constraint_type(list, b, MOI.Integer)
    _add_constraint_type(list, b, MOI.ZeroOne)
    return list
end

function MOI.get(
    b::VariablesContainer,
    ::MOI.ListOfConstraintIndices{MOI.VariableIndex,S},
) where {S}
    list = MOI.ConstraintIndex{MOI.VariableIndex,S}[]
    flag = _single_variable_flag(S)
    for (index, mask) in enumerate(b.set_mask)
        if !iszero(mask & flag)
            push!(list, MOI.ConstraintIndex{MOI.VariableIndex,S}(index))
        end
    end
    return list
end

# These `MOI.ModelLike` fallbacks have to be redefined here as
# `VariablesContainer` is not a subtype of `MOI.ModelLike`

function MOI.throw_if_not_valid(b::VariablesContainer, index)
    if !MOI.is_valid(b, index)
        throw(MOI.InvalidIndex(index))
    end
end

function MOI.set(
    ::VariablesContainer,
    ::MOI.ConstraintFunction,
    ci::MOI.ConstraintIndex,
    func,
)
    return throw(MOI.FunctionTypeMismatch{MOI.func_type(ci),typeof(func)}())
end

function MOI.set(
    ::VariablesContainer,
    ::MOI.ConstraintSet,
    ci::MOI.ConstraintIndex{MOI.VariableIndex},
    set,
)
    return throw(MOI.SetTypeMismatch{MOI.set_type(ci),typeof(set)}())
end

function MOI.supports_constraint(
    ::VariablesContainer,
    ::Type{<:MOI.AbstractFunction},
    ::Type{<:MOI.AbstractSet},
)
    return false
end

function MOI.add_constraint(
    ::VariablesContainer,
    func::MOI.AbstractFunction,
    set::MOI.AbstractSet,
)
    return throw(MOI.UnsupportedConstraint{typeof(func),typeof(set)}())
end

###
### Hyperrectangle
###

"""
    struct Hyperrectangle{T} <: AbstractVectorBounds
        lower::Vector{T}
        upper::Vector{T}
    end

A struct for the .constants field in MatrixOfConstraints.
"""
struct Hyperrectangle{T} <: AbstractVectorBounds
    lower::Vector{T}
    upper::Vector{T}
end

Hyperrectangle{T}() where {T} = Hyperrectangle{T}(T[], T[])

function Base.:(==)(a::Hyperrectangle, b::Hyperrectangle)
    return a.lower == b.lower && a.upper == b.upper
end

function Base.empty!(b::Hyperrectangle)
    empty!(b.lower)
    empty!(b.upper)
    return b
end

function Base.resize!(b::Hyperrectangle, n)
    resize!(b.lower, n)
    resize!(b.upper, n)
    return
end

function load_constants(
    b::Hyperrectangle{T},
    offset,
    set::SUPPORTED_VARIABLE_SCALAR_SETS{T},
) where {T}
    flag = _single_variable_flag(typeof(set))
    if iszero(flag & _LOWER_BOUND_MASK)
        b.lower[offset+1] = _no_lower_bound(T)
    else
        b.lower[offset+1] = _lower_bound(set)
    end
    if iszero(flag & _UPPER_BOUND_MASK)
        b.upper[offset+1] = _no_upper_bound(T)
    else
        b.upper[offset+1] = _upper_bound(set)
    end
    return
end

function_constants(::Hyperrectangle{T}, row) where {T} = zero(T)


using DataFrames
using JLD2
using Base.Test

opt = get_options()

info("Load metadata from excel table stations_metadata.xlsx")

df_sel = load_metadata(opt["stat_sel"])

info("Clean the metadata table and save to stat_selected.csv")

df_nice = clean_metadata(df_sel)

writetable(Pkg.dir("IcanProj", "data", "stat_selected.csv"), df_nice)

info("Get data and save in WatershedData struct to wsh_info.jld2")

wsh_info = get_watershed_data(df_sel)

@save Pkg.dir("IcanProj", "data", "wsh_info.jld2") wsh_info

# Run some basic tests on wsh_info

for wsh_single in wsh_info

    @show wsh_single.name

    @test all(wsh_single.elev .>= 0)
    @test all(wsh_single.elev .<= 3000)
    @test size(wsh_single.gridcel) == (50, 50)

end


using Plots
gr(xlims=[-1,6],ylims=[-4,6],border=:none)

plot([0,1],-[4,4],fill=(0,:royalblue1),fillalpha=0.7)
plot!([1,2],-[4,4],fill=(0,:darkorange1),fillalpha=0.7)
plot!([4,5],[3,3],fill=(0,:darkorange1),fillalpha=0.7)
plot!([4,5],[2,2],fill=(0,:white),fillalpha=1)
plot!([-0.1,2.1],-[2,2],fill=(0,:white),lw=0.1,fillalpha=1)
plot!([0,5],[0,0],seriescolor=:black,lw=3,legend=nothing)

for i = 1:5
    plot!([0,5],[i,i],seriescolor=:black,lw=3,legend=nothing)
end
for i = 0:5
    plot!([i,i],[0,5],seriescolor=:black,lw=3,legend=nothing)
end

for i = 0:5
    plot!([i,i],[-2,-4],seriescolor=:black,lw=3,legend=nothing)
end
for i = 2:4
    plot!([0,5],[-i,-i],seriescolor=:black,lw=3,legend=nothing)
end

current()

savefig("GPU_3.png")


__precompile__()

module FEBOLPlots

using FEBOL

using PyPlot:
    imshow,
    xlabel,
    ylabel,
    contour,
    figure,
    pause,
    hold,
    axis,
    title,
    scatter,
    gcf,
    savefig,
    matplotlib,
    rc,
    tick_params,
    clf,
    cla

using PyPlot
import PyPlot.plot

using Reel      # for gifs

#rc("font", family="serif")
#rc("text", usetex=true)
rc("font", family="Times New Roman", size=16)
println("rc_context = ", rc_context())

export
    visualize,
    gif,
    plot_vehicle,
    plot,
    hold,
    title,
    pause,
    figure,
    cla,
    gcf,
    savefig,
    Frames

const LocTuple = NTuple{2, Float64}

include("visualize.jl")
include("plotting.jl")
include("gif.jl")
include("special.jl")

end # module


module Rounding
include("fenv_constants.jl")

export
    RoundingMode, RoundNearest, RoundToZero, RoundUp, RoundDown, RoundFromZero,
    get_rounding, set_rounding, with_rounding

## rounding modes ##
immutable RoundingMode
    code::Int
    RoundingMode(c::Integer) = new(c)
end
const RoundNearest = RoundingMode(0)
const RoundToZero = RoundingMode(1)
const RoundUp = RoundingMode(2)
const RoundDown = RoundingMode(3)
const RoundFromZero = RoundingMode(4)

function to_fenv(r::RoundingMode)
    if r === RoundNearest
        JL_FE_TONEAREST
    elseif r === RoundToZero
        JL_FE_TOWARDZERO
    elseif r === RoundUp
        JL_FE_UPWARD
    elseif r === RoundDown
        JL_FE_DOWNWARD
    elseif r === RoundFromZero
        error("unsupported rounding mode")
    else
        error("invalid rounding mode")
    end
end

function from_fenv(r::Integer)
    if r == JL_FE_TONEAREST
        return RoundNearest
    elseif r == JL_FE_DOWNWARD
        return RoundDown
    elseif r == JL_FE_UPWARD
        return RoundUp
    elseif r == JL_FE_TOWARDZERO
        return RoundToZero
    else
        error("invalid rounding mode code")
    end
end

set_rounding{T<:Union(Float32,Float64)}(::Type{T},r::RoundingMode) = ccall(:fesetround, Cint, (Cint,), to_fenv(r))
get_rounding{T<:Union(Float32,Float64)}(::Type{T}) = from_fenv(ccall(:fegetround, Cint, ()))

function with_rounding{T}(f::Function, ::Type{T}, rounding::RoundingMode)
    old_rounding = get_rounding(T)
    set_rounding(T,rounding)
    try
        return f()
    finally
        set_rounding(T,old_rounding)
    end
end

end #module


using Printf

Jstr(J::Rational) = isone(denominator(J)) ? "$(numerator(J))" : "$(numerator(J))/$(denominator(J))"
Jstr(J::Integer) = "$J"
Jstr(J::AbstractString) = J
Jstr(J::AbstractVector) = join(Jstr.(J), ",")

function term_str(term, J)
    i = findfirst(!isnumeric, term)
    mult = i > 1 ? parse(Int, term[1:i-1]) : 0
    isoddterm = term[end] == '*'
    termsym = isoddterm ? term[i:end-1] : term[i:end]

    terms = (i > 1 ? "\$^{$mult}\$$(termsym)" : termsym)*(isoddterm ? "\$^{\\textrm{o}}\$" : "")
    ismissing(J) ? terms : "$(terms)\$_{$(Jstr(J))}\$"
end

function cfg_str(cfg)
    subshells = map(split(cfg, ".")) do s
        if s[1] == '('
            i = findfirst(isequal('<'), s)
            term = strip(isnothing(i) ? s[2:end-1] : s[2:i-1])
            J = isnothing(i) ? missing : strip(s[i+1:end-2])
            "($(term_str(term, J)))"
        else
            i = findlast(!isnumeric, s)
            i < length(s) ? "$(s[1:i])\$^{$(s[i+1:end])}\$" : s
        end
    end
    join(subshells, " ")
end

limit_str(::Missing) = ""

function limit_str(cfg)
    m = match(r"(.*)\((.*)\)", cfg)

    if isnothing(m)
        cfg
    else
        a,b = split(m[2], " ")
        cfgs = cfg_str("$(a).($(b))")
        "$(m[1]) [$(cfg_str(cfgs))]"
    end
end

function latex_table(io::IO, data, name; offset=0u"eV")
    no_offset = iszero(NoUnits(offset/u"eV"))

    write(io, "\\section{$name}\n\\footnotesize\n")
    write(io, "\\tablehead{%\n")
    write(io, "\\textbf{Configuration}&\\textbf{Term}&\\textbf{Energy}")
    no_offset || write(io, "&\\textbf{Offset energy}")
    write(io, "\\\\\n\\hline\n")
    write(io, "}\n")
    write(io, "\\begin{supertabular}{l|l|l$(no_offset ? "" : "|l")}\n")

    cfgs = data[!,1]
    terms = data[!,2]
    Js = data[!,3]
    Es = data[!,4]
    for (cfg,term,J,E) in zip(cfgs,terms,Js,Es)
        Eoff = NoUnits((E + offset)/u"eV")
        ismissing(term) && continue
        if term == "Limit"
            write(io, "\\hline\n")
            write(io, @sprintf("%s & Limit & %08.5f eV", limit_str(cfg), NoUnits(E/u"eV")))
            no_offset || write(io, @sprintf("& %08.5f eV", Eoff))
            write(io, "\\\\\n")
            continue
        end
        write(io, @sprintf("%s & %s & %08.5f eV", cfg_str(cfg), term_str(term,J), NoUnits(E/u"eV")))
            no_offset || write(io, @sprintf("& %08.5f eV", Eoff))
        write(io, "\\\\\n")

    end

    write(io, "\\end{supertabular}\n\n")
end

function latex_table(fun::Function, filename::String)
    open(filename, "w") do file
        write(file, raw"""\documentclass[11pt]{article}
\usepackage[margin=2cm]{geometry}
\usepackage{multicol}
\usepackage{supertabular}
\usepackage{amsmath}
\usepackage{amssymb}
\begin{document}
\twocolumn

""")
        fun(file)
        write(file, "\n\n\\end{document}\n")
    end
end

latex_table(filename::String, data, name; kwargs...) =
    latex_table(file -> latex_table(file, data, name; kwargs...), filename)

export latex_table


@testset "Filter" begin
    # Defining our missing datasets
    a = allowmissing(1.0:1.0:20.0)
    a[[2, 3, 7]] .= missing
    mask = map(!ismissing, a)

    # We call collect to not have a wrapper type that references the same data.
    m = collect(reshape(a, 5, 4))
    m_no_missing = reshape(1.0:1.0:20.0, 5, 4)

    aa = AxisArray(
        deepcopy(m),
        Axis{:time}(DateTime(2017, 6, 5, 5):Hour(1):DateTime(2017, 6, 5, 9)),
        Axis{:id}(1:4)
    )

    table = DataFrame(
        :sin => allowmissing(sin.(1.0:1.0:20.0)),
        :cos => allowmissing(sin.(1.0:1.0:20.0)),
    )

    table.sin[[2, 3, 7, 12, 19]] .= missing

    @test Filter() == Filter()

    @testset "Vector" begin
        result = apply(a, Filter())
        expected = deleteat!(deepcopy(a), [2, 3, 7])

        @test result == expected
        @test result == Impute.filter(a)

        a2 = deepcopy(a)
        Impute.filter!(a2)
        @test a2 == expected
    end

    @testset "Matrix" begin
        @test_throws UndefKeywordError apply(m, Filter())
        @test_throws UndefKeywordError Impute.filter(m)
        @test_throws MethodError Impute.filter!(m)

        @testset "no missing" begin
            result = apply(m_no_missing, Filter())
            @test isequal(result, m_no_missing)

            result = apply(m_no_missing, Filter(); kw=:foo)  # kwargs unused
            @test isequal(result, m_no_missing)
        end

        @testset "rows" begin
            result = apply(m, Filter(); dims=:rows)
            expected = m[[1, 4, 5], :]

            @test isequal(result, expected)
            @test isequal(result, (Impute.filter(m; dims=:rows)))
            @test isequal(collect(result'), Impute.filter(collect(m'); dims=:cols))
        end

        @testset "cols" begin
            result = apply(m, Filter(); dims=:cols)
            expected = copy(m)[:, 3:4]

            @test isequal(result, expected)
            @test isequal(result, Impute.filter(m; dims=:cols))
            @test isequal(collect(result'), Impute.filter(collect(m'); dims=:rows))
        end
    end

    @testset "Tables" begin
        @testset "DataFrame" begin
            df = deepcopy(table)

            @test_throws UndefKeywordError apply(df, Filter())
            @test_throws UndefKeywordError Impute.filter(df)
            @test_throws MethodError Impute.filter!(df)

            @testset "rows" begin
                result = apply(df, Filter(); dims=:rows)
                expected = dropmissing(df)

                @test isequal(result, expected)
                @test isequal(result, Impute.filter(df; dims=:rows))
            end

            @testset "cols" begin
                result = apply(df, Filter(); dims=:cols)
                expected = select(df, :cos)

                @test isequal(result, expected)
                @test isequal(result, Impute.filter(df; dims=:cols))
            end
        end

        @testset "Column Table" begin
            coltab = Tables.columntable(table)

            @test_throws UndefKeywordError apply(coltab, Filter())
            @test_throws UndefKeywordError Impute.filter(coltab)
            @test_throws MethodError Impute.filter!(coltab)

            @testset "rows" begin
                result = apply(coltab, Filter(); dims=:rows)
                expected = Tables.columntable(dropmissing(table))

                @test isequal(result, expected)
                @test isequal(result, Impute.filter(coltab; dims=:rows))
            end

            @testset "cols" begin
                result = apply(coltab, Filter(); dims=:cols)
                expected = Tables.columntable(TableOperations.select(coltab, :cos))

                @test isequal(result, expected)
                @test isequal(result, Impute.filter(coltab; dims=:cols))
            end
        end

        @testset "Row Table" begin
            @testset "rows" begin
                rowtab = Tables.rowtable(table)

                result = apply(rowtab, Filter(); dims=:rows)
                expected = Tables.rowtable(dropmissing(table))

                @test isequal(result, expected)
                @test isequal(result, Impute.filter(rowtab; dims=:rows))
                @test isequal(result, Impute.filter(rowtab))

                rowtab_ = Impute.filter!(rowtab)
                @test isequal(rowtab, expected)
                @test isequal(rowtab_, expected)
            end

            @testset "cols" begin
                rowtab = Tables.rowtable(table)

                @test_throws ArgumentError Impute.filter!(rowtab; dims=:cols)
                result = apply(rowtab, Filter(); dims=:cols)
                expected = Tables.rowtable(TableOperations.select(rowtab, :cos))

                @test isequal(result, expected)
                @test isequal(result, Impute.filter(rowtab; dims=:cols))
            end

        end
    end

    @testset "AxisArray" begin
        @test_throws UndefKeywordError apply(aa, Filter())
        @test_throws UndefKeywordError Impute.filter(aa)
        @test_throws MethodError Impute.filter!(aa)

        @testset "rows" begin
            result = apply(aa, Filter(); dims=:rows)
            expected = aa[[1, 4, 5], :]

            @test isa(result, AxisArray)
            @test isequal(result, expected)
            @test isequal(result, Impute.filter(aa; dims=:rows))
            @test isequal(collect(result'), Impute.filter(collect(aa'); dims=:cols))
        end

        @testset "cols" begin
            result = apply(aa, Filter(); dims=:cols)
            expected = copy(aa)[:, 3:4]

            @test isa(result, AxisArray)
            @test isequal(result, expected)
            @test isequal(result, Impute.filter(aa; dims=:cols))
            @test isequal(collect(result'), Impute.filter(collect(aa'); dims=:rows))
        end
    end

    @testset "functional" begin
        expected = deleteat!(deepcopy(a), [2, 3, 7])
        @test Impute.filter(!ismissing, a) == expected

        b = deepcopy(a)
        @test Impute.filter!(!ismissing, b) == expected
        @test b == expected
    end
end


using InteractiveUtils
using InteractiveUtils: typesof
using Core: Typeof
import Base: copy!
import Base.Broadcast: broadcasted, materialize!

mutable struct Context <: AContext
  cache::Union{IdDict{Any,Any},Nothing}
end

Context() = Context(nothing)

cache(cx::Context) = cx.cache === nothing ? (cx.cache = IdDict()) : cx.cache

struct Pullback{S,T}
  t::T
end

Pullback{S}(x) where S = Pullback{S,typeof(x)}(x)

struct CompileError
  T
  e
end

function Base.showerror(io::IO, e::CompileError)
  print(io, "Compiling $(e.T): ")
  showerror(io, e.e)
end

# interface2.jl

# Wrappers
_pullback(f, args...) = _pullback(Context(), f, args...)

tailmemaybe(::Nothing) = nothing
tailmemaybe(x::Tuple) = Base.tail(x)

function pullback(f, args...)
  y, back = _pullback(f, args...)
  y, Δ -> tailmemaybe(back(Δ))
end

sensitivity(y::Number) = one(y)
sensitivity(y::Complex) = error("Output is complex, so the gradient is not defined.")
sensitivity(y::AbstractArray) = error("Output is an array, so the gradient is not defined. Perhaps you wanted jacobian.")
sensitivity(y) = error("Output should be scalar; gradients are not defined for output $(repr(y))")

"""
    gradient(f, args...)

Returns a tuple containing `∂f/∂x` for each argument `x`,
the derivative (for scalar `x`) or the gradient.

`f(args...)` must be a real number, see [`jacobian`](@ref) for array output.

See also [`withgradient`](@ref) to keep the value `f(args...)`,
and [`pullback`](@ref) for value and back-propagator.

```jldoctest; setup=:(using Zygote)
julia> gradient(*, 2.0, 3.0, 5.0)
(15.0, 10.0, 6.0)

julia> gradient(x -> sum(abs2,x), [7.0, 11.0, 13.0])
([14.0, 22.0, 26.0],)

julia> gradient([7, 11], 0, 1) do x, y, d
         p = size(x, d)
         sum(x.^p .+ y)
       end
([14.0, 22.0], 2.0, nothing)
```
"""
function gradient(f, args...)
  y, back = pullback(f, args...)
  grad = back(sensitivity(y))
  isnothing(grad) ? nothing : map(_project, args, grad)
end

# Base.adjoint(f::Function) = x -> gradient(f, x)[1]  # piracy!
Base.adjoint(f::Function) = x -> begin  # still piracy! avoids projection for legacy reasons
  y, back = pullback(f, x)
  back(sensitivity(y))[1]
end

"""
    withgradient(f, args...)
    withgradient(f, ::Params)

Returns both the value of the function and the [`gradient`](@ref),
as a named tuple.

```jldoctest; setup=:(using Zygote)
julia> y, ∇ = withgradient(/, 1, 2)
(val = 0.5, grad = (0.5, -0.25))

julia> ∇ == gradient(/, 1, 2)
true
```
"""
function withgradient(f, args...)
  y, back = pullback(f, args...)
  grad = back(sensitivity(y))
  results = isnothing(grad) ? map(_ -> nothing, args) : map(_project, args, grad)
  (val=y, grad=results)
end

# Param-style wrappers

"""
    gradient(() -> loss(), ps::Params) -> Grads

Gradient with implicit parameters. Takes a zero-argument function,
and returns a dictionary-like container, whose keys are arrays `x in ps`.

```jldoctest; setup=:(using Zygote)
julia> x = [1 2 3; 4 5 6]; y = [7, 8]; z = [1, 10, 100];

julia> g = gradient(Params([x, y])) do
         sum(x .* y .* z')
       end
Grads(...)

julia> g[x]
2×3 Matrix{Float64}:
 7.0  70.0  700.0
 8.0  80.0  800.0

julia> haskey(g, z)  # only x and y are parameters
false
```
"""
gradient

"""
    Params([A, B])

Container for implicit parameters, used when differentiating
a zero-argument funtion `() -> loss(A, B)` with respect to `A, B`.
"""
struct Params
  order::Buffer # {Any, Vector{Any}}
  params::IdSet{Any} # TODO store ids only
end

Params() = Params(Buffer([], false), IdSet())
Params(xs) = Params(Buffer(xs, false), IdSet(xs))
Params(ps::Params) = ps
Params(xs::Tuple) = Params(collect(xs))

@forward Params.order Base.iterate, Base.length, Base.getindex
@forward Params.params Base.in

Base.map(::typeof(_project), args::Tuple{Params}, grad) = grad  # skip _project in gradient(f, ::Params)

function Base.union!(ps::Params, itrs...)
  foreach(itr -> foreach(x -> push!(ps, x), itr), itrs)
  return ps
end

Base.copy(ps::Params) = union!(Params(), ps)
Base.union(ps::Params, itrs...) = union!(copy(ps), itrs...)
Base.issetequal(ps1::Params, ps2::Params) = issetequal(ps1.params, ps2.params)
Base.issetequal(ps1::Params, x::Base.AbstractSet) = issetequal(ps1.params, x)
Base.issetequal(x::Base.AbstractSet, ps1::Params) = issetequal(x, ps1.params)

function Base.intersect!(ps::Params, itrs...)
  for itr in itrs
    for x in collect(ps)
      x ∉ itr && delete!(ps, x)
    end
  end
  return ps
end

Base.intersect(ps::Params, itrs...) = intersect!(copy(ps), itrs...)

function Base.push!(ps::Params, x)
  if !(x in ps.params)
    push!(ps.order, x)
    push!(ps.params, x)
  end
  return ps
end

Base.push!(ps::Params, x...) = (foreach(x -> push!(ps, x), x); ps)

function Base.delete!(ps::Params, x)
  if x in ps.params
    delete!(ps.params, x)
    i = findfirst(y -> y === x, ps.order)
    deleteat!(ps.order, i)
  end
  return ps
end

Base.Broadcast.broadcasted(f, ps::Params) = broadcasted(f, ps.order)

@adjoint function Broadcast.broadcasted(f::Function, ps::Params)
  f.(ps), _ -> throw(ArgumentError("Zygote.Params does not support broadcasting within gradients, try iteration `for p in ps`"))
end

Base.:(==)(x::Params, y::Params) = x.order.data == y.order.data

function Base.show(io::IO, ps::Params)
  print(io, "Params([")
  join(io, ps.order, ", ")
  print(io, "])")
end

"""
    copy!(ps::Params, x::AbstractVector)
    copy!(x::AbstractVector, ps::Params)

Copies the content of array `x` into the parameters `ps` or viceversa.
The length of `x` has to be equal to the sum of the lengths
of all parameters.
"""
function copy!(ps::Params, x::AbstractVector)
  @assert length(x) == sum(length(p) for p in ps)
  i = 0
  for p in ps
    p .= reshape(x[i+1:i+length(p)], size(p))
    i += length(p)
  end
  ps
end

function copy!(x::AbstractVector, ps::Params)
  @assert length(x) == sum(length(p) for p in ps)
  i = 0
  for p in ps
    x[i+1:i+length(p)] .= vec(p)
    i += length(p)
  end
  ps
end

"""
    Grads(...)

Dictionary-like container returned when taking gradients with
respect to implicit parameters. For an array `W`, appearing
within `Params([W, A, B...])`, the gradient is `g[W]`.
"""
struct Grads
  grads::IdDict{Any,Any}
  params::Params
end

Base.show(io::IO, ps::Grads) = print(io, "Grads(...)")

@forward Grads.grads  Base.setindex!
@forward Grads.params  Base.length

const ADictOrGrads = Union{AbstractDict, Grads}

# Dictionary interface.
# Don't use the IdDict directly since it may contain some spurious pairs.
Base.haskey(gs::Grads, x) = x ∈ gs.params
Base.keys(gs::Grads) = gs.params
Base.values(gs::Grads) = (gs.grads[p] for p in gs.params)

function Base.iterate(gs::Grads, state...)
  res = iterate(gs.params, state...)
  isnothing(res) && return nothing
  p, next_state = res
  return gs[p], next_state
end

function Base.getindex(gs::Grads, x)
  isbits(x) && error("Only reference types can be differentiated with `Params`.")
  return gs.grads[x]
end

"""
    copy!(gs::Grads, x::AbstractVector)
    copy!(x::AbstractVector, gs::Grads)

Copies the content of array `x` into the gradient object `gs` or vice versa. The
length of `x` has to be equal to the sum of the lengths of all gradients.
"""
function copy!(gs::Grads, x::AbstractVector)
  i = 0
  for p in gs.params
    gs[p] .= reshape(x[i+1:i+length(p)], size(p))
    i += length(p)
  end
  x
end

function copy!(x::AbstractVector,  gs::Grads)
  i = 0
  for p in gs.params
    x[i+1:i+length(p)] .= vec(gs[p])
    i += length(p)
  end
  x
end

function Base.merge!(gs_dst::Grads, gs_srcs::Grads...)
  for gs_src in gs_srcs
    union!(gs_dst.params, gs_src.params)
    merge!(gs_dst.grads, gs_src.grads)
  end
  gs_dst
end

function Base.copy(gs::Grads)
  gs_new = Grads(IdDict(), gs.params)
  merge!(gs_new, gs)
end

broadcasted(f, gs::Grads, gss::ADictOrGrads...) = map(f, gs, gss...)

broadcasted(f, a::Numeric, gs::Grads) = map(x -> f(a, x), gs)
broadcasted(f, gs::Grads, a::Numeric) = map(x -> f(x, a), gs)

function materialize!(gs1::Grads, gs2::Grads)
  issetequal(gs1.params, gs2.params) ||
    throw(ArgumentError("Expected Grads objects with the same Params."))
  for p in gs1.params
    gs1[p] = gs2[p]
  end
  return gs1
end

function Base.map(f, gs1::Grads, gss::ADictOrGrads...)
  gsout = Grads(IdDict{Any,Any}(), Params(gs1.params))
  return map!(f, gsout, gs1, gss...)
end

function Base.map!(f, gsout::Grads, gss::ADictOrGrads...)
  all(issetequal(gsout.params, keys(gs)) for gs in gss) ||
    throw(ArgumentError("map! expects Grads objects with the same Params."))
  for p in gsout.params
    gsout[p] = f((_getformap(gs, p) for gs in gss)...)
  end
  return gsout
end

function _getformap(gs, p)
  g = gs[p]
  isnothing(g) ? fill!(similar(p), 0) : g
end

function pullback(f, ps::Params)
  cx = Context()
  y, back = _pullback(cx, f)
  y, function (Δ)
    for p in ps
      cache(cx)[p] = nothing
    end
    back(Δ)
    Grads(cx.cache, ps) # TODO make a copy
  end
end

# Code Reflection

function code_ir(f, T)
  m = meta(Tuple{Typeof(f),T.parameters...})
  return IR(m)
end

function code_irm(ex)
  isexpr(ex, :call) || error("@code_ir f(args...)")
  f, args = ex.args[1], ex.args[2:end]
  :(code_ir($(esc(f)), typesof($(esc.(args)...))))
end

macro code_ir(ex)
  code_irm(ex)
end

macro code_adjoint(ex)
  :(Adjoint($(code_irm(ex)), varargs = varargs($(esc(:($InteractiveUtils.@which $ex))), length(($(esc.(ex.args)...),)))))
end


function allocate_coefficients(mesh::Union{TreeMesh, CurvedMesh, UnstructuredQuadMesh}, equations, dg::DG, cache)
  # We must allocate a `Vector` in order to be able to `resize!` it (AMR).
  # cf. wrap_array
  zeros(eltype(cache.elements), nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache))
end

@inline function wrap_array(u_ode::AbstractVector, mesh::Union{TreeMesh{1},CurvedMesh{1}}, equations, dg::DG, cache)
  @boundscheck begin
    @assert length(u_ode) == nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache)
  end
  # We would like to use
  #   reshape(u_ode, (nvariables(equations), nnodes(dg), nnodes(dg), nelements(dg, cache)))
  # but that results in
  #   ERROR: LoadError: cannot resize array with shared data
  # when we resize! `u_ode` during AMR.

  # The following version is fast and allows us to `resize!(u_ode, ...)`.
  # OBS! Remember to `GC.@preserve` temporaries such as copies of `u_ode`
  #      and other stuff that is only used indirectly via `wrap_array` afterwards!
  unsafe_wrap(Array{eltype(u_ode), ndims(mesh)+2}, pointer(u_ode),
              (nvariables(equations), nnodes(dg), nelements(dg, cache)))
end

function compute_coefficients!(u, func, t, mesh::Union{TreeMesh{1},CurvedMesh{1}}, equations, dg::DG, cache)

  @threaded for element in eachelement(dg, cache)
    for i in eachnode(dg)
      x_node = get_node_coords(cache.elements.node_coordinates, equations, dg, i, element)
      u_node = func(x_node, t, equations)
      set_node_vars!(u, u_node, equations, dg, i, element)
    end
  end
end

@inline function wrap_array(u_ode::AbstractVector, mesh::Union{TreeMesh{2},CurvedMesh{2},UnstructuredQuadMesh}, equations, dg::DG, cache)
  @boundscheck begin
    @assert length(u_ode) == nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache)
  end
  # We would like to use
  #   reshape(u_ode, (nvariables(equations), nnodes(dg), nnodes(dg), nelements(dg, cache)))
  # but that results in
  #   ERROR: LoadError: cannot resize array with shared data
  # when we resize! `u_ode` during AMR.

  # The following version is fast and allows us to `resize!(u_ode, ...)`.
  # OBS! Remember to `GC.@preserve` temporaries such as copies of `u_ode`
  #      and other stuff that is only used indirectly via `wrap_array` afterwards!
  unsafe_wrap(Array{eltype(u_ode), ndims(mesh)+2}, pointer(u_ode),
              (nvariables(equations), nnodes(dg), nnodes(dg), nelements(dg, cache)))
end

function compute_coefficients!(u, func, t, mesh::Union{TreeMesh{2},CurvedMesh{2},UnstructuredQuadMesh}, equations, dg::DG, cache)

  @threaded for element in eachelement(dg, cache)
    for j in eachnode(dg), i in eachnode(dg)
      x_node = get_node_coords(cache.elements.node_coordinates, equations, dg, i, j, element)
      u_node = func(x_node, t, equations)
      set_node_vars!(u, u_node, equations, dg, i, j, element)
    end
  end
end

@inline function wrap_array(u_ode::AbstractVector, mesh::Union{TreeMesh{3},CurvedMesh{3}}, equations, dg::DG, cache)
  @boundscheck begin
    @assert length(u_ode) == nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache)
  end
  # We would like to use
  #   reshape(u_ode, (nvariables(equations), nnodes(dg), nnodes(dg), nnodes(dg), nelements(dg, cache)))
  # but that results in
  #   ERROR: LoadError: cannot resize array with shared data
  # when we resize! `u_ode` during AMR.

  # The following version is fast and allows us to `resize!(u_ode, ...)`.
  # OBS! Remember to `GC.@preserve` temporaries such as copies of `u_ode`
  #      and other stuff that is only used indirectly via `wrap_array` afterwards!
  unsafe_wrap(Array{eltype(u_ode), ndims(mesh)+2}, pointer(u_ode),
              (nvariables(equations), nnodes(dg), nnodes(dg), nnodes(dg), nelements(dg, cache)))
end

function compute_coefficients!(u, func, t, mesh::Union{TreeMesh{3},CurvedMesh{3}}, equations, dg::DG, cache)

  @threaded for element in eachelement(dg, cache)
    for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
      x_node = get_node_coords(cache.elements.node_coordinates, equations, dg, i, j, k, element)
      u_node = func(x_node, t, equations)
      set_node_vars!(u, u_node, equations, dg, i, j, k, element)
    end
  end
end


module TestModelSearch

# using Revise
using Test
using MLJModels
import MLJBase.info

pca = info("PCA", pkg="MultivariateStats")
cnst = info("ConstantRegressor", pkg="MLJModels")

@test_throws ArgumentError info("Julia")

@info(ConstantRegressor) == cnst
@test info(Standardizer()) == info("Standardizer", pkg="MLJModels")

@testset "localmodels" begin
    tree = info("DecisionTreeRegressor")
    @test cnst in localmodels(modl=TestModelSearch)
    @test !(tree in localmodels(modl=TestModelSearch))
    import MLJModels
    import DecisionTree
    import MLJModels.DecisionTree_.DecisionTreeRegressor
    @test tree in localmodels(modl=TestModelSearch)
end

@testset "models() and localmodels()" begin
    t(model) = model.is_pure_julia
    mods = models(t)
    @test pca in mods
    @test cnst in mods
    @test !(info("SVC") in mods)
    mods = localmodels(t, modl=TestModelSearch)
    @test cnst in mods
    @test !(pca in mods)
    u(model) = !(model.is_supervised)
    @test pca in models(u, t)
    @test !(cnst in models(u, t))
end

end
true


"""
This module defines the following data types:
* `Geom` : acquisition geometry, i.e., positions of supersources, sources and receivers
* `Src` : source related acquisition parameters, e.g., source wavelet
It also provides methods that either does operations on these data type or
help their construction.
"""
module Acquisition

import GeoPhyInv.Models
import GeoPhyInv.Utils
using Distributions
using DataFrames
using Random
using SparseArrays
using CSV

"""
Acquisiton has supersources, sources and receivers.
Each supersource has `ns` multiple sources that are
injected (or active) simultaneously.
For each supersource,
a set of `nr` receivers are
recording waves.

# Fields

* `sx::Vector{Vector{Float64,1},1}` : ``x`` positions of sources
* `sz::Vector{Vector{Float64,1},1}` : ``z`` positions of sources
* `rx::Vector{Vector{Float64,1},1}` : ``x`` positions of receivers
* `rz::Vector{Vector{Float64,1},1}` : ``z`` positions of receivers
* `nss::Int64` : number of supersources
* `ns::Vector{Int64,1}` : number of sources for every supersource
* `nr::Vector{Int64,1}` : number of receivers for every supersource
"""
mutable struct Geom
	sx::Vector{Vector{Float64}}
	sz::Vector{Vector{Float64}}
	rx::Vector{Vector{Float64}}
	rz::Vector{Vector{Float64}}
	nss::Int64 # change to +ve only integer later
	ns::Vector{Int64} # change to +ve only integer later
	nr::Vector{Int64} # change to +ve only integer later
	"adding conditions that are to be false while construction"
	Geom(sx, sz, rx, rz, nss, ns, nr) =
		any([
                  broadcast(length, sx) != ns,
                  broadcast(length, sz) != ns,
                  broadcast(length, rx) != nr,
                  broadcast(length, rz) != nr,
       		  length(sx) != nss,
       		  length(sz) != nss,
       		  length(rz) != nss,
       		  length(rx) != nss
		  ]) ?
		error("Geom construct") : new(sx, sz, rx, rz, nss, ns, nr)
end # type

"Compare if two `TD`'s  are equal"
function Base.isequal(geom1::Geom, geom2::Geom, attrib=:all)
	srcfnames = [:sx, :sz, :nss, :ns]
	recfnames = [:rx, :rz, :nr, :nss]
	srcvec=[(isequal(getfield(geom1, name),getfield(geom2, name))) for name in srcfnames]
	recvec=[(isequal(getfield(geom1, name),getfield(geom2, name))) for name in recfnames]
	if(attrib == :sources)
		return all(srcvec)
	elseif(attrib == :receivers)
		return all(recvec)
	else
		return all(vcat(srcvec, recvec))
	end
end

"""
Print information about `Geom`
"""
function Base.print(geom::Geom, name::String="")
	println("\tAcquisition Geometry:\t",name)
	println("\t> number of supersources:\t",geom.nss)
	println("\t> sources per supersource:\t","min\t",minimum(geom.ns[:]), "\tmax\t", maximum(geom.ns[:]))
	println("\t> receivers per supersource:\t","min\t",minimum(geom.nr[:]), "\tmax\t", maximum(geom.nr[:]))
	println("\t> number of unique positions of:\t","sources\t",Geom_get([geom],:nus), "\treceivers\t", Geom_get([geom],:nur))
end

"""
Return some derived fields of `Geom`

# Arguments

* `acq::Vector{Geom}` : a vector of `Geom`
* `attrib::Symbol` : attribute to determine the return object
  * `=:nus` number of unique source positions in acquisition
  * `=:nur` number of unique receiver positions in acquisition
  * `=:uspos` a tuple of x and z positions of all the unique sources
  * `=:urpos` a tuple of x and z position of all the unique receivers
  * `=:geomurpos` a `Geom` vector as if all the unique receiver positions are used for each supersource
  * `=:geomuspos` a `Geom` vector as if all the unique source positions are used for each supersource
"""
function Geom_get(acq::Vector{Geom}, attrib::Symbol)
	ngeom = length(acq)
	sposx = Geom_getvec(acq,:sx);	sposz = Geom_getvec(acq,:sz)
	isequal(length(sposx), length(sposz)) ?
		spos = [[sposz[is],sposx[is]] for is in 1:length(sposx)] : error("input acq corrupt")

	rposx = Geom_getvec(acq,:rx);	rposz = Geom_getvec(acq,:rz)
	isequal(length(rposx), length(rposz)) ?
		rpos = [[rposz[ir],rposx[ir]] for ir in 1:length(rposx)] : error("input acq corrupt")

	uspos = unique(spos); nus=length(uspos)
	urpos = unique(rpos); nur=length(urpos)
	uspost = ([uspos[iu][1] for iu in 1:nus], [uspos[iu][2] for iu in 1:nus]);
        urpost = ([urpos[iu][1] for iu in 1:nur], [urpos[iu][2] for iu in 1:nur]);

	if(attrib == :nus)
		return nus
	elseif(attrib == :nur)
		return nur
	elseif(attrib == :uspos)
		return uspost
	elseif(attrib == :urpos)
		return urpost
	elseif(attrib == :geomurpos)
		return [Geom(acq[igeom].sx, acq[igeom].sz,
	       fill(urpost[2],acq[igeom].nss),
	       fill(urpost[1],acq[igeom].nss), acq[igeom].nss, acq[igeom].ns,
	       fill(nur,acq[igeom].nss)) for igeom=1:ngeom]
	elseif(attrib == :geomuspos)
		return [Geom(fill(uspost[2],acq[igeom].nss),
	       fill(uspost[1],acq[igeom].nss),
	       acq[igeom].rx, acq[igeom].rz,
	       acq[igeom].nss, fill(nus,acq[igeom].nss),
	       acq[igeom].nr) for igeom=1:ngeom]
	else
		error("invalid attrib")
	end
end

"""
Given receiver positions `rpos` and `rpos0`.
Returns an array Int indices of the dimension of number of supersources
with `true` at indices, if the waves due to that particular source are
recorded.
"""
function Geom_find(acq::Geom; rpos::Array{Float64,1}=nothing, rpos0::Array{Float64,1}=nothing)
	rpos==nothing ? error("need rpos") : nothing
	sson = Array{Vector{Int64}}(undef,acq.nss);
	for iss = 1:acq.nss
		rvec = [[acq.rz[iss][ir],acq.rx[iss][ir]] for ir=1:acq.nr[iss]]
		ir=findfirst(rvec, rpos)
		if(rpos0==nothing)
			ir0=1;
		else
			ir0=findfirst(rvec, rpos0);
		end
		sson[iss] = ((ir != 0) && (ir0 != 0)) ? [ir,ir0] : [0]
	end
	return sson
end

#=
"""
Modify input `Geom` such that the output `Geom` has
either sources or receivers on the boundary of
`mgrid`.

# Arguments

* `acqgeom::Geom` : input geometry
* `mgridi` : grid to determine the boundary
* `attrib::Symbol` : decide return
  * `=:srcborder` sources on boundary (useful for back propagation)
  * `=:recborder` receivers on boundary
"""
function Geom_boundary(acqgeom::Geom,
	      mgrid
	      attrib::Symbol
	     )

	if((attrib == :recborder) | (attrib == :srcborder))
		bz, bx, nb = border(mgrid, 3, :inner)
	end
	"change the position of receivers to the boundary"
	if(attrib == :recborder)
		nr = nb; 		sx = acqgeom.sx;
		sz = acqgeom.sz;		ns = acqgeom.ns;
		nss = acqgeom.nss;
		return Geom(sx, sz, fill(bx,nss), fill(bz,nss), nss, ns, fill(nr,nss))
	"change the position of source (not supersources) to the boundary for back propagation"
	elseif(attrib == :srcborder)
		ns = nb;		rx = acqgeom.rx;
		rz = acqgeom.rz;		nr = acqgeom.nr;
		nss = acqgeom.nss;
		return Geom(fill(bx,nss), fill(bz,nss), rx, rz, nss, fill(ns,nss), nr)
	else
		error("invalid attrib")
	end
end
=#

"""
Check if all the sources and receivers in `Geom` are within the model

# Return

* `true` if all the positions within the model, `false` otherwise
"""
function Geom_check(geom::Geom, mgrid)
	xmin, zmin, xmax, zmax = mgrid[2][1], mgrid[1][1], mgrid[2][end], mgrid[1][end]

	checkvec=fill(false, geom.nss)
	for iss=1:geom.nss
		checkvec[iss] = any(
		      vcat(
		     (((geom.sx[iss] .- xmin).*(xmax .- geom.sx[iss])) .< 0.0),
		     (((geom.sz[iss] .- zmin).*(zmax .- geom.sz[iss])) .< 0.0),
		     (((geom.rx[iss] .- xmin).*(xmax .- geom.rx[iss])) .< 0.0),
		     (((geom.rz[iss] .- zmin).*(zmax .- geom.rz[iss])) .< 0.0),
		     ) )
	end
	return !(any(checkvec))
end

"""
Check if the input acquisition geometry is fixed spread.
"""
function Geom_isfixed(geom::Geom)
	isfixed=false
	# can have different source positions, but also different number of sources?
	for field in [:rx, :rz, :nr, :ns]
		f=getfield(geom, field)
		if(all(f .== f[1:1]))
			isfixed=true
		else
			isfixed=false
			return isfixed
		end
	end
	return isfixed
end

"""
A fixed spread acquisition has same set of sources and
receivers for each supersource.
This method constructs a
fixed spread acquisition geometry using either a
horizontal or vertical array of supersources/ receivers.
Current implementation has only one source for every supersource.

# Arguments

* `ssmin::Float64` : minimum coordinate for sources
* `ssmax::Float64` : maximum coordinate for sources
* `ss0::Float64` : consant coordinate for sources
* `rmin::Float64` : minimum coordinate for receivers
* `rmax::Float64` : maximum coordinate for receivers
* `r0::Float64` : consant coordinate for receivers
* `nss::Int64` : number of supersources
* `nr::Int64` : number of receivers
* `ssattrib::Symbol=:horizontal` : supersource array kind
  `=:vertical` : vertical array of supersources
  `=:horizontal` horizontal array of supersources
* `rattrib::Symbol=:horizontal` : receiver array kind
  `=:vertical` : vertical array of receivers
  `=:horizontal` horizontal array of receivers
* `rand_flags::Vector{Bool}=[false, false]` : decide placement of supersources and receivers
  `=[true, false]` : randomly place supersources for regularly spaced receivers
  `=[true, true]` : randomly place supersources and receivers
  `=[false, false]` : regularly spaced supersources and receivers
  `=[false, true]` : randomly place receivers for regularly spaced supersources

# Return
* a fixed spread acquisition geometry `Geom`
"""
function Geom_fixed(
	      ssmin::Real,
	      ssmax::Real,
	      ss0::Real,
	      rmin::Real,
	      rmax::Real,
	      r0::Real,
	      nss::Int64,
	      nr::Int64,
	      ssattrib::Symbol=:horizontal,
	      rattrib::Symbol=:horizontal,
 	      rand_flags::Vector{Bool}=[false, false];
	      ssα::Real=0.0,
	      rα::Real=0.0,
	      ns::Vector{Int64}=ones(Int,nss),
	      srad::Real=0.0
	     )

	ssmin=Float64(ssmin); rmin=Float64(rmin)
	ssmax=Float64(ssmax); rmax=Float64(rmax)
	ss0=Float64(ss0); r0=Float64(r0)
	ssα=Float64(ssα)*pi/180.
	rα=Float64(rα)*pi/180.

	ssarray = isequal(nss,1) ? fill(ssmin,1) : (rand_flags[1] ?
					     Random.rand(Uniform(minimum([ssmin,ssmax]),maximum([ssmin,ssmax])),nss) : range(ssmin,stop=ssmax,length=nss))
	if(ssattrib==:horizontal)
		ssz = ss0.+(ssarray.-minimum(ssarray)).*sin(ssα)/cos(ssα); ssx=ssarray
	elseif(ssattrib==:vertical)
		ssx = ss0.+(ssarray.-minimum(ssarray)).*sin(ssα)/cos(ssα); ssz=ssarray
	else
		error("invalid ssattrib")
	end

	rarray = isequal(nr,1) ? fill(rmin,1) : (rand_flags[2] ?
					  Random.rand(Uniform(minimum([rmin,rmax]),maximum([rmin,rmax])),nr) : range(rmin,stop=rmax,length=nr))
	if(rattrib==:horizontal)
		rz = r0.+(rarray.-minimum(rarray)).*sin(rα)/cos(rα); rx = rarray
	elseif(rattrib==:vertical)
		rx = r0.+(rarray.-minimum(rarray)).*sin(rα)/cos(rα); rz = rarray
	else
		error("invalid rattrib")
	end
	rxall = [rx for iss=1:nss];
	rzall = [rz for iss=1:nss];
	ssxall=[zeros(ns[iss]) for iss=1:nss];
	sszall=[zeros(ns[iss]) for iss=1:nss];
	for iss in 1:nss
		for is in 1:ns[iss]
			θ=Random.rand(Uniform(-Float64(pi),Float64(pi)))
			r = iszero(srad) ? 0.0 : Random.rand(Uniform(0,srad))
			x=r*cos(θ)
			z=r*sin(θ)
			ssxall[iss][is]=ssx[iss]+x
			sszall[iss][is]=ssz[iss]+z
		end
	end
	return Geom(ssxall, sszall, rxall, rzall, nss, ns, fill(nr,nss))
end

"""
Randomly place source and receivers in a given model
"""
function Geom_fixed(mod::Models.Seismic,
	      nss::Int64=5,
	      nr::Int64=5,
	      ns::Vector{Int64}=ones(Int,nss))

	xmin=mod.mgrid[2][1]
	xmax=mod.mgrid[2][end]
	zmin=mod.mgrid[1][1]
	zmax=mod.mgrid[1][end]

	rxall = [Random.rand(Uniform(xmin,xmax),nr) for iss=1:nss];
	rzall = [Random.rand(Uniform(xmin,xmax),nr) for iss=1:nss];
	ssxall = [Random.rand(Uniform(xmin,xmax),ns[iss]) for iss=1:nss];
	sszall = [Random.rand(Uniform(xmin,xmax),ns[iss]) for iss=1:nss];

	return Geom(ssxall, sszall, rxall, rzall, nss, ns, fill(nr,nss))
end

"""
Circular acquisition. The sources and receivers can also be placed on a circle of radius
`rad`. The origin of the circle is at `loc`.
This geometry is unrealistic, but useful for testing.
Receivers are placed such that the limits
of the angular offset are given by `θlim`

# Arguments

* `nss::Int64=10` : number of supersources
* `nr::Int64=10` : number receivers for each super source
* `loc::Vector{Float64}=[0.,0.]` : location of origin
* `rad::Vector{Float64}=[100.,100.]` : radius for source and receiver circles, for example,
  * `=[0.,100.]` for sources at the center of circle
* `θlim::Vector{Float64}=[0.,π]` : acquisition is limited to these angular offsets between 0 and π

# Return

* a circular acquisition geometry `Geom`
"""
function Geom_circ(;
		   nss::Int64=10,
		   nr::Int64=10,
		   loc::Vector=[0.,0.],
		   rad::Vector=[100.,100.],
		   θlim::Vector=[0.,π]
	       )

	# modify nr such that approximately nr receivers are present in θlim
	nra = nr * convert(Int64, floor(π/(abs(diff(θlim)[1]))))

	(nra==0) && error("θlim should be from 0 to π")

	sxx, szz = circ_coord(loc..., nss, rad[1])
	rxxa, rzza = circ_coord(loc..., nra, rad[2])
	sx=Array{Array{Float64}}(undef,nss)
	sz=Array{Array{Float64}}(undef,nss)
	rx=Array{Array{Float64}}(undef,nss)
	rz=Array{Array{Float64}}(undef,nss)
	ns=Array{Int64}(undef,nss)
	nr=Array{Int64}(undef,nss)

	for iss=1:nss
		rangles = [angle(complex(rxxa[ira]-loc[2], rzza[ira]-loc[1])) for ira in 1:nra]
		sangles = angle(complex(sxx[iss]-loc[2], szz[iss]-loc[1]))
		diff = abs.(rangles.-sangles)
		diff = [min(diff[ira], (2π - diff[ira])) for ira in 1:nra]
		angles = findall(minimum(θlim) .<= diff .<= maximum(θlim))
		sx[iss] = [sxx[iss]]
		sz[iss] = [szz[iss]]
		ns[iss] = 1
		nr[iss] = length(angles)
		rx[iss] = rxxa[angles]
		rz[iss] = rzza[angles]
	end
	geom =  Geom(sx, sz, rx, rz, nss, ns, nr)

	print(geom, "circular acquisition geometry")
	return geom
end

function circ_coord(z, x, n, rad)
	θ = (n==1) ? [0.0] : collect(range(0, stop=2π, length=n+1))
	xx = (rad .* cos.(θ) .+ x)[1:n]
	zz = (rad .* sin.(θ) .+ z)[1:n]
	return xx, zz
end

"""
Appends the input a vector of acquisition geometries.
"Merge" `Geom` objects in `geomv` to `geom1`.
Typically you have one acquisition,
you change it and want to create the macro acquisition
where the supersources of both acquisitions are appended.
"""
function Geom_add!(geom1::Geom, geomv::Vector{Geom})
	for igeom=1:length(geomv)
		append!(geom1.sx, geomv[igeom].sx)
		append!(geom1.sz, geomv[igeom].sz)
		append!(geom1.rx, geomv[igeom].rx)
		append!(geom1.rz, geomv[igeom].rz)
		append!(geom1.ns, geomv[igeom].ns)
		append!(geom1.nr, geomv[igeom].nr)
		geom1.nss += geomv[igeom].nss
	end
	return geom1
end

"""
Adds input positions as either sources or receivers of every supershot.
"""
function Geom_add!(geom::Geom, zpos::Vector{Float64}, xpos::Vector{Float64}, attrib::Symbol)
	length(xpos) == length(zpos) ? nothing : error("same length vectors needed")
	for iss=1:geom.nss
		if(attrib == :receivers)
			geom.rx[iss] = vcat(geom.rx[iss][:], xpos)
			geom.rz[iss] = vcat(geom.rz[iss][:], zpos)
			geom.nr[iss] += length(xpos)
		elseif(attrib == :sources)
			geom.sx[iss] = vcat(geom.sx[iss][:], xpos)
			geom.sz[iss] = vcat(geom.sz[iss][:], zpos)
			geom.ns[iss] += length(xpos)
		end
	end
end

"""
Advance either source or receiver array in an acquisition geometry in horizontal or vertical directions.

# Arguments
* `geom::Geom` : acquisition geometry that is updated
* `advances::Vector{Float64}=[[0.,0.], [0.,0.,]]` : source and receiver advancements
"""
function Geom_advance(geom::Geom, advances::Vector{Vector{Float64}}=[[0.,0.], [0.,0.]])
	geom_out=deepcopy(geom)
	for iss=1:geom.nss
		geom_out.sx[iss][:] .+= advances[1][2]
		geom_out.sz[iss][:] .+= advances[1][1]
		geom_out.rx[iss][:] .+= advances[2][2]
		geom_out.rz[iss][:] .+= advances[2][1]
	end
	return geom_out
end

"""
return a vector of the order
"""
function Geom_getvec(geom::Vector{Geom}, field::Symbol)
	np = length(geom);
	vect = [getfield(geom[ip],field) for ip=1:np]
	return vec(vcat(vcat(vect...)...))
end

"""
Data type for the source related parameters during acquisiton.

# Fields
* `nss::Int64` : number of supersources
* `ns::Array{Int64}` : number of sources for each supersource
* `fields::Vector{Symbol}` : number of fields
* `wav::Array{Float64}` : wavelets in time domain
* `tgrid` : time grid
"""
mutable struct Src
	nss::Int64
	ns::Vector{Int64}
	fields::Vector{Symbol}
	wav::Array{Array{Float64,2},2}
	tgrid::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}
	Src(nss, ns, fields, wav, tgrid) =
		any([
		  any([fields[ifield] ∉ [:P, :Vx, :Vz] for ifield in 1:length(fields)]),
		  length(fields) == 0,
		  broadcast(size,wav) ≠ [(length(tgrid),ns[iss]) for iss=1:nss, ifield=1:length(fields)]
		  ]) ?
		error("error in Src construction") : new(nss, ns, fields, wav, tgrid)
end

function Base.isapprox(src1::Src, src2::Src)
	vec=([(src1.nss==src2.nss), (src1.fields==src2.fields), (src1.ns==src2.ns),
       		(isequal(src1.tgrid, src2.tgrid)),
		])
	return all(vec)
end

function Base.copyto!(srco::Src, src::Src)
	if(isapprox(srco, src))
		for f in [:wav]
			srcowav=getfield(srco, f)
			srcwav=getfield(src, f)
			for iss=1:src.nss, ifield=1:length(src.fields)
				srcowavv=srcowav[iss,ifield]
				srcwavv=srcwav[iss,ifield]
				copyto!(srcowavv, srcwavv)
			end
		end
		return srco
	else
		error("attempt to copy dissimilar sources")
	end
end

"""
Allocate `Src` with zeros depending on the acquisition geometry.
"""
function Src_zeros(acqgeom::Geom,  fields::Vector{Symbol}, tgrid::StepRangeLen)
	wavsrc = [zeros(length(tgrid),acqgeom.ns[iss]) for iss=1:acqgeom.nss, ifield=1:length(fields)]
	return Src(acqgeom.nss, acqgeom.ns, fields, wavsrc, deepcopy(tgrid))
end

"""
Priiiint information about `Src`
"""
function Base.print(src::Src, name::String="")
	println("\tSource Acquisition:\t",name)
	println("\t> number of supersources:\t",src.nss)
	println("\t> sources per supersource:\t","min\t",minimum(src.ns[:]), "\tmax\t", maximum(src.ns[:]))

	freqmin, freqmax, freqpeak = freqs(src)
	println("\t> frequency:\t","min\t",freqmin, "\tmax\t",freqmax,"\tpeak\t",freqpeak)
	println("\t> time:\t","min\t",src.tgrid[1], "\tmax\t", src.tgrid[end])
	println("\t> samples:\t",length(src.tgrid))
end

"""
Return minimum, maximum and peak frequencies of `Src`
"""
function freqs(src::Src)
	freqmin = minimum([Utils.findfreq(src.wav[i,j][:,:],src.tgrid,attrib=:min) for i in 1:src.nss, j in 1:length(src.fields)])
	freqmax = maximum([Utils.findfreq(src.wav[i,j][:,:],src.tgrid,attrib=:max) for i in 1:src.nss, j in 1:length(src.fields)])
	freqpeak = mean([Utils.findfreq(src.wav[i,j][:,:],src.tgrid,attrib=:peak) for i in 1:src.nss, j in 1:length(src.fields)])
	return freqmin, freqmax, freqpeak
end

"""
Constructor for `Src` data type.
Uses same source wavelet, i.e., `wav` for all sources and supersources

# Arguments

* `nss::Int64` : number of supersources
* `ns::Int64` : number of sources
* `fields::Vector{Symbol}` : number of fields the sources are exciting
* `wav::Array{Float64}` : a source wavelet that is used for all sources and supersources
* `tgrid` : time grid for the wavelet
"""
function Src_fixed(nss::Int64,
	     ns::Int64,
	     fields::Vector{Symbol},
	     wav::Array{Float64},
	     tgrid::StepRangeLen
	     )

	wavsrc = [repeat(wav,inner=(1,ns)) for iss=1:nss, ifield=1:length(fields)]
	return Src(nss, fill(ns, nss), fields, wavsrc, deepcopy(tgrid))
end

"""
Constructor of `Src`, which is typical for a input model such that
the model has `nλ` wavelengths.

# Arguments

* `nss::Int64` : number of supersources
* `ns::Int64` : number of source per each supersource
* `fields::Vector{Symbol}` :

# Keyword Arguments

* `mod::Models.Seismic` :
* `nλ::Int64=10` : number of wavelengths in the mod
* `wav_func::Function=(fqdom, tgrid)->Utils.Wavelets.ricker(fqdom,tgrid)` : which wavelet to generate, see Utils.Wavelets.jl
* `tmaxfrac::Float64=1.0` : by default the maximum modelling time is computed using the average velocity and the diagonal distance of the model, use this fraction to increase of reduce the maximum time
"""
function Src_fixed_mod(
		nss::Int64, ns::Int64, fields::Vector{Symbol};
		mod::Models.Seismic=nothing,
		nλ::Int64=10,
		wav_func::Function=(fqdom, tgrid)->Utils.Wavelets.ricker(fqdom,tgrid),
		tmaxfrac::Float64=1.0
		)

	x=mod.mgrid[2]; z=mod.mgrid[1]
	# dominant wavelength using mod dimensions
	λdom=mean([(abs(x[end]-x[1])), (abs(z[end]-z[1]))])/real(nλ)
	# average P velocity
	vavg=Models.χ([mean(mod.χvp)], mod.ref.vp, -1)[1]

	fqdom = vavg/λdom

	# maximum distance the wave travels
	d = sqrt((x[1]-x[end]).^2+(z[1]-z[end]).^2)

	# use two-way maximum distance to get tmax
	tmax=2.0*d*inv(vavg)*tmaxfrac

	# choose sampling interval to obey max freq of source wavelet
	δmin = minimum([step(mod.mgrid[2]), step(mod.mgrid[1])])
	vpmax = mod.vp0[2]
	δt=0.5*δmin/vpmax

	# check if δt is reasonable
	#(δt > 0.1/fqdom) : error("decrease spatial sampling or nλ")

	tgrid=range(0.0, stop=tmax, step=δt)
	wav = wav_func(fqdom, tgrid);

	wavsrc = [hcat([wav_func(fqdom, tgrid) for is in 1:ns]...) for iss=1:nss, ifield=1:length(fields)]
	src=Src(nss, fill(ns, nss), fields, wavsrc, deepcopy(tgrid))
	print(src)
	# choose ricker waveletes of fdom
	return src
end

"""
Generate band-limited random source signals
"""
function Src_fixed_random(nss::Int64, ns::Int64, fields::Vector{Symbol};
			  distvec=[Normal() for iss in 1:nss],
			  sparsepvec=[1. for iss in 1:nss],
			  fmin::Float64=0.0,
			  fmax::Float64=0.0,
			  tgrid::StepRangeLen=nothing,
			  tmaxfrac::Float64=1.0)
	wavsrc = [repeat(zeros(length(tgrid)),inner=(1,ns)) for iss=1:nss, ifield=1:length(fields)]
	for ifield in 1:length(fields), iss in 1:nss, is in 1:ns
		wavsrc[iss, ifield][:,is] = Utils.get_tapered_random_tmax_signal(tgrid, fmin=fmin, fmax=fmax, tmaxfrac=tmaxfrac, dist=distvec[iss],
								 sparsep=sparsepvec[iss])
	end
	src=Src(nss, fill(ns, nss), fields, wavsrc, deepcopy(tgrid))
	print(src)
	return src
end

"""
Function that returns Src after time reversal
"""
function Src_tr(src::Src)
	return Src(src.nss,src.ns,src.fields,
	    [reverse(src.wav[i,j],dims=1) for i in 1:src.nss, j in 1:length(src.fields)],deepcopy(src.tgrid))
end

"""
Pad `Src`
tgrids should be same in all Src
"""
function Src_uspos(src::Vector{Src}, acq::Vector{Geom})
	np = length(src) == length(acq) ? length(src) : error("unequal sizez")

	# unique source positions
	nus=Geom_get(acq,:nus)
	uspos=Geom_get(acq,:uspos)
	# all zeros for all unique positions
	wavout = [[zeros(src[ip].length(tgrid),nus) for iss=1:src[ip].nss, ifield=1:length(src[ip].fields)] for ip=1:np]
	# fill source wavelets when necessary
	for ip=1:np
		for ifield=1:length(src[ip].fields), iss=1:src[ip].nss, is=1:acq[ip].ns[iss]
			is0=findall([[uspos[1][i]-acq[ip].sz[iss][is],
		      uspos[2][i]-acq[ip].sx[iss][is]] == [0., 0.,] for i in 1:nus])

			wavout[ip][iss, ifield][:,is0] = src[ip].wav[iss, ifield][:,is]
		end
	end
	# output src
	return [Src(src[ip].nss,fill(nus, src[ip].nss),
	     src[ip].fields,wavout[ip],src[ip].tgrid) for ip=1:np]
end

"""
Return a sparse ACQ matrix, for a given acqgeom
data=ACQ*snapshot
"""
function ACQmat(acqgeom::Geom,mgrid,iss=1)
	nz,nx=length.(mgrid)
	ACQ=spzeros(acqgeom.nr[iss],prod(length.(mgrid)))
	for ir = 1:acqgeom.nr[iss]
		irx=argmin(abs.(mgrid[2].-acqgeom.rx[iss][ir]))
		irz=argmin(abs.(mgrid[1].-acqgeom.rz[iss][ir]))
		ACQ[ir,irz+(irx-1)*nz]=1.0
	end
	return ACQ
end

"""
return a vector of the order

"""
function Src_getvec(src::Vector{Src}, field::Symbol)
	np = length(src);
	vect = [getfield(src[ip],field) for ip=1:np]
	return vec(hcat(hcat(vect...)...))
end

function save(geom, folder)
	!(isdir(folder)) && error("invalid directory")
	for iss in 1:geom.nss
		file=joinpath(folder, string("s", iss, ".csv"))
		CSV.write(file,DataFrame(hcat(geom.sx[iss], geom.sz[iss])))
		file=joinpath(folder, string("r", iss, ".csv"))
		CSV.write(file,DataFrame(hcat(geom.rx[iss], geom.rz[iss])))
	end
end

end # module


using CleanEOSTable
using Base.Test
using HDF5

mktempdir() do path
    println("Downloading...")
    download("https://stellarcollapse.org/EOS/LS220_234r_136t_50y_analmu_20091212_SVNr26.h5.bz2",
             "$path/LS220_234r_136t_50y_analmu_20091212_SVNr26.h5.bz2")
    println("Unpacking...")
    run(`bunzip2 $path/LS220_234r_136t_50y_analmu_20091212_SVNr26.h5.bz2`)

    tablename = joinpath(path, "$path/LS220_234r_136t_50y_analmu_20091212_SVNr26.h5")
    new_tablename = joinpath(path, "$path/LS220_234r_136t_50y_analmu_20091212_SVNr26_corrected.h5")
    cleanEOSTable(tablename, new_tablename)

    println("Comparing tables...")
    table = h5open(tablename, "r")
    table_logenergy = read(table, "logenergy")
    new_table = h5open(new_tablename, "r")
    new_table_logenergy = read(new_table, "logenergy")

    @test count(new_table_logenergy .!= table_logenergy) == 741
end