text:
module RudeOilTests
using RudeOil
using FactCheck: facts, context, @fact, greater_than, not, @fact_throws

import Base: ismatch
ismatch(regex::String) = x -> ismatch(Regex(regex), x)
contains(item) = x -> item ∈ x

#= include("machine.jl") =#

const machine = Machine("TemporaryRUDEOILTestMachine")
try
  # Use same machine for all tests.
  # Cos takes time to create
  start_machine(machine)

  include("env.jl")
  include("image.jl")
  include("container.jl")

finally
  remove!(machine)
end
end


text:
export schedule, Schedulers
"""
    schedule(model) → ids
Return an iterator over the scheduled IDs using the model's scheduler.
"""
schedule(model::ABM) = model.scheduler(model)

# Notice how the above lines are *outside* the submodule

"""
    Schedulers
Submodule containing all predefined schedulers of Agents.jl and the scheduling API.
Schedulers have a very simple interface. They are functions that take as an input the ABM and
return an iterator over agent IDs. Notice that this iterator can be a "true" iterator
(non-allocated) or can be just a standard vector of IDs. You can define your own scheduler
according to this API and use it when making an [`AgentBasedModel`](@ref).
You can also use the function `schedule(model)` to obtain the scheduled ID list,
if you prefer to write your own `step!`-like loop.

See also [Advanced scheduling](@ref) for making more advanced schedulers.

Notice that schedulers can be given directly to model creation, and thus become the
"default" scheduler a model uses, but they can just as easily be incorporated in a
`model_step!` function as shown in [Advanced stepping](@ref).
"""
module Schedulers
using Agents
using Random: shuffle!, randsubseq

export randomly, by_id, fastest, partially, by_property, by_type

####################################
# Schedulers
####################################
"""
    Schedulers.fastest
A scheduler that activates all agents once per step in the order dictated by the
agent's container, which is arbitrary (the keys sequence of a dictionary).
This is the fastest way to activate all agents once per step.
"""
fastest(model::ABM) = keys(model.agents)

"""
    Schedulers.by_id
A scheduler that activates all agents agents at each step according to their id.
"""
function by_id(model::ABM)
    agent_ids = sort(collect(keys(model.agents)))
    return agent_ids
end

"""
    Schedulers.randomly
A scheduler that activates all agents once per step in a random order.
Different random ordering is used at each different step.
"""
function randomly(model::ABM)
    order = shuffle!(model.rng, collect(keys(model.agents)))
end

"""
    Schedulers.partially(p)
A scheduler that at each step activates only `p` percentage of randomly chosen agents.
"""
function partially(p::Real)
    function partial(model::ABM)
        ids = collect(keys(model.agents))
        return randsubseq(model.rng, ids, p)
    end
    return partial
end

"""
    Schedulers.by_property(property)
A scheduler that at each step activates the agents in an order dictated by their `property`,
with agents with greater `property` acting first. `property` can be a `Symbol`, which
just dictates which field of the agents to compare, or a function which inputs an agent
and outputs a real number.
"""
function by_property(p)
    function property(model::ABM)
        ids = collect(keys(model.agents))
        properties = [Agents.get_data(model[id], p) for id in ids]
        s = sortperm(properties)
        return ids[s]
    end
end

"""
    Schedulers.by_type(shuffle_types::Bool, shuffle_agents::Bool)
A scheduler useful only for mixed agent models using `Union` types.
- Setting `shuffle_types = true` groups by agent type, but randomizes the type order.
Otherwise returns agents grouped in order of appearance in the `Union`.
- `shuffle_agents = true` randomizes the order of agents within each group, `false` returns
the default order of the container (equivalent to [`Schedulers.fastest`](@ref)).
"""
function by_type(shuffle_types::Bool, shuffle_agents::Bool)
    function by_union(model::ABM{S,A}) where {S,A}
        types = Agents.union_types(A)
        sets = [Integer[] for _ in types]
        for agent in allagents(model)
            idx = findfirst(t -> t == typeof(agent), types)
            push!(sets[idx], agent.id)
        end
        shuffle_types && shuffle!(model.rng, sets)
        shuffle_agents && [shuffle!(model.rng, set) for set in sets]
        vcat(sets...)
    end
end

"""
    Schedulers.by_type((C, B, A), shuffle_agents::Bool)
A scheduler that activates agents by type in specified order (since `Union`s are not order
preserving). `shuffle_agents = true` randomizes the order of agents within each group.
"""
function by_type(order::Tuple{Type,Vararg{Type}}, shuffle_agents::Bool)
    function by_ordered_union(model::ABM{S,A}) where {S,A}
        types = Agents.union_types(A)
        if order !== nothing
            @assert length(types) == length(order) "Invalid dimension for `order`"
            types = order
        end
        sets = [Integer[] for _ in types]
        for agent in allagents(model)
            idx = findfirst(t -> t == typeof(agent), types)
            push!(sets[idx], agent.id)
        end
        shuffle_agents && [shuffle!(model.rng, set) for set in sets]
        vcat(sets...)
    end
end

end # Schedulers submodule

ext: jl


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

# Linear algebra functions for dense matrices in column major format

## BLAS cutoff threshold constants

const SCAL_CUTOFF = 2048
const DOT_CUTOFF = 128
const ASUM_CUTOFF = 32
const NRM2_CUTOFF = 32

# Generic cross-over constant based on benchmarking on a single thread with an i7 CPU @ 2.5GHz
# L1 cache: 32K, L2 cache: 256K, L3 cache: 6144K
# This constant should ideally be determined by the actual CPU cache size
const ISONE_CUTOFF = 2^21 # 2M

function rmul!(X::Array{T}, s::T) where T<:BlasFloat
    s == 0 && return fill!(X, zero(T))
    s == 1 && return X
    if length(X) < SCAL_CUTOFF
        generic_rmul!(X, s)
    else
        BLAS.scal!(length(X), s, X, 1)
    end
    X
end

lmul!(s::T, X::Array{T}) where {T<:BlasFloat} = rmul!(X, s)

rmul!(X::Array{T}, s::Number) where {T<:BlasFloat} = rmul!(X, convert(T, s))
function rmul!(X::Array{T}, s::Real) where T<:BlasComplex
    R = typeof(real(zero(T)))
    GC.@preserve X BLAS.scal!(2*length(X), convert(R,s), convert(Ptr{R},pointer(X)), 1)
    X
end

function isone(A::StridedMatrix)
    m, n = size(A)
    m != n && return false # only square matrices can satisfy x == one(x)
    if sizeof(A) < ISONE_CUTOFF
        _isone_triacheck(A, m)
    else
        _isone_cachefriendly(A, m)
    end
end

@inline function _isone_triacheck(A::StridedMatrix, m::Int)
    @inbounds for i in 1:m, j in i:m
        if i == j
            isone(A[i,i]) || return false
        else
            iszero(A[i,j]) && iszero(A[j,i]) || return false
        end
    end
    return true
end

# Inner loop over rows to be friendly to the CPU cache
@inline function _isone_cachefriendly(A::StridedMatrix, m::Int)
    @inbounds for i in 1:m, j in 1:m
        if i == j
            isone(A[i,i]) || return false
        else
            iszero(A[j,i]) || return false
        end
    end
    return true
end

"""
    isposdef!(A) -> Bool

Test whether a matrix is positive definite (and Hermitian) by trying to perform a
Cholesky factorization of `A`, overwriting `A` in the process.
See also [`isposdef`](@ref).

# Examples
```jldoctest
julia> A = [1. 2.; 2. 50.];

julia> isposdef!(A)
true

julia> A
2×2 Array{Float64,2}:
 1.0  2.0
 2.0  6.78233
```
"""
isposdef!(A::AbstractMatrix) = ishermitian(A) && isposdef(cholfact!(Hermitian(A)))

"""
    isposdef(A) -> Bool

Test whether a matrix is positive definite (and Hermitian) by trying to perform a
Cholesky factorization of `A`.
See also [`isposdef!`](@ref)

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

julia> isposdef(A)
true
```
"""
isposdef(A::AbstractMatrix) = ishermitian(A) && isposdef(cholfact(Hermitian(A)))
isposdef(x::Number) = imag(x)==0 && real(x) > 0

# the definition of strides for Array{T,N} is tuple() if N = 0, otherwise it is
# a tuple containing 1 and a cumulative product of the first N-1 sizes
# this definition is also used for StridedReshapedArray and StridedReinterpretedArray
# which have the same memory storage as Array
function stride(a::Union{DenseArray,StridedReshapedArray,StridedReinterpretArray}, i::Int)
    if i > ndims(a)
        return length(a)
    end
    s = 1
    for n = 1:(i-1)
        s *= size(a, n)
    end
    return s
end
strides(a::Union{DenseArray,StridedReshapedArray,StridedReinterpretArray}) = size_to_strides(1, size(a)...)

function norm(x::StridedVector{T}, rx::Union{UnitRange{TI},AbstractRange{TI}}) where {T<:BlasFloat,TI<:Integer}
    if minimum(rx) < 1 || maximum(rx) > length(x)
        throw(BoundsError(x, rx))
    end
    GC.@preserve x BLAS.nrm2(length(rx), pointer(x)+(first(rx)-1)*sizeof(T), step(rx))
end

vecnorm1(x::Union{Array{T},StridedVector{T}}) where {T<:BlasReal} =
    length(x) < ASUM_CUTOFF ? generic_vecnorm1(x) : BLAS.asum(x)

vecnorm2(x::Union{Array{T},StridedVector{T}}) where {T<:BlasFloat} =
    length(x) < NRM2_CUTOFF ? generic_vecnorm2(x) : BLAS.nrm2(x)

"""
    triu!(M, k::Integer)

Return the upper triangle of `M` starting from the `k`th superdiagonal,
overwriting `M` in the process.

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

julia> triu!(M, 1)
5×5 Array{Int64,2}:
 0  2  3  4  5
 0  0  3  4  5
 0  0  0  4  5
 0  0  0  0  5
 0  0  0  0  0
```
"""
function triu!(M::AbstractMatrix, k::Integer)
    m, n = size(M)
    if !(-m + 1 <= k <= n + 1)
        throw(ArgumentError(string("the requested diagonal, $k, must be at least ",
            "$(-m + 1) and at most $(n + 1) in an $m-by-$n matrix")))
    end
    idx = 1
    for j = 0:n-1
        ii = min(max(0, j+1-k), m)
        for i = (idx+ii):(idx+m-1)
            M[i] = zero(M[i])
        end
        idx += m
    end
    M
end

triu(M::Matrix, k::Integer) = triu!(copy(M), k)

"""
    tril!(M, k::Integer)

Return the lower triangle of `M` starting from the `k`th superdiagonal, overwriting `M` in
the process.

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

julia> tril!(M, 2)
5×5 Array{Int64,2}:
 1  2  3  0  0
 1  2  3  4  0
 1  2  3  4  5
 1  2  3  4  5
 1  2  3  4  5
```
"""
function tril!(M::AbstractMatrix, k::Integer)
    m, n = size(M)
    if !(-m - 1 <= k <= n - 1)
        throw(ArgumentError(string("the requested diagonal, $k, must be at least ",
            "$(-m - 1) and at most $(n - 1) in an $m-by-$n matrix")))
    end
    idx = 1
    for j = 0:n-1
        ii = min(max(0, j-k), m)
        for i = idx:(idx+ii-1)
            M[i] = zero(M[i])
        end
        idx += m
    end
    M
end
tril(M::Matrix, k::Integer) = tril!(copy(M), k)

"""
    fillband!(A::AbstractMatrix, x, l, u)

Fill the band between diagonals `l` and `u` with the value `x`.
"""
function fillband!(A::AbstractMatrix{T}, x, l, u) where T
    m, n = size(A)
    xT = convert(T, x)
    for j in 1:n
        for i in max(1,j-u):min(m,j-l)
            @inbounds A[i, j] = xT
        end
    end
    return A
end

function diagind(m::Integer, n::Integer, k::Integer=0)
    if !(-m <= k <= n)
        throw(ArgumentError(string("requested diagonal, $k, must be at least $(-m) and ",
            "at most $n in an $m-by-$n matrix")))
    end
    k <= 0 ? range(1-k, step=m+1, length=min(m+k, n)) : range(k*m+1, step=m+1, length=min(m, n-k))
end

"""
    diagind(M, k::Integer=0)

An `AbstractRange` giving the indices of the `k`th diagonal of the matrix `M`.

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

julia> diagind(A,-1)
2:4:6
```
"""
diagind(A::AbstractMatrix, k::Integer=0) = diagind(size(A,1), size(A,2), k)

"""
    diag(M, k::Integer=0)

The `k`th diagonal of a matrix, as a vector.

See also: [`diagm`](@ref)

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

julia> diag(A,1)
2-element Array{Int64,1}:
 2
 6
```
"""
diag(A::AbstractMatrix, k::Integer=0) = A[diagind(A,k)]

"""
    diagm(kv::Pair{<:Integer,<:AbstractVector}...)

Construct a square matrix from `Pair`s of diagonals and vectors.
Vector `kv.second` will be placed on the `kv.first` diagonal.
`diagm` constructs a full matrix; if you want storage-efficient
versions with fast arithmetic, see [`Diagonal`](@ref), [`Bidiagonal`](@ref)
[`Tridiagonal`](@ref) and [`SymTridiagonal`](@ref).

# Examples
```jldoctest
julia> diagm(1 => [1,2,3])
4×4 Array{Int64,2}:
 0  1  0  0
 0  0  2  0
 0  0  0  3
 0  0  0  0

julia> diagm(1 => [1,2,3], -1 => [4,5])
4×4 Array{Int64,2}:
 0  1  0  0
 4  0  2  0
 0  5  0  3
 0  0  0  0
```
"""
function diagm(kv::Pair{<:Integer,<:AbstractVector}...)
    A = diagm_container(kv...)
    for p in kv
        inds = diagind(A, p.first)
        for (i, val) in enumerate(p.second)
            A[inds[i]] += val
        end
    end
    return A
end
function diagm_container(kv::Pair{<:Integer,<:AbstractVector}...)
    T = promote_type(map(x -> eltype(x.second), kv)...)
    n = mapreduce(x -> length(x.second) + abs(x.first), max, kv)
    return zeros(T, n, n)
end
function diagm_container(kv::Pair{<:Integer,<:BitVector}...)
    n = mapreduce(x -> length(x.second) + abs(x.first), max, kv)
    return falses(n, n)
end

function tr(A::Matrix{T}) where T
    n = checksquare(A)
    t = zero(T)
    for i=1:n
        t += A[i,i]
    end
    t
end

"""
    kron(A, B)

Kronecker tensor product of two vectors or two matrices.

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

julia> B = [im 1; 1 -im]
2×2 Array{Complex{Int64},2}:
 0+1im  1+0im
 1+0im  0-1im

julia> kron(A, B)
4×4 Array{Complex{Int64},2}:
 0+1im  1+0im  0+2im  2+0im
 1+0im  0-1im  2+0im  0-2im
 0+3im  3+0im  0+4im  4+0im
 3+0im  0-3im  4+0im  0-4im
```
"""
function kron(a::AbstractMatrix{T}, b::AbstractMatrix{S}) where {T,S}
    R = Matrix{promote_op(*,T,S)}(undef, size(a,1)*size(b,1), size(a,2)*size(b,2))
    m = 1
    for j = 1:size(a,2), l = 1:size(b,2), i = 1:size(a,1)
        aij = a[i,j]
        for k = 1:size(b,1)
            R[m] = aij*b[k,l]
            m += 1
        end
    end
    R
end

kron(a::Number, b::Union{Number, AbstractVecOrMat}) = a * b
kron(a::AbstractVecOrMat, b::Number) = a * b
kron(a::AbstractVector, b::AbstractVector) = vec(kron(reshape(a ,length(a), 1), reshape(b, length(b), 1)))
kron(a::AbstractMatrix, b::AbstractVector) = kron(a, reshape(b, length(b), 1))
kron(a::AbstractVector, b::AbstractMatrix) = kron(reshape(a, length(a), 1), b)

# Matrix power
(^)(A::AbstractMatrix, p::Integer) = p < 0 ? power_by_squaring(inv(A), -p) : power_by_squaring(A, p)
function (^)(A::AbstractMatrix{T}, p::Integer) where T<:Integer
    # make sure that e.g. [1 1;1 0]^big(3)
    # gets promotes in a similar way as 2^big(3)
    TT = promote_op(^, T, typeof(p))
    return power_by_squaring(convert(AbstractMatrix{TT}, A), p)
end
function integerpow(A::AbstractMatrix{T}, p) where T
    TT = promote_op(^, T, typeof(p))
    return (TT == T ? A : copyto!(similar(A, TT), A))^Integer(p)
end
function schurpow(A::AbstractMatrix, p)
    if istriu(A)
        # Integer part
        retmat = A ^ floor(p)
        # Real part
        if p - floor(p) == 0.5
            # special case: A^0.5 === sqrt(A)
            retmat = retmat * sqrt(A)
        else
            retmat = retmat * powm!(UpperTriangular(float.(A)), real(p - floor(p)))
        end
    else
        S,Q,d = schur(complex(A))
        # Integer part
        R = S ^ floor(p)
        # Real part
        if p - floor(p) == 0.5
            # special case: A^0.5 === sqrt(A)
            R = R * sqrt(S)
        else
            R = R * powm!(UpperTriangular(float.(S)), real(p - floor(p)))
        end
        retmat = Q * R * Q'
    end

    # if A has nonpositive real eigenvalues, retmat is a nonprincipal matrix power.
    if isreal(retmat)
        return real(retmat)
    else
        return retmat
    end
end
function (^)(A::AbstractMatrix{T}, p::Real) where T
    n = checksquare(A)

    # Quicker return if A is diagonal
    if isdiag(A)
        TT = promote_op(^, T, typeof(p))
        retmat = copy_oftype(A, TT)
        for i in 1:n
            retmat[i, i] = retmat[i, i] ^ p
        end
        return retmat
    end

    # For integer powers, use power_by_squaring
    isinteger(p) && return integerpow(A, p)

    # If possible, use diagonalization
    if issymmetric(A)
        return (Symmetric(A)^p)
    end
    if ishermitian(A)
        return (Hermitian(A)^p)
    end

    # Otherwise, use Schur decomposition
    return schurpow(A, p)
end
(^)(A::AbstractMatrix, p::Number) = exp(p*log(A))

# Matrix exponential

"""
    exp(A::AbstractMatrix)

Compute the matrix exponential of `A`, defined by

```math
e^A = \\sum_{n=0}^{\\infty} \\frac{A^n}{n!}.
```

For symmetric or Hermitian `A`, an eigendecomposition ([`eigfact`](@ref)) is
used, otherwise the scaling and squaring algorithm (see [^H05]) is chosen.

[^H05]: Nicholas J. Higham, "The squaring and scaling method for the matrix exponential revisited", SIAM Journal on Matrix Analysis and Applications, 26(4), 2005, 1179-1193. [doi:10.1137/090768539](https://doi.org/10.1137/090768539)

# Examples
```jldoctest
julia> A = Matrix(1.0I, 2, 2)
2×2 Array{Float64,2}:
 1.0  0.0
 0.0  1.0

julia> exp(A)
2×2 Array{Float64,2}:
 2.71828  0.0
 0.0      2.71828
```
"""
exp(A::StridedMatrix{<:BlasFloat}) = exp!(copy(A))
exp(A::StridedMatrix{<:Union{Integer,Complex{<:Integer}}}) = exp!(float.(A))

## Destructive matrix exponential using algorithm from Higham, 2008,
## "Functions of Matrices: Theory and Computation", SIAM
function exp!(A::StridedMatrix{T}) where T<:BlasFloat
    n = checksquare(A)
    if ishermitian(A)
        return copytri!(parent(exp(Hermitian(A))), 'U', true)
    end
    ilo, ihi, scale = LAPACK.gebal!('B', A)    # modifies A
    nA   = norm(A, 1)
    Inn    = Matrix{T}(I, n, n)
    ## For sufficiently small nA, use lower order Padé-Approximations
    if (nA <= 2.1)
        if nA > 0.95
            C = T[17643225600.,8821612800.,2075673600.,302702400.,
                     30270240.,   2162160.,    110880.,     3960.,
                           90.,         1.]
        elseif nA > 0.25
            C = T[17297280.,8648640.,1995840.,277200.,
                     25200.,   1512.,     56.,     1.]
        elseif nA > 0.015
            C = T[30240.,15120.,3360.,
                    420.,   30.,   1.]
        else
            C = T[120.,60.,12.,1.]
        end
        A2 = A * A
        P  = copy(Inn)
        U  = C[2] * P
        V  = C[1] * P
        for k in 1:(div(size(C, 1), 2) - 1)
            k2 = 2 * k
            P *= A2
            U += C[k2 + 2] * P
            V += C[k2 + 1] * P
        end
        U = A * U
        X = V + U
        LAPACK.gesv!(V-U, X)
    else
        s  = log2(nA/5.4)               # power of 2 later reversed by squaring
        if s > 0
            si = ceil(Int,s)
            A /= convert(T,2^si)
        end
        CC = T[64764752532480000.,32382376266240000.,7771770303897600.,
                1187353796428800.,  129060195264000.,  10559470521600.,
                    670442572800.,      33522128640.,      1323241920.,
                        40840800.,           960960.,           16380.,
                             182.,                1.]
        A2 = A * A
        A4 = A2 * A2
        A6 = A2 * A4
        U  = A * (A6 * (CC[14]*A6 + CC[12]*A4 + CC[10]*A2) +
                  CC[8]*A6 + CC[6]*A4 + CC[4]*A2 + CC[2]*Inn)
        V  = A6 * (CC[13]*A6 + CC[11]*A4 + CC[9]*A2) +
                   CC[7]*A6 + CC[5]*A4 + CC[3]*A2 + CC[1]*Inn

        X = V + U
        LAPACK.gesv!(V-U, X)

        if s > 0            # squaring to reverse dividing by power of 2
            for t=1:si; X *= X end
        end
    end

    # Undo the balancing
    for j = ilo:ihi
        scj = scale[j]
        for i = 1:n
            X[j,i] *= scj
        end
        for i = 1:n
            X[i,j] /= scj
        end
    end

    if ilo > 1       # apply lower permutations in reverse order
        for j in (ilo-1):-1:1; rcswap!(j, Int(scale[j]), X) end
    end
    if ihi < n       # apply upper permutations in forward order
        for j in (ihi+1):n;    rcswap!(j, Int(scale[j]), X) end
    end
    X
end

## Swap rows i and j and columns i and j in X
function rcswap!(i::Integer, j::Integer, X::StridedMatrix{<:Number})
    for k = 1:size(X,1)
        X[k,i], X[k,j] = X[k,j], X[k,i]
    end
    for k = 1:size(X,2)
        X[i,k], X[j,k] = X[j,k], X[i,k]
    end
end

"""
    log(A{T}::StridedMatrix{T})

If `A` has no negative real eigenvalue, compute the principal matrix logarithm of `A`, i.e.
the unique matrix ``X`` such that ``e^X = A`` and ``-\\pi < Im(\\lambda) < \\pi`` for all
the eigenvalues ``\\lambda`` of ``X``. If `A` has nonpositive eigenvalues, a nonprincipal
matrix function is returned whenever possible.

If `A` is symmetric or Hermitian, its eigendecomposition ([`eigfact`](@ref)) is
used, if `A` is triangular an improved version of the inverse scaling and squaring method is
employed (see [^AH12] and [^AHR13]). For general matrices, the complex Schur form
([`schur`](@ref)) is computed and the triangular algorithm is used on the
triangular factor.

[^AH12]: Awad H. Al-Mohy and Nicholas J. Higham, "Improved inverse  scaling and squaring algorithms for the matrix logarithm", SIAM Journal on Scientific Computing, 34(4), 2012, C153-C169. [doi:10.1137/110852553](https://doi.org/10.1137/110852553)

[^AHR13]: Awad H. Al-Mohy, Nicholas J. Higham and Samuel D. Relton, "Computing the Fréchet derivative of the matrix logarithm and estimating the condition number", SIAM Journal on Scientific Computing, 35(4), 2013, C394-C410. [doi:10.1137/120885991](https://doi.org/10.1137/120885991)

# Examples
```jldoctest
julia> A = Matrix(2.7182818*I, 2, 2)
2×2 Array{Float64,2}:
 2.71828  0.0
 0.0      2.71828

julia> log(A)
2×2 Array{Float64,2}:
  1.0  0.0
 -0.0  1.0
```
"""
function log(A::StridedMatrix)
    # If possible, use diagonalization
    if ishermitian(A)
        logHermA = log(Hermitian(A))
        return isa(logHermA, Hermitian) ? copytri!(parent(logHermA), 'U', true) : parent(logHermA)
    end

    # Use Schur decomposition
    n = checksquare(A)
    if istriu(A)
        return triu!(parent(log(UpperTriangular(complex(A)))))
    else
        if isreal(A)
            SchurF = schurfact(real(A))
        else
            SchurF = schurfact(A)
        end
        if !istriu(SchurF.T)
            SchurS = schurfact(complex(SchurF.T))
            logT = SchurS.Z * log(UpperTriangular(SchurS.T)) * SchurS.Z'
            return SchurF.Z * logT * SchurF.Z'
        else
            R = log(UpperTriangular(complex(SchurF.T)))
            return SchurF.Z * R * SchurF.Z'
        end
    end
end

"""
    sqrt(A::AbstractMatrix)

If `A` has no negative real eigenvalues, compute the principal matrix square root of `A`,
that is the unique matrix ``X`` with eigenvalues having positive real part such that
``X^2 = A``. Otherwise, a nonprincipal square root is returned.

If `A` is symmetric or Hermitian, its eigendecomposition ([`eigfact`](@ref)) is
used to compute the square root. Otherwise, the square root is determined by means of the
Björck-Hammarling method [^BH83], which computes the complex Schur form ([`schur`](@ref))
and then the complex square root of the triangular factor.

[^BH83]:

    Åke Björck and Sven Hammarling, "A Schur method for the square root of a matrix",
    Linear Algebra and its Applications, 52-53, 1983, 127-140.
    [doi:10.1016/0024-3795(83)80010-X](https://doi.org/10.1016/0024-3795(83)80010-X)

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

julia> sqrt(A)
2×2 Array{Float64,2}:
 2.0  0.0
 0.0  2.0
```
"""
function sqrt(A::StridedMatrix{<:Real})
    if issymmetric(A)
        return copytri!(parent(sqrt(Symmetric(A))), 'U')
    end
    n = checksquare(A)
    if istriu(A)
        return triu!(parent(sqrt(UpperTriangular(A))))
    else
        SchurF = schurfact(complex(A))
        R = triu!(parent(sqrt(UpperTriangular(SchurF.T)))) # unwrapping unnecessary?
        return SchurF.vectors * R * SchurF.vectors'
    end
end
function sqrt(A::StridedMatrix{<:Complex})
    if ishermitian(A)
        sqrtHermA = sqrt(Hermitian(A))
        return isa(sqrtHermA, Hermitian) ? copytri!(parent(sqrtHermA), 'U', true) : parent(sqrtHermA)
    end
    n = checksquare(A)
    if istriu(A)
        return triu!(parent(sqrt(UpperTriangular(A))))
    else
        SchurF = schurfact(A)
        R = triu!(parent(sqrt(UpperTriangular(SchurF.T)))) # unwrapping unnecessary?
        return SchurF.vectors * R * SchurF.vectors'
    end
end

function inv(A::StridedMatrix{T}) where T
    checksquare(A)
    S = typeof((one(T)*zero(T) + one(T)*zero(T))/one(T))
    AA = convert(AbstractArray{S}, A)
    if istriu(AA)
        Ai = triu!(parent(inv(UpperTriangular(AA))))
    elseif istril(AA)
        Ai = tril!(parent(inv(LowerTriangular(AA))))
    else
        Ai = inv!(lufact(AA))
        Ai = convert(typeof(parent(Ai)), Ai)
    end
    return Ai
end

"""
    cos(A::AbstractMatrix)

Compute the matrix cosine of a square matrix `A`.

If `A` is symmetric or Hermitian, its eigendecomposition ([`eigfact`](@ref)) is used to
compute the cosine. Otherwise, the cosine is determined by calling [`exp`](@ref).

# Examples
```jldoctest
julia> cos(fill(1.0, (2,2)))
2×2 Array{Float64,2}:
  0.291927  -0.708073
 -0.708073   0.291927
```
"""
function cos(A::AbstractMatrix{<:Real})
    if issymmetric(A)
        return copytri!(parent(cos(Symmetric(A))), 'U')
    end
    T = complex(float(eltype(A)))
    return real(exp!(T.(im .* A)))
end
function cos(A::AbstractMatrix{<:Complex})
    if ishermitian(A)
        return copytri!(parent(cos(Hermitian(A))), 'U', true)
    end
    T = complex(float(eltype(A)))
    X = exp!(T.(im .* A))
    @. X = (X + $exp!(T(-im*A))) / 2
    return X
end

"""
    sin(A::AbstractMatrix)

Compute the matrix sine of a square matrix `A`.

If `A` is symmetric or Hermitian, its eigendecomposition ([`eigfact`](@ref)) is used to
compute the sine. Otherwise, the sine is determined by calling [`exp`](@ref).

# Examples
```jldoctest
julia> sin(fill(1.0, (2,2)))
2×2 Array{Float64,2}:
 0.454649  0.454649
 0.454649  0.454649
```
"""
function sin(A::AbstractMatrix{<:Real})
    if issymmetric(A)
        return copytri!(parent(sin(Symmetric(A))), 'U')
    end
    T = complex(float(eltype(A)))
    return imag(exp!(T.(im .* A)))
end
function sin(A::AbstractMatrix{<:Complex})
    if ishermitian(A)
        return copytri!(parent(sin(Hermitian(A))), 'U', true)
    end
    T = complex(float(eltype(A)))
    X = exp!(T.(im .* A))
    Y = exp!(T.(.-im .* A))
    @inbounds for i in eachindex(X)
        x, y = X[i]/2, Y[i]/2
        X[i] = Complex(imag(x)-imag(y), real(y)-real(x))
    end
    return X
end

"""
    sincos(A::AbstractMatrix)

Compute the matrix sine and cosine of a square matrix `A`.

# Examples
```jldoctest
julia> S, C = sincos(fill(1.0, (2,2)));

julia> S
2×2 Array{Float64,2}:
 0.454649  0.454649
 0.454649  0.454649

julia> C
2×2 Array{Float64,2}:
  0.291927  -0.708073
 -0.708073   0.291927
```
"""
function sincos(A::AbstractMatrix{<:Real})
    if issymmetric(A)
        symsinA, symcosA = sincos(Symmetric(A))
        sinA = copytri!(parent(symsinA), 'U')
        cosA = copytri!(parent(symcosA), 'U')
        return sinA, cosA
    end
    T = complex(float(eltype(A)))
    c, s = reim(exp!(T.(im .* A)))
    return s, c
end
function sincos(A::AbstractMatrix{<:Complex})
    if ishermitian(A)
        hermsinA, hermcosA = sincos(Hermitian(A))
        sinA = copytri!(parent(hermsinA), 'U', true)
        cosA = copytri!(parent(hermcosA), 'U', true)
        return sinA, cosA
    end
    T = complex(float(eltype(A)))
    X = exp!(T.(im .* A))
    Y = exp!(T.(.-im .* A))
    @inbounds for i in eachindex(X)
        x, y = X[i]/2, Y[i]/2
        X[i] = Complex(imag(x)-imag(y), real(y)-real(x))
        Y[i] = x+y
    end
    return X, Y
end

"""
    tan(A::AbstractMatrix)

Compute the matrix tangent of a square matrix `A`.

If `A` is symmetric or Hermitian, its eigendecomposition ([`eigfact`](@ref)) is used to
compute the tangent. Otherwise, the tangent is determined by calling [`exp`](@ref).

# Examples
```jldoctest
julia> tan(fill(1.0, (2,2)))
2×2 Array{Float64,2}:
 -1.09252  -1.09252
 -1.09252  -1.09252
```
"""
function tan(A::AbstractMatrix)
    if ishermitian(A)
        return copytri!(parent(tan(Hermitian(A))), 'U', true)
    end
    S, C = sincos(A)
    S /= C
    return S
end

"""
    cosh(A::AbstractMatrix)

Compute the matrix hyperbolic cosine of a square matrix `A`.
"""
function cosh(A::AbstractMatrix)
    if ishermitian(A)
        return copytri!(parent(cosh(Hermitian(A))), 'U', true)
    end
    X = exp(A)
    @. X = (X + $exp!(float(-A))) / 2
    return X
end

"""
    sinh(A::AbstractMatrix)

Compute the matrix hyperbolic sine of a square matrix `A`.
"""
function sinh(A::AbstractMatrix)
    if ishermitian(A)
        return copytri!(parent(sinh(Hermitian(A))), 'U', true)
    end
    X = exp(A)
    @. X = (X - $exp!(float(-A))) / 2
    return X
end

"""
    tanh(A::AbstractMatrix)

Compute the matrix hyperbolic tangent of a square matrix `A`.
"""
function tanh(A::AbstractMatrix)
    if ishermitian(A)
        return copytri!(parent(tanh(Hermitian(A))), 'U', true)
    end
    X = exp(A)
    Y = exp!(float.(.-A))
    @inbounds for i in eachindex(X)
        x, y = X[i], Y[i]
        X[i] = x - y
        Y[i] = x + y
    end
    X /= Y
    return X
end

"""
    acos(A::AbstractMatrix)

Compute the inverse matrix cosine of a square matrix `A`.

If `A` is symmetric or Hermitian, its eigendecomposition ([`eigfact`](@ref)) is used to
compute the inverse cosine. Otherwise, the inverse cosine is determined by using
[`log`](@ref) and [`sqrt`](@ref).  For the theory and logarithmic formulas used to compute
this function, see [^AH16_1].

[^AH16_1]: Mary Aprahamian and Nicholas J. Higham, "Matrix Inverse Trigonometric and Inverse Hyperbolic Functions: Theory and Algorithms", MIMS EPrint: 2016.4. [https://doi.org/10.1137/16M1057577](https://doi.org/10.1137/16M1057577)

# Examples
```jldoctest
julia> acos(cos([0.5 0.1; -0.2 0.3]))
2×2 Array{Complex{Float64},2}:
  0.5-5.55112e-17im  0.1-2.77556e-17im
 -0.2+2.498e-16im    0.3-3.46945e-16im
```
"""
function acos(A::AbstractMatrix)
    if ishermitian(A)
        acosHermA = acos(Hermitian(A))
        return isa(acosHermA, Hermitian) ? copytri!(parent(acosHermA), 'U', true) : parent(acosHermA)
    end
    SchurF = schurfact(complex(A))
    U = UpperTriangular(SchurF.T)
    R = triu!(parent(-im * log(U + im * sqrt(I - U^2))))
    return SchurF.Z * R * SchurF.Z'
end

"""
    asin(A::AbstractMatrix)

Compute the inverse matrix sine of a square matrix `A`.

If `A` is symmetric or Hermitian, its eigendecomposition ([`eigfact`](@ref)) is used to
compute the inverse sine. Otherwise, the inverse sine is determined by using [`log`](@ref)
and [`sqrt`](@ref).  For the theory and logarithmic formulas used to compute this function,
see [^AH16_2].

[^AH16_2]: Mary Aprahamian and Nicholas J. Higham, "Matrix Inverse Trigonometric and Inverse Hyperbolic Functions: Theory and Algorithms", MIMS EPrint: 2016.4. [https://doi.org/10.1137/16M1057577](https://doi.org/10.1137/16M1057577)

# Examples
```jldoctest
julia> asin(sin([0.5 0.1; -0.2 0.3]))
2×2 Array{Complex{Float64},2}:
  0.5-4.16334e-17im  0.1-5.55112e-17im
 -0.2+9.71445e-17im  0.3-1.249e-16im
```
"""
function asin(A::AbstractMatrix)
    if ishermitian(A)
        asinHermA = asin(Hermitian(A))
        return isa(asinHermA, Hermitian) ? copytri!(parent(asinHermA), 'U', true) : parent(asinHermA)
    end
    SchurF = schurfact(complex(A))
    U = UpperTriangular(SchurF.T)
    R = triu!(parent(-im * log(im * U + sqrt(I - U^2))))
    return SchurF.Z * R * SchurF.Z'
end

"""
    atan(A::AbstractMatrix)

Compute the inverse matrix tangent of a square matrix `A`.

If `A` is symmetric or Hermitian, its eigendecomposition ([`eigfact`](@ref)) is used to
compute the inverse tangent. Otherwise, the inverse tangent is determined by using
[`log`](@ref).  For the theory and logarithmic formulas used to compute this function, see
[^AH16_3].

[^AH16_3]: Mary Aprahamian and Nicholas J. Higham, "Matrix Inverse Trigonometric and Inverse Hyperbolic Functions: Theory and Algorithms", MIMS EPrint: 2016.4. [https://doi.org/10.1137/16M1057577](https://doi.org/10.1137/16M1057577)

# Examples
```jldoctest
julia> atan(tan([0.5 0.1; -0.2 0.3]))
2×2 Array{Complex{Float64},2}:
  0.5+1.38778e-17im  0.1-2.77556e-17im
 -0.2+6.93889e-17im  0.3-4.16334e-17im
```
"""
function atan(A::AbstractMatrix)
    if ishermitian(A)
        return copytri!(parent(atan(Hermitian(A))), 'U', true)
    end
    SchurF = schurfact(complex(A))
    U = im * UpperTriangular(SchurF.T)
    R = triu!(parent(log((I + U) / (I - U)) / 2im))
    return SchurF.Z * R * SchurF.Z'
end

"""
    acosh(A::AbstractMatrix)

Compute the inverse hyperbolic matrix cosine of a square matrix `A`.  For the theory and
logarithmic formulas used to compute this function, see [^AH16_4].

[^AH16_4]: Mary Aprahamian and Nicholas J. Higham, "Matrix Inverse Trigonometric and Inverse Hyperbolic Functions: Theory and Algorithms", MIMS EPrint: 2016.4. [https://doi.org/10.1137/16M1057577](https://doi.org/10.1137/16M1057577)
"""
function acosh(A::AbstractMatrix)
    if ishermitian(A)
        acoshHermA = acosh(Hermitian(A))
        return isa(acoshHermA, Hermitian) ? copytri!(parent(acoshHermA), 'U', true) : parent(acoshHermA)
    end
    SchurF = schurfact(complex(A))
    U = UpperTriangular(SchurF.T)
    R = triu!(parent(log(U + sqrt(U - I) * sqrt(U + I))))
    return SchurF.Z * R * SchurF.Z'
end

"""
    asinh(A::AbstractMatrix)

Compute the inverse hyperbolic matrix sine of a square matrix `A`.  For the theory and
logarithmic formulas used to compute this function, see [^AH16_5].

[^AH16_5]: Mary Aprahamian and Nicholas J. Higham, "Matrix Inverse Trigonometric and Inverse Hyperbolic Functions: Theory and Algorithms", MIMS EPrint: 2016.4. [https://doi.org/10.1137/16M1057577](https://doi.org/10.1137/16M1057577)
"""
function asinh(A::AbstractMatrix)
    if ishermitian(A)
        return copytri!(parent(asinh(Hermitian(A))), 'U', true)
    end
    SchurF = schurfact(complex(A))
    U = UpperTriangular(SchurF.T)
    R = triu!(parent(log(U + sqrt(I + U^2))))
    return SchurF.Z * R * SchurF.Z'
end

"""
    atanh(A::AbstractMatrix)

Compute the inverse hyperbolic matrix tangent of a square matrix `A`.  For the theory and
logarithmic formulas used to compute this function, see [^AH16_6].

[^AH16_6]: Mary Aprahamian and Nicholas J. Higham, "Matrix Inverse Trigonometric and Inverse Hyperbolic Functions: Theory and Algorithms", MIMS EPrint: 2016.4. [https://doi.org/10.1137/16M1057577](https://doi.org/10.1137/16M1057577)
"""
function atanh(A::AbstractMatrix)
    if ishermitian(A)
        return copytri!(parent(atanh(Hermitian(A))), 'U', true)
    end
    SchurF = schurfact(complex(A))
    U = UpperTriangular(SchurF.T)
    R = triu!(parent(log((I + U) / (I - U)) / 2))
    return SchurF.Z * R * SchurF.Z'
end

for (finv, f, finvh, fh, fn) in ((:sec, :cos, :sech, :cosh, "secant"),
                                 (:csc, :sin, :csch, :sinh, "cosecant"),
                                 (:cot, :tan, :coth, :tanh, "cotangent"))
    name = string(finv)
    hname = string(finvh)
    @eval begin
        @doc """
            $($name)(A::AbstractMatrix)

        Compute the matrix $($fn) of a square matrix `A`.
        """ ($finv)(A::AbstractMatrix{T}) where {T} = inv(($f)(A))
        @doc """
            $($hname)(A::AbstractMatrix)

        Compute the matrix hyperbolic $($fn) of square matrix `A`.
        """ ($finvh)(A::AbstractMatrix{T}) where {T} = inv(($fh)(A))
    end
end

for (tfa, tfainv, hfa, hfainv, fn) in ((:asec, :acos, :asech, :acosh, "secant"),
                                       (:acsc, :asin, :acsch, :asinh, "cosecant"),
                                       (:acot, :atan, :acoth, :atanh, "cotangent"))
    tname = string(tfa)
    hname = string(hfa)
    @eval begin
        @doc """
            $($tname)(A::AbstractMatrix)
        Compute the inverse matrix $($fn) of `A`. """ ($tfa)(A::AbstractMatrix{T}) where {T} = ($tfainv)(inv(A))
        @doc """
            $($hname)(A::AbstractMatrix)
        Compute the inverse matrix hyperbolic $($fn) of `A`. """ ($hfa)(A::AbstractMatrix{T}) where {T} = ($hfainv)(inv(A))
    end
end

"""
    factorize(A)

Compute a convenient factorization of `A`, based upon the type of the input matrix.
`factorize` checks `A` to see if it is symmetric/triangular/etc. if `A` is passed
as a generic matrix. `factorize` checks every element of `A` to verify/rule out
each property. It will short-circuit as soon as it can rule out symmetry/triangular
structure. The return value can be reused for efficient solving of multiple
systems. For example: `A=factorize(A); x=A\\b; y=A\\C`.

| Properties of `A`          | type of factorization                          |
|:---------------------------|:-----------------------------------------------|
| Positive-definite          | Cholesky (see [`cholfact`](@ref))  |
| Dense Symmetric/Hermitian  | Bunch-Kaufman (see [`bkfact`](@ref)) |
| Sparse Symmetric/Hermitian | LDLt (see [`ldltfact`](@ref))      |
| Triangular                 | Triangular                                     |
| Diagonal                   | Diagonal                                       |
| Bidiagonal                 | Bidiagonal                                     |
| Tridiagonal                | LU (see [`lufact`](@ref))            |
| Symmetric real tridiagonal | LDLt (see [`ldltfact`](@ref))      |
| General square             | LU (see [`lufact`](@ref))            |
| General non-square         | QR (see [`qrfact`](@ref))            |

If `factorize` is called on a Hermitian positive-definite matrix, for instance, then `factorize`
will return a Cholesky factorization.

# Examples
```jldoctest
julia> A = Array(Bidiagonal(fill(1.0, (5, 5)), :U))
5×5 Array{Float64,2}:
 1.0  1.0  0.0  0.0  0.0
 0.0  1.0  1.0  0.0  0.0
 0.0  0.0  1.0  1.0  0.0
 0.0  0.0  0.0  1.0  1.0
 0.0  0.0  0.0  0.0  1.0

julia> factorize(A) # factorize will check to see that A is already factorized
5×5 Bidiagonal{Float64,Array{Float64,1}}:
 1.0  1.0   ⋅    ⋅    ⋅
  ⋅   1.0  1.0   ⋅    ⋅
  ⋅    ⋅   1.0  1.0   ⋅
  ⋅    ⋅    ⋅   1.0  1.0
  ⋅    ⋅    ⋅    ⋅   1.0
```
This returns a `5×5 Bidiagonal{Float64}`, which can now be passed to other linear algebra functions
(e.g. eigensolvers) which will use specialized methods for `Bidiagonal` types.
"""
function factorize(A::StridedMatrix{T}) where T
    m, n = size(A)
    if m == n
        if m == 1 return A[1] end
        utri    = true
        utri1   = true
        herm    = true
        sym     = true
        for j = 1:n-1, i = j+1:m
            if utri1
                if A[i,j] != 0
                    utri1 = i == j + 1
                    utri = false
                end
            end
            if sym
                sym &= A[i,j] == A[j,i]
            end
            if herm
                herm &= A[i,j] == conj(A[j,i])
            end
            if !(utri1|herm|sym) break end
        end
        ltri = true
        ltri1 = true
        for j = 3:n, i = 1:j-2
            ltri1 &= A[i,j] == 0
            if !ltri1 break end
        end
        if ltri1
            for i = 1:n-1
                if A[i,i+1] != 0
                    ltri &= false
                    break
                end
            end
            if ltri
                if utri
                    return Diagonal(A)
                end
                if utri1
                    return Bidiagonal(diag(A), diag(A, -1), :L)
                end
                return LowerTriangular(A)
            end
            if utri
                return Bidiagonal(diag(A), diag(A, 1), :U)
            end
            if utri1
                if (herm & (T <: Complex)) | sym
                    try
                        return ldltfact!(SymTridiagonal(diag(A), diag(A, -1)))
                    end
                end
                return lufact(Tridiagonal(diag(A, -1), diag(A), diag(A, 1)))
            end
        end
        if utri
            return UpperTriangular(A)
        end
        if herm
            cf = cholfact(A)
            if cf.info == 0
                return cf
            else
                return factorize(Hermitian(A))
            end
        end
        if sym
            return factorize(Symmetric(A))
        end
        return lufact(A)
    end
    qrfact(A, Val(true))
end
factorize(A::Adjoint)   =   adjoint(factorize(parent(A)))
factorize(A::Transpose) = transpose(factorize(parent(A)))

## Moore-Penrose pseudoinverse

"""
    pinv(M[, tol::Real])

Computes the Moore-Penrose pseudoinverse.

For matrices `M` with floating point elements, it is convenient to compute
the pseudoinverse by inverting only singular values above a given threshold,
`tol`.

The optimal choice of `tol` varies both with the value of `M` and the intended application
of the pseudoinverse. The default value of `tol` is
`eps(real(float(one(eltype(M)))))*minumum(size(M))`, which is essentially machine epsilon
for the real part of a matrix element multiplied by the larger matrix dimension. For
inverting dense ill-conditioned matrices in a least-squares sense,
`tol = sqrt(eps(real(float(one(eltype(M))))))` is recommended.

For more information, see [^issue8859], [^B96], [^S84], [^KY88].

# Examples
```jldoctest
julia> M = [1.5 1.3; 1.2 1.9]
2×2 Array{Float64,2}:
 1.5  1.3
 1.2  1.9

julia> N = pinv(M)
2×2 Array{Float64,2}:
  1.47287   -1.00775
 -0.930233   1.16279

julia> M * N
2×2 Array{Float64,2}:
 1.0          -2.22045e-16
 4.44089e-16   1.0
```

[^issue8859]: Issue 8859, "Fix least squares", https://github.com/JuliaLang/julia/pull/8859

[^B96]: Åke Björck, "Numerical Methods for Least Squares Problems",  SIAM Press, Philadelphia, 1996, "Other Titles in Applied Mathematics", Vol. 51. [doi:10.1137/1.9781611971484](http://epubs.siam.org/doi/book/10.1137/1.9781611971484)

[^S84]: G. W. Stewart, "Rank Degeneracy", SIAM Journal on Scientific and Statistical Computing, 5(2), 1984, 403-413. [doi:10.1137/0905030](http://epubs.siam.org/doi/abs/10.1137/0905030)

[^KY88]: Konstantinos Konstantinides and Kung Yao, "Statistical analysis of effective singular values in matrix rank determination", IEEE Transactions on Acoustics, Speech and Signal Processing, 36(5), 1988, 757-763. [doi:10.1109/29.1585](https://doi.org/10.1109/29.1585)
"""
function pinv(A::StridedMatrix{T}, tol::Real) where T
    m, n = size(A)
    Tout = typeof(zero(T)/sqrt(one(T) + one(T)))
    if m == 0 || n == 0
        return Matrix{Tout}(undef, n, m)
    end
    if istril(A)
        if istriu(A)
            maxabsA = maximum(abs.(diag(A)))
            B = zeros(Tout, n, m)
            for i = 1:min(m, n)
                if abs(A[i,i]) > tol*maxabsA
                    Aii = inv(A[i,i])
                    if isfinite(Aii)
                        B[i,i] = Aii
                    end
                end
            end
            return B
        end
    end
    SVD         = svdfact(A, full = false)
    Stype       = eltype(SVD.S)
    Sinv        = zeros(Stype, length(SVD.S))
    index       = SVD.S .> tol*maximum(SVD.S)
    Sinv[index] = one(Stype) ./ SVD.S[index]
    Sinv[findall(.!isfinite.(Sinv))] = zero(Stype)
    return SVD.Vt' * (Diagonal(Sinv) * SVD.U')
end
function pinv(A::StridedMatrix{T}) where T
    tol = eps(real(float(one(T))))*min(size(A)...)
    return pinv(A, tol)
end
function pinv(x::Number)
    xi = inv(x)
    return ifelse(isfinite(xi), xi, zero(xi))
end

## Basis for null space

"""
    nullspace(M[, tol::Real])

Computes a basis for the nullspace of `M` by including the singular
vectors of A whose singular have magnitude are greater than `tol*σ₁`,
where `σ₁` is `A`'s largest singular values. By default, the value of
`tol` is the smallest dimension of `A` multiplied by the [`eps`](@ref)
of the [`eltype`](@ref) of `A`.

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

julia> nullspace(M)
3×1 Array{Float64,2}:
 0.0
 0.0
 1.0

julia> nullspace(M, 2)
3×3 Array{Float64,2}:
 0.0  1.0  0.0
 1.0  0.0  0.0
 0.0  0.0  1.0
```
"""
function nullspace(A::StridedMatrix, tol::Real = min(size(A)...)*eps(real(float(one(eltype(A))))))
    m, n = size(A)
    (m == 0 || n == 0) && return Matrix{T}(I, n, n)
    SVD = svdfact(A, full=true)
    indstart = sum(SVD.S .> SVD.S[1]*tol) + 1
    return copy(SVD.Vt[indstart:end,:]')
end
nullspace(a::StridedVector, tol::Real = min(size(a)...)*eps(real(float(one(eltype(a)))))) = nullspace(reshape(a, length(a), 1), tol)

"""
    cond(M, p::Real=2)

Condition number of the matrix `M`, computed using the operator `p`-norm. Valid values for
`p` are `1`, `2` (default), or `Inf`.
"""
function cond(A::AbstractMatrix, p::Real=2)
    if p == 2
        v = svdvals(A)
        maxv = maximum(v)
        return maxv == 0.0 ? oftype(real(A[1,1]),Inf) : maxv / minimum(v)
    elseif p == 1 || p == Inf
        checksquare(A)
        return _cond1Inf(A, p)
    end
    throw(ArgumentError("p-norm must be 1, 2 or Inf, got $p"))
end
_cond1Inf(A::StridedMatrix{<:BlasFloat}, p::Real) = _cond1Inf(lufact(A), p, norm(A, p))
_cond1Inf(A::AbstractMatrix, p::Real)             = norm(A, p)*norm(inv(A), p)

## Lyapunov and Sylvester equation

# AX + XB + C = 0

"""
    sylvester(A, B, C)

Computes the solution `X` to the Sylvester equation `AX + XB + C = 0`, where `A`, `B` and
`C` have compatible dimensions and `A` and `-B` have no eigenvalues with equal real part.

# Examples
```jldoctest
julia> A = [3. 4.; 5. 6]
2×2 Array{Float64,2}:
 3.0  4.0
 5.0  6.0

julia> B = [1. 1.; 1. 2.]
2×2 Array{Float64,2}:
 1.0  1.0
 1.0  2.0

julia> C = [1. 2.; -2. 1]
2×2 Array{Float64,2}:
  1.0  2.0
 -2.0  1.0

julia> X = sylvester(A, B, C)
2×2 Array{Float64,2}:
 -4.46667   1.93333
  3.73333  -1.8

julia> A*X + X*B + C
2×2 Array{Float64,2}:
  2.66454e-15  1.77636e-15
 -3.77476e-15  4.44089e-16
```
"""
function sylvester(A::StridedMatrix{T},B::StridedMatrix{T},C::StridedMatrix{T}) where T<:BlasFloat
    RA, QA = schur(A)
    RB, QB = schur(B)

    D = -(adjoint(QA) * (C*QB))
    Y, scale = LAPACK.trsyl!('N','N', RA, RB, D)
    rmul!(QA*(Y * adjoint(QB)), inv(scale))
end
sylvester(A::StridedMatrix{T}, B::StridedMatrix{T}, C::StridedMatrix{T}) where {T<:Integer} = sylvester(float(A), float(B), float(C))

sylvester(a::Union{Real,Complex}, b::Union{Real,Complex}, c::Union{Real,Complex}) = -c / (a + b)

# AX + XA' + C = 0

"""
    lyap(A, C)

Computes the solution `X` to the continuous Lyapunov equation `AX + XA' + C = 0`, where no
eigenvalue of `A` has a zero real part and no two eigenvalues are negative complex
conjugates of each other.

# Examples
```jldoctest
julia> A = [3. 4.; 5. 6]
2×2 Array{Float64,2}:
 3.0  4.0
 5.0  6.0

julia> B = [1. 1.; 1. 2.]
2×2 Array{Float64,2}:
 1.0  1.0
 1.0  2.0

julia> X = lyap(A, B)
2×2 Array{Float64,2}:
  0.5  -0.5
 -0.5   0.25

julia> A*X + X*A' + B
2×2 Array{Float64,2}:
 0.0          6.66134e-16
 6.66134e-16  8.88178e-16
```
"""
function lyap(A::StridedMatrix{T}, C::StridedMatrix{T}) where {T<:BlasFloat}
    R, Q = schur(A)

    D = -(adjoint(Q) * (C*Q))
    Y, scale = LAPACK.trsyl!('N', T <: Complex ? 'C' : 'T', R, R, D)
    rmul!(Q*(Y * adjoint(Q)), inv(scale))
end
lyap(A::StridedMatrix{T}, C::StridedMatrix{T}) where {T<:Integer} = lyap(float(A), float(C))
lyap(a::T, c::T) where {T<:Number} = -c/(2a)

max_stars_repo_name: djsegal/julia-fork

max_issues_repo_name: djsegal/julia-fork

max_forks_repo_name: djsegal/julia-fork


text:
using GenerativeModels
using POMDPs
using POMDPToolbox
using ProfileView
using Multilane

#Set up problem configuration
nb_lanes = 2
pp = PhysicalParam(nb_lanes,lane_length=200.) #2.=>col_length=8
r_crash = -1.
accel_cost = -1e-5
decel_cost = -5e-6
invalid_cost = -1e-5
lineride_cost = -1e-5
lanechange_cost = -2e-5
_discount = 1.
nb_cars=3

rng = MersenneTwister(9)

rmodel = OriginalRewardModel(r_crash,
                             accel_cost,
                             decel_cost,
                             invalid_cost,
                             lineride_cost,
                             lanechange_cost)

dmodel = IDMMOBILModel(nb_cars, pp)

mdp = OriginalMDP(dmodel, rmodel, _discount)

policy = RandomPolicy(mdp)

sim = RolloutSimulator(max_steps=100)

@time simulate(sim, mdp, policy, initial_state(mdp, rng))
@time for i in 1:100
    simulate(sim, mdp, policy, initial_state(mdp, rng))
end


text:
println("===================================================================")
println("Beginning GMHExamples.jl/ode/predatorprey/scripts/PredatorPrey1.jl")
println("===================================================================")

import GeneralizedMetropolisHastings #pre-include to avoid race condition when running in parallel
import GMHExamples

println("Number of parallel processes running: ",nprocs())
println("Use addprocs() in the REPL if you want to run on more processes")

@everywhere using GeneralizedMetropolisHastings
@everywhere using GMHExamples
using PyPlot

#Standard M-H for nproposals == 1
#Generalized M-H for nproposals > 1
nproposals = 200

#MCMC iteration specifications
nburnin = 500
niterations = 1000
ntunerperiod = 20

###Initial conditions for the ODE (prey and predator populations)
initial = [50.0,5.0]

###Default values of the parameters and prior boundaries
#defaults = [0.36,107.38,0.87,53.36,0.64,0.27]
defaults = [0.4,107.0,0.9,53.0,0.7,0.3]
lows = zeros(6)
highs = 150*ones(6)

###The variance of the noise on the input data
variance = sqrt(10.0)*ones(2)

println("==========================================")
println("Simulation parameters defined successfully")
println("==========================================")

###Create a FitzHugh-Nagumo model with measurement data, ODE function and parameters with default values and priors
m = predatorpreymodel(initial,variance,lows,highs,defaults)

###Show the model
println("==========================")
println("Model defined successfully")
println("==========================")
show(m)

###Plot the measurement data (simmulated data + noise)
figure(string("PredatorPrey1-",nprocs())) ; clf()
subplot(221)
plot(dataindex(m),measurements(m)[:,1];label="Prey")
plot(dataindex(m),measurements(m)[:,2];label="Predator")
xlabel("Time")
ylabel("Amplitude")
title("Predator-Prey measurement data")
grid("on")
legend(loc="upper right",fancybox="true")

###Create a Metropolis Sampler with normal proposal density
s1 = sampler(:mh,:normal,1.0,[0.01,1.0,0.01,1.0,0.01,0.01])
s2 = sampler(:adaptive,0.01,numparas(m))
println("=============================")
println("Samplers defined successfully")
println("=============================")
show(s1)
show(s2)

###Create a tuner that scales the proposal density
t1 = tuner(:scale,ntunerperiod,0.5,:erf)
t2 = tuner(:monitor,ntunerperiod)
println("======+====================")
println("Tuners defined successfully")
println("=======+===================")
show(t1)
show(t2)

###Create a Generalized Metropolis-Hastings runner (which will default to Standard MH when nproposals=1)
p = policy(:mh,nproposals;initialize=:default)
r = runner(p,niterations,nproposals;numburnin = nburnin)
println("===========================")
println("Runner defined successfully")
println("===========================")
show(r)

###Run the MCMC (can take quite a bit of time)
println("=======================")
println("Run the MCMC simulation")
println("=======================")
@time c1 = run!(r,m,s1,t1)
@time c2 = run!(r,m,s2,t2)
println("==========================")
println("Completed MCMC simulation")
println("=========================")

###Show the results of the simulations
show(c1)
show(c2)

nparas = numparas(m)
meanparamvals1 = mean(samples(c1),2)
stdparamvals1 = std(samples(c1),2)
meanparamvals2 = mean(samples(c2),2)
stdparamvals2 = std(samples(c2),2)

println("Results of the MH MCMC simulation:")
for i=1:nparas
    println(" parameter $(parameters(m)[i].key):  mean = ",meanparamvals1[i]," std = ",stdparamvals1[i])
end

println("Results of the Adaptive MCMC simulation:")
for i=1:nparas
    println(" parameter $(parameters(m)[i].key):  mean = ",meanparamvals2[i]," std = ",stdparamvals2[i])
end

println("================")
println("Plotting results")
println("================")

###Plot the loglikelihood values across samples
###After an initial few low values, this should remain relatively high
subplot(222)
plot(1:numsamples(c1),logposterior(c1,m))
title("Normal Metropolis-Hastings")
xlabel("Samples")
ylabel("Log-Posterior")

subplot(224)
plot(1:numsamples(c2),logposterior(c2,m))
title("Adaptive Metropolis-Hastings")
xlabel("Samples")
ylabel("Log-Posterior")

subplot(223)
modeldata1 = evaluate!(m,vec(meanparamvals1))
modeldata2 = evaluate!(m,vec(meanparamvals2))
plot(dataindex(m),measurements(m)[:,1];label="Prey")
plot(dataindex(m),measurements(m)[:,2];label="Predator")
plot(dataindex(m),modeldata1[:,1];label="MH Prey")
plot(dataindex(m),modeldata1[:,2];label="MH Predator")
plot(dataindex(m),modeldata2[:,1];label="AMH Prey")
plot(dataindex(m),modeldata2[:,2];label="AMH Predator")
xlabel("Time")
ylabel("Amplitude")
title("Predator-Prey model data")
grid("on")
legend(loc="upper right",fancybox="true")

###Plot the histograms of parameter values
figure(string("PredatorPrey-MH-",nprocs())) ; clf()
for i=1:nparas
    subplot(230 + i)
    h = PyPlot.plt[:hist](vec(getindex(samples(c1),i,:)),20)
    grid("on")
    title("Parameter $(parameters(m)[i].key)")
    xlabel("Values")
    ylabel("Samples")
end

figure(string("PredatorPrey-AMH-",nprocs())) ; clf()
for i=1:nparas
    subplot(230 + i)
    h = PyPlot.plt[:hist](vec(getindex(samples(c2),i,:)),20)
    grid("on")
    title("Parameter $(parameters(m)[i].key)")
    xlabel("Values")
    ylabel("Samples")
end


text:
# # [Wind- and convection-driven mixing in an ocean surface boundary layer](@id gpu_example)
#
# This example simulates mixing by three-dimensional turbulence in an ocean surface
# boundary layer driven by atmospheric winds and convection. It demonstrates:
#
#   * How to set-up a grid with varying spacing in the vertical direction
#   * How to use the `SeawaterBuoyancy` model for buoyancy with a linear equation of state.
#   * How to use a turbulence closure for large eddy simulation.
#   * How to use a function to impose a boundary condition.
#
# ## Install dependencies
#
# First let's make sure we have all required packages installed.

# ```julia
# using Pkg
# pkg"add Oceananigans, CairoMakie"
# ```

# We start by importing all of the packages and functions that we'll need for this
# example.

using Random
using Printf
using CairoMakie

using Oceananigans
using Oceananigans.Units: minute, minutes, hour

# ## The grid
#
# We use 32²×24 grid points with 2 m grid spacing in the horizontal and
# varying spacing in the vertical, with higher resolution closer to the
# surface. Here we use a stretching function for the vertical nodes that
# maintains relatively constant vertical spacing in the mixed layer, which
# is desirable from a numerical standpoint:

Nz = 24          # number of points in the vertical direction
Lz = 32          # (m) domain depth

refinement = 1.2 # controls spacing near surface (higher means finer spaced)
stretching = 12  # controls rate of stretching at bottom

## Normalized height ranging from 0 to 1
h(k) = (k - 1) / Nz

## Linear near-surface generator
ζ₀(k) = 1 + (h(k) - 1) / refinement

## Bottom-intensified stretching function
Σ(k) = (1 - exp(-stretching * h(k))) / (1 - exp(-stretching))

## Generating function
z_faces(k) = Lz * (ζ₀(k) * Σ(k) - 1)

grid = RectilinearGrid(size = (32, 32, Nz),
                          x = (0, 64),
                          y = (0, 64),
                          z = z_faces)

# We plot vertical spacing versus depth to inspect the prescribed grid stretching:

fig = Figure(resolution=(1200, 800))
ax = Axis(fig[1, 1], ylabel = "Depth (m)", xlabel = "Vertical spacing (m)")
lines!(ax, grid.Δzᵃᵃᶜ[1:grid.Nz], grid.zᵃᵃᶜ[1:grid.Nz])
scatter!(ax, grid.Δzᵃᵃᶜ[1:Nz], grid.zᵃᵃᶜ[1:Nz])

save("ocean_wind_mixing_convection_grid_spacing.svg", fig)
nothing #hide

# ![](ocean_wind_mixing_convection_grid_spacing.svg)

# ## Buoyancy that depends on temperature and salinity
#
# We use the `SeawaterBuoyancy` model with a linear equation of state,

buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(thermal_expansion = 2e-4,
                                                                    haline_contraction = 8e-4))

# ## Boundary conditions
#
# We calculate the surface temperature flux associated with surface heating of
# 200 W m⁻², reference density `ρₒ`, and heat capacity `cᴾ`,

Qʰ = 200.0  # W m⁻², surface _heat_ flux
ρₒ = 1026.0 # kg m⁻³, average density at the surface of the world ocean
cᴾ = 3991.0 # J K⁻¹ kg⁻¹, typical heat capacity for seawater

Qᵀ = Qʰ / (ρₒ * cᴾ) # K m s⁻¹, surface _temperature_ flux

# Finally, we impose a temperature gradient `dTdz` both initially and at the
# bottom of the domain, culminating in the boundary conditions on temperature,

dTdz = 0.01 # K m⁻¹

T_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵀ),
                                bottom = GradientBoundaryCondition(dTdz))

# Note that a positive temperature flux at the surface of the ocean
# implies cooling. This is because a positive temperature flux implies
# that temperature is fluxed upwards, out of the ocean.
#
# For the velocity field, we imagine a wind blowing over the ocean surface
# with an average velocity at 10 meters `u₁₀`, and use a drag coefficient `cᴰ`
# to estimate the kinematic stress (that is, stress divided by density) exerted
# by the wind on the ocean:

u₁₀ = 10    # m s⁻¹, average wind velocity 10 meters above the ocean
cᴰ = 2.5e-3 # dimensionless drag coefficient
ρₐ = 1.225  # kg m⁻³, average density of air at sea-level

Qᵘ = - ρₐ / ρₒ * cᴰ * u₁₀ * abs(u₁₀) # m² s⁻²

# The boundary conditions on `u` are thus

u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))

# For salinity, `S`, we impose an evaporative flux of the form

@inline Qˢ(x, y, t, S, evaporation_rate) = - evaporation_rate * S # [salinity unit] m s⁻¹
nothing # hide

# where `S` is salinity. We use an evporation rate of 1 millimeter per hour,

evaporation_rate = 1e-3 / hour # m s⁻¹

# We build the `Flux` evaporation `BoundaryCondition` with the function `Qˢ`,
# indicating that `Qˢ` depends on salinity `S` and passing
# the parameter `evaporation_rate`,

evaporation_bc = FluxBoundaryCondition(Qˢ, field_dependencies=:S, parameters=evaporation_rate)

# The full salinity boundary conditions are

S_bcs = FieldBoundaryConditions(top=evaporation_bc)

# ## Model instantiation
#
# We fill in the final details of the model here: upwind-biased 5th-order
# advection for momentum and tracers, 3rd-order Runge-Kutta time-stepping,
# Coriolis forces, and the `AnisotropicMinimumDissipation` closure
# for large eddy simulation to model the effect of turbulent motions at
# scales smaller than the grid scale that we cannot explicitly resolve.

model = NonhydrostaticModel(; grid, buoyancy,
                            advection = UpwindBiasedFifthOrder(),
                            timestepper = :RungeKutta3,
                            tracers = (:T, :S),
                            coriolis = FPlane(f=1e-4),
                            closure = AnisotropicMinimumDissipation(),
                            boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs))

# Notes:
#
# * To use the Smagorinsky-Lilly turbulence closure (with a constant model coefficient) rather than
#   `AnisotropicMinimumDissipation`, use `closure = SmagorinskyLilly()` in the model constructor.
#
# * To change the `architecture` to `GPU`, replace `architecture = CPU()` with
#   `architecture = GPU()`.

# ## Initial conditions
#
# Our initial condition for temperature consists of a linear stratification superposed with
# random noise damped at the walls, while our initial condition for velocity consists
# only of random noise.

## Random noise damped at top and bottom
Ξ(z) = randn() * z / model.grid.Lz * (1 + z / model.grid.Lz) # noise

## Temperature initial condition: a stable density gradient with random noise superposed.
Tᵢ(x, y, z) = 20 + dTdz * z + dTdz * model.grid.Lz * 1e-6 * Ξ(z)

## Velocity initial condition: random noise scaled by the friction velocity.
uᵢ(x, y, z) = sqrt(abs(Qᵘ)) * 1e-3 * Ξ(z)

## `set!` the `model` fields using functions or constants:
set!(model, u=uᵢ, w=uᵢ, T=Tᵢ, S=35)

# ## Setting up a simulation
#
# We set-up a simulation with an initial time-step of 10 seconds
# that stops at 40 minutes, with adaptive time-stepping and progress printing.

simulation = Simulation(model, Δt=10.0, stop_time=40minutes)

# The `TimeStepWizard` helps ensure stable time-stepping
# with a Courant-Freidrichs-Lewy (CFL) number of 1.0.

wizard = TimeStepWizard(cfl=1.0, max_change=1.1, max_Δt=1minute)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

# Nice progress messaging is helpful:

## Print a progress message
progress_message(sim) = @printf("Iteration: %04d, time: %s, Δt: %s, max(|w|) = %.1e ms⁻¹, wall time: %s\n",
                                iteration(sim), prettytime(sim), prettytime(sim.Δt),
                                maximum(abs, sim.model.velocities.w), prettytime(sim.run_wall_time))

simulation.callbacks[:progress] = Callback(progress_message, IterationInterval(20))

# We then set up the simulation:

# ## Output
#
# We use the `JLD2OutputWriter` to save ``x, z`` slices of the velocity fields,
# tracer fields, and eddy diffusivities. The `prefix` keyword argument
# to `JLD2OutputWriter` indicates that output will be saved in
# `ocean_wind_mixing_and_convection.jld2`.

## Create a NamedTuple with eddy viscosity
eddy_viscosity = (; νₑ = model.diffusivity_fields.νₑ)

filename = "ocean_wind_mixing_and_convection"

simulation.output_writers[:slices] =
    JLD2OutputWriter(model, merge(model.velocities, model.tracers, eddy_viscosity),
                     filename = filename * ".jld2",
                     indices = (:, grid.Ny/2, :),
                     schedule = TimeInterval(1minute),
                     overwrite_existing = true)

# We're ready:

run!(simulation)

# ## Turbulence visualization
#
# We animate the data saved in `ocean_wind_mixing_and_convection.jld2`.
# We prepare for animating the flow by loading the data into
# FieldTimeSeries and defining functions for computing colorbar limits.

filepath = filename * ".jld2"

time_series = (w = FieldTimeSeries(filepath, "w"),
               T = FieldTimeSeries(filepath, "T"),
               S = FieldTimeSeries(filepath, "S"),
               νₑ = FieldTimeSeries(filepath, "νₑ"))

## Coordinate arrays
xw, yw, zw = nodes(time_series.w)
xT, yT, zT = nodes(time_series.T)

""" Return colorbar levels equispaced between `(-clim, clim)` and encompassing the extrema of `c`. """
function divergent_levels(c, clim, nlevels=21)
    cmax = maximum(abs, c)
    levels = clim > cmax ? range(-clim, stop=clim, length=nlevels) : range(-cmax, stop=cmax, length=nlevels)

    return (levels[1], levels[end]), levels
end

""" Return colorbar levels equispaced between `clims` and encompassing the extrema of `c`."""
function sequential_levels(c, clims, nlevels=20)
    levels = range(clims[1], stop=clims[2], length=nlevels)
    cmin, cmax = minimum(c), maximum(c)
    cmin < clims[1] && (levels = vcat([cmin], levels))
    cmax > clims[2] && (levels = vcat(levels, [cmax]))

    return clims, levels
end
nothing # hide

# We start the animation at ``t = 10minutes`` since things are pretty boring till then:

times = time_series.w.times
intro = searchsortedfirst(times, 10minutes)

# We are now ready to animate using Makie. We use Makie's `Observable` to animate
# the data. To dive into how `Observable`s work we refer to
# [Makie.jl's Documentation](https://makie.juliaplots.org/stable/documentation/nodes/index.html).

n = Observable(intro)

 wₙ = @lift interior(time_series.w[$n],  :, 1, :)
 Tₙ = @lift interior(time_series.T[$n],  :, 1, :)
 Sₙ = @lift interior(time_series.S[$n],  :, 1, :)
νₑₙ = @lift interior(time_series.νₑ[$n], :, 1, :)

fig = Figure(resolution = (1000, 500))

axis_kwargs = (xlabel="x (m)",
               ylabel="z (m)",
               aspect = AxisAspect(grid.Lx/grid.Lz),
               limits = ((0, grid.Lx), (-grid.Lz, 0)))

ax_w  = Axis(fig[2, 1]; title = "Vertical velocity", axis_kwargs...)
ax_T  = Axis(fig[2, 3]; title = "Temperature", axis_kwargs...)
ax_S  = Axis(fig[3, 1]; title = "Salinity", axis_kwargs...)
ax_νₑ = Axis(fig[3, 3]; title = "Eddy viscocity", axis_kwargs...)

title = @lift @sprintf("t = %s", prettytime(times[$n]))

wlims = (-0.05, 0.05)
Tlims = (19.7, 19.99)
Slims = (35, 35.005)
νₑlims = (1e-6, 5e-3)

hm_w = heatmap!(ax_w, xw, zw, wₙ; colormap = :balance, colorrange = wlims)
Colorbar(fig[2, 2], hm_w; label = "m s⁻¹")

hm_T = heatmap!(ax_T, xT, zT, Tₙ; colormap = :thermal, colorrange = Tlims)
Colorbar(fig[2, 4], hm_T; label = "ᵒC")

hm_S = heatmap!(ax_S, xT, zT, Sₙ; colormap = :haline, colorrange = Slims)
Colorbar(fig[3, 2], hm_S; label = "g / kg")

hm_νₑ = heatmap!(ax_νₑ, xT, zT, νₑₙ; colormap = :thermal, colorrange = νₑlims)
Colorbar(fig[3, 4], hm_νₑ; label = "m s⁻²")

fig[1, 1:4] = Label(fig, title, textsize=24, tellwidth=false)

# And now record a movie.

frames = intro:length(times)

@info "Making a motion picture of ocean wind mixing and convection..."

record(fig, filename * ".mp4", frames, framerate=8) do i
    msg = string("Plotting frame ", i, " of ", frames[end])
    print(msg * " \r")
    n[] = i
end
nothing #hide

# ![](ocean_wind_mixing_and_convection.mp4)


text:
module Equations

    export Equation
    export AbstractEquationODE, AbstractEquationPODE,
           AbstractEquationDAE, AbstractEquationPDAE,
           AbstractEquationSDE, AbstractEquationPSDE

    export periodicity

    export ODE, IODE, PODE, HODE, VODE, SODE
    export DAE, IDAE, PDAE, HDAE, VDAE
    export SDE, PSDE, SPSDE

    include("equations/equations.jl")

    include("equations/ode.jl")
    include("equations/hode.jl")
    include("equations/iode.jl")
    include("equations/pode.jl")
    include("equations/vode.jl")

    include("equations/dae.jl")
    include("equations/hdae.jl")
    include("equations/idae.jl")
    include("equations/pdae.jl")
    include("equations/vdae.jl")

    include("equations/sode.jl")

    include("equations/sde.jl")
    include("equations/psde.jl")
    include("equations/spsde.jl")

end


text:
using MakieCore
using Documenter

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

makedocs(;
    modules=[MakieCore],
    authors="Simon Danisch",
    repo="https://github.com/SimonDanisch/MakieCore.jl/blob/{commit}{path}#{line}",
    sitename="MakieCore.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://SimonDanisch.github.io/MakieCore.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/SimonDanisch/MakieCore.jl",
)


text:
include("./imports.jl")

# TODO: test apply!(Sequential)

a = randn(5, 5)
item = ArrayItem(a)

tfm = MapElem(x -> x + 1)

buf = apply(tfm, item)

apply!(buf, tfm, item)

DataAugmentation.copyitemdata!([item], [buf])
item
buf

@testset ExtendedTestSet "apply!(buf, ::Map, ...)" begin
    newitem() = ArrayItem(randn(5, 5))
    tfm = MapElem(x -> x + 1)

    buf = apply(tfm, newitem())
    buf_copy = deepcopy(buf)

    @test buf.data ≈ buf_copy.data
    @test_nowarn apply!(buf, tfm, newitem())
    @test !(buf.data ≈ buf_copy.data)
end

@testset ExtendedTestSet "copyitemdata!" begin
    newitem() = ArrayItem(randn(5, 5))
    @testset ExtendedTestSet "" begin
        a = newitem()
        b = newitem()
        @test !(a.data ≈ b.data)
        DataAugmentation.copyitemdata!(a, b)
        @test a.data ≈ b.data
    end

    @testset ExtendedTestSet "" begin
        a = newitem()
        b = newitem()
        @test !(a.data ≈ b.data)
        DataAugmentation.copyitemdata!([a], [b])
        @test a.data ≈ b.data
    end

    @testset ExtendedTestSet "" begin
        a = newitem()
        b = newitem()
        @test !(a.data ≈ b.data)
        DataAugmentation.copyitemdata!((a,), (b,))
        @test a.data ≈ b.data
    end
end

@testset ExtendedTestSet "`Buffered`" begin
    newitem() = ArrayItem(randn(5, 5))
    # buffer should be created
    tb = Buffered(MapElem(x -> x + 1))
    @test isnothing(tb.buffer)
    @test_nowarn apply(tb, newitem())
    @test !isnothing(tb.buffer)

    # buffer should change
    buf = deepcopy(tb.buffer)
    @test_nowarn apply(tb, newitem())
    @test !(buf[1].data ≈ tb.buffer[1].data)

    # inplace version
    buf2 = deepcopy(buf)
    @test_nowarn apply!(buf2, tb, newitem())
    @test !(buf2[1].data ≈ tb.buffer[1].data)
end

@testset ExtendedTestSet "`BufferedThreadsafe`" begin
    newitem() = ArrayItem(randn(5, 5))
    # buffer should be created
    tbt = BufferedThreadsafe(MapElem(x -> x + 1))
    tb = tbt.buffereds[1]
    @test isnothing(tb.buffer)
    @test_nowarn apply(tbt, newitem())
    @test !isnothing(tb.buffer)

    # buffer should change
    buf = deepcopy(tb.buffer)
    @test_nowarn apply(tbt, newitem())
    @test !(buf[1].data ≈ tb.buffer[1].data)

    # inplace version
    buf2 = deepcopy(buf)
    @test_nowarn apply!(buf2, tbt, newitem())
    @test !(buf2[1].data ≈ tb.buffer[1].data)
end


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

name = "SEAL"
version = v"3.5.6"

# Collection of sources required to complete build
sources = [
    ArchiveSource("https://github.com/microsoft/SEAL/archive/v$(version).tar.gz",
                  "13674a39a48c0d1c6ff544521cf10ee539ce1af75c02bfbe093f7621869e3406"),
    DirectorySource("./bundled")
]

# Bash recipe for building across all platforms
script = raw"""
cd $WORKSPACE/srcdir/SEAL-*

# Apply patches that fix some cross-compilation issues
for f in ${WORKSPACE}/srcdir/patches/*.patch; do
    atomic_patch -p1 ${f}
done

# The last three '-DSEAL_USE__*' flags are required to circumvent
# cross-compilation issues
if [[ "${target}" == *-darwin* ]]; then
  # C++17 is disabled on MacOS due to the environment being too old.
  cmake . \
    -DCMAKE_INSTALL_PREFIX=$prefix \
    -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN%.*}_clang.cmake \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
    -DBUILD_SHARED_LIBS=ON \
    -DSEAL_BUILD_SEAL_C=ON \
    -DSEAL_USE___BUILTIN_CLZLL=OFF \
    -DSEAL_USE__ADDCARRY_U64=OFF \
    -DSEAL_USE__SUBBORROW_U64=OFF \
    -DSEAL_USE_CXX17=OFF
else
  cmake . \
    -DCMAKE_INSTALL_PREFIX=$prefix \
    -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN%.*}_clang.cmake \
    -DCMAKE_BUILD_TYPE=Release \
    -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
    -DBUILD_SHARED_LIBS=ON \
    -DSEAL_BUILD_SEAL_C=ON \
    -DSEAL_USE___BUILTIN_CLZLL=OFF \
    -DSEAL_USE__ADDCARRY_U64=OFF \
    -DSEAL_USE__SUBBORROW_U64=OFF
fi
make -j${nproc}
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    Linux(:x86_64, libc=:musl),
    MacOS(:x86_64),
    Linux(:x86_64, libc=:glibc),
    Linux(:aarch64, libc=:glibc),
    FreeBSD(:x86_64),
    Linux(:powerpc64le, libc=:glibc),
    Linux(:aarch64, libc=:musl)
]

# Fix incompatibilities across the GCC 4/5 version boundary due to std::string,
# as suggested by the wizard
platforms = expand_cxxstring_abis(platforms)

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

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
]

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


text:
s = open("input.txt") do file
    readlines(file);
end

function ParseRule(rule)
    m = match(r"(\w+) (\w+) bags contain", rule);

    color = m.captures[1] * " " * m.captures[2];

    bags = [];
    counts = [];
    for m in eachmatch(r"(\d+) (\w+) (\w+) bag", rule)
        push!(bags, join(m.captures[2:3], " "))
        push!(counts, parse(Int, m.captures[1]))
    end

    return (color, bags, counts);
end

function CreateColorGraph(lines)
    bagGraph = Dict();
    for line in lines
        rule = ParseRule(line);

        for i = 1:length(rule[2])
            color = rule[2][i];
            count = rule[3][i];
            push!(get!(bagGraph, color, []), (rule[1], count));
        end
    end

    return bagGraph;
end

graph = CreateColorGraph(s);

colorStack = ["shiny gold"];
seen = Set();

while length(colorStack) > 0
    curColor = pop!(colorStack);
    parent = graph[curColor];
    for parentColor in parent
        if in(parentColor[1], seen)
            continue;
        end

        push!(seen, parentColor[1]);

        if !haskey(graph, parentColor[1])
            continue
        end

        push!(colorStack, parentColor[1]);
    end
end

println(length(seen));


text:
module VXI11

include("ONCRPC.jl")
using .ONCRPC
include("rpc.jl")

const LOCK_TIMEOUT = 10000
const IO_TIMEOUT = 10000

mutable struct Link
    c::ONCRPC.Channel
    lid::UInt32
    maxRecvSize::UInt32
end

close(l::Link) = destroy_link(l.c, l.lid)

function Link(c::ONCRPC.Channel, device="inst0")
    link = create_link(c, Create_LinkParms(rand(Int32), false, 10000, device))
    @assert link.error == 0
    l = Link(c, link.link, link.maxRecvSize)
    finalizer(l, close)
    l
end

function Link(ip::Base.IPAddr, port, device="inst0")
    Link(ONCRPC.Channel(connect(ip, port)), device)
end

function Base.write(l::Link, data)
    buf = IOBuffer()
    write(buf, data)
    data = take!(buf)
    resp = device_write(l.c, Device_WriteParms(l.lid, UInt32(IO_TIMEOUT), UInt32(LOCK_TIMEOUT),
        UInt32(1) << 3, data))
    @assert resp.error == 0
    @assert resp.size == sizeof(data)
    sizeof(data)
end

function Base.read(l::Link)
    rr = device_read(l.c, Device_ReadParms(l.lid, UInt32(l.maxRecvSize), UInt32(IO_TIMEOUT), UInt32(LOCK_TIMEOUT), UInt32(0), UInt32(0)))
    rr.data
end

end


text:
struct ConstantPolicy <: Algorithm end

function train(learner::Learner{ConstantPolicy})
    train_samples = learner.train_data.samples
    best_r = 0.0
    best_r_reward = -Inf
    for _sample in train_samples
        r = _sample.bid_price
        r_reward = 0.0
        for sample in train_samples
            pay_price = sample.pay_price
            bid_price = sample.bid_price
            if r <= pay_price
                r_reward += pay_price
            elseif pay_price < r <= bid_price
                r_reward += r
            else
                @assert r > bid_price
            end
        end
        if r_reward > best_r_reward
            best_r = r
            best_r_reward = r_reward
        end
    end
    linear_model = LinearModel(zeros(feature_count(learner.train_data)), best_r)
    return TrainingResults(learner, linear_model)
end

max_stars_repo_name: joehuchette/reserve-price-optimization

max_issues_repo_name: joehuchette/reserve-price-optimization

max_forks_repo_name: joehuchette/reserve-price-optimization


text:
"""
Provides types and functions to work with Markdown syntax trees.

The module is similar to the [Markdown standard library](https://docs.julialang.org/en/v1/stdlib/Markdown/),
but aims to be stricter and provide a more well-defined API.

!!! note
    Markdown2 does not provide a parser, just a data structure to represent Markdown ASTs.

# Markdown nodes

The types in this module represent the different types of nodes you can have in a Markdown
abstract syntax tree (AST). Currently it supports all the nodes necessary to represent Julia
flavored Markdown. But having this as a separate module from the Markdown standard library
allows us to consistently extend the node type we support (e.g. to support the raw HTML
nodes from [CommonMark](https://spec.commonmark.org/0.29/#raw-html), or strikethrough text
from [GitHub Flavored Markdown](https://github.github.com/gfm/#strikethrough-extension-)).

Markdown nodes split into to two different classes: [block nodes and inline
nodes](https://spec.commonmark.org/0.29/#blocks-and-inlines). Generally, the direct children
of a particular node can only be either inline or block (e.g. paragraphs contain inline
nodes, admonitions contain block nodes as direct children).

In Markdown2, this is represented using a simple type hierarchy. All Markdown nodes are
subtypes of either the [`MarkdownBlockNode`](@ref) or the [`MarkdownInlineNode`](@ref)
abstract type. Both of these abstract types themselves are a subtype of the
[`MarkdownNode`](@ref).

# Additional methods

* The `Base.convert(::Type{Markdown2.MD}, md::Markdown.MD)` method can be used to convert
  the Julia Markdown standard libraries ASTs into Markdown2 ASTs.
* The [`walk`](@ref) function can be used for walking over a [`Markdown2.MD`](@ref) tree.
"""
module Markdown2
using DocStringExtensions
import Markdown

# Abstract type hierarchy for Markdown nodes
# ==========================================
"""
    abstract type MarkdownNode

Supertype for all Markdown nodes.
"""
abstract type MarkdownNode end

"""
    abstract type MarkdownBlockNode <: MarkdownNode

Supertype for all block-level Markdown nodes.
"""
abstract type MarkdownBlockNode <: MarkdownNode end

"""
    abstract type MarkdownInlineNode <: MarkdownNode

Supertype for all inline Markdown nodes.
"""
abstract type MarkdownInlineNode <: MarkdownNode end

# Concrete types representing markdown nodes
# ==========================================
"""
    struct MD

The root node of a Markdown document. Its children are a list of top-level block-type nodes.
Note that `MD` is not a subtype of `MarkdownNode`.
"""
struct MD
    nodes :: Vector{MarkdownBlockNode}

    MD(content::AbstractVector) = new(content)
end
MD() = MD([])

# Forward some array methods
Base.push!(md::MD, x) = push!(md.nodes, x)
Base.getindex(md::MD, args...) = md.nodes[args...]
Base.setindex!(md::MD, args...) = setindex!(md.nodes, args...)
Base.lastindex(md::MD) = endof(md.nodes)
Base.length(md::MD) = length(md.nodes)
Base.isempty(md::MD) = isempty(md.nodes)

# Block nodes
# -----------
"""
    struct ThematicBreak <: MarkdownBlockNode

A block node represeting a thematic break (a `<hr>` tag).
"""
struct ThematicBreak <: MarkdownBlockNode end

struct Heading <: MarkdownBlockNode
    level :: Int
    nodes :: Vector{MarkdownInlineNode}

    function Heading(level::Integer, nodes::Vector{MarkdownInlineNode})
        @assert 1 <= level <= 6 # TODO: error message
        new(level, nodes)
    end
end

struct CodeBlock <: MarkdownBlockNode
    language :: String
    code :: String
end

#struct HTMLBlock <: MarkdownBlockNode end # the parser in Base does not support this currently
#struct LinkDefinition <: MarkdownBlockNode end # the parser in Base does not support this currently

"""
    struct Paragraph <: MarkdownBlockNode

Represents a paragraph block-type node. Its children are inline nodes.
"""
struct Paragraph <: MarkdownBlockNode
    nodes :: Vector{MarkdownInlineNode}
end

## Container blocks
struct BlockQuote <: MarkdownBlockNode
    nodes :: Vector{MarkdownBlockNode}
end

"""
    struct List <: MarkdownBlockNode

If `.orderedstart` is `nothing` then the list is unordered. Otherwise is specifies the first
number in the list.
"""
struct List <: MarkdownBlockNode
    tight :: Bool
    orderedstart :: Union{Int, Nothing}
    items :: Vector{Vector{MarkdownBlockNode}} # TODO: Better types?
end

# Non-Commonmark extensions
struct DisplayMath <: MarkdownBlockNode
    formula :: String
end

struct Footnote <: MarkdownBlockNode
    id :: String
    nodes :: Vector{MarkdownBlockNode} # Footnote is a container block
end

struct Table <: MarkdownBlockNode
    align :: Vector{Symbol}
    cells :: Array{Vector{MarkdownInlineNode}, 2} # TODO: better type?
    # Note: Table is _not_ a container type -- the cells can only contan inlines.
end

struct Admonition <: MarkdownBlockNode
    category :: String
    title :: String
    nodes :: Vector{MarkdownBlockNode} # Admonition is a container block
end

# Inline nodes
# ------------

struct Text <: MarkdownInlineNode
    text :: String
end

struct CodeSpan <: MarkdownInlineNode
    code :: String
end

struct Emphasis <: MarkdownInlineNode
    nodes :: Vector{MarkdownInlineNode}
end

struct Strong <: MarkdownInlineNode
    nodes :: Vector{MarkdownInlineNode}
end

struct Link <: MarkdownInlineNode
    destination :: String
    #title :: String # the parser in Base does not support this currently
    nodes :: Vector{MarkdownInlineNode}
end

struct Image <: MarkdownInlineNode
    destination :: String
    description :: String
    #title :: String # the parser in Base does not support this currently
    #nodes :: Vector{MarkdownInlineNode} # the parser in Base does not parse the description currently
end
#struct InlineHTML <: MarkdownInlineNode end # the parser in Base does not support this currently
struct LineBreak <: MarkdownInlineNode end

# Non-Commonmark extensions
struct InlineMath <: MarkdownInlineNode
    formula :: String
end

struct FootnoteReference <: MarkdownInlineNode
    id :: String
end

# Conversion methods
# ==================
"""
    convert(::Type{MD}, md::Markdown.MD) -> Markdown2.MD

Converts a Markdown standard library AST into a Markdown2 AST.
"""
function Base.convert(::Type{MD}, md::Markdown.MD)
    nodes = map(_convert_block, md.content)
    MD(nodes)
end

_convert_block(xs::Vector) = MarkdownBlockNode[_convert_block(x) for x in xs]
_convert_block(b::Markdown.HorizontalRule) = ThematicBreak()
function _convert_block(b::Markdown.Header{N}) where N
    text = _convert_inline(b.text)
    # Empty headings have just an empty String as text
    nodes = isa(text, AbstractVector) ? text : MarkdownInlineNode[text]
    Heading(N, nodes)
end
_convert_block(b::Markdown.Code) = CodeBlock(b.language, b.code)
_convert_block(b::Markdown.Paragraph) = Paragraph(_convert_inline(b.content))
_convert_block(b::Markdown.BlockQuote) = BlockQuote(_convert_block(b.content))
function _convert_block(b::Markdown.List)
    tight = all(isequal(1), length.(b.items))
    orderedstart = (b.ordered == -1) ? nothing : b.ordered
    List(tight, orderedstart, _convert_block.(b.items))
end

# Non-Commonmark extensions
_convert_block(b::Markdown.LaTeX) = DisplayMath(b.formula)
_convert_block(b::Markdown.Footnote) = Footnote(b.id, _convert_block(b.text))
function _convert_block(b::Markdown.Table)
    @assert all(isequal(length(b.align)), length.(b.rows)) # TODO: error
    cells = [_convert_inline(b.rows[i][j]) for i = 1:length(b.rows), j = 1:length(b.align)]
    Table(
        b.align,
        [_convert_inline(b.rows[i][j]) for i = 1:length(b.rows), j = 1:length(b.align)]
    )
end
_convert_block(b::Markdown.Admonition) = Admonition(b.category, b.title, _convert_block(b.content))

# Fallback
function _convert_block(x)
    @debug "Strange inline Markdown node (typeof(x) = $(typeof(x))), falling back to repr()" x
    Paragraph([Text(repr(x))])
end

_convert_inline(xs::Vector) = MarkdownInlineNode[_convert_inline(x) for x in xs]
_convert_inline(s::String) = Text(s)
function _convert_inline(s::Markdown.Code)
    @assert isempty(s.language) # TODO: error
    CodeSpan(s.code)
end
_convert_inline(s::Markdown.Bold) = Strong(_convert_inline(s.text))
_convert_inline(s::Markdown.Italic) = Emphasis(_convert_inline(s.text))
function _convert_inline(s::Markdown.Link)
    text = _convert_inline(s.text)
    # Autolinks (the `<URL>` syntax) yield Link objects where .text is just a String
    nodes = isa(text, AbstractVector) ? text : [text]
    Link(s.url, nodes)
end
_convert_inline(s::Markdown.Image) = Image(s.url, s.alt)
# struct InlineHTML <: MarkdownInlineNode end # the parser in Base does not support this currently
_convert_inline(::Markdown.LineBreak) = LineBreak()

# Non-Commonmark extensions
_convert_inline(s::Markdown.LaTeX) = InlineMath(s.formula)
function _convert_inline(s::Markdown.Footnote)
    @assert s.text === nothing # footnote references should not have any content, TODO: error
    FootnoteReference(s.id)
end

# Fallback
function _convert_inline(x)
    @debug "Strange inline Markdown node (typeof(x) = $(typeof(x))), falling back to repr()" x
    Text(repr(x))
end

# walk() function
# ===============
"""
    walk(f, element)

Calls `f(element)` on `element` and any of its child elements. The elements are assumed to
be [`Markdown2`](@ref) elements.
"""
function walk end

function walk(f, node::T) where {T <: Union{MarkdownNode, MD}}
    f(node) || return
    if :nodes in fieldnames(T)
        walk(f, node.nodes)
    end
    return
end

function walk(f, nodes::Vector)
    for node in nodes
        walk(f, node)
    end
end

function walk(f, list::List)
    f(list) || return
    for item in list.items
        walk(f, item)
    end
end

function walk(f, table::Table)
    f(table) || return
    for cell in table.cells
        walk(f, cell)
    end
end
end


text:
# Function for sparse arrays
# value!, deriv!
# `output` can have more dimensions than `target`, in which case do broadcasting

# TODO: find way to use normal broadcast for this.
# probably with the new changes in MLMetric and compare modes

@inline function value(loss::MarginLoss, target::AbstractSparseArray, output::AbstractArray)
    buffer = similar(output)
    value!(buffer, loss, target, output)
end

##
@generated function value!(
        buffer::AbstractArray,
        loss::MarginLoss,
        target::AbstractSparseArray{Q,Ti,M},
        output::AbstractArray{T,N}
    ) where {T,N,Q,Ti,M}
    M > N && throw(ArgumentError("target has more dimensions than output; broadcasting not supported in this direction."))
    quote
      @dimcheck size(buffer) == size(output)
      @nexprs $M (n)->@dimcheck(size(target,n) == size(output,n))
      zeroQ = zero(Q)
      negQ = Q(-1)
      @simd for I in CartesianIndices(size(output))
          @nexprs $N n->(i_n = I[n])
          tgt = @nref($M,target,i)
          if tgt == zeroQ
              # convention is that zeros in a sparse array are interpreted as negative one
              @inbounds @nref($N,buffer,i) = value(loss, negQ, @nref($N,output,i))
          else
              @inbounds @nref($N,buffer,i) = value(loss, tgt, @nref($N,output,i))
          end
      end
      buffer
    end
end


text:
# An example of link prediction using negative and positive samples.
# Ported from https://docs.dgl.ai/tutorials/blitz/4_link_predict.html#sphx-glr-tutorials-blitz-4-link-predict-py
# See the comparison paper https://arxiv.org/pdf/2102.12557.pdf for more details

using Flux
using Flux: onecold, onehotbatch
using Flux.Losses: logitbinarycrossentropy
using GraphNeuralNetworks
using MLDatasets: PubMed
using Statistics, Random, LinearAlgebra
using CUDA
CUDA.allowscalar(false)

# arguments for the `train` function
Base.@kwdef mutable struct Args
    η = 1f-3             # learning rate
    epochs = 200          # number of epochs
    seed = 17             # set seed > 0 for reproducibility
    usecuda = false      # if true use cuda (if available)
    nhidden = 64        # dimension of hidden features
    infotime = 10 	     # report every `infotime` epochs
end

# We define our own edge prediction layer but could also
# use GraphNeuralNetworks.DotDecoder instead.
struct DotPredictor end

function (::DotPredictor)(g, x)
    z = apply_edges((xi, xj, e) -> sum(xi .* xj, dims=1), g, xi=x, xj=x)
    # z = apply_edges(xi_dot_xj, g, xi=x, xj=x) # Same with built-in method
    return vec(z)
end

function train(; kws...)
    args = Args(; kws...)

    args.seed > 0 && Random.seed!(args.seed)

    if args.usecuda && CUDA.functional()
        device = gpu
        args.seed > 0 && CUDA.seed!(args.seed)
        @info "Training on GPU"
    else
        device = cpu
        @info "Training on CPU"
    end

    ### LOAD DATA
    data = PubMed.dataset()
    g = GNNGraph(data.adjacency_list)

    # Print some info
    @info g
    @show is_bidirected(g)
    @show has_self_loops(g)
    @show has_multi_edges(g)
    @show mean(degree(g))
    isbidir = is_bidirected(g)

    # Move to device
    g = g |> device
    X = data.node_features |> device

    #### TRAIN/TEST splits
    # With bidirected graph, we make sure that an edge and its reverse
    # are in the same split
    train_pos_g, test_pos_g = rand_edge_split(g, 0.9, bidirected=isbidir)
    test_neg_g = negative_sample(g, num_neg_edges=test_pos_g.num_edges, bidirected=isbidir)

    ### DEFINE MODEL #########
    nin, nhidden = size(X,1), args.nhidden

    # We embed the graph with positive training edges in the model
    model = WithGraph(GNNChain(GCNConv(nin => nhidden, relu),
                               GCNConv(nhidden => nhidden)),
                      train_pos_g) |> device

    pred = DotPredictor()

    ps = Flux.params(model)
    opt = ADAM(args.η)

    ### LOSS FUNCTION ############

    function loss(pos_g, neg_g = nothing; with_accuracy=false)
        h = model(X)
        if neg_g === nothing
            # We sample a negative graph at each training step
            neg_g = negative_sample(pos_g, bidirected=isbidir)
        end
        pos_score = pred(pos_g, h)
        neg_score = pred(neg_g, h)
        scores = [pos_score; neg_score]
        labels = [fill!(similar(pos_score), 1); fill!(similar(neg_score), 0)]
        l = logitbinarycrossentropy(scores, labels)
        if with_accuracy
            acc = 0.5 * mean(pos_score .>= 0) + 0.5 * mean(neg_score .< 0)
            return l, acc
        else
            return l
        end
    end

    ### LOGGING FUNCTION
    function report(epoch)
        train_loss, train_acc = loss(train_pos_g, with_accuracy=true)
        test_loss, test_acc = loss(test_pos_g, test_neg_g, with_accuracy=true)
        println("Epoch: $epoch  $((; train_loss, train_acc))  $((; test_loss, test_acc))")
    end

    ### TRAINING
    report(0)
    for epoch in 1:args.epochs
        gs = Flux.gradient(() -> loss(train_pos_g), ps)
        Flux.Optimise.update!(opt, ps, gs)
        epoch % args.infotime == 0 && report(epoch)
    end
end

# train()


text:
# This file is a part of AstroLib.jl. License is MIT "Expat".

function imf(mass::AbstractVector{T}, expon::AbstractVector{T},
             mass_range::AbstractVector{T}) where {T<:AbstractFloat}
    ne_comp = length(expon)
    if length(mass_range) != ne_comp + 1
        error("Length of array mass_range is not one more than that of expon")
    end
    integ = Vector{T}(undef, ne_comp)
    joint = ones(T, ne_comp)
    for i = 1:ne_comp
        if expon[i] != -1
            integ[i] = (mass_range[i+1]^(1 + expon[i]) - mass_range[i]^(1 + expon[i]))/
                       (1 + expon[i])
        else
            integ[i] = log(mass_range[i+1]/mass_range[i])
        end
        if i != 1
            joint[i] = joint[i-1]*(mass_range[i]^(expon[i-1] - expon[i]))
        end
    end
    norm = joint./(dot(integ, joint))
    psi = fill!(similar(mass), 0)
    for i = 1:ne_comp
        test = findall(mass_range[i].< mass.<mass_range[i+1])
        if length(test) !=0
            psi[test] = norm[i].*(mass[test].^expon[i])
        end
    end
    return psi
end

"""
    imf(mass, expon, mass_range) -> psi

### Purpose ###

Compute an N-component power-law logarithmic initial mass function (IMF).

### Explanation ###

The function is normalized so that the total mass distribution equals
one solar mass.

### Arguments ###

* `mass`: mass in units of solar mass, vector.
* `expon`: power law exponent, vector. The number of values in expon equals
  the number of different power-law components in the IMF.
* `mass_range`: vector containing the mass upper and lower limits of the
  IMF and masses where the IMF exponent changes. The number of values in
  mass_range should be one more than in expon. The values in mass_range
  should be monotonically increasing and positive.

### Output ###

* `psi`: mass function, number of stars per unit logarithmic mass interval
  evaluated for supplied masses.

### Example ###

Show the number of stars per unit mass interval at 3 Msun for a Salpeter
(expon = -1.35) IMF, with a mass range from 0.1 MSun to 110 Msun.

```jldoctest
julia> using AstroLib

julia> imf([3], [-1.35], [0.1, 110]) / 3
1-element Array{Float64,1}:
 0.01294143518151214
```

### Notes ###

Code of this function is based on IDL Astronomy User's Library.
"""
imf(mass::AbstractVector{<:Real}, expon::AbstractVector{<:Real}, mass_range::AbstractVector{<:Real}) =
    imf(float(mass), float(expon), float(mass_range))


text:
include("gibbssampling.jl")
include("hmcsampling.jl")

function log_gp_prior(gp::_MCGP,f::AbstractVector)
    logpdf(MvNormal(gp.μ₀,gp.K),f)
end

function log_joint_model(model::MCGP{T,L,<:SamplingInference},x) where {T,L}
    fs = unsqueeze(model,x)
    log_likelihood(model.likelihood,get_y(model),fs) + sum(log_gp_prior.(model.f),fs)
end

function unsqueeze(model::MCGP,f)
    n = model.nFeatures
    Tuple(f[((i-1)*n+1):(i*n)] for i in 1:model.nLatent)
end

function grad_log_joint_model(model::MCGP{T,L,<:SamplingInference},x) where {T,L}
    fs = unsqueeze(model,x)
    vcat(grad_log_likelihood(model.likelihood,get_y(model),fs)...) + vcat(grad_log_gp_prior.(model.f,fs)...)
end

log_likelihood(l::Likelihood,y::AbstractVector,f::Tuple{<:AbstractVector{T}}) where {T<:Real} = logpdf.(l,y,first(f))

function grad_log_likelihood(l::Likelihood,y::AbstractVector,f::Tuple{<:AbstractVector{T}}) where {T}
    grad_log_pdf.(l,y,f)
end

function grad_log_gp_prior(gp,f)
    gp.K/(f-gp.μ₀)
end