using LinearAlgebra

import LinearAlgebra: mul!, rmul!

lmul!(G::LinearAlgebra.Givens, ::Nothing) = nothing
rmul!(::Nothing, G::LinearAlgebra.Givens) = nothing

function svdvals2x2(d1, d2, e)
    d1sq = d1*d1
    d2sq = d2*d2
    esq  = e*e
    b = d1sq + d2sq + esq
    D = sqrt(abs2((d1 + d2)*(d1 - d2)) + esq*(2*(d1sq + d2sq) + esq))
    D2 = b + D
    λ1 = 2*d1sq*d2sq/D2
    λ2 = D2/2
    return minmax(sqrt(λ1), sqrt(λ2))
end

function svdIter!(B::Bidiagonal{T}, n1, n2, shift, U = nothing, Vᴴ = nothing) where T<:Real

    if B.uplo === 'U'

        d = B.dv
        e = B.ev

        G, r = givens(d[n1] - abs2(shift)/d[n1], e[n1], n1, n1 + 1)
        lmul!(G, Vᴴ)

        ditmp       = d[n1]
        ei1         = e[n1]
        di          = ditmp*G.c + ei1*G.s
        ei1         = -ditmp*G.s + ei1*G.c
        di1         = d[n1 + 1]
        bulge       = di1*G.s
        di1        *= G.c

        for i = n1:n2 - 2
            G, r      = givens(di, bulge, i, i + 1)
            rmul!(U, G')
            d[i]      = G.c*di + G.s*bulge
            ei        = G.c*ei1 + G.s*di1
            di1       = -G.s*ei1 + G.c*di1
            ei1       = e[i + 1]
            bulge     = G.s*ei1
            ei1      *= G.c

            G, r      = givens(ei, bulge, i + 1, i + 2)
            lmul!(G, Vᴴ)
            e[i]      = ei*G.c + bulge*G.s
            di        = di1*G.c + ei1*G.s
            ei1       = -di1*G.s + ei1*G.c
            di2       = d[i + 2]
            bulge     = di2*G.s
            di1       = di2*G.c
        end

        G, r      = givens(di, bulge, n2 - 1, n2)
        rmul!(U, G')
        d[n2 - 1] = G.c*di + G.s*bulge
        e[n2 - 1] = G.c*ei1 + G.s*di1
        d[n2]     = -G.s*ei1 + G.c*di1

    else
        throw(ArgumentError("lower bidiagonal version not implemented yet"))
    end

    return B
end

# See LAWN 3
function svdDemmelKahan!(B::Bidiagonal{T}, n1, n2, U = nothing, Vᴴ = nothing) where T<:Real

    if B.uplo === 'U'

        d = B.dv
        e = B.ev

        oldcs = one(T)
        G     = LinearAlgebra.Givens(1, 2, one(T), zero(T))
        Gold  = G

        for i = n1:n2 - 1
            G, r  = givens(d[i] * G.c, e[i], i, i + 1)
            lmul!(G, Vᴴ)
            if i != n1
                e[i - 1] = Gold.s * r
            end

            Gold, d[i] = givens(Gold.c * r, d[i + 1] * G.s, i, i + 1)
            rmul!(U, Gold')
        end

        h         = d[n2] * G.c
        e[n2 - 1] = h * Gold.s
        d[n2]     = h * Gold.c

    else
        throw(ArgumentError("lower bidiagonal version not implemented yet"))
    end

    return B
end

# Recurrence to estimate smallest singular value from LAWN3 Lemma 1
function estimate_σ⁻(dv::AbstractVector, ev::AbstractVector, n1::Integer, n2::Integer)
    λ  = abs(dv[n2])
    B∞ = λ
    for j in (n2 - 1):-1:n1
        λ = abs(dv[j])*(λ/(λ + abs(ev[j])))
        B∞ = min(B∞, λ)
    end

    μ  = abs(dv[n1])
    B1 = μ
    for j in n1:(n2 - 1)
        μ = abs(dv[j + 1])*(μ*(μ + abs(ev[j])))
        B1 = min(B1, μ)
    end

    return min(B∞, B1)
end

# The actual SVD solver routine
# Notice that this routine doesn't adjust the sign and sorts the values
function __svd!(B::Bidiagonal{T}, U = nothing, Vᴴ = nothing; tol = 100eps(T), debug = false) where T<:Real

    n = size(B, 1)
    n1 = 1
    n2 = n
    d = B.dv
    e = B.ev
    count = 0

    # See LAWN3 page 6 and 22
    σ⁻ = estimate_σ⁻(d, e, n1, n2)
    fudge = n
    thresh = tol*σ⁻

    if B.uplo === 'U'
        while true

            # Search for biggest index for non-zero off diagonal value in e
            for n2i = n2:-1:1
                if n2i == 1
                    # We are done!
                    return nothing
                else
                    if abs(e[n2i - 1]) > thresh
                        n2 = n2i
                        break
                    end
                end
            end

            # Search for largest sub-bidiagonal matrix ending at n2
            for _n1 = (n2 - 1):-1:1
                if _n1 == 1
                    n1 = _n1
                    break
                elseif abs(e[_n1 - 1]) < thresh
                    n1 = _n1
                    break
                end
            end

            debug && println("n1=", n1, ", n2=", n2, ", d[n1]=", d[n1], ", d[n2]=", d[n2], ", e[n1]=", e[n1], " e[n2-1]=", e[n2-1], " thresh=", thresh)

            # See LAWN 3 p 18 and
            # We approximate the smallest and largest singular values of the
            # current block to determine if it's safe to use shift or if
            # the zero shift algorithm is required to maintain high relative
            # accuracy
            σ⁻ = estimate_σ⁻(d, e, n1, n2)
            σ⁺ = max(maximum(view(d, n1:n2)), maximum(view(e, n1:(n2 - 1))))

            if fudge*tol*σ⁻ <= eps(σ⁺)
                svdDemmelKahan!(B, n1, n2, U, Vᴴ)
            else
                shift = svdvals2x2(d[n2 - 1], d[n2], e[n2 - 1])[1]
                if shift^2 < eps(B.dv[n1]^2)
                    # Shift is too small to affect the iteration so just use
                    # zero-shift algorithm anyway
                    svdDemmelKahan!(B, n1, n2, U, Vᴴ)
                else
                    svdIter!(B, n1, n2, shift, U, Vᴴ)
                end
            end

            count += 1
            debug && println("count=", count)
        end
    else
        # Just transpose the matrix
        error("SVD for lower triangular bidiagonal matrices isn't implemented yet.")
    end
end

function _svdvals!(B::Bidiagonal{T}; tol = eps(T), debug = false) where T<:Real
    __svd!(B, tol = tol, debug = debug)
    return sort(abs.(diag(B)), rev = true)
end

function _sort_and_adjust!(U, s, Vᴴ)
    n = length(s)

    # Adjust sign of singular values if necessary
    if any(!isposdef, s)
        rmul!(U, Diagonal(sign.(s)))
        map!(abs, s, s)
    end

    # Sort singular values if necessary
    # FixMe! Try to do this without (much) allocation
    p = sortperm(s, by = abs, rev = true)
    if p != 1:n
        U[:]  = U[:,p]
        s[:]  = s[p]
        Vᴴ[:] = Vᴴ[p,:]
    end

    return nothing
end

function _svd!(B::Bidiagonal{T}; tol = eps(T), debug = false) where T<:Real
    n  = size(B, 1)

    U  = Matrix{T}(I, n, n)
    Vᴴ = Matrix{T}(I, n, n)

    __svd!(B, U, Vᴴ, tol = tol, debug = debug)

    s = diag(B)

    # Sort and adjust the singular values. The algorithm used in __svd! might produce negative singular values and it will return the values in no particular ordering.
    _sort_and_adjust!(U, s, Vᴴ)

    return SVD(U, s, Vᴴ)
end

struct BidiagonalFactorization{T,S,U,V}
    bidiagonal::Bidiagonal{T,S}
    reflectors::U
    τl::V
    τr::V
end

function bidiagonalize!(A::AbstractMatrix)
    m, n = size(A)
    τl, τr = eltype(A)[], eltype(A)[]

    if m >= n
        # tall case: upper bidiagonal
        for i in 1:min(m, n)
            x = view(A, i:m, i)
            τi = LinearAlgebra.reflector!(x)
            push!(τl, τi)
            LinearAlgebra.reflectorApply!(x, τi, view(A, i:m, (i + 1):n))
            if i < n
                x = view(A, i, i + 1:n)
                conj!(x)
                τi = LinearAlgebra.reflector!(x)
                push!(τr, τi)
                LinearAlgebra.reflectorApply!(view(A, (i + 1):m, (i + 1):n), x, τi)
            end
        end

        bd = Bidiagonal(real(diag(A)), real(diag(A, 1)), :U)

        return BidiagonalFactorization{eltype(bd),typeof(bd.dv),typeof(A),typeof(τl)}(bd, A, τl, τr)
    else
        # wide case: lower bidiagonal
        for i in 1:min(m, n)
            x = view(A, i, i:n)
            conj!(x)
            τi = LinearAlgebra.reflector!(x)
            push!(τr, τi)
            LinearAlgebra.reflectorApply!(view(A, (i + 1):m, i:n), x, τi)
            if i < m
                x = view(A, i + 1:m, i)
                τi = LinearAlgebra.reflector!(x)
                push!(τl, τi)
                LinearAlgebra.reflectorApply!(x, τi, view(A, (i + 1):m, (i + 1):n))
            end
        end

        bd = Bidiagonal(real(diag(A)), real(diag(A, -1)), :L)

        return BidiagonalFactorization{eltype(bd),typeof(bd.dv),typeof(A),typeof(τl)}(bd, A, τl, τr)
    end
end

# This one is currently very type unstable but it is useful to be able to reuse
# existing elementary reflectors types such as QRPackedQ, LQPackedQ, and HessenbergQ
# such that multiplication methods are available for free.
function Base.getproperty(F::BidiagonalFactorization, s::Symbol)
    BD = getfield(F, :bidiagonal)
    R  = getfield(F, :reflectors)
    τl = getfield(F, :τl)
    τr = getfield(F, :τr)
    if BD.uplo === 'U'
        if s === :leftQ
            return LinearAlgebra.QRPackedQ(R, τl)
        elseif s === :rightQ
            if VERSION < v"1.3.0-DEV.243"
                return LinearAlgebra.HessenbergQ(copy(transpose(R[1:size(R,2),:])), τr)
            else
                factors = copy(transpose(R[1:size(R,2),:]))
                return LinearAlgebra.HessenbergQ{eltype(factors),typeof(factors),typeof(τr),false}('U', factors, τr)
            end
        else
            return getfield(F, s)
        end
    else
        if s === :leftQ
            if VERSION < v"1.3.0-DEV.243"
                return LinearAlgebra.HessenbergQ(R[:,1:size(R,1)], τl)
            else
                factors = R[:,1:size(R,1)]
                return LinearAlgebra.HessenbergQ{eltype(factors),typeof(factors),typeof(τr),false}('U', factors, τl)
            end
        elseif s === :rightQ
            # return transpose(LinearAlgebra.LQPackedQ(R, τr)) # FixMe! check that this shouldn't be adjoint
            LinearAlgebra.QRPackedQ(copy(transpose(R)), τr)
        else
            return getfield(F, s)
        end
    end
end

# For now, we define a generic lmul! and rmul! for HessenbergQ here
function lmul!(Q::LinearAlgebra.HessenbergQ, B::AbstractVecOrMat)

    m, n = size(B, 1), size(B, 2)

    if m != size(Q, 2)
        throw(DimensionMismatch(""))
    end

    for j in 1:n
        for l in length(Q.τ):-1:1
            τl = Q.τ[l]
            ṽ  = view(Q.factors, (l + 2):m, l)
            b  = view(B, (l + 2):m, j)
            vᴴb = B[l + 1, j]
            if length(ṽ) > 0
                vᴴb += ṽ'b
            end
            B[l + 1, j] -= τl*vᴴb
            b .-= ṽ .* τl .* vᴴb
        end
    end

    return B
end

function lmul!(adjQ::Adjoint{<:Any,<:LinearAlgebra.HessenbergQ}, B::AbstractVecOrMat)

    Q = parent(adjQ)
    m, n = size(B, 1), size(B, 2)

    if m != size(adjQ, 2)
        throw(DimensionMismatch(""))
    end

    for j in 1:n
        for l in 1:length(Q.τ)
            τl = Q.τ[l]
            ṽ  = view(Q.factors, (l + 2):m, l)
            b  = view(B, (l + 2):m, j)
            vᴴb = B[l + 1, j]
            if length(ṽ) > 0
                vᴴb += ṽ'b
            end
            B[l + 1, j] -= τl'*vᴴb
            b .-= ṽ .* τl' .* vᴴb
        end
    end

    return B
end

function rmul!(A::AbstractMatrix, Q::LinearAlgebra.HessenbergQ)

    m, n = size(A)

    if n != size(Q, 1)
        throw(DimensionMismatch(""))
    end

    for i in 1:m
        for l in 1:(n - 1)
            τl = Q.τ[l]
            ṽ  = view(Q.factors, (l + 2):n, l)
            aᵀ = transpose(view(A, i, (l + 2):n))
            aᵀv = A[i, l + 1]
            if length(ṽ) > 0
                aᵀv += aᵀ*ṽ
            end
            A[i, l + 1] -= aᵀv*τl
            aᵀ .-= aᵀv .* τl .* ṽ'
        end
    end

    return A
end

function rmul!(A::AbstractMatrix, adjQ::Adjoint{<:Any,<:LinearAlgebra.HessenbergQ})

    m, n = size(A)
    Q = parent(adjQ)

    if n != size(adjQ, 1)
        throw(DimensionMismatch(""))
    end

    for i in 1:m
        for l in (n - 1):-1:1
            τl = Q.τ[l]
            ṽ  = view(Q.factors, (l + 2):n, l)
            aᵀ = transpose(view(A, i, (l + 2):n))
            aᵀv = A[i, l + 1]
            if length(ṽ) > 0
                aᵀv += aᵀ*ṽ
            end
            A[i, l + 1] -= aᵀv*τl'
            aᵀ .-= aᵀv .* τl' .* ṽ'
        end
    end

    return A
end

function rmul!(A::AbstractMatrix, Q::LinearAlgebra.LQPackedQ)

    m, n = size(A)

    if n != size(Q, 1)
        throw(DimensionMismatch(""))
    end

    for i in 1:m
        for l in length(Q.τ):-1:1
            τl = Q.τ[l]
            ṽ  = view(Q.factors, l, (l + 1):n)
            aᵀ = transpose(view(A, i, (l + 1):n))
            aᵀv = A[i, l]
            if length(ṽ) > 0
                aᵀv += aᵀ*ṽ
            end
            A[i, l] -= aᵀv*τl
            aᵀ .-= aᵀv .* τl .* ṽ'
        end
    end

    return A
end

# Overload LinearAlgebra methods

LinearAlgebra.svdvals!(B::Bidiagonal{T}; tol = eps(T), debug = false) where T<:Real = _svdvals!(B, tol = tol, debug = debug)

"""
    svdvals!(A [, tol, debug])

Generic computation of singular values.
```jldoctest
julia> using LinearAlgebra, GenericLinearAlgebra, Quaternions

julia> n = 20;

julia> H = [big(1)/(i + j - 1) for i in 1:n, j in 1:n]; # The Hilbert matrix

julia> Float64(svdvals(H)[end]/svdvals(Float64.(H))[end] - 1) # The relative error of the LAPACK based solution in 64 bit floating point.
-0.9999999999447275

julia> A = qr([Quaternion(randn(4)...) for i in 1:3, j in 1:3]).Q *
           Diagonal([3, 2, 1]) *
           qr([Quaternion(randn(4)...) for i in 1:3, j in 1:3]).Q'; # A quaternion matrix with the singular value 1, 2, and 3.

julia> svdvals(A) ≈ [3, 2, 1]
true
```
"""
function LinearAlgebra.svdvals!(A::StridedMatrix; tol = eps(real(eltype(A))), debug = false)
    B = bidiagonalize!(A).bidiagonal
    # It doesn't matter that we transpose the bidiagonal matrix when we are only interested in the values
    return svdvals!(Bidiagonal(B.dv, B.ev, :U), tol = tol, debug = debug)
end

# FixMe! The full keyword is redundant for Bidiagonal and should be removed from Base
LinearAlgebra.svd!(B::Bidiagonal{T};
    tol = eps(T),
    full = false,
    # To avoid breaking on <Julia 1.3, the `alg` keyword doesn't do anything. Once we drop support for Julia 1.2
    # and below, we can make the keyword argument work correctly
    alg = nothing,
    debug = false) where T<:Real = _svd!(B, tol = tol, debug = debug)

"""
    svd!(A[, tol, full, debug])::SVD

A generic singular value decomposition (SVD). The implementation only uses Julia functions so the SVD can be computed for any element type provided that the necessary arithmetic operations are supported by the element type.

- `tol`: The relative tolerance for determining convergence. The default value is `eltype(T)` where `T` is the element type of the input matrix bidiagonal (i.e. after converting the matrix to bidiagonal form).

- `full`: Sepcifies if all the left and right singular vectors be returned or if only the vectors us to the number of rows and columns of the input matrix `A` should be returned (the default).

- `debug`: A Boolean flag to activate debug information during the executions of the algorithm. The default is false.

# Algorithm
...tomorrow

# Example

```jldoctest
julia> svd(big.([1 2; 3 4]))
SVD{BigFloat, BigFloat, Matrix{BigFloat}}
U factor:
2×2 Matrix{BigFloat}:
 -0.404554   0.914514
 -0.914514  -0.404554
singular values:
2-element Vector{BigFloat}:
 5.464985704219042650451188493284182533042584640492784181017488774646871847029449
 0.3659661906262578204229643842614005434788136943931877734325179702209382149672422
Vt factor:
2×2 Matrix{BigFloat}:
 -0.576048  -0.817416
 -0.817416   0.576048
```
"""
function LinearAlgebra.svd!(A::StridedMatrix{T};
    tol = eps(real(eltype(A))),
    full = false,
    # To avoid breaking on <Julia 1.3, the `alg` keyword doesn't do anything. Once we drop support for Julia 1.2
    # and below, we can make the keyword argument work correctly
    alg = nothing,
    debug = false) where T

    m, n = size(A)

    # Convert the input matrix A to bidiagonal form and return a BidiagonalFactorization object
    BF = bidiagonalize!(A)

    # The location of the super/sub-diagonal of the bidiagonal matrix depends on the aspect ratio of the input. For tall and square matrices, the bidiagonal matrix is upper whereas it is lower for wide input matrices as illustrated below. The 'O' entried indicate orthogonal (unitary) matrices.

    #      A      =   Q_l  *   B   * Q_r^H

    #     |x x x| = |O O O| |x x  | |O O O|
    #     |x x x|   |O O O|*|  x x|*|O O O|
    #     |x x x|   |O O O| |    x| |O O O|
    #     |x x x|   |O O O|
    #     |x x x|   |O O O|

    # |x x x x x| = |O O O| |x    | |O O O O O|
    # |x x x x x|   |O O O|*|x x  |*|O O O O O|
    # |x x x x x|   |O O O| |  x x| |O O O O O|

    _B = BF.bidiagonal
    B  = _B.uplo === 'U' ? _B : Bidiagonal(_B.dv, _B.ev, :U)

    # Compute the SVD of the bidiagonal matrix B
    F = _svd!(B, tol = tol, debug = debug)

    # Form the matrices U and Vᴴ by combining the singular vector matrices of the bidiagonal SVD with the Householder reflectors from the bidiagonal factorization.
    if _B.uplo === 'U'
        U  = Matrix{T}(I, m, full ? m : n)
        U[1:n,1:n] = F.U
        lmul!(BF.leftQ, U)
        Vᴴ = F.Vt*BF.rightQ'
    else
        U  = BF.leftQ*F.V
        Vᴴ = Matrix{T}(I, full ? n : m, n)
        Vᴴ[1:m,1:m] = F.U'
        rmul!(Vᴴ, BF.rightQ')
    end

    s = F.S

    return SVD(U, s, Vᴴ)
end


using Knet: @diff, Param, value, grad, params
using Knet: sigm, tanh, softmax

mutable struct Layer
    wai::Param
    was::Param
    ba::Param
    wri::Param
    wrs::Param
    br::Param
    wsi::Param
    bs::Param
    state
end

Layer(in_size, out_size) =
begin
    sq = sqrt(2/(in_size+out_size))
    wai = Param(2*sq .* randn(Float32, in_size, out_size) .- sq)
    was = Param(2*sq .* randn(Float32, out_size, out_size) .- sq)
    ba  = Param(zeros(Float32, 1, out_size))
    wri = Param(2*sq .* randn(Float32, in_size, out_size) .- sq)
    wrs = Param(2*sq .* randn(Float32, out_size, out_size) .- sq)
    br  = Param(zeros(Float32, 1, out_size))
    wsi = Param(2*sq .* randn(Float32, in_size, out_size) .- sq)
    bs  = Param(zeros(Float32, 1, out_size))
    state = zeros(Float32, 1, out_size)
    layer = Layer(wai, was, ba, wri, wrs, br, wsi, bs, state)
layer
end

(layer::Layer)(in) =
begin
    focus  = sigm.(in * layer.wai + layer.state * layer.was + layer.ba)
    keep   = sigm.(in * layer.wri + layer.state * layer.wrs + layer.br)
    interm = tanh.(in * layer.wsi + layer.state .* focus + layer.bs)
    layer.state = keep .* interm + (1 .- keep) .* layer.state
end

mk_def(in_size, layers, out_size) =
begin
    hm_layers = length(layers)+1

    model_def = "struct Model\n"
    for i in 1:hm_layers
        model_def *= "l$i::Layer\n"
    end
    model_def *= "end"

    eval(Meta.parse(model_def))

    model_call = "(model::Model)(io) =\n"
    model_call *= "begin\n"
    for i in 1:hm_layers
        model_call *= "io = model.l$i(io)\n"
    end
    model_call *= "end"

    eval(Meta.parse(model_call))

end

mk_layers(in_size, layers, out_size) =
begin
    hm_layers = length(layers)+1

    hm_layers = length(layers) +1
    internal = Layer[]
    for i in 1:hm_layers
        if     i == 1         layer = Layer(in_size, layers[i])
        elseif i == hm_layers layer = Layer(layers[end], out_size)
        else                  layer = Layer(layers[i-1], layers[i])
        end
        push!(internal, layer)
    end
internal
end


# Continuous interpolation of discrete frequency-domain data
using Interpolations
import Interpolations: interpolate

function interpolate(s::Spectrum)
    l,h = limits(s)
    fn = extrapolate(interpolate(s[:], BSpline(Cubic(Natural())), OnGrid()), Flat())
    # Interpolate.jl only supports increasing ranges, so will have to
    # use flip `h` and `l` here and invert the axis in the next line.
    scaled = scale(fn, range(l; stop=h, length=length(s)))
    δ -> scaled[l + (h - δ)]
end

function resample(s::Spectrum, shifts::AbstractArray)
    intp = interpolate(s)
    intp.(shifts)
end


using BcdV3
using StaticArrays
using MimuConstants
using LinearAlgebra
using Printf
using Dates

# Noise less simulations
TT = Float64
active_imus = [3,4,5,6, 27,28, 29, 30]
r_all = get_InitialPosMIMU4444BT_corrected()
r_m = r_all[:, active_imus]
r_ref = r_m[:,1]
r_m .-= r_ref
r_all .-= r_ref

sig_a = get_rms_noise_acc()
sig_g = get_rms_noise_gyro()
Na = size(r_m, 2)
Ng = 8

N = 5000 # Dyn
Ns = []
show_every = 100
N_bcd  = 1000
tol_dlogp = 1.0e-9

Nt = Na + Ng
g_mag = 9.81

b = [@SVector randn(3) for k = 1:Nt]
T = [I + 1.0e-3*@SMatrix randn(3,3) for k = 1:Nt]
T[1] = I + 1.0e-3*SMatrix{3,3, Float64}(triu(randn(3,3)))
r = [SVector{3,TT}(r_m[:,k]) for k = 1:Na]

w = [deg2rad(1000.0) * @SVector randn(3) for _ = 1:N]
w_dot = [deg2rad(40e3) * @SVector randn(3) for _ = 1:N]
s = [20.0 * @SVector randn(3) for _ = 1:N]
w_m = hcat(Array.(w)...)

y_d_true = Array{SVector{3,TT}}(undef, Nt,N)
y_d_true_2 = Array{SVector{3,TT}}(undef, Nt,N)
for n = 1:N
    for k = 1:Na
        y_d_true[k,n] = T[k]*(Ω(w[n])^2*r[k] + Ω(w_dot[n])*r[k] + s[n]) + b[k]
        y_d_true_2[k,n] = T[k]*(Ω(2*w[n])^2*r[k] + Ω(2*w_dot[n])*r[k] + 2*s[n]) + b[k]
    end
    for k in range(Na+1, length=Ng)
        y_d_true[k,n] = T[k]*w[n] + b[k]
        y_d_true_2[k,n] = T[k]*w[n]*2 + b[k]
    end
end

Q_inv = [SMatrix{3,3,Float64}(I)/( k <= Na ? sig_a^2 :  sig_g^2) for k = 1:Nt]

b0 = [b_n + 1.0e-2*randn(3) for  b_n in b]
T0 = [T[k] + 1.0e-3*(k == 1 ? triu(randn(3,3)) : randn(3,3))  for k = 1:Nt]
r0_m = copy(r_m)
r0_m[:,2:end] += 1.0e-3*randn(3,Na-1)
r0 = [SVector{3}(r0_m[:,k]) for k = 1:Na]

y = Array{SVector{3,TT}}(undef, Nt,N)
y2 = Array{SVector{3,TT}}(undef, Nt,N)
for n = 1:N
    for k = 1:Na
        e = sig_a * randn(3)
        y[k,n] = y_d_true[k,n] + e
        y2[k,n] = y_d_true_2[k,n] + e
    end
    for k in range(Na+1, length=Ng)
        e = sig_g * randn(3)
        y[k,n] = y_d_true[k,n] + e
        y2[k,n] = y_d_true_2[k,n] + e
    end
end
dw = zeros(3, N)
dw2 = zeros(3, N)

w0 = deepcopy(w)
r_hat, T_hat, b_hat, eta = bcd!(
    copy(r), copy(T), copy(b), y, Q_inv, Ns, g_mag, Val{3*(Na+Ng)}(),Val{Na}(),
    1.0e-12, 100, 1.0e-8, 200, tol_dlogp, 1.0-14, 1; warn = false, w0 = w0
)
w_hat_m = hcat(Array.(eta[:w])...)
dw .+= (w_m  - w_hat_m).^2
mse1 = sum(dw)/length(dw)
rmse1 = sqrt(mse1)

r_hat, T_hat, b_hat, eta = bcd!(
    copy(r), copy(T), copy(b), y2, Q_inv, Ns, g_mag, Val{3*(Na+Ng)}(),Val{Na}(),
    1.0e-12, 100, 1.0e-8, 200, tol_dlogp, 1.0-14, 1; warn = false, w0 = w0.*2
)
w_hat_m = hcat(Array.(eta[:w])...)
dw2 .+= (w_m.*2  - w_hat_m).^2
mse2 = sum(dw2)/length(dw2)
rmse2 = sqrt(mse2)

# rmse is lower for second


export  ProPolyHermite, Cpro, degree, ProPolyH, prohermite_coeffmatrix,
        FamilyProPolyHermite, FamilyScaledProPolyHermite,
        derivative,
        evaluate!, evaluate,
        vander!, vander

"""
ProPolyHermite <: Hermite

An immutable structure for probabilistic Hermite polynomials

## Fields
-  `m` : order of the polynomial
-  `P` : probabilistic Hermite polynomial of order m
- `scaled` : with rescaling to have unitary norm
"""
struct ProPolyHermite <: Hermite
    m::Int64
    P::ImmutablePolynomial{Float64}
    scaled::Bool
end

function Base.show(io::IO, P::ProPolyHermite)
    println(io, string(P.m)*"-th order probabilistic Hermite polynomial "*string(P.P)*", scaled = "*string(P.scaled))
end

# Hen(x)  = (-1)ⁿ*exp(x²/2)dⁿ/dxⁿ exp(-x²/2)
# Hen′(x) = n*Hen-1(x)
# Hen″(x) = n*(n-1)*Hen-1(x)

"""
$(TYPEDSIGNATURES)

Normalization constant of the m-th order probabilistic Hermite polynomial
"""
# Cpro(m::Int64) =sqrt(sqrt(2*π) * gamma(m+1))
Cpro(m::Int64) =sqrt(gamma(m+1)) #sqrt(2π) is already accounted for by sampling

"""
$(TYPEDSIGNATURES)

Normalization constant of the m-th order probabilistic Hermite polynomial
"""
Cpro(P::ProPolyHermite) = Cpro(P.m)

"""
$(TYPEDSIGNATURES)

Return the degree of the probabilistic Hermite polynomial `P`.
"""
degree(P::ProPolyHermite) = P.m

# Adapted https://people.sc.fsu.edu/~jburkardt/m_src/hermite_polynomial/h_polynomial_coefficients.m

#    Input, integer N, the highest order polynomial to compute.
#    Note that polynomials 0 through N will be computed.
#
#    Output, real C(1:N+1,1:N+1), the coefficients of the Hermite
#    polynomials.

# Recurrence relation:
# Hen+1 = x*Hen - n*Hen-1

"""
$(TYPEDSIGNATURES)

Computes the coefficients of the probabilistic Hermite polynomials up to order m.
"""
function prohermite_coeffmatrix(m::Int64)

    if m < 0
        return Float64[]
    end

    coeff = zeros(m+1, m+1)
    coeff[1, 1] = 1.0

    if m == 0
        return coeff
    end

    coeff[2, 2] = 1.0

    if m==1
        return coeff
    end

    for i  = 2:m
        coeff[i+1, 1]      =  -      (i - 1) * coeff[i-1, 1]
        coeff[i+1, 2:i-1]  =                   coeff[i  , 1:i-2] -
                                     (i - 1) * coeff[i-1, 2:i-1]
        coeff[i+1, i]      =                   coeff[i  , i-1]
        coeff[i+1, i+1]    =                   coeff[i  , i]
    end
    return coeff
end

const ProPolyH = prohermite_coeffmatrix(CstMaxDegree)

function ProPolyHermite(m::Int64;scaled::Bool= false)
    @assert m>=0 "The order of the polynomial should be >=0"
    if scaled ==false
        return ProPolyHermite(m, ImmutablePolynomial(view(ProPolyH,m+1,1:m+1)), scaled)
    else
        C = 1/Cpro(m)
        return ProPolyHermite(m, ImmutablePolynomial(C*view(ProPolyH,m+1,1:m+1)), scaled)
    end
end

"""
$(TYPEDSIGNATURES)

Evaluates the probabilistic Hermite polynomial `P`.
"""
(P::ProPolyHermite)(x) = P.P(x)

const FamilyProPolyHermite = ntuple(i->ProPolyHermite(i-1), CstMaxDegree+1)
const FamilyScaledProPolyHermite = ntuple(i->ProPolyHermite(i-1; scaled = true), CstMaxDegree+1)

"""
$(TYPEDSIGNATURES)

Computes the k-th derivative of the probabilistic Hermite polynomial `P`.
"""
function derivative(P::ProPolyHermite, k::Int64)
    m = P.m
    @assert k>=0 "This function doesn't compute anti-derivatives of Hermite polynomials"
    if m>=k
        factor = exp(loggamma(m+1) - loggamma(m+1-k))
        if P.scaled == false
            return ImmutablePolynomial(factor*ProPolyH[m-k+1,1:m-k+1])
        else
            C = 1/Cpro(m)
            return ImmutablePolynomial(factor*C*ProPolyH[m-k+1,1:m-k+1])
        end
    else
        return ImmutablePolynomial((0.0))
    end
end

"""
$(TYPEDSIGNATURES)

Evaluates in-place the probabilistic Hermite polynomial `P`.
"""
function evaluate!(dV, P::ProPolyHermite, x)
    m = P.m
    N = size(x,1)
    @assert size(dV) == (N, m+1) "Wrong dimension of the Vander matrix"

    # H_0(x) = 1, H_1(x) = x
    col = view(dV,:,1)
    @avx @. col = 1.0

    if P.scaled
        rmul!(col, 1/Cpro(0))
    end
    if m == 0
        return dV
    end

    col = view(dV,:,2)
    @avx @. col = x

    if P.scaled
        rmul!(col, 1/Cpro(1))
    end

    if m == 1
        return dV
    end

    if P.scaled
        @inbounds for i=2:m
            colp1 = view(dV,:,i+1)
            col   = view(dV,:,i)
            colm1 = view(dV,:,i-1)
            @avx @. colp1 = Cpro(i-1)* x * col - Cpro(i-2)*(i-1)*colm1
            # dV[:,i+1] = 2.0*Cphy(i-1)*x .* dV[:,i] - 2.0*Cphy(i-2)*i*dV[:,i-1]
            rmul!(colp1, 1/Cpro(i))
        end
    else
        @inbounds for i=2:m
            colp1 = view(dV,:,i+1)
            col   = view(dV,:,i)
            colm1 = view(dV,:,i-1)
            @avx @. colp1 = x * col - (i-1)*colm1
        end
    end
    return dV
end

"""
$(TYPEDSIGNATURES)

Evaluates the probabilistic Hermite polynomial `P`.
"""
evaluate(P::ProPolyHermite, x::Array{Float64,1}) = evaluate!(zeros(size(x,1), P.m+1), P, x)

# He_{n}^(k)(x) =  n!/(n-k)! H_{n-k}(x)

"""
$(TYPEDSIGNATURES)

Computes in-place the Vandermonde matrix of the k-th derivative of the probabilistic Hermite polynomial `P`.
"""
function vander!(dV, P::ProPolyHermite, k::Int64, x)
    m = P.m

    if k==0
        evaluate!(dV, P, x)
    # Derivative
    elseif k==1
        # Use recurrence relation H′_{n+1} = (1 + 1/n)*(x*H′_{n} - n*H′_{n-1})
        N = size(x,1)
        @assert size(dV) == (N, m+1) "Wrong dimension of the Vander matrix"

        col = view(dV,:,1)
        @avx @. col = 0.0
        if m==0
            return dV
        end

        col = view(dV,:,2)
        @avx @. col = 1.0

        if m==1
            return dV
        end

        col = view(dV,:,3)
        @avx @. col = 2.0*x

        if m==2
            return dV
        end

        @inbounds for i=3:m
            colp1 = view(dV,:,i+1)
            col   = view(dV,:,i)
            colm1 = view(dV,:,i-1)

            @avx @. colp1 = (1.0 + 1.0/(i-1.0))*(x * col - (i-1.0) * colm1)
        end

        if P.scaled
            @inbounds for i=1:m
                colp1 = view(dV,:,i+1)
                factor = exp(loggamma(i+1) - loggamma(i+1-k))
                rmul!(colp1, 1/(Cpro(i)*sqrt(factor)))
            end

        end
        return dV
    # Second Derivative
    elseif k==2
        # Use recurrence relation He″_{n+1} = (n+1)/(n-1)*(x*He″_{n} - 2*n*He″_{n-1})
        N = size(x,1)
        @assert size(dV) == (N, m+1) "Wrong dimension of the Vander matrix"

        col = view(dV,:,1)
        @avx @. col = 0.0
        if m==0
            return dV
        end

        col = view(dV,:,2)
        @avx @. col = 0.0

        if m==1
            return dV
        end

        col = view(dV,:,3)
        @avx @. col = 2.0

        if m==2
            return dV
        end

        col = view(dV,:,4)
        @avx @. col = 6.0 * x

        if m==3
            return dV
        end

        @inbounds for i=4:m
            colp1 = view(dV,:,i+1)
            col   = view(dV,:,i)
            colm1 = view(dV,:,i-1)

            @avx @. colp1 = ((i)/(i-2.0))*(x * col - (i-1.0) * colm1)
        end

        if P.scaled
            @inbounds for i=1:m
                colp1 = view(dV,:,i+1)
                factor = exp(loggamma(i+1.0) - loggamma(i+1.0-k))
                @avx rmul!(colp1, 1/(Cpro(i)*sqrt(factor)))
            end

        end
        return dV
    end
end

"""
$(TYPEDSIGNATURES)

Computes the Vandermonde matrix of the k-th derivative of the probabilistic Hermite polynomial `P`.
"""
vander(P::ProPolyHermite, k::Int64, x::Array{Float64,1}) = vander!(zeros(size(x,1), P.m+1), P, k, x)


# function testlinearomega()
#   x = uniform(0.0, 1.0)
#   rand(x; ΩT = LinearΩ)
# end

# function x_(ω)
#   a = rand(ω)
#   b = bernoulli(ω, 0.3, Bool)
#   c = normal(ω, 1.0, 2.0)
#   a + b + c
# end

function testlinear()
  x = normal(0.0, 1.0, (1,2,3))
  w = Omega.LinearΩ()
  x(w)
end

# x = ciid(x_)
# ω = LinearΩ()
# x(ω)


using ImageTransformations: center
using Photometry
using Distributions
using Statistics
using HCIToolbox: get_annulus_segments
using StaticArrays
using Rotations

"""
    detectionmap([method=snr], data, fwhm; fill=0)

Parallel implementation of arbitrary detection mapping applied to each pixel in the input image. Any invalid values will be set to `fill`.

The following methods are provided in the [`Metrics`](@ref) module:
* [`snr`](@ref) - signal-to-noise ratio (S/N) using student-t statistics to account for small sample penalty.
* [`significance`](@ref) - Gaussian signifance using student-t statistics to account for samll sample penalty.

!!! tip
    This code is automatically multi-threaded, so be sure to set `JULIA_NUM_THREADS` before loading your runtime to take advantage of it!
"""
function detectionmap(method, data::AbstractMatrix{T}, fwhm; fill=zero(T)) where T
    out = fill!(similar(data), fill)
    width = minimum(size(data)) / 2 - 1.5 * fwhm

    masked = get_annulus_segments(data, fwhm/2 + 2, width, mode=:apply)
    coords = findall(!iszero, masked)

    Threads.@threads for coord in coords
        val = method(data, coord, fwhm)
        @inbounds out[coord] = isfinite(val) ? val : fill
    end

    return out
end

detectionmap(data, fwhm) = detectionmap(snr, data, fwhm)

"""
    snr(data, position, fwhm)

Calculate the signal to noise ratio (SNR, S/N) for a test point at `position` using apertures of diameter `fwhm` in a residual frame.

Uses the method of Mawet et al. 2014 which includes penalties for small sample statistics. These are encoded by using a student's t-test for calculating the SNR.

!!! note
    SNR is not equivalent to significance, use [`significance`](@ref) instead
"""
function snr(data::AbstractMatrix, position, fwhm)
    x, y = position
    cy, cx = center(data)
    separation = sqrt((x - cx)^2 + (y - cy)^2)
    separation > fwhm / 2 + 1 || return NaN
    r = fwhm / 2

    N = floor(Int, 2 * π  * separation / fwhm)
    dθ = 2π / N
    R = RotMatrix{2}(dθ)

    # initial points
    rx = x - cx
    ry = y - cy

    fluxes = similar(data, N)

    for idx in eachindex(fluxes)
        ap_x = rx + cx
        ap_y = ry + cy
        rx, ry = R * SA[rx, ry]
        ap = CircularAperture(ap_x, ap_y, r)
        fluxes[idx] = photometry(ap, data).aperture_sum
    end

    other_elements = @view fluxes[2:end]
    bkg_σ = std(other_elements) # ddof = 1 by default
    return (fluxes[1] - mean(other_elements)) / (bkg_σ * sqrt(1 + 1/(N - 1)))
end

snr(data::AbstractMatrix, idx::CartesianIndex, fwhm) = snr(data, reverse(idx.I), fwhm)

"""
    significance(data, position, fwhm)

Calculates the Gaussian significance from the signal-to-noise ratio (SNR, S/N) for a test point at `position` using apertures of diameter `fwhm` in a residual frame.

The Gaussian signifiance is calculated from converting the SNR confidence limit from a student t distribution to a Gaussian via

``\\text{sig}(\\text{SNR}) = \\Phi^{-1}\\left[\\int_0^\\text{SNR}{t_\\nu(x)dx}\\right]``

where the degrees of freedom ``\\nu`` is given as ``2\\pi r / \\Gamma - 2`` where r is the radial distance of each pixel from the center of the frame.

# See Also
[`snr`](@ref)
"""
function significance(data::AbstractMatrix, position, fwhm)
    x, y = position
    cy, cx = center(data)
    separation = sqrt((x - cx)^2 + (y - cy)^2)
    _snr = snr(data, position, fwhm)
    # put in one line to allow broadcast fusion
    return snr_to_sig(_snr, separation, fwhm)
end
significance(data::AbstractMatrix, idx::CartesianIndex, fwhm) = significance(data, reverse(idx.I), fwhm)

function snr_to_sig(snr, separation, fwhm)
    dof = 2 * π * separation / fwhm - 2
    dof > 0 || return NaN
    return quantile(Normal(), cdf(TDist(dof), float(snr)))
end
function sig_to_snr(sig, separation, fwhm)
    dof = 2 * π * separation / fwhm - 2
    dof > 0 || return NaN
    return quantile(TDist(dof), cdf(Normal(), float(sig)))
end

"""
    noise(data, position, fwhm)

Calculate the statistical noise for a test point at `position` using apertures of diameter `fwhm` in a residual frame.

Uses the standard deviation of the apertures in the entire annulus. This is distinct from the [`snr`](@ref) noise calculation, which defines a confidence interval using student-t statistics. This means you cannot simply create a noise map and divide it from the signal to create an equivalent S/N map.
"""
function noise(data::AbstractMatrix, position, fwhm)
    x, y = position
    cy, cx = center(data)
    separation = sqrt((x - cx)^2 + (y - cy)^2)
    separation > fwhm / 2 + 1 || return NaN
    r = fwhm / 2

    N = floor(Int, 2 * π  * separation / fwhm)
    dθ = 2π / N
    R = RotMatrix{2}(dθ)

    # initial points
    rx = x - cx
    ry = y - cy

    fluxes = similar(data, N)

    for idx in eachindex(fluxes)
        ap_x = rx + cx
        ap_y = ry + cy
        rx, ry = R * SA[rx, ry]
        ap = CircularAperture(ap_x, ap_y, r)
        fluxes[idx] = photometry(ap, data).aperture_sum
    end

    return std(fluxes)
end

noise(data::AbstractMatrix, idx::CartesianIndex, fwhm) = noise(data, reverse(idx.I), fwhm)


# Modified from Flux.jl to accept the td error instead of two inputs
function Flux.huber_loss(y; δ=eltype(y)(1))
    abs_error = abs.(y)
    temp = abs_error .< δ
    x = eltype(y)(0.5)
    hub_loss = sum(((abs_error.^2) .* temp) .* x .+ δ*(abs_error .- x*δ) .* (1 .- temp)) * 1 // length(y)
end


"""
    AbstractSampling

Abstract type for sparse sampling.

Encodes the "basis transformation" of a propagator from the truncated IR
basis coefficients `G_ir[l]` to time/frequency sampled on sparse points
`G(x[i])` together with its inverse, a least squares fit:

         ________________                   ___________________
        |                |    evaluate     |                   |
        |     Basis      |---------------->|     Value on      |
        |  coefficients  |<----------------|  sampling points  |
        |________________|      fit        |___________________|

"""
abstract type AbstractSampling{T,Tmat,F<:SVD} end

function cond(sampling::AbstractSampling)
    return first(sampling.matrix_svd.S) / last(sampling.matrix_svd.S)
end

function Base.show(io::IO, smpl::S) where {S<:AbstractSampling}
    println(io, S)
    print(io, "Sampling points: ")
    return println(io, smpl.sampling_points)
end

"""
    TauSampling <: AbstractSampling

Sparse sampling in imaginary time.

Allows the transformation between the IR basis and a set of sampling points
in (scaled/unscaled) imaginary time.
"""
struct TauSampling{T,Tmat,F<:SVD} <: AbstractSampling{T,Tmat,F}
    sampling_points::Vector{T}
    matrix::Matrix{Tmat}
    matrix_svd::F
end

"""
    TauSampling(basis[, sampling_points])

Construct a `TauSampling` object. If not given, the `sampling_points` are chosen as
the extrema of the highest-order basis function in imaginary time. This turns out
to be close to optimal with respect to conditioning for this size (within a few percent).
"""
function TauSampling(
    basis::AbstractBasis, sampling_points=default_tau_sampling_points(basis)
)
    matrix = eval_matrix(TauSampling, basis, sampling_points)
    sampling = TauSampling(sampling_points, matrix, svd(matrix))

    if iswellconditioned(basis) && cond(sampling) > 1e8
        @warn "Sampling matrix is poorly conditioned (cond = $(cond(sampling)))."
    end

    return sampling
end

const TauSampling64 = @static if VERSION < v"1.9-"
    TauSampling{Float64,Float64,SVD{Float64,Float64,Matrix{Float64}}}
else
    TauSampling{Float64,Float64,SVD{Float64,Float64,Matrix{Float64},Vector{Float64}}}
end

"""
    MatsubaraSampling <: AbstractSampling

Sparse sampling in Matsubara frequencies.

Allows the transformation between the IR basis and a set of sampling points
in (scaled/unscaled) imaginary frequencies.
"""
struct MatsubaraSampling{T,Tmat,F<:SVD} <: AbstractSampling{T,Tmat,F}
    sampling_points::Vector{T}
    matrix::Matrix{Tmat}
    matrix_svd::F
end

const MatsubaraSampling64 = @static if VERSION < v"1.9-"
    MatsubaraSampling{Int64,ComplexF64,SVD{ComplexF64,Float64,Matrix{ComplexF64}}}
else
    MatsubaraSampling{
        Int64,ComplexF64,SVD{ComplexF64,Float64,Matrix{ComplexF64},Vector{Float64}}
    }
end

"""
    MatsubaraSampling(basis[, sampling_points])

Construct a `MatsubaraSampling` object. If not given, the `sampling_points` are chosen as
the (discrete) extrema of the highest-order basis function in Matsubara. This turns out
to be close to optimal with respect to conditioning for this size (within a few percent).
"""
function MatsubaraSampling(
    basis::AbstractBasis, sampling_points=default_matsubara_sampling_points(basis)
)
    matrix = eval_matrix(MatsubaraSampling, basis, sampling_points)
    sampling = MatsubaraSampling(sampling_points, matrix, svd(matrix))

    if iswellconditioned(basis) && cond(sampling) > 1e8
        @warn "Sampling matrix is poorly conditioned (cond = $(cond(sampling)))."
    end

    return sampling
end

"""
    eval_matrix(T, basis, x)

Return evaluation matrix from coefficients to sampling points. `T <: AbstractSampling`.
"""
eval_matrix(::Type{TauSampling}, basis, x) = permutedims(basis.u(x))
eval_matrix(::Type{MatsubaraSampling}, basis, x) = permutedims(basis.uhat(x))

"""
    evaluate(sampling, al; dim=1)

Evaluate the basis coefficients `al` at the sparse sampling points.
"""
function evaluate(
    smpl::AbstractSampling{S,Tmat}, al::AbstractArray{T,N}; dim=1
) where {S,Tmat,T,N}
    if size(smpl.matrix, 2) ≠ size(al, dim)
        msg = "Number of columns (got $(size(smpl.matrix, 2))) has to match al's size in dim (got $(size(al, dim)))"
        throw(DimensionMismatch(msg))
    end
    bufsize = (size(al)[1:(dim - 1)]..., size(smpl.matrix, 1), size(al)[(dim + 1):end]...)
    buffer = Array{promote_type(Tmat, T),N}(undef, bufsize)
    return evaluate!(buffer, smpl, al; dim)
end

"""
    evaluate!(buffer::AbstractArray{T,N}, sampling, al; dim=1) where {T,N}

Like [`evaluate`](@ref), but write the result to `buffer`.
Please use dim = 1 or N to avoid allocating large temporary arrays internally.
"""
function evaluate!(buffer::AbstractArray, smpl::AbstractSampling, al; dim=1)
    bufsize = (size(al)[1:(dim - 1)]..., size(smpl.matrix, 1), size(al)[(dim + 1):end]...)
    if size(buffer) ≠ bufsize
        msg = "Buffer has the wrong size (got $(size(buffer)), expected $bufsize)"
        throw(DimensionMismatch(msg))
    end
    return matop_along_dim!(buffer, smpl.matrix, al, dim, mul!)
end

"""
    fit(sampling, al::AbstractArray{T,N}; dim=1)

Fit basis coefficients from the sparse sampling points
Please use dim = 1 or N to avoid allocating large temporary arrays internally.
"""
function fit(
    smpl::AbstractSampling{S,Tmat}, al::AbstractArray{T,N}; dim=1
) where {S,Tmat,T,N}
    if size(smpl.matrix, 1) ≠ size(al, dim)
        msg = "Number of rows (got $(size(smpl.matrix, 1))) has to match al's size in dim (got $(size(al, dim)))"
        throw(DimensionMismatch(msg))
    end
    bufsize = (size(al)[1:(dim - 1)]..., size(smpl.matrix, 2), size(al)[(dim + 1):N]...)
    buffer = Array{promote_type(Tmat, T),N}(undef, bufsize)
    return fit!(buffer, smpl, al; dim)
end

"""
    workarrlength(smpl::AbstractSampling, al; dim=1)

Return length of workarr for `fit!`.
"""
function workarrlength(smpl::AbstractSampling, al::AbstractArray; dim=1)
    return length(smpl.matrix_svd.S) * (length(al) ÷ size(al, dim))
end

"""
    fit!(buffer, sampling, al::Array{T,N}; dim=1)

Like [`fit`](@ref), but write the result to `buffer`.
Please use dim = 1 or N to avoid allocating large temporary arrays internally.
The length of `workarry` cannot be smaller than the returned value of `workarrlengthfit`.
"""
function fit!(
    buffer::Array{S,N}, smpl::AbstractSampling, al::Array{T,N};
    dim=1, workarr::Vector{S}=Vector{S}(undef, workarrlength(smpl, al; dim)),
) where {S,T,N}
    bufsize = (size(al)[1:(dim - 1)]..., size(smpl.matrix, 2), size(al)[(dim + 1):end]...)
    if size(buffer) ≠ bufsize
        msg = "Buffer has the wrong size (got $(size(buffer)), expected $bufsize)"
        throw(DimensionMismatch(msg))
    end
    length(workarr) ≥ workarrlength(smpl, al; dim) ||
        throw(ArgumentError("workarr too small"))
    return div_noalloc!(buffer, smpl.matrix_svd, al, workarr, dim)
end

"""
    movedim(arr::AbstractArray, src => dst)

Move `arr`'s dimension at `src` to `dst` while keeping the order of the remaining
dimensions unchanged.
"""
function movedim(arr::AbstractArray{T,N}, dims::Pair) where {T,N}
    src, dst = dims
    src == dst && return arr
    return permutedims(arr, getperm(N, dims))
end

function getperm(N, dims::Pair)
    src, dst = dims
    perm = collect(1:N)
    deleteat!(perm, src)
    insert!(perm, dst, src)
    return perm
end

"""
    matop_along_dim!(buffer, mat, arr::AbstractArray, dim::Integer, op)

Apply the operator `op` to the matrix `mat` and to the array `arr` along the dimension
`dim`, writing the result to `buffer`.
"""
function matop_along_dim!(buffer, mat, arr::AbstractArray{T,N}, dim, op) where {T,N}
    1 ≤ dim ≤ N || throw(DomainError(dim, "Dimension must be in [1, $N]"))

    if dim == 1
        matop!(buffer, mat, arr, op, 1)
    elseif dim != N
        # Move the target dim to the first position
        perm = getperm(N, dim => 1)
        arr_perm = permutedims(arr, perm)
        buffer_perm = permutedims(buffer, perm)
        matop!(buffer_perm, mat, arr_perm, op, 1)
        permutedims!(buffer, buffer_perm, getperm(N, 1 => dim))
    else
        # Apply the operator to the last dimension
        matop!(buffer, mat, arr, op, N)
    end
    return buffer
end

"""
    matop!(buffer, mat, arr::AbstractArray, op, dim)

Apply the operator `op` to the matrix `mat` and to the array `arr` along the first
dimension (dim=1) or the last dimension (dim=N).
"""
function matop!(
    buffer::AbstractArray{S,N}, mat, arr::AbstractArray{T,N}, op, dim
) where {S,T,N}
    if dim == 1
        flatarr = reshape(arr, (size(arr, 1), :))
        flatbuffer = reshape(buffer, (size(buffer, 1), :))
        op(flatbuffer, mat, flatarr)
    elseif dim == N
        flatarr = reshape(arr, (:, size(arr, N)))
        flatbuffer = reshape(buffer, (:, size(buffer, N)))
        op(flatbuffer, flatarr, transpose(mat))
    else
        throw(DomainError("Dimension must be 1 or $N"))
    end
    return buffer
end

function div_noalloc!(
    buffer::AbstractArray{S,N}, mat::SVD, arr::AbstractArray{T,N}, workarr, dim
) where {S,T,N}
    1 ≤ dim ≤ N || throw(DomainError(dim, "Dimension must be in [1, $N]"))

    if dim == 1
        flatarr = reshape(arr, (size(arr, 1), :))
        flatbuffer = reshape(buffer, (size(buffer, 1), :))
        ldiv_noalloc!(flatbuffer, mat, flatarr, workarr)
    elseif dim != N
        # Move the target dim to the first position
        arr_perm = movedim(arr, dim => 1)
        buffer_perm = movedim(buffer, dim => 1)
        flatarr = reshape(arr_perm, (size(arr_perm, 1), :))
        flatbuffer = reshape(buffer_perm, (size(buffer_perm, 1), :))
        ldiv_noalloc!(flatbuffer, mat, flatarr, workarr)
        buffer .= movedim(buffer_perm, 1 => dim)
    else
        flatarr = reshape(arr, (:, size(arr, N)))
        flatbuffer = reshape(buffer, (:, size(buffer, N)))
        rdiv_noalloc!(flatbuffer, flatarr, mat, workarr)
    end
    return buffer
end

function ldiv_noalloc!(Y::AbstractMatrix, A::SVD, B::AbstractMatrix, workarr)
    # Setup work space
    worksize = (size(A.U, 2), size(B, 2))
    worklength = prod(worksize)
    length(workarr) ≥ worklength ||
        throw(DimensionMismatch("size(workarr)=$(size(workarr)), min worksize=$worklength"))
    workarr_view = reshape(view(workarr, 1:worklength), worksize)

    mul!(workarr_view, A.U', B)
    workarr_view ./= A.S
    return mul!(Y, A.V, workarr_view)
end

function rdiv_noalloc!(Y::AbstractMatrix, A::AbstractMatrix, B::SVD, workarr)
    # Setup work space
    worksize = (size(A, 1), size(B.U, 2))
    worklength = prod(worksize)
    length(workarr) ≥ worklength ||
        throw(DimensionMismatch("size(workarr)=$(size(workarr)), min worksize=$worklength"))
    workarr_view = reshape(view(workarr, 1:worklength), worksize)

    # Note: conj creates a temporary matrix
    mul!(workarr_view, A, conj(B.U))
    workarr_view ./= reshape(B.S, 1, :)
    return mul!(Y, workarr_view, conj(B.Vt))
end


export RQ, rq, rq!

"""
    RQ <: Factorization

An RQ matrix factorization stored in a packed format, typically obtained from
[`rq`](@ref). If ``A`` is an `m`×`n` matrix, then

```math
A = [0 R] Q
```

where ``Q`` is an orthogonal/unitary matrix and ``R`` is upper triangular (trapezoidal
if `m < n`).
The matrix ``Q`` is stored as a sequence of Householder reflectors ``v_i``
and coefficients ``\\tau_i`` where:

```math
Q = \\prod_{i=1}^{\\min(m,n)} (I - \\tau_i v_i v_i^T).
```

Iterating the decomposition produces the components `R` and `Q`.

The object has two fields:

* `factors` is an `m`×`n` matrix.

  - The upper triangular part contains the elements of ``R``, that is `R =
    triu(F.factors)` for a `QR` object `F`.

  - The subdiagonal part contains the reflectors ``v_i`` stored in a packed format where
    ``v_i`` is the ``i``th column of the matrix `V = I + tril(F.factors, -1)`.

* `τ` is a vector  of length `min(m,n)` containing the coefficients ``\tau_i``.

"""
struct RQ{T,S<:AbstractMatrix{T}} <: Factorization{T}
    factors::S
    τ::Vector{T}

    function RQ{T,S}(factors, τ) where {T,S<:AbstractMatrix{T}}
        require_one_based_indexing(factors)
        new{T,S}(factors, τ)
    end
end
RQ(factors::AbstractMatrix{T}, τ::Vector{T}) where {T} = RQ{T,typeof(factors)}(factors, τ)
function RQ{T}(factors::AbstractMatrix, τ::AbstractVector) where {T}
    RQ(convert(AbstractMatrix{T}, factors), convert(Vector{T}, τ))
end

# iteration for destructuring into components
Base.iterate(S::RQ) = (S.R, Val(:Q))
Base.iterate(S::RQ, ::Val{:Q}) = (S.Q, Val(:done))
Base.iterate(S::RQ, ::Val{:done}) = nothing

Base.size(F::RQ, dim::Integer) = size(getfield(F, :factors), dim)
Base.size(F::RQ) = size(getfield(F, :factors))

function Base.getproperty(F::RQ, d::Symbol)
    m, n = size(F)
    if d === :R
        if m <= n
            return triu!(getfield(F, :factors)[1:m, n-m+1:n])
        else
            return triu(getfield(F, :factors),min(0,n-m))
        end
    elseif d === :Q
        return RQPackedQ(getfield(F, :factors), F.τ)
    else
        getfield(F, d)
    end
end

"""
    RQPackedQ <: AbstractMatrix

The orthogonal/unitary ``Q`` matrix of an RQ factorization stored in [`RQ`](@ref) format.
"""
struct RQPackedQ{T,S<:AbstractMatrix{T}} <: AbstractQ{T}
    factors::S
    τ::Vector{T}

    function RQPackedQ{T,S}(factors, τ) where {T,S<:AbstractMatrix{T}}
        require_one_based_indexing(factors)
        new{T,S}(factors, τ)
    end
end

# Conversions
RQ{T}(A::RQ) where {T} = RQ(convert(AbstractMatrix{T}, A.factors), convert(AbstractVector{T}, A.τ))
Factorization{T}(A::RQ{T}) where {T} = A
Factorization{T}(A::RQ) where {T} = RQ{T}(A)
AbstractMatrix(F::RQ) = F.R * F.Q
AbstractArray(F::RQ) = AbstractMatrix(F)
Matrix(F::RQ) = Array(AbstractArray(F))
Array(F::RQ) = Matrix(F)

function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, F::RQ)
    summary(io, F); println(io)
    println(io, "R factor:")
    show(io, mime, F.R)
    println(io, "\nQ factor:")
    show(io, mime, F.Q)
end

RQPackedQ(factors::AbstractMatrix{T}, τ::Vector{T}) where {T} = RQPackedQ{T,typeof(factors)}(factors, τ)
function RQPackedQ{T}(factors::AbstractMatrix, τ::AbstractVector) where {T}
    RQPackedQ(convert(AbstractMatrix{T}, factors), convert(Vector{T}, τ))
end
RQPackedQ{T}(Q::RQPackedQ) where {T} = RQPackedQ(convert(AbstractMatrix{T}, Q.factors), convert(Vector{T}, Q.τ))
AbstractMatrix{T}(Q::RQPackedQ{T}) where {T} = Q
AbstractMatrix{T}(Q::RQPackedQ) where {T} = RQPackedQ{T}(Q)
# this makes the rectangular version, alas.
# function Matrix{T}(Q::RQPackedQ{T}) where {T<:BlasFloat}
#     m,n = size(Q.factors)
#     i0= m > n ? (m-n+1) : 1
#     LAPACK.orgrq!(Q.factors[i0:m,1:n], Q.τ)
# end

Matrix{T}(Q::RQPackedQ{S}) where {T,S} = Matrix{T}(lmul!(Q, Matrix{S}(I, size(Q,2), size(Q,2))))
Base.size(Q::RQPackedQ, dim::Integer) = size(getfield(Q, :factors), dim == 1 ? 2 : dim)

function lmul!(A::RQPackedQ{T,S}, B::StridedVecOrMat{T}) where {T<:BlasFloat, S<:StridedMatrix}
    m,n = size(A.factors)
    if m > n
        return LAPACK.ormrq!('L','N',view(A.factors, m-n+1:m, 1:n),A.τ,B)
    else
        return LAPACK.ormrq!('L','N',A.factors,A.τ,B)
    end
end
function lmul!(A::RQPackedQ, B::AbstractVecOrMat)
    require_one_based_indexing(B)
    mA, nA = size(A.factors)
    mB, nB = size(B,1), size(B,2)
    if nA != mB
        throw(DimensionMismatch("matrix A has dimensions ($nA,$nA) but B has dimensions ($mB, $nB)"))
    end
    Afactors = A.factors
    mm = min(mA,nA)
    @inbounds begin
        for k = mm:-1:1
            krow = mA-mm+k
            kcol = nA-mm+k
            for j = 1:nB
                vBj = B[kcol,j]
                for i = 1:kcol-1
                    # vBj += conj(Afactors[i,k])*B[i,j]
                    vBj += (Afactors[krow,i])*B[i,j]
                end
                vBj = conj(A.τ[k])*vBj
                B[kcol,j] -= vBj
                for i = 1:kcol-1
                    # B[i,j] -= Afactors[i,k]*vBj
                    B[i,j] -= conj(Afactors[krow,i])*vBj
                end
            end
        end
    end
    B
end

function (*)(B::StridedMatrix, Q::RQPackedQ)
    TBQ = promote_type(eltype(Q), eltype(B))
    Qnew = convert(AbstractMatrix{TBQ}, Q)
    if size(Q.factors, 2) == size(B, 2)
        Bnew = copy_oftype(B, TBQ)
    elseif size(Q.factors, 1) == size(B, 2)
        # People will expect this, misguided though it seems.
        Bnew = hcat(zeros(TBQ, size(B,1), size(Q.factors, 2) - size(B,2)),B)
    else
        throw(DimensionMismatch("first dimension of matrix must have size either $(size(Q.factors, 1)) or $(size(Q.factors, 2))"))
    end
    rmul!(Bnew, Qnew)
end

function lmul!(adjA::Adjoint{<:Any,<:RQPackedQ{T,S}}, B::StridedVecOrMat{T}
               ) where {T<:BlasReal,S<:StridedMatrix}
    A = adjA.parent
    m,n = size(A.factors)
    if m > n
        return LAPACK.ormrq!('L','T',view(A.factors, m-n+1:m, 1:n),A.τ,B)
    else
        return LAPACK.ormrq!('L','T',A.factors,A.τ,B)
    end
end
function lmul!(adjA::Adjoint{<:Any,<:RQPackedQ{T,S}}, B::StridedVecOrMat{T}
               ) where {T<:BlasComplex,S<:StridedMatrix}
    A = adjA.parent
    m,n = size(A.factors)
    if m > n
        return LAPACK.ormrq!('L','C',view(A.factors, m-n+1:m, 1:n),A.τ,B)
    else
        return LAPACK.ormrq!('L','C',A.factors,A.τ,B)
    end
end
function lmul!(adjA::Adjoint{<:Any,<:RQPackedQ}, B::AbstractVecOrMat)
    require_one_based_indexing(B)
    A = adjA.parent
    mA, nA = size(A.factors)
    mB, nB = size(B,1), size(B,2)
    if nA != mB
        throw(DimensionMismatch("matrix A has dimensions ($nA,$nA) but B has dimensions ($mB, $nB)"))
    end
    Afactors = A.factors
    mm = min(mA,nA)
    @inbounds begin
        for k = 1:mm
            krow = mA-mm+k
            kcol = nA-mm+k
            for j = 1:nB
                vBj = B[kcol,j]
                for i = 1:kcol-1
                    # vBj += conj(Afactors[i,k])*B[i,j]
                    vBj += (Afactors[krow,i])*B[i,j]
                end
                vBj = A.τ[k]*vBj
                B[kcol,j] -= vBj
                for i = 1:kcol-1
                    # B[i,j] -= Afactors[i,k]*vBj
                    B[i,j] -= conj(Afactors)[krow,i]*vBj
                end
            end
        end
    end
    B
end
function rmul!(A::StridedVecOrMat{T}, B::RQPackedQ{T,S}
               ) where {T<:BlasFloat,S<:StridedMatrix}
    m,n = size(B.factors)
    if m > n
        return LAPACK.ormrq!('R','N',view(B.factors, m-n+1:m, 1:n),B.τ,A)
    else
        return LAPACK.ormrq!('R','N',B.factors,B.τ,A)
    end
end
function rmul!(A::StridedMatrix,Q::RQPackedQ)
    mQ, nQ = size(Q.factors)
    mA, nA = size(A,1), size(A,2)
    if nA != nQ
        throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but matrix Q has dimensions ($nQ, $nQ)"))
    end
    Qfactors = Q.factors
    mm = min(mQ, nQ)
    @inbounds begin
        for k = 1:min(mQ,nQ)
            krow = mQ-mm+k
            kcol = nQ-mm+k
            for i = 1:mA
                vAi = A[i,kcol]
                for j = 1:kcol-1
                    # vAi += A[i,j]*Qfactors[j,k]
                    vAi += A[i,j]*conj(Qfactors[krow,j])
                end
                vAi = vAi*conj(Q.τ[k])
                A[i,kcol] -= vAi
                for j = 1:kcol-1
                    # A[i,j] -= vAi*conj(Qfactors[j,k])
                    A[i,j] -= vAi*(Qfactors[krow,j])
                end
            end
        end
    end
    A
end
function rmul!(A::StridedVecOrMat{T}, adjB::Adjoint{<:Any,<:RQPackedQ{T}}
               ) where {T<:BlasReal}
    B = adjB.parent
    m,n = size(B.factors)
    if m > n
        return LAPACK.ormrq!('R','T',view(B.factors, m-n+1:m, 1:n),B.τ,A)
    else
        return LAPACK.ormrq!('R','T',B.factors,B.τ,A)
    end
end
function rmul!(A::StridedVecOrMat{T}, adjB::Adjoint{<:Any,<:RQPackedQ{T}}
               ) where {T<:BlasComplex}
    B = adjB.parent
    m,n = size(B.factors)
    if m > n
        return LAPACK.ormrq!('R','C',view(B.factors, m-n+1:m, 1:n),B.τ,A)
    else
        return LAPACK.ormrq!('R','C',B.factors,B.τ,A)
    end
end
function rmul!(A::StridedMatrix, adjQ::Adjoint{<:Any,<:RQPackedQ})
    Q = adjQ.parent
    mQ, nQ = size(Q.factors)
    mA, nA = size(A,1), size(A,2)
    if nA != nQ
        throw(DimensionMismatch("matrix A has dimensions ($mA,$nA) but matrix Q has dimensions ($nQ, $nQ)"))
    end
    Qfactors = Q.factors
    mm = min(mQ, nQ)
    @inbounds begin
        for k = min(mQ,nQ):-1:1
            krow = mQ-mm+k
            kcol = nQ-mm+k
            for i = 1:mA
                vAi = A[i,kcol]
                for j = 1:kcol-1
                    # vAi += A[i,j]*Qfactors[j,k]
                    vAi += A[i,j]*conj(Qfactors[krow,j])
                end
                vAi = vAi*Q.τ[k]
                A[i,kcol] -= vAi
                for j = 1:kcol-1
                    # A[i,j] -= vAi*conj(Qfactors[j,k])
                    A[i,j] -= vAi*(Qfactors[krow,j])
                end
            end
        end
    end
    A
end

ldiv!(A::RQ{T}, B::StridedMatrix{S}) where {T,S} = lmul!(adjoint(A.Q), ldiv!(UpperTriangular(A.R), B))
ldiv!(A::RQ{T}, B::StridedVector{S}) where {T,S} = lmul!(adjoint(A.Q), ldiv!(UpperTriangular(A.R), B))

# apply shifted reflector from right
@inline function reflectorYlppa!(x::AbstractVector, τ::Number, A::StridedMatrix)
    require_one_based_indexing(x)
    m, n = size(A)
    if length(x) != n
        throw(DimensionMismatch("reflector has length $(length(x)), which must match the second dimension of matrix A, $n"))
    end
    @inbounds begin
        for j = 1:m
            # dot
            Ajv = A[j, n]
            for i = 1:n-1
                Ajv += A[j, i] * x[i+1]
            end

            Ajv = τ * Ajv

            # ger
            A[j, n] -= Ajv
            for i = 1:n-1
                A[j, i] -= Ajv * x[i+1]'
            end
        end
    end
    return A
end

function _rqfactUnblocked!(A::AbstractMatrix{T}) where {T}
    require_one_based_indexing(A)
    m, n = size(A)
    τ = zeros(T, min(m,n))
    v = zeros(T, max(m,n))
    mm = min(m,n)
    for k = mm:-1:1
        krow = m-mm+k
        kcol = n-mm+k
        v[1] = A[krow,kcol]'
        for i = 2:kcol
            v[i] = A[krow,i-1]'
        end
        x = view(v, 1:kcol)
        if kcol==1 && T<:Real
            τk = zero(T)
        else
            τk = reflector!(x)
        end
        τ[k] = τk
        reflectorYlppa!(x, τk, view(A, 1:krow-1, 1:kcol))
        A[krow,kcol] = v[1]'
        for i = 2:kcol
            A[krow,i-1] = v[i]'
        end
    end
    RQ(A, τ)
end

"""
    rq(A) -> F

Compute the RQ factorization of the matrix `A`: an orthogonal/unitary matrix `Q` and
and upper triangular or trapezoidal matrix `R` such that
```math
A = [0 R] Q
```
The returned object `F` stores the factorization in a packed format [`RQ`](@ref).
Unlike the `QR` factorization this does not handle least squares problems.
If `A` is `m`×`n` then `Q` is `n`×`n` and `R` is `m`×`min(m,n)`.
"""
function rq end

"""
    rq!(A) -> F

The same as [`rq`](@ref), but saves space by overwriting the input `A`.
"""
rq!(A::StridedMatrix{<:BlasFloat}) = RQ(LAPACK.gerqf!(A)...)
rq!(A::StridedMatrix) = _rqfactUnblocked!(A)

rq(x::Number) = rq(fill(x,1,1))
function rq(v::AbstractVector)
    require_one_based_indexing(v)
    rq(reshape(v, (length(v), 1)))
end

# There is a local version of this:
# using LinearAlgebra: _qreltype

function rq(A::AbstractMatrix{T}) where {T}
    require_one_based_indexing(A)
    AA = similar(A, LinearAlgebra._qreltype(T), size(A))
    copyto!(AA, A)
    return rq!(AA)
end


function evaluate_kinetics_model(t,x,data_dictionary)

    # get stuff from the data_dictionary -
    stochiometric_matrix = data_dictionary["stochiometric_matrix"]
    number_of_compartments = data_dictionary["number_of_compartments"]
    number_of_states_per_compartment = data_dictionary["number_of_states_per_compartment"]
    number_of_rates_per_compartment = data_dictionary["number_of_rates_per_compartment"]
    compartment_volume_array = data_dictionary["compartment_volume_array"]

    # what is my kinetics function pointer?
    kinetics_function_pointer = data_dictionary["kinetics_function_pointer"]

    # assemble the overall rate vector -
    overall_kinetics_vector = Array{Float64,1}()
    for compartment_index = 1:number_of_compartments

        # calc the rates -
        rates_in_this_compartment = kinetics_function_pointer(t,x,compartment_index,data_dictionary)
        volume = compartment_volume_array[compartment_index]
        rate_array = stochiometric_matrix*(rates_in_this_compartment*volume)

        # package -
        for value in rate_array
            push!(overall_kinetics_vector,value)
        end
    end

    # return -
    return overall_kinetics_vector
end

function evaluate_transport_model(t,x,data_dictionary)

    # get stuff from the data_dictionary -
    number_of_compartments = data_dictionary["number_of_compartments"]
    number_of_states_per_compartment = data_dictionary["number_of_states_per_compartment"]

    # what is my transport function pointer?
    transport_function_pointer = data_dictionary["transport_function_pointer"]

    # assemble transport terms -
    overall_transport_vector = Array{Float64,1}()
    for compartment_index = 1:number_of_compartments

        # calc the transport terms -
        compartment_transport_terms = transport_function_pointer(t,x,compartment_index,data_dictionary)
        for value in compartment_transport_terms
            push!(overall_transport_vector,value)
        end
    end

    # return -
    return overall_transport_vector
end

function balances(dx,x,data_dictionary,t)

    # what is my system size?
    compartment_volume_array = data_dictionary["compartment_volume_array"]
    number_of_states_per_compartment = data_dictionary["number_of_states_per_compartment"]
    number_of_compartments = data_dictionary["number_of_compartments"]
    number_of_states = number_of_states_per_compartment*number_of_compartments

    # compute mass matrix -
    mass_vector = Array{Float64,1}()
    for compartment_index = 1:number_of_compartments
        volume = compartment_volume_array[compartment_index]
        for state_index = 1:number_of_states_per_compartment
            push!(mass_vector,volume)
        end
    end

    # compute the transport, and kinetics terms -
    transport_terms = evaluate_transport_model(t,x,data_dictionary)
    kinetic_terms = evaluate_kinetics_model(t,x,data_dictionary)

    # compute the right hand side of the balances -
    dxdt = transport_terms + kinetic_terms

    # package -
    for index = 1:number_of_states
        dx[index] = mass_vector[index]*dxdt[index]
    end
end


"""
    Farneback(Args...)

A method for dense optical flow estimation developed by Gunnar Farneback. It
computes the optical flow for all the points in the frame using the polynomial
representation of the images. The idea of polynomial expansion is to approximate
the neighbourhood of a point in a 2D function with a polynomial. Displacement
fields are estimated from the polynomial coefficients depending on how the
polynomial transforms under translation.

The different arguments are:

 -  flow_est          =  Array of SVector{2} containing estimate flow values for all points in the frame
 -  iterations        =  Number of iterations the displacement estimation algorithm is run at each
                         point
 -  window_size       =  Size of the search window at each pyramid level; the total size of the
                         window used is 2*window_size + 1
 -  σw                =  Standard deviation of the Gaussian weighting filter
 -  neighbourhood     =  size of the pixel neighbourhood used to find polynomial expansion for each pixel;
                         larger values mean that the image will be approximated with smoother surfaces,
                         yielding more robust algorithm and more blurred motion field
 -  σp                =  standard deviation of the Gaussian that is used to smooth derivatives used as a
                         basis for the polynomial expansion (Applicability)
 -  est_flag          =  true -> Use flow_est as initial estimate
                         false -> Assume zero initial flow values
 -  gauss_flag        =  false -> use box filter
                         true -> use gaussian filter instead of box filter of the same size for optical flow
                         estimation; usually, this option gives more accurate flow than with a box filter,
                         at the cost of lower speed (Default Value)

## References

Farnebäck G. (2003) Two-Frame Motion Estimation Based on Polynomial Expansion. In: Bigun J.,
Gustavsson T. (eds) Image Analysis. SCIA 2003. Lecture Notes in Computer Science, vol 2749. Springer, Berlin,
Heidelberg

Farnebäck, G.: Polynomial Expansion for Orientation and Motion Estimation. PhD thesis, Linköping University,
Sweden, SE-581 83 Linköping, Sweden (2002) Dissertation No 790, ISBN 91-7373-475-6.

"""

struct Farneback{F <: Float64, I <: Int, V <: Bool} <: OpticalFlowAlgo
    flow_est::Array{SVector{2, F}, 2}
    iterations::I
    window_size::I
    σw::F
    neighbourhood::I
    σp::F
    est_flag::V
    gauss_flag::V
end

Farneback(flow_est::Array{SVector{2, F}, 2}, iterations::I, window_size::I, σw::F, neighbourhood::I, σp::F, est_flag::V) where {F <: Float64, I <: Int, V <: Bool} = LK{F, I, V}(flow_est, iterations, window_size, σw, neighbourhood, σp, est_flag, true)

function optflow(first_img::AbstractArray{T, 2}, second_img::AbstractArray{T,2}, algo::Farneback{}) where T <: Gray
    if algo.est_flag
        @assert size(first_img) == size(algo.flow_est)
    end
    A1, B1, _ = polynomial_expansion(first_img, algo.neighbourhood, algo.σp)
    A2, B2, _ = polynomial_expansion(second_img, algo.neighbourhood, algo.σp)

    AtA = similar(A1)
    AtB = similar(B1)

    disp = Array{SVector{2, Float64}, 2}(size(first_img))
    for i = 1:size(first_img)[1]
        for j = 1:size(first_img)[2]
            disp[i,j] = SVector{2, Float64}(0.0,0.0)
        end
    end

    for k = 1:algo.iterations
        for i = 1:size(first_img)[1]
            for j = 1:size(first_img)[2]
                d = disp[i,j]

                mod_i = round(Int, i + d[2])
                mod_j = round(Int, j + d[1])

                mod_i = clamp(mod_i, 1, size(first_img)[1])
                mod_j = clamp(mod_j, 1, size(first_img)[2])

                sub_A1 = SMatrix{2,2,Float64}(@view A1[i,j,:,:])
                sub_B1 = SVector{2,Float64}(@view B1[i,j,:])
                sub_A2 = SMatrix{2,2,Float64}(@view A2[mod_i,mod_j,:,:])
                sub_B2 = SVector{2,Float64}(@view B2[mod_i,mod_j,:])

                A = (sub_A1 + sub_A2)/2
                dB = A*d - ((sub_B2 - sub_B1)/2)

                AtA[i,j,:,:] .= A'*A
                AtB[i,j,:] .= A'*dB
            end
        end

        len = floor(Int, algo.window_size/2)
        w = ImageFiltering.KernelFactors.gaussian((algo.σw),(algo.window_size))[-len:len]

        for i = 1:2
            for j = 1:2
                if algo.gauss_flag
                    AtA[:,:,i,j] = @view conv2(w, w, AtA[:,:,i,j])[1+len:size(first_img)[1]+len,1+len:size(first_img)[2]+len]
                else
                    AtA[:,:,i,j] =  @view conv2(ones(algo.window_size)./algo.window_size, ones(algo.window_size)./algo.window_size, AtA[:,:,i,j])[1+len:size(first_img)[1]+len,1+len:size(first_img)[2]+len]
                end
            end
            if algo.gauss_flag
                AtB[:,:,i] = @view conv2(w, w, AtB[:,:,i])[1+len:size(first_img)[1]+len,1+len:size(first_img)[2]+len]
            else
                AtB[:,:,i,j] = @view conv2(ones(algo.window_size)./algo.window_size, ones(algo.window_size)./algo.window_size, AtB[:,:,i,j])[1+len:size(first_img)[1]+len,1+len:size(first_img)[2]+len]
            end
        end

        for i = 1:size(first_img)[1]
            for j = 1:size(first_img)[2]
                disp[i,j] = pinv2x2(SMatrix{2,2,Float64}(AtA[i,j,:,:]))*SVector{2,Float64}(AtB[i,j,:])
            end
        end
    end

    return disp

end

"""
```
A, B, C = polynomial_expansion(img, neighbourhood, σ)
```

Returns the polynomial coefficients of the approximation of the neighbourhood of a
point in a 2D function with a polynomial.
The expansion is: f(x) = x'Ax + B'x + C

Parameters:

-  img               =  Grayscale image who polynomial expansion is to be found
-  neighbourhood     =  Size of the pixel neighbourhood used to find polynomial expansion for each pixel;
                        larger values mean that the image will be approximated with smoother surfaces,
                        yielding more robust algorithm and more blurred motion field
-  σ                 =  Standard deviation of the Gaussian that is used to smooth derivatives used as a
                        basis for the polynomial expansion (Applicability)

## References

Farnebäck, G.: Polynomial Expansion for Orientation and Motion Estimation. PhD thesis, Linköping University,
Sweden, SE-581 83 Linköping, Sweden (2002) Dissertation No 790, ISBN 91-7373-475-6.

"""

function polynomial_expansion(img::AbstractArray{T, 2}, neighbourhood::Int, σ::Real) where T <: Gray
    len = floor(Int, neighbourhood/2)
    a = ImageFiltering.Kernel.gaussian((σ,σ), (neighbourhood,neighbourhood))[-len:len, -len:len]

    c = ones(size(img))
    padded_c = parent(padarray(c, Fill(0, (len, len))))
    padded_img = parent(padarray(img, Fill(0, (len, len))))

    B0 = zeros(neighbourhood*neighbourhood, 6)
    for x = -len:len
        for y = -len:len
            B0[neighbourhood*(x + len) + (y + len) + 1, :] = [1, x, y, x*x, y*y, x*y]
        end
    end

    B = SMatrix{neighbourhood*neighbourhood, 6}(B0)
    Wa = SDiagonal{neighbourhood^2,Float64}( @view a[:])
    BtWa = B'*Wa

    # If we assume that the confidence is constant we can move several
    # matrix multiplications out of the main loop.
    Wc = SDiagonal{neighbourhood^2,Float64}(ones(1,neighbourhood^2)...)
    BtWaWc = BtWa*Wc
    G = BtWaWc*B
    G_inv = SMatrix{6,6,Float64}(pinv(convert(Array,G)))
    G_invBtWaWc = G_inv*BtWaWc

    poly_A = zeros(size(img)[1], size(img)[2], 2, 2)
    poly_B = zeros(size(img)[1], size(img)[2], 2)
    poly_C = zeros(size(img))

    for i = (len + 1):(size(img)[1] - len)
        for j = (len + 1):(size(img)[2] - len)
            sub_img = @view padded_img[(i - len):(i + len), (j - len):(j + len)]
            r = G_invBtWaWc*SVector{neighbourhood^2,Float64}(@view sub_img[:])

            poly_A[i-len, j-len, 1, 1] =  r[4]
            poly_A[i-len, j-len, 2, 1] =  r[6]/2
            poly_A[i-len, j-len, 1, 2] =  r[6]/2
            poly_A[i-len, j-len, 2, 2] =  r[5]

            poly_B[i-len, j-len, 1] = r[2]
            poly_B[i-len, j-len, 2] = r[3]

            poly_C[i-len, j-len] = r[1]
        end
    end

    return (poly_A, poly_B, poly_C)

end


using LangrethsRules
using Test

function integrate(x::AbstractVector, y::AbstractVector)
    if isone(length(y))
        return zero(first(y))
    end
    @inbounds retval = (x[2] - x[1]) * (y[1] + y[2])
    @inbounds @fastmath @simd for i in 2:(length(y)-1)
        retval += (x[i+1] - x[i]) * (y[i] + y[i+1])
    end
    return 1 // 2 * retval
end

function compute(A, B, ts)
    l = zero(lesser(A))
    for i in 1:size(l, 1)
        for j in 1:i
            l[i,j] += integrate(ts, greater(A)[i, 1:i] .* lesser(B)[1:i, j])
            l[i,j] -= integrate(ts, lesser(A)[i, 1:j] .* greater(B)[1:j, j])
            l[i,j] -= integrate(ts[j:i], lesser(A)[i, j:i] .* lesser(B)[j:i, j])

            l[j,i] = -conj(l[i,j])
        end
        l[i,i] = eltype(l) <: Real ? 0.0 : im * imag(l[i,i])
    end

    g = zero(greater(A))
    for i in 1:size(g, 1)
        for j in 1:i
            g[i,j] -= integrate(ts, lesser(A)[i, 1:i] .* greater(B)[1:i, j])
            g[i,j] += integrate(ts, greater(A)[i, 1:j] .* lesser(B)[1:j, j])
            g[i,j] += integrate(ts[j:i], greater(A)[i, j:i] .* greater(B)[j:i, j])

            g[j,i] = -conj(g[i,j])
        end
        g[i,i] = eltype(g) <: Real ? 0.0 : im * imag(g[i,i])
    end

    TimeOrderedGF(l, g, ts)
end

N = 100

# Suppose a non-equidistant time-grid
ts = sort(N*rand(N));

# And 2 time-ordered GFs defined on that grid
a = let
    x = rand(ComplexF64,N,N)
    y = rand(ComplexF64,N,N)
    TimeOrderedGF(x - x', y - y', ts) # Skew-symmetric L and G components
end
b = let
    x = rand(ComplexF64,N,N)
    y = rand(ComplexF64,N,N)
    TimeOrderedGF(x - x', y - y', ts) # Skew-symmetric L and G components
end

@test greater(a ⋆ b) ≈ greater(compute(a, b, ts))
@test lesser(a ⋆ b) ≈ lesser(compute(a, b, ts))
@test retarded(a ⋆ b) ≈ adjoint(advanced(a ⋆ b))
# @test retarded(a ⋆ b ⋆ a) ≈ retarded(a ⋆ (b ⋆ a))


using QuantumLab
using Base.Test

if (!isdefined(:indent))
  indent = ""
end

info("$(now())  $(indent)READING:   h2o.xyz -> h2o::Geometry ")
h2o = readGeometryXYZ("h2o.xyz")

# test BasisSetExchange
# as this takes quite some time we primarily only want to do this during continuous integration (travis-ci) and when we haven-t checked this before
if (!isfile("STO-3G.tx93"))
  info("$(now())  $(indent)OBTAINING:   STO-3G.tx93")
  bseEntries = obtainBasisSetExchangeEntries()
  display(bseEntries)
  stoEntry   = computeBasisSetExchangeEntry("sto-3g",bseEntries)[3]
  downloadBasisSetBasisSetExchange(stoEntry,"STO-3G.tx93")
  @test_throws ErrorException computeBasisSetExchangeEntry("NotDefined",bseEntries)
end

info("$(now())  $(indent)READING:   STO-3G.tx93 -> sto3g::BasisSet")
sto3g = readBasisSetTX93("STO-3G.tx93")

info("$(now())  $(indent)COMPUTING:   sto3g,h2o -> bas")
bas = computeBasis(sto3g,h2o)

info("$(now())  $(indent)COMPUTING:   bas,h2o -> matrixOverlap, matrixKinetic, matrixNuclearAttraction, ERIs")
matrixOverlap = computeMatrixOverlap(bas)
matrixKinetic = computeMatrixKinetic(bas)
matrixNuclearAttraction = computeMatrixNuclearAttraction(bas,h2o)
ERIs = computeTensorElectronRepulsionIntegrals(bas)

info("$(now())  $(indent)COMPUTING:   h2o -> matrixSADguess")
matrixSADguess = computeDensityGuessSAD("HF","STO-3G",h2o)

info("$(now())  $(indent)COMPUTING:   sto3g, h2o -> shells, shells_native")
shells = computeBasisShellsLibInt2(sto3g,h2o)
shells_native = computeBasisShells(sto3g,h2o)

info("$(now())  $(indent)COMPUTING:   (HartreeFock)  shells, h2o, matrixSADguess -> density")
density = evaluateSCF(shells,h2o,mean(matrixSADguess),info=false,detailedinfo=false)[3]

info("$(now())  $(indent)COMPUTING:   density, matrixKinetic, matrixNuclearAttraction, ERIs -> matrixFock")
matrixFock = computeMatrixFock(density,matrixKinetic,matrixNuclearAttraction,ERIs)

densityvirt = inv(matrixOverlap) - density

info("$(now())  $(indent)COMPUTING:   shells, density -> S, T, J, K, Fock")
S = computeMatrixOverlap(shells)
T = computeMatrixKinetic(shells)
J = computeMatrixCoulomb(shells,density)
K = computeMatrixExchange(shells,density)
V = computeMatrixNuclearAttraction(shells,h2o)
fock = computeMatrixFock(T,V,J,K)
moenergies = eigvals(Symmetric(fock),Symmetric(S))


export iauPvdpv
"""
Inner (=scalar=dot) product of two pv-vectors.

This function is part of the International Astronomical Union's
SOFA (Standards Of Fundamental Astronomy) software collection.

Status:  vector/matrix support function.

Given:
   a        double[2][3]      first pv-vector
   b        double[2][3]      second pv-vector

Returned:
   adb      double[2]         a . b (see note)

Note:

   If the position and velocity components of the two pv-vectors are
   ( ap, av ) and ( bp, bv ), the result, a . b, is the pair of
   numbers ( ap . bp , ap . bv + av . bp ).  The two numbers are the
   dot-product of the two p-vectors and its derivative.

Called:
   iauPdp       scalar product of two p-vectors

This revision:  2013 June 18

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# void iauPvdpv(double a[2][3], double b[2][3], double adb[2])
function iauPvdpv(a::AbstractMatrix{<:Real}, b::AbstractMatrix{<:Real})
   # Preallocate return values
   adb = zeros(Float64, 2)

   ccall((:iauPvdpv, libsofa_c), Cvoid,
         (Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}),
          convert(Array{Float64, 2}, a'),
          convert(Array{Float64, 2}, b'),
          adb)

   return SVector{2}(adb)
end


# Tests automatic differentiation on sn.jl:
include("../src/sn_jacobian.jl")
using PyPlot

@testset "sn_jacobian" begin

function test_sn_jacobian2(l_max)
r0 = [0.01,100.0]
nb = 10
#l_max = 20
#l_max = 5
n_max = l_max^2+2*l_max

# Now, carry out finite-difference derivative computation:
function sn_jac_num(l_max::Int64,r::T,b::T) where {T <: Real}
  dq = big(1e-24)
  n_max = l_max^2+2*l_max
# Allocate an array for s_n:
  sn_big = zeros(BigFloat,n_max+1)
# Make BigFloat versions of r & b:
  r_big = big(r); b_big = big(b)
# Compute s_n to BigFloat precision:
  s_n!(l_max,r_big,b_big,sn_big)
# Now, compute finite differences:
  sn_jac_big= zeros(BigFloat,n_max+1,2)
  sn_plus = zeros(BigFloat,n_max+1)
  s_n!(l_max,r_big+dq,b_big,sn_plus)
  sn_minus = zeros(BigFloat,n_max+1)
  s_n!(l_max,r_big-dq,b_big,sn_minus)
  sn_jac_big[:,1] = (sn_plus-sn_minus)*.5/dq
  s_n!(l_max,r_big,b_big+dq,sn_plus)
  s_n!(l_max,r_big,b_big-dq,sn_minus)
  sn_jac_big[:,2] = (sn_plus-sn_minus)*.5/dq
return convert(Array{Float64,1},sn_big),convert(Array{Float64,2},sn_jac_big)
end

#sn_jacobian
#convert(Array{Float64,2},sn_jac_big)-sn_jacobian

get_cmap("plasma")
epsilon = 1e-10; delta = 1e-3
for i=1:2
  fig,axes = subplots(1,1)
  r=r0[i]
  if r < 1.0
    b = [linearspace(epsilon,delta,nb); linearspace(delta,r-delta,nb);
     -logarithmspace(log10(delta),log10(epsilon),nb) .+r; linearspace(r-epsilon,r+epsilon,nb); logarithmspace(log10(epsilon),log10(delta),nb) .+r;
     linearspace(r+delta,1-r-delta,nb); -logarithmspace(log10(delta),log10(epsilon),nb) .+(1-r); linearspace(1-r-epsilon,1-r+epsilon,nb);
     logarithmspace(log10(epsilon),log10(delta),nb).+(1-r); linearspace(1-r+delta,1+r-delta,nb); -logarithmspace(log10(delta),log10(epsilon),nb) .+(1+r)]
     nticks = 14
     xticknames=[L"$10^{-15}$",L"$10^{-12}$",L"$10^{-3}$",L"$r-10^{-3}$",L"$r-10^{-12}$",L"$r+10^{-12}$",L"$r+10^{-3}$",
     L"$1-r-10^{-3}$",L"$1-r-10^{-12}$",L"$1-r+10^{-12}$",L"$1-r+10^{-3}$",L"$1+r-10^{-3}$",L"$1+r-10^{-12}$",L"$1+r-10^{-15}$"]
  else
    b = [logarithmspace(log10(epsilon),log10(delta),nb).+(r-1); linearspace(r-1+delta,r-delta,nb);
     -logarithmspace(log10(delta),log10(epsilon),nb) .+r; linearspace(r-epsilon,r+epsilon,nb); logarithmspace(log10(epsilon),log10(delta),nb) .+r;
     linearspace(r+delta,r+1-delta,nb); -logarithmspace(log10(delta),log10(epsilon),nb) .+(r+1)]
     nticks = 8
     xticknames=[L"$r-1+10^{-12}$",L"$r-1+10^{-3}$",L"$r-10^{-3}$",L"$r-10^{-12}$",L"$r+10^{-12}$",L"$r+10^{-3}$",
     L"$r+1-10^{-3}$",L"$r+1-10^{-12}$"]
  end
#  if r < 1.0
#    b = [linearspace(1e-8,epsilon,nb); linearspace(epsilon,delta,nb); linearspace(delta,r-delta,nb);
##     linearspace(r-delta,r-epsilon,nb); linearspace(r-epsilon,r,nb); linearspace(r,r+epsilon,nb); linearspace(r+epsilon,r+delta,nb);
#     linearspace(r-delta,r-epsilon,nb); linearspace(r-epsilon,r+epsilon,nb); linearspace(r+epsilon,r+delta,nb);
#     linearspace(r+delta,1-r-delta,nb); linearspace(1-r-delta,1-r-epsilon,nb); linearspace(1-r-epsilon,1-r+epsilon,nb);
#     linearspace(1-r+epsilon,1-r+delta,nb); linearspace(1-r+delta,1+r-delta,nb); linearspace(1+r-delta,1+r-epsilon,nb);linearspace(1+r-epsilon,1+r-1e-8,nb)]
#  else
#    b = [linearspace(r-1+1e-8,r-1+epsilon,nb); linearspace(r-1+epsilon,r-1+delta,nb); linearspace(r-1+delta,r-delta,nb);
##     linearspace(r-delta,r-epsilon,nb); linearspace(r-epsilon,r,nb); linearspace(r,r+epsilon,nb); linearspace(r+epsilon,r+delta,nb);
#     linearspace(r-delta,r-epsilon,nb); linearspace(r-epsilon,r+epsilon,nb); linearspace(r+epsilon,r+delta,nb);
#     linearspace(r+delta,r+1-delta,nb); linearspace(r+1-delta,r+1-epsilon,nb); linearspace(r+1-epsilon,r+1-1e-8,nb)]
#  end
  k2 = (1.-(r-b).^2)./(4.*b.*r)
  k2H = ((b+1).^2-r.^2)./(4b)
  sn_jac_grid = zeros(length(b),n_max+1,2)
  sn_jac_array= zeros(n_max+1,2)
  sn_grid = zeros(length(b),n_max+1)
  sn_grid_big = zeros(length(b),n_max+1)
  sn_array= zeros(n_max+1)
  sn_jac_grid_num = zeros(length(b),n_max+1,2)
  for j=1:length(b)
#    println("r: ",r," b: ",b[j])
    sn_array,sn_jac_array= sn_jac(l_max,r,b[j])
    s_n!(l_max,r,b[j],sn_array)
    sn_grid[j,:]=sn_array
    sn_jac_grid[j,:,:]=sn_jac_array
    sn_array,sn_jac_array = sn_jac_num(l_max,r,b[j])
    sn_jac_grid_num[j,:,:]=sn_jac_array
    sn_grid_big[j,:]=sn_array
  end
# Now, make plots:
  ax = axes
  m=0;l=0
  for n=0:n_max
    if m==0
#      ax[:plot](abs.(sn_jac_grid[:,n+1,1]-sn_jac_grid_num[:,n+1,1]),color=cmap(l/(l_max+2)),lw=1,label=string("l=",l))
#      ax[:semilogy](abs.(asinh.(sn_jac_grid[:,n+1,1])-sn_jac_grid_num[:,n+1,1]),lw=1,label=string("l=",l))
#      ax[:semilogy](b,abs.(asinh.(sn_jac_grid[:,n+1,1])-asinh.(sn_jac_grid_num[:,n+1,1])),lw=1,label=string("l=",l))
#      ax[:semilogy](abs.(asinh.(sn_jac_grid[:,n+1,1])-asinh.(sn_jac_grid_num[:,n+1,1])),lw=1,label=string("l=",l))
      ax[:semilogy](abs.(asinh.(sn_jac_grid[:,n+1,1])-asinh.(sn_jac_grid_num[:,n+1,1])),lw=1,label=string("l=",l))
#      semilogy(abs.(asinh.(sn_jac_grid[:,n+1,1])-asinh.(sn_jac_grid_num[:,n+1,1])),lw=1,label=string("l=",l))
    else
#      ax[:plot](abs.(sn_jac_grid[:,n+1,1]-sn_jac_grid_num[:,n+1,1]),color=cmap(l/(l_max+2)),lw=1)
#      ax[:semilogy](abs.(sn_jac_grid[:,n+1,1]-sn_jac_grid_num[:,n+1,1]),lw=1)
      ax[:semilogy](abs.(asinh.(sn_jac_grid[:,n+1,1])-asinh.(sn_jac_grid_num[:,n+1,1])),lw=1)
#      semilogy(abs.(asinh.(sn_jac_grid[:,n+1,1])-asinh.(sn_jac_grid_num[:,n+1,1])),lw=1)
#      ax[:semilogy](b,abs.(asinh.(sn_jac_grid[:,n+1,1])-asinh.(sn_jac_grid_num[:,n+1,1])),lw=1)
    end
#    ax[:plot](abs.(sn_jac_grid[:,n+1,2]-sn_jac_grid_num[:,n+1,2]),color=cmap(l/(l_max+2)),lw=1)
#    ax[:semilogy](abs.(sn_jac_grid[:,n+1,2]-sn_jac_grid_num[:,n+1,2]),lw=1)
#    ax[:semilogy](b,abs.(asinh.(sn_jac_grid[:,n+1,2])-asinh.(sn_jac_grid_num[:,n+1,2])),lw=1)
    ax[:semilogy](abs.(asinh.(sn_jac_grid[:,n+1,2])-asinh.(sn_jac_grid_num[:,n+1,2])),lw=1)
    ax[:semilogy](abs.(asinh.(sn_grid[:,n+1])-asinh.(sn_grid_big[:,n+1])),lw=1)
#    semilogy(abs.(asinh.(sn_jac_grid[:,n+1,2])-asinh.(sn_jac_grid_num[:,n+1,2])),lw=1)
    println("n: ",n," m: ",m," l: ",l," mu: ",l-m," nu: ",l+m," max dsn/dr: ",maximum(abs.(asinh.(sn_jac_grid[:,n+1,1])-asinh.(sn_jac_grid_num[:,n+1,1]))),
      " maximum dsn/db: ",maximum(abs.(asinh.(sn_jac_grid[:,n+1,2])-asinh.(sn_jac_grid_num[:,n+1,2]))))
    m +=1
    if m > l
      l += 1
      m = -l
    end
  end
  ax[:legend](loc="upper right",fontsize=6)
  ax[:set_xlabel]("b values")
  ax[:set_ylabel]("Derivative Error")
  ax[:axis]([0,length(b),1e-16,1])
  ax[:set_xticks](nb*linearspace(0,nticks-1,nticks))
  ax[:set_xticklabels](xticknames,rotation=45)
  if i==1
    ax[:set_title]("r = 0.01")
  else
    ax[:set_title]("r = 100")
  end
#  ax[:axis]([minimum(b),maximum(b),1e-16,1])
#  ax[:axis]([0,length(b),1e-16,1])
#  read(STDIN,Char)
  clf()
  for n=0:n_max
      plot(b,sn_grid[:,n+1])
  end
#  read(STDIN,Char)
## Loop over n and see where the differences between the finite-difference
## and AutoDiff are greater than the derivative value:
l=0; m=0
for n=0:n_max
  diff = sn_jac_grid[:,n+1,:]-sn_jac_grid_num[:,n+1,:]
  for ib=1:length(b)
    test1 = isapprox(sn_jac_grid[ib,n+1,1],sn_jac_grid_num[ib,n+1,1],atol=1e-8)
    @test test1
    if ~test1
      println("b: ",b[ib]," n: ",n," dsn/dr: ",sn_jac_grid[ib,n+1,1]," dsn/dr_num: ",sn_jac_grid_num[ib,n+1,1]," diff: ",sn_jac_grid[ib,n+1,1]-sn_jac_grid_num[ib,n+1,1])
    end
    test2 = isapprox(sn_jac_grid[ib,n+1,2],sn_jac_grid_num[ib,n+1,2],atol=1e-8)
    @test test2
    if ~test2
      println("b: ",b[ib]," n: ",n," dsn/db: ",sn_jac_grid[ib,n+1,2]," dsn/db_num: ",sn_jac_grid_num[ib,n+1,2]," diff: ",sn_jac_grid[ib,n+1,2]-sn_jac_grid_num[ib,n+1,2])
    end
  end
  mask = (abs.(diff) .> 1e-3*abs.(sn_jac_grid[:,n+1,:])) .& (abs.(sn_jac_grid[:,n+1,:]) .> 1e-5)
  if sum(mask) > 0 || mod(l-m,4) == 0
    println("n: ",n," max dsn/dr: ",maximum(abs.(diff)))
    println("b: ",b[mask[:,1]],b[mask[:,2]]," k^2_H: ",k2H[mask[:,1]],k2H[mask[:,2]])
# end
#end
# Now, plot derivatives for each value of n:
#for n=0:n_max
    clf()
    plot(b,sn_grid[:,n+1])
    plot(b,sn_jac_grid[:,n+1,1])
    plot(b[mask[:,1]],sn_jac_grid[mask[:,1],n+1,1],"o")
    plot(b,sn_jac_grid[:,n+1,2])
    plot(b[mask[:,2]],sn_jac_grid[mask[:,2],n+1,2],"o")
    plot(b,sn_jac_grid_num[:,n+1,1],linestyle="--",linewidth=2)
    plot(b,sn_jac_grid_num[:,n+1,2],linestyle="--",linewidth=2)
    println("n: ",n," l: ",l," m: ",m," mu: ",l-m," nu: ",l+m)
#    read(STDIN,Char)
  end
  m +=1
  if m > l
    l += 1
    m = -l
  end
end
#
#
##end
#
##loglog(abs.(reshape(sn_jac_grid,length(b)*(n_max+1)*2)),abs.(reshape(sn_jac_grid-sn_jac_grid_num,length(b)*(n_max+1)*2)),".")
##loglog(abs.(reshape(sn_jac_grid,length(b)*(n_max+1)*2)),abs.(reshape(sn_jac_grid_num,length(b)*(n_max+1)*2)),".")
end
return
end

test_sn_jacobian2(10)
end


using Distributions
type SmallFiniteHorizonMDP
    P::Array{Float64, 4}
    atbls::Array{Distributions.AliasTable,3}
    R::Array{Float64, 3}
    p0::Categorical
    s::UInt64
    t::UInt64
end
function SmallFiniteHorizonMDP(P, R, p0, s, t)
    S, A, H = size(R)
    atbls = Array{Distributions.AliasTable, 3}(S, A, H)
    for t=1:H, a=1:A, s=1:S
        atbls[s,a,t] = Distributions.AliasTable(P[:, s, a, t])
    end
    SmallFiniteHorizonMDP(P, atbls, R, p0, s, t)
end

horizon(mdp::SmallFiniteHorizonMDP) = size(mdp.R, 3)
nS(mdp::SmallFiniteHorizonMDP) = size(mdp.R, 1)
nA(mdp::SmallFiniteHorizonMDP) = size(mdp.R, 2)
rewards(mdp::SmallFiniteHorizonMDP) = mdp.R
function restart!(mdp::SmallFiniteHorizonMDP)
    s = rand(mdp.p0)
    mdp.s = s
    mdp.t = 1
    s
end
function step!(mdp::SmallFiniteHorizonMDP, a)
    r = mdp.R[mdp.s, a, mdp.t]
    sn = rand(mdp.atbls[mdp.s, a, mdp.t])
    mdp.s = sn
    mdp.t += 1
    stop = mdp.t > horizon(mdp)
    sn, r, stop
end

function expected_return(mdp::SmallFiniteHorizonMDP, policy)
    H = horizon(mdp)
    S = nS(mdp)
    A = nA(mdp)

    Vold = zeros(S)
    Vnew = zeros(S)

    for t=H:-1:1
        for s=1:S
            a = policy(s, t)
            Vnew[s] = mdp.R[s, a, t] + mdp.P[:, s, a, t] ⋅ Vold
        end
        Vnew, Vold = Vold, Vnew
    end
    mdp.p0.p ⋅ Vold
end

function optimal_policy(mdp::SmallFiniteHorizonMDP)
    H = horizon(mdp)
    S = nS(mdp)
    A = nA(mdp)
    policy = zeros(UInt64, S, H)
    Vold = zeros(S)
    Qnew = zeros(S, A)

    for t=H:-1:1
        Qnew = mdp.R[:, :, t] + squeeze(sum(Vold .* mdp.P[:, :, :, t], 1), 1)
        for s=1:S
            policy[s, t] = indmax(Qnew[s, :])
            Vold[s] = Qnew[s, policy[s, t]]
        end
    end
    (s, t) -> policy[s, t]
end

function value_fun(mdp::SmallFiniteHorizonMDP, policy)
    H = horizon(mdp)
    S = nS(mdp)
    A = nA(mdp)

    V = zeros(S, H+1)

    for t=H:-1:1
        for s=1:S
            a = policy(s, t)
            V[s, t] = mdp.R[s, a, t] + mdp.P[:, s, a, t] ⋅ V[:, t+1]
        end
    end
    V
end

gammarand(a) = if a >= 1
            rand(Distributions.GammaGDSampler(Gamma(a)))
        else
            rand(Distributions.GammaGSSampler(Gamma(a)))
        end

function dirrand!(res, αs)
    ss = 0.
    for (i, α) in enumerate(αs)
        b = gammarand(α)
        ss += b
        res[i] = b
    end
    res ./= ss
end
function dirrand(αs)
    res = zeros(length(αs))
    dirrand!(res, αs)
    res
end
function randomMDP(S, A, H, α =.1)
    dir = Dirichlet(S, α)
    P = zeros(S, S, A, H)
    R = rand(S, A, H)
    R[rand(S, A, H) .<= 0.85] = 0.
    p0 = Categorical(dirrand(repeated(α, S)))
    for s=1:S, a=1:A, t=1:H
        P[:, s, a, t] = dirrand(repeated(α, S))
    end
    SmallFiniteHorizonMDP(P, R, p0, 1, 1)
end
function stateDepTestMDP(N, ϵ=1e-2)
    S = 2N + 2
    A = 2
    H = 2
    pp0 = zeros(S)
    pp0[1] = 1
    p0 = Categorical(pp0)
    R = zeros(S, A, H)
    R[(S - N + 1):end, :, 2] = 1
    R[2, :, 2] = 0.5 + ϵ
    P = zeros(S, S, A, H)
    for s=1:S, a=1:A
        P[s, s, a, 2] = 1.
    end
    P[3:end, 1, 1, 1] = 1 / 2 / N
    P[2, 1, 2, 1] = 1.
    SmallFiniteHorizonMDP(P, R, p0, 1, 1)
end

type ChainMDP
    P::Array{Float64, 3}
    R::Vector{Float64}
    N::Int64
    s::UInt64
    t::UInt64
    r0::Float64
end
nS(m::ChainMDP) = m.N
nA(m::ChainMDP) = 2
horizon(m::ChainMDP) = m.N
function rewards(m::ChainMDP)
    R = zeros(nS(m), nA(m), horizon(m))
    R .= m.R
    R
end

function ChainMDP(N, r0=0.)
    P = zeros(N, N, 2)
    for s=1:N
        sn = max(s - 1, 1)
        P[sn, s, 1] = 1.
        P[sn, s, 2] = 1./N
        sn = min(s+1, N)
        P[sn, s, 2] = 1 - 1./N
    end

    R = zeros(N)
    R[end] = 1.
    R[1] = r0
    ChainMDP(P, R, N, 1, 1, r0)
end
function restart!(mdp::ChainMDP)
    mdp.s = 1
    mdp.t = 1
    1
end
function step!(mdp::ChainMDP, a)
    r = 0.
    if mdp.s == mdp.N
        r = 1.
    elseif mdp.s == 1
        r = mdp.r0
    end
    sn = max(mdp.s-1, 1)
    if a == 2 && rand() > 1/mdp.N
        # go right
        sn = min(mdp.s+1,mdp.N)
    end
    mdp.s = sn
    mdp.t += 1
    stop = mdp.t > horizon(mdp)
    sn, r, stop
end

function expected_return(mdp::ChainMDP, policy)
    H = horizon(mdp)
    S = nS(mdp)
    A = nA(mdp)

    Vold = zeros(S)
    Vnew = zeros(S)

    for t=H:-1:1
        for s=1:S
            a = policy(s, t)
            Vnew[s] = mdp.R[s] + mdp.P[:, s, a] ⋅ Vold
        end
        Vnew, Vold = Vold, Vnew
    end
    Vold[1]
end

function optimal_policy(mdp::ChainMDP)
    H = horizon(mdp)
    S = nS(mdp)
    A = nA(mdp)
    policy = zeros(UInt64, S, H)
    Vold = zeros(S)
    Qnew = zeros(S, A)

    for t=H:-1:1
        Qnew = reshape(mdp.R, S, 1) .+ squeeze(sum(Vold .* mdp.P, 1), 1)
        for s=1:S
            policy[s, t] = indmax(Qnew[s, :])
            Vold[s] = Qnew[s, policy[s, t]]
        end
    end
    (s, t) -> policy[s, t]
end


function duration(vs,fs,d) #That function cut samples length of any sound according to their duration.
    if(d<=size(vs)[1]/fs-1/fs)
    k=int64(d*fs+1)
    Vs=zeros(k,size(vs)[2])
    for j=1:size(vs)[2]
        for i=1:size(vs)[1]
            if(i<k+1)
        Vs[i,j]=vs[i,j]
            end
        end
    end
    return Vs
    else
        error("Tiempo excedido")
    end
end
#Inicio del programa

dur=5; # duration of sound signal
using WAV #That package let me read any sound signal and convert that signal in samples
using PyPlot #Plot Pakage
using DSP #function for Digital Signal Processing
close("all")

archivo1="C:\\JULIA\\Grab6.wav"
vs1, fs1 = wavread(archivo1)
Vs1=duration(vs1,fs1,dur)
Vs1
length((Vs1))
archivo2="C:\\JULIA\\Grab7.wav"
vs2, fs2 = wavread(archivo2)
Vs2=duration(vs2,fs2,dur)
Vs2
archivo3="C:\\JULIA\\Grab8.wav"
vs3, fs3 = wavread(archivo3)
Vs3=duration(vs3,fs3,dur)
Vs3

dt=1/fs1;
#te=[dt:dt:dt*(size(Vs1)[1])]
t=linspace(dt,dt*size(Vs1)[1],size(Vs1)[1])
plot(t,Vs1[:,1])
##plot(te,Vs1[:,1])
ruido1=0.006/20*rand(size(t))
#Señales con ruido
x1=Vs1[:,1]+ruido1
x2=Vs2[:,1]+ruido1
x3=Vs3[:,1]+ruido1

f=linspace(-1*fs1/2,1*fs1/2,length(t))
figure(1)
espectro1=abs(fftshift(fft(x1)/length(t)))
espectro2=abs(fftshift(fft(x2)/length(t)))
espectro3=abs(fftshift(fft(x3)/length(t)))
subplot(3,2,1)
plot(f,espectro1)
title("Primera señal")
subplot(3,2,3)
plot(f,espectro2)
title("Segunda señal")
subplot(3,2,5)
plot(f,espectro3)
title("Tercera señal")
#filtrado de señal con ruido

#Filtracion del ruido

    #Primera señal
response1=Lowpass(3000,fs=fs1)
orden=Butterworth(10);
filtrada1=filt(digitalfilter(response1,orden),x1)
FILTRADA1=abs(fftshift(fft(filtrada1)/length(t)))
subplot(3,2,2)
plot(f,FILTRADA1)
title("Primera señal filtrada 3Khz")

    #Segunda señal
response2=Lowpass(3000,fs=fs2)
orden=Butterworth(10);
filtrada2=filt(digitalfilter(response2,orden),x2)
FILTRADA2=abs(fftshift(fft(filtrada2)/length(t)))
subplot(3,2,4)
plot(f,FILTRADA2)
title("Segunda señal filtrada 3Khz")

    #Tercera señal
response3=Lowpass(4000,fs=fs3)
orden=Butterworth(10);
filtrada3=filt(digitalfilter(response3,orden),x3)
FILTRADA3=abs(fftshift(fft(filtrada3)/length(t)))
subplot(3,2,6)
plot(f,FILTRADA3)
title("Tercera señal filtrada 4Khz")

    #ModulacionAM Doble Banda lateral
figure(2)
title("ModulacionAM")
dbl2=filtrada2.*cos(2*pi*3500*t)
DBL2=abs(fftshift(fft(dbl2)/length(t)))
subplot(2,2,1)
plot(f,DBL2)
title("Modulacion 3.5Khz segunda señal (DBL)")
response22=Bandpass(3500,6500,fs=fs2)
blu2=filt(digitalfilter(response22,orden),dbl2)
BLU2=abs(fftshift(fft(blu2)/length(t)))
subplot(2,2,2)
plot(f,BLU2)
title("BLU segunda señal")
dbl3=filtrada3.*cos(2*pi*7000*t)
DBL3=abs(fftshift(fft(dbl3)/length(t)))
subplot(2,2,3)
plot(f,DBL3)
title("Modulacion 7Khz tercera señal (DBL)")
response33=Bandpass(7000,11000,fs=fs3)
blu3=filt(digitalfilter(response33,orden),dbl3)
BLU3=abs(fftshift(fft(blu3)/length(t)))
subplot(2,2,4)
plot(f,BLU3)
title("BLU tercera señal")
figure(3)
plot(f,FILTRADA1)
plot(f,BLU2)
plot(f,BLU3)
title("Señal FDM")
    #Señal FDM con ruido
suma=filtrada1+blu2+blu3+ruido1
FDM=abs(fftshift(fft(suma)/length(t)))
figure(4)
subplot(3,1,1)
plot(f,FDM)
title("Señal FDM con ruido")
    #Filtracion de señal FDM
response_FDM=Lowpass(11000,fs=fs1)
orden_FDM=Butterworth(15)
fil_FDM=filt(digitalfilter(response_FDM,orden_FDM),suma)
FIL_FDM=abs(fftshift(fft(fil_FDM)/length(t)))
subplot(3,1,2)
plot(f,FIL_FDM)
title("Señal FDM filtrada")
    #ModulacionFDM
mod_FDM=fil_FDM.*cos(2*pi*11100*t)
MOD_FDM=abs(fftshift(fft(mod_FDM)/length(t)))
subplot(3,1,3)
plot(f,MOD_FDM)
title("FDM modulada a 11.2Khz")


using JSON, LinearAlgebra, JLD2

# read the data from json
tempCMEParams = Dict()
open("src/CMEParamsData/iltcme.json", "r") do f #iltcme.json is from Horvath, Telek, see CMEParams_metadata, or fullData.zip if you have it
    global tempCMEParams
    tempCMEParams=JSON.parse(f)  # parse and transform data
end

# put the parameters in CMEParams with keys corresponding the ME order
CMEParams = Dict()
for n in keys(tempCMEParams)
    if !in(2*tempCMEParams[n]["n"]+1,keys(CMEParams)) # if no already in dict add it
        CMEParams[2*tempCMEParams[n]["n"]+1] = tempCMEParams[n]
    elseif tempCMEParams[n]["cv2"]<CMEParams[2*tempCMEParams[n]["n"]+1]["cv2"]
        # if its already in there, add only if it has the smallest CV
        CMEParams[2*tempCMEParams[n]["n"]+1] = tempCMEParams[n]
    end
end

# a function to numerically approximate D
function integrateD(evals,params)
    # evals is an integer specifying how many points to eval the function at
    # params is a CMEParams dictionary entry, i.e. CMEParams[3]
    N = 2*params["n"]+1 # ME order

    α = zeros(N)
    α[1] = params["c"]
    a = params["a"]
    b = params["b"]
    ω =  params["omega"]
    for k in 1:params["n"]
        kω = k*ω
        α[2*k] = (1/2)*( a[k]*(1+kω) - b[k]*(1-kω) )/(1+kω^2)
        α[2*k+1] = (1/2)*( a[k]*(1-kω) + b[k]*(1+kω) )/(1+kω^2)
    end

    period = 2*π/ω # the orbit repeats after this time
    edges = range(0,period,length=evals+1) # points at which to evaluate the fn
    h = period/(evals)

    orbit_LHS = α
    orbit_RHS = zeros(N)
    v_RHS = zeros(N)
    v_RHS[1] = 1
    v_LHS = ones(N)
    D = zeros(N,N)
    for t in edges[2:end]
        orbit_RHS[1] = α[1]
        for k in 1:params["n"]
            kωt = k*ω*t
            idx = 2*k
            idx2 = idx+1
            temp_cos = cos(kωt)
            temp_sin = sin(kωt)
            orbit_RHS[idx] = α[idx]*temp_cos + α[idx2]*temp_sin
            orbit_RHS[idx2] = -α[idx]*temp_sin + α[idx2]*temp_cos
            v_RHS[idx] = temp_cos - temp_sin
            v_RHS[idx2] = temp_sin + temp_cos
        end
        orbit_RHS = orbit_RHS./sum(orbit_RHS)
        orbit = (orbit_LHS+orbit_RHS)./2

        v = exp(-(t-h))*(v_LHS - exp(-h)*v_RHS)

        Dᵢ = v*orbit'
        D += Dᵢ

        orbit_LHS = copy(orbit_RHS)
        v_LHS = copy(v_RHS)
    end
    D = (1/(1-exp(-period)))*D
    return D
end

# numerical approximation of D
k = 100_000_000
T=[]
N=[]
for n in keys(CMEParams)
    evals = Int(ceil(k/(n/3)))# number of function evals
    display(n)
    display(evals)
    CMEParams[n]["D"], t = @timed integrateD(evals,CMEParams[n])
    CMEParams[n]["intDevals"] = evals
    display(t)
    push!(N,n)
    push!(T,t)
end

scatter(N,T)

# add order 1 (exponential) to params
CMEParams[1] = Dict(
  "n"       => 0,
  "c"       => 1,
  "b"       => Any[],
  "mu2"     => 1,
  "a"       => Any[],
  "omega"   => 1,
  "phi"     => [],
  "mu1"     => 1,
  "D"       => [1.0],
  "cv2"     => 1,
  "optim"   => "full",
  "lognorm" => [],
)

open("dev/CMEParams.json","w") do f
    JSON.print(f, CMEParams)
end

@save "dev/CMEParams.jld2" CMEParams

let
    CMEKeys = sort(collect(keys(CMEParams)))
    a = 0
    filecounter = 1
    tempDict = Dict()
    for key in CMEKeys
        a += key^2*8
        tempDict[key] = CMEParams[key]
        if a > 2e7
            open(pwd()*"/dev/CMEParamsData/CMEParams"*string(filecounter)*".json","w") do f
                JSON.print(f, tempDict)
            end
            @save pwd()*"/dev/CMEParamsData/CMEParams"*string(filecounter)*".jld2" tempDict
            tempDict = Dict()
            a = 0
            filecounter += 1
        end
    end
end