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

# QR and Hessenberg Factorizations

struct QR{T,S<:AbstractMatrix} <: Factorization{T}
    factors::S
    τ::Vector{T}
    QR{T,S}(factors::AbstractMatrix{T}, τ::Vector{T}) where {T,S<:AbstractMatrix} = new(factors, τ)
end
QR(factors::AbstractMatrix{T}, τ::Vector{T}) where {T} = QR{T,typeof(factors)}(factors, τ)
# Note. For QRCompactWY factorization without pivoting, the WY representation based method introduced in LAPACK 3.4
struct QRCompactWY{S,M<:AbstractMatrix} <: Factorization{S}
    factors::M
    T::Matrix{S}
    QRCompactWY{S,M}(factors::AbstractMatrix{S}, T::AbstractMatrix{S}) where {S,M<:AbstractMatrix} = new(factors, T)
end
QRCompactWY(factors::AbstractMatrix{S}, T::AbstractMatrix{S}) where {S} = QRCompactWY{S,typeof(factors)}(factors, T)

struct QRPivoted{T,S<:AbstractMatrix} <: Factorization{T}
    factors::S
    τ::Vector{T}
    jpvt::Vector{BlasInt}
    QRPivoted{T,S}(factors::AbstractMatrix{T}, τ::Vector{T}, jpvt::Vector{BlasInt}) where {T,S<:AbstractMatrix} =
        new(factors, τ, jpvt)
end
QRPivoted(factors::AbstractMatrix{T}, τ::Vector{T}, jpvt::Vector{BlasInt}) where {T} =
    QRPivoted{T,typeof(factors)}(factors, τ, jpvt)

function qrfactUnblocked!{T}(A::AbstractMatrix{T})
    m, n = size(A)
    τ = zeros(T, min(m,n))
    for k = 1:min(m - 1 + !(T<:Real), n)
        x = view(A, k:m, k)
        τk = reflector!(x)
        τ[k] = τk
        reflectorApply!(x, τk, view(A, k:m, k + 1:n))
    end
    QR(A, τ)
end

# Find index for columns with largest two norm
function indmaxcolumn(A::StridedMatrix)
    mm = norm(view(A, :, 1))
    ii = 1
    for i = 2:size(A, 2)
        mi = norm(view(A, :, i))
        if abs(mi) > mm
            mm = mi
            ii = i
        end
    end
    return ii
end

function qrfactPivotedUnblocked!(A::StridedMatrix)
    m, n = size(A)
    piv = collect(UnitRange{BlasInt}(1,n))
    τ = Array{eltype(A)}(min(m,n))
    for j = 1:min(m,n)

        # Find column with maximum norm in trailing submatrix
        jm = indmaxcolumn(view(A, j:m, j:n)) + j - 1

        if jm != j
            # Flip elements in pivoting vector
            tmpp = piv[jm]
            piv[jm] = piv[j]
            piv[j] = tmpp

            # Update matrix with
            for i = 1:m
                tmp = A[i,jm]
                A[i,jm] = A[i,j]
                A[i,j] = tmp
            end
        end

        # Compute reflector of columns j
        x = view(A, j:m, j)
        τj = LinAlg.reflector!(x)
        τ[j] = τj

        # Update trailing submatrix with reflector
        LinAlg.reflectorApply!(x, τj, view(A, j:m, j+1:n))
    end
    return LinAlg.QRPivoted{eltype(A), typeof(A)}(A, τ, piv)
end

# LAPACK version
qrfact!(A::StridedMatrix{<:BlasFloat}, ::Type{Val{false}}) = QRCompactWY(LAPACK.geqrt!(A, min(minimum(size(A)), 36))...)
qrfact!(A::StridedMatrix{<:BlasFloat}, ::Type{Val{true}}) = QRPivoted(LAPACK.geqp3!(A)...)
qrfact!(A::StridedMatrix{<:BlasFloat}) = qrfact!(A, Val{false})

# Generic fallbacks

"""
    qrfact!(A, pivot=Val{false})

`qrfact!` is the same as [`qrfact`](@ref) when `A` is a subtype of
`StridedMatrix`, but saves space by overwriting the input `A`, instead of creating a copy.
An [`InexactError`](@ref) exception is thrown if the factorization produces a number not
representable by the element type of `A`, e.g. for integer types.
"""
qrfact!(A::StridedMatrix, ::Type{Val{false}}) = qrfactUnblocked!(A)
qrfact!(A::StridedMatrix, ::Type{Val{true}}) = qrfactPivotedUnblocked!(A)
qrfact!(A::StridedMatrix) = qrfact!(A, Val{false})

"""
    qrfact(A, pivot=Val{false}) -> F

Computes the QR factorization of `A`. The return type of `F` depends on the element type of
`A` and whether pivoting is specified (with `pivot==Val{true}`).

| Return type   | `eltype(A)`     | `pivot`      | Relationship between `F` and `A` |
|:--------------|:----------------|:-------------|:---------------------------------|
| `QR`          | not `BlasFloat` | either       | `A==F[:Q]*F[:R]`                 |
| `QRCompactWY` | `BlasFloat`     | `Val{false}` | `A==F[:Q]*F[:R]`                 |
| `QRPivoted`   | `BlasFloat`     | `Val{true}`  | `A[:,F[:p]]==F[:Q]*F[:R]`        |

`BlasFloat` refers to any of: `Float32`, `Float64`, `Complex64` or `Complex128`.

The individual components of the factorization `F` can be accessed by indexing:

| Component | Description                               | `QR`            | `QRCompactWY`      | `QRPivoted`     |
|:----------|:------------------------------------------|:----------------|:-------------------|:----------------|
| `F[:Q]`   | `Q` (orthogonal/unitary) part of `QR`     | ✓ (`QRPackedQ`) | ✓ (`QRCompactWYQ`) | ✓ (`QRPackedQ`) |
| `F[:R]`   | `R` (upper right triangular) part of `QR` | ✓               | ✓                  | ✓               |
| `F[:p]`   | pivot `Vector`                            |                 |                    | ✓               |
| `F[:P]`   | (pivot) permutation `Matrix`              |                 |                    | ✓               |

The following functions are available for the `QR` objects: [`size`](@ref)
and [`\\`](@ref). When `A` is rectangular, `\\` will return a least squares
solution and if the solution is not unique, the one with smallest norm is returned.

Multiplication with respect to either thin or full `Q` is allowed, i.e. both `F[:Q]*F[:R]`
and `F[:Q]*A` are supported. A `Q` matrix can be converted into a regular matrix with
[`full`](@ref) which has a named argument `thin`.

# Example

```jldoctest
julia> A = [3.0 -6.0; 4.0 -8.0; 0.0 1.0]
3×2 Array{Float64,2}:
 3.0  -6.0
 4.0  -8.0
 0.0   1.0

julia> F = qrfact(A)
Base.LinAlg.QRCompactWY{Float64,Array{Float64,2}} with factors Q and R:
[-0.6 0.0 0.8; -0.8 0.0 -0.6; 0.0 -1.0 0.0]
[-5.0 10.0; 0.0 -1.0]

julia> F[:Q] * F[:R] == A
true
```

!!! note
    `qrfact` returns multiple types because LAPACK uses several representations
    that minimize the memory storage requirements of products of Householder
    elementary reflectors, so that the `Q` and `R` matrices can be stored
    compactly rather as two separate dense matrices.

    The data contained in `QR` or `QRPivoted` can be used to construct the
    `QRPackedQ` type, which is a compact representation of the rotation matrix:

    ```math
    Q = \\prod_{i=1}^{\\min(m,n)} (I - \\tau_i v_i v_i^T)
    ```

    where ``\\tau_i`` is the scale factor and ``v_i`` is the projection vector
    associated with the ``i^{th}`` Householder elementary reflector.

    The data contained in `QRCompactWY` can be used to construct the
    `QRCompactWYQ` type, which is a compact representation of the rotation
    matrix

    ```math
    Q = I + Y T Y^T
    ```

    where `Y` is ``m \\times r`` lower trapezoidal and `T` is ``r \\times r``
    upper triangular. The *compact WY* representation [^Schreiber1989] is not
    to be confused with the older, *WY* representation [^Bischof1987]. (The
    LAPACK documentation uses `V` in lieu of `Y`.)

    [^Bischof1987]: C Bischof and C Van Loan, "The WY representation for products of Householder matrices", SIAM J Sci Stat Comput 8 (1987), s2-s13. [doi:10.1137/0908009](http://dx.doi.org/10.1137/0908009)

    [^Schreiber1989]: R Schreiber and C Van Loan, "A storage-efficient WY representation for products of Householder transformations", SIAM J Sci Stat Comput 10 (1989), 53-57. [doi:10.1137/0910005](http://dx.doi.org/10.1137/0910005)
"""
function qrfact{T}(A::AbstractMatrix{T}, arg)
    AA = similar(A, typeof(zero(T)/norm(one(T))), size(A))
    copy!(AA, A)
    return qrfact!(AA, arg)
end
function qrfact{T}(A::AbstractMatrix{T})
    AA = similar(A, typeof(zero(T)/norm(one(T))), size(A))
    copy!(AA, A)
    return qrfact!(AA)
end
qrfact(x::Number) = qrfact(fill(x,1,1))

"""
    qr(A, pivot=Val{false}; thin::Bool=true) -> Q, R, [p]

Compute the (pivoted) QR factorization of `A` such that either `A = Q*R` or `A[:,p] = Q*R`.
Also see [`qrfact`](@ref).
The default is to compute a thin factorization. Note that `R` is not
extended with zeros when the full `Q` is requested.
"""
qr(A::Union{Number, AbstractMatrix}, pivot::Union{Type{Val{false}}, Type{Val{true}}}=Val{false}; thin::Bool=true) =
    _qr(A, pivot, thin=thin)
function _qr(A::Union{Number, AbstractMatrix}, ::Type{Val{false}}; thin::Bool=true)
    F = qrfact(A, Val{false})
    full(getq(F), thin=thin), F[:R]::Matrix{eltype(F)}
end
function _qr(A::Union{Number, AbstractMatrix}, ::Type{Val{true}}; thin::Bool=true)
    F = qrfact(A, Val{true})
    full(getq(F), thin=thin), F[:R]::Matrix{eltype(F)}, F[:p]::Vector{BlasInt}
end

"""
    qr(v::AbstractVector) -> w, r

Computes the polar decomposition of a vector.
Returns `w`, a unit vector in the direction of `v`, and
`r`, the norm of `v`.

See also [`normalize`](@ref), [`normalize!`](@ref),
and [`LinAlg.qr!`](@ref).

# Example

```jldoctest
julia> v = [1; 2]
2-element Array{Int64,1}:
 1
 2

julia> w, r = qr(v)
([0.447214, 0.894427], 2.23606797749979)

julia> w*r == v
true
```
"""
function qr(v::AbstractVector)
    nrm = norm(v)
    if !isempty(v)
        vv = copy_oftype(v, typeof(v[1]/nrm))
        return __normalize!(vv, nrm), nrm
    else
        T = typeof(zero(eltype(v))/nrm)
        return T[], oneunit(T)
    end
end

"""
    LinAlg.qr!(v::AbstractVector) -> w, r

Computes the polar decomposition of a vector. Instead of returning a new vector
as `qr(v::AbstractVector)`, this function mutates the input vector `v` in place.
Returns `w`, a unit vector in the direction of `v` (this is a mutation of `v`),
and `r`, the norm of `v`.

See also [`normalize`](@ref), [`normalize!`](@ref),
and [`qr`](@ref).
"""
function qr!(v::AbstractVector)
    nrm = norm(v)
    __normalize!(v, nrm), nrm
end

# Conversions
convert{T}(::Type{QR{T}}, A::QR) = QR(convert(AbstractMatrix{T}, A.factors), convert(Vector{T}, A.τ))
convert{T}(::Type{Factorization{T}}, A::QR{T}) = A
convert{T}(::Type{Factorization{T}}, A::QR) = convert(QR{T}, A)
convert{T}(::Type{QRCompactWY{T}}, A::QRCompactWY) = QRCompactWY(convert(AbstractMatrix{T}, A.factors), convert(AbstractMatrix{T}, A.T))
convert{T}(::Type{Factorization{T}}, A::QRCompactWY{T}) = A
convert{T}(::Type{Factorization{T}}, A::QRCompactWY) = convert(QRCompactWY{T}, A)
convert(::Type{AbstractMatrix}, F::Union{QR,QRCompactWY}) = F[:Q] * F[:R]
convert(::Type{AbstractArray}, F::Union{QR,QRCompactWY}) = convert(AbstractMatrix, F)
convert(::Type{Matrix}, F::Union{QR,QRCompactWY}) = convert(Array, convert(AbstractArray, F))
convert(::Type{Array}, F::Union{QR,QRCompactWY}) = convert(Matrix, F)
full(F::Union{QR,QRCompactWY}) = convert(AbstractArray, F)
convert{T}(::Type{QRPivoted{T}}, A::QRPivoted) = QRPivoted(convert(AbstractMatrix{T}, A.factors), convert(Vector{T}, A.τ), A.jpvt)
convert{T}(::Type{Factorization{T}}, A::QRPivoted{T}) = A
convert{T}(::Type{Factorization{T}}, A::QRPivoted) = convert(QRPivoted{T}, A)
convert(::Type{AbstractMatrix}, F::QRPivoted) = (F[:Q] * F[:R])[:,invperm(F[:p])]
convert(::Type{AbstractArray}, F::QRPivoted) = convert(AbstractMatrix, F)
convert(::Type{Matrix}, F::QRPivoted) = convert(Array, convert(AbstractArray, F))
convert(::Type{Array}, F::QRPivoted) = convert(Matrix, F)
full(F::QRPivoted) = convert(AbstractArray, F)

function show(io::IO, F::Union{QR, QRCompactWY, QRPivoted})
    println(io, "$(typeof(F)) with factors Q and R:")
    show(io, F[:Q])
    println(io)
    show(io, F[:R])
end

function getindex(A::QR, d::Symbol)
    m, n = size(A)
    if d == :R
        return triu!(A.factors[1:min(m,n), 1:n])
    elseif d == :Q
        return getq(A)
    else
        throw(KeyError(d))
    end
end
function getindex(A::QRCompactWY, d::Symbol)
    m, n = size(A)
    if d == :R
        return triu!(A.factors[1:min(m,n), 1:n])
    elseif d == :Q
        return getq(A)
    else
        throw(KeyError(d))
    end
end
function getindex{T}(A::QRPivoted{T}, d::Symbol)
    m, n = size(A)
    if d == :R
        return triu!(A.factors[1:min(m,n), 1:n])
    elseif d == :Q
        return getq(A)
    elseif d == :p
        return A.jpvt
    elseif d == :P
        p = A[:p]
        n = length(p)
        P = zeros(T, n, n)
        for i in 1:n
            P[p[i],i] = one(T)
        end
        return P
    else
        throw(KeyError(d))
    end
end

# Type-stable interface to get Q
getq(A::QRCompactWY) = QRCompactWYQ(A.factors,A.T)
getq(A::Union{QR, QRPivoted}) = QRPackedQ(A.factors,A.τ)

struct QRPackedQ{T,S<:AbstractMatrix} <: AbstractMatrix{T}
    factors::S
    τ::Vector{T}
    QRPackedQ{T,S}(factors::AbstractMatrix{T}, τ::Vector{T}) where {T,S<:AbstractMatrix} = new(factors, τ)
end
QRPackedQ(factors::AbstractMatrix{T}, τ::Vector{T}) where {T} = QRPackedQ{T,typeof(factors)}(factors, τ)

struct QRCompactWYQ{S, M<:AbstractMatrix} <: AbstractMatrix{S}
    factors::M
    T::Matrix{S}
    QRCompactWYQ{S,M}(factors::AbstractMatrix{S}, T::Matrix{S}) where {S,M<:AbstractMatrix} = new(factors, T)
end
QRCompactWYQ(factors::AbstractMatrix{S}, T::Matrix{S}) where {S} = QRCompactWYQ{S,typeof(factors)}(factors, T)

convert{T}(::Type{QRPackedQ{T}}, Q::QRPackedQ) = QRPackedQ(convert(AbstractMatrix{T}, Q.factors), convert(Vector{T}, Q.τ))
convert{T}(::Type{AbstractMatrix{T}}, Q::QRPackedQ{T}) = Q
convert{T}(::Type{AbstractMatrix{T}}, Q::QRPackedQ) = convert(QRPackedQ{T}, Q)
convert{S}(::Type{QRCompactWYQ{S}}, Q::QRCompactWYQ) = QRCompactWYQ(convert(AbstractMatrix{S}, Q.factors), convert(AbstractMatrix{S}, Q.T))
convert{S}(::Type{AbstractMatrix{S}}, Q::QRCompactWYQ{S}) = Q
convert{S}(::Type{AbstractMatrix{S}}, Q::QRCompactWYQ) = convert(QRCompactWYQ{S}, Q)
convert{T}(::Type{Matrix}, A::Union{QRPackedQ{T},QRCompactWYQ{T}}) = A_mul_B!(A, eye(T, size(A.factors, 1), minimum(size(A.factors))))
convert(::Type{Array}, A::Union{QRPackedQ,QRCompactWYQ}) = convert(Matrix, A)

"""
    full(A::Union{QRPackedQ,QRCompactWYQ}; thin::Bool=true) -> Matrix

Converts an orthogonal or unitary matrix stored as a `QRCompactWYQ` object, i.e. in the
compact WY format [^Bischof1987], or in the `QRPackedQ` format, to a dense matrix.

Optionally takes a `thin` Boolean argument, which if `true` omits the columns that span the
rows of `R` in the QR factorization that are zero. The resulting matrix is the `Q` in a thin
QR factorization (sometimes called the reduced QR factorization). If `false`, returns a `Q`
that spans all rows of `R` in its corresponding QR factorization.
"""
function full{T}(A::Union{QRPackedQ{T},QRCompactWYQ{T}}; thin::Bool = true)
    if thin
        convert(Array, A)
    else
        A_mul_B!(A, eye(T, size(A.factors, 1)))
    end
end

size(A::Union{QR,QRCompactWY,QRPivoted}, dim::Integer) = size(A.factors, dim)
size(A::Union{QR,QRCompactWY,QRPivoted}) = size(A.factors)
size(A::Union{QRPackedQ,QRCompactWYQ}, dim::Integer) = 0 < dim ? (dim <= 2 ? size(A.factors, 1) : 1) : throw(BoundsError())
size(A::Union{QRPackedQ,QRCompactWYQ}) = size(A, 1), size(A, 2)

function getindex(A::Union{QRPackedQ,QRCompactWYQ}, i::Integer, j::Integer)
    x = zeros(eltype(A), size(A, 1))
    x[i] = 1
    y = zeros(eltype(A), size(A, 2))
    y[j] = 1
    return dot(x, A_mul_B!(A, y))
end

## Multiplication by Q
### QB
A_mul_B!{T<:BlasFloat}(A::QRCompactWYQ{T}, B::StridedVecOrMat{T}) = LAPACK.gemqrt!('L','N',A.factors,A.T,B)
A_mul_B!{T<:BlasFloat}(A::QRPackedQ{T}, B::StridedVecOrMat{T}) = LAPACK.ormqr!('L','N',A.factors,A.τ,B)
function A_mul_B!(A::QRPackedQ, B::AbstractVecOrMat)
    mA, nA = size(A.factors)
    mB, nB = size(B,1), size(B,2)
    if mA != mB
        throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but B has dimensions ($mB, $nB)"))
    end
    Afactors = A.factors
    @inbounds begin
        for k = min(mA,nA):-1:1
            for j = 1:nB
                vBj = B[k,j]
                for i = k+1:mB
                    vBj += conj(Afactors[i,k])*B[i,j]
                end
                vBj = A.τ[k]*vBj
                B[k,j] -= vBj
                for i = k+1:mB
                    B[i,j] -= Afactors[i,k]*vBj
                end
            end
        end
    end
    B
end

function (*)(A::Union{QRPackedQ,QRCompactWYQ}, b::StridedVector)
    TAb = promote_type(eltype(A), eltype(b))
    Anew = convert(AbstractMatrix{TAb}, A)
    if size(A.factors, 1) == length(b)
        bnew = copy_oftype(b, TAb)
    elseif size(A.factors, 2) == length(b)
        bnew = [b; zeros(TAb, size(A.factors, 1) - length(b))]
    else
        throw(DimensionMismatch("vector must have length either $(size(A.factors, 1)) or $(size(A.factors, 2))"))
    end
    A_mul_B!(Anew, bnew)
end
function (*)(A::Union{QRPackedQ,QRCompactWYQ}, B::StridedMatrix)
    TAB = promote_type(eltype(A), eltype(B))
    Anew = convert(AbstractMatrix{TAB}, A)
    if size(A.factors, 1) == size(B, 1)
        Bnew = copy_oftype(B, TAB)
    elseif size(A.factors, 2) == size(B, 1)
        Bnew = [B; zeros(TAB, size(A.factors, 1) - size(B,1), size(B, 2))]
    else
        throw(DimensionMismatch("first dimension of matrix must have size either $(size(A.factors, 1)) or $(size(A.factors, 2))"))
    end
    A_mul_B!(Anew, Bnew)
end

### QcB
Ac_mul_B!{T<:BlasReal}(A::QRCompactWYQ{T}, B::StridedVecOrMat{T}) = LAPACK.gemqrt!('L','T',A.factors,A.T,B)
Ac_mul_B!{T<:BlasComplex}(A::QRCompactWYQ{T}, B::StridedVecOrMat{T}) = LAPACK.gemqrt!('L','C',A.factors,A.T,B)
Ac_mul_B!{T<:BlasReal}(A::QRPackedQ{T}, B::StridedVecOrMat{T}) = LAPACK.ormqr!('L','T',A.factors,A.τ,B)
Ac_mul_B!{T<:BlasComplex}(A::QRPackedQ{T}, B::StridedVecOrMat{T}) = LAPACK.ormqr!('L','C',A.factors,A.τ,B)
function Ac_mul_B!(A::QRPackedQ, B::AbstractVecOrMat)
    mA, nA = size(A.factors)
    mB, nB = size(B,1), size(B,2)
    if mA != mB
        throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but B has dimensions ($mB, $nB)"))
    end
    Afactors = A.factors
    @inbounds begin
        for k = 1:min(mA,nA)
            for j = 1:nB
                vBj = B[k,j]
                for i = k+1:mB
                    vBj += conj(Afactors[i,k])*B[i,j]
                end
                vBj = conj(A.τ[k])*vBj
                B[k,j] -= vBj
                for i = k+1:mB
                    B[i,j] -= Afactors[i,k]*vBj
                end
            end
        end
    end
    B
end
function Ac_mul_B(Q::Union{QRPackedQ,QRCompactWYQ}, B::StridedVecOrMat)
    TQB = promote_type(eltype(Q), eltype(B))
    return Ac_mul_B!(convert(AbstractMatrix{TQB}, Q), copy_oftype(B, TQB))
end

### QBc/QcBc
for (f1, f2) in ((:A_mul_Bc, :A_mul_B!),
                 (:Ac_mul_Bc, :Ac_mul_B!))
    @eval begin
        function ($f1)(Q::Union{QRPackedQ,QRCompactWYQ}, B::StridedVecOrMat)
            TQB = promote_type(eltype(Q), eltype(B))
            Bc = similar(B, TQB, (size(B, 2), size(B, 1)))
            ctranspose!(Bc, B)
            return ($f2)(convert(AbstractMatrix{TQB}, Q), Bc)
        end
    end
end

### AQ
A_mul_B!{T<:BlasFloat}(A::StridedVecOrMat{T}, B::QRCompactWYQ{T}) = LAPACK.gemqrt!('R','N', B.factors, B.T, A)
A_mul_B!{T<:BlasFloat}(A::StridedVecOrMat{T}, B::QRPackedQ{T}) = LAPACK.ormqr!('R', 'N', B.factors, B.τ, A)
function A_mul_B!(A::StridedMatrix,Q::QRPackedQ)
    mQ, nQ = size(Q.factors)
    mA, nA = size(A,1), size(A,2)
    if nA != mQ
        throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but matrix Q has dimensions ($mQ, $nQ)"))
    end
    Qfactors = Q.factors
    @inbounds begin
        for k = 1:min(mQ,nQ)
            for i = 1:mA
                vAi = A[i,k]
                for j = k+1:mQ
                    vAi += A[i,j]*Qfactors[j,k]
                end
                vAi = vAi*Q.τ[k]
                A[i,k] -= vAi
                for j = k+1:nA
                    A[i,j] -= vAi*conj(Qfactors[j,k])
                end
            end
        end
    end
    A
end

function (*)(A::StridedMatrix, Q::Union{QRPackedQ,QRCompactWYQ})
    TAQ = promote_type(eltype(A), eltype(Q))
    return A_mul_B!(copy_oftype(A, TAQ), convert(AbstractMatrix{TAQ}, Q))
end

### AQc
A_mul_Bc!{T<:BlasReal}(A::StridedVecOrMat{T}, B::QRCompactWYQ{T}) = LAPACK.gemqrt!('R','T',B.factors,B.T,A)
A_mul_Bc!{T<:BlasComplex}(A::StridedVecOrMat{T}, B::QRCompactWYQ{T}) = LAPACK.gemqrt!('R','C',B.factors,B.T,A)
A_mul_Bc!{T<:BlasReal}(A::StridedVecOrMat{T}, B::QRPackedQ{T}) = LAPACK.ormqr!('R','T',B.factors,B.τ,A)
A_mul_Bc!{T<:BlasComplex}(A::StridedVecOrMat{T}, B::QRPackedQ{T}) = LAPACK.ormqr!('R','C',B.factors,B.τ,A)
function A_mul_Bc!(A::AbstractMatrix,Q::QRPackedQ)
    mQ, nQ = size(Q.factors)
    mA, nA = size(A,1), size(A,2)
    if nA != mQ
        throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but matrix Q has dimensions ($mQ, $nQ)"))
    end
    Qfactors = Q.factors
    @inbounds begin
        for k = min(mQ,nQ):-1:1
            for i = 1:mA
                vAi = A[i,k]
                for j = k+1:mQ
                    vAi += A[i,j]*Qfactors[j,k]
                end
                vAi = vAi*conj(Q.τ[k])
                A[i,k] -= vAi
                for j = k+1:nA
                    A[i,j] -= vAi*conj(Qfactors[j,k])
                end
            end
        end
    end
    A
end
function A_mul_Bc(A::AbstractMatrix, B::Union{QRCompactWYQ,QRPackedQ})
    TAB = promote_type(eltype(A),eltype(B))
    BB = convert(AbstractMatrix{TAB}, B)
    if size(A,2) == size(B.factors, 1)
        AA = similar(A, TAB, size(A))
        copy!(AA, A)
        return A_mul_Bc!(AA, BB)
    elseif size(A,2) == size(B.factors,2)
        return A_mul_Bc!([A zeros(TAB, size(A, 1), size(B.factors, 1) - size(B.factors, 2))], BB)
    else
        throw(DimensionMismatch("matrix A has dimensions $(size(A)) but matrix B has dimensions $(size(B))"))
    end
end
@inline A_mul_Bc(rowvec::RowVector, B::Union{LinAlg.QRCompactWYQ,LinAlg.QRPackedQ}) = ctranspose(B*ctranspose(rowvec))

### AcQ/AcQc
for (f1, f2) in ((:Ac_mul_B, :A_mul_B!),
                 (:Ac_mul_Bc, :A_mul_Bc!))
    @eval begin
        function ($f1)(A::StridedVecOrMat, Q::Union{QRPackedQ,QRCompactWYQ})
            TAQ = promote_type(eltype(A), eltype(Q))
            Ac = similar(A, TAQ, (size(A, 2), size(A, 1)))
            ctranspose!(Ac, A)
            return ($f2)(Ac, convert(AbstractMatrix{TAQ}, Q))
        end
    end
end

A_ldiv_B!{T<:BlasFloat}(A::QRCompactWY{T}, b::StridedVector{T}) = (A_ldiv_B!(UpperTriangular(A[:R]), view(Ac_mul_B!(A[:Q], b), 1:size(A, 2))); b)
A_ldiv_B!{T<:BlasFloat}(A::QRCompactWY{T}, B::StridedMatrix{T}) = (A_ldiv_B!(UpperTriangular(A[:R]), view(Ac_mul_B!(A[:Q], B), 1:size(A, 2), 1:size(B, 2))); B)

# Julia implementation similarly to xgelsy
function A_ldiv_B!{T<:BlasFloat}(A::QRPivoted{T}, B::StridedMatrix{T}, rcond::Real)
    mA, nA = size(A.factors)
    nr = min(mA,nA)
    nrhs = size(B, 2)
    if nr == 0 return zeros(T, 0, nrhs), 0 end
    ar = abs(A.factors[1])
    if ar == 0 return zeros(T, nr, nrhs), 0 end
    rnk = 1
    xmin = ones(T, 1)
    xmax = ones(T, 1)
    tmin = tmax = ar
    while rnk < nr
        tmin, smin, cmin = LAPACK.laic1!(2, xmin, tmin, view(A.factors, 1:rnk, rnk + 1), A.factors[rnk + 1, rnk + 1])
        tmax, smax, cmax = LAPACK.laic1!(1, xmax, tmax, view(A.factors, 1:rnk, rnk + 1), A.factors[rnk + 1, rnk + 1])
        tmax*rcond > tmin && break
        push!(xmin, cmin)
        push!(xmax, cmax)
        for i = 1:rnk
            xmin[i] *= smin
            xmax[i] *= smax
        end
        rnk += 1
    end
    C, τ = LAPACK.tzrzf!(A.factors[1:rnk,:])
    A_ldiv_B!(UpperTriangular(C[1:rnk,1:rnk]),view(Ac_mul_B!(getq(A),view(B, 1:mA, 1:nrhs)),1:rnk,1:nrhs))
    B[rnk+1:end,:] = zero(T)
    LAPACK.ormrz!('L', eltype(B)<:Complex ? 'C' : 'T', C, τ, view(B,1:nA,1:nrhs))
    B[1:nA,:] = view(B, 1:nA, :)[invperm(A[:p]::Vector{BlasInt}),:]
    return B, rnk
end
A_ldiv_B!{T<:BlasFloat}(A::QRPivoted{T}, B::StridedVector{T}) = vec(A_ldiv_B!(A,reshape(B,length(B),1)))
A_ldiv_B!{T<:BlasFloat}(A::QRPivoted{T}, B::StridedVecOrMat{T}) = A_ldiv_B!(A, B, maximum(size(A))*eps(real(float(one(eltype(B))))))[1]
function A_ldiv_B!{T}(A::QR{T}, B::StridedMatrix{T})
    m, n = size(A)
    minmn = min(m,n)
    mB, nB = size(B)
    Ac_mul_B!(A[:Q], view(B, 1:m, :))
    R = A[:R]
    @inbounds begin
        if n > m # minimum norm solution
            τ = zeros(T,m)
            for k = m:-1:1 # Trapezoid to triangular by elementary operation
                x = view(R, k, [k; m + 1:n])
                τk = reflector!(x)
                τ[k] = τk'
                for i = 1:k - 1
                    vRi = R[i,k]
                    for j = m + 1:n
                        vRi += R[i,j]*x[j - m + 1]'
                    end
                    vRi *= τk
                    R[i,k] -= vRi
                    for j = m + 1:n
                        R[i,j] -= vRi*x[j - m + 1]
                    end
                end
            end
        end
        Base.A_ldiv_B!(UpperTriangular(view(R, :, 1:minmn)), view(B, 1:minmn, :))
        if n > m # Apply elementary transformation to solution
            B[m + 1:mB,1:nB] = zero(T)
            for j = 1:nB
                for k = 1:m
                    vBj = B[k,j]
                    for i = m + 1:n
                        vBj += B[i,j]*R[k,i]'
                    end
                    vBj *= τ[k]
                    B[k,j] -= vBj
                    for i = m + 1:n
                        B[i,j] -= R[k,i]*vBj
                    end
                end
            end
        end
    end
    return B
end
A_ldiv_B!(A::QR, B::StridedVector) = A_ldiv_B!(A, reshape(B, length(B), 1))[:]
function A_ldiv_B!(A::QRPivoted, b::StridedVector)
    A_ldiv_B!(QR(A.factors,A.τ), b)
    b[1:size(A.factors, 2)] = view(b, 1:size(A.factors, 2))[invperm(A.jpvt)]
    b
end
function A_ldiv_B!(A::QRPivoted, B::StridedMatrix)
    A_ldiv_B!(QR(A.factors, A.τ), B)
    B[1:size(A.factors, 2),:] = view(B, 1:size(A.factors, 2), :)[invperm(A.jpvt),:]
    B
end

# convenience methods
## return only the solution of a least squares problem while avoiding promoting
## vectors to matrices.
_cut_B(x::AbstractVector, r::UnitRange) = length(x)  > length(r) ? x[r]   : x
_cut_B(X::AbstractMatrix, r::UnitRange) = size(X, 1) > length(r) ? X[r,:] : X

## append right hand side with zeros if necessary
function _append_zeros(b::AbstractVector, T::Type, n)
    if n > length(b)
        x = zeros(T, n)
        return copy!(x, b)
    else
        return copy_oftype(b, T)
    end
end
function _append_zeros(B::AbstractMatrix, T::Type, n)
    if n > size(B, 1)
        X = zeros(T, (n, size(B, 2)))
        X[1:size(B,1), :] = B
        return X
    else
        return copy_oftype(B, T)
    end
end

function (\){TA,TB}(A::Union{QR{TA},QRCompactWY{TA},QRPivoted{TA}}, B::AbstractVecOrMat{TB})
    S = promote_type(TA,TB)
    m, n = size(A)
    m == size(B,1) || throw(DimensionMismatch("left hand side has $m rows, but right hand side has $(size(B,1)) rows"))

    AA = convert(Factorization{S}, A)

    X = A_ldiv_B!(AA, _append_zeros(B, S, n))

    return _cut_B(X, 1:n)
end

# With a real lhs and complex rhs with the same precision, we can reinterpret the complex
# rhs as a real rhs with twice the number of columns.

# convenience methods to compute the return size correctly for vectors and matrices
_ret_size(A::Factorization, b::AbstractVector) = (max(size(A, 2), length(b)),)
_ret_size(A::Factorization, B::AbstractMatrix) = (max(size(A, 2), size(B, 1)), size(B, 2))

function (\){T<:BlasReal}(A::Union{QR{T},QRCompactWY{T},QRPivoted{T}}, BIn::VecOrMat{Complex{T}})
    m, n = size(A)
    m == size(BIn, 1) || throw(DimensionMismatch("left hand side has $m rows, but right hand side has $(size(BIn,1)) rows"))

# |z1|z3|  reinterpret  |x1|x2|x3|x4|  transpose  |x1|y1|  reshape  |x1|y1|x3|y3|
# |z2|z4|      ->       |y1|y2|y3|y4|     ->      |x2|y2|     ->    |x2|y2|x4|y4|
#                                                 |x3|y3|
#                                                 |x4|y4|
    B = reshape(transpose(reinterpret(T, BIn, (2, length(BIn)))), size(BIn, 1), 2*size(BIn, 2))

    X = A_ldiv_B!(A, _append_zeros(B, T, n))

# |z1|z3|  reinterpret  |x1|x2|x3|x4|  transpose  |x1|y1|  reshape  |x1|y1|x3|y3|
# |z2|z4|      <-       |y1|y2|y3|y4|     <-      |x2|y2|     <-    |x2|y2|x4|y4|
#                                                 |x3|y3|
#                                                 |x4|y4|
    XX = reinterpret(Complex{T}, transpose(reshape(X, div(length(X), 2), 2)), _ret_size(A, BIn))
    return _cut_B(XX, 1:n)
end

##TODO:  Add methods for rank(A::QRP{T}) and adjust the (\) method accordingly
##       Add rcond methods for Cholesky, LU, QR and QRP types
## Lower priority: Add LQ, QL and RQ factorizations

# FIXME! Should add balancing option through xgebal


text:
#############################################################################
#
#    Will run every benchmarking unit in directory 'benchunits'
#
#    Each julia file in 'benchunits' should return a benchmark DataFrame
#     in a variable called 'res'
#
#    Results are appended to 'benchlog.csv'
#
#############################################################################

using DataFrames
using Benchmark
using MCMC

# cd to MCMC package dir to get correct git commit id
mcmcdir = joinpath(dirname(Base.find_in_path("MCMC")), "..")
cd(mcmcdir) do
	if readall(`git status -s`) != ""
		println("\nThere are untracked/uncommited changes to the package :")
		run(`git status`)
		println("\nProceed anyway with benchmarking ? (yes/NO)")
		c = readline(STDIN)
		ismatch(r"^y|Y", c) || error("Benchmarking aborted")
	end

	unitdir = joinpath(mcmcdir, "benchmarks", "benchunits")
	benchlog = joinpath(mcmcdir, "benchmarks", "benchlog.csv")

	fb = open(benchlog, "a+")
	for u in split(readall((`ls $unitdir`)))  #  u = split(readall((`ls $unitdir`)))[1]
		println("benchmarking '$u'")
		include(joinpath(unitdir, u))

		# printtable(fb, res, ',', '"', false)
		printtable(fb, res)
	end
	close(fb)
end

ext: jl


text:
using Test
using QuantumLattices.Prerequisites.SimpleTrees

@simpletree struct ATree end
@simpletree struct BTree end {<:AbstractString, ::Int}

@testset "AbstractSimpleTree" begin
    tree = ATree{String, Int}()

    @test eltype(tree) == Pair{String, Int}
    @test keytype(tree) == String
    @test valtype(tree) == Int
    @test root(tree) === nothing

    push!(tree, "L0", 0)
    push!(tree, "L0", "L1-1", 2)
    push!(tree, "L0", "L1-2", 3)
    push!(tree, "L1-1", "L2-1", 4)
    push!(tree, "L1-1", "L2-2", 5)
    push!(tree, "L1-2", "L2-3", 6)
    push!(tree, "L1-2", "L2-4", 7)
    tree["L0"] = 1

    @test root(tree) == "L0"
    @test length(tree) == 7
    @test collect(keys(tree, simpletreedepth)) == ["L0", "L1-1", "L2-1", "L2-2", "L1-2", "L2-3", "L2-4"]
    @test collect(keys(tree, simpletreewidth)) == ["L0", "L1-1", "L1-2", "L2-1", "L2-2", "L2-3", "L2-4"]
    @test collect(values(tree, simpletreedepth)) == [1, 2, 4, 5, 3, 6, 7]
    @test collect(values(tree, simpletreewidth)) == [1, 2, 3, 4, 5, 6, 7]
    @test collect(pairs(tree, simpletreedepth)) == [("L0", 1), ("L1-1", 2), ("L2-1", 4), ("L2-2", 5), ("L1-2", 3), ("L2-3", 6), ("L2-4", 7)]
    @test collect(pairs(tree, simpletreewidth)) == [("L0", 1), ("L1-1", 2), ("L1-2", 3), ("L2-1", 4), ("L2-2", 5), ("L2-3", 6), ("L2-4", 7)]
    @test [isleaf(tree, node) for node in keys(tree, simpletreewidth)] == [false, false, false, true, true, true, true]
    @test [level(tree, node) for node in keys(tree, simpletreewidth)] == [1, 2, 2, 3, 3, 3, 3]
    @test [ancestor(tree, "L2-1", i) for i = 0:2] == ["L2-1", "L1-1", "L0"]
    @test [ancestor(tree, "L2-3", i) for i = 0:2] == ["L2-3", "L1-2", "L0"]
    @test descendants(tree, "L0", 0) == ["L0"]
    @test descendants(tree, "L0", 1) == ["L1-1", "L1-2"]
    @test descendants(tree, "L0", 2) == ["L2-1", "L2-2", "L2-3", "L2-4"]
    @test siblings(tree, "L0") == []
    @test siblings(tree, "L1-1") == ["L1-2"]
    @test siblings(tree, "L2-1") == ["L2-2"]
    @test leaves(tree) == ["L2-1", "L2-2", "L2-3", "L2-4"]

    backup = deepcopy(tree)

    sub = subtree(tree, "L0")
    empty!(tree)
    append!(tree, sub)
    @test tree == backup
    @test isequal(tree, backup)

    sub = subtree(tree, "L1-1")
    delete!(tree, "L1-1")
    append!(tree, "L0", sub)
    move!(tree, "L1-2", "L0")
    @test tree == backup
    @test isequal(tree, backup)
end

@testset "SimpleTree" begin
    tree = SimpleTree{String, Int}()
    @test eltype(tree)  == eltype(typeof(tree)) == Pair{String, Int}
    @test keytype(tree) == keytype(typeof(tree)) == String
    @test valtype(tree) == valtype(typeof(tree)) == Int
    @test root(tree) === nothing
end

@testset "@simpletree" begin
    @test valtype(BTree) == Int
    @test keytype(BTree{String}) == String
    @test valtype(BTree{String}) == Int
    @test BTree{String}().TREECORE == SimpleTreeCore{String, Int}()
end


text:
Base.@kwdef mutable struct LindseyMethod{EF,ST}
    ef::EF
    discretizer::ST = Empirikos.Discretizer()
end

function StatsBase.fit(lindsey::LindseyMethod, Xs::AbstractVector{<:Real}, args...)
    lindsey = Empirikos.set_defaults(lindsey, Xs)
    return _fit(lindsey, Xs, args...)
end

function StatsBase.fit(lindsey::LindseyMethod, Xs, args...)
    return _fit(lindsey, Xs, args...)
end

function _fit(lindsey::LindseyMethod, Xs::AbstractVector{<:Real}, args...)
    Xs_summary = summarize(Xs, args..., lindsey.discretizer)
    return _fit(lindsey, Xs_summary)
end

function _fit(lindsey::LindseyMethod, Xs::AbstractVector{<:Interval}, args...)
    Xs_summary = summarize(Xs, args...)
    return _fit(lindsey, Xs_summary)
end

function _fit(lindsey::LindseyMethod, Xs::Empirikos.MultinomialSummary)
    mdpts = (first.(keys(Xs)) .+ last.(keys(Xs))) ./ 2

    ef = lindsey.ef

    #poisson_predictor = cefm.Q.(collect(mdpts))
    poisson_predictor = vcat(ef.Q.(mdpts)'...)
    poisson_predictor = hcat(fill(1.0, length(mdpts)), poisson_predictor)
    poisson_response = Empirikos.multiplicity(Xs)
    poisson_offset = pdf.(ef.basemeasure, mdpts)

    poisson_fit = fit(
        GeneralizedLinearModel,
        poisson_predictor,
        poisson_response,
        Poisson();
        offset=poisson_offset,
    )

    αs = coef(poisson_fit)[2:end]
    return ef(αs)
end


text:
using Documenter, WaterLily

makedocs(
    modules = [WaterLily],
    repo="https://github.com/weymouth/WaterLily.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", nothing) == "true",
        canonical="https:/weymouth.github.io/WaterLily.jl/",
        assets=String[],
    ),
    authors = "Gabriel Weymouth",
    sitename = "WaterLily.jl",
    pages = ["index.md"]
    # strict = true,
    # clean = true,
    # checkdocs = :exports,
)

deploydocs(
    repo = "github.com/weymouth/WaterLily.jl.git",
    target = "build",
    branch = "gh-pages",
    versions = ["stable" => "v^", "v#.#" ],
)


text:
# Workaround for GR warnings
ENV["GKSwstype"] = "100"

using Documenter, GeoStats
using DocumenterTools: Themes

istravis = "TRAVIS" ∈ keys(ENV)

Themes.compile(joinpath(@__DIR__,"src/assets/geostats-light.scss"), joinpath(@__DIR__,"src/assets/themes/documenter-light.css"))
Themes.compile(joinpath(@__DIR__,"src/assets/geostats-dark.scss"), joinpath(@__DIR__,"src/assets/themes/documenter-dark.css"))

makedocs(
  format = Documenter.HTML(
    assets=["assets/favicon.ico", asset("https://fonts.googleapis.com/css?family=Montserrat|Source+Code+Pro&display=swap", class=:css)],
    prettyurls = istravis,
    mathengine = KaTeX(Dict(
      :macros => Dict(
        "\\x" => "\\boldsymbol{x}",
        "\\z" => "\\boldsymbol{z}",
        "\\l" => "\\boldsymbol{\\lambda}",
        "\\c" => "\\boldsymbol{c}",
        "\\C" => "\\boldsymbol{C}",
        "\\g" => "\\boldsymbol{g}",
        "\\G" => "\\boldsymbol{G}",
        "\\f" => "\\boldsymbol{f}",
        "\\F" => "\\boldsymbol{F}",
        "\\R" => "\\mathbb{R}",
        "\\1" => "\\mathbb{1}"
      )
    ))
  ),
  sitename = "GeoStats.jl",
  authors = "Júlio Hoffimann",
  pages = [
    "Home" => "index.md",
    "Basic workflow" => "workflow.md",
    "Reference guide" => [
      "Data" => "data.md",
      "Domains" => "domains.md",
      "Problems" => "problems.md",
      "Solvers" => "solvers.md",
      "Validation" => "validation.md",
      "Statistics" => "statistics.md",
      "Clustering" => "clustering.md",
      "Weighting" => "weighting.md",
      "Miscellaneous" => "miscellaneous.md",
      "Variography" => [
        "variography/empirical.md",
        "variography/theoretical.md",
        "variography/fitting.md"
      ],
      "Kriging" => [
        "kriging/estimators.md",
        "kriging/solver.md"
      ],
      "Gaussian simulation" => [
        "gausssim/lu.md",
        "gausssim/fft.md",
        "gausssim/sgs.md"
      ],
      "Point Patterns" => [
        "pointpatterns/pointprocs.md",
        "pointpatterns/pointops.md"
      ],
      "Plotting" => "plotting.md"
    ],
    "Contributing" => [
      "contributing/guidelines.md",
      "Developer guide" => [
        "contributing/solvers.md"
      ]
    ],
    "Resources" => [
      "resources/education.md",
      "resources/publications.md",
      "resources/ecosystem.md"
    ],
    "About" => [
      "Community" => "about/community.md",
      "License" => "about/license.md",
      "Citing" => "about/citing.md"
    ],
    "Index" => "links.md"
  ]
)

deploydocs(repo="github.com/JuliaEarth/GeoStats.jl.git")


text:
##Details of reactions in a convenient format for later use
immutable Stoichiometry
    nspecies::Int ##How many species in total
    nreactions::Int ##How many reactions in total
    input_length::Int ##Length of the input arrays below
    ##Reaction input_reaction[i] requires input_index[i] units of species input_species[i]
    input_reactions::Array{Int, 1}
    input_species::Array{Int, 1}
    input_indices::Array{Int, 1}
    ##change[i,j] represents net change in number of units of species i from reaction j
    change::Array{Int, 2}

    function Stoichiometry(nspecies, nreactions, input_length, input_reactions, input_species, input_indices, change)
        if (nspecies < 1)
            error("Must have at least 1 species")
        elseif (nreactions < 1)
            error("Must have at least 1 reaction")
        elseif (length(input_reactions) != input_length)
            error("input_length inconsistent with input_reactions length")
        elseif (length(input_species) != input_length)
            error("input_length inconsistent with input_species length")
        elseif (length(input_indices) != input_length)
            error("input_length inconsistent with input_indices length")
        elseif (maximum(input_reactions) > nreactions)
            error("input_reactions values must not be greater than nreactions")
        elseif (minimum(input_reactions) <= 0)
            error("input_reactions values must be positive")
        elseif (maximum(input_species) > nspecies)
            error("input_species values must not be greater than nspecies")
        elseif (minimum(input_species) <= 0)
            error("input_species values must be positive")
        elseif (minimum(input_indices) <= 0)
            error("input_indices values must be positive")
        elseif (size(change) != (nspecies, nreactions))
            error("change matrix dimensions incorrect")
        end
        new(nspecies, nreactions, input_length, input_reactions, input_species, input_indices, change)
    end
end

##Stoichiometry constructor using P (inputs) and Q (outputs) matrices as in Owen et at 2014. (i.e. P[i,j] is number of species j required for reaction i, and Q[i,j] is number of species j output by reaction i).
function Stoichiometry(P::Array{Int, 2}, Q::Array{Int, 2})
    if (size(P) != size(Q))
        error("Dimensions of P and Q must be the same")
    elseif (minimum(P) < 0)
        error("All P elements must be non-negative")
    elseif (minimum(Q) < 0)
        error("All Q elements must be non-negative")
    end
    (nreactions, nspecies) = size(P)
    input_length = 0
    input_reactions = Array{Int}(0)
    input_species = Array{Int}(0)
    input_indices = Array{Int}(0)
    ##Store the non-zero elements of P
    for this_reaction in 1:nreactions
        for this_species in 1:nspecies
            this_index = P[this_reaction, this_species]
            if (this_index > 0)
                push!(input_reactions, this_reaction)
                push!(input_species, this_species)
                push!(input_indices, this_index)
                input_length += 1
            end
        end
    end
    Stoichiometry(nspecies, nreactions, input_length, input_reactions, input_species, input_indices, (Q-P)')
end

##Simulate forward 1 event
##Using stoichiometry "s", current state "state" and constants "θ"
##Return tuple of time till next event (can be Inf) and event number
function gillespie_step(s::Stoichiometry, state::Array{Int, 1}, θ::Array{Float64, 1})
    rates = copy(θ)
    expdist = Exponential()
    for i in 1:s.input_length
        rates[s.input_reactions[i]] *= state[s.input_species[i]]^s.input_indices[i]
    end
    total_rate = sum(rates)
    Δt = rand(expdist) / total_rate
    if (total_rate == 0.0)
        event = 1 ##An arbitrary value as event type irrelevant
    else
        event = rand(Categorical(rates/total_rate))
    end
    (Δt, event)
end

##Perform a simulation of the Gillespie algorithm
##with stoichiometry "s", initial state "state0" and constants "θ", observation times "obs_times", maximum iterations to perform maxit
##A tuple (success, observations) is returned
##success is true if all the observation times were reached in the specified number of iterations
##observations is a matrix with entry [:,i] the state at time obs_times[i] (if success=false, unobserved entries are zero)
function gillespie_partial_sim(s::Stoichiometry, state0::Array{Int, 1}, θ::Array{Float64, 1}, obs_times::Array{Float64, 1}, maxit::Integer)
    if (length(state0) != s.nspecies)
        error("Number of species given in initial state and stoichiometry do not match")
    elseif (length(θ) != s.nreactions)
        error("Number of reactions given in θ and stoichiometry do not match")
    elseif (minimum(θ) < 0)
        error("Negative values in θ")
    elseif (!issorted(obs_times, lt=((x,y)->x<=y)))
        error("obs_times must be strictly increasing")
    elseif (obs_times[1] < 0)
        error("obs_times must be non-negative")
    end
    nobs = length(obs_times)
    observations = zeros(Int, (s.nspecies, nobs))
    obs_count = 1 ##Which observation we are currently looking for
    if (obs_times[1] == 0)
        observations[:,1] = state0
        obs_count += 1
    end
    t_curr = 0.0
    state_curr = copy(state0)
    t_next = obs_times[obs_count]
    it_count = 0
    success = false
    while (it_count < maxit)
        (Δt, event) = gillespie_step(s, state_curr, θ)
        it_count += 1
        if (Δt > t_next - t_curr)
            observations[:,obs_count] = state_curr
            t_curr = t_next
            obs_count += 1
            if (obs_count > nobs)
                success = true
                break
            else
                t_next = obs_times[obs_count]
            end
        else
            state_curr += s.change[:,event]
            t_curr += Δt
        end
    end
    (success, observations)
end

##Perform a simulation of the Gillespie algorithm
##with stoichiometry "s", initial state "state0" and constants "θ"
##Until (a) time maxt is reached (b) maxit iterations have been performed (c) No further events possible.
##The output is a vector of event times and a matrix whose columns are corresponding states
function gillespie_sim(s::Stoichiometry, state0::Array{Int, 1}, θ::Array{Float64, 1}, maxt::AbstractFloat, maxit::Integer)
    if (length(state0) != s.nspecies)
        error("Number of species given in initial state and stoichiometry do not match")
    elseif (length(θ) != s.nreactions)
        error("Number of reactions given in θ and stoichiometry do not match")
    elseif (minimum(θ) < 0)
        error("Negative values in θ")
    end
    event_times = [0.0]
    observations = Array{Int}(s.nspecies, 1)
    observations[:,1] = state0
    obs_count = 2 ##Which observation we are currently looking for
    t_curr = 0.0
    state_curr = copy(state0)
    while (obs_count <= maxit)
        (Δt, event) = gillespie_step(s, state_curr, θ)
        t_curr += Δt
        if (t_curr == Inf || t_curr > maxt)
            break
        end
        state_curr += s.change[:,event]
        push!(event_times, t_curr)
        observations = hcat(observations, copy(state_curr))
        obs_count += 1
    end
    (event_times, observations)
end


text:
using Random

mutable struct RandGen
    stdgen::MersenneTwister
    is_seeded::Bool

    function RandGen()

        stdgen = MersenneTwister()
        is_seeded = false

        new(
            stdgen,
            is_seeded,
        )
    end
end

function randint(rand_gen::RandGen, low::Int, high::Int)

    x = rand(rand_gen.stdgen, UInt32)
    range = high - low
    return low + (x % range)
end

function randn(rand_gen::RandGen, high::Int)

    x = rand(rand_gen.stdgen, UInt32)
    return x % high
end

function rand01(rand_gen::RandGen)

    x = rand(rand_gen.stdgen, UInt32)
    return Float32(x / (typemax(UInt32) + 1))
end

function randbool(rand_gen::RandGen)

    return rand01(rand_gen) > 0.5
end

function randint(rand_gen::RandGen)

    return Int(rand(rand_gen.stdgen, UInt32))
end

function seed(rand_gen::RandGen, seed::Int)

    rand_gen.stdgen = Random.seed!(rand_gen.stdgen, seed)
    rand_gen.is_seeded = true
    return nothing
end


text:
using LinearAlgebra
using SparseArrays
using Plots
using ForwardDiff

uexact(x) = log(2+sin(pi*x))
dudx_exact(x) = ForwardDiff.derivative(uexact,x)
f(x) = dudx_exact(x)-ForwardDiff.derivative(dudx_exact,x) # du/dx - d^2u/dx^2 = f(x)
α = uexact(-1)
β = uexact(1)

function solve(m,f,α,β)
    x = LinRange(-1,1,m+2) # add 2 endpoints
    xint = x[2:end-1]
    h = x[2]-x[1]
    A = (1/h^2) * spdiagm(0=>2*ones(m),-1=>-ones(m-1),1=>-ones(m-1))
    Q = (1/(2*h)) * spdiagm(-1=>-ones(m-1),1=>ones(m-1))
    b = f.(xint)
    b[1] += α/h^2 + α/(2*h)
    b[end] += β/h^2 - β/(2*h)
    u = (Q+A)\b
    return u,x,xint,h
end

function solve(m,ϵ,f,α,β)
    x = LinRange(-1,1,m+2) # add 2 endpoints
    xint = x[2:end-1]
    h = x[2]-x[1]
    A = (1/h^2) * spdiagm(0=>2*ones(m),-1=>-ones(m-1),1=>-ones(m-1))
    Q = (1/(2*h)) * spdiagm(-1=>-ones(m-1),1=>ones(m-1))
    # Q = (1/(h)) * spdiagm(0=>-ones(m-1),1=>ones(m-1))
    b = f.(xint)
    b[1] += ϵ*α/h^2 + α/(2*h)
    b[end] += ϵ*β/h^2 - β/(2*h)
    u = (Q+ϵ*A)\b
    return u,x,xint,h
end

ϵ = .0001
u,x,xint,h = solve(100,ϵ,x->1,0,0)
plot(xint,u,leg=false)

# mvec = 2 .^ (2:9)
# hvec = zeros(length(mvec))
# err_inf = zeros(length(mvec))
# for (i,m) in enumerate(mvec)
#
#     ϵ = 1.0
#     u,x,xint,h = solve(m,ϵ,f,α,β)
#     hvec[i] = h
#     err_inf[i] = maximum(@. abs(uexact(xint) - u))
# end
#
# plot(hvec,err_inf,marker=:dot,label="Max error")
# plot!(hvec,hvec,linestyle=:dash,label="O(h)")
# plot!(hvec,hvec.^2,linestyle=:dash,label="O(h^2)")
# plot!(xaxis=:log,yaxis=:log,legend=:bottomright)


text:
module PIC
using StaticArrays
using Interpolations
using Plots
import Base

gr()

include("error.jl")
include("debug.jl")
include("num.jl")
include("constants.jl")
include("interp.jl")
include("grid.jl")
include("fields.jl")
include("particles.jl")
include("poisson.jl")
include("timestepping.jl")
include("plotting.jl")

include("test.jl")

#
# positions = zeros(2, 1)
# positions[:, 1] = [0.0, 0.0]
# particles = Particles{2}(positions, 1.0)
# particles.velocities[:] = [0.2, 0.07]
# grid = Grid((30, 30), (10.0, 10.0))
# charge_density = zeros(size(grid))
# potential = zeros(size(grid))
#
# n_steps = 100
# n_skips = 1
# anim = @animate for _ = 1:n_steps
#     charge_density .= 0
#     #sample_charge_density!(charge_density, grid, particles)
#     sample_current_density!(charge_density, grid, particles, DIM1)
#     positions .+= particles.velocities
#     @show sum(charge_density) * cell_volume(grid)
#     # potential .= 0
#     # solve_poisson!(
#     #     potential,
#     #     charge_density,
#     #     grid,
#     #     Free,
#     #     Jacobi(max_deviation = 1e-6),
#     # )
#     scalarheatmap(charge_density)#, clims = (0, 8))#(-0.3, 0.3))
# end every n_skips
# gif(anim, "test.gif", fps = 15)

#
# grid = Grid((30, 30, 3), (10.0, 10.0, 0.1))
# charge_density = zeros(size(grid))
# positions = []
# pt = [60, 60, 2]
# charge_density[pt...] = 1.0

# potential = zeros(size(grid))
# solve_poisson!(
#     potential,
#     charge_density,
#     grid,
#     Free,
#     Jacobi(max_deviation = 1e-6),
# )
# p1 = scalarheatmap(potential, clims = (-1e-3, 0))
# savefig(p1, "jac.png")
#
# p6 = plot(lower_edges(grid)[1], potential[1:end, pt[2]], ylims = (-1e-3, 0))
# savefig(p6, "jac_1d.png")

# p2 = scalarheatmap(potential[:, :, 2], clims = (-1e-3, 0))
# savefig(p2, "gs.png")
#
# p5 = plot(
#     lower_edges(grid)[1],
#     potential[1:end, pt[2:end]...],
#     ylims = (-1e-3, 0),
# )
# savefig(p5, "gs_1d.png")
#
# potential = zeros(size(grid))
# source_indices = findall(!iszero, charge_density)
# for k = 1:size(grid)[3]
#     for j = 1:size(grid)[2]
#         for i = 1:size(grid)[1]
#             potential[i, j, k] = compute_free_potential(
#                 SVector(
#                     lower_edges(grid)[1][i],
#                     lower_edges(grid)[2][j],
#                     lower_edges(grid)[3][k],
#                 ),
#                 charge_density,
#                 cell_volume(grid),
#                 lower_edges(grid),
#                 source_indices,
#             )
#         end
#     end
# end
# p2 = scalarheatmap(potential[:, :, 2], clims = (-1e-3, 0))
# savefig(p2, "correct.png")
#
# p3 = plot(
#     lower_edges(grid)[1],
#     potential[1:end, pt[2:end]...],
#     ylims = (-1e-3, 0),
# )
# savefig(p3, "correct_1d.png")
#
# analytical(charge, ::Dim{3}) =
#     -charge ./ (4π .* abs.(lower_edges(grid)[1] .- lower_edges(grid)[1][pt[1]]))
#
# analytical(charge, ::Dim{2}) =
#     charge .* log.(abs.(lower_edges(grid)[1] .- lower_edges(grid)[1][pt[1]])) ./
#     2π
#
# p4 = plot(
#     lower_edges(grid)[1],
#     analytical(
#         charge_density[pt...] * cell_volume(grid),
#         Dim{length(size(grid))}(),
#     ),
#     ylims = (-1e-3, 0),
# )
# savefig(p4, "analytical_1d.png")

# potential = zeros(size(grid)[1])
# analytical1 = analytical(charge_density[pt...] * cell_extents(grid)[1])
# potential[1] = analytical1[1]
# potential[end] = analytical1[end]
# for n = 1:10000
#     for i = 2:length(potential)-1
#         potential[i] =
#             0.5 * (
#                 potential[i-1] + potential[i+1] -
#                 cell_extents(grid)[1]^2 * charge_density[i, pt[2]]
#             )
#     end
# end
# p7 = plot(lower_edges(grid)[1], potential)
# savefig(p7, "gs1_1d.png")
#
# p8 = plot(lower_edges(grid)[1], analytical1)
# savefig(p8, "analytical1_1d.png")

# n_steps = 20
# n_skips = 1
# set_poisson_boundaries!(potential, charge_density, grid, Free)
# anim = @animate for _ = 1:n_steps
#     @show update_poisson!(
#         potential,
#         charge_density,
#         grid,
#         Free,
#         GaussSeidel(relaxation_factor = 1.5),
#     ) / (prod(extents(grid)) * sum(abs.(charge_density)))
#     scalarheatmap(potential)
# end every n_skips
# gif(anim, "test.gif", fps = 15)

end

max_stars_repo_name: lars-frogner/PIC

max_issues_repo_name: lars-frogner/PIC

max_forks_repo_name: lars-frogner/PIC


text:
using ContinuousPOMDPTreeSearchExperiments
using Multilane
using POMCPOW
using POMDPs
using MCTS
using ARDESPOT
using DataFrames
using CSV

file_contents = readstring(@__FILE__())

@show cor = 0.0
@show lambda = 2.0

@show N = 1000
@show n_iters = 1_000_000
@show max_time = 1.0
@show max_depth = 40
@show val = SimpleSolver()
alldata = DataFrame()

dpws = DPWSolver(depth=max_depth,
                 n_iterations=n_iters,
                 max_time=max_time,
                 exploration_constant=8.0,
                 k_state=4.5,
                 alpha_state=1/10.0,
                 enable_action_pw=false,
                 check_repeat_state=false,
                 estimate_value=RolloutEstimator(val)
                 # estimate_value=val
                )

solvers = Dict{String, Solver}(
    "qmdp" => QBSolver(dpws),
    "pftdpw" => begin
        m = 15
        wup = WeightUpdateParams(smoothing=0.0, wrong_lane_factor=0.05)
        rng = MersenneTwister(123)
        up = BehaviorParticleUpdater(nothing, m, 0.0, 0.0, wup, rng)
        BBMDPSolver(dpws, up)
    end,
    "pftdpw_5" => begin
        m = 5
        wup = WeightUpdateParams(smoothing=0.0, wrong_lane_factor=0.05)
        rng = MersenneTwister(123)
        up = BehaviorParticleUpdater(nothing, m, 0.0, 0.0, wup, rng)
        BBMDPSolver(dpws, up)
    end,
    "pftdpw_10" => begin
        m = 10
        wup = WeightUpdateParams(smoothing=0.0, wrong_lane_factor=0.05)
        rng = MersenneTwister(123)
        up = BehaviorParticleUpdater(nothing, m, 0.0, 0.0, wup, rng)
        BBMDPSolver(dpws, up)
    end,
    "pftdpw_20" => begin
        m = 20
        wup = WeightUpdateParams(smoothing=0.0, wrong_lane_factor=0.05)
        rng = MersenneTwister(123)
        up = BehaviorParticleUpdater(nothing, m, 0.0, 0.0, wup, rng)
        BBMDPSolver(dpws, up)
    end,
    "pomcpow" => begin
        wup = WeightUpdateParams(smoothing=0.0, wrong_lane_factor=0.05)
        POMCPOWSolver(tree_queries=n_iters,
                               criterion=MaxUCB(8.0),
                               max_depth=max_depth,
                               max_time=max_time,
                               enable_action_pw=false,
                               k_observation=4.5,
                               alpha_observation=1/10.0,
                               estimate_value=FORollout(val),
                               # estimate_value=val,
                               check_repeat_obs=false,
                               node_sr_belief_updater=BehaviorPOWFilter(wup)
                              )
    end,
    "despot" => begin
        rng = MersenneTwister(13)
        b = IndependentBounds(DefaultPolicyLB(val), 1.02, check_terminal=true)
        sol = DESPOTSolver(lambda=0.01,
                     K=100,
                     D=max_depth,
                     max_trials=1_000_000,
                     T_max=max_time,
                     bounds=b,
                     random_source=MemorizingSource(100, max_depth, rng, min_reserve=25),
                     default_action=ReportWhenUsed(MLAction(0.0, 0.0)),
                     rng=rng)
    end,
    "simple" => SimpleSolver()
)

behaviors = standard_uniform(correlation=cor)
pp = PhysicalParam(4, lane_length=100.0)
dmodel = NoCrashIDMMOBILModel(10, pp,
                              behaviors=behaviors,
                              p_appear=1.0,
                              lane_terminate=true,
                              max_dist=1000.0
                             )
rmodel = SuccessReward(lambda=lambda)
pomdp = NoCrashPOMDP{typeof(rmodel), typeof(behaviors)}(dmodel, rmodel, 0.95, false)
mdp = NoCrashMDP{typeof(rmodel), typeof(behaviors)}(dmodel, rmodel, 0.95, false)

for (k, solver) in solvers
    @show k

    sims = []
    for i in 1:N
        if k == "qmdp"
            planner = deepcopy(solve(solver, mdp))
        else
            planner = deepcopy(solve(solver, pomdp))
        end
        srand(planner, i+50_000)
        wup = WeightUpdateParams(smoothing=0.0, wrong_lane_factor=0.05)
        filter = BehaviorParticleUpdater(pomdp, 5000, 0.0, 0.0, wup, MersenneTwister(i+50_000))
        # filter = AggressivenessUpdater(pomdp, 5000, 0.0, 0.0, wup, MersenneTwister(i+50_000))

        rng = MersenneTwister(i+70_000)
        is = initial_state(pomdp, rng)
        ips = MLPhysicalState(is)

        md = Dict(:solver=>k, :i=>i)
        sim = Sim(deepcopy(pomdp),
                  planner,
                  filter,
                  ips,
                  is,
                  rng=rng,
                  max_steps=100,
                  metadata=md
                 )

        push!(sims, sim)
    end

    data = run_parallel(sims)
    # data = run(sims)

    if isempty(alldata)
        alldata = data
    else
        alldata = vcat(alldata, data)
    end

    rs = data[:reward]
    println(@sprintf("reward: %6.3f ± %6.3f", mean(rs), std(rs)/sqrt(length(rs))))
end

datestring = Dates.format(now(), "E_d_u_HH_MM")
copyname = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "data", "multilane_table_$(datestring).jl")
write(copyname, file_contents)
filename = Pkg.dir("ContinuousPOMDPTreeSearchExperiments", "data", "multilane_$(datestring).csv")
println("saving to $filename...")
CSV.write(filename, alldata)
println("done.")


text:
mutable struct NitscheHH3{T} <: MaxwellOperator3D
    gamma::T
end

function quaddata(operator::NitscheHH3,
    localtestbasis::LagrangeRefSpace,
    localtrialbasis::LagrangeRefSpace,
    testelements, trialelements)

  tqd = quadpoints(localtestbasis,  testelements,  (10,))
  bqd = quadpoints(x -> localtrialbasis(x, Val{:withcurl}), trialelements, (8,))

  #return QuadData(tqd, bqd)
  return (tpoints=tqd, bpoints=bqd)
end

function quadrule(op::NitscheHH3, g::LagrangeRefSpace, f::LagrangeRefSpace, i, τ, j, σ, qd)
    DoubleQuadStrategy(
        qd.tpoints[1,i],
        qd.bpoints[1,j]
    )
end

function integrand(op::NitscheHH3, kernel, test_vals, test_point, trial_vals, trial_point)
    Gxy = kernel.green
    @assert length(test_point.patch.tangents) == 1
    tx = normalize(tangents(test_point, 1))
    gx = test_vals[1]
    curlfy = trial_vals[2]
    return gx*dot(tx, Gxy * curlfy)
end


text:
#---------------------------------------------------------
# # [Option yflip](@id 2-yflip)
#---------------------------------------------------------

# Examples illustrating the `yflip` option in `MIRTjim`.

# Packages needed:

using MIRTjim: jim, prompt
using ImageGeoms: ImageGeom
using ImagePhantoms: shepp_logan, SheppLoganEmis, phantom

# The following line is helpful when running this example.jl file as a script;
# this way it will prompt user to hit a key after each image is displayed.

isinteractive() && jim(:prompt, true);

# ### Simple 2D image and 3D stack of images

i1 = shepp_logan(128, SheppLoganEmis()) # 2D image
i1 = i1[12:117,:] # non-square for illustration
i2 = cat(dims=3, i1, 2i1, 3i1)/3 # 3D stack of images
size(i2)

# This figure illustrates how the `yflip` option affects
# the image orientation in the "usual" case where y=1:N.
# This is a typical convention in "image processing" of digital images
# that lack any physical coordinates for their axes.

ji = (args...; kwargs...) -> jim(args...; kwargs..., prompt=false)
jim(
    ji( i1, "2D default"),
    ji( i1, yflip=false, "2D yflip=false"),
    ji( i1, yflip=true, "2D yflip=true"),
    ji( i2, ncol=2, "3D default"),
    ji( i2, ncol=2, yflip=false, "3D yflip=false"),
    ji( i2, ncol=2, yflip=true, "3D yflip=true"),
)

# ### 2D image and 3D stack of images with specified x,y coordinates

# Now consider the case where we think of the image as a function `f(x,y)`
# in the typical case where 0,0 is in the "center" of the image.
# This convention is more common in physics-based imaging and inverse problems
# where the image axes have meaningful values (and often have units).

objects = shepp_logan(SheppLoganEmis(); fovs=(120,120))
ig = ImageGeom(dims=(50,64), deltas=(2,2), offsets=(0.5,0.5))
xy = axes(ig)
i3 = phantom(xy..., objects)
i4 = cat(dims=3, i3, 2i3, 3i3) # 3D stack of images

jim(
    ji(xy..., i3, "2D default", xlabel="x", ylabel="y"),
    ji(xy..., i3, yflip=false, "2D yflip=false"),
    ji(xy..., i3, yflip=true, "2D yflip=true"),
    ji(xy..., i4, ncol=2, "3D default"),
    ji(xy..., i4, ncol=2, yflip=false, "3D yflip=false"),
    ji(xy..., i4, ncol=2, yflip=true, "3D yflip=true"),
    prompt = false,
)

# In both cases the "default" option for `yflip`
# has the typical desired behavior,
# but one can modify the default to achieve other behaviors.


text:
export git_otype, message, target_id, target, tagger, lookup_tag

Base.show(io::IO, t::GitTag) = begin
    write(io, "GitTag(\"$(name)\" @ $(target_id(t))")
end

name(t::GitTag) = bytestring(ccall((:git_tag_name, libgit2), Ptr{UInt8}, (Ptr{Void},), t))

function message(t::GitTag)
    msgptr = ccall((:git_tag_message, libgit2), Ptr{UInt8}, (Ptr{Void},), t)
    msgptr == C_NULL && return nothing
    return bytestring(msgptr)
end

function target_id(t::GitTag)
    idptr = ccall((:git_tag_target_id, libgit2), Ptr{UInt8}, (Ptr{Void},), t)
    return Oid(idptr)
end

function target(t::GitTag)
    tptr = Ptr{Void}[0]
    @check ccall((:git_tag_target, libgit2), Cint,
                 (Ptr{Ptr{Void}}, Ptr{Void}), tptr, t)
    return gitobj_from_ptr(tptr[1])
end

function tagger(t::GitTag)
    # returns a reference to an existing struct so don't free
    sig_ptr = ccall((:git_tag_tagger, libgit2), Ptr{SignatureStruct}, (Ptr{Void},), t)
    sig_ptr == C_NULL && return nothing
    return Signature(sig_ptr)
end


text:
# LibTOAST.jl: interface to the TOAST++ library
# Copyright (C) 2019 Samuel Powell

@testset "SystemMatrix" begin

  # Construct an empty system matrix, make sure its empty
  S = SystemMatrix(mesh2D)
  @test LibTOAST.ncols(S) == 3511
  @test LibTOAST.nrows(S) == 3511

  nvals = length(S.mesh.colidx)
  @test nvals == 24211

  sysmatvals = unsafe_wrap(Array,LibTOAST.valptr(S), nvals )
  @test all(sysmatvals .== 0)

  # Do some matrix assembly and make sure it comes out the same
  assemble!(S, FF)

  jlsmatvals = sparse(S).nzval
  @test sysmatvals == jlsmatvals

end


text:
function part1(input)
    program = parse_program(input)
    accumulator, halts = run_program(program)
    @assert !halts
    return accumulator
end

function part2(input)
    program = parse_program(input)
    for i = 1:length(program)
        opcode = program[i].opcode
        if opcode == :nop
            program[i].opcode = :jmp
        elseif opcode == :jmp
            program[i].opcode = :nop
        else
            continue
        end
        accumulator, halts = run_program(program)
        halts && return accumulator
        program[i].opcode = opcode
    end
end

mutable struct Instruction
    opcode::Symbol
    value::Int
end

function parse_program(input)
    program = Instruction[]
    for line in readlines(input)
        name, value = split(line, " ")
        push!(program, Instruction(Symbol(name), parse(Int, value)))
    end
    return program
end

function run_program(program)
    accumulator = 0
    program_counter = 1
    visited_instructions = Set{Int}()
    while true
        program_counter == length(program) + 1 && return accumulator, true
        program_counter in visited_instructions && return accumulator, false
        push!(visited_instructions, program_counter)
        instruction = program[program_counter]
        if instruction.opcode == :acc
            accumulator += instruction.value
            program_counter += 1
        elseif instruction.opcode == :jmp
            program_counter += instruction.value
        elseif instruction.opcode == :nop
            program_counter += 1
        end
    end
end


text:
export orLLi!

# --------- likelihoods ----------

"likelihood of data"
function orLL!(grad::AbstractVector{T}, hess::AbstractMatrix{T}, tmpgrad::AbstractVector, η::AbstractVector, y::AbstractVector{<:Integer}, X::AbstractMatrix, β::AbstractVector, γ::AbstractVector, model::Type{Val{D}}, sgn::Real=-one(T)) where {D,T}

  # # ---------- SAFETY CHECKS ---------------
  D ∈ (:probit, :logit) || throw(error("Can only do :probit or :logit model"))
  k,n = size(X)
  L = num_categories(y)
  length(β) == k    || throw(DimensionMismatch())
  length(γ) == L-1  || throw(DimensionMismatch())
  length(η) == n    || throw(DimensionMismatch())
  all(X[:,1] .== 1.0) && throw(error("X cannot have intercept"))
  # ----------------------------------------

  # ensure intercepts don't cross
  any(diff(γ) .<= 0.0)  &&  return -Inf

  # reset grad + hess
  length(grad) > 0 && (grad .= 0.0)
  length(hess) > 0 && (hess .= 0.0)

  mul!(η, transpose(X), β)  # update η

  LL = zero(T)
  for j in Base.OneTo(length(y))
      LL += orLLi!(grad, hess, tmpgrad, y[j], η[j], @view(X[:,j]), γ, model)
  end

  # flip signs of these
  if sign(sgn) < zero(sgn)
      length(grad) > 0  && (grad .*= -one(T))
      length(hess) > 0  && (hess .*= -one(T))
      return -LL
  end

  return LL
end

"observation-specific likelihood"
function orLLi!(grad::AbstractVector{T}, hess::AbstractMatrix{T}, tmpgrad::AbstractVector, l::Integer, η::Real, x::AbstractVector, γ::AbstractVector{T}, model::Type) where {T}
  any(diff(γ) .< zero(T))  &&  return -Inf
  L = length(γ)
  k = length(x)
  η1 = l == 1   ? -Inf : γ[l-1] - η
  η2 = l == L+1 ?  Inf : γ[l]   - η

  # LL takes a short-cut to improve accuracy if l == 1 or L+1
  if l == 1
      F = cdf(η2, model)
      LL = logcdf(η2, model)
  elseif l == L+1
      F = ccdf(η1, model)
      LL = logccdf(η1, model)
  else
      F = cdf(η2, model) - cdf(η1, model)
      LL = log(F)
  end

# https://github.com/scipy/scipy/blob/a2ffe09aa751749f2372aa13c19c61b2dec5266f/scipy/stats/_continuous_distns.py
# https://github.com/JuliaStats/Distributions.jl/blob/master/src/truncated/normal.jl
# https://github.com/cossio/TruncatedNormal.jl/blob/master/notes/normal.pdf

  if length(grad) > 0
      dlogF1 = pdf(η1, model) / F
      dlogF2 = pdf(η2, model) / F
      dlogF  = dlogcdf_trunc(η1,η2,model) # (pdf(η2, model) - pdf(η1, model)) / F  # try using _F1 from truncated/normal.jl

      # so that we can update grad directly if no hessian
      if length(hess) == 0
          tmpgrad2 = grad
      else
          tmpgrad2 = tmpgrad
          tmpgrad .= 0.0
      end

      tmpgrad2[1:k] .-= dlogF .* x
      l > 1    && (tmpgrad2[k+l-1] -= dlogF1)
      l < L+1  && (tmpgrad2[k+l]   += dlogF2)

      # update actual gradient vector if computing hessian
      length(hess) > 0 && (grad .+= tmpgrad)
  end

  if length(hess) > 0
      d2logF1 = dpdf(η1, model) / F
      d2logF2 = dpdf(η2, model) / F
      d2logF  = d2logF2 - d2logF1

      LinearAlgebra.BLAS.ger!(-one(T), tmpgrad, tmpgrad, hess)
      LinearAlgebra.BLAS.ger!(d2logF, x, x, @view(hess[1:k,1:k]))
      if l > 1
          hess[k+l-1, 1:k  ] .+= d2logF1 * x
          hess[1:k  , k+l-1] .+= d2logF1 * x
          hess[k+l-1, k+l-1]  -= d2logF1
      end
      if l < L+1
          hess[k+l, 1:k] .-= d2logF2 * x
          hess[1:k, k+l] .-= d2logF2 * x
          hess[k+l, k+l]  += d2logF2
      end
  end

  return LL
end


text:
using GLMakie: Node, @lift, labelslidergrid!, set_close_to!, connect!
using GLMakie: Figure, heatmap, contourf
using GaussianBeamLens

z = GaussianBeamLens.z
x = GaussianBeamLens.x

# Sliders
x0_index = Node(1)
w0_index = Node(2)
f_index = Node(3)
Lensz_index = Node(4)

# Use the values from the Sliders as inputs to get Intensity of the field at every point in x and z
@time I = @lift(GaussianBeamLensPropagate($x0_index, $w0_index, $f_index, $Lensz_index))

# Initialize figure
fig = Figure(resolution = (2304, 768))

@time heatmap(fig[1, 1], z, x, I,
	colormap = :thermal,
	axis =(aspect=3,
		title = "Gaussian Beam: Lens Transmission",
        xlabel="Propagation Direction (z)",
        ylabel="Lateral Offset (x)"))
#Other colormap options: deep

# Define a grid with slider label, value ranges
@time lsgrid = labelslidergrid!(
    fig,
    ["lateral offset", "input beam width", "focal length", "lens z-position"],
    [-240:10:240, 2:1:50, 0:10:2000, 0:10:2000];
    formats = [x -> "$x λ"],
    width = 555,
    tellheight = false)
fig[1, 2] = lsgrid.layout

# Set default value for each slider
set_close_to!(lsgrid.sliders[1], 60)
set_close_to!(lsgrid.sliders[2], 32)
set_close_to!(lsgrid.sliders[3], 1000)
set_close_to!(lsgrid.sliders[4],1000)

# Connect sliders to values
connect!(x0_index, lsgrid.sliders[1].value)
connect!(w0_index, lsgrid.sliders[2].value)
connect!(f_index, lsgrid.sliders[3].value)
connect!(Lensz_index, lsgrid.sliders[4].value)

# show figure
fig

max_stars_repo_name: mkitti/GaussianBeamLens

max_issues_repo_name: mkitti/GaussianBeamLens

max_forks_repo_name: mkitti/GaussianBeamLens


text:
module InferableTasks

export @iasync, @ispawn

struct InferableTask{T}
    task::Task
end

Base.fetch(t::InferableTask{T}) where {T} = fetch(t.task)::T
Base.wait(t::InferableTask) = wait(t.task)

macro ispawn(ex)
    inferable(ex) do ex
        :($Threads.@spawn $ex)
    end |> esc
end

macro iasync(ex)
    inferable(ex) do ex
        :($Base.@async $ex)
    end |> esc
end

# Use the world age that would be used from
# `Core.Compiler.return_type(f, t)`.  This is different from what
# `Base.get_world_counter` returns.  A big thanks to Chris Foster!:
# https://discourse.julialang.org/t/can-we-have-inferable-fetch-task/45541/2
_get_world_counter() = ccall(:jl_get_tls_world_age, UInt, ())

# I'm hoping this would make sure that the caller of `@ispawn` would
# be invalidated if the function `$f` has to be invalidated.  If the
# caller is invalidated, it'll get a new world age.  This in turn,
# invalidates `$f` called in the task because it is now called from a
# new world age (via `invoke_in_world`).
# --- https://discourse.julialang.org/t/can-we-have-inferable-fetch-task/45541/4
const __NEVER = Ref(false)

function inferable(spawn_macro, ex)
    @gensym f T world
    quote
        local $f, $T, $world
        $Base.@noinline $f() = $ex
        if $__NEVER[]
            $f()
        end
        $T = $Core.Compiler.return_type($f, $Tuple{})
        $world = $_get_world_counter()
        $InferableTask{$T}($(spawn_macro(:($Base.invoke_in_world($world, $f)))))
    end
end

end


text:
"""
    splitdims(array, [dims])

Eagerly split the contents of `array` into a nested array of arrays. The outermost
array contains the specified dimension(s) `dims`, which may be an integer, a
tuple of integers, or defaults to the final array dimension. The nested arrays
will contain all the remaining dimensions (in ascending order).

See also `splitdimsview`, which performs a similar operation lazily.

#### Examples:

```julia
julia> splitdims([1 2; 3 4])
2-element Array{Array{Int64,1},1}:
 [1, 3]
 [2, 4]

julia> splitdims([1 2; 3 4], 1)
2-element Array{Array{Int64,1},1}:
 [1, 2]
 [3, 4]
```
"""
@inline splitdims(a::AbstractArray, i::Int) = _splitdims(a, Val((i,))) # @inline forces constant propagation
@inline splitdims(a::AbstractArray) = _splitdims(a, Val((ndims(a),)))
@inline splitdims(a::AbstractArray{<:Any, 0}) = _splitdims(a, Val(()))
@inline splitdims(a::AbstractArray, dims::Tuple{Vararg{Int}}) = _splitdims(a, Val(dims))

function _splitdims(a::AbstractArray{<:Any, n}, ::Val{dims}) where {n, dims}
    innerdims = default_innerdims(dims, n)
    check_dims(Val(n), Val(dims), Val(innerdims))

    allaxes = axes(a)
    outeraxes = getindices(allaxes, dims)
    inneraxes = getindices(allaxes, innerdims)

    local out
    alloc = false
    for i in CartesianIndices(outeraxes)
        inds = slice_inds(i, Val(dims), Val(n))
        tmp = getindices(a, inds...)
        if !alloc
            out = similar(a, typeof(tmp), outeraxes)
            alloc = true
        end
        out[i] = tmp
    end
    return out
end

@pure function default_innerdims(dims::Tuple{Vararg{Int}}, n::Int)
    t = ()
    for i in 1:n
        if !(i in dims)
            t = (i, t...)
        end
    end
    return t
end

@generated function check_dims(::Val{n}, ::Val{dims}, ::Val{innerdims}) where {n, dims, innerdims}
    d = (dims..., innerdims...)
    if length(d) != n
        error("Incorrect dimensions")
    end
    for i in 1:n
        if !(i in d)
            error("Incorrect dimensions")
        end
    end
    return nothing
end

@generated function slice_inds(i::CartesianIndex, ::Val{dims}, ::Val{n}) where {dims, n}
    out = []
    for j in 1:n
        k = findfirst(equalto(j), dims)
        if k == 0
            out = [out..., :]
        else
            out = [out..., :(i[$k])]
        end
    end
    return :(tuple($(out...)))
end

## Lazy version

struct SplitDimsArray{T, N, Dims, A} <: AbstractArray{T, N}
    parent::A
end

Base.parent(a::SplitDimsArray) = a.parent

Base.axes(a::SplitDimsArray{T, N, Dims}) where {T, N, Dims} = getindices(axes(parent(a)), Dims)
Base.IndexStyle(::SplitDimsArray) = Base.IndexCartesian()
@propagate_inbounds function Base.getindex(a::SplitDimsArray{T, N, Dims}, i::Int...) where {T, N, Dims}
    return view(parent(a), slice_inds(CartesianIndex(i), Val(Dims), Val(ndims(parent(a))))...)
end

"""
    splitdimsview(array, [dims])

Lazily split the contents of `array` into a nested array of arrays. The outermost
array contains the specified dimension(s) `dims`, which may be an integer, a
tuple of integers, or defaults to the final array dimension. The nested arrays
will contain all the remaining dimensions (in ascending order).

See also `splitdims`, which performs a similar operation eagerly.

#### Examples:

```julia
julia> splitdims([1 2; 3 4])
2-element SplitDimsArray{SubArray{Int64,1,Array{Int64,2},Tuple{Base.Slice{Base.OneTo{Int64}},Int64},true},1,(2,),Array{Int64,2}}:
 [1, 3]
 [2, 4]

julia> splitdims([1 2; 3 4], 1)
2-element SplitDimsArray{SubArray{Int64,1,Array{Int64,2},Tuple{Int64,Base.Slice{Base.OneTo{Int64}}},true},1,(1,),Array{Int64,2}}:
 [1, 2]
 [3, 4]
```
"""
@inline splitdimsview(a::AbstractArray) = SplitDimsArray{new_eltype(typeof(a), Val((ndims(a),))), 1, (ndims(a),), typeof(a)}(a)
@inline splitdimsview(a::AbstractArray{<:Any, 0}) = SplitDimsArray{new_eltype(typeof(a), Val(())), 0, (), typeof(a)}(a)
@inline splitdimsview(a::AbstractArray, i::Int) = SplitDimsArray{new_eltype(typeof(a), Val((i,))), 1, (i,), typeof(a)}(a)
@inline function splitdimsview(a::AbstractArray{<:Any, N}, dims::NTuple{M, Int}) where {N, M}
    SplitDimsArray{new_eltype(typeof(a), Val(dims)), M, dims, typeof(a)}(a)
end

@pure _subtract(N::Int, M::Int) = N - M

function new_eltype(::Type{A}, ::Val{Dims}) where {A, Dims}
    return Core.Inference.return_type(view, splat_inds(Tuple{A, Core.Inference.return_type(slice_inds, Tuple{CartesianIndex{length(Dims)}, Val{Dims}, Val{ndims(A)}})}))
end

@pure function splat_inds(::Type{T}) where {T <: Tuple}
    a = T.parameters[1]
    b = T.parameters[2].parameters
    return Tuple{a, b...}
end


text:
module ChartmetricScraper

using HTTP,
    Dates
import JSON

export Token,
    newtoken!,
    albumrequest,
    playlistrequest,
    curatorrequest,
    Request,
    dorequest,
    parseresponse,
    getparameters,
    setparameters!,
    setparameter!,
    writestate,
    readstate!,
    buildrequesturl,
    getparsed,
    hasvalue,
    search

include("token.jl")
include("request.jl")
include("helpers.jl")
include("requests_helpers.jl")
end # module


text:
const CMat = Ptr{Cvoid}

"""
    A Petsc matrix, actually just a pointer to the actual C matrix
"""
struct PetscMat
    ptr::Ref{CMat}
    comm::MPI.Comm

    PetscMat(comm::MPI.Comm) = new(Ref{CMat}(), comm)
end

# allows us to pass PetscMat objects directly into CMat ccall signatures
Base.cconvert(::Type{CMat}, mat::PetscMat) = mat.ptr[]

"""
MatSetValue(mat::PetscMat, i::PetscInt, j::PetscInt, v::PetscScalar, mode::InsertMode)

Wrapper to `MatSetValue`. Indexing starts at 0 (as in PETSc).

# Implementation
For an unknow reason, calling PETSc.MatSetValue leads to an "undefined symbol: MatSetValue" error.
So this wrapper directly call MatSetValues (anyway, this is what is done in PETSc...)
"""
function MatSetValue(mat::PetscMat, i::PetscInt, j::PetscInt, v::PetscScalar, mode::InsertMode)
    MatSetValues(mat, PetscIntOne, [i], PetscIntOne, [j], [v], mode)
end

function MatSetValue(mat::PetscMat, i::Integer, j::Integer, v::Number, mode::InsertMode)
    MatSetValue(mat, PetscInt(i), PetscInt(j), PetscScalar(v), mode)
end

"""
MatSetValues(mat::PetscMat, nI::PetscInt, I::Vector{PetscInt}, nJ::PetscInt, J::Vector{PetscInt}, V::Array{PetscScalar}, mode::InsertMode)

Wrapper to `MatSetValues`. Indexing starts at 0 (as in PETSc)
"""
function MatSetValues(mat::PetscMat, nI::PetscInt, I::Vector{PetscInt}, nJ::PetscInt, J::Vector{PetscInt}, V::Array{PetscScalar}, mode::InsertMode)
    error = ccall((:MatSetValues, libpetsc), PetscErrorCode,
        (CMat, PetscInt, Ptr{PetscInt}, PetscInt, Ptr{PetscInt}, Ptr{PetscScalar}, InsertMode),
        mat, nI, I, nJ, J, V, mode)
    @assert iszero(error)
end

function MatSetValues(mat::PetscMat, I::Vector{PetscInt}, J::Vector{PetscInt}, V::Array{PetscScalar}, mode::InsertMode)
    MatSetValues(mat, PetscInt(length(I)), I, PetscInt(length(J)), J, V, mode)
end

function MatSetValues(mat::PetscMat, I, J, V, mode::InsertMode)
    MatSetValues(mat, PetscInt.(I), PetscInt.(collect(J)), PetscScalar.(V), mode)
end

"""
    MatCreate(comm::MPI.Comm, mat::PetscMat)

Wrapper to `MatCreate`
"""
function MatCreate(comm::MPI.Comm = MPI.COMM_WORLD)
    mat = PetscMat(comm)
    error = ccall((:MatCreate, libpetsc), PetscErrorCode, (MPI.MPI_Comm, Ptr{CMat}), comm, mat.ptr)
    @assert iszero(error)
    return mat
end

"""
    MatCompositeAddMat(mat::PetscMat, smat::PetscMat)

Wrapper to `MatCompositeAddMat`.
"""
function MatCompositeAddMat(mat::PetscMat, smat::PetscMat)
    error = ccall((:MatCompositeAddMat, libpetsc), PetscErrorCode, (CMat, CMat), mat, smat)
    @assert iszero(error)
end

"""
    MatCreateDense(comm::MPI.Comm, m::PetscInt, n::PetscInt, M::PetscInt, N::PetscInt)

Wrapper to `MatCreateDense`. Last argument `data` is not supported yet (NULL is passed).
"""
function MatCreateDense(comm::MPI.Comm, m::PetscInt, n::PetscInt, M::PetscInt, N::PetscInt)
    mat = PetscMat(comm)
    error = ccall((:MatCreateDense, libpetsc), PetscErrorCode,
        (MPI.MPI_Comm, PetscInt, PetscInt, PetscInt, PetscInt, Ptr{PetscScalar}, Ptr{CMat}),
        comm, m, n, M, N, C_NULL, mat.ptr)
    @assert iszero(error)
    return mat
end

function MatCreateDense(comm::MPI.Comm,
    m::Integer = PETSC_DECIDE,
    n::Integer = PETSC_DECIDE,
    M::Integer = PETSC_DECIDE,
    N::Integer = PETSC_DECIDE)

    return MatCreateDense(comm, PetscInt(m), PetscInt(n), PetscInt(M), PetscInt(N))
end

"""
    MatGetSize(mat::PetscMat)

Wrapper to `MatGetSize`. Return the number of rows and cols of the matrix (global number).
"""
function MatGetSize(mat::PetscMat)
    nrows_glo = Ref{PetscInt}(0)
    ncols_glo = Ref{PetscInt}(0)

    error = ccall((:MatGetSize, libpetsc), PetscErrorCode, (CMat, Ref{PetscInt}, Ref{PetscInt}), mat, nrows_glo, ncols_glo)
    @assert iszero(error)

    return nrows_glo[], ncols_glo[]
end

"""
    MatGetLocalSize(mat::PetscMat)

Wrapper to `MatGetLocalSize`. Return the number of local rows and cols of the matrix (i.e on the processor).
"""
function MatGetLocalSize(mat::PetscMat)
    nrows_loc = Ref{PetscInt}(0)
    ncols_loc = Ref{PetscInt}(0)

    error = ccall((:MatGetLocalSize, libpetsc), PetscErrorCode, (CMat, Ref{PetscInt}, Ref{PetscInt}), mat, nrows_loc, ncols_loc)
    @assert iszero(error)

    return nrows_loc[], ncols_loc[]
end

"""
    MatMult(mat::PetscMat, x::PetscVec, y::PetscVec)

Wrapper to `MatMult`
"""
function MatMult(mat::PetscMat, x::PetscVec, y::PetscVec)
    error = ccall((:MatMult, libpetsc), PetscErrorCode, (CMat, CVec, CVec), mat, x, y)
    @assert iszero(error)
end

"""
    MatSetSizes(mat::PetscMat, nrows_loc::PetscInt, ncols_loc::PetscInt, nrows_glo::PetscInt, ncols_glo::PetscInt)

Wrapper to MatSetSizes
"""
function MatSetSizes(mat::PetscMat, nrows_loc::PetscInt, ncols_loc::PetscInt, nrows_glo::PetscInt, ncols_glo::PetscInt)
    error = ccall((:MatSetSizes, libpetsc),
                PetscErrorCode,
                (CMat, PetscInt, PetscInt, PetscInt, PetscInt),
                mat, nrows_loc, ncols_loc, nrows_glo, ncols_glo
            )
    @assert iszero(error)
end

function MatSetSizes(mat::PetscMat, nrows_loc::Integer, ncols_loc::Integer, nrows_glo::Integer, ncols_glo::Integer)
    MatSetSizes(mat, PetscInt(nrows_loc), PetscInt(ncols_loc), PetscInt(nrows_glo), PetscInt(ncols_glo))
end

"""
    MatSetType(mat::PetscMat, type::String)

Wrapper for `MatSetType`. Values for `type` alors available here:
https://petsc.org/release/docs/manualpages/Mat/MatType.html#MatType
"""
function MatSetType(mat::PetscMat, type::String)
    error = ccall((:MatSetType, libpetsc), PetscErrorCode, (CMat, Cstring), mat, type)
    @assert iszero(error)
end

"""
    MatSetUp(mat::PetscMat)

Wrapper to MatSetUp
"""
function MatSetUp(mat::PetscMat)
    error = ccall((:MatSetUp, libpetsc), PetscErrorCode, (CMat,), mat)
    @assert iszero(error)
end

"""
    MatSetFromOptions(mat::PetscMat)

Wrapper to MatSetFromOptions
"""
function MatSetFromOptions(mat::PetscMat)
    error = ccall((:MatSetFromOptions, libpetsc), PetscErrorCode, (CMat,), mat)
    @assert iszero(error)
end

"""
    MatGetOwnershipRange(mat::PetscMat)

Wrapper to `MatGetOwnershipRange`

The result `(rstart, rend)` is a Tuple indicating the rows handled by the local processor.

# Warning
`PETSc` indexing starts at zero (so `rstart` may be zero) and `rend-1` is the last row
handled by the local processor.
"""
function MatGetOwnershipRange(mat::PetscMat)
    rstart = Ref{PetscInt}(0)
    rend = Ref{PetscInt}(0)

    error = ccall((:MatGetOwnershipRange, libpetsc), PetscErrorCode, (CMat, Ref{PetscInt}, Ref{PetscInt}), mat, rstart, rend)
    @assert iszero(error)

    return rstart[], rend[]
end

"""
    Wrapper to `MatAssemblyBegin`
"""
function MatAssemblyBegin(mat::PetscMat, type::MatAssemblyType)
    error = ccall((:MatAssemblyBegin, libpetsc), PetscErrorCode, (CMat, MatAssemblyType), mat, type)
    @assert iszero(error)
end

"""
    Wrapper to `MatAssemblyEnd`
"""
function MatAssemblyEnd(mat::PetscMat, type::MatAssemblyType)
    error = ccall((:MatAssemblyEnd, libpetsc), PetscErrorCode, (CMat, MatAssemblyType), mat, type)
    @assert iszero(error)
end

"""
    Wrapper to MatCreateVecs
"""
function MatCreateVecs(mat::PetscMat, vecr::PetscVec, veci::PetscVec)
    error = ccall((:MatCreateVecs, libpetsc), PetscErrorCode, (CMat, Ptr{CVec}, Ptr{CVec}), mat, vecr.ptr, veci.ptr)
    @assert iszero(error)
end

function MatCreateVecs(mat::PetscMat)
    vecr = PetscVec(mat.comm); veci = PetscVec(mat.comm)
    MatCreateVecs(mat, vecr, veci)
    return vecr, veci
end

"""
    MatMPIAIJSetPreallocation(mat::PetscMat, dnz::PetscInt, dnnz::Vector{PetscInt}, onz::PetscInt, onnz::Vector{PetscInt})

Wrapper to `MatMPIAIJSetPreallocation`. `dnnz` and `onnz` not tested yet.
"""
function MatMPIAIJSetPreallocation(mat::PetscMat, dnz::PetscInt, dnnz::Vector{PetscInt}, onz::PetscInt, onnz::Vector{PetscInt})
    error = ccall((:MatMPIAIJSetPreallocation, libpetsc), PetscErrorCode,
        (CMat, PetscInt, Ptr{PetscInt}, PetscInt, Ptr{PetscInt}),
        mat, dnz, dnnz, onz, onnz)
    @assert iszero(error)
end

"""
    MatSeqAIJSetPreallocation(mat::PetscMat, nz::PetscInt, nnz::Vector{PetscInt})

Wrapper to `MatSeqAIJSetPreallocation`. `nnz` not tested yet.
"""
function MatSeqAIJSetPreallocation(mat::PetscMat, nz::PetscInt, nnz::Vector{PetscInt} = C_NULL)
    error = ccall((:MatSeqAIJSetPreallocation, libpetsc), PetscErrorCode,
        (CMat, PetscInt, Ptr{PetscInt}), mat, nz, nnz)
    @assert iszero(error)
end

"""
    MatView(mat::PetscMat, viewer::PetscViewer = PetscViewerStdWorld())

Wrapper to `MatView`
"""
function MatView(mat::PetscMat, viewer::PetscViewer = PetscViewerStdWorld())
    error = ccall((:MatView, libpetsc), PetscErrorCode, (CMat, CViewer), mat, viewer)
    @assert iszero(error)
end

"""
    MatDestroy(mat::PetscMat)

Wrapper to MatDestroy
"""
function MatDestroy(mat::PetscMat)
    error = ccall((:MatDestroy, libpetsc), PetscErrorCode, (Ptr{CMat},), mat.ptr)
    @assert iszero(error)
end

max_stars_repo_name: bmxam/JuliaPetsc

max_issues_repo_name: bmxam/JuliaPetsc

max_forks_repo_name: bmxam/JuliaPetsc


text:
using Pkg

using Documenter, GenomicFeatures

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

format = Documenter.HTML(
    edit_link = "develop"
)

makedocs(
    modules = [GenomicFeatures],
    checkdocs = :all,
    linkcheck = true,
    format = format,
    sitename = "GenomicFeatures.jl",
    authors = replace(join(Pkg.TOML.parsefile("Project.toml")["authors"], ", "), r" <.*?>" => "" ) * ", The BioJulia Organisation, and other contributors.",
    pages = [
        "Home" => "index.md",
        "Intervals" => "man/intervals.md",
        "API Reference" => "api/api.md"
    ],
)

deploydocs(
    repo = "github.com/BioJulia/GenomicFeatures.jl.git",
    devbranch = "develop",
    push_preview = true
)


text:
module ArbitrageFreeSABR
using LinearAlgebra
using SparseArrays

include("BlackScholes.jl")
include("ImpliedVolatility.jl")
include("BachelierVolatility.jl")
include("Bachelier.jl")

abstract type SABRModel end
struct ArbitrageFreeSABRModel <: SABRModel end
struct FreeBoundarySABRModel <: SABRModel end

"Structure for the parameters of the SABR model for a given option maturity"
struct SABRMaturity{T<:Real,M<:SABRModel}
  α::T
  β::T
  ρ::T
  ν::T
  forward::T
  time::T
  model::M
end

"Probability density in the finite difference grid, at a given time, expressed in the transformed coordinate theta"
struct TransformedDensity{T<:Real,M<:SABRModel}
  maturity::SABRMaturity{T,M}
  zmin::T
  zmax::T
  h::T
  zm::Vector{T}
  ϑ::Vector{T}
  PL::T
  PR::T
end

export ArbitrageFreeSABRModel,
       FreeBoundarySABRModel,
       SABRMaturity,
       makeTransformedDensityLawsonSwayne,
       priceTransformedDensity

function solveStep(ϑ::Vector{G}, PL::G, PR::G, h::G, Fm::Vector{G}, Cm::Vector{G}, Em::Vector{G}, dt::G) where {G}
  frac = dt / (2 * h)
  M = length(ϑ)
  B = zeros(M)
  A = zeros(M - 1)
  C = zeros(M - 1)
  @. B[2:M-1] = 1.0 + frac * (Cm[2:M-1] * Em[2:M-1] * (1.0 / (Fm[3:M] - Fm[2:M-1]) + 1.0 / (Fm[2:M-1] - Fm[1:M-2])))
  @. C[2:M-1] = -frac * Cm[3:M] * Em[3:M] / (Fm[3:M] - Fm[2:M-1])
  @. A[1:M-2] = -frac * Cm[1:M-2] * Em[1:M-2] / (Fm[2:M-1] - Fm[1:M-2])
  B[1] = Cm[1] / (Fm[2] - Fm[1]) * Em[1]
  C[1] = Cm[2] / (Fm[2] - Fm[1]) * Em[2]
  B[M] = Cm[M] / (Fm[M] - Fm[M-1]) * Em[M]
  A[M-1] = Cm[M-1] / (Fm[M] - Fm[M-1]) * Em[M-1]
  tri = Tridiagonal(A, B, C)
  p1 = ϑ[1]
  pm = ϑ[M]
  ϑ[1] = 0
  ϑ[M] = 0
  ϑ0 = tri \ ϑ
  PL0 = PL + dt * Cm[2] / (Fm[2] - Fm[1]) * Em[2] * ϑ0[2]
  PR0 = PR + dt * Cm[M-1] / (Fm[M] - Fm[M-1]) * Em[M-1] * ϑ0[M-1]
  ϑ[1] = p1
  ϑ[M] = pm
  return ϑ0, PL0, PR0
end

function computeBoundaries(m::SABRMaturity{G,ArbitrageFreeSABRModel}, nd) where {G}
  zmin = -nd * sqrt(m.time)
  zmax = -zmin
  if (m.β < 1)
    ybar = -m.forward^(1 - m.β) / (1 - m.β)
    zbar = -1 / m.ν * log((sqrt(1 - m.ρ^2 + (m.ρ + m.ν * ybar / m.α)^2) - m.ρ - m.ν * ybar / m.α) / (1 - m.ρ))
    if (zbar > zmin)
      zmin = zbar
    end
  end
  return zmin, zmax
end

function computeBoundaries(m::SABRMaturity{G,FreeBoundarySABRModel}, nd) where {G}
  zmin = -nd * sqrt(m.time)
  zmax = -zmin
  return zmin, zmax
end

function yofz(m::SABRMaturity{G,M}, zm) where {G,M<:SABRModel}
  return m.α / m.ν * (sinh.(m.ν * zm) + m.ρ * (cosh.(m.ν * zm) .- 1))
end
function fofy(m::SABRMaturity{G,ArbitrageFreeSABRModel}, ym) where {G}
  expr = m.forward^(1 - m.β) .+ (1 - m.β) * ym
  return (max.(expr, 0)).^(1 / (1 - m.β))
end
function fofy(m::SABRMaturity{G,FreeBoundarySABRModel}, ym) where {G}
  u = sign(m.forward) * abs(m.forward)^(1 - m.β) .+ (1 - m.β) * ym
  return sign.(u) .* ((abs.(u)).^(1 / (1 - m.β)))
end
function cofy(m::SABRMaturity{G,ArbitrageFreeSABRModel}, ym, Fm) where {G}
  return sqrt.(m.α^2 .+ 2 * m.ρ * m.α * m.ν * ym + m.ν^2 * ym.^2) .* Fm.^(m.β)
end
function cofy(m::SABRMaturity{G,FreeBoundarySABRModel}, ym, Fm) where {G}
  return sqrt.(m.α^2 .+ 2 * m.ρ * m.α * m.ν * ym + m.ν^2 * ym.^2) .* abs.(Fm).^(m.β)
end
function gammaOfF(m::SABRMaturity{G,ArbitrageFreeSABRModel}, Fm, j0::Int) where {G}
  Gamma = (Fm.^m.β .- m.forward^m.β) ./ (Fm .- m.forward)
  Gamma[j0+1] = m.β / m.forward.^(1 - m.β)
  return Gamma
end

function gammaOfF(m::SABRMaturity{G,FreeBoundarySABRModel}, Fm, j0::Int) where {G}
  Gamma = (abs.(Fm).^m.β .- abs(m.forward)^m.β) ./ (Fm .- m.forward)
  return Gamma
end
function yOfStrike(m::SABRMaturity{T,ArbitrageFreeSABRModel}, strike) where {T}
  return (strike.^(1 - m.β) .- m.forward^(1 - m.β)) / (1 - m.β)
end

function yOfStrike(m::SABRMaturity{T,FreeBoundarySABRModel}, strike) where {T}
  return (sign.(strike) .* abs.(strike).^(1 - m.β) .- m.forward^(1 - m.β)) / (1 - m.β)
end

function makeForward(m::SABRMaturity{T,M}, z) where {T,M<:SABRModel}
  y = yofz(m, z)
  return fofy(m, y)
end

function makeTransformedDensityLawsonSwayne(
  maturity::SABRMaturity{G,M},
  N::Int,
  timesteps::Int,
  nd::Real
) where {G,M<:SABRModel}
  α = maturity.α
  β = maturity.β
  ν = maturity.ν
  ρ = maturity.ρ
  forward = maturity.forward
  T = maturity.time
  zmin, zmax = computeBoundaries(maturity, nd)
  J = N - 2
  h0 = (zmax - zmin) / J
  j0 = ceil((0 - zmin) / h0)
  h = (0 - zmin) / (j0 - 0.5)
  z = collect(0:(J+1)) * h .+ zmin
  zmax = z[J+1] #z[J+2] is outside as we shift the grid
  zm = z .- 0.5 * h
  ym = yofz(maturity, zm)
  ym[1] = ym[2] #avoid negative
  Fm = fofy(maturity, ym)
  Cm = cofy(maturity, ym, Fm)
  Gammam = gammaOfF(maturity, Fm, Int(j0))
  ymax = yofz(maturity, zmax)
  ymin = yofz(maturity, zmin)
  Fmax = fofy(maturity, ymax)
  Fmin = fofy(maturity, ymin)
  Fm[1] = 2 * Fmin - Fm[2]
  Fm[J+2] = 2 * Fmax - Fm[J+1]
  Cm[1] = Cm[2]
  Cm[J+2] = Cm[J+1]
  dt = T / timesteps
    #Ptotal = sum(h*ϑ[2:J+1])+PL+PR
    #Ftotal = Fm[2:J+1]*ϑ[2:J+1]*h+Fmin*PL+Fmax*PR
  b = 1 - sqrt(2) / 2 #Lawson Swayne param
  dt1 = dt * b
  dt2 = dt * (1 - 2 * b)
  Em = ones(J + 2)
  Emdt1 = exp.(ρ * ν * α * Gammam * dt1)
  Emdt1[1] = Emdt1[2]
  Emdt1[J+2] = Emdt1[J+1]
  Emdt2 = exp.(ρ * ν * α * Gammam * dt2)
  Emdt2[1] = Emdt2[2]
  Emdt2[J+2] = Emdt2[J+1]
  PL = 0.0
  PR = 0.0
  ϑ = zeros(J + 2)
  ϑ[Int(j0)+1] = 1.0 / h
  for t = 1:timesteps
    @. Em *= Emdt1
    ϑ1, PL1, PR1 = solveStep(ϑ, PL, PR, h, Fm, Cm, Em, dt1)
    @. Em *= Emdt1
    ϑ2, PL2, PR2 = solveStep(ϑ1, PL1, PR1, h, Fm, Cm, Em, dt1)
    @. ϑ = (sqrt(2) + 1) * ϑ2 - sqrt(2) * ϑ1
    PL = (sqrt(2) + 1) * PL2 - sqrt(2) * PL1
    PR = (sqrt(2) + 1) * PR2 - sqrt(2) * PR1
    @. Em *= Emdt2
      #Ptotal = sum(h*ϑ[2:J+1])+PL+PR
      #Ftotal = Fm[2:J+1]*ϑ[2:J+1]+h+Fmin*PL+Fmax*PR
  end
  density = TransformedDensity(maturity, zmin, zmax, h, zm, ϑ, PL, PR)
  return density
end

@enum DensitySmoothing none = 1 linear = 2 midpoint = 3 linearc0 = 4

function computeLinearValues(density::TransformedDensity)
  n = length(density.ϑ)
  Fm = makeForward(density.maturity, density.zm) #FIXME check Fm[1] != Fm[2] (should be negative, is it possible? eventually, just mirror it, the choice is free)
  Fm[1] = -Fm[2]
  F = makeForward(density.maturity, density.zm .+ 0.5 * density.h)
  s = zeros(n)
  d = zeros(n)
  dl = zeros(n - 1)
  du = zeros(n - 1)
  rhs = zeros(n)
  firstrow = zeros(n)
  lastrow = zeros(n)
  for i = 2:n-1
    d[i] = F[i] - F[i-1] - 0.5 * (F[i] - Fm[i])^2 / (Fm[i+1] - Fm[i]) + 0.5 * (F[i-1] - Fm[i])^2 / (Fm[i-1] - Fm[i])
    du[i] = 0.5 * (F[i] - Fm[i])^2 / (Fm[i+1] - Fm[i])
    dl[i-1] = -0.5 * (F[i-1] - Fm[i])^2 / (Fm[i-1] - Fm[i])
    rhs[i] = density.ϑ[i] * density.h
    lastrow[i] += 0.5 * (F[i] - Fm[i])^2 - 0.5 * (F[i-1] - Fm[i])^2 - (F[i] - Fm[i])^3 / (Fm[i+1] - Fm[i]) / 3.0 +
                  (F[i-1] - Fm[i])^3 / (Fm[i-1] - Fm[i]) / 3.0
    lastrow[i-1] += -(F[i-1] - Fm[i])^3 / (Fm[i-1] - Fm[i]) / 3.0
    lastrow[i+1] += (F[i] - Fm[i])^3 / (Fm[i+1] - Fm[i]) / 3.0
      #rhs[0] += Fm[i]*density.ϑ[i]*density.h
      ##use F*a + b expression for F^2, eventually for F as well.
  end
  rhs[n] = 0
  firstrow[1] = Fm[2] # Fm[2]*s[1] + Fm[1]*s[2] = 0
  firstrow[2] = -Fm[1]
  rhs[1] = 0
  tri = Tridiagonal(dl, d, du)
  #then add top row and bottom row to matrix
  stri = sparse(tri)
  stri[1, :] = firstrow
  stri[n, :] = lastrow
  coeff = stri \ rhs
  return coeff
end

function priceTransformedDensity(
  density::TransformedDensity,
  isCall::Bool,
  strike::G,
  smoothing::DensitySmoothing,
  smoothingParameters = nothing
) where {G}
  maturity = density.maturity
  α = maturity.α
  β = maturity.β
  ν = maturity.ν
  ρ = maturity.ρ
  forward = maturity.forward
  ystrike = yOfStrike(maturity, strike)
  zstrike = -1 / ν * log.((sqrt.(1 .- ρ^2 .+ (ρ + ν * ystrike / α).^2) .- ρ - ν * ystrike / α) / (1 - ρ))
  if (zstrike <= density.zmin)
    p = forward .- strike
  else
    if (zstrike >= density.zmax)
      p = 0
    else
      Fmax = makeForward(maturity, density.zmax)
      p = (Fmax .- strike) * density.PR
      k0 = Int.(ceil.((zstrike .- density.zmin - eps()*1000) / density.h))
      ztilde = density.zmin .+ k0 * density.h
      ftilde = makeForward(maturity, ztilde)
      term = ftilde .- strike
      if smoothing == linearc0
        Fm = makeForward(maturity, density.zm[k0:end]) #zm[k0+1] = zmin+k0*h
        if k0 == 1
          Fm[1] = -Fm[2]
        end
        F = makeForward(maturity, density.zm[k0:end] .+ 0.5 * density.h)
        coeff = convert(Array{Float64}, smoothingParameters)[k0:end]
        p += sum((Fm[3:end-1] .- strike) * density.h .* density.ϑ[k0+2:end-1])
        for i = 3:length(Fm)-1
          p += (0.5 * (F[i] - Fm[i])^2 - 0.5 * (F[i-1] - Fm[i])^2) * coeff[i] +
               (coeff[i+1] - coeff[i]) * (F[i] - Fm[i])^3 / (Fm[i+1] - Fm[i]) / 3.0 +
               (coeff[i] - coeff[i-1]) * (F[i-1] - Fm[i])^3 / (Fm[i-1] - Fm[i]) / 3.0
        end
        if strike < Fm[2]
          p += (F[2] - Fm[2])^3 * (coeff[3] - coeff[2]) / (Fm[3] - Fm[2]) / 3.0 + 0.5 * (F[2] - Fm[2])^2 * coeff[2]+Fm[2]*(0.5*(F[2]-Fm[2])^2*(coeff[3]-coeff[2])/(Fm[3]-Fm[2])+coeff[2]*(F[2]-Fm[2]))
          p -= strike * ((F[2]^2 - Fm[2]^2) * (coeff[3] - coeff[2]) / (Fm[3] - Fm[2]) / 2.0 + (F[2] - Fm[2]) * (coeff[2]-Fm[2]*(coeff[3] - coeff[2]) / (Fm[3] - Fm[2])))
          p -= (strike-Fm[2])^3 * (coeff[1] - coeff[2]) / (Fm[1] - Fm[2]) / 3.0 + 0.5 * (strike - Fm[2])^2 * coeff[2]+Fm[2]*(0.5*(strike-Fm[2])^2*(coeff[1]-coeff[2])/(Fm[1]-Fm[2])+coeff[2]*(strike-Fm[2]))
#          p += (Fm[2]^3 - strike^3) * (coeff[2] - coeff[1]) / (Fm[2] - Fm[1]) / 3.0 +
          #     0.5 * (Fm[2]^2 - strike^2) *( coeff[1] -Fm[1]* (coeff[2] - coeff[1]) / (Fm[2] - Fm[1]))
          p -= strike *
               ((Fm[2]^2 - strike^2) * (coeff[2] - coeff[1]) / (Fm[2] - Fm[1]) / 2.0 + (Fm[2] - strike) * (coeff[1] -Fm[1]* (coeff[2] - coeff[1]) / (Fm[2] - Fm[1])))
        else
          p += (F[2]^3 - strike^3) * (coeff[3] - coeff[2]) / (Fm[3] - Fm[2]) / 3.0 +
               0.5 * (F[2]^2 - strike^2) * (coeff[2]-Fm[2]*(coeff[3] - coeff[2]) / (Fm[3] - Fm[2]))
          p -= strike *
               ((F[2]^2 - strike^2) * (coeff[3] - coeff[2]) / (Fm[3] - Fm[2]) / 2.0 + (F[2] - strike) *( coeff[2]-Fm[2]*(coeff[3] - coeff[2]) / (Fm[3] - Fm[2])))
        end
      else
        Fm = makeForward(maturity, density.zm[k0+1:end]) #zm[k0+1] = zmin+k0*h
        if term > 0 && (smoothing == linear || smoothing == midpoint)
          dFdz = 0.0
          ϑ0 = density.ϑ[k0+1]
          if smoothing == linear
            ftildem = makeForward(maturity, ztilde - density.h)
            bk = (2 * Fm[1] - ftildem - ftilde) / (ftilde - ftildem)
            dFdz = (ftilde - ftildem) / density.h / (1 + bk * (ftilde + 2 * strike - 3 * ftildem) / (ftilde - ftildem))
          elseif smoothing == midpoint
            ftildem = makeForward(maturity, ztilde - density.h)
          #dFdz = (ftilde - ftildem) / density.h      #preserves zeroth
            dFdz = 2 * (ftilde - Fm[1]) / density.h   #equivalent todFdz = (ftilde - Fm[1]) / (ztilde - density.zm[k0+1])
          #dFdz = (ftilde-ftildem)^2 / (2*density.h*(Fm[1]-ftildem)) #preserve price but not zeroth
          end
          p += 0.5 * term * term * ϑ0 / dFdz
        end
        p += sum((Fm[2:end-1] .- strike) * density.h .* density.ϑ[k0+2:end-1])
      end
    end
  end
  if !isCall
    p = p - (forward - strike) # Call-Put = forward - strike
  end

  return p
end

function cumulativeDensity(density::TransformedDensity, isCall::Bool, strike::G, smoothing::DensitySmoothing) where {G}
  maturity = density.maturity
  α = maturity.α
  β = maturity.β
  ν = maturity.ν
  ρ = maturity.ρ
  forward = maturity.forward
  ystrike = yOfStrike(maturity, strike)
  zstrike = -1 / ν * log.((sqrt.(1 .- ρ^2 .+ (ρ + ν * ystrike / α).^2) .- ρ - ν * ystrike / α) / (1 - ρ))
  if (zstrike <= density.zmin)
    p = -1.0
  else
    if (zstrike >= density.zmax)
      p = 0.0
    else
      Fmax = makeForward(maturity, density.zmax)
      p = -density.PR
      k0 = Int.(ceil.((zstrike .- density.zmin - eps()*1000) / density.h))
      ztilde = density.zmin .+ k0 * density.h
      ftilde = makeForward(maturity, ztilde)
      term = ftilde .- strike
      Fm = makeForward(maturity, density.zm[k0+1:end]) #zm[k0+1] = zmin+k0*h
      if smoothing != none && term > 0
        dFdz = 0.0
        ϑ0 = density.ϑ[k0+1]
        if smoothing == linear
          ftildem = makeForward(maturity, ztilde - density.h)
          bk = (2 * Fm[1] - ftildem - ftilde) / (ftilde - ftildem)
          dFdz = (ftilde - ftildem) / density.h / (1 + bk * (ftilde + 2 * strike - 3 * ftildem) / (ftilde - ftildem))
        elseif smoothing == midpoint
          ftildem = makeForward(maturity, ztilde - density.h)
          #dFdz = (ftilde - ftildem) / density.h      #preserves zeroth
          dFdz = 2 * (ftilde - Fm[1]) / density.h   #equivalent todFdz = (ftilde - Fm[1]) / (ztilde - density.zm[k0+1])
          #dFdz = (ftilde-ftildem)^2 / (2*density.h*(Fm[1]-ftildem)) #preserve price but not zeroth
        end
        p += -term * ϑ0 / dFdz
      end
      p += -density.h .* sum(density.ϑ[k0+2:end-1])
    end
  end
  if !isCall
    p = p + 1 # Call-Put = forward - strike
  end
  return p
end

end # module


text:
# This uses reduce instead of sum because sum can't handle empty collections properly.
# There is no universal zero-element, so it has to be specified manually and cannot
# be determined automatically. Otherwise, n=0 will cause an error.
sum_of_squares(n::Int) = reduce(+, i^2 for i in 1:n; init=0)

# However, sum{T<:Real}(r::Range{T}) can handle "empty" ranges.
square_of_sum(n::Int) = sum(1:n)^2

difference(n::Int) = square_of_sum(n) - sum_of_squares(n)


text:
"""
    _checked_rrule

like `rrule` but throws an error if the `rrule` is not defined.
Rather than returning `nothing`
"""
function _checked_rrule(f, args...; kwargs...)
    r = rrule(f, args...; kwargs...)
    r isa Nothing && _throw_checked_rrule_error(f, args...; kwargs...)
    return r
end

@noinline function _throw_checked_rrule_error(f, args...; kwargs...)
    io = IOBuffer()
    print(io, "can't differentiate `", f, '(')
    join(io, map(arg->string("::", typeof(arg)), args), ", ")
    if !isempty(kwargs)
        print(io, ";")
        join(io, map(((k, v),)->string(k, "=", v), kwargs), ", ")
    end
    print(io, ")`; no matching `rrule` is defined")
    throw(ArgumentError(String(take!(io))))
end