# Tensor factorization
#----------------------
const OFA = OrthogonalFactorizationAlgorithm

import LinearAlgebra: svd!, svd
const SVDAlg = Union{SVD, SDD}

Base.@deprecate(
    svd(t::AbstractTensorMap, leftind::IndexTuple, rightind::IndexTuple;
            trunc::TruncationScheme = notrunc(), p::Real = 2, alg::SVDAlg = SDD()),
    tsvd(t, leftind, rightind; trunc = trunc, p = p, alg = alg))
Base.@deprecate(
    svd(t::AbstractTensorMap;
            trunc::TruncationScheme = notrunc(), p::Real = 2, alg::SVDAlg = SDD()),
    tsvd(t; trunc = trunc, p = p, alg = alg))
Base.@deprecate(
    svd!(t::AbstractTensorMap;
            trunc::TruncationScheme = notrunc(), p::Real = 2, alg::SVDAlg = SDD()),
    tsvd(t; trunc = trunc, p = p, alg = alg))

"""
    tsvd(t::AbstractTensorMap, leftind::Tuple, rightind::Tuple;
        trunc::TruncationScheme = notrunc(), p::Real = 2, alg::Union{SVD, SDD} = SDD())
        -> U, S, V, ϵ

Compute the (possibly truncated)) singular value decomposition such that
`norm(permute(t, leftind, rightind) - U * S *V) ≈ ϵ`, where `ϵ` thus represents the truncation error.

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in
`t` to be destroyed/overwritten, by using `tsvd!(t, trunc = notrunc(), p = 2)`.

A truncation parameter `trunc` can be specified for the new internal dimension, in which
case a truncated singular value decomposition will be computed. Choices are:
*   `notrunc()`: no truncation (default);
*   `truncerr(η::Real)`: truncates such that the p-norm of the truncated singular values is
    smaller than `η` times the p-norm of all singular values;
*   `truncdim(χ::Int)`: truncates such that the equivalent total dimension of the internal
    vector space is no larger than `χ`;
*   `truncspace(V)`: truncates such that the dimension of the internal vector space is
    smaller than that of `V` in any sector.
*   `truncbelow(χ::Real)`: truncates such that every singular value is larger then `χ` ;

The method `tsvd` also returns the truncation error `ϵ`, computed as the `p` norm of the
singular values that were truncated.

THe keyword `alg` can be equal to `SVD()` or `SDD()`, corresponding to the underlying LAPACK
algorithm that computes the decomposition (`_gesvd` or `_gesdd`).

Orthogonality requires `spacetype(t)<:InnerProductSpace`, and `svd(!)` is currently
only implemented for `spacetype(t)<:EuclideanSpace`.
"""
tsvd(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple; kwargs...) =
    tsvd!(permute(t, p1, p2; copy = true); kwargs...)

"""
    leftorth(t::AbstractTensorMap, leftind::Tuple, rightind::Tuple;
                alg::OrthogonalFactorizationAlgorithm = QRpos()) -> Q, R

Create orthonormal basis `Q` for indices in `leftind`, and remainder `R` such that
`permute(t, leftind, rightind) = Q*R`.

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in `t`
to be destroyed/overwritten, by using `leftorth!(t, alg = QRpos())`.

Different algorithms are available, namely `QR()`, `QRpos()`, `SVD()` and `Polar()`. `QR()`
and `QRpos()` use a standard QR decomposition, producing an upper triangular matrix `R`.
`Polar()` produces a Hermitian and positive semidefinite `R`. `QRpos()` corrects the
standard QR decomposition such that the diagonal elements of `R` are positive. Only
`QRpos()` and `Polar()` are uniqe (no residual freedom) so that they always return the same
result for the same input tensor `t`.

Orthogonality requires `spacetype(t)<:InnerProductSpace`, and `leftorth(!)` is currently
only implemented for `spacetype(t)<:EuclideanSpace`.
"""
leftorth(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple; kwargs...) =
    leftorth!(permute(t, p1, p2; copy = true); kwargs...)

"""
    rightorth(t::AbstractTensorMap, leftind::Tuple, rightind::Tuple;
                alg::OrthogonalFactorizationAlgorithm = LQpos()) -> L, Q

Create orthonormal basis `Q` for indices in `rightind`, and remainder `L` such that
`permute(t, leftind, rightind) = L*Q`.

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in `t`
to be destroyed/overwritten, by using `rightorth!(t, alg = LQpos())`.

Different algorithms are available, namely `LQ()`, `LQpos()`, `RQ()`, `RQpos()`, `SVD()` and
`Polar()`. `LQ()` and `LQpos()` produce a lower triangular matrix `L` and are computed using
a QR decomposition of the transpose. `RQ()` and `RQpos()` produce an upper triangular
remainder `L` and only works if the total left dimension is smaller than or equal to the
total right dimension. `LQpos()` and `RQpos()` add an additional correction such that the
diagonal elements of `L` are positive. `Polar()` produces a Hermitian and positive
semidefinite `L`. Only `LQpos()`, `RQpos()` and `Polar()` are uniqe (no residual freedom) so
that they always return the same result for the same input tensor `t`.

Orthogonality requires `spacetype(t)<:InnerProductSpace`, and `rightorth(!)` is currently
only implemented for `spacetype(t)<:EuclideanSpace`.
"""
rightorth(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple; kwargs...) =
    rightorth!(permute(t, p1, p2; copy = true); kwargs...)

"""
    leftnull(t::AbstractTensor, leftind::Tuple, rightind::Tuple;
                alg::OrthogonalFactorizationAlgorithm = QRpos()) -> N

Create orthonormal basis for the orthogonal complement of the support of the indices in
`leftind`, such that `N' * permute(t, leftind, rightind) = 0`.

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in `t`
to be destroyed/overwritten, by using `leftnull!(t, alg = QRpos())`.

Different algorithms are available, namely `QR()` (or equivalently, `QRpos()`), `SVD()` and
`SDD()`. The first assumes that the matrix is full rank and requires `iszero(atol)` and
`iszero(rtol)`. With `SVD()` and `SDD()`, `rightnull` will use the corresponding singular
value decomposition, and one can specify an absolute or relative tolerance for which
singular values are to be considered zero, where `max(atol, norm(t)*rtol)` is used as upper
bound.

Orthogonality requires `spacetype(t)<:InnerProductSpace`, and `leftnull(!)` is currently
only implemented for `spacetype(t)<:EuclideanSpace`.
"""
leftnull(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple; kwargs...) =
    leftnull!(permute(t, p1, p2; copy = true); kwargs...)

"""
    rightnull(t::AbstractTensor, leftind::Tuple, rightind::Tuple;
                alg::OrthogonalFactorizationAlgorithm = LQ(),
                atol::Real = 0.0,
                rtol::Real = eps(real(float(one(eltype(t)))))*iszero(atol)) -> N

Create orthonormal basis for the orthogonal complement of the support of the indices in
`rightind`, such that `permute(t, leftind, rightind)*N' = 0`.

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in `t`
to be destroyed/overwritten, by using `rightnull!(t, alg = LQpos())`.

Different algorithms are available, namely `LQ()` (or equivalently, `LQpos`), `SVD()` and
`SDD()`. The first assumes that the matrix is full rank and requires `iszero(atol)` and
`iszero(rtol)`. With `SVD()` and `SDD()`, `rightnull` will use the corresponding singular
value decomposition, and one can specify an absolute or relative tolerance for which
singular values are to be considered zero, where `max(atol, norm(t)*rtol)` is used as upper
bound.

Orthogonality requires `spacetype(t)<:InnerProductSpace`, and `rightnull(!)` is currently
only implemented for `spacetype(t)<:EuclideanSpace`.
"""
rightnull(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple; kwargs...) =
    rightnull!(permute(t, p1, p2; copy = true); kwargs...)

"""
    eigen(t::AbstractTensor, leftind::Tuple, rightind::Tuple; kwargs...) -> D, V

Compute eigenvalue factorization of tensor `t` as linear map from `rightind` to `leftind`.

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in `t`
to be destroyed/overwritten, by using `eigen!(t)`. Note that the permuted tensor on which
`eigen!` is called should have equal domain and codomain, as otherwise the eigenvalue
decomposition is meaningless and cannot satisfy
```
permute(t, leftind, rightind) * V = V * D
```

Accepts the same keyword arguments `scale`, `permute` and `sortby` as `eigen` of dense
matrices. See the corresponding documentation for more information.

See also `eig` and `eigh`
"""
LinearAlgebra.eigen(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple; kwargs...) =
    eigen!(permute(t, p1, p2; copy = true); kwargs...)

"""
    eig(t::AbstractTensor, leftind::Tuple, rightind::Tuple; kwargs...) -> D, V

Compute eigenvalue factorization of tensor `t` as linear map from `rightind` to `leftind`.
The function `eig` assumes that the linear map is not hermitian and returns type stable
complex valued `D` and `V` tensors for both real and complex valued `t`. See `eigh` for
hermitian linear maps

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in
`t` to be destroyed/overwritten, by using `eig!(t)`. Note that the permuted tensor on
which `eig!` is called should have equal domain and codomain, as otherwise the eigenvalue
decomposition is meaningless and cannot satisfy
```
permute(t, leftind, rightind) * V = V * D
```

Accepts the same keyword arguments `scale`, `permute` and `sortby` as `eigen` of dense matrices. See the corresponding documentation for more information.

See also `eigen` and `eigh`.
"""
eig(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple; kwargs...) =
    eig!(permute(t, p1, p2; copy = true); kwargs...)

"""
    eigh(t::AbstractTensorMap{<:EuclideanSpace}, leftind::Tuple, rightind::Tuple) -> D, V

Compute eigenvalue factorization of tensor `t` as linear map from `rightind` to `leftind`.
The function `eigh` assumes that the linear map is hermitian and `D` and `V` tensors with
the same `eltype` as `t`. See `eig` and `eigen` for non-hermitian tensors. Hermiticity
requires that the tensor acts on inner product spaces, and the current implementation
requires `spacetyp(t) <: EuclideanSpace`.

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in
`t` to be destroyed/overwritten, by using `eigh!(t)`. Note that the permuted tensor on
which `eigh!` is called should have equal domain and codomain, as otherwise the eigenvalue
decomposition is meaningless and cannot satisfy
```
permute(t, leftind, rightind) * V = V * D
```

See also `eigen` and `eig`.
"""
eigh(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple) =
    eigh!(permute(t, p1, p2; copy = true))

"""
    isposdef(t::AbstractTensor{<:EuclideanSpace}, leftind::Tuple, rightind::Tuple) -> ::Bool

Test whether a tensor `t` is positive definite as linear map from `rightind` to `leftind`.

If `leftind` and `rightind` are not specified, the current partition of left and right
indices of `t` is used. In that case, less memory is allocated if one allows the data in
`t` to be destroyed/overwritten, by using `isposdef!(t)`. Note that the permuted tensor on
which `isposdef!` is called should have equal domain and codomain, as otherwise it is
meaningless

Accepts the same keyword arguments `scale`, `permute` and `sortby` as `eigen` of dense
matrices. See the corresponding documentation for more information.
"""
LinearAlgebra.isposdef(t::AbstractTensorMap, p1::IndexTuple, p2::IndexTuple) =
    isposdef!(permute(t, p1, p2; copy = true))

tsvd(t::AbstractTensorMap; trunc::TruncationScheme = NoTruncation(),
                            p::Real = 2, alg::Union{SVD, SDD} = SDD()) =
    tsvd!(copy(t); trunc = trunc, p = p, alg = alg)
leftorth(t::AbstractTensorMap; alg::OFA = QRpos(), kwargs...) =
    leftorth!(copy(t); alg = alg, kwargs...)
rightorth(t::AbstractTensorMap; alg::OFA = LQpos(), kwargs...) =
    rightorth!(copy(t); alg = alg, kwargs...)
leftnull(t::AbstractTensorMap; alg::OFA = QR(), kwargs...) =
    leftnull!(copy(t); alg = alg, kwargs...)
rightnull(t::AbstractTensorMap; alg::OFA = LQ(), kwargs...) =
    rightnull!(copy(t); alg = alg, kwargs...)
LinearAlgebra.eigen(t::AbstractTensorMap; kwargs...) = eigen!(copy(t); kwargs...)
eig(t::AbstractTensorMap; kwargs...) = eig!(copy(t); kwargs...)
eigh(t::AbstractTensorMap; kwargs...) = eigh!(copy(t); kwargs...)
LinearAlgebra.isposdef(t::AbstractTensorMap) = isposdef!(copy(t))

# Orthogonal factorizations (mutation for recycling memory):
# only correct if Euclidean inner product
#------------------------------------------------------------------------------------------
leftorth!(t::AdjointTensorMap{S}; alg::OFA = QRpos()) where {S<:EuclideanSpace} =
    map(adjoint, reverse(rightorth!(adjoint(t); alg = alg')))

rightorth!(t::AdjointTensorMap{S}; alg::OFA = LQpos()) where {S<:EuclideanSpace} =
    map(adjoint, reverse(leftorth!(adjoint(t); alg = alg')))

leftnull!(t::AdjointTensorMap{S}; alg::OFA = QR(), kwargs...) where {S<:EuclideanSpace} =
    adjoint(rightnull!(adjoint(t); alg = alg', kwargs...))

rightnull!(t::AdjointTensorMap{S}; alg::OFA = LQ(), kwargs...) where {S<:EuclideanSpace} =
    adjoint(leftnull!(adjoint(t); alg = alg', kwargs...))

function tsvd!(t::AdjointTensorMap{S};
                trunc::TruncationScheme = NoTruncation(),
                p::Real = 2,
                alg::Union{SVD, SDD} = SDD()) where {S<:EuclideanSpace}
    u, s, vt, err = tsvd!(adjoint(t); trunc = trunc, p = p, alg = alg)
    return adjoint(vt), adjoint(s), adjoint(u), err
end

function leftorth!(t::TensorMap{<:EuclideanSpace};
                    alg::Union{QR, QRpos, QL, QLpos, SVD, SDD, Polar} = QRpos(),
                    atol::Real = zero(float(real(eltype(t)))),
                    rtol::Real = (alg ∉ (SVD(), SDD())) ? zero(float(real(eltype(t)))) :
                    eps(real(float(one(eltype(t)))))*iszero(atol))
    if !iszero(rtol)
        atol = max(atol, rtol*norm(t))
    end
    I = sectortype(t)
    S = spacetype(t)
    A = storagetype(t)
    Qdata = SectorDict{I, A}()
    Rdata = SectorDict{I, A}()
    dims = SectorDict{I, Int}()
    for c in blocksectors(domain(t))
        isempty(block(t,c)) && continue
        Q, R = _leftorth!(block(t, c), alg, atol)
        Qdata[c] = Q
        Rdata[c] = R
        dims[c] = size(Q, 2)
    end
    V = S(dims)
    if alg isa Polar
        @assert V ≅ domain(t)
        W = domain(t)
    elseif length(domain(t)) == 1 && domain(t) ≅ V
        W = domain(t)
    elseif length(codomain(t)) == 1 && codomain(t) ≅ V
        W = codomain(t)
    else
        W = ProductSpace(V)
    end
    return TensorMap(Qdata, codomain(t)←W), TensorMap(Rdata, W←domain(t))
end

function leftnull!(t::TensorMap{<:EuclideanSpace};
                    alg::Union{QR, QRpos, SVD, SDD} = QRpos(),
                    atol::Real = zero(float(real(eltype(t)))),
                    rtol::Real = (alg ∉ (SVD(), SDD())) ? zero(float(real(eltype(t)))) :
                    eps(real(float(one(eltype(t)))))*iszero(atol))
    if !iszero(rtol)
        atol = max(atol, rtol*norm(t))
    end
    I = sectortype(t)
    S = spacetype(t)
    A = storagetype(t)
    V = codomain(t)
    Ndata = SectorDict{I, A}()
    dims = SectorDict{I, Int}()
    for c in blocksectors(V)
        N = _leftnull!(block(t, c), alg, atol)
        Ndata[c] = N
        dims[c] = size(N, 2)
    end
    W = S(dims)
    return TensorMap(Ndata, V←W)
end

function rightorth!(t::TensorMap{<:EuclideanSpace};
                    alg::Union{LQ, LQpos, RQ, RQpos, SVD, SDD, Polar} = LQpos(),
                    atol::Real = zero(float(real(eltype(t)))),
                    rtol::Real = (alg ∉ (SVD(), SDD())) ? zero(float(real(eltype(t)))) :
                    eps(real(float(one(eltype(t)))))*iszero(atol))
    if !iszero(rtol)
        atol = max(atol, rtol*norm(t))
    end
    I = sectortype(t)
    S = spacetype(t)
    A = storagetype(t)
    Ldata = SectorDict{I, A}()
    Qdata = SectorDict{I, A}()
    dims = SectorDict{I, Int}()
    for c in blocksectors(codomain(t))
        isempty(block(t,c)) && continue
        L, Q = _rightorth!(block(t, c), alg, atol)
        Ldata[c] = L
        Qdata[c] = Q
        dims[c] = size(Q, 1)
    end
    V = S(dims)
    if alg isa Polar
        @assert V ≅ codomain(t)
        W = codomain(t)
    elseif length(codomain(t)) == 1 && codomain(t) ≅ V
        W = codomain(t)
    elseif length(domain(t)) == 1 && domain(t) ≅ V
        W = domain(t)
    else
        W = ProductSpace(V)
    end
    return TensorMap(Ldata, codomain(t)←W), TensorMap(Qdata, W←domain(t))
end

function rightnull!(t::TensorMap{<:EuclideanSpace};
                    alg::Union{LQ, LQpos, SVD, SDD} = LQpos(),
                    atol::Real = zero(float(real(eltype(t)))),
                    rtol::Real = (alg ∉ (SVD(), SDD())) ? zero(float(real(eltype(t)))) :
                    eps(real(float(one(eltype(t)))))*iszero(atol))
    if !iszero(rtol)
        atol = max(atol, rtol*norm(t))
    end
    I = sectortype(t)
    S = spacetype(t)
    A = storagetype(t)
    V = domain(t)
    Ndata = SectorDict{I, A}()
    dims = SectorDict{I, Int}()
    for c in blocksectors(V)
        N = _rightnull!(block(t, c), alg, atol)
        Ndata[c] = N
        dims[c] = size(N, 1)
    end
    W = S(dims)
    return TensorMap(Ndata, W←V)
end

function tsvd!(t::TensorMap{<:EuclideanSpace};
                trunc::TruncationScheme = NoTruncation(),
                p::Real = 2,
                alg::Union{SVD, SDD} = SDD())
    S = spacetype(t)
    I = sectortype(t)
    A = storagetype(t)
    Ar = similarstoragetype(t, real(eltype(t)))
    Udata = SectorDict{I, A}()
    Σmdata = SectorDict{I, Ar}() # this will contain the singular values as matrix
    Vdata = SectorDict{I, A}()
    dims = SectorDict{sectortype(t), Int}()
    if isempty(blocksectors(t))
        W = S(dims)
        truncerr = zero(real(eltype(t)))
        return TensorMap(Udata, codomain(t)←W), TensorMap(Σmdata, W←W),
                    TensorMap(Vdata, W←domain(t)), truncerr
    end
    for (c, b) in blocks(t)
        U, Σ, V = _svd!(b, alg)
        Udata[c] = U
        Vdata[c] = V
        if @isdefined Σdata # cannot easily infer the type of Σ, so use this construction
            Σdata[c] = Σ
        else
            Σdata = SectorDict(c=>Σ)
        end
        dims[c] = length(Σ)
    end
    if !isa(trunc, NoTruncation)
        Σdata, truncerr = _truncate!(Σdata, trunc, p)
        truncdims = SectorDict{I, Int}()
        for c in blocksectors(t)
            truncdim = length(Σdata[c])
            if truncdim != 0
                truncdims[c] = truncdim
                if truncdim != dims[c]
                    Udata[c] = Udata[c][:, 1:truncdim]
                    Vdata[c] = Vdata[c][1:truncdim, :]
                end
            else
                delete!(Udata, c)
                delete!(Vdata, c)
                delete!(Σdata, c)
            end
        end
        dims = truncdims
        W = S(dims)
    else
        W = S(dims)
        if length(domain(t)) == 1 && domain(t)[1] ≅ W
            W = domain(t)[1]
        elseif length(codomain(t)) == 1 && codomain(t)[1] ≅ W
            W = codomain(t)[1]
        end
        truncerr = abs(zero(eltype(t)))
    end
    for (c, Σ) in Σdata
        Σmdata[c] = copyto!(similar(Σ, length(Σ), length(Σ)), Diagonal(Σ))
    end
    return TensorMap(Udata, codomain(t)←W), TensorMap(Σmdata, W←W),
            TensorMap(Vdata, W←domain(t)), truncerr
end

function LinearAlgebra.ishermitian(t::TensorMap)
    domain(t) == codomain(t) || return false
    spacetype(t) <: EuclideanSpace || return false # hermiticity only defined for euclidean
    for (c, b) in blocks(t)
        ishermitian(b) || return false
    end
    return true
end

LinearAlgebra.eigen!(t::TensorMap) = ishermitian(t) ? eigh!(t) : eig!(t)

function eigh!(t::TensorMap{<:EuclideanSpace}; kwargs...)
    domain(t) == codomain(t) ||
        throw(SpaceMismatch("`eigh!` requires domain and codomain to be the same"))
    S = spacetype(t)
    I = sectortype(t)
    A = storagetype(t)
    Ar = similarstoragetype(t, real(eltype(t)))
    Ddata = SectorDict{I, Ar}()
    Vdata = SectorDict{I, A}()
    dims = SectorDict{I, Int}()
    for (c, b) in blocks(t)
        values, vectors = eigen!(Hermitian(b); kwargs...)
        d = length(values)
        Ddata[c] = copyto!(similar(values, (d, d)), Diagonal(values))
        Vdata[c] = vectors
        dims[c] = d
    end
    if length(domain(t)) == 1
        W = domain(t)[1]
    else
        W = S(dims)
    end
    return TensorMap(Ddata, W←W), TensorMap(Vdata, domain(t)←W)
end

function eig!(t::TensorMap; kwargs...)
    domain(t) == codomain(t) ||
        throw(SpaceMismatch("`eig!` requires domain and codomain to be the same"))
    S = spacetype(t)
    I = sectortype(t)
    T = complex(eltype(t))
    Ac = similarstoragetype(t, T)
    Ddata = SectorDict{I, Ac}()
    Vdata = SectorDict{I, Ac}()
    dims = SectorDict{I, Int}()
    for (c, b) in blocks(t)
        values, vectors = eigen!(b; kwargs...)
        d = length(values)
        Ddata[c] = copyto!(similar(values, T, (d, d)), Diagonal(values))
        if eltype(vectors) == T
            Vdata[c] = vectors
        else
            Vdata[c] = copyto!(similar(vectors, T), vectors)
        end
        dims[c] = d
    end
    if length(domain(t)) == 1
        W = domain(t)[1]
    else
        W = S(dims)
    end
    return TensorMap(Ddata, W←W), TensorMap(Vdata, domain(t)←W)
end

function LinearAlgebra.isposdef!(t::TensorMap)
    domain(t) == codomain(t) ||
        throw(SpaceMismatch("`isposdef` requires domain and codomain to be the same"))
    spacetype(t) <: EuclideanSpace || return false
    for (c, b) in blocks(t)
        isposdef!(b) || return false
    end
    return true
end


function test_energy(mesh, sbp, eqn, opts)
  energy_final = calcNorm(eqn, eqn.q_vec)
  q_initial = zeros(eqn.q_vec)
  ICfunc = AdvectionEquationMod.ICDict[opts["IC_name"]]
  ICfunc(mesh, sbp, eqn, opts, q_initial)
  energy_initial = calcNorm(eqn, q_initial)

  @test  abs(energy_initial - energy_final)  < 1e-12
end

"""
  Test energy stability in 2 and 3 dimensions.
"""
function test_energy_serial()
  @testset "----- Testing Energy Stability -----" begin

    start_dir = pwd()
    cd(dirname(@__FILE__))

    fname = "input_vals_periodic.jl"
    cd("./2dp1")
    mesh, sbp, eqn, opts = solvePDE(fname)
    test_energy(mesh, sbp, eqn, opts)

    cd("../3dp1")
    mesh, sbp, eqn, opts = solvePDE(fname)
    test_energy(mesh, sbp, eqn, opts)

    cd(start_dir)

  end

  return nothing
end

#test_energy_serial()
add_func1!(AdvectionTests, test_energy_serial, [TAG_SHORTTEST])


###
### logistic and logit
###

logistic(x::Real) = inv(one(x) + exp(-x))

function logistic_logjac(x::Real)
    mx = -abs(x)
    mx - 2*log1p(exp(mx))
end

logit(x::Real) = log(x / (one(x) - x))

logit_logjac(y) = -log(y) - log1p(-y)

###
### calculations
###

"""
    $SIGNATURES

Number of elements (strictly) above the diagonal in an ``n×n`` matrix.
"""
unit_triangular_dimension(n::Int) = n * (n-1) ÷ 2

###
### type calculations
###

"""
    $(SIGNATURES)

Extend element type of argument so that it is closed under the algebra used by this package.

Pessimistic default for non-real types.
"""
function extended_eltype(::Type{S}) where S
    T = eltype(S)
    T <: Real ? typeof(√(one(T))) : Any
end

extended_eltype(x::T) where T = extended_eltype(T)

###
### macros
###

"""
$(SIGNATURES)

Workaround for https://github.com/JuliaLang/julia/issues/14919 to make
transformation types callable.

TODO: remove when this issue is closed, also possibly remove MacroTools as a
dependency if not used elsewhere.
"""
macro calltrans(ex)
    if @capture(ex, struct T1_ fields__ end)
        @capture T1 (T2_ <: S_|T2_)
        @capture T2 (T3_{params__}|T3_)
        quote
            Base.@__doc__ $(esc(ex))
            (t::$(esc(T3)))(x) = transform(t, x)
        end
    else
        throw(ArgumentError("can't find anything to make callable in $(ex)"))
    end
end

####
#### random values
####

"Shared part of docstrings for keyword arguments of or passed to [`random_reals`](@ref)."
const _RANDOM_REALS_KWARGS_DOC = """
A standard multivaritate normal or Cauchy is used, depending on `cauchy`, then scaled with
`scale`. `rng` is the random number generator used.
"""

_random_reals_scale(rng::AbstractRNG, scale::Real, cauchy::Bool) =
    cauchy ? scale / abs2(randn(rng)) : scale * 1.0

"""
$(SIGNATURES)

Random real number.

$(_RANDOM_REALS_KWARGS_DOC)
"""
random_real(; scale::Real = 1, cauchy::Bool = false, rng::AbstractRNG = GLOBAL_RNG) =
    randn(rng) * _random_reals_scale(rng, scale, cauchy)

"""
$(SIGNATURES)

Random vector in ``ℝⁿ``.

$(_RANDOM_REALS_KWARGS_DOC)
"""
random_reals(n::Integer; scale::Real = 1, cauchy::Bool = false,
             rng::AbstractRNG = GLOBAL_RNG) =
                 randn(rng, n) .* _random_reals_scale(rng, scale, cauchy)


# Starting point for Arrays example file for Learning Julia

# create an array using [] syntax
arr1 = [1, 2, 3, 4, 5]
println(arr1)
println("Type of arr1 is ", typeof(arr1))
println("Size of arr1 is ", length(arr1))

# access an array element - array indexes are 1-based
println("Element at position 2 of this array is ", arr1[2])
println()

# arrays can hold different types of values
arr2 = ["one", 2, 3.0, -4]
println(arr2)
println("Type of arr2 is ", typeof(arr2))
println("Size of arr2 is ", length(arr2))
println()

# declare an array with a particular type
arr3 = Float64[]
push!(arr3, 1.0, 2.0, 3.0) # the ! symbol indicates that we are modifying a defined variable
println(arr3)
println("Type of arr3 is ", typeof(arr3))
println("Size of arr3 is ", length(arr3))
println()

# populate an array - fill wth default values
arr4 = zeros(4)
println(arr4)
println()

arr5 = ones(4)
println(arr5)
println()

arr6 = fill(5, 7) # fill 7 instances of 5
println(arr6)
println()

# Array sorting using the sort function
arr7 = sort(arr1, rev = true)
println("arr1 reverse sorted is ", arr7)
println()

# sort! modifies the original array in place
sort!(arr7) # again, ! modify in place
println("arr1 reverse reverse sorted is ", arr7)
println()

# Convert to string with a delimiter
strVal = join(arr1, "-")
println(strVal)


smooth(g::AbstractMatrix, k::IdentityKernel) = g
smooth(g::AbstractMatrix, k::SmoothingKernel) = k.smoother(g, k.ξ)

nobs(mf::MomentFunction) = mf.nobs
nmom(mf::MomentFunction) = mf.nmom
npar(mf::MomentFunction) = mf.npar

function get_mom_deriv(g::Function, dtype::Symbol, k::SmoothingKernel, nobs, nmom, npar)

  s(θ::Vector)             = smooth(g(θ), k)
  gn(θ::Vector, p::Vector) = s(θ)'*p
  g1(θ::Vector, λ::Vector) = s(θ)*λ

  g2(θ::Vector, p::Vector, λ::Vector) = (p'*s(θ)*λ)[1]

  gn_closure(θ::Vector) = gn(θ, __p)
  g1_closure(θ::Vector) = g1(θ, __λ)
  g2_closure(θ::Vector) = g2(θ, __p, __λ)

  gn_closure!(θ::Vector, gg) = gg[:] = gn(θ, __p)
  g1_closure!(θ::Vector, gg) = gg[:] = g1(θ, __λ)

  if dtype==:typed
  	∂gn  = ForwardDiff.forwarddiff_jacobian(gn_closure, Float64, fadtype=:typed)
  	∂g1  = ForwardDiff.forwarddiff_jacobian(g1_closure, Float64, fadtype=:typed)
  	∂²g2 = ForwardDiff.forwarddiff_hessian(g2_closure, Float64, fadtype=:typed)
  elseif dtype==:dual
  	∂gn  = ForwardDiff.forwarddiff_jacobian(gn_closure!, Float64, fadtype=:dual, n = npar, m = nmom)
  	∂g1  = ForwardDiff.forwarddiff_jacobian(g1_closure!, Float64, fadtype=:dual, n = npar, m = nobs)
  	∂²g2 = ForwardDiff.forwarddiff_hessian(g2_closure, Float64, fadtype=:typed)
  elseif dtype==:diff
  	∂gn(θ::Vector)  = Calculus.jacobian(gn_closure, θ, :central)
  	∂g1(θ::Vector)  = Calculus.jacobian(g1_closure, θ, :central)
  	∂²g2(θ::Vector) = Calculus.hessian(g2_closure, θ, :central)
  end
  return (g, s, gn, ∂gn, ∂g1, ∂²g2)
end

function MomentFunction(g::Function, dtype::Symbol; nobs = Nothing, npar = Nothing, nmom = Nothing)
  ## Default is no smoothing
  MomentFunction(get_mom_deriv(g, dtype, IdentityKernel(), nobs, nmom, npar)...,
                                IdentityKernel(), nobs, nmom, npar)
end

function MomentFunction(g::Function, dtype::Symbol, k::SmoothingKernel;
                         nobs = Nothing, npar = Nothing, nmom = Nothing)
  ## Default is no smoothing
  MomentFunction(get_mom_deriv(g, dtype, k, nobs, nmom, npar)..., k, nobs, nmom, npar)
end


using FRGRealTime.MathFuns

function massflow1(k, m2, T, lambda, Nf=4)
    (1 / 6) * Nf * lambda * TF.Fb1(k, m2, T)
end

function massflow2(k, m2, T, lambda, Nf=4)
    (1 / 6) * (Nf + 2) * lambda * TF.Fb1(k, m2, T)
end

function lambdaflow1(k, m2, T, lambda)
    -(1 / 3) * lambda^2 * TF.Fb2(k, m2, T)
end

function lambdaflow2(k, m2, T, lambda, Nf=4)
    -(1 / 6) * lambda^2 * (Nf + 8) * TF.Fb2(k, m2, T)
end

function Imlambdaflow1(p0, k, m2, lambda, T, Imlambda, Relambda, ϵ=0.1)
    -(1 / 3) * (
        -Imlambda^2 * TF.ImFb2(p0, k, m2, T, ϵ) +
        2 * TF.ReFb2(p0, k, m2, T, ϵ) * Relambda * Imlambda +
        TF.ImFb2(p0, k, m2, T, ϵ) * Relambda^2
    )
end

function Relambdaflow1(p0, k, m2, lambda, T, Imlambda, Relambda, ϵ=0.1)
    -(1 / 3) * (
        -Imlambda^2 * TF.ReFb2(p0, k, m2, T, ϵ) -
        2 * TF.ImFb2(p0, k, m2, T, ϵ) * Relambda * Imlambda +
        TF.ReFb2(p0, k, m2, T, ϵ) * Relambda^2
    )
end

function Imlambdaflow2(p0, k, m2, lambda, T, Imlambda, Relambda, ϵ=0.1)
    -(2 / 3) * (
        -Imlambda^2 * TF.ImFb2(p0, k, m2, T, ϵ) +
        2 * TF.ReFb2(p0, k, m2, T, ϵ) * Relambda * Imlambda +
        TF.ImFb2(p0, k, m2, T, ϵ) * Relambda^2
    )
end

function Relambdaflow2(
    p0,
    k,
    m2,
    lambda,
    T,
    Imlambda,
    Relambda,
    ϵ=0.1,
    Nf=4,
)
    -(2 / 3) * (
        -Imlambda^2 * TF.ReFb2(p0, k, m2, T, ϵ) -
        2 * TF.ImFb2(p0, k, m2, T, ϵ) * Relambda * Imlambda +
        TF.ReFb2(p0, k, m2, T, ϵ) * Relambda^2
    ) - (1 / 6) * lambda^2 * (Nf + 4) * TF.Fb2(k, m2, T)
end

function Imlambdaflow3(p0, q0, k, m2, lambda, T, Imlambda, Relambda, ϵ=0.1)
    -(1 / 3) * (
        -Imlambda^2 * TF.ImFb2(p0 + q0, k, m2, T, ϵ) +
        2 *
        TF.ReFb2(p0 + q0, k, m2, T, ϵ) *
        Relambda *
        Imlambda +
        TF.ImFb2(p0 + q0, k, m2, T, ϵ) * Relambda^2
    ) -
    (1 / 3) * (
        -Imlambda^2 * TF.ImFb2(p0 - q0, k, m2, T, ϵ) +
        2 *
        TF.ReFb2(p0 - q0, k, m2, T, ϵ) *
        Relambda *
        Imlambda +
        TF.ImFb2(p0 - q0, k, m2, T, ϵ) * Relambda^2
    )
end

function Imlambdaflow3_ac(p0, q0, k, m2, lambda, T, Imlambda, Relambda, ϵ=0.1)
    -(1 / 3) * (
        -Imlambda^2 * TF.ImFb2(p0 + q0, k, m2, T, ϵ) +
        2 *
        TF.ReFb2_ac(p0 + q0, k, m2, T, ϵ) *
        Relambda *
        Imlambda +
        TF.ImFb2(p0 + q0, k, m2, T, ϵ) * Relambda^2
    ) -
    (1 / 3) * (
        -Imlambda^2 * TF.ImFb2(p0 - q0, k, m2, T, ϵ) +
        2 *
        TF.ReFb2_ac(p0 - q0, k, m2, T, ϵ) *
        Relambda *
        Imlambda +
        TF.ImFb2(p0 - q0, k, m2, T, ϵ) * Relambda^2
    )
end
function Relambdaflow3_ac(
    p0,
    q0,
    k,
    m2,
    lambda,
    T,
    Imlambda,
    Relambda,
    ϵ=0.1,
    Nf=4,
)
    -(1 / 3) * (
        -Imlambda^2 * TF.ReFb2_ac(p0 + q0, k, m2, T, ϵ) -
        2 *
        TF.ImFb2(p0 + q0, k, m2, T, ϵ) *
        Relambda *
        Imlambda + TF.ReFb2_ac(p0 + q0, k, m2, T, ϵ) * Relambda^2
    ) -
    (1 / 3) * (
        -Imlambda^2 * TF.ReFb2_ac(p0 - q0, k, m2, T, ϵ) -
        2 *
        TF.ImFb2(p0 - q0, k, m2, T, ϵ) *
        Relambda *
        Imlambda + TF.ReFb2_ac(p0 - q0, k, m2, T, ϵ) * Relambda^2
    ) - (1 / 6) * lambda^2 * (Nf + 4) * TF.Fb2(k, m2, T)
end

function Imlambdaflow4(p0, q0, k, m2, lambda, T, Imlambda, Relambda, ϵ = 0.1)
    -(1 / 3) * (TF.ImFb2(p0 + q0, k, m2, T, ϵ) * Relambda^2) -
    (1 / 3) * (TF.ImFb2(p0 - q0, k, m2, T, ϵ) * Relambda^2)
end

function Imlambdaflow4_ac(p0, q0, k, m2, lambda, T, Imlambda, Relambda, ϵ = 0.1)
    -(1 / 3) * (TF.ImFb2(p0 + q0, k, m2, T, ϵ) * Relambda^2) -
    (1 / 3) * (TF.ImFb2(p0 - q0, k, m2, T, ϵ) * Relambda^2)
end

function Relambdaflow4(
    p0,
    q0,
    k,
    m2,
    lambda,
    T,
    Imlambda,
    Relambda,
    ϵ = 0.1,
    Nf = 4,
)
    -(1 / 3) * (TF.ReFb2(p0 + q0, k, m2, T, ϵ) * Relambda^2) -
    (1 / 3) * (TF.ReFb2(p0 - q0, k, m2, T, ϵ) * Relambda^2) -
    (1 / 6) * lambda^2 * (Nf + 4) * TF.Fb2(k, m2, T)
end

function Relambdaflow4_ac(
    p0,
    q0,
    k,
    m2,
    lambda,
    T,
    Imlambda,
    Relambda,
    ϵ = 0.1,
    Nf = 4,
)
    -(1 / 3) * (TF.ReFb2_ac(p0 + q0, k, m2, T, ϵ) * Relambda^2) -
    (1 / 3) * (TF.ReFb2_ac(p0 - q0, k, m2, T, ϵ) * Relambda^2) -
    (1 / 6) * lambda^2 * (Nf + 4) * TF.Fb2(k, m2, T)
end

function twopointflow(p0, k, m2, T, lambda)
    -1 / 48 * (
        k^4 *
        csch(sqrt(k^2 + m2) / (2 * T))^2 *
        (
            lambda(k, abs(sqrt(k^2 + m2) - p0)) +
            lambda(k, abs(sqrt(k^2 + m2) + p0))
        ) *
        (sqrt(k^2 + m2) + T * sinh(sqrt(k^2 + m2) / T))
    ) / ((k^2 + m2)^(3 / 2) * pi^2 * T)
end

function twopointflow2(p0, k, m2, T, lambda)
    -1 / 24 * (
        k^4 *
        csch(sqrt(k^2 + m2) / (2 * T))^2 *
        lambda(k) *
        (sqrt(k^2 + m2) + T * sinh(sqrt(k^2 + m2) / T))
    ) / ((k^2 + m2)^(3 / 2) * pi^2 * T)
end

function twopointflow3(p0, k, m2, T, lambda)
    -1 / 24 * (
        k^4 *
        csch(sqrt(k^2 + m2) / (2 * T))^2 *
        lambda(k) *
        (sqrt(k^2 + m2) + T * sinh(sqrt(k^2 + m2) / T))
    ) / ((k^2 + m2)^(3 / 2) * pi^2 * T)
end

function RelambdaflowMktest(
    p0,
    q0,
    k,
    m2,
    lambda,
    T,
    Imlambda,
    RelambdaMk,
    ϵ = 0.1,
    Nf = 4,
)
    -(1 / 3) * (ReFb2Mkpositivetest(p0, k, m2, T, ϵ) * RelambdaMk^2)
end

function ReFb2Mkpositivetest(p0, k, m2, T, ϵ = 0.1)
    -1 / 48 * (k^4 * (ϵ - 4 * p0 * Eb(k, m2) + 8 * Eb(k, m2)^2)) / (
        pi^2 *
        Eb(k, m2)^3 *
        (p0 + 2 * Eb(k, m2)) *
        (p0^2 + ϵ + 4 * Eb(k, m2) * (-p0 + Eb(k, m2)))
    )
end


println("""
    Create a function that takes any real matrix as an input and returns
    the element type of the matrix:
""")

realmatrixeltype(x::Matrix{<:Real}) = eltype(x)
display(realmatrixeltype)

println("""
    Matrix m filled with random integers from 0 to 9:
""")

m = rand(0:9, 3, 3)
display(m)

println("""
    Element type of m:
""")

display(realmatrixeltype(m))


using Hecke, DelimitedFiles

boundcond = ARGS[1]
gtype = ARGS[2]

_gtype = replace(replace(replace(replace(replace(gtype, '[' => '_'), ']' => '_'), ' ' => '_'), ',' => '_'), "__" => "_")

file = "conductors_" * boundcond * "_" * _gtype;

bound = fmpz(Meta.eval(Meta.parse(boundcond)))
gtype = convert(Vector{Int}, Meta.eval(Meta.parse(gtype)))

sprint_formatted(fmt, args...) = @eval @sprintf($fmt, $(args...))

gtype_right = snf(abelian_group(gtype))[1].snf
g_exp = gtype_right[end]
wild_primes = factor(g_exp).fac
for (p, v) in wild_primes
  if p == 2
    wild_primes[p] = v+2
  else
    wild_primes[p] = v+1
  end
end
conductors = Vector{Int}()
for i = 3:bound
  lf = factor(i)
  admissible = true
  for (p, v) in lf
    if p in keys(wild_primes)
      if v > 1 && v <= wild_primes[p]
        continue
      elseif v > wild_primes[p]
        admissible = false
        break
      else
        #we need to check that p-1 is not coprime to the other factors
        if length(wild_primes) == 1
          admissible = false
          break
        else
          t = p-1
          found = false
          for q in keys(wild_primes)
            if !iscoprime(t, q)
              found = true
              break
            end
          end
          if !found
            admissible = false
            break
          end
        end
      end
    else
      if v > 1
        admissible = false
        break
      end
      #p is tamely ramified...
      t = p-1
      found = false
      for q in keys(wild_primes)
        if !iscoprime(t, q)
          found = true
          break
        end
      end
      if !found
        admissible = false
        break
      end
    end
  end
  if admissible
    push!(conductors, i)
  end
end

println("DONE")

if length(ARGS) == 3
  file = ARGS[3]
else
  file = file * "_" * string(length(conductors))
end

if isfile(file)
  error("File $file already exists")
end

open(file, "w") do f
  print("Writing conductors to $file ... ")
  writedlm(f, conductors)
  println("DONE")
end

@assert countlines(file) == length(conductors)


using DataFrames
using GLM
using RCall
using Test

@testset "Cubic spline matches at points" begin
    using Glissa
    x = collect(0:0.1:pi)
    f = cos.(x)
    cs = cubic_spline(x, f, [0, 0]; slope=FreeSlope())
    for i1 in 1:length(x)
        @test abs(cs(x[i1]) - f[i1]) < 1e-9
    end
end

function points_slope(logx, logy)
    df = DataFrame(logx = logx, logy = logy)
    model = @formula(logy ~ logx)
    regress = lm(model, df)
    slope = coef(regress)[2]
    slope
end

@testset "Cubic spline error reduces as dx reduces" begin
    using Glissa
    Δ = (1e-5) * 2 .^ (0:14)
    maxerr = zeros(Float64, length(Δ))
    for j in 1:length(Δ)
        x = collect(0:(Δ[j]):(pi/2))
        f = cos.(x)
        cs = cubic_spline(x, f, [0, 0]; slope=FreeSlope())

        xx = collect(0:(Δ[j]/10):(pi/2))
        mm = zero(Float64)
        for x1 in xx
            mm = max(mm, abs(cs(x1) - cos(x1)))
        end
        maxerr[j] = mm
    end
    slope = points_slope(log.(Δ), log.(maxerr))
    @test 0.7 < slope
    @test slope < 1.0
end

@testset "Cubic spline is continuous" begin
    using Glissa
    Δ = 0.1
    x = collect(0:Δ:pi)
    f = sin.(x)
    cs = cubic_spline(x, f, [0, 0]; slope=FreeSlope())
    for i1 in 1:(length(x) - 1)
        d = Δ / 10
        prev = 100
        # Assert values approach the expected value.
        for xl in range(x[i1] + Δ / 10, length = 5, stop = x[i1] + 1e-9)
            v = abs(sin(x[i1]) - cs(xl))
            @test v <= prev
            prev = v
        end
        # Assert they get close-enough for the final value.
        @test prev < 1e-5
    end
    for i1 in 2:length(x)
        d = Δ / 10
        prev = 100
        # Look at approach from below.
        for xl in range(x[i1] - Δ / 10, length = 5, stop = x[i1] - 1e-9)
            v = abs(sin(x[i1]) - cs(xl))
            @test v <= prev
            prev = v
        end
        @test prev < 1e-5
    end
end

@testset "Cubic spline is continuous for uneven x" begin
    using Glissa
    Δ = 0.1
    x = vcat([0.0, 0.1, 0.2, 0.4, 0.5, 1], 2:5)
    f = sin.(x)
    cs = cubic_spline(x, f, [0, 0]; slope=FreeSlope())
    for i1 in 1:(length(x) - 1)
        d = Δ / 10
        prev = 100
        # Assert values approach the expected value.
        for xl in range(x[i1] + Δ / 10, length = 5, stop = x[i1] + 1e-9)
            v = abs(sin(x[i1]) - cs(xl))
            @test v <= prev
            prev = v
        end
        # Assert they get close-enough for the final value.
        @test prev < 1e-5
    end
    for i1 in 2:length(x)
        d = Δ / 10
        prev = 100
        # Look at approach from below.
        for xl in range(x[i1] - Δ / 10, length = 5, stop = x[i1] - 1e-9)
            v = abs(sin(x[i1]) - cs(xl))
            @test v <= prev
            prev = v
        end
        @test prev < 1e-5
    end
end

@testset "Cubic spline overshoot" begin
    using Glissa
    # RPN 15A data from Fritsch-Carlson
    x = [7.99, 8.09, 8.19, 8.7, 9.2, 10, 12, 15, 20]
    f = [0, 2.76429e-5, 4.37498e-2, 0.169183, 0.469428, 0.943740, 0.998636, 0.999919, 0.999994]
    cs = cubic_spline(x, f, [0.0, 0.0]; slope = Monotonic())
    xx = 8.08:0.0001:8.1
    # plot(xx, cs.(xx))
    ϵ = 1e-8
    for i in 2:(length(x) - 1)
        left = (cs(x[i] - ϵ) - cs(x[i] - 2ϵ)) / ϵ
        right = (cs(x[i] + 2ϵ) - cs(x[i] + ϵ))/ϵ
        d = abs(left - right)
        @test d < 1e-6
    end
    ϵ = 1e-5
    second_derivative = zeros(Float64, length(x) - 2)
    for i in 2:(length(x) - 1)
        left = (cs(x[i] - 3ϵ) - 2*cs(x[i] - 2ϵ) + cs(x[i] - ϵ)) / ϵ^2
        right = (cs(x[i] + ϵ) - 2*cs(x[i] + 2ϵ) + cs(x[i] + 3ϵ)) / ϵ^2
        second_derivative[i - 1] = abs(left - right)
    end
    @test second_derivative[1] > 5 * second_derivative[5]
end

@testset "Hyman criterion matches splinefun in R" begin
    using Glissa
    using RCall
    # this hyman filter comes from the R splinefun source code.
    # x is the knots. y is values. b is gradient.
    R"""hyman_filter <- function(x, y, b)
    {
        n <- length(x)
        ss <- diff(y) / diff(x)
        S0 <- c(ss[1L], ss)
        S1 <- c(ss, ss[n-1L])
        t1 <- pmin(abs(S0), abs(S1))
        sig <- b
        ind <- S0*S1 > 0
        sig[ind] <- S1[ind]
        ind <- sig >= 0
        if(sum(ind)) b[ind] <- pmin(pmax(0, b[ind]), 3*t1[ind])
        ind <- !ind
        if(sum(ind)) b[ind] <- pmax(pmin(0, b[ind]), -3*t1[ind])
        b
    }
    """
    x = collect(0:0.1:π)
    y = sin.(x)
    yp = cos.(x)
    hf = collect(R"""hyman_filter($x, $y, $yp)""")
    Glissa.project_to_monotonicity!(x, y, yp)
    @test maximum(abs.(yp .- hf)) < 1e-12

    # This example is constructed by hand to need monotonic smoothing.
    x = [7.99, 8.09, 8.19, 8.7, 9.2, 10, 12, 15, 20]
    f = [0, 2.76429e-5, 4.37498e-2, 0.169183, 0.469428, 0.943740, 0.998636, 0.999919, 0.999994]
    fp = similar(f)
    Glissa.global_derivatives!(x, f, fp)
    hf = collect(R"""hyman_filter($x, $f, $fp)""")
    # Assert that the uncorrected slopes are different from corrected ones.
    @test maximum(abs.(hf - fp)) > 1e-10
    Glissa.project_to_monotonicity!(x, f, fp)
    # After correction, they agree with the R version.
    @test maximum(abs.(fp .- hf)) < 1e-12

    R"""spl_coef_conv <- function(x, y, b)
    {
        n <- length(x)
        h <- diff(x); y <- -diff(y)
        b0 <- b[-n]; b1 <- b[-1L]
        cc <- -(3*y + (2*b0 + b1)*h) / h^2
        c1 <- (3*y[n-1L] + (b0[n-1L] + 2*b1[n-1L])*h[n-1L]) / h[n-1L]^2
        c <- c(cc, c1)
        dd <- (2*y/h + b0 + b1) / h^2
        d <- c(dd, dd[n-1L])
        list(c, d)
    }
    """
    cd = R"""spl_coef_conv($x, $f, $fp)"""
    c = collect(cd[1])
    d = collect(cd[2])
    C = zeros(Float64, 4, length(x) - 1)
    Glissa.hyman_coefficients!(x, f, fp, C)
    # The R functions define vectors that are 1 longer.
    @test maximum(abs.(C[3, :] - c[1:(end-1)])) < 1e-12
    @test maximum(abs.(C[4, :] - d[1:(end-1)])) < 1e-12
end

@testset "Cubic spline is compatible with forward diff" begin
    using Glissa
    using ForwardDiff
    x = [7.99, 8.09, 8.19, 8.7, 9.2, 10, 12, 15, 20]
    x2 = collect(8:0.01:10)
    y = [0, 2.76429e-5, 4.37498e-2, 0.169183, 0.469428, 0.943740, 0.998636, 0.999919, 0.999994]

    function g(f::AbstractVector{T}) where {T}
        cs = Glissa.cubic_spline(x, f, zeros(T, 2); slope = Monotonic())
        cs.(x2)
    end

    interp = g(y)
    jacob = ForwardDiff.jacobian(g, y)
    @test size(jacob) == (201, length(x))
    @test jacob[1, 1] > 0
    @test abs(jacob[1, 9]) < 1e-7
end

@testset "Cubic spline is compatible with reverse diff" begin
    using Glissa
    using ReverseDiff
    x = [7.99, 8.09, 8.19, 8.7, 9.2, 10, 12, 15, 20]
    x2 = collect(8:0.01:10)
    f = [0, 2.76429e-5, 4.37498e-2, 0.169183, 0.469428, 0.943740, 0.998636, 0.999919, 0.999994]
    function g(f::AbstractVector{T}, x1, x2) where {T}
        cs = Glissa.cubic_spline(x1, f, zeros(T, 2); slope = Monotonic())
        cs.(x2)
    end
    interp = g(f, x, x2)
    f_tape = ReverseDiff.GradientTape(g, (f, x, x2))
    compiled_f_tape = ReverseDiff.compile(f_tape)
    inputs = (f, x, x2)
    cfg = ReverseDiff.GradientConfig(inputs)
    results = similar(interp)
    # This doesn't work, and I'd like it to work.
    # ReverseDiff.gradient!(results, compiled_f_tape, inputs)
end

@testset "Cubic spline integrates a flat line" begin
    using Glissa
    using Random
    using Distributions
    x = Float64[0, 1, 3, 5]
    y = Float64[1, 1, 1, 1]
    cs = Glissa.cubic_spline(x, y, [0.0, 0.0]; slope = Monotonic())
    @test abs(cs(2.5) - 1) < 1e-6
    exact(z) = z
    rng = MersenneTwister(9234724)
    for j in 1:100
        a = rand(rng, Uniform(0, 5))
        b = rand(rng, Uniform(a, 5))
        found = Glissa.integrate(cs, a, b)
        @test abs(found - exact(b) + exact(a)) < 1e-9
    end
end

@testset "Cubic spline integrates a sloped line" begin
    using Glissa
    using Random
    using Distributions
    x = Float64[0, 1, 3, 5]
    y = Float64[0, 1, 3, 5]
    cs = Glissa.cubic_spline(x, y, [1.0, 1.0]; slope = Monotonic())
    @test abs(cs(2.5) - 2.5) < 1e-6
    exact(z) = 0.5 * z^2
    rng = MersenneTwister(9292734)
    for j in 1:100
        a = rand(rng, Uniform(0, 5))
        b = rand(rng, Uniform(a, 5))
        found = Glissa.integrate(cs, a, b)
        @test abs(found - exact(b) + exact(a)) < 1e-9
    end
end

@testset "Cubic spline integrates a quadratic line" begin
    using Glissa
    using Random
    using Distributions
    Δ = 0.1
    x = collect(0:Δ:5)
    y = x.^2
    cs = Glissa.cubic_spline(x, y, [0.0, 10.0]; slope = Monotonic())
    for q in [0.3, 1.2, 4.7, 4.99]
        @test abs(cs(q) - q^2) < Δ^3
    end
    exact(z) = z^3 / 3
    rng = MersenneTwister(89192734)
    for j in 1:100
        a = rand(rng, Uniform(0, 5))
        b = rand(rng, Uniform(a, 5))
        found = Glissa.integrate(cs, a, b)
        @test abs(found - exact(b) + exact(a)) < 8*Δ^3
    end
end

@testset "Cubic spline integrates a cosine" begin
    using Glissa
    using Random
    using Distributions
    Δ = π / 30
    x = collect(0:Δ:π)
    y = cos.(x)
    cs = Glissa.cubic_spline(x, y, [0.0, 0.0]; slope = Monotonic())
    for q in [0.3, 1.2, 1.7, 2.99]
        @test abs(cs(q) - cos(q)) < Δ^3
    end
    exact(z) = sin(z)
    rng = MersenneTwister(89192734)
    for j in 1:100
        a = rand(rng, Uniform(0, 3))
        b = rand(rng, Uniform(a, 3))
        found = Glissa.integrate(cs, a, b)
        @test abs(found - exact(b) + exact(a)) < Δ^3
    end
end

@testset "Cubic spline derivative" begin
    using Glissa
    cs = CubicSpline{Int,Float64}([0, 1, 2],
        Float64[
            1 2;
            3 7;
            4 0;
            1 -1
        ]
    )
    csprime = Glissa.derivative(cs)
    @test size(csprime.c) == (3, 2)
    @test csprime.c ≈ Float64[
        3 7;
        8 0;
        3 -3
    ]
end

@testset "Cubic spline multiplication" begin
    using Glissa
    cs1 = CubicSpline{Int,Float64}([0, 1, 2],
        Float64[
            1 2;
            3 7;
            4 0
        ]
    )
    using Glissa
    cs2 = CubicSpline{Int,Float64}([0, 1, 2],
        Float64[
            2 2;
            1 -2
        ]
    )
    cs3 = cs1 * cs2
    @test cs3.c ≈ Float64[
        2 4;
        7 10;
        11 -14;
        4 0
    ]
end


multidims = zeros(Int8, 2, 3)
tupdims = zeros(Int8, (2, 3))
defaulttype = zeros((2, 3))

println(multidims)
println(tupdims)
println(defaulttype)


# Script to obtain the distribution of separations for random particles

push!(LOAD_PATH, "./models")
push!(LOAD_PATH, ".")
push!(LOAD_PATH, "./SparseInverseProblems/src")
using SparseInverseProblemsMod
using SuperResModels
using TestCases
using Utils

using PyCall
@pyimport numpy as np
@pyimport matplotlib.pyplot as plt

### Manual: This code is to generate files that contain the distribution of the
### distances of the generated particles. To have an efficient Sampling for plotting.
### To modify the interval of interest, change the lowLimit and upperLimit variables.

### Static parameters and container for use.
f_c = 20
x_max = 1.0
num_x = 10*f_c
filter_x =  ones(2*f_c+1);
model_static = SuperResModels.Fourier1d(f_c, x_max, filter_x, num_x)
### Dynamic parameters and container for use.
K = 2
tau = 1.0/(K*1.0)
# the maximum velocity on the domain is given by x_max/2 + K tau v_max = 1
v_max = 0.5
num_v = 10
model_dynamic = SuperResModels.DynamicFourier1d(model_static, v_max, tau, K, num_v*K)

# How particles are generated
minimumNumPartic = 4
maximumNumPartic = 10

test_case = () -> TestCases.cloud_1d_full(x_max, v_max, 1, 1, K, tau, rand(minimumNumPartic:maximumNumPartic))

### We generate particles and we evaluate their separation
N = 1000000
separations = zeros(N)
for i = 1:N
    if mod(i,100000)== 0
	println(i)
    end
    (thetas, weights) = test_case()
    separations[i] = Utils.separation_val(thetas,K,tau,x_max)
end

#plt.figure()
#plt.hist(separations,1000)
#plt.show()

# Spatial resolution of this approximation
M = 500
if x_max != 1.0
    warn(" Beware, the interval where we decided to consider the separations was tuned
    for the specific case of x_max equal to 1, since you modified this value, check if
        upper and lowerLimit is adequate for your use, just plot an histogram of the 		separations.")
end
upperLimit = 0.06
lowerLimit = 0.001
binss = linspace(lowerLimit,upperLimit,M)
vals = np.histogram(separations, bins = binss )
vals = vals[1]/(binss[2]-binss[1])/N

if minimum(vals) <= 0.00000001
error(" There are bin intervals without particles ")
end

plt.figure()
plt.plot(binss[1:end-1], vals)
plt.show()

np.save("data/1Dsimulations/separationDistribVal_06interval.npy", vals)
np.save("data/1Dsimulations/separationDistribBins_06interval.npy", binss)


#################################################################
#
#    1st order derivation testing for functions
#
#################################################################

tf(x) = sum(x)
@compare tf v0ref
@compare tf v1ref
@compare tf v2ref

rosenbrock(x) = (1 - x[1])^2 + 100(x[2] - x[1]^2)^2   # function to be derived
@compare rosenbrock [1., 1.]
@compare rosenbrock [0., 3.]
@compare rosenbrock [10., -10.]

function tf(x)
    a=zeros(1+3)
    for i in 1:4
        t = 4+3+2
        a[i] += b[i]*x+t
    end
    sum(a)
end
@compare tf v0ref
@compare tf 10.

function tf(x)
    a=0
    for i in 1:10
        for j in 1:10
            a += log(x) * sin(j)
        end
    end
    a
end
@compare tf 10.

function tf(x)
    a = zeros(2)
    c = 0
    for i in 1:2
        a[i] = x
        a[i] += b[i+1]
        c = c * sin(2i)
    end
    sum(a)
end
@compare tf v0ref

function tf(x)
    a = zeros(2)
    c = 0
    for i in 1:2
        d = 3
        for j in 1:2
            a[i] = x
            a[i] += b[i+1]
            d = d * sin(2i)
        end
        c += d
    end
    z = c * sum(a)
    z - sum(a)
end
@compare tf v0ref

function tf(x)
    z = 0
    for i in 1:length(x)
        z = i * x[i]
    end
    z
end
@compare tf v1ref
@compare tf v2ref


@testset "ZeroOne" begin
    @test @inferred(classify(0.2f0, LabelEnc.ZeroOne(0.3))) === 0.
    @test @inferred(classify(0.4f0, LabelEnc.ZeroOne(0.3))) === 1.
    @test classify.([0.2,0.4,0.6], LabelEnc.ZeroOne(0.3)) == [0,1,1]
    @test eltype(classify.([0.2,0.4,0.6], LabelEnc.ZeroOne(0.3))) <: Float64
    @test @inferred(classify([0.2,0.4,0.6], LabelEnc.ZeroOne(0.3))) == [0,1,1]
    @test eltype(classify([0.2,0.4,0.6], LabelEnc.ZeroOne(0.3))) <: Float64
    for T in (Float16, Float32, Float64)
        @test @inferred(classify(T(0.4), 0.5)) === zero(T)
        @test @inferred(classify(T(0.5), 0.5)) === one(T)
        @test @inferred(classify(T(0.4), 0.3)) === one(T)
        @test @inferred(classify(T(0.4), LabelEnc.ZeroOne)) === zero(T)
        @test @inferred(classify(T(0.6), LabelEnc.ZeroOne)) === one(T)
        # broadcast
        @test classify.(T[0.4,0.6], 0.5) == [0,1]
        @test classify.(T[0.4,0.6], LabelEnc.ZeroOne) == [0,1]
        @test @inferred(classify(T[0.4,0.6], 0.5)) == [0,1]
        @test @inferred(classify(T[0.4,0.6], LabelEnc.ZeroOne)) == [0,1]
        @test eltype(classify.(T[0.4,0.6], 0.5)) <: T
        @test eltype(classify.(T[0.4,0.6], LabelEnc.ZeroOne)) <: T
        @test eltype(classify(T[0.4,0.6], 0.5)) <: T
        @test eltype(classify(T[0.4,0.6], LabelEnc.ZeroOne)) <: T
        buffer = zeros(2)
        @test @inferred(classify!(buffer, [0.4,0.6], 0.5)) == [0,1]
        @test buffer == [0,1]
        buffer = zeros(2)
        @test @inferred(classify!(buffer, [0.4,0.6], LabelEnc.ZeroOne)) == [0,1]
        @test buffer == [0,1]
    end
    for T in (Float16, Float32, Float64, UInt8, Int32, Int64)
        @test @inferred(classify(0.2f0, LabelEnc.ZeroOne(T,0.3))) === zero(T)
        @test @inferred(classify(0.4f0, LabelEnc.ZeroOne(T,0.3))) === one(T)
        @test @inferred(classify(0.2f0, LabelEnc.ZeroOne(T))) === zero(T)
        @test @inferred(classify(0.4f0, LabelEnc.ZeroOne(T))) === zero(T)
        @test @inferred(classify(0.6f0, LabelEnc.ZeroOne(T))) === one(T)
        # broadcast
        @test @inferred(classify([0.4,0.6], LabelEnc.ZeroOne(T))) == [0,1]
        @test classify.([0.4,0.6], LabelEnc.ZeroOne(T)) == [0,1]
        @test classify.([0.2,0.4,0.6], LabelEnc.ZeroOne(T,0.3)) == [0,1,1]
        @test eltype(classify([0.4,0.6], LabelEnc.ZeroOne(T))) <: T
        @test eltype(classify.([0.4,0.6], LabelEnc.ZeroOne(T))) <: T
        @test eltype(classify.([0.2,0.4,0.6], LabelEnc.ZeroOne(T,0.3))) <: T
        buffer = zeros(2)
        @test @inferred(classify!(buffer, [0.4,0.6], LabelEnc.ZeroOne(T))) == [0,1]
        @test buffer == [0,1]
    end
end

@testset "MarginBased" begin
    for T in (Float16, Float32, Float64)
        @test @inferred(classify(T(0.2),  LabelEnc.MarginBased)) === one(T)
        @test @inferred(classify(T(0.6),  LabelEnc.MarginBased)) === one(T)
        @test @inferred(classify(T(0.0),  LabelEnc.MarginBased)) === one(T)
        @test @inferred(classify(T(-0.1), LabelEnc.MarginBased)) === -one(T)
        @test @inferred(classify(T(-5.1), LabelEnc.MarginBased)) === -one(T)
        @test classify.(T[0,-.1,0.2,3,-4], LabelEnc.MarginBased) == [1,-1,1,1,-1]
        @test @inferred(classify(T[0,-.1,0.2,3,-4], LabelEnc.MarginBased)) == [1,-1,1,1,-1]
        @test eltype(classify.(T[0,-.1,0.2,3,-4], LabelEnc.MarginBased)) <: T
        @test eltype(classify(T[0,-.1,0.2,3,-4], LabelEnc.MarginBased)) <: T
        buffer = zeros(5)
        @test @inferred(classify!(buffer, T[0,-.1,0.2,3,-4], LabelEnc.MarginBased)) == [1,-1,1,1,-1]
        @test buffer == [1,-1,1,1,-1]
    end
    for T in (Int16, Int32, Int64)
        @test @inferred(classify(T(12), LabelEnc.MarginBased)) === one(T)
        @test @inferred(classify(T(2),  LabelEnc.MarginBased)) === one(T)
        @test @inferred(classify(T(0),  LabelEnc.MarginBased)) === one(T)
        @test @inferred(classify(T(-1), LabelEnc.MarginBased)) === -one(T)
        @test @inferred(classify(T(-5), LabelEnc.MarginBased)) === -one(T)
        @test classify.(T[0,-1,2,3,-4], LabelEnc.MarginBased) == [1,-1,1,1,-1]
        @test @inferred(classify(T[0,-1,2,3,-4], LabelEnc.MarginBased)) == [1,-1,1,1,-1]
        @test eltype(classify.(T[0,-1,2,3,-4], LabelEnc.MarginBased)) <: T
        @test eltype(classify(T[0,-1,2,3,-4], LabelEnc.MarginBased)) <: T
    end
    for T in (Float16, Float32, Float64, Int16, Int32, Int64)
        @test @inferred(classify(12,  LabelEnc.MarginBased(T))) === one(T)
        @test @inferred(classify(2f0, LabelEnc.MarginBased(T))) === one(T)
        @test @inferred(classify(0,   LabelEnc.MarginBased(T))) === one(T)
        @test @inferred(classify(-1., LabelEnc.MarginBased(T))) === -one(T)
        @test @inferred(classify(-5,  LabelEnc.MarginBased(T))) === -one(T)
        @test classify.([0,-.1,0.2,3,-4], LabelEnc.MarginBased(T)) == [1,-1,1,1,-1]
        @test @inferred(classify([0,-.1,0.2,3,-4], LabelEnc.MarginBased(T))) == [1,-1,1,1,-1]
        @test eltype(classify.([0,-.1,0.2,3,-4], LabelEnc.MarginBased(T))) <: T
        @test eltype(classify([0,-.1,0.2,3,-4], LabelEnc.MarginBased(T))) <: T
    end
end

@testset "OneOfK" begin
    for T in (Float16, Float32, Float64)
        @test @inferred(classify(T[0.4,0.2,0.9,.1], LabelEnc.OneOfK)) === 3
        @test @inferred(classify(T[0.4,0.2,0.9,.1], LabelEnc.OneOfK{T})) === 3
        @test @inferred(classify(T[0.4,0.2,0.9,.1], LabelEnc.OneOfK(Val{4}))) === 3
    end
    for T in (Int16, Int32, Int64)
        @test @inferred(classify(T[4,2,9,10,2], LabelEnc.OneOfK)) === 4
        @test @inferred(classify(T[4,2,9,10,2], LabelEnc.OneOfK{T})) === 4
        @test @inferred(classify(T[4,2,9,10,2], LabelEnc.OneOfK(Val{4}))) === 4
    end
    for T in (Float16, Float32, Float64, Int16, Int32, Int64)
        A = T[2 1 0 2 1 0 1;
              1 0 2 1 2 1 2;
              0 2 1 0 0 2 0]
        @test @inferred(classify(A, LabelEnc.OneOfK)) == [1,3,2,1,2,3,2]
        @test @inferred(classify(A, LabelEnc.OneOfK, ObsDim.Last())) == [1,3,2,1,2,3,2]
        @test @inferred(classify(A, LabelEnc.OneOfK, obsdim=2)) == [1,3,2,1,2,3,2]
        @test @inferred(classify(A, LabelEnc.OneOfK(Val{3}))) == [1,3,2,1,2,3,2]
        @test @inferred(classify(A, LabelEnc.OneOfK(Val{3}), ObsDim.Last())) == [1,3,2,1,2,3,2]
        @test @inferred(classify(A, LabelEnc.OneOfK(Val{3}), obsdim=2)) == [1,3,2,1,2,3,2]
        @test eltype(classify(A, LabelEnc.OneOfK)) <: Int
        @test eltype(classify(A, LabelEnc.OneOfK, ObsDim.Last())) <: Int
        @test eltype(classify(A, LabelEnc.OneOfK, obsdim=2)) <: Int
        @test eltype(classify(A, LabelEnc.OneOfK(Val{3}))) <: Int
        @test eltype(classify(A, LabelEnc.OneOfK(Val{3}), ObsDim.Last())) <: Int
        @test eltype(classify(A, LabelEnc.OneOfK(Val{3}), obsdim=2)) <: Int
        buffer = zeros(7)
        @test @inferred(classify!(buffer, A, LabelEnc.OneOfK(3), ObsDim.Last())) == [1,3,2,1,2,3,2]
        @test buffer == [1,3,2,1,2,3,2]
        buffer = zeros(7)
        @test @inferred(classify!(buffer, A, LabelEnc.OneOfK(3); obsdim=:last)) == [1,3,2,1,2,3,2]
        @test buffer == [1,3,2,1,2,3,2]
        At = A'
        @test @inferred(classify(At, LabelEnc.OneOfK, ObsDim.First())) == [1,3,2,1,2,3,2]
        @test @inferred(classify(At, LabelEnc.OneOfK, obsdim=1)) == [1,3,2,1,2,3,2]
        @test @inferred(classify(At, LabelEnc.OneOfK(Val{3}), ObsDim.First())) == [1,3,2,1,2,3,2]
        @test @inferred(classify(At, LabelEnc.OneOfK(Val{3}), obsdim=1)) == [1,3,2,1,2,3,2]
        @test eltype(classify(At, LabelEnc.OneOfK, ObsDim.First())) <: Int
        @test eltype(classify(At, LabelEnc.OneOfK, obsdim=1)) <: Int
        @test eltype(classify(At, LabelEnc.OneOfK(Val{3}))) <: Int
        @test eltype(classify(At, LabelEnc.OneOfK(Val{3}), ObsDim.First())) <: Int
        @test eltype(classify(At, LabelEnc.OneOfK(Val{3}), obsdim=1)) <: Int
        buffer = zeros(7)
        @test @inferred(classify!(buffer, At, LabelEnc.OneOfK(3), ObsDim.First())) == [1,3,2,1,2,3,2]
        @test buffer == [1,3,2,1,2,3,2]
        buffer = zeros(7)
        @test @inferred(classify!(buffer, At, LabelEnc.OneOfK, obsdim=1)) == [1,3,2,1,2,3,2]
        @test buffer == [1,3,2,1,2,3,2]
    end
end


using Gadfly

set_default_plot_size(6inch, 3inch)

plot(x=[0.236033, 0.346517, 0.312707, 0.00790928, 0.488613, 0.210968, 0.951916, 0.999905, 0.251662, 0.986666],
     y=[0.555751, 0.437108, 0.424718, 0.773223, 0.28119, 0.209472, 0.251379, 0.0203749, 0.287702, 0.859512],
     color=[0.0770269, 0.685074, 0.988958, 0.2822, 0.824017, 0.69052, 0.077905, 0.953607, 0.0857399, 0.581857],
     Scale.color_asinh)


using LoopVectorization
using MuladdMacro

export jacobi, jacobi_series, jacobi_sum, jacobi_weight
"""
    jacobi_weight(α, β, x)

Compute the weighting function (1-x)^α * (1-x)^β .
"""
function jacobi_weight(α, β, x)
	return (1 - x)^α * (1 - x)^β
end

# calcac_startb computes the a and c terms of the recurrence relation for the
# jacobi polynomials.  It also computes three pieces of b which do not depend on
# x.  This is not a formal docstring because documeter complains if I document
# the function and do not export it.

function calcac_startb(n, α, β)
    a = (2n) * (n + α + β) * (2n + α + β - 2)
    c = 2 * (n + α - 1) * (n + β - 1) * (2n + α + β)
    b1 = (2n + α + β - 1)
    b2 = (2n + α + β)
    b2 = b2 * (b2 - 2)
    b3 = α^2 - β^2
    return a, c, b1, b2, b3
end

"""
    jacobi(n, α, β, x)

Compute the Jacobi polynomial of order n and weights α,β at point x.

This function uses a recurrence relation and is numerical stable to very high
order.  The computation time is linear w.r.t. the order n.  jacobi(n, a, b, x)
is also linear w.r.t. the size of argument x.  x should be passed as an array,
and jacobi should not be called with dot syntax, for best performance.

See also: [`jacobi_series`](@ref), [`jacobi_sum`](@ref)
"""
function jacobi(n, α, β, x)
    # the body of this function has largely been copy/pasted to jacobi_sequence
    # and jacobi_sum.  Changes should be propagated.
	if n == 0
        return ones(size(x))
    elseif n == 1
        return (α + 1) .+ (α + β + 2) .* ((x .- 1) ./ 2)
    end
	Pnm1 = (α + 1) .+ (α + β + 2) .* ((x .- 1) ./ 2)
	a, c, b1, b2, b3 = calcac_startb(2, α, β)
	inva = 1 / a
	# .- c .* Pnm2, but Pnm2 == 1, so drop it
	Pn = similar(Pnm1)
    Pn .= @muladd (b1 .* (b2 .* x .+ b3) .* Pnm1 .- c) .* inva
    if n == 2
        return Pn
    end
	Pnm2 = similar(Pn)
	for i = 3:n
		# assign Pnm2 to Pn to reuse a buffer
		Pnm2, Pnm1, Pn = Pnm1, Pn, Pnm2
		a, c, b1, b2, b3 = calcac_startb(i, α, β)
        inva = 1 / a
		@avx unroll = 4 for i in eachindex(Pn)
			Pn[i] = (b1 * (b2 * x[i] + b3) * Pnm1[i] - c * Pnm2[i]) * inva
		end
    end
	return Pn
end

"""
    jacobi_series(ns, α, β, x)

Compute a series of Jacobi polynomials of orders n.  Returns an array with shape
(size(x)..., length(ns)).  That is, the _final_ dimension contains the modes
and the first dimension(s) are spatial.

See also: [`jacobi`](@ref), [`jacobi_sum`](@ref)
"""
function jacobi_series(ns, α, β, x)
    # the body of this function is largely copied from jacobi.  Changes to
    # either should be propagated.
    # allocate the output buffer
    out = Array{eltype(x), ndims(x)+1}(undef, size(x)..., length(ns))
    plane_idx = ndims(x)+1
    lowest_plane = 1
    if ns[1] == 0
        # set the first element of out to ones, if the first n the user wants
        # is n=0
        fill!(selectdim(out, plane_idx, lowest_plane), 1.)
        lowest_plane += 1
    end

    if lowest_plane > length(ns)
        return out
    end
    # seed the recurrence relation
	Pnm1 = (α + 1) .+ (α + β + 2) .* ((x .- 1) ./ 2)
	a, c, b1, b2, b3 = calcac_startb(2, α, β)
	inva = 1 / a

    # unlike the regular jacobi function, we do not need to allocate
    # a scratch buffer, since we can use the lowest plane of the output
    # bufffer with impunity.  This only works because there is no weighting
    # so the lowest plane is guaranteed to contain a "pure" Jacobi polynomial
    # the exception is one buffer for Pnm2 to start up since we don't know if
    # the user opted to keep it

    # .- c .* Pnm2, but Pnm2 == 1, so drop it
    Pn = @muladd (b1 .* (b2 .* x .+ b3) .* Pnm1 .- c) .* inva
    Pnm2 = similar(Pn)

    if ns[lowest_plane] == 1
        view = selectdim(out, plane_idx, lowest_plane)
        view .= Pnm1
        lowest_plane += 1
    end
    if lowest_plane > length(ns)
        return out
    end
	if ns[lowest_plane] == 2
        view = selectdim(out, plane_idx, lowest_plane)
        view .= Pn
        lowest_plane += 1
    end
    if lowest_plane > length(ns)
        return out
    end

    # now do the loop that avoids recurrence
	for i = 3:ns[end]
        # Pnm2, Pnm1, Pn are purposfully not part of out
		Pnm2, Pnm1, Pn = Pnm1, Pn, Pnm2
		a, c, b1, b2, b3 = calcac_startb(i, α, β)
        inva = 1 / a
		@avx unroll = 4 for i in eachindex(Pn)
			Pn[i] = (b1 * (b2 * x[i] + b3) * Pnm1[i] - c * Pnm2[i]) * inva
        end
        for j in eachindex(ns)
			if ns[j] == i
                # this order is relevant to the sum
                view = selectdim(out, plane_idx, lowest_plane)
                view .= Pn
                lowest_plane += 1
                break
			end
		end
    end
	return out
end

"""
    jacobi_sum(ns, weights, α, β, x)

Compute a sum of Jacobi polynomial of order n weighted by weights.

See also: [`jacobi`](@ref), [`jacobi_series`](@ref)
"""
function jacobi_sum(ns, weights, α, β, x)
    # the body of this function is largely copied from jacobi.  Changes to
    # either should be propagated.

	# start by allocating the output array, be a little clever
	# around the initialization value
	out = similar(x)
	if ns[1] == 0
		fill!(out, weights[1])
	else
		fill!(out, 0)
	end

	# Pnm1 = P_1, Pnm2= P_0, == 1
	# this block of setup is the same as before
	Pnm1 = (α + 1) .+ (α + β + 2) .* ((x.-1)./2)
	a, c, b1, b2, b3 = calcac_startb(2, α, β)
	inva = 1 / a
	# .- c .* Pnm2, but Pnm2 == 1, so drop it
    Pn = @muladd (b1 .* (b2 .* x .+ b3) .* Pnm1 .- c) .* inva

	# add P_2 if we can.  Two iffs is faster than
	# a searchsorted
	if ns[1] == 1
		out .+= weights[1] .* Pnm1
	elseif ns[2] == 1
		out .+= weights[2] .* Pnm1
	end
	if ns[1] == 2
		out .+= weights[1] .* Pn
	elseif ns[2] == 2
		out .+= weights[2] .* Pn
	end

	Pnm2 = similar(Pn);
	for i = 3:ns[end]
		# assign Pnm2 to Pn to reuse a buffer
		Pnm2, Pnm1, Pn = Pnm1, Pn, Pnm2
		a, c, b1, b2, b3 = calcac_startb(i, α, β)
        inva = 1 / a
		@avx unroll=4 for i in eachindex(Pn)
			Pn[i] = (b1 * (b2 * x[i] + b3) * Pnm1[i] - c * Pnm2[i]) * inva
		end
		# i is the order that was just computed
		for j in eachindex(ns)
			if ns[j] == i
				# this order is relevant to the sum
				out .+= weights[j] .* Pn
				break
			end
		end
    end
	return out
end


load(::Day{6}, path) = parse.(Int, split(readline(path), ","))

# Part 1 & 2 use the same code
function count_lanternfish(init_fish, days)
    # Keep Julia's 1-based indexing in mind: counts[i] = fish with i-1 days left to live
    counts = map(i -> count(==(i), init_fish), 0:8)
    for _ in 1:days
        born = counts[1]
        counts[1:(end - 1)] = counts[2:end]
        counts[7] += born
        counts[9] = born
    end
    return sum(counts)
end

solve1(::Day{6}, data) = count_lanternfish(data, 80)
solve2(::Day{6}, data) = count_lanternfish(data, 256)

testresult1(::Day{6}) = 5934
testresult2(::Day{6}) = 26984457539


using PredicateComposition
using Test

@testset "PredicateComposition.jl" begin
    tstdata = [(1:i).^2 for i=1:10]
    @test filter(length ≺ 3, tstdata) == [[1],[1,4]]
    @test filter((length ⪰ 3) ⩓ (length ⪯ 5), tstdata) == [[1,4,9],[1,4,9,16],[1,4,9,16,25]]

    counteven(A) = count(iseven,A)
    countodd(A) = count(isodd,A)
    @test filter(MIN(length,counteven) ≺ 3, tstdata) == [[1],[1,4],[1,4,9],[1,4,9,16],[1,4,9,16,25]]
    @test filter(MAX(countodd,counteven) ≺ 3, tstdata) == [[1],[1,4],[1,4,9],[1,4,9,16]]

    @test map(SUM(length,counteven), tstdata) == [1, 3, 4, 6, 7, 9, 10, 12, 13, 15]
    @test SUM(length,counteven).(tstdata) == [1, 3, 4, 6, 7, 9, 10, 12, 13, 15]

    @test filter(length ≣ counteven, tstdata) == []
    @test filter(length ≺ counteven, tstdata) == []
    @test filter(length ⪯ counteven, tstdata) == []
    @test filter(length ≻ counteven, tstdata) == tstdata
    @test filter(length ⪰ counteven, tstdata) == tstdata

    @test filter(length ≣ countodd, tstdata) == [[1]]
    @test filter(length ≺ countodd, tstdata) == []
    @test filter(length ⪯ countodd, tstdata) == [[1]]
    @test filter(length ≻ countodd, tstdata) == tstdata[2:end]
    @test filter(length ⪰ countodd, tstdata) == tstdata

    @test count(iseven ⩔ isodd, 1:100) == 100
    @test count(!iseven ⩓ !isodd, 1:100) == 0
    @test count(iscntrl ⩔ isnumeric, Char.(0:100)) == 42
    @test count(iscntrl ⩔ isnumeric, Char.(0:10000)) == 598
    @test count(isletter ⩔ !isascii, Char.(0:100)) == 30
    @test count(isletter ⩔ !isascii, Char.(0:10000)) == 9925
end


using Unwrap
using Test

@time @testset "Unwrap 1D" begin include("test_unwrap_1D.jl") end
@time @testset "Unwrap 2D" begin include("test_unwrap_2D.jl") end
@time @testset "Unwrap 3D" begin include("test_unwrap_3D.jl") end


using Distributions
using MLBase

"Generate a correlation matrix of a given size and density of network edges."
function randcor(K, netDensity)
    IC = diagm(abs(randn(K)))
    for i in 1:K, j in 1:i-1
        IC[i,j] = rand() < netDensity ? randn()-1 : 0.0
        IC[j,i] = IC[i,j]
    end
    mineval = minimum(eig(IC)[1])
    if mineval < 0.01
        IC -= eye(K)*mineval*1.1
    end
    C = inv(IC)

    Base.cov2cor!(C, sqrt(diag(C)))
end

function upper(X::AbstractMatrix)
    x = Float64[]
    for i in 1:size(X)[1], j in i+1:size(X)[2]
        push!(x, X[i,j])
    end
    x
end
function area_under_pr(truth::AbstractVector, predictor::AbstractVector; resolution=4000)
    rocData = MLBase.roc(round(Int64, truth), float(invperm(sortperm(predictor))), resolution)
    vals = collect(map(x->(recall(x), -precision(x)), rocData))
    sort!(vals)
    xvals = map(x->x[1], vals)
    yvals = map(x->-x[2], vals)
    # println(xvals)
    # println(yvals)
    area_under_curve([0.0; xvals], [yvals[1]; yvals]) # make sure we extend all the way to zero
end
function area_under_curve(x, y) # must be sorted by increasing x
    area = 0.0
    lastVal = NaN
    for i in 2:length(x)
        v = (y[i-1]+y[i])/2 * (x[i]-x[i-1])
        if !isnan(v)
            area += v
            lastVal = v
        elseif !isnan(lastVal)
            area += lastVal
        end
    end
    area
end
function area_under_pr(truth::AbstractMatrix, predictor::AbstractMatrix)
    area_under_pr(abs(upper(truth)) .> 0.01, abs(upper(predictor)))
end


using Test
using Random
using LinearAlgebra
using Combinatorics
using TensorKit
using TensorKit: ProductSector, fusiontensor
using TensorOperations
TensorOperations.disable_cache() # avoids memory overflow during CI?
using TupleTools
using TupleTools: StaticLength
using Base.Iterators: take, product

Random.seed!(1234)

smallset(::Type{G}) where {G<:Sector} = take(values(G), 6)
function smallset(::Type{ProductSector{Tuple{G1,G2}}}) where {G1,G2}
    iter = product(smallset(G1),smallset(G2))
    s = collect(i × j for (i,j) in iter if dim(i)*dim(j) <= 6)
    return length(s) > 6 ? rand(s, 6) : s
end
function smallset(::Type{ProductSector{Tuple{G1,G2,G3}}}) where {G1,G2,G3}
    iter = product(smallset(G1),smallset(G2),smallset(G3))
    s = collect(i × j × k for (i,j,k) in iter if dim(i)*dim(j)*dim(k) <= 6)
    return length(s) > 6 ? rand(s, 6) : s
end
function randsector(::Type{G}) where {G<:Sector}
    s = collect(smallset(G))
    a = rand(s)
    while a == one(a) # don't use trivial label
        a = rand(s)
    end
    return a
end
function hasfusiontensor(G::Type{<:Sector})
    try
        fusiontensor(one(G), one(G), one(G))
        return true
    catch
        return false
    end
end

include("sectors.jl")
include("fusiontrees.jl")
include("spaces.jl")
include("tensors.jl")


doc"""
    Categorical(p)
A *Categorical distribution* is parameterized by a probability vector `p` (of length `K`).
$P(X = k) = p[k]  \quad \text{for } k = 1, 2, \ldots, K.$
```julia
Categorical(p)   # Categorical distribution with probability vector p
params(d)        # Get the parameters, i.e. (p,)
probs(d)         # Get the probability vector, i.e. p
ncategories(d)   # Get the number of categories, i.e. K
```
Here, `p` must be a real vector, of which all components are nonnegative and sum to one.
**Note:** The input vector `p` is directly used as a field of the constructed distribution, without being copied.
External links:
* [Categorical distribution on Wikipedia](http://en.wikipedia.org/wiki/Categorical_distribution)
"""

immutable Categorical{T<:Real} <: DiscreteUnivariateDistribution
    K::Int
    p::Vector{T}

    Categorical(p::Vector{T}, ::NoArgCheck) = new(length(p), p)

    function Categorical(p::Vector{T})
        @check_args(Categorical, isprobvec(p))
        new(length(p), p)
    end

    function Categorical(k::Integer)
        @check_args(Categorical, k >= 1)
        new(k, fill(1/k, k))
    end
end

Categorical{T<:Real}(p::Vector{T}, ::NoArgCheck) = Categorical{T}(p, NoArgCheck())
Categorical{T<:Real}(p::Vector{T}) = Categorical{T}(p)
Categorical(k::Integer) = Categorical{Float64}(k)

@distr_support Categorical 1 d.K

### Conversions

convert{T<:Real, S<:Real}(::Type{Categorical{T}}, p::Vector{S}) = Categorical(Vector{T}(p))
convert{T<:Real, S<:Real}(::Type{Categorical{T}}, d::Categorical{S}) = Categorical(Vector{T}(d.p))

### Parameters

ncategories(d::Categorical) = d.K
probs(d::Categorical) = d.p
params(d::Categorical) = (d.p,)
@inline partype{T<:Real}(d::Categorical{T}) = T

### Statistics

function categorical_mean{T<:Real}(p::AbstractArray{T})
    k = length(p)
    s = zero(T)
    for i = 1:k
        @inbounds s += p[i] * i
    end
    s
end

mean(d::Categorical) = categorical_mean(d.p)

function median{T<:Real}(d::Categorical{T})
    k = ncategories(d)
    p = probs(d)
    cp = zero(T)
    i = 0
    while cp < 1/2 && i <= k
        i += 1
        @inbounds cp += p[i]
    end
    i
end

function var{T<:Real}(d::Categorical{T})
    k = ncategories(d)
    p = probs(d)
    m = categorical_mean(p)
    s = zero(T)
    for i = 1:k
        @inbounds s += abs2(i - m) * p[i]
    end
    s
end

function skewness{T<:Real}(d::Categorical{T})
    k = ncategories(d)
    p = probs(d)
    m = categorical_mean(p)
    s = zero(T)
    for i = 1:k
        @inbounds s += (i - m)^3 * p[i]
    end
    v = var(d)
    s / (v * sqrt(v))
end

function kurtosis{T<:Real}(d::Categorical{T})
    k = ncategories(d)
    p = probs(d)
    m = categorical_mean(p)
    s = zero(T)
    for i = 1:k
        @inbounds s += (i - m)^4 * p[i]
    end
    s / abs2(var(d)) - 3
end

entropy(d::Categorical) = entropy(d.p)

function mgf{T<:Real}(d::Categorical{T}, t::Real)
    k = ncategories(d)
    p = probs(d)
    s = zero(T)
    for i = 1:k
        @inbounds s += p[i] * exp(t)
    end
    s
end

function cf{T<:Real}(d::Categorical{T}, t::Real)
    k = ncategories(d)
    p = probs(d)
    s = zero(T) + zero(T)*im
    for i = 1:k
        @inbounds s += p[i] * cis(t)
    end
    s
end

mode(d::Categorical) = indmax(probs(d))

function modes(d::Categorical)
    K = ncategories(d)
    p = probs(d)
    maxp = maximum(p)
    r = Array(Int, 0)
    for k = 1:K
        @inbounds if p[k] == maxp
            push!(r, k)
        end
    end
    r
end

### Evaluation

function cdf{T<:Real}(d::Categorical{T}, x::Int)
    k = ncategories(d)
    p = probs(d)
    x < 1 && return zero(T)
    x >= k && return one(T)
    c = p[1]
    for i = 2:x
        @inbounds c += p[i]
    end
    return c
end

pdf{T<:Real}(d::Categorical{T}, x::Int) = insupport(d, x) ? d.p[x] : zero(T)

logpdf(d::Categorical, x::Int) = insupport(d, x) ? log(d.p[x]) : -Inf

pdf(d::Categorical) = copy(d.p)

function _pdf!{T<:Real}(r::AbstractArray, d::Categorical{T}, rgn::UnitRange)
    vfirst = round(Int, first(rgn))
    vlast = round(Int, last(rgn))
    vl = max(vfirst, 1)
    vr = min(vlast, d.K)
    p = probs(d)
    if vl > vfirst
        for i = 1:(vl - vfirst)
            r[i] = zero(T)
        end
    end
    fm1 = vfirst - 1
    for v = vl:vr
        r[v - fm1] = p[v]
    end
    if vr < vlast
        for i = (vr - vfirst + 2):length(rgn)
            r[i] = zero(T)
        end
    end
    return r
end

function quantile(d::Categorical, p::Float64)
    0 <= p <= 1 || throw(DomainError())
    k = ncategories(d)
    pv = probs(d)
    i = 1
    v = pv[1]
    while v < p && i < k
        i += 1
        @inbounds v += pv[i]
    end
    i
end

# sampling

sampler(d::Categorical) = AliasTable(d.p)

### sufficient statistics

immutable CategoricalStats <: SufficientStats
    h::Vector{Float64}
end

function add_categorical_counts!{T<:Integer}(h::Vector{Float64}, x::AbstractArray{T})
    for i = 1 : length(x)
        @inbounds xi = x[i]
        h[xi] += 1.   # cannot use @inbounds, as no guarantee that x[i] is in bound
    end
    h
end

function add_categorical_counts!{T<:Integer}(h::Vector{Float64}, x::AbstractArray{T}, w::AbstractArray{Float64})
    n = length(x)
    if n != length(w)
        throw(ArgumentError("Inconsistent array lengths."))
    end
    for i = 1 : n
        @inbounds xi = x[i]
        @inbounds wi = w[i]
        h[xi] += wi   # cannot use @inbounds, as no guarantee that x[i] is in bound
    end
    h
end

function suffstats{T<:Integer}(::Type{Categorical}, k::Int, x::AbstractArray{T})
    CategoricalStats(add_categorical_counts!(zeros(k), x))
end

function suffstats{T<:Integer}(::Type{Categorical}, k::Int, x::AbstractArray{T}, w::AbstractArray{Float64})
    CategoricalStats(add_categorical_counts!(zeros(k), x, w))
end

typealias CategoricalData Tuple{Int, AbstractArray}

suffstats(::Type{Categorical}, data::CategoricalData) = suffstats(Categorical, data...)
suffstats(::Type{Categorical}, data::CategoricalData, w::AbstractArray{Float64}) = suffstats(Categorical, data..., w)

# Model fitting

function fit_mle(::Type{Categorical}, ss::CategoricalStats)
    Categorical(pnormalize!(ss.h))
end

function fit_mle{T<:Integer}(::Type{Categorical}, k::Integer, x::AbstractArray{T})
    Categorical(pnormalize!(add_categorical_counts!(zeros(k), x)), NoArgCheck())
end

function fit_mle{T<:Integer}(::Type{Categorical}, k::Integer, x::AbstractArray{T}, w::AbstractArray{Float64})
    Categorical(pnormalize!(add_categorical_counts!(zeros(k), x, w)), NoArgCheck())
end

fit_mle(::Type{Categorical}, data::CategoricalData) = fit_mle(Categorical, data...)
fit_mle(::Type{Categorical}, data::CategoricalData, w::AbstractArray{Float64}) = fit_mle(Categorical, data..., w)

fit_mle{T<:Integer}(::Type{Categorical}, x::AbstractArray{T}) = fit_mle(Categorical, maximum(x), x)
fit_mle{T<:Integer}(::Type{Categorical}, x::AbstractArray{T}, w::AbstractArray{Float64}) = fit_mle(Categorical, maximum(x), x, w)

fit(::Type{Categorical}, data::CategoricalData) = fit_mle(Categorical, data)
fit(::Type{Categorical}, data::CategoricalData, w::AbstractArray{Float64}) = fit_mle(Categorical, data, w)


__precompile__()
module Random
using Compat #xor
export ParisiRapuano, randperm!, getRNG,
        randchoice, randchoice!, randint,
        DiscreteDistribution

using Base.Random
import Base: rand, srand

include("parisi_rapuano.jl")
include("discrete_distribution.jl")

# Knuth shuffles (see Wikipedia)
"""
    randperm!([rng,] v::AbstractArray)

In-place random permutation of the elements of `v`.
"""
function randperm!(rng::AbstractRNG, v::AbstractArray)
    n = length(v)
    for i=1:n-1
        r = rand(rng, 0:n-i)
        v[i], v[i+r] = v[i+r], v[i]
    end
end

randperm!(v::AbstractArray) = randperm!(GLOBAL_RNG, v)

"""
randchoice!([rng,] a, k; exclude=nothing)

Sample `k` elements from array `a` without repetition and eventually excluding elements in `exclude`.
Pay attention, it has complecity `O(length(a))`, since it creates a copy of `a`
internally.
Use [`randchoice!`](@ref) if you want a `O(k)` sampler.
"""
randchoice(rng::AbstractRNG, a::AbstractArray{T}, k::Integer; exclude = T[]) where {T} = randchoice!(rng, collect(a), k; exclude = exclude)
randchoice(a::AbstractArray{T}, k::Integer; exclude = T[]) where {T} = randchoice(GLOBAL_RNG, a, k; exclude = exclude)

"""
randchoice!([rng,] a, k; exclude=nothing)

Sample `k` elements from array `a` without repetition and eventually excluding elements in `exclude`.
Pay attention, it changes the order of the elements in `a`.
Use [`randchoice`](@ref) for the non-mutating version.
"""
function randchoice!(rng::AbstractRNG, a::AbstractArray{T}, k::Integer; exclude::Vector{T} = T[]) where T
    length(a) < k + length(exclude) && error("Array too short.")
    res = Vector{eltype(a)}()
    sizehint!(res, k)
    n = length(a)
    i = 1
    while length(res) < k
        r = rand(rng, 1:n-i+1)
        if !(a[r] in exclude)
            push!(res, a[r])
            a[r], a[n-i+1] = a[n-i+1], a[r]
            i += 1
        end
    end
    return res
end

randchoice!(a::AbstractArray{T}, k::Integer; exclude = T[]) where {T} = randchoice!(GLOBAL_RNG, a, k; exclude = exclude)

"""
    randchoice([rng,] itr)

Returns a random element of `itr`. This implementation
is more efficient than the one in Base for `rand(itr)`.
"""
randchoice(rng, itr) = nth(itr, _randchoice(rng, length(itr)))
randchoice(itr) = randchoice(GLOBAL_RNG, itr)

_randchoice(rng, n::Int) = floor(Int, rand(rng) * n) + 1

#
#    nth(xs, n::Integer)
#
#Return the n'th element of xs. Mostly useful for non indexable collections.
function nth(xs, n::Integer)
    n > 0 || throw(BoundsError(xs, n))
    # catch, if possible
    applicable(length, xs) && (n ≤ length(xs) || throw(BoundsError(xs, n)))
    s = start(xs)
    i = 0
    while !done(xs, s)
        (val, s) = next(xs, s)
        i += 1
        i == n && return val
    end
    # catch iterators with no length but actual finite size less then n
    throw(BoundsError(xs, n))
end

nth(xs::AbstractArray, n::Integer) = xs[n]

"""
    getRNG(seed)

Returns `MersenneTwister(seed)` if `seed` is greater than `0`, `GLOBAL_RNG`
otherwise.
"""
getRNG(seed::Integer) = seed > 0 ? MersenneTwister(seed) : GLOBAL_RNG

end #submodule


## This script makes the example graphs we'll use in subsequent studies

##
using Plots
using GraphRecipes
using NearestNeighbors
using Random
using SparseArrays
using LinearAlgebra
##
""" This function is weird because it only gets half of the knn, but it makes
a useful picture, so I'm keeping this bug. DO NOT USE FOR ANTYHING REAL. """
function gnk(n,k)
  xy = rand(2,n)
  T = KDTree(xy)
  idxs = knn(T, xy, k)[1]
  # form the edges for sparse
  ei = Int[]
  ej = Int[]
  for i=1:n
    for j=idxs[i]
      if i > j
        push!(ei,i)
        push!(ej,j)
      end
    end
  end
  return xy, ei, ej
end
#Random.seed!(3)
#xy, ei, ej = gnk(25,5)
Random.seed!(1)
xy, ei, ej = gnk(15,4)
gr(size=(300,300))
graphplot(ei, ej, x =xy[1,:], y=xy[2,:],
  markercolor=:black, markerstrokecolor=:white,
  markersize=4, linecolor=1, linealpha=0.8, linewidth=0.7,
  axis_buffer=0.02, background=nothing)

##
for i=1:size(xy,2)
  annotate!(xy[1,i],xy[2,i],"$i")
end
gui()
## We are going to remove edge 14,15 for simplicity
n = size(xy,2)
A = sparse(ei,ej,1,n,n)
A = A + A' # make it symmetric
A[14,15] = 0
A[15,14] = 0
A[6,7] = 1
A[7,6] = 1
dropzeros!(A)
## Remove vertex 4 as it's disconnected
verts = collect(1:15)
deleteat!(verts,4)
A = A[verts,verts]
xy = xy[:,verts]
##
ei,ej = findnz(triu(A))[1:2]
graphplot(ei, ej, x =xy[1,:], y=xy[2,:],
  markercolor=:black, markerstrokecolor=:white,
  markersize=4, linecolor=1, linealpha=0.8, linewidth=0.7,
  axis_buffer=0.02, background=nothing)
##
using DelimitedFiles
writedlm("undir-graph-14.xy", xy')
writedlm("undir-graph-14.edges", [ei ej])
## Save a few figures
pyplot(size=(300,300))
p = graphplot(ei, ej, x =xy[1,:], y=xy[2,:],
  markercolor=:white, markerstrokecolor=:black,
  markersize=8, linecolor=3, linealpha=1.0, linewidth=1.0,
  axis_buffer=0.02, background=:transparent,dpi=300)
savefig("undir-graph-14.pdf")

## Save a few figures
pyplot(size=(300,300))
p = graphplot(ei, ej, x =xy[1,:], y=xy[2,:],
  markercolor=:white, markerstrokecolor=:black,
  markersize=14, linecolor=3, linealpha=1.0, linewidth=1.0,
  axis_buffer=0.02, background=:transparent,dpi=300)
for i=1:size(A,1)
  annotate!(xy[1,i],xy[2,i],text("$i",9,:black))
end
savefig("undir-graph-14-labels.pdf")


a = create_var(1, 1)
b = eye(3)
c = a + b
test_expr(c, (a,))

p = b + c
test_expr(p, (a,))

x = create_var(3, 4)
d = a - x
test_expr(d, (a, x))

f = x + a
test_expr(f, (a, x))

g = x - a
test_expr(g, (a, x))

info("All add/subtract tests passed")


# ------------------------------------------------------------------------------

PairIntTuples = Tuple{Tuple{Vararg{Int}},Tuple{Vararg{Int}}}

function det_minor_expansion_inner(
        m::MatElem{<: T},
        discarded::PairIntTuples,
        cache::Dict{PairIntTuples,T}
    ) where T <: RingElem
    n = size(m, 1);
    if length(discarded[1]) == n
        return 1;
    end
    if discarded in keys(cache)
        return cache[discarded];
    end
    result = 0;
    row = minimum(setdiff(Set(1:n), Set(discarded[1])));
    dis_rows = Tuple(sort([[i for i in discarded[1]]; row]));
    sign = 1;
    for col = 1:n
        if !(col in discarded[2])
            dis_cols = Tuple(sort([[i for i in discarded[2]]; col]));
            result += sign * m[row, col] * det_minor_expansion_inner(m, (dis_rows, dis_cols), cache);
            sign = -1 * sign;
        end
    end
    if length(discarded[1]) > 1
        cache[discarded] = result;
    end
    if length(discarded[1]) < 3
        @debug "Discarded: $discarded; $(Dates.now())"
        flush(stdout)
    end
    return result;
end

function det_minor_expansion(m::MatElem{T}) where T <: RingElem
    cache = Dict{PairIntTuples,T}();
    return det_minor_expansion_inner(m, (Tuple{}(), Tuple{}()), cache);
end

# ------------------------------------------------------------------------------

"""
    Bezout_matrix(f, g, var_elim)

Compute the Bezout matrix of two polynomials f, g with respect to var_elim

Inputs:
- `f` - first polynomial
- `g` - second polynomial
- `var_elim` - variable, of which f and g are considered as polynomials

Output:
- `M::MatrixElem` - The Bezout matrix
"""
function Bezout_matrix(f::P, g::P, var_elim::P) where P <: MPolyElem
    parent_ring = parent(f)
    deg_f = Nemo.degree(f, var_elim)
    deg_g = Nemo.degree(g, var_elim)
    n = max(deg_f, deg_g)
    coeffs_f = [coeff(f, [var_elim], [i]) for i in 0:n]
    coeffs_g = [coeff(g, [var_elim], [i]) for i in 0:n]
    GL = AbstractAlgebra.MatrixSpace(parent_ring, n, n)
    M = zero(GL)
    for i in 1:n
        for j in 1:n
            M[i, j] = sum(coeffs_f[j + k + 1] * coeffs_g[i - k] - coeffs_g[j + k + 1] * coeffs_f[i - k] for k in 0:min(i - 1, n - j))
        end
    end
    return M
end

# ------------------------------------------------------------------------------

"""
    Sylvester_matrix(f, g, var_elim)

Compute the Sylvester matrix of two polynomials f, g with respect to var_elim
Inputs:
- `f` - first polynomial
- `g` - second polynomial
- `var_elim` - variable, of which f and g are considered as polynomials
Output:
- `M::MatrixElem` - The Sylvester matrix
"""
function Sylvester_matrix(f::P, g::P, var_elim::P) where P <: MPolyElem
    parent_ring = parent(f)
    deg_f = Nemo.degree(f, var_elim)
    deg_g = Nemo.degree(g, var_elim)
    n = deg_f + deg_g
    GL = AbstractAlgebra.MatrixSpace(parent_ring, n, n)
    M = zero(GL)
    for i in 1:deg_f
        for j in 0:deg_g
            M[i, j + i] = coeff(g, [var_elim], [j])
        end
    end
    for i in 1:deg_g
        for j in 0:deg_f
            M[i + deg_f, j + i] = coeff(f, [var_elim], [j])
        end
    end

    return M
end

# ------------------------------------------------------------------------------

"""
    simplify_matrix(M)

Eliminate GCD of entries of every row and column

Input:
- `M::MatrixElem` - matrix to be simplified

Output:
- `M::MatrixElem` - Simplified matrix
- `extra_factors::Vector{AbstractAlgebra.MPolyElem}` - array of GCDs eliminated from `M`.
"""
function simplify_matrix(M::MatElem{P}) where P <: MPolyElem
    """
    An auxiliary function taking a list of coordinates of cells
    and dividing them by their gcd.
    Returns the gcd
    """
    function _simplify_range(coords::Array{Tuple{Int,Int},1})
        gcd_temp = M[coords[1]...]
        for c in coords[2:end]
            gcd_temp = gcd(gcd_temp, M[c...])
        end
        if gcd_temp != 1
            for c in coords
                M[c...] = divexact(M[c...], gcd_temp)
            end
        end
        return gcd_temp
    end

    extra_factors = Array{P,1}()
    rows_cols = Array{Array{Tuple{Int,Int},1},1}()
    # adding all rows
    for i in 1:nrows(M)
        push!(rows_cols, [(i, j) for j in 1:ncols(M)])
    end
    # adding all columns
    for j in 1:ncols(M)
        push!(rows_cols, [(i, j) for i in 1:nrows(M)])
    end
    # eliminating factors
    for r in rows_cols
        gcd_temp = _simplify_range(r)
        if gcd_temp != 1
            push!(extra_factors, gcd_temp)
        end
    end

    return M, extra_factors
end

# ------------------------------------------------------------------------------

# Definition: for an irreducible polynomial P in n variables, we will call
# an iterator Julia object *generic point generator* if it generates
# an infinite random sequence of points in C^n such that
#    1. P vanishes at each of the points
#    2. for every other irreducible polynomial, the probability of
#    vanishing at each of these points is zero
#
# Each returned point is dictionary from the variables to the values

# ------------------------------------------------------------------------------

abstract type PointGenerator{P} end

mutable struct ODEPointGenerator{P} <: PointGenerator{P}
    ode::ODE{P}
    big_ring::MPolyRing
    precision::Int
    cached_points::Array{Dict{P,<: FieldElem},1}
    number_type::Type

    function ODEPointGenerator{P}(ode::ODE{P}, big_ring::MPolyRing) where P <: MPolyElem
        prec = length(ode.x_vars) + 1
        number_type = typeof(one(base_ring(big_ring)))
        return new(ode, big_ring, prec, Array{Dict{P,number_type}}[], number_type)
    end
end

# ------------------------------------------------------------------------------

function Base.iterate(gpg::ODEPointGenerator{P}, i::Int=1) where P <: MPolyElem{<: FieldElem}
    if i > length(gpg.cached_points)
        @debug "Generating new point on the variety"
        sample_max = i * 50
        result = undef
        while true
            @debug "Preparing initial condition"
            flush(stdout)
            base_field = base_ring(gpg.big_ring)
            param_values = Dict{P,Int}(p => rand(1:sample_max) for p in gpg.ode.parameters)
            initial_conditions = Dict{P,Int}(x => rand(1:sample_max) for x in gpg.ode.x_vars)
            input_values = Dict{P,Array{Int,1}}(u => [rand(1:sample_max) for _ in 1:gpg.precision] for u in gpg.ode.u_vars)
            @debug "Computing a power series solution"
            flush(stdout)
            ps_solution = undef
            try
                ps_solution = power_series_solution(gpg.ode, param_values, initial_conditions, input_values, gpg.precision)
            catch e
                @debug "$e"
                flush(stdout)
                continue
            end

            @debug "Constructing the point"
            flush(stdout)
            result = Dict{P,gpg.number_type}(switch_ring(p, gpg.big_ring) => base_field(c) for (p, c) in param_values)
            for u in gpg.ode.u_vars
                for i in 0:(gpg.precision - 1)
                    result[str_to_var(var_to_str(u) * "_$i", gpg.big_ring)] = coeff(ps_solution[u], i) * factorial(i)
                end
            end
            for y in gpg.ode.y_vars
                for j in 0:(gpg.precision - 1)
                    result[str_to_var(var_to_str(y) * "_$j", gpg.big_ring)] = coeff(ps_solution[y], j) * factorial(j)
                end
            end
            for x in gpg.ode.x_vars
                result[switch_ring(x, gpg.big_ring)] = coeff(ps_solution[x], 0)
                result[str_to_var(var_to_str(x) * "_dot", gpg.big_ring)] = coeff(ps_solution[x], 1)
            end
            break
        end
        push!(gpg.cached_points, result)
    end
    return (gpg.cached_points[i], i + 1)
end

# ------------------------------------------------------------------------------

"""
    choose(polys, generic_point_generator)

Input:
- `polys` - an array of distinct irreducible polynomials in the same ring
- `generic_point_generator` - a generic point generator as described above for one of polys

Output:
- the polynomial that vanishes at the `generic_point_generator`
"""
function choose(polys::Array{P,1}, generic_point_generator) where P <: MPolyElem{<: FieldElem}
    vars = gens(parent(polys[1]))
    for p in generic_point_generator
        if length(polys) <= 1
            break
        end
        # get accounts for the fact that the big ring may contain some auxiliary variables, e.g. rand_proj_var
        point = [get(p, v, zero(base_ring(parent(polys[1])))) for v in vars]
        polys = filter(e -> (evaluate(e, point) == 0), polys)
        flush(stdout)
    end
    return polys[1]
end

# ------------------------------------------------------------------------------

"""
    eliminate_var(f, g, var_elim, generic_point_generator)

Eliminate variable from a pair of polynomials

Input:
- `f` and `g` - polynomials
- `var_elim` - variable to be eliminated
- `generic_point_generator` - a generic point generator object for the factor
      of the resultant of `f` and `g` of interest

Output:
- `polynomial` - the desired factor of the resultant of `f` and `g`
"""
function eliminate_var(f::P, g::P, var_elim::P, generic_point_generator) where P <: MPolyElem{<: FieldElem}
    # Linear comb
    while f != 0 && g != 0
        if Nemo.degree(f, var_elim) > Nemo.degree(g, var_elim)
            f, g = g, f
        end
        lf = coeff(f, [var_elim], [Nemo.degree(f, var_elim)])
        lg = coeff(g, [var_elim], [Nemo.degree(g, var_elim)])
        (flag, q) = divides(lg, lf)
        if flag
            @debug "\t Decreasing degree with linear combination $(Dates.now())"
            flush(stdout)
            g = g - q * f * var_elim^(Nemo.degree(g, var_elim) - Nemo.degree(f, var_elim))
        elseif (Nemo.degree(g, var_elim) == Nemo.degree(f, var_elim))
            (flag, q) = divides(lf, lg)
            if flag
                @debug "\t Decreasing degree with linear combination $(Dates.now())"
                flush(stdout)
                f = f - q * g
            else
                break
            end
        else
            break
        end
    end

    if f == 0 || g == 0
        return f + g
    end

    # Case (f(v^d), g(v^d)):
    list_deg = []
    for d in 0:Nemo.degree(f, var_elim)
        if coeff(f, [var_elim], [d]) != 0
            push!(list_deg, d)
        end
    end
    for d in 0:Nemo.degree(g, var_elim)
        if coeff(g, [var_elim], [d]) != 0
            push!(list_deg, d)
        end
    end
    gcd_deg = list_deg[1]
    for ele in list_deg[2:end]
        gcd_deg = gcd(gcd_deg, ele)
    end
    if gcd_deg > 1
        f = sum(coeff(f, [var_elim], [gcd_deg * i]) * (var_elim^i) for i in 0:(Nemo.degree(f, var_elim) ÷ gcd_deg))
        g = sum(coeff(g, [var_elim], [gcd_deg * i]) * (var_elim^i) for i in 0:(Nemo.degree(g, var_elim) ÷ gcd_deg))
    end

    resultant = undef
    matrix_factors = []
    if Nemo.degree(f, var_elim) == 0
        resultant = f
    else
        if Nemo.degree(f, var_elim) > 1
            @debug "Calculating Bezout Matrix $(Dates.now())"
            flush(stdout)
            M = Bezout_matrix(f, g, var_elim)
        else
            @debug "Calculating Sylvester matrix"
            flush(stdout)
            M = Sylvester_matrix(f, g, var_elim)
        end
        @debug "Simplifying the matrix $(Dates.now())"
        flush(stdout)
        M_simp, matrix_factors = simplify_matrix(M)
        @debug "Removed factors $(map(length, matrix_factors))"
        M_size = zero(Nemo.MatrixSpace(Nemo.ZZ, ncols(M_simp), ncols(M_simp)))
        for i in 1:ncols(M_simp)
            for j in 1:ncols(M_simp)
                M_size[i,j] = length(M_simp[i,j])
            end
        end
        @debug "\t Matrix size: \n $M_size"
        @debug "\t Computing determinant $(Dates.now())"
        flush(stdout)
        resultant = det_minor_expansion(M_simp)
    end
    # Step 4: Eliminate extra factors
    factors = fast_factor(resultant)
    for mfac in matrix_factors
        for fac in fast_factor(mfac)
            if !(fac in factors)
                push!(factors, fac)
            end
        end
    end
    res = choose(factors, generic_point_generator)
    for f in factors
        if f != res
            @debug "\t \t Size of extra factor: $(length(f)); $(Dates.now())"
            @debug "\t \t It is $f"
            flush(stdout)
        end
    end
    return res
end


"""
    outputType(s, e::Env)

Determine the output type for the sentential term `s` in the environment `e`
containing the type information. If `s` is not properly typed return nothing.
"""
outputType(s::Variable, e::Env) = get(e.userTypes, s.name, nothing)[1]
function outputType(s::Union{FunctionTerm,PredicateTerm}, e::Env)
  if isa(s, PredicateTerm) && s.name.name == "=" # treat equality PredicateTerm
    oTypes = outputType.(s.args, Ref(e))
    all(oTypes .== Ref(oTypes[1])) ? BoolType() : nothing
  else
    if haskey(e, s)
      ts = e[s]    # the type of s
      if length(ts) == length(s.args)+1 && all(outputType.(s.args, Ref(e)) .== ts[1:end-1])
        ts[end]    # args are well-typed so return the output type
      else
        nothing    # args are not well-typed so return nothing
      end
    else           # can't find type in environment
      nothing
    end
  end
end
outputType(s::Union{NegationTerm,QuantifierTerm}, e::Env) =
  outputType(s.scope, e) == BoolType() ? BoolType() : nothing
outputType(j::JunctionTerm, e::Env) =
  all(outputType.(j.juncts, Ref(e)) .== Ref(BoolType())) ? BoolType() : nothing

"""
    isWellTyped(s, e::Env)

`true/false` depending on whether `s` is well-typed/ill-typed in environment
`e`.
"""
isWellTyped(s::SententialTerm, e::Env) = outputType(s, e) == BoolType()
isWellTyped(f::FunctionTerm, e::Env) =
  haskey(e,f) ? outputType(f, e) == e[f][end] : false


using Test, EllZeroTrendFiltering

t = range(0, stop=4π, length=50)
g = sin

@time ℓ = compute_transition_costs(g, t);
cost_last = QuadraticPolynomial(0.0, 0.0, 0.0)

@time Λ =  pwq_dp_constrained(ℓ, cost_last, 7);

I_sols = [      [[1, 50]],
                [[1, 34, 50], [1, 17, 50]],
                [[1, 21, 30, 50]],
                [[1, 8, 19, 30, 50], [1, 21, 32, 43, 50]],
                [[1, 8, 19, 32, 43, 50]]
         ]

f_costs = [5.328255648628215, 4.8783936257642315, 1.7402065022125042, 0.9196880458290417, 0.09174442455649423]

@testset "Optimal Fit m=$m" for m = 1:5
    I, _, f = recover_solution(Λ[m, 1], ℓ, cost_last)

    @test I ∈ I_sols[m]
    @test f ≈ f_costs[m]    atol = 1e-8
end

##
@testset "Brute force m=$m" for m = 1:4
    I_bf, _, f_bf = brute_force_search(ℓ, cost_last, m);

    @test I_bf ∈  I_sols[m]
    @test f_bf ≈ f_costs[m]    atol = 1e-8
end

##

@testset "Regularize ζ=$ζ" for ζ in 0.1:0.1:2

    Λ_reg = pwq_dp_regularized(ℓ, cost_last, ζ)

    # recover_optimal_index_set returns the cost inclusive the regularization penality,
    # revober optimal solution does not do so. It is arguably more interesting
    # to test cost including regularization.
    I, _, f_reg = recover_optimal_index_set(Λ_reg[1])

    # Use the costs above to find out how many segments the solution should contain
    m_expected = argmin([f_costs[m] + m*ζ for m=1:5])

    @test m_expected == length(I) - 1
    @test I ∈ I_sols[m_expected]
    @test f_reg ≈ f_costs[m_expected] + ζ*m_expected
end