"""
    svdsolve(A::AbstractMatrix, [howmany = 1, which = :LR, T = eltype(A)]; kwargs...)
    svdsolve(f, m::Int, n::Int, [howmany = 1, which = :LR, T = Float64]; kwargs...)
    svdsolve(f, x₀, y₀, [howmany = 1, which = :LM]; kwargs...)
    svdsolve(f, x₀, y₀, howmany, which, algorithm)

Compute `howmany` singular values from the linear map encoded in the matrix `A` or by the
function `f`. Return singular values, left and right singular vectors and a
`ConvergenceInfo` structure.

### Arguments:
The linear map can be an `AbstractMatrix` (dense or sparse) or a general function or callable
object. Since both the action of the linear map and its adjoint are required in order to compute
singular values, `f` can either be a tuple of two callable objects (each accepting a single argument),
representing the linear map and its adjoint respectively, or, `f` can be a single callable object
that accepts two input arguments, where the second argument is a flag that indicates whether
the adjoint or the normal action of the linear map needs to be computed. The latter form still
combines well with the `do` block syntax of Julia, as in
```julia
vals, lvecs, rvecs, info = svdsolve(x₀, y₀, howmany, which; kwargs...) do (x, flag)
    if flag
        # y = compute action of adjoint map on x
    else
        # y = compute action of linear map on x
    end
    return y
end
```

For a general linear map encoded using either the tuple or the two-argument form, the best approach
is to provide a start vector `x₀` (in the domain of the linear map). Alternatively, one can specify
the number `n` of columns of the linear map, in which case `x₀ = rand(T, n)` is used, where the
default value of `T` is `Float64`, unless specified differently. If an `AbstractMatrix` is used,
a starting vector `x₀` does not need to be provided; it is chosen as `rand(T, size(A,1))`.

The next arguments are optional, but should typically be specified. `howmany` specifies how many
singular values and vectors should be computed; `which` specifies which singular values should
be targetted. Valid specifications of `which` are
*   `LR`: largest singular values
*   `SR`: smallest singular values
However, the largest singular values tend to converge more rapidly.

### Return values:
The return value is always of the form `vals, lvecs, rvecs, info = svdsolve(...)` with
*   `vals`: a `Vector{<:Real}` containing the singular values, of length at least `howmany`,
    but could be longer if more singular values were converged at the same cost.
*   `lvecs`: a `Vector` of corresponding left singular vectors, of the same length as `vals`.
*   `rvecs`: a `Vector` of corresponding right singular vectors, of the same length as `vals`.
    Note that singular vectors are not returned as a matrix, as the linear map could act on any
    custom Julia type with vector like behavior, i.e. the elements of the lists `lvecs` (`rvecs`)
    are objects that are typically similar to the starting guess `y₀` (`x₀`), up to a possibly
    different `eltype`. When the linear map is a simple `AbstractMatrix`, `lvecs` and `rvecs`
    will be `Vector{Vector{<:Number}}`.
*   `info`: an object of type [`ConvergenceInfo`], which has the following fields
    -   `info.converged::Int`: indicates how many singular values and vectors were actually
        converged to the specified tolerance `tol` (see below under keyword arguments)
    -   `info.residual::Vector`: a list of the same length as `vals` containing the residuals
        `info.residual[i] = A * rvecs[i] - vals[i] * lvecs[i]`.
    -   `info.normres::Vector{<:Real}`: list of the same length as `vals` containing the norm
        of the residual `info.normres[i] = norm(info.residual[i])`
    -   `info.numops::Int`: number of times the linear map was applied, i.e. number of times
        `f` was called, or a vector was multiplied with `A` or `A'`.
    -   `info.numiter::Int`: number of times the Krylov subspace was restarted (see below)
!!! warning "Check for convergence"
    No warning is printed if not all requested eigenvalues were converged, so always check
    if `info.converged >= howmany`.

### Keyword arguments:
Keyword arguments and their default values are given by:
*   `krylovdim`: the maximum dimension of the Krylov subspace that will be constructed.
    Note that the dimension of the vector space is not known or checked, e.g. `x₀` should not
    necessarily support the `Base.length` function. If you know the actual problem dimension
    is smaller than the default value, it is useful to reduce the value of `krylovdim`, though
    in principle this should be detected.
*   `tol`: the requested accuracy according to `normres` as defined above. If you work in e.g.
    single precision (`Float32`), you should definitely change the default value.
*   `maxiter`: the number of times the Krylov subspace can be rebuilt; see below for further
    details on the algorithms.
*   `orth`: the orthogonalization method to be used, see [`Orthogonalizer`](@ref)

### Algorithm
The last method, without default values and keyword arguments, is the one that is finally called,
and can also be used directly. Here the algorithm is specified, though currently only [`GKL`](@ref)
is available. `GKL` refers to the the partial Golub-Kahan-Lanczos bidiagonalization which forms
the basis for computing the approximation to the singular values. This factorization is dynamically
shrunk and expanded (i.e. thick restart) similar to the Krylov-Schur factorization for eigenvalues.
"""
function svdsolve end

function svdsolve(A::AbstractMatrix, howmany::Int = 1, which::Selector = :LR, T::Type = eltype(A); kwargs...)
    svdsolve(A, rand(T, size(A,1)), howmany, which; kwargs...)
end
function svdsolve(f, n::Int, howmany::Int = 1, which::Selector = :LR, T::Type = Float64; kwargs...)
    svdsolve(f, rand(T, n), howmany, which; kwargs...)
end
function svdsolve(f, x₀, howmany::Int = 1, which::Symbol = :LR; kwargs...)
    which == :LR || which == :SR || error("invalid specification of which singular values to target: which = $which")
    alg = GKL(; kwargs...)
    svdsolve(f, x₀, howmany, which, alg)
end

function svdsolve(A, x₀, howmany::Int, which::Symbol, alg::GKL)
    krylovdim = alg.krylovdim
    maxiter = alg.maxiter
    howmany > krylovdim && error("krylov dimension $(krylovdim) too small to compute $howmany singular values")

    ## FIRST ITERATION: setting up
    numiter = 1
    # Compute Lanczos factorization
    iter = GKLIterator(svdfun(A), x₀, alg.orth)
    fact = initialize(iter)
    numops = 2
    sizehint!(fact, krylovdim)
    β = normres(fact)
    tol::eltype(β) = alg.tol
    while length(fact) < krylovdim
        fact = expand!(iter, fact)
        numops += 2
        normres(fact) < tol && length(fact) >= howmany && break
    end

    # Process
    # allocate storage
    HH = fill(zero(eltype(fact)), krylovdim+1, krylovdim)
    PP = fill(zero(eltype(fact)), krylovdim, krylovdim)
    QQ = fill(zero(eltype(fact)), krylovdim, krylovdim)

    # initialize
    β = normres(fact)
    m = length(fact)
    P = copyto!(view(PP, 1:m, 1:m), I)
    Q = copyto!(view(QQ, 1:m, 1:m), I)
    f = view(HH, m+1, 1:m)
    B = rayleighquotient(fact) # Bidiagional (lower)

    # compute singular value decomposition
    P, S, Q = bidiagsvd!(B, P, Q)
    if which == :SR
        reversecols!(P)
        reverserows!(S)
        reverserows!(Q)
    elseif which != :LR
        error("invalid specification of which singular values to target: which = $which")
    end
    mul!(f, view(Q', m, :), β)

    converged = 0
    while converged < length(fact) && abs(f[converged+1]) < tol
        converged += 1
    end

    ## OTHER ITERATIONS: recycle
    while numiter < maxiter && converged < howmany
        numiter += 1

        # Determine how many to keep
        keep = div(3*krylovdim + 2*converged, 5) # strictly smaller than krylovdim since converged < howmany <= krylovdim, at least equal to converged

        # Update basis by applying P and Q using Householder reflections
        U = basis(fact, :U)
        basistransform!(U, view(P, :, 1:keep))
        # for j = 1:m
        #     h, ν = householder(P, j:m, j)
        #     lmul!(h, view(P, :, j+1:krylovdim))
        #     rmul!(U, h')
        # end
        V = basis(fact, :V)
        basistransform!(V, view(Q', :, 1:keep))
        # for j = 1:m
        #     h, ν = householder(Q, j, j:m)
        #     rmul!(view(Q, j+1:krylovdim, :), h)
        #     rmul!(V, h)
        # end

        # Shrink GKL factorization (no longer strictly GKL)
        r = residual(fact)
        U[keep+1] = rmul!(r, 1/normres(fact))
        H = fill!(view(HH, 1:keep+1, 1:keep), 0)
        @inbounds for j = 1:keep
            H[j,j] = S[j]
            H[keep+1,j] = f[j]
        end

        # Restore bidiagonal form in the first keep columns
        @inbounds for j = keep:-1:1
            h, ν = householder(H, j+1, 1:j, j)
            H[j+1,j] = ν
            H[j+1,1:j-1] .= 0
            rmul!(view(H, 1:j, :), h')
            rmul!(V, h')
            h, ν = householder(H, 1:j, j, j)
            H[j,j] = ν
            @inbounds H[1:j-1,j] .= 0
            lmul!(h, view(H, :, 1:j-1))
            rmul!(U, h')
        end
        @inbounds for j = 1:keep
            fact.αs[j] = H[j,j]
            fact.βs[j] = H[j+1,j]
        end
        fact = shrink!(fact, keep)

        # GKL factorization: recylce fact
        while length(fact) < krylovdim
            fact = expand!(iter, fact)
            numops += 2
            normres(fact) < tol && length(fact) >= howmany && break
        end

        # post process
        β = normres(fact)
        m = length(fact)
        P = copyto!(view(PP, 1:m, 1:m), I)
        Q = copyto!(view(QQ, 1:m, 1:m), I)
        f = view(HH, m+1, 1:m)
        B = rayleighquotient(fact) # Bidiagional (lower)

        # compute singular value decomposition
        P, S, Q = bidiagsvd!(B, P, Q)
        if which == :SR
            reversecols!(P)
            reverserows!(S)
            reverserows!(Q)
        elseif which != :LR
            error("invalid specification of which singular values to target: which = $which")
        end
        mul!(f, view(Q', m, :), β)
        converged = 0
        while converged < length(fact) && abs(f[converged+1]) < tol
            converged += 1
        end
    end

    if converged > howmany
        howmany = converged
    end
    values = S[1:howmany]

    # Compute schur vectors
    Pv = view(P,:, 1:howmany)
    Qv = view(Q, 1:howmany, :)

    # Compute convergence information
    leftvectors = let U = basis(fact, :U)
        [U*v for v in cols(Pv)]
    end
    rightvectors = let V = basis(fact, :V)
        [V*v for v in cols(Qv')]
    end
    residuals = let r = residual(fact)
        [r*last(v) for v in cols(Qv')]
    end
    normresiduals = let f = f
        map(i->abs(f[i]), 1:howmany)
    end

    return values, leftvectors, rightvectors, ConvergenceInfo(converged, residuals, normresiduals, numiter, numops)
end

svdfun(A::AbstractMatrix) = (x,flag) -> flag ? A'*x : A*x
svdfun((f,fadjoint)::Tuple{Any,Any}) = (x,flag) -> flag ? fadjoint(x) : f(x)
svdfun(f) = f


using ConfigEnv
using AlphaVantage
using DataFrames, CSV
using Dates
using Indicators
using Setfield
using Plots
using BenchmarkTools

dotenv();

########################################
# Downloading data
########################################
df = digital_currency_daily("BTC", datatype = "csv", parser = x -> CSV.File(x.body)) |> DataFrame

df = filter(:timestamp => x -> Date(2019, 1, 1) <= x <= Date(2020, 1, 1), df)
price = df."close (USD)"

sma1 = sma(price; n = 10)
sma2 = sma(price; n = 20)

########################################
# Converting raw data to signal generation
########################################

struct MAH{T}
    ma::T
    mah::T
end

function build_mah(v)
    T = eltype(v)
    res = Vector{MAH{T}}(undef, length(v))
    @inbounds res[1] = MAH(v[1], T(NaN))
    @inbounds for i in 2:length(v)
        res[i] = MAH(v[i], v[i - 1])
    end
    return res
end

sma1 = build_mah(sma1)
sma2 = build_mah(sma2)

function cross_signal(mah1, mah2)
    (mah1.mah < mah2.mah) & (mah1.ma > mah2.ma) && return 1
    (mah1.mah > mah2.mah) & (mah1.ma < mah2.ma) && return -1
    return 0
end

########################################
# Broker and single simulation run
########################################

struct Broker
    cash::Float64
    shares::Int
end

function run(broker, price, sma1, sma2)
    @inbounds for i in eachindex(sma1)
        signal = cross_signal(sma1[i], sma2[i])
        if signal > 0
            p = price[i]
            shares = floor(Int, broker.cash/p)
            @set! broker.shares = shares
            @set! broker.cash -= shares*p
        elseif signal < 0
            broker.shares > 0 || continue # no shorting
            p = price[i]
            @set! broker.cash += broker.shares*p
            @set! broker.shares = 0
        end
    end

    # Cash out
    p = price[end]
    @set! broker.cash += broker.shares*p
    @set! broker.shares = 0

    return broker
end

broker = Broker(1_000_000, 0)

@btime run($broker, $price, $sma1, $sma2)
  # 352.389 ns (0 allocations: 0 bytes)

########################################
# Grid run
########################################

smas = map(n -> sma(price; n = n), 10:100) .|> build_mah

struct SimInfo
    i::Int
    j::Int
    x::Float64
end

function run(broker, price, smas)
    res = SimInfo[]
    for i in 1:length(smas)
        for j in i+1:length(smas)
            sma1 = smas[i]
            sma2 = smas[j]
            out = run(broker, price, sma1, sma2)
            push!(res, SimInfo(i, j, out.cash))
        end
    end
    return res
end

@btime run($broker, $price, $smas);
  # 1.635 ms (12 allocations: 192.55 KiB)

########################################
# Final heatmap
########################################

res = run(broker, price, smas)

res_dual = map(x -> SimInfo(x.j, x.i, x.x), res)

res_diag = map(i -> SimInfo(i, i, broker.cash), 1:length(smas))

res = vcat(res, res_dual, res_diag)
sort!(res, by = x -> (x.i, x.j))
vals = map(x -> x.x/broker.cash - 1, reshape(res, :, length(smas)))
heatmap(vals, c=cgrad([:blue,:gray,:red]))
savefig("BTC_DMAC.png")


cd("INSERT_PATH")

using Distributions
# using Polynomials
using PolynomialRoots
using DataFrames

include("toolBox.jl")
include("polyDNRR.jl")

#########################################################
# Parameter Declaration
#########################################################

ss = 270

## used data sets
# w1, lam1 = 0.6, 1.8e04 # f_{2,1}, (ss200), dataseeds: 23,24,21,32 , etas = 3.5,4.9,6.3,7.5 (sampseeds=94,95,93,96)
# w1, lam1 = 0.6, 3e04 # f_{2,1}, (ss200), dataseed: 44 , eta = 2.7 (sampseed=99)

w1, lam1 = 0.6, 1e05 # f_{2,1}, (ss200), dataseeds: 23,24,21,32 , etas = 3.5,4.9,6.3,7.5 (sampseeds=94,95,93,96)

lam2 = 1e-02 * lam1

# map parameters and initial condition
theta, x0 = [0.05, 2.55, 0., -0.99], 1.
# theta, x0 = [1., 0., -1.71], 0.1

# corresponding deterministic orbit
xdet = zeros(ss)
xdet[1] = polyMap(theta, x0)
for i in 2:ss
  xdet[i] = polyMap(theta, xdet[i - 1])
end

# orbit contaminated with dynamical noise
# first argument is  type of noise: 0 is parametric, 1-3 nonparametric
dataSeed = 2
data = genData(3, ss, theta, x0, w1, lam1, lam2, dataSeed)
eta = sqrt(w1*1/lam1+(1-w1)*1/lam2)/std(data)

#########################################################
# GSB Reconstruction - Denoising
#########################################################

# Dynamical Noise Only

degree, k, gibbsIter, burnIn, dela, delb, wdist, gEps1, gEps2, thLow, thUp, zLow, zUp, papr, pbpr, thin, samplerSeed =
  5, ss, 150000, 50000, 1e05, 1e-03, 1e-03, 1e-03, 1e-03, -10.0, 10.0, -10.0, 10.0, 1., 1., 10, 999;

thetas, x0s, noise, clusters, ps, deltas, ys = 0, 0, 0, 0, 0, 0, 0;
gc()

filename = "/Results"
savelocation = string(pwd(), filename, "/seed$samplerSeed")
mkpath(savelocation)
writedlm(string(savelocation,"/xdet.txt"), xdet)

@time thetas = polyDNRR(data, degree, k, gibbsIter, burnIn, dela, delb, wdist, gEps1, gEps2, thLow, thUp, zLow, zUp, papr, pbpr, thin, samplerSeed);

thetas = 0;
gc()


using EmpiricalRisks
using Base.Test

## Auxiliary functions

function verify_multipred(pred::PredictionModel{1,0}, θ, X::DenseMatrix)
    n = size(X, 2)
    @test ninputs(pred, X) == n
    rr = zeros(n)
    for i = 1:n
        rr[i] = predict(pred, θ, X[:,i])
    end
    @test_approx_eq predict(pred, θ, X) rr

    buffer = zeros(n)
    @test_approx_eq predict!(buffer, pred, θ, X) rr
end

function verify_multipred(pred::PredictionModel{1,1}, θ, X::DenseMatrix)
    n = size(X, 2)
    @test ninputs(pred, X) == n
    p = length(predict(pred, θ, X[:,1]))
    rr = zeros(p, n)
    for i = 1:n
        rr[:,i] = predict(pred, θ, X[:,i])
    end
    @test_approx_eq predict(pred, θ, X) rr

    buffer = zeros(p, n)
    @test_approx_eq predict!(buffer, pred, θ, X) rr
end

## Data preparation

d = 5
k = 3
n = 12

X = randn(d, n)

## Predictors

# LinearPred

a = randn()
w = randn(d)
wa = [w; a]
b = 2.5

pred = LinearPred(d)
@test inputlen(pred) == d
@test inputsize(pred) == (d,)
@test outputlen(pred) == 1
@test outputsize(pred) == ()
@test paramlen(pred) == d
@test paramsize(pred) == (d,)
@test isvalidparam(pred, w)

for i = 1:n
    x_i = X[:,i]
    @test_approx_eq predict(pred, w, x_i) dot(w, x_i)
end
verify_multipred(pred, w, X)

# AffinePred

pred = AffinePred(d, b)
@test inputlen(pred) == d
@test inputsize(pred) == (d,)
@test outputlen(pred) == 1
@test outputsize(pred) == ()
@test paramlen(pred) == d + 1
@test paramsize(pred) == (d+1,)
@test isvalidparam(pred, wa)

for i = 1:n
    x_i = X[:,i]
    @test_approx_eq predict(pred, wa, x_i) dot(w, x_i) + a * b
end
verify_multipred(pred, wa, X)

# MvLinearPred

a = randn(k)
W = randn(k, d)
Wa = [W a]
b = 2.5

pred = MvLinearPred(d, k)
@test inputlen(pred) == d
@test inputsize(pred) == (d,)
@test outputlen(pred) == k
@test outputsize(pred) == (k,)
@test paramlen(pred) == d * k
@test paramsize(pred) == (k, d)
@test isvalidparam(pred, W)

for i = 1:n
    x_i = X[:,i]
    @test_approx_eq predict(pred, W, x_i) W * x_i
end
verify_multipred(pred, W, X)

# MvAffinePred

pred = MvAffinePred(d, k, b)
@test inputlen(pred) == d
@test inputsize(pred) == (d,)
@test outputlen(pred) == k
@test outputsize(pred) == (k,)
@test paramlen(pred) == (d+1) * k
@test paramsize(pred) == (k, d+1)
@test isvalidparam(pred, Wa)

for i = 1:n
    x_i = X[:,i]
    @test_approx_eq predict(pred, Wa, x_i)  W * x_i + a * b
end
verify_multipred(pred, Wa, X)


using Test, UncertainData

@testset "StrictlyDecreasing" begin

    # Create some uncertain data with decreasing magnitude and zero overlap between values,
    # so we're guaranteed that a strictly decreasing sequence through the dataset exists.
    N = 10
    t = [i <= N/2 ? CertainValue(float(i)) : UncertainValue(Normal, i, 1) for i = N:-1:1]
    T = UncertainIndexDataset(t)
    iv = UncertainIndexValueDataset(t, t)

    test_cs = [
        NoConstraint(),
        TruncateLowerQuantile(0.2),
        TruncateUpperQuantile(0.2),
        TruncateQuantiles(0.2, 0.8),
        TruncateMaximum(50),
        TruncateMinimum(-50),
        TruncateRange(-50, 50),
        TruncateStd(1)
    ]
    test_seqs = [
        StrictlyDecreasing(StartToEnd())
    ]

    @testset "$(test_seqs[i])" for i in 1:length(test_seqs)
        @test resample(t, StrictlyDecreasing()) isa Vector{Float64}
        @test resample(T, StrictlyDecreasing()) isa Vector{Float64}

        iv_draw = resample(iv, StrictlyDecreasing())
        @test iv_draw isa Tuple{Vector{Float64}, Vector{Float64}}

        @testset "$(test_cs[i])" for i in 1:length(test_cs)
            c = test_cs[i] # sequential + single constraint
            cs = sample(test_cs, N) # sequential + multiple constraints
            @test resample(t, StrictlyDecreasing(), c) isa Vector{Float64}
            @test resample(t, StrictlyDecreasing(), cs) isa Vector{Float64}
            @test resample(T, StrictlyDecreasing(), c) isa Vector{Float64}
            @test resample(T, StrictlyDecreasing(), cs) isa Vector{Float64}

            # Single extra constraint
            iv_draw = resample(iv, StrictlyDecreasing(), c)
            @test iv_draw isa Tuple{Vector{Float64}, Vector{Float64}}
            @test all(diff(iv_draw[1]) .< 0)

            iv_draw = resample(iv, StrictlyDecreasing(), c, c)
            @test iv_draw isa Tuple{Vector{Float64}, Vector{Float64}}
            @test all(diff(iv_draw[1]) .< 0)

            # Multiple extra constraints
            iv_draw = resample(iv, StrictlyDecreasing(), cs)
            #@test iv_draw isa Tuple{Vector{Float64}, Vector{Float64}}
            #@test all(diff(iv_draw[1]) .< 0)

            # iv_draw = resample(iv, StrictlyDecreasing(), cs, cs)
            # @test iv_draw isa Tuple{Vector{Float64}, Vector{Float64}}
            # @test all(diff(iv_draw[1]) .< 0)

            # iv_draw = resample(iv, StrictlyDecreasing(), c, cs)
            # @test iv_draw isa Tuple{Vector{Float64}, Vector{Float64}}
            # @test all(diff(iv_draw[1]) .< 0)

            # iv_draw = resample(iv, StrictlyDecreasing(), cs, c)
            # @test iv_draw isa Tuple{Vector{Float64}, Vector{Float64}}
            # @test all(diff(iv_draw[1]) .> 0)
        end
    end
end


module GP

using Distances, LinearAlgebra, Optim

# structs
struct hyperparameters
    σ::AbstractVector
    ℓ::AbstractVector
    σ_n::Float64
    ℓx::Vector{Int64}
end

struct hyperparameters_multi
    σ::AbstractVector
    ℓ::AbstractVector
    σ_n::Float64
    ℓx::Vector{Int64}
end

struct GP_obj
    X::AbstractArray
    Lxx::AbstractArray
    α::AbstractArray
    Θ::hyperparameters
end

struct GP_obj_multi
    X::AbstractArray
    Lxx::AbstractArray
    α::AbstractArray
    Θ::hyperparameters_multi
end

struct SS
    A::AbstractArray
    B::AbstractArray
    C::AbstractArray
    D::AbstractArray
end

# squared exponential kernel
function k_SE(R, σ)
    K = (σ^2) * exp.(-(1/2) * R)
    return K
end

# calculate distance matrix
function r(x₁, x₂)
    r = pairwise(SqEuclidean(), x₁, x₂, dims=2)
    return r
end

# multiply kernels
function r_SE_mul(X₁, X₂, ℓ)
    R = zeros(size(X₁,1), size(X₂,1))
    for i in 1:size(X₁,2)
        R += r(reshape(X₁[:, i], 1, size(X₁,1)), reshape(X₂[:, i], 1, size(X₂,1))) ./ (ℓ[i]^2)
    end
    return R
end

# analytical derivative of posterior mean for dimension d
function δμδxsf(GP_obj,Xs,Σs,k,d)
    return -((GP_obj.Θ.ℓ[GP_obj.Θ.ℓx[d]].^(-2))*(GP_obj.X[:,d] .- Xs[:,d]') .* k(GP_obj.X,Xs,GP_obj.Θ,d))'*GP_obj.α
end

function εf(GP_obj,Xs,Σs,k)
    δμδxs = zeros(size(Xs))
    for d in 1:size(Xs,2)
        δμδxs[:,d] = δμδxsf(GP_obj,Xs,Σs,k,d)
    end
    return [(δμδxs[i,:]' * + diagm(Σs[i,:]) * δμδxs[i,:]) for i in 1:size(Xs,1)]
end

function μf(X, xs, k, Θ, α)
    return (k(X,xs,Θ)'*α)[]
end

# calculate posterior distribution
function posterior(GP_obj, Xs, Σs, k, calc_σ)
    Kxs = k(GP_obj.X,Xs,GP_obj.Θ)
    Kss = k(Xs,Xs,GP_obj.Θ)
    μ = (Kxs'*GP_obj.α)
    # uncertainty propagation
    if calc_σ==true
        σ = sqrt.(diag(Kss) - sum((GP_obj.Lxx \ Kxs).^2, dims=1)' + εf(GP_obj,Xs,Σs,k))
    else
        σ = zeros(size(Xs,1))
    end
    return μ, σ
end

# training kernel cholesky factorization
function calcLxx(X, y, k, Θ::hyperparameters)
    # kernel
    Kxx = k(X,X,Θ) + Θ.σ_n*I
    Cxx = cholesky(Kxx)
    Lxx = Cxx.L
    Lxx_y = (Lxx \ y)
    α = Lxx' \ Lxx_y
    return Lxx, α
end

# negative log marginal likelihood (for optimization of hyperparameters)
function nlml(X, y, k, Θₒₚₜ, Θ₀::hyperparameters)
    nσ, nℓ = size(Θ₀.σ,1), size(Θ₀.ℓ,1)
    Θ = hyperparameters(Θₒₚₜ[1:nσ], Θₒₚₜ[nσ+1:nσ+nℓ], Θ₀.σ_n, Θ₀.ℓx)
    # cholesky factorization of k(X,X')
    Lxx, α = calcLxx(X, y, k, Θ)
    # log marginal likelihood
    n = size(X,1)
    logml = -(1/2) * y' * α  - sum(log.(diag(Lxx))) - (n/2) * log(2*pi)
    print(".")
    return -logml[1]
end

# greedy learning of sparse approximation
function sparse_subset(X, y, ix, k, Θ, m)
    # dimensions of input
    d = size(X,2)
    # induced points
    Xm = X[1,:]'
    ym = y[1]
    Θm = Θ
    ixm = ix[1]
    # not-yet-chosen points
    Xn = X[2:end, :]
    yn = y[2:end]

    print("[")
    for j = 1:(m-1)
        Hmax=-1e5; xmax=zeros(size(X,2)); ymax=0; imax=0; xs=zeros(size(X,2)); ys=0
        print(".")
        for i in 1:size(yn,1)
            # new test point
            xs = Xn[i, :]'
            ys = yn[i]
            Xms = vcat(Xm, xs)
            yms = vcat(ym, ys)
            # calculate the entropy
            Lxx, α = calcLxx(Xms, yms, k, Θ)
            Hs = (1/2)*sum(log.(((2*pi*exp(1))^d)*(diag(Lxx))))
            #Kxx = k(Xms,Xms,Θ) + 0.001*I
            #Hs = (1/2)*sum(log.(((2*pi*exp(1))^d)*(det(Kxx))))
            # use point that maximizes differential entropy
            if Hs>Hmax
                #println("$Hs > $Hmax")
                Hmax = Hs
                xmax = xs
                ymax = ys
                imax = i
            end
        end
        if imax!=0
            Xm = vcat(Xm, xmax)
            ym = vcat(ym, ymax)
            ixm = vcat(ixm,ix[imax])
            Xn = Xn[[i for i in 1:size(Xn,1) if i!=imax],:]
            yn = yn[[i for i in 1:size(Xn,1) if i!=imax]]
        end
    end
    println("]")
    return Xm, ym, ixm
end

# optimize hyperparameters
function optimize_hyperparameters(X, y, k, Θ₀::hyperparameters)
    Θ₀_opt = vcat(Θ₀.σ, Θ₀.ℓ)
    nσ, nℓ = length(Θ₀.σ), length(Θ₀.ℓ)
    lower = vcat(zeros(nσ), zeros(nℓ))
    upper = [Inf for i in 1:(nσ+nℓ)]

    print("[")
    results = optimize(Θ -> nlml(X, y, k, Θ, Θ₀), lower, upper, Θ₀_opt, Fminbox(LBFGS())) #LBFGS()
    println("]")
    print(results)

    Θ₁ = results.minimizer
    Θp = hyperparameters(Θ₁[1:nσ], Θ₁[nσ+1:nσ+nℓ], Θ₀.σ_n, Θ₀.ℓx)

    return Θp
end

end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ 5ca81b3a-7926-11eb-38e4-37053a873b78
begin
	using Markdown
	using InteractiveUtils
	using Images
	using Plots
	using PlutoUI
end

# ╔═╡ 8d041392-7926-11eb-308b-55a47a287cea
using JImages

# ╔═╡ f590466a-7926-11eb-18d7-3109f09d4f85

# ╔═╡ cab9623c-7926-11eb-024e-0fe16e3170d6
begin
	path = string(datadir, "/", "lena.h5")
	println("path to Lena is $path")
	lena = load_h5(path, "lena")
end

# ╔═╡ 02f670b8-7927-11eb-3d2e-8bb684cc0747
kg1 = Kernel.gaussian(1)

# ╔═╡ 6ae9cd0a-7927-11eb-256d-adfecbd5f27e
@time lena_g1 = imfilter(lena, reflect(centered(kg1)), Fill(0))

# ╔═╡ 85b1ec0a-7927-11eb-0d1e-8f132863920d
Gray.(lena)

# ╔═╡ aea0934e-7927-11eb-1025-ef5ebadbfb4a
iter = 100

# ╔═╡ b27aed28-7927-11eb-3e8a-17e5edee2584
md" Running RL with $iter iterations"

# ╔═╡ 7da73da2-7928-11eb-08cd-39c737d26622
@time lena_d, DFR = richardson_lucy(lena_g1, kg1,
                                    iterations=iter,
                                    view_window =1,
                                    fom_window=64,
                                    debug_level=quiet)

# ╔═╡ b2a9fe70-7928-11eb-1e88-d9749038d9ab
DFR

# ╔═╡ 3808c506-7933-11eb-3b76-1bcc9c18b823
plot(DFR[3:end])

# ╔═╡ a11f9b24-7928-11eb-0259-53dc3facfaae
md"DFR shows the variation of the relative difference. About 50 iterations are enough"

# ╔═╡ c23a1db0-7933-11eb-19bb-e7670600f387
Gray.(lena_d)

# ╔═╡ ccc3275e-7933-11eb-3b6d-8f88a26c5390
md"Running the fast version"

# ╔═╡ fef6fd5c-7933-11eb-2c9c-977ec2ea4a55
@time lena_d2 = frichardson_lucy(lena_g1, kg1, iterations=iter)

# ╔═╡ 962d09b2-7933-11eb-0cc0-f7f5cdf38116
md"## An alternative deconvolution"

# ╔═╡ a58bf18e-7d0c-11eb-2b5b-f52987f7c47a
md"### Gaussian kernel"

# ╔═╡ f60fab46-7d0c-11eb-1fc3-05956cd1a30b
function pprint(x)
	with_terminal() do
    	println(x)
	end
end

# ╔═╡ b7b4b8d0-7d0c-11eb-257e-8d2f1c78c9bd
begin
	G(x,y, σ=0.5) = 1 / sqrt(2πσ^2) * e^(-(x^2 + y^2) / (2σ^2) )
	G(I::CartesianIndex{2}, σ=0.5) = 1 / sqrt(2πσ^2) * e^(-(x^2 + y^2) / (2σ^2) )
end

# ╔═╡ 8c12e61c-7d0d-11eb-0955-f35ae7f18e89
zz = [[10., 30.] [20.,40.] ]

# ╔═╡ 12906938-7d1a-11eb-1560-9bb1872423e5
begin
	t(x,y) = x + y
	TT(I::CartesianIndex{2}) = t(I.I...)
end

# ╔═╡ 55d11c48-7d0f-11eb-22a3-d129c2a7e5e1
begin
	for i in CartesianIndices(zz)
          println(" cartesian index i = ", i ," value = ", zz[i])
		  println(" i.I = ", i.I... ," value = ", zz[i])
			println(" i.I = ", i.I... ," t(i.I...) = ", t(i.I...))
	end
end

# ╔═╡ 23619db8-7d1a-11eb-2144-c9766bda5ae5
begin
	for i in CartesianIndices(zz)
          println(" cartesian index i = ", i ," value = ", zz[i])
			println(" TT($i) = ", TT(i))
	end
end

# ╔═╡ cff0e784-7d27-11eb-1c8a-c73c55c769cd
l = 4 * ceil(Int, 0.5) + 1

# ╔═╡ 02791abe-7d28-11eb-00f0-a90f55eee923
w = 5÷2

# ╔═╡ 1bb6d316-7d28-11eb-18e7-fb4a5fcc031e
CartesianIndices((-w:w, -w:w))

# ╔═╡ 30c12eb4-7d28-11eb-35e0-83ae2b97ea51
bar(a,x...) = (a,b,x)

# ╔═╡ d5ba26e4-7d27-11eb-2b71-81b1385785b3

# ╔═╡ Cell order:
# ╠═5ca81b3a-7926-11eb-38e4-37053a873b78
# ╠═8d041392-7926-11eb-308b-55a47a287cea
# ╠═f590466a-7926-11eb-18d7-3109f09d4f85
# ╠═cab9623c-7926-11eb-024e-0fe16e3170d6
# ╠═02f670b8-7927-11eb-3d2e-8bb684cc0747
# ╠═6ae9cd0a-7927-11eb-256d-adfecbd5f27e
# ╠═85b1ec0a-7927-11eb-0d1e-8f132863920d
# ╠═aea0934e-7927-11eb-1025-ef5ebadbfb4a
# ╠═b27aed28-7927-11eb-3e8a-17e5edee2584
# ╠═7da73da2-7928-11eb-08cd-39c737d26622
# ╠═b2a9fe70-7928-11eb-1e88-d9749038d9ab
# ╠═3808c506-7933-11eb-3b76-1bcc9c18b823
# ╠═a11f9b24-7928-11eb-0259-53dc3facfaae
# ╠═c23a1db0-7933-11eb-19bb-e7670600f387
# ╠═ccc3275e-7933-11eb-3b6d-8f88a26c5390
# ╠═fef6fd5c-7933-11eb-2c9c-977ec2ea4a55
# ╠═962d09b2-7933-11eb-0cc0-f7f5cdf38116
# ╠═a58bf18e-7d0c-11eb-2b5b-f52987f7c47a
# ╠═f60fab46-7d0c-11eb-1fc3-05956cd1a30b
# ╠═b7b4b8d0-7d0c-11eb-257e-8d2f1c78c9bd
# ╠═8c12e61c-7d0d-11eb-0955-f35ae7f18e89
# ╠═12906938-7d1a-11eb-1560-9bb1872423e5
# ╠═55d11c48-7d0f-11eb-22a3-d129c2a7e5e1
# ╠═23619db8-7d1a-11eb-2144-c9766bda5ae5
# ╠═cff0e784-7d27-11eb-1c8a-c73c55c769cd
# ╠═02791abe-7d28-11eb-00f0-a90f55eee923
# ╠═1bb6d316-7d28-11eb-18e7-fb4a5fcc031e
# ╠═30c12eb4-7d28-11eb-35e0-83ae2b97ea51
# ╠═d5ba26e4-7d27-11eb-2b71-81b1385785b3


# Functions

import MutableArithmetics

"""
    AbstractFunction

Abstract supertype for function objects.
"""
abstract type AbstractFunction <: MutableArithmetics.AbstractMutable end

"""
    output_dimension(f::AbstractFunction)

Return 1 if `f` has a scalar output and the number of output components if `f`
has a vector output.
"""
function output_dimension end

"""
    AbstractScalarFunction

Abstract supertype for scalar-valued function objects.
"""
abstract type AbstractScalarFunction <: AbstractFunction end
output_dimension(::AbstractScalarFunction) = 1

Base.broadcastable(f::AbstractScalarFunction) = Ref(f)
Base.ndims(::Type{<:AbstractScalarFunction}) = 0
Base.ndims(::AbstractScalarFunction) = 0

"""
    AbstractVectorFunction

Abstract supertype for vector-valued function objects.
"""
abstract type AbstractVectorFunction <: AbstractFunction end

"""
    SingleVariable(variable)

The function that extracts the scalar variable referenced by `variable`, a `VariableIndex`.
This function is naturally be used for single variable bounds or integrality constraints.
"""
struct SingleVariable <: AbstractScalarFunction
    variable::VariableIndex
end

"""
    VectorOfVariables(variables)

The function that extracts the vector of variables referenced by `variables`, a `Vector{VariableIndex}`.
This function is naturally be used for constraints that apply to groups of variables, such as an "all different" constraint, an indicator constraint, or a complementarity constraint.
"""
struct VectorOfVariables <: AbstractVectorFunction
    variables::Vector{VariableIndex}
end
output_dimension(f::VectorOfVariables) = length(f.variables)

"""
    struct ScalarAffineTerm{T}
        coefficient::T
        variable_index::VariableIndex
    end

Represents ``c x_i`` where ``c`` is `coefficient` and ``x_i`` is the variable
identified by `variable_index`.
"""
struct ScalarAffineTerm{T}
    coefficient::T
    variable_index::VariableIndex
end

# Note: ScalarAffineFunction is mutable because its `constant` field is likely of an immutable
# type, while its `terms` field is of a mutable type, meaning that creating a `ScalarAffineFunction`
# allocates, and it is desirable to provide a zero-allocation option for working with
# ScalarAffineFunctions. See https://github.com/jump-dev/MathOptInterface.jl/pull/343.
"""
    ScalarAffineFunction{T}(terms, constant)

The scalar-valued affine function ``a^T x + b``, where:
* ``a`` is a sparse vector specified by a list of
  [`ScalarAffineTerm`](@ref) structs.
* ``b`` is a scalar specified by `constant::T`

Duplicate variable indices in `terms` are accepted, and the corresponding
coefficients are summed together.
"""
mutable struct ScalarAffineFunction{T} <: AbstractScalarFunction
    terms::Vector{ScalarAffineTerm{T}}
    constant::T
end

"""
    struct VectorAffineTerm{T}
        output_index::Int64
        scalar_term::ScalarAffineTerm{T}
    end

A `ScalarAffineTerm` plus its index of the output component of a
`VectorAffineFunction` or `VectorQuadraticFunction`.
`output_index` can also be interpreted as a row index into a sparse matrix,
where the `scalar_term` contains the column index and coefficient.
"""
struct VectorAffineTerm{T}
    output_index::Int64
    scalar_term::ScalarAffineTerm{T}
end

function VectorAffineTerm(
    output_index::Base.Integer,
    scalar_term::ScalarAffineTerm,
)
    return VectorAffineTerm(convert(Int64, output_index), scalar_term)
end

"""
    VectorAffineFunction{T}(terms, constants)

The vector-valued affine function ``A x + b``, where:
* ``A`` is a sparse matrix specified by a list of `VectorAffineTerm` objects.
* ``b`` is a vector specified by `constants`

Duplicate indices in the ``A`` are accepted, and the corresponding coefficients
are summed together.
"""
struct VectorAffineFunction{T} <: AbstractVectorFunction
    terms::Vector{VectorAffineTerm{T}}
    constants::Vector{T}
end
output_dimension(f::VectorAffineFunction) = length(f.constants)

"""
    struct ScalarQuadraticTerm{T}
        coefficient::T
        variable_index_1::VariableIndex
        variable_index_2::VariableIndex
    end

Represents ``c x_i x_j`` where ``c`` is `coefficient`, ``x_i`` is the variable
identified by `variable_index_1` and ``x_j`` is the variable identified by
`variable_index_2`.
"""
struct ScalarQuadraticTerm{T}
    coefficient::T
    variable_index_1::VariableIndex
    variable_index_2::VariableIndex
end

# Note: ScalarQuadraticFunction is mutable because its `constant` field is likely of an immutable
# type, while its other fields are of mutable types, meaning that creating a `ScalarQuadraticFunction`
# allocates, and it is desirable to provide a zero-allocation option for working with
# ScalarQuadraticFunctions. See https://github.com/jump-dev/MathOptInterface.jl/pull/343.
"""
    ScalarQuadraticFunction{T}(affine_terms, quadratic_terms, constant)

The scalar-valued quadratic function ``\\frac{1}{2}x^TQx + a^T x + b``, where:
* ``a`` is a sparse vector specified by a list of `ScalarAffineTerm` structs.
* ``b`` is a scalar specified by `constant`.
* ``Q`` is a symmetric matrix specified by a list of `ScalarQuadraticTerm`
  structs.

Duplicate indices in ``a`` or ``Q`` are accepted, and the corresponding
coefficients are summed together. "Mirrored" indices `(q,r)` and `(r,q)` (where
`r` and `q` are `VariableIndex`es) are considered duplicates; only one need be
specified.

For example, for two scalar variables ``y, z``, the quadratic expression
``yz + y^2`` is represented by the terms
`ScalarQuadraticTerm.([1.0, 2.0], [y, y], [z, y])`.
"""
mutable struct ScalarQuadraticFunction{T} <: AbstractScalarFunction
    affine_terms::Vector{ScalarAffineTerm{T}}
    quadratic_terms::Vector{ScalarQuadraticTerm{T}}
    constant::T
end

"""
    struct VectorQuadraticTerm{T}
        output_index::Int64
        scalar_term::ScalarQuadraticTerm{T}
    end

A [`ScalarQuadraticTerm`](@ref) plus its index of the output component of a
`VectorQuadraticFunction`. Each output component corresponds to a
distinct sparse matrix ``Q_i``.
"""
struct VectorQuadraticTerm{T}
    output_index::Int64
    scalar_term::ScalarQuadraticTerm{T}
end

function VectorQuadraticTerm(
    output_index::Base.Integer,
    scalar_term::ScalarQuadraticTerm,
)
    return VectorQuadraticTerm(convert(Int64, output_index), scalar_term)
end

"""
    VectorQuadraticFunction{T}(affine_terms, quadratic_terms, constants)

The vector-valued quadratic function with i`th` component ("output index")
defined as ``\\frac{1}{2}x^TQ_ix + a_i^T x + b_i``, where:
* ``a_i`` is a sparse vector specified by the `VectorAffineTerm`s with
  `output_index == i`.
* ``b_i`` is a scalar specified by `constants[i]`
* ``Q_i`` is a symmetric matrix specified by the `VectorQuadraticTerm` with
  `output_index == i`.

Duplicate indices in ``a_i`` or ``Q_i`` are accepted, and the corresponding
coefficients are summed together. "Mirrored" indices `(q,r)` and `(r,q)` (where
`r` and `q` are `VariableIndex`es) are considered duplicates; only one need be
specified.
"""
struct VectorQuadraticFunction{T} <: AbstractVectorFunction
    affine_terms::Vector{VectorAffineTerm{T}}
    quadratic_terms::Vector{VectorQuadraticTerm{T}}
    constants::Vector{T}
end
output_dimension(f::VectorQuadraticFunction) = length(f.constants)

# Function modifications

"""
    AbstractFunctionModification

An abstract supertype for structs which specify partial modifications to functions, to be used for making small modifications instead of replacing the functions entirely.
"""
abstract type AbstractFunctionModification end

"""
    ScalarConstantChange{T}(new_constant::T)

A struct used to request a change in the constant term of a scalar-valued function.
Applicable to `ScalarAffineFunction` and `ScalarQuadraticFunction`.
"""
struct ScalarConstantChange{T} <: AbstractFunctionModification
    new_constant::T
end

"""
    VectorConstantChange{T}(new_constant::Vector{T})

A struct used to request a change in the constant vector of a vector-valued function.
Applicable to `VectorAffineFunction` and `VectorQuadraticFunction`.
"""
struct VectorConstantChange{T} <: AbstractFunctionModification
    new_constant::Vector{T}
end

"""
    ScalarCoefficientChange{T}(variable::VariableIndex, new_coefficient::T)

A struct used to request a change in the linear coefficient of a single variable
in a scalar-valued function.
Applicable to `ScalarAffineFunction` and `ScalarQuadraticFunction`.
"""
struct ScalarCoefficientChange{T} <: AbstractFunctionModification
    variable::VariableIndex
    new_coefficient::T
end

# Note: MultiRowChange is mutable because its `variable` field of an immutable
# type, while `new_coefficients` is of a mutable type, meaning that creating a `MultiRowChange`
# allocates, and it is desirable to provide a zero-allocation option for working with
# MultiRowChanges. See https://github.com/jump-dev/MathOptInterface.jl/pull/343.
"""
    MultirowChange{T}(variable::VariableIndex, new_coefficients::Vector{Tuple{Int64, T}})

A struct used to request a change in the linear coefficients of a single
variable in a vector-valued function. New coefficients are specified by
`(output_index, coefficient)` tuples. Applicable to `VectorAffineFunction` and
`VectorQuadraticFunction`.
"""
mutable struct MultirowChange{T} <: AbstractFunctionModification
    variable::VariableIndex
    new_coefficients::Vector{Tuple{Int64,T}}
end

function MultirowChange(
    variable::VariableIndex,
    new_coefficients::Vector{Tuple{Ti,T}},
) where {Ti<:Base.Integer,T}
    return MultirowChange(
        variable,
        [(convert(Int64, i), j) for (i, j) in new_coefficients],
    )
end

# Implementation of comparison for functions
function Base.:(==)(f::VectorOfVariables, g::VectorOfVariables)
    return f.variables == g.variables
end

function Base.isapprox(
    f::Union{SingleVariable,VectorOfVariables},
    g::Union{SingleVariable,VectorOfVariables};
    kwargs...,
)
    return f == g
end

# For affine and quadratic functions, terms are compressed in a dictionary using `_dicts` and then the dictionaries are compared with `dict_compare`
function dict_compare(d1::Dict, d2::Dict{<:Any,T}, compare::Function) where {T}
    return all(kv -> compare(kv.second, Base.get(d2, kv.first, zero(T))), d1)
end

# Build a dictionary where the duplicate keys are summed
function sum_dict(kvs::Vector{Pair{K,V}}) where {K,V}
    d = Dict{K,V}()
    for kv in kvs
        key = kv.first
        d[key] = kv.second + Base.get(d, key, zero(V))
    end
    return d
end

"""
    coefficient(t::Union{ScalarAffineTerm, ScalarQuadraticTerm
                         VectorAffineTerm, VectorQuadraticTerm})

Finds the coefficient stored in the term `t`.
"""
function coefficient end

function coefficient(t::Union{ScalarAffineTerm,ScalarQuadraticTerm})
    return t.coefficient
end
function coefficient(t::Union{VectorAffineTerm,VectorQuadraticTerm})
    return t.scalar_term.coefficient
end

"""
    term_indices(t::Union{ScalarAffineTerm, ScalarQuadraticTerm,
                         VectorAffineTerm, VectorQuadraticTerm})

Returns the indices of the input term `t` as a tuple of `Int`s.

* For `t::ScalarAffineTerm`, this is a 1-tuple of the variable index.
* For `t::ScalarQuadraticTerm`, this is a 2-tuple of the variable indices
  in non-decreasing order.
* For `t::VectorAffineTerm`, this is a 2-tuple of the row/output and
  variable indices.
* For `t::VectorQuadraticTerm`, this is a 3-tuple of the row/output and
  variable indices in non-decreasing order.
"""
term_indices(t::ScalarAffineTerm) = (t.variable_index.value,)
function term_indices(t::ScalarQuadraticTerm)
    return minmax(t.variable_index_1.value, t.variable_index_2.value)
end
function term_indices(t::Union{VectorAffineTerm,VectorQuadraticTerm})
    return (t.output_index, term_indices(t.scalar_term)...)
end

"""
    term_pair(t::Union{ScalarAffineTerm, ScalarQuadraticTerm,
                       VectorAffineTerm, VectorQuadraticTerm})

Returns the pair [`term_indices`](@ref) `=>` [`coefficient`](@ref) of the term.
"""
function term_pair(
    t::Union{
        ScalarAffineTerm,
        ScalarQuadraticTerm,
        VectorAffineTerm,
        VectorQuadraticTerm,
    },
)
    return term_indices(t) => coefficient(t)
end

function _dicts(f::Union{ScalarAffineFunction,VectorAffineFunction})
    return (sum_dict(term_pair.(f.terms)),)
end

function _dicts(f::Union{ScalarQuadraticFunction,VectorQuadraticFunction})
    return (
        sum_dict(term_pair.(f.affine_terms)),
        sum_dict(term_pair.(f.quadratic_terms)),
    )
end

"""
    constant(f::Union{ScalarAffineFunction, ScalarQuadraticFunction})

Returns the constant term of the scalar function
"""
constant(f::Union{ScalarAffineFunction,ScalarQuadraticFunction}) = f.constant

"""
    constant(f::Union{VectorAffineFunction, VectorQuadraticFunction})

Returns the vector of constant terms of the vector function
"""
constant(f::Union{VectorAffineFunction,VectorQuadraticFunction}) = f.constants

function Base.isapprox(
    f::F,
    g::G;
    kwargs...,
) where {
    F<:Union{
        ScalarAffineFunction,
        ScalarQuadraticFunction,
        VectorAffineFunction,
        VectorQuadraticFunction,
    },
    G<:Union{
        ScalarAffineFunction,
        ScalarQuadraticFunction,
        VectorAffineFunction,
        VectorQuadraticFunction,
    },
}
    return isapprox(constant(f), constant(g); kwargs...) && all(
        dict_compare.(
            _dicts(f),
            _dicts(g),
            (α, β) -> isapprox(α, β; kwargs...),
        ),
    )
end

function constant(
    f::Union{ScalarAffineFunction,ScalarQuadraticFunction},
    T::DataType,
)
    return constant(f)
end
function constant(
    f::Union{VectorAffineFunction,VectorQuadraticFunction},
    T::DataType,
)
    return constant(f)
end

"""
    constant(f::SingleVariable, T::DataType)

The constant term of a `SingleVariable` function is
the zero value of the specified type `T`.
"""
constant(f::SingleVariable, T::DataType) = zero(T)

"""
    constant(f::VectorOfVariables, T::DataType)

The constant term of a `VectorOfVariables` function is a
vector of zero values of the specified type `T`.
"""
constant(f::VectorOfVariables, T::DataType) = zeros(T, length(f.variables))

# isbits type, nothing to copy
Base.copy(func::SingleVariable) = func

Base.copy(func::VectorOfVariables) = VectorOfVariables(copy(func.variables))

"""
    copy(func::Union{ScalarAffineFunction, VectorAffineFunction})

Return a new affine function with a shallow copy of the terms and constant(s)
from `func`.
"""
function Base.copy(
    func::F,
) where {F<:Union{ScalarAffineFunction,VectorAffineFunction}}
    return F(copy(func.terms), copy(constant(func)))
end

"""
    copy(func::Union{ScalarQuadraticFunction, VectorQuadraticFunction})

Return a new quadratic function with a shallow copy of the terms and constant(s)
from `func`.
"""
function Base.copy(
    func::F,
) where {F<:Union{ScalarQuadraticFunction,VectorQuadraticFunction}}
    return F(
        copy(func.affine_terms),
        copy(func.quadratic_terms),
        copy(constant(func)),
    )
end

# Define shortcuts for
# SingleVariable -> ScalarAffineFunction
function ScalarAffineFunction{T}(f::SingleVariable) where {T}
    return ScalarAffineFunction([ScalarAffineTerm(one(T), f.variable)], zero(T))
end
# VectorOfVariables -> VectorAffineFunction
function VectorAffineFunction{T}(f::VectorOfVariables) where {T}
    n = length(f.variables)
    return VectorAffineFunction(
        map(
            i -> VectorAffineTerm(i, ScalarAffineTerm(one(T), f.variables[i])),
            1:n,
        ),
        zeros(T, n),
    )
end

# Conversion between scalar functions
# Conversion to SingleVariable
function Base.convert(::Type{SingleVariable}, f::ScalarAffineFunction)
    if (
        !iszero(f.constant) ||
        !isone(length(f.terms)) ||
        !isone(f.terms[1].coefficient)
    )
        throw(InexactError(:convert, SingleVariable, f))
    end
    return SingleVariable(f.terms[1].variable_index)
end

function Base.convert(
    ::Type{SingleVariable},
    f::ScalarQuadraticFunction{T},
) where {T}
    return convert(SingleVariable, convert(ScalarAffineFunction{T}, f))
end

# Conversion to ScalarAffineFunction
function Base.convert(::Type{ScalarAffineFunction{T}}, α::T) where {T}
    return ScalarAffineFunction{T}(ScalarAffineTerm{T}[], α)
end

function Base.convert(
    ::Type{ScalarAffineFunction{T}},
    f::SingleVariable,
) where {T}
    return ScalarAffineFunction{T}(f)
end

function Base.convert(
    ::Type{ScalarAffineTerm{T}},
    t::ScalarAffineTerm{T},
) where {T}
    return t
end

function Base.convert(
    ::Type{ScalarAffineTerm{T}},
    t::ScalarAffineTerm,
) where {T}
    return ScalarAffineTerm{T}(t.coefficient, t.variable_index)
end

function Base.convert(
    ::Type{ScalarAffineFunction{T}},
    f::ScalarAffineFunction{T},
) where {T}
    return f
end

function Base.convert(
    ::Type{ScalarAffineFunction{T}},
    f::ScalarAffineFunction,
) where {T}
    return ScalarAffineFunction{T}(f.terms, f.constant)
end

function Base.convert(
    ::Type{ScalarAffineFunction{T}},
    f::ScalarQuadraticFunction{T},
) where {T}
    if !Base.isempty(f.quadratic_terms)
        throw(InexactError(:convert, ScalarAffineFunction{T}, f))
    end
    return ScalarAffineFunction{T}(f.affine_terms, f.constant)
end

# Conversion to ScalarQuadraticFunction
function Base.convert(::Type{ScalarQuadraticFunction{T}}, α::T) where {T}
    return ScalarQuadraticFunction{T}(
        ScalarAffineTerm{T}[],
        ScalarQuadraticTerm{T}[],
        α,
    )
end

function Base.convert(
    ::Type{ScalarQuadraticFunction{T}},
    f::SingleVariable,
) where {T}
    return convert(
        ScalarQuadraticFunction{T},
        convert(ScalarAffineFunction{T}, f),
    )
end

function Base.convert(
    ::Type{ScalarQuadraticFunction{T}},
    f::ScalarAffineFunction{T},
) where {T}
    return ScalarQuadraticFunction{T}(
        f.terms,
        ScalarQuadraticTerm{T}[],
        f.constant,
    )
end

function Base.convert(::Type{VectorOfVariables}, g::SingleVariable)
    return VectorOfVariables([g.variable])
end

function Base.convert(
    ::Type{VectorAffineFunction{T}},
    g::SingleVariable,
) where {T}
    return VectorAffineFunction{T}(
        [VectorAffineTerm(1, ScalarAffineTerm(one(T), g.variable))],
        [zero(T)],
    )
end

function Base.convert(
    ::Type{VectorQuadraticFunction{T}},
    g::SingleVariable,
) where {T}
    return VectorQuadraticFunction{T}(
        [VectorAffineTerm(1, ScalarAffineTerm(one(T), g.variable))],
        VectorQuadraticTerm{T}[],
        [zero(T)],
    )
end

function Base.convert(
    ::Type{VectorAffineFunction{T}},
    g::ScalarAffineFunction,
) where {T}
    return VectorAffineFunction{T}(
        VectorAffineTerm{T}[VectorAffineTerm(1, term) for term in g.terms],
        [g.constant],
    )
end

function Base.convert(
    ::Type{VectorQuadraticFunction{T}},
    g::ScalarAffineFunction,
) where {T}
    return VectorQuadraticFunction{T}(
        VectorAffineTerm{T}[VectorAffineTerm(1, term) for term in g.terms],
        VectorQuadraticTerm{T}[],
        [g.constant],
    )
end

function Base.convert(
    ::Type{VectorQuadraticFunction{T}},
    g::ScalarQuadraticFunction,
) where {T}
    return VectorQuadraticFunction{T}(
        VectorAffineTerm{T}[
            VectorAffineTerm(1, term) for term in g.affine_terms
        ],
        VectorQuadraticTerm{T}[
            VectorQuadraticTerm(1, term) for term in g.quadratic_terms
        ],
        [g.constant],
    )
end


function FVMIBL(w::Array{Float64,2}, U::Array{Float64,1}, Ut::Array{Float64,1}, Ux::Array{Float64,1})

    n = Int(length(w)/2)
    dx = Float64(1.0/n)

    # correlate the unknown values from the del and E values
    del , E, FF ,B, S, dfde = correlate(w)
    #del , E, FF ,B, S, dfde  = calc_shapes(n,w)

    fL, fR, UipL ,UipR, FFipL, FFipR, dfdeipL, dfdeipR, wipL, wipR = fluxReconstruction(w , U, FF, dfde, del, E)

    lamb1L ,lamb2L = eigenlamb(UipL, dfdeipL, FFipL, wipL)
    lamb1R ,lamb2R = eigenlamb(UipR, dfdeipR, FFipR, wipR)

    lamb1 ,lamb2 = eigenlamb(U, dfde, FF, w)

    dt = calc_Dt(lamb1 ,lamb2, 0.8, dx)

    #dt = max(dtL,dtR)
    #dt = 0.001

    # two step forward Euler methods for adding the source term to the right hand-side
    # of the transport equations.

    #z = RHSSource(U,B, del,Ut, Ux, FF, E, S)

    # step 1 : assuming this as a homogeneous equation and advanced half a step
    w1 = w .+ ((fL - fR) .* ((dt)/(dx)))

    del , E, FF ,B, S, dfde = correlate(w1)
    #del , E, FF ,B, S, dfde  = calc_shapes(n,w1)

    fL, fR, UipL ,UipR, FFipL, FFipR, dfdeipL, dfdeipR, wipL, wipR = fluxReconstruction(w1 , U, FF, dfde, del, E)

    z = RHSSource(U,B, del,Ut, Ux, FF, E, S)

    #dtL = calc_Dt(UipL, dfdeipL, FFipL, wipL, 0.8, dx)
    #dtR = calc_Dt(UipR, dfdeipR, FFipR, wipR, 0.8, dx)

    #dt = max(dtL,dtR)

    j1, j2 = sepeartionJ(lamb1, lamb2, dt, dx)

    # step 2 : by considering source terms advanced a full step using 2nd order midpoint rule
    w2 =1/2*(w1+w) .+ (dt).*z

    return w2,dt, j1, j2
end

function init(n)

    # initial conditions of IBVP
    E_init = 0.4142 * ones(n,1)
    E = E_init
    B = 131.9*E_init.^3 - 167.32*E_init.^2 + 76.642.*E_init .- 11.068
    del = sqrt.(B.*0.005);
    F= 4.8274.*E_init.^4 - 5.9816*E_init.^3 + 4.0274*E_init.^2 + 0.23247.*E_init .+ 0.15174;

    return del, E, F ,B
end

function eigenlamb(U::Array{Float64,1}, dfde::Array{Float64,1}, FF::Array{Float64,1}, w::Array{Float64,2})
    lamb1 =  U.*((dfde .- 1.0)
            + sqrt.(1.0 .+ 4.0*FF .- 2.0.* dfde .- (4.0*((w[:,2]./w[:,1]) .-1.0).*dfde) .+ dfde.^2))
    lamb2 = U.*((dfde .- 1.0)
            - sqrt.(1.0 .+ 4.0*FF .- 2.0.* dfde .- (4.0*((w[:,2]./w[:,1]) .-1.0).*dfde) .+ dfde.^2))

    return lamb1, lamb2
end

function calc_Dt(lamb1::Array{Float64,1}, lamb2::Array{Float64,1}, cfl::Float64, dx::Float64 )

    # calculate time step values based on eigenvalues

    dti = cfl.*(dx./(max.(lamb1,lamb2)))
    dt = minimum(dti)
    return dt
end

function fluxReconstruction(w::Array{Float64,2}, U::Array{Float64,1}, FF::Array{Float64,1}, dfde::Array{Float64,1}, del::Array{Float64,1} , E::Array{Float64,1})

    delF = U.*(w[:,2]- w[:,1])
    EF = U.* FF.* w[:,1]
    F = hcat(delF, EF)
    # first-order approximation of left and right side of the i+1/2
    # cell interface

    # 1) approximation of edge velocity
    UR = U
    UL = U
    UipL = UR
    UipR = [UL[2:end]; UL[1]]

    # 2) approximation of FF
    FFR = FF
    FFL = FF
    FFipL = FFR
    FFipR = [FFL[2:end]; FFL[1]]

    # 3) approximation of dfde
    dfdeR = dfde
    dfdeL = dfde
    dfdeipL = dfdeR
    dfdeipR = [dfdeL[2:end]; dfdeL[1]]

    # 4) approximation w (a two-dimensional array)
    wR = w
    wL = w
    wipL = wR
    wipR = [wL[2:end,:]; wL[1:1,:]]

    # calculating left and right flux based on calculated left and right quantities
    delFluxL = UipL.*(wipL[:,2] - wipL[:,1])
    delFluxR = UipR.*(wipR[:,2] - wipR[:,1])
    EFluxL = UipL.* FFipL.* wipL[:,1]
    EFluxR = UipR.* FFipR.* wipR[:,1]

    fpL = hcat(delFluxL, EFluxL)
    fpR = hcat(delFluxR, EFluxR)

    # calculating wave speed at the left and right interaces
    wsL = maxWaveSpeed(UipL, wipL, dfdeipL, FFipL)
    wsR = maxWaveSpeed(UipR, wipR, dfdeipR, FFipR)
    # selecting the maximum wave speed
    ws = max.(wsL+wsR)
    ww = hcat(ws,ws)

    # flux reconstruction of left and right side of the i+1/2 interface
    fR = 0.5*((fpL + fpR) + ww.* (wipL - wipR))
    fL = [fR[end:end,:];fR[1:end-1,:]]

    # specifying the boundary conditions using internal extrapolation from the calculated flux of the
    # neighbouring cell centers
    #fL[1,:] = [F[1,1]; F[1,2]]
    fL[1,:] = 0.5*((F[1,:]) - wsR[1,:].* (wipR[1,:]))
    #fL[1,:] = [0; 0]

    fR[end,:] = [F[end,1];F[end,2]]

    #fR[end,:] = [0.0;0.0]

    return fL, fR, UipL ,UipR, FFipL, FFipR, dfdeipL, dfdeipR, wipL, wipR

end

function maxWaveSpeed(Uip::Array{Float64,1}, wip::Array{Float64,2}, dfdeip::Array{Float64,1}, FFip::Array{Float64,1})

    # calculate wave speed at the interfaces of the cell
    ws = abs.(Uip).*((dfdeip .- 1.0)
        + sqrt.(1.0 .+ 4.0*FFip .- 2.0.* dfdeip .- (4.0*((wip[:,2]./wip[:,1]) .-1.0).*dfdeip) .+ dfdeip.^2))

    return ws
end

function RHSSource(U::Array{Float64,1} ,B::Array{Float64,1}, del::Array{Float64,1},Ut::Array{Float64,1}, Ux::Array{Float64,1}, FF::Array{Float64,1}, E::Array{Float64,1}, S::Array{Float64,1} )

    # the source terms of the system of equations
    z1 = B./(2.0*del) .- del.* (Ut./U) .- (E.+ 1.0).*del.*Ux
    z2 = S./del .- 2.0*E.*del.* (Ut./U) .- 2.0*FF.*del.*Ux
    z = hcat(z1,z2)

    return z

end

function calc_shapes(ncell::Int64, sol::Array{Float64})

    E = zeros(ncell); B = zeros(ncell); F = zeros(ncell); S = zeros(ncell)
    dfde = zeros(ncell); del = zeros(ncell)

    for i = 1:ncell
        del[i] = sol[i,1]
        E[i] = sol[i,2]/del[i] - 1.
        F[i] = 4.8274*E[i]^4 - 5.9816*E[i]^3 + 4.0274*E[i]^2 + 0.23247*E[i] + 0.15174
        if E[i] < -0.0616
            B[i] = -225.86*E[i]^3 - 3016.6*E[i]^2 - 208.68*E[i] - 17.915
        elseif E[i] > -0.0395
            B[i] = 131.9*E[i]^3 - 167.32*E[i]^2 + 76.642*E[i] - 11.068
        else
            B[i] = 0.5*(-225.86*E[i]^3 - 3016.6*E[i]^2 - 208.68*E[i] - 17.915
                        + 131.9*E[i]^3 - 167.32*E[i]^2 + 76.642*E[i] - 11.068)
        end
        if E[i] < -0.0582
            S[i] = 451.55*E[i]^3 + 2010*E[i]^2 + 138.96*E[i] + 11.296
        elseif E[i] > -0.042
            S[i] = -96.739*E[i]^3 + 117.74*E[i]^2 - 46.432*E[i] + 6.8074
        else
            S[i] = 0.5*(451.55*E[i]^3 + 2010*E[i]^2 + 138.96*E[i] + 11.296
                        - 96.739*E[i]^3 + 117.74*E[i]^2 - 46.432*E[i] + 6.8074)
        end
        dfde[i] = 4*4.8274*E[i]^3 - 3*5.9816*E[i]^2 + 2*4.0274*E[i] + 0.23247
    end

    return del, E, F ,B, S, dfde
end

function sepeartionJ(lamb1::Array{Float64,1}, lamb2::Array{Float64,1}, dt::Float64, dx::Float64)

    N1 = length(lamb1)
    lambj1 = sum(lamb1)/N1
    J1Sep = (dt)/(dx) .* (lamb1[1:end-1] - lamb1[2:end])./ (lambj1*N1)

    for i=2:N1-1
    if lamb1[i].<0.0
    J1Sep[i] =  (dt)/(dx) .* (lamb1[i+1] - lamb1[i-1])./ (lambj1*N1)
    end
    end

    N2 = length(lamb2)
    lambj2 = sum(lamb2)/N2
    J2Sep = (dt)/(dx) .* (lamb2[1:end-1] - lamb2[2:end])./ (lambj2*N2)

    for i=2:N2-1
    if lamb2[i].<0.0
    J2Sep[i] =  (dt)/(dx) * (lamb2[i+1] - lamb2[i-1]) / (lambj2*N2)
    end
    end

    return J1Sep, J2Sep

end


@testset "Capsule" begin
    @testset "constructors" begin
        o, extr, r = Point2f0(1, 2), Point2f0(3, 4), 5f0
        s = Capsule(o, extr, r)
        @test typeof(s) == Capsule{2,Float32}
        @test typeof(s) == Capsule2{Float32}
        @test origin(s) == o
        @test extremity(s) == extr
        @test radius(s) == r
        #@test abs(height(s)- norm([1,2]-[3,4]))<1e-5
        h = norm(o - extr)
        @test isapprox(height(s), h)
        #@test norm(direction(s) - Point{2,Float32}([2,2]./norm([1,2]-[3,4])))<1e-5
        @test isapprox(direction(s), Point2f0(2, 2) ./ h)
        v1 = rand(Point{3, Float64}); v2 = rand(Point{3, Float64}); R = rand()
        s = Capsule(v1, v2, R)
        @test typeof(s) == Capsule{3, Float64}
        @test typeof(s) == Capsule3{Float64}
        @test origin(s) == v1
        @test extremity(s) == v2
        @test radius(s) == R
        @test height(s) == norm(v2 - v1)
        #@test norm(direction(s) - Point{3,Float64}((v2-v1)./norm(v2-v1)))<1e-10
        @test isapprox(direction(s), (v2-v1) ./ norm(v2 .- v1))
      end

    @testset "decompose" begin
        o, extr, r = Point2f0(1, 2), Point2f0(3, 4), 5f0
        s = Capsule(o, extr, r)
        positions = Point{3, Float32}[
          (4.535534, -1.5355341, 0.0),
          (2.9134173, -2.619398, 0.0),
          (1.0000002, -3.0, 0.0),
          (-0.913417, -2.6193976, 0.0),
          (-2.5355341, -1.5355337, 0.0),
          (-3.6193976, 0.08658302, 0.0),
          (-4.0, 2.0000002, 0.0),
          (-3.619398, 3.9134173, 0.0),
          (-2.535534, 5.5355344, 0.0),
          (-0.53553426, 7.535534, 0.0),
          (1.0865824, 8.619398, 0.0),
          (2.9999998, 9.0, 0.0),
          (4.913417, 8.619397, 0.0),
          (6.535534, 7.535534, 0.0),
          (7.619397, 5.913417, 0.0),
          (8.0, 3.9999998, 0.0),
          (7.619398, 2.0865827, 0.0),
          (6.535534, 0.4644655, 0.0)
        ]
        @test decompose(Point3f0, s, 4) ≈ positions
        FT = Face{3, Int}
        faces = FT[
          (1, 3, 2),
          (1, 4, 3),
          (1, 5, 4),
          (1, 6, 5),
          (1, 7, 6),
          (1, 8, 7),
          (1, 9, 8),
          (1, 10, 9),
          (1, 11, 10),
          (1, 12, 11),
          (1, 13, 12),
          (1, 14, 13),
          (1, 15, 14),
          (1, 16, 15),
          (1, 17, 16),
          (1, 18, 17)
        ]
        @test faces == decompose(FT, s, 4)

        v1 = Point{3, Float64}(1,2,3); v2 = Point{3, Float64}(4,5,6); R = 5.0
        s = Capsule(v1, v2, R)
        positions = Point{3,Float64}[
            (4.535533905932738, -1.5355339059327373, 3.0),
            (3.1616954987389683, -2.3711193256429137, 1.8952880865480273),
            (1.4587585476806852, -2.5412414523193148, 0.9587585476806852),
            (-0.31402023360539255, -2.020000734336378, 0.332989516029115),
            (7.535533905932738, 1.4644660940672627, 6.0),
            (8.371119325642914, 2.8383045012610317, 7.104711913451973),
            (8.541241452319316, 4.541241452319315, 8.041241452319316),
            (8.020000734336378, 6.314020233605392, 8.667010483970884),
            (3.0412414523193148, 4.041241452319315, -1.0824829046386295),
            (1.7811492852594561, 2.7811492852594557, -1.8764343108748296),
            (0.4021342206547498, 1.4021342206547494, -1.927992798267443),
            (-0.8858611987114287, 0.11413880128857112, -1.2293090544897973),
            (6.041241452319315, 7.041241452319315, 1.9175170953613705),
            (6.990573112163402, 7.990573112163402, 3.3329895160291154),
            (7.484617125293379, 8.48461712529338, 5.1544901063711865),
            (7.448159769230341, 8.44815976923034, 7.104711913451973),
            (-2.535533905932737, 5.535533905932738, 2.9999999999999996),
            (-3.3711193256429137, 4.161695498738968, 1.8952880865480268),
            (-3.5412414523193148, 2.4587585476806852, 0.9587585476806848),
            (-3.020000734336378, 0.6859797663946074, 0.332989516029115),
            (0.46446609406726314, 8.535533905932738, 6.0),
            (1.8383045012610322, 9.371119325642914, 7.104711913451972),
            (3.5412414523193148, 9.541241452319316, 8.041241452319316),
            (5.314020233605392, 9.020000734336378, 8.667010483970884),
            (-1.0412414523193152, -0.04124145231931431, 7.0824829046386295),
            (-1.990573112163402, -0.9905731121634007, 5.667010483970884),
            (-2.4846171252933793, -1.4846171252933784, 3.845509893628814),
            (-2.4481597692303416, -1.4481597692303412, 1.8952880865480275),
            (1.9587585476806848, 2.9587585476806857, 10.08248290463863),
            (3.2188507147405434, 4.218850714740545, 10.87643431087483),
            (4.59786577934525, 5.597865779345251, 10.927992798267443),
            (5.885861198711429, 6.885861198711429, 10.229309054489796),
            (-1.8867513459481287, -0.8867513459481287, 0.11324865405187134),
            (6.886751345948129, 7.886751345948129, 8.886751345948129)
        ]
        @test decompose(Point3{Float64},s,(4, 4)) ≈ positions
        faces = Face[
          (13, 9, 1),
          (5, 13, 1),
          (4, 12, 33),
          (16, 8, 34),
          (9, 10, 1),
          (10, 2, 1),
          (14, 13, 5),
          (14, 5, 6),
          (10, 11, 2),
          (11, 3, 2),
          (15, 14, 6),
          (15, 6, 7),
          (11, 12, 3),
          (12, 4, 3),
          (16, 15, 7),
          (16, 7, 8),
          (21, 17, 9),
          (13, 21, 9),
          (12, 20, 33),
          (24, 16, 34),
          (17, 18, 9),
          (18, 10, 9),
          (22, 21, 13),
          (22, 13, 14),
          (18, 19, 10),
          (19, 11, 10),
          (23, 22, 14),
          (23, 14, 15),
          (19, 20, 11),
          (20, 12, 11),
          (24, 23, 15),
          (24, 15, 16),
          (29, 25, 17),
          (21, 29, 17),
          (20, 28, 33),
          (32, 24, 34),
          (25, 26, 17),
          (26, 18, 17),
          (30, 29, 21),
          (30, 21, 22),
          (26, 27, 18),
          (27, 19, 18),
          (31, 30, 22),
          (31, 22, 23),
          (27, 28, 19),
          (28, 20, 19),
          (32, 31, 23),
          (32, 23, 24),
          (5, 1, 25),
          (29, 5, 25),
          (28, 4, 33),
          (8, 32, 34),
          (1, 2, 25),
          (2, 26, 25),
          (6, 5, 29),
          (6, 29, 30),
          (2, 3, 26),
          (3, 27, 26),
          (7, 6, 30),
          (7, 30, 31),
          (3, 4, 27),
          (4, 28, 27),
          (8, 7, 31),
          (8, 31, 32)
        ]
        @test faces == decompose(Face{3, Int}, s, (4, 4))
        m = GLPlainMesh(s, (4, 4))
        @test m.faces == faces
        @test m.vertices ≈ positions
        m = GLNormalMesh(s)# just test that it works without explicit resolution parameter
        @test m isa GLNormalMesh

    end
end


using FileIO
using IterTools
using Makie
using MakieRipserer
using MeshIO
using Ripserer
using SparseArrays
# import Pkg; Pkg.add(url="https://github.com/mtsch/TopologicalDatasets.jl")
using TopologicalDatasets

torus = sample(Torus(; r1=3), 1000)
res = ripserer(Alpha(torus); reps=true, dim_max=2)

# barcode
bcplt = plot_barcode(res; infinity=Observable(5))
save(joinpath(@__DIR__, "assets", "barcode.png"), bcplt)

# diagram
dgplt = plot(res; infinity=Observable(5))
save(joinpath(@__DIR__, "assets", "diagram.png"), dgplt)

# cocycle
repplt = plot(torus)
plot!(repplt, res[end][end - 1].representative, torus)
save(joinpath(@__DIR__, "assets", "cocycle.png"), repplt)

# app demo

time = Observable(0.0)
torus_app = MakieRipserer.app(torus; time=time)
rng = range(0, 4.5; length=200)
rng = vcat(rng, reverse(rng[2:(end - 1)]))
record(torus_app, joinpath(@__DIR__, "assets", "torus.gif"), rng) do r
    println("time: ", r)
    time[] = r
end

# sublevel cat
cat = load(joinpath(@__DIR__, "assets", "cat.obj"))
cat_pts = cat.position
cat_index = Dict([cat_pts[i] => i for i in 1:length(cat_pts)])
dist = zeros(length(cat_pts), length(cat_pts))
for t in cat
    for (p1, p2) in IterTools.subsets(t, Val(2))
        i = cat_index[p1]
        j = cat_index[p2]
        dist[i, j] = dist[j, i] = max(p1[3], p2[3])
    end
end
for point in cat_pts
    i = cat_index[point]
    dist[i, i] = point[3]
end
cat_flt = Rips(sparse(dist))
time = Observable(-1.0)
cat_app = MakieRipserer.app(cat_pts, cat_flt; time=time, dim_max=2)
rng = range(
    minimum(Ripserer.births(cat_flt)), maximum(Ripserer.births(cat_flt)); length=200
)
rng = vcat(rng, reverse(rng[2:(end - 1)]))
record(cat_app, joinpath(@__DIR__, "assets", "cat.gif"), rng) do r
    println("time: ", r)
    time[] = r
end


# ---
# title: 1334. Find the City With the Smallest Number of Neighbors at a Threshold Distance
# id: problem1334
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Graph
# link: <https://leetcode.com/problems/find-the-city-with-the-smallest-number-of-neighbors-at-a-threshold-distance/description/>
# hidden: true
# ---
#
# There are `n` cities numbered from `0` to `n-1`. Given the array `edges` where
# `edges[i] = [fromi, toi, weighti]` represents a bidirectional and weighted
# edge between cities `fromi` and `toi`, and given the integer
# `distanceThreshold`.
#
# Return the city with the smallest number **** of cities that are reachable
# through some path and whose distance is **at most** `distanceThreshold`, If
# there are multiple such cities, return the city with the greatest number.
#
# Notice that the distance of a path connecting cities _**i**_ and _**j**_ is
# equal to the sum of the edges' weights along that path.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/01/16/find_the_city_01.png)
#
#
#
#     Input: n = 4, edges = [[0,1,3],[1,2,1],[1,3,4],[2,3,1]], distanceThreshold = 4
#     Output: 3
#     Explanation: The figure above describes the graph.
#     The neighboring cities at a distanceThreshold = 4 for each city are:
#     City 0 -> [City 1, City 2]
#     City 1 -> [City 0, City 2, City 3]
#     City 2 -> [City 0, City 1, City 3]
#     City 3 -> [City 1, City 2]
#     Cities 0 and 3 have 2 neighboring cities at a distanceThreshold = 4, but we have to return city 3 since it has the greatest number.
#
#
# **Example 2:**
#
# **![](https://assets.leetcode.com/uploads/2020/01/16/find_the_city_02.png)**
#
#
#
#     Input: n = 5, edges = [[0,1,2],[0,4,8],[1,2,3],[1,4,2],[2,3,1],[3,4,1]], distanceThreshold = 2
#     Output: 0
#     Explanation: The figure above describes the graph.
#     The neighboring cities at a distanceThreshold = 2 for each city are:
#     City 0 -> [City 1]
#     City 1 -> [City 0, City 4]
#     City 2 -> [City 3, City 4]
#     City 3 -> [City 2, City 4]
#     City 4 -> [City 1, City 2, City 3]
#     The city 0 has 1 neighboring city at a distanceThreshold = 2.
#
#
#
#
# **Constraints:**
#
#   * `2 <= n <= 100`
#   * `1 <= edges.length <= n * (n - 1) / 2`
#   * `edges[i].length == 3`
#   * `0 <= fromi < toi < n`
#   * `1 <= weighti, distanceThreshold <= 10^4`
#   * All pairs `(fromi, toi)` are distinct.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


## MEtropolis
export Metropolis
struct Metropolis <: MCMCRule
    num_switches::Int

    # Inner constructor to assert that the number of switches is always odd.
    Metropolis(i) = begin
        if iseven(i)
            @warn """
                Metropolis sampling with an even number of switches per step
                is known to be non-ergodic, because it can potentially never
                change the configuration.
                Setting the number of switches to $(i+1).
                """
            i += 1
        end
        return new(i)
    end
end

"""
    Metropolis([n_switches=1])

Metropolis rule for Markov Chain Monte Carlo sampling where at every step
`n_switches` sites are switched.

The number of switches must be odd, otherwise there might
"""
Metropolis() = Metropolis(1)

function markov_chain_step!(σ, s::MCMCSampler{Metropolis}, net::Union{MatrixNet,KetNet}, c)
    # the denominator
    logψ_σ = 2*real(net(σ))

    for i=1:s.rule.num_switches
        # 1- Flip a spin
        sites_to_flip = 1:nsites(σ)
        flipat = rand(c.rng, sites_to_flip)
        old_val, new_val =flipat!(c.rng, σ, flipat)

        # 2- Compute the probability to be in that state
        logψ_σp = 2*real(net(σ)) #log_prob_ψ(net, σ)
        log_ratio = logψ_σp - logψ_σ
        prob = exp(log_ratio)
        rv = rand(c.rng)

        # If the move is rejected, revert it
        if !(rv < prob)
            setat!(σ, flipat, old_val)
        else # else accept it
            logψ_σ = logψ_σp
            c.steps_accepted += 1
        end
    end
    c.steps_done += 1

    return true
end

function markov_chain_step!(σ::LUState, s::MCMCSampler{Metropolis}, net::Union{MatrixNet,KetNet}, c)
    # Switch num_switches times
    for i=1:s.rule.num_switches
        # 1- Flip a spin
        sites_to_flip = 1:nsites(σ)
        flipat = rand(c.rng, sites_to_flip) #1:nsites(σ))
        old_val, new_val = flipat_fast!(c.rng, σ, flipat)

        # 2- Compute the probability to be in that state
        log_ratio = 2*real(logψ_Δ(net, σ))
        prob = exp(log_ratio)
        rv = rand(c.rng)
        # If the move is rejected, revert it
        if !(rv < prob)
            clear_changes!(σ)
        else
            update_lut!(σ, net)
            c.steps_accepted += 1
        end
    end
    c.steps_done += 1

    return true
end


# Lebesgue measure

export Lebesgue

struct Lebesgue{X} <: PrimitiveMeasure end

function Pretty.tile(::Lebesgue{X}) where {X}
    result = Pretty.literal("Lebesgue(")
    result *= Pretty.tile(X)
    result *= Pretty.literal(")")
end

Lebesgue(X) = Lebesgue{X}()

sampletype(::Lebesgue{ℝ}) = Float64
sampletype(::Lebesgue{ℝ₊}) = Float64
sampletype(::Lebesgue{𝕀}) = Float64

testvalue(::Lebesgue{ℝ}) = 0.0
testvalue(::Lebesgue{𝕀}) = 0.5
testvalue(::Lebesgue{ℝ₊}) = 1.0
testvalue(::Lebesgue{<:Real}) = 0.0

logdensity(::Lebesgue, x) = zero(x)

Base.:∘(::typeof(basemeasure), ::Type{Lebesgue}) = Lebesgue


#-----------------------------------------------------------------# partitionplot
@userplot struct PartitionPlot
    args
end

@recipe function f(o::PartitionPlot; nparts=100, stat=Extrema(), by=nothing, dropmissing=false)
    t = o.args[1]
    sel = map(x -> lowerselection(t, x), o.args[2:end])
    # Create the (Indexed)Partition object
    o = if length(sel) == 1
        Partition(stat, nparts)
    elseif length(sel) == 2
        T = typeof(collect(rows(t)[1])[sel[1]])  # Type of first selection
        IndexedPartition(T, stat, nparts)
    else
        throw(ArgumentError("too many arguments for partitionplot"))
    end
    # Wrap the (Indexed)Partition in an FTSeries to possibly remove missing values
    s = FTSeries(o; filter = dropmissing ? !IndexedTables._ismissing : x -> true)
    if by === nothing
        reduce(s, t; select=sel)
    else
        grp = groupreduce(s, t, by; select=sel)
        grp = isa(t, DDataset) ? collect(grp) : grp
        for row in rows(grp)
            @series begin row[2] end
        end
    end
end

"""
    partitionplot(table, y;    stat=Extrema(), nparts=100, by=nothing, dropmissing=false)
    partitionplot(table, x, y; stat=Extrema(), nparts=100, by=nothing, dropmissing=false)

Plot a summary of variable `y` against `x` (`1:length(y)` if not specified).  Using `nparts`
approximately-equal sections along the x-axis, the data in `y` over each section is
summarized by `stat`.
"""
partitionplot


# Parameters:
#	N::Int
#	- the problem size
#	K::Int
#	- the problem is homogenous at N-K
#	P::Array{Array{Array{Int, 1}, 1}, 1}
#	- output1 from partitions_bl()
#	ZFI::Array{Int, 1}
#	- output2 from partitions_bl()
#	WI::Array{Int, 2}
#	- output3 from partitions_bl()
# Return Values:
#	PT::Array{Array{Array{Int, 1}, 1}, 1} (Partition Tree)
#	- for each value, i, in PT[n][j], P[n][j] decomposes into P[n-1][i]
#	- length(PT[n]) = 0 for n = 1:(N - K)
#	- length(PT[n][p]) = 0 for p <= ZFI[n]
function partition_tree_bl(N::Int, K::Int, P::Array{Array{Array{Int, 1}, 1}, 1}, ZFI::Array{Int, 1}, WI::Array{Int, 2})
	PT = Array(Array{Array{Int, 1}, 1}, N)
	for n = 1:(N - K)
		PT[n] = Array(Array{Int, 1}, 0)
	end
	for n = (N - K + 1):N
		Pd = P[n - 1]
		Pn = P[n]
		PTn = Array(Array{Int, 1} , length(Pn))
		for p = 1:ZFI[n]
			PTn[p] = Array(Int, 0)
		end
		for p = (ZFI[n] + 1):length(Pn)
			PDA = pda(Pn[p])
			if length(PDA) == 1 || PDA[1][1] == PDA[2][1]
				lb = 1
				if PDA[1][1] != 1
					lb = WI[n - 1, PDA[1][1] - 1] + 1
				end
				PTn[p] = dia(PDA, Pd, lb)
			else
				lb1 = 1
				if PDA[1][1] != 1
					lb1 = WI[n - 1, PDA[1][1] - 1] + 1
				end
				lb2 = WI[n - 1, PDA[2][1] - 1]
				PTn[p] = dia(PDA, Pd, lb1, lb2)
			end
		end
		PT[n] = PTn
	end
	return PT
end


const stft_docs = """
The short-time fourier transform can be thought of as a filterbank, with each
filter centered at each fft bin frequency. The window function determines the
frequency response of each bandpass filter (each filter is an FIR filter of the
window function modulated by the center frequency).

This implementation has a few differences from the `stft` implementation in
DSP.jl that I think make it more consistent with the filterbank interpretation,
and make the output easier to reason about:

1. The first window is centered at the first sample of `x`, and in general the
   nth window is centered at `(n-1)*hop`. This also ensures that no input samples
   are multiplied by zero, which would make them unrecoverable.
2. `x` is zero-padded to ensure that the center of the last window is no
   earlier than the last sample of `x`. Similarly this ensures that no information
   from `x` is lost, and also makes the width of the result easier to reason about.
   (it is only a function of the hop size, not the window size).
3. Each window is centered at index 1 before taking the FFT, avoiding a linear
   phase shift in each frame. One benefit of this is that when the same frequency
   component present in `x` is picked up by multiple filters, the phase will be
   coherent across the bands (assuming it is not demodulated).
4. It provides the `demod` option to work with demodulated filter bank outputs.
   This is useful because the phase progresses much more slowly (proportional to
   the difference between the signal frequency and the center frequency, rather
   than at the signal frequency itself).
5. An `istft2` implementation is provided. Note that `istft2` is currently not
   fully compatible with DSP.jl's `stft` function. In order to make it compatible
   we'd need to match the window offset and also perform an `fftshift`.

`stft2` and `istft2` should satisfy the property that for any real-valued signal
`x` and `n`:

    istft2(stft2(x, n), n)[1:length(x)] ≈ x

Note that the round-trip may introduce some zero-padding at the end.
"""

# TODO: support windows longer than nfft - alternatively support windows with
# padding in the middle where the window length is still nfft. They should be
# equivilent ways to do the same thing, though maybe `stft` and `istft` will need
# to handle them specially...
# TODO: add options for boundary conditions so you don't get big bursts
# in the STFT at the beginning and end for truncated signals
"""
    stft2(x, nfft, hop=nfft÷2;
          window=rect,
          demod=false)

Perform a short-time fourier transform on a real-valued time signal `x`. The
result is a complex-valued matrix of dimension:

    (nfft÷2+1, cld(length(x)-1, hop)+1)

- `nfft` is the FFT size.
- `hop` is the hop size from frame to frame. You can think of this as the
   sampling period of the continuous filterbank.
- `window` adds an analysis window. It can be an `AbstractArray`
  of length `nfft`, or a function that can be called as
  `f(nfft; zerophase=true)` to generate such a window. Note that generally the
  window function should be periodic and centered around index `1`, which is
  what the `zerophase` option in DSP.jl's window functions does. See the DSP.jl
  documentation for more details.
- `demod` should be set to `true` if the STFT phase should be demodulated (i.e
  a "sliding window" stft). This is equivilent to de-modulating the output of
  each filter by its center frequency.

$stft_docs
"""
function stft2(x, nfft, hop=nfft÷2;
               window=rect,
               demod=false)
    # we could handle this, but it's simpler not to, and this is the common case
    @assert iseven(nfft)
    @assert hop > 0

    L = length(x)
    N = nfft÷2+1
    # the first window will be centered at the first sample. The last window
    # should be centered at the last sample or later, so that COLA will apply,
    # and we don't need to worry about numerical issues because some samples
    # got scaled way down by the window function.
    M = cld(L-1, hop)+1

    X = zeros(Complex{eltype(x)}, N, M)

    win = if window isa Function
        window(nfft; zerophase=true)
    elseif window isa AbstractArray
        @assert length(window) == nfft
        window
    else
        throw(ArgumentError("window must be an Array or Function"))
    end

    xbuf = zeros(eltype(x), nfft)
    Xbuf = similar(X, N)

    p = plan_rfft(xbuf)

    # each time we copy into the buffer for the FFT, we have to copy the
    # positive part (including the window center) and the negative part
    # separately. We want the window center to be at the beginning of the fft
    # buffer, or else we add a half-buffer circular delay to each frame

    # we could simplify this (at the expense of more copying) by just copying
    # the data in one block and calling `ifftshift`.

    xidx = 1
    nfft2 = nfft÷2

    for c in 1:M
        fill!(xbuf, zero(eltype(xbuf)))
        # copy the positive half to the beginning of the fft buffer
        # note the center of the window could be past the last sample
        # of `x`, in which case `npos` should be 0.
        npos = max(0, min(nfft2, L-xidx+1))
        copyto!(xbuf, 1, x, xidx, npos)
        # copy the negative half to the end of the fft buffer
        nneg = min(nfft2, xidx-1)
        # note that if xidx is past the end of `x`, we don't actually have
        # nneg samples available to copy. it's still useful to compute the
        # indices for the copying though
        copystart = xidx-nneg
        ncopy = max(0, min(nneg, L-copystart+1))
        copyto!(xbuf, nfft-nneg+1, x, copystart, ncopy)
        # zero-out any parts of the buffer that weren't filled. This should
        # only happen at the beginning and end of the process
        # xbuf[npos+1:end-nneg] .= zero(eltype(xbuf))
        if win !== nothing
            xbuf .*= win
        end
        mul!(Xbuf, p, xbuf)
        copyto!(X, (c-1)*N+1, Xbuf, 1, N)
        xidx += hop
    end

    if demod
        # we want to convert to a sliding-window STFT  by multiplying by
        # exp(-jnω). This is equivilent to demodulating each band of the STFT
        # filter bank
        n = (0:M-1)*hop
        ω = range(0, π, length=N)
        # TODO we could do this without allocating a whole copy of X by doing
        # the multipliation frame-wise
        X .*= exp.(-im .* n' .* ω)
    end

    X
end

"""
    istft2(X, nfft, hop=nfft÷2;
           out_nfft=nfft,
           out_hop=Int(out_nfft * hop / nfft),
           window=rect,
           demod=false)
Perform an inverse short-time fourier transform on the complex-valued matrix
`X`.

- `nfft` is the FFT size that was used to compute the STFT.
- `hop` is the hop size of the frames in `X`
- `out_nfft` is the desired iFFT size to use for reconstruction.
- `out_hop` is the hop size for the resynthesized windows. You can set this
  to something different from `hop` for basic time stretching (though you will
  likely also need to manually update the phase values).
- `window` adds a resynthesis window. It can be an `AbstractArray`
  of length `nfft`, or a function that can be called as
  `f(nfft; zerophase=true)` to generate such a window. We currently only handle
  the painless case where the round-trip window (elementwise product of the
  analysis and synthesis windows) has the constant-overlap-add property. Note
  that generally the window function should be periodic and centered around
  index `1`, which is what the `zerophase` option in DSP.jl's window functions
  does. See the DSP.jl documentation for more details.
- `demod` should be set to `true` if the STFT phase has been demodulated (i.e
  a "sliding window" stft).

$stft_docs
"""
function istft2(X, nfft, hop=nfft÷2;
                out_nfft=nfft,
                # out_hop will throw an error if the given sizes aren't
                # compatible. TODO: give a better error message or handle this
                # more gracefully
                out_hop=Int(out_nfft * hop / nfft),
                window=rect,
                demod=false)
    # we could handle this case, but it's simpler not to
    @assert iseven(nfft)
    @assert iseven(out_nfft)
    @assert hop > 0
    @assert out_hop > 0

    N = size(X, 1)
    M = size(X, 2)

    # the output signal starts in the center of the first window, and
    # ends no later than the center of the last window
    L = (M-1) * out_hop + 1
    out = zeros(real(eltype(X)), L)

    # TODO: if the element type is e.g. Float32, we should use the same for
    # the window
    # TODO: scaling should probably be handled more automatically, and
    # also we should handle the non-COLA case. Also we're not scaling
    # the windows above because we're assuming they're COLA and add to 1
    # ALSO there are still some artifacts at the very beginning and end,
    # where the samples don't get the right number of windows adding up,
    # so the scale factor is off.
    win = if window isa Function
        window(out_nfft, zerophase=true)
    elseif window isa AbstractArray
        @assert length(window) == out_nfft
        window
    else

    end

    if demod
        # we want to convert to a sliding-signal STFT (which is what most stft
        # implementations do natively) by multiplying by exp(jnω). This is
        # equivilent to modulating the STFT filter bank by the band's center
        # frequency.
        n = (0:M-1)*hop
        ω = range(0, π, length=N)
        # TODO we could do this without allocating a whole copy of X by doing
        # the multipliation frame-wise
        X .*= exp.(im .* n' .* ω)
    end

    # this is probably pretty close to supporting nfft != N, but not really
    # tested in that case
    Xbuf = similar(X, out_nfft÷2+1)

    # fill with zeros, so that if we partially fill it later, it ends up being
    # zero-padded
    fill!(Xbuf, zero(eltype(X)))
    outbuf = similar(out, out_nfft)
    p = plan_irfft(Xbuf, out_nfft)
    copylen = min(length(Xbuf), N)

    nfft2 = out_nfft÷2

    outoffset = 0
    for c in 1:M
        copyto!(Xbuf, 1, X, (c-1)*N+1, copylen)
        mul!(outbuf, p, Xbuf)
        outbuf .*= win
        npos = min(nfft2, L-outoffset)
        @simd for i in 1:npos
            @inbounds out[outoffset+i] += outbuf[i]
        end
        nneg = min(nfft2, outoffset)
        @simd for i in 1:nneg
            @inbounds out[outoffset-nneg+i] += outbuf[end-nneg+i]
        end
        outoffset += out_hop
    end

    out
end


load("kernels.jl")

dest1 = Array(Int, (2,3))
dest2 = Array(Int, (2,3))
B = [1 2 3
     4 5 6]
C = [1 0 1
     0 1 0]
D = [ 0  0  0
     10 10 10]

# code = quote
#     A = B.*C + D
#     dest1[...] = A
#     dest2[...] = A + C
# end
# ex = code_kernel(code)
# eval(ex)

@kernel let
    A = B.*C + D
    dest1[...] = A
    dest2[...] = A + C
end

function f(B,C,D)
    A = B.*C + D
    dest1 = A
    dest2 = A + C
    (dest1, dest2)
end
dest1f, dest2f = f(B,C,D)

println("kernel:     dest1 =\n$dest1")
println("Array ops:  dest1 =\n$dest1f")
println()
println("kernel:     dest2 =\n$dest2")
println("Array ops:  dest2 =\n$dest2f")
println()


using LinearAlgebra
using Statistics: mean, std

Base.signbit(z::Complex{T}) where T<:Number = signbit(real(z))

"""
    integrate is the main entry point

    input:
    ------
    eq: a Symbolics expression to integrate
    x: the independent variable (optional)

    abstol: the desired tolerance
    num_steps: the number of different steps with expanding basis to be tried
    num_trials: the number of trials in each step (no changes to the basis)
    radius: the radius of the disk in the complex plane to generate random test points
    show_basis: if true, the basis (list of candidate terms) is printed
    opt: the sparse regression optimizer
    bypass: if true, do not integrate terms separately but consider all at once
    symbolic: try symbolic integration first
    max_basis: the maximum number of candidate terms to consider
    verbose: print a detailed report
    complex_plane: generate random test points on the complex plane (if false, the points will be on real axis)
    homotomy: use the homotopy algorithm to generat the basis

    output:
    -------
    solved, unsolved, err

    solved is the solved integral and unsolved is the residual unsolved portion of the input
    err is the numerical error in reaching the solution
"""
function integrate(eq, x=nothing; abstol=1e-6, num_steps=2, num_trials=5, radius=1.0,
                   show_basis=false, opt = STLSQ(exp.(-10:1:0)), bypass=false,
                   symbolic=true, max_basis=100, verbose=false, complex_plane=true,
                   homotopy=true)
    eq = expand(eq)
    eq = apply_div_rule(eq)

    if x == nothing
        x = var(eq)
        if x == nothing
            @syms 𝑥
            x = 𝑥
        end
    end

    l = Logger(verbose)

    # eq is a constant
    if !isdependent(eq, x)
        return x * eq, 0, 0
    end

    # homotopy = homotopy && !bypass

    return integrate_sum(eq, x, l; bypass, abstol, num_trials, num_steps,
                              radius, show_basis, opt, symbolic,
                              max_basis, verbose, complex_plane, homotopy)
end

"""
    integrate_sum applies the integral summation rule ∫ Σᵢ fᵢ(x) dx = Σᵢ ∫ fᵢ(x) dx

    eq: the integrand
    x: the indepedent variable
    l: a logger

    output is the same as integrate
"""
function integrate_sum(eq, x, l; bypass=false, kwargs...)
    solved = 0
    unsolved = 0
    ϵ₀ = 0
    ts = bypass ? [eq] : terms(eq)

    if length(ts) > 1
        inform(l, "Integrating sum", ts)
    end

    for p in ts
        s, u, ϵ = integrate_term(p, x, l; kwargs...)
        solved += s
        unsolved += u
        ϵ₀ = max(ϵ₀, ϵ)
    end

    if !isequal(unsolved, 0)
        eq = factor_rational(simplify_trigs(unsolved))

        if !isequal(eq, unsolved)
            eq = expand(eq)
            unsolved = 0
            ϵ₀ = 0
            ts = bypass ? [eq] : terms(eq)

            if length(ts) > 1
                inform(l, "Integrating transformed sum", ts)
            else
                inform(l, "Transforming the expression", ts[1])
            end

            for p in ts
                s, u, ϵ = integrate_term(p, x, l; kwargs...)
                solved += s
                unsolved += u
                ϵ₀ = max(ϵ₀, ϵ)

                if !isequal(u, 0)   # premature termination on the first failure
                    return 0, eq, ϵ₀
                end

            end
        end
    end

    return expand(solved), unsolved, ϵ₀
end

"""
    integrate_term is the central part of the code that tries different
    methods to integrate eq

    eq: the integrand
    x: the indepedent variable
    l: a logger

    output is the same as integrate
"""
function integrate_term(eq, x, l; kwargs...)
    args = Dict(kwargs)
    abstol, num_steps, num_trials, show_basis, symbolic, verbose, max_basis,
        radius, homotopy = args[:abstol], args[:num_steps],
        args[:num_trials], args[:show_basis], args[:symbolic], args[:verbose],
        args[:max_basis], args[:radius], args[:homotopy]

    attempt(l, "Integrating term", eq)

    if is_number(eq)
        y = eq * x
        result(l, "Successful", y)
        return y, 0, 0
    end

    basis = generate_basis(eq, x; homotopy)

    if show_basis
        inform(l, "Generating basis (|β| = $(length(basis)))", basis)
    end

    if length(basis) > max_basis
        result(l, "|β| = $(length(basis)) is too large")
        return 0, eq, Inf
    end

    D = Differential(x)
    ϵ₀ = Inf
    y₀ = 0

    # rescue
    ϵᵣ = Inf
    yᵣ = 0

    for i = 1:num_steps
        if length(basis) > max_basis break end

        Δbasis = [expand_derivatives(D(f)) for f in basis]

        if symbolic
            y, ϵ = try_symbolic(Float64, eq, x, basis, Δbasis; kwargs...)

            if !isequal(y, 0) && accept_solution(eq, x, y, radius) < abstol
                result(l, "Successful symbolic", y)
                return y, 0, 0
            else
                inform(l, "Failed symbolic")
            end
        end

        for j = 1:num_trials
            r = radius #*sqrt(2)^j
            y, ϵ = try_integrate(Float64, eq, x, basis, Δbasis, r; kwargs...)

            ϵ = accept_solution(eq, x, y, r)
            if ϵ < abstol
                result(l, "Successful numeric (attempt $j out of $num_trials)", y)
                return y, 0, ϵ
            elseif ϵ < ϵᵣ
                ϵᵣ = ϵ
                yᵣ = y
            end
        end

        inform(l, "Failed numeric")

        if i < num_steps
            basis = expand_basis(basis, x)
            if show_basis
                inform(l, "Expanding the basis (|β| = $(length(basis)))", basis)
            elseif verbose
                inform(l, "Expanding the basis (|β| = $(length(basis)))")
            end
        end
    end

    if ϵᵣ < abstol * 10
        result(l, "Accepting numeric (rescued)", yᵣ)
        return yᵣ, 0, ϵᵣ
    else
        result(l, "Unsucessful", eq)
        return 0, eq, ϵ₀
    end
end

###############################################################################

"""
    the core of the randomized parameter-fitting algorithm

    `try_integrate` tries to find a linear combination of the basis, whose
    derivative is equal to eq

    output
    -------
    integral, error
"""
function try_integrate(T, eq, x, basis, Δbasis, radius; kwargs...)
    args = Dict(kwargs)
    abstol, opt, complex_plane, verbose =
        args[:abstol], args[:opt], args[:complex_plane], args[:verbose]

    basis = basis[2:end]    # remove 1 from the beginning
    Δbasis = Δbasis[2:end]
    n = length(basis)

    # A is an nxn matrix holding the values of the fragments at n random points
    A = zeros(Complex{T}, (n, n))
    X = zeros(Complex{T}, n)

    init_basis_matrix!(T, A, X, x, eq, Δbasis, radius, complex_plane; abstol)

    y₁, ϵ₁ = sparse_fit(T, A, x, basis, Δbasis, opt; abstol)
    if ϵ₁ < abstol
        return y₁, ϵ₁
    end

    y₂, ϵ₂ = find_singlet(T, A, basis; abstol)
    if ϵ₂ < abstol
        return y₂, ϵ₂
    end

    if n < 8    # here, 8 is arbitrary and signifies a small basis
        y₃, ϵ₃ = find_dense(T, A, basis; abstol)
        if ϵ₃ < abstol
            return y₃, ϵ₃
        end
    end

    # moving toward the poles
    ∂eq = expand_derivatives(Differential(x)(eq))
    modify_basis_matrix!(T, A, X, x, eq, ∂eq, Δbasis, radius; abstol)
    y₄, ϵ₄ = sparse_fit(T, A, x, basis, Δbasis, opt; abstol)

    if ϵ₄ < abstol || ϵ₄ < ϵ₁
        return y₄, ϵ₄
    else
        return y₁, ϵ₁
    end
end

function init_basis_matrix!(T, A, X, x, eq, Δbasis, radius, complex_plane; abstol=1e-6)
    n = size(A, 1)
    k = 1
    i = 1

    while k <= n
        try
            x₀ = test_point(complex_plane, radius)
            X[k] = x₀
            d = Dict(x => x₀)

            b₀ = Complex{T}(substitute(eq, d))
            if is_proper(b₀)
                for j = 1:n
                    A[k,j] = Complex{T}(substitute(Δbasis[j], d)) / b₀
                end
                if all(is_proper, A[k,:])
                    k += 1
                end
            end
        catch e
            println("Error from init_basis_matrix!: ", e)
        end
    end
end

function modify_basis_matrix!(T, A, X, x, eq, ∂eq, Δbasis, radius; abstol=1e-6)
    n = size(A, 1)
    k = 1
    for k = 1:n
        d = Dict(x => X[k])
        # One Newton iteration toward the poles
        # note the + sign instead of the usual - in Newton-Raphson's method. This is
        # because we are moving toward the poles and not zeros.
        x₀ = X[k] + Complex{T}(substitute(eq, d)) / Complex{T}(substitute(∂eq, d))
        X[k] = x₀
        d = Dict(x => x₀)
        b₀ = Complex{T}(substitute(eq, d))
        for j = 1:n
            A[k,j] = Complex{T}(substitute(Δbasis[j], d)) / b₀
        end
    end
end

function sparse_fit(T, A, x, basis, Δbasis, opt; abstol=1e-6)
    n = length(basis)
    # find a linearly independent subset of the basis
    l = find_independent_subset(A; abstol)
    A, basis, Δbasis, n = A[l,l], basis[l], Δbasis[l], sum(l)

    try
        b = ones(n)
        # q₀ = A \ b
        q₀ = Optimize.init(opt, A, b)
        @views Optimize.sparse_regression!(q₀, A, permutedims(b)', opt, maxiter = 1000)
        ϵ = rms(A * q₀ - b)
        q = nice_parameter.(q₀)
        if sum(iscomplex.(q)) > 2 return nothing, Inf end   # eliminating complex coefficients
        return sum(q[i]*basis[i] for i = 1:length(basis) if q[i] != 0; init=zero(x)), abs(ϵ)
    catch e
        println("Error from sparse_fit", e)
        return nothing, Inf
    end
end

function find_singlet(T, A, basis; abstol)
    σ = vec(std(A; dims=1))
    μ = vec(mean(A; dims=1))
    l = (σ .< abstol) .* (abs.(μ) .> abstol)
    if sum(l) == 1
        k = findfirst(l)
        return nice_parameter(1/μ[k]) * basis[k], σ[k]
    else
        return nothing, Inf
    end
end

function find_dense(T, A, basis; abstol=1e-6)
    n = size(A, 1)
    b = ones(T, n)

    try
        q = A \ b
        if minimum(abs.(q)) > abstol
            ϵ = maximum(abs.(A*q .- b))
            if ϵ < abstol
                y = sum(nice_parameter.(q) .* basis)
                return y, ϵ
            end
        end
    catch e
        #
    end
    return nothing, Inf
end


using Plots
using Omega
import Omega.Space
using Omega.Space: getdim

"""
Contour Plot of two dimensions of Ω

```
x = normal(0.0, 1.0)
y = normal(0.0, 1.0)
c1 = err(x ==ₛ y)
c2 = err(x >ₛ y)
ωcontour(c2)
```

"""
function ωcontour(xrv::RandVar;
                  ΩT = defΩ(),
                  xdim = 1,
                  ydim = 2,
                  xrng = 0:0.005:1,
                  yrng = 0:0.005:1,
                  plt = plot(),
                  fill = true,
                  kwargs...)
  ω = ΩT()
  xrv(ω)
  function f(x, y)
    ω = Space.update(ω, xdim, x)
    ω = Space.update(ω, ydim, y)
    xrv(ω)
  end
  contour!(plt, xrng, yrng, f; fill = fill, kwargs...)
end

# "x ∈ [0, 1]?"
# isunit(x) = 0.0 <= x <= 1.0

"Plot path in Omega Space from sequence of samples"
function ωtrace(ωs::Vector;
                   xdim = 1,
                   ydim = 2,
                   plt = plot())
  xs, ys = ntranspose(map(ω -> (getdim(ω, xdim), getdim(ω, ydim)), ωs))
  plot!(plt, xs, ys, arrow = :arrow, linealpha = 0.5, legend = false)
end

# function plottraces(qpdata, plt = plot())
#   for qp in qpdata
#     plottrace(qp, plt)
#   end
#   plt
# end

# function testcb2(;alg = HMC, n = 1000, kwargs...)
#   μ = normal(0.0, 1.0)
#   x = normal(μ, 1.0)
#   y = (x ==ₛ 0.0)
#   # y = (x == 0.0) | (μ < 0.0)
#   cb, cbdata = Omega.tracecb(Omega.QP, deepcopy)
#   # cb = [default_cbs(n); cb]
#   rand(μ, y, n; alg = alg, cb = cb, kwargs...)
#   qpdata = cbdata[2]
#   plt = plot()
#   ωcontour(y, x.id, μ.id, Omega.defΩ(HMC)(), plt = plt)
#   print(plottraces(qpdata, plt))
#   qpdata, plt
# end

# testcb(nsteps = 20, stepsize = 0.01, n = 10000)

# function testcb2(;kwargs...)
#   μ = normal(0.0, 1.0)
#   x = normal(μ, 1.0)
#   # y = (x == 0.0) | (μ < 0.0)
#   y = (x^2 + μ^2 == 1.0)
#   cb, cbdata = Omega.tracecb(Omega.QP)
#   cb = [default_cbs(n); cb]
#   rand(μ, y, HMC; n = n, cb = cb, kwargs...)
#   qpdata = cbdata[2]
#   plt = plot()
#   ωcontour(y, x.id, μ.id, Omega.defΩ(HMC)(), plt = plt)
#   plottraces(qpdata, plt)
#   qpdata, plt
# end
# # qpdata, plt = testcb2(nsteps = 20, stepsize = 0.01, n=5000)

# function to_output_space(q1, q2)
#   n1  = Distributions.Normal(0, 1)
#   μ = quantile(n1, q1)
#   n2 = Distributions.Normal(μ, 1)
#   μ, quantile(n2, q2)
# end

# # function plotoutput(data, i=0,plt = plot())
# #   transf(d) = to_output_space(d.q |> Omega.inv_transform)...)[2]
# #   ys = map(transf, data)
# #   k = 1.0/(data |> length)
# #   xs = i:k:(i+1-k)
# #   plot!(plt, xs, ys, arrow = :arrow, linealpha = 0.5, legend=false)
# # end

# function plotoutputs(data, plt=plot())
#   for (i, qp) in enumerate(data)
#     plotoutput(qp, i, plt)
#   end
#   plt
# end

# function plotoutputXY(data, plt = plot())
#   xs, ys = zip([g((d.q |> Omega.inv_transform)...)
#                 for d in data]...) |> collect
#   plot!(plt, xs |> collect, ys |> collect, arrow = :arrow, linealpha = 0.5, legend=false)
# end

# function plotoutputsXY(data, plt=plot())
#   foreach(data) do qp
#     plotoutputXY(qp, plt)
#   end
#   plt
# end

# function filtered_qp(qpdata)
#   indexes = filter(1:length(qpdata)-1) do i
#     qpdata[i][1].q != qpdata[i+1][1].q
#   end
#   qpdata[indexes]
# end