text:
# # Quadrotor
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/Quadrotor.ipynb)
#
#md # !!! note "Overview"
#md #     System type: polynomial continuous system\
#md #     State dimension: 2\
#md #     Application domain: Chemical kinetics
#
# ## Model description
#
# We study the dynamics of a quadrotor as derived in [xxx]. Let us first introduce
# the variables required to describe the model: the inertial (north) position
# ``x_1``, the inertial (east) position ``x_2``, the altitude ``x_3``, the
# longitudinal velocity ``x_4``, the lateral velocity ``x_5``, the vertical
# velocity ``x_6``, the roll angle ``x_7``, the pitch angle ``x_8``, the yaw
# angle ``x_9``, the roll rate ``x_{10}``, the pitch rate ``x_{11}``, and the
# yaw rate ``x_{12}``. We further  require the following parameters: gravity
# constant ``g = 9.81`` [m/s``^2``], radius of center mass ``R = 0.1`` [m],
# distance of motors to center mass ``l = 0.5`` [m], motor mass
# ``M_{rotor} = 0.1`` [kg], center mass ``M = 1`` [kg], and total mass
# ``m = M + 4M_{rotor}``.
#
# From the above parameters we can compute the moments of inertia as
#
# ```math
# \begin{array}{lcll}
# J_x & = & \frac{2}{5}, M, R^2 + 2, l^2, M_{rotor}, \\
# J_y & = & J_x, \\
# J_z & = & \frac{2}{5}, M, R^2 + 4, l^2, M_{rotor}.
# \end{array}
# ```
#
# Finally, we can write the set of ordinary differential equations for the
# quadrotor according to [xxx]:
# ```math
# \left\{
# \begin{array}{lcl}
# \dot{x}_1 & = & \cos(x_8)\cos(x_9)x_4 + \Big(\sin(x_7)\sin(x_8)\cos(x_9) - \cos(x_7)\sin(x_9)\Big)x_5 \\
# & & + \Big(\cos(x_7)\sin(x_8)\cos(x_9) + \sin(x_7)\sin(x_9)\Big)x_6 \\
# \dot{x}_2 & = & \cos(x_8)\sin(x_9)x_4 + \Big(\sin(x_7)\sin(x_8)\sin(x_9) + \cos(x_7)\cos(x_9)\Big)x_5 \\
# & & + \Big(\cos(x_7)\sin(x_8)\sin(x_9) - \sin(x_7)\cos(x_9)\Big)x_6 \\
# \dot{x}_3 & = & \sin(x_8)x_4 - \sin(x_7)\cos(x_8)x_5 - \cos(x_7)\cos(x_8)x_6 \\
# \dot{x}_4 & = & x_{12}x_5 - x_{11}x_6 - g\sin(x_8) \\
# \dot{x}_5 & = & x_{10}x_6 - x_{12}x_4 + g\cos(x_8)\sin(x_7) \\
# \dot{x}_6 & = & x_{11}x_4 - x_{10}x_5 + g\cos(x_8)\cos(x_7) - \frac{F}{m} \\
# \dot{x}_7 & = & x_{10} + \sin(x_7)\tan(x_8)x_{11} + \cos(x_7)\tan(x_8)x_{12} \\
# \dot{x}_8 & = & \cos(x_7)x_{11} - \sin(x_7)x_{12} \\
# \dot{x}_9 & = & \frac{\sin(x_7)}{\cos(x_8)}x_{11} + \frac{\cos(x_7)}{\cos(x_8)}x_{12} \\
# \dot{x}_{10} & = & \frac{J_y - J_z}{J_x}x_{11}x_{12} + \frac{1}{J_x}\tau_\phi \\
# \dot{x}_{11} & = & \frac{J_z - J_x}{J_y}x_{10}x_{12} + \frac{1}{J_y}\tau_\theta \\
# \dot{x}_{12} & = & \frac{J_x - J_y}{J_z}x_{10}x_{11} + \frac{1}{J_z}\tau_\psi
# \end{array}
# \right.
# ```
#
# To check interesting control specifications, we stabilize the quadrotor using
# simple PD controllers for height, roll, and pitch. The inputs to the
# controller are the desired values for height, roll, and pitch ``u_1``, ``u_2``,
# and ``u_3``, respectively. The equations of the controllers are:
# ```math
# \begin{array}{lcll}
# F & = & m \, g - 10(x_3 - u_1) + 3x_6 \; & (\text{height control}), \\
# \tau_\phi & = & -(x_7 - u_2) - x_{10} & (\text{roll control}), \\
# \tau_\theta & = & -(x_8 - u_3) - x_{11} & (\text{pitch control}).
# \end{array}
# ```

# We leave the heading uncontrolled so that we set ``\tau_\psi = 0``.

using ReachabilityAnalysis
using ReachabilityAnalysis: is_intersection_empty

# ## Reachability settings
#
#The task is to change the height from ``0``~[m] to ``1``~[m] within ``5``~[s].
# A goal region ``[0.98,1.02]`` of the height ``x_3`` has to be reached within
# ``5``~[s] and the height has to stay below ``1.4`` for all times. After
# ``1``~[s] the height should stay above ``0.9``~[m]. The initial value for the
# position and velocities (i.e., from ``x_1`` to ``x_6``) is uncertain and given
# by ``[-\Delta,\Delta]``~[m], with ``\Delta=0.4``. All other variables are
# initialized to ``0``. This preliminary analysis must be followed by a
# corresponding evolution for ``\Delta = 0.1`` and ``\Delta= 0.8`` while keeping
# all the settings the same. No goals are specified for these cases: the
# objective instead is to understand the scalability of each tool with fixed
# settings.

# parameters of the model
const g = 9.81           # gravity constant in m/s^2
const R = 0.1            # radius of center mass in m
const l = 0.5            # distance of motors to center mass in m
const Mrotor = 0.1       # motor mass in kg
const M = 1.0            # center mass in kg
const m = M + 4*Mrotor   # total mass in kg
const mg = m*g

# moments of inertia
const Jx = (2/5)*M*R^2 + 2*l^2*Mrotor
const Jy = Jx
const Jz = (2/5)*M*R^2 + 4*l^2*Mrotor
const Cyzx = (Jy - Jz)/Jx
const Czxy = (Jz - Jx)/Jy
const Cxyz = 0.0 #(Jx - Jy)/Jz

# considering the control parameters as *parameters*
const u₁ = 1.0
const u₂ = 0.0
const u₃ = 0.0

const Tspan = (0.0, 5.0)
const v3 = LazySets.SingleEntryVector(3, 12, 1.0)

@inline function quad_property(solz)
    tf = tend(solz)

    #Condition: b1 = (x[3] < 1.4) for all time
    unsafe1 = HalfSpace(-v3, -1.4) # unsafe: -x3 <= -1.4
    b1 =  all([is_intersection_empty(unsafe1, set(R)) for R in solz(0.0 .. tf)])
    #b1 = ρ(v3, solz) < 1.4

    #Condition: x[3] > 0.9 for t ≥ 1.0
    unsafe2 = HalfSpace(v3, 0.9) # unsafe: x3 <= 0.9
    b2 = all([is_intersection_empty(unsafe2, set(R)) for R in solz(1.0 .. tf)])

    #Condition: x[3] ⊆ Interval(0.98, 1.02) for t ≥ 5.0 (t=5 is `tf`)
    b3 = set(project(solz[end], vars=(3))) ⊆ Interval(0.98, 1.02)

    return b1 && b2 && b3
end

@taylorize function quadrotor!(dx, x, params, t)
    #unwrap the variables and the controllers; the last three are the controllers
    #x₁, x₂, x₃, x₄, x₅, x₆, x₇, x₈, x₉, x₁₀, x₁₁, x₁₂, u₁, u₂, u₃ = x
    x₁  = x[1]
    x₂  = x[2]
    x₃  = x[3]
    x₄  = x[4]
    x₅  = x[5]
    x₆  = x[6]
    x₇  = x[7]
    x₈  = x[8]
    x₉  = x[9]
    x₁₀ = x[10]
    x₁₁ = x[11]
    x₁₂ = x[12]

    #equations of the controllers
    F = (mg - 10*(x₃ - u₁)) + 3*x₆  # height control
    τϕ = -(x₇ - u₂) - x₁₀            # roll control
    τθ = -(x₈ - u₃) - x₁₁            # pitch control
    local τψ = 0.0                   # heading is uncontrolled

    Tx = τϕ/Jx
    Ty = τθ/Jy
    Tz = τψ/Jz
    F_m = F/m

    #Some abbreviations
    sx7 = sin(x₇)
    cx7 = cos(x₇)
    sx8 = sin(x₈)
    cx8 = cos(x₈)
    sx9 = sin(x₉)
    cx9 = cos(x₉)

    sx7sx9 = sx7*sx9
    sx7cx9 = sx7*cx9
    cx7sx9 = cx7*sx9
    cx7cx9 = cx7*cx9
    sx7cx8 = sx7*cx8
    cx7cx8 = cx7*cx8
    sx7_cx8 = sx7/cx8
    cx7_cx8 = cx7/cx8

    x4cx8 = cx8*x₄

    p11 = sx7_cx8*x₁₁
    p12 = cx7_cx8*x₁₂
    xdot9 = p11 + p12

    #differential equations for the quadrotor

    dx[1] = (cx9*x4cx8 + (sx7cx9*sx8 - cx7sx9)*x₅) + (cx7cx9*sx8 + sx7sx9)*x₆
    dx[2] = (sx9*x4cx8 + (sx7sx9*sx8 + cx7cx9)*x₅) + (cx7sx9*sx8 - sx7cx9)*x₆
    dx[3] = (sx8*x₄ - sx7cx8*x₅) - cx7cx8*x₆
    dx[4] = (x₁₂*x₅ - x₁₁*x₆) - g*sx8
    dx[5] = (x₁₀*x₆ - x₁₂*x₄) + g*sx7cx8
    dx[6] = (x₁₁*x₄ - x₁₀*x₅) + (g*cx7cx8 - F_m)
    dx[7] = x₁₀ + sx8*xdot9
    dx[8] = cx7*x₁₁ - sx7*x₁₂
    dx[9] = xdot9
    dx[10] = Cyzx * (x₁₁ * x₁₂) + Tx
    dx[11] = Czxy * (x₁₀ * x₁₂) + Ty
    dx[12] = Cxyz * (x₁₀ * x₁₁) + Tz

    return dx
end

function quadrotor(; T=5.0, plot_vars=[0, 3],
                property=quad_property,
                project_reachset=true,
                Wpos = 0.4, Wvel = 0.4)

    #initial conditions
    X0c = zeros(12)
    ΔX0 = [Wpos, Wpos, Wpos, Wvel, Wvel, Wvel, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    X0 = Hyperrectangle(X0c, ΔX0)

    #initial-value problem
    prob = @ivp(x' = quadrotor!(x), dim: 12, x(0) ∈ X0);

    return prob
end

#-

# ## Results

cases = ["Δ=0.1", "Δ=0.4", "Δ=0.8"];

# ----------------------------------------
#  Case 1: smaller uncertainty
# ----------------------------------------
Wpos = 0.1
Wvel = 0.1
prob = quadrotor(project_reachset=false, Wpos=Wpos, Wvel=Wvel)
alg = TMJets(abstol=1e-7, orderT=5, orderQ=1, adaptive=false)

# Warm-up run
sol1 = solve(prob, tspan=Tspan, alg=alg);
solz1 = overapproximate(sol1, Zonotope);

# ----------------------------------------
# Case 2: intermediate uncertainty
# ----------------------------------------
Wpos = 0.4
Wvel = 0.4
prob = quadrotor(project_reachset=false, Wpos=Wpos, Wvel=Wvel)
alg = TMJets(abstol=1e-7, orderT=5, orderQ=1, adaptive=false)

# warm-up run
sol2 = solve(prob, tspan=Tspan, alg=alg);
solz2 = overapproximate(sol2, Zonotope);

# verify that specification holds
#property = quad_property(solz2)
#println("Validate property, case $(cases[2]) : $(property)")

# ----------------------------------------
# Case 3: large uncertainty
# ----------------------------------------
Wpos = 0.8
Wvel = 0.8
prob = quadrotor(project_reachset=false, Wpos=Wpos, Wvel=Wvel)
alg = TMJets(abstol=1e-7, orderT=5, orderQ=1, adaptive=false)

# Warm-up run
sol3 = solve(prob, tspan=Tspan, alg=alg);
solz3 = overapproximate(sol3, Zonotope);

#-

using Plots

Plots.plot(solz3,  vars=(0, 3), linecolor="green",  color=:green,  alpha=0.8)
Plots.plot!(solz2, vars=(0, 3), linecolor="blue",   color=:blue,   alpha=0.8)
Plots.plot!(solz1, vars=(0, 3), linecolor="yellow", color=:yellow, alpha=0.8,
    xlab="t", ylab="x3",
    xtick=[0., 1., 2., 3., 4., 5.], ytick=[-1., -0.5, 0., 0.5, 1., 1.5],
    xlims=(0., 5.), ylims=(-1., 1.5))


text:
sha(s) = bytes2hex(sha256(s))
path2sha(path::AbstractString) = sha(read(path))

"""
    State(input_sha::String, julia_version::String)

State obtained from a Pluto notebook file (".jl") where

- `input_sha`: SHA checksum calculated over the file
- `julia_version`: \\\$VERSION
"""
struct State
    input_sha::String
    julia_version::String
end

"""
    State(text::AbstractString)

Create a new State from a Pluto notebook file (".jl").
"""
State(text::AbstractString) = State(sha(text), string(VERSION))

function n_cache_lines()
    # Determine length by using `string(state::State)`.
    state = State("a", "b")
    lines = split(string(state), '\n')
    return length(lines)
end

const STATE_IDENTIFIER = "[PlutoStaticHTML.State]"

function string(state::State)::String
    return """
        <!--
            # This information is used for caching.
            $STATE_IDENTIFIER
            input_sha = "$(state.input_sha)"
            julia_version = "$(state.julia_version)"
        -->
        """
end

"Extract State from a HTML file which contains a State as string somewhere."
function extract_state(html::AbstractString)::State
    sep = '\n'
    lines = split(html, sep)
    start = findfirst(contains(STATE_IDENTIFIER), lines)::Int
    stop = start + 2
    info = join(lines[start:stop], sep)
    entries = parsetoml(info)["PlutoStaticHTML"]["State"]
    return State(entries["input_sha"], entries["julia_version"])
end

"Convert a notebook at `path` to a State."
function path2state(path::AbstractString)::State
    @assert endswith(path, ".jl")
    code = read(path, String)
    State(code)
end

ext: jl


text:
module Tools3D

using PyPlot, Random, TransD_GP, GP, Statistics

function geomprogdepth(n, dy, c)
    dy*(1.0-c^n)/(1-c)
end

function getn(z, dy, c)
    log(1 - z/dy*(1-c))/log(c)
end

function nicenup(g::PyPlot.Figure;fsize=16)
    for ax in gcf().axes
        ax.tick_params("both",labelsize=fsize)
        ax.xaxis.label.set_fontsize(fsize)
        ax.yaxis.label.set_fontsize(fsize)
        any(keys(ax) .== :zaxis) && ax.zaxis.label.set_fontsize(fsize)
        ax.title.set_fontsize(fsize)

        if typeof(ax.get_legend_handles_labels()[1]) != Array{Any,1}
            ax.legend(loc="best", fontsize=fsize)
        end
    end
    g.tight_layout()
end

function meshgrid(xs, ys)
    xx, yy = [xs[i] for i in 1:length(xs), j in 1:length(ys)], [ys[j] for i in 1:length(xs), j in 1:length(ys)]
    return xx, yy
end

function makezn(;n          = 20,
                 dz         = 1.6,
                 extendfrac = 1.245)
    @assert typeof(n) == Int
    @assert dz > 0.0
    @assert extendfrac > 1.0
    znrange            = 1.0:n
    zboundaries        = geomprogdepth.(znrange, dz, extendfrac)
    thickness          = [zboundaries[1]; diff(zboundaries)[1:end-1]]
    zall               = [zboundaries[1]/2; 0.5*(zboundaries[1:end-1] + zboundaries[2:end])]
    znall              = getn.(zall, dz, extendfrac)

    figure()
    plot(znall, zall)
    xlabel("depth index")
    ylabel("depth associated km")
    grid()
    nicenup(gcf())

    f, ax = plt.subplots(1, 2, figsize=(10,5))
    ax[1].stem(zboundaries[1:end-1], zboundaries[1:end-1], markerfmt="")
    ax[1].stem(zall, zall, "k--", markerfmt=" ")
    ax[1].set_xlabel("depth km")
    ax[1].set_ylabel("depth km")

    ax[2].stem(znrange[1:end-1], znrange[1:end-1], markerfmt="")
    ax[2].stem(znall, znall, "k--", markerfmt=" ")
    ax[2].set_ylabel("depth index")
    ax[2].set_xlabel("depth index")
    nicenup(gcf())

    f, ax = plt.subplots(1, 2, figsize=(10,5), sharey=true)
    ax[1].stem(zall[1:end-1],thickness, "k--", markerfmt=" ")
    ax[1].set_xlabel("depth km")
    ax[1].set_ylabel("thickness km")
    ax[1].yaxis.grid(which="major")
    ax[2].stem(znall[1:end-1],thickness, "k--", markerfmt=" ")
    ax[2].set_xlabel("depth index")
    ax[2].yaxis.grid(which="major")
    nicenup(f)
    return zall, znall
end

function makeopt(;nmin      = 2,
                  nmax      = 400,
                  λ         = [100.0, 200.0, 1.0],
                  δ         = 0.1,
                  fbounds   = [0.5 4],
                  demean    = true,
                  sdev_prop = [0.4],
                  sdev_pos  = [80.0;80.0;1.0],
                  stat_window = 500,
                  pnorm     = 2.0,
                  λx        = 1600.0,
                  λy        = 1600.0,
                  dxfrac    = 0.025,
                  dyfrac    = 0.025,
                  rseed     = 4,
                  timebirth = true,
                  znall     = nothing
                )
    @assert znall != nothing
    dx = dxfrac*λx
    dy = dyfrac*λy
    x = 0:(dx):λx-dx
    y = 0:(dy):λy-dy
    xall = zeros(3,length(x)*length(y)*length(znall))
    for i in 1:size(xall,2)
        xid, yid, zid = Tuple(CartesianIndices((length(x),length(y),length(znall)))[i])
        xall[:,i] = [x[xid]; y[yid]; znall[zid]]
    end
    xbounds = zeros(Float64,size(xall,1),2)
    for dim in 1:size(xall, 1)
        xbounds[dim,:] = [minimum(xall[dim,:]), maximum(xall[dim,:])]
    end
    ## Initialize a model using these options
    Random.seed!(rseed)
    opt = TransD_GP.OptionsStat(nmin = nmin,
                            nmax = nmax,
                            xbounds = xbounds,
                            fbounds = fbounds,
                            xall = xall,
                            λ = λ,
                            δ = δ,
                            demean = demean,
                            sdev_prop = sdev_prop,
                            sdev_pos = sdev_pos,
                            stat_window = stat_window,
                            pnorm = pnorm,
                            quasimultid = false,
                            updatenonstat = false,
                            needλ²fromlog = false
                            )

    optdummy = TransD_GP.OptionsNonstat(opt,
                            nmin = 2,
                            nmax = 3,
                            fbounds = fbounds,
                            δ = δ,
                            demean = demean,
                            sdev_prop = sdev_prop,
                            sdev_pos = sdev_pos,
                            pnorm = pnorm,
                            K = opt.K
                            )
    return opt, optdummy, x, y
end

function slicemodel(m::TransD_GP.Model,
                    opt::TransD_GP.Options;
                    slicesx    = [],
                    slicesy    = [],
                    slicesz    = [],
                    dz         = nothing,
                    extendfrac = nothing,
                    figsize    = (7,9)
                    )
    @assert !any((dz, extendfrac) .== nothing)
    @assert any([slicesx !=  [], slicesy != [], slicesz !=[]])
    slicemodel(m.fstar, m.n, m.xtrain, m.ftrain, opt::TransD_GP.Options; slicesx=slicesx, slicesy=slicesy, slicesz=slicesz,
                dz=dz, extendfrac=extendfrac, figsize=figsize)
end

function slicemodel(fstar::Array{Float64}, npoints::Int,
                    xtrain::Array{Float64, 2}, ftrain::Array{Float64},
                    opt::TransD_GP.Options;
                    slicesx    = [],
                    slicesy    = [],
                    slicesz    = [],
                    dz         = nothing,
                    extendfrac = nothing,
                    figsize    = (7,9)
                    )
    x = unique(opt.xall[1,:])
    y = unique(opt.xall[2,:])
    z = geomprogdepth.(unique(opt.xall[3,:]), dz, extendfrac)
    @assert !any((dz, extendfrac) .== nothing)
    @assert any([slicesx !=  [], slicesy != [], slicesz !=[]])
    nplots = sum([slicesx != [], slicesy != [], slicesz != []])
    v = reshape(fstar,length(x), length(y), length(z))
    fig = figure(figsize=(nplots*figsize[1],figsize[2]))
    using3D()
    ax = Array{Any,1}(undef, nplots)
    iplot = 0
    if slicesz!=[]
        xx,yy = meshgrid(x,y)
        l = slicesz
        iplot+=1
        ax[iplot]=fig.add_subplot(1, nplots, iplot, projection="3d")
        for i in l
            c=(v[:,:,i].-minimum(v))./(maximum(v)-minimum(v))
            ax[iplot].plot_surface(xx, yy, z[i]*ones(size(yy)), facecolors=plt.cm.jet_r(c), shade=false)
        end
    end
    if slicesx!=[]
        yy,zz = meshgrid(y,z)
        l = slicesx
        iplot+=1
        ax[iplot]=fig.add_subplot(1, nplots, iplot, projection="3d")
        for i in l
            c=(v[i,:,:].-minimum(v))./(maximum(v)-minimum(v))
            ax[iplot].plot_surface(x[i]*ones(size(yy)), yy, zz, facecolors=plt.cm.jet_r(c), shade=false)
        end
    end
   if slicesy!=[]
        xx,zz = meshgrid(x,z)
        l = slicesy
        iplot+=1
        ax[iplot]=fig.add_subplot(1, nplots, iplot, projection="3d")
        for i in l
            c=(v[:,i,:].-minimum(v))./(maximum(v)-minimum(v))
            ax[iplot].plot_surface(xx, y[i]*ones(size(zz)), zz, facecolors=plt.cm.jet_r(c), shade=false)
        end
    end
    for a in ax
        if npoints>0
            a.scatter3D(xtrain[1,1:npoints], xtrain[2,1:npoints], geomprogdepth.(xtrain[3,1:npoints], dz, extendfrac),
                c=ftrain[1:npoints], vmin=minimum(v), vmax=maximum(v), s=50, cmap="jet_r")
        end
        a.set_xlim(extrema(x))
        a.set_ylim(extrema(y))
        a.set_zlim(extrema(z))
        a.set_xlabel("x km")
        a.set_ylabel("y km")
        a.set_zlabel("depth km")
        a.invert_zaxis()
    end
    nicenup(gcf(),fsize=14)
    fig.subplots_adjust(right=0.9)
    cbar_ax = fig.add_axes([0.92, 0.15, 0.025, 0.7])
    cbar = fig.colorbar(plt.cm.ScalarMappable(cmap="jet_r",norm=matplotlib.colors.Normalize(vmin=minimum(v), vmax=maximum(v))), cax=cbar_ax)
    cbar.set_label(L"\log_{10} \rho", fontsize=14)
    cbar.ax.tick_params(labelsize=14)
end

function makecubemodel(opt::TransD_GP.Options;
                        dz         = nothing,
                        extendfrac = nothing,
                        z1         = 10.0,
                        z2         = 400.0,
                        yleft      = 500.0,
                        yright     = 1100.0,
                        xup        = 980.0,
                        xdown      = 580.0,
                        zup        = 70.0,
                        zdown      = 250.0,
                        ρ0         = 2,
                        ρ1         = 4,
                        ρ2         = 0.5,
                        ρanom      = 0.8,
                        ntosmooth  = 1,
                        )
    @assert extendfrac != nothing
    @assert dz != nothing
    @assert z1<zup<zdown<z2
    @assert yleft<yright
    @assert xdown<xup

    ρ = zeros(size(opt.xall, 2))
    zmin, zmax = geomprogdepth.(extrema(opt.xall[3,:]), dz, extendfrac)
    for i = 1:size(opt.xall, 2)
        x, y, zn = opt.xall[:,i]
        z = geomprogdepth(zn, dz, extendfrac)
        if z<z1
            ρ[i] = ρ0
        elseif (zup<z<zdown) & (yleft<y<yright) & (xdown<x<xup)
            ρ[i] = ρanom*(ρ1 + (z-zmin)/(zmax-z1)*(ρ2-ρ1)) #ρanom
        elseif z>z2
            ρ[i] = ρ2
        else
            ρ[i] = ρ1 + (z-zmin)/(zmax-z1)*(ρ2-ρ1)
        end
    end
    nx, ny, nz = length(unique(opt.xall[1,:])), length(unique(opt.xall[2,:])), length(unique(opt.xall[3,:]))
    ρ .= boxcarn(reshape(ρ, nx, ny, nz), ntosmooth)[:]
    opt.fbounds = [0.98*minimum(ρ) 1.02*maximum(ρ)]
    return ρ
end

function get_training_data(fstar::Array{Float64, 1},
                           opt::TransD_GP.Options;
                           dz = nothing,
                           extendfrac = nothing,
                           sdmaxfrac = 0.05,
                           rseed     = 12,
                           zbreak    = 12321.0,
                           takeevery = 4,
                           fractrain = 0.05
                          )
    @assert extendfrac != nothing
    @assert dz != nothing
    @assert zbreak != 12321.0
    @assert sdmaxfrac > 0 && sdmaxfrac < 1
    Random.seed!(rseed)
    x, y =  unique(opt.xall[1,:]), unique(opt.xall[2,:])
    z = geomprogdepth.(unique(opt.xall[3,:]), dz, extendfrac)
    @assert minimum(z) < zbreak < maximum(z)
    noisyd = NaN .+ zeros(Float64, size(fstar))
    ntrain = round(Int, fractrain*length(fstar))
    linidx = randperm(length(fstar))[1:ntrain]
    δtry = sdmaxfrac*maximum(abs.(fstar))
    for l in linidx
        row, col, zid = Tuple(CartesianIndices((length(x),length(y),length(z)))[l])
        if z[zid] > zbreak || (rem(row, takeevery) == 0 && rem(col, takeevery) == 0 && rem(zid, takeevery) == 0)
            noisyd[l] = fstar[l] + δtry*randn()
        end
    end
    f = figure(figsize=(10,10))
    scatter3D(opt.xall[1,:], opt.xall[2,:], geomprogdepth.(opt.xall[3,:], dz, extendfrac),
                          c=noisyd, vmin=minimum(fstar), vmax=maximum(fstar), cmap="jet_r")
    xlim(extrema(x))
    ylim(extrema(y))
    zlim(extrema(z))
    xlabel("x km")
    ylabel("y km")
    zlabel("depth km")
    cbar = f.colorbar(plt.cm.ScalarMappable(cmap="jet_r",norm=matplotlib.colors.Normalize(vmin=minimum(fstar), vmax=maximum(fstar))),ax=gca())
    cbar.set_label(L"\log_{10} \rho")
    gca().zaxis.label.set_fontsize(14); nicenup(gcf())
    gca().invert_zaxis()

    noisyd, δtry
end

function plot_last_target_model(opt_in::TransD_GP.Options, d::Array{Float64, 1}, sd::Float64;
                                slicesx    = [],
                                slicesy    = [],
                                slicesz    = [],
                                dz         = nothing,
                                extendfrac = nothing,
                                fsize      = 14
                               )
    @assert !any((dz, extendfrac) .== nothing)
    @assert any([slicesx !=  [], slicesy != [], slicesz !=[]])
    Tacrosschains = gettargtemps(opt_in)
    last_target_model_idx = findall(abs.(Tacrosschains[end,:] .-1.0) .< 1e-12)
    for idx in last_target_model_idx
        opt_in.fstar_filename = "models_"*opt_in.fdataname*"_$idx.bin"
        m_last = TransD_GP.history(opt_in, stat=:fstar)[end]
        slicemodel(m_last, 0,[0. 0.],[0.], opt_in, slicesx=slicesx, slicesy=slicesy, slicesz=slicesz, dz=dz, extendfrac=extendfrac)
        sd > 0.0 && calc_simple_RMS(d, m_last, sd)
    end

end

function gettargtemps(opt_in::TransD_GP.Options)
    nchains = length(filter( x -> occursin(r"misfits.*bin", x), readdir(pwd()) )) # my terrible regex
    @info "Number of chains is $nchains"
    # now look at any chain to get how many iterations
    costs_filename = "misfits_"*opt_in.fdataname
    opt_in.costs_filename    = costs_filename*"_1.bin"
    iters          = TransD_GP.history(opt_in, stat=:iter)
    niters         = length(iters)
    @info "McMC has run for $(iters[end]) iterations"
    # then create arrays of unsorted by temperature T
    Tacrosschains  = zeros(Float64, niters, nchains)
    # get the values into the arrays
    for ichain in 1:nchains
        opt_in.costs_filename = costs_filename*"_$ichain.bin"
        Tacrosschains[:,ichain] = TransD_GP.history(opt_in, stat=:T)
    end
    Tacrosschains
end

function assembleTat1(opt::TransD_GP.Options, irange::StepRange{Int})
    Tacrosschains = gettargtemps(opt)
    iters = TransD_GP.history(opt, stat=:iter)
    @assert 0<irange.start<=length(iters)
    @assert irange.stop<=length(iters)
    @info "obtaining models $(iters[irange.start]) to $(iters[irange.stop])"
    nat1 = length(findall(abs.(Tacrosschains[end,:] .-1.0) .< 1e-12))
    @info "No. of chains at 1 is $nat1"
    niters = length(irange)
    mat1 = zeros(Float64, size(opt.xall,2), niters, nat1)
    cumT = cumsum(Tacrosschains .== 1,dims=1)
    @info size(mat1)
    for (istep, stp) in enumerate(irange)
        @info "$(istep/niters*100)% done"
        for (itemp, idx) in enumerate(findall(abs.(Tacrosschains[stp,:] .-1.0) .< 1e-12))
            opt.fstar_filename = "models_"*opt.fdataname*"_$idx.bin"
            mat1[:,istep,itemp] = TransD_GP.history(opt, stat=:fstar)[cumT[stp,idx]]
        end
    end
    mat1
end

function assembleTat1(opt::TransD_GP.Options; burninfrac=0.5)
    @assert 0.0<1.0burninfrac<1.0
    Tacrosschains = gettargtemps(opt)
    iters = TransD_GP.history(opt, stat=:iter)
    start = round(Int, length(iters)*burninfrac)
    @info "obtaining models $(iters[start]) to $(iters[end])"
    nmodels = sum((Tacrosschains[start:end,:] .== 1))
    # storing models column wise in nmodels rows
    mat1 = zeros(Float64, size(opt.xall,2), nmodels)
    imodel = 0
    for ichain in 1:size(Tacrosschains, 2)
        @info "chain $ichain"
        at1idx = findall(Tacrosschains[:,ichain].==1).>= start
        ninchain = sum(at1idx)
        ninchain == 0 && continue
        opt.fstar_filename = "models_"*opt.fdataname*"_$ichain.bin"
        mat1[:,imodel+1:imodel+ninchain] .= reduce(hcat, TransD_GP.history(opt, stat=:fstar)[at1idx])
        imodel += ninchain
    end
    mat1
end

function savepngs(opt::TransD_GP.Options, irange::StepRange{Int};
                  slicesx    = [],
                  slicesy    = [],
                  slicesz    = [],
                  dz         = nothing,
                  extendfrac = nothing,
                  fsize      = 14
                 )
    @assert !any((dz, extendfrac) .== nothing)
    @assert any([slicesx !=  [], slicesy != [], slicesz !=[]])
    Tacrosschains = gettargtemps(opt)
    iters = TransD_GP.history(opt, stat=:iter)
    @assert 0<irange.start<=length(iters)
    @assert irange.stop<=length(iters)
    nat1 = length(findall(abs.(Tacrosschains[end,:] .-1.0) .< 1e-12))
    @info "No. of chains at 1 is $nat1"
    @info "obtaining models $(iters[irange.start]) to $(iters[irange.stop])"
    M = assembleTat1(opt,irange)
    for (iridx, ir) in enumerate(irange), ic in 1:nat1
        Tools3D.slicemodel(M[:,iridx,ic],0,[0. 0.],[0.], opt,
                            slicesx=slicesx, slicesy=slicesy, slicesz=slicesz,
                            dz=dz, extendfrac=extendfrac)
        @info "saving target model $(iters[ir]) chain $ic"
        gcf().suptitle("model_"*lpad(iters[ir],7,'0')*"_chain_"*lpad(ic,3,'0'))
        savefig("model_"*lpad(iters[ir],7,'0')*"_chain_"*lpad(ic,3,'0')*".png", dpi=300)
        close()
    end
    nothing
end

function calc_simple_RMS(d::AbstractArray, fstar::Array{Float64, 1}, sd::Float64)

    select = .!isnan.(d)
    r = (d[select] - fstar[select])/sd
    n = sum(select)
    @info "χ^2 error is $(r'*r) for $n points RMS: $(sqrt(r'*r/n))"
    nothing
end

function boxcarn(A::AbstractArray, ntosmooth::Int)
    out = zeros(eltype(A), size(A))
    R = CartesianIndices(A)
    Ifirst, Ilast = first(R), last(R)
    I1 = oneunit(Ifirst)
    for I in R
        n, s = 0, zero(eltype(out))
        for J in max(Ifirst, I-ntosmooth*I1):min(Ilast, I+ntosmooth*I1)
            s += A[J]
            n += 1
        end
        out[I] = s/n
    end
    out
end

end


text:
module TestUtils
    using ManifoldLearning
    using Test

    @testset "Utils" begin

	k = 3
	X = rand(3, 15)

	# test k-nn graph
	D, E = ManifoldLearning.find_nn(X, k, excluding=false)
	@test size(X,2) == size(D,2) && size(D, 1)-1 == k
	@test size(X,2) == size(E,2) && size(E, 1)-1 == k

	D, E = ManifoldLearning.find_nn(X, k)
	@test size(X,2) == size(D,2) && size(D, 1) == k
	@test size(X,2) == size(E,2) && size(E, 1) == k

	# test connected components
	CC = ManifoldLearning.components(E)
	@test length(CC) > 0

	# test shortest path
        D = transpose([7 9 14; 7 10 15; 9 10 2; 15 6 0; 6 9 11; 14 2 9])
        E = transpose([2 3 6; 1 3 4; 1 2 6; 2 5 4; 3 4 6; 1 3 5])
	P, PD = ManifoldLearning.dijkstra(D, E, 1)
	@test PD[5] == 20. && PD[6] == 11.
    end
end


text:
using WiltonInts84
using StaticArrays
using Test

p1 = SVector(0.0, 0.0, 0.0)
p2 = SVector(1.0, 0.0, 0.0)
p3 = SVector(0.0, 1.0, 0.0)
h = 0.1

x = -0.5 : 0.01 : 1.5
y = -0.5 : 0.01 : 1.5

r, R = 0.0, 0.15
fails = 0
M = zeros(typeof(h), length(x), length(y))
for i in eachindex(x)
    for j in eachindex(y)
        _c = SVector(x[i],y[j],h)
        try
            _I, _K = wiltonints(p1,p2,p3,_c,r,R,Val{0})
            M[i,j] = _I[2]
        catch
            # @show i j _c
            global fails += 1
        end
    end
end

# There are still some cases where the package fails,
# but at least lets make sure that number does not increase...
@test fails == 20


text:
module day22

include("inputs.jl")

ex = inputs.ex_d22()
input = inputs.input_d22()

function parse_inputs(input)
    p1blob, p2blob = split(input, "\n\n")
    map([p1blob, p2blob]) do deck
        parse.(Int, split(deck, "\n")[2:end])
    end
end

# Part 1
function play(hands)
    while !any(isempty.(hands))
        cards = popfirst!.(hands)
        if cards[1] > cards[2]
            push!(hands[1], cards...)
        else
            push!(hands[2], reverse(cards)...)
        end
    end
    return hands
end

function score(hands)
    score = 0
    for hand in hands; if !isempty(hand)
        score += sum([card*ind for (ind, card) in enumerate(reverse(hand))])
    end; end
    return score
end

part1(input) = parse_inputs(input) |> play |> score
@assert part1(ex) == 306
@info "Day 22, Part 1 answer: $(part1(input))"

# Part 2
function apply_win!(hands, cards, winner)
    push!(hands[winner], popat!(cards, winner))
    push!(hands[winner], pop!(cards))
end

function play_rc(hands)
    seen_before = Set()

    while !any(isempty.(hands))
        if hash(hands) ∈ seen_before
            return (1, hands)
        else
            push!(seen_before, hash(hands))
        end

        cards = popfirst!.(hands)   # draw cards

        if all(length.(hands) .>= cards)
            newhands = deepcopy(getindex.(hands, UnitRange.(1, cards)))
            (winner, _) = play_rc(newhands)
        else
            winner = cards[1] > cards[2] ? 1 : 2
        end

        apply_win!(hands, cards, winner)

    end

    winner = isempty(hands[1]) ? 2 : 1
    return (winner, hands)
end

part2(input) = parse_inputs(input) |> play_rc |> last |> score
@assert part2(ex) == 291
@info "Day 22, Part 2 answer: $(part2(input))"

end


text:
module tensorflow
  const _ProtoBuf_Top_ = @static isdefined(parentmodule(@__MODULE__), :_ProtoBuf_Top_) ? (parentmodule(@__MODULE__))._ProtoBuf_Top_ : parentmodule(@__MODULE__)
  include("tensor_shape_pb.jl")
  include("types_pb.jl")
  module tf2xla
    const _ProtoBuf_Top_ = @static isdefined(parentmodule(@__MODULE__), :_ProtoBuf_Top_) ? (parentmodule(@__MODULE__))._ProtoBuf_Top_ : parentmodule(@__MODULE__)
    include("host_compute_metadata_pb.jl")
  end
  module tpu
    const _ProtoBuf_Top_ = @static isdefined(parentmodule(@__MODULE__), :_ProtoBuf_Top_) ? (parentmodule(@__MODULE__))._ProtoBuf_Top_ : parentmodule(@__MODULE__)
    include("topology_pb.jl")
  end
end


text:
function advance!(pac::T, pap) where {T<:P_common{FdtdAcoustic}}
    for ipw in pac.activepw
        # store p for the last two steps
        # pppppp!(pap[ipw],pac.attrib_mod)
        advance_kernel!(pap[ipw], pac)
    end
    return nothing
end

function advance_kernel!(pap, pac::T) where {T<:P_common{FdtdAcoustic,2}}

    w1t = pap.w1[:t]
    mw = pap.memory_pml
    pml = pac.pml

    #compute dpdx and dpdz at [it-1] for all propagating fields
    @parallel compute_dp!(w1t[:p], w1t[:dpdx], w1t[:dpdz], pac.fc[:dzI], pac.fc[:dxI])
    memoryx!(mw[:dpdx], w1t[:dpdx], pml[:dpdx][:a], pml[:dpdx][:b], pml[:dpdx][:kI])
    memoryz!(mw[:dpdz], w1t[:dpdz], pml[:dpdz][:a], pml[:dpdz][:b], pml[:dpdz][:kI])

    #update velocity at [it-1/2] using
    #velocity at [it-3/2] and dpdx and dpdz at [it-1]
    @parallel compute_v!(
        w1t[:vx],
        w1t[:vz],
        pac.mod[:rhoI],
        w1t[:dpdx],
        w1t[:dpdz],
        pac.fc[:dt],
    )#

    #rigid boundary conditions
    @parallel (1:pac.ic[:nz]) dirichletx!(w1t[:vx], w1t[:vz], pac.ic[:nx])
    @parallel (1:pac.ic[:nx]) dirichletz!(w1t[:vz], w1t[:vx], pac.ic[:nz])

    @parallel compute_dv!(
        w1t[:vx],
        w1t[:vz],
        w1t[:dvxdx],
        w1t[:dvzdz],
        pac.fc[:dxI],
        pac.fc[:dzI],
    )#
    memoryx!(mw[:dvxdx], w1t[:dvxdx], pml[:dvxdx][:a], pml[:dvxdx][:b], pml[:dvxdx][:kI])
    memoryz!(mw[:dvzdz], w1t[:dvzdz], pml[:dvzdz][:a], pml[:dvzdz][:b], pml[:dvzdz][:kI])

    #compute pressure at [it] using p at [it-1] and dvxdx
    #and dvzdz at [it-1/2]
    @parallel compute_p!(w1t[:p], w1t[:dvxdx], w1t[:dvzdz], pac.mod[:K], pac.fc[:dt])

end

function advance_kernel!(pap, pac::T) where {T<:P_common{FdtdAcoustic,3}}

    w1t = pap.w1[:t]
    mw = pap.memory_pml
    pml = pac.pml

    #compute dpdx and dpdz at [it-1] for all propagating fields
    @parallel compute_dp!(
        w1t[:p],
        w1t[:dpdx],
        w1t[:dpdy],
        w1t[:dpdz],
        pac.fc[:dzI],
        pac.fc[:dyI],
        pac.fc[:dxI],
    )
    memoryx!(mw[:dpdx], w1t[:dpdx], pml[:dpdx][:a], pml[:dpdx][:b], pml[:dpdx][:kI])
    memoryy!(mw[:dpdy], w1t[:dpdy], pml[:dpdy][:a], pml[:dpdy][:b], pml[:dpdy][:kI])
    memoryz!(mw[:dpdz], w1t[:dpdz], pml[:dpdz][:a], pml[:dpdz][:b], pml[:dpdz][:kI])

    #update velocity at [it-1/2] using
    #velocity at [it-3/2] and dpdx and dpdz at [it-1]
    @parallel compute_v!(
        w1t[:vx],
        w1t[:vy],
        w1t[:vz],
        pac.mod[:rhoI],
        w1t[:dpdx],
        w1t[:dpdy],
        w1t[:dpdz],
        pac.fc[:dt],
    )#

    #rigid boundary conditions
    @parallel (1:pac.ic[:nz], 1:pac.ic[:ny]) dirichletx!(
        w1t[:vx],
        w1t[:vz],
        w1t[:vy],
        pac.ic[:nx],
    )
    @parallel (1:pac.ic[:nz], 1:pac.ic[:nx]) dirichlety!(
        w1t[:vy],
        w1t[:vz],
        w1t[:vx],
        pac.ic[:ny],
    )
    @parallel (1:pac.ic[:ny], 1:pac.ic[:nx]) dirichletz!(
        w1t[:vz],
        w1t[:vy],
        w1t[:vx],
        pac.ic[:nz],
    )

    @parallel compute_dv!(
        w1t[:vx],
        w1t[:vy],
        w1t[:vz],
        w1t[:dvxdx],
        w1t[:dvydy],
        w1t[:dvzdz],
        pac.fc[:dxI],
        pac.fc[:dyI],
        pac.fc[:dzI],
    )#
    memoryx!(mw[:dvxdx], w1t[:dvxdx], pml[:dvxdx][:a], pml[:dvxdx][:b], pml[:dvxdx][:kI])
    memoryy!(mw[:dvydy], w1t[:dvydy], pml[:dvydy][:a], pml[:dvydy][:b], pml[:dvydy][:kI])
    memoryz!(mw[:dvzdz], w1t[:dvzdz], pml[:dvzdz][:a], pml[:dvzdz][:b], pml[:dvzdz][:kI])

    #compute pressure at [it] using p at [it-1] and dvxdx
    #and dvzdz at [it-1/2]
    @parallel compute_p!(
        w1t[:p],
        w1t[:dvxdx],
        w1t[:dvydy],
        w1t[:dvzdz],
        pac.mod[:K],
        pac.fc[:dt],
    )

end

"""
Exchange pointers (i.e., set names of NamedArray) instead of copying arrays around
"""
function pppppp!(pap, attrib_mod)
    w2 = pap.w1
    names_old = names(w2)[1]
    names_new = vcat(circshift(names_old[1:3], -1), names_old[4:end])
    setnames!(w2, names_new, 1)
    # (old method), use for debugging
    #copyto!.(w2[:tpp],w2[:tp])
    #copyto!.(w2[:tp],w2[:t])
end

# these relative indices of the arrays point to same location
# [ix,iy,iz]     --> pressure
# [ix-1/2,iy,iz]       --> vx
# [ix,iy-1/2,iz]       --> vy
# [ix,iy,iz-1/2]       --> vz

@parallel function compute_dp!(p, dpdx, dpdz, dzI, dxI)
    @all(dpdx) = @d_xi(p) * dxI #
    @all(dpdz) = @d_zi(p) * dzI #
    return
end

@parallel function compute_dp!(p, dpdx, dpdy, dpdz, dzI, dyI, dxI)
    @all(dpdx) = @d_xi(p) * dxI #
    @all(dpdy) = @d_yi(p) * dyI #
    @all(dpdz) = @d_zi(p) * dzI #
    return
end

@parallel function compute_v!(vx, vz, rhoI, dpdx, dpdz, dt)#
    @inn(vx) = @inn(vx) + dt * @av_xi(rhoI) * @all(dpdx)
    @inn(vz) = @inn(vz) + dt * @av_zi(rhoI) * @all(dpdz)
    return
end

@parallel function compute_v!(vx, vy, vz, rhoI, dpdx, dpdy, dpdz, dt)#
    @inn(vx) = @inn(vx) + dt * @av_xi(rhoI) * @all(dpdx)
    @inn(vy) = @inn(vy) + dt * @av_yi(rhoI) * @all(dpdy)
    @inn(vz) = @inn(vz) + dt * @av_zi(rhoI) * @all(dpdz)
    return
end

#
@parallel function compute_dv!(vx, vz, dvxdx, dvzdz, dxI, dzI)#
    @all(dvxdx) = @d_xa(vx) * dxI #
    @all(dvzdz) = @d_za(vz) * dzI #
    return
end

@parallel function compute_dv!(vx, vy, vz, dvxdx, dvydy, dvzdz, dxI, dyI, dzI)#
    @all(dvxdx) = @d_xa(vx) * dxI #
    @all(dvydy) = @d_ya(vy) * dyI #
    @all(dvzdz) = @d_za(vz) * dzI #
    return
end

# no attenuation (no memory in stress-strain relation)
@parallel function compute_p!(p, dvxdx, dvzdz, K, dt)
    @all(p) = @all(p) + @all(K) * (@all(dvxdx) + @all(dvzdz)) * dt
    return
end

# no attenuation (no memory in stress-strain relation)
@parallel function compute_p!(p, dvxdx, dvydy, dvzdz, K, dt)
    @all(p) = @all(p) + @all(K) * (@all(dvxdx) + @all(dvzdz) + @all(dvydy)) * dt
    return
end

# #
#  function compute_p!(pap,pac,::T) where {T<:Union{FdtdAcousticVisco}}
# 	dvxdx=pap.w1[:dx][:vx]
# 	dvzdz=pap.w1[:dz][:vz]
# 	p=pap.w1[:t][:p]
# 	pp=pap.w1[:tp][:p]
# 	r=pap.w2[:t][:r]
# 	rp=pap.w2[:tp][:r]
# 	prvzvx!(p,pp,r,rp,dvzdz,dvxdx,pac.mod[:K],pac.mod3[:memcoeff1],pac.mod3[:memcoeff2],pac.fc[:dt],pac.ic[:nz],pac.ic[:nx],pac.ic[:nsls])
# end
# # viscoacoustic modeling (memory in stress-strain relation)
#  function prvzvx!(p,pp,r,rp,dvxdx,dvzdz,mod,mod31,mod32,fc1,nz,nx,nsls)#pw,dpdxw,dpdzw,modK,nz,nx,dt, ::FdtdAcousticVisco)
# 	for ix=1:nx  # see limits above
# 	@simd for iz=1:nz

# 		 # use the Auxiliary Differential Equation form,
# 		 # which is second-order accurate in time if implemented
# 		 #following eq (14) of Robertsson, Blanch and Symes, Geophysics, vol.
# 		 #	  59(9), pp 1444-1456 (1994), which is what we do here

# 		 # loop over standard linear solids
# 		 sumr=0.0
# 		 for isls = 1:nsls
# 			# this average of the two terms comes from eq (14) of
# 			#Robertsson, Blanch and Symes, Geophysics, vol. 59(9), pp 1444-1456 (1994)
# 			# central finite-difference around the time of vx (not the staggerred time grid!); therefore the factor 0.5
# 			r[isls,iz,ix] = r[isls,iz,ix] + ((dvxdx[iz,ix] + dvzdz[iz,ix]) * mod[iz,ix] * mod32[isls,iz,ix] - (mod31[isls,iz,ix] * r[isls,iz,ix])) * fc1 * 0.5
# 			sumr += r[isls,iz,ix]
# 		 end

# 		# this average of the two terms comes from eq (13) of
# 		#Robertsson, Blanch and Symes, Geophysics, vol. 59(9), pp 1444-1456 (1994)
# 		#pressure(i,j) = pressure(i,j) + (- kappa_half_x * (value_dvx_dx + value_dvy_dy) +
# 		#  0.5d0 * sum_of_memory_variables_kappa) * DELTAT
# 		# Here sumr should be at the time same as dvxdx
# 		@inbounds p[iz,ix] = pp[iz,ix] + ((mod[iz,ix] * (dvxdx[iz,ix] + dvzdz[iz,ix])) + (sumr)) * fc1
# 	end
# 	end
# end

# ====================================
# not used at the moment
# ====================================

# function dtsample(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 = 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, modrhovxI, modrhovzI, modK)
#     nz = size(p, 1)
#     nx = size(p, 2)
#     # boundary conditions
#     for i = 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

#     dx24I = 1.0
#     dz24I = 1.0
#     dt = 1.0

#     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)

#     modrhovxI = ones(nz, nx)
#     modrhovzI = ones(nz, nx)
#     modK = 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 = 1:ntimes

#         advance_kernel!(
#             pw,
#             dpdxw,
#             dpdzw,
#             dx24I,
#             dz24I,
#             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,
#             dt,
#             modrhovxI,
#             modrhovzI,
#             memory_dvx_dxw,
#             memory_dvz_dzw,
#             b_x,
#             b_z,
#             a_x,
#             a_z,
#             k_xI,
#             k_zI,
#             modK,
#         )
#     end

#     return pw
# end


text:
# Hydra primitives

named(arg) = isexpr(arg, :(::)) && length(arg.args) == 1 ? :($(gensym())::$(arg.args[1])) : arg

typeless(x) = MacroTools.prewalk(x -> isexpr(x, :(::)) ? x.args[1] : x, x)

macro spmd(ex)
  @capture(shortdef(ex), (name_(args__) = body_) |
                         (name_(args__) where {Ts__} = body_)) || error("Need a function definition")
  f, T = isexpr(name, :(::)) ?
    (length(name.args) == 1 ? (esc(gensym()), esc(name.args[1])) : esc.(name.args)) :
    (esc(gensym()), :(typeof($(esc(name)))))
  Ts = esc.(Ts == nothing ? [] : Ts)
  args = esc.(named.(args))
  argnames = typeless.(args)
  cx = :($(esc(:__mask__))::Mask)
  fargs = [cx, :($f::$T), args...]
  MacroTools.@q begin
    @inline Hydra.spmd($(fargs...)) where $(Ts...) = $(esc(body))
  end
end

# Hydra execution

mask(n) = vect(ntuple(_ -> true, n)...)

lane1(xs::AbstractVec) = xs[1]
lane1(x) = x

macro spmd(n, ex)
  :(lane1(spmd(mask($(esc(n))), () -> $(esc(ex)))))
end


text:
#
# public.jl --
#
# Implement public interface of ScientificCameras.
#
#-------------------------------------------------------------------------------
#
# This file is part of "AndorCameras.jl" released under the MIT license.
#
# Copyright (C) 2017-2021, Éric Thiébaut.
#

function open(::Type{Camera}, dev::Integer)
    handle = check(AT.Open(dev))
    cam = Camera()
    cam.state = 1
    cam.model = _UNKNOWN_MODEL
    cam.handle = handle
    if isimplemented(cam, CameraModel)
        model = cam[CameraModel]
        if model == "SIMCAM CMOS"
            cam.model = _SIM_CAM_MODEL
        elseif model == "ZYLA-4.2P-USB3"
            cam.model = _ZYLA_USB_MODEL
        end
    end
    return finalizer(_close, cam)
end

close(cam::Camera) = _close(cam; throwerrors=true)

function stop(cam::Camera)
    if cam.state != 2
        @warn "not acquiring"
        return nothing
    end
    check(_stop(cam))
    check(AT.Flush(cam))
    cam.state = 1
    return nothing
end

abort(cam::Camera) = stop(cam)

getfullwidth(cam::Camera) = Int(cam[SensorWidth])

getfullheight(cam::Camera) = Int(cam[SensorHeight])

#getroistride(cam::Camera) =
#    (isimplemented(cam, AOIStride) ? cam[AOIStride] :
#     div(getroiwidth(cam)*getbitsperpixel(cam) + 7, 8)) :: Int

getbitsperpixel(cam::Camera) = round(Int, 8*cam[BytesPerPixel])

function getbinning(cam::Camera) :: Tuple{Int,Int}
    if isimplemented(cam, AOIHBin) && isimplemented(cam, AOIVBin)
        return (cam[AOIHBin], cam[AOIVBin])
    elseif isimplemented(cam, AOIBinning)
        return parsebinning(repr(cam, AOIBinning))
    else
        return (1, 1)
    end
end

function parsebinning(str::AbstractString) :: Tuple{Int,Int}
    @noinline failure(str::AbstractString) =
        error("unknown binning ($str)")
    i = something(findfirst(isequal('x'), str), -1)
    if i > 0
        try
            return (parse(Int, str[1:i-1]), parse(Int, str[i+1:end]))
        catch
            nothing
        end
    end
    failure(str)
end

function getroi(cam::Camera)
    xsub, ysub = getbinning(cam)
    width = (isimplemented(cam, AOIWidth) ? cam[AOIWidth]
             : div(cam[SensorWidth], xsub))
    height = (isimplemented(cam, AOIHeight) ? cam[AOIHeight]
             : div(cam[SensorHeight], ysub))
    xoff = (isimplemented(cam, AOILeft) ? cam[AOILeft] - 1 : 0)
    yoff = (isimplemented(cam, AOITop)  ? cam[AOITop]  - 1 : 0)
    return ROI(xsub, ysub, xoff, yoff, width, height)
end

function setroi!(cam::Camera, roi::ROI)
    # Check parameters.
    checkstate(cam, 1, throwerrors=true)
    fullwidth = getfullwidth(cam)
    fullheight = getfullheight(cam)
    checkroi(roi, fullwidth, fullheight)

    # The ROI features must be configured in the order listed below as
    # features towards the top of the list will override the values below
    # them if the values are incompatible.
    #
    #  - `AOIHBin`, `AOIVBin` or `AOIBinning`
    #  - `AOIWidth` is in macro-pixel units
    #  - `AOILeft` is in pixel units
    #  - `AOIHeight` is in macro-pixel units
    #  - `VerticallyCenterAOI`
    #  - `AOITop` is in pixel units
    #
    if isimplemented(cam, AOIHBin) && isimplemented(cam, AOIVBin)
        if cam[AOIHBin] != roi.xsub
            cam[AOIHBin] = roi.xsub
        end
        if cam[AOIVBin] != roi.ysub
            cam[AOIVBin] = roi.ysub
        end
    elseif isimplemented(cam, AOIBinning)
        (xsub, ysub) = parsebinning(repr(cam, AOIBinning))
        if roi.xsub != xsub || roi.ysub != ysub
            str = @sprintf("%dx%d", roi.xsub, roi.ysub)
            cam[AOIBinning] = str
        end
    elseif roi.xsub != 1 || roi.ysub != 1
        error("pixel binning is not supported")
    end
    if isimplemented(cam, AOIWidth)
        if cam[AOIWidth] != roi.width
            cam[AOIWidth] = roi.width
        end
    elseif roi.xoff + roi.width*roi.xsub != fullwidth
        error("setting horizontal size of ROI is not supported")
    end
    if isimplemented(cam, AOILeft)
        left = roi.xoff + 1
        if cam[AOILeft] != left
            cam[AOILeft] = left
        end
    elseif roi.xoff != 0
        error("setting horizontal offset of ROI is not supported")
    end

    if isimplemented(cam, AOIHeight)
        if cam[AOIHeight] != roi.height
            cam[AOIHeight] = roi.height
        end
    elseif roi.yoff + roi.height*roi.ysub != fullheight
        error("setting vertical size of ROI is not supported")
    end
    if isimplemented(cam, VerticallyCenterAOI)
        cam[VerticallyCenterAOI] = false
    end
    if isimplemented(cam, AOITop)
        top = roi.yoff + 1
        if cam[AOITop] != top
            cam[AOITop] = top
        end
    elseif roi.yoff != 0
        error("setting vertical offset of ROI is not supported")
    end
    return nothing
end

# "Packed" encodings mean that the pixel values are packed on consecutive
# bytes.  Some exotic encodings only exists for the "SimCam" (simulated
# camera) model so their exact definition is not critical in practice.
const _PIXEL_ENCODING_TYPES = Dict{String,DataType}(
    "Mono8" => Monochrome{8},  # SimCam only
    "Mono12" => Monochrome{16}, # 16 bits but only 12 are significant
    "Mono16" => Monochrome{16},
    "Mono32" => Monochrome{32},
    "Mono12Packed" => Monochrome{12},
    "RGB8Packed" => RGB{8}, # SimCam only
    "Mono12Coded" => Monochrome{16}, # probably inexact but SimCam only
    "Mono12CodedPacked" => Monochrome{12}, # probably inexact but SimCam only
    "Mono22Parallel"  => Monochrome{32}, # probably inexact but SimCam only
    "Mono22PackedParallel" => Monochrome{22}) # probably inexact but SimCam only

const _PIXEL_ENCODING_NAMES = Dict{DataType,String}()
for (key, val) in _PIXEL_ENCODING_TYPES
    _PIXEL_ENCODING_NAMES[val] = key
end

# FIXME: As part of the initialisation of the camera, we could create fast
#        tables via the index number of the supproted formats.
function supportedpixelformats(cam::Camera)
    U = Union{}
    for i in 1:maximum(cam, PixelEncoding)
        if isimplemented(cam, PixelEncoding, i)
            str = repr(cam, PixelEncoding, i)
            U = Union{U, _PIXEL_ENCODING_TYPES[str]}
        end
    end
    return U
end

getpixelformat(cam::Camera) =
    get(_PIXEL_ENCODING_TYPES, repr(cam, PixelEncoding), Nothing)

function setpixelformat!(cam::Camera, ::Type{T}) where {T<:PixelFormat}
    checkstate(cam, 1, throwerrors=true)
    if T == getpixelformat(cam)
        return T
    elseif haskey(_PIXEL_ENCODING_NAMES, T)
        cam[PixelEncoding] = _PIXEL_ENCODING_NAMES[T]
        return T
    end
    error("unsupported pixel encoding")
end

getspeed(cam::Camera) =
    (cam[FrameRate], cam[ExposureTime])

function setspeed!(cam::Camera, fps::Float64, exp::Float64)
    # First reduce frame rate (if fps is smaller than actual value), then set
    # exposure time, then augment frame rate (if fps is larger than actual
    # value).  This strategy is to avoid havin incompatible settings at any
    # time.
    if cam[FrameRate] > fps
        cam[FrameRate] = fps
    end
    if cam[ExposureTime] != exp
        cam[ExposureTime] = exp
    end
    if cam[FrameRate] < fps
        cam[FrameRate] = fps
    end
    return nothing
end

getgain(cam::Camera) = 1.0

setgain!(cam::Camera, gain::Float64) =
    gain == 1.0 || error("invalid gain factor")

getbias(cam::Camera) = 0.0

setbias!(cam::Camera, bias::Float64) =
    bias == 0.0 || error("invalid bias level")

getgamma(cam::Camera) = 1.0

setgamma!(cam::Camera, gamma::Float64) =
    gamma == 1.0 || error("invalid gamma factor")

# Size of metadata fields.
const LENGTH_FIELD_SIZE = 4
const CID_FIELD_SIZE = 4
const TIMESTAMP_FIELD_SIZE = 8

# Size of metadata with a timestamp.
const METADATA_SIZE = (LENGTH_FIELD_SIZE + CID_FIELD_SIZE
                       + TIMESTAMP_FIELD_SIZE)

# Extend method.
function getcapturebitstype(cam::Camera)
    T = equivalentbitstype(getpixelformat(cam))
    return (T == Nothing ? UInt8 : T)
end

# Check timeout and convert it in milliseconds.
function timeout2ms(sec::Real)
    sec > 0 || throw(ArgumentError("invalid timeout"))
    msec = sec*1_000
    return (msec > typemax(AT.INFINITE) ? AT.INFINITE :
            round(typeof(AT.INFINITE), msec))
end

@noinline _warntimeout(cnt::Integer) =
    @warn "Acquisition timeout after $cnt image(s)"

# Extend method.
function read(cam::Camera, ::Type{T}, num::Int;
              nbufs::Integer = 4,
              skip::Integer = 0,
              timeout::Real = defaulttimeout(cam),
              truncate::Bool = false,
              ignoretimeouts::Bool = false,
              quiet::Bool = false) where {T}

    # Check timeout and convert it in milliseconds.
    ms = timeout2ms(timeout)

    # Allocate vector of images.
    imgs = Vector{Array{T,2}}(undef, num)

    # Start the acquisition.
    start(cam, T, nbufs)

    # Acquire all images.
    cnt = 0
    while cnt < num
        try
            _wait(cam, ms, skip > 0, quiet)
            if skip > zero(skip)
                skip -= one(skip)
            else
                cnt += 1
                imgs[cnt] = copy(cam.lastimg)
            end
        catch err
            solved = false
            if isa(err, TimeoutError)
                if ignoretimeouts
                    # Pretend the problem has been solved.
                    solved = true
                elseif truncate
                    # Truncate the image sequence.
                    num = cnt
                    resize!(imgs, num)
                    solved = true
                end
                if solved
                    # The error has been solved, print a warning unless quiet
                    # is true.
                    quiet || _warntimeout(cnt)
                end
            end
            if !solved
                # The error has not been solved, abort the acquisition and
                # rethrow the exception.
                abort(cam)
                rethrow(err)
            end
        end
    end

    # Stop the acquisition and return the sequence of images.
    abort(cam)
    return imgs
end

function read(cam::Camera, ::Type{T};
              skip::Integer = 0,
              nbufs::Integer = 1,
              timeout::Real = defaulttimeout(cam),
              ignoretimeouts::Bool = false,
              quiet::Bool = false) where {T}

    # Check timeout and convert it in milliseconds.
    ms = timeout2ms(timeout)

    # Start the acquisition.
    start(cam, T, nbufs)

    # Acquire a single image.
    cnt = 0
    while cnt < 1
        try
            _wait(cam, ms, skip > 0, quiet)
            if skip > zero(skip)
                skip -= one(skip)
            else
                cnt += 1
            end
        catch err
            if isa(err, TimeoutError) && ignoretimeouts
                # Print a warning unless quiet is true.
                quiet || _warntimeout(cnt)
            else
                # Abort the acquisition and rethrow the exception.
                abort(cam)
                rethrow(err)
            end
        end
    end

    # Stop the acquisition and return the image.
    abort(cam)
    return cam.lastimg
end

# Extend method.
function start(cam::Camera, ::Type{T}, nbufs::Int) where {T}

    # Check arguments.
    checkstate(cam, 1, throwerrors=true)
    nbufs ≥ 1 || throw(ArgumentError("invalid number of acquisition buffers"))

    # Make sure no buffers are currently in use.
    check(AT.Flush(cam))

    # Turn on metadata (this must be done before getting the frame size).
    if (isimplemented(cam, MetadataEnable) &&
        isimplemented(cam, MetadataTimestamp) &&
        isimplemented(cam, TimestampClockFrequency))
        cam[MetadataEnable] = true
        cam[MetadataTimestamp] = true
        cam.clockfrequency = cam[TimestampClockFrequency]
    else
        cam.clockfrequency = 0
    end

    # Get size of buffers.
    framesize = cam[ImageSizeBytes]

    # Create queue of frame buffers.
    if length(cam.bufs) != nbufs
        resize!(cam.bufs, nbufs)
    end
    for i in 1:nbufs
        if ! isassigned(cam.bufs, i) || sizeof(cam.bufs[i]) != framesize
            cam.bufs[i] = Vector{UInt8}(undef, framesize)
        end
        status = AT.QueueBuffer(cam, cam.bufs[i])
        if isfailure(status)
            AT.Flush(cam)
            error(status)
        end
    end

    # Allocate array to store last image.
    width = (isimplemented(cam, AOIWidth) ? cam[AOIWidth]
             : cam[SensorWidth])
    height = (isimplemented(cam, AOIHeight) ? cam[AOIHeight]
              : cam[SensorHeight])
    stride = div(framesize - (cam.clockfrequency > 0 ? METADATA_SIZE : 0),
                 height)
    cam.lastimg = Array{T,2}(undef, width, height)
    cam.mono12packed = (repr(cam, PixelEncoding) == "Mono12Packed") # FIXME:
    if isimplemented(cam, AOIStride)
        cam.bytesperline = cam[AOIStride]
        if cam.bytesperline != stride
            @warn("computed stride ($(stride) bytes) is not equal to "*
                  "AOIStride ($(cam.bytesperline) bytes)")
        end
    else
        cam.bytesperline = stride
    end

    # Set the camera to continuously acquires frames.
    cam[CycleMode] = "Continuous"
    cam[TriggerMode] = "Internal"

    # Start the acquisition.
    check(AT.Command(cam, AcquisitionStart))
    cam.state = 2
    return nothing
end

# Extend method.
function wait(cam::Camera, sec::Float64 = 0.0)
    ticks = _wait(cam, timeout2ms(sec), false)
    return cam.lastimg, ticks
end

# Extend method.
function release(cam::Camera)
    checkstate(cam, 2, throwerrors=true)
    return nothing
end

function _wait(cam::Camera, ms::Integer, skip::Bool=false, quiet::Bool=false)
    # Check arguments.
    checkstate(cam, 2, throwerrors=true)

    # Sleep in this thread until data is ready.  Note that the timeout argument
    # is pretended to be `Cint` because `AT.INFINITE` is `-1` whereas it is
    # `Cuint`.  This limits the maximum allowed timeout to about 24.9 days
    # which should be sufficient!
    status, bufptr, bufsiz = AT.WaitBuffer(cam, ms)
    if isfailure(status)
        if status.code != AT.ERR_TIMEDOUT
            error(status)
        end
        if cam.model == _ZYLA_USB_MODEL
            # Reset USB connection.
            dev = find_zyla()
            if dev == ""
                quiet || @warn "Timeout! USB connection not found..."
            else
                reset_usb(dev)
                AT.Command(cam, AcquisitionStop)  # FIXME:
                AT.Command(cam, AcquisitionStart) # FIXME:
                quiet || @warn "Timeout! USB connection has been reset..."
            end
        end
        throw(TimeoutError())
    end

    # Get the timestamp.
    local ticks::Float64 # to enforce type stability
    if cam.clockfrequency > 0
        tickscnt = unsafe_load(Ptr{UInt64}(bufptr + bufsiz - METADATA_SIZE))
        if ENDIAN_BOM == 0x01020304
            tickscnt = bswap(tickscnt)
        end
        ticks = tickscnt/cam.clockfrequency
    else
        ticks = time()
    end

    # Find buffer index.
    for buf in cam.bufs
        if pointer(buf) == bufptr
            # Extract buffer data and requeue buffer.
            if ! skip
                if cam.mono12packed
                    extractmono12packed!(cam.lastimg, buf, cam.bytesperline)
                else
                    extract!(cam.lastimg, buf, cam.bytesperline)
                end
            end
            status = AT.QueueBuffer(cam, bufptr, bufsiz)
            if isfailure(status)
                # Stop acquisition and report error.
                _stop(cam)
                AT.Flush(cam)
                cam.state = 1
                error(status)
            end
            return ticks
        end
    end
    error("bad buffer address")
end


text:
using Pkg

ssh = true
if ssh
    git = "git@github.com:SebastianAment/"
else
    git = "https://github.com/SebastianAment/"
end

add(x) = Pkg.add(Pkg.PackageSpec(url = git * x * ".git"))

add("LazyInverse.jl")
add("LinearAlgebraExtensions.jl")
add("KroneckerProducts.jl")
add("WoodburyIdentity.jl")
add("OptimizationAlgorithms.jl")
add("CovarianceFunctions.jl")

# not reachable without ssh
if ssh
    add("GaussianDistributions.jl")
end


text:
function test_data(d, l)
    x = simulate_random_snparray(undef, 1000, 1000)
    xla = SnpLinAlg{Float64}(x, model=ADDITIVE_MODEL, center=true, scale=true)
    z = ones(1000, 1)
    y = rand(1000)
    v = IHTVariable(xla, z, y, 1, 10, d, l, Int[], Float64[], :none, falses(1), false)
    MendelIHT.init_iht_indices!(v, false, trues(1000))
    return x, z, y, v
end

function make_IHTvar(d, μ, y)
    n = length(μ)
    v = IHTVariable(rand(n, 10), rand(n, 1), y, 1, 10, d, IdentityLink(),
        Int[], Float64[], :none, falses(1), false)
    MendelIHT.init_iht_indices!(v, false, trues(n))
    v.μ = μ
    return v
end

@testset "loglikelihood" begin
    Random.seed!(2019)
    d = Normal()
    μ = rand(1000000)
    y = [rand(Normal(μi, 1)) for μi in μ]
    v = make_IHTvar(d, μ, y) # create IHTVariable for testing

    # since we estimate σ with sum((y - E(y))^2)/n, there is small bias, so atol is large
    @test isapprox(MendelIHT.loglikelihood(v), sum(logpdf.(Normal.(μ, 1.0), y)), atol=1)

    d = Bernoulli()
    p = rand(10000)
    y = rand(0.0:1.0, 10000)
    v = make_IHTvar(d, p, y) # create IHTVariable for testing

    @test isapprox(MendelIHT.loglikelihood(v), sum(logpdf.(Bernoulli.(p), y)), atol=1e-8)

    d = Poisson()
    λ = rand(10000)
    y = Float64.(rand(1:100, 10000))
    v = make_IHTvar(d, λ, y) # create IHTVariable for testing

    @test isapprox(MendelIHT.loglikelihood(v), sum(logpdf.(Poisson.(λ), y)), atol=1e-8)

    d = NegativeBinomial()
    p = rand(1000000)
    r = ones(1000000)
    y = Float64.([rand(NegativeBinomial(1.0, p_i)) for p_i in p])
    v = make_IHTvar(d, p, y) # create IHTVariable for testing

    @test isapprox(MendelIHT.loglikelihood(v), sum(logpdf.(NegativeBinomial.(r, r ./ (p .+ r)), y)), atol=1e-5)
end

@testset "deviance" begin
    Random.seed!(2019)
    # Need more test other than normal distribution!!
    d = Normal()
    y = randn(1000)
    μ = randn(1000)
    v = make_IHTvar(d, μ, y) # create IHTVariable for testing
    @test sum(abs2, y .- μ) ≈ MendelIHT.deviance(v)
end

@testset "update_μ!" begin
    Random.seed!(2019)
    d = Normal
    μ = zeros(1000)
    xb = rand(1000)
    MendelIHT.update_μ!(μ, xb, canonicallink(d()))

    @test all(isapprox(μ, xb, atol=1e-8))

    d = Bernoulli
    p = zeros(1000)
    xb = rand(1000)
    MendelIHT.update_μ!(p, xb, canonicallink(d()))

    @test all(isapprox(p, 1 ./ (1 .+ exp.(-1.0 .* xb)), atol=1e-8))

    d = Poisson
    λ = zeros(1000)
    xb = rand(1000)
    MendelIHT.update_μ!(λ, xb, canonicallink(d()))

    @test all(isapprox(λ, (exp.(xb)), atol=1e-8))

    d = NegativeBinomial
    p = zeros(1000)
    xb = rand(1000)
    MendelIHT.update_μ!(p, xb, LogLink()) #use loglink for NegativeBinomial

    @test all(isapprox(p, exp.(xb), atol=1e-8))
end

@testset "_choose!" begin
    J = 1
    k = 10

    Random.seed!(1111)
    x, z, y, v = test_data(Normal(), IdentityLink())
    fill!(v.idx, false)
    fill!(v.idc, false)
    v.idx[1:10] .= trues(10)
    MendelIHT._choose!(v)

    @test all(v.idx[1:10] .== trues(10))
    @test v.idc[1] == false

    v.idx[1:11] .= trues(11)
    MendelIHT._choose!(v)

    @test sum(v.idx[1:11]) == 10

    # set k = 1
    x, z, y, v = test_data(Normal(), IdentityLink())
    fill!(v.idx, false)
    v.idc .= true
    v.k = 1
    MendelIHT._choose!(v)

    @test v.idc[1] == true
    @test sum(v.idc) == 1
    @test sum(v.idx) == 0

    x, z, y, v = test_data(Normal(), IdentityLink())
    fill!(v.idx, false)
    v.idc .= true
    v.idx[1:10] .= trues(10)
    MendelIHT._choose!(v)

    @test sum(v.idx) + sum(v.idc) == 10
end

@testset "_iht_backtrack" begin
    @test MendelIHT._iht_backtrack_(-150.0, -100.0, 1, 10) == true
    @test MendelIHT._iht_backtrack_(-50.0, -100.0, 1, 10) == false
    @test MendelIHT._iht_backtrack_(150.0, 100.0, 1, 10) == false
    @test MendelIHT._iht_backtrack_(50.0, 100.0, 1, 10) == true
    @test MendelIHT._iht_backtrack_(50.0, 100.0, 0, 3) == true
    @test MendelIHT._iht_backtrack_(50.0, 100.0, 11, 10) == false
    @test MendelIHT._iht_backtrack_(50.0, 100.0, 11, 11) == false
end

@testset "standardize!" begin
	Random.seed!(1017)
    z = rand(1000, 1000)
    standardize!(z)
    @test isapprox(0.0, mean(z[:, 1]), atol=1e-8)
    @test isapprox(0.0, mean(z[:, 10]), atol=1e-8)
    @test isapprox(0.0, mean(z[:, 100]), atol=1e-8)
    @test isapprox(1.0, std(z[:, 1]), atol=1e-8)
    @test isapprox(1.0, std(z[:, 10]), atol=1e-8)
    @test isapprox(1.0, std(z[:, 100]), atol=1e-8)

    z = rand(1000, 1000)
    col1_μ = mean(z[:, 1])
    col1_σ = std(z[:, 1])
    standardize!(@view(z[:, 2:end]))
    @test isapprox(col1_μ, mean(z[:, 1]), atol=1e-8)
    @test isapprox(0.0, mean(z[:, 10]), atol=1e-8)
    @test isapprox(0.0, mean(z[:, 100]), atol=1e-8)
    @test isapprox(col1_σ, std(z[:, 1]), atol=1e-8)
    @test isapprox(1.0, std(z[:, 10]), atol=1e-8)
    @test isapprox(1.0, std(z[:, 100]), atol=1e-8)
end

@testset "project_k!" begin
    x = rand(100000)
    k = 100
    p = sortperm(x, rev = true)
    top_k_index = p[1:k]
    last_k_index = p[k+1:end]
    project_k!(x, k)

    @test all(x[top_k_index] .!= 0.0)
    @test all(x[last_k_index] .== 0.0)
end

# Note: since depending on RNG the groups and x,y vectors will be different,
#       we do not test specific values, but only test properties of result.
@testset "project_group_sparse!" begin
    #2 active groups, 3 active predictors per group, 10 total predictors
    m, n, k = 2, 3, 10
    y = randn(k);
    group = rand(1:5, k);
    x = copy(y)
    project_group_sparse!(x, group, m, n)

    non_zero_position = findall(!iszero, x)
    non_zero_entries = x[non_zero_position]
    @test all(non_zero_entries .== y[non_zero_position])
    @test length(non_zero_position) ≤ 6

    # 2 active groups, 50% of active predictor per group, variable group size, 9 total predictors
    J, n = 2, 15
    k = [1, 1, 2, 2, 3]
    y = 5rand(n)
    y_copy = copy(y)
    group = [1; 2; 2; 3; 3; 3; 4; 4; 4; 4; 5; 5; 5; 5; 5]
    project_group_sparse!(y, group, J, k)
    non_zero_position = findall(!iszero, y)
    non_zero_entries = y[non_zero_position]
    @test all(non_zero_entries .== y_copy[non_zero_position])
    @test length(non_zero_position) ≤ 2 + 3

    # test project_group_sparse! is equivalent to project_k! when max group = 1
    m, n, k = 1, 10, 100000
    group = ones(Int, k) #everybody is in the same group
    y = randn(k)
    x = copy(y)
    project_k!(x, n)
    project_group_sparse!(y, group, m, n)
    @test all(x .== y)
end

@testset "maf_weights" begin
    Random.seed!(33)
    x = simulate_random_snparray(undef, 1000, 10000)
    m = SnpArrays.maf(x)
    p = maf_weights(x)

    @test eltype(p) <: AbstractFloat
    @test all(p .>= 1.0)
    @test p[1] ≈ 1 / (2.0sqrt(m[1] * (1 - m[1])))
    @test p[2] ≈ 1 / (2.0sqrt(m[2] * (1 - m[2])))

    p = maf_weights(x, max_weight=2.0)
    @test all(1.0 .<= p .<= 2)
    @test p[1] ≈ 1 / (2.0sqrt(m[1] * (1 - m[1])))
end