truncation.jl

const docs_truncation_kwargs = """
* `atol::Real`       : Absolute tolerance for the truncation
* `rtol::Real`       : Relative tolerance for the truncation
* `maxrank::Integer` : Maximal rank for the truncation
* `minrank::Integer` : Minimal rank for the truncation
* `maxerror::Real`   : Maximal truncation error.
* `filter`           : Custom filter to select truncated values.
"""

const docs_truncation_strategies = """
- [`notrunc`](@ref)
- [`truncrank`](@ref)
- [`trunctol`](@ref)
- [`truncerror`](@ref)
- [`truncfilter`](@ref)
"""

const docs_null_truncation_kwargs = """
* `atol::Real`       : Absolute tolerance for the truncation
* `rtol::Real`       : Relative tolerance for the truncation
* `maxnullity::Real` : Maximal rank for the truncation
"""

"""
    TruncationStrategy(; kwargs...)

Select a truncation strategy based on the provided keyword arguments.

## Keyword arguments
The following keyword arguments are all optional, and their default value (`nothing`)
will be ignored. It is also allowed to combine multiple of these, in which case the kept
values will consist of the intersection of the different truncated strategies (except
`minrank`, which uses union semantics to guarantee a lower bound on the number of kept values).

$docs_truncation_kwargs
"""
function TruncationStrategy(;
        atol::Union{Real, Nothing} = nothing,
        rtol::Union{Real, Nothing} = nothing,
        maxrank::Union{Integer, Nothing} = nothing,
        minrank::Union{Integer, Nothing} = nothing,
        maxerror::Union{Real, Nothing} = nothing,
        filter = nothing,
    )
    strategy = notrunc()

    if !isnothing(atol) || !isnothing(rtol)
        atol = @something atol 0
        rtol = @something rtol 0
        strategy &= trunctol(; atol, rtol)
    end

    isnothing(maxrank) || (strategy &= truncrank(maxrank))
    isnothing(maxerror) || (strategy &= truncerror(; atol = maxerror))
    isnothing(filter) || (strategy &= truncfilter(filter))

    # union constraint: guarantee a lower bound on number of kept values
    # special-case NoTruncation: keeping everything already satisfies any minrank
    if !isnothing(minrank) && !(strategy isa NoTruncation)
        strategy |= truncrank(minrank)
    elseif !isnothing(minrank)
        strategy = truncrank(minrank)
    end

    return strategy
end

"""
    null_truncation_strategy(; kwargs...)

Select a nullspace truncation strategy based on the provided keyword arguments.

## Keyword arguments
The following keyword arguments are all optional, and their default value (`nothing`)
will be ignored. It is also allowed to combine multiple of these, in which case the
discarded values will consist of the intersection of the different truncated strategies.

$docs_null_truncation_kwargs
"""
function null_truncation_strategy(; atol = nothing, rtol = nothing, maxnullity = nothing)
    if isnothing(maxnullity) && isnothing(atol) && isnothing(rtol)
        return notrunc()
    end
    atol = @something atol 0
    rtol = @something rtol 0
    trunc = trunctol(; atol, rtol, keep_below = true)
    return !isnothing(maxnullity) ? trunc & truncrank(maxnullity; rev = false) : trunc
end

"""
    NoTruncation()

Trivial truncation strategy that keeps all values, mostly for testing purposes.
See also [`notrunc()`](@ref).
"""
struct NoTruncation <: TruncationStrategy end

"""
    notrunc()

Truncation strategy that does nothing, and keeps all the values.
"""
notrunc() = NoTruncation()

# TODO: Base.Ordering?
"""
    TruncationByOrder(howmany::Int, by::Function, rev::Bool)

Truncation strategy to keep the first `howmany` values when sorted according to `by` in increasing (decreasing) order if `rev` is false (true).

See also [`truncrank`](@ref).
"""
struct TruncationByOrder{F} <: TruncationStrategy
    howmany::Int
    by::F
    rev::Bool
end

"""
    truncrank(howmany::Integer; by=abs, rev::Bool=true)

Truncation strategy to keep the first `howmany` values when sorted according to `by` or the last `howmany` if `rev` is true.
"""
function truncrank(howmany::Integer; by = abs, rev::Bool = true)
    return TruncationByOrder(howmany, by, rev)
end

"""
    TruncationByFilter(filter::Function)

Truncation strategy to keep the values for which `filter` returns true.

See also [`truncfilter`](@ref).
"""
struct TruncationByFilter{F} <: TruncationStrategy
    filter::F
end

"""
    truncfilter(filter)

Truncation strategy to keep the values for which `filter` returns true.
"""
truncfilter(f) = TruncationByFilter(f)

"""
    TruncationByValue(atol::Real, rtol::Real, p::Real, by, keep_below::Bool=false)

Truncation strategy to keep the values that satisfy `by(val) > max(atol, rtol * norm(values, p))`.
If `keep_below = true`, discard these values instead.
See also [`trunctol`](@ref)
"""
struct TruncationByValue{T <: Real, P <: Real, F} <: TruncationStrategy
    atol::T
    rtol::T
    p::P
    by::F
    keep_below::Bool
end
function TruncationByValue(atol::Real, rtol::Real, p::Real = 2, by = abs, keep_below::Bool = true)
    return TruncationByValue(promote(atol, rtol)..., p, by, keep_below)
end

"""
    trunctol(; atol::Real=0, rtol::Real=0, p::Real=2, by=abs, keep_below::Bool=false)

Truncation strategy to keep the values that satisfy `by(val) > max(atol, rtol * norm(values, p))`.
If `keep_below = true`, discard these values instead.
"""
function trunctol(; atol::Real = 0, rtol::Real = 0, p::Real = 2, by = abs, keep_below::Bool = false)
    return TruncationByValue(atol, rtol, p, by, keep_below)
end

"""
    TruncationByError(; atol::Real, rtol::Real, p::Real)

Truncation strategy to discard values until the error caused by the discarded values exceeds some tolerances.
See also [`truncerror`](@ref).
"""
struct TruncationByError{T <: Real, P <: Real} <: TruncationStrategy
    atol::T
    rtol::T
    p::P
end
function TruncationError(atol::Real, rtol::Real, p::Real = 2)
    return TruncationError(promote(atol, rtol)..., p)
end

"""
    truncerror(; atol::Real=0, rtol::Real=0, p::Real=2)

Truncation strategy for truncating values such that the error in the factorization
is smaller than `max(atol, rtol * norm)`, where the error is determined using the `p`-norm.
"""
function truncerror(; atol::Real = 0, rtol::Real = 0, p::Real = 2)
    return TruncationByError(promote(atol, rtol)..., p)
end

"""
    TruncationIntersection(trunc::TruncationStrategy, truncs::TruncationStrategy...)

Truncation strategy that composes multiple truncation strategies, keeping values that are
common between them.
"""
struct TruncationIntersection{T <: Tuple{Vararg{TruncationStrategy}}} <: TruncationStrategy
    components::T
end
function TruncationIntersection(trunc::TruncationStrategy, truncs::TruncationStrategy...)
    return TruncationIntersection((trunc, truncs...))
end

function Base.:&(trunc1::TruncationStrategy, trunc2::TruncationStrategy)
    return TruncationIntersection((trunc1, trunc2))
end

# flatten components
function Base.:&(trunc1::TruncationIntersection, trunc2::TruncationIntersection)
    return TruncationIntersection((trunc1.components..., trunc2.components...))
end
function Base.:&(trunc1::TruncationIntersection, trunc2::TruncationStrategy)
    return TruncationIntersection((trunc1.components..., trunc2))
end
function Base.:&(trunc1::TruncationStrategy, trunc2::TruncationIntersection)
    return TruncationIntersection((trunc1, trunc2.components...))
end

# drop notrunc
Base.:&(::NoTruncation, trunc::TruncationStrategy) = trunc
Base.:&(trunc::TruncationStrategy, ::NoTruncation) = trunc
Base.:&(::NoTruncation, ::NoTruncation) = notrunc()

# disambiguate
Base.:&(::NoTruncation, trunc::TruncationIntersection) = trunc
Base.:&(trunc::TruncationIntersection, ::NoTruncation) = trunc

"""
    TruncationUnion(trunc::TruncationStrategy, truncs::TruncationStrategy...)

Truncation strategy that composes multiple truncation strategies, keeping values that are
present in any of them.
"""
struct TruncationUnion{T <: Tuple{Vararg{TruncationStrategy}}} <: TruncationStrategy
    components::T
end
function TruncationUnion(trunc::TruncationStrategy, truncs::TruncationStrategy...)
    return TruncationUnion((trunc, truncs...))
end

function Base.:|(trunc1::TruncationStrategy, trunc2::TruncationStrategy)
    return TruncationUnion((trunc1, trunc2))
end

# flatten components
function Base.:|(trunc1::TruncationUnion, trunc2::TruncationUnion)
    return TruncationUnion((trunc1.components..., trunc2.components...))
end
function Base.:|(trunc1::TruncationUnion, trunc2::TruncationStrategy)
    return TruncationUnion((trunc1.components..., trunc2))
end
function Base.:|(trunc1::TruncationStrategy, trunc2::TruncationUnion)
    return TruncationUnion((trunc1, trunc2.components...))
end

# NoTruncation is the absorbing element for | (union with "keep all" = keep all)
Base.:|(::NoTruncation, ::TruncationStrategy) = notrunc()
Base.:|(::TruncationStrategy, ::NoTruncation) = notrunc()
Base.:|(::NoTruncation, ::NoTruncation) = notrunc()

# disambiguate
Base.:|(::NoTruncation, ::TruncationUnion) = notrunc()
Base.:|(::TruncationUnion, ::NoTruncation) = notrunc()

@doc """
    truncation_error(values, ind)
Compute the truncation error as the 2-norm of the values that are not kept by `ind`.
""" truncation_error, truncation_error!