text:
# draw Gaussian on circle

using Distributions
# using Fontconfig
using Cairo
using Gadfly, Colors
using Random

using ApproxManifoldProducts
using KernelDensityEstimate

using TransformUtils

const TU = TransformUtils

setForceEvalDirect!(true)

##

pd = Normal(-pi, 0.25)
pd2 = Normal(0.0, 0.75)
pd3 = MixtureModel([Normal(1.0, 0.3);Normal(-1.0, 0.2)], [0.6;0.4])

arr = Function[
  (x)->pdf(pd, x);
  (x)->pdf(pd2, x);
  (x)->pdf(pd3, x)
]

pl = plotCircBeliefs(arr)

## look at known failures

pts = [-pi/2; pi/2+1e-10]
pd3 = MixtureModel([Normal(pts[1], 0.1);Normal(pts[2], 0.1)], [0.5;0.5])

arr = Function[
  (x)->pdf(pd3, x);
]

pl = plotCircBeliefs(arr)

# pl |> PDF("/tmp/test.pdf")
# @async run(`evince /tmp/test.pdf`)

## Plot a KDE on circle

pts = rand(Uniform(-pi,pi), 100)

pc = kde!(pts, [0.1;])

gg = (x)->pc([x;])[1]

arr = [gg;]

pl = plotCircBeliefs(arr)

## plot two modes on circular manifold (using tree based eval)

# these do work
pts = [-pi/2; pi/2]
pts = [-pi/2; pi/2-1e-1]
pts = TU.wrapRad.([-pi/2; pi+0.1])
pts = [-pi/2+1e-1; pi/2]

# these dont work with difftheta, but fine with `-`
pts = [-pi/2; pi/2+1e-10]
pts = [-pi/2; pi/2+1.0]
pts = [-pi/2-1e-10; pi/2]

p = kde!(pts, [0.2], (addtheta,), (difftheta,))
# p.bt.centers

# circular plot does not work
plotKDECircular(p)

# linear plot works with similar eval
x=reshape(collect(range(-pi,stop=pi,length=1000)), 1,:)

# should work, but definitely doesn't
y = p( x, false, 1e-3, (addtheta,),(difftheta,) )
Gadfly.plot(x=x, y=y, Geom.line)

# shouldnt necessarily work, but does
# y = p( x, false, 1e-3, (+,),(-,) )
# Gadfly.plot(x=x, y=y, Geom.line)

gg = (x)->p([x;], false, 1e-3, (addtheta,), (difftheta,))[1]

arr = [gg;]

pl = plotCircBeliefs(arr)

## debugging the circular plotting function

arr = []

for bd in [p;]
  gg = (x::Float64)->bd([x;])[1]
  push!(arr, gg)
end

N=1000
th = range(-pi, stop=pi-1e-15, length=N)
c=["green"; "blue"; "deepskyblue"; "magenta"; "cyan"]
logpdf=true

c = cat(["red";], c, dims=1)
beliefs = Dict{Int, Function}()
beliefs[1] = (x)->1.0
global j = 1
for ar in arr
  global j
  j += 1
  beliefs[j] = logpdf ? (x)->log(ar(x)+1.0) : (x)->ar(x) #
end

PL = []
for j in 1:length(beliefs)
  obj = zeros(N)
  X = zeros(N)
  Y = zeros(N)
  for i in 1:N
    @show i, th
    @show t = th[i]
    obj[i] = 1.0 # beliefs[1](t) # = 1.0
    obj[i] += (j != 1 ? beliefs[j](t) : 0.0) # TODO: error must be here?
    # second term directional statistics for plotting (not evaluation)
    # obj[i] += (j != 1 ? beliefs[j](t-2pi) : 0.0 )
    # obj[i] += (j != 1 ? beliefs[j](t+2pi) : 0.0 )
    @show t, obj[i]
    X[i] = cos(t)*obj[i]
    Y[i] = sin(t)*obj[i]
  end
  push!(PL, Gadfly.layer(x=deepcopy(X), y=deepcopy(Y), Gadfly.Geom.path, Gadfly.Theme(default_color=parse(Colorant,c[j])))[1] )
end

Gadfly.plot(PL...)

# pts = [0.0001*randn(N2).-pi/2; 0.0001*randn(N2).+pi/2]
# shuffle!(pts)

## repeat with linear plot (still using same eval)

#

ext: jl


text:
module SpecialFunctions
using LinearAlgebra
using OffsetArrays
using ..Constants

export cricbessel, crichankel, twobytwoinverse

const ci = 1.0im

"""
    ricbessel(n, ds, ϵ, Ψ)

Compute the Ricatti-Bessel function.

References:

- [Riccati-Bessel Functions](https://mathworld.wolfram.com/Riccati-BesselFunctions.html) at Wolfram MathWorld
"""
function ricbessel(n::Int64, ds::Float64, ϵ::Float64)
    Ψ = OffsetArray(zeros(n + 1), 0:n)

    if ds < n
        ns = floor(Int, ds + 4(ds^0.3333) + 17)
        ns = max(n + 10, ns)
        dns = 0
        for i in (ns - 1):-1:n
            sn = (i + 1) / ds
            dns = sn - 1 / (dns + sn)
        end
        Ψ[n] = dns
        Ψ[n - 1] = n / ds - 1 / (dns + (n) / ds)
        for i in (n - 2):-1:1
            sn = (i + 1) / ds
            Ψ[i] = sn - 1 / (Ψ[i + 1] + sn)
        end
        Ψ_t = sin(ds)
        Ψ[0] = Ψ_t
        ds2 = ds * ds
        sum = Ψ_t * Ψ_t / ds2
        for i in 1:n
            Ψ_t = Ψ_t / (i / ds + Ψ[i])
            sum = sum + (2i + 1) * Ψ_t * Ψ_t / ds2
            err = abs(1 - sum)
            Ψ[i] = Ψ_t
            if err < ϵ
                # nmax = i
                return Ψ
            end
        end
        # nmax = n
    else
        Ψ[0] = sin(ds)
        Ψ[1] = Ψ[0] / ds - cos(ds)
        for i in 1:(n - 1)
            sn = (2i + 1) / ds
            Ψ[i + 1] = sn * Ψ[i] - Ψ[i - 1]
        end
        # nmax = n
    end

    return Ψ
end

function richankel(n::Int64, ds::Float64)
    Ξ = OffsetArray(zeros(ComplexF64, n + 1), 0:n)

    if ds < n
        ns = floor(Int, ds + 4(ds^0.3333) + 17)
        ns = max(n + 10, ns)
        dns = 0
        for i in (ns - 1):-1:n
            sn = (i + 1) / ds
            dns = sn - 1 / (dns + sn)
        end
        Ξ[n] = dns
        Ξ[n - 1] = (n) / ds - 1 / (dns + (n) / ds)
        for i in (n - 2):-1:1
            sn = (i + 1) / ds
            Ξ[i] = sn - 1 / (Ξ[i + 1] + sn)
        end
        chi0 = -cos(ds)
        psi = sin(ds)
        chi1 = chi0 / ds - psi
        Ξ[0] = ComplexF64(psi, chi0)
        for i in 1:n
            chi2 = (2i + 1) / ds * chi1 - chi0
            psi = psi / ((i) / ds + Ξ[i])
            Ξ[i] = ComplexF64(psi, chi1)
            chi0 = chi1
            chi1 = chi2
        end
    else
        chi0 = -cos(ds)
        psi0 = sin(ds)
        chi1 = chi0 / ds - psi0
        psi1 = psi0 / ds + chi0
        Ξ[0] = ComplexF64(psi0, chi0)
        Ξ[1] = ComplexF64(psi1, chi1)
        for i in 1:(n - 1)
            sn = (2i + 1) / ds
            Ξ[i + 1] = sn * Ξ[i] - Ξ[i - 1]
        end
    end
    return Ξ
end

"""
    cricbessel(n, ds)

Compute the Riccati-Bessel function, return complex values.
"""
function cricbessel(n::Int64, ds::ComplexF64)
    Ψ, Χ = cspherebessel(n, ds)
    for i in 0:n
        Ψ[i] *= ds
    end
    return Ψ
end

function crichankel(n::Int64, ds::ComplexF64)
    Ξ = OffsetArray(zeros(ComplexF64, n + 1), 0:n)
    Ξ[0] = -ci * exp(ci * ds)
    Ψ_0 = sin(ds)
    if abs(Ξ[0]) / abs(Ψ_0) < 1e-6
        Ξ[1] = -exp(ci * ds) * (ci + ds) / ds
        for i in 2:n
            Ξ[i] = (2i + 1) / ds * Ξ[i - 1] - Ξ[i - 2]
        end
    else
        Ψ, Χ = cspherebessel(n, ds)
        for i in 1:n
            Ξ[i] = (Ψ[i] + ci * Χ[i]) * ds
        end
    end
    return Ξ
end

"""
    cspherebessel(n, z)

Compute spherical Bessel functions jn(z) & yn(z)
"""
function cspherebessel(n::Int64, z::ComplexF64)
    a0 = abs(z)
    nm = n
    csj = OffsetArray(zeros(ComplexF64, n + 1), 0:n)
    csy = OffsetArray(zeros(ComplexF64, n + 1), 0:n)
    if a0 < 1e-60
        fill!(csy, ComplexF64(-1e300, 0))
        csy[0] = ComplexF64(1, 0)
        return csj, csy
    end
    csj[0] = sin(z) / z
    csj[1] = (csj[0] - cos(z)) / z
    if n >= 2
        csa = csj[0]
        csb = csj[1]
        m = msta1(a0, 200)
        if m < n
            nm = m
        else
            m = msta2(a0, n, 15)
        end
        cf = 0.0
        cf0 = 0.0
        cf1 = 1.0 - 100
        for k in m:-1:0
            cf = (2.0 * k + 3.0) * cf1 / z - cf0
            if k <= nm
                csj[k] = cf
            end
            cf0 = cf1
            cf1 = cf
        end
        if (abs(csa) > abs(csb))
            cs = csa / cf
        else
            cs = csb / cf0
        end
        for k in 0:min(nm, n)
            csj[k] = cs * csj[k]
        end
    end
    fill!(csy, ComplexF64(1e200, 0))
    csy[0] = -cos(z) / z
    csy[1] = (csy[0] - sin(z)) / z
    for k in 2:min(nm, n)
        if abs(csj[k - 1]) > abs(csj[k - 2])
            csy[k] = (csj[k] * csy[k - 1] - 1.0 / (z * z)) / csj[k - 1]
        else
            csy[k] = (csj[k] * csy[k - 2] - (2.0 * k - 1.0) / z^3) / csj[k - 2]
        end
    end

    return csj, csy
end

"""
    msta1(x, mp)

Determine the starting point for backward recurrence such that the magnitude of Jn(x) at that point is about 10^(-MP).
"""
function msta1(x::Float64, mp::Int64)
    a0 = abs(x)
    n0 = floor(Int, 1.1 * a0) + 1
    f0 = envj(n0, a0) - mp
    n1 = n0 + 5
    f1 = envj(n1, a0) - mp
    nn = 0
    for it in 1:20
        nn = floor(Int, n1 - (n1 - n0) / (1 - f0 / f1))
        f = envj(nn, a0) - mp
        if abs(nn - n1) < 1
            break
        end
        n0 = n1
        f0 = f1
        n1 = nn
        f1 = f
    end
    return nn
end

function msta2(x::Float64, n::Int64, mp::Int64)
    a0 = abs(x)
    hmp = 0.5mp
    ejn = envj(n, a0)
    if ejn <= hmp
        obj = mp
        n0 = floor(Int, 1.1 * a0)
    else
        obj = hmp + ejn
        n0 = n
    end
    f0 = envj(n0, a0) - obj
    n1 = n0 + 5
    f1 = envj(n1, a0) - obj
    nn = 0
    for it in 1:20
        nn = floor(Int, n1 - (n1 - n0) / (1.0 - f0 / f1))
        f = envj(nn, a0) - obj
        if abs(nn - n1) < 1
            break
        end
        n0 = n1
        f0 = f1
        n1 = nn
        f1 = f
    end
    return nn + 10
end

function envj(n::Int64, x::Float64)
    n = max(1, abs(n))
    return 0.5log10(6.28n) - n * log10(1.36 * x / n)
end

function vcfunc(ctx::ConstantContext, m::Int64, n::Int64, k::Int64, l::Int64)
    vcn = OffsetArray(zeros(n + l + 1), 0:(n + l))
    vcfunc!(ctx, m, n, k, l, vcn)
    return vcn
end

function vcfunc!(ctx::ConstantContext, m::Int64, n::Int64, k::Int64, l::Int64, vcn::OffsetArray{Float64,1})
    init!(ctx, n + l)

    bcof, fnr, _, _ = get_offset_constants(ctx)
    wmax = n + l
    wmin = max(abs(n - l), abs(m + k))
    fill!(vcn, 0)
    vcn[wmax] = bcof[n + m, l + k] * bcof[n - m, l - k] / bcof[n + n, l + l]
    if wmin == wmax
        return
    end
    vcn[wmax - 1] =
        vcn[wmax] * (l * m - k * n) * fnr[2(l + n) - 1] / fnr[l] / fnr[n] / fnr[n + l + m + k] / fnr[n + l - m - k]
    if wmin == wmax - 1
        return
    end
    mk = m + k
    vcmax = abs(vcn[wmax]) + abs(vcn[wmax - 1])

    # a downwards recurrence is used initially
    for w in wmax:-1:(wmin + 2)
        t1::Float64 =
            2w * fnr[2w + 1] * fnr[2w - 1] /
            (fnr[w + mk] * fnr[w - mk] * fnr[n - l + w] * fnr[l - n + w] * fnr[n + l - w + 1] * fnr[n + l + w + 1])
        t2 = ((m - k) * w * (w - 1) - mk * n * (n + 1) + mk * l * (l + 1)) / (2w * (w - 1))
        t3::Float64 =
            fnr[w - mk - 1] *
            fnr[w + mk - 1] *
            fnr[l - n + w - 1] *
            fnr[n - l + w - 1] *
            fnr[n + l - w + 2] *
            fnr[n + l + w] / ((2(w - 1)) * fnr[2w - 3] * fnr[2w - 1])
        vcn[w - 2] = (t2 * vcn[w - 1] - vcn[w] / t1) / t3
        if (wmax - w) % 2 == 1
            vctest = abs(vcn[w - 2]) + abs(vcn[w - 1])
            vcmax = max(vcmax, vctest)
            rat = vctest / vcmax
            # if/when the coefficients start to decrease in magnitude, an upwards recurrence takes over
            if rat < 0.01 && w - 2 > wmin
                wmax = w - 3
                vcfuncuprec!(ctx, m, n, k, l, wmax, vcn)
                break
            end
        end
    end
end

"""
    vcfuncuprec(ctx, m, n, k, l, wmax, vcn)

Upwards VC coefficient recurrence, used internally by vcfunc(ctx, m, n, k, l).
"""
function vcfuncuprec!(
    ctx::ConstantContext,
    m::Int64,
    n::Int64,
    k::Int64,
    l::Int64,
    wmax::Int64,
    vcn::OffsetArray{Float64,1},
)
    bcof, fnr, monen, _ = get_offset_constants(ctx)

    mk = abs(m + k)
    nl = abs(n - l)
    if nl >= mk
        w = nl
        if n >= l
            m1 = m
            n1 = n
            l1 = l
            k1 = k
        else
            m1 = k
            n1 = l
            k1 = m
            l1 = n
        end
        vc1 = monen[k1 + l1] * bcof[l1 + k1, w - m1 - k1] * bcof[l1 - k1, w + m1 + k1] / bcof[2l1, 2w + 1]
    else
        w = mk
        if m + k >= 0
            vc1 = monen[n + m] * bcof[n - l + w, l - k] * bcof[l - n + w, n - m] / bcof[2w + 1, n + l - w]
        else
            vc1 = monen[l + k] * bcof[n - l + w, l + k] * bcof[l - n + w, n + m] / bcof[2w + 1, n + l - w]
        end
    end

    w1 = w
    vcn[w] = vc1
    w = w1 + 1
    mk = m + k
    w2 = min(wmax, n + l)
    if w2 > w1
        t1 =
            2 * w * fnr[2w + 1] * fnr[2w - 1] /
            (fnr[w + mk] * fnr[w - mk] * fnr[n - l + w] * fnr[l - n + w] * fnr[n + l - w + 1] * fnr[n + l + w + 1])
        if w1 == 0
            t2 = 0.5 * (m - k)
        else
            t2 = ((m - k) * w * (w - 1) - mk * n * (n + 1) + mk * l * (l + 1)) / (2w * (w - 1))
        end
        vcn[w] = t1 * t2 * vcn[w1]
    end

    for w in (w1 + 2):w2
        t1 =
            2w * fnr[2w + 1] * fnr[2w - 1] /
            (fnr[w + mk] * fnr[w - mk] * fnr[n - l + w] * fnr[l - n + w] * fnr[n + l - w + 1] * fnr[n + l + w + 1])
        t2 = ((m - k) * w * (w - 1) - mk * n * (n + 1) + mk * l * (l + 1)) / (2w * (w - 1))
        t3 =
            fnr[w - mk - 1] *
            fnr[w + mk - 1] *
            fnr[l - n + w - 1] *
            fnr[n - l + w - 1] *
            fnr[n + l - w + 2] *
            fnr[n + l + w] / ((2 * (w - 1)) * fnr[2w - 3] * fnr[2w - 1])
        vcn[w] = t1 * (t2 * vcn[w - 1] - t3 * vcn[w - 2])
    end

    return
end

function normalizedlegendre(ctx::ConstantContext, cbe::Float64, mmax::Int64, nmax::Int64)
    init!(ctx, mmax + nmax)

    bcof, fnr, monen, _ = get_offset_constants(ctx)
    dc = OffsetArray(zeros(2mmax + 1, nmax + 1), (-mmax):mmax, 0:nmax)

    # cbe must ∈ [-1, 1] otherwise sbe will be a complex number
    sbe = √((1 + cbe) * (1 - cbe))
    for m in 0:mmax
        dc[m, m] = monen[m] * (0.5 * sbe)^m * bcof[m, m]
        if m == nmax
            break
        end
        dc[m, m + 1] = fnr[2m + 1] * cbe * dc[m, m]
        for n in (m + 1):(nmax - 1)
            dc[m, n + 1] =
                (-fnr[n - m] * fnr[n + m] * dc[m, n - 1] + (2n + 1) * cbe * dc[m, n]) /
                (fnr[n + 1 - m] * fnr[n + 1 + m])
        end
    end

    for m in 1:mmax
        for n in m:nmax
            dc[-m, n] = monen[m] * dc[m, n]
        end
    end

    return dc
end

function rotcoef(ctx::ConstantContext, cbe::Float64, kmax::Int64, nmax::Int64)
    init!(ctx, kmax + nmax)
    bcof, fnr, monen, _ = get_offset_constants(ctx)

    dc = OffsetArray(zeros(2kmax + 1, nmax * (nmax + 2) + 1), (-kmax):kmax, 0:(nmax * (nmax + 2)))
    dk0 = OffsetArray(zeros(2nmax + 3), (-nmax - 1):(nmax + 1))
    dk01 = OffsetArray(zeros(2nmax + 3), (-nmax - 1):(nmax + 1))

    # cbe must ∈ [-1, 1] otherwise sbe will be a complex number
    sbe = √((1 + cbe) * (1 - cbe))
    cbe2 = 0.5(1 + cbe)
    sbe2 = 0.5(1 - cbe)

    dk0[0] = 1
    sben = 1.0
    dc[0, 0] = 1
    dk01[0] = 0

    for n in 1:nmax
        knmax = min(n, kmax)
        nn1 = n * (n + 1)
        sben *= sbe / 2
        dk0[n] = monen[n] * sben * bcof[n, n]
        dk0[-n] = monen[n] * dk0[n]
        dk01[n] = 0
        dk01[-n] = 0
        dc[0, nn1 + n] = dk0[n]
        dc[0, nn1 - n] = dk0[-n]

        for k in (-n + 1):(n - 1)
            kn = nn1 + k
            dkt = dk01[k]
            dk01[k] = dk0[k]
            dk0[k] = (cbe * (2n - 1) * dk01[k] - fnr[n - k - 1] * fnr[n + k - 1] * dkt) / (fnr[n + k] * fnr[n - k])
            dc[0, kn] = dk0[k]
        end

        for m in 1:knmax
            fmn = 1 / fnr[n - m + 1] / fnr[n + m]
            m1 = m - 1
            dkm0 = 0.0

            for k in (-n):n
                kn = nn1 + k
                dkm1 = dkm0
                dkm0 = dc[m1, kn]
                dkn1 = k == n ? 0.0 : dc[m1, kn + 1]
                dc[m, kn] =
                    (
                        fnr[n + k] * fnr[n - k + 1] * cbe2 * dkm1 - fnr[n - k] * fnr[n + k + 1] * sbe2 * dkn1 -
                        k * sbe * dc[m1, kn]
                    ) * fmn
                dc[-m, nn1 - k] = dc[m, kn] * monen[k + m]
            end
        end
    end

    return dc
end

"""
    taufunc(ctx, cb, nmax)

Normalized vector spherical harmonic function.
"""
function taufunc(ctx::ConstantContext, cb::Float64, nmax::Int64)
    drot = rotcoef(ctx, cb, 1, nmax)
    τ = OffsetArray(zeros(nmax + 2, nmax, 2), 0:(nmax + 1), 1:nmax, 1:2)
    for n in 1:nmax
        nn1 = n * (n + 1)
        fnm = √((2n + 1) / 2) / 4

        for m in (-n):-1
            mn = nn1 + m
            τ[n + 1, -m, 1] = -fnm * (-drot[-1, mn] + drot[1, mn])
            τ[n + 1, -m, 2] = -fnm * (drot[-1, mn] + drot[1, mn])
        end

        for m in 0:n
            mn = nn1 + m
            τ[m, n, 1] = -fnm * (-drot[-1, mn] + drot[1, mn])
            τ[m, n, 2] = -fnm * (drot[-1, mn] + drot[1, mn])
        end
    end
    return τ
end

"""
    pifunc(cb, ephi, nmax, ndim)

Vector spherical harmonic function.

TODO: `ndim` seems to be useless. Can it be safely deleted?
"""
function pifunc(ctx::ConstantContext, cb::Float64, E_ϕ::ComplexF64, nmax::Int64, ndim::Int64)
    drot = rotcoef(ctx, cb, 1, nmax)

    _, fnr, _, _ = get_offset_constants(ctx)

    E_ϕm = OffsetArray(zeros(ComplexF64, 2nmax + 1), (-nmax):nmax)
    E_ϕm[0] = 1.0

    # ndim should be as large as or larger than nmax
    π_vec = OffsetArray(zeros(ComplexF64, ndim + 2, ndim, 2), 0:(ndim + 1), 1:ndim, 1:2)

    for m in 1:nmax
        E_ϕm[m] = E_ϕ * E_ϕm[m - 1]
        E_ϕm[-m] = conj(E_ϕm[m])
    end

    for n in 1:nmax
        cin = (-1im)^(n) * fnr[2n + 1]
        nn1 = n * (n + 1)
        for m in n:-1:1
            π_vec[n + 1, m, 1] = cin * drot[1, nn1 - m] * E_ϕm[-m]
            π_vec[n + 1, m, 2] = cin * drot[-1, nn1 - m] * E_ϕm[-m]
        end

        for m in 0:n
            π_vec[m, n, 1] = cin * drot[1, nn1 + m] * E_ϕm[m]
            π_vec[m, n, 2] = cin * drot[-1, nn1 + m] * E_ϕm[m]
        end
    end

    return π_vec
end

function planewavecoef(ctx::ConstantContext, α::Float64, β::Float64, nodr::Int64)
    init!(ctx, nodr)

    _, _, monen, _ = get_offset_constants(ctx)

    τ_lr = OffsetArray(zeros(nodr + 2, nodr, 2), 0:(nodr + 1), 1:nodr, 1:2)
    pmnp0 = OffsetArray(zeros(ComplexF64, nodr + 2, nodr, 2, 2), 0:(nodr + 1), 1:nodr, 1:2, 1:2)

    cb = cos(β)
    sb = √((1 - cb) * (1 + cb))
    ca = cos(α)
    sa = sin(α)
    E_α = ComplexF64(ca, sa)
    τ = taufunc(ctx, cb, nodr)
    τ_lr[:, :, 1] = (τ[:, :, 1] .+ τ[:, :, 2]) * 0.5
    τ_lr[:, :, 2] = (τ[:, :, 1] .- τ[:, :, 2]) * 0.5
    E_αm = ephicoef(E_α, nodr)

    for n in 1:nodr
        cin = 4 * ci^(n + 1)
        for p in 1:2
            sp = -monen[p]
            for m in (-n):-1
                pmnp0[n + 1, -m, p, 1] = -cin * τ_lr[n + 1, -m, p] * E_αm[-m]
                pmnp0[n + 1, -m, p, 2] = sp * ci * cin * τ_lr[n + 1, -m, p] * E_αm[-m]
            end

            for m in 0:n
                pmnp0[m, n, p, 1] = -cin * τ_lr[m, n, p] * E_αm[-m]
                pmnp0[m, n, p, 2] = sp * ci * cin * τ_lr[m, n, p] * E_αm[-m]
            end
        end
    end

    return pmnp0
end

"""
    gaussianbeamcoef(ctx, α, β, cbeam, nodr)

Regular VSWF expansion coefficients for a gaussian beam, localized approximation.

cbeam = 1/(k omega)
"""
function gaussianbeamcoef(ctx::ConstantContext, α::Float64, β::Float64, cbeam::Float64, nodr::Int64)
    pmnp0 = planewavecoef(ctx, α, β, nodr)
    for n in 1:nodr
        gbn = exp(-((n + 0.5) * cbeam)^2)
        for p in 1:2
            for k in 1:2
                for m in (-n):-1
                    pmnp0[n + 1, -m, p, k] *= gbn
                end

                for m in 0:n
                    pmnp0[m, n, p, k] *= gbn
                end
            end
        end
    end
    return pmnp0
end

function ephicoef(E_ϕ::ComplexF64, nodr::Int64)
    E_ϕm = OffsetArray(zeros(ComplexF64, 2nodr + 1), (-nodr):nodr)
    E_ϕm[0] = 1.0
    for m in 1:nodr
        E_ϕm[m] = E_ϕ * E_ϕm[m - 1]
        E_ϕm[-m] = conj(E_ϕm[m])
    end
    return E_ϕm
end

"""
    sphereplanewavecoef(ctx, nsphere, neqns, nodr, nodrmax, α, β, rpos, hostsphere, numberfieldexp, rimedium)

Plane wave expansion coefficients at sphere origins. A phase shift is used.
"""
function sphereplanewavecoef(
    ctx::ConstantContext,
    nodr::Array{Int64,1},
    α::Float64,
    β::Float64,
    rpos::Array{Float64,2},
    hostsphere::Array{Int64,1},
    numberfieldexp::Array{Int64,1},
    rimedium::Array{ComplexF64,1},
)
    nsphere = length(nodr)
    nodrmax = maximum(nodr)

    # Valid the input
    @assert length(hostsphere) == nsphere
    @assert length(numberfieldexp) == nsphere
    @assert size(rpos) == (3, nsphere)
    @assert length(rimedium) == 2

    rib = 2 / (1 / rimedium[1] + 1 / rimedium[2])
    pmnp0 = planewavecoef(ctx, α, β, nodrmax)
    cb = cos(β)
    sb = √((1 - cb) * (1 + cb))
    ca = cos(α)
    sa = sin(α)

    neqns = 2 * sum(numberfieldexp[i] * nodr[i] * (nodr[i] + 2) for i in 1:nsphere)
    pmnp = zeros(ComplexF64, 2neqns)

    l = 0
    for i in 1:nsphere
        phasefac = exp(ci * rib * ((ca * rpos[1, i] + sa * rpos[2, i]) * sb + rpos[3, i] * cb))
        for j in 1:numberfieldexp[i]
            for k in 1:2
                for p in 1:2
                    for n in 1:nodr[i]
                        for m in 0:(nodr[i] + 1)
                            l += 1
                            if hostsphere[i] == 0 && j == 1
                                pmnp[l] = phasefac * pmnp0[m, n, p, k]
                            else
                                pmnp[l] = 0
                            end
                        end
                    end
                end
            end
        end
    end

    return pmnp
end

function axialtrancoefrecurrence(
    ctx::ConstantContext,
    itype::Int64,
    r::Float64,
    ri::Array{ComplexF64,1},
    nmax::Int64,
    lmax::Int64,
)
    @assert itype == 1 || itype == 3
    @assert length(ri) == 2

    nlmax = max(nmax, lmax)
    nlmin = min(nmax, lmax)

    nlmax0 = 0
    if nlmax > 0
        nlmax0 = nlmax
        axialtrancoefinit!(ctx, nlmax)
    end

    ndim = atcdim(nmax, lmax)
    ac = zeros(ComplexF64, ndim)
    act = zeros(ComplexF64, nmax, lmax, 2)
    actt = zeros(ComplexF64, 2, 2)

    if iszero(r)
        if itype != 1
            return ac
        end
        iadd0 = 0
        for ma in 0:nlmin
            m1 = max(1, ma)
            for m in (-ma):(2m1):ma
                blockdim = 2(nmax - m1 + 1) * (lmax - m1 + 1)
                iadd1 = iadd0 + blockdim
                for l in m1:nlmax
                    act[l, l, 1] = act[l, l, 2] = 1
                end
                ac[(iadd0 + 1):iadd1] = reshape(act[m1:nmax, m1:lmax, 1:2], (blockdim,))
                iadd0 = iadd1
            end
        end
        return ac
    end

    z = r * ri
    Ξ = OffsetArray(zeros(ComplexF64, nmax + lmax + 1, 2), 0:(nmax + lmax), 1:2)
    for p in 1:2
        Ξ[:, p] = itype == 1 ? cricbessel(nmax + lmax, z[p]) / z[p] : crichankel(nmax + lmax, z[p]) / z[p]
    end

    vcc, fnm1, fn, fnp1 = get_offset_axial_constants(ctx)

    lm1 = lmax - 1
    iadd0 = 0
    for ma in 0:nlmin
        m1 = max(1, ma)
        lp1 = m1 + 1
        for m in (-ma):(2m1):ma
            blockdim = 2(nmax - m1 + 1) * (lmax - m1 + 1)
            iadd1 = iadd0 + blockdim
            n = m1
            for l in m1:lmax
                wmin = abs(n - l)
                wmax = n + l
                ml = l * (l + 1) + m
                for p in 1:2
                    actt[1, p] = actt[2, p] = 0
                    for i in wmin:2:wmax
                        actt[1, p] += vcc[n, ml, i] * Ξ[i, p]
                    end
                    for i in (wmin + 1):2:(wmax - 1)
                        actt[2, p] += vcc[n, ml, i] * Ξ[i, p]
                    end
                    actt[2, p] *= ci
                end
                act[n, l, 1] = actt[1, 1] + actt[2, 1]
                act[n, l, 2] = actt[1, 2] - actt[2, 2]
            end
            l = lmax
            ml = l * (l + 1) + m
            for n in (m1 + 1):nmax
                wmin = abs(n - l)
                wmax = n + l
                for p in 1:2
                    actt[1, p] = actt[2, p] = 0
                    for i in wmin:2:wmax
                        actt[1, p] += vcc[n, ml, i] * Ξ[i, p]
                    end
                    for i in (wmin + 1):2:(wmax - 1)
                        actt[2, p] += vcc[n, ml, i] * Ξ[i, p]
                    end
                    actt[2, p] *= ci
                end
                act[n, l, 1] = actt[1, 1] + actt[2, 1]
                act[n, l, 2] = actt[1, 2] - actt[2, 2]
            end

            if m1 < nlmin
                for n in m1:(nmax - 1)
                    np1 = n + 1
                    nm1 = n - 1
                    for p in 1:2
                        sp = -(-1)^p
                        for i in m1:lm1
                            act[np1, i, p] =
                                -act[n, i + 1, p] * fnp1[m, i] + (fn[m, i] - fn[m, n]) * act[n, i, p] * sp * ci
                        end
                        for i in lp1:lm1
                            act[np1, i, p] += act[n, i - 1, p] * fnm1[m, i]
                        end
                        if n > m1
                            for i in m1:lm1
                                act[np1, i, p] += act[nm1, i, p] * fnm1[m, n]
                            end
                        end
                        for i in m1:lm1
                            act[np1, i, p] /= fnp1[m, n]
                        end
                    end
                end
            end
            ac[(iadd0 + 1):iadd1] = reshape(act[m1:nmax, m1:lmax, 1:2], (blockdim,))
            iadd0 = iadd1
        end
    end

    return ac
end

function axialtrancoefinit!(ctx::ConstantContext, nmax::Int64)
    init!(ctx, nmax)

    ctx.vcc = zeros(nmax, nmax * (nmax + 2), 2nmax + 1)
    ctx.fnm1 = zeros(2nmax + 1, nmax)
    ctx.fn = zeros(2nmax + 1, nmax)
    ctx.fnp1 = zeros(2nmax + 1, nmax)

    _, fnr, monen, _ = get_offset_constants(ctx)
    vcc, fnm1, fn, fnp1 = get_offset_axial_constants(ctx)
    vc1 = OffsetArray(zeros(2nmax + 1), 0:(2nmax))
    vc2 = OffsetArray(zeros(2nmax + 1), 0:(2nmax))
    for n in 1:nmax
        n21 = 2n + 1
        for l in 1:nmax
            c1 = fnr[n21] * fnr[2l + 1]
            ll1 = l * (l + 1)
            vcfunc!(ctx, -1, n, 1, l, vc2)
            wmin = abs(n - l)
            wmax = n + l
            nlmin = min(n, l)
            for m in (-nlmin):nlmin
                ml = ll1 + m
                c2 = -c1 * monen[m]
                vcfunc!(ctx, -m, n, m, l, vc1)
                for w in wmin:wmax
                    inlw = ci^(n - l + w)
                    vcc[n, ml, w] = c2 * vc1[w] * vc2[w] * (real(inlw) + imag(inlw))
                end
            end
        end
    end

    for m in (-nmax):nmax
        for l in max(1, abs(m)):nmax
            lp1 = l + 1
            lm1 = l - 1
            fnm1[m, l] = fnr[lm1] * fnr[lp1] * fnr[l - m] * fnr[l + m] / fnr[lm1 + l] / fnr[lp1 + l] / l
            fn[m, l] = m / l / lp1
            fnp1[m, l] = fnr[l] * fnr[l + 2] * fnr[lp1 - m] * fnr[lp1 + m] / fnr[l + lp1] / fnr[2l + 3] / lp1
        end
    end
end

"""
    tranordertest(ctx, r, ri, lmax, ϵ)

Test to determine convergence of regular VSWF addition theorem for max order lmax and translation distance r w/ refractive index ri.
"""
function tranordertest(ctx::ConstantContext, r::Float64, ri::ComplexF64, lmax::Int64, ϵ::Float64)
    if iszero(r)
        return lmax
    end

    nlim = 200
    init!(ctx, nlim + lmax)
    _, fnr, _, _ = get_offset_constants(ctx)

    z = r * ri
    vc1 = OffsetArray(zeros(nlim + lmax + 1), 0:(nlim + lmax))
    s = 0.0

    for n in 1:nlim
        Ξ = cricbessel(n + lmax, z)
        for l in 0:(n + lmax)
            Ξ[l] = Ξ[l] / z * ci^l
        end
        n21 = 2n + 1
        l = lmax
        c = fnr[n21] * fnr[2l + 1] * ci^(n - l)
        vcfunc!(ctx, -1, n, 1, l, vc1)
        wmin = abs(n - l)
        wmax = n + l
        m = 1
        a = 0.0im
        b = 0.0im
        for w in wmin:wmax
            alnw = vc1[w]^2
            if (n + l + w) & 1 == 0
                a += alnw * Ξ[w]
            else
                b += alnw * Ξ[w]
            end
        end
        a *= c
        b *= c
        s += a * conj(a) + b * conj(b)
        if abs(1.0 - s) < ϵ
            return max(n, lmax)
        end
    end

    return max(nlim, lmax)
end

function atcadd(m::Int64, n::Int64, ntot::Int64)
    return n - ntot + (max(1, m) * (1 + 2ntot - max(1, m))) ÷ 2 + ntot * min(1, m)
end

function atcdim(ntot::Int64, ltot::Int64)
    nmin = min(ntot, ltot)
    nmax = max(ntot, ltot)
    return 2(nmin * (1 - nmin * nmin + 3nmax * (nmin + 2))) ÷ 3
end

function moffset(m::Int64, ntot::Int64, ltot::Int64)
    if m == 0
        return 0
    elseif m < 0
        return 2 * (-((1 + m) * (2 + m) * (3 + 2m + 3ntot)) - 3ltot * (2 + ntot + m * (3 + m + 2ntot))) ÷ 3
    else
        return 2 * (-3ltot * (m - 1)^2 + 6ltot * m * ntot + (m - 1) * (m * (-4 + 2m - 3ntot) + 3(ntot + 1))) ÷ 3
    end
end

"""
    gentrancoef(ctx, itype, xptran, ri, nrow0, nrow1, ncol0, ncol1, iaddrow0, iaddcol0)

Calculates the vwh translation coefficients for a general translation from one origin to another.
"""
function gentrancoef(
    ctx::ConstantContext,
    itype::Int64,
    xptran::Array{Float64,1},
    ri::Array{ComplexF64,1},
    nrow0::Int64,
    nrow1::Int64,
    ncol0::Int64,
    ncol1::Int64,
    iaddrow0::Int64,
    iaddcol0::Int64,
)
    @assert itype == 1 || itype == 3
    @assert length(xptran) == 3
    @assert length(ri) == 2
    @assert 1 <= nrow0 <= nrow1
    @assert 1 <= ncol0 <= ncol1

    ntot = nrow1 + ncol1
    jnc = OffsetArray(zeros(ComplexF64, ntot + 1, 2), 0:ntot, 1:2)
    r, ct, E_ϕ = cartosphere(xptran)
    nblkr0 = (nrow0 - 1) * (nrow0 + 1)
    nblkr1 = nrow1 * (nrow1 + 2)
    nblkc0 = (ncol0 - 1) * (ncol0 + 1)
    nblkc1 = ncol1 * (ncol1 + 2)

    ac = zeros(ComplexF64, 2, nblkr1 - nblkr0 + iaddrow0, nblkc1 - nblkc0 + iaddcol0)

    if iszero(r)
        for n in (nblkr0 + 1):nblkr1
            mn = n - nblkr0 + iaddrow0
            if n > nblkc0 && n <= nblkc1 && itype == 1
                ac[1, mn, n - nblkc0 + iaddcol0] = 1
                ac[2, mn, n - nblkc0 + iaddcol0] = 1
            end
        end
        return ac
    end

    kmax = 0
    ct0 = ct
    drot = rotcoef(ctx, ct0, kmax, ntot)
    E_ϕm = ephicoef(E_ϕ, ntot)
    z = ri * r

    for p in 1:2
        jnc[:, p] = itype == 1 ? cricbessel(ntot, z[p]) / z[p] : crichankel(ntot, z[p]) / z[p]
        for n in 0:ntot
            jnc[n, p] *= ci^n
        end
    end

    _, fnr, monen, _ = get_offset_constants(ctx)

    vc1 = OffsetArray(zeros(ntot + 1), 0:ntot)
    vc2 = OffsetArray(zeros(ntot + 1), 0:ntot)
    atc = zeros(ComplexF64, 2, 2)
    for l in ncol0:ncol1
        ll1 = l * (l + 1)
        for n in nrow0:nrow1
            nn1 = n * (n + 1)
            wmax = n + l
            vcfunc!(ctx, -1, n, 1, l, vc2)
            c = -ci^(n - l) * fnr[2n + 1] * fnr[2l + 1]
            for k in (-l):l
                kl = ll1 + k - nblkc0 + iaddcol0
                for m in (-n):n
                    m1m = monen[m]
                    mn = nn1 + m - nblkr0 + iaddrow0
                    v = k - m
                    vcfunc!(ctx, -m, n, k, l, vc1)
                    a = 0.0
                    b = 0.0
                    wmin = max(abs(v), abs(n - l))
                    for p in 1:2
                        for w in wmax:-1:wmin
                            vw = w * (w + 1) + v
                            if (wmax - w) & 1 == 0
                                a += vc1[w] * vc2[w] * jnc[w, p] * drot[0, vw]
                            else
                                b += vc1[w] * vc2[w] * jnc[w, p] * drot[0, vw]
                            end
                        end
                        atc[p, 1] = a
                        atc[p, 2] = b
                    end
                    ac[1, mn, kl] = (atc[1, 1] + atc[1, 2]) * c * m1m * E_ϕm[v]
                    ac[2, mn, kl] = (atc[2, 1] - atc[2, 2]) * c * m1m * E_ϕm[v]
                end
            end
        end
    end

    return ac
end

function cartosphere(xp::Array{Float64,1})
    @assert length(xp) == 3

    r = xp[1]^2 + xp[2]^2 + xp[3]^2
    if iszero(r)
        return r, 1.0, 1.0 + 0.0im
    end

    r = √r
    ct = xp[3] / r
    E_ϕ = (iszero(xp[1]) && iszero(xp[2])) ? 1.0 + 0.0im : ComplexF64(xp[1], xp[2]) / √(xp[1]^2 + xp[2]^2)

    return r, ct, E_ϕ
end

"""
    eulerrotation(xp)

Euler rotation of a point (xp[1], xp[2], xp[3]).
"""
function eulerrotation(xp::Array{Float64,1}, eulerangf::Array{Float64,1}, dir::Int64)
    @assert length(xp) == 3
    @assert length(eulerangf) == 3

    eulerang = dir == 1 ? eulerangf : -eulerangf[3:-1:1]

    cang = [cos(x) for x in eulerang]
    sang = [sin(x) for x in eulerang]

    mat1 = [[cang[1], -sang[1], 0.0] [sang[1], cang[1], 0.0] [0.0, 0.0, 1.0]]
    mat2 = [[cang[2], 0.0, sang[2]] [0.0, 1.0, 0.0] [-sang[2], 0.0, cang[2]]]
    mat3 = [[cang[3], -sang[3], 0.0] [sang[3], cang[3], 0.0] [0.0, 0.0, 1.0]]

    return mat3 * (mat2 * (mat1 * xp))
end

function planewavetruncationorder(r::Float64, rimedium::Array{ComplexF64}, ϵ::Float64)
    @assert length(rimedium) == 2

    rib = 2.0 / (1.0 / rimedium[1] + 1.0 / rimedium[2])
    n1 = max(10, floor(Int, 3r + 1))
    rri = r * rib
    jn = cricbessel(n1, rri) / rri
    eir = exp(-ci * rri)
    s = jn[0] * eir

    for n in 1:n1
        s += ci^n * (2n + 1) * jn[n] * eir
        err = abs(1.0 - s)
        if err < ϵ
            return n
        end
    end

    return n1
end

function vwhcalc(ctx::ConstantContext, rpos::Array{Float64,1}, ri::Array{ComplexF64,1}, nodr::Int64, itype::Int64)
    @assert itype == 1 || itype == 3
    @assert length(rpos) == 3
    @assert length(ri) == 2

    a1vec = OffsetArray(zeros(ComplexF64, 2nodr + 1), (-nodr):nodr)
    b1vec = OffsetArray(zeros(ComplexF64, 2nodr + 1), (-nodr):nodr)
    z1vec = OffsetArray(zeros(ComplexF64, 2nodr + 1), (-nodr):nodr)
    a2vec = OffsetArray(zeros(ComplexF64, 2nodr + 1), (-nodr):nodr)
    b2vec = OffsetArray(zeros(ComplexF64, 2nodr + 1), (-nodr):nodr)
    z2vec = OffsetArray(zeros(ComplexF64, 2nodr + 1), (-nodr):nodr)
    umn = OffsetArray(zeros(ComplexF64, 2nodr + 5, nodr + 2, 2), (-nodr - 2):(nodr + 2), 0:(nodr + 1), 1:2)
    hn = OffsetArray(zeros(ComplexF64, nodr + 2, 2), 0:(nodr + 1), 1:2)

    nodrmax = 0
    if nodr > nodrmax
        nodrmax = nodr
        init!(ctx, nodr + 2)
    end

    r, ct, E_ϕ = cartosphere(rpos)
    vwh = zeros(ComplexF64, 3, 2, nodr * (nodr + 2))
    _, fnr, monen, vwh_coef = get_offset_constants(ctx)
    if r <= 1e-4
        if itype == 3
            return vwh
        end
        for p in 1:2
            vwh[1, p, 1] = 0.5fnr[2] / fnr[3]
            vwh[2, p, 3] = vwh[2, p, 1] = -0.5ci * fnr[2] / fnr[3]
            vwh[3, p, 2] = fnr[2] / fnr[6]
            vwh[1, p, 3] = -.5fnr[2] / fnr[3]
        end
        return vwh
    end
    nodrp1 = nodr + 1
    nodrm1 = nodr - 1

    a = ri * r
    for p in 1:2
        hn[:, p] = (itype == 1 ? cricbessel(nodrp1, a[p]) : crichankel(nodrp1, a[p])) / a[p]
    end

    pmn = rotcoef(ctx, ct, 0, nodrp1)
    E_ϕm = ephicoef(E_ϕ, nodrp1)

    for p in 1:2
        umn[0, 0, p] = hn[0, p] * fnr[2]
        for n in 1:nodrp1
            nn1 = n * (n + 1)
            coeff = fnr[2] * hn[n, p]
            for i in (-n):n
                umn[i, n, p] = pmn[0, nn1 + i] * E_ϕm[i] * coeff
            end
        end
    end

    for p in 1:2
        sp = -monen[p]
        for n in 1:nodr
            nn1 = n * (n + 1)
            np1 = n + 1
            nm1 = n - 1
            for i in (-n):n
                a1vec[i] = vwh_coef[i, n, 1, 1] * umn[i + 1, np1, p] + vwh_coef[i, n, 1, -1] * umn[i + 1, nm1, p]
                b1vec[i] = vwh_coef[i, n, -1, 1] * umn[i - 1, np1, p] + vwh_coef[i, n, -1, -1] * umn[i - 1, nm1, p]
                z1vec[i] = vwh_coef[i, n, 0, 1] * umn[i, np1, p] + vwh_coef[i, n, 0, -1] * umn[i, nm1, p]
                a2vec[i] = vwh_coef[i, n, 1, 0] * umn[i + 1, n, p]
                b2vec[i] = vwh_coef[i, n, -1, 0] * umn[i - 1, n, p]
                z2vec[i] = vwh_coef[i, n, 0, 0] * umn[i, n, p]
                vwh[1, p, nn1 + i] = -0.5(a1vec[i] + b1vec[i]) - 0.5ci * sp * (a2vec[i] + b2vec[i])
                vwh[2, p, nn1 + i] = -0.5ci * (-a1vec[i] + b1vec[i]) - 0.5sp * (a2vec[i] - b2vec[i])
                vwh[3, p, nn1 + i] = -z1vec[i] - sp * ci * z2vec[i]
            end
        end
    end

    return vwh
end

function vwhaxialcalc(ctx::ConstantContext, rpos::Array{Float64,1}, ri::Array{ComplexF64,1}, nodr::Int64, itype::Int64)
    @assert itype == 1 || itype == 3
    @assert length(rpos) == 3
    @assert length(ri) == 2

    umn = OffsetArray(zeros(ComplexF64, 5, nodr + 2, 2), -2:2, 0:(nodr + 1), 1:2)
    hn = OffsetArray(zeros(ComplexF64, nodr + 2, 2), 0:(nodr + 1), 1:2)

    nodrmax = 0
    if nodr > nodrmax
        nodrmax = nodr
        init!(ctx, nodr + 2)
    end

    r, ct, E_ϕ = cartosphere(rpos)
    vwh = zeros(ComplexF64, 3, 2, 2, nodr)
    _, fnr, monen, vwh_coef = get_offset_constants(ctx)
    if r <= 1e-4
        if itype == 3
            return vwh
        end
        vwh[1, 1, 1, 1] = 0.5fnr[2] / fnr[3]
        vwh[2, 1, 1, 1] = -.5ci * fnr[2] / fnr[3]
        vwh[1, 1, 2, 1] = -.5fnr[2] / fnr[3]
        vwh[2, 1, 2, 1] = -.5ci * fnr[2] / fnr[3]
        return vwh
    end
    nodrp1 = nodr + 1
    nodrm1 = nodr - 1

    a = ri * r
    for p in 1:2
        hn[:, p] = (itype == 1 ? cricbessel(nodrp1, a[p]) : crichankel(nodrp1, a[p])) / a[p]
    end

    pmn = normalizedlegendre(ctx, ct, 2, nodrp1)
    E_ϕm = ephicoef(E_ϕ, 2)
    umn[0, 0, 1:2] = hn[0, 1:2] * fnr[2]
    for n in 1:nodrp1
        p = min(n, 2)
        for m in (-p):p
            umn[m, n, 1:2] = hn[n, 1:2] * (fnr[2] * pmn[m, n] * E_ϕm[m])
        end
    end

    for p in 1:2
        sp = -monen[p]
        for n in 1:nodr
            np1 = n + 1
            nm1 = n - 1
            for m in [-1, 1]
                k = (m + 1) ÷ 2 + 1
                mp1 = m + 1
                mm1 = m - 1
                a1 = vwh_coef[m, n, 1, 1] * umn[mp1, np1, p] + vwh_coef[m, n, 1, -1] * umn[mp1, nm1, p]
                b1 = vwh_coef[m, n, -1, 1] * umn[mm1, np1, p] + vwh_coef[m, n, -1, -1] * umn[mm1, nm1, p]
                z1 = vwh_coef[m, n, 0, 1] * umn[m, np1, p] + vwh_coef[m, n, 0, -1] * umn[m, nm1, p]
                a2 = vwh_coef[m, n, 1, 0] * umn[mp1, n, p]
                b2 = vwh_coef[m, n, -1, 0] * umn[mm1, n, p]
                z2 = vwh_coef[m, n, 0, 0] * umn[m, n, p]
                vwh[1, p, k, n] = -0.5(a1 + b1) - 0.5ci * sp * (a2 + b2)
                vwh[2, p, k, n] = -0.5ci * (-a1 + b1) - 0.5sp * (a2 - b2)
                vwh[3, p, k, n] = -z1 - sp * ci * z2
            end
        end
    end

    return vwh
end

function twobytwoinverse(mat::Array{ComplexF64,2})
    @assert size(mat) == (2, 2)
    @assert det(mat) != 0

    return inv(mat)
end

end # module SpecialFunctions


text:
include("type_define.jl")
include("mesh_medium.jl")
include("boundary.jl")
include("fd_coefficient.jl")
include("mesh_wavefield.jl")
include("mesh_wavelet.jl")
# include("wavefield_parameter.jl")


text:
using Cassette

Cassette.@context Ctx

mutable struct Callback
    f::Any
end

function Cassette.execute(ctx::Ctx, ::typeof(println), args...)
    previous = ctx.metadata.f
    ctx.metadata.f = () -> (previous(); println(args...))
    return nothing
end

a = rand(3)
b = rand(3)
function add(a, b)
    println("I'm about to add $a + $b")
    c = a + b
    println("c = $c")
    return c
end

add(a, b)

ctx = Ctx(metadata = Callback(() -> nothing));

c = Cassette.overdub(ctx, add, a, b)

ctx.metadata.f()


text:
"""
Subroutine dsel2

Given an array x, this subroutine permutes the elements of the
array keys so that

  abs(x(keys(i))) <= abs(x(keys(k))),  1 <= i <= k,
  abs(x(keys(k))) <= abs(x(keys(i))),  k <= i <= n.

In other words, the smallest k elements of x in absolute value are
x(keys(i)), i = 1,...,k, and x(keys(k)) is the kth smallest element.

MINPACK-2 Project. March 1998.
Argonne National Laboratory.
William D. Kastak, Chih-Jen Lin, and Jorge J. More'.

Revised October 1999. Length of x was incorrectly set to n.
"""
function dsel2(n,x,keys,k)
    if (n <= 1 || k <= 0 || k > n)
        return
    end

    u = n
    l = 1
    lc = n
    lp = 2*n

    while (l < u)

        # Choose the partition as the median of the elements in
        # positions l+s*(u-l) for s = 0, 0.25, 0.5, 0.75, 1.
        # Move the partition element into position l.

        p1 = floor(Int,(u+3*l)/4)
        p2 = floor(Int,(u+l)/2)
        p3 = floor(Int,(3*u+l)/4)

        # Order the elements in positions l and p1.

        if abs(x[keys[l]]) > abs(x[keys[p1]])
            swap = keys[l]
            keys[l] = keys[p1]
            keys[p1] = swap
        end

        # Order the elements in positions p2 and p3.

        if abs(x[keys[p2]]) > abs(x[keys[p3]])
            swap = keys[p2]
            keys[p2] = keys[p3]
            keys[p3] = swap
        end

        # Swap the larger of the elements in positions p1
        # and p3, with the element in position u, and reorder
        # the first two pairs of elements as necessary.

        if abs(x[keys[p3]]) > abs(x[keys[p1]])
            swap = keys[p3]
            keys[p3] = keys[u]
            keys[u] = swap
            if abs(x[keys[p2]]) > abs(x[keys[p3]])
                swap = keys[p2]
                keys[p2] = keys[p3]
                keys[p3] = swap
            end
        else
            swap = keys[p1]
            keys[p1] = keys[u]
            keys[u] = swap
            if abs(x[keys[l]]) > abs(x[keys[p1]])
                swap = keys[l]
                keys[l] = keys[p1]
                keys[p1] = swap
            end
        end

        # If we define a(i) = abs(x(keys(i))) for i = 1,...,n, we have
        # permuted keys so that
        #
        #  a(l) <= a(p1), a(p2) <= a(p3), max(a(p1),a(p3)) <= a(u).
        #
        # Find the third largest element of the four remaining
        # elements (the median), and place in position l.

        if abs(x[keys[p1]]) > abs(x[keys[p3]])
            if abs(x[keys[l]]) <= abs(x[keys[p3]])
                swap = keys[l]
                keys[l] = keys[p3]
                keys[p3] = swap
            end
        else
            if abs(x[keys[p2]]) <= abs(x[keys[p1]])
                swap = keys[l]
                keys[l] = keys[p1]
                keys[p1] = swap
            else
                swap = keys[l]
                keys[l] = keys[p2]
                keys[p2] = swap
            end
        end

        # Partition the array about the element in position l.

        m = l
        for i=l+1:u
            if abs(x[keys[i]]) < abs(x[keys[l]])
                m = m + 1
                swap = keys[m]
                keys[m] = keys[i]
                keys[i] = swap
            end
        end

        # Move the partition element into position m.

        swap = keys[l]
        keys[l] = keys[m]
        keys[m] = swap

        # Adjust the values of l and u.

        if k >= m
            l = m + 1
        end
        if k <= m
            u = m - 1
        end

        # Check fo multiple medians if the length of the subarray
        # has not decreased by 1/3 after two consecutive iterations.

        if (3*(u-l) > 2*lp) && (k > m)

            # Partition the remaining elements into those elements
            # equal to x(m), and those greater than x(m). Adjust
            # the values of l and u.

            p = m
            for i=m+1:u
                if abs(x[keys[i]]) == abs(x[keys[m]])
                    p = p + 1
                    swap = keys[p]
                    keys[p] = keys[i]
                    keys[i] = swap
                end
            end
            l = p + 1
            if k <= p
                u = p - 1
            end
        end

        # Update the length indicators for the subarray.

        lp = lc
        lc = u-l
    end

    return
end


text:
using Printf
using JLD2
using Oceananigans.Utils

"""
    JLD2OutputWriter{F, I, O, IF, IN, KW} <: AbstractOutputWriter

An output writer for writing to JLD2 files.
"""
mutable struct JLD2OutputWriter{F, I, O, IF, IN, KW} <: AbstractOutputWriter
        filepath :: String
         outputs :: O
        interval :: I
       frequency :: F
            init :: IF
       including :: IN
        previous :: Float64
            part :: Int
    max_filesize :: Float64
           async :: Bool
           force :: Bool
         verbose :: Bool
         jld2_kw :: KW
end

noinit(args...) = nothing

"""
    JLD2OutputWriter(model, outputs; interval=nothing, frequency=nothing, dir=".",
                     prefix="", init=noinit, including=[:grid, :coriolis, :buoyancy, :closure],
                     part=1, max_filesize=Inf, force=false, async=false, verbose=false, jld2_kw=Dict{Symbol, Any}())

Construct a `JLD2OutputWriter` that writes `label, func` pairs in `outputs` (which can be a `Dict` or `NamedTuple`)
to a JLD2 file, where `label` is a symbol that labels the output and `func` is a function of the form `func(model)`
that returns the data to be saved.

Keyword arguments
=================
- `frequency::Int`   : Save output every `n` model iterations.
- `interval::Int`    : Save output every `t` units of model clock time.
- `dir::String`      : Directory to save output to. Default: "." (current working directory).
- `prefix::String`   : Descriptive filename prefixed to all output files. Default: "".
- `init::Function`   : A function of the form `init(file, model)` that runs when a JLD2 output file is initialized.
                       Default: `noinit(args...) = nothing`.
- `including::Array` : List of model properties to save with every file. By default, the grid, equation of state,
                       Coriolis parameters, buoyancy parameters, and turbulence closure parameters are saved.
- `part::Int`        : The starting part number used if `max_filesize` is finite. Default: 1.
- `max_filesize::Int`: The writer will stop writing to the output file once the file size exceeds `max_filesize`, and
                       write to a new one with a consistent naming scheme ending in `part1`, `part2`, etc. Defaults to
                       `Inf`.
- `force::Bool`      : Remove existing files if their filenames conflict. Default: `false`.
- `async::Bool`      : Write output asynchronously. Default: `false`.
- `verbose::Bool`    : Log what the output writer is doing with statistics on compute/write times and file sizes.
                       Default: `false`.
- `jld2_kw::Dict`    : Dict of kwargs to be passed to `jldopen` when data is written.
"""
function JLD2OutputWriter(model, outputs; interval=nothing, frequency=nothing, dir=".", prefix="",
                          init=noinit, including=[:grid, :coriolis, :buoyancy, :closure],
                          part=1, max_filesize=Inf, force=false, async=false, verbose=false,
                          jld2_kw=Dict{Symbol, Any}())

    validate_interval(interval, frequency)

    mkpath(dir)
    filepath = joinpath(dir, prefix * ".jld2")
    force && isfile(filepath) && rm(filepath, force=true)

    jldopen(filepath, "a+"; jld2_kw...) do file
        init(file, model)
        saveproperties!(file, model, including)
    end

    return JLD2OutputWriter(filepath, outputs, interval, frequency, init, including,
                            0.0, part, max_filesize, async, force, verbose, jld2_kw)
end

function write_output(model, fw::JLD2OutputWriter)
    verbose = fw.verbose
    verbose && @info @sprintf("Calculating JLD2 output %s...", keys(fw.outputs))
    tc = Base.@elapsed data = Dict((name, f(model)) for (name, f)
                                    in zip(keys(fw.outputs), values(fw.outputs)))
    verbose && @info "Calculation time: $(prettytime(tc))"

    iter = model.clock.iteration
    time = model.clock.time

    filesize(fw.filepath) >= fw.max_filesize && start_next_file(model, fw)

    path = fw.filepath
    verbose && @info "Writing JLD2 output $(keys(fw.outputs)) to $path..."
    t0, sz = time_ns(), filesize(path)

    if fw.async
        @async remotecall(jld2output!, 2, path, iter, time, data, fw.jld2_kw)
    else
        jld2output!(path, iter, time, data, fw.jld2_kw)
    end

    t1, newsz = time_ns(), filesize(path)

    verbose && @info @sprintf("Writing done: time=%s, size=%s, Δsize=%s",
                              prettytime((t1-t0)/1e9), pretty_filesize(newsz), pretty_filesize(newsz-sz))

    return nothing
end

"""
    jld2output!(path, iter, time, data, kwargs)

Write the (name, value) pairs in `data`, including the simulation
`time`, to the JLD2 file at `path` in the `timeseries` group,
stamping them with `iter` and using `kwargs` when opening
the JLD2 file.
"""
function jld2output!(path, iter, time, data, kwargs)
    jldopen(path, "r+"; kwargs...) do file
        file["timeseries/t/$iter"] = time
        for (name, datum) in data
            file["timeseries/$name/$iter"] = datum
        end
    end
    return nothing
end

function start_next_file(model, fw::JLD2OutputWriter)
    verbose = fw.verbose
    sz = filesize(fw.filepath)
    verbose && @info begin
        "Filesize $(pretty_filesize(sz)) has exceeded maximum file size $(pretty_filesize(fw.max_filesize))."
    end

    if fw.part == 1
        part1_path = replace(fw.filepath, r".jld2$" => "_part1.jld2")
        verbose && @info "Renaming first part: $(fw.filepath) -> $part1_path"
        mv(fw.filepath, part1_path, force=fw.force)
        fw.filepath = part1_path
    end

    fw.part += 1
    fw.filepath = replace(fw.filepath, r"part\d+.jld2$" => "part" * string(fw.part) * ".jld2")
    fw.force && isfile(fw.filepath) && rm(fw.filepath, force=true)
    verbose && @info "Now writing to: $(fw.filepath)"

    jldopen(fw.filepath, "a+"; fw.jld2_kw...) do file
        fw.init(file, model)
        saveproperties!(file, model, fw.including)
    end
end

"""
    FieldOutput([return_type=Array], field)

Returns a `FieldOutput` type intended for use with the `JLD2OutputWriter`.
Calling `FieldOutput(model)` returns `return_type(field.data.parent)`.
"""
struct FieldOutput{O, F}
    return_type :: O
          field :: F
end

FieldOutput(field) = FieldOutput(Array, field) # default
(fo::FieldOutput)(model) = fo.return_type(fo.field.data.parent)

"""
    FieldOutputs(fields)

Returns a dictionary of `FieldOutput` objects with key, value
pairs corresponding to each name and value in the `NamedTuple` `fields`.
Intended for use with `JLD2OutputWriter`.

Examples
========

```julia
julia> output_writer = JLD2OutputWriter(model, FieldOutputs(model.velocities), frequency=1);

julia> keys(output_writer.outputs)
Base.KeySet for a Dict{Symbol,FieldOutput{UnionAll,F} where F} with 3 entries. Keys:
  :w
  :v
  :u
```
"""
function FieldOutputs(fields)
    names = propertynames(fields)
    nfields = length(fields)
    return Dict((names[i], FieldOutput(fields[i])) for i in 1:nfields)
end


text:
#!/usr/bin/julia

for i = 1:100
    str = i % 3 == 0 ? "Fizz" : ""
    str *= i % 5 == 0 ? "Buzz" : ""
    if isempty(str)
        str = i
    end
    println(str)
end


text:
function Δtsample(p)
	p[1] .= 0.0  # it is necessary to have these boundary conditions
	nt=length(p)
	dpdt=zero(p)
	pout=zero(p)
		for it=1:nt-1 # use these bounds later
			dpdt[it] = p[it+1]-p[it]
		end

		for it=1:nt-1 # use these exact bounds later
			pout[it+1] = dpdt[it+1]-dpdt[it]
		end
	return pout
end

"""
Simple test to see staggerred grid spatial derivatives that are used satisfy adjoint test.
"""
function Δxsample(p)
	nx=length(p)
	# boundary conditions
	p[1:2].=0.0
	p[nx-1:nx].=0.0
	dpdx=zero(p)
	pout=zero(p)
	for i in 1:100 # applying operators a 100 times, say
		for ix=2:nx-2 # use these bounds later
			dpdx[ix] = (27.e0*p[ix+1]-27.e0*p[ix]-p[ix+2]+p[ix-1]) * 1.0
		end

		for ix=3:nx-1 # use these exact bounds later
			pout[ix] = (27.e0*dpdx[ix]-27.e0*dpdx[ix-1]-dpdx[ix+1]+dpdx[ix-2]) * 1.0
		end
	end

	return pout
end

"""
advance by one time step for adjoint test
"""
function advance_sample(ntimes, p, modrrvx, modrrvz, modttI)
	nz=size(p,1)
	nx=size(p,2)
	# boundary conditions
	for i in 1:3
		p[1:2,:,i] .= 0.0
		p[:,1:2,i] .= 0.0
		p[nz-1:nz,:,i] .= 0.0
		p[:,nx-1:nx,i] .= 0.0
	end

	δx24I=1.; δz24I=1.; δt=1.

	a_x=zeros(nx)
	a_x_half=zeros(nx)
	b_x=zeros(nx);	b_x_half=zeros(nx)
	k_x=ones(nx);	k_xI=ones(nx)
	k_x_half=ones(nx);	k_x_halfI=ones(nx)

	k_z=ones(nz);	k_zI=ones(nz)
	k_z_half=ones(nz);	k_z_halfI=ones(nz)
	a_z=zeros(nz);	a_z_half=zeros(nz)
	b_z=zeros(nz);	b_z_half=zeros(nz)

	modrrvx=ones(nz,nx)
	modrrvz=ones(nz,nx)
	modttI=ones(nz,nx)

	pw=zeros(nz,nx,3)
	#input
	copyto!(pw, p)
#=
	for iz in 1:nz
		for ix in 1:nx
			pw[iz,ix,3]=p[iz,ix,3]
		end
	end
=#
	dpdxw=zeros(nz,nx,3); 	dpdzw=zeros(nz,nx,3)
	ppw=zeros(nz,nx,3); 	pppw=zeros(nz,nx,3)

	memory_dvx_dxw=zeros(nz,nx)
	memory_dvx_dzw=zeros(nz,nx)
	memory_dvz_dxw=zeros(nz,nx)
	memory_dvz_dzw=zeros(nz,nx)
	memory_dp_dxw=zeros(nz,nx)
	memory_dp_dzw=zeros(nz,nx)

	for i in 1:ntimes

		advance_kernel!(pw, dpdxw, dpdzw, δx24I, δz24I,
			 memory_dp_dxw,memory_dp_dzw,
			 b_x_half,b_z_half,a_x_half,a_z_half,k_x_halfI,k_z_halfI,
			 nx,nz,δt,modrrvx,modrrvz,memory_dvx_dxw,memory_dvz_dzw,
			 b_x,b_z,a_x,a_z,k_xI,k_zI,modttI)
	end

	return pw
end

@inbounds @fastmath function advance!(pac, pap)
	# aliases
	p=pap.p; pp=pap.pp; ppp=pap.ppp;
	dpdx=pap.dpdx; dpdz=pap.dpdz;
	memory_dp_dx=pap.memory_dp_dx; memory_dp_dz=pap.memory_dp_dz;
	memory_dvx_dx=pap.memory_dvx_dx; memory_dvz_dz=pap.memory_dvz_dz
	modttI=pac.modttI; modrrvx=pac.modrrvx; modrrvz=pac.modrrvz
	δx24I=pac.δx24I; δz24I=pac.δz24I; δt=pac.δt
	nx=pac.nx; nz=pac.nz
	a_x=pac.a_x; b_x=pac.b_x; k_xI=pac.k_xI; a_x_half=pac.a_x_half; b_x_half=pac.b_x_half; k_x_halfI=pac.k_x_halfI
	a_z=pac.a_z; b_z=pac.b_z; k_zI=pac.k_zI; a_z_half=pac.a_z_half; b_z_half=pac.b_z_half; k_z_halfI=pac.k_z_halfI

	for ipw in pac.activepw
		pw=p[ipw]
		ppw=pp[ipw]
		pppw=ppp[ipw]
		dpdxw=dpdx[ipw]
		dpdzw=dpdz[ipw]
		memory_dp_dxw=memory_dp_dx[ipw]
		memory_dp_dzw=memory_dp_dz[ipw]
		memory_dvx_dxw=memory_dvx_dx[ipw]
		memory_dvz_dzw=memory_dvz_dz[ipw]

		# store p for the last two steps
		pppppp!(pw,ppw,pppw)

		advance_kernel!(pw, dpdxw, dpdzw, δx24I, δz24I,
			 memory_dp_dxw, memory_dp_dzw,
			 b_x_half,b_z_half,a_x_half,a_z_half,k_x_halfI,k_z_halfI,
			 nx,nz,δt,modrrvx,modrrvz,memory_dvx_dxw,memory_dvz_dzw,
			 b_x,b_z,a_x,a_z,k_xI,k_zI,modttI)

	end

	return nothing
end

function advance_kernel!(pw, dpdxw, dpdzw, δx24I, δz24I,
			 memory_dp_dxw,memory_dp_dzw,
			 b_x_half,b_z_half,a_x_half,a_z_half,k_x_halfI,k_z_halfI,
			 nx,nz,δt,modrrvx,modrrvz,memory_dvx_dxw,memory_dvz_dzw,
			 b_x,b_z,a_x,a_z,k_xI,k_zI,modttI)
	#compute dpdx and dpdz at [it-1] for all propagating fields
	update_dpdx!(pw, dpdxw, δx24I, memory_dp_dxw, b_x_half, a_x_half, k_x_halfI, nx, nz)
	update_dpdz!(pw, dpdzw, δz24I, memory_dp_dzw, b_z_half, a_z_half, k_z_halfI, nx, nz)

	#update velocity at [it-1/2] using
	#velocity at [it-3/2] and dpdx and dpdz at [it-1]
	update_vx!(pw, dpdxw, δt, modrrvx, nx, nz)
	update_vz!(pw, dpdzw, δt, modrrvz, nx, nz)

	#rigid boundary conditions (not currently implemented, need to do adjoint test)
	#such that:
	#p[3,.] is where velocities are zero and
	#p[.,n-1] is where velocities are zero
#	rigid_boundary_left!(p,nx,nz)
#	rigid_boundary_right!(p,nx,nz)
#	rigid_boundary_top!(p,nx,nz)
#	rigid_boundary_bottom!(p,nx,nz)

	dvdx!(dpdxw,pw,memory_dvx_dxw,b_x,a_x,k_xI,nz,nx,δx24I)
	dvdz!(dpdzw,pw,memory_dvz_dzw,b_z,a_z,k_zI,nz,nx,δz24I)

	#compute pressure at [it] using p at [it-1] and dvxdx
	#and dvzdz at [it-1/2]
	pvzvx!(pw,dpdxw,dpdzw,modttI,nz,nx,δt)

end

function pppppp!(pw,ppw,pppw)
	copyto!(pppw,ppw)
	copyto!(ppw,pw)
end

#
@inbounds @fastmath function dvdx!(dpdxw,pw,memory_dvx_dxw,b_x,a_x,k_xI,nz,nx,δx24I)
	for ix=3:nx-1 # ix=1,2,nx are useless, rigid boundary conditions apply on ix=3,nx-1
	@simd for iz=1:nz
		@inbounds dpdxw[iz,ix,2] = (27.e0*pw[iz,ix,2]-27.e0*pw[iz,ix-1,2]-pw[iz,ix+1,2]+pw[iz,ix-2,2]) * (δx24I)
		@inbounds memory_dvx_dxw[iz,ix] = b_x[ix] * memory_dvx_dxw[iz,ix] + a_x[ix] * dpdxw[iz,ix,2] # pml
		@inbounds dpdxw[iz,ix,2] = dpdxw[iz,ix,2] * k_xI[ix] + memory_dvx_dxw[iz,ix] # pml
	end
	end
end

@inbounds @fastmath function dvdz!(dpdzw,pw,memory_dvz_dzw,b_z,a_z,k_zI,nz,nx,δz24I)
	for ix=1:nx
	@simd for iz=3:nz-1 # iz=1,2 are usless, boundary conditions on iz=3,nz-1
		@inbounds dpdzw[iz,ix,3] = (27.e0*pw[iz,ix,3]-27.e0*pw[iz-1,ix,3]-pw[iz+1,ix,3]+pw[iz-2,ix,3]) * (δz24I)
		@inbounds memory_dvz_dzw[iz,ix] = b_z[iz] * memory_dvz_dzw[iz,ix] + a_z[iz] * dpdzw[iz,ix,3] # pml
		@inbounds dpdzw[iz,ix,3] = dpdzw[iz,ix,3] * k_zI[iz] + memory_dvz_dzw[iz,ix] # pml
	end
	end
end

@inbounds @fastmath function pvzvx!(pw,dpdxw,dpdzw,modttI,nz,nx,δt)
	for ix=1:nx  # see limits above
	@simd for iz=1:nz
		@inbounds pw[iz,ix,1] += (modttI[iz,ix] * (dpdxw[iz,ix,2] + dpdzw[iz,ix,3])) * δt #* boundary_p(iz,ix)
	end
	end
end

# update for ix=[3,...,nx-2]
# ix=1,2 are virtual nodes (apply rigid boundary condition later)
# ix=nx-1,nx are virtual nodes (apply rigid boundary condition later)
@inbounds @fastmath function update_dpdx!(pw, dpdxw, δx24I, memory_dp_dxw, b_x_half, a_x_half, k_x_halfI, nx, nz)
	for ix=2:nx-2
	@simd for iz=1:nz
		@inbounds dpdxw[iz,ix,1] = (27.e0*pw[iz,ix+1,1]-27.e0*pw[iz,ix,1]-pw[iz,ix+2,1]+pw[iz,ix-1,1]) * (δx24I)
		@inbounds memory_dp_dxw[iz,ix] = b_x_half[ix] * memory_dp_dxw[iz,ix] + a_x_half[ix] * dpdxw[iz,ix,1] # pml
		@inbounds dpdxw[iz,ix,1] = dpdxw[iz,ix,1] * k_x_halfI[ix] + memory_dp_dxw[iz,ix] # pml
	end
	end
end

# update for iz=[2,...,nz-2]
# iz=1,2 are virtual node
# iz=nz-1,nz are virtual nodes
@inbounds @fastmath function update_dpdz!(pw, dpdzw, δz24I, memory_dp_dzw, b_z_half, a_z_half, k_z_halfI, nx, nz)
	for ix=1:nx
	@simd for iz=2:nz-2
		@inbounds dpdzw[iz,ix,1] = (27.e0*pw[iz+1,ix,1]-27.e0*pw[iz,ix,1]-pw[iz+2,ix,1]+pw[iz-1,ix,1]) * (δz24I)
		@inbounds memory_dp_dzw[iz,ix] = b_z_half[iz] * memory_dp_dzw[iz,ix] + a_z_half[iz] * dpdzw[iz,ix,1] # pml
		@inbounds dpdzw[iz,ix,1] = dpdzw[iz,ix,1] * k_z_halfI[iz] + memory_dp_dzw[iz,ix] # pml
	end
	end
end

# dpdx is previously computed from [2,...,nx-2]
@inbounds @fastmath function update_vx!(pw, dpdxw, δt, modrrvx,  nx, nz)
	for ix=1:nx # see dpdxw computation above
	@simd for iz=1:nz
		@inbounds pw[iz,ix,2] += (dpdxw[iz,ix,1]) * δt * modrrvx[iz,ix] #* boundary_vx(iz,ix)
	end
	end
end

# dpdz is previously computed from [2,...,nz-2]
@inbounds @fastmath function update_vz!(pw, dpdzw, δt, modrrvz, nx, nz)
	for ix=1:nx
	@simd for iz=1:nz # see dpdzw computation above
		@inbounds pw[iz,ix,3] +=  (dpdzw[iz,ix,1]) * δt * modrrvz[iz,ix] #* boundary_vz(iz,ix)
	end
	end
end

# rigid boundary conditions on the left side of the model
function rigid_boundary_left!(pw,nx,nz)
	@simd for iz=1:nz
		pw[iz,1,2] = -1.0*pw[iz,4,2]  # ix=1 is the virtual node for vx
		pw[iz,2,2] = -1.0*pw[iz,3,2]  # ix=2 is the virtual node for vx
		pw[iz,3,3] = 0.0 # ix=3, is where the vz should be zero
	end
end

function rigid_boundary_right!(pw,nx,nz)
	@simd for iz=1:nz
		pw[iz,nx-1,2] = -1.0*pw[iz,nx-2,2]  # ix=nx-1 is the virtual node for vx
		pw[iz,nx,2] = -1.0*pw[iz,nx-3,2]  # ix=nx is the virtual node for vx
		pw[iz,nx-1,3] = 0.0 # ix=nx-1, is where the vz should be zero
	end
end

# rigid boundary conditions on the top side of the model
function rigid_boundary_top!(pw,nx,nz)
	@simd for ix=1:nx
		pw[1,ix,3] = -1.0*pw[4,ix,3]  # iz=1 is the virtual node for vz
		pw[2,ix,3] = -1.0*pw[3,ix,3]  # iz=2 is the virtual node for vz
		pw[3,ix,2] = 0.0 # iz=3 is where the vx should be zero
	end
end

function rigid_boundary_bottom!(pw,nx,nz)
	@simd for ix=1:nx
		pw[nz-1,ix,3] = -1.0*pw[nz-2,ix,3]  # iz=nz-1 is the virtual node for vz
		pw[nz,ix,3] = -1.0*pw[nz-3,ix,3]  # iz=nz-1 is the virtual node for vz
		pw[nz-1,ix,2] = 0.0 # iz=nz-1, is where the vx should be zero
	end
end


text:
using PaddedViews

"""
    scale(cube::AbstractArray{T,3}, scales)

Linearly stretch each frame in `cube` by the corresponding scale in `scales`. Uses bilinear interpolaiton internally. The output size can be specified or else we will choose the smallest size that contains the largest stretch.
"""
function scale(cube::AbstractArray{T,3}, scales) where T
    out = similar(cube)
    Threads.@threads for i in axes(cube, 3)
        frame = @view cube[:, :, i]
        out[:, :, i] .= scale(frame, scales[i])
    end
    return out
end

"""
    scale(frame::AbstractMatrix, scale)

Linearly stretch `frame` with the ratio `scale`. Uses bilinear interpolation internally.
"""
scale(frame::AbstractMatrix, scale) = cropview(imresize(frame; ratio=scale), size(frame); force=true, verbose=false)

"""
    invscale(cube::AbstractArray{T,3}, scales)

Linearly contract each frame in `cube` by the corresponding scale in `scales`. Uses bilinear interpolaiton internally.
"""
function invscale(cube::AbstractArray{T,3}, scales) where T
    out = similar(cube)
    Threads.@threads for i in axes(cube, 3)
        out[:, :, i] .=  invscale(cube[:, :, i], scales[i])
    end
    return out
end

"""
    invscale(frame::AbstractMatrix, scale)

Linearly contract `frame` with the ratio `scale`. Uses bilinear interpolation internally.
"""
function invscale(frame::AbstractMatrix{T}, scale) where T
    out = imresize(frame; ratio=inv(scale))
    size(out) == size(frame) && return out
    offset = div.(size(frame) .- size(out), 2) .+ 1
    return PaddedView(zero(T), out, size(frame), offset)
end

"""
    scale_and_stack(spcube::AbstractArray{T,4}, scales)

Given a 4-D spectral ADI (SDI) tensor this function scales each temporal slice according to `scales` and then concatenates into a cube with `nλ * nf` frames.
"""
function scale_and_stack(spcube::AbstractArray{T,4}, scales) where T
    nx, ny, nλ, nf = size(spcube)
    out = similar(spcube, nx, ny, nλ * nf)
    Threads.@threads for n in axes(spcube, 4)
        slice = (n - 1) * nλ + firstindex(spcube, 4):n * nλ
        frame = @view spcube[:, :, :, n]
        out[:, :, slice] .= scale(frame, scales)
    end

    return out
end

"""
    invscale_and_collapse(stack_cube::AbstractArray{T,3}, scales; kwargs...)

Given an SDI tensor that has been stacked into a cube, invscales each spectral slice and combines with [`collapse`](@ref). The output will be cropped to `size`.
"""
function invscale_and_collapse(stack_cube::AbstractArray{T,3}, scales; kwargs...) where T
    nλ = length(scales)
    n = Base.size(stack_cube, 3) ÷ nλ
    out = similar(stack_cube, size(stack_cube)[1:2]..., n)
    Threads.@threads for n in axes(out, 3)
        slice = (n - 1) * nλ + firstindex(out, 3):n * nλ
        cube = @view stack_cube[:, :, slice]
        out[:, :, n] .= collapse(invscale(cube, scales); kwargs...)
    end
    return out
end

"""
    scale_list(wavelengths)

Returns a list of scaling factors for aligning SDI tensors from a list of wavelengths.

# Examples

```jldoctest
julia> scale_list([0.5, 2, 4])
3-element Vector{Float64}:
 8.0
 2.0
 1.0
```
"""
scale_list(wavelengths) = maximum(wavelengths) ./ wavelengths


text:
# Autogenerated wrapper script for libnode_jll for x86_64-linux-gnu-cxx03
export libnode, node

JLLWrappers.@generate_wrapper_header("libnode")
JLLWrappers.@declare_library_product(libnode, "libnode.so.83")
JLLWrappers.@declare_executable_product(node)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libnode,
        "lib/libnode.so.83",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_executable_product(
        node,
        "bin/node",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
if is_linux()
    path = joinpath(Pkg.dir("Dyno", "deps", "usr", "bin"))
    mkpath(path)
    filename = joinpath(path, "dynare_m")
    download("https://github.com/albop/demo_dyno/raw/master/dynare_m", filename)
    run(`chmod +x $filename`)
end

# using BinDeps
# @BinDeps.setup
# dynare = library_dependency("dynare")
# provides(Binaries, URI("https://github.com/albop/demo_dyno/raw/master/dynare_m"), dynare, os=:linux)
# @BinDeps.install


text:
using RanDAGen
using Documenter

makedocs(;
    modules=[RanDAGen],
    authors="Trystan Kaes <trystanblkaes@gmail.com>",
    repo="https://github.com/TrystanKaes/RanDAGen.jl/blob/{commit}{path}#L{line}",
    sitename="RanDAGen.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://RanDAGen.github.io/RanDAGen.jl",
        assets=String[],
    ),
    # devbranch = "master",
    # devurl = "dev",
    versions = ["stable" => "v^", "v#.#", "dev" => "dev"],
    pages=[
        "Home" => "index.md",
    ],
    push_preview=true,
)

deploydocs(;
    repo="github.com/TrystanKaes/RanDAGen.jl",
)


text:
using DataDeps

register(DataDep(
  "thyroid-disease-new-thyroid",
  "http://archive.ics.uci.edu/ml/datasets/Thyroid+Disease",
  "http://archive.ics.uci.edu/ml/machine-learning-databases/thyroid-disease/new-thyroid.data",
  "0760039b928ecba5170f377b8fd96caaf55ed26f9861511c52f81395a9d8227a",
  post_fetch_method=(path -> begin
    UCIData.process_dataset(path,
        target_index=1,
        feature_indices=2:6,
    )
  end),
))

max_stars_repo_name: JackDunnNZ/uci-data

max_issues_repo_name: JackDunnNZ/uci-data

max_forks_repo_name: JackDunnNZ/uci-data


text:
# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control

"""
    module Simulate_PendulumWithoutMacro (examples/withoutMacros)

Simulate Pendulum as ODE model (version that uses Variables, but no macros)
"""
module Simulate_Pendulum

using ..ModiaMath

mutable struct PendulumWithoutMacro <: ModiaMath.AbstractComponentWithVariables
    _internal::ModiaMath.ComponentInternal

    # Parameters
    L::Float64
    m::Float64
    d::Float64
    g::Float64
    phi0_deg::Float64

    # Variables
    phi::ModiaMath.RealScalar
    w::ModiaMath.RealScalar
    a::ModiaMath.RealScalar

    function PendulumWithoutMacro(;L=1.0, m=1.0, d=0.1, g=9.81, phi0_deg=90)
        this = new(ModiaMath.ComponentInternal(:PendulumWithoutMacro, nothing))
        this.L = L
        this.m = m
        this.d = d
        this.g = g
        this.phi0_deg = phi0_deg

        @assert(L > 0.0)
        @assert(m > 0.0)
        @assert(d >= 0.0)

        phi = ModiaMath.RealScalar(start=phi0_deg * pi / 180, unit="rad", fixed=true, info="Relative rotation angle", numericType=ModiaMath.XD_EXP)
        ModiaMath.initComponent!(this, phi, :phi)

        w   = ModiaMath.RealScalar(start=0.0, unit="rad/s", fixed=true, info="Relative angular velocity", integral=phi, numericType=ModiaMath.XD_EXP)
        ModiaMath.initComponent!(this, w, :w)

        a   = ModiaMath.RealScalar(unit="rad/s^2", info="Relative angular acceleration", integral=w, numericType=ModiaMath.DER_XD_EXP)
        ModiaMath.initComponent!(this, a, :a)

        return this
    end
end

function ModiaMath.computeVariables!(p::PendulumWithoutMacro, sim::ModiaMath.SimulationState)
    L = p.L
    m = p.m
    d = p.d
    g = p.g
    phi = p.phi.value
    w   = p.w.value

    p.a.value = (-m * g * L * sin(phi) - d * w) / (m * L^2)
end

simulationModel = ModiaMath.SimulationModel(PendulumWithoutMacro(L=0.8, m=0.5, d=0.2), stopTime=5.0)
#ModiaMath.print_ModelVariables(simulationModel)
result = ModiaMath.simulate!(simulationModel, log=true)

ModiaMath.plot(result, [:phi, :w, :a])

end


text:
dependencies = Dict{Symbol,Symbol}()

function calc_first(prd::Production)
	subprods = prd.subprods
	father_idx = prd.enum

  if length(prd.firsts) != 0
      return prd.firsts
  end

	for subprods_idx=1:length(subprods)

    current_prod_first = []
    for sp_prod_idx =1:length(subprods[subprods_idx])
        has_eps = false

        #verificar se a producao atual tem EPS caso tenha chamar first da proxima
        if subprods[subprods_idx][sp_prod_idx] isa Symbol
            for eps_it in getProd(subprods[subprods_idx][sp_prod_idx]).subprods
                if eps_it[1] == EPS
                    has_eps = true
                end
            end
        end

        if typeof(subprods[subprods_idx][sp_prod_idx]) == Symbol
			      #println("Começa com Produção")

			      son_idx = getProd_idx(subprods[subprods_idx][sp_prod_idx])

            if length(grammar[son_idx].firsts) == 0
			          calc_first(grammar[son_idx])
            end

            son_firsts = grammar[son_idx].firsts
			      for i in son_firsts
				        #grammar[father_idx].firsts = vcat(grammar[father_idx].firsts, i)

                # so adicionar EPS a regra caso não tenha mais nenhuma outra e.g
                # A -> BCD, first(B) = FIRST(C) - eps U First(D)
                # B -> C | eps
                # D -> d
                # so se D -> d | eps, então First(B) conteria EPS
                if i[1] == Int(EPS) && sp_prod_idx == length(subprods[subprods_idx])
                    current_prod_first = vcat(current_prod_first, i)
                elseif i[1] == Int(EPS)
                    continue
                else
                    current_prod_first = vcat(current_prod_first, i)
                end
			      end
		    else
			      #println("Começa com Token")
  	        #grammar[father_idx].firsts = vcat(grammar[father_idx].firsts, Int(subprods[subprods_idx][1]))
            current_prod_first = vcat(current_prod_first, Int(subprods[subprods_idx][sp_prod_idx]))
		    end
        if !has_eps
            break
        end
    end

      current_prod_first = unique(current_prod_first)
      push!(grammar[father_idx].firsts, current_prod_first)
	end

	#grammar[father_idx].firsts = unique(grammar[father_idx].firsts)
	return grammar[father_idx].firsts
end

#=
for p in grammar
	calc_first(p)
end
=#

max_stars_repo_name: NaelsonDouglas/namelesscompiler

max_issues_repo_name: NaelsonDouglas/namelesscompiler

max_forks_repo_name: NaelsonDouglas/namelesscompiler


text:
const HTSRecordFlagBits = (
    # the read is paired in sequencing, no matter whether it is mapped in a pair
    paired=UInt16(1),
    # the read is mapped in a proper pair
    proper_pair=UInt16(2),
    # the read itself is unmapped; conflictive with BAM_FPROPER_PAIR
    unmap=UInt16(4),
    # the mate is unmapped
    munmap=UInt16(8),
    # the read is mapped to the reverse strand
    reverse=UInt16(16),
    # the mate is mapped to the reverse strand
    mreverse=UInt16(32),
    # this is read1
    read1=UInt16(64),
    # this is read2
    read2=UInt16(128),
    # not primary alignment
    secondary=UInt16(256),
    # QC failure
    qcfail=UInt16(512),
    # optical or PCR duplicate
    dup=UInt16(1024),
    # supplementary alignment
    supplementary=UInt16(2048),
)

"""
Access HTSRecord bitwise flags.
"""
struct HTSRecordFlag
    value::UInt16
end

function Base.convert(::Type{UInt16}, flag::HTSRecordFlag)
    getfield(flag, :value)
end

@inline function Base.getproperty(flag::HTSRecordFlag, name::Symbol)
    bit = getproperty(HTSRecordFlagBits, name)
    (convert(UInt16, flag) & bit) > 0
end

function Base.propertynames(::HTSRecordFlag)
    propertynames(HTSRecordFlagBits)
end

# TODO: simplify this
function Base.summary(flag::HTSRecordFlag)
    str = [
        flag.paired ? "paired" : "unpaired",
        flag.proper_pair ? "proper_pair" : "improper_pair",
        flag.unmap ? "unmapped" : "mapped",
        flag.munmap ? "mate_unmapped" : "mate_mapped",
        flag.reverse ? "rev_strand" : "pos_strand",
        flag.mreverse ? "mate_rev_strand" : "mate_pos_strand",
        flag.read1 ? "firstread" : "not_firstread",
        flag.read2 ? "lastread" : "not_lastread",
        flag.secondary ? "secondary" : "primary",
        flag.qcfail ? "qcfail" : "not_qcfail",
        flag.dup ? "dup" : "not_dup",
        flag.supplementary ? "supplementary" : "not_supplementary",
    ]
    join(str, ",")
end


text:
"""
    benchmark_result_types(benchmark_result::BenchmarkResult)
"""
function benchmark_result_types(::BenchmarkResult{TA, TB, TC}) where {TA, TB, TC}
    return (TA, TB, TC)
end

"""
    benchmark_result_df(benchmark_result::BenchmarkResult)
"""
function benchmark_result_df(benchmark_result::BenchmarkResult)
    return deepcopy(benchmark_result.df)
end

function Base.show(io::IO, br::BenchmarkResult{TA, TB, TC}) where {TA, TB, TC}
    println(io, "Benchmark Result of Matrix{$TC}=Matrix{$TA}×Matrix{$TB}")
    println(io, br.df)
    return nothing
end


text:
# These tests have been adapted from Criterim clojure library

facts("outlier effect") do
    C = Criterion
    @fact C.outlier_effect(0.009) => C.Unaffected()
    @fact C.outlier_effect(0.09)  => C.Slight()
    @fact C.outlier_effect(0.45)  => C.Moderate()
    @fact C.outlier_effect(0.51)  => C.Severe()
end

facts("outliers") do
    C = Criterion
    @fact C.classify_outliers([1,2,5,7,8])        => C.Outliers(5, 0, 0, 0, 0)
    @fact C.classify_outliers([1,2,3,5,7,8])      => C.Outliers(6, 0, 0, 0, 0)
    @fact C.classify_outliers([-100,1,2,5,7,8,9]) => C.Outliers(7, 1, 0, 0, 0)
    @fact C.classify_outliers([-10,1,2,5,7,8,9])  => C.Outliers(7, 0, 1, 0, 0)
    @fact C.classify_outliers([1,2,3,5,7,8,22])   => C.Outliers(7, 0, 0, 1, 0)
    @fact C.classify_outliers([1,1,2,5,7,8,100])  => C.Outliers(7, 0, 0, 0, 1)
end

facts("outlier significance") do
    # http://www.ellipticgroup.com/misc/article_supplement.pdf, p22
    mean_est = 1.395522860870968
    std_est  = sqrt(0.0013859776344426547 * 0.0013859776344426547)
    niter    = 67108864
    o_var = Criterion.outlier_variance(mean_est, std_est, niter)
    @fact o_var.frac => 0.9960022873987793
end


text:
function generate(io::IO, libCvts::Union{Pair{<:AbstractString, Vector{Converted}}, Pair{Symbol, Vector{Converted}}, Pair{Nothing, Vector{Converted}}}...; relto::String = ".", comments::Bool = true, exports::Bool = true)
	println(io, "# This file is automatically generated.  Do not edit.")
	println(io, "using CBinding: 𝐣𝐥")
	foreach(libCvts) do (lib, cvts)
		println(io)
		lib = isnothing(lib) ? " " : " $(lib isa AbstractString ? repr(String(lib)) : String(lib)) "
		println(io, "𝐣𝐥.@cbindings", lib, "begin")
		generate(io, cvts, 1, relto = relto, comments = comments, exports = exports)
		println(io, "end")
	end
end

generate(io::IO, cvts::Vector{Converted}, indent::Int; relto::String = ".", comments::Bool = true, exports::Bool = true) = foreach(enumerate(cvts)) do (ind, cvt)
	ind > 1 && println(io, repeat('\t', indent))
	generate(io, cvt, indent, relto = relto, comments = comments, exports = exports)
end

function generate(io::IO, cvt::Converted, indent::Int; relto::String = ".", comments::Bool = true, exports::Bool = true)
	tabs = repeat('\t', indent)

	if comments
		for (name, comment) in cvt.comments
			isempty(name) && continue

			name = startswith(name, '@') ? ":("*name*")" : name

			comment = """
			𝐣𝐥.@doc \"\"\"
			$(rstrip(replace(replace(string(comment, expr = (startswith(cvt.expr, "𝐣𝐥.@cextern") ? cvt.expr : ""), relto = relto), '\\' => "\\\\"), '\"' => "\\\""), '\n'))
			\"\"\" $(name)
			"""

			println(io, tabs, rstrip(replace(comment, '\n' => "\n$(tabs)")))
		end
	end

	if exports
		exps = sort(collect(filter(str -> !isempty(str) && !startswith(str, '_'), keys(cvt.comments))))
		isempty(exps) || println(io, tabs, "export $(join(exps, ", "))")
	end

	println(io, tabs, replace(cvt.expr, '\n' => "\n$(tabs)"))
end


text:
function clear()
    begin
	   # Use system to process commands
	   run(@unix ? `clear` : `cmd /c cls`)
	end
end


text:
#!/usr/bin/env julia

using EchoJulia
using MAT

# Finds raw files and converts them to MATLAB files optionally using
# an Echoview calibration supplement for correction.

function main(args)

    ecs_filename = nothing

    if args[1] == "-c" # Echoview calibration supplement
        ecs_filename = args[2]
        filenames = args[3:end]
    else
        filenames = args
    end

    if length(filenames) == 1
        dir = filenames[1]
        if isdir(dir)
            filenames =  filter(x->endswith(x,".raw"), readdir(dir))
            filenames = ["$(dir)$(x)" for x in filenames]
        end
    end

    for filename in filenames

        out ="$(splitext(basename(filename))[1]).mat"

        if isfile(out)
            info("Skipping $out.")
            continue
        end

        try
            info("Processing $filename ...")
            raw = SimradRaw.load(filename)
            if ecs_filename == nothing
                cal = nothing
            else
                cal = EchoviewEcs.load(ecs_filename)
            end

            dict = transform(raw, calibration=cal)

            info("Writing $out ...")
            matwrite(out, dict)
        catch ex
            warn(ex)
        end

    end

end

main(ARGS)


text:
function pythagorean_triple_with_sum(sum)
    if sum < 11
        error("The smallest possible sum of a pythagorean triple is 11.")
    end

    ms = 2:(sum + 1)

    ns = 1:sum

    for m in ms
        for n in ns
            if m > n
                a = (m * m) - (n * n)

                b = 2 * m * n

                c = (m * m) + (n * n)

                if (a + b + c == sum) && ((a * a) + (b * b) == (c * c))
                    return [a, b, c]
                end
            end
        end
    end

    return []
end

println(prod(pythagorean_triple_with_sum(1000)))

max_stars_repo_name: seaneshbaugh/rosetta-euler

max_issues_repo_name: seaneshbaugh/rosetta-euler

max_forks_repo_name: seaneshbaugh/rosetta-euler


text:
export
    UAVActions,
    UAV,
    get_action,
    belief_state_from_slave,
    get_polar_state,
    get_belief_state,
    get_other_uav_id,
    apply_sensor_error

#=
This file defines the uav type, which uses different strategies to
make decisions based on the current state of the system.
=#

# type to hold the actions and signals emitted from uavs
type UAVActions
    actions::Vector{Float64}
    # signals sent by the uavs to coordinate next actions
    signals::Vector{Symbol}
    function UAVActions(num_uavs::Int64)
        actions = Array(Float64, num_uavs)
        signals = [:no_signal, :no_signal]
        return new(actions, signals)
    end
end

# type for uav
type UAV
    # id gives index into tables for this uav
    uav_id::Int64
    # actions
    actions::Vector{Float64}
    num_actions::Int64
    joint_actions::Matrix{Float64}
    num_joint_actions::Int64
    # tables
    policy_table::PolicyTable
    coordination_table::CoordinationTable
    # type of policy
    is_policy_joint::Bool
    # strategy
    strategy::Function
    action_selection_method::String
    # sensor accuracy information
    # we now need to pass the random number generator to the uav:
    rng::MersenneTwister
    # gaussian sensors, mean = 0, std dev these values
    sensor_range_std::Float64
    sensor_theta_std::Float64
    sensor_psi_std::Float64
    sensor_velocity_std::Float64
    use_sensor_noise::Bool
    function UAV(uav_id::Int64, actions::Vector{Float64}, joint_actions::Matrix{Float64},
            policy_table::PolicyTable, coordination_table::CoordinationTable, is_policy_joint::Bool,
            strategy::Function, action_selection_method::String, rng::MersenneTwister,
            sensor_range_std::Float64, sensor_theta_std::Float64, sensor_psi_std::Float64,
            sensor_velocity_std::Float64, use_sensor_noise::Bool)
        return new(uav_id, actions, length(actions), joint_actions, size(joint_actions, 2),
                policy_table, coordination_table, is_policy_joint, strategy, action_selection_method,
                rng, sensor_range_std, sensor_theta_std, sensor_psi_std, sensor_velocity_std, use_sensor_noise)
    end
end

#=
Description:
Delegates to the strategy of the uav to decide which action to take

Parameters:
- uav: the uav deciding which action to take
- state: state object which contains all variable information
    needed by the uav to decide which action to take

Return Value:
- action: float value that is the action of this uav
- signal: float value indiciating coordination action
- utilities: the utilities used to derive the actions
- belief_state: the belief_state used to index the policy table
- different_action: whether or not the coordination table prevented
    the taking of a different action from the greedy case and
    the coordinationt table action was coc. This is used to
    wake the pilot.
=#
function get_action(uav::UAV, state::State)

    # compute belief state
    belief_state = get_belief_state(uav, state)

    # get the baseline utilities from the policy table
    utilities = get_utilities(uav.policy_table, belief_state)

    # delegate to the strategy to decide which action to take
    # as well as whether to signal the other plane
    action, signals, different_action = uav.strategy(uav, state, utilities)

    return action, signals, utilities, belief_state, different_action
end

#=
Description:
Converts the state from the master point of view to the slave point of view
to select the correct action.
=#
function belief_state_from_slave(belief_state::Array{Float64,1})
    v0 = belief_state[4]
    v1 = belief_state[5]
    slave_belief_state = belief_state
    # [1] stays the same
    # [2] changes - requires calc
    # [3] changes - requires calc
    # [4] becomes [5] - no calculation
    slave_belief_state[4] = v1
    # [5] becomes [4] - no calculation
    slave_belief_state[5] = v0
    # [6] stays the same

    # solve for x and y coordinates
    range = belief_state[1]
    theta = belief_state[2] # rad
    psi = belief_state[3] # rad
    p1 = [0;0]
    p2 = [range*cos(theta);range*sin(theta)]

    # move the point 1 and rotate
    p1 = p1-p2
    # rotate point 1
    rot = [cos(psi) sin(psi);-sin(psi) cos(psi)]
    p1 = rot*p1

    # calculate new theta
    thetaNew = atan2(p1[2],p1[1]) # rad
    while thetaNew > pi
        thetaNew = thetaNew - 2*pi
    end
    while thetaNew < -pi
        thetaNew = thetaNew + 2*pi
    end

    # calculate new psi
    psiNew = 2*pi-psi # rad
    # ensure new psi is between 0 and 2*pi
    while psiNew < 0
        psiNew = psiNew+2*pi
    end

    slave_belief_state[2] = thetaNew # rad
    slave_belief_state[3] = psiNew # rad

    return slave_belief_state
end

#=
Description:
Converts the state to polar format and check that it is in bounds of the policy

Parameters:
- uav: the uav
- state: state object to convert to polar

Return Value:
- polar_state: the state converted to polar coordinates, which consists of
    [xr, yr, pr, vown, vint, resp, resp] to [rho, theta, psi, vown, vint, resp, resp]
=#
function get_polar_state(uav::UAV, state::State)
    # delegate to state for conversion
    if uav.is_policy_joint # joint policy
        polar_state = to_polar(state)
    else
        polar_state = to_polar_non_joint(state)
    end

    # get bounds of grid
    max_r = maximum(uav.policy_table.grid.cutPoints[1])
    min_t = minimum(uav.policy_table.grid.cutPoints[2])
    max_t = maximum(uav.policy_table.grid.cutPoints[2])

    # check that the new polar state is inside of the bounds
    # and if not return the terminal state
    r = polar_state[1]
    theta = polar_state[2]
    if r > max_r || theta < min_t || theta > max_t
        polar_state = TERMINAL_POLAR_STATE
    end

    return polar_state
end

#=
Description:
Populates a belief state of the uav

Parameters:
- uav: the uav, the belief state of which is being calculated
- state: the state converted to belief state

Return Value:
- belief_state: probability distribution over possible states
=#
function get_belief_state(uav::UAV, state::State)
    # initialize belief as sparse vector over all possible states
    belief_state = spzeros(uav.policy_table.num_states, 1)

    # convert state to polar
    polar_state = get_polar_state(uav, state)

    if !uav.is_policy_joint
        if uav.uav_id != 1
            if polar_state != TERMINAL_POLAR_STATE
                polar_state = belief_state_from_slave(polar_state)
            end
        end
    end

    # add sensor error if using noisy sensors
    if uav.use_sensor_noise
        sensor_error_state = apply_sensor_error(uav, polar_state)
    else
        sensor_error_state = polar_state
    end

    # populate belief_state
    if polar_state == TERMINAL_POLAR_STATE
        belief_state[end] = 1.0
    else
        indices, weights = interpolants(uav.policy_table.grid, sensor_error_state)
        belief_state[indices] = weights
    end

    return belief_state
end

#=
Description:
Returns the id value of the other uav

Parameters:
- uav: the current uav

Return Value:
- id: int value of the other uavs id
=#
function get_other_uav_id(uav::UAV)
    if uav.uav_id == 1
        return 2
    else
        return 1
    end
end

function apply_sensor_error(uav::UAV, polar_state::Vector{Float64})
    # copy polar state because it might be used
    # elsewhere in the above function and it's not
    # obvious whether we might want to use the
    # original or a noisy version above
    sensor_error_state = deepcopy(polar_state)

    # range
    sensor_error_state[1] = (sensor_error_state[1]
        + randn(uav.rng) * uav.sensor_range_std)
    sensor_error_state[2] = (sensor_error_state[2]
        + randn(uav.rng) * uav.sensor_theta_std)
    sensor_error_state[3] = (sensor_error_state[3]
        + randn(uav.rng) * uav.sensor_psi_std)

    # apply noise to other planes velocity
    if uav.uav_id == 1
        other_uav_velocity_idx = 5
    else
        other_uav_velocity_idx = 4
    end
    sensor_error_state[other_uav_velocity_idx] = (
        sensor_error_state[other_uav_velocity_idx]
        + randn(uav.rng) * uav.sensor_velocity_std)

    return sensor_error_state
end


text:
using DiffEqBase
import DiffEqBase: AbstractODESolution, HermiteInterpolation,AbstractDiffEqInterpolation
using DiffEqSensitivity
using ForwardDiff

abstract type AbstractSimulation end
export AbstractSimulation

struct Simulation{Q<:AbstractODESolution,W<:AbstractDomain,L<:AbstractArray,G<:Function,G2<:AbstractArray} <: AbstractSimulation
    sol::Q
    domain::W
    names::L
    N::G
    Ns::G2
end

function Simulation(sol::Q,domain::W) where {Q<:AbstractODESolution,W<:AbstractDomain}
    names = getfield.(domain.phase.species,:name)
    Ns = sum(hcat(sol.interp.u...)[domain.indexes[1]:domain.indexes[2],:],dims=1)
    Nderivs = sum(hcat(sol.interp.du...)[domain.indexes[1]:domain.indexes[2],:],dims=1)
    N = HermiteInterpolation(sol.interp.t,Ns,Nderivs)
    F(t::T) where {T<:Real} = N(t,nothing,Val{0},sol.prob.p,:left)
    return Simulation(sol,domain,names,F,Ns)
end

export Simulation

length(p::T) where {T<:AbstractSimulation} = 1
export length

iterate(p::T) where {T<:AbstractSimulation} = p
export iterate

Broadcast.broadcastable(p::T) where {T<:AbstractSimulation} = Ref(p)
export broadcastable

spcindex(bsol::Z,name::Q) where {Z<:Simulation,Q<:AbstractString} = findfirst(isequal(name),getfield.(bsol.domain.phase.species,:name))
export spcindex

function molefractions(bsol::Q,name::W,t::E) where {Q<:AbstractSimulation, W<:String, E<:Real}
    @assert name in bsol.names
    ind = findfirst(isequal(name),bsol.names)
    return bsol.sol(t)[ind]/bsol.N(t)
end

function molefractions(bsol::Q, t::E) where {Q<:AbstractSimulation,E<:Real}
    return bsol.sol(t)[bsol.domain.indexes[1]:bsol.domain.indexes[2]]./bsol.N(t)
end

function molefractions(bsol::Q) where {Q<:AbstractSimulation}
    @views return hcat(bsol.sol.u...)[bsol.domain.indexes[1]:bsol.domain.indexes[2],:]./bsol.Ns
end

export molefractions

getT(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantTPDomain,ConstantTVDomain},K<:Real,Q,G,L} = bsol.domain.T
getT(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantVDomain,ParametrizedVDomain,ConstantPDomain,ParametrizedPDomain},K<:Real,Q,G,L} = bsol.sol(t)[bsol.domain.indexes[3]]
getT(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ParametrizedTConstantVDomain,ParametrizedTPDomain},K<:Real,Q,G,L} = bsol.domain.T(t)
export getT
getV(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantVDomain,ConstantTVDomain,ParametrizedTConstantVDomain},K<:Real,Q,G,L} = bsol.domain.V
getV(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedVDomain,K<:Real,Q,G,L} = bsol.domain.V(t)
getV(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantTPDomain,ParametrizedTPDomain},K<:Real,Q,G,L} = bsol.sol(t)[bsol.domain.indexes[3]]
getV(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ParametrizedPDomain,ConstantPDomain},K<:Real,Q,G,L} = bsol.sol(t)[bsol.domain.indexes[4]]
export getV
getP(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantTPDomain,ConstantPDomain},K<:Real,Q,G,L} = bsol.domain.P
getP(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantTVDomain,ParametrizedTConstantVDomain},K<:Real,Q,G,L} = 1.0e6
getP(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ParametrizedTPDomain,ParametrizedPDomain},K<:Real,Q,G,L} = bsol.domain.P(t)
getP(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantVDomain, ParametrizedVDomain},K<:Real,Q,G,L} = bsol.sol(t)[bsol.domain.indexes[4]]
export getP
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ConstantTPDomain,K<:Real,Q,G,L} = bsol.domain.P/(R*bsol.domain.T)
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:Union{ConstantVDomain,ConstantTVDomain,ParametrizedTConstantVDomain},K<:Real,Q,G,L} = bsol.N(t)/bsol.domain.V
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedVDomain,K<:Real,Q,G,L} = bsol.N(t)/bsol.domain.V(t)
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedTPDomain,K<:Real,Q,G,L} = bsol.domain.P(t)/(R*bsol.domain.T(t))
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ConstantPDomain,K<:Real,Q,G,L} = bsol.domain.P/(R*getT(bsol,t))
getC(bsol::Simulation{Q,W,L,G}, t::K) where {W<:ParametrizedPDomain,K<:Real,Q,G,L} = bsol.domain.P(t)/(R*getT(bsol,t))
export getC
"""
calculates the rates of production/loss at a given time point
this outputs a sparse matrix of  num reactions xnum species containing the production/loss
rate of that species associated with that reaction
"""
function rops(bsol::Q,t::X) where {Q<:Simulation,X<:Real}
    ropmat = zeros(length(bsol.domain.phase.reactions),length(bsol.domain.phase.species))
    cs,kfs,krevs = calcthermo(bsol.domain,bsol.sol(t),t)[[2,9,10]]
    @simd for i in 1:length(bsol.domain.phase.reactions)
        rxn = bsol.domain.phase.reactions[i]
        R = getrate(rxn,cs,kfs,krevs)
        for ind in rxn.productinds
            ropmat[i,ind] += R
        end
        for ind in rxn.reactantinds
            ropmat[i,ind] -= R
        end
    end
    return ropmat
end

"""
calculates the rates of production/loss at a given time point for a single species
this outputs a sparse vector of length num reactions containing the production/loss
rate associated with that reaction for the given species
"""
function rops(bsol::Y,name::X,t::Z) where {Y<:Simulation, X<:AbstractString, Z<:Real}
    rop = spzeros(length(bsol.domain.phase.reactions))
    cs,kfs,krevs = calcthermo(bsol.domain,bsol.sol(t),t)[[2,9,10]]
    ind = findfirst(isequal(name),getfield.(bsol.domain.phase.species,:name))
    @assert !isa(ind,Nothing) "species $name not in species array"
    for (i,rxn) in enumerate(bsol.domain.phase.reactions)
        c = 0
        R = getrate(rxn,cs,kfs,krevs)
        c -= count(isequal(ind),rxn.reactantinds)
        c += count(isequal(ind),rxn.productinds)
        if c != 0
            rop[i] = c*R
        end
    end
    return rop
end
export rops

"""
Calculates sensitivities with respect to `target` at the time point at the end of the simulation
The returned sensitivities are the normalized values

By default uses the InterpolatingAdjoint algorithm with vector Jacobian products calculated with ReverseDiffVJP(true)
this assumes no changes in code branching during simulation, if that were to become no longer true, the Tracker
based alternative algorithm is slower, but avoids this concern.
"""
function getadjointsensitivities(bsol::Q,target::String,solver::W;sensalg::W2=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true)),abstol::Float64=1e-6,reltol::Float64=1e-3,kwargs...) where {Q,W,W2}
    @assert target in bsol.names || target in ["T","V"]
    if target in ["T","V"]
        ind = bsol.domain.indexes[end]
    else
        ind = findfirst(isequal(target),bsol.names)
    end
    function g(y::X,p::Array{Y,1},t::Z) where {Q,V,X,Y<:Float64,Z}
        dy = similar(y,length(y))
        return dydtreactor!(dy,y,t,bsol.domain,[],p=p)[ind]
    end
    function g(y::Array{X,1},p::Y,t::Z) where {Q,V,X<:Float64,Y,Z}
        dy = similar(p,length(y))
        return dydtreactor!(dy,y,t,bsol.domain,[],p=p)[ind]
    end
    function g(y::Array{X,1},p::Array{Y,1},t::Z) where {Q,V,X<:ForwardDiff.Dual,Y<:ForwardDiff.Dual,Z}
        dy = similar(y,length(y))
        return dydtreactor!(dy,y,t,bsol.domain,[],p=p)[ind]
    end
    dgdu(out, y, p, t) = ForwardDiff.gradient!(out, y -> g(y, p, t), y)
    dgdp(out, y, p, t) = ForwardDiff.gradient!(out, p -> g(y, p, t), p)
    du0,dpadj = adjoint_sensitivities(bsol.sol,solver,g,nothing,(dgdu,dgdp);sensealg=sensalg,abstol=abstol,reltol=reltol,kwargs...)
    dpadj[length(bsol.domain.phase.species)+1:end] .*= bsol.domain.p[length(bsol.domain.phase.species)+1:end]
    if !(target in ["T","V"])
        dpadj ./= bsol.sol(bsol.sol.t[end])[ind]
    end
    return dpadj
end
export getadjointsensitivities

function getconcentrationsensitivity(bsol::Simulation{Q,W,L,G}, numerator::String, denominator::String, t::K) where {W<:Union{ConstantVDomain,ConstantTVDomain,ParametrizedTConstantVDomain},K<:Real,Q,G,L}
    @assert numerator in bsol.names
    @assert denominator in bsol.names
    indnum = findfirst(isequal(numerator),bsol.names)
    inddeno = findfirst(isequal(denominator),bsol.names)
    Nvars = bsol.domain.indexes[end]-bsol.domain.indexes[1]+1
    Nrxns = length(bsol.domain.phase.reactions)
    x,dp = extract_local_sensitivities(bsol.sol,t)
    s = dp[inddeno][indnum]
    val = s/bsol.sol(t)[indnum] #constant volume
    if t == 0
        return 0.0
    else
        return val
    end
end

function getconcentrationsensitivity(bsol::Simulation{Q,W,L,G}, numerator::String, denominator::String, t::K) where {W<:Union{ConstantTPDomain,ParametrizedTPDomain,ConstantPDomain,ParametrizedPDomain,ParametrizedVDomain},K<:Real,Q,G,L}
    @assert numerator in bsol.names
    @assert denominator in bsol.names
    indnum = findfirst(isequal(numerator),bsol.names)
    inddeno = findfirst(isequal(denominator),bsol.names)
    Nvars = bsol.domain.indexes[end]-bsol.domain.indexes[1]+1
    Nrxns = length(bsol.domain.phase.reactions)
    x,dp = extract_local_sensitivities(bsol.sol,t)
    svals = dp[inddeno]
    s = svals[indnum]
    V = getV(bsol,t)
    c = bsol.sol(t)[indnum]/V
    val =  (s-c*sum(svals)*R*bsol.domain.T/bsol.domain.P)/(c*V) #known T and P
    if t == 0
        return 0.0
    else
        return val
    end
end

function getconcentrationsensitivity(bsol::Simulation{Q,W,L,G}, numerator::String, denominator::Z, t::K) where {W<:Union{ConstantVDomain,ConstantTVDomain,ParametrizedTConstantVDomain},K<:Real,Z<:Integer,Q,G,L}
    @assert numerator in bsol.names
    indnum = findfirst(isequal(numerator),bsol.names)
    inddeno = denominator
    Nvars = bsol.domain.indexes[end]-bsol.domain.indexes[1]+1
    Nrxns = length(bsol.domain.phase.reactions)
    x,dp = extract_local_sensitivities(bsol.sol,t)
    s = dp[inddeno+length(bsol.domain.phase.species)][indnum]
    T = getT(bsol,t)
    P = getP(bsol,t)
    C = getC(bsol,t)
    k = bsol.domain.p[inddeno+length(bsol.domain.phase.species)]
    val = s*k/bsol.sol(t)[indnum] #constant volume
    if t == 0
        return 0.0
    else
        return val
    end
end

function getconcentrationsensitivity(bsol::Simulation{Q,W,L,G}, numerator::String, denominator::Z, t::K) where {W<:Union{ConstantTPDomain,ParametrizedTPDomain,ConstantPDomain,ParametrizedPDomain,ParametrizedVDomain},K<:Real,Z<:Integer,Q,G,L}
    @assert numerator in bsol.names
    indnum = findfirst(isequal(numerator),bsol.names)
    inddeno = denominator
    Nvars = bsol.domain.indexes[end]-bsol.domain.indexes[1]+1
    Nrxns = length(bsol.domain.phase.reactions)
    x,dp = extract_local_sensitivities(bsol.sol,t)
    svals = dp[inddeno+length(bsol.domain.phase.species)]
    s = svals[indnum]
    V = getV(bsol,t)
    T = getT(bsol,t)
    P = getP(bsol,t)
    C = getC(bsol,t)
    c = bsol.sol(t)[indnum]/V
    k = bsol.domain.p[inddeno+length(bsol.domain.phase.species)]
    val = k*(s-c*sum(svals)*R*bsol.domain.T/bsol.domain.P)/(c*V) #known T and P
    if t == 0
        return 0.0
    else
        return val
    end
end

export getconcentrationsensitivity

"""
calculate the rates of all reactions at time t
"""
function rates(bsol::Q,t::X) where {Q<:Simulation,X<:Real}
    cs,kfs,krevs = calcthermo(bsol.domain,bsol.sol(t),t)[[2,9,10]]
    return [getrate(rxn,cs,kfs,krevs) for rxn in bsol.domain.phase.reactions]
end

"""
calculate the rates of all reactions at given times ts
defaults to using bsol.sol.t if ts is not supplied
"""
function rates(bsol::Q;ts::X=Array{Float64,1}()) where {Q<:Simulation,X<:AbstractArray}
    if length(ts) == 0
        ts = bsol.sol.t
    end
    return hcat([rates(bsol,t) for t in ts]...)
end

export rates


text:
# # Test with manually weighted ICN
icn = ICN(param=true)
@test max_icn_length() == 18
show_layers(icn)
icn.weigths = vcat(trues(18), falses(6))
@test CN.is_viable(icn)
@test length(icn) == 31

compo = compose(icn)
@test code(compo; name = "test_composition") == "test_composition = identity ∘ sum ∘ sum ∘ [param_minus_val, val_minus_param" *
    ", count_bounding_param, count_g_param, count_l_param, count_eq_param," *
    " contiguous_vals_minus_rev, contiguous_vals_minus, count_l_right, count_g_right" * ", count_l_left, count_g_left, count_lesser, count_greater, count_eq_right, " * "count_eq_left, count_eq, identity]"

v = [1,2,4,3]
@test composition(compo)(v; param=2, dom_size=4) == 67

CN.generate_weights(icn)

## Test GA and exploration
domains = [domain([1,2,3,4]) for i in 1:4]
compo, _ = explore_learn_compose(domains, allunique)
@test composition(compo)([1,2,3,3], dom_size = 4) > 0.0

## Test export to file
composition_to_file!(compo, "test_dummy.jl", "all_different")


text:
# PROX_SUM_SQUARE    Proximal operator of sum-of-squares.
#
#   prox_sum_square(v,lambda) is the proximal operator of
#   (1/2)||.||_2^2 with parameter lambda.
function prox_sum_square(v, lambda)
    x = (1 / (1 + lambda)) * v
    return x
end


text:
using TerminalLoggers
using Test

using Logging:
    LogLevel, BelowMinLevel, Debug, Info, Warn, Error, AboveMaxLevel,
    shouldlog, handle_message, min_enabled_level, catch_exceptions

using ProgressLogging: Progress
using UUIDs: UUID

include("TerminalLogger.jl")
include("StickyMessages.jl")