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


nrange(::Model"ising_extended") = 3
# H = -J Σⱼ XⱼXⱼ₊₁ - h Σⱼ Zⱼ - J₂ Σⱼ XⱼZⱼ₊₁Xⱼ₊₂
function ITensors.MPO(::Model"ising_extended", s; J=1.0, h=1.0, J₂=0.0)
  N = length(s)
  a = OpSum()
  if J != 0
    for n in 1:(N - 1)
      a += -J, "X", n, "X", n + 1
    end
  end
  if J₂ != 0
    for n in 1:(N - 2)
      a += -J₂, "X", n, "Z", n + 1, "X", n + 2
    end
  end
  if h != 0
    for n in 1:N
      a += -h, "Z", n
    end
  end
  return splitblocks(linkinds, MPO(a, s))
end

# H = -J Σⱼ XⱼXⱼ₊₁ - h Σⱼ Zⱼ - J₂ Σⱼ XⱼZⱼ₊₁Xⱼ₊₂
function ITensors.OpSum(::Model"ising_extended", n1, n2; J=1.0, h=1.0, J₂=0.0)
  opsum = OpSum()
  if J != 0
    opsum += -J / 2, "X", n1, "X", n2
    opsum += -J / 2, "X", n2, "X", n2 + 1
  end
  if J₂ != 0
    opsum += -J₂, "X", n1, "Z", n2, "X", n2 + 1
  end
  opsum += -h / 3, "Z", n1
  opsum += -h / 3, "Z", n2
  opsum += -h / 3, "Z", n2 + 1
  return opsum
end

function ITensors.ITensor(::Model"ising_extended", s1, s2, s3; J=1.0, h=1.0, J₂=0.0)
  opsum = OpSum(Model"ising_extended"(), 1, 2; J=J, h=h, J₂=J₂)
  return prod(MPO(opsum, [s1, s2, s3]))
end

function reference(::Model"ising_extended", ::Observable"energy"; J=1.0, h=1.0, J₂=0.0)
  f(k) = sqrt((J * cos(k) + J₂ * cos(2k) - h)^2 + (J * sin(k) + J₂ * sin(2k))^2)
  return -1 / 2π * ITensorInfiniteMPS.∫(k -> f(k), -π, π)
end


module CryptoGroups

include("utils.jl")

function spec end
function specialize end
function order end
function bitlength end
function <| end
function tobits end

function validate end # isvalid
function value end
function modulus end

function cofactor end

function a end
function b end

# In most cases probably using sizeof directly is more appropriate
bitlength(::Type{T}) where T <: Integer = sizeof(T) * 8

function bitlength(p::Integer)

    bits = bitstring(p)
    start = findfirst(x -> x == '1', bits)
    N = length(bits) - start + 1

    return N
end

function bitlength(p::BigInt)

    bytes = int2bytes(p)
    bits = bitstring(bytes[end])
    start = findfirst(x -> x == '1', bits)
    N = length(bytes) * 8  - (start - 1)

    return N
end

include("Fields/Fields.jl")
include("Curves/Curves.jl")
include("Specs/Specs.jl")

include("groups.jl")
include("elgamal.jl")
include("spec.jl")

export @bin_str

export spec, specialize, order, bitlength

import .Fields: Field, BinaryField, FP, F2GNB, F2PB, PrimeField
export Field, FP, F2GNB, F2PB # frombits and tobits method shall be dealt with convert

import .Curves: AbstractPoint, ECPoint, AffinePoint, BinaryCurve, validate, gx, gy, a, b, oncurve
export AbstractPoint, ECPoint, AffinePoint, BinaryCurve, validate #, gx, gy, a, b, oncurve

import .Specs: MODP, Koblitz, ECP, EC2N, PB, GNB, generator, Hash, PRG, RO, ROPRG
export MODP, Koblitz, ECP, EC2N, PB, GNB, generator, Hash, PRG, RO, ROPRG

export ElGamal

end


struct SupportSample
    s::Any
    x::Int
end

struct CategoricalWithSupport <: Distribution{SupportSample} end

"""
    categorical(probs::AbstractArray{U, 1}) where {U <: Real}

Given a vector of probabilities `probs` where `sum(probs) = 1`, and a vector of support, sample an element  `support[i]` from the support with probability `probs[i]`.
"""
const categorical_with_support = CategoricalWithSupport()

function logpdf(::CategoricalWithSupport, ss::SupportSample, probs::AbstractArray{U,1},support) where {U <: Real}
    log(probs[ss.x])
end

function logpdf_grad(::CategoricalWithSupport, ss::SupportSample, probs::AbstractArray{U,1},support)  where {U <: Real}
    grad = zeros(length(probs))
    grad[ss.x] = 1.0 / probs[ss.x]
    (nothing, grad)
end

function random(::CategoricalWithSupport, probs::AbstractArray{U,1},support) where {U <: Real}
    idx = rand(Distributions.Categorical(probs))

    SupportSample(support[idx],idx)
end

(::CategoricalWithSupport)(probs,support) = random(CategoricalWithSupport(), probs,support)

has_output_grad(::CategoricalWithSupport) = false
has_argument_grads(::CategoricalWithSupport) = (true,)

export categorical_with_support


struct Params_SV{ComplexArray,Total,KF}
    # boundary conditions
    periodic::Bool
    # normalization
    total::Total

    # algorithm-specific
    # gradient kernel
    kernelfactor::KF

    # fft buffers
    complex_surface::ComplexArray
    complex_void::ComplexArray
end

function Params_SV(img; periodic::Bool=true, kernelfactor=KernelFactors.sobel)
    box = size(img)
    complex_box = periodic ? box : box .* 2

    # scale factor
    total = cnt_total(img, periodic)

    p = Params_SV(
        periodic,
        total,
        kernelfactor,
        similar(img, ComplexF32, complex_box),
        similar(img, ComplexF32, complex_box)
    )
    cf_type = :full
    p, cf_type
end

"""
    correllation_function!(res, img, params::Params_SV)

Compute surface-void correlation function in positive directions
"""
function correllation_function!(res, img, params::Params_SV)
    ix = CartesianIndices(img)

    f = params.complex_surface .= 0
    g = params.complex_void .= 0
    v_f = view(f, ix)
    v_g = view(g, ix)
    v_f .= gradient_norm(img, params.kernelfactor)
    v_g .= img .== 0

    cross_correlation!(f, g)

    res .= real.(v_f) ./ params.total
end

"""
    surfvoid(image; periodic = false)

Calculate `Fsv(x)` (surface-void) correlation function map
for the N-dimensional image and return a `CFMap` object.

The `image` contains the probability of the voxel being in the correct phase.

# Examples
```jldoctest
julia> surfvoid([1 0; 0 1]; periodic=true).result
3×3 Matrix{Float32}:
 0.176777  0.176777  0.176777
 0.176777  0.176777  0.176777
 0.176777  0.176777  0.176777
```
"""
function surfvoid(image; periodic::Bool=false, kernelfactor=KernelFactors.sobel)
    corr_function_map(image, Params_SV; periodic, kernelfactor)
end

function surfvoid(image, phase::Int; periodic::Bool=false, kernelfactor=KernelFactors.sobel)
    one_phase_img = image .== phase
    corr_function_map(one_phase_img, Params_SV; periodic, kernelfactor)
end


# thanks @tduretz
module test_o2

# Test correct handling of second order points

using Triangulate

function test()
    triin=Triangulate.TriangulateIO()
    triin.pointlist=Cdouble[0.0 0.0 ; 1.0 0.0 ; 1.0  1.0; 0.0 1.0]'
    triin.segmentlist=Cint[1 2 ; 2 3 ; 3 4 ; 4 1 ]'
    triin.segmentmarkerlist=Cint[1, 2, 3, 4]

    # Test P1 element creation
    (triout, vorout)=triangulate("pqQ", triin)
    test1 = triout.pointlist==[0.0  1.0  1.0  0.0;
                               0.0  0.0  1.0  1.0]

    test1 = test1 && triout.trianglelist==[4  2;
                                           1  3;
                                           2  4]

    # Test P2 element creation
    (triout, vorout)=triangulate("pqo2Q", triin)
    test2 = triout.pointlist==[0.0  1.0  1.0  0.0  0.0  0.5  0.5  1.0  0.5;
                               0.0  0.0  1.0  1.0  0.5  0.0  0.5  0.5  1.0]

    test2 = test2 && triout.trianglelist==[4  2;
                                           1  3;
                                           2  4;
                                           6  9;
                                           7  7;
                                           5  8]

    test1 && test2
end

end


# general method given a reachability algorithm for the linear system
function reach_CARLIN_alg(X0, F1, F2; alg, resets, N, T, Δt0, bloat, compress)
    if resets == 0
        reach_CARLIN(X0, F1, F2; alg, N, T, Δt=Δt0, bloat, compress)
    else
        reach_CARLIN_resets(X0, F1, F2, resets; alg, N, T, Δt=Δt0, bloat, compress)
    end
end

function reach_CARLIN(X0, F1, F2; alg, N, T, Δt, bloat, compress, A=nothing)
    error_bound_func = error_bound_specabs

    # lift initial states
    if compress
        ŷ0 = lift_vector(X0, N) # see CarlemanLinearization/linearization.jl
    else
        ŷ0 = kron_pow_stack(X0, N) |> box_approximation
    end

    # solve continuous ODE
    if isnothing(A)
        A = build_matrix(F1, F2, N; compress)
    end
    prob = @ivp(ŷ' = Aŷ, ŷ(0) ∈ ŷ0)
    sol = solve(prob, T=T, alg=alg, Δt0=Δt)

    # projection onto the first n variables
    n = dim(X0)
    πsol_1n = _project(sol, 1:n)

    if !bloat
        return πsol_1n
    end

    # compute errors
    errfunc = error_bound_func(X0, Matrix(F1), Matrix(F2), N=N)

    # evaluate error bounds for each reach-set in the solution
    E = [errfunc.(tspan(R)) for R in sol]

    # if the interval is always > 0 then we can just take max(Ei)

    # symmetrize intervals
    E_rad = [symmetric_interval_hull(Interval(ei)) for ei in E]
    E_ball = [BallInf(zeros(n), max(ei)) for ei in E_rad]

    # sum the solution with the error
    fp_bloated = [ReachSet(set(Ri) ⊕ Ei, tspan(Ri)) for (Ri, Ei) in zip(πsol_1n, E_ball)] |> Flowpipe

    return fp_bloated
end

function _compute_resets(resets::Int, T)
    return mince(0 .. T, resets+1)
end

function _compute_resets(resets::Vector{Float64}, T)
    # assumes initial time is 0
    aux = vcat(0.0, resets, T)
    return [interval(aux[i], aux[i+1]) for i in 1:length(aux)-1]
end

function reach_CARLIN_resets(X0, F1, F2, resets; alg, N, T, Δt, bloat, compress)

    # build state matrix (remains unchanged upon resets)
    A = build_matrix(F1, F2, N; compress)

    # time intervals to compute
    time_intervals = _compute_resets(resets, T)

    # compute until first chunk
    T1 = sup(first(time_intervals))
    sol_1 = reach_CARLIN(X0, F1, F2; alg, N, T=T1, Δt=interval(0)+Δt, bloat, compress, A=A)

    # preallocate output flowpipe
    fp_1 = flowpipe(sol_1)
    out = Vector{typeof(fp_1)}()
    push!(out, fp_1)

    # approximate final reach-set
    Rlast = sol_1[end]
    X0 = box_approximation(set(Rlast))

    # compute remaining chunks
    for i in 2:length(time_intervals)
        T0 = T1
        Ti = sup(time_intervals[i])
        sol_i = reach_CARLIN(X0, F1, F2; alg, N, T=Ti-T0, Δt=interval(T0)+Δt, bloat, compress, A=A)
        push!(out, flowpipe(sol_i))
        X0 = box_approximation(set(sol_i[end]))
        T1 = Ti
    end

    return MixedFlowpipe(out)
end


function samplepoints(streamlinelength, λ_min , λ_max, n)
    x0 = rand() * 1000.0m
    xs = range(x0, x0 + streamlinelength, length = n)
    no = noise_between_wavelengths(λ_min, λ_max, xs, normalize = false)
end

"""
    z_transform_contributions(vs::T, z::Complex)
        -> Complex
Unilateral z-transform
where
    vs are the samples
    z is a point in the complex plane

"""
function z_transform_contributions(vs::T, z::Complex) where {T<:Union{AbstractRange, Vector}}
    halfn = div(length(vs), 2)
    [s * z^-(i - halfn) for (i, s) in enumerate(vs)]
end
z_transform_contributions_points(vs, z) = [EM * Point(real(ŝ), imag(ŝ)) for ŝ in z_transform_contributions(vs, z)]


export mult_gauss, mult_gauss!, mult_gauss_chn, mult_gauss_chn!

@inline function comb_mult_gauss(x, n)
    return x * n
end

mult_gauss!(X, σ=0.1, μ=1; clip=false) = apply_noise!(comb_mult_gauss, f_awg(σ, μ, complex_copy(σ), complex_copy(μ)), X, clip)
mult_gauss(X, σ=0.1, μ=1; clip=false) = mult_gauss!(copy(X), σ, μ, clip=clip)

"""
    mult_gauss(X; clip=false[, σ=0.1, μ=1])

Returns the array `X` with the array value multiplied with a gauss distribution (standard deviation `σ` and mean `μ`) .
`σ` and `μ` are optional arguments representing standard deviation and mean of gauss.
If keyword argument `clip` is provided the values are clipped to be in [0, 1].
If `X` is a RGB{Normed} or Gray{Normed} image, then the values will be automatically clipped and the keyword
`clip` is meaningless.

If `X<:Complex`, `μ` and `σ` are applied to the imaginary in the same way as for the real part.
If you want to have different behaviour for real and imaginary part, simply
choose `μ` or `σ` complex.
"""
mult_gauss

mult_gauss_chn(X, σ=0.1, μ=1; clip=false) = mult_gauss_chn!(copy(X), σ, μ, clip=clip)
mult_gauss_chn!(X, σ=0.1, μ=1; clip=false) = apply_noise_chn!(comb_mult_gauss, f_awg(σ, μ, complex_copy(σ), complex_copy(μ)), X, clip)

"""
    mult_gauss_chn(X; clip=false[, σ=0.1, μ=0.0])

Returns the RGB image `X` with the values of the pixel multiplied with a
gauss distribution (standard deviation `σ` and mean `μ`) pixelwise.
However, every channel of one pixel receives the same amount of noise.
The noise therefore acts roughly as intensity - but not color - changing noise.
If keyword argument `clip` is provided the values are clipped to be in [0, 1].
`σ` and `μ` are optional arguments representing standard deviation and mean of gauss.

If `X<:Complex`, `μ` and `σ` are applied to the imaginary in the same way as for the real part.
If you want to have different behaviour for real and imaginary part, simply
choose `μ` or `σ` complex.
"""
mult_gauss_chn


const R  = 6372.1  # Radius of the Earth [km]

function haversine_distance(lat1, lon1, lat2, lon2)
    d = 0
    # Your code goes here!
    return d
end


"""
*MultipleTesting* package

* Adjusting p-values for multiple testing
* Estimating the fraction π₀ of tests under the null hypothesis
* Combination of p-values
* Higher criticism
"""
module MultipleTesting

using StatsBase
import StatsBase: fit

using Distributions
import Distributions: estimate

import SpecialFunctions: digamma

export PValues,
    adjust,
    PValueAdjustment,
    Bonferroni,
    BenjaminiHochberg,
    BenjaminiHochbergAdaptive,
    BenjaminiYekutieli,
    BenjaminiLiu,
    Hochberg,
    Holm,
    Hommel,
    Sidak,
    ForwardStop,
    BarberCandes,
    estimate,
    Pi0Estimator,
    Storey,
    StoreyBootstrap,
    LeastSlope,
    Oracle,
    TwoStep,
    RightBoundary,
    CensoredBUM,
    CensoredBUMFit,
    BUM,
    BUMFit,
    FlatGrenander,
    ConvexDecreasing,
    ConvexDecreasingFit,
    isin,
    fit,
    BetaUniformMixtureModel,
    PValueCombination,
    combine,
    Fisher,
    Logit,
    Stouffer,
    Tippett,
    Simes,
    Wilkinson,
    Minimum,
    HigherCriticismScores,
    HigherCriticismThreshold

include("types.jl")
include("utils.jl")
include("pval-adjustment.jl")
include("pi0-estimators.jl")
include("model.jl")
include("combinations.jl")
include("higher-criticism.jl")

end


function likelihood(k, model::AbstractBinomialModel, observation)
    return mean(
                exp.(-0.5f0 .* ((observation .- model.q .* k) ./ model.σ).^2)
                ./ (sqrt(2*Float32(pi)) .* model.σ)
           , dims = 2
           )[:,1]
end

function kernel_likelihood_indices!(u, v, idxT, kT, q, σ, observation, r, randstates)
    # use column-index first (gives a 3x speedup)
    # kT stands for the transpose of k
    # idxT stands for the transpose of idx
    j = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    M_in, M_out = size(kT)
    @inbounds if j <= M_out
        vsum = 0f0
        CurMax = 1f0
        for i in 1:M_in
            # omitting normalization constant here; it is only needed for u
            vi      = CUDA.exp(-0.5f0 *((observation - q[j] * kT[i,j]) / σ[j])^2)
            vsum   += vi
            v[i, j] = vsum
            # sample descending sequence of sorted random numbers
            # r[i,M_in] >= ... >= r[i,2] >= r[i,1]
            # Algorithm by:
            # Bentley & Saxe, ACM Transactions on Mathematical Software, Vol 6, No 3
            # September 1980, Pages 359--364
            mirrori = M_in - i + 1 # mirrored index i
            CurMax *= CUDA.exp(CUDA.log(GPUArrays.gpu_rand(Float32, CUDA.CuKernelContext(), randstates)) / mirrori)
            r[mirrori, j] = CurMax
        end
        # compute average likelihood across inner particles
        # (with normalization constant that was omitted from v for speed)
        u[j] = vsum / (M_in * CUDA.sqrt(2*Float32(pi)) * σ[j])
        # O(n) binning algorithm for sorted samples
        bindex = 1 # bin index
        @inbounds for i in 1:M_in
            # scale random numbers (this is equivalent to normalizing v)
            rsample = r[i, j] * vsum
            # checking bindex <= M_in - 1 is redundant since
            # v[M_in, j] = vsum
            while rsample > v[bindex, j]
                bindex += 1
            end
            idxT[i, j] = bindex
        end
    end
    return nothing
end

function likelihood_indices(k, model::AbstractBinomialModel, observation)
    M_out, M_in = size(k)
    r           = CuArray{Float32}(undef, M_in, M_out)
    u           = CuArray{Float32}(undef, M_out)
    v           = CuArray{Float32}(undef, M_in, M_out)
    idx         = CuArray{Int}(undef, M_out, M_in)

    rng = GPUArrays.default_rng(CuArray)

    kernel  = @cuda launch=false kernel_likelihood_indices!(
                u, v,
                idx', k',
                model.q, model.σ,
                Float32(observation),
                r,
                rng.state
              )
    config  = launch_configuration(kernel.fun)
    threads = Base.min(M_out, config.threads, 256)
    blocks  = cld(M_out, threads)
    kernel(
        u, v,
        idx', k', # pass transposes for column-major indexing (gives a 3x speedup)
        model.q, model.σ,
        observation,
        r,
        rng.state
        ;
        threads=threads, blocks=blocks
    )
    return u, idx
end

function inner_resample_helper!(in, out, idx)
    function kernel(in, out, idx)
        j = (blockIdx().x - 1) * blockDim().x + threadIdx().x
        @inbounds if j <= size(out, 2)
            @inbounds for i in 1:size(out, 1)
                out[i, j] = in[i, idx[i, j]]
            end
        end#if
        return nothing
    end
    kernel  = @cuda launch=false kernel(in, out, idx)
    config  = launch_configuration(kernel.fun)
    threads = Base.min(size(out, 2), config.threads, 256)
    blocks  = cld(size(out, 2), threads)
    kernel(in, out, idx; threads=threads, blocks=blocks)
    return out
end

function inner_resample_helper!(in, idx)
    out = similar(in)
    inner_resample_helper!(in, out, idx)
    in .= out
    return in
end

function likelihood_resample!(state::BinomialState, model, observation::BinomialObservation)
    u, idx = likelihood_indices(state.k, model, observation.EPSP)
    M_out  = size(state.n, 1)
    inner_resample_helper!(state.n, idx)
    inner_resample_helper!(state.k, idx)
    return u
end


"""
    getOverlap(MPSA, MPSB)

calculates the between to MPS encoded in MPSA and MPSB: < MPSB | MPSA >

# Arguments
- MPSA: vector containing the ket of the overlap
- MPSB: vector containing the bra of the overlap

return:

"""
function getOverlap(MPSA::Vector{<:Any}, MPSB::Vector{<:Any})
	Length = Int(size(MPSA)[1])

        @tensor T_multi[α_bra, α_ket, β_bra, β_ket] := (MPSA[1][α_ket, d, β_ket] *
							conj(MPSB[1][α_bra, d, β_bra]))

        for site = 2:Length
        	#println("add ",site)
        	@tensor begin

        		E_local[a_in',a_in,a_out',a_out] := (MPSA[site][a_in, d, a_out] *
        						     conj(MPSB[site][a_in', d, a_out']))

        		T_multi[α_bra, α_ket, β_bra, β_ket] := (T_multi[α_bra, α_ket, γ_bra, γ_ket] *
        							E_local[γ_bra, γ_ket, β_bra, β_ket])
        	end
        end

	return T_multi #(α_bra, α_ket, β_bra, β_ket)
end

"""
    getExpValue(MPSvec, MPOvec)

calculates the expectation value of an MPO : < MPSvec | MPOvec | MPSvec >

# Arguments
- MPSvec: vector of M's
- MPOvec: vector of MPO's

return:

"""
function getExpValue(MPSvec::Vector{<:Any}, MPOvec::Vector{MPO})
	OpString = stringMPO(MPOvec)
	expValue = 0
	for (i, x) in enumerate(OpString[1])
		test = applyTM_OP(MPSvec, MPSvec, x, MPOvec, Id(size(MPSvec[1])[1]))
		expValue += tr(test)
	end

	return expValue
end

"""
    singleSiteExpValue(MPSvec, C, Op, site; can_left = true)

calculates the expectation value of a single site Operator Op given
the local MPS-matrix MPSvec[site] and singular-values between the next(previous) bond
depending on the left(right) canonical condition of M specified by can_left.

# Arguments
- MPSvec: containing the local MPS-matrices
- C: singular values missing to bring M into mixed canonical form
- Op: local single site operator
- site: site on which Op should act
- can_left = true : is MPSvec[site] in left or right can. form.

return:

	expectation values of < MPSvec | Op_site | MPSvec >

"""
singleSiteExpValue(MPSvec::Vector{<:Any}, C::AbstractArray{<:Number},
		   Op::AbstractArray{<:Number}, site::Int; can_left=true) = singleSiteExpValue(MPSvec[site], C, Op, can_left = can_left)

function singleSiteExpValue(M::AbstractArray{<:T}, C::AbstractArray{<:T}, Op::AbstractArray{<:T}; can_left = true) where T<:Number
	if can_left
		@tensor tmp[β_ket,α_ket,d] := C[γ, β_ket] * M[α_ket, d, γ]
	else
		@tensor tmp[β_ket,α_ket,d] := C[α_ket, γ] * M[γ, d, β_ket]
	end

	@tensor tmp[d_ket,d_bra] := tmp[β,α,d_ket] * conj(tmp[β,α,d_bra])
	@tensor expValue[] := Op[d_bra,d_ket] * tmp[d_ket,d_bra]

	#expVlaue = 2.0
	return expValue
end

function getCoef(MPSvec)

	# to do get Coefficient vector
end


"""
Copyright Marcus Greiff 2020

Tests for the power distribution functionality
"""

# Include modules and recompile if running test standalone
if !isdefined(Main, :runAllTests)
    using Test, AerialVehicleControl, LinearAlgebra
    recompile();
end

# Testing fixture
function fixture()
    J  = rand(3,3); J  = J + J'; J  = J + I*maximum(abs.(eigvals(J)));
    kR, kc, kw = 1.0, 1.0, 1.0;
    a, b, c, d = 0.35, 0.26, 1.0+rand(), 1.0+rand();
    C  = con_state_qw_fsf_t(J, kR, kc, kw; a=a, b=b, c=c, d=d);
    C.thrust = 0.5;
    t = rand(3,1);
    C.torques = ntuple(i->Cdouble(t[i,1]), 3)
    Cc = deepcopy(C);
    d, dc, dm = get_matrix_and_copy(rand(4,1))
    return C, Cc, t, d, dc, dm;
end

@verbose(1, "Testing the power distribution...")
@testset "Test the power distribution" begin
    @testset "Good inputs" begin
        # Fixture
        C, Cc, t, d, dc, dm = fixture();

        apar = C.param_a;
        bpar = C.param_b;
        cpar = C.param_c;
        dpar = C.param_d;
        T = C.thrust;

        # Modified configuration for the razor drone
        A = [ +1.00  +1.00  +1.00  +1.00;
              -dpar  -dpar  +dpar  +dpar;
              -dpar  +dpar  +dpar  -dpar;
              +cpar  -cpar  +cpar  -cpar]

        #A = [+1    +1    +1    +1   ;
        #     +0    -dpar +0    +dpar;
        #     +dpar +0    -dpar +0   ;
        #     -cpar +cpar -cpar +cpar]

        f_julia = A \ ([T; t]);

        f_julia[f_julia .<0.0] .= 0.0

        d_julia = -(apar/(2.0*bpar)) .+ sqrt.((apar/(2*bpar)).^2 .+ f_julia./bpar);

        # Test
        status = ccall((:compute_power_distribution, AerialVehicleControl.CONT_LIB_PATH),
            Cint,
            (Ref{con_state_qw_fsf_t}, Ref{matrix_double_t}),
            C,
            dm)

        # Assert
        @test isapprox(status, 1);
        for i = 1:4
            @test isapprox(d_julia[i,1], d[i,1]);
        end
    end
    @testset "Bad inputs (negative thrust)" begin
        # Fixture
        C, Cc, t, d, dc, dm = fixture();
        C.thrust = -C.thrust;
        Cc = deepcopy(C);

        # Test
        status = ccall((:compute_power_distribution, AerialVehicleControl.CONT_LIB_PATH),
            Cint,
            (Ref{con_state_qw_fsf_t}, Ref{matrix_double_t}),
            C,
            dm)

        # Assert
        @test isapprox(status, 0);
    end
    @testset "Bad inputs (negative parameter value)" begin
        for i = 1:4
            # Fixture
            C, Cc, t, d, dc, dm = fixture();
            C.param_a = -C.param_a;
            Cc = deepcopy(C);

            # Test
            status = ccall((:compute_power_distribution, AerialVehicleControl.CONT_LIB_PATH),
                Cint,
                (Ref{con_state_qw_fsf_t}, Ref{matrix_double_t}),
                C,
                dm)

            # Assert
            @test isapprox(status, 0);
        end
    end
    @testset "Bad inputs (wrong duty cycle row dimensions)" begin
        # Fixture
        C, Cc, t, _, _, _ = fixture();
        d, dc, dm = get_matrix_and_copy(rand(3,1));

        # Test
        status = ccall((:compute_power_distribution, AerialVehicleControl.CONT_LIB_PATH),
            Cint,
            (Ref{con_state_qw_fsf_t}, Ref{matrix_double_t}),
            C,
            dm)
        # Assert
        @test isapprox(status, 0);
    end
    @testset "Bad inputs (wrong duty cycle column dimensions)" begin
        # Fixture
        C, Cc, t, _, _, _ = fixture();
        d, dc, dm = get_matrix_and_copy(rand(4,2));

        # Test
        status = ccall((:compute_power_distribution, AerialVehicleControl.CONT_LIB_PATH),
            Cint,
            (Ref{con_state_qw_fsf_t}, Ref{matrix_double_t}),
            C,
            dm)
        # Assert
        @test isapprox(status, 0);
    end
end


################################################################################
# Matrix containing the number of observations and the condition
# number per pixel

@doc raw"""
    mutable struct NobsMatrixElement

This structure encodes the condition matrix for a pixel in the map. It
is essentially matrix M_p in Eq. (10) of KurkiSuonio2009, but with
steroids, as it implements the following fields:

- `m`: the 3×3 matrix in Eq. (10)
- `invm`: the 3×3 inverse of `m`
- `invcond`: the inverse of the condition number *rcond* for `m`; this
  field is useful to check whether the samples in a TOD and the attack
  angles are sufficient to constrain a unique solution for I, Q, and U
  for that pixel.
- `neglected`: a Boolean flag telling whether the pixel was skipped or
  not during map-making (this usually happens if `invcond` is too
  small)

A typical use case is an array of `NobsMatrixElement` objects in an
array that get updated via [`update_nobs_matrix!`](@ref), like in the
following example:

```julia
NPIX = 10   # Number of pixels in the map
nobs_matrix = [NobsMatrixElement() for i in 1:NPIX]

# The variables psi_angle, sigma_squared, pixidx, and flagged
# must have been defined somewhere else; they contain the TOD
update_nobs_matrix!(nobs_matrix, psi_angle, sigma_squared, pixidx, flagged)
```

"""
mutable struct NobsMatrixElement{T <: Real}
    # We do not define "m" as symmetric, as we want to update its
    # fields one by one; however, we only consider the upper part
    m::Array{T, 2}
    invm::Symmetric{T, Array{T, 2}}
    invcond::T

    # If "true", do not consider this pixel in the solution, as it was
    # not covered enough by the scanning strategy
    neglected::Bool

    NobsMatrixElement{T}() where {T <: Real} = new(
        Symmetric(zeros(T, 3, 3)),
        Symmetric(zeros(T, 3, 3)),
        0.0,
        true,
    )
end

function Base.show(io::IO, nobsmatr::NobsMatrixElement{T}) where {T <: Real}
    print(io, "NobsMatrixElement(invcond => $(nobsmatr.invcond), neglected => $(nobsmatr.neglected))")
end

function Base.show(io::IO, ::MIME"text/plain", nobsmatr::NobsMatrixElement{T}) where {T <: Real}
    println(
        io,
        @sprintf(
            """NobsMatrixElement:
Inverse condition number.......... %e
Neglected......................... %s
Matrix:
    %.5f   %.5f   %.5f
    %.5f   %.5f   %.5f
    %.5f   %.5f   %.5f
Inverse matrix:
    %.5f   %.5f   %.5f
    %.5f   %.5f   %.5f
    %.5f   %.5f   %.5f
    """,
            nobsmatr.invcond,
            nobsmatr.neglected,
            nobsmatr.m[1, 1], nobsmatr.m[1, 2], nobsmatr.m[1, 3],
            nobsmatr.m[2, 1], nobsmatr.m[2, 2], nobsmatr.m[2, 3],
            nobsmatr.m[3, 1], nobsmatr.m[3, 2], nobsmatr.m[3, 3],
            nobsmatr.invm[1, 1], nobsmatr.invm[1, 2], nobsmatr.invm[1, 3],
            nobsmatr.invm[2, 1], nobsmatr.invm[2, 2], nobsmatr.invm[2, 3],
            nobsmatr.invm[3, 1], nobsmatr.invm[3, 2], nobsmatr.invm[3, 3],
        )
    )
end

@doc raw"""
    update_nobs!(nobs::NobsMatrixElement{T}; threshold = 1e-7)

This function makes sure that all the elements in `nobs` are
synchronized. It should be called whenever the field `nobs.m` (matrix
M_p in Eq. 10 of KurkiSuonio2009) has changed.

"""
function update_nobs!(nobs::NobsMatrixElement{T}; threshold = 1e-7) where {T <: Real}
    c = cond(nobs.m)

    if isfinite(c)
        nobs.invm = inv(Symmetric(nobs.m))
        nobs.invcond = 1 / c
        nobs.neglected = nobs.invcond < threshold
    else
        nobs.invcond = 0.0
        nobs.neglected = true
    end
end

@doc raw"""
    update_nobs_matrix!(nobs_matrix::Vector{NobsMatrixElement{T}}, psi_angle, sigma_squared, pixidx, flagged)

Apply Eq. (10) of KurkiSuonio2009 iteratively on the samples of a TOD
to update matrices ``M_p`` in `nobs_matrix`. The meaning of the
parameters is the following:

- `nobs_matrix` is the structure that gets updated by this function
- `psi_angle` is an array of `N` elements, containing the polarization
  angles (in radians)
- `sigma_squared` is an array of `N` elements, each containing the
  value of σ^2 for the samples in the TOD
- `pixidx` is an array of `N` elements, containing the pixel index
  observed by the TOD
- `flagged` is a Boolean array of `N` elements; `true` means that the
  sample in the TOD should not be used to produce the map. This can be
  used to produce jackknives and to neglect moving objects in the TOD

"""
function update_nobs_matrix!(
    nobs_matrix::Vector{NobsMatrixElement{T}},
    psi_angle,
    sigma_squared,
    pixidx,
    flagged;
    comm = nothing,
) where {T <: Real}

    @assert length(psi_angle) == length(sigma_squared)
    @assert length(psi_angle) == length(pixidx)
    @assert length(psi_angle) == length(flagged)

    for (idx, curpix, curpsi, curflagged) in zip(
        1:length(pixidx),
        pixidx,
        psi_angle,
        flagged,
    )
        curflagged && continue

        constant = 1 / sigma_squared[idx]

        sin_term, cos_term = sincos(2 * curpsi)
        nobs_matrix[curpix].m[1, 1] += constant
        nobs_matrix[curpix].m[1, 2] += cos_term * constant
        nobs_matrix[curpix].m[1, 3] += sin_term * constant

        nobs_matrix[curpix].m[2, 2] += cos_term^2 * constant
        nobs_matrix[curpix].m[2, 3] += cos_term * sin_term * constant

        nobs_matrix[curpix].m[3, 3] += sin_term^2 * constant
    end

    if use_mpi() && comm != nothing
        for curpix in eachindex(nobs_matrix)
            # TODO: repeated calls to MPI_AllReduce are tremendously slow!
            MPI.Allreduce!(MPI.IN_PLACE, nobs_matrix[curpix].m, +, comm)
        end
    end
end


@testset "discreteshape" begin
    img = generate_surface(SynthSphere(), 0.5, [0,0,1], radius = 50, background = false)
    Z, p, q = retrieve_surface(DiscreteShape(), img, smoothness=1000)
    @test Z[38,38] ≈ 0.0003182796154911338
    @test Z[1,1] ≈ 2.0548315482140964e-6
    @test isapprox(Z[76,76], 1.2021784812891086e-19; atol = 2e-15)
    @test Z[45,45] ≈ 0.00040018054317404224
    @test Z[50,18] ≈ 0.00014000747445970892
    @test Z[11, 68] ≈ 0.00018127096811786046
    @test Z[42,42] ≈ 0.0003883024349324123
end


import ..UncertainValues:
	TheoreticalDistributionScalarValue,
    AbstractUncertainTwoParameterScalarValue,
    AbstractUncertainThreeParameterScalarValue,
    UncertainScalarTheoreticalTwoParameter,
    UncertainScalarTheoreticalThreeParameter,
    UncertainScalarNormallyDistributed,
    UncertainScalarUniformlyDistributed,
    UncertainScalarBetaDistributed,
    UncertainScalarBetaPrimeDistributed,
    UncertainScalarBetaBinomialDistributed,
    UncertainScalarBinomialDistributed,
    UncertainScalarGammaDistributed,
    UncertainScalarFrechetDistributed
import Distributions

# Resample for generic
resample(uv::TheoreticalDistributionScalarValue) = rand(uv.distribution)
resample(uv::TheoreticalDistributionScalarValue, n::Int) = rand(uv.distribution, n)

# Custom resample methods for each type of uncertain scalars based on
# distributions (in case we want to implement custom sampling for some of them)
resample(uv::UncertainScalarTheoreticalThreeParameter) = rand(uv.distribution)
resample(uv::UncertainScalarTheoreticalTwoParameter) = rand(uv.distribution)
resample(uv::UncertainScalarNormallyDistributed) = rand(uv.distribution)
resample(uv::UncertainScalarUniformlyDistributed) = rand(uv.distribution)
resample(uv::UncertainScalarBetaDistributed) = rand(uv.distribution)
resample(uv::UncertainScalarBetaPrimeDistributed) = rand(uv.distribution)
resample(uv::UncertainScalarBetaBinomialDistributed) = rand(uv.distribution)
resample(uv::UncertainScalarGammaDistributed) = rand(uv.distribution)
resample(uv::UncertainScalarFrechetDistributed) = rand(uv.distribution)
resample(uv::UncertainScalarBinomialDistributed) = rand(uv.distribution)

resample(uv::UncertainScalarTheoreticalThreeParameter, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarTheoreticalTwoParameter, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarNormallyDistributed, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarUniformlyDistributed, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarBetaDistributed, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarBetaPrimeDistributed, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarBetaBinomialDistributed, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarGammaDistributed, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarFrechetDistributed, n::Int) = rand(uv.distribution, n)
resample(uv::UncertainScalarBinomialDistributed, n::Int) = rand(uv.distribution, n)

resample(x::Distributions.Truncated) = rand(x)
resample(x::Distributions.Truncated, n::Int) = rand(x, n)

resample(x::Distributions.Distribution) = rand(x)
resample(x::Distributions.Distribution, n::Int) = rand(x, n)


myfunc(x) = log(sum(x))

@testset "others" begin
    test_compare(myfunc; x=rand(2))
end


include("../src/SplinePlotting.jl")

function nrbtestcrv()
    pnts = [0.5 1.5 4.5 3.0 7.5 6.0 8.5;
        3.0 5.5 5.5 1.5 1.5 4.0 4.5;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0]
    crv = nrbmak(pnts, [[0, 0, 0, 1/4, 1/2, 3/4, 3/4, 1, 1, 1]])
    return crv
end

function nrbtestcrv3d()
    coefs =[ 6.0  0.0  6.0  1;
        -5.5  0.5  5.5  1;
        -5.0  1.0 -5.0  1;
         4.5  1.5 -4.5  1;
         4.0  2.0  4.0  1;
        -3.5  2.5  3.5  1;
        -3.0  3.0 -3.0  1;
         2.5  3.5 -2.5  1;
         2.0  4.0  2.0  1;
        -1.5  4.5  1.5  1;
        -1.0  5.0 -1.0  1;
         0.5  5.5 -0.5  1;
         0.0  6.0  0.0  1]'
    knots = [0, 0, 0, 0, .1, .2, .3, .4, .5, .6, .7, .8, .9, 1, 1, 1, 1]
    crv3d = nrbmak(coefs, [knots])
    return crv3d
end

function nrbtestsrf()
    pnts = zeros(3, 5, 5)

    pnts[:,:,1] = [ 0.0  3.0  5.0  8.0 10.0;     # w*x
                0.0  0.0  0.0  0.0  0.0;         # w*y
                2.0  2.0  7.0  7.0  8.0]         # w*z

    pnts[:,:,2] = [ 0.0  3.0  5.0  8.0 10.0;
                    3.0  3.0  3.0  3.0  3.0;
                    0.0  0.0  5.0  5.0  7.0]

    pnts[:,:,3] = [ 0.0  3.0  5.0  8.0 10.0;
                    5.0  5.0  5.0  5.0  5.0;
                    0.0  0.0  5.0  5.0  7.0]

    pnts[:,:,4] = [ 0.0  3.0  5.0  8.0 10.0;
                    8.0  8.0  8.0  8.0  8.0;
                    5.0  5.0  8.0  8.0 10.0]

    pnts[:,:,5] = [ 0.0  3.0  5.0  8.0 10.0;
                   10.0 10.0 10.0 10.0 10.0;
                    5.0  5.0  8.0  8.0 10.0]

    uKnots = [0, 0, 0, 1/3, 2/3, 1, 1, 1]
    vKnots = [0, 0, 0, 1/3, 2/3, 1, 1, 1]

    srf = nrbmak(pnts, [uKnots, vKnots])
    return srf
end

function nrbtestvol()
    knots = [[0, 0, 0, 1/2, 1, 1, 1], [0, 0, 0, 1, 1, 1],
             [0, 0, 0, 1/6, 2/6, 1/2, 1/2, 4/6, 5/6, 1, 1, 1]];

    coefs = [-1.0000   -0.9734   -0.7071    1.4290    1.0000    3.4172
             0    2.4172         0    0.0148   -2.0000   -1.9734
             0    2.0000    4.9623    9.4508    4.0000    2.0000
        1.0000    1.0000    0.7071    1.0000    1.0000    1.0000
       -0.8536         0   -0.6036    1.9571    1.2071    3.5000
        0.3536    2.5000    0.2500    0.5429   -1.7071   -1.0000
             0    2.0000    4.4900    8.5444    3.4142    2.0000
        0.8536    1.0000    0.6036    1.0000    0.8536    1.0000
       -0.3536   -4.0000   -0.2500   -1.2929    1.7071    1.0000
        0.8536         0    0.6036   -2.7071   -1.2071   -5.0000
             0    2.0000    4.4900   10.0711    3.4142    2.0000
        0.8536    1.0000    0.6036    1.0000    0.8536    1.0000
             0   -4.0000         0    0.7071    2.0000    5.0000
        1.0000    4.0000    0.7071   -0.7071   -1.0000   -5.0000
             0    2.0000    4.9623   14.4142    4.0000    2.0000
        1.0000    1.0000    0.7071    1.0000    1.0000    1.0000
       -2.5000   -4.0000   -1.7678    0.7071    1.0000    5.0000
             0    4.0000         0   -0.7071   -3.5000   -5.0000
             0    2.0000    6.0418   14.4142    4.0000    2.0000
        1.0000    1.0000    0.7071    1.0000    1.0000    1.0000
       -2.4379         0   -1.7238    2.7071    1.9527    5.0000
        0.9527    4.0000    0.6737    1.2929   -3.4379   -1.0000
             0    2.0000    6.6827   10.0711    4.0000    2.0000
        1.0000    1.0000    0.7071    1.0000    1.0000    1.0000
       -0.9734   -1.0000   -0.6883    0.7071    3.4172    1.0000
        2.4172         0    1.7092   -1.4142   -1.9734   -2.0000
             0    4.0000    6.6827    4.9623    4.0000         0
        1.0000    1.0000    0.7071    0.7071    1.0000    1.0000
             0   -0.8536         0    0.8536    3.5000    1.2071
        2.5000    0.3536    1.7678   -1.2071   -1.0000   -1.7071
             0    3.4142    6.0418    4.4900    4.0000         0
        1.0000    0.8536    0.7071    0.6036    1.0000    0.8536
       -4.0000   -0.3536   -2.8284    1.2071    1.0000    1.7071
             0    0.8536         0   -0.8536   -5.0000   -1.2071
             0    3.4142    7.1213    4.4900    4.0000         0
        1.0000    0.8536    0.7071    0.6036    1.0000    0.8536
       -4.0000         0   -2.8284    1.4142    5.0000    2.0000
        4.0000    1.0000    2.8284   -0.7071   -5.0000   -1.0000
             0    4.0000   10.1924    4.9623    4.0000         0
        1.0000    1.0000    0.7071    0.7071    1.0000    1.0000
       -4.0000   -2.5000   -2.8284    0.7071    5.0000    1.0000
        4.0000         0    2.8284   -2.4749   -5.0000   -3.5000
             0    4.0000   10.1924    6.0418    4.0000         0
        1.0000    1.0000    0.7071    0.7071    1.0000    1.0000
             0   -2.4379         0    1.3808    5.0000    1.9527
        4.0000    0.9527    2.8284   -2.4309   -1.0000   -3.4379
             0    4.0000    7.1213    6.6827    4.0000         0
        1.0000    1.0000    0.7071    0.7071    1.0000    1.0000
       -1.0000   -0.9734    0.2071    2.4163    1.0000    3.4172
             0    2.4172   -1.2071   -1.3954   -2.0000   -1.9734
        2.0000    4.0000    7.0178    6.6827    2.0000         0
        1.0000    1.0000    1.0000    0.7071    1.0000    1.0000
       -0.8536         0    0.3536    2.4749    1.2071    3.5000
        0.3536    2.5000   -0.8536   -0.7071   -1.7071   -1.0000
        1.7071    4.0000    6.3498    6.0418    1.7071         0
        0.8536    1.0000    0.8536    0.7071    0.8536    1.0000
       -0.3536   -4.0000    0.8536    0.7071    1.7071    1.0000
        0.8536         0   -0.3536   -3.5355   -1.2071   -5.0000
        1.7071    4.0000    6.3498    7.1213    1.7071         0
        0.8536    1.0000    0.8536    0.7071    0.8536    1.0000
             0   -4.0000    1.2071    3.5355    2.0000    5.0000
        1.0000    4.0000   -0.2071   -3.5355   -1.0000   -5.0000
        2.0000    4.0000    7.0178   10.1924    2.0000         0
        1.0000    1.0000    1.0000    0.7071    1.0000    1.0000
       -2.5000   -4.0000   -0.5429    3.5355    1.0000    5.0000
             0    4.0000   -1.9571   -3.5355   -3.5000   -5.0000
        2.0000    4.0000    8.5444   10.1924    2.0000         0
        1.0000    1.0000    1.0000    0.7071    1.0000    1.0000
       -2.4379         0   -0.0355    3.5355    1.9527    5.0000
        0.9527    4.0000   -1.4497   -0.7071   -3.4379   -1.0000
        2.0000    4.0000    9.4508    7.1213    2.0000         0
        1.0000    1.0000    1.0000    0.7071    1.0000    1.0000]

    coefs = reshape(coefs, 4, 4, 3, 9)
    horseshoe = nrbmak(coefs, knots)

    return horseshoe
end


@testset "Mass of hermitian lattices" begin
  Qx, x = PolynomialRing(FlintQQ, "x")
  f = x-1
  K, a = NumberField(f, "a", cached = false)
  Kt, t = PolynomialRing(K, "t")
  g = t^2+(2)
  E, b = NumberField(g, "b", cached = false)
  D = matrix(E, 3, 3, [(1), 0, 0, 0, (1), 0, 0, 0, (1)])
  gens = Vector{Hecke.NfRelElem{nf_elem}}[map(E, [(-2)*b+(-2), b+(6), 0]), map(E, [0, (1), (1)]), map(E, [b+(-6), (-6)*b+(6), 0])]
  L = hermitian_lattice(E, generators = gens, gram_ambient_space = D)
  m = @inferred mass(L)
  @test m == 1//32

  Qx, x = PolynomialRing(FlintQQ, "x")
  f = x^2-3
  K, a = NumberField(f, "a", cached = false)
  Kt, t = PolynomialRing(K, "t")
  g = t^2+(1)
  E, b = NumberField(g, "b", cached = false)
  D = matrix(E, 3, 3, [(1), 0, 0, 0, (1), 0, 0, 0, (1)])
  gens = Vector{Hecke.NfRelElem{nf_elem}}[map(E, [(-1//2*a-1//2)*b+(1//2*a+5//2), (-1), 0]), map(E, [(a-9)*b+(15*a+3), 0, (-a)*b+(8*a+13)]), map(E, [(-1610*a+2990)*b+(2070*a-2990), 0, (-68*a+393)*b+(19*a+462)])]
  L = hermitian_lattice(E, generators = gens, gram_ambient_space = D)
  m = @inferred mass(L)
  @test m == 1//108

  Qx, x = PolynomialRing(FlintQQ, "x")
  f = x^4-x^3-4*x^2+4*x+1
  K, a = NumberField(f, "a", cached = false)
  Kt, t = PolynomialRing(K, "t")
  g = t^2+(a^3 - 1*a^2 - 4*a + 5)
  E, b = NumberField(g, "b", cached = false)
  D = matrix(E, 3, 3, [(1), 0, 0, 0, (1), 0, 0, 0, (1)])
  gens = Vector{Hecke.NfRelElem{nf_elem}}[map(E, [(1), 0, 0]), map(E, [0, (1), 0]), map(E, [0, 0, (1)])]
  L = hermitian_lattice(E, generators = gens, gram_ambient_space = D)
  m = @inferred mass(L)
  @test m == 1//10125
end


"""
    struct ConeMantle{T,TR,TP,D} <: AbstractCurvedSurfacePrimitive{T}

Surface primitive describing the mantle of a [`Cone`](@ref).

## Parametric types
* `T`: Precision type.
* `TR`: Type of the radius `r`.
    * `TR == T`: CylinderMantle (constant radius `r` at all `z`).
    * `TR == Tuple{T, T}`: VaryingCylinderMantle (bottom radius at `r[1]`, top radius at `r[2]`).
* `TP`: Type of the angular range `φ`.
    * `TP == Nothing`: Full 2π Cone.
    * `TP == Tuple{T, T}`: Partial Cone ranging from `φ[1]` to `φ[2]`.
* `D`: Direction in which the normal vector points (`:inwards` or `:outwards`).

## Fields
* `r::TR`: Definition of the radius of the `ConeMantle` (in m).
* `φ::TP`: Range in polar angle `φ` over which the `ConeMantle` extends (in radians).
* `hZ::T`: Half of the height of the `ConeMantle` (in m).
* `origin::CartesianPoint{T}`: Origin of the `Cone` which has this `ConeMantle` as surface.
* `rotation::SMatrix{3,3,T,9}`: Rotation matrix of the `Cone` which has this `ConeMantle` as surface.
"""
@with_kw struct ConeMantle{T,TR,TP,D} <: AbstractCurvedSurfacePrimitive{T}
    r::TR = 1
    φ::TP = nothing
    hZ::T = 1

    # axis::Line{T} = Line{T}(CartesianPoint{T}(zero(T), zero(T), -hZ), CartesianVector{T}(zero(T), zero(T), 2hZ))
    origin::CartesianPoint{T} = zero(CartesianPoint{T})
    rotation::SMatrix{3,3,T,9} = one(SMatrix{3, 3, T, 9})
end

radius_at_z(hZ::T, rBot::T, rTop::T, z::T) where {T} = iszero(hZ) ? rBot : rBot + (hZ+z)*(rTop - rBot)/(2hZ)
radius_at_z(cm::ConeMantle{T,T}, z::T) where {T} = cm.r
radius_at_z(cm::ConeMantle{T,Tuple{T,T}}, z::T) where {T} = radius_at_z(cm.hZ, cm.r[1], cm.r[2], z)

get_φ_limits(cm::ConeMantle{T,<:Any,Tuple{T,T}}) where {T} = cm.φ[1], cm.φ[2]
get_φ_limits(cm::ConeMantle{T,<:Any,Nothing}) where {T} = T(0), T(2π)

function normal(cm::ConeMantle{T,T}, pt::CartesianPoint{T}) where {T}
    pto = _transform_into_object_coordinate_system(pt, cm)
    cyl = CylindricalPoint(pto)
    return CartesianVector(_transform_into_global_coordinate_system(
            CartesianPoint(CylindricalPoint{T}(cyl.r, cyl.φ, zero(T))), cm))
end
function normal(cm::ConeMantle{T,Tuple{T,T},<:Any,:inwards}, pt::CartesianPoint{T}) where {T}
    pto = _transform_into_object_coordinate_system(pt, cm)
    cyl = CylindricalPoint(pto)
    Δr = cm.r[2] - cm.r[1]
    Δz = 2cm.hZ
    return CartesianVector(_transform_into_global_coordinate_system(
            CartesianPoint(CylindricalPoint{T}(-one(T), cyl.φ, Δr / Δz)), cm))
end
function normal(cm::ConeMantle{T,Tuple{T,T},<:Any,:outwards}, pt::CartesianPoint{T}) where {T}
    pto = _transform_into_object_coordinate_system(pt, cm)
    cyl = CylindricalPoint(pto)
    Δr = cm.r[2] - cm.r[1]
    Δz = 2cm.hZ
    return CartesianVector(_transform_into_global_coordinate_system(
            CartesianPoint(CylindricalPoint{T}( one(T), cyl.φ, -Δr / Δz)), cm))
end

function vertices(cm::ConeMantle{T}, n_arc::Int64)::Vector{CartesianPoint{T}} where {T}
    φMin, φMax = get_φ_limits(cm)
    n_arc = _get_n_points_in_arc_φ(cm, n_arc)
    φ = range(φMin, φMax, length = n_arc+1)
    rbot = radius_at_z(cm,-cm.hZ)
    rtop = radius_at_z(cm,cm.hZ)
    botcircle = [_transform_into_global_coordinate_system(CartesianPoint{T}(rbot*cos(φ), rbot*sin(φ), -cm.hZ), cm) for φ in φ]
    topcircle = [_transform_into_global_coordinate_system(CartesianPoint{T}(rtop*cos(φ), rtop*sin(φ), cm.hZ), cm) for φ in φ]
    append!(botcircle, topcircle)
end

function sample(cm::ConeMantle{T}, spacing::T)::Vector{CartesianPoint{T}} where {T}
    φMin, φMax = get_φ_limits(cm)
    Δφ = abs(φMax - φMin)
    full2π = isnothing(cm.φ)
    rbot = radius_at_z(cm,-cm.hZ)
    rtop = radius_at_z(cm,cm.hZ)
    l = hypot(2cm.hZ, rtop - rbot)
    z = range(-cm.hZ, cm.hZ, length = max(2,1+Int(ceil(l/spacing))))
    [_transform_into_global_coordinate_system(CartesianPoint{T}((radius_at_z(cm,z) .* reverse(sincos(φ)))..., z), cm) for z in z for r in [radius_at_z(cm,z)] for φ in (r == 0 ? [φMin] : range(φMin, φMax - full2π*spacing/radius_at_z(cm,z), length = max(2,1+Int(ceil(Δφ*radius_at_z(cm,z)/spacing)))))]
end

function connections(cm::ConeMantle, n_arc::Int64)::Vector{Vector{Int64}}
    n_arc = _get_n_points_in_arc_φ(cm, n_arc)
    [[i,i+1,i+n_arc+2,i+n_arc+1] for i in 1:n_arc]
end

function connections(cm::ConeMantle, n_arc::Int64, n_vert_lines::Int64)::Vector{Vector{Int64}}
    n_arc = _get_n_points_in_arc_φ(cm, n_arc)
    verts = [[i, i + n_arc + 1] for i in _get_vert_lines_range(cm,n_arc,n_vert_lines)]
    circ1 =  [[i, i + 1] for i in 1:n_arc]
    circ2 =  [[i, i + 1] for i in n_arc+2:2*n_arc+1]
    append!(verts, circ1, circ2)
end

get_label_name(::ConeMantle) = "Cone Mantle"

const FullConeMantle{T,D} = ConeMantle{T,Tuple{T,T},Nothing,D} # ugly name but works for now, should just be `ConeMantle`...
const PartialConeMantle{T,D} = ConeMantle{T,Tuple{T,T},Tuple{T,T},D}

extremum(cm::FullConeMantle{T}) where {T} = hypot(cm.hZ, max(cm.r...))

extremum(cm::PartialConeMantle{T}) where {T} = hypot(cm.hZ, max(cm.r...))

function lines(sp::FullConeMantle{T}; n = 2) where {T}
    bot_origin = _transform_into_global_coordinate_system(CartesianPoint{T}(zero(T),zero(T), -sp.hZ), sp)
    bot_ellipse = Circle{T}(r = sp.r[1], φ = sp.φ, origin = bot_origin, rotation = sp.rotation)
    top_origin = _transform_into_global_coordinate_system(CartesianPoint{T}(zero(T),zero(T), +sp.hZ), sp)
    top_ellipse = Circle{T}(r = sp.r[2], φ = sp.φ, origin = top_origin, rotation = sp.rotation)
    φs = range(T(0), step = T(2π) / n, length = n)
    edges = [ Edge{T}(_transform_into_global_coordinate_system(CartesianPoint(CylindricalPoint{T}(sp.r[1], φ, -sp.hZ)), sp),
                      _transform_into_global_coordinate_system(CartesianPoint(CylindricalPoint{T}(sp.r[2], φ, +sp.hZ)), sp)) for φ in φs ]
    bot_ellipse, top_ellipse, edges
end
function lines(sp::PartialConeMantle{T}; n = 2) where {T}
    bot_origin = _transform_into_global_coordinate_system(CartesianPoint{T}(zero(T), zero(T), -sp.hZ), sp)
    bot_ellipse = PartialCircle{T}(r = sp.r[1], φ = sp.φ, origin = bot_origin, rotation = sp.rotation)
    top_origin = _transform_into_global_coordinate_system(CartesianPoint{T}(zero(T), zero(T), +sp.hZ), sp)
    top_ellipse = PartialCircle{T}(r = sp.r[2], φ = sp.φ, origin = top_origin, rotation = sp.rotation)
    φs = range(sp.φ[1], stop = sp.φ[2], length = n)
    edges = [ Edge{T}(_transform_into_global_coordinate_system(CartesianPoint(CylindricalPoint{T}(sp.r[1], φ, -sp.hZ)), sp),
                      _transform_into_global_coordinate_system(CartesianPoint(CylindricalPoint{T}(sp.r[2], φ, +sp.hZ)), sp)) for φ in φs ]
    bot_ellipse, top_ellipse, edges
end

"""
    intersection(cm::ConeMantle{T,Tuple{T,T}}, l::Line{T}) where {T}

Calculates the intersections of a `Line` with a `ConeMantle`.

## Arguments
* `cm::ConeMantle{T,Tuple{T,T}}`: The `ConeMantle`.
* `l::Line{T}`: The `Line`.

!!! note
    The function will always return 2 CartesianPoint's.
    If the line just touches the mantle, the two points will be the same.
    If the line does not touch the mantle at all, the two points will have NaN's as there coordinates.
    If the line crosses the mantle only once, two points will be returned. The two points will be the same point (the intersection).
    If the line lies inside the mantle and is parallel to it. The same point will be returned which is the origin of the line.
"""
function intersection(cm::ConeMantle{T,Tuple{T,T}}, l::Line{T}) where {T}
    obj_l = _transform_into_object_coordinate_system(l, cm) # direction is not normalized

    L1 = obj_l.origin.x
    L2 = obj_l.origin.y
    L3 = obj_l.origin.z
    D1 = obj_l.direction.x
    D2 = obj_l.direction.y
    D3 = obj_l.direction.z

    hZ = cm.hZ
    R0 = cm.r[1]

    S = (cm.r[2] - cm.r[1]) / 2hZ # r slope

    f1 = D1^2 + D2^2 - D3^2*S^2
    # f1 is zero if there is only one intersection.

    λ1, λ2 = if f1 == 0 # one intersection
        term1 = -hZ^2*S^2 - 2hZ*L3*S^2 - 2*hZ*R0*S + L1^2 + L2^2 - L3^2*S^2 - 2L3*R0*S - R0^2
        term2 = L1*sqrt(D3^2*S^2 - D2^2) - D2*L2 + D3*hZ*S^2 + D3*L3*S^2 + D3*R0*S
        λ = term1 / (2term2)
        λ, λ
    else # two or no intersections
        λ = inv(f1)
        term1 = (2D1*L1 + 2D2*L2 - 2D3*hZ*S^2 - 2D3*L3*S^2 - 2D3*R0*S)^2
        term2 = -hZ^2*S^2 - 2hZ*L3*S^2 - 2hZ*R0*S + L1^2 + L2^2 - L3^2*S^2 - 2L3*R0*S - R0^2
        term3 = -D1*L1 - D2*L2 + D3*hZ*S^2 + D3*L3*S^2 + D3*R0*S
        term4 = term1 - 4*f1*term2
        # sq::T = term4 < 0 ? T(NaN) : sqrt(term4)
        sq::T = sqrt(abs(term4))

        # @show term4
        λ1 = λ * (-sq/2 + term3)
        λ2 = λ * (+sq/2 + term3)
        λ1, λ2
    end
    ints1 = obj_l.origin + λ1 * obj_l.direction
    ints2 = obj_l.origin + λ2 * obj_l.direction
    return _transform_into_global_coordinate_system(ints1, cm),
           _transform_into_global_coordinate_system(ints2, cm)
end

# """
#     intersection(cm::ConeMantle{T,T}, l::Line{T}) where {T}

# The function will always return 2 CartesianPoint's.
# If the line just touches the mantle, the two points will be the same.
# If the line does not touch the mantle at all, the two points will have NaN's as there coordinates.
# """
# function intersection(cm::ConeMantle{T,T}, l::Line{T}) where {T}
#     obj_l = _transform_into_object_coordinate_system(l, cm) # direction is not normalized

#     L1 = obj_l.origin.x
#     L2 = obj_l.origin.y
#     L3 = obj_l.origin.z
#     D1 = obj_l.direction.x
#     D2 = obj_l.direction.y
#     D3 = obj_l.direction.z

#     f1 = D1^2 + D2^2
#     λ = inv(f1) # f1 is only 0 if obj_l is parallel to the axis of the cone
#                 # (here eZ -> D1 = D2 = 0)
#                 # We assume here that this is not the case ->
#                 # We check this in choosing the sample / evaluating dimensions in `paint!`
#     hZ = cm.hZ
#     R0 = cm.r

#     term1 = (2D1*L1 + 2D2*L2)^2
#     term2 = L1^2 + L2^2 - R0^2
#     term3 = -D1*L1 - D2*L2
#     term4 = term1 - 4*f1*term2
#     sq::T = term4 < 0 ? T(NaN) : sqrt(term1 - 4*f1*term2) # if this

#     λ1 = λ * (-sq/2 + term3)
#     λ2 = λ * (+sq/2 + term3)

#     ints1 = obj_l.origin + λ1 * obj_l.direction
#     ints2 = obj_l.origin + λ2 * obj_l.direction
#     return _transform_into_global_coordinate_system(ints1, cm),
#            _transform_into_global_coordinate_system(ints2, cm)
# end

# function get_2d_grid_ticks_and_proj(cm::ConeMantle{T}, t) where {T}
#     pts = extreme_points(cm)

#     t_idx_range_x, t_idx_range_y, t_idx_range_z = get_min_max_index_ranges((pts, t))
#     ls = (length(t_idx_range_x), length(t_idx_range_y), length(t_idx_range_z))
#     ls = (
#         ls[1] == 1 ? typemax(eltype(ls)) : ls[1],
#         ls[2] == 1 ? typemax(eltype(ls)) : ls[2],
#         ls[3] == 1 ? typemax(eltype(ls)) : ls[3]
#     )
#     eX = CartesianVector{T}(one(T),zero(T),zero(T))
#     eY = CartesianVector{T}(zero(T),one(T),zero(T))
#     eZ = CartesianVector{T}(zero(T),zero(T),one(T))
#     axis = cm.rotation * CartesianVector{T}(zero(T), zero(T), one(T))

#     proj, t1, t2 = if (axis × eZ != zero(CartesianVector{T}))
#         Val{:xy}(), t_idx_range_x, t_idx_range_y
#     elseif (axis × eY != zero(CartesianVector{T}))
#         Val{:xz}(), t_idx_range_x, t_idx_range_z
#     elseif (axis × eX != zero(CartesianVector{T}))
#         Val{:yz}(), t_idx_range_y, t_idx_range_z
#     else
#         error("Sampling Error. Have to extend cases")
#         # Should never happen. This else case can be removed after testing.
#     end

#     return t1, t2, proj
# end

# function distance_to_surface(pt::AbstractCoordinatePoint{T}, c::ConeMantle{T, <:Any, Nothing, <:Any})::T where {T}
#     pcy = CylindricalPoint(pt)
#     distance_to_line(PlanarPoint{T}(pcy.r,pcy.z), LineSegment(c))
# end

# function distance_to_surface(pt::AbstractCoordinatePoint{T}, c::ConeMantle{T, <:Any, <:AbstractInterval, <:Any})::T where {T}
#     pcy = CylindricalPoint(pt)
#     φMin::T, φMax::T, _ = get_φ_limits(c)
#     if _in_φ(pcy, c.φ)
#         return distance_to_line(PlanarPoint{T}(pcy.r,pcy.z), LineSegment(c))
#     else
#         return distance_to_line(CartesianPoint(pt), LineSegment(c, _φNear(pcy.φ, φMin, φMax)))
#     end
# end


const eta0 = 4π*299792458e-7

#plane wave
function u2α(k, u::PlaneWave, centers::Array{T,2}, P) where T <: Real
	#build incoming coefficients for plane wave incident field
	α = Array{Complex{Float64}}(undef, size(centers,1)*(2P + 1))
	for ic = 1:size(centers,1)
		phase = exp(1.0im*k*(cos(u.θi)*centers[ic,1] + sin(u.θi)*centers[ic,2])) #phase shift added to move cylinder coords
		for p = -P:P
			α[(ic-1)*(2P+1) + p + P + 1] = phase*exp(1.0im*p*(π/2-u.θi))
		end
	end
	α
end

"""
    uinc(k, points, u::Einc)

Computes incident field produced by `u` with outer wavenumber `k` at `points`.
"""
function uinc(k, points::Array{T,2}, u::PlaneWave) where T <: Real
	exp.(1.0im*k*(cos(u.θi)*points[:,1] + sin(u.θi)*points[:,2]))
end

function uinc(k, points::Array{T,1}, u::PlaneWave) where T <: Real
	exp(1.0im*k*(cos(u.θi)*points[1] + sin(u.θi)*points[2]))
end

#line (filament) source
function u2α(k, u::LineSource, centers::Array{T,2}, P) where T <: Real
	α = Array{Complex{Float64}}(undef, size(centers,1)*(2P + 1))
	for ic = 1:size(centers,1)
		R = hypot(centers[ic,1] - u.x, centers[ic,2] - u.y)
		θ = atan(centers[ic,2] - u.y, centers[ic,1] - u.x)
		α[(ic-1)*(2P+1) + 0 + P + 1] = 0.25im*besselh(0, 1, k*R)
		for p = 1:P
			α[(ic-1)*(2P+1) - p + P + 1] = 0.25im*exp(1im*p*θ)*besselh(p, 1, k*R)
			α[(ic-1)*(2P+1) + p + P + 1] = 0.25im*exp(-1im*p*θ)*besselh(-p, 1, k*R)
		end
	end
	α
end

function uinc(k, points::Array{T,2}, u::LineSource) where T <: Real
	r = hypot.(u.x .- points[:,1], u.y .- points[:,2])
	if any(r .== 0)
		@warn("uinc: encountered singularity in incident field, returned NaN")
		r[r.==0] = NaN
	end
	0.25im*besselh.(0, 1, k*r)
end

function uinc(k, points::Array{T,1}, u::LineSource) where T <: Real
	r = hypot(u.x - points[1], u.y - points[2])
	if r == 0
		@warn("uinc: encountered singularity in incident field, returned NaN")
		return NaN
	end
	0.25im*besselh(0, 1, k*r)
end

#current source TODO: simple avg->trap rule
function u2α(k, u::CurrentSource, centers::Array{T,2}, P) where T <: Real
	α = zeros(Complex{Float64}, size(centers,1)*(2P + 1))
	c = 0.25im*u.len/length(u.σ)
	for ic = 1:size(centers,1)
		for is = 1:length(u.σ)
			R = hypot(centers[ic,1] - u.p[is,1], centers[ic,2] - u.p[is,2])
			θ = atan(centers[ic,2] - u.p[is,2], centers[ic,1] - u.p[is,1])

			α[(ic-1)*(2P+1) + P + 1] +=	c*u.σ[is]*besselh(0, 1, k*R)
			for p = 1:P
				H = besselh(p, 1, k*R)
				α[(ic-1)*(2P+1) - p + P + 1] +=	c*u.σ[is]*H*exp(1im*p*θ)
				α[(ic-1)*(2P+1) + p + P + 1] +=	c*u.σ[is]*(-1)^p*H*exp(-1im*p*θ)
			end
		end
	end
	α
end

function uinc(k, points::Array{T,2}, u::CurrentSource) where T <: Real
	res = Array{Complex{Float64}}(undef, size(points, 1))
	r = Array{Float64}(undef, size(u.p, 1))
	for i = 1:size(points, 1)
		r .= hypot.(u.p[:,1] - points[i,1], u.p[:,2] - points[i,2])
		if any(r .== 0) #point is on source
			@warn("uinc: encountered singularity in incident field, returned NaN")
			res[i] = NaN
		end
		res[i] = (0.25im*u.len/length(u.σ))*sum(besselh.(0, 1, k*r).*u.σ)
	end
	res
end

function uinc(k, points::Array{T,1}, u::CurrentSource) where T <: Real
	r = hypot.(u.p[:,1] - points[1], u.p[:,2] - points[2])
	if any(r .== 0) #point is on source
		@warn("uinc: encountered singularity in incident field, returned NaN")
		return NaN
	end
	(0.25im*u.len/length(u.σ))*sum(besselh.(0, 1, k*r).*u.σ)
end

#compute error in field approximation (on multipole disk)
function err_α(k, ui::Einc, P, shape, center)
	#first assert shape and field source do not intersect
	if isa(ui, LineSource)
		hypot(center[1] - ui.x, center[2] - ui.y) <= shape.R && error("LineSource is in shape")
	elseif isa(ui, CurrentSource)
		#assuming straight line, see mathworld.wolfram.com/Circle-LineIntersection.html
		x1 = ui.p[1,1] - center[1]
		y1 = ui.p[1,2] - center[2]
		x2 = ui.p[end,1] - center[1]
		y2 = ui.p[end,2] - center[2]
		dr2 = (x2 - x1)^2 + (y2 - y1)^2
		D2 = (x1*y2 - x2*y1)^2
		shape.R^2*dr2 - D2 < 0 || error("CurrentSource is in shape")
	end

	Nt = 10_000
	θ = range(0, stop=2π, length=Nt)[1:end-1]
	α = u2α(k, ui, center, P)
	u_exact = uinc(k, center .+ shape.R*[cos.(θ) sin.(θ)], ui)
	u_α = sum(α[p + P + 1]*besselj(p, k*shape.R)*exp.(1im*p*θ) for p = -P:P)
	norm(u_α - u_exact)/norm(u_exact)
end

function hxinc(k, p, u::PlaneWave)
    (sin(u.θi)/eta0)*exp.(1.0im*k*(cos(u.θi)*p[:,1] + sin(u.θi)*p[:,2]))
end

function hyinc(k, p, u::PlaneWave)
    (-cos(u.θi)/eta0)*exp.(1.0im*k*(cos(u.θi)*p[:,1] + sin(u.θi)*p[:,2]))
end

function hxinc(k, p, u::LineSource)
    R = sqrt.((view(p,:,1) .- u.x).^2 + (view(p,:,2) .- u.y).^2)
    if any(R .== 0)
		@warn("Hinc: encountered singularity in incident field, returned NaN")
		R[R.==0] = NaN
	end
    h = (0.25/eta0./R).*besselh.(1,k*R).*(u.y .- view(p,:,2))
end

function hyinc(k, p, u::LineSource)
    R = sqrt.((view(p,:,1) .- u.x).^2 + (view(p,:,2) .- u.y).^2)
    if any(R .== 0)
		@warn("Hinc: encountered singularity in incident field, returned NaN")
		R[R.==0] = NaN
	end
    h = (0.25/eta0./R).*besselh.(1,k*R).*(view(p,:,1) .- u.x)
end

#these untested
function hxinc(k, points, u::CurrentSource)
    res = Array{Complex{Float64}}(undef, size(points, 1))
	r = Array{Float64}(undef, size(u.p, 1))
    y = Array{Float64}(undef, size(u.p, 1))
	for i = 1:size(points, 1)
		y[:] = points[i,2] - view(u.p,:,2)
        r[:] = hypot.(points[i,1] .- view(u.p,:,1), y)
		if any(r .== 0) #point is on source
			@warn("Hinc: encountered singularity in incident field, returned NaN")
			res[i] = NaN
		end
		res[i] = (-0.25*u.len/length(u.σ)/eta0)*sum(besselh.(1, 1, k*r).*u.σ.*y./r)
	end
	res
end
function hyinc(k, points, u::CurrentSource)
    res = Array{Complex{Float64}}(undef, size(points, 1))
	r = Array{Float64}(undef, size(u.p, 1))
    x = Array{Float64}(undef, size(u.p, 1))
	for i = 1:size(points, 1)
		x[:] = points[i,1] .- view(u.p,:,1)
        r[:] = hypot.(x, points[i,2] .- view(u.p,:,2))
		if any(r .== 0) #point is on source
			@warn("Hinc: encountered singularity in incident field, returned NaN")
			res[i] = NaN
		end
		res[i] = (0.25*u.len/length(u.σ)/eta0)*sum(besselh.(1, 1, k*r).*u.σ.*x./r)
	end
	res
end


using Makie

function main()
    s = Scene()
    x = -10:0.01:-0.1
    lines!(s, x, 1 ./ x)
    x = 0.1:0.01:10
    lines!(s, x, 1 ./ x)
end

main()


using KernelAbstractions

@kernel function darcyKernel!(uⁿ⁺¹, vⁿ⁺¹, wⁿ⁺¹, h, Δx, Δy, Δz, K)
    i, j, k = @index(Global, NTuple)

    @inbounds begin
        uⁿ⁺¹[i,j,k] = -K[i,j,k] * ∇3dᶜ(h, i, j, k, (Δx, Δy, Δz))[1]    # u
        vⁿ⁺¹[i,j,k] = -K[i,j,k] * ∇3dᶜ(h, i, j, k, (Δx, Δy, Δz))[2]    # v
        wⁿ⁺¹[i,j,k] = -K[i,j,k] * ∇3dᶜ(h, i, j, k, (Δx, Δy, Δz))[3]    # w
    end
end