"""
Functions for finding piecewise linear approximations to concave
    functions at integer points.

    These functions solve the sparsest reduction problem from the paper.
"""

function check_cb_submodular(w::Vector{Float64},tol = 1e-10)
    """
    Given a set of weights defining a cardinality-based function,
    check whether it is submodular.

    w(i) = w_{i-1}
    """
    for i = 2:length(w)-1
        if 2*w[i] - w[i-1] - w[i+1] < -tol
            return false
        end
    end
    return true
end

function check_cb_symmetric(w::Vector{Float64})
    """
    Check whether weights defining a cardinality-based function
    define a symmetric splitting function.

    If there are k weights, this is for a function defined on
    k-1 nodes, since we include w_0.
    """
    k = length(w)
    for i = 1:round(Int64,floor(k/2))
        if w[i] != w[k-i+1]
            return false
        end
    end
    return true
end

function random_CombGad(k::Int64)
    """
    This generates a random combined cardinality-based gadget.

    Returns the parameters as well as the integer
    evaluations.
    """
    J = rand(round(Int64,k/2):k-1)
    a = rand(J)
    b = sort((k-1)*rand(J))
    z0 = 2*rand()
    zk = 2*rand()
    w = CombGad(z0,zk,k,a,b)
    return z0, zk, a, b, w
end

function random_scb_function(k::Int64,prob1=1.0,prob2=1.0)
    """
    Generate a random cardinality-based submodular function.
        k = number of nodes the function is defined on.

    This generates a random combined cardinality-based gadget.

    prob1 = probability that z0 > 0
    prob2 = probability that zk > 0
    """
    J = rand(round(Int64,k/2):k-1)
    a = round.(rand(J),digits = 2)
    while minimum(a) == 0
        a = round.(rand(J),digits = 2)
    end
    b = round.(sort((k-1)*rand(J)),digits = 2)
    while minimum(b) == 0
        b = round.(sort((k-1)*rand(J)),digits = 2)
    end

    if rand() < prob1
        z0 = 2*round.(rand(),digits = 2)
    else
        z0 = 0.0
    end
    if rand() < prob2
        zk = 2*round.(rand(),digits = 2)
    else
        zk = 0.0
    end
    w = CombGad(z0,zk,k,a,b)
    w = w./maximum(w)*3
end

function CombGad(z0,zk,k::Int64,a,b)
    """
    This is a k-node Combined Gadget function.
    Returs the integer values.
    """
    J = length(a)
    @assert(length(a) == length(b))
    @assert(is_increasing(b))

    @assert(b[J] < k)
    @assert(minimum(a) > 0)
    @assert(minimum(b) > 0)
    I = collect(0:k)
    w = z0*(k .- I) + zk*I + sum( a[j]*min.( I*(k-b[j]), (k .- I)*b[j]) for j = 1:J)

    return w
end

function CombGad_func(z0,zk,k::Int64,a,b)
    """
    This is a k-node Combined Gadget function.
    It returns a callable function for any x
    """
    J = length(a)
    if J > 0
        @assert(length(a) == length(b))
        @assert(is_increasing(b))
        @assert(b[J] < k)
        @assert(minimum(a) > 0)
        @assert(minimum(b) > 0)
        f = x->(z0*(k .- x) + zk*x + sum( a[j]*min.( x*(k-b[j]), (k .- x)*b[j]) for j = 1:J))
    else
        f = x->(z0*(k .- x) + zk*x)
    end
    return f
end

function is_increasing(b)
    for i = 1:length(b)-1
        if b[i+1] < b[i]
            return false
        end
    end
    return true
end

is_decreasing(b) = is_increasing(reverse(b))

function plf_to_CombGad(m::Vector{Float64},b::Vector{Float64},k::Int64,f0::Float64)
    """
    Find the Combined Gadget weights for a concave, piecewise linear function f.

    m = slope vector
    b = breakpoints
    k = defines interval [0,k]
    f0 = f(0)

    Output:
    z0, zk, a, b = parameters for Combined Gadget Function
    f = vector of function evaluations at breakpoints.

    More detail
    ------------
    A nonnegative, piecewise linear function f on [0,k] is uniquely determined
    by its breakpoints, slopes, and a single point (0,f(0)).

    Given the slopes m and breakpoints b, and point f(0) of such a function,
    which is additionally concave (meaning m is a decreasing vector).

    We return the parameters z0, zk, a, and b that define the CombGad function
    whose continuous extension is exactly the function f.

    Let J+1 be the number of linear pieces.

    We'll also return the function values at breakpoints: (b_i, f_i)
    """
    @assert(is_decreasing(m))
    J = length(m)-1
    z0 = f0/k
    a = zeros(J)
    f = zeros(J+2)
    f[1] = f0
    bext = [0;b;k]
    for j = 1:J
        a[j] = 1/k*(m[j]-m[j+1])
        f[j+1] = m[j]*(bext[j+1]-bext[j]) + f[j]
    end
    f[J+2] = m[J+1]*(bext[J+2]-bext[J+1])+f[J+1]
    zk = f[end]/k
    return z0, zk, a, b, f[2:end-1]
end

function breakpoints_to_slopes(b,f)
    """
    Given breakpoints and function evaluations at these breakpoints
        for a nonnegative, piecewise linear function on [0,k],
        return the slopes of the curves.

        b[i] = is the ith breakpoint
        f[i] = the piecewise linear function evaluated at b[i]

        Although not technically breakpoints, the first and last entry of b
        should be 0 and k for some integer k, so that we can define the function.
    """
    @assert(length(b) == length(f))
    @assert(minimum(b) == 0)
    @assert(is_increasing(b))
    @assert(b[1] == 0)
    k = b[end]
    @assert(k == round(k))

    J = length(b)-1
    m = zeros(J)
    for i = 1:J
        m[i] = (f[i+1]-f[i])/(b[i+1]-b[i])
    end
    return m

end

# functions for the greedy search
function line_through(x1,y1,x2,y2)
    """
    Return the line through points (x1,y1) and (x2,y2).
    Gives the slope as well as the line as a callable function.
    """
    m = (y2-y1)/(x2-x1)
    inter = y1-m*x1
    l_fun = x->(m*x .+ inter)
    return m, l_fun
end

function line_intersection(m1,m2,x1,x2,y1,y2)
    """
    Given two lines, defined by slopws m1, m2 and
    points (x1,y1), (x2,y2), find the point of intersection.
    """
    if m1 == m2
        println("Lines are parallel")
        return
    end
    b1 = y1 - m1*x1
    b2 = y2 - m2*x2
    x = (b2-b1)/(m1-m2)
    y = m1*x+b1
    return x,y
end

function Refined_SCB_to_CGF(w,epsi=1e-14)
    """
    This allows you to find the CGF with minimum number of linear pieces
    AND the smallest (or approximately the smallest) approximation guarantee
    you can get using that number of pieces.

    It refines the output from one call to the original algoirhtm
    for finding the minimum number of pieces.
    """
    z0, zk, a, b, cgf = SubCardFun_to_CGF_weights(w,epsi,false)
    Jtrue = length(a) + 1
    J = Jtrue
    best_eps = epsi
    # Now find the same number of linear pieces, but better
    # approximations, if possible
    while Jtrue == J && epsi > 1e-10
        epsi = epsi/2
        z0_t, zk_t, a_t, b_t, cgf_t = SubCardFun_to_CGF_weights(w,epsi,false)
        J = length(a_t) + 1
        if J == Jtrue
            z0 = z0_t
            zk = zk_t
            a = a_t
            b = b_t
            cgf = cgf_t
            best_eps = epsi
        end
    end
    return z0, zk, a, b, cgf, best_eps
end

function SubCardFun_to_CGF_weights(w,epsi=1e-14,verbose = false)
    """
    Given function evaluations for a submodular cardinality-based
    function w, return the parameters needed to construct a
    combination of CB-gadgets that approximately model that function.
    """
    # even if the goal is to perfectly fit the function,
    # need to fit it to within a small tolerance for numerical issues
    epsi = max(epsi,1e-14)
    k = length(w)-1
    Lines, Slopes, X, Y, Ranges = get_pwl_approx(w,epsi,verbose)
    z0, zk, a, b, f = Lines_to_CombGadget(Slopes,X,Y,k)
    cgf = CombGad_func(z0,zk,k,a, b)
    return z0, zk, a, b, cgf
end

function Lines_to_CombGadget(Slopes,X,Y,k)
    """
    Given information about a set of lines,
    (namely, the slope and a point for each line)
    return the parameters that describe the piecewise
    linear function as a Combined Gadget Function.
    """
    b1 = Y[1] - Slopes[1]*X[1]
    f0 = b1
    b = zeros(length(X)-1)
    for t = 1:length(X)-1
        x,y = line_intersection(Slopes[t],Slopes[t+1],X[t],X[t+1],Y[t],Y[t+1])
        b[t] = x
    end

    z0, zk, a, b, f = plf_to_CombGad(Slopes,b,k,f0)

    return z0, zk, a, b, f
end

function SymmetricSCB_to_Gadget(w,epsi,verbose = false)
    @assert(is_increasing(w))

    # in the symmetric case we also assume w_0 = 0
    @assert(w[1] == 0)
    r = length(w) - 1
    Lines, M, X, Y, Ranges = get_pwl_approx(w,epsi,false)
    mlast = M[end]

    # In the symmetric case, penalties are increasing
    # and we need to ensure the last linear piece is flat
    if mlast < 0
        # negative sloped lines not needed
        # replace with flat line
        M[end] = 0
        X[end] = r
        Y[end] = w[end]
        # pop!(Lines)
    elseif Ranges[end][1] == r-1 && Ranges[end-1][2] == r-1 && mlast > 0
        # Covered point r-1 twice. Get rid of it
        M[end] = 0
        X[end] = r
        Y[end] = w[end]
        # pop!(Lines)
    elseif mlast > 0
        push!(M,0)
        push!(X,r)
        push!(Y,w[end])
    end
    # push!(Lines, x->w[end]*ones(length(x)))

    b = zeros(length(X)-1)
    for t = 1:length(X)-1
        if M[t] == M[t+1]
            @show r, M
            @show X, Y
            @show mlast
        end
        x,y = line_intersection(M[t],M[t+1],X[t],X[t+1],Y[t],Y[t+1])
        b[t] = x
    end

    a = zeros(length(M)-1)
    for i = 1:length(a)
        a[i] = M[i]- M[i+1]
    end

    return a,b
end

function get_pwl_approx(w,epsi,verbose = false)
    """
    Given penalties w for a submodular cardinality based function,
    find a minimum sized set of linear pieces that provides
    an upper bounding (1+epsi) approximation to w everywhere.
    """
    k = length(w)-1
    @assert(epsi >= 0)
    Lines = Vector()
    Slopes = Vector{Float64}()
    X = Vector{Float64}()
    Y = Vector{Float64}()
    new_i = 1
    Ranges = Vector{Tuple{Int64,Int64}}()

    while new_i <= k+1
        i = new_i
        x2, w2, m, l_fun, new_i = next_line(w,i, epsi; verbose=verbose)

        if i == k || i == k+1
            push!(Ranges,(k-1,k))
        else
            push!(Ranges,(i-1,new_i-2))
        end
        # Save information about the lines
        push!(Lines,l_fun)
        push!(Slopes,m)
        push!(X,x2)
        push!(Y,w2)

    end
    # push!(Ranges,k+1)
    return Lines,Slopes, X, Y, Ranges
end

function next_line(w,i,epsi;verbose = false)
    """
    w is the submodular cardinality-based penalties we're trying to cover.
    i is the next index in w that currently isn't covered by a linear piece.
    Next index i means next integer we need to cover is x[i] = i-1
    """
    k = length(w)-1
    x = collect(0:k)

    # Check for end case 1: x[i] == k-1 or x[i] == k
    if x[i] == k-1 || x[i] == k
        # Add line covering (k-1,w_{k-1}) and (k,w_k)
        xi = x[end-1]
        wi = w[end-1]
        x2 = x[end]
        w2 = w[end]
        m, l_fun = line_through(xi,wi,x2,w2)
        if verbose println("Last line through ($xi,$wi) and ($x2,$w2)") end
        return x2, w2, m, l_fun, k+2
    end

    # Otherwise, start a new line at point (xi,li).
    # This is the worst approx at xi we will accept, maximizing reach.
    xi = x[i]
    li = (1+epsi)*w[i]
    if verbose println("Start a line at ($(xi),$li)") end

    # Visit next point to define a new line
    current = i+1
    x2 = x[current]
    w2 = w[current]
    m, l_fun = line_through(xi,li,x2,w2)
    if verbose println("Candidate line L: through ($x2,$w2)") end

    current += 1

    # Visit points one at a time, greedily finding best cover
    while true
        x_current = x[current]          # integer value we consider
        w_current = w[current]          # true penalty at that value
        l_current = l_fun(x_current)    # approximation provided by current line

        if verbose
            println("Visiting i = $x_current")
            println("(w[$x_current],L[$x_current],(1+eps)w[$x_current]) = ($(w_current),$(l_current),$((1+epsi)*w_current))")
        end

        if l_current < w_current

            if verbose println("Undershot, need new line.") end

            # These are the points defining the new line
            x2 = x_current
            w2 = w_current
            m, l_fun = line_through(xi,li,x2,w2)
            if verbose println("Candidate line L: through ($x2,$w2)") end

            # If this covers last point, we're done
            if x2 == k
                return x2, w2, m, l_fun, k+2
            else
                current += 1
            end

        elseif l_current > (1+epsi)*w_current
            if verbose println("Overshot, return") end

            # we can return the next line, slope m going through point (x2, w2)
            # and also the next point that is now uncovered

            return x2, w2, m, l_fun, current

        elseif l_current <= (1+epsi)*w_current && x_current == k

            # We're done, there's nothing left to cover.
            return x2, w2, m, l_fun, k+2

        else
            # We've covered a point, but there are more points to go.
            current += 1
            x_current = x[current]
        end
    end
end

function clique_round_bound(epsilon)
    """
    Return the upper bound on the number of linear pieces
    needed to bound the clique concave function from [1,k/2]
    for any k.
    """
    return ceil.(log2.(log2.(1 ./ epsilon))) .* ceil.(2*(1 .+ epsilon) ./ sqrt.(2*epsilon))
end


function slicerm(g, dim, x0, xargs; N=1000, w=0.5, m=10)
    #= multi-D slice sampler, g: log of distribution function,
    dim number of dimensions
    x0 start point,
    xargs extra arguments for distribution, N*dim samples to return,
        w "out step" distance, m max number of outward steps.

        Note that you only have to worry about log scaling in
	'vertical' comparisons
        in these Monte Carlo chain calculations.
    =#
    xs = zeros(N*dim,dim)  # array that will be returned
    xs[1,:] = x0  #x0 should be vector of length dim
    x1 = zeros(1,dim)
    L = zeros(1,dim)
    R = zeros(1,dim)
    way = zeros(1,dim)  # which axis to go along in space
    for i in 2:N*dim
        for d in 1:dim
            way = 0*way #clear it
            way[d] = 1.0 #set nonzero in direction we go for slicing
            y0 = g(x0,xargs)  #height of distribution at x0
            y = y0+log(rand()) # height for slice
            #start stepping out
            U = rand()   # between 0 and 1
            L = x0 - w*way*U
            R = L + w*way
            V = rand()
            J = floor(m*V)
            K = (m - 1) - J
            while J > 0 && y < g(L,xargs)
                L = L - w*way
                J = J - 1
            end
            while K > 0 && y < g(R,xargs)
                R = R + w*way
                K = K - 1
            end
            #now should be stepped out beyond distribution at slice level
            # work back in if no value found:
            Lbar, Rbar = L, R
            while true
                U = rand()
                x1 = Lbar + U * (Rbar - Lbar)  # vector subtraction should be
						# correct dir
                if y < g(x1,xargs)
                    break # exit while loop
                end
                if x1[d] < x0[d]
                    Lbar = x1
                    else
                    Rbar = x1
                end
            end
            xs[i,:] = x1 # found an acceptable point, record in chain (a row)
            x0 = x1 # set initial to new point for next round.
        end # for d
    end  #for N
    return xs
end


"""
function chroma(stft::Matrix{Complex{Cdouble}},
	fs::Cdouble, wsize::Cuint, γ::Cdouble, ϵ::Cdouble
)::Matrix{Cdouble}

Compute the chroma features of a signal from its Short Time Fourier Transform.

The transform must be provided as a matrix in which each column represents a
frame and with time evolving along the second dimension. The return will be a
matrix of size 12xN where N is the number of frames of the STFT and each column
is a vector of normalized chroma features.

**Parameters:**

`stft`: The Short Time Fourier Transform of the signal from which the features
must be computed.

`fs`: The sampling rate of the signal.

`wsize`: The number of samples of the window used in the STFT.

`γ`: The compression factor to be used in the chroma features extraction.

`ϵ`: The smallest norm that a chroma vector needs to have to not be replaced by
a stub one.
"""
function chroma(
	stft::Matrix{Complex{Cdouble}},
	fs::Cdouble,
	wsize::Cuint,
	γ::Cdouble,
	ϵ::Cdouble
)::Matrix{Cdouble}

	# The number of chroma vectors is defined by the frames in the STFT
	nofframes = size(stft)[2]
	pitchspectrum = Matrix{Cdouble}(undef, 128, nofframes)

	# Relation of a pitch with a frequency defined by the MIDI standard
	pitch2freq(p) = round(Cint, (wsize/fs)*2^((p - 69)/12)*440)

	# The bound frequencies of each one of the pitch classes
	bounds = [1 + pitch2freq(p - 0.5) for p in 0:128]

	for n in 1:nofframes

		# The pitch spectrum is a log spaced spectrum where a component is
		# proportional to the sum of the power of all stft components inside the
		# bounds of a pitch class
		pitchspectrum[:, n] =
			[sum(abs2.(stft[bounds[i]:bounds[i + 1],n])) for i in 1:128]

	end

	chroma = Matrix{Cdouble}(undef, 12, nofframes)

	for n in 1:nofframes

		# A chroma vector is given by the sum of the pitch classes of the same
		# note in all octaves
		chroma[:, n] =
			[sum(pitchspectrum[(class % 12 + 1):12:end, n]) for class in 0:11]

	end

	# Compresses the chroma vectors
	chroma = log10.(1 .+ γ .* chroma)

	stubFeature = ones(12)./√12

	# Evaluates the norm of each chroma vector...
	for i in 1:nofframes

		# ...and if it's smaller than the ϵ parameter normalizes the vector
		if ((norm = √(sum(chroma[:, i].^2))) > ϵ)
			chroma[:, i] ./= norm
		# Otherwise replaces the vector with a stub one
		else
			chroma[:, i] = stubFeature
		end

	end

	return chroma

end


using BNBKnapsack
using Test

@testset "BNBKnapsack.jl" begin
    # Execute the BNB
    #@test BNBKnapsack.execute(BNBKnapsack.Problem(4,[50, 40, 45, 55],[5, 4, 3, 4],13,1e-4, 1e-3)) == 0
end

n_objects = 4               # Also the number of variables.
values = [16, 22, 12, 8]
weights = [5, 7, 4, 3]
max_weight = 14

integer_precision = 1e-4    # Integers are represented as floating-point values: if they are close enough to an integer, they are considered integer.
gap_precision = 1e-3

problem = BNBKnapsack.Problem(n_objects,values,weights,max_weight,integer_precision,gap_precision)
BNBKnapsack.execute(problem)


######## Fractional Brownian Motion ########
"""
Fractional Brownian motion.

# Members
- hurst: the Hurst exponent
"""
struct FractionalBrownianMotion <: SSSIProcess
    hurst::Real

    function FractionalBrownianMotion(hurst::Real)
        0. < hurst < 1. || error("Hurst exponent must be bounded in 0 and 1.")
        new(hurst)
    end
end

ss_exponent(X::FractionalBrownianMotion) = X.hurst

@doc raw""" fBm_autocov(t,s,H)

Autocovariance function of standard fBm:
    $\frac 1 2 (|t|^{2H} + |s|^{2H} - |t-s|^{2H})$
"""
fBm_autocov = (t::Real,s::Real,H::Real) -> 1/2 * (abs(t)^(2H) + abs(s)^(2H) - abs(t-s)^(2H))

autocov(X::FractionalBrownianMotion, t::Real, s::Real) = fBm_autocov(t,s,X.hurst)

"""
    fBm_covmat(G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, H::Real)

Compute the covariance matrix of fBm with hurst `H` between two grids `G1` and `G2`.
This gives the same result as `covmat(FractionalBrownianMotion(H), G1, G2)`.
"""
function fBm_covmat(G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, H::Real)
    Σ = zeros(length(G1),length(G2))
    for (c,s) in enumerate(G2), (r,t) in enumerate(G1)
        Σ[r,c] = fBm_autocov(t, s, H)
    end
    return Σ
end

fBm_covmat(G, H) = Matrix(Symmetric(fBm_covmat(G, G, H)))

# Moving average kernels of fBm
"""
K_+ kernel
"""
function Kplus(x::Real, t::Real, H::Real)
    # @assert t>0
    p::Real = H-1/2
    v::Real = 0
    if x<0
        v = (t-x)^p - (-x)^p
    elseif 0<=x<t
        v = (t-x)^p
    else
        v = 0
    end
    return v
end

"""
K_- kernel
"""
function Kminus(x::Real, t::Real, H::Real)
    p::Real = H-1/2
    v::Real = 0
    if x<=0
        v = 0
    elseif 0<x<=t
        v = -(x)^p
    else
        v = (x-t)^p - (x)^p
    end
    return v
end

"""
K_+ + K_- kernel
"""
Kppm(x, t, H) = Kplus(x, t, H) + Kminus(x, t, H)

"""
K_+ - K_- kernel
"""
Kpmm(x, t, H) = Kplus(x, t, H) - Kminus(x, t, H)

######## Fractional Gaussian Noise ########

@doc raw"""
Fractional Gaussian noise (fGn) is the discrete time version of the continuous time differential process of a fBm: $ B(t+\delta) - B(t) $.  It is defined as $ B(n+l) - B(n) $, where `l` is the lag.

# Note
- The anti-causal convention adopted by the definition here has no impact since the process is stationary.
"""
struct FractionalGaussianNoise <: IncrementProcess{FractionalBrownianMotion}
    parent_process::FractionalBrownianMotion
    lag::Integer

    function FractionalGaussianNoise(hurst::Real, lag::Integer=1)
        lag >= 1 || error("Lag must be >= 1.")
        new(FractionalBrownianMotion(hurst), lag)
    end
end

ss_exponent(X::FractionalGaussianNoise) = X.parent_process.hurst
lag(X::FractionalGaussianNoise) = X.lag
step(X::FractionalGaussianNoise) = 1

@doc raw"""
    fGn_autocov(t::Real,H::Real,δ::Real)

Autocovariance function of stardard (continuous time) fGn: $ 1/2 (|t+δ|^{2H} + |t-δ|^{2H} - 2|t|^{2H}) $
"""
fGn_autocov = (t::Real,H::Real,δ::Real) -> 1/2 * (abs(t+δ)^(2H) + abs(t-δ)^(2H) - 2*abs(t)^(2H))

function autocov(X::FractionalGaussianNoise, n::Integer)
    return fGn_autocov(n, ss_exponent(X), lag(X))
end

"""
    fGn_covmat(G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, H::Real, δ::Real)

Compute the covariance matrix of standard (continuous time) fractional Gaussian noise of Hurst exponent `H` and time lag `δ` between the grid `G1` and `G2`.

# Notes
The special case of discrete regular grids with integer lag `d` is equivalent to `covmat(FractionalGaussianNoise(H, d), 1:N, 1:M)`.
"""
function fGn_covmat(G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, H::Real, δ::Real)
    Σ = zeros(length(G1),length(G2))
    for (c,s) in enumerate(G2), (r,t) in enumerate(G1)
        Σ[r,c] = fGn_autocov(t-s, H, δ)
    end
    return Σ
end

fGn_covmat(G::AbstractVector{<:Real}, H::Real, δ::Real) = Matrix(Symmetric(fGn_covmat(G, G, H, δ)))

fGn_covmat(N::Integer, H::Real, d::Integer) = fGn_covmat(1:N, H, d)

######## Fractional Wavelet noise (fWn) ########

@doc raw"""
A fractional Wavelet noise (fWn) is discrete time version of the high-pass filtered process of fBm:
$$
W(t) = \sum_{n=0}^{p-1} B(t-n\delta) \psi[n]
$$

It is defined as
$$
W[k] = \sum_{n=0}^{p-1} B(k-n) \psi[n]
$$

# Notes
- fGn is a special case of fWn.
```julia
d = 1  # lag for fGn
fGn = FracFin.FractionalGaussianNoise(0.7, d)

filter = vcat(1,zeros(d-1),-1)
fWn = FracFin.FractionalWaveletNoise(0.7, filter)

G = FracFin.covmat(fGn, 10)
W = FracFin.covmat(fWn, 10)
maximum(abs.(G-W))  # close to 0
```
"""
struct FractionalWaveletNoise <: FilteredProcess{DiscreteTime, FractionalBrownianMotion}
    parent_process::FractionalBrownianMotion
    filter::AbstractVector{<:Real}
    coeff::AbstractVector{<:Real}  # coefficients for the summation of |t + k|^(2H)
    supp::AbstractVector{<:Integer}  # support (index range) of the coefficients

    function FractionalWaveletNoise(hurst::Real, filter::AbstractVector{<:Real})
        @assert isapprox(sum(filter), 0; atol=1e-8)  "Filter must not contain DC."

        # whatever the mode of transform the support is always 1-length(filter):length(filter)-1, and the coefficients satisfy coeffs[k] = coeffs[-k]
        new(FractionalBrownianMotion(hurst), filter, fWn_cov_coeff(filter), 1-length(filter):length(filter)-1)
    end
end

ss_exponent(X::FractionalWaveletNoise) = X.parent_process.hurst
filter(X::FractionalWaveletNoise) = X.filter
step(X::FractionalWaveletNoise) = 1

"""
# Notes
- `fWn_cov_coeff(u,v)` is the reverse of `fWn_cov_coeff(v,u)`.
"""
function fWn_cov_coeff(u::AbstractVector{<:Real}, v::AbstractVector{<:Real})
    m = length(u)
    n = length(v)

    return -1 * native_conv(u, reverse(v))
end

fWn_cov_coeff(u::AbstractVector{<:Real}) = fWn_cov_coeff(u, u)

"""
Autocovariance function of a continuous time fWn.
"""
function fWn_autocov(X::FractionalWaveletNoise, t::Real, δ::Real; mode::Symbol=:causal)
    return 1/2 * sum(X.coeff .* abs.(t .+ δ * X.supp).^(2*ss_exponent(X)))
end

autocov(X::FractionalWaveletNoise, n::Integer) = fWn_autocov(X, n, 1)

"""
    covmat(X::FractionalWaveletNoise, G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, δ::Real=1)

Compute the covariance matrix between the grid `G1` and `G2` of a continuous time fWn `X` with step `δ`.

# Notes
- Like `fGn_covmat()`, this function is equivalent to `covmat()` in case of regular grids.
"""
function covmat(X::FractionalWaveletNoise, G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, δ::Real=1)
    Σ = zeros(length(G1),length(G2))
    for (c,t) in enumerate(G2), (r,s) in enumerate(G1)
        Σ[r,c] = fWn_autocov(X, s-t, δ)
    end

    return Σ
end

covmat(X::FractionalWaveletNoise, G::AbstractVector{<:Real}, δ::Real=1) = Matrix(Symmetric(covmat(X, G, G, δ)))

## Redundant functions
function fWn_covmat(G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, filter::AbstractVector{<:Real}, H::Real, δ::Real)
    coeff = fWn_cov_coeff(filter)
    n = length(filter)
    Σ = zeros(length(G1),length(G2))

    for (c,t) in enumerate(G2), (r,s) in enumerate(G1)
        Σ[r,c] = 1/2 * sum(coeff .* abs.((s-t) .+ δ*(1-n:n-1)).^(2H))
    end

    return Σ
end

fWn_covmat(G::AbstractVector{<:Real}, f::AbstractVector{<:Real}, H::Real, δ::Real) = Matrix(Symmetric(fWn_covmat(G, G, f, H, δ)))

######## fWn bank (fWnb) ########

"""
A fractional Wavelet noise bank (fWnb) is a fBm filtered by a bank of filters, in other words, a collection of fWn.

# Notes
- This process is multivariate hence can not be properly handled in the scalar framework as fWn.
"""
struct FractionalWaveletNoiseBank <: AbstractRandomField # <: FilteredProcess{DiscreteTime, FractionalBrownianMotion}
    parent_process::FractionalBrownianMotion
    filters::AbstractVector
    coeffs::AbstractMatrix  # 2-tuple: pair coefficients for the summation of |t + k|^(2H) and the range of index
    supps::AbstractMatrix  # supports (index range) of the coefficients

    function FractionalWaveletNoiseBank(hurst::Real, filters::AbstractVector{<:AbstractVector{<:Real}}, mode::Symbol)
        @assert all(isapprox(sum(f), 0; atol=1e-8) for f in filters)  "Filters must not contain DC."

        # Gotcha! Double for-loops in array comprehension has reversed order than the usual double for-loop, e.g.
        # [(r,c) for r=1:3, c=1:2] is a 3-by-2 matrix
        coeffs = [fWn_cov_coeff(u,v) for u in filters, v in filters]  # the first for-loop variable `u` is the row dimension

        # Supports, following the same convention as in `convmask()`, but as correlation not convolution.
        supps = if mode in [:causal, :left]  # <-> :anticausal in `convmask`
            [(1-length(u):length(v)-1) for u in filters, v in filters]
        elseif mode in [:anticausal, :right]
            [(1-length(v):length(u)-1) for u in filters, v in filters]
        else
            [(1-floor(Int,length(v)/2)-ceil(Int,length(u)/2)):(floor(Int,length(u)/2)+ceil(Int,length(v)/2)-1) for u in filters, v in filters]
        end

        new(FractionalBrownianMotion(hurst), filters, coeffs, supps)
    end
end

ss_exponent(X::FractionalWaveletNoiseBank) = X.parent_process.hurst
filter(X::FractionalWaveletNoiseBank) = X.filters
step(X::FractionalWaveletNoiseBank) = 1

function fWn_autocov(A::AbstractMatrix{<:AbstractVector{<:Real}}, K::AbstractMatrix{<:AbstractVector{<:Real}}, H::Real, t::Real, δ::Real; partial::Bool=false, kwargs...)
    @assert size(A) == size(K) && size(A,1) == size(A,2)  "Mismatched dimensions."
    @assert 0 < H < 1  "Invalid Hurst exponent."

    J = size(A, 1)
    Σ = zeros(Real, (J,J))

    if partial
        for r=1:J
            Σ[r,r] = 1/2 * sum(A[r,r] .* abs.(t .+ δ*K[r,r]).^(2H))
        end
    else
        for c=1:J, r=1:J
            Σ[r,c] = 1/2 * sum(A[r,c] .* abs.(t .+ δ*K[r,c]).^(2H))
        end
    end
    return Σ
end

"""
Covariance of fWnb. This covariance is a matrix since fWnb is multivariate process.

# Args
- X: a fWnb process
- t: time
- δ: sampling step

# Notes
- filter at [r,c] is the reverse of that at [c,r]
- range of index at [r,c] is reverse of that at [c,r] times -1
By consequent the final matrix Σ(t) = Σ'(-t).
"""
function autocov(X::FractionalWaveletNoiseBank, t::Real, δ::Real; kwargs...)
    fWn_autocov(X.coeffs, X.supps, ss_exponent(X), t, δ; kwargs...)
end

# function autocov(X::FractionalWaveletNoiseBank, t::Real, δ::Real)
#     J = length(X.filters)
#     H = ss_exponent(X)
#     Σ = zeros(Real, (J,J))

#     for c=1:J, r=1:J
#         # X.coeffs[r,c] is (filter, range of index)
#         Σ[r,c] = 1/2 * sum(X.coeffs[r,c] .* abs.(t .+ δ*X.supps[r,c]).^(2H))
#     end
#     return Σ
# end

function fWn_covmat(A::AbstractMatrix{<:AbstractVector{<:Real}}, K::AbstractMatrix{<:AbstractVector{<:Real}}, H::Real, G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}; kwargs...)
    @assert size(A) == size(K) && size(A,1) == size(A,2)  "Dimension mismatches."

    J = size(A,1)
    l1, l2 = length(G1), length(G2)
    Σ = zeros(l1*J, l2*J)

    for c=0:l2-1,r=0:l1-1  # inner (row) iteration is on `r`
        Σ[(r*J+1):(r*J+J), (c*J+1):(c*J+J)] = fWn_autocov(A, K, H, G1[r+1]-G2[c+1], δ; kwargs...)
    end
    return Σ  # not forcing symmetry
end

function fWn_covmat(A::AbstractMatrix{<:AbstractVector{<:Real}}, K::AbstractMatrix{<:AbstractVector{<:Real}}, H::Real, G::AbstractVector{<:Real}; kwargs...)
    @assert size(A) == size(K) && size(A,1) == size(A,2)  "Dimension mismatches."

    J = size(A,1)
    l = length(G)
    Σ = zeros(l*J, l*J)

    for c=0:l-1,r=0:l-1
        # use the property:
        # autocov(x, t) = transpose(autocov(x, -t))
        Σ[(r*J+1):(r*J+J), (c*J+1):(c*J+J)] = (r>=c) ? fWn_autocov(A, K, H, G[r+1]-G[c+1], δ; kwargs...) : transpose(Σ[(c*J+1):(c*J+J), (r*J+1):(r*J+J)])
    end

    return Matrix(Symmetric(Σ))  # forcing symmetry
end

function fWn_covmat(A::AbstractMatrix{<:AbstractVector{<:Real}}, K::AbstractMatrix{<:AbstractVector{<:Real}}, H::Real, G::AbstractVector{<:Integer}; kwargs...)
    @assert size(A) == size(K) && size(A,1) == size(A,2)  "Dimension mismatches."

    J = size(A,1)
    l = length(G)
    Σ = zeros(l*J, l*J)

    Σs = [fWn_autocov(A, K, H, G[n]-G[1], 1; kwargs...) for n=1:length(G)]

    for c=0:l-1,r=0:l-1
        Σ[(r*J+1):(r*J+J), (c*J+1):(c*J+J)] = (r>=c) ? Σs[r-c+1] : transpose(Σs[c-r+1])
    end

    return Matrix(Symmetric(Σ))  # forcing symmetry
end

"""
Compute the covariance matrix of standard (continuous time) fractional Wavelet noise bank of Hurst exponent `H` and time lag `δ` between the grid `G1` and `G2`.
"""
function covmat(X::FractionalWaveletNoiseBank, G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, δ::Real; kwargs...)
    return fWn_covmat(X.coeffs, X.supps, ss_exponent(X), G1, G2, δ; kwargs...)
end

function covmat(X::FractionalWaveletNoiseBank, G::AbstractVector{<:Real}, δ::Real; kwargs...)
    return fWn_covmat(X.coeffs, X.supps, ss_exponent(X), G, δ; kwargs...)
end

function covmat(X::FractionalWaveletNoiseBank, G::AbstractVector{<:Integer}; kwargs...)
    return fWn_covmat(X.coeffs, X.supps, ss_exponent(X), G; kwargs...)
end

# """
# Compute the covariance matrix of standard (continuous time) fractional Wavelet noise bank of Hurst exponent `H` and time lag `δ` between the grid `G1` and `G2`.
# """
# function covmat(X::FractionalWaveletNoiseBank, G1::AbstractVector{<:Real}, G2::AbstractVector{<:Real}, δ::Real)
#         J = length(X.filters)
#     l1, l2 = length(G1), length(G2)
#     Σ = zeros(l1*J, l2*J)

#     for c=0:l2-1,r=0:l1-1  # inner (row) iteration is on `r`
#         Σ[(r*J+1):(r*J+J), (c*J+1):(c*J+J)] = autocov(X, G1[r+1]-G2[c+1], δ)
#     end
#     return Σ  # not forcing symmetry
# end

# function covmat(X::FractionalWaveletNoiseBank, G::AbstractVector{<:Real}, δ::Real)
#     J = length(X.filters)
#     l = length(G)
#     Σ = zeros(l*J, l*J)

#     for c=0:l-1,r=0:l-1
#         # use the property:
#         # autocov(x, t) = transpose(autocov(x, -t))
#         Σ[(r*J+1):(r*J+J), (c*J+1):(c*J+J)] = (r>=c) ? autocov(X, G[r+1]-G[c+1], δ) : transpose(Σ[(c*J+1):(c*J+J), (r*J+1):(r*J+J)])
#     end

#     return Matrix(Symmetric(Σ))  # forcing symmetry
# end

# function covmat(X::FractionalWaveletNoiseBank, G::AbstractVector{<:Integer})
#     J = length(X.filters)
#     l = length(G)
#     Σ = zeros(l*J, l*J)

#     Σs = [autocov(X, G[n]-G[1], 1) for n=1:length(G)]

#     for c=0:l-1,r=0:l-1
#         Σ[(r*J+1):(r*J+J), (c*J+1):(c*J+J)] = (r>=c) ? Σs[r-c+1] : transpose(Σs[c-r+1])
#     end

#     return Matrix(Symmetric(Σ))  # forcing symmetry
# end

# covmat(X::FractionalWaveletNoiseBank, l::Integer) = covmat(X, 1:l)


import GeometricFlux: message, update, propagate

in_channel = 10
out_channel = 5
N = 6
adj = [0. 1. 0. 0. 0. 0.;
       1. 0. 0. 1. 1. 1.;
       0. 0. 0. 0. 0. 1.;
       0. 1. 0. 0. 1. 0.;
       0. 1. 0. 1. 0. 1.;
       0. 1. 1. 0. 1. 0.]
ne = [[2], [1,4,5,6], [6], [2,5], [2,4,6], [2,3,5]]

struct NewLayer <: MessagePassing
    adjlist::AbstractVector{<:AbstractVector}
    weight
end

NewLayer(adjm::AbstractMatrix, m, n) = NewLayer(neighbors(adjm), randn(m,n))

(l::NewLayer)(X) = propagate(l, X=X, aggr=:add)
message(n::NewLayer; x_i=zeros(0), x_j=zeros(0)) = x_j
update(::NewLayer; X=zeros(0), M=zeros(0)) = M

X = Array(reshape(1:N*in_channel, in_channel, N))
l = NewLayer(adj, in_channel, out_channel)

message(n::NewLayer; x_i=zeros(0), x_j=zeros(0)) = n.weight' * x_j

@testset "msgpass" begin
    Y = l(X)
    @test size(Y) == (out_channel, N)
end


# FunctionAlgebra.jl

# [[file:~/Documents/Julia/scrap.org::*FunctionAlgebra.jl][FunctionAlgebra.jl:1]]
function Base.:+(f1::Union{Function,Operator}, f2::Union{Function,Operator})
    if f1 == f2
	t -> 2*f1(t)
    else
	t -> f1(t) + f2(t)
    end
end

Base.:+(a::Number, f::Union{Function,Operator}) = t -> a + f(t)
Base.:+(f::Union{Function,Operator}, a::Number) = t -> f(t) + a

#_____________________________________________
# Subtraction
function Base.:-(f1::Union{Function,Operator}, f2::Union{Function,Operator})
    if f1 == f2
	function (t)
	    0
	end
    else
	function (t)
	    f1(t) - f2(t)
	end
    end
end

Base.:-(a::Number, f::Union{Function,Operator}) = t -> a - f(t)
Base.:-(f::Union{Function,Operator}, a::Number) = t -> f(t) - a

#_____________________________________________
# Multiplication
function Base.:*(f1::Function, f2::Function)
    if f1 == f2
	t -> f1(t)*f1(t)
    else
	t -> f1(t)*f2(t)
    end
end

function Base.:*(f1::Operator, f2::Operator)
    t -> (f1(f2))(t)
end

Base.:*(a::Number, f::Union{Function,Operator}) = t -> a*f(t)
Base.:*(f::Union{Function,Operator}, a::Number) = *(a, f)

#_____________________________________________
# Division

function Base.:/(f1::Function, f2::Function)
    if f1 == f2
	t -> 1
    else
	t -> f1(t)/f2(t)
    end
end

Base.:/(a::Number, f::Function) = t -> a/f(t)
Base.:/(f::Function, a::Number) = t -> f(t)/a

#_____________________________________________
# Exponents

function Base.:^(a::Operator, b::Integer)
    function (t)
	foldl((x,y)->a(x),1:b,init=t)
    end
end
# FunctionAlgebra.jl:1 ends here


module Consts

using LinearAlgebra

const I2 = [1 0;0 1]
const X = [0 1;1 0]
const Z = [1 0;0 -1]
const G_X = 0.5 * (kron(X, X, X, X) + I(2^4))
const G_Z = 0.5 * (kron(Z, Z, Z, Z) + I(2^4))
const O = G_X * G_Z

end

function transfer_matrix(A, O = Consts.O)
    A_dag = conj(A)
    uM = reshape(A, :, size(A, 3))
    lM = reshape(A_dag, :, size(A, 3))

    T = lM * O * transpose(uM)
    shape = ntuple(4) do idx
        isodd(idx) ? size(A, 1) : size(A, 2)
    end
    T2 = reshape(T, shape)
    T3 = permutedims(T2, (1,3,2,4))
    return reshape(T3, size(A, 1)^2, size(A, 2)^2)
end

function boundary_vector(A, O)
    T = conj(A) * O * transpose(A)
    return vec(T)
end

function expectation(As, Os)
    lA, rA = first(As), last(As)
    lO, rO = first(Os), last(Os)
    lT = boundary_vector(lA, lO)
    rT = boundary_vector(rA, rO)
    T = rT
    for idx in length(As)-1:-1:2
        A, O = As[idx], Os[idx]
        T = transfer_matrix(A, O) * T
    end
    return transpose(lT) * T
end

function sample_paulistring(L::Int)
    labels_a = rand(0:1, L)
    labels_b = rand(0:1, L)

    return map(zip(labels_a, labels_b)) do (a, b)
        if a == 0 && b == 0
            Consts.I2
        elseif a == 0 && b == 1
            Consts.Z
        elseif a == 1 && b == 0
            Consts.X
        else # 11
            Consts.X * Consts.Z
        end
    end
end

function renyi_entropy(As; nsamples::Int=10^5)
    L = length(As)
    C = 2^L # 4^L / 2^L
    estimate = ThreadsX.sum(1:nsamples) do _
        Os = sample_paulistring(L)
        expectation(As, Os)^4
    end / nsamples * C
    return -log(estimate)
end

function enumerate_paulistring(L::Int)
    labels_a = Iterators.product([0:1 for _ in 1:L]...)
    labels_b = Iterators.product([0:1 for _ in 1:L]...)

    return map(Iterators.product(labels_a, labels_b)) do (as, bs)
        map(zip(as, bs)) do (a, b)
            if a == 0 && b == 0
                Consts.I2
            elseif a == 0 && b == 1
                Consts.Z
            elseif a == 1 && b == 0
                Consts.X
            else # 11
                Consts.X * Consts.Z
            end
        end
    end |> vec
end

function exact_renyi_entropy(As)
    expect = sum(expectation(As, Os)^4 for Os in enumerate_paulistring(length(As)))
    return -log(expect / 2^length(As))
end


struct RecurrencePlan{T}
    A::Vector{T}
    B::Vector{T}
    C::Vector{T}
    rf₀::Vector{T}
    rf₁::Vector{T}
    rf₀inv::Vector{T}
    rf₁inv::Vector{T}
    rb₀::Vector{T}
    rb₁::Vector{T}
    rb₀inv::Vector{T}
    rb₁inv::Vector{T}
end

function RecurrencePlan(α::T,β::T,N::Int) where T
    A  = T[recA(α,β,k) for k=0:N-1]
    B  = T[recB(α,β,k) for k=0:N-1]
    C  = T[recC(α,β,k) for k=1:N  ]
    rf₀ = T[rf0(α,β,k) for k=0:N-1]
    rf₁ = T[rf1(α,β,k) for k=0:N-1]
    rf₀inv = T[inv(rf₀[k]) for k=1:N]
    rf₁inv = T[inv(rf₁[k]) for k=1:N]
    rb₀ = T[rb0(α,β,k) for k=1:N-1]
    rb₁ = T[rb1(α,β,k) for k=1:N-1]
    rb₀inv = T[inv(rb₀[k]) for k=1:N-1]
    rb₁inv = T[inv(rb₁[k]) for k=1:N-1]
    RecurrencePlan(A,B,C,rf₀,rf₁,rf₀inv,rf₁inv,rb₀,rb₁,rb₀inv,rb₁inv)
end

recA(α::T,β::T,k::Int) where T = k == 0 && ((α+β==0)||(α+β==-1)) ? (α+β)/2+1 : (2k+α+β+1)*(2k+α+β+2)/(2*(k+1)*(k+α+β+1))
recB(α::T,β::T,k::Int) where T = k==0&&((α+β==0)||(α+β==-1)) ? (α-β)/2 : (α-β)*(α+β)*(2k+α+β+1)/(2*(k+1)*(k+α+β+1)*(2k+α+β))
recC(α::T,β::T,k::Int) where T = (k+α)*(k+β)*(2k+α+β+2)/((k+1)*(k+α+β+1)*(2k+α+β))

rf0(α::T,β::T,k::Int) where T = (k+α+one(T))/(k+one(T))
rf1(α::T,β::T,k::Int) where T = -(k+β+one(T))/(k+one(T))

rb0(α::T,β::T,k::Int) where T = (k+one(T))/k*(α+β+k+1)/(k+β)*(α+β+2k)/(α+β+2k+2)
rb1(α::T,β::T,k::Int) where T = -(k+one(T))/k*(α+β+k+1)/(k+α)*(α+β+2k)/(α+β+2k+2)

# Forward polynomial recurrence. Select modified algorithms for x≈±1.

function forward_recurrence!(p::AbstractVector,N::Int,θ::Number,plan::RecurrencePlan{T}) where T
    if θ ≤ 1/4
        reinsch_f0!(p,N,sinpi(θ/2),plan)
    elseif θ ≥ 3/4
        reinsch_f1!(p,N,cospi(θ/2),plan)
    else#if 1/4 < θ < 3/4
        orthogonal_polynomial_recurrence!(p,N,cospi(θ),plan)
    end
end

forward_recurrence!(p::AbstractVector,θ::Number,plan::RecurrencePlan{T}) where T =
    forward_recurrence!(p,length(p),θ,plan)

function forward_recurrence!(p::AbstractVector,N::Int,θ::Number,cpθ::Number,spθ::Number,plan::RecurrencePlan{T}) where T
    if θ ≤ 1/4
        reinsch_f0!(p,N,spθ,plan)
    elseif θ ≥ 3/4
        reinsch_f1!(p,N,cpθ,plan)
    else#if 1/4 < θ < 3/4
        orthogonal_polynomial_recurrence!(p,N,cospi(θ),plan)
    end
end

# Forward recurrence

function orthogonal_polynomial_recurrence!(p::AbstractVector,N::Int,x::Number,plan::RecurrencePlan{V}) where V
    A=plan.A
    B=plan.B
    C=plan.C

    T = promote_type(eltype(p),typeof(x),V)
    pk1, pk = one(T), muladd(A[1],x,B[1])
    p[1],p[2] = pk1, pk

    @inbounds for k = 2:N-1
        Ak, Bk, Ck = A[k], B[k], C[k-1]
        pk, pk1 = muladd(muladd(Ak,x,Bk),pk,-Ck*pk1), pk
        p[k+1] = pk
    end
end

# Modified forward recurrence algorithm for θ near 0
# spθ = sinpi(θ/2)

function reinsch_f0!(p::AbstractVector,N::Int,spθ::Number,plan::RecurrencePlan{V}) where V
    A=plan.A
    C=plan.C
    r=plan.rf₀
    rinv=plan.rf₀inv

    T = promote_type(eltype(p),typeof(spθ),V)
    xm1 = -2spθ^2
    p[1] = one(T)
    dk = A[1]*xm1*rinv[1]
    p[2] = (p[1]+dk)*r[1]

    @inbounds for k = 2:N-1
        Ak, Ck, rk, rkinv = A[k], C[k-1], r[k], rinv[k]
        temp = muladd(Ak*xm1,p[k],Ck*dk)
        dk, p[k+1] = temp*rkinv, muladd(rk,p[k],temp)
    end
end

# Modified forward recurrence for θ near 1
# cpθ = cospi(θ/2)

function reinsch_f1!(p::AbstractVector,N::Int,cpθ::Number,plan::RecurrencePlan{V}) where V
    A=plan.A
    C=plan.C
    r=plan.rf₁
    rinv=plan.rf₁inv

    T = promote_type(eltype(p),typeof(cpθ),V)
    xp1 = 2cpθ^2
    p[1] = one(T)
    dk = A[1]*xp1*rinv[1]
    p[2] = (p[1]+dk)*r[1]

    @inbounds for k = 2:N-1
        Ak, Ck, rk, rkinv = A[k], C[k-1], r[k], rinv[k]
        temp = muladd(Ak*xp1,p[k],Ck*dk)
        dk, p[k+1] = temp*rkinv, muladd(rk,p[k],temp)
    end
end

# Backward polynomial recurrence. Select modified algorithms for x≈±1.

function backward_recurrence(c::AbstractVector,N::Int,θ::Number,plan::RecurrencePlan{T}) where T
    if θ ≤ 1/4
        reinsch_b0(c,N,sinpi(θ/2),plan)
    elseif θ ≥ 3/4
        reinsch_b1(c,N,cospi(θ/2),plan)
    else#if 1/4 < θ < 3/4
        clenshaw(c,N,cospi(θ),plan)
    end
end

backward_recurrence(c::AbstractVector,θ::Number,plan::RecurrencePlan{T}) where T =
    backward_recurrence(c,length(c),θ,plan)

function backward_recurrence(c::AbstractVector,N::Int,θ::Number,cpθ::Number,spθ::Number,plan::RecurrencePlan{T}) where T
    if θ ≤ 1/4
        reinsch_b0(c,N,spθ,plan)
    elseif θ ≥ 3/4
        reinsch_b1(c,N,cpθ,plan)
    else#if 1/4 < θ < 3/4
        clenshaw(c,N,cospi(θ),plan)
    end
end

# Clenshaw-Smith algorithm

function clenshaw(c::AbstractVector,N::Int,x::Number,plan::RecurrencePlan{V}) where V
    A=plan.A
    B=plan.B
    C=plan.C

    T = promote_type(eltype(c),typeof(x),V)
    bk, bk1, bk2 = zero(T), zero(T), zero(T)

    @inbounds for k = N:-1:2
        ck, Ak, Bk, Ck = c[k], A[k], B[k], C[k]
        bk = muladd(muladd(Ak,x,Bk),bk1,muladd(-Ck,bk2,ck))
        bk2, bk1, bk = bk1, bk, bk2
    end

    @inbounds ck, Ak, Bk, Ck = c[1], A[1], B[1], C[1]
    muladd(muladd(Ak,x,Bk),bk1,muladd(-Ck,bk2,ck))
end

# Modified Clenshaw-Smith algorithm for θ near 0
# spθ = sinpi(θ/2)

function reinsch_b0(c::AbstractVector,N::Int,spθ::Number,plan::RecurrencePlan{V}) where V
    A=plan.A
    C=plan.C
    r=plan.rb₀
    rinv=plan.rb₀inv

    T = promote_type(eltype(c),typeof(spθ),V)
    xm1 = -2spθ^2
    uk, dk = zero(T), zero(T)

    @inbounds for k=N:-1:2
        ck, Ak, Ck, rk, rkinv = c[k], A[k], C[k], r[k-1], rinv[k-1]
        temp = muladd(Ak*xm1,uk,muladd(Ck,dk,ck))
        dk, uk = temp*rk, muladd(rkinv,uk,temp)
    end

    @inbounds ck, Ak, Ck = c[1], A[1], C[1]
    muladd(Ak*xm1,uk,muladd(Ck,dk,ck))
end

# Modified Clenshaw-Smith algorithm for θ near 1
# cpθ = cospi(θ/2)

function reinsch_b1(c::AbstractVector,N::Int,cpθ::Number,plan::RecurrencePlan{V}) where V
    A=plan.A
    C=plan.C
    r=plan.rb₁
    rinv=plan.rb₁inv

    T = promote_type(eltype(c),typeof(cpθ),V)
    xp1 = 2cpθ^2
    uk, dk = zero(T), zero(T)

    @inbounds for k=N:-1:2
        ck, Ak, Ck, rk, rkinv = c[k], A[k], C[k], r[k-1], rinv[k-1]
        temp = muladd(Ak*xp1,uk,muladd(Ck,dk,ck))
        dk, uk = temp*rk, muladd(rkinv,uk,temp)
    end

    @inbounds ck, Ak, Ck = c[1], A[1], C[1]
    muladd(Ak*xp1,uk,muladd(Ck,dk,ck))
end


using OffsetArrays

"""
Legendre polynomial on the interval [-1..+1].
"""
@inline function _legendre(j::Int, x::T) where {T}
    if j <= 0
        return one(T)
    elseif j == 1
        return x
    else
        return ( (2j-1) * _legendre(j-1, x) * x - (j-1) * _legendre(j-2, x) ) / j
    end
end

"""
Legendre basis on the interval [0..+1].
"""
struct Legendre{T, LT} <: Basis{T}
    b::LT
    n::Int

    function Legendre{T}(n::Integer) where {T}
        p = n-1
        b = OffsetArray([y -> _legendre(i, 2y-1) * sqrt(T(2i+1)) for i in 0:p], 0:p)
        new{T, typeof(b)}(b, n)
    end

end

Legendre(::Type{T}, n::Integer) where {T} = Legendre{T}(n)
Legendre(n::Integer) = Legendre(Float64, n)

function _eval(L::Legendre{LT}, x::DT, j::Int) where {LT,DT}
    local T = promote_type(LT, DT)
    @assert j ≥ 0 && j < L.n
    _legendre(j, 2x-1) * sqrt(T(2j+1))
end

(L::Legendre)(x::Number, j::Integer) = L.b[j](x)

basis(L::Legendre) = L.b
nbasis(L::Legendre) = L.n
order(L::Legendre) = nbasis(L)
degree(L::Legendre) = nbasis(L) - 1

Base.eltype(::Legendre{T}) where {T} = T
Base.eachindex(L::Legendre) = eachindex(L.b)
Base.axes(L::Legendre) = (Inclusion(0..1), eachindex(L))

Base.hash(L::Legendre, h::UInt) = hash(L.n, h)
Base.:(==)(L1::Legendre, L2::Legendre) = (L1.n == L2.n)
Base.isequal(L1::Legendre{T1}, L2::Legendre{T2}) where {T1,T2} = (T1 == T2 && L1 == L2)

Base.getindex(L::Legendre, x::Number, j::Integer) = L(x,j)
Base.getindex(L::Legendre, x::Number,  ::Colon) = [b(x) for b in L.b]
Base.getindex(L::Legendre, X::AbstractVector, j::Integer) = L.(X,j)
Base.getindex(L::Legendre, X::AbstractVector,  ::Colon) = [b(x) for x in X, b in L.b]

## Derivative

"""
Derivative of Legendre polynomial on the interval [-1..+1].
"""
@inline function _legendre_derivative(j::Int, x::T) where {T}
    if j <= 0
        return zero(T)
    elseif j == 1
        return one(T)
    else
        return ( (2j-1) * _legendre(j-1, x) + (2j-1) * _legendre_derivative(j-1, x) * x - (j-1) * _legendre_derivative(j-2, x) ) / j
    end
end

@simplify *(D::Derivative, L::Legendre) = Mul(D,L)

const LegendreDerivative = QMul2{<:Derivative,<:Legendre}

"""
Evaluate derivative of Legendre polynomial on the interval [0..+1].
"""
function _eval(D::LegendreDerivative, x::DT, j::Int) where {DT}
    @assert j ≥ 0 && j < nbasis(D.B)
    _legendre_derivative(j, promote_type(eltype(D.B), DT)(2x-1)) * 2
end

Base.getindex(D::LegendreDerivative, x::Number, j::Integer) = _eval(D, x, j)
Base.getindex(D::LegendreDerivative, x::Number,  ::Colon) = [_eval(D, x, j) for j in eachindex(D.B)]
Base.getindex(D::LegendreDerivative, X::AbstractVector, j::Integer) = [_eval(D, x, j) for x in X]
Base.getindex(D::LegendreDerivative, X::AbstractVector,  ::Colon) = [_eval(D, x, j) for x in X, j in eachindex(D.B)]

Base.adjoint(L::Legendre) = Derivative(axes(L,1)) * L


function edit_distance(a, b, submat::AbstractSubstitutionMatrix{T}, insertion_cost::T, deletion_cost::T) where T
    m = length(a)
    n = length(b)
    trace = Matrix{Trace}(undef, (m + 1, n + 1))
    D = Vector{T}(undef, m + 1)
    D[1] = T(0)
    for i in 1:m
        D[i+1] = i * insertion_cost
        trace[i+1,1] = TRACE_INSERT
    end
    for j in 1:n
        d_diag = D[1]
        D[1] = j * deletion_cost
        trace[1,j+1] = TRACE_DELETE
        for i in 1:m
            ins = D[i]   + insertion_cost
            del = D[i+1] + deletion_cost
            mat = d_diag + submat[a[i],b[j]]
            # find the best score and its trace
            best = min(del, ins, mat)
            t = TRACE_NONE
            del == best && (t |= TRACE_DELETE)
            ins == best && (t |= TRACE_INSERT)
            mat == best && (t |= TRACE_MATCH)
            d_diag = D[i+1]
            D[i+1] = best
            trace[i+1,j+1] = t
        end
    end
    return D[m+1], trace, (m, n)
end

function edit_traceback(a, b, trace, endpos)
    anchors = Vector{AlignmentAnchor}()
    i, j = endpos
    @start_traceback
    while i ≥ 1 || j ≥ 1
        t = trace[i+1,j+1]
        if t & TRACE_MATCH > 0
            if a[i] == b[j]
                @anchor OP_SEQ_MATCH
            else
                @anchor OP_SEQ_MISMATCH
            end
        elseif t & TRACE_DELETE > 0
            @anchor OP_DELETE
        elseif t & TRACE_INSERT > 0
            @anchor OP_INSERT
        else
            error("failed to trace back")
        end
    end
    @finish_traceback
    return AlignedSequence(a, anchors)
end


# K. A. Fujita, Y. Toyoshima, S. Uda, Y. I. Ozaki, H. Kubota, and S. Kuroda
# Decoupling of receptor and downstream signals in the Akt pathway by its low-pass filter characteristics
# https://doi.org/10.1126/scisignal.2000810

using Logging

using StructuralIdentifiability

logger = Logging.SimpleLogger(stdout, Logging.Debug)
global_logger(logger)

ode = @ODEmodel(
    EGFR'(t) = EGFR_turnover * pro_EGFR(t) + EGF_EGFR(t) * reaction_1_k2 - EGFR(t) * EGFR_turnover - EGF_EGFR(t) * reaction_1_k1,
    pEGFR'(t) = EGF_EGFR(t) * reaction_9_k1 - pEGFR(t) * reaction_4_k1 + pEGFR_Akt(t) * reaction_2_k2 + pEGFR_Akt(t) * reaction_3_k1 - Akt(t) * pEGFR(t) * reaction_2_k1,
    pEGFR_Akt'(t) = Akt(t) * pEGFR(t) * reaction_2_k1 - pEGFR_Akt(t) * reaction_3_k1 - pEGFR_Akt(t) * reaction_2_k2,
    Akt'(t) = pAkt(t) * reaction_7_k1 + pEGFR_Akt(t) * reaction_2_k2 - Akt(t) * pEGFR(t) * reaction_2_k1,
    pAkt'(t) = pAkt_S6(t) * reaction_5_k2 - pAkt(t) * reaction_7_k1 + pAkt_S6(t) * reaction_6_k1 + pEGFR_Akt(t) * reaction_3_k1 - S6(t) * pAkt(t) * reaction_5_k1,
    S6'(t) = pAkt_S6(t) * reaction_5_k2 + pS6(t) * reaction_8_k1 - S6(t) * pAkt(t) * reaction_5_k1,
    pAkt_S6'(t) = S6(t) * pAkt(t) * reaction_5_k1 - pAkt_S6(t) * reaction_6_k1 - pAkt_S6(t) * reaction_5_k2,
    pS6'(t) = pAkt_S6(t) * reaction_6_k1 - pS6(t) * reaction_8_k1,
    EGF_EGFR'(t) = EGF_EGFR(t) * reaction_1_k1 - EGF_EGFR(t) * reaction_9_k1 - EGF_EGFR(t) * reaction_1_k2,
    y1(t) = a1 * (pEGFR(t) + pEGFR_Akt(t)),
    y2(t) = a2 * (pAkt(t) + pAkt_S6(t)),
    y3(t) = a3 * pS6(t)
)

@time println(assess_identifiability(ode))


export zcom_NO
function zcom_NO(x)::Float64
    """
    Computes the z component of the center of mass of the NO projectile
    """
    z_N = x[1, 3]
    z_O = x[2, 3]
    zcom = (z_N * m_N + z_O*m_O)/(m_N + m_O)
    return zcom
end


#module GraphicReasoning

using Calculus

#puts a semiCircle on the graph starting at x=a,y=fa, and going until x=b.
#integral is (fa+pi/2)*(b-a)
function semiCircle(a::Int,fa,b::Int)
	r = (b-a)/2.0
	x0 = (b+a)/2.0

	flip = rand([true false])
	if flip
		x -> fa - sqrt((r)^2-(x-(x0))^2)
	else
		x -> fa + sqrt((r)^2-(x-(x0))^2)
	end
end

function makeLineEnd(exc)
	y = range(-5,5,length=11)
	let cond = true
		while(cond)
		fb = Int(rand(y))
		if fb != exc
			return fb
		end
		end
	end
end

#integral is line(b)*(b-a)
function line(a::Int,fa,b::Int)
	y = range(-5,5,length=11)
	cond = true
	fb = makeLineEnd(fa)

	m = (fb-fa)/(b-a)

	x -> m*(x-a)+fa
end

function plotQuestion()
	p = plot()

	d = [2,3,4]
	y_start = 0
	pos = -5
	ys = []
	while pos<=5
		ds = rand(d)
		is_line = rand([true,false])
			if is_line
				f = line(pos,y_start,pos+ds)
				y_start = Int(f(pos+ds))
				txt = "($(pos+ds),$(y_start))"
				plot!(f,pos,pos+ds;annotation=(pos+ds,y_start,text(txt,8)))
				push!(ys,y_start)
			else
				f = semiCircle(pos,y_start,pos+ds)
				plot!(f,pos,pos+ds)
				annotate!([(pos+ds,y_start,Plots.text("($(pos+ds),$(y_start))",:left,8))])
				annotate!([(pos+(ds/2),f(pos+(ds/2)),Plots.text("($pos,$(f(pos+(ds/2))))",:bottom,8))])
				push!(ys,f(pos+(ds/2)))
			end
			if pos+ds>5
				push!(ys,pos+ds)
				break
			end
			pos = pos+ds
	end
	#going through x=-5 to x=5
	y0 = min(minimum(ys),-5)
	yf = max(maximum(ys),5)

	plot!(
	xlims=[y0,yf],
	ylims=[y0,yf],
	title="graph of f",
	minorgrid=true,
	size=(600,600),
	legend=false
	)

	savefig("question.png")

end

#end


using Random

@testset "predict" begin
    Random.seed!(100)

    @model function linear_reg(x, y, σ = 0.1)
        β ~ Normal(0, 1)

        for i ∈ eachindex(y)
            y[i] ~ Normal(β * x[i], σ)
        end
    end

    @model function linear_reg_vec(x, y, σ = 0.1)
        β ~ Normal(0, 1)
        y ~ MvNormal(β .* x, σ)
    end

    f(x) = 2 * x + 0.1 * randn()

    Δ = 0.1
    xs_train = 0:Δ:10; ys_train = f.(xs_train);
    xs_test = [10 + Δ, 10 + 2 * Δ]; ys_test = f.(xs_test);

    # Infer
    m_lin_reg = linear_reg(xs_train, ys_train);
    chain_lin_reg = sample(m_lin_reg, NUTS(100, 0.65), 200);

    # Predict on two last indices
    m_lin_reg_test = linear_reg(xs_test, fill(missing, length(ys_test)));
    predictions = Turing.Inference.predict(m_lin_reg_test, chain_lin_reg)

    ys_pred = vec(mean(Array(group(predictions, :y)); dims = 1))

    @test sum(abs2, ys_test - ys_pred) ≤ 0.1

    # Ensure that `rng` is respected
    predictions1 = let rng = MersenneTwister(42)
        predict(rng, m_lin_reg_test, chain_lin_reg[1:2])
    end
    predictions2 = let rng = MersenneTwister(42)
        predict(rng, m_lin_reg_test, chain_lin_reg[1:2])
    end
    @test all(Array(predictions1) .== Array(predictions2))

    # Predict on two last indices for vectorized
    m_lin_reg_test = linear_reg_vec(xs_test, missing);
    predictions_vec = Turing.Inference.predict(m_lin_reg_test, chain_lin_reg)
    ys_pred_vec = vec(mean(Array(group(predictions_vec, :y)); dims = 1))

    @test sum(abs2, ys_test - ys_pred_vec) ≤ 0.1

    # Multiple chains
    chain_lin_reg = sample(m_lin_reg, NUTS(100, 0.65), MCMCThreads(), 200, 2);
    m_lin_reg_test = linear_reg(xs_test, fill(missing, length(ys_test)));
    predictions = Turing.Inference.predict(m_lin_reg_test, chain_lin_reg)

    @test size(chain_lin_reg, 3) == size(predictions, 3)

    for chain_idx in MCMCChains.chains(chain_lin_reg)
        ys_pred = vec(mean(Array(group(predictions[:, :, chain_idx], :y)); dims = 1))
        @test sum(abs2, ys_test - ys_pred) ≤ 0.1
    end

    # Predict on two last indices for vectorized
    m_lin_reg_test = linear_reg_vec(xs_test, missing);
    predictions_vec = Turing.Inference.predict(m_lin_reg_test, chain_lin_reg)

    for chain_idx in MCMCChains.chains(chain_lin_reg)
        ys_pred_vec = vec(mean(
            Array(group(predictions_vec[:, :, chain_idx], :y));
            dims = 1
        ))
        @test sum(abs2, ys_test - ys_pred_vec) ≤ 0.1
    end
end


uniform(dim, n) = rand(dim, n)

true_uniform(n, m) = reduce(hcat, collect.(Iterators.product((1:n) / n .- 1 / 2n, (1:m) / m .- 1 / 2m)))

function net(nline, npoint)
    dx = 1 / npoint
    xs = range(dx / 2, length = npoint, step = dx)
    dy = 1 / nline
    ys = range(dy / 2, length = nline, step = dy)

    a = [[x, y] for x in xs, y in ys]
    b = [[y, x] for x in xs, y in ys]
    reduce(hcat, [a[:]; b[:]])
end

function circle_cluster(n)
    ws = range(0., length=n, step=2π / n)
    [f(w) for f in (cos, sin), w in ws]
end

normal_cluster(dim, n) = randn(dim, n)

function cross_cluster(w)
    T = [
        cos(w) sin(w)
        -sin(w) cos(w)
    ]
    x = [
        -2 -1 0 1 2 0 0 0 0
        0 0 0 0 0 -2 -1 1 2
    ]
    mapreduce(x -> T * x, hcat, eachcol(x))
end

clusters(centers, cluster) = mod.(reduce(hcat, [cluster() .+ c for c in eachcol(centers)]), 1.)

function voronoi_metric(x, support)
    a = colwise(PeriodicEuclidean(ones(2)), support, x) |> sort
    abs(a[2] - a[1])
end

function voronoi(n, support)
    ndims = size(support, 1)
    all_points = rand(ndims, 1_000_000)

    a = [voronoi_metric(x, support) for x in eachcol(all_points)]

    ix = sortperm(a)
    points = all_points[:, ix[1:n]]

    # scatter(points[1, :], points[2, :], markersize=2)
    # scatter!(support[1, :], support[2, :]) |>  display

    points
end


@inline M{T}(i::T)           = (i == 1) ? T(1) : T(2)^(i - T(1)) + T(1)
@inline cc_iM(i::Int)           = (i == 1) ? 1 : Int(log2(i - 1) + 1)
@inline dM{T}(i::T)           = (i < T(3)) ? T(2^(i - 1)) : T(2^(i - 2))
@inline cc_g(i::Int, j::Int)     = (i == 1) ?  0.5  : (j - 1) / (M(i) - 1.0)
@inline cc_g(i::Int)            = Float64[cc_g(i, j) for j = 1:M(i)]
@inline cc_dg(i::Int, j::Int)    = i == 2 ? [0.0, 1.0][j] : cc_g(i, 2j)
@inline cc_dg(i::Int)           = Float64[cc_dg(i, j) for j = 1:dM(i)]

@inline function cc_bf_l(x::Float64, xij::Float64, mi)
    if (mi == 1)
        return 1.0
    end
    dx = (1.0 - (mi - 1.0) * abs(x - xij))
    return (dx > 0.0) ? dx : 0.0
end

cc_bf_l(x::Float64, xij::Float64) = cc_bf_l(x, xij, M(level(xij)))
cc_bf_l(x::Vector{Float64}, xij::Vector{Float64}) = prod(map(cc_bf_l, x, xij))

@inline function cc_bf_q(x::Float64, xij::Float64, mi)
	if (mi == 1)
		return 1.0
	end
	dx = 1.0 - ((mi - 1.0) * (x - xij))^2
	return (dx > 0.0) ?  dx : 0.0
end

cc_bf_q(x::Float64, xij::Float64) = cc_bf_q(x, xij, cc_M(level(xij)))
cc_bf_q(x::Vector{Float64}, xij::Vector{Float64}) = map(cc_bf_q, x, xij)

function cc_simpsonsw(i::Int, j::Int)
    @assert 1 ≤ j ≤ M(i)
    if i == 1
        return 1.0
    end
    w = 1 / (M(i) - 1) / 3
    if j == 1 || j == M(i)
        return w
    else
        w *= mod(j, 2) == 0 ? 4 : 2
    end
    return w
end

cc_simpsonsw(i::Int) = Float64[cc_simpsonsw(i, j) for j = 1:M(i)]

function cc_dsimpsonsw(i::Int)
    if i == 1
        return cc_simpsonsw(i)
    elseif i == 2
        return cc_simpsonsw(i) - [0;1.0;0]
    else
        return cc_simpsonsw(i) - vec(Float64[cc_simpsonsw(i - 1)'zeros(1, M(i - 1))])[1:M(i - 1) - 1]
    end
end


module Monotones

export relative_entropy
import RTFOOL.Context

function relative_entropy(ps::AbstractVector{Float64}, qs::AbstractVector{Float64})
    if length(ps) != length(qs)
        error("states have different sizes")
    elseif !(sum(ps) ≈ 1.0) || !(sum(qs) ≈ 1.0)
        error("state does not sum to 1.0")
    end
    sum(map((p,q) -> (p > 0.0) ? p*log(p/q) : 0.0, ps, qs))
end
relative_entropy(ctx::Context) = relative_entropy(ctx.system_state, ctx.bath_state)

end


using Revise
using Test, QuHamiltonian, Lattices

# set up something useful
using LinearAlgebra, SparseArrays
⊗ = kron
II = Matrix{Float64}(I, 2, 2)
##########################

@testset "vertex sum" begin

    h = @sum Z[@vertex(i)]
    ltc = Chain(4)
    H = h(ltc)

    A = Z ⊗ II ⊗ II ⊗ II + II ⊗ Z ⊗ II ⊗ II + II ⊗ II ⊗ Z ⊗ II + II ⊗ II ⊗ II ⊗ Z

    lhs = Bit[0, 0, 0, 0]
    rhs = Bit[0, 0, 0, 0]

    for i = 1:2^4
        for j = 1:2^4
            @test H[lhs, rhs] == A[convert(Int, lhs) + 1, convert(Int, rhs) + 1]
            lhs << 1
        end
        rhs << 1
    end

    @test Matrix(H) == A
    @test sparse(H) == sparse(A)
end

@testset "edge sum" begin

    @nearest i, j
    h = @sum Z[i]Z[j]
    ltc = Chain(4)
    H = h(ltc)

    lhs = Bit[0, 0, 0, 0]
    rhs = Bit[0, 0, 0, 0]

    A = Z ⊗ Z ⊗ II ⊗ II + II ⊗ Z ⊗ Z ⊗ II + II ⊗ II ⊗ Z ⊗ Z + Z ⊗ II ⊗ II ⊗ Z
    Matrix(H)
    for i = 1:1<<4
        for j = 1:1<<4
            @test H[lhs, rhs] == A[convert(Int, lhs) + 1, convert(Int, rhs) + 1]
            lhs << 1
        end
        rhs << 1
    end

    @test Matrix(H) == A
    @test sparse(H) == sparse(A)
end

@testset "add/sub" begin

    @nearest i, j
    @vertex l
    h = @sum(X[i]X[j]) - 0.1 * @sum(Z[l])
    ltc = Chain(4)
    H = h(ltc)

    lhs = Bit[0, 0, 0, 0]
    rhs = Bit[0, 0, 0, 0]

    H[lhs, rhs]
end

# @testset "multiple vertex sum" begin
#
#     @vertex i, j
#     h = @sum Z[i]X[j] + @sum (Z * X)[i]
#     ltc = Chain(4)
#     H = h(ltc)
#
#     A =
#     # (2, 1)
#     X ⊗ Z ⊗ II ⊗ II +
#     # (3, 1)
#     X ⊗ II ⊗ Z ⊗ II +
#     # (4, 1)
#     X ⊗ II ⊗ II ⊗ Z +
#
#     # (1, 2)
#     Z ⊗ X ⊗ II ⊗ II +
#     # (3, 2)
#     II ⊗ X ⊗ Z ⊗ II +
#     # (4, 2)
#     II ⊗ X ⊗ II ⊗ Z +
#
#     # (1, 3)
#     Z ⊗ II ⊗ X ⊗ II +
#     # (2, 3)
#     II ⊗ Z ⊗ X ⊗ II +
#     # (4, 3)
#     II ⊗ II ⊗ X ⊗ Z +
#
#     # (1, 4)
#     Z ⊗ II ⊗ II ⊗ X +
#     # (2, 4)
#     II ⊗ Z ⊗ II ⊗ X +
#     # (3, 4)
#     II ⊗ II ⊗ Z ⊗ X +
#
#     # (1, 1)
#     (Z * X) ⊗ II ⊗ II ⊗ II +
#     # (2, 2)
#     II ⊗ (Z * X) ⊗ II ⊗ II +
#     # (3, 3)
#     II ⊗ II ⊗ (Z * X) ⊗ II +
#     # (4, 4)
#     II ⊗ II ⊗ II ⊗ (Z * X)
#
#     Matrix(H) == A
#
# end


export AddAttention

doc"""
    AddAttention(x1, x2)

```math
f(x_{i},x_{j}) = W^{T} \sigma (W_{1}x_{i} + W_{2}x_{j} + b) + b
```
"""
struct AddAttention
    l1
    l2
end

function AddAttention(::Type{T}, insize::Int, outsize::Int) where T
    l1 = Linear(T, 2insize, insize)
    l2 = Linear(T, insize, outsize)
    AddAttention(l1, l2)
end

function (f::AddAttention)(x1::Var, x2::Var)
    h = pairwise(x1, x2)
    h = f.l1(h)
    h = relu(h)
    f.l2(h)
end


#На прямоугольном поле произвольных размеров расставить маркеры в виде "шахматных" клеток, начиная с юго-западного угла поля, когда каждая отдельная "шахматная" клетка имеет размер n x n клеток поля (n - это параметр функции).
#Начальное положение Робота - произвольное, конечное - совпадает с начальным. Клетки на севере и востоке могут получаться "обрезанными" - зависит от соотношения размеров поля и "шахматных" клеток.
function chess(r, n)
    z = 0
    m = 0
    z, m = corner(r)
    side = Ost;
    x = 0;
    y = 0;
    if ((x ÷ n) + (y ÷ n)) % 2 == 0
        putmarker!(r)
    end
    while isborder(r,Nord) == false
        x, y = chess_line!(side, n, x, y, r)
        move!(r,Nord);
        y = y + 1;
        side = revers(side);
        put!(r, x, y, n)
    end
    x, y = chess_line!(side, n, x, y, r)
    go_home!(r, z, m)
end
function movements_mark!(r,side)
    while (isborder(r,side) ==  false) && (ismarker(r) == true)
        move!(r,side)
    end
end
function revers(side)
    side = HorizonSide((Int(side)+2)%4)
    return side
end
function count_move!(r,side,count)
    for i = 1:count
        move!(r,side)
    end
end
function movements_to_side!(r,side)
    count = 0
    while isborder(r,side) ==  false
        move!(r,side)
        count = count + 1
    end
    return count
end
function corner(r)
    x = 0
    y = 0
    while (isborder(r, Sud) == false) || (isborder(r, West) == false)
        y = y + movements_to_side!(r, Sud)
        x = x + movements_to_side!(r, West)
    end
    return x, y
end
function go_home!(r, x, y)
    corner(r)
    count_move!(r, Nord, y)
    count_move!(r, Ost, x)
end
function chess_line!(side, n, x, y, r)
    while isborder(r, side) == false
        move!(r, side)
        if side == Ost
            x = x + 1;
        else
            x = x - 1;
        end

        put!(r, x, y, n)
    end
    return x, y
end
function put!(r, x, y, n)
    if ((x ÷ n) + (y ÷ n)) % 2 == 0
        putmarker!(r)
    end
end
n = 3
chess(r,n)


using Gadfly

x1 = rand(40)
y1 = 4.*x1 .+ 2 .+randn(40)
x2 = rand(40)
y2 = -6.*x2 .+ 3 .+ randn(40)
x  = [x1;x2]
y  = [y1;y2]
col = [fill("Slope 4",40); fill("Slope -6",40)]
plot(x=x,y=y,colour=col, Geom.point, Geom.smooth(method=:lm))


using Base.Test
using JointPosteriors

model = ( ProbabilityVector(3), )

struct BinaryClassificationData
  X::Array{Int64,1}
  freq::Array{Int64,1}
  NmX::Array{Int64, 1}
  α_m_m1::Float64
  β_m_m1::Float64
  α_p_m1::Float64
  β_p_m1::Float64
  α_τ_m1::Float64
  β_τ_m1::Float64
end

function log_density(data, p)
    @inbounds log_π = data.α_m_m1 * log(p[2]) + data.β_m_m1 * log(1 - p[2]) + data.α_p_m1*log(p[3]) + data.β_p_m1 * log(1 - p[3]) + data.α_τ_m1 * log(p[1]) + data.β_τ_m1 * log(1 - p[1])

    @inbounds @simd for i ∈ eachindex(data.X)
        log_π += data.freq[i] * log( p[1] * (1 - p[2])^data.X[i] * p[2]^(data.NmX[i]) + (1-p[1]) * p[3]^data.X[i] * (1-p[3])^data.NmX[i] )
    end
    log_π
end

function BinaryClassificationData(X::Array{Int, 1}, freq::Array{Int,1}, n::Int; αm::Real = 1, βm::Real = 1, αp::Real = 1, βp::Real = 1, ατ::Real = 1, βτ::Real = 1)
  BinaryClassificationData(X, freq, n .- X, αm - 1, βm - 1, αp - 1, βp - 1, ατ - 1, βτ - 1)
end
X = [0, 1, 2, 3, 4, 7, 8, 9];
freq = [10, 2, 2, 1, 2, 3, 2, 16];
data = BinaryClassificationData(X, freq, 9, βm = 2, βp = 2);

#mod = Model(model)
#md = LogDensities.ModelDiff(log_density, mod.diff_buffer, data);
#jp = fit(mod, data)

function run_tests(model, data, μ, σ, quants, f, tol = 1e-2)
    builds = (SparseQuadratureGrids.AdaptiveRaw, SparseQuadratureGrids.Adaptive, SparseQuadratureGrids.SmolyakRaw, SparseQuadratureGrids.Smolyak)
    q_rules = (SparseQuadratureGrids.GenzKeister, SparseQuadratureGrids.KronrodPatterson)
    @testset begin
        @testset "Build: $b, Q rule: $q." for b ∈ builds, q ∈ q_rules
            mod = Model(model, b{q})
            fit(mod, data)
            jp = fit(mod, data)
            m_grid = marginal(jp, f)
            m_norm = marginal(jp, f, Normal)
            @test isapprox(m_norm.μ, μ, rtol = tol)
            @test isapprox(m_norm.σ, σ, rtol = tol)
            @test isapprox(quantile(m_grid, .025), quants[1], rtol = tol)
            @test isapprox(quantile(m_grid, .25), quants[2], rtol = tol)
            @test isapprox(quantile(m_grid, .5), quants[3], rtol = tol)
            @test isapprox(quantile(m_grid, .75), quants[4], rtol = tol)
            @test isapprox(quantile(m_grid, .975), quants[5], rtol = tol)
            println("\n")
            println(quantile.(m_norm, [.025,.25,.5,.75,.975]))
            println("\n")
            @test isapprox(quantile(m_norm, .025), quants[1], rtol = tol)
            @test isapprox(quantile(m_norm, .25), quants[2], rtol = tol)
            @test isapprox(quantile(m_norm, .5), quants[3], rtol = tol)
            @test isapprox(quantile(m_norm, .75), quants[4], rtol = tol)
            @test isapprox(quantile(m_norm, .975), quants[5], rtol = tol)
        end
    end
end

quants = [0.391, 0.495, 0.55, 0.602, 0.696]
run_tests(model, data, 0.5504, 0.077, quants, p -> p[1], 10^-1.5)


# generate a tridiagonal matrix
function tridiagm(a, b, c, n)
       dd, du = ones(n), ones(n - 1)
       b*diagm(0 => dd) + a*diagm(-1 => du) + c*diagm(1 => du)
end

SUITE["Balls"] = BenchmarkGroup()
SUITE["Linear map"] = BenchmarkGroup()
begin
    for set_type in (Ball1, Ball2, BallInf)
        for n in (1, 10, 100, 1000)
            X = set_type(ones(n), 0.1)
            d = [i/n for i in 1:n]
            SUITE["Balls"][string(set_type), n] = @benchmarkable σ($d, $X)

            A = tridiagm(1, -2, 0.5, n)
            Y = A * X
            SUITE["Linear map"][string(set_type), n] = @benchmarkable σ($d, $Y)
        end
    end
end


#!/usr/bin/env julia
#

using Printf

e_str1 = "2.718"
e = parse(Float64, e_str1)
println(5e)
#> 13.5914
num_15 = parse(Int, "15")
println(3num_15)
#> 45

@printf "e = %0.2f\n" e
#> 2.718
e_str2 = @sprintf("%0.3f", e)

# show that the 2 strings are the same
println("e_str1 == e_str2: $(e_str1 == e_str2)")
#> e_str1 == e_str2: true

# available number format characters are f, e, g, c, s, p, d
# (pi is a predefined constant; however, since its type is
# "MathConst" it has to be converted to a float to be formatted)
@printf "fix trailing precision: %0.3f\n" float(pi)
#> fix trailing precision: 3.142
@printf "scientific form: %0.6e\n" 1000pi
#> scientific form: 3.141593e+03
# g is not implemented yet
@printf "a character: %c\n" 'α'
#> a character: α
@printf "a string: %s\n" "look I'm a string!"
#> a string: look I'm a string!
@printf "right justify a string: %50s\n" "width 50, text right justified!"
#> right justify a string:                    width 50, text right justified!
@printf "a pointer: %p\n" 100000000
#> a pointer: 0x0000000005f5e100
@printf "print a integer: %d\n" 1e10
#> print an integer: 10000000000


using Findpeaks

using DelimitedFiles

using Plots
default(legend=false)

data = readdlm("../assets/example_spectrum.txt")
x = data[:, 1]
y = data[:, 2]

peaks = findpeaks(y, x, min_prom=1000.)

plot(x, y, title="Prominent peaks")
scatter!(x[peaks], y[peaks])
savefig("prom_peaks.png")

sep = 0.2

peaks = findpeaks(y, x, min_dist=sep)

plot(x, y, title="Peaks at least $sep units apart")
scatter!(x[peaks], y[peaks])
savefig("dist_peaks.png")


export incrementer, multiplier, inverse_decay, cached_inverse_decay, sample_avg, cached_sample_avg, exp_decay

"""
    incrementer(;start=1, step=1)

# Example
```julia
julia> f = incrementer();

julia> [f() for _ in 1:3]
3-element Array{Int64,1}:
 1
 2
 3
```
"""
function incrementer(;start=1, step=1)
    i = start - step
    () -> i += step
end

"""
    multiplier(;start=1.0, ratio=1.0)

# Example
```julia
julia> f = multiplier(ratio=0.9);

julia> [f() for _ in 1:3]
3-element Array{Float64,1}:
 1.0
 0.9
 0.81
```
"""
function multiplier(;start=1.0, ratio=1.0)
    x = start / ratio
    () -> x *= ratio
end

"""
    inverse_decay()

# Example
```julia
julia> f = inverse_decay()
#25 (generic function with 1 method)

julia> [f() for _ in 1:5]
5-element Array{Float64,1}:
 1.0
 0.5
 0.3333333333333333
 0.25
 0.2
```
"""
function inverse_decay()
    x = 0
    (args...) -> begin
        x += 1
        1/x
    end
end

"""
    sample_avg()

# Example
```julia
julia> f = sample_avg();

julia> f(2)
2.0

julia> f(3) # (2+3)/2
2.5

julia> f(5) # (2+3+5)/3
3.3333333333333335
```
"""
function sample_avg()
    t = 0
    avg = 0
    (x) -> begin
        t += 1
        avg += (x-avg)/t
    end
end

"""
    cached_inverse_decay()

# Example
```julia
julia> f = cached_inverse_decay();

julia> f(:a) # cached!
1.0

julia> f(:a) # call again
0.5

julia> f(:a) # call again
0.3333333333333333

julia> f(:b) # a new cache
1.0
```
"""
function cached_inverse_decay()
    cache = Dict()
    function (x)
        if !haskey(cache, x)
            cache[x] = inverse_decay()
        end
        cache[x]()
    end
end

"""
    cached_sample_avg()

# Example
```julia
julia> f = cached_sample_avg();

julia> f(:a, 3)  # cache :a
3.0

julia> f(:a, 5)  # calculate avg value of  :a (3 + 5) / 2
4.0

julia> f(:a, 8)  # calculate avg value of  :a (3 + 5 + 8) /3
5.333333333333333

julia> f(:b, 0)  # cache another value
0.0
```
"""
function cached_sample_avg()
    cache = Dict()
    function (k, x)
        if !haskey(cache, k)
            cache[k] = sample_avg()
        end
        cache[k](x)
    end
end

"""
    exp_decay(init=1.0, λ=0.1, decay_step=1000, clip=1e-4)

See [Exponential Decay](https://en.wikipedia.org/wiki/Exponential_decay)
"""
function exp_decay(;init=1.0, λ=0.1, decay_step=1000, clip=1e-4)
    i = -1
    function f()
        i += 1
        max(init * exp(- λ * i / decay_step), clip)
    end
end


module tools

export unfolding, folding, mode_n_mult, tten, hosvd
export get_xi, get_rhs_sin, get_rhs_norm, fullten
export rhs_sin, rhs_norm

export aca, cur, cur_fun
export unfolding_fun, aca_fun, approx_aca

using LinearAlgebra

#### New interface for right-hand side
# This can be used more like functions.
# For any other function on the right hand side,
# just make a non-abstract subtype of rhs and
# implement the desired function as a callable.

# The functions (r::T)(i,j,k) where T <: rhs is some subtype of rhs
# make the objects callable, i.e. one can write
# > r = rhs_sin(100)
# > r(1,35,83)
# this makes the code very convenient.  The Base.getindex
# implementation for rhs is for a more natural indexign via
# > r[1,35,83]
# which returns the same value as the call above.

abstract type rhs end

struct rhs_sin <: rhs
    n::Int
    den::Float64
    rhs_sin(n::Int) = new(n, 1/(n+1))
end

function (r::rhs_sin)(i::Int,j::Int,k::Int)
    # multiplication with reciprocal denominator is much faster than
    # division
    sin((i+j+k)*r.den)
end

struct rhs_norm <: rhs
    n::Int
    den::Float64
    rhs_norm(n::Int) = new(n, 1/(n+1))
end

function (r::rhs_norm)(i::Int,j::Int,k::Int)
    sqrt( (i*r.den)^2 + (j*r.den)^2 + (k*r.den)^2)
end

function Base.getindex(r::rhs, i::Int, j::Int, k::Int)
    # cheap boundscheck; note this does not use the Julia convetnions
    # of bounds checking with several nested function calls
    # check may be removed for better performance
    correct = 1<=i<=r.n && 1<=j<=r.n && 1<=k<=r.n
    if ~correct
        throw(BoundsError(r,(i,j,k)))
    end
    return r(i,j,k)
end

# alternative to getindex for callables This makes arrayindexing with
# Fortran syntax possible, i.e. a(i,j,k)
#### use with care!!
#### uncomment only if you know what you do
# function (a::Array)(i::Int...)
#     a[i...]
# end

"""
    unfolding_fun(a, n, dims) -> a_u(i,j)

Returns a function for evaluating the unfolding (matricisation,
flattening) of a tensor that is given only via a function.

"""
function unfolding_fun_expr(a, n, dims)
    # this is just for getting the first index after permutation, it
    # is assumed that the rank of the tensors is not huge
    # This is a bit easy here, because I assume tensors of fixed order
    # 3 only. Then products just vanish.
    temp = dims[1:end .!= n]
    ind = Tuple(temp)

    # here one-based indexing is a PITA. Maybe this can be done
    # easier, but at least this works. It has been tested with the
    # array based version.
    if n == 1
        ret = :( (i,j) ->
                 # $a(i, ind2sub($ind,j)...)
                 $a(i, CartesianIndices($ind)[j][1], CartesianIndices($ind)[j][2])
                 # if rem(j,$ind)==0
                 # a(i,$ind,div(j,$ind))
                 # else
                 # a(i,rem(j,$ind),1+div(j,$ind))
                 # end
                 )
    elseif n == 2
        ret = :( (i,j) ->
                 $a(CartesianIndices($ind)[j][1], i, CartesianIndices($ind)[j][2])
                 # if rem(j,$ind)==0
                 # a($ind,i,div(j,$ind))
                 # else
                 # a(rem(j,$ind),i,1+div(j,$ind))
                 # end
                 )
    elseif n==3
        ret = :( (i,j) ->
                 $a(CartesianIndices($ind)[j][1], CartesianIndices($ind)[j][2], i)
                 # if rem(j,$ind)==0
                 # a($ind,div(j,$ind),i)
                 # else
                 # a(rem(j,$ind),1+div(j,$ind),i)
                 # end
                 )
    end
    # show(ret)                   # for seeing what substitutions are made
    return ret
end

function unfolding_fun(a, n, dims)
    return eval(unfolding_fun_expr(a, n, dims))
end

"""
    get_xi(n::Int)
returns a vector xi of length n with entries
```julia
xi[i] = 1/(n+1)
```
"""
function get_xi(n::Int)
    collect(1:n)/(n+1)
end

@inline function ξ(i, n)
    i/(n+1)
end

function get_rhs_sin(xi::Array{Float64,1})
    n = length(xi)
    rhs = Array{Float64,3}(undef, (n,n,n))
    @inbounds for i in 1:n
        for j in 1:n
            @simd for k in 1:n
                rhs[k,j,i] = sin(xi[k]+xi[j]+xi[i])
            end
        end
    end
    return rhs
end

function get_rhs_sin(n::Int)
    rhs = Array{Float64,3}(undef, (n,n,n))
    @inbounds for i in 1:n
        for j in 1:n
            @simd for k in 1:n
                rhs[k,j,i] = sin(ξ(k,n)+ξ(j,n)+ξ(i,n))
            end
        end
    end
    return rhs
end

function get_rhs_sin_opt(n::Int)
    # this function retruns different values
    # this is due to machine arithmetic and
    # the precomputation of the denominator
    rhs = Array{Float64,3}(undef, (n,n,n))
    den::Float64 = 1.0/(n+1)
    @inbounds for i in 1:n
        for j in 1:n
            @simd for k in 1:n
                rhs[k,j,i] = sin((k+j+i)*den)
            end
        end
    end
    return rhs
end

function get_rhs_norm(n::Int)
    rhs = Array{Float64,3}(undef, (n,n,n))
    @inbounds for i in 1:n
        for j in 1:n
            @simd for k in 1:n
                rhs[k, j, i] = sqrt(ξ(k,n)*ξ(k,n)+ξ(j,n)*ξ(j,n)+ξ(i,n)*ξ(i,n))
            end
        end
    end
    return rhs
end

function get_rhs_norm(xi::Array{Float64,1})
    n = length(xi)
    rhs = Array{Float64,3}(undef, (n,n,n))
    @inbounds for i in 1:n
        for j in 1:n
            @simd for k in 1:n
                rhs[k,j,i] = sqrt(xi[k]*xi[k] + xi[j]*xi[j] + xi[i]*xi[i])
            end
        end
    end
    return rhs
end

function approx_sin(n)
    rhs = get_rhs_sin(n)
    hosvd(rhs, 1e-4)
end

function approx_norm(n)
    rhs = get_rhs_norm(n)
    hosvd(rhs, 1e-4)
end

"""
n-mode unfolding or matricisation of a tensor a
"""
function unfolding(a, n)
    dims = size(a)
    indices = setdiff(1:ndims(a),n)
    reshape(permutedims(a,[n;indices]),dims[n],prod(dims[indices]))
end

"""
fold matrix into tensor by mode n
reverse of unfolding
"""
function folding(m, n, dims)
    indices = setdiff(1:length(dims), n)
    a = reshape(m, (dims[n],dims[indices]...))
    permutedims(a,invperm([n;indices]))
end

"""
mode-n multiplication of tensor a with matrix m

# inputs
- a :  Array
- n :  mode for multiplication
- m :  matrix that gets multiplied with a in mode n
"""
function mode_n_mult(a, n, m)
    dims = size(a)
    # new_dims = [dims...]
    new_dims = collect(dims)
    new_dims[n] = size(m,1)
    b = m*unfolding(a, n)
    folding(b,n,new_dims)
end

"""
Type for Tucker tensor
core:   Array, this is the core tensor of size (r₁, r₂, r₃, ...)
frames: Array of Arrays, these are the mode frames
"""
struct tten{T<:AbstractFloat}
    core::Array{T}
    frames::Array{Array{T,2},1}

end

"""
    hosvd(a, ϵ)

Computes the Higher Order Singular Value Decomposition

The computation is done via sequentially truncated HOSVD, this is much faster,
as the tensor gets compressed early.
Returns a `tten` object. The tten is a tucker tensor consisting of a
core tensor and an array of mode frames.

input:
- a: full tensor
- ϵ: relative error bound

output:
- `tten` object representing the Tucker approximation
"""
function hosvd(a, ϵ::T=1e-4, sv::Bool=false) where {T<:AbstractFloat}
    d = ndims(a)
    # dims = [size(a)...]
    dims = collect(size(a))
    ϵ = ϵ/sqrt(d) * norm(a[:],2) # divide relative error equal to all dimensions
    ϵ² = ϵ*ϵ                     # use the sqare for cheaper comparison

    frames = Array{Array{Float64,2},1}(undef, d)
    if sv
        svs = Array{Array{Float64,1},1}(undef, d)
    end
    for i in 1:d
        U,S,V = svd(unfolding(a,i))
        rᵢ = 1
        temp = 0.0
        for j in length(S):-1:1
            temp += S[j]*S[j]
            if temp > ϵ²
                # take the one satisfying the error bound, but stay
                # within the bounds => check whether on bound or not
                rᵢ = j==length(S) ? j : (j)
                break
            end
        end
        frames[i] = U[:, 1:rᵢ]
        #a = mode_n_mult(a,i,frames[i]')
        # use a more compact representation of the remaining matrix
        # the implementation of mode_n_mult copys the data around more than once
        # this works as the remainder S*V is just the same as
        # U'*A = U'*U*A*V', as U is orthogonal
        dims[i] = rᵢ
        a = folding(diagm(0=>S[1:rᵢ])*V[:,1:rᵢ]',i,dims);
        if sv
            svs[i] = S
        end
        # garbage collection makes life easier and code faster
        GC.gc()
    end
    if sv
        return tten(a,frames), svs
    else
        return tten(a,frames)
    end
end

# the hosvd2 functions are worse because they reduce the core at the
# end when done during the first loop, simultaneous to svd's, it
# reduces the size and thus cost much
function hosvd2(a, ranks::Vector{T}) where {T<:Int}
    d = ndims(a)
    @assert(d==length(ranks), "dimension mismatch")

    frames = Array{Array{Float64,2},1}(undef, d)
    for i in 1:d
        U = svd(unfolding(a,i))
        frames[i] = U[:, 1:ranks[i]]
    end

    for i in 1:d
        a = mode_n_mult(a,i,frames[i]')
    end
    return tten(a, frames)

end

function hosvd2(a, ranks::Tuple{Vararg{T}}) where {T<:Int}
    # rank_arr = [ranks...]
    rank_arr = collect(ranks)
    hosvd(a, rank_arr)
end

function hosvd2(a, tol::T) where {T<:AbstractFloat}
    d = ndims(a)
    tol = tol/sqrt(d) * norm(a[:],2)

    frames = Array{Array{Float64,2},1}(undef, d)
    for i in 1:d
        U,S = svd(unfolding(a,i))
        r_i = length(S)
        for j in 1:length(S)
            # rework this
            # this is not exactly what we want
            # better is ||S[j:end]|| < tol or even
            # sum(S[j:end].^2) < tol^2
            if S[j] < tol
                r_i = j
                break
            end
        end
        frames[i] = U[:, 1:r_i]
        gc()
    end

    for i in 1:d
        a = mode_n_mult(a,i,frames[i]')
        gc()
    end
    return tten(a,frames)
end

function fullten(a::tten{T}) where {T<:AbstractFloat}
    res = a.core
    for i in 1:length(a.frames)
        res = mode_n_mult(res, i, a.frames[i])
        GC.gc()
    end
    return res
end

# """
#     approx_aca(a, dims, ϵ)

# Compute an approximation of a in Tucker format. dim is an array
# holding the size of a.
# """
# function approx_aca(a, dims, ϵ::Float64=1e-4)
#     d = length(dims)
#     frames =  Array{Array{Float64,2},1}(d)

#     a_u = unfolding_fun(a, 1, dims)
#     dims_u = [dims[1], prod(dims[2:end]) ]
#     I,J = aca_fun(a_u, dims_u, ϵ)
#     C,U,R = cur_fun(a_u, dims_u, I, J)
#     frames[1] = C
#     dims[1] = length(I)
#     core = folding(U*R, 1, dims)

#     for i = 2:3
#         core_u = unfolding(core, i)
#         I,J = aca(core_u, ϵ)
#         C, U, R = cur(core_u, I, J)
#         frames[i] = C
#         dims[i] = length(I)
#         core = folding(U*R, i, dims)
#         gc()
#     end

#     return tten(core, frames)
# end

"""
    aca(a, ϵ)

Computes adaptive cross approximation of matrix a with relative accuracy ϵ.

This implementation uses partial pivoting, this implies the array has
to be evaluated along a row or column. For the computation not the
whole array is needed.

"""
function aca(a::Array{Float64,2},ϵ::Float64=1e-4)
    Rk = a
    maxrk = min(size(a)...)
    I = Int[]
    J = Int[]
    I_good = Int[]
    J_good = Int[]
    k = 1
    i = 1
    # i = 70
    normk2 = 0.0
    us = []
    vs = []
    ϵ = ϵ*ϵ
    # use ↓ this crappy construct to simulate a do-while-loop
    while true
        # maxval is not really used here but comes for free in mymax
        # function
        # maxval,j = mymax(abs.(Rk[i,:]),J)
        # use better pivoting compared to pseudo-alg from lecture
        i,j = find_pivot(Rk, i, I, J)
        # reorder execution compared to algorith from lecture
        # now this is not really important anymore, but it does not
        # matter where this is done. For an early (not working)
        # version this was important
        append!(I, i)
        append!(J, j)

        δ = Rk[i,j]
        # improve the error here. ϵ is maybe not optimal, as it is
        # just the sqared error bound for the overall
        # approximation. This just works for the tested cases.
        if abs(δ) < ϵ
            if length(I) == maxrk - 1
                # if this happens, no proper LRA has been calculated
                @warn("Return with maximal rank => no Low Rank Approximation")
                return I_good,J_good
            end

            # try again, if still bad exit note: this is a trade off
            # between accuracy and speed.
            # It may happen, that there are still good candidates, but
            # those cannot be reached with the search used for the
            # pivot elements. E.g. for block diagonal matrices this
            # easily happens.
            i,j = find_pivot(Rk, i, I, J)
            δ = Rk[i,j]
            if abs(δ) < ϵ
                @warn("Could not find good pivot anymore.")
                return I_good,J_good
            end
        else
            uk = Rk[:,j]
            vk = Rk[i,:]'/δ
            Rk = Rk - uk*vk
            k = k + 1
            # use only the indices that contribute
            append!(I_good, i)
            append!(J_good, j)
            push!(us, uk)
            push!(vs, vk)

            # stopping criterion
            # calculate the parts seperately
            # first interactions of uk, vk with old u and v then the
            # sqared norm, at the end add everything up
            temp = 0
            # k-2 due to early incrementation of k
            for l in 1:k-2
                temp += uk'*us[l] * vs[l]*vk'
            end
            #                         ↓ sqared norms
            normk2 = normk2 + 2temp + (uk'*uk)*(vk*vk')

            # check for stopping criterion
            # also if it is not fulfilled check for high rank, if rank
            # gets too high just return silently and don't complain
            ### maybe it is a good idea to add some special warning or
            ### throw an error if rank exceeds the bounds, after all
            ### this is for _low rank_ approximation
            if uk'*uk * vk*vk' <= ϵ*normk2 || k==maxrk
                return I_good,J_good
            end
        end
        # if abs(δ)<2eps() does not make a
        # difference, if uk and vk are not updated, the error estimate
        # remains the same. So this can be checked in the else
        # statement

        # needed when directly implementing pseudo-algo from lecture
        # that is not the best choice
        # maxval, i = mymax(abs.(uk), I)
    end
end

function find_pivot(a, i, I, J, n::Int=3)
    j = 1
    for l = 1:n
        maxval, j = mymax(abs.(a[i,:]), J)
        maxval, i = mymax(abs.(a[:,j]), I)
    end
    return i,j
end

"""
    mymax(v, I) -> max,i

Computes the maximal value and the index of v, where i ∉ I. I is some
indexset, every element of I is a valid index for v. (The
implementation assumes that at least the minimal index of I is valid.)
"""
function mymax(v, I)
    indices = setdiff(1:length(v),I)
    ind = min(indices...)
    maxval = v[ind]
    for i ∈ indices
        if v[i] > maxval
            maxval = v[i]
            ind = i
        end
    end
    return maxval, ind
end

"""
    update_rk(a, I, J) -> r

takes a function a(i,j), arrays if indices I, J and returns function
calculating the residual a - u*s*v'. Where u are columns J of a, v are
rows I of a, and s is inverse of a(I,J).
"""
function update_rk(a, I, J)
    # check this again! Not sure this works
    #
    # Idea: compute s = inv(a(I,J))
    # then compute a - u*s*v'
    # for the last computation use only the rows and columns that are
    # required
    # this function returns a closure, if not used to it looks strange
    # at first

    if length(I) != length(J)
        error("I and J have to be of equal size")
        return
    end

    s = zeros(length(I),length(J))
    for (j,jj) ∈ enumerate(J)
        for (i,ii) ∈ enumerate(I)
            s[i,j] = Base.invokelatest(a, ii,jj)
        end
    end
    s = inv(s)

    function u(ii)
        # check sizes!
        ret = zeros(length(J))
        for (j,jj) ∈ enumerate(J)
            ret[j] = Base.invokelatest(a, ii,jj)
        end
        return ret
    end

    function v(jj)
        ret = zeros(length(I))
        for (i,ii) ∈ enumerate(I)
            ret[i] = Base.invokelatest(a, ii,jj)
        end
        return ret
    end

    # define the closure
    function rk(ii,jj)
        return Base.invokelatest(a, ii,jj) - u(ii)'*s*v(jj)
    end

    return rk
end

"""
   find_pivot_fun(a, dims, i, I, J)

retunr indices i,j that maximise a(i,j) using a heuristic (search
along axes of a). dims array contains the size of a, i.e.
a ∈ ℝ^(dims[1],dims[2]).

"""
function find_pivot_fun(a, dims, i, I, J)
    j = 1
    for n = 1:3
        a1 = jj -> abs(a(i,jj))
        # maxval, j = mymax_fun( jj-> abs(a(i,jj)), dims[2], J)
        maxval, j = mymax_fun(a1, dims[2], J)
        a2 = ii -> abs(a(ii,j))
        # maxval, i = mymax_fun( ii-> abs(a(ii,j)), dims[1], I)
        maxval, i = mymax_fun(a2, dims[1], I)
    end
    return i,j
end

"""
    mymax_fun(v, n, I)

return maximum value of function v and index ∈ {1...n}∖I
"""
function mymax_fun(v, n, I)
    indices = setdiff(1:n, I)
    ind = min(indices...)
    # maxval = v(ind)
    maxval = Base.invokelatest(v, ind)
    for i ∈ indices
        # if v(i) > maxval
        if Base.invokelatest(v, i) > maxval
            # maxval = v(i)
            maxval = Base.invokelatest(v, i)
            ind = i
        end
    end
    return maxval, ind
end

function aca_fun(a, dims, ϵ::Float64=1e-4)
    #R = Array{Any,1}([a])
    # R = []
    # push!(R, a)
    # Rk = a
    maxrk = min(dims...)
    I = Int[]
    J = Int[]
    I_good = Int[]
    J_good = Int[]
    k = 1
    i = 1
    normk2 = 0.0
    us = []
    vs = []
    Rk = []
    push!(Rk, a)
    ϵ = ϵ*ϵ

    while true
        i,j = find_pivot_fun(Rk[k], dims, i, I, J)
        append!(I, i)
        append!(J, j)

        # δ = R[k](i,j)
        δ = Base.invokelatest(Rk[k], i,j)
        if abs(δ) < ϵ
            # if abs(δ) == 0
            #     return sort!(I),sort!(J)
            # end
            if length(I) == maxrk-1
                @warn("Return with maximal rank")
                return I_good, J_good
            end

            i,j = find_pivot_fun(Rk[k], dims, i, I, J)
            δ = Rk[k](i,j)
            if abs(δ) < ϵ
                @warn("Could not find good pivot aymore.")
                return I_good, J_good
            end
        else
            # uk = (ii)->R[k](ii,j)
            # vk = (jj)->R[k](i,jj)/δ
            # uk = (ii)->Rk[k](ii,j)
            # vk = (jj)->Rk[k](i,jj)/δ
            uk = (ii)->Base.invokelatest(Rk[k], ii,j)
            vk = (jj)->Base.invokelatest(Rk[k],i,jj)/δ
            push!(us, uk)
            push!(vs, vk)

            # FIX: calculate the whole function in every step
            # this costs a loop of length k in each step
            # push!(R, (ii,jj)->R[k](ii,jj)-uk(ii)*vk(jj) )
            # Rk = update_rk(Rk, uk, vk)
            # Rk = update_rk(a, us, vs)

            # update later, otherwise uk and vk get zero!!!!!
            # Rk = update_rk(a, I, J)
            # k = k+1

            temp = 0.0
            for l in 1:k-2
                # temp += prod(map( ii->uk(ii)*us[l](ii), 1:dims[1])) *
                #     prod(map( jj->ns[l](jj)*vk(jj), 1:dims[2]))
                # use the handy |> piping syntax
                temp += sum(1:dims[1] .|> ii->uk(ii)*us[l](ii)) *
                    sum(1:dims[2] .|> jj->vs[l](jj)*vk(jj))
            end

            # calculate product of sqared norms
            # sqnormprod = prod(map( ii->uk(ii)*uk(ii), 1:dims[1])) * prod(map( jj->vk(jj)*vk(jj), 1:dims[2]))
            sqnormprod = sum(1:dims[1] .|> ii->uk(ii)*uk(ii)) * sum(1:dims[2] .|> jj->vk(jj)*vk(jj))
            normk2 = normk2 + 2temp + sqnormprod

            if sqnormprod <= ϵ*normk2 || k==maxrk
                return I_good, J_good
            end

            # use only the indices that contribute to the solution
            append!(I_good, i)
            append!(J_good, j)

            # update here, otherwise uk and vk would be zero
            # k already increased
            k = k+1
            push!(Rk, update_rk(a, I_good, J_good))
        end

    end
end

"""
    cur_fun(a, dims, I, J) -> c, u, r

returns CUR decomposition of `a` given as function a(i,j). dims array
defines sizes of a, i.e. a ∈ ℝ^(dims[1],dims[2]). I, J are the row and
columns indices as returned from aca_fun.

When using aca with functions for huge sizes be careful. This function creates **full matrices**!

"""
function cur_fun(a, dims, I, J)
    m = dims[1]
    n = dims[2]
    c = zeros(m, length(J))
    u = zeros(length(I),length(J))
    r = zeros(length(I), n)

    # set c to coloumns of a
    for (j,jj) ∈ enumerate(J)
        for i in 1:m
            c[i,j] = Base.invokelatest(a, i,jj)
        end
    end

    # set r to rows of a
    for j in 1:n
        for (i,ii) ∈ enumerate(I)
            r[i,j] = Base.invokelatest(a, ii,j)
        end
    end

    # first get a[I,J], then invert it
    for (j,jj) ∈ enumerate(J)
        for (i,ii) ∈ enumerate(I)
            u[i,j] = Base.invokelatest(a, ii,jj)
        end
    end
    u = inv(u)

    return c, u, r
end

"""
   cur(a, i, j) -> C,U,R

Compute the CUR approximation given the original matrix a and the
indiex sets i and j.
"""
function cur(a, i, j)
    return a[:,j], inv(a[i,j]), a[i,:]
end

"""
    approx_aca(a, dims, ϵ)

Compute an approximation of a in Tucker format. dim is an array
holding the size of a.
"""
function approx_aca(a, dims, ϵ::Float64=1e-4)
    d = length(dims)
    frames =  Array{Array{Float64,2},1}(undef, d)

    a_u = unfolding_fun(a, 1, dims)
    dims_u = [dims[1], prod(dims[2:end]) ]
    I,J = aca_fun(a_u, dims_u, ϵ)
    C,U,R = cur_fun(a_u, dims_u, I, J)
    frames[1] = C
    dims[1] = length(I)
    core = folding(U*R, 1, dims)

    for i = 2:3
        core_u = unfolding(core, i)
        I,J = aca(core_u, ϵ)
        C, U, R = cur(core_u, I, J)
        frames[i] = C
        dims[i] = length(I)
        core = folding(U*R, i, dims)
        gc()
    end

    return tten(core, frames)
end

end


using DiscreteElasticRods
DER = DiscreteElasticRods
using Zygote
using ForwardDiff
using FiniteDiff
using ReverseDiff
using BenchmarkTools
using IterativeSolvers
using SparseDiffTools
using DelimitedFiles

struct AHesVec{F,uType}
    f::F
    u::uType
end

Base.size(L::AHesVec) = (length(L.u),length(L.u))
Base.size(L::AHesVec,i::Int) = length(L.u)
Base.:*(L::AHesVec,x::AbstractVector) =  SparseDiffTools.autoback_hesvec(L.f, L.u, x)

import LinearAlgebra
function LinearAlgebra.mul!(du::AbstractVector,L::AHesVec,v::AbstractVector)
    du .= L*v
end

# Rod sizes
n_array = [10, 20, 30, 40, 50]
grad_array = zeros(3, length(n_array))
hess_array = zeros(5, length(n_array))
minres_array = zeros(2, length(n_array))

#for k = 1:length(n_array)
k = 1
    # Setup an energy functional
n = n_array[k]
ref_rod = DER.random_rod(n)
ref_strains = DER.full_kinematics(ref_rod)
rod_props = DER.elliptical_stiffness(10, 0.3, 0.01, 0.04)

function total_energy(q)
    Δx = reshape(q[1:3*n+6], 3, n+2)
    Δθ = reshape(q[3*n+7:end], 1, n+1)
    rod_coords = DER.rod_delta(Δx, Δθ)

    new_rod = DER.rod_update(ref_rod, rod_coords)
    strains = DER.full_kinematics(new_rod)
    DER.elastic_energy(ref_strains, strains, rod_props)
end

function total_energy_array(q)
    Δx = reshape(q[1:3*n+6], 3, n+2)
    Δθ = reshape(q[3*n+7:end], 1, n+1)
    x, d = DER.rod_update(ref_rod.x, ref_rod.d, Δx, Δθ)
    l, κ, τ = DER.full_kinematics(x, d)
    strains = DER.rod_strains(l, κ, τ)
    DER.elastic_energy(ref_strains, strains, rod_props)
end

nq = 4*n+7
q = rand(nq)
# tape = ReverseDiff.GradientTape(total_energy_array, q)
# compiled_tape = ReverseDiff.compile(tape)
#
# grad_reverse!(out, q) = ReverseDiff.gradient!(out, compiled_tape, q)
# grad_forward(q) = ForwardDiff.gradient(total_energy, q)
# grad_zygote(q) = first(Zygote.gradient(total_energy, q))
#
# # Timing the Gradients
# out = similar(q)
# grad_reverse!(out,q)
# grad_forward(q)
# grad_zygote(q)
# rbench = @benchmark grad_reverse!($out, $q)
# fbench = @benchmark grad_forward($q)
# zbench = @benchmark grad_zygote($q)
#
# grad_array[1,k] = minimum(fbench).time/1e9
# grad_array[2,k] = minimum(rbench).time/1e9
# grad_array[3,k] = minimum(zbench).time/1e9

# q = zeros(nq)
# v = rand(nq)
# nh = @benchmark num_hesvec(total_energy, $q, $v)
# nah = @benchmark numauto_hesvec(total_energy, $q, $v)
# anh = @benchmark autonum_hesvec(total_energy, $q, $v)
# nbh = @benchmark numback_hesvec(total_energy, $q, $v)
# abh = @benchmark autoback_hesvec(total_energy, $q, $v)
#
# hess_array[1,k] = minimum(nh).time/1e9
# hess_array[2,k] = minimum(nah).time/1e9
# hess_array[3,k] = minimum(anh).time/1e9
# hess_array[4,k] = minimum(nbh).time/1e9
# hess_array[5,k] = minimum(abh).time/1e9

#writedlm("hessian_timing.txt", hess_array, ',')

grd = rand(nq)

Hvec1 = HesVec(total_energy_array, q; autodiff=true)
out = zeros(length(q))
mna = @benchmark minres!($out, $Hvec1, $grd)

AHVec = AHesVec(total_energy, q)
out = zeros(length(q))
mab = @benchmark minres!($out, $AHVec, $grd)

minres_array[1,k] = minimum(mna).time/1e9
minres_array[2,k] = minimum(mab).time/1e9
writedlm("minres_timing.txt", minres_array, ',')

print("Benchmarked n = $n\n")
#end

writedlm("gradient_timing.txt", grad_array, ',')
writedlm("hessian_timing.txt", hess_array, ',')
writedlm("minres_timing.txt", minres_array, ',')

# ------------------------------------------------------------------------------
# Minres Timing
# ------------------------------------------------------------------------------


# ---
# title: 632. Smallest Range Covering Elements from K Lists
# id: problem632
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Hash Table, Two Pointers, String
# link: <https://leetcode.com/problems/smallest-range-covering-elements-from-k-lists/description/>
# hidden: true
# ---
#
# You have `k` lists of sorted integers in **non-decreasing  order**. Find the
# **smallest** range that includes at least one number from each of the `k`
# lists.
#
# We define the range `[a, b]` is smaller than range `[c, d]` if `b - a < d - c`
# **or** `a < c` if `b - a == d - c`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [[4,10,15,24,26],[0,9,12,20],[5,18,22,30]]
#     Output: [20,24]
#     Explanation:
#     List 1: [4, 10, 15, 24,26], 24 is in range [20,24].
#     List 2: [0, 9, 12, 20], 20 is in range [20,24].
#     List 3: [5, 18, 22, 30], 22 is in range [20,24].
#
#
# **Example 2:**
#
#
#
#     Input: nums = [[1,2,3],[1,2,3],[1,2,3]]
#     Output: [1,1]
#
#
# **Example 3:**
#
#
#
#     Input: nums = [[10,10],[11,11]]
#     Output: [10,11]
#
#
# **Example 4:**
#
#
#
#     Input: nums = [[10],[11]]
#     Output: [10,11]
#
#
# **Example 5:**
#
#
#
#     Input: nums = [[1],[2],[3],[4],[5],[6],[7]]
#     Output: [1,7]
#
#
#
#
# **Constraints:**
#
#   * `nums.length == k`
#   * `1 <= k <= 3500`
#   * `1 <= nums[i].length <= 50`
#   * `-105 <= nums[i][j] <= 105`
#   * `nums[i]` is sorted in **non-decreasing** order.
#
#
## @lc code=start
using LeetCode

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