text:
module NormalizeQuantiles

export normalizeQuantiles

export sampleRanks
export qnTiesMethods,tmMin,tmMax,tmOrder,tmReverse,tmRandom,tmAverage

using Distributed
using SharedArrays
using Random
using Statistics

@enum qnTiesMethods tmMin tmMax tmOrder tmReverse tmRandom tmAverage

function checkForNotANumber(x::Any)
    (!isa(x,Integer) && !isa(x,Real)) || isnan(x)
end

@doc "
### qnmatrix::Array{Float64} function normalizeQuantiles(matrix::AbstractArray)
Calculate the quantile normalized data for the input matrix

Parameter:
    matrix::AbstractArray
The input data as an array of values interpreted as matrix(rows,columns)

Return value:
    qnmatrix::Array{Float64}
The quantile normalized data as Array{Float64}

Example:

    using NormalizeQuantiles

    array = [ 3.0 2.0 1.0 ; 4.0 5.0 6.0 ; 9.0 7.0 8.0 ; 5.0 2.0 8.0 ]
    qn = normalizeQuantiles(array)

    row = 2
    column = 2
    array=convert(Array{Any},array)
    array[row,column] = missing
    qn = normalizeQuantiles(array)
"
function normalizeQuantiles(matrix::AbstractArray)
    if ndims(matrix) > 2
        throw(ArgumentError("normalizeQuantiles expects an array of dimension 2"))
    end
    nrows=size(matrix,1)
    ncols=size(matrix,2)
    # preparing the result matrix
    qnmatrix=SharedArray{Float64}(nrows,ncols)
    if ncols>0 && nrows>0
        # foreach column: sort the values without NAs; put NAs (if any) back into sorted list
        NormalizeQuantiles.sortColumns!(matrix,qnmatrix)
        # foreach row: set all values to the mean of the row, except NAs
        NormalizeQuantiles.meanRows!(qnmatrix)
        # foreach column: equal values in original column should all be mean of normalized values
        # foreach column: reorder the values back to the original order
        NormalizeQuantiles.equalValuesInColumnAndOrderToOriginal!(matrix,qnmatrix,nrows)
    end
    convert(Array{Float64},qnmatrix)
end

function sortColumns!(matrix::AbstractArray,qnmatrix::SharedArray{Float64})
    tcol=1
    @inbounds @sync @distributed for scolumn in eachindex(matrix[end,:])
        sortcol=[ NormalizeQuantiles.checkForNotANumber(x) ? NaN : Float64(x) for x in matrix[:,scolumn] ]
        missingIndices=findall(isnan.(sortcol))
        sort!(sortcol)
        #putting original NaNs back into place:
        nMissPos=lastindex(missingIndices)
        delta=nMissPos
        index=lastindex(sortcol)
        while nMissPos>0
            missingPos=missingIndices[nMissPos]
            if index==missingPos
                sortcol[index]=NaN
                nMissPos-=1
                delta-=1
            else
                sortcol[index]=sortcol[index-delta]
            end
            index-=1
        end
        qnmatrix[:,tcol]=sortcol
        tcol+=1
    end
end

function meanRows!(qnmatrix::SharedArray{Float64})
    @inbounds @sync @distributed for srow = eachindex(qnmatrix[:,end])
        goodIndices=.!isnan.(qnmatrix[srow,:])
        rowView=view(qnmatrix,srow,goodIndices)
        #rowmean=mean(rowView)
        #
        #in julia function sum the order of summing up the elements can vary on different
        #types of AbstractArray (e.g. qnmatrix[row,goodIndices] and rowView). Summing over
        #a large number of floats in different order yield different results because of
        #floating point precision, so we calculate the sum in a well defined way:
        lrow=length(rowView)
        sum=0.0
        for i in eachindex(rowView)
            sum+=rowView[i]
        end
        rowmean=sum/lrow
        rowView.=rowmean
    end
end

function equalValuesInColumnAndOrderToOriginal!(matrix::AbstractArray,qnmatrix::SharedArray{Float64},nrows)
    tcol=1
    @inbounds @sync @distributed for scolumn in eachindex(matrix[end,:])
        goodIndices=.!isnan.(qnmatrix[:,tcol])
        #matrix can be an OffsetArray
        colview=view(matrix,collect(eachindex(matrix[:,scolumn]))[goodIndices],scolumn)
        sortp=sortperm(Float64.(colview))
        if length(sortp)>0
            NormalizeQuantiles.setMeanForEqualOrigValues!(colview[sortp],qnmatrix,tcol,goodIndices)
        end
        qnmatrix[(1:nrows)[goodIndices][sortp],tcol]=qnmatrix[goodIndices,tcol]
        tcol+=1
    end
end

function setMeanForEqualOrigValues!(sortedArrayNoNAs::AbstractArray,qnmatrix::SharedArray{Float64},column::Int,goodIndices)
    nrows=length(sortedArrayNoNAs)
    foundIndices=zeros(Int,nrows)
    goodIndices2=findall(goodIndices)
    count=1
    lastValue=sortedArrayNoNAs[1]
    for i in 2:nrows
        nextValue=sortedArrayNoNAs[i]
        foundIndices[count]=goodIndices2[i-1]
        if nextValue==lastValue
            count+=1
        else
            if count>1
                qnmatrix[foundIndices[1:count],column].=mean(qnmatrix[foundIndices[1:count],column])
                count=1
            end
        end
        lastValue=nextValue
    end
    foundIndices[count]=goodIndices2[nrows]
    if count>1
        qnmatrix[foundIndices[1:count],column].=mean(qnmatrix[foundIndices[1:count],column])
    end
end

@doc "
### (Array{Union{Missing,Int}},Dict{Int,Array{Int}}) sampleRanks(array::AbstractArray;tiesMethod::qnTiesMethods=tmMin,naIncreasesRank=false,resultMatrix=false)

Calculate ranks of the values of a given vector.

Parameters:

    array: the input array
    tiesMethod: the method how ties (equal values) are treated
       possible values: tmMin tmMax tmOrder tmReverse tmRandom tmAverage
       default is tmMin
    naIncreasesRank: if true than any NA increases the following ranks by 1
    resultMatrix: if true than return a dictionary of rank keys and array of indices values

Example:

    using NormalizeQuantiles

    a = [ 5.0 2.0 4.0 3.0 1.0 ]

    (r,m)=sampleRanks(a)

    (r,m)=sampleRanks(a,tiesMethod=tmMin,naIncreasesRank=false,resultMatrix=true)

r is the vector of ranks.

m is a dictionary with rank as keys and as value the indices of all values of this rank.

"
function sampleRanks(array::AbstractArray;tiesMethod::qnTiesMethods=tmMin,naIncreasesRank::Bool=false,resultMatrix::Bool=false)
    nrows=length(array)
    naCounts=zeros(Int,nrows)
    goodIndices=falses(nrows)
    naCount=0
    reducedIndex=1
    goodIndex=1
    for arrayIndex in eachindex(array)
        if NormalizeQuantiles.checkForNotANumber(array[arrayIndex])
            naCount+=1
        else
            goodIndices[goodIndex]=true
            naCounts[reducedIndex]=naCount
            reducedIndex+=1
        end
        goodIndex+=1
    end
    reducedArraySorted=Float64.(array[goodIndices])
    reducedArraySortedIndices=sortperm(reducedArraySorted)
    reducedArraySorted=reducedArraySorted[reducedArraySortedIndices]
    result=Array{Union{Missing,Int}}(missing,nrows)
    rankMatrix=Dict{Int,Array{Int}}()
    group = Array{Int,1}()
    firstFound = true
    lastFound = NaN
    nextRank=1
    rankIncrement=0
    doIncreaseRank=false
    headingNA=false
    groupIndex=1
    offset=firstindex(array)-groupIndex
    for resultIndex in eachindex(array)
        if goodIndices[resultIndex]
            if headingNA
                headingNA=false
                if doIncreaseRank
                    nextRank+=rankIncrement
                    rankIncrement=0
                    doIncreaseRank=false
                end
            end
            reducedArrayIndex=reducedArraySortedIndices[groupIndex]
            if !firstFound && lastFound != reducedArraySorted[groupIndex]
                nextRank = NormalizeQuantiles.setRank!(group,nextRank,offset,tiesMethod,result,resultMatrix,rankMatrix)
                group = [naCounts[reducedArrayIndex]+reducedArrayIndex+offset]
                if doIncreaseRank
                    nextRank+=rankIncrement
                    rankIncrement=0
                    doIncreaseRank=false
                end
            else
                push!(group,naCounts[reducedArrayIndex]+reducedArrayIndex+offset)
            end
            lastFound = reducedArraySorted[groupIndex]
            firstFound = false
            groupIndex+=1
        else
            if naIncreasesRank
                doIncreaseRank=true
                rankIncrement+=1
            end
            if firstFound
                headingNA=true
            end
        end
    end
    NormalizeQuantiles.setRank!(group,nextRank,offset,tiesMethod,result,resultMatrix,rankMatrix)
    (result,rankMatrix)
end

function setRank!(group::Array{Int},nextRank::Int,offset::Int,tiesMethod::qnTiesMethods,result::Array{Union{Missing,Int}},resultMatrix::Bool,rankMatrix::Dict{Int,Array{Int}})
    ranksCount=length(group)
    ranks=nextRank:(nextRank+ranksCount-1)
    if tiesMethod==tmMin
        minRank=minimum(ranks)
        result[group.-offset].=minRank
        if resultMatrix
            rankMatrix[minRank]=group
        end
        nextRank+=1
    elseif tiesMethod==tmMax
        maxRank=maximum(ranks)
        result[group.-offset].=maxRank
        if resultMatrix
            rankMatrix[maxRank]=group
        end
        nextRank=maxRank+1
    elseif tiesMethod==tmOrder
        maxRank=maximum(ranks)
        result[group.-offset]=ranks
        if resultMatrix
            for rankIndex in 1:length(ranks)
                rankMatrix[ranks[rankIndex]]=[group[rankIndex]]
            end
        end
        nextRank=maxRank+1
    elseif tiesMethod==tmReverse
        maxRank=maximum(ranks)
        ranks=reverse(ranks,dims=1)
        result[group.-offset]=ranks
        if resultMatrix
            for rankIndex in 1:length(ranks)
                rankMatrix[ranks[rankIndex]]=[group[rankIndex]]
            end
        end
        nextRank=maxRank+1
    elseif tiesMethod==tmRandom
        maxRank=maximum(ranks)
        ranks=ranks[randperm(ranksCount)]
        result[group.-offset]=ranks
        if resultMatrix
            for rankIndex in 1:length(ranks)
                rankMatrix[ranks[rankIndex]]=[group[rankIndex]]
            end
        end
        nextRank=maxRank+1
    elseif tiesMethod==tmAverage
        rankMean=round(Int,mean(ranks))
        result[group.-offset].=rankMean
        if resultMatrix
            rankMatrix[rankMean]=group
        end
        nextRank=rankMean+1
    end
    nextRank
end

end # module


text:
using InducingPoints
using Documenter

ENV["GKSwstype"] = "100"

makedocs(;
    modules=[InducingPoints],
    authors="Theo Galy-Fajou <theo.galyfajou@gmail.com> and contributors",
    repo="https://github.com/JuliaGaussianProcesses/InducingPoints.jl/blob/{commit}{path}#L{line}",
    sitename="InducingPoints.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://JuliaGaussianProcesses.github.io/InducingPoints.jl",
        assets=String[],
    ),
    strict=true,
    pages=["Home" => "index.md", "Algorithms" => "algorithms.md", "API" => "api.md"],
)

deploydocs(; repo="github.com/JuliaGaussianProcesses/InducingPoints.jl", push_preview=true)

ext: jl


text:
matern12(x1, x2; ρ=ρ, σ²=σ²) = MaternCovariance_12(abs(x1-x2); ρ=ρ, σ²=σ²)

function MaternCovariance_12(d; ρ=ρ, σ²=σ²)

    @assert(d >= 0)

    σ² * exp(-d/ρ)

end

matern32(x1, x2; ρ=ρ, σ²=σ²) = MaternCovariance_32(abs(x1-x2); ρ=ρ, σ²=σ²)

function MaternCovariance_32(d; ρ=ρ, σ²=σ²)

    @assert(d >= 0)

    σ² * (1 + sqrt(3)*d/ρ) * exp(- sqrt(3)*d/ρ)

end

matern52(x1, x2; ρ=ρ, σ²=σ²) = MaternCovariance_52(abs(x1-x2); ρ=ρ, σ²=σ²)

function MaternCovariance_52(d; ρ=ρ, σ²=σ²)

    @assert(d >= 0)

    σ² * (1 + sqrt(5)*d/ρ + 5*d^2/(3*ρ^2)) * exp(- sqrt(5)*d/ρ)

end

rbf(x1, x2; ρ=ρ, σ²=σ²) = RBFCovariance(abs(x1-x2); ρ=ρ, σ²=σ²)

function RBFCovariance(d; ρ=ρ, σ²=σ²)

    @assert(d >= 0)

    σ² * exp(-d^2 / (2ρ))

end

linear(x1, x2; c=c, σ²=σ²) = σ² * (x1 - c) * (x2 - c)

function calculatekernelmatrix(k, ρ, σ², x)

    calculatekernelmatrix(k, ρ, σ², x, x)

end

function calculatekernelmatrix(k, ρ, σ², x1, x2)

    [k(a, b; ρ = ρ, σ² = σ²) for a in x1, b in x2]

end


text:
"""
    tune_δ(algorithm, time_horizon, [precision], [δ], [options])

Find the threshold step size for a given property.

### Input

- `algorithm`        -- the function that is used to evaluate ``δ``;
                        the call should be: `algorithm(N, δ)`
- `time_horizon`     -- time horizon
- `precision`        -- (optional, default: 1e-5) precision of result (s. a.)
- `δ`                -- (optional, default: 1e-1) initial guess for the time
                        step
- `rel_step_size`    -- (optional, default: 0.1) relative step size
- `biggest_working`  -- (optional, default: 0.) biggest known working ``δ``
- `smallest_failing` -- (optional, default: `time_horizon`) smallest known
                        failing ``δ``

### Algorithm

Finds a value for ``δ`` such that `algorithm`

  1. returns `true` for ``δ`` and
  2. returns `false` for (``δ`` + `precision`).
"""
function tune_δ(algorithm::Function,
                time_horizon::Float64,
                precision::Float64=1e-5,
                δ::Float64=1e-1,
                rel_step_size::Float64=0.1;
                biggest_working=0.,
                smallest_failing=time_horizon
                )::Float64
    biggest_working_modified = biggest_working != 0.
    smallest_failing_modified = smallest_failing != time_horizon

    step = 0.
    runtime_best = -1.

    if δ <= 0. || precision <= 0.
        throw(DomainError())
    end

    i = 0
    while true
        i += 1
        info("Iteration $i")
        info("δ: $δ")

        @assert (δ >= biggest_working && δ <= smallest_failing) "invalid δ: " *
            "$δ (not in [$biggest_working, $smallest_failing])"

        # choose N according to δ
        N = ceil(Int, time_horizon / δ)

        # check property
        runtime = @elapsed begin
            answer = algorithm(N, δ)
        end
        if answer
            # success
            biggest_working = δ
            biggest_working_modified = true
            runtime_best = runtime
        else
            # failure
            smallest_failing = δ
            smallest_failing_modified = true
        end

        # choose new δ
        if smallest_failing - biggest_working <= precision
            # reached precision window -> terminate
            break
        elseif biggest_working_modified && smallest_failing_modified
            # choose δ in the center of the uncertainty interval
            δ_new = smallest_failing + (biggest_working - smallest_failing) / 2.
        elseif biggest_working_modified
            # no upper bound yet, choose δ that is bigger wrt. the relative ball
            δ_new = min(δ * (1. + rel_step_size), time_horizon)
        else
            @assert smallest_failing_modified
            # no lower bound yet, choose δ that is bigger wrt. the relative ball
            δ_new = max(δ * (1. - rel_step_size), 0.)
        end

        # round δ down wrt. precision
        inverse_precision = round(Int, 1. / precision)
        δ_new = floor(δ_new * inverse_precision) / inverse_precision
        if δ_new >= smallest_failing || δ_new <= biggest_working
            # reached precision window -> terminate
            break
        end

        δ = δ_new
    end
    info("biggest working δ found: $biggest_working, runs in $runtime_best sec)")
    info("smallest failing δ found: " *
         (smallest_failing_modified ? "$smallest_failing" : "--"))
    return δ
end


text:
const TableauHSPARKsecondary{DT} = AbstractTableauSPARK{:hspark_secondary,DT}
const ParametersHSPARKsecondary = AbstractParametersHSPARK{:hspark_secondary}

@doc raw"""
Specialised Partitioned Additive Runge-Kutta integrator for Hamiltonian systems
subject to Dirac constraints with projection on secondary constraint.

This integrator solves the following system of equations for the internal stages,
```math
\begin{align}
Q_{n,i} &= q_{n} + h \sum \limits_{j=1}^{s} a_{ij} V_{n,j} + h \sum \limits_{j=1}^{r} \alpha_{ij} U_{n,j} , & i &= 1, ..., s , \\
P_{n,i} &= p_{n} + h \sum \limits_{j=1}^{s} a_{ij} F_{n,j} + h \sum \limits_{j=1}^{r} \alpha_{ij} G_{n,j} , & i &= 1, ..., s , \\
\tilde{Q}_{n,i} &= q_{n} + h \sum \limits_{j=1}^{s} \tilde{a}_{ij} V_{n,j} + h \sum \limits_{j=1}^{r} \tilde{\alpha}_{ij} U_{n,j} , & i &= 1, ..., r , \\
\tilde{P}_{n,i} &= p_{n} + h \sum \limits_{j=1}^{s} \tilde{a}_{ij} F_{n,j} + h \sum \limits_{j=1}^{r} \tilde{\alpha}_{ij} G_{n,j} , & i &= 1, ..., r , \\
0 &= \sum \limits_{j=1}^{r} \omega_{ij} \tilde{\Phi}_{n,j} , & i &= 1, ..., r-1 ,
\end{align}
```
with definitions
```math
\begin{align}
V_{n,i} &= \hphantom{-} \frac{\partial H}{\partial p} (Q_{n,i}, P_{n,i}) , & i &= 1, ..., s , \\
F_{n,i} &=           -  \frac{\partial H}{\partial q} (Q_{n,i}, P_{n,i}) , & i &= 1, ..., s , \\
U_{n,i} &= \hphantom{-} \frac{\partial \phi}{\partial p} (\tilde{Q}_{n,i}, \tilde{P}_{n,i})^{T} \Lambda_{n,i} , & i &= 1, ..., r , \\
G_{n,i} &=           -  \frac{\partial \phi}{\partial q} (\tilde{Q}_{n,i}, \tilde{P}_{n,i})^{T} \Lambda_{n,i} , & i &= 1, ..., r , \\
\tilde{\Phi}_{n,i} &= \phi(\tilde{Q}_{n,i}, \tilde{P}_{n,i}) , & i &= 1, ..., r ,
\end{align}
```
and update rule
```math
\begin{align}
q_{n+1} &= q_{n} + h \sum \limits_{i=1}^{s} b_{i} V_{n,i} + h \sum \limits_{i=1}^{r} \beta_{i} U_{n,i} , \\
p_{n+1} &= p_{n} + h \sum \limits_{i=1}^{s} b_{i} F_{n,i} + h \sum \limits_{i=1}^{r} \beta_{i} G_{n,i} , \\
0 &= \phi (q_{n+1}, p_{n+1}) .
\end{align}
```
"""
struct IntegratorHSPARKsecondary{DT, TT, ET <: PDAE{DT,TT},
                                       PT <: ParametersHSPARKsecondary{DT,TT},
                                       ST <: NonlinearSolver{DT},
                                       IT <: InitialGuessPODE{DT,TT}} <: AbstractIntegratorHSPARK{DT, TT}
    equation::ET
    tableau::TableauHSPARKsecondary{TT}

    params::PT
    solver::ST
    iguess::IT
end

function IntegratorHSPARKsecondary(equation::VDAE{DT,TT},
                                   tableau::TableauVSPARKsecondary{TT}, Δt::TT) where {DT,TT}
    D = equation.d
    S = tableau.s
    Σ = tableau.σ

    # println(tableau)

    N = 2*D*S + 4*D*Σ

    if isdefined(tableau, :d) && length(tableau.d) > 0
        N += D
    end

    # create params
    params = ParametersHSPARKsecondary{DT,D,S,Σ}(equation.ϑ, equation.f, equation.g, equation.g̅, equation.ϕ, equation.ψ, Δt, tableau)

    # create solver
    solver = create_nonlinear_solver(DT, N, params)

    # create initial guess
    iguess = InitialGuessPODE(get_config(:ig_interpolation), equation, Δt)

    # create integrator
    IntegratorHSPARKsecondary{DT, TT, typeof(equation), typeof(params), typeof(solver), typeof(iguess)}(
                              equation, tableau, params, solver, iguess)
end

pstages(int::IntegratorHSPARKsecondary) = int.tableau.σ

function compute_stages!(x::Vector{ST}, cache::IntegratorCacheSPARK{ST,TT,D,S,Σ},
                                        params::ParametersHSPARKsecondary{DT,TT,D,S,Σ}) where {ST,DT,TT,D,S,Σ}
    local t::TT

    for i in 1:S
        for k in 1:D
            # copy x to Y, Z
            cache.Yi[i][k] = x[2*(D*(i-1)+k-1)+1]
            cache.Zi[i][k] = x[2*(D*(i-1)+k-1)+2]

            # compute Q and P
            cache.Qi[i][k] = params.q[k] + params.Δt * cache.Yi[i][k]
            cache.Pi[i][k] = params.p[k] + params.Δt * cache.Zi[i][k]
        end

        # compute f(X)
        t = params.t + params.Δt * params.tab.p.c[i]
        params.f_v(t, cache.Qi[i], cache.Pi[i], cache.Vi[i])
        params.f_f(t, cache.Qi[i], cache.Pi[i], cache.Fi[i])
    end

    for i in 1:Σ
        for k in 1:D
            # copy y to Y, Z and Λ
            cache.Yp[i][k] = x[2*D*S+4*(D*(i-1)+k-1)+1]
            cache.Zp[i][k] = x[2*D*S+4*(D*(i-1)+k-1)+2]
            cache.Up[i][k] = x[2*D*S+4*(D*(i-1)+k-1)+3]
            cache.Λp[i][k] = x[2*D*S+4*(D*(i-1)+k-1)+4]

            # compute Q and V
            cache.Qp[i][k] = params.q[k] + params.Δt * cache.Yp[i][k]
            cache.Pp[i][k] = params.p[k] + params.Δt * cache.Zp[i][k]
        end

        # compute f(X)
        t = params.t + params.Δt * params.tab.p̃.c[i]
        params.f_v(t, cache.Qp[i], cache.Pp[i], cache.Vp[i])
        params.f_f(t, cache.Qp[i], cache.Pp[i], cache.Fp[i])

        params.f_g(t, cache.Qp[i], cache.Pp[i], cache.Up[i], cache.Gp[i])
        params.f_g(t, cache.Qp[i], cache.Pp[i], cache.Λp[i], cache.G̅p[i])

        params.f_ϕ(t, cache.Qp[i], cache.Pp[i], cache.Φp[i])
        params.f_ψ(t, cache.Qp[i], cache.Pp[i], cache.Vp[i], cache.Fp[i], cache.Ψp[i])
    end

    # if length(params.tab.d) > 0
    #     for k in 1:D
    #         cache.μ[k] = x[2*D*S+4*D*Σ+k]
    #     end
    # end

    # compute q and p
    cache.q̃ .= params.q
    cache.p̃ .= params.p
    for i in 1:S
        cache.q̃ .+= params.Δt .* params.tab.q.b[1][i] .* cache.Vi[i]
        cache.p̃ .+= params.Δt .* params.tab.p.b[1][i] .* cache.Fi[i]
    end
    for i in 1:Σ
        cache.q̃ .+= params.Δt .* params.tab.q.b[2][i] .* cache.Up[i]
        cache.q̃ .+= params.Δt .* params.tab.q.b[3][i] .* cache.Λp[i]
        cache.p̃ .+= params.Δt .* params.tab.p.b[2][i] .* cache.Gp[i]
        cache.p̃ .+= params.Δt .* params.tab.p.b[3][i] .* cache.G̅p[i]
    end

    # compute ϕ(q,p)
    t = params.t + params.Δt
    params.f_ϕ(t, cache.q̃, cache.p̃, cache.ϕ̃)
end

"Compute stages of specialised partitioned additive Runge-Kutta methods for variational systems."
@generated function function_stages!(y::Vector{ST}, b::Vector{ST}, params::ParametersHSPARKsecondary{DT,TT,D,S,Σ}) where {ST,DT,TT,D,S,Σ}
    cache = IntegratorCacheSPARK{ST,TT,D,S,Σ}()

    quote
        compute_stages!(y, $cache, params)

        # compute b = - [(Y-AV-AU), (Z-AF-AG), Φ]
        for i in 1:S
            for k in 1:D
                b[2*(D*(i-1)+k-1)+1] = - $cache.Yi[i][k]
                b[2*(D*(i-1)+k-1)+2] = - $cache.Zi[i][k]
                for j in 1:S
                    b[2*(D*(i-1)+k-1)+1] += params.tab.q.a[1][i,j] * $cache.Vi[j][k]
                    b[2*(D*(i-1)+k-1)+2] += params.tab.p.a[1][i,j] * $cache.Fi[j][k]
                end
                for j in 1:Σ
                    b[2*(D*(i-1)+k-1)+1] += params.tab.q.a[2][i,j] * $cache.Up[j][k]
                    b[2*(D*(i-1)+k-1)+1] += params.tab.q.a[3][i,j] * $cache.Λp[j][k]
                    b[2*(D*(i-1)+k-1)+2] += params.tab.p.a[2][i,j] * $cache.Gp[j][k]
                    b[2*(D*(i-1)+k-1)+2] += params.tab.p.a[3][i,j] * $cache.G̅p[j][k]
                end
            end
        end

        # compute b = - [(Y-AV-AU), (Z-AF-AG), Φ, ωΨ]
        for i in 1:Σ
            for k in 1:D
                b[2*D*S+4*(D*(i-1)+k-1)+1] = - $cache.Yp[i][k]
                b[2*D*S+4*(D*(i-1)+k-1)+2] = - $cache.Zp[i][k]
                b[2*D*S+4*(D*(i-1)+k-1)+3] = - $cache.Φp[i][k]
                # b[2*D*S+4*(D*(i-1)+k-1)+4] = - $cache.Ψp[i][k]
                b[2*D*S+4*(D*(i-1)+k-1)+4] = 0
                for j in 1:S
                    b[2*D*S+4*(D*(i-1)+k-1)+1] += params.tab.q̃.a[1][i,j] * $cache.Vi[j][k]
                    b[2*D*S+4*(D*(i-1)+k-1)+2] += params.tab.p̃.a[1][i,j] * $cache.Fi[j][k]
                end
                for j in 1:Σ
                    b[2*D*S+4*(D*(i-1)+k-1)+1] += params.tab.q̃.a[2][i,j] * $cache.Up[j][k]
                    b[2*D*S+4*(D*(i-1)+k-1)+1] += params.tab.q̃.a[3][i,j] * $cache.Λp[j][k]
                    b[2*D*S+4*(D*(i-1)+k-1)+2] += params.tab.p̃.a[2][i,j] * $cache.Gp[j][k]
                    b[2*D*S+4*(D*(i-1)+k-1)+2] += params.tab.p̃.a[3][i,j] * $cache.G̅p[j][k]
                end
                for j in 1:Σ
                    b[2*D*S+4*(D*(i-1)+k-1)+4] -= params.tab.ω[i,j] * $cache.Ψp[j][k]
                end
                b[2*D*S+4*(D*(i-1)+k-1)+4] -= params.tab.ω[i,Σ+1] * $cache.ϕ̃[k]
            end
        end

        # if length(params.tab.d) > 0
        #     for i in 1:Σ
        #         for k in 1:D
        #             b[2*D*S+4*(D*(i-1)+k-1)+2] -= $cache.μ[k] * params.tab.d[i] / params.tab.p.b[2][i]
        #         end
        #     end
        #
        #     for k in 1:D
        #         b[2*D*S+4*D*Σ+k] = 0
        #         for i in 1:Σ
        #             b[2*D*S+4*D*Σ+k] -= $cache.Vp[i][k] * params.tab.d[i]
        #         end
        #     end
        # end
    end
end

function initial_guess!(int::IntegratorHSPARKsecondary, cache::IntegratorCacheSPARK)
    for i in 1:nstages(int)
        evaluate!(int.iguess, cache.q, cache.p, cache.v, cache.f,
                              cache.q̅, cache.p̅, cache.v̅, cache.f̅,
                              cache.q̃, cache.p̃, cache.ṽ, cache.f̃,
                              tableau(int).q.c[i], tableau(int).p.c[i])

        for k in 1:ndims(int)
            int.solver.x[2*(ndims(int)*(i-1)+k-1)+1] = (cache.q̃[k] - cache.q[k])/timestep(int)
            int.solver.x[2*(ndims(int)*(i-1)+k-1)+2] = (cache.p̃[k] - cache.p[k])/timestep(int)
        end
    end

    for i in 1:pstages(int)
        evaluate!(int.iguess, cache.q, cache.p, cache.v, cache.f,
                              cache.q̅, cache.p̅, cache.v̅, cache.f̅,
                              cache.q̃, cache.p̃, cache.ṽ, cache.f̃,
                              tableau(int).q̃.c[i], tableau(int).p̃.c[i])

        for k in 1:ndims(int)
            int.solver.x[2*ndims(int)*nstages(int)+4*(ndims(int)*(i-1)+k-1)+1] = (cache.q̃[k] - cache.q[k])/timestep(int)
            int.solver.x[2*ndims(int)*nstages(int)+4*(ndims(int)*(i-1)+k-1)+2] = (cache.p̃[k] - cache.p[k])/timestep(int)
            int.solver.x[2*ndims(int)*nstages(int)+4*(ndims(int)*(i-1)+k-1)+3] = 0
            int.solver.x[2*ndims(int)*nstages(int)+4*(ndims(int)*(i-1)+k-1)+4] = 0
        end
    end

    if length(tableau(int).d) > 0
        for k in 1:ndims(int)
            int.solver.x[2*ndims(int)*nstages(int)+4*ndims(int)*pstages(int)+k] = 0
        end
    end
end

function update_solution!(int::IntegratorHSPARKsecondary{DT,TT}, sol::AtomicSolutionPDAE{DT,TT}) where {DT,TT}
    # compute final update
    update_solution!(sol.q, sol.q̃, int.cache.Vi, int.params.tab.q.b[1], timestep(int))
    update_solution!(sol.p, sol.p̃, int.cache.Fi, int.params.tab.p.b[1], timestep(int))

    # compute projection
    update_solution!(sol.q, sol.q̃, int.cache.Up, int.params.tab.q.b[2], timestep(int))
    update_solution!(sol.q, sol.q̃, int.cache.Λp, int.params.tab.q.b[3], timestep(int))
    update_solution!(sol.p, sol.p̃, int.cache.Gp, int.params.tab.p.b[2], timestep(int))
    update_solution!(sol.p, sol.p̃, int.cache.G̅p, int.params.tab.p.b[3], timestep(int))
end


text:
# This example demonstrates one-sided communication,
# specifically passive Remote Memory Access (RMA)

using MPI

MPI.Init()
const world_sz = MPI.Comm_size(MPI.COMM_WORLD)
const rank = MPI.Comm_rank(MPI.COMM_WORLD)

# allocate memory
all_ranks = fill(-1, world_sz)
# create RMA window on all ranks
win = MPI.Win_create(all_ranks, MPI.COMM_WORLD)

# let each rank write its rank number into window
if rank != 0
    # lock window (MPI.LOCK_SHARED works as well)
    MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 0, 0, win)
    # each rank writes to exposed windows of rank 0
    # Signature: obj, target_rank, target_displacement, window
    MPI.Put(rank, 0, rank, win)
    # finish the communication epoch
    MPI.Win_unlock(0, win)
else
    all_ranks[1] = 0
end

# wait with printing
MPI.Win_fence(0, win)

# print window content on all ranks
if rank == 0
    println("After Put with lock / unlock, window content on rank 0:")
    @show all_ranks
end

# free window
MPI.free(win)


text:
module CorpusUtils
using DataStructures, DynamicData

include("./types.jl")
#Lexicon basic methods
Base.length(l::CorpusUtils.Lexicon) = Base.length(l.words)

function addword(l::CorpusUtils.Lexicon, s::AbstractString, allow_growing = true)
	if haskey(l.string_dict, s)
		w = l.words[l.string_dict[s]]
		w.freq +=1
		return w
	end
	new_word = Word(length(l)+1, s)
	push!(l.words, new_word)
	l.string_dict[s] = new_word.id
	new_word.freq += 1
	new_word
end

function add_text_file(c::CorpusUtils.Corpus, filename::String, timestamp::DateTime; stop_word_file::String="data/stopwords/stop_snowball.txt")
	lines = readlines(filename)
	content = ""
	for line in lines
		content *= line
	end
	d = Document()
	content = replace(content, r"-.?\n", "")
	stopwords = String[]
	try
		stopwords = read_stop_word_file(stop_word_file)
	catch end
	try
		for line in split(content, '\n')
			words = split(line, ' ')
			for word in words
				word = strip(word)
				word = replace(word,r"[[:punct:]]","")
				word = replace(word,['\n','\t']," ")
				word = lowercase(word)
				if word == "" || length(word) ≤ 3 || !isalpha(word) || word ∈ stopwords
					continue
				end
				word_obj = addword(c.lexicon, word)
				if haskey(d, word_obj)
					d[word_obj] += 1
				else
					d[word_obj] = 1
				end
				# println("$word ($(length(word)), $(length(strip(word))))")
			end
		end
	catch UnicodeError
		println("current word is no unicode")
	end

    dp = DynamicData.add(c.documents, d, timestamp)

	for w_obj in collect(keys(d))
		w_obj.doc_freq += 1
		push!(w_obj.used_in_docs, dp)
	end
end

function add_text_file_directory(c::Corpus, path::String, timestamp::DateTime; stop_word_file::String="data/stopwords/stop_snowball.txt")
	for file in readdir(path)
		cur_path = joinpath(path, file)
		if isfile(cur_path)
			add_text_file(c, cur_path, timestamp, stop_word_file=stop_word_file)
		end
	end
end

function read_stop_word_file(filename::String)
	stopwords = String[]
	f = open(filename, "r")
	while !eof(f)
		cur_line = readline(f)
		cur_line = chomp(cur_line)
		if length(cur_line) == 0
			continue
		end
		w = strip(split(cur_line, "|")[1])
		w = replace(w,r"[[:punct:]]","")
		w = replace(w,['\n','\t']," ")
		w = lowercase(w)

		push!(stopwords, w)
	end
	stopwords = collect(filter(w->length(w) != 0, stopwords))
	close(f)
	stopwords
end

#pruning methods
function prune_dictionary!(c::CorpusUtils.Corpus, remove_crit::Function, remove_words_only::Bool=false)
	words_exclude = filter(remove_crit, c.lexicon.words)
	println("removing $(length(words_exclude)) words from corpus")
	c.lexicon.words = setdiff(c.lexicon.words, words_exclude)
	# c.lexicon.words = words_include
	empty!(c.lexicon.string_dict)
	for (id, word) in enumerate(c.lexicon.words)
		c.lexicon.string_dict[word.value] = id
		word.id = id
	end
	if !remove_words_only
		for w in words_exclude
			[delete!(doc.value,w) for doc in w.used_in_docs]
		end
		prune_documents!(c, 1)
	end
end

function prune_documents!(corpus::Corpus, minimum_length::Int64=10)
	docs_remove = filter(x->length(x.value) < minimum_length, corpus.documents)
	# docs_remove = setdiff(collect(1:length(corpus.documents)), docs_keep)
	println("removing $(length(collect(docs_remove))) documents")
	for (i, doc) in enumerate(docs_remove)
		# t_doc = DynamicData.get_timestamp_index(corpus.documents, corpus.documents[d_id].timestamp)
		(words, freqs) = get_words_and_counts(doc.value)
		#first update word stats for each of the words in the document
		for (wi, w) in enumerate(words)
			corpus.lexicon.words[w].freq -= freqs[wi]
			corpus.lexicon.words[w].doc_freq -= 1
			corpus.lexicon.words[w].used_in_docs = setdiff(corpus.lexicon.words[w].used_in_docs, [doc])
			# deleteat!(, findin(corpus.lexicon.words[w].used_in_docs, [doc])[1])
		end
		#remove doc
		DynamicData.remove!(corpus.documents, doc)
	end
	#finally check if we have zero freq words
	prune_unused_words!(corpus)
end

function prune_stop_words!(c::Corpus, stop_words_file)
	words = read_stop_word_file(stop_words_file)
	prune_dictionary!(c, w -> w.value ∈ words)
end

function prune_word_doc_ratio!(c::Corpus, max_doc_ratio::Float64)
	all_docs = length(c.documents)
	prune_dictionary!(c, w -> (length(w.used_in_docs)/all_docs) > max_doc_ratio)
end

prune_word_length!(c::Corpus, min_length::Int64) = prune_dictionary!(c, w -> length(w.value) < min_length)
prune_word_frequency!(c::Corpus, min_freq::Int64) = prune_dictionary!(c, w -> w.freq < min_freq)
prune_word_term!(c::Corpus, term::String) = prune_dictionary!(c, w -> w.value == term)
prune_unused_words!(c::Corpus) = prune_dictionary!(c, w -> w.freq < 1, true)
prune_sparse_words!(c::Corpus, ratio::Float64) = prune_dictionary!(c, w -> ratio > length(w.used_in_times)/length(DynamicData.get_unique_timestamps(c.documents)))

function prune_words_tfidf!(c::Corpus, score_threshold::Float64 = .05)
	all_words = sum(map(x->x.freq, c.lexicon.words))
	all_docs = length(c.documents)
	scores = Dict{Any, Float64}()
	[scores[w] = w.freq/all_words * log(all_docs/w.doc_freq) for w in c.lexicon.words]
	max_value = maximum(collect(values(scores)))
	prune_dictionary!(c, w -> (w.freq/all_words * log(all_docs/w.doc_freq))/max_value < score_threshold)
end
"""
Partitions the dataset by rounding the timestamps to the provided granularity.
Returns a new DynamicDateSet with rounded timestamps (only works for DateTime timestamps).
"""
function get_partitioning(source::CorpusUtils.Corpus, granularity::Union{Type{Dates.Year}, Type{Dates.Month}, Type{Dates.Week}, Type{Dates.Day}})
	target = DynamicData.DynamicDataSet{Document, DateTime}()
	target_lexicon = Lexicon()
	for w in source.lexicon.words
		w_new = Word(w.id, w.value)
		w_new.freq = w.freq
		w_new.doc_freq = w.doc_freq
		w_new.used_in_docs = Vector{DynamicData.DataPoint{Document, DateTime}}()
		push!(target_lexicon.words, w_new)
	end
	target_lexicon.string_dict = deepcopy(source.lexicon.string_dict)
	for d in source.documents
		dp = DynamicData.add(target, d.value, round(d.timestamp, granularity), false)
		[push!(w.used_in_docs, dp) for w in keys(dp.value)]
	end
	DynamicData.rebuild_indexes(target)
	Corpus{DynamicData.DynamicDataSet{Document, DateTime}}(target, target_lexicon)
end

#i/o

function write_sequential_corpus(corpus, filename_corpus, filename_lexicon)
	all_times = DynamicData.get_unique_timestamps(corpus.documents)
	f = open(filename_corpus,"w")
	@printf(f, "%d\n", length(all_times))
	for ts in all_times
		all_docs_t = DynamicData.get_datapoints_for_timestamp(corpus.documents, ts)
		@printf(f, "%f\n%d\n", Dates.datetime2unix(ts), length(all_docs_t))
		for doc in all_docs_t
			(words, freqs) = get_words_and_counts(doc.value)
			@printf(f, "%d", length(words))
			for (i,w) in enumerate(words)
				@printf(f, " %d:%d", w-1, freqs[i])
			end
			@printf(f,"\n")
		end
	end
	close(f)

	f = open(filename_lexicon, "w")
	for w in corpus.lexicon.words
		@printf(f, "%s\n", w.value)
	end
	close(f)
end

function write_corpus(corpus, filename_corpus, filename_lexicon)
	f = open(filename_corpus, "w")
	for doc in corpus.documents
		(words, freqs) = get_words_and_counts(doc.value)
		@printf(f, "%d", length(words))
		for (i,w) in enumerate(words)
			@printf(f, " %d:%d", w-1, freqs[i])
		end
		@printf(f,"\n")
	end
	close(f)
	f = open(filename_lexicon, "w")
	for w in corpus.lexicon.words
		@printf(f, "%s\n", w.value)
	end
	close(f)
end

#read sequential num corpus
function read_sequential_corpus(f, l::Lexicon)
	stream = open(f,"r")
	numTimes = parse(Int64, readline(stream))
	# assume seconds after epoch

	# DynamicData
	c = DynamicData.DynamicDataSet{Document, DateTime}()
	for t=1:numTimes
	    if t % 100 == 0
	        print("t=", t)
	        print("\n")
	    end
		cur_date = toDate(readline(stream))
		# cur_date.num_val = t
		numDocsT = parse(Int64,readline(stream))
		for i in 1:numDocsT
			doc = Document()
			#get a doc id
			dp = DynamicData.add(c, doc, cur_date, false)
			t_id = DynamicData.get_timestamp_index(c, cur_date)
			#add the content
			addToDocument(dp, split(readline(stream), "#")[1], l, t_id)
		end
	end
	close(stream)
	DynamicData.rebuild_indexes(c)
	c
end

function toDate(line::AbstractString)
	Dates.unix2datetime(round(Int64, parse(Float64, chomp(line))))
end

function addToDocument(doc::DynamicData.DataPoint, line::AbstractString, l::Lexicon, t_id::Int64)
	for term in split(chomp(line))[2:end]
		(w, doc_freq) = splitTermFreq(term, l)
		push!(w.used_in_docs, doc)
		push!(doc.value, w => doc_freq)
	end
end

function splitTermFreq(term::AbstractString, l::Lexicon)
	parts = split(term, ":")
	w = l.words[parse(Int64, parts[1])+1]
	doc_freq = parse(Int64,parts[2])
	w.freq += doc_freq
	w.doc_freq += 1
	Pair(w, doc_freq)
end

#read lexicon
function read_lexicon(stream)
	lines = readlines(stream)
	Lexicon(map(chomp, lines))
end

function get_words_and_counts(doc::Document)
	(map(w->w.id, collect(keys(doc))), collect(values(doc)))
end

function create_corpus(lexicon_file::String, corpus_file::String)
	l = CorpusUtils.read_lexicon(lexicon_file)
	c = CorpusUtils.read_sequential_corpus(corpus_file, l)
	CorpusUtils.Corpus{DynamicData.DynamicDataSet{Document, DateTime}}(c,l)
end

function used_in_times(c::Corpus, w::Word)
	DynamicData.get_timestamp_indexes_for_subset(c.documents, convert(Vector{DynamicData.DataPoint{Document, DateTime}}, w.used_in_docs))
end
end # module


text:
# convert matrix to band storage
function band(A::AbstractMatrix,kl,ku)
    m, n = size(A)
    AB = zeros(eltype(A),kl+ku+1,n)
    for j = 1:n
        for i = max(1,j-ku):min(m,j+kl)
            AB[ku+1-j+i,j] = A[i,j]
        end
    end
    return AB
end

# convert band storage to general matrix
function unband(AB::AbstractMatrix,m,kl,ku)
    bm, n = size(AB)
    A = zeros(eltype(AB),m,n)
    for j = 1:n
        for i = max(1,j-ku):min(m,j+kl)
            A[i,j] = AB[ku+1-j+i,j]
        end
    end
    return A
end

# zero out elements not on matrix bands
function bandex(A::AbstractMatrix,kl,ku)
    m, n = size(A)
    AB = band(A,kl,ku)
    B = unband(AB,m,kl,ku)
    return B
end

const CublasFloat = Union{Float64,Float32,ComplexF64,ComplexF32}
const CublasReal = Union{Float64,Float32}
const CublasComplex = Union{ComplexF64,ComplexF32}


text:
"""
    AbstractSchedule

All schedules should inherit from this type or its subtypes.
Concrete subtypes are expected to conform to the
[iteration interface](https://docs.julialang.org/en/v1/manual/interfaces/#man-interface-iteration)
and implement [`Base.getindex`](https://docs.julialang.org/en/v1/manual/interfaces/#Indexing).
"""
abstract type AbstractSchedule end

Base.firstindex(schedule::AbstractSchedule) = 1

"""
    ScheduleIterator{T<:AbstractSchedule, S}
    ScheduleIterator(schedule::T)

Create a stateful iterator around `schedule`.
See also [`next!`](#)
"""
mutable struct ScheduleIterator{T<:AbstractSchedule, S}
    schedule::T
    state::Union{S, Nothing}
end
function ScheduleIterator(schedule::T) where T<:AbstractSchedule
    _, state = iterate(schedule)

    ScheduleIterator{T, typeof(state)}(schedule, nothing)
end

"""
    next!(iter::ScheduleIterator)

Advance `iter` by one iteration and return the next value.
See also [`ScheduleIterator`](#)
"""
function next!(iter::ScheduleIterator)
    val, iter.state = isnothing(iter.state) ? iterate(iter.schedule) : iterate(iter.schedule, iter.state)

    return val
end

"""
    DecaySchedule <: AbstractSchedule

An abstract type for all decay schedules.
Such schedules conform to a formula:
```
s(t) = λ * g(t)
```
where `s(t)` is the schedule output, `λ` is the base value, and `g(t)` is the decay function.
Concrete subtypes must implment [`basevalue`](#) and [`decay`](#).
"""
abstract type DecaySchedule <: AbstractSchedule end

# decay interface
"""
    basevalue(s::T) where T<:DecaySchedule

Concrete subtypes of [`DecaySchedule`](#) must implement this function for `T`.
Given `s`, the function should return the base value of `s`.
"""
function basevalue end
"""
    decay(s::T, t) where T<:DecaySchedule

Concrete subtypes of [`DecaySchedule`](#) must implement this function for `T`.
Given `s`, `decay(s, t)` should return the value of the decay function at iteration `t`.
"""
function decay end

Base.getindex(schedule::DecaySchedule, t::Integer) = basevalue(schedule) * decay(schedule, t)

Base.iterate(schedule::DecaySchedule, t = 1) = (schedule[t], t + 1)

"""
    CyclicSchedule <: AbstractSchedule

An abstract type for all cyclic schedules.
Such schedules conform to a formula:
```
s(t) = abs(λ0 - λ1) * g(t) + min(λ0, λ1)
```
where `s(t)` is the schedule output, `λ0` is the start value,
`λ1` is the end value, and `g(t)` is the cycle function.
Concrete subtypes must implment [`startvalue`](#), [`endvalue`](#) and [`cycle`](#).
"""
abstract type CyclicSchedule <: AbstractSchedule end

# cyclic interface
"""
    startvalue(s::T) where T<:CyclicSchedule

Concrete subtypes of [`CyclicSchedule`](#) must implement this function for `T`.
Given `s`, the function should return the start value of `s`.
"""
function startvalue end
"""
    endvalue(s::T) where T<:CyclicSchedule

Concrete subtypes of [`CyclicSchedule`](#) must implement this function for `T`.
Given `s`, the function should return the end value of `s`.
"""
function endvalue end
"""
    cycle(s::T, t) where T<:CyclicSchedule

Concrete subtypes of [`CyclicSchedule`](#) must implement this function for `T`.
Given `s`, `cycle(s, t)` should return the value of the cycle function at iteration `t`.
"""
function cycle end

function Base.getindex(schedule::CyclicSchedule, t::Integer)
    k0, k1 = startvalue(schedule), endvalue(schedule)

    return abs(k0 - k1) * cycle(schedule, t) + min(k0, k1)
end

Base.iterate(schedule::CyclicSchedule, t = 1) = (schedule[t], t + 1)

"""
    Lambda{T} <: AbstractSchedule
    Lambda(;f)

Wrap an arbitrary function `f` into a schedule.
The schedule output at iteration `t` is `f(t)`.
"""
struct Lambda{T} <: AbstractSchedule
    f::T
end
Lambda(;f) = Lambda(f)

Base.getindex(schedule::Lambda, t) = schedule.f(t)

Base.iterate(schedule::Lambda, t = 1) = (schedule[t], t + 1)
Base.IteratorEltype(::Type{<:Lambda}) = Base.EltypeUnknown()
Base.IteratorSize(::Type{<:Lambda}) = Base.SizeUnkown()

"""
    reverse(f, period)

Return a reverse function such that `reverse(f, period)(t) == f(period - t)`.
"""
reverse(f, period) = t -> f(period - t)
"""
    symmetric(f, period)

Return a symmetric function such that for `t ∈ [1, period / 2)`,
the symmetric function evaluates to `f(t)`, and when `t ∈ [period / 2, period)`,
the symmetric functions evaluates to `f(period - t)`.
"""
symmetric(f, period) = t -> (t < period / 2) ? f(t) : f(period - t)

at(x::Number, t) = x
at(x::AbstractSchedule, t) = x[t]

"""
    Sequence{T<:AbstractVector, S<:Integer} <: AbstractSchedule
    Sequence(;schedules, step_sizes)

A sequence of schedules.
The output of this schedule is the concatenation of `schedules` where each
schedule is evaluated for each step size in `step_sizes`.

Note that `schedules` can also be a vector of numbers (not just schedules).

# Arguments
- `schedules::AbstractVector`: a vector of schedules or numbers
- `step_sizes::Vector{<:Integer}`: a vector of iteration lengths for each schedule
"""
struct Sequence{T<:AbstractVector, S<:Integer} <: AbstractSchedule
    schedules::T
    step_sizes::Vector{S}
end
Sequence(;schedules, step_sizes) = Sequence(schedules, step_sizes)

function Base.getindex(schedule::Sequence, t::Integer)
    accum_steps = cumsum(schedule.step_sizes)
    i = findlast(x -> t > x, accum_steps)
    i = isnothing(i) ? 1 :
            (i >= length(schedule.schedules)) ? length(schedule.schedules) : i + 1
    toffset = (i > 1) ? t - accum_steps[i - 1] : t

    return at(schedule.schedules[i], toffset)
end
function Base.iterate(schedule::Sequence, state = (1, 1, 1))
    t, i, t0 = state
    if (i <= length(schedule.step_sizes)) && (t >= t0 + schedule.step_sizes[i])
        # move onto next step range
        i += 1
        t0 = t
    end

    return at(schedule.schedules[i], t - t0 + 1), (t + 1, i, t0)
end
Base.IteratorSize(::Type{<:Sequence}) = Base.SizeUnknown()

"""
    Loop{T<:AbstractSchedule, S<:Integer} <: AbstractSchedule
    Loop(;f, period)

Create a schedule that loops `f` every `period` iterations.
Note that `f` must be a subtype of [`AbstractSchedule`](#).
To loop arbitrary functions, wrap them in [`Lambda`](#).

# Arguments
- `f::AbstractSchedule`: the schedule to loop
- `period::Integer`: how often to loop
"""
struct Loop{T<:AbstractSchedule, S<:Integer} <: AbstractSchedule
    cycle_func::T
    period::S
end
Loop(;f, period) = Loop(f, period)

function Base.getindex(schedule::Loop, t)
    tcycle = mod(t, schedule.period)
    tcycle = (tcycle == 0) ? schedule.period : tcycle

    return schedule.cycle_func[tcycle]
end

Base.iterate(schedule::Loop, t = 1) = (schedule[t], t + 1)

Base.eltype(::Type{<:Loop{T}}) where T<:Union{<:DecaySchedule, <:CyclicSchedule} = eltype(T)
Base.IteratorSize(::Type{<:Loop}) = Base.IsInfinite()


text:
export BoundedStateSpace
export Identity, VectorView, OutputMatrix
export sample_space, state2workspace, volume, dim
export steering_control, propagate, collision_waypoints, waypoints, plot_waypoints
export in_state_space, is_free_state, is_free_motion, is_free_path, defaultNN, setup_steering, controltype

#====================== Implementation Notes ======================
TODO: the information below is outdated.

To specify a new state space (with associated distance function M,
a Metric or QuasiMetric) the user must specify:
1. evaluate(M, v, w)
2. waypoints(M, v, w)
3. time_waypoint(M, v, w, t)    (if time-based discretization is desired)
4. cost_waypoint(M, v, w, c)    (if cost-based discretization is desired)

Optionally, to improve performance, the user may specify: ## TODO: OUTDATED BELOW
1. defaultNN(M, init)
2. workspace_waypoints(v, w, M, s2w)
3. collision_waypoints(v, w, M, s2w)    (i.e. a sparser version of workspace_waypoints)

In rare circumstances, the user might need to specify:
1. setup_steering(SS, r)
2. controltype(M)

==================================================================#

### State Space Definitions
struct BoundedStateSpace{S<:State,M<:PreMetric,W<:State2Workspace} <: StateSpace{S}
    lo::S
    hi::S
    dist::M
    s2w::W
end
changeprecision{T<:AbstractFloat}(::Type{T}, SS::BoundedStateSpace) =
    BoundedStateSpace(map(x -> changeprecision(T,x), (SS.lo, SS.hi, SS.dist, SS.s2w))...)
changeprecision{T<:AbstractFloat}(::Type{T}, s2w::State2Workspace) = s2w

### Sampling
sample_space(SS::BoundedStateSpace) = (SS.lo + rand(typeof(SS.lo)).*(SS.hi-SS.lo))
volume(SS::BoundedStateSpace) = prod(SS.hi-SS.lo)    # TODO: only makes sense for a Euclidean metric
dim{S}(SS::BoundedStateSpace{S}) = length(S)         # TODO: this should really be a function of the metric dist

### State2Workspace (configuration space -> collision checking workspace)
struct Identity <: State2Workspace end
struct VectorView{N} <: State2Workspace
    inds::NTuple{N,Int}
end
VectorView(inds::Int...) = VectorView(inds)
VectorView(inds) = VectorView(Tuple(inds))
struct OutputMatrix{M,N,T<:AbstractFloat,L} <: State2Workspace
    C::SMatrix{M,N,T,L}
end
OutputMatrix(C::Matrix) = OutputMatrix(SMatrix{size(C,1),size(C,2)}(C))
changeprecision{T<:AbstractFloat}(::Type{T}, s2w::OutputMatrix) = OutputMatrix(changeprecision(T, s2w.C))

state2workspace(v::State, SS::StateSpace) = state2workspace(v, SS.s2w)
state2workspace(v::State, s2w::Identity) = v
state2workspace(v::AbstractVector, s2w::VectorView) = v[s2w.inds]
state2workspace(v::AbstractVector, s2w::OutputMatrix) = s2w.C*SVector(v)

workspace2state(w::AbstractVector, SS::StateSpace) = workspace2state(w, SS, SS.s2w)
workspace2state(w::AbstractVector, SS::StateSpace, s2w::Identity) = w
workspace2state(w::AbstractVector, SS::StateSpace, s2w::State2Workspace) = workspace2state(w, sample_space(SS), s2w)
function workspace2state(w::AbstractVector, v::AbstractVector, s2w::VectorView)
    mv = MVector(v)
    mv[s2w.inds] = w
    typeof(v)(mv)
end
workspace2state(w::AbstractVector, v::AbstractVector, s2w::OutputMatrix) = v + typeof(v)(Matrix(s2w.C) \ (w - s2w.C*v))

### Propagation and Waypoints
setup_steering(SS::StateSpace, r) = setup_steering(SS.dist, r)
setup_steering(d::PreMetric, r) = nothing
setup_steering(d::ChoppedPreMetric, r) = (d.chopval = r)
controltype(d::PreMetric) = NullControl  # TODO: not sure if I even use this anymore, but will be relevant for caching

## Generic propagation methods
function propagate{T}(d::Union{Metric, QuasiMetric}, v::State, u::StepControl{T}, s::AbstractFloat)
    s <= 0 ? v : s >= duration(u) ? propagate(d, v, u) : propagate(d, v, StepControl(T(s),u.u))
end
function propagate(d::Union{Metric, QuasiMetric}, v::State, us::ControlSequence)
    for u in us
        v = propagate(d, v, u)
    end
    v
end
function propagate(d::Union{Metric, QuasiMetric}, v::State, us::ControlSequence, s::AbstractFloat)
    s <= 0 && return v
    t = zero(s)
    for u in us
        if s >= t + duration(u)
            v = propagate(d, v, u)
            t += duration(u)
        else
            return propagate(d, v, u, s - t)
        end
    end
    v
end
function propagate{T<:AbstractFloat}(d::Union{Metric, QuasiMetric}, v::State, u::ControlInfo, s::AbstractVector{T})
    [propagate(d, v, u, t) for t in s]
end
function propagate{T<:AbstractFloat}(d::Union{Metric, QuasiMetric}, v::State, u::ControlSequence, s::AbstractVector{T})
    issorted(s) || error("Times should be sorted as input to propagate.")
    tf = duration(u)
    s = clamp(s, zero(tf), tf)
    path = typeof(v)[]
    t0, i = zero(tf), 1
    for t in s
        while t >= t0 + duration(u[i]) && i < length(u)    # second clause necessary because of numerical error
            v = propagate(d, v, u[i])
            t0 += duration(u[i])
            i += 1
        end
        push!(path, propagate(d, v, u[i], t-t0))
    end
    path
end

## Generic waypoints methods
function waypoints(d::Union{Metric, QuasiMetric}, v::State, w::State, n::Int)
    u = steering_control(d, v, w)
    t = duration(u)
    propagate(d, v, u, linspace(typeof(t)(0), t, n))
end
function waypoints(d::Union{Metric, QuasiMetric}, v::State, w::State, dt::AbstractFloat)
    u = steering_control(d, v, w)
    t = duration(u)
    propagate(d, v, u, [typeof(t)(0):dt:t; t])
end
function collision_waypoints(d::Union{Metric, QuasiMetric}, v::State, us::ControlSequence)
    path = typeof(v)[]
    for u in us
        append!(path, collision_waypoints(d, v, u))
        v = propagate(d, v, u)
    end
    path
end
collision_waypoints(d::Union{Metric, QuasiMetric}, v::State, w::State) = push!(collision_waypoints(d, v, steering_control(d,v,w)), w)

## StateSpace and CLBM
for f in (:propagate, :collision_waypoints, :waypoints, :steering_control)
    @eval $f(SS::StateSpace, args...) = $f(SS.dist, args...)
    @eval $f(CLBM::ChoppedPreMetric, args...) = $f(CLBM.m, args...)
end
steering_control(SS::StateSpace, V::State...) = vcat([steering_control(SS.dist,V[i],V[i+1]) for i in 1:length(V)-1]...)

### Validity Checking
in_state_space(v::AbstractVector, SS::StateSpace) = @all [SS.lo[i] <= v[i] <= SS.hi[i] for i in 1:length(v)]
is_free_state(v::State, CC::CollisionChecker, SS::StateSpace) =
    in_state_space(v, SS) && is_free_state(state2workspace(v, SS), CC)
function is_free_motion(v::State, w::State, CC::CollisionChecker, SS::StateSpace)
    wps = collision_waypoints(SS, v, w)
    @all [(in_state_space(wps[i], SS) &&
           is_free_motion(state2workspace(wps[i], SS.s2w), state2workspace(wps[i+1], SS.s2w), CC))
           for i in 1:length(wps)-1]
end
is_free_path(p::Path, CC::CollisionChecker, SS::StateSpace) =
    @all [is_free_motion(p[i], p[i+1], CC, SS) for i in 1:length(p)-1]

### Near Neighbor Setup (should be extended for best performance)
defaultNN(SS::StateSpace, init) = defaultNN(SS.dist, init)
defaultNN(d::Metric, init) = MetricNN(typeof(init)[init], d, init)
defaultNN(d::QuasiMetric, init) = QuasiMetricNN(typeof(init)[init], d, init)
function defaultNN{M}(d::ChoppedPreMetric{M}, init)
    M <: Metric && return MetricNN(typeof(init)[init], d, init)
    M <: QuasiMetric && return QuasiMetricNN(typeof(init)[init], d, init)
    error("Unsupported ChoppedPreMetric")
end
# pairwise_distances(d::PreMetric, V) = # TODO: implement this in a reasonable manner with ControlInfo

### Plotting
plot(SS::StateSpace) = plot_bounds(SS.lo, SS.hi)
plot_waypoints(SS::StateSpace, v::State, w::State) =
    map(x -> state2workspace(x, SS.s2w), isa(SS.dist, Euclidean) ? collision_waypoints(SS,v,w) : waypoints(SS,v,w,10))
function plot_path(p::Path, SS::StateSpace; kwargs...)
    plt.scatter(zip([state2workspace(x, SS)[1:2] for x in p]...)...; kwargs...)
    wps = hcat([statevec2mat(plot_waypoints(SS, p[i], p[i+1])) for i in 1:length(p)-1]...)
    plot_path(wps; kwargs...)
end
function plot_tree(V::Path, A::Vector{Int}, SS::StateSpace; kwargs...)    # V is a vector of states, i.e. a Path
    VW = [state2workspace(v, SS) for v in V]
    pts = statevec2mat(VW[find(A)])
    plt.scatter(pts[1,:], pts[2,:], zorder=1; kwargs...)
    X = vcat([[statevec2mat(plot_waypoints(SS, V[A[w]], V[w]))[1,:]; nothing] for w in find(A)]...)
    Y = vcat([[statevec2mat(plot_waypoints(SS, V[A[w]], V[w]))[2,:]; nothing] for w in find(A)]...)
    plt.plot(X, Y, linewidth=.5, linestyle="-", zorder=1; kwargs...)
end

include("statespaces/geometric.jl")
include("statespaces/linearquadratic.jl")
include("statespaces/simplecars.jl")


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

# Checked integer arithmetic

import Base: checked_abs, checked_neg, checked_add, checked_sub, checked_mul,
             checked_div, checked_rem, checked_fld, checked_mod, checked_cld,
             add_with_overflow, sub_with_overflow, mul_with_overflow

# checked operations

@testset for T in (Int8, Int16, Int32, Int64, Int128)
    # regular cases
    @test checked_abs(T(0)) === T(0)
    @test checked_neg(T(0)) === T(0)
    @test checked_add(T(0)) === T(0)
    @test checked_mul(T(0)) === T(0)

    for s in (-1, +1)
        @test checked_abs(T(3s)) === T(abs(3s))
        @test checked_neg(T(3s)) === T(-3s)
        @test checked_add(T(3s)) === T(3s)
        @test checked_mul(T(3s)) === T(3s)
        @test checked_abs(T(s*typemax(T))) === typemax(T)
        @test checked_neg(T(s*typemax(T))) === T(-s*typemax(T))
        @test checked_add(T(s*typemax(T))) === T(s*typemax(T))
        @test checked_mul(T(s*typemax(T))) === T(s*typemax(T))
    end

    # corner cases
    @test_throws OverflowError checked_abs(typemin(T))
    @test_throws OverflowError checked_neg(typemin(T))

    # regular cases
    for s1 in (-1, +1), s2 in (-1, +1)
        @test checked_add(T(4s1), T(3s2)) === T(4s1 + 3s2)
        @test checked_sub(T(4s1), T(3s2)) === T(4s1 - 3s2)
        @test checked_mul(T(4s1), T(3s2)) === T(4s1 * 3s2)
        @test checked_div(T(4s1), T(3s2)) === T(div(4s1, 3s2))
        @test checked_rem(T(4s1), T(3s2)) === T(rem(4s1, 3s2))
        @test checked_fld(T(4s1), T(3s2)) === T(fld(4s1, 3s2))
        @test checked_mod(T(4s1), T(3s2)) === T(mod(4s1, 3s2))
        @test checked_cld(T(4s1), T(3s2)) === T(cld(4s1, 3s2))
    end

    # corner cases
    halfmax = T(typemax(T)÷2)
    halfmax_plus1 = T(halfmax+1)
    halfmin = T(typemin(T)÷2)
    halfmin_minus1 = T(halfmin-1)
    sqrtmax = T(1) << T(sizeof(T)*4)
    half_sqrtmax = sqrtmax >> T(1)
    half_sqrtmax_plus1 = half_sqrtmax + T(1)

    @test checked_add(typemax(T), T(-1)) === T(typemax(T) - 1)
    @test_throws OverflowError checked_add(typemin(T), T(-1))
    @test checked_add(typemax(T), T(0)) === typemax(T)
    @test checked_add(typemin(T), T(0)) === typemin(T)
    @test_throws OverflowError checked_add(typemax(T), T(1))
    @test checked_add(typemin(T), T(1)) === T(typemin(T) + 1)
    @test checked_add(T(-1), typemax(T)) === T(typemax(T) - 1)
    @test_throws OverflowError checked_add(T(-1), typemin(T))
    @test checked_add(T(0), typemax(T)) === typemax(T)
    @test checked_add(T(0), typemin(T)) === typemin(T)
    @test_throws OverflowError checked_add(T(1), typemax(T))
    @test checked_add(T(1), typemin(T)) === T(typemin(T) + 1)
    @test checked_add(typemax(T), typemin(T)) === T(-1)
    @test checked_add(typemin(T), typemax(T)) === T(-1)
    @test_throws OverflowError checked_add(halfmax_plus1, halfmax_plus1)
    @test_throws OverflowError checked_add(halfmin, halfmin_minus1)

    @test add_with_overflow(typemax(T), T(-1)) === (T(typemax(T) - 1), false)
    @test add_with_overflow(typemin(T), T(-1)) === (T(typemax(T)), true)
    @test add_with_overflow(typemax(T), T(0))  === (typemax(T), false)
    @test add_with_overflow(typemin(T), T(0))  === (typemin(T), false)
    @test add_with_overflow(typemax(T), T(1))  === (typemin(T), true)
    @test add_with_overflow(typemin(T), T(1))  === (T(typemin(T) + 1), false)
    @test add_with_overflow(T(-1), typemax(T)) === (T(typemax(T) - 1), false)
    @test add_with_overflow(T(-1), typemin(T)) === (typemax(T), true)
    @test add_with_overflow(T(0), typemax(T))  === (typemax(T), false)
    @test add_with_overflow(T(0), typemin(T))  === (typemin(T), false)
    @test add_with_overflow(T(1), typemax(T))  === (typemin(T), true)
    @test add_with_overflow(T(1), typemin(T))  === (T(typemin(T) + 1), false)
    @test add_with_overflow(typemax(T), typemin(T)) === (T(-1), false)
    @test add_with_overflow(typemin(T), typemax(T)) === (T(-1), false)
    @test add_with_overflow(halfmax_plus1, halfmax_plus1) === (typemin(T), true)
    @test add_with_overflow(halfmin, halfmin_minus1) === (typemax(T), true)

    @test_throws OverflowError checked_sub(typemax(T), T(-1))
    @test checked_sub(typemax(T), T(0)) === typemax(T)
    @test checked_sub(typemax(T), T(1)) === T(typemax(T) - 1)
    @test checked_sub(typemin(T), T(-1)) === T(typemin(T) + 1)
    @test checked_sub(typemin(T), T(0)) === typemin(T)
    @test_throws OverflowError checked_sub(typemin(T), T(1))
    @test checked_sub(T(0), typemax(T)) === T(typemin(T) + 1)
    @test checked_sub(T(1), typemax(T)) === T(typemin(T) + 2)
    @test checked_sub(T(-1), typemin(T)) === typemax(T)
    @test_throws OverflowError checked_sub(T(0), typemin(T))
    @test checked_sub(typemax(T), typemax(T)) === T(0)
    @test checked_sub(typemin(T), typemin(T)) === T(0)
    @test checked_sub(halfmax, T(-halfmin)) === T(-1)
    @test_throws OverflowError checked_sub(halfmin, T(-halfmin_minus1))

    @test sub_with_overflow(typemax(T), T(-1)) === (typemin(T), true)
    @test sub_with_overflow(typemax(T), T(0))  === (typemax(T), false)
    @test sub_with_overflow(typemax(T), T(1))  === (T(typemax(T) - 1), false)
    @test sub_with_overflow(typemin(T), T(-1)) === (T(typemin(T) + 1), false)
    @test sub_with_overflow(typemin(T), T(0))  === (typemin(T), false)
    @test sub_with_overflow(typemin(T), T(1))  === (typemax(T), true)
    @test sub_with_overflow(T(0), typemax(T))  === (T(typemin(T) + 1), false)
    @test sub_with_overflow(T(1), typemax(T))  === (T(typemin(T) + 2), false)
    @test sub_with_overflow(T(-1), typemin(T)) === (typemax(T), false)
    @test sub_with_overflow(T(0), typemin(T))  === (typemin(T), true)
    @test sub_with_overflow(typemax(T), typemax(T)) === (T(0), false)
    @test sub_with_overflow(typemin(T), typemin(T)) === (T(0), false)
    @test sub_with_overflow(halfmax, T(-halfmin)) === (T(-1), false)
    @test sub_with_overflow(halfmin, T(-halfmin_minus1)) === (typemax(T), true)

    @test checked_mul(typemax(T), T(0)) === T(0)
    @test checked_mul(typemin(T), T(0)) === T(0)
    @test checked_mul(typemax(T), T(1)) === typemax(T)
    @test checked_mul(typemin(T), T(1)) === typemin(T)
    @test_throws OverflowError checked_mul(sqrtmax, half_sqrtmax)
    @test checked_mul(sqrtmax, -half_sqrtmax) === T(sqrtmax * -half_sqrtmax)
    @test_throws OverflowError checked_mul(sqrtmax, -half_sqrtmax_plus1)
    @test checked_mul(-sqrtmax, half_sqrtmax) === T(-sqrtmax * half_sqrtmax)
    @test_throws OverflowError checked_mul(-sqrtmax, half_sqrtmax_plus1)
    @test_throws OverflowError checked_mul(-sqrtmax, -half_sqrtmax)

    @test mul_with_overflow(typemax(T), T(0)) === (T(0), false)
    @test mul_with_overflow(typemin(T), T(0)) === (T(0), false)
    @test mul_with_overflow(typemax(T), T(1)) === (typemax(T), false)
    @test mul_with_overflow(typemin(T), T(1)) === (typemin(T), false)
    @test mul_with_overflow(sqrtmax, half_sqrtmax) === (typemin(T), true)
    @test mul_with_overflow(sqrtmax, -half_sqrtmax) === (T(sqrtmax * -half_sqrtmax), false)
    @test mul_with_overflow(sqrtmax, -half_sqrtmax_plus1) === (T(sqrtmax * -half_sqrtmax_plus1), true)
    @test mul_with_overflow(-sqrtmax, half_sqrtmax) === (T(-sqrtmax * half_sqrtmax), false)
    @test mul_with_overflow(-sqrtmax, half_sqrtmax_plus1) === (T(-sqrtmax * half_sqrtmax_plus1), true)
    @test mul_with_overflow(-sqrtmax, -half_sqrtmax) === (T(-sqrtmax * -half_sqrtmax), true)

    @test checked_div(typemax(T), T(1)) === typemax(T)
    @test_throws DivideError checked_div(typemax(T), T(0))
    @test checked_div(typemax(T), T(-1)) === T(-typemax(T))
    @test checked_div(typemin(T), T(1)) === typemin(T)
    @test_throws DivideError checked_div(typemin(T), T(0))
    @test_throws DivideError checked_div(typemin(T), T(-1))
    @test checked_rem(typemax(T), T(1)) === T(0)
    @test_throws DivideError checked_rem(typemax(T), T(0))
    @test checked_rem(typemax(T), T(-1)) === T(0)
    @test checked_rem(typemin(T), T(1)) === T(0)
    @test_throws DivideError checked_rem(typemin(T), T(0))
    @test checked_rem(typemin(T), T(-1)) === T(0)
    @test checked_fld(typemax(T), T(1)) === typemax(T)
    @test_throws DivideError checked_fld(typemax(T), T(0))
    @test checked_fld(typemax(T), T(-1)) === T(-typemax(T))
    @test checked_fld(typemin(T), T(1)) === typemin(T)
    @test_throws DivideError checked_fld(typemin(T), T(0))
    @test_throws DivideError checked_fld(typemin(T), T(-1))
    @test checked_mod(typemax(T), T(1)) === T(0)
    @test_throws DivideError checked_mod(typemax(T), T(0))
    @test checked_mod(typemax(T), T(-1)) === T(0)
    @test checked_mod(typemin(T), T(1)) === T(0)
    @test_throws DivideError checked_mod(typemin(T), T(0))
    @test checked_mod(typemin(T), T(-1)) === T(0)
    @test_throws DivideError checked_cld(typemax(T), T(0))
    @test checked_cld(typemax(T), T(-1)) === T(-typemax(T))
    @test checked_cld(typemin(T), T(1)) === typemin(T)
    @test_throws DivideError checked_cld(typemin(T), T(0))
    @test_throws DivideError checked_cld(typemin(T), T(-1))
end

@testset for T in (UInt8, UInt16, UInt32, UInt64, UInt128)
    # regular cases
    @test checked_abs(T(0)) === T(0)
    @test checked_neg(T(0)) === T(0)
    @test checked_add(T(0)) === T(0)
    @test checked_mul(T(0)) === T(0)

    @test checked_abs(T(3)) === T(3)
    @test_throws OverflowError checked_neg(T(3))
    @test checked_add(T(3)) === T(3)
    @test checked_mul(T(3)) === T(3)

    # regular cases
    @test checked_add(T(4), T(3)) === T(7)
    @test checked_sub(T(4), T(3)) === T(1)
    @test checked_mul(T(4), T(3)) === T(12)

    # corner cases
    halfmax = T(typemax(T)÷2)
    halfmax_plus1 = T(halfmax+1)
    sqrtmax = T(1) << T(sizeof(T)*4)

    @test checked_add(typemax(T), T(0)) === typemax(T)
    @test checked_add(T(0), T(0)) === T(0)
    @test_throws OverflowError checked_add(typemax(T), T(1))
    @test checked_add(T(0), T(1)) === T(T(0) + 1)
    @test checked_add(T(0), typemax(T)) === typemax(T)
    @test checked_add(T(0), T(0)) === T(0)
    @test_throws OverflowError checked_add(T(1), typemax(T))
    @test checked_add(T(1), T(0)) === T(T(0) + 1)
    @test checked_add(typemax(T), T(0)) === typemax(T)
    @test checked_add(T(0), typemax(T)) === typemax(T)
    @test_throws OverflowError checked_add(halfmax_plus1, halfmax_plus1)

    @test add_with_overflow(typemax(T), T(0)) === (typemax(T), false)
    @test add_with_overflow(T(0), T(0))       === (T(0), false)
    @test add_with_overflow(typemax(T), T(1)) === (T(0), true)
    @test add_with_overflow(T(0), T(1))       === (T(T(0) + 1), false)
    @test add_with_overflow(T(0), typemax(T)) === (typemax(T), false)
    @test add_with_overflow(T(0), T(0))       === (T(0), false)
    @test add_with_overflow(T(1), typemax(T)) === (T(0), true)
    @test add_with_overflow(T(1), T(0))       === (T(T(0) + 1), false)
    @test add_with_overflow(typemax(T), T(0)) === (typemax(T), false)
    @test add_with_overflow(T(0), typemax(T)) === (typemax(T), false)
    @test add_with_overflow(halfmax_plus1, halfmax_plus1) === (T(0), true)

    @test checked_sub(typemax(T), T(0)) === typemax(T)
    @test checked_sub(typemax(T), T(1)) === T(typemax(T) - 1)
    @test checked_sub(T(0), T(0)) === T(0)
    @test_throws OverflowError checked_sub(T(0), T(1))
    @test_throws OverflowError checked_sub(T(0), typemax(T))
    @test_throws OverflowError checked_sub(T(1), typemax(T))
    @test checked_sub(T(0), T(0)) === T(0)
    @test checked_sub(typemax(T), typemax(T)) === T(0)

    @test sub_with_overflow(typemax(T), T(0)) === (typemax(T), false)
    @test sub_with_overflow(typemax(T), T(1)) === (T(typemax(T) - 1), false)
    @test sub_with_overflow(T(0), T(0))       === (T(0), false)
    @test sub_with_overflow(T(0), T(1))       === (typemax(T), true)
    @test sub_with_overflow(T(0), typemax(T)) === (T(1), true)
    @test sub_with_overflow(T(1), typemax(T)) === (T(2), true)
    @test sub_with_overflow(T(0), T(0))       === (T(0), false)
    @test sub_with_overflow(typemax(T), typemax(T)) === (T(0), false)

    @test checked_mul(typemax(T), T(0)) === T(0)
    @test checked_mul(T(0), T(0)) === T(0)
    @test checked_mul(typemax(T), T(1)) === typemax(T)
    @test checked_mul(T(0), T(1)) === T(0)
    @test_throws OverflowError checked_mul(sqrtmax, sqrtmax)

    @test mul_with_overflow(typemax(T), T(0)) === (T(0), false)
    @test mul_with_overflow(T(0), T(0)) === (T(0), false)
    @test mul_with_overflow(typemax(T), T(1)) === (typemax(T), false)
    @test mul_with_overflow(T(0), T(1)) === (T(0), false)
    @test mul_with_overflow(sqrtmax, sqrtmax) === (T(0), true)

    @test checked_div(typemax(T), T(1)) === typemax(T)
    @test_throws DivideError checked_div(typemax(T), T(0))
    @test checked_rem(typemax(T), T(1)) === T(0)
    @test_throws DivideError checked_rem(typemax(T), T(0))
    @test checked_fld(typemax(T), T(1)) === typemax(T)
    @test_throws DivideError checked_fld(typemax(T), T(0))
    @test checked_mod(typemax(T), T(1)) === T(0)
    @test_throws DivideError checked_mod(typemax(T), T(0))
    @test checked_cld(typemax(T), T(1)) === typemax(T)
    @test_throws DivideError checked_cld(typemax(T), T(0))
end

@testset "Boolean" begin
    @test checked_add(false) === 0
    @test checked_add(true) === 1
    @test checked_neg(false) === 0
    @test checked_neg(true) === -1
    @test checked_abs(true) === true
    @test checked_abs(false) === false
    @test checked_mul(false) === false
    @test checked_mul(true) === true

    @test checked_add(true, true) === 2
    @test checked_add(true, false) === 1
    @test checked_add(false, false) === 0
    @test checked_add(false, true) === 1

    @test checked_sub(true, true) === 0
    @test checked_sub(true, false) === 1
    @test checked_sub(false, false) === 0
    @test checked_sub(false, true) === -1

    @test checked_mul(true, false) === false
    @test checked_mul(false, false) === false
    @test checked_mul(true, true) === true
    @test checked_mul(false, true) === false

    @test checked_div(true, true) === true
    @test checked_div(false, true) === false
    @test_throws DivideError checked_div(true, false)
    @test checked_rem(true, true) === false
    @test checked_rem(false, true) === false
    @test_throws DivideError checked_rem(true, false)
    @test checked_fld(true, true) === true
    @test checked_fld(false, true) === false
    @test_throws DivideError checked_fld(true, false)
    @test checked_mod(true, true) === false
    @test checked_mod(false, true) === false
    @test_throws DivideError checked_mod(true, false)
    @test checked_cld(true, true) === true
    @test checked_cld(false, true) === false
    @test_throws DivideError checked_cld(true, false)
end
@testset "BigInt" begin
    @test checked_abs(BigInt(-1)) == BigInt(1)
    @test checked_abs(BigInt(1)) == BigInt(1)
    @test checked_neg(BigInt(-1)) == BigInt(1)
    @test checked_neg(BigInt(1)) == BigInt(-1)
    @test checked_add(BigInt(1), BigInt(1)) == BigInt(2)
    @test checked_sub(BigInt(1), BigInt(2)) == BigInt(-1)
    @test checked_mul(BigInt(2), BigInt(10)) == BigInt(20)
    @test checked_div(BigInt(10), BigInt(2)) == BigInt(5)
    @test checked_rem(BigInt(9), BigInt(4)) == BigInt(1)
    @test checked_fld(BigInt(10), BigInt(3)) == BigInt(3)
    @test checked_mod(BigInt(9), BigInt(4)) == BigInt(1)
    @test checked_cld(BigInt(10), BigInt(3)) == BigInt(4)
end

@testset "Additional tests" begin
    # test promotions
    @test checked_add(UInt(4), UInt8(3)) === UInt(7)
    @test checked_sub(UInt(4), UInt8(3)) === UInt(1)
    @test checked_mul(UInt(4), UInt8(3)) === UInt(12)
    @test checked_div(UInt(4), UInt8(2)) === UInt(2)
    @test checked_mod(UInt(4), UInt8(2)) === UInt(0)
    @test checked_rem(UInt(5), UInt8(2)) === UInt(1)
    @test checked_fld(UInt(5), UInt8(1)) === UInt(5)
    @test checked_cld(UInt(5), UInt8(1)) === UInt(5)

    @test checked_sub(UInt(4), UInt(3)) === UInt(1)
    @test_throws OverflowError checked_sub(UInt(5), UInt(6))
    @test checked_mul(UInt(4), UInt(3)) === UInt(12)

    @test checked_sub(Int128(-1),Int128(-2)) === Int128(1)

    @test_throws OverflowError checked_mul(UInt32(2)^30, UInt32(2)^2)
    @test_throws OverflowError checked_mul(UInt64(2)^62, UInt64(2)^2)

    @test checked_add(UInt128(1), UInt128(2)) === UInt128(3)
    @test_throws OverflowError checked_add(UInt128(2)^127, UInt128(2)^127)

    @test checked_sub(UInt128(2), UInt128(1)) === UInt128(1)
    @test_throws OverflowError checked_sub(UInt128(3), UInt128(4))

    @test checked_mul(UInt128(3), UInt128(4)) === UInt128(12)
    @test_throws OverflowError checked_mul(UInt128(2)^127, UInt128(2))

end

@testset "Multiple arguments for add, mul" begin
    @test checked_add(1, 2, 3) === 6
    @test checked_add(1, 2, 3, 4) === 10
    @test checked_add(1, 2, 3, 4, 5) === 15
    @test checked_add(1, 2, 3, 4, 5, 6) === 21
    @test checked_add(1, 2, 3, 4, 5, 6, 7) === 28
    @test checked_add(1, 2, 3, 4, 5, 6, 7, 8) === 36

    @test checked_mul(1, 2, 3) === 6
    @test checked_mul(1, 2, 3, 4) === 24
    @test checked_mul(1, 2, 3, 4, 5) === 120
    @test checked_mul(1, 2, 3, 4, 5, 6) === 720
    @test checked_mul(1, 2, 3, 4, 5, 6, 7) === 5040
    @test checked_mul(1, 2, 3, 4, 5, 6, 7, 8) === 40320
end

max_stars_repo_name: vanillajonathan/julia

max_issues_repo_name: vanillajonathan/julia

max_forks_repo_name: vanillajonathan/julia


text:
module LinearMaps

export LinearMap
export ⊗, kronsum, ⊕
export FillMap

using LinearAlgebra
import LinearAlgebra: mul!
using SparseArrays

using Base: require_one_based_indexing

abstract type LinearMap{T} end

const MapOrVecOrMat{T} = Union{LinearMap{T}, AbstractVecOrMat{T}}
const MapOrMatrix{T} = Union{LinearMap{T}, AbstractMatrix{T}}
const RealOrComplex = Union{Real, Complex}

Base.eltype(::LinearMap{T}) where {T} = T

abstract type MulStyle end

struct FiveArg <: MulStyle end
struct ThreeArg <: MulStyle end

MulStyle(::FiveArg, ::FiveArg) = FiveArg()
MulStyle(::ThreeArg, ::FiveArg) = ThreeArg()
MulStyle(::FiveArg, ::ThreeArg) = ThreeArg()
MulStyle(::ThreeArg, ::ThreeArg) = ThreeArg()
MulStyle(::LinearMap) = ThreeArg() # default
MulStyle(::AbstractVecOrMat) = FiveArg()
MulStyle(A::LinearMap, As::LinearMap...) = MulStyle(MulStyle(A), MulStyle(As...))

Base.isreal(A::LinearMap) = eltype(A) <: Real
LinearAlgebra.issymmetric(::LinearMap) = false # default assumptions
LinearAlgebra.ishermitian(A::LinearMap{<:Real}) = issymmetric(A)
LinearAlgebra.ishermitian(::LinearMap) = false # default assumptions
LinearAlgebra.isposdef(::LinearMap) = false # default assumptions

Base.ndims(::LinearMap) = 2
Base.size(A::LinearMap, n) =
    (n == 1 || n == 2 ? size(A)[n] : error("LinearMap objects have only 2 dimensions"))
Base.axes(A::LinearMap, n::Integer) =
    (n == 1 || n == 2 ? axes(A)[n] : error("LinearMap objects have only 2 dimensions"))
Base.length(A::LinearMap) = prod(size(A))

# check dimension consistency for multiplication A*B
_iscompatible((A, B)) = size(A, 2) == size(B, 1)
function check_dim_mul(A, B)
    _iscompatible((A, B)) ||
        throw(DimensionMismatch("second dimension of left factor, $(size(A, 2)), " *
            "does not match first dimension of right factor, $(size(B, 1))"))
    return nothing
end
# check dimension consistency for multiplication C = A*B
function check_dim_mul(C, A, B)
    mA, nA = size(A) # A always has two dimensions
    mB, nB = size(B, 1), size(B, 2)
    (mB == nA && size(C, 1) == mA && size(C, 2) == nB) ||
        throw(DimensionMismatch("A has size ($mA,$nA), B has size ($mB,$nB), C has size $(size(C))"))
    return nothing
end

# conversion of AbstractVecOrMat to LinearMap
convert_to_lmaps_(A::AbstractVecOrMat) = LinearMap(A)
convert_to_lmaps_(A::LinearMap) = A
convert_to_lmaps() = ()
convert_to_lmaps(A) = (convert_to_lmaps_(A),)
@inline convert_to_lmaps(A, B, Cs...) =
    (convert_to_lmaps_(A), convert_to_lmaps_(B), convert_to_lmaps(Cs...)...)

# The (internal) multiplication logic is as follows:
#  - `*(A, x)` calls `mul!(y, A, x)` for appropriately-sized y
#  - `mul!` checks consistency of the sizes, and calls `_unsafe_mul!`,
#    which does not check sizes, but potentially one-based indexing if necessary
#  - by default, `_unsafe_mul!` is redirected back to `mul!`
#  - custom map types only need to implement 3-arg (vector) `mul!`, and
#    everything else (5-arg multiplication, application to matrices,
#    conversion to matrices) will just work

"""
    *(A::LinearMap, x::AbstractVector)::AbstractVector

Compute the action of the linear map `A` on the vector `x`.

!!! compat "Julia 1.3"
    In Julia versions v1.3 and above, objects `L` of any subtype of `LinearMap`
    are callable in the sense that `L(x) = L*x` for `x::AbstractVector`.

## Examples
```jldoctest; setup=(using LinearAlgebra, LinearMaps)
julia> A=LinearMap([1.0 2.0; 3.0 4.0]); x=[1.0, 1.0];

julia> A*x
2-element Array{Float64,1}:
 3.0
 7.0

julia> A(x)
2-element Array{Float64,1}:
 3.0
 7.0
```
"""
function Base.:(*)(A::LinearMap, x::AbstractVector)
    check_dim_mul(A, x)
    T = promote_type(eltype(A), eltype(x))
    y = similar(x, T, axes(A)[1])
    return mul!(y, A, x)
end

(L::LinearMap)(x::AbstractVector) = L*x

"""
    mul!(Y::AbstractVecOrMat, A::LinearMap, B::AbstractVector) -> Y
    mul!(Y::AbstractMatrix, A::LinearMap, B::AbstractMatrix) -> Y

Calculates the action of the linear map `A` on the vector or matrix `B` and stores the
result in `Y`, overwriting the existing value of `Y`. Note that `Y` must not be aliased
with either `A` or `B`.

## Examples
```jldoctest; setup=(using LinearAlgebra, LinearMaps)
julia> A=LinearMap([1.0 2.0; 3.0 4.0]); B=[1.0, 1.0]; Y = similar(B); mul!(Y, A, B);

julia> Y
2-element Array{Float64,1}:
 3.0
 7.0

julia> A=LinearMap([1.0 2.0; 3.0 4.0]); B=[1.0 1.0; 1.0 1.0]; Y = similar(B); mul!(Y, A, B);

julia> Y
2×2 Array{Float64,2}:
 3.0  3.0
 7.0  7.0
```
"""
function mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector)
    check_dim_mul(y, A, x)
    return _unsafe_mul!(y, A, x)
end

"""
    mul!(C::AbstractVecOrMat, A::LinearMap, B::AbstractVector, α, β) -> C
    mul!(C::AbstractMatrix, A::LinearMap, B::AbstractMatrix, α, β) -> C

Combined inplace multiply-add ``A B α + C β``. The result is stored in `C` by overwriting
it. Note that `C` must not be aliased with either `A` or `B`.

## Examples
```jldoctest; setup=(using LinearAlgebra, LinearMaps)
julia> A=LinearMap([1.0 2.0; 3.0 4.0]); B=[1.0, 1.0]; C=[1.0, 3.0];

julia> mul!(C, A, B, 100.0, 10.0) === C
true

julia> C
2-element Array{Float64,1}:
 310.0
 730.0

julia> A=LinearMap([1.0 2.0; 3.0 4.0]); B=[1.0 1.0; 1.0 1.0]; C=[1.0 2.0; 3.0 4.0];

julia> mul!(C, A, B, 100.0, 10.0) === C
true

julia> C
2×2 Array{Float64,2}:
 310.0  320.0
 730.0  740.0
```
"""
function mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector, α::Number, β::Number)
    check_dim_mul(y, A, x)
    return _unsafe_mul!(y, A, x, α, β)
end

function _generic_mapvec_mul!(y, A, x, α, β)
    # this function needs to call mul! for, e.g.,  AdjointMap{...,<:CustomMap}
    if isone(α)
        iszero(β) && return mul!(y, A, x)
        z = A * x
        if isone(β)
            y .+= z
        else
            y .= y.*β .+ z
        end
        return y
    elseif iszero(α)
        iszero(β) && return fill!(y, zero(eltype(y)))
        isone(β) && return y
        # β != 0, 1
        return rmul!(y, β)
    else # α != 0, 1
        iszero(β) && return rmul!(mul!(y, A, x), α)
        z = A * x
        if isone(β)
            y .+= z .* α
        else
            y .= y .* β .+ z .* α
        end
        return y
    end
end

# the following is of interest in, e.g., subspace-iteration methods
function mul!(Y::AbstractMatrix, A::LinearMap, X::AbstractMatrix)
    check_dim_mul(Y, A, X)
    return _unsafe_mul!(Y, A, X)
end
function mul!(Y::AbstractMatrix, A::LinearMap, X::AbstractMatrix, α::Number, β::Number)
    check_dim_mul(Y, A, X)
    return _unsafe_mul!(Y, A, X, α, β)
end

function _generic_mapmat_mul!(Y, A, X, α=true, β=false)
    for (Xi, Yi) in zip(eachcol(X), eachcol(Y))
        mul!(Yi, A, Xi, α, β)
    end
    return Y
end

_unsafe_mul!(y, A::MapOrVecOrMat, x) = mul!(y, A, x)
_unsafe_mul!(y, A::AbstractVecOrMat, x, α, β) = mul!(y, A, x, α, β)
function _unsafe_mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector, α, β)
    return _generic_mapvec_mul!(y, A, x, α, β)
end
function _unsafe_mul!(y::AbstractMatrix, A::LinearMap, x::AbstractMatrix)
    return _generic_mapmat_mul!(y, A, x)
end
function _unsafe_mul!(y::AbstractMatrix, A::LinearMap, x::AbstractMatrix, α, β)
    return _generic_mapmat_mul!(y, A, x, α, β)
end

const LinearMapTuple = Tuple{Vararg{LinearMap}}

include("left.jl") # left multiplication by a transpose or adjoint vector
include("transpose.jl") # transposing linear maps
include("wrappedmap.jl") # wrap a matrix of linear map in a new type, thereby allowing to alter its properties
include("uniformscalingmap.jl") # the uniform scaling map, to be able to make linear combinations of LinearMap objects and multiples of I
include("linearcombination.jl") # defining linear combinations of linear maps
include("scaledmap.jl") # multiply by a (real or complex) scalar
include("composition.jl") # composition of linear maps
include("functionmap.jl") # using a function as linear map
include("blockmap.jl") # block linear maps
include("kronecker.jl") # Kronecker product of linear maps
include("fillmap.jl") # linear maps representing constantly filled matrices
include("conversion.jl") # conversion of linear maps to matrices
include("show.jl") # show methods for LinearMap objects

"""
    LinearMap(A::LinearMap; kwargs...)::WrappedMap
    LinearMap(A::AbstractVecOrMat; kwargs...)::WrappedMap
    LinearMap(J::UniformScaling, M::Int)::UniformScalingMap
    LinearMap{T=Float64}(f, [fc,], M::Int, N::Int = M; kwargs...)::FunctionMap

Construct a linear map object, either from an existing `LinearMap` or `AbstractVecOrMat` `A`,
with the purpose of redefining its properties via the keyword arguments `kwargs`;
a `UniformScaling` object `J` with specified (square) dimension `M`; or
from a function or callable object `f`. In the latter case, one also needs to specify
the size of the equivalent matrix representation `(M, N)`, i.e., for functions `f` acting
on length `N` vectors and producing length `M` vectors (with default value `N=M`).
Preferably, also the `eltype` `T` of the corresponding matrix representation needs to be
specified, i.e., whether the action of `f` on a vector will be similar to, e.g., multiplying
by numbers of type `T`. If not specified, the devault value `T=Float64` will be assumed.
Optionally, a corresponding function `fc` can be specified that implements the adjoint
(=transpose in the real case) of `f`.

The keyword arguments and their default values for the function-based constructor are:
*   `issymmetric::Bool = false` : whether `A` or `f` act as a symmetric matrix
*   `ishermitian::Bool = issymmetric & T<:Real` : whether `A` or `f` act as a Hermitian
    matrix
*   `isposdef::Bool = false` : whether `A` or `f` act as a positive definite matrix.
For existing linear maps or matrices `A`, the default values will be taken by calling
internal functions `_issymmetric`, `_ishermitian` and `_isposdef` on the existing object `A`.
These in turn dispatch to (overloads of) `LinearAlgebra`'s `issymmetric`, `ishermitian`,
and `isposdef` methods whenever these checks are expected to be computationally cheap or even
known at compile time as for certain structured matrices, but return `false` for generic
`AbstractMatrix` types.

For the function-based constructor, there is one more keyword argument:
*   `ismutating::Bool` : flags whether the function acts as a mutating matrix multiplication
    `f(y,x)` where the result vector `y` is the first argument (in case of `true`),
    or as a normal matrix multiplication that is called as `y=f(x)` (in case of `false`).
    The default value is guessed by looking at the number of arguments of the first
    occurrence of `f` in the method table.
"""
LinearMap(A::MapOrVecOrMat; kwargs...) = WrappedMap(A; kwargs...)
LinearMap(J::UniformScaling, M::Int) = UniformScalingMap(J.λ, M)
LinearMap(f, M::Int; kwargs...) = LinearMap{Float64}(f, M; kwargs...)
LinearMap(f, M::Int, N::Int; kwargs...) = LinearMap{Float64}(f, M, N; kwargs...)
LinearMap(f, fc, M::Int; kwargs...) = LinearMap{Float64}(f, fc, M; kwargs...)
LinearMap(f, fc, M::Int, N::Int; kwargs...) = LinearMap{Float64}(f, fc, M, N; kwargs...)

LinearMap{T}(A::MapOrMatrix; kwargs...) where {T} = WrappedMap{T}(A; kwargs...)
LinearMap{T}(f, args...; kwargs...) where {T} = FunctionMap{T}(f, args...; kwargs...)

end # module


text:
import Optim

# Transformation from C to R
split_(x)  = x[1:Int(length(x)/2)],x[Int(length(x)/2) + 1:end]

"BFGS using MoreThuente"
function optim( z, alpha, c, lb=nothing; linesearch = Optim.LineSearches.MoreThuente() )
    function grad!(stor,x)
        stor[:] = vcat(Fgradient(split_(x)..., c)...)
    end

    result = Optim.optimize( x->residual(split_(x)...,c),
                             grad!,
                             vcat(z, alpha),
                             Optim.BFGS(linesearch = linesearch), Optim.Options(
                                iterations = 14000,
                                allow_f_increases = true))

    return split_(Optim.minimizer(result))
end


text:
"""
    abstract type LearningTask

Represents a concrete approach for solving a learning task.

A `LearningTask` defines how data is processed encoded and decoded
before and after going through a model.

## Extending

It is recommended to use [`AbstractBlockTask`](#)s like [`BlockTask`](#)
and [`SupervisedTask`](#) to construct tasks, but you may subtype
`LearningTask` for lower-level control.

There is a core interface that will allow you to train models and
perform inference (for supervised tasks). It consists of

- [`encodesample`](#)
- [`encodeinput`](#)
- [`decodeypred`](#)

You can optionally implement additional interfaces to get support for
higher-level features of the library.

- Training interface: [`tasklossfn`](#), [`taskmodel`](#)
- Testing interface: [`mocksample`](#), [`mockinput`](#), [`mocktarget`](#),
    [`mockmodel`](#)
- Batching: [`shouldbatch`](#)

"""
abstract type LearningTask end

# ## Encoding contexts

"""
    abstract type Context

Represents a context in which a data transformation
is made. This allows using dispatching for varying behavior,
for example, to apply augmentations only during training or
use non-destructive cropping during inference.

See [`Training`](#), [`Validation`](#) and
[`Inference`](#).
"""
abstract type Context end

"""
    Training <: Context

A context for applying data transformations during training. [`Encoding`](#)s and
[`LearningTask`](#)s can dispatch on this when certain transformations,
like random augmentations, should only be applied during training.
"""
struct Training <: Context end

"""
    Validation <: Context

A context for applying data transformations during validation/testing.
[`Encoding`](#)s and [`LearningTask`](#)s can dispatch on this when
certain transformations, like random augmentations, should not be applied
during validation, only in training.
"""
struct Validation <: Context end

struct Inference <: Context end

# ## Encoding interface

function encodesample end

function encodeinput end

function decodeypred end
const decodeŷ = decodeypred

decodey(args...; kwargs...) = decodeypred(args...; kwargs...)

# ## Buffered encoding interface
#
# If not overwritten, applies the non-buffering method.

encodesample!(buf, task, ctx, sample) = encodesample(task, ctx, sample)
encodeinput!(buf, task, ctx, sample) = encodeinput(task, ctx, sample)
decodeypred!(buf, task, ctx, ypred) = decodeypred(task, ctx, ypred)
const decodeŷ! = decodeypred!
decodey!(args...; kwargs...) = decodeypred!(args...; kwargs...)

# ## Training interface

"""
    tasklossfn(task)

Default loss function to use when training models for `task`.
"""
function tasklossfn end

"""
    taskmodel(task, backbone)

Construct a model for `task` from a backbone architecture, for example
by attaching a task-specific head model.
"""
function taskmodel end

# ## Testing interface

"""
    mocksample(task)

Generate a random `sample` compatible with `task`.
"""
mocksample(task) = (mockinput(task), mocktarget(task))

"""
    mockinput(task)

Generate a random `input` compatible with `task`.
"""
function mockinput end

"""
    mocktarget(task)

Generate a random `target` compatible with `task`.
"""
function mocktarget end

"""
    mockmodel(task)

Generate a `model` compatible with `task` for testing.
"""
function mockmodel end

# ## Batching interface

"""
    shouldbatch(::LearningTask)

Define whether encoded samples for a learning task should be
batched. The default is `true`.
"""
shouldbatch(::LearningTask) = true


text:
## estimate MAP by first estimating MAP for each chain
MAP(x::T, data) where T <: Fit = MAP([MAP(ch, data) for ch in x.chains], data)

MAP(x::Chain, data) = x.theta[last(findmax(sum(x.loglik_p, dims=2)[:,1]))]

function MAP(theta::Vector{T}, data) where T <: AbstractParam
    M = length(theta)
    ll = -Inf
    tmax = 1

    for t in 1:M
        new_ll = loglik(theta[t], data)

        if new_ll > ll
            ll = new_ll
            tmax = t
        end
    end

    theta[tmax]
end

weibull_loghaz(time, shape, scale) = log(scale) .+ log(shape) .+ (shape-1) .* log.(time)
weibull_cumhaz(time, shape, scale) = scale .* time.^shape
##inv_weibull_cumhaz(p, shape, scale) = (p ./ scale).^(1/shape)

function loglik_p_surv(T, E, Xe, shape::Tx, scale::Tx, delta::Vector{Tx}) where Tx <: Real
    if (shape .<= 0.0) | (scale .<= 0.0)
        -Inf
    else
        Xed = Xe*delta
        sum(E.*(weibull_loghaz(T .- 1, shape, scale) .+ Xed) .- weibull_cumhaz(T .- 1, shape, scale).*exp.(Xed))
    end
end

function loglik_p_surv(T, E, Xe, haz_basis, cumhaz_basis, w, delta)
    Xed = Xe*delta
    ew = exp.(w)
    sum(E.*(log.(haz_basis*ew) .+ Xed) .- cumhaz_basis*ew.*exp.(Xed))
end

##invquad(L, x) = dot(x, L' \ (L \ x))
##mvn_logpdf(x, mu, L) = -length(mu)/2 * log(2*pi) - log(det(L)) - invquad(L, x-mu)/2

function loglik_p_long(b, beta, mu, Sigma, sigma, Y, times, n, Xl, Z)
    target = zero(eltype(mu))
    d = Normal(0, sigma)

    for j in 1:n
        target += logpdf(d, Y[j] - dot(Z[j,:], b) - dot(Xl[j,:], beta))
        ##target += logpdf(Normal(zero(eltype(sigma)), sigma), Y[j] - dot(Z[j,:], u) - dot(Xl[j,:], beta))
        ## loglikelihood(Normal(0, sigma), Y[1:n] .- Z[1:n,:]*u - Xl[1:n,:]*beta)
    end

    target + logpdf(MvNormal(mu, Sigma), b) ##mvn_logpdf(b, zero(b), Sigma)
end

function loglik_p_long(beta, mu, Sigma, sigma, Y, times, n, Xl, Z)
    m = Z*mu + Xl*beta
    C = Array(Symmetric(Z*Sigma*Z' + sigma^2*diagm(0 => ones(Float64, n))))
    issuccess(cholesky(C; check=false)) || return -Inf
    logpdf(MvNormal(m, C), Y)
end

function loglik_p(theta::Param, data)
    @unpack T, E, Xe, Xl, Xp, Z, Y, times, n = data
    @unpack xi, cind, delta, bl_shape, bl_scale, b, beta, mu, S, L, sigma = theta

    cp = class_prob(xi, Xp)

    N = length(T)
    G = size(beta)[1]

    loglik = zeros(Float64, N)

    Sigma = [build_cov_matrix(S[g,:], L[g,:,:]) for g in 1:G]

    for i in 1:N
        g = cind[i]
        lll_vec = loglik_p_long(b[i,:], beta[g,:], mu[g,:], Sigma[g], sigma, Y[i,1:n[i]], times[i,1:n[i]], n[i], Xl[i,1:n[i],:], Z[i,1:n[i],:])
        lle_vec = loglik_p_surv(T[i], E[i], Xe[i,:]', bl_shape[g], bl_scale[g], delta[g,:])
        loglik[i] = lll_vec + lle_vec + log(cp[i,g])
    end

    loglik
end

function loglik_p(theta::ParamSpline, data)
    @unpack T, E, Xe, Xl, Xp, Z, Y, times, n, haz_basis, cumhaz_basis = data
    @unpack xi, cind, delta, w, b, beta, mu, S, L, sigma = theta

    cp = class_prob(xi, Xp)

    N = length(T)
    G = size(beta)[1]

    ll = zeros(Float64, N)

    Sigma = [build_cov_matrix(S[g,:], L[g,:,:]) for g in 1:G]

    for i in 1:N
        g = cind[i]
        lll_vec = loglik_p_long(b[i,:], beta[g,:], mu[g,:], Sigma[g], sigma, Y[i,1:n[i]], times[i,1:n[i]], n[i], Xl[i,1:n[i],:], Z[i,1:n[i],:])
        lle_vec = loglik_p_surv(T[i], E[i], Xe[i,:]', haz_basis[i,:]', cumhaz_basis[i,:]', w[g,:], delta[g,:])
        ll[i] = lll_vec + lle_vec + log(cp[i,g])
    end
    ll
end

loglik(theta, data) = sum(loglik_p(theta, data))

loglik_array(x::T, data::Data) where T <: AbstractFit = loglik_array(x.chains, data)

function loglik_array(x::AbstractVector{T}, data::Data) where T <: Chain
    M, K, N = length(x[1].theta), length(x), length(data.T)

    ##lp = SharedArray{Float64,3}((M, K, N))
    lp = Array{Float64,3}(undef, M, K, N) ##SharedArray{Float64,3}((M, K, N))

    ##@sync @distributed
    for m in 1:M
        for k in 1:K
            lp[m,k,:] = loglik_p(x[k].theta[m], data)
        end
    end
    lp
end