#!/usr/bin/julia

# This code was written following the algorithm presented in
#
# â”‚ FernÃ¡ndez RodrÃ­guez, A., de Santiago Rodrigo, L., LÃ³pez GuillÃ©n, E. et al.
# â”‚ â€œCoding Pronyâ€™s method in MATLAB and applying it to biomedical signal filteringâ€�
# â”‚ BMC Bioinformatics 19, 451 (2018)
# â”‚ https://doi.org/10.1186/s12859-018-2473-y

using PolynomialRoots, LinearAlgebra

"""
    predictionmatrix(samples::AbstractVector, p::Int)

Return the "forward linear prediction matrix"

    â�› y[p]    y[p-1]  â‹¯  y[1]   â��
    â�œ y[p+1]  y[p]    â‹¯  y[2]   â�Ÿ
    â�œ  â‹®       â‹®      â‹±   â‹®     â�Ÿ
    â�� y[N-1]  y[N-2]  â‹¯  y[N-p] â� 

with `y = samples`.
"""
function predictionmatrix(samples::AbstractVector, p::Int)
    N = length(samples)
    mat = zeros(eltype(samples), N - p, p)
    for rown in 1:(N - p), coln in 1:p
        mat[rown, coln] = samples[p + (rown - 1) - (coln - 1)]
    end
    return mat
end

"""
    observationvector(samples::AbstractVector, p::Int)

Construct the observation vector

     â�› y[p+1] â��
    -â�œ y[p+2] â�Ÿ
     â�œ  â‹®     â�Ÿ
     â�� y[N]   â� 

with `y = samples`.
"""
function observationvector(samples::AbstractVector, p::Int)
    N = length(samples)
    v = zeros(eltype(samples), N - p)
    for i in 1:(N - p)
        v[i] = -samples[p + i]
    end
    return v
end

"""
    cproots(samples::AbstractVector, p::Int)

Find the coefficients ``aâ‚–``, ``kâˆˆ{1,â€¦,p}`` of the characteristic polynomial
``Ï•(z) = âˆ�â‚–â‚Œâ‚�áµ– (z-zâ‚–) = âˆ‘â‚–â‚Œâ‚€áµ– aâ‚–záµ–â�»áµ�``.

They are the (approximate) solution to the linear system (`y = samples`)

    â�› y[p]    y[p-1]  â‹¯  y[1]   â�� â�› a[1] â��    â�› y[p+1] â��  
    â�œ y[p+1]  y[p]    â‹¯  y[2]   â�Ÿ â�œ a[2] â�Ÿ â‰ˆ -â�œ y[p+2] â�Ÿ ;
    â�œ  â‹®       â‹®      â‹±   â‹®     â�Ÿ â�œ  â‹®   â�Ÿ    â�œ â‹®      â�Ÿ 
    â�� y[N-1]  y[N-2]  â‹¯  y[N-p] â�  â�� a[p] â�     â�� y[N]   â�  

`a[0]` is always 1.
"""
function cproots(samples::AbstractVector, p::Int)
    # A \ b returns the exact solution of Ax=b if A is square. For rectangular A
    # the result is the minimum-norm least squares solution.
    a = predictionmatrix(samples, p) \ observationvector(samples, p)
    # Look at the definition of Ï•: aâ‚– is the coefficient of the (p-k)-th power.
    # Since `roots` from PolynomialRoots requires that the roots from the
    # lowest coefficients to the highest, the list must be reversed.
    reverse!(a)
    # The last coefficient, aâ‚€=1 (which multiplies záµ–), is then added to the list.
    push!(a, one(eltype(a)))
    return roots(a)
end

function amplitudes(roots::AbstractVector, samples::AbstractVector)
    N = length(samples)
    p = length(roots)
    # Now that the roots zâ‚–, kâˆˆ{1,â€¦,p} of Ï• are known, the constants hâ‚– can be
    # found from the equation
    #
    #   y[n] = âˆ‘â‚–â‚Œâ‚�áµ– hâ‚–zâ‚–â�¿â�»Â¹
    #
    # which can be written in matrix form as
    #
    #   â�› 1     1     â‹¯  1     â�� â�› hâ‚� â��   â�› y[1] â��
    #   â�œ zâ‚�    zâ‚‚    â‹¯  zâ‚š    â�Ÿ â�œ hâ‚‚ â�Ÿ   â�œ y[2] â�Ÿ
    #   â�œ zâ‚�Â²   zâ‚‚Â²   â‹¯  zâ‚šÂ²   â�Ÿ â�œ hâ‚ƒ â�Ÿ = â�œ y[3] â�Ÿ .
    #   â�œ â‹®     â‹®     â‹±  â‹®     â�Ÿ â�œ â‹®  â�Ÿ   â�œ  â‹®   â�Ÿ
    #   â�� zâ‚�á´ºâ�»Â¹ zâ‚‚á´ºâ�»Â¹ â‹¯  zâ‚šá´ºâ�»Â¹ â�  â�� hâ‚š â�    â�� y[N] â�  
    #
    Z = zeros(eltype(samples), N, p)
    for rown in 1:N, coln in 1:p
        Z[rown, coln] = roots[coln]^(rown - 1)
    end
    return Z \ samples
end

"""
    prony(f::Function, xstart::Real, xend::Real, xstep::Real, p::Int)

Apply Prony's algorithm to find the linear combination ``x â†¦ âˆ‘â±¼â‚Œâ‚�áµ– hâ±¼ exp(Î»â±¼x)``
that best approximates `f` on [`xstart`, `xend`], sampled with step size
`xstep`.

Returns a Pair `(h,Î»)` of lists that contain the coefficients.
"""
function prony(f::Function, xstart::Real, xend::Real, xstep::Real, p::Int)
    X = xstart:xstep:xend
    Y = f.(X)
    if 2p > length(X)
        error("2p must be smaller than the number of data points.")
    end
    z = cproots(Y, p)
    h = amplitudes(z, Y)
    # We have successfully approximated Y = {yâ‚™}â‚™â‚Œâ‚�á´º as
    #   
    #   yâ‚™ = âˆ‘â‚–â‚Œâ‚�áµ– hâ‚– zâ‚–â�¿â�»Â¹ (n â‰¥ 1)
    #
    # where n âˆˆ eachindex(X). If we define zâ‚– := exp(Î»â‚–) then
    #   
    #   yâ‚™ = âˆ‘â‚–â‚Œâ‚�áµ– hâ‚– exp(Î»â‚–(n-1)).
    #
    Î» = log.(z)
    # Now n must be traced back to the values in X, which is assumed to be a
    # list of equidistant, ordered, values xâ‚™. If Ï„ is the step size of X, i.e.
    #
    #   X = { nÏ„ : n âˆˆ ğ��� and n < Ì„n }
    #   
    # then n-1 = Ï„â�»Â¹ Ï„(n-1) = xâ‚™/Ï„, therefore
    #
    #   y(xâ‚™) = âˆ‘â‚–â‚Œâ‚�áµ– hâ‚– exp(Î»â‚–xâ‚™/Ï„).
    #
    # The coefficients we want are then hâ‚– and Î»â‚–/Ï„. It's okay if one of them
    # depends on Ï„: after all this parameter determines which values yâ‚™ are
    # picked up by the algorithm. Hopefully the final result doesn't depend
    # on the particular choice of Ï„ (it seems like it doesn't...).
    return h, Î» ./ xstep
end
