text:
module Kalman

using HDF5
using LinearAlgebra
using Logging
using Statistics

import Base:
    write, read

export AR1KalmanFilter, reset!, advance!, observe!, generate, whiten, log_likelihood

mutable struct AR1KalmanFilter
    ypred::Float64
    vypred::Float64
    mu::Float64
    var::Float64
    tau::Float64
end

struct AR1KalmanPosterior
    ts::Array{Float64, 1}
    ys::Array{Float64, 1}
    dys::Array{Float64, 1}
    mu::Float64
    sigma::Float64
    dtmin::Float64
    T::Float64
end

function AR1KalmanPosterior(ts, ys, dys)
    AR1KalmanPosterior(ts, ys, dys, mean(ys), std(ys), minimum(diff(ts)), ts[end]-ts[1])
end

struct AR1KalmanParams
    mu::Float64
    sigma::Float64
    nu::Float64
    tau::Float64
end

function to_params(post::AR1KalmanPosterior, x::Array{Float64, 1})
    AR1KalmanParams(x[1], x[2], x[3], x[4])
end

function to_array(post::AR1KalmanPosterior, p::AR1KalmanParams)
    Float64[p.mu, p.sigma, p.nu, p.tau]
end

function log_prior(post::AR1KalmanPosterior, x::Array{Float64, 1})
    log_prior(post, to_params(post, x))
end

function rand_between(a, b)
    a + (b-a)*rand()
end

function rand_between(a, b, n)
    a .+ (b-a).*rand(n)
end

function init(post::AR1KalmanPosterior, n)
    pts = zeros(4,n)

    mu = post.mu
    sig = post.sigma

    dtmin = post.dtmin
    T = post.T

    pts[1,:] = rand_between((mu-10*sig), (mu+10*sig), n)
    pts[2,:] = exp.(rand_between(log(sig/10.0),log(sig*10.0), n))
    pts[3,:] = exp.(rand_between(log(0.1), log(10.0), n))
    pts[4,:] = exp.(rand_between(log(0.1*dtmin),log(10.0*T), n))

    pts
end

function log_prior(post::AR1KalmanPosterior, p::AR1KalmanParams)
    mu = post.mu
    sig = post.sigma

    dtmin = post.dtmin
    T = post.T

    if p.mu < mu - 10*sig || p.mu > mu + 10*sig
        return -Inf
    end

    if p.sigma < sig/10.0 || p.sigma > sig*10.0
        return -Inf
    end

    if p.nu < 0.1 || p.nu > 10.0
        return -Inf
    end

    if p.tau < dtmin/10.0 || p.tau > 10.0*T
        return -Inf
    end

    # Flat in log(sigma), log(tau), log(nu); Flat in mu
    -log(p.sigma) - log(p.tau) - log(p.nu)
end

function log_likelihood(post::AR1KalmanPosterior, x::Array{Float64,1})
    log_likelihood(post, to_params(post, x))
end

function log_likelihood(post::AR1KalmanPosterior, p::AR1KalmanParams)
    filt = make_filter(post, p)

    log_likelihood(filt, post.ts, post.ys, p.nu*post.dys)
end

function make_filter(post, x::Array{Float64, 1})
    make_filter(post, to_params(post, x))
end

function make_filter(post::AR1KalmanPosterior, p::AR1KalmanParams)
    AR1KalmanFilter(p.mu, p.sigma*p.sigma, p.tau)
end

""" Produce a Kalman filter for an AR(1) process with the given mean,
variance, and timescale. """
function AR1KalmanFilter(mu, var, tau)
    AR1KalmanFilter(mu, var, mu, var, tau)
end

""" Reset the given filter to its initial state.  """
function reset!(filt::AR1KalmanFilter)
    filt.ypred = filt.mu
    filt.vypred = filt.var
end

""" Advance the given filter by a step `dt`.  """
function advance!(filt::AR1KalmanFilter, dt::Float64)
    lfac = exp(-dt/filt.tau)

    filt.ypred = filt.mu + (filt.ypred - filt.mu)*lfac
    filt.vypred = lfac*(filt.vypred - filt.var)*lfac + filt.var
end

""" Incorporate the given observation with measurement uncertainty (1
s.d.) into the filter estimate. """
function observe!(filt::AR1KalmanFilter, y::Float64, dy::Float64)
    vy = dy*dy + filt.vypred

    gain = filt.vypred / vy

    filt.ypred = filt.ypred + (y - filt.ypred)*gain
    filt.vypred = filt.vypred - gain*gain*vy
end

function predict(filt::AR1KalmanFilter)
    filt.ypred, filt.vypred
end

""" Generate the GP represented by the given filter, sampled at the
given times with the given observational uncertainties.  """
function generate(filt::AR1KalmanFilter, ts::Array{Float64,1}, dys::Array{Float64,1})
    reset!(filt)

    n = size(ts, 1)
    ys = zeros(n)

    ytrue = sqrt(filt.var)*randn() + filt.mu
    ys[1] = ytrue + dys[1]*randn()

    for i in 2:n
        observe!(filt, ytrue, 0.0)
        advance!(filt, ts[i]-ts[i-1])
        ytrue = filt.ypred + sqrt(filt.vypred)*randn()
        ys[i] = ytrue + dys[i]*randn()
    end

    ys
end

""" Whiten the observations assuming they are drawn from the GP
represented by the given filter.  The outputs will be independent
N(0,1) distributed. """
function whiten(filt::AR1KalmanFilter, ts::Array{Float64,1}, ys::Array{Float64, 1}, dys::Array{Float64,1})
    n = size(ts, 1)

    xs = zeros(n)

    reset!(filt)

    xs[1] = (ys[1] - filt.ypred)/sqrt(filt.vypred + dys[1]*dys[1])
    for i in 2:n
        observe!(filt, ys[i-1], dys[i-1])
        advance!(filt, ts[i]-ts[i-1])
        xs[i] = (ys[i] - filt.ypred) / sqrt(filt.vypred + dys[i]*dys[i])
    end

    xs
end

square(x) = x*x

function log_likelihood_term(filt, y, dy)
    var = filt.vypred + dy*dy
    -0.5*log(2.0*pi) - 0.5*log(var) - 0.5*square(y-filt.ypred)/var
end

""" The log-likelihood function for the given data assuming that it
comes from the GP represented by the filter. """
function log_likelihood(filt::AR1KalmanFilter, ts::Array{Float64,1}, ys::Array{Float64,1}, dys::Array{Float64,1})
    n = size(ts, 1)

    ll = zero(ys[1])

    reset!(filt)

    ll += log_likelihood_term(filt, ys[1], dys[1])
    for i in 2:n
        observe!(filt, ys[i-1], dys[i-1])
        advance!(filt, ts[i]-ts[i-1])
        ll += log_likelihood_term(filt, ys[i], dys[i])
    end

    ll
end

function psd(filt::AR1KalmanFilter, fs::Array{Float64, 1})
    4.0*filt.tau ./ (1.0 + (2.0*pi*filt.tau*fs).^2)
end

mutable struct CARMAKalmanFilter
    mu::Float64
    x::Array{ComplexF64,1}
    vx::Array{ComplexF64, 2}
    K::Array{ComplexF64, 1}
    lambda::Array{ComplexF64, 1}
    vxtemp::Array{ComplexF64, 2}
    v::Array{ComplexF64, 2}
    arroots::Array{ComplexF64, 1}
    b::Array{ComplexF64, 2}
    tscale::Float64
end

function reset!(filt::CARMAKalmanFilter)
    p = size(filt.x, 1)

    filt.x = zeros(ComplexF64, p)
    filt.vx = copy(filt.v)
    filt.K = zeros(ComplexF64, p)

    filt
end

"""Construct a polynomial from the given roots.  Returns an array of
coefficients, `c`, represting the polynomial as `p(x) =
sum(c[i]*x^(i-1))`.

"""
function poly(roots::Array{T,1}) where {T <: Number}
    n = size(roots, 1) + 1

    if n == 1
        return T[one(T)]
    else
        poly = zeros(T, n)

        poly[2] = one(T)
        poly[1] = -roots[1]

        for i in 2:size(roots,1)
            r = roots[i]
            for j in n:-1:2
                poly[j] = poly[j-1] - r*poly[j]
            end
            poly[1] = -poly[1]*r
        end

        poly
    end
end

function polyeval(roots::Array{T, 1}, x::T) where {T <: Number}
    p = one(x)

    for i in 1:size(roots,1)
        p = p*(x - roots[i])
    end

    p
end

function CARMAKalmanFilter(mu::Float64, sigma::Float64, arroots::Array{ComplexF64, 1}, maroots::Array{ComplexF64, 1})
    p = size(arroots, 1)
    q = size(maroots, 1)

    @assert q < p "q must be less than p: q = $q, p = $p"
    @assert all(real(arroots) .< 0.0) "AR roots must have negative real part: $arroots"
    @assert all(real(maroots) .< 0.0) "MA roots must have negative real part: $maroots"

    tscale = exp(-mean(log.(abs.(arroots)))) # Try to stabilise the variance computation

    arroots = arroots.*tscale
    maroots = maroots.*tscale

    beta = poly(maroots)
    beta /= beta[1]
    b = cat(beta, zeros(p-q-1), dims=1)
    b = b'

    U = zeros(ComplexF64, (p,p))
    for j in 1:p
        for i in 1:p
            U[i,j] = arroots[j]^(i-1)
        end
    end

    # Rotated observation vector
    b = b*U

    e = zeros(ComplexF64, p)
    e[end] = one(ComplexF64)

    J = U \ e

    V = zeros(ComplexF64, (p,p))
    for j in 1:p
        for i in 1:p
            V[i,j] = -J[i]*conj(J[j])/(arroots[i] + conj(arroots[j]))
        end
    end

    s2 = sigma*sigma/(b*V*b')[1]
    V = V*s2

    CARMAKalmanFilter(mu, zeros(ComplexF64, p), V, zeros(ComplexF64, p), zeros(ComplexF64, p), zeros(ComplexF64, (p,p)), copy(V), copy(arroots), b, tscale)
end

@inbounds function advance!(filt::CARMAKalmanFilter, dt::Float64)
    p = size(filt.x, 1)

    for i in 1:p
        x = filt.arroots[i]*dt/filt.tscale
        filt.lambda[i] = exp(x)
    end
    lam = filt.lambda

    for i in 1:p
        filt.x[i] = lam[i]*filt.x[i]
    end

    for j in 1:p
        for i in 1:p
            filt.vx[i,j] = lam[i]*conj(lam[j])*(filt.vx[i,j] - filt.v[i,j]) + filt.v[i,j]
        end
    end

    filt
end

@inbounds function observe!(filt::CARMAKalmanFilter, y::Float64, dy::Float64)
    p = size(filt.x, 1)

    ey, vy = predict(filt)
    vy += dy*dy

    for i in 1:p
        filt.K[i] = zero(filt.K[i])
        for j in 1:p
            filt.K[i] += filt.vx[i,j]*conj(filt.b[j])/vy
        end
    end

    for i in 1:p
        filt.x[i] = filt.x[i] + (y - ey)*filt.K[i]
    end

    for j in 1:p
        for i in 1:p
            filt.vx[i,j] = filt.vx[i,j] - vy*filt.K[i]*conj(filt.K[j])
        end
    end

    filt
end

@inbounds function predict(filt::CARMAKalmanFilter)
    p = size(filt.x,1)

    yp = filt.mu
    for i in 1:p
        yp += real(filt.b[1,i]*filt.x[i])
    end

    vyp = 0.0
    for i in 1:p
        for j in 1:p
            vyp += real(filt.b[1,i]*filt.vx[i,j]*conj(filt.b[1,j]))
        end
    end

    yp, vyp
end

function whiten(filt::CARMAKalmanFilter, ts, ys, dys)
    n = size(ts, 1)

    reset!(filt)
    zs = zeros(n)

    for i in 1:n
        yp, vyp = predict(filt)

        zs[i] = (ys[i] - yp)/sqrt(vyp + dys[i]*dys[i])

        observe!(filt, ys[i], dys[i])

        if i < n
            advance!(filt, ts[i+1]-ts[i])
        end
    end

    zs
end

function draw_and_collapse!(filt::CARMAKalmanFilter)
    nd = size(filt.x, 1)
    try
        for i in 1:nd
            filt.vx[i,i] = real(filt.vx[i,i]) # Fix a roundoff error problem?
        end
        L = cholesky(Hermitian(filt.vx)).L
        filt.x = filt.x + L*randn(nd)
        filt.vx = zeros(ComplexF64, (nd, nd))
    catch e
        if isa(e, PosDefException)
            F = eigfact(filt.vx)
            for i in eachindex(F[:values])
                l = real(F[:values][i])
                v = F[:vectors][:,i]

                if l < 0.0
                    l = 0.0
                end

                filt.x = filt.x + sqrt(l)*randn().*v
            end
            filt.vx = zeros(ComplexF64, (nd, nd))
        else
            rethrow()
        end
    end
end

function generate(filt::CARMAKalmanFilter, ts, dys)
    n = size(ts, 1)
    nd = size(filt.x, 1)

    ys = zeros(n)

    reset!(filt)

    for i in 1:n
        # Draw a new state
        draw_and_collapse!(filt)

        y, _ = predict(filt)

        ys[i] = y + dys[i]*randn()

        if i < n
            advance!(filt, ts[i+1]-ts[i])
        end
    end

    ys
end

function log_likelihood(filt, ts, ys, dys)
    n = size(ts, 1)

    ll = -0.5*n*log(2.0*pi)

    reset!(filt)
    for i in 1:n
        yp, vyp = predict(filt)

        if vyp < 0.0
            warn("Kalman filter has gone unstable!")
            return -Inf
        end

        dy = ys[i] - yp
        vy = vyp + dys[i]*dys[i]

        ll -= 0.5*log(vy)
        ll -= 0.5*dy*dy/vy

        observe!(filt, ys[i], dys[i])

        if i < n
            advance!(filt, ts[i+1]-ts[i])
        end
    end

    ll
end

function carmacovariance(ts, sigma, arroots, maroots)
    n = size(ts, 1)
    p = size(arroots, 1)
    q = size(maroots, 1)

    @assert q < p

    beta = poly(maroots)
    beta /= beta[1]

    cm = zeros((n,n))

    for i in 1:n
        for j in i:n
            for k in 1:p
                r = arroots[k]
                br = zero(beta[1])
                bmr = zero(beta[1])
                rprod = one(arroots[1])
                for l in 1:q+1
                    br += beta[l]*r^(l-1)
                    bmr += beta[l]*(-r)^(l-1)
                end
                for l in 1:p
                    if l != k
                        rprod *= (arroots[l] - r)*(conj(arroots[l]) + r)
                    end
                end

                cm[i,j] += real(br*bmr*exp(r*abs(ts[j]-ts[i]))/(-2.0*real(r)*rprod))
            end
            cm[j,i] = cm[i,j]
        end
    end

    sfact = sigma*sigma/cm[1,1]
    for i in 1:n
        for j in 1:n
            cm[i,j] *= sfact
        end
    end

    cm
end

function carmagenerate(ts, dys, mu, sigma, arroots, maroots)
    n = size(ts, 1)

    cm = carmacovariance(ts, sigma, arroots, maroots)
    for i in 1:n
        cm[i,i] += dys[i]*dys[i]
    end

    L = chol(cm, Val{:L})

    ys = mu + L*randn(n)

    ys
end

function raw_carma_log_likelihood(ts, ys, dys, mu, sigma, arroots, maroots)
    n = size(ts, 1)

    cm = carmacovariance(ts, sigma, arroots, maroots)
    for i in 1:n
        cm[i,i] += dys[i]*dys[i]
    end

    zs = ys - mu

    F = cholfact(cm)
    L = F[:L]

    logdet = 0.0
    for i in 1:n
        logdet += log(L[i,i])
    end

    -0.5*n*log(2*pi) - logdet - 0.5*dot(zs, F \ zs)
end

struct CARMAKalmanPosterior
    ts::Array{Float64, 1}
    ys::Array{Float64, 1}
    dys::Array{Float64,1}
    p::Int
    q::Int
end

function write(f::Union{HDF5File, HDF5Group}, p::CARMAKalmanPosterior)
    f["ts", "compress", 3, "shuffle", ()] = p.ts
    f["ys", "compress", 3, "shuffle", ()] = p.ys
    f["dys", "compress", 3, "shuffle", ()] = p.dys
    f["p"] = p.p
    f["q"] = p.q
end

function read(f::Union{HDF5File, HDF5Group}, ::Type{CARMAKalmanPosterior})
    CARMAKalmanPosterior(f)
end

function CARMAKalmanPosterior(f::Union{HDF5File, HDF5Group})
    CARMAKalmanPosterior(read(f, "ts"), read(f, "ys"), read(f, "dys"), read(f, "p"), read(f, "q"))
end

function nparams(post::CARMAKalmanPosterior)
    post.p + post.q + 3
end

struct CARMAPosteriorParams
    mu::Float64
    sigma::Float64
    nu::Float64
    arroots::Array{ComplexF64, 1}
    maroots::Array{ComplexF64, 1}
end

function to_roots(x::Array{Float64, 1})
    n = size(x, 1)
    r = zeros(ComplexF64, n)

    if n == 0
        r
    elseif n == 1
        r[1] = -exp(x[1])
        r
    else
        for i in 1:2:n-1
            logb = x[i]
            logc = x[i+1]

            b = exp(logb)
            c = exp(logc)

            d = b*b - 4*c
            if d < 0.0
                sd = sqrt(-d)
                r[i] = -0.5*(b + sd*1im)
                r[i+1] = -0.5*(b - sd*1im)
            else
                sd = sqrt(d)
                r[i] = -0.5*(b + sd)
                r[i+1] = -0.5*(b - sd)
            end
        end

        if n % 2 == 1
            r[end] = -exp(x[end])
        end

        r
    end
end

function to_rparams(x::Array{ComplexF64, 1})
    n = size(x, 1)
    rp = zeros(n)

    if n == 0
        rp
    elseif n == 1
        rp = log(-real(x[1]))
    else
        for i in 1:2:n-1
            r1 = x[i]
            r2 = x[i+1]

            b = real(-(r1 + r2))
            c = real(r1*r2)

            rp[i] = log(b)
            rp[i+1] = log(c)
        end

        if n % 2 == 1
            rp[end] = log(-real(x[end]))
        end

        sort_root_params(rp)
    end
end

function root_params_ordered(rp::Array{Float64, 1})
    n = size(rp, 1)

    if n == 0
        true
    elseif n == 1
        true
    else
        x = -Inf
        for i in 1:2:n-1
            if rp[i] > x
                x = rp[i]
            else
                return false
            end
        end
        true
    end
end

function sort_root_params(rp::Array{Float64, 1})
    n = size(rp, 1)

    if n == 0
        rp
    elseif n == 1
        rp
    else
        logbs = rp[1:2:n-1]
        logcs = rp[2:2:n]
        inds = sortperm(logbs)

        rps = zeros(n)
        for i in 1:2:n-1
            rps[i] = logbs[inds[div(i,2)+1]]
            rps[i+1] = logcs[inds[div(i,2)+1]]
        end

        if n % 2 == 1
            rps[end] = rp[end]
        end

        rps
    end
end

function to_params(post::CARMAKalmanPosterior, x::Array{Float64,1})
    @assert size(x,1)==nparams(post)

    CARMAPosteriorParams(x[1], x[2], x[3], to_roots(x[4:3+post.p]), to_roots(x[4+post.p:end]))
end

function to_array(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    x = zeros(nparams(post))

    x[1] = p.mu
    x[2] = p.sigma
    x[3] = p.nu
    x[4:3+post.p] = to_rparams(p.arroots)
    x[4+post.p:end] = to_rparams(p.maroots)

    x
end

function rmin_rmax(post::CARMAKalmanPosterior)
    dt = minimum(diff(post.ts))
    T = post.ts[end] - post.ts[1]

    min_r = 1.0/T
    max_r = 1.0/dt

    min_r, max_r
end

function log_prior(post::CARMAKalmanPosterior, x::Array{Float64, 1})
    if !root_params_ordered(x[4:3+post.p]) || !root_params_ordered(x[4+post.p:end])
        return -Inf
    end
    log_prior(post, to_params(post, x))
end

function log_prior(post::CARMAKalmanPosterior, x::CARMAPosteriorParams)
    min_r, max_r = rmin_rmax(post)

    mu = mean(post.ys)
    sig = std(post.ys)

    if x.mu < mu - 10*sig || x.mu > mu + 10*sig
        return -Inf
    end

    if x.sigma < sig / 10.0 || x.sigma > sig * 10.0
        return -Inf
    end

    if x.nu < 0.1 || x.nu > 10.0
        return -Inf
    end

    for i in 1:post.p
        r = x.arroots[i]
        if real(r) < -max_r || real(r) > -min_r
            return -Inf
        end

        if imag(r) < -max_r || imag(r) > max_r
            return -Inf
        end
    end

    for i in 1:post.q
        r = x.maroots[i]
        if real(r) < -max_r || real(r) > -min_r
            return -Inf
        end

        if imag(r) < -max_r || imag(r) > max_r
            return -Inf
        end
    end

    -log(x.sigma) - log(x.nu)
end

function make_filter(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    CARMAKalmanFilter(post, p)
end

function CARMAKalmanFilter(post::CARMAKalmanPosterior, x::Array{Float64, 1})
    CARMAKalmanFilter(post, to_params(post, x))
end

function CARMAKalmanFilter(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    CARMAKalmanFilter(p.mu, p.sigma, p.arroots, p.maroots)
end

function generate(post::CARMAKalmanPosterior, x::Array{Float64, 1})
    generate(post, to_params(post, x))
end

function generate(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    filt = CARMAKalmanFilter(post, p)

    dy = post.dys*p.nu

    y = generate(filt, post.ts, dy)

    y
end

function log_likelihood(post::CARMAKalmanPosterior, x::Array{Float64, 1})
    log_likelihood(post, to_params(post, x))
end

function log_likelihood(post::CARMAKalmanPosterior, x::CARMAPosteriorParams)
    filt = nothing
    try
        filt = CARMAKalmanFilter(post, x)
    catch e
        if isa(e, DomainError)
            warn("Could not construct Kalman filter for parameters $(x)")
            return -Inf
        else
            rethrow()
        end
    end

    dys = post.dys * x.nu

    log_likelihood(filt, post.ts, post.ys, dys)
end

function log_posterior(post::CARMAKalmanPosterior, x::Array{Float64, 1})
    log_posterior(post, to_params(post, x))
end

function log_posterior(post::CARMAKalmanPosterior, x::CARMAPosteriorParams)
    lp = log_prior(post, x)

    if lp == -Inf
        lp
    else
        lp + log_likelihood(post, x)
    end
end

function randroots(rmin, rmax, n)
    p = zeros(n)

    bmin = 2*rmin
    bmax = 2*rmax

    cmin = rmin*rmin
    cmax = 2.0*rmax*rmax

    logbmin = log(bmin)
    logbmax = log(bmax)

    logcmin = log(cmin)
    logcmax = log(cmax)

    for i in 1:2:n-1
        logb = 0.0
        logc = 0.0
        while true
            logb = logbmin + (logbmax-logbmin)*rand()
            logc = logcmin + (logcmax-logcmin)*rand()

            rs = to_roots(Float64[logb, logc])

            r1 = rs[1]
            r2 = rs[2]

            if real(r1) < -rmin && real(r1) > -rmax && imag(r1) > -rmax && imag(r1) < rmax && real(r2) < -rmin && real(r2) > -rmax && imag(r2) > -rmax && imag(r2) < rmax
                break
            end
        end

        p[i] = logb
        p[i+1] = logc
    end

    if n % 2 == 1
        p[end] = log(rmin) + (log(rmax)-log(rmin))*rand()
    end

    to_roots(p)
end

function init(post, n)
    rmin, rmax = rmin_rmax(post)

    mu0 = mean(post.ys)
    sig0 = std(post.ys)

    xs = zeros((nparams(post), n))

    for i in 1:n
        mu = mu0-10*sig0 + 20*sig0*rand()
        sig = exp(log(0.1*sig0) + rand()*(log(10.0*sig0) - log(0.1*sig0)))
        nu = exp(log(0.1) + rand()*(log(10.0)-log(0.1)))

        arroots = randroots(rmin, rmax, post.p)
        maroots = randroots(rmin, rmax, post.q)

        p = CARMAPosteriorParams(mu, sig, nu, arroots, maroots)
        xs[:,i] = to_array(post, p)
    end

    xs
end

function whiten(post::CARMAKalmanPosterior, p::Array{Float64, 1})
    whiten(post, to_params(post, p))
end

function whiten(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    filt = CARMAKalmanFilter(post, p)

    whiten(filt, post.ts, post.ys, post.dys*p.nu)
end

""" Returns `(r, dr)` the residual and standard deviation on the
residual for the given parameters.
"""
function residuals(post::CARMAKalmanPosterior, p::Array{Float64, 1})
    residuals(post, to_params(post, p))
end

function residuals(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    filt = CARMAKalmanFilter(post, p)

    residuals(filt, post.ts, post.ys, post.dys*p.nu)
end

function residuals(filt::CARMAKalmanFilter, ts::Array{Float64, 1}, ys::Array{Float64, 1}, dys::Array{Float64, 1})
    reset!(filt)

    rs = zeros(size(ts, 1))
    drs = zeros(size(ts, 1))

    for i in 1:size(ts, 1)
        y, vy = predict(filt)

        rs[i] = ys[i] - y
        drs[i] = sqrt(vy + dys[i]*dys[i])

        observe!(filt, ys[i], dys[i])

        if i < size(ts, 1)
            advance!(filt, ts[i+1]-ts[i])
        end
    end

    rs, drs
end

function psdfreq(post::CARMAKalmanPosterior; nyquist_factor=1.0, oversample_factor=1.0)
    T = maximum(post.ts) - minimum(post.ts)
    dt_med = median(diff(post.ts))

    fmax = nyquist_factor/(2.0*dt_med)
    df = 1.0/(oversample_factor*T)

    collect(df:df:fmax)
end

""" Returns the one-sided PSD of the CARMA process in `post` described
by parameters `x` at the frequencies `fs`.

The PSD is normalised so that ``\\int_{0}^{\\infty} df \\, P(f) =
\\sigma^2``, the variance of the process.  """
function psd(post::CARMAKalmanPosterior, x::Array{Float64, 2}, fs::Array{Float64, 1})
    psds = [psd(post, x[:,i], fs) for i in 1:size(x,2)]
    hcat(psds...)
end

function psd(post::CARMAKalmanPosterior, x::Array{Float64, 1}, fs::Array{Float64, 1})
    psd(post, to_params(post, x), fs)
end

function psd(post::CARMAKalmanPosterior, p::CARMAPosteriorParams, fs::Array{Float64, 1})
    psd = zeros(size(fs, 1))

    filt = CARMAKalmanFilter(post, p)

    ar = p.arroots
    ma = p.maroots

    pp = size(ar, 1)
    q = size(ma, 1)

    # Formula for autocovariance at zero lag; see, e.g. [Kelly+
    # (2014)](https://dx.doi.org/10.1088/0004-637X/788/1/33).
    s = zero(ar[1]) # Complex zero
    pnorm = polyeval(ma, zero(ar[1]))
    for k in 1:pp
        rk = ar[k]
        t1 = polyeval(ma, rk)
        t2 = polyeval(ma, -rk)
        t3 = prod([(ar[l]-rk)*(conj(ar[l]) + rk) for l in 1:pp if l != k])
        s += t1*t2/(-2.0*real(rk)*t3*pnorm*pnorm)
    end

    sig = p.sigma/real(sqrt(s))

    for i in 1:size(fs, 1)
        f = fs[i]
        tpif = 2.0*pi*1.0im*f

        numer = polyeval(p.maroots, tpif) / polyeval(p.maroots, zero(tpif))
        denom = polyeval(p.arroots, tpif)

        psd[i] = 2.0*sig*sig*abs2(numer)/abs2(denom) # 2.0 for one-sided
    end

    psd
end

function frequencies(post::CARMAKalmanPosterior, x::Array{Float64, 1})
    frequencies(post, to_params(post, x))
end

function frequencies(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    imag(p.arroots)/(2.0*pi)
end

function frequencies(post::CARMAKalmanPosterior, x::Array{Float64, 2})
    fs = [frequencies(post, x[:,i]) for i in 1:size(x,2)]
    hcat(fs...)
end

function drates(post::CARMAKalmanPosterior, x::Array{Float64, 1})
    drates(post, to_params(post, x))
end

function drates(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    -real(p.arroots)
end

function drates(post::CARMAKalmanPosterior, x::Array{Float64, 2})
    drs = [drates(post, x[:,i]) for i in 1:size(x,2)]
    hcat(drs...)
end

function qfactors(post::CARMAKalmanPosterior, x::Array{Float64, 1})
    qfactors(post, to_params(post, x))
end

function qfactors(post::CARMAKalmanPosterior, p::CARMAPosteriorParams)
    fs = frequencies(post, p)
    dr = drates(post, p)

    pi*fs ./ dr
end

function qfactors(post::CARMAKalmanPosterior, x::Array{Float64, 2})
    qs = [qfactors(post, x[:,i]) for i in 1:size(x,2)]
    hcat(qs...)
end

""" Return the predicted state of the underlying process and the
variance in the prediction at the given times.  """
function predict(post::CARMAKalmanPosterior, x::Array{Float64, 1}, t::Float64)
    predict(post, to_params(post, x), t)
end

# This is now generic over the type of filt; should work for AR1 and CARMA.
function filter_prediction(filt, ts::Array{Float64, 1}, ys::Array{Float64, 1}, dys::Array{Float64, 1}, obs::Array{Bool, 1})
    n = size(ys, 1)

    yps = zeros(n)
    vyps = zeros(n)

    reset!(filt)

    for i in 1:n
        y,vy = predict(filt)

        yps[i] = y
        vyps[i] = vy

        if obs[i]
            observe!(filt, ys[i], dys[i])
        end

        if i < n
            advance!(filt, ts[i+1]-ts[i])
        end
    end

    yps, vyps
end

function predict(post, x::Array{Float64, 1}, ts::Array{Float64, 1})
    predict(post, to_params(post, x), ts)
end

function predict(post, p, ts::Array{Float64, 1})
    all_ts = vcat(post.ts, ts)
    all_obs = convert(Array{Bool, 1}, vcat(trues(size(post.ts,1)), falses(size(ts,1))))
    all_ys = vcat(post.ys, zeros(size(ts,1)))
    all_dys = vcat(post.dys, zeros(size(ts, 1)))

    iforward = sortperm(all_ts)
    ibackward = reverse(iforward)

    filt = make_filter(post, p)

    ysf, vysf = filter_prediction(filt, all_ts[iforward], all_ys[iforward], all_dys[iforward], all_obs[iforward])
    ysb, vysb = filter_prediction(filt, -all_ts[ibackward], all_ys[ibackward], all_dys[ibackward], all_obs[ibackward])

    ysb = reverse(ysb)
    vysb = reverse(vysb)

    ys = (vysf.*ysb + vysb.*ysf)./(vysf + vysb)
    vys = 1.0./(1.0./vysb + 1.0./vysf)

    ys_out = zeros(size(all_ts,1))
    vys_out = zeros(size(all_ts,1))

    for i in 1:size(all_ts, 1)
        ys_out[iforward[i]] = ys[i]
        vys_out[iforward[i]] = vys[i]
    end

    ys_out[size(post.ts,1)+1:end], vys_out[size(post.ts,1)+1:end]
end

"""
    draw_extrapolation(post, p_or_x, ts, dys)

Produce an extrapolation of the process with parameters `p_or_x`
fitted to the data in `post` at the times `ts` and with reported
observational errors `dys`.

Note that the true observational uncertainty is scaled by the `nu`
parameter: `dys_true = dys*p.nu`.

`p_or_x` can be either an array or a parameter object.
"""
function draw_extrapolation(post::CARMAKalmanPosterior, x::Array{Float64, 1}, ts::Array{Float64, 1}, dys::Array{Float64, 1})
    draw_extrapolation(post, to_params(post, x), ts, dys)
end

function draw_extrapolation(post::CARMAKalmanPosterior, p::CARMAPosteriorParams, ts::Array{Float64, 1}, dys::Array{Float64, 1})
    @assert all(ts .> maximum(post.ts))
    @assert issorted(ts)

    filt = CARMAKalmanFilter(post, p)

    reset!(filt)

    # First, fit the filter to the observed data; by the end of the
    # loop, the filter has incorporated each data point in post.
    for i in eachindex(post.ts)
        observe!(filt, post.ys[i], p.nu*post.dys[i])

        if i < size(post.ts,1)
            advance!(filt, post.ts[i+1]-post.ts[i])
        end
    end

    # Now run the filter forward through the ts
    tcurrent = post.ts[end]
    ys_out = zeros(size(ts, 1))
    for i in eachindex(ts)
        advance!(filt, ts[i] - tcurrent)
        tcurrent = ts[i]

        draw_and_collapse!(filt)

        y, _ = predict(filt)

        ys_out[i] = y + p.nu*dys[i]*randn()
    end

    ys_out
end

struct MultiSegmentCARMAKalmanPosterior
    ts::Array{Array{Float64, 1}, 1}
    ys::Array{Array{Float64, 1}, 1}
    dys::Array{Array{Float64, 1}, 1}
    p::Int
    q::Int

    function MultiSegmentCARMAKalmanPosterior(ts, ys, dys, p, q)
        @assert p>0
        @assert q<p
        @assert q>=0

        new(ts, ys, dys, p, q)
    end
end

function nparams(post::MultiSegmentCARMAKalmanPosterior)
    post.p + post.q + 1 + 2*size(post.ts,1)
end

function nsegments(post::MultiSegmentCARMAKalmanPosterior)
    size(post.ts, 1)
end

struct MultiSegmentCARMAPosteriorParams
    mu::Array{Float64, 1}
    sigma::Float64
    nu::Array{Float64, 1}
    arroots::Array{ComplexF64, 1}
    maroots::Array{ComplexF64, 1}
end

function to_params(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    @assert size(x, 1)==nparams(post)

    ns = nsegments(post)

    MultiSegmentCARMAPosteriorParams(x[1:ns], x[ns+1], x[ns+2:2*ns+1], to_roots(x[2*ns+2:2*ns+1+post.p]), to_roots(x[2*ns+2+post.p:end]))
end

function to_array(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    x = zeros(nparams(post))

    ns = nsegments(post)

    x[1:ns] = p.mu
    x[ns+1] = p.sigma
    x[ns+2:2*ns+1] = p.nu
    x[2*ns+2:2*ns+1+post.p] = to_rparams(p.arroots)
    x[2*ns+2+post.p:end] = to_rparams(p.maroots)

    x
end

function rmin_rmax(post::MultiSegmentCARMAKalmanPosterior)
    allts = sort(vcat(post.ts...))

    dt = minimum(diff(allts))
    T = allts[end] - allts[1]

    min_r = 1.0/T
    max_r = 1.0/dt

    min_r, max_r
end

function log_prior(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    ns = nsegments(post)
    if !root_params_ordered(x[2*ns+2:2*ns+1+post.p]) || !root_params_ordered(x[2*ns+2+post.p:end])
        return -Inf
    end
    log_prior(post, to_params(post, x))
end

function log_prior(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    min_r, max_r = rmin_rmax(post)

    mus = Float64[mean(y) for y in post.ys]
    sigs = Float64[std(y) for y in post.ys]

    if any(p.mu .< mus - 10.0*sigs) || any(p.mu .> mus + 10.0*sigs)
        return -Inf
    end

    sigwts = Float64[size(y,1)-1 for y in post.ys]
    sig_tot = sqrt(sum(sigwts.*sigs.*sigs)/sum(sigwts))

    if p.sigma < sig_tot / 10.0 || p.sigma > sig_tot*10.0
        return -Inf
    end

    if any(p.nu .< 0.1) || any(p.nu .> 10.0)
        return -Inf
    end

    for i in 1:post.p
        r = p.arroots[i]
        if real(r) < -max_r || real(r) > -min_r
            return -Inf
        end

        if imag(r) < -max_r || imag(r) > max_r
            return -Inf
        end
    end

    for i in 1:post.q
        r = p.maroots[i]
        if real(r) < -max_r || real(r) > -min_r
            return -Inf
        end

        if imag(r) < -max_r || imag(r) > max_r
            return -Inf
        end
    end

    -log(p.sigma) - sum(log.(p.nu))
end

function make_filter(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    CARMAKalmanFilter(post, p)
end

function CARMAKalmanFilter(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    CARMAKalmanFilter(post, to_params(post, x))
end

function CARMAKalmanFilter(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    CARMAKalmanFilter(0.0, p.sigma, p.arroots, p.maroots)
end

function log_likelihood(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    log_likelihood(post, to_params(post, x))
end

function alltsysdys(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    allts = vcat(post.ts...)

    ys = Array{Float64, 1}[ys .- mu for (ys, mu) in zip(post.ys, p.mu)]
    dys = Array{Float64, 1}[dys.*nu for (dys, nu) in zip(post.dys, p.nu)]

    allys = vcat(ys...)
    alldys = vcat(dys...)

    inds = sortperm(allts)
    allrinds = invperm(inds)  # Gives the index in the allts array
                           # corresponding to each data point
    allts = allts[inds]
    allys = allys[inds]
    alldys = alldys[inds]

    rinds = Array{Int, 1}[]
    i = 1
    for t in post.ts
        n = size(t, 1)
        push!(rinds, allrinds[i:i+n-1])
        i = i + n
    end

    (allts, allys, alldys, rinds)
end

function log_likelihood(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    filt = nothing
    try
        filt = CARMAKalmanFilter(post, p)
    catch e
        if isa(e, DomainError)
            warn("Could not construct Kalman filter for parameters $(p)")
            return -Inf
        else
            rethrow()
        end
    end

    allts, allys, alldys, rinds = alltsysdys(post, p)

    log_likelihood(filt, allts, allys, alldys)
end

function log_posterior(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    log_posterior(post, to_params(post, x))
end

function log_posterior(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    lp = log_prior(post, p)

    if lp == -Inf
        lp
    else
        lp + log_likelihood(post, p)
    end
end

function init(post::MultiSegmentCARMAKalmanPosterior, n)
    rmin, rmax = rmin_rmax(post)

    mu0s = Float64[mean(y) for y in post.ys]
    sig0s = Float64[std(y) for y in post.ys]

    sigwts = Float64[size(y,1)-1 for y in post.ys]
    total_sig = sqrt(sum(sigwts.*sig0s.*sig0s)/sum(sigwts))

    xs = zeros((nparams(post), n))

    for i in 1:n
        mus = mu0s-10.0*sig0s + 20.0*sig0s.*rand(size(sig0s,1))
        sig = exp(log(0.1*total_sig) + rand()*(log(10.0*total_sig) - log(0.1*total_sig)))
        nus = exp.(log(0.1) .+ rand(size(mu0s,1)).*(log(10.0) - log(1.0)))

        arroots = randroots(rmin, rmax, post.p)
        maroots = randroots(rmin, rmax, post.q)

        p = MultiSegmentCARMAPosteriorParams(mus, sig, nus, arroots, maroots)
        xs[:,i] = to_array(post,p)
    end

    xs
end

function psdfreq(post::MultiSegmentCARMAKalmanPosterior; nyquist_factor=1.0, oversample_factor=1.0)
    T = post.ts[end][end] - post.ts[1][1]
    dt_med = median(diff(sort(vcat(post.ts...))))

    fmax = nyquist_factor/(2.0*dt_med)
    df = 1.0/(oversample_factor*T)

    collect(df:df:fmax)
end

function psd(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 2}, fs::Array{Float64, 1})
    psds = [psd(post, x[:,i], fs) for i in 1:size(x,2)]
    hcat(psds...)
end

function psd(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1}, fs::Array{Float64, 1})
    psd(post, to_params(post, x), fs)
end

function psd(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams, fs::Array{Float64, 1})
    psd = zeros(size(fs, 1))

    filt = CARMAKalmanFilter(post, p)

    for i in 1:size(fs, 1)
        f = fs[i]
        tpif = 2.0*pi*1.0im*f

        numer = polyeval(p.maroots, tpif) / polyeval(p.maroots, 0.0+0.0im)
        denom = polyeval(p.arroots, tpif)

        psd[i] = 2.0*filt.sig*filt.sig*abs2(numer)/abs2(denom)
    end

    psd
end

function frequencies(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    frequencies(post, to_params(post, x))
end

function frequencies(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    imag(p.arroots)/(2.0*pi)
end

function frequencies(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 2})
    fs = [frequencies(post, x[:,i]) for i in 1:size(x,2)]
    hcat(fs...)
end

function drates(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    drates(post, to_params(post, x))
end

function drates(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    -real(p.arroots)
end

function drates(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 2})
    drs = [drates(post, x[:,i]) for i in 1:size(x,2)]
    hcat(drs...)
end

function qfactors(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    qfactors(post, to_params(post, x))
end

function qfactors(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    fs = frequencies(post, p)
    dr = drates(post, p)

    pi*fs ./ dr
end

function qfactors(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 2})
    qs = [qfactors(post, x[:,i]) for i in 1:size(x,2)]
    hcat(qs...)
end

function residuals(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1})
    residuals(post, to_params(post, x))
end

function residuals(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams)
    filt = CARMAKalmanFilter(post, p)

    allts, allys, alldys, rinds = alltsysdys(post, p)
    allrs, alldrs = residuals(filt, allts, allys, alldys)

    rs = [[allrs[i] for i in ri] for ri in rinds]
    drs = [[alldrs[i] for i in ri] for ri in rinds]

    (rs, drs)
end

function predict(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1}, ts::Array{Float64, 1})
    predict(post, to_params(post, x), ts)
end

function predict(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams, ts::Array{Float64, 1})
    allts, allys, alldys, rinds = alltsysdys(post, p)

    singlepost = CARMAKalmanPosterior(allts, allys, alldys, post.p, post.q)
    singleparams = to_params(singlepost, zeros(nparams(singlepost)))

    singleparams.mu = 0.0
    singleparams.sigma = p.sigma
    singleparams.nu = 1.0
    singleparams.arroots = p.arroots
    singleparams.maroots = p.maroots

    predict(singlepost, singleparams, ts)
end

function draw_extrapolation(post::MultiSegmentCARMAKalmanPosterior, x::Array{Float64, 1}, ts::Array{Float64, 1}, dys::Array{Float64, 1})
    draw_extrapolation(post, to_params(post, x), ts, dys)
end

"""
If `post` is a `MultiSegmentCARMAKalmanPosterior`, then the returned
extrapolation will have `mu = 0` and `nu = 1`.  If a different mean
and error scaling are desired they should be applied separately.
"""
function draw_extrapolation(post::MultiSegmentCARMAKalmanPosterior, p::MultiSegmentCARMAPosteriorParams, ts::Array{Float64, 1}, dys::Array{Float64, 1})
    allts, allys, alldys, rinds = alltsysdys(post, p)

    singlepost = CARMAKalmanPosterior(allts, allys, alldys, post.p, post.q)
    singleparams = to_params(singlepost, zeros(nparams(singlepost)))

    singleparams.mu = 0.0
    singleparams.sigma = p.sigma
    singleparams.nu = 1.0
    singleparams.arroots = p.arroots
    singleparams.maroots = p.maroots

    draw_extrapolation(singlepost, singleparams, ts, dys)
end

end


text:
mutable struct LocalNNFunctionApproximator{NT<:NNTree, V<:AbstractVector{Float64}} <: LocalFunctionApproximator
    nntree::NT
    nnpoints::AbstractVector{V}
    nnvalues::Vector{Float64}
    knnK::Union{Int64,Nothing}
    rnnR::Union{Float64,Nothing}
end
# Two default constructors for k and r
LocalNNFunctionApproximator(nntree::NT, nnpts::AbstractVector{V}, knnK::Int64) where {NT <: NNTree,V<:AbstractVector{Float64}} =
    LocalNNFunctionApproximator(nntree, nnpts, zeros(length(nntree.indices)), knnK, nothing)
LocalNNFunctionApproximator(nntree::NT, nnpts::AbstractVector{V}, rnnR::Float64) where {NT <: NNTree,V<:AbstractVector{Float64}} =
    LocalNNFunctionApproximator(nntree, nnpts, zeros(length(nntree.indices)), nothing, rnnR)

################ INTERFACE FUNCTIONS ################
function n_interpolating_points(nnfa::LocalNNFunctionApproximator)
    return length(nnfa.nntree.indices)
end

function get_all_interpolating_points(nnfa::LocalNNFunctionApproximator)
  return nnfa.nnpoints
end

function get_all_interpolating_values(nnfa::LocalNNFunctionApproximator)
  return nnfa.nnvalues
end

function get_interpolating_nbrs_idxs_wts(nnfa::LocalNNFunctionApproximator, v::AbstractVector{Float64})

    @assert (nnfa.knnK != nothing || nnfa.rnnR != nothing)

    if nnfa.knnK != nothing
        # Do k-NN lookup to get data and distances
        idxs, dists = knn(nnfa.nntree, v, nnfa.knnK)
    else
        # Do inrange lookup to get data
        # Then use metric to get dists between query and each nearest neighbor
        idxs = inrange(nnfa.nntree, v, nnfa.rnnR)
        dists = zeros(length(idxs))
        for (i,idx) in enumerate(idxs)
            dists[i] = Distances.evaluate(nnfa.nntree.metric, v, nnfa.nnpoints[idx])
        end
    end

    weights = zeros(length(dists))

    # If exactly one point, set that probability of that idx to 1
    # and that of all others to 0.
    if minimum(dists) < eps(Float64)
        weights[argmin(dists)] = 1.0
    else
        for (i,d) in enumerate(dists)
            weights[i] = 1.0/d
        end
        weights /= sum(weights)
    end

    return (idxs, weights)
end

function compute_value(nnfa::LocalNNFunctionApproximator, v::AbstractVector{Float64})

    idxs, wts = get_interpolating_nbrs_idxs_wts(nnfa, v)

    # Do a weighted average of values
    value::Float64 = 0.0
    wtsum::Float64 = 0.0
    for (i,idx) in enumerate(idxs)
        value += wts[i]*nnfa.nnvalues[idx]
        wtsum += wts[i]
    end
    value /= wtsum

    return value
end

function compute_value(nnfa::LocalNNFunctionApproximator, v_list::AbstractVector{V}) where V <: AbstractVector{Float64}
    @assert length(v_list) > 0
    vals = [compute_value(nnfa, pt) for pt in v_list]
    return vals
end

function set_all_interpolating_values(nnfa::LocalNNFunctionApproximator, nnvalues::AbstractVector{Float64})
  nnfa.nnvalues = copy(nnvalues)
end

ext: jl


text:
@testset "Enmap slicing" begin
    shape0, wcs0 = fullsky_geometry(π/180)
    m = Enmap(rand(shape0...), wcs0)

    # regular slicing
    m_sliced = m[5:10, 1:end]
    wcs = getwcs(m_sliced)

    @test [-1.0, 1.0] ≈ collect(wcs.cdelt)
    @test [176.5, 91.0] ≈ collect(wcs.crpix)
    @test [0.5, 0.0] ≈ collect(wcs.crval)
    @test (m.data)[5:10, 1:end] ≈ m_sliced

    # backwards slicing
    m_sliced = m[1:12, end:-1:begin]
    wcs = getwcs(m_sliced)
    @test [-1.0, -1.0] ≈ collect(wcs.cdelt)
    @test [180.5, 91.0] ≈ collect(wcs.crpix)
    @test [0.5, 0.0] ≈ collect(wcs.crval)
    @test (m.data)[1:12, end:-1:begin] ≈ m_sliced

    # non-unit steps
    m_sliced = m[3:5:24, 39:-3:2]
    wcs = getwcs(m_sliced)
    @test [-5.0, -3.0] ≈ collect(wcs.cdelt)
    @test [36.1, -16.666666666666668] ≈ collect(wcs.crpix)
    @test [0.5, 0.0] ≈ collect(wcs.crval)
    @test (m.data)[3:5:24, 39:-3:2] ≈ m_sliced

    # removing the first two axes
    @test typeof(m[1,:]) == Array{Float64,1}
    @test typeof(m[:,1]) == Array{Float64,1}
end

##
@testset "Enmap view slicing" begin
    shape0, wcs0 = fullsky_geometry(π/180)
    m = Enmap(rand(shape0...), wcs0)

    # regular slicing
    m_sliced = @view m[5:10, 1:end]
    wcs = getwcs(m_sliced)
    @test [-1.0, 1.0] ≈ collect(wcs.cdelt)
    @test [176.5, 91.0] ≈ collect(wcs.crpix)
    @test [0.5, 0.0] ≈ collect(wcs.crval)
    @test (m.data)[5:10, 1:end] ≈ m_sliced

    # backwards slicing
    m_sliced = @view m[1:12, end:-1:begin]
    wcs = getwcs(m_sliced)
    @test [-1.0, -1.0] ≈ collect(wcs.cdelt)
    @test [180.5, 91.0] ≈ collect(wcs.crpix)
    @test [0.5, 0.0] ≈ collect(wcs.crval)
    @test (m.data)[1:12, end:-1:begin] ≈ m_sliced

    # non-unit steps
    m_sliced = @view m[3:5:24, 39:-3:2]
    wcs = getwcs(m_sliced)
    @test [-5.0, -3.0] ≈ collect(wcs.cdelt)
    @test [36.1, -16.666666666666668] ≈ collect(wcs.crpix)
    @test [0.5, 0.0] ≈ collect(wcs.crval)
    @test (m.data)[3:5:24, 39:-3:2] ≈ m_sliced
end

## WCS should be not be shared under deepcopy, broadcasting, or broadcasted assignment
@testset "Enmap copying behavior" begin
    for copy_op in (copy, deepcopy, similar)  # these should all do the same thing: NEVER keep the same WCS
        shape0, wcs0 = fullsky_geometry(Pixell.WCS.WCSTransform, π/180)
        m = Enmap(rand(shape0...), wcs0)
        m2 = copy_op(m)
        @test !(m.wcs === m2.wcs)
        m2.wcs.cdelt = collect([99., 99.])
        @test !(m.wcs.cdelt ≈ collect([99., 99.]))

        m = Enmap(rand(shape0...), wcs0)
        m2 = m.^2
        @test !(m.wcs === m2.wcs)
        m2.wcs.cdelt = collect([99., 99.])
        @test !(m.wcs.cdelt ≈ collect([99., 99.]))

        m = Enmap(rand(shape0...), wcs0)
        m2 = copy_op(m)
        m2 .= m
        @test !(m.wcs === m2.wcs)
        m2.wcs.cdelt = collect([99., 99.])  # also make sure sub-arrays aren't shared
        @test !(m.wcs.cdelt ≈ collect([99., 99.]))
    end
end

##
@testset "Enmap broadcasting" begin
    shape, wcs = fullsky_geometry(deg2rad(1); dims=(3,))
    A, B = rand(shape...), rand(shape...)
    ma = Enmap(A, wcs)
    mb = Enmap(B, wcs)
    @test A .+ B == ma .+ mb
    @test A .+ B == ma .+ B
    @test A .+ B == A .+ mb
    @test A .+ B .* sin.(A.^2) == (ma .+ mb .* sin.(ma.^2))

    ma .= 1.0
    @test all(ma .≈ 1.0)
    ma .= mb
    @test all(ma .≈ mb)
    ma[1,:,3] .= 2.0
    @test all(ma[1,:,3] .≈ 2.0)
    ma[:,end,3] .= 3.0
    @test all(ma[:,end,3] .≈ 3.0)
    ma[:,1,:] .= mb[:,2,:]
    @test all(ma.data[:,1,:] .≈ mb.data[:,2,:])

    A, B = rand(shape...), rand(shape...)
    ma = Enmap(A, wcs)
    mb = Enmap(B, wcs)
    mb[:,:,1] .= ma[:,:,1]
    @test all(ma.data[:,:,1] .≈ mb.data[:,:,1])
    @test !all(ma.data[:,:,2] .≈ mb.data[:,:,2])
    @test !all(ma.data[:,:,3] .≈ mb.data[:,:,3])
    mb[:,:,:] .= ma[:,:,:]
    @test all(ma.data .≈ mb.data)

    mv = @view ma[1,:,1]
    @test Pixell.getwcs(mv) == Pixell.NoWCS()
    mv = ma[1,:,1]
    @test Pixell.getwcs(mv) == Pixell.NoWCS()
    mv = ma[1,:,:]
    @test Pixell.getwcs(mv) == Pixell.NoWCS()
    mv = ma[1:5,:,1]
    @test Pixell.getwcs(mv) != Pixell.NoWCS()
    mv = ma[1:5,:,:]
    @test Pixell.getwcs(mv) != Pixell.NoWCS()
end

@testset "Enmap WCS props" begin
    shape, wcs = fullsky_geometry(π/180)
    m = Enmap(zeros(shape), wcs)
    @test stride(m,1) == 1
    @test stride(m,2) == stride(m.data, 2)

    imap = read_map("data/test.fits"; trim=false)
    @test Pixell.getunit(imap.wcs) ≈ π/180
    @test sprint(show, imap.wcs) == "WCSTransform(naxis=2,cdelt=[-1.0, 1.0],crval=[0.5, 0.0],crpix=[180.5, 91.0])"

    @test collect(Pixell.getcrval(imap.wcs)) == imap.wcs.crval
    @test collect(Pixell.getcrpix(imap.wcs)) == imap.wcs.crpix
    @test collect(Pixell.getcdelt(imap.wcs)) == imap.wcs.cdelt

    wcs = convert(CarClenshawCurtis{Float64}, imap.wcs)
    @test Pixell.getcrval(imap.wcs) == wcs.crval
    @test Pixell.getcrpix(imap.wcs) == wcs.crpix
    @test Pixell.getcdelt(imap.wcs) == wcs.cdelt

    imap.wcs.cunit = ["rad", "rad"]
    @test Pixell.getunit(imap.wcs) ≈ 1.0
    imap.wcs.cunit = ["arcmin", "arcmin"]
    @test Pixell.getunit(imap.wcs) ≈ π / 180 / 60
    imap.wcs.cunit = ["arcsec", "arcsec"]
    @test Pixell.getunit(imap.wcs) ≈ π / 180 / 60 / 60
    imap.wcs.cunit = ["mas", "mas"]
    @test Pixell.getunit(imap.wcs) ≈ π / 180 / 60 / 60 / 1000
end


text:
include("testfunctions.jl")
using DiferentialEvolution
using PlotlyJS

#Plots
#=
len=500
x = LinRange(-600,600,len)
y = LinRange(-600,600,len)
X = ones(len)*x'
Y = y*ones(len)'
Z = Eggholder.(X,Y)
Eggholder(-2.5153506020027042e64, 1.79281509087981e64)
@time critic1 = DE_JADE(Eggholder, 500, 25, [400,400], [600,600], 0.2, 1/20)
@time critic2 = DE_dither(Eggholder, 1000, 500, [-512,-512], [512,512],0.9,0.5)
p_init = contour(x =x, y= y,z = Z)
crit = scatter(x=[critic[2][1]],y=[critic[2][2]], mode="markers",color=2)
plot([p_init,crit])
=#

f(x) = sinc(x)
x = -5:0.1:5
y = f.(x)
g(par...) = par*[@. x^i for i in 1:length(par)]
error(a,b,c,d,e,f,h,i,k,l,m,n) = sum(abs.(y-g(a,b,c,d,e,f,h,i,k,l,m,n)))
res = DiferentialEvolution.DE_JADE(error,1000,500,[-100,-100,-100,-100,-100,-100,-100,-100,-100,-100,-100,-100],[100,100,100,100,100,100,100,100,100,100,100,100],0.15,0.2)

a1 = scatter(x=x, y=y)
a2 = scatter(x=x, y=g(res[2]...))
plot([a1,a2])


text:
function jsrender(session::Session, content::Vector)
    DOM.um("span", jsrender.((session,), content)...)
end

jsrender(session::Session, md::Markdown.MD) = JSServe.jsrender(session, jsrender(session, md.content))

function jsrender(session::Session, header::Markdown.Header{l}) where l
    DOM.um("h$l", htmlinline(session, header.text))
end

function jsrender(session::Session, code::Markdown.Code)
    maybe_lang = !isempty(code.language) ? Any[:class=>"language-$(code.language)"] : []
    DOM.um("pre", DOM.um("code", maybe_lang..., code.code))
end

function jsrender(session::Session, md::Markdown.Paragraph)
    return DOM.p(htmlinline(session, md.content))
end

function jsrender(session::Session, md::Markdown.BlockQuote)
    DOM.um("blockquote", jsrender(session, md.content))
end

function jsrender(session::Session, f::Markdown.Footnote)
    return DOM.div(
        class = "footnote", id = "footnote-$(f.id)",
        DOM.p(class = "footnote-title", f.id),
        jsrender(session, f.text)
    )
end

function jsrender(session::Session, md::Markdown.Admonition)
    return DOM.div(class = "admonition $(md.category)") do
        DOM.p(class = "admonition-title", md.title)
        jsrender(session, md.content)
    end
end

function jsrender(session::Session, md::Markdown.List)
    maybe_attr = md.ordered > 1 ? Any[:start => string(md.ordered)] : []
    tag = Markdown.isordered(md) ? "ol" : "ul"
    DOM.um(
        tag, maybe_attr...,
        map(md.items) do item
            DOM.um("li", jsrender(session, item))
        end...
    )
end

function jsrender(session::Session, md::Markdown.HorizontalRule)
    return DOM.m_unesc("hr")
end

function htmlinline(session::Session, content::Vector)
    DOM.um("div", htmlinline.((session,), content)...)
end

function htmlinline(session::Session, code::Markdown.Code)
    DOM.um("code", code.code)
end

function htmlinline(session::Session, md::Union{Symbol, AbstractString})
    return md
end

function htmlinline(session::Session, md::Markdown.Bold)
    DOM.um("strong", htmlinline(session, md.text))
end

function htmlinline(session::Session, md::Markdown.Italic)
    DOM.um("em", htmlinline(session, md.text))
end

function htmlinline(session::Session, md::Markdown.Image)
    DOM.um("img", src = md.url, alt = md.alt)
end

function htmlinline(session::Session, f::Markdown.Footnote)
    DOM.um("a", href = "#footnote-$(f.id)", class = "footnote", string(session, "[", f.id, "]"))
end

function htmlinline(session::Session, link::Markdown.Link)
    return DOM.um("a", href = link.url, htmlinline(session, link.text))
end

function htmlinline(session::Session, br::Markdown.LineBreak)
    DOM.um("br")
end

htmlinline(session::Session, x) = jsrender(session, x)


text:
module heat_juafem_examples
using FinEtools
using FinEtools.MeshExportModule
using Compat.Test

function heat_juafem_example()
    println("""Heat conduction example from JuAFEM.""")
    t0 = time()

    A = 2.0
    thermal_conductivity =  [i==j ? one(FFlt) : zero(FFlt) for i=1:2, j=1:2]; # conductivity matrix
    function getsource!(forceout::FFltVec, XYZ::FFltMat, tangents::FFltMat, fe_label::FInt)
        forceout[1] = 1.0; #heat source
    end
    N = 1000;

    println("Mesh generation")
    @time fens,fes = Q4block(A, A, N, N)
    fens.xyz[:,1] .-= A/2
    fens.xyz[:,2] .-= A/2

    geom = NodalField(fens.xyz)
    Temp = NodalField(zeros(size(fens.xyz,1),1))

    println("Searching nodes  for BC")
    @time l1 = selectnode(fens; box=[-A/2 -A/2 -A/2 A/2], inflate = 1.0/N/100.0)
    @time l2 = selectnode(fens; box=[A/2 A/2 -A/2 A/2], inflate = 1.0/N/100.0)
    @time l3 = selectnode(fens; box=[-A/2 A/2 -A/2 -A/2], inflate = 1.0/N/100.0)
    @time l4 = selectnode(fens; box=[-A/2 A/2 A/2 A/2], inflate = 1.0/N/100.0)
    List = vcat(l1, l2, l3, l4);
    setebc!(Temp, List, true, 1, 0.0)
    applyebc!(Temp)

    numberdofs!(Temp)

    t1 = time()

    m = MatHeatDiff(thermal_conductivity)
    femm = FEMMHeatDiff(IntegData(fes, GaussRule(2, 2)), m)

    println("Conductivity")
    @time K=conductivity(femm, geom, Temp)

    println("Internal heat generation")
    fi = ForceIntensity(FFlt, 1, getsource!);
    @time F1 = distribloads(femm, geom, Temp, fi, 3);

    println("Factorization")
    @time K = cholesky(K)
    println("Solution of the factorized system")
    @time U =  K\(F1)
    scattersysvec!(Temp, U[:])

    println("Total time elapsed = $(time() - t0) [s]")
    println("Solution time elapsed = $(time() - t1) [s]")

    println("Maximum temperature = $(maximum(Temp.values)) ")
    # using MeshExportModule

    # File =  "a.vtk"
    # MeshExportModule.vtkexportmesh(File, fes.conn, hcat(geom.values,  Temp.values), MeshExportModule.Q4; scalars=[("Temperature", Temp.values)])
    # @async run(`"paraview.exe" $File`)
end # heat_juafem_example

function allrun()
    #println("#####################################################")
    println("# heat_juafem_example ")
    heat_juafem_example()
    return true
end # function allrun

end # module heat_juafem_examples


text:
#Wave equation Differential Equation
function Zilndroa_Diffeq!(df,f,param,t)

    #Get Operators and variables
    Opps        = param[1]
    vars        = param[2]

    r = Opps.rcoord
    θ = Opps.tcoord

    RHS = Opps.DiffO_1
    Dr  = Opps.Dr
    Dθ  = Opps.Dθ

    dr_ψ  = Opps.aux_vex_1
    dθ_ψ  = Opps.aux_vex_2

    #Fields of our problem (Julia changes the rhs if you change the lhs)
    ψ   = f.x[1]
    ϕ   = f.x[2]

    dψ  = df.x[1]
    dϕ  = df.x[2]

    mul!(dr_ψ , Dr , ψ)
    mul!(dθ_ψ , Dθ , ψ)

    #Differential equation
    mul!(dψ , I ,  ϕ)
    mul!(dϕ , RHS ,  ψ)

    Nr      = vars.rnodes
    Ntheta  = vars.thnodes
    Ω       = vars.Ω
    #Boundary conditions

    for i in 1:Ntheta

        #Inner BC (Impedance)
        idx = (i-1)*Nr+1
        dψ[idx] = Z(t , θ[i] ,  vars ) * ( dr_ψ[idx] - 0.5 * ψ[idx] / r[1] ) - Ω * dθ_ψ[idx]
        #dψ[idx] = 0

        #Outer BC (reflective)
        idx = i*Nr
        dψ[idx] = 0
    end

end

function solve_zilindroa_system(p_in::Param)

    #
    # Setup Routine
    #

    #Welcome message
    welcome_message();

    #Start time
    init_time = time()
    max_runtime = time_string_to_seconds(p_in.max_runtime)

    #Setup routine setup check
    if( !setup_folder(p_in) )
        print_exit_message(  time() - init_time )
        return false
    end

    #Read parameters
    p = p_in
    print_parameters_to_screen(p)
    println("OK")

    #Storage file
    println("Creating HDF5 file...")

    fdata   = "./data"
    fname   = fdata * p.folder *  "$(p.fname).h5"
    fid = h5open(fname, "w")

    save_parameters_to_file( p , fid )
    println("OK")

    #Create operators structures
    println("Setting up Operators...")
    Operators = set_up_operator(p)
    println("OK")

    #Time variables
    tmin = p.t_sim_init
    tmax = p.t_sim_final
    tspan       = ( tmin , tmax )

    iter        = 0
    iter_save   = 0
    every       = floor(Int , p.out_every_t/p.deltat + 1.0)
    save_time   = 0

    #Initial Conditions
    println("Setting up inicial configuration...")
    ψ , dψ = get_gaussian_pulse( p , Operators)

    U = ArrayPartition( ψ , dψ )
    println("OK")

    my_params = ( Operators , p)

    #Problem ODE and integrator
    prob = ODEProblem(  Zilndroa_Diffeq! , U , tspan , my_params )

    integrator = init( prob , RK4() , save_everystep=false , dt=p.deltat , adaptive=false )

    #Create a group for the functions
    g1 =  create_group(fid , "ψ");
    g2 =  create_group(fid , "dψ");

    #Save first configuration of field

    dset_name = string("iteration_",iter_save)

    data1 = integrator.u.x[1]
    dset1, dtype1  = create_dataset(g1, dset_name , data1 )
    write_dataset(dset1, dtype1 , data1)

    data2 = integrator.u.x[2]
    dset2, dtype2  = create_dataset(g2, dset_name , data2 )
    write_dataset(dset2, dtype2 , data2)

    #Add HDF5.attributes time to both datasets
    HDF5.attributes(dset1)["time"] =  integrator.t
    HDF5.attributes(dset2)["time"] =  integrator.t

    #Progress bar and time evolution
    println("Starting up simulation...")
    prog1 = Progress( Int(floor(p.t_sim_final-p.t_sim_init))*100 )

    for (u,t) in tuples(integrator)

        #Increase iteration counter
        iter += 1

        #Step integrator
        DifferentialEquations.step!(integrator)

        #Update bar
        ProgressMeter.update!(prog1, Int(floor(integrator.t*100)))

        #If save iteration
        if iter % every == 0

            iter_save += 1
            save_time = integrator.t

            #data set names
            dset_name = string("iteration_",iter_save)

            #data set for both displacement and derivative
            data1 = integrator.u.x[1]
            dset1, dtype1  = create_dataset(g1, dset_name , data1 )
            write_dataset(dset1, dtype1 , data1)

            data2 = integrator.u.x[2]
            dset2, dtype2  = create_dataset(g2, dset_name , data2 )
            write_dataset(dset2, dtype2 , data2)

            #Add HDF5.attributes time to both datasets
            HDF5.attributes(dset1)["time"] =  integrator.t
            HDF5.attributes(dset2)["time"] =  integrator.t

        end

        #Check if time has exceeded
        if( time_since(init_time) > max_runtime )
            println("\nMaximum allowed runtime reached , Ending Simulation!")
            break
        end

    end

    #Add time information to file
    HDF5.attributes(fid)["max_iter"]    = iter_save
    HDF5.attributes(fid)["t_min"]       = p.t_sim_init
    HDF5.attributes(fid)["t_max"]       = save_time
    HDF5.attributes(fid)["out_every_t"] = p.out_every_t
    HDF5.attributes(fid)["dt"]          = p.deltat

    close(fid)
    println("Saving Files...")

    print_exit_message(  time() - init_time )

    return true

end


text:
@testset "Testing kalman.jl" begin

    # set up
    A = [.95 0; 0. .95]
    Q = eye(2) .* 0.5
    G = eye(2) .* .5
    R = eye(2) .* 0.2
    kf = Kalman(A, G, Q, R)

    rough_kwargs = Dict(:atol => 1e-2, :rtol => 1e-4)

    sig_inf, kal_gain = stationary_values(kf)

    # Compute the Kalman gain and sigma infinity according to the
    # recursive equations and compare
    mat_inv = inv(G * sig_inf * G' + R)

    kal_recursion = A * sig_inf * (G') * mat_inv
    sig_recursion = A * sig_inf * A' - kal_recursion * G * sig_inf * A' + Q
    # test stationary
    @test isapprox(sig_inf, sig_recursion; rough_kwargs...)
    @test isapprox(kal_gain, kal_recursion; rough_kwargs...)

    # test_update_using_stationary
    set_state!(kf, [0.0 0.0]', sig_inf)
    update!(kf, [0.0 0.0]')
    @test isapprox(kf.cur_sigma, sig_inf; rough_kwargs...)
    @test isapprox(kf.cur_x_hat, [0.0 0.0]'; rough_kwargs...)

    # test update nonstationary
    curr_x, curr_sigma = ones(2, 1), eye(2) .* .75
    y_observed = fill(0.75, 2, 1)
    set_state!(kf, curr_x, curr_sigma)
    update!(kf, y_observed)
    mat_inv = inv(G * curr_sigma * G' + R)
    curr_k = A * curr_sigma * (G') * mat_inv
    new_sigma = A * curr_sigma * A' - curr_k * G * curr_sigma * A' + Q
    new_xhat = A * curr_x + curr_k * (y_observed - G * curr_x)
    @test isapprox(kf.cur_sigma, new_sigma; rough_kwargs...)
    @test isapprox(kf.cur_x_hat, new_xhat; rough_kwargs...)

end  # @testset


text:
#Look up Table
include("lut/mt.jl")

"""
    tetIx(tIx, vals, iso::Real)

Determines which case in the triangle table we are dealing with
"""
@inline function tetIx(tIx, vals, iso::Real)
    v1 = vals[subTets[tIx][1]]
    v2 = vals[subTets[tIx][2]]
    v3 = vals[subTets[tIx][3]]
    v4 = vals[subTets[tIx][4]]
    ix = v1 < iso ? 0x01 : 0x00
    (v2 < iso) && (ix |= 0x02)
    (v3 < iso) && (ix |= 0x04)
    (v4 < iso) && (ix |= 0x08)
    ix
end

"""
    vertId(e, x, y, z, nx, ny)

Determines a unique integer ID associated with the edge. This is used
as a key in the vertex dictionary. It needs to be both unambiguous (no
two edges get the same index) and unique (every edge gets the same ID
regardless of which of its neighboring voxels is asking for it) in order
for vertex sharing to be implemented properly.
"""
@inline function vertId(e, x, y, z, nx, ny)
    dx, dy, dz = voxCrnrPosInt[voxEdgeCrnrs[e][1]]
    voxEdgeDir[e]+7*(x-0x01+dx+nx*(y-0x01+dy+ny*(z-0x01+dz)))
end

"""
    vertPos(e, x, y, z, vals::NTuple{8,T}, iso::Real, eps::Real, ::Type{VertType}) where {T<:Real, VertType}

Assuming an edge crossing, determines the point in space at which it
occurs.
eps represents the "bump" factor to keep vertices away from voxel
corners (thereby preventing degeneracies).
"""
@inline function vertPos(e, x, y, z, vals::NTuple{8,T}, iso::Real, eps::Real, ::Type{VertType}) where {T<:Real, VertType}

    ixs     = voxEdgeCrnrs[e]
    srcVal  = vals[ixs[1]]
    tgtVal  = vals[ixs[2]]
    a       = min(max((iso-srcVal)/(tgtVal-srcVal), eps), one(T)-eps)
    b       = one(T)-a
    c1 = voxCrnrPos(VertType)[ixs[1]]
    c2 = voxCrnrPos(VertType)[ixs[2]]

    VertType(x,y,z) + c1 .* b + c2.* a
end

"""
    getVertId(e, x, y, z, nx, ny,
                vals, iso::Real,
                vts::Dict,
                vtsAry::Vector,
                eps::Real)

Gets the vertex ID, adding it to the vertex dictionary if not already
present.
"""
@inline function getVertId(e, x, y, z, nx, ny,
                           vals, iso::Real,
                           vts::Dict,
                           vtsAry::Vector,
                           eps::Real)

    VertType = eltype(vtsAry)
    vId = vertId(e, x, y, z, nx, ny)
    # TODO we can probably immediately construct the vertex array here and use vert id to map to sequential ordering
    if !haskey(vts, vId)
        v = vertPos(e, x, y, z, vals, iso, eps, VertType)
        push!(vtsAry, v)
        vts[vId] = length(vtsAry)
    end
    vts[vId]
end

"""
    voxEdgeId(subTetIx, tetEdgeIx)
Given a sub-tetrahedron case and a tetrahedron edge ID, determines the
corresponding voxel edge ID.
"""
@inline function voxEdgeId(subTetIx, tetEdgeIx)
    srcVoxCrnr = subTets[subTetIx][tetEdgeCrnrs[tetEdgeIx][1]]
    tgtVoxCrnr = subTets[subTetIx][tetEdgeCrnrs[tetEdgeIx][2]]
    return voxEdgeIx[srcVoxCrnr][tgtVoxCrnr]
end

"""
    procVox(vals, iso::Real, x, y, z, nx, ny,
                    vts::Dict, vtsAry::Vector, fcs::Vector,
                    eps::Real)

Processes a voxel, adding any new vertices and faces to the given
containers as necessary.
"""
function procVox(vals, iso::Real, x, y, z, nx, ny,
                 vts::Dict, vtsAry::Vector, fcs::Vector,
                 eps::Real)
    VertType = eltype(vtsAry)
    FaceType = eltype(fcs)
    # check each sub-tetrahedron in the voxel
    @inbounds for i = 1:6
        tIx = tetIx(i, vals, iso)
        (tIx == 0x00 || tIx == 0x0f) && continue

        e = tetTri[tIx]

        # add the face to the list
        push!(fcs, FaceType(
                    getVertId(voxEdgeId(i, e[1]), x, y, z, nx, ny, vals, iso, vts, vtsAry, eps),
                    getVertId(voxEdgeId(i, e[2]), x, y, z, nx, ny, vals, iso, vts, vtsAry, eps),
                    getVertId(voxEdgeId(i, e[3]), x, y, z, nx, ny, vals, iso, vts, vtsAry, eps)))

        # bail if there are no more faces
        iszero(e[4]) && continue
        push!(fcs, FaceType(
                    getVertId(voxEdgeId(i, e[4]), x, y, z, nx, ny, vals, iso, vts, vtsAry, eps),
                    getVertId(voxEdgeId(i, e[5]), x, y, z, nx, ny, vals, iso, vts, vtsAry, eps),
                    getVertId(voxEdgeId(i, e[6]), x, y, z, nx, ny, vals, iso, vts, vtsAry, eps)))
    end
end

"""
Given a 3D array and an isovalue, extracts a mesh represention of the
an approximate isosurface by the method of marching tetrahedra.
"""
function isosurface(sdf::AbstractArray{T, 3}, method::MarchingTetrahedra, ::Type{VertType}=SVector{3,Float32}, ::Type{FaceType}=SVector{3, Int};
                    origin=VertType(-1,-1,-1), widths=VertType(2,2,2)) where {T, VertType, FaceType}
    vts        = Dict{Int, Int}()
    fcs        = FaceType[]
    vtsAry = VertType[]
    sizehint!(vts, div(length(sdf),8))
    sizehint!(vtsAry, div(length(sdf),8))
    sizehint!(fcs, div(length(sdf),4))
    # process each voxel
    nx::Int, ny::Int, nz::Int = size(sdf)
    @inbounds for k = 1:nz-1, j = 1:ny-1, i= 1:nx-1
        vals = (sdf[i, j, k],
                sdf[i, j+1, k],
                sdf[i+1, j+1, k],
                sdf[i+1, j, k],
                sdf[i, j, k+1],
                sdf[i, j+1, k+1],
                sdf[i+1, j+1, k+1],
                sdf[i+1, j, k+1])
        cubeindex = _get_cubeindex(vals,method.iso)
        if cubeindex != 0x00 && cubeindex != 0xff
            procVox(vals, method.iso, i, j, k, nx, ny, vts, vtsAry, fcs, method.eps)
        end
    end

    _correct_vertices!(vtsAry, size(sdf), origin, widths, VertType)

    vtsAry,fcs
end

"""
    _correct_vertices!(vts, size, origin, widths, VertType)

The marchingTetrahedra function returns vertices on the (1-based) indices of the
SDF's data, ignoring its actual bounds. This function adjusts the vertices in
place so that they correspond to points within the SDF bounds.
"""
function _correct_vertices!(vts, size, origin, widths, ::Type{VertType}) where {VertType}
    s = widths ./ VertType(size .- 1)  # subtract 1 because an SDF with N points per side has N-1 cells
    for i in eachindex(vts)
        vts[i] = (vts[i] .- 1) .* s .+ origin  # subtract 1 to fix 1-indexing
    end
end


text:
export LaplaceApproximation, laplace

laplace() = LaplaceApproximation()

struct LaplaceApproximation <: AbstractApproximationMethod end

approximation_name(::LaplaceApproximation)       = "LaplaceApproximation"
approximation_short_name(::LaplaceApproximation) = "LP"

using ForwardDiff
using Optim

function getweights(::LaplaceApproximation, mean::AbstractVector{T}, covariance::AbstractMatrix{T}) where { T <: Real }
    return getweights(srcubature(), mean, covariance)
end

function getpoints(::LaplaceApproximation, mean::AbstractVector{T}, covariance::AbstractMatrix{T}) where { T <: Real }
    return getpoints(srcubature(), mean, covariance)
end

function approximate_meancov(::LaplaceApproximation, g::Function, distribution)

    logg = (z) -> log(g(z))
    logd = (z) -> logpdf(distribution, z)

    logf   = (z) -> logg(z) + logd(z)
    d_logf = (z) -> ForwardDiff.gradient(logf, z)

    result = optimize((d) -> -(logf(d)), mean(distribution), LBFGS())
    if !Optim.converged(result)
        error("LaplaceApproximation: convergence failed")
    end

    m = Optim.minimizer(result)
    c = -cholinv(ForwardDiff.hessian(logf, m))

    return m, c
end

function approximate_kernel_expectation(::LaplaceApproximation, g::Function, distribution)
    return approximate_kernel_expectation(srcubature(), g, distribution)
end


text:
#=
Popis
  Funkcia zisti ci graf obsahuje paralelne hrany.

Parametre
  g : datova struktura grafu :: GenericGraph

Navratova hodnota
  : parameter reprezentujuci najdenie paralelnej hrany(True=najdena/False=nenajdena) :: Bool

Poznamky
  Funkcia prejde cez vsetky hrany grafu a zisti, ci sa v nom nachadzaju
  paralelne hrany. Funkcia je ukoncena ihned ako najde prvy vyskyt paralelnosti hran pricom
  vrati hodnotu True. Ak paralelne hrany neboli najdene navratova hodnota je False.
=#

function contains_parallel_edges(g::GenericGraph)

  parallel_check = Dict{Int64, Array{Tuple{Int64, Int64},1}}()
  parallel_check = [v.index => Tuple{Int64, Int64}[] for v=vertices(g)]

  for e in edges(g)

    if in((target(e).index, source(e).index), parallel_check[target(e).index]) ||
       in((source(e).index, target(e).index), parallel_check[source(e).index])

      return true

    else

      push!(parallel_check[source(e).index], (source(e).index, target(e).index))

    end

  end

  return false

end

max_stars_repo_name: mpuk/GraphAnalysis

max_issues_repo_name: mpuk/GraphAnalysis

max_forks_repo_name: mpuk/GraphAnalysis


text:
"""
The data assimilation system (4D Var)
Load:
    x_a_init.txt
    y_o.txt
Save:
    x_b_4d.txt
    x_a_4d.txt
"""

include("setting.jl");
using DelimitedFiles;
using LinearAlgebra;
using OrdinaryDiffEq;
using Optim, LineSearches;

# load initial condition
x_a_init = readdlm("data/x_a_init.txt");

# load observations
y_o_save = readdlm("data/y_o.txt");
y_o_time = y_o_save[:, 1];
y_o_data = y_o_save[:, 2:end];

# initial x_a: from x_a_init
x_a_save = fill(NaN, nT + 1, N)
x_a_save[1, :] = x_a_init';

# initial x_b: no values at the initial time (assign NaN)
x_b_save = fill(NaN, nT + 1, N)

# initial background error variance
σ = 0.2
R = convert(Matrix, Diagonal(repeat([σ^2], N)));
B_init = R
B_init = readdlm("data/b_nmc.txt");

# observation operator
function observation_operators(y_o::Matrix{Float64})::Vector{Matrix{Float64}}
    len = size(y_o, 1)
    hs = Vector{Matrix{Float64}}()
    for i = 1:len
        row = map(x -> isnan(x) ? 0.0 : 1.0, y_o[i, :])
        push!(hs, Diagonal(row))
    end
    hs
end

H = observation_operators(y_o_data)

# cost function
function J(
    x✶::Vector{Float64},
    x_b::Vector{Float64},
    B✶::Matrix{Float64},
    ts::Vector{Float64},
    y_o::Matrix{Float64},
    H::Vector{Matrix{Float64}},
    R::Vector{Matrix{Float64}})::Float64
    @assert length(H) == length(R)
    @assert length(R) == size(y_o, 1)
    n = length(R)
    xs, _ = traj(x✶, ts)
    db = x_b - x✶
    J1 = db' * inv(B✶) * db
    @assert all(!, isnan.(J1))
    J2 = 0.0
    for i = 1:n
        @inbounds dy = replace(y -> isnan(y) ? 0.0 : y, y_o[i, :]) - H[i] * xs[i]
        @inbounds J2 += dy' * inv(R[i]) * dy
        @assert all(!, isnan.(J2)) "j2 nan in loop $i"
    end
    0.5 * J1 + 0.5 * J2
end

# gradient of cost function, s[:] notation is mandatory
function J_grad!(
    s::Vector{Float64},
    x✶::Vector{Float64},
    x_b::Vector{Float64},
    B✶::Matrix{Float64},
    ts::Vector{Float64},
    y_o::Matrix{Float64},
    H::Vector{Matrix{Float64}},
    R::Vector{Matrix{Float64}})
    n = length(R)
    xs, M = traj(x✶, ts)
    Jg1 = inv(B✶) * (x✶ - x_b)
    @assert all(!, isnan.(Jg1))

    Jg2 = zeros(size(Jg1))
    for i = 1:n
        @inbounds dgy = replace(y -> isnan(y) ? 0.0 : y, y_o[i, :]) - H[i] * xs[i]
        @inbounds Jg2 += transpose(M[i]) * transpose(H[i]) * inv(R[i]) * dgy
        @assert all(!, isnan.(Jg2))
    end
    @inbounds s[:] = Jg1 - Jg2
end

# jacobian of lorenz96
function jac_lorenz(x::Vector{Float64})::Matrix{Float64}
    # jac[i, j] = -1 when j = i
    # jac[i, j] = x[i-1] when j = i+1
    # jac[i, j] = x[i+1] - x[i-2] when j = i-1
    # jac[i, j] = -x[i-1] when j = i-2
    len = length(x)
    lenm = len - 1
    jac = zeros(len, len)

    for i = 1:len
        ip1 = (i + 1 + lenm) % len + 1
        im1 = (i - 1 + lenm) % len + 1
        im2 = (i - 2 + lenm) % len + 1
        @inbounds jac[i, i] = -1
        @inbounds jac[i, ip1] = x[im1]
        @inbounds jac[i, im1] = x[ip1] - x[im2]
        @inbounds jac[i, im2] = -x[im1]
    end
    jac
end

function traj(x✶::Vector{Float64}, times::Vector{Float64})::Tuple{Vector{Vector},Vector{Matrix{Float64}}}
    # some missing time step
    # assumption of constant dU might fail
    nw = length(times) - 1
    nt = length(times)
    # this assertion should not exist
    @assert nw != 0
    ms = Vector{Matrix}(undef, nt)
    @inbounds ms[1] = Matrix(1.0I, N, N)
    xs = Vector{Vector}(undef, nt)
    @inbounds xs[1] = x✶
    dU = diff(times)
    fine_step = 100
    xt = x✶

    for i = 1:nw
        @inbounds ms[i+1] = ms[i]
        if dU[i] > 1e-12
            @inbounds ddU = dU[i] / fine_step
            for _ = 1:fine_step
                L = I + ddU * jac_lorenz(xt)
                xt += ddU * lorenz96(xt, Missing, Missing)
                @inbounds ms[i+1] = L * ms[i+1]
            end
        end
        @inbounds xs[i+1] = xt
    end

    xs[2:end], ms[2:end]
end

tbias = 1.0 * dT
ibias = convert(Integer, cld(tbias, dT))
ua = 1
ue = length(y_o_time)

for tt = 1:nT-ibias
    tts = tt - 1
    tti = tt + 1
    Ts = tts * dT
    Ta = tt * dT
    Wa = Ta + tbias
    # assimilation window from Ts to Ta
    global us = ua + 1
    for i = us:ue
        global ua = i
        if abs(y_o_time[i] - Wa) < dT / 100.0
            break
        elseif Wa < y_o_time[i]
            break
        end
    end
    println("Cycle = ", tt, ", Ts = ", round(Ts; digits = 3), ", Ta = ", round(Ta; digits = 3))

    #--------------
    # forecast step
    #--------------
    prob = ODEProblem(lorenz96, x_a_save[tt, :], (Ts, Ta))
    sol = solve(prob, Tsit5())
    x_b_save[tti, :] = sol[end]

    #--------------
    # analysis step
    #--------------

    # background
    x_b = sol[end]
    x_b_time = Ta

    # observation
    y_o = y_o_data[us:ua, :]
    h_o = H[us:ua]
    r_o = [repeat(R, 1, 1, ua - us + 1)[:, :, i] for i in 1:(ua-us+1)]

    # analysis scheme (4dvar)
    traj_time = [x_b_time; y_o_time[us:ua]]

    opt = optimize(x -> J(x, x_b, B_init, traj_time, y_o, h_o, r_o),
        (s, x) -> J_grad!(s, x, x_b, B_init, traj_time, y_o, h_o, r_o),
        x_b,
        method = ConjugateGradient(),
        g_tol = 1e-12,
        iterations = 100
    )
    # @show opt
    @assert Optim.converged(opt)

    x_a = opt.minimizer

    x_a_save[tti, :] = x_a
end # end tt loop

open("data/x_a_4d.txt", "w") do xa
    writedlm(xa, x_a_save)
end

open("data/x_b_4d.txt", "w") do xb
    writedlm(xb, x_b_save)
end

max_stars_repo_name: Typas/Data-Assimilation-Project

max_issues_repo_name: Typas/Data-Assimilation-Project

max_forks_repo_name: Typas/Data-Assimilation-Project