text:
using PermutedArrays
using Documenter

DocMeta.setdocmeta!(PermutedArrays, :DocTestSetup, :(using PermutedArrays); recursive=true)

makedocs(;
    modules=[PermutedArrays],
    authors="Chad Scherrer <chad.scherrer@gmail.com> and contributors",
    repo="https://github.com/cscherrer/PermutedArrays.jl/blob/{commit}{path}#{line}",
    sitename="PermutedArrays.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://cscherrer.github.io/PermutedArrays.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/cscherrer/PermutedArrays.jl",
)

max_stars_repo_name: cscherrer/PermutedArrays

max_issues_repo_name: cscherrer/PermutedArrays

max_forks_repo_name: cscherrer/PermutedArrays


text:
struct S a end
function f(s::S) s.a = 1 end

ext: jl

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia


text:
@testset "454.4-sum-ii.jl" begin
    @test four_sum_count([1, 2], [-2, -1], [-1, 2], [0, 2]) == 2
end


text:
# Subsampling space can definitely be necessary
# Subsampling time may not be, maybe rather use DESolution methods.

# Supposing adaptive time, specify "saveat" -- simple.
# Optional, allow specifying "tstops" which forces non-interpolation.

import Base: to_indices, _maybetail, @_inline_meta, tail, getindex

abstract type AbstractSubsampler{D} end
function coordinate_indices(lattice::LATTICE,
        subsampler_arr::AbstractArray{<:AbstractSubsampler}) where {LATTICE<:AbstractLattice}
     idxs_arr = map(subsampler_arr) do subsampler
        coordinate_indices(lattice, subsampler)
    end
    all_keep = ones(Bool, size(lattice)...)
    for idxs in idxs_arr
        keep = zeros(Bool, size(lattice)...)
        keep[idxs] .= true
        all_keep .&= keep
    end
    all_indices = CartesianIndices(lattice)
    return all_indices[all_keep]
end
function subsample(lattice::LATTICE, sub) where {LATTICE<:AbstractLattice}
    idxs = coordinate_indices(lattice, sub)
    LATTICE(lattice.arr[idxs])
end
subsample(space::AbstractLattice, ::Nothing) = space
function getindex(lattice::LATTICE, sub) where {T,D,LATTICE<:AbstractLattice{T,D}}
    idxs = coordinate_indices(lattice, sub)
    LATTICE(lattice.arr[idxs])
end

@with_kw struct IndexSubsampler{D} <: AbstractSubsampler{D}
	strides::NTuple{D,Int}
end
@with_kw struct RightCutFromValue{D,T} <: AbstractSubsampler{D}
	cut::NTuple{D,T}
end
export RightCutFromValue

@with_kw struct RightCutProportionFromValue{D,T} <: AbstractSubsampler{D}
    cut::NTuple{D,T}
    proportion::NTuple{D,T}
end
export RightCutProportionFromValue

"IndexInfo specifies the index of 0 and the stride (`Δ`) between indices."
@with_kw struct IndexInfo{D,N}
	Δ::NTuple{N,D}
	origin_idx::CartesianIndex{N}
end

"""
	StrideToEnd(stride, start=1)

StrideToEnd is a custom index type that acts like `start:stride:end`, to circumvent the fact that you can't put `start:stride:end` into a variable.
"""
struct StrideToEnd{STOP}
    stride::Int
	start::Int
    stop::STOP
    StrideToEnd{STOP}(stride::Int, start::Int=1, stop::STOP=nothing) where {STOP<:Union{Nothing,Int}} = new(stride, start,stop)
end
StrideToEnd(stride, start, stop::S=nothing) where S = StrideToEnd{S}(stride, start, stop)
to_indices(A, inds, I::Tuple{StrideToEnd{Int}, Vararg{Any}})=
(@_inline_meta; (I[1].start:I[1].stride:I[1].stop, to_indices(A, _maybetail(inds), tail(I))...))
to_indices(A, inds, I::Tuple{StrideToEnd{Nothing}, Vararg{Any}})=
	(@_inline_meta; (I[1].start:I[1].stride:inds[1][end], to_indices(A, _maybetail(inds), tail(I))...))
to_indices(A, inds, I::Tuple{NTuple{N,StrideToEnd}, Vararg{Any}}) where N = to_indices(A, inds, (I[1]..., _maybetail(I)...))
getindex(A::AbstractArray, S::StrideToEnd) = getindex(A, to_indices(A, (S,))...)

function coordinate_indices(lattice::AbstractLattice{T}, subsampler::RightCutFromValue{D,T}) where {T,D}
    axis_dxs = map(zip(subsampler.cut, coordinate_axes(lattice))) do (value, axis)
        findfirst(axis .>= value)
    end
    stride = 1
    cuts = Tuple(axis_dxs)
    return StrideToEnd.(stride, cuts)
end
function coordinate_indices(lattice::AbstractLattice{T}, subsampler::RightCutProportionFromValue{D,T}) where {T,D}
    start_stop = map(zip(subsampler.cut, subsampler.proportion, coordinate_axes(lattice))) do (value, proportion, axis)
        axis_begin = findfirst(axis .>= value)
        remaining_span = length(axis) - axis_begin
        axis_end = floor(Int, axis_begin + remaining_span * proportion)
        (axis_begin, axis_end)
    end
    axis_starts, axis_stops = zip(start_stop...)
    return StrideToEnd.(1, axis_starts, axis_stops)
end

function coordinate_indices(::Any, subsampler::IndexSubsampler{D}) where D
    return StrideToEnd.(subsampler.strides, 1)
end

function coordinate_indices(lattice, sub_dxs::AbstractArray{<:CartesianIndex})
    #@warn "subsampling one pop"
    #FIXME don't subsample one pop
    idxs = population(sub_dxs, 1)
end


text:
include("viral_load_infectivity_testpos.jl")

const scen_names = ["(b) Status Quo","(c1) Fortnightly concurrent PCR","(c2) Fortnightly random PCR", "(d) 3 LFDs per week","(e) 2 LFDs per week","(f) Daily LFDs","(g) Daily LFDs + PCR","(h) 3 LFDs + PCR",
    "(a) No testing"]

const shift_pattern = [true, true, true, true, false, false, false,
                       true, true, true, true, true, false, false]
const twolfd_pattern = [true, false, true, false, false, false, false,
                           true, false, false, true, false, false, false]
const threelfd_pattern = [true, false, true, true, false, false, false,
                           true, false, true, false, true, false, false]

function scenario_1_setup(Ndays::Int, shift_pattern_start::Int)   #2 LFDs + 1 concurrent PCR
    #PCRs happen from first Monday
    #infection day is DayOfWeek[shift_pattern_start]
    #Next Monday is 1 + mod(15 - shift_pattern_start,7)
    #i.e. if shift_pattern_start = 1, infection day is a Monday, so PCR starts on day 1
    #if shift_pattern_start = 2, infection day is a Tuesday, so PCR starts on day 7 etc.
    test_day0 = 1 + mod(15 - shift_pattern_start,7)

    Nrepeats = Int64(ceil((Ndays - 13 + shift_pattern_start)/14))
    LFD_test_days_bool = vcat(twolfd_pattern[shift_pattern_start:end], repeat(twolfd_pattern,Nrepeats))
    LFD_test_days_bool = LFD_test_days_bool[1:Ndays]

    LFD_test_days = collect(1:length(LFD_test_days_bool))[LFD_test_days_bool]
    PCR_test_days = collect(test_day0:7:Ndays)
    PCR_delays = draw_PCR_delays(length(PCR_test_days))
    PCR_result_days = PCR_test_days .+ PCR_delays
    TestDays = vcat(LFD_test_days, PCR_test_days)
    ResultDays = vcat(LFD_test_days, PCR_result_days)
    TestTypes = vcat(ones(length(LFD_test_days)), zeros(length(PCR_test_days)))
    itd = sortperm(ResultDays)
    return TestDays[itd], TestTypes[itd], ResultDays[itd]
end

function scenario_2_setup(Ndays::Int, shift_pattern_start::Int)   #2 LFDs + 1 concurrent fortnightly PCR
    test_day0 = 1 + mod(15 - shift_pattern_start,14)

    Nrepeats = Int64(ceil((Ndays - 13 + shift_pattern_start)/14))
    LFD_test_days_bool = vcat(twolfd_pattern[shift_pattern_start:end], repeat(twolfd_pattern,Nrepeats))
    LFD_test_days_bool = LFD_test_days_bool[1:Ndays]

    LFD_test_days = collect(1:length(LFD_test_days_bool))[LFD_test_days_bool]
    PCR_test_days = collect(test_day0:14:Ndays)
    PCR_delays = draw_PCR_delays(length(PCR_test_days))
    PCR_result_days = PCR_test_days .+ PCR_delays
    TestDays = vcat(LFD_test_days, PCR_test_days)
    ResultDays = vcat(LFD_test_days, PCR_result_days)
    TestTypes = vcat(ones(length(LFD_test_days)), zeros(length(PCR_test_days)))
    itd = sortperm(ResultDays)

    return TestDays[itd], TestTypes[itd], ResultDays[itd]
end

function scenario_3_setup(Ndays::Int, shift_pattern_start::Int)  #2 LFDs + 1 random fortnightly PCR
    Nrepeats = Int64(ceil((Ndays - 13 + shift_pattern_start)/14))
    shifts = vcat(shift_pattern[shift_pattern_start:end], repeat(shift_pattern,Nrepeats))
    shift_days = collect(1:length(shifts))[shifts]
    test_day0 = rand(shift_days[shift_days .< 15]) #random fortnightly PCR on work day

    LFD_test_days_bool = vcat(twolfd_pattern[shift_pattern_start:end], repeat(twolfd_pattern,Nrepeats))
    LFD_test_days_bool = LFD_test_days_bool[1:Ndays]

    LFD_test_days = collect(1:length(LFD_test_days_bool))[LFD_test_days_bool]
    PCR_test_days = collect(test_day0:14:Ndays)
    PCR_delays = draw_PCR_delays(length(PCR_test_days))
    PCR_result_days = PCR_test_days .+ PCR_delays
    TestDays = vcat(LFD_test_days, PCR_test_days)
    ResultDays = vcat(LFD_test_days, PCR_result_days)
    TestTypes = vcat(ones(length(LFD_test_days)), zeros(length(PCR_test_days)))
    itd = sortperm(ResultDays)

    return TestDays[itd], TestTypes[itd], ResultDays[itd]
end

function scenario_4_setup(Ndays::Int, shift_pattern_start::Int)  #3 LFDs
    Nrepeats = Int64(ceil((Ndays - 13 + shift_pattern_start)/14))
    LFD_test_days_bool = vcat(threelfd_pattern[shift_pattern_start:end], repeat(threelfd_pattern,Nrepeats))
    LFD_test_days_bool = LFD_test_days_bool[1:Ndays]

    LFD_test_days = collect(1:length(LFD_test_days_bool))[LFD_test_days_bool]
    TestDays = copy(LFD_test_days)
    ResultDays = copy(LFD_test_days)
    TestTypes = ones(length(LFD_test_days))
    itd = sortperm(ResultDays)

    return TestDays[itd], TestTypes[itd], ResultDays[itd]
end

function scenario_5_setup(Ndays::Int, shift_pattern_start::Int)  #2 LFDs
    Nrepeats = Int64(ceil((Ndays - 13 + shift_pattern_start)/14))
    LFD_test_days_bool = vcat(twolfd_pattern[shift_pattern_start:end], repeat(twolfd_pattern,Nrepeats))
    LFD_test_days_bool = LFD_test_days_bool[1:Ndays]

    LFD_test_days = collect(1:length(LFD_test_days_bool))[LFD_test_days_bool]
    TestDays = copy(LFD_test_days)
    TestTypes = ones(length(LFD_test_days))

    return TestDays, TestTypes, TestDays
end

function scenario_6_setup(Ndays::Int, shift_pattern_start::Int)  #Daily LFDs
    Nrepeats = Int64(ceil((Ndays - 13 + shift_pattern_start)/14))
    LFD_test_days_bool = vcat(shift_pattern[shift_pattern_start:end], repeat(shift_pattern,Nrepeats))
    LFD_test_days_bool = LFD_test_days_bool[1:Ndays]

    TestDays = collect(1:Ndays)[LFD_test_days_bool]
    TestTypes = ones(length(TestDays))
    return TestDays, TestTypes, TestDays #result days same as test days
end

function scenario_7_setup(Ndays::Int, shift_pattern_start::Int)   #Daily LFDs + 1 concurrent PCR
    test_day0 = 1 + mod(15 - shift_pattern_start,7)

    Nrepeats = Int64(ceil((Ndays - 13 + shift_pattern_start)/14))
    LFD_test_days_bool = vcat(shift_pattern[shift_pattern_start:end], repeat(shift_pattern,Nrepeats))
    LFD_test_days_bool = LFD_test_days_bool[1:Ndays]
    LFD_test_days = collect(1:Ndays)[LFD_test_days_bool]

    PCR_test_days = collect(test_day0:7:Ndays)
    PCR_delays = draw_PCR_delays(length(PCR_test_days))
    PCR_result_days = PCR_test_days .+ PCR_delays
    TestDays = vcat(LFD_test_days, PCR_test_days)
    ResultDays = vcat(LFD_test_days, PCR_result_days)
    TestTypes = vcat(ones(length(LFD_test_days)), zeros(length(PCR_test_days)))
    itd = sortperm(ResultDays)

    return TestDays[itd], TestTypes[itd], ResultDays[itd]
end

function scenario_8_setup(Ndays::Int, shift_pattern_start::Int)   #3 LFDs + 1 concurrent PCR
    test_day0 = 1 + mod(15 - shift_pattern_start,7)

    Nrepeats = Int64(ceil((Ndays - 13 + shift_pattern_start)/14))
    LFD_test_days_bool = vcat(threelfd_pattern[shift_pattern_start:end], repeat(threelfd_pattern,Nrepeats))
    LFD_test_days_bool = LFD_test_days_bool[1:Ndays]

    LFD_test_days = collect(1:length(LFD_test_days_bool))[LFD_test_days_bool]
    PCR_test_days = collect(test_day0:7:Ndays)
    PCR_delays = draw_PCR_delays(length(PCR_test_days))
    PCR_result_days = PCR_test_days .+ PCR_delays
    TestDays = vcat(LFD_test_days, PCR_test_days)
    ResultDays = vcat(LFD_test_days, PCR_result_days)
    TestTypes = vcat(ones(length(LFD_test_days)), zeros(length(PCR_test_days)))
    itd = sortperm(ResultDays)
    return TestDays[itd], TestTypes[itd], ResultDays[itd]
end

"""
    init_testing_random!(sim::Dict, testing_params::Dict, i_day::Int, Ndays::Int)

Function to initialise test positive profiles and test isolation probabilities

## Arguments:
`sim` = Dict generated by `init_VL_and_infectiousness(Ntot::Int, Pisol::Float64)`

`testing_params` = Dict containing testing options, must have:

    "scenario" => `String` = Choose from '2_LFD_+_concPCR', '3_LFD', '2_LFD_+_Fort_concPCR'
                                          '2_LFD_+_Fort_randPCR', '2_LFD'

    "comply_prob" => `Float64` = Chance of performing next LFD test (random)

`i_day` = Day to start simulation on (counting from 1)

`Ndays` = Total number of days to simulate (counting from 1)

## Returns:
`Array{Int64,1}` = Days when testing will occur

`Int64` = Next testing day index (from `i_day`)

## See also:
`init_VL_and_infectiousness(Ntot::Int, Pisol::Float64)`
"""
function init_testing_random!(sim::Dict, testing_params::Dict, Conf_PCR::Bool)
    #add test positivity profiles to simulation Dict, i_day is start day, Ndays is total length of sim
    #testing params contains "tperiod" (days between test)
    #testing params contains "protocol" (LFD_mass_protocol or PCR_mass_protocol)
    #optional: sens_rel: relative sensitivity of LFD as VL -> infinity
    Ndays = length.(sim["VL_profiles"])
    sim["will_isolate_with_test"] = ones(Bool,sim["Ntot"])
    sim["testing_paused"] = zeros(Bool,sim["Ntot"])
    sim["resume_testing"] = -ones(Int64,sim["Ntot"])
    sim["test_days"] = Array{Array{Int64,1},1}(undef,sim["Ntot"])
    sim["test_types"] = Array{Array{Int64,1},1}(undef,sim["Ntot"])
    sim["test_result_days"] = Array{Array{Int64,1},1}(undef,sim["Ntot"])
    sim["conf_PCR"] = Array{Array{Bool,1},1}(undef,sim["Ntot"])
    TestOutput = []
    if testing_params["scenario"] == scen_names[1]
        TestOutput = scenario_1_setup.(Ndays, sim["shift_pattern_start_day"])
    elseif testing_params["scenario"] == scen_names[2]
        TestOutput = scenario_2_setup.(Ndays, sim["shift_pattern_start_day"])
    elseif testing_params["scenario"] == scen_names[3]
        TestOutput = scenario_3_setup.(Ndays, sim["shift_pattern_start_day"])
    elseif testing_params["scenario"] == scen_names[4]
        TestOutput = scenario_4_setup.(Ndays, sim["shift_pattern_start_day"])
    elseif testing_params["scenario"] == scen_names[5]
        TestOutput = scenario_5_setup.(Ndays, sim["shift_pattern_start_day"])
    elseif testing_params["scenario"] == scen_names[6]
        TestOutput = scenario_6_setup.(Ndays, sim["shift_pattern_start_day"])
    elseif testing_params["scenario"] == scen_names[7]
        TestOutput = scenario_7_setup.(Ndays, sim["shift_pattern_start_day"])
    elseif testing_params["scenario"] == scen_names[8]
        TestOutput = scenario_8_setup.(Ndays, sim["shift_pattern_start_day"])
    end
    for i in 1:sim["Ntot"]
        sim["test_days"][i] = TestOutput[i][1]
        sim["test_types"][i] = TestOutput[i][2]
        sim["test_result_days"][i] = TestOutput[i][3]
        sim["conf_PCR"][i] = zeros(Bool,length(sim["test_days"][i]))
        if Conf_PCR
            sim["conf_PCR"][i][sim["test_types"][i] .== 1] .= true
        end
    end

    sim["test_pos_prob"] = get_test_probs.(sim["VL_profiles"], sim["test_days"],
        sim["test_types"]; Do_LFD_test_prob = testing_params["comply_prob"],
        all_or_nothing = testing_params["aon_compliance"])
    #assuming random compliance
end

function get_no_testing_scenario(Ntot::Int, Pisol::Float64;
        Day5release_bool::Bool = false)
    sim_baseline = init_VL_and_infectiousness(Ntot, Pisol)
    sim_baseline["shift_pattern_start_day"] = rand(1:length(shift_pattern),Ntot)
    sim_baseline["test_pos_prob"] = Array{Array{Float64,1},1}(undef,Ntot)
    sim_baseline["test_result_days"]  = Array{Array{Int64,1},1}(undef,Ntot)
    Conf_PCR_h = fill(zeros(Bool,0),sim_baseline["Ntot"])
    sim_baseline["isol_days"] = Array{Array{Int64,1},1}(undef,Ntot)
    sim_baseline["tot_isol_days"] = zeros(Int64,Ntot)
    for j in 1:Ntot
        sim_baseline["test_pos_prob"][j] = zeros(0)
        sim_baseline["test_result_days"][j] = zeros(Int64,0)
        sim_baseline["isol_days"][j] = zeros(Int64,0)
        sim_baseline["tot_isol_days"][j] = run_testing_scenario!(sim_baseline["isol_days"][j],
            sim_baseline["infection_profiles"][j], sim_baseline["test_pos_prob"][j],
            sim_baseline["test_result_days"][j], sim_baseline["symp_day"][j] .+ 1,
            sim_baseline["will_isolate"][j], sim_baseline["VL_profiles"][j],
            Conf_PCR_h[j]; Day5release=Day5release_bool)
    end

    return sim_baseline
end

function run_testing_scenarios_vs_baseline(sim_baseline::Dict, LFD_comply::Float64,
                                      Conf_PCR::Bool; Day5release_bool::Bool = false,
                                      LFD_AllorNone::Bool = false)
    Ntot = sim_baseline["Ntot"]
    Nscens = length(scen_names)
    sim_scens = Array{Dict,1}(undef,Nscens)
    for i in 1:(Nscens-1)
        sim_scens[i] = copy(sim_baseline)
        init_testing_random!(sim_scens[i], Dict("scenario"=>scen_names[i],
                 "comply_prob"=>LFD_comply, "aon_compliance"=>LFD_AllorNone), Conf_PCR)

        sim_scens[i]["isol_days"] = copy.(sim_baseline["isol_days"])
        sim_scens[i]["tot_isol_days"] = run_testing_scenario!.(sim_scens[i]["isol_days"],
            sim_scens[i]["infection_profiles"],  sim_scens[i]["test_pos_prob"], sim_scens[i]["test_result_days"],
            sim_scens[i]["symp_day"] .+ 1,  sim_scens[i]["will_isolate"], sim_scens[i]["VL_profiles"],
            sim_scens[i]["conf_PCR"]; Day5release=Day5release_bool)
    end
    sim_scens[Nscens] = copy(sim_baseline)
    for i in 1:Nscens
        sim_scens[i]["inf_days"] = zeros(Int64, Ntot)
        sim_scens[i]["tot_work_isol_days"] = zeros(Int64, Ntot)
        sim_scens[i]["inf_profile_isolation"] = copy(sim_scens[i]["infection_profiles"])
        Ndays = length.(sim_scens[i]["VL_profiles"])
        Nrepeats = Int64.(ceil.((Ndays .- 13 .+ sim_scens[i]["shift_pattern_start_day"]) ./ 14))
#         print("Scenario ",i,'\n')
        for j in 1:Ntot
            if length(sim_scens[i]["isol_days"][j]) > 0
                maxisolday = max(sim_scens[i]["isol_days"][j]...)
                if maxisolday > Ndays[j]
                    Nrepeats[j] = Int64(ceil((maxisolday - 13 + sim_scens[i]["shift_pattern_start_day"][j]) / 14))
                end
            end
            shifts = vcat(shift_pattern[sim_scens[i]["shift_pattern_start_day"][j]:end],
                          repeat(shift_pattern,Nrepeats[j]))
#             print("Shifts: ", shifts, '\n')
            inf_shifts = shifts[1:Ndays[j]]
#             print("Inf Shifts: ", inf_shifts, '\n')
            #count people as not infectious if not at work
#             print("Infectious profile raw: ", sim_scens[i]["infection_profiles"][j], '\n')
            sim_scens[i]["infection_profiles"][j][inf_shifts .== false] .= 0.0
#             print("Infectious profile shifts removed: ", sim_scens[i]["infection_profiles"][j], '\n')
            #extract isolation days that occur in infectious window
#             print("Isolation days: ", sim_scens[i]["isol_days"][j], '\n')
            inf_isol_days = sim_scens[i]["isol_days"][j][sim_scens[i]["isol_days"][j] .<=
                                            length(sim_scens[i]["infection_profiles"][j])]
#             print("Infectious Isolation days: ", inf_isol_days, '\n')
            #copy over infectious profile with shifts removed
            sim_scens[i]["inf_profile_isolation"][j] = copy(sim_scens[i]["infection_profiles"][j])
            #set isolation days to zero in modified inf profile
            sim_scens[i]["inf_profile_isolation"][j][inf_isol_days] .= 0.0
#             print("Infectious profile isolation removed: ", sim_scens[i]["inf_profile_isolation"][j], '\n')
            sim_scens[i]["inf_days"][j] = sum((sim_scens[i]["inf_profile_isolation"][j] .> 0))
#             print("Infectious days worked: ", sim_scens[i]["inf_days"][j], '\n')
            #number of isol days in work
            sim_scens[i]["tot_work_isol_days"][j] = sum(shifts[sim_scens[i]["isol_days"][j]])
#             print("Work days missed due to isolation: ", sim_scens[i]["tot_work_isol_days"][j], '\n')
#             print("\n\n")
        end
    end

    return sim_scens, scen_names
end


text:
# implementation of derivative and gradient kernel
########################### Gradient kernel ####################################
# f ∼ GP(μ, k)
# ∂f ∼ GP(∂μ, dkxy) # gradient kernel
struct GradientKernel{T, K, IT<:InputTrait} <: MultiKernel{T}
    k::K
    input_trait::IT
    function GradientKernel{T, K}(k, it = input_trait(k)) where {T, K}
        new{T, K, typeof(it)}(k, it)
    end
end
GradientKernel(k::AbstractKernel{T}) where T = GradientKernel{T, typeof(k)}(k)
GradientKernel(k) = GradientKernel{Float64, typeof(k)}(k)
@inline input_trait(G::GradientKernel) = G.input_trait
# gramian_eltype(G::GradientKernel) = eltype(G.k)
Base.eltype(G::GradientKernel) = eltype(G.k)

function elsize(G::Gramian{<:AbstractMatrix, <:GradientKernel}, i::Int)
    i ≤ 2 ? length(G.x[1]) : 1
end

(G::GradientKernel)(x, y) = gradient_kernel(G.k, x, y, G.input_trait)

# O(d²) fallback for non-structured case
function gradient_kernel(k, x, y, ::GenericInput)
    K = allocate_gradient_kernel(k, x, y, GenericInput())
    gradient_kernel!(K, k, x, y, GenericInput())
end
function allocate_gradient_kernel(k, x, y, ::GenericInput)
    zeros(gramian_eltype(k, x, y), length(x), length(y))
end

function gradient_kernel!(K::AbstractMatrix, k, x, y, ::GenericInput)
    value = zeros(eltype(K), length(x)) # necessary?
    derivs = K # jacobian (higher order terms)
    result = DiffResults.DiffResult(value, K)
    g(y) = ForwardDiff.gradient(z->k(z, y), x)
    ForwardDiff.jacobian!(result, g, y)
    return K
end

# specialized lazy representation of element of gradient kernel gramian
# increases performance dramatically, compared to generic Woodbury implementation
struct GradientKernelElement{T, K, X, Y, IT} <: Factorization{T}
    k::K
    x::X
    y::Y
    input_trait::IT
end

function GradientKernelElement(k, x, y, IT::InputTrait = input_trait(k))
    T = gramian_eltype(k, x, y)
    GradientKernelElement{T, typeof(k), typeof(x), typeof(y), typeof(IT)}(k, x, y, IT)
end
Base.size(K::GradientKernelElement) = (length(K.x), length(K.y))
Base.size(K::GradientKernelElement, i::Int) = 1 ≤ i ≤ 2 ? size(K)[i] : 1
Base.eltype(K::GradientKernelElement{T}) where T = T
# gradient kernel element only used for sparsely representable elements
Base.Matrix(K::GradientKernelElement) = K * I(size(K, 1))

function Base.:*(G::GradientKernelElement, a)
    T = promote_type(eltype(G), eltype(a))
    b  = zeros(T, size(a))
    mul!(b, G, a)
end

const GenericGradientKernelElement{T, K, X, Y} = GradientKernelElement{T, K, X, Y, <:GenericInput}

const IsotropicGradientKernelElement{T, K, X, Y} = GradientKernelElement{T, K, X, Y, IsotropicInput}
function IsotropicGradientKernelElement{T}(k, x, y) where T
    IsotropicGradientKernelElement{T, typeof(k), typeof(x), typeof(y)}(k, x, y, IsotropicInput())
end

# isotropic kernel
function LinearAlgebra.mul!(b, G::IsotropicGradientKernelElement, a, α::Number = 1, β::Number = 0) #, ::IsotropicInput = G.input_trait)
    r = difference(G.x, G.y) # other implementation appears to be faster for d = 128 since r does not have to be computed twice?
    r² = sum(abs2, r)
    k1, k2 = derivative_laplacian(G.k, r²) # IDEA: these could be computed once and stored in G
    dot_r_a = r'a # dot(r, a)
    @. b = α * -2(k1 * a + 2*k2 * r * dot_r_a) + β * b
end

# sparse but not completely lazy representation
function WoodburyFactorizations.Woodbury(K::IsotropicGradientKernelElement)
    k, x, y = K.k, K.x, K.y
    r = x - y
    r = reshape(r, :, 1) # do these work without reshape?
    r² = sum(abs2, r)
    d = length(x)
    k1, k2 = derivative_laplacian(k, r²)
    D = (-2k1*I)(d)
    C = MMatrix{1, 1}(-4k2)
    return K = Woodbury(D, r, C, r')
end

function gradient_kernel!(K::IsotropicGradientKernelElement, k, x, y, ::IsotropicInput)
    typeof(K)(k, x, y, IsotropicInput())
end

function gradient_kernel(k, x, y, ::IsotropicInput)
    T = gramian_eltype(k, x, y)
    IsotropicGradientKernelElement{T}(k, x, y)
end

const DotProductGradientKernelElement{T, K, X, Y} = GradientKernelElement{T, K, X, Y, DotProductInput}
function DotProductGradientKernelElement{T}(k, x, y) where T
    DotProductGradientKernelElement{T, typeof(k), typeof(x), typeof(y)}(k, x, y, DotProductInput())
end
function LinearAlgebra.mul!(b, K::DotProductGradientKernelElement, a, α::Number = 1, β::Number = 0)
    k, x, y = K.k, K.x, K.y
    d² = dot(x, y)
    k1, k2 = derivative_laplacian(k, d²)
    dot_x_a = x'a # dot(x, a)
    @. b = α * (k1 * a + k2 * y * dot_x_a) + β * b
end

# sparse but not completely lazy representation
function WoodburyFactorizations.Woodbury(K::DotProductGradientKernelElement)
    k, x, y = K.k, K.x, K.y
    d² = dot(x, y)
    k1, k2 = derivative_laplacian(k, d²)
    D = (k1*I)(length(x))
    C = MMatrix{1, 1}(k2)
    return K = Woodbury(D, copy(y), C, copy(x)')
end

function gradient_kernel!(K::DotProductGradientKernelElement, k, x, y, ::DotProductInput)
    typeof(K)(k, x, y, DotProductInput())
end

function gradient_kernel(k, x, y, ::DotProductInput)
    T = gramian_eltype(k, x, y)
    DotProductGradientKernelElement{T}(k, x, y)
end

const LinearFunctionalGradientKernelElement{T, K, X, Y} = GradientKernelElement{T, K, X, Y, StationaryLinearFunctionalInput}
function LinearFunctionalGradientKernelElement{T}(k, x, y) where T
    LinearFunctionalGradientKernelElement{T, typeof(k), typeof(x), typeof(y)}(k, x, y, StationaryLinearFunctionalInput())
end

function LinearAlgebra.mul!(b, K::LinearFunctionalGradientKernelElement, a, α::Number = 1, β::Number = 0)
    k, x, y = K.k, K.x, K.y
    r = difference(x, y)
    cr = dot(k.c, r)
    k1, k2 = derivative_laplacian(k, cr)
    dot_c_a = k.c'a # dot(c, a)
    @. b = α * -k2 * k.c * dot_c_a + β * b
end

# Base.Matrix(K::LinearFunctionalGradientKernelElement) = Matrix(LazyMatrixProduct(K)) # or more generally: K * I(size(K, 1))
# sparse but not completely lazy representation
function LazyMatrixProduct(K::LinearFunctionalGradientKernelElement)
    k, x, y = K.k, K.x, K.y
    r = difference(x, y)
    cr = dot(k.c, r)
    k1, k2 = derivative_laplacian(k, cr)
    c = zeros(eltype(k.c), length(x), 1)
    @. c = k.c
    c2 = -k2*c
    return LazyMatrixProduct(c, c2')
end

# is this necessary?
function gradient_kernel!(K::LinearFunctionalGradientKernelElement, k, x, y, ::StationaryLinearFunctionalInput)
    typeof(K)(k, x, y, StationaryLinearFunctionalInput())
end

function gradient_kernel(k, x, y, ::StationaryLinearFunctionalInput)
    T = gramian_eltype(k, x, y)
    LinearFunctionalGradientKernelElement{T}(k, x, y)
end

function evaluate_block!(Gij, k::GradientKernel, x, y, IT = input_trait(k))
    gradient_kernel!(Gij, k.k, x, y, IT)
end

############################# Constant Kernel ##################################
# efficient specialization for constants
# IDEA: have special constant input trait, since it can combine with any kernel
gradient_kernel(k::Constant, x, y, ::IsotropicInput) = gradient_kernel(k, x, y)
gradient_kernel!(K, k::Constant, x, y, ::IsotropicInput) = gradient_kernel!(K, k, x, y)
function gradient_kernel(k::Constant, x, y)
    Zeros(length(x), length(y))
end
gradient_kernel!(K::Zeros, k::Constant, x, y) = K
gradient_kernel!(K::AbstractMatrix, k::Constant, x, y) = (K .= 0)

function gramian(k::GradientKernel{<:Any, <:Constant}, x::AbstractVector, y::AbstractVector)
    d = length(x[1])
    Zeros(length(x) * d, length(y) * d)
end
########################### Neural Network Kernel ##############################
# specialization of gradient nn kernel, written before more general AD-implementation,
# which obviates the need for special derivations like these
function allocate_gradient_kernel(k::NeuralNetwork, x, y, ::GenericInput)
    x isa SVector && (x = MVector(x))
    y isa SVector && (y = MVector(y))
    U = hcat(x, y) # this preserves static arrays
    D = Diagonal(zero(x))
    # C = similar(x, 2, 2)
    elty = typeof(k(x, y))
    C = zero(MMatrix{2, 2, elty})
    K = Woodbury(D, U, C, U')
end

# specialization for neural network kernel
# NOTE: assumes K is allocated with allocate_gradient_kernel
# IDEA: take care of σ constant of NN kernel (variance of bias term)
function gradient_kernel!(K::Woodbury, k::NeuralNetwork, x, y, ::GenericInput = GenericInput())
    Nx, Ny = dot(x, x) + 1, dot(y, y) + 1
    dxy, Nxy = dot(x, y), Nx * Ny
    d = dxy / sqrt(Nxy)
    f0(x) = 2/π * asin(x)
    f1(x) = 2/π * inv(sqrt(1-x^2)) # first derivative
    f2(x) = 2/π * inv(sqrt(1-x^2))^3 * x # second derivative
    k0, k1, k2 = f0(d), f1(d), f2(d) # kernel value and derivatives
    k1 /= sqrt(Nxy)
    k2 /= Nxy # incorporate normalization to simplify following expression

    # gradient
    # gradx = k1 / sqrt(Nxy) * (y - dxy / Nx * x) # this appears correct

    # second order derivatives
    @. K.A.diag = k1
    @. K.U[:, 1] = x
    @. K.U[:, 2] = y
    k1k2dxy = (k1 + k2 * dxy)
    K.C[1, 1] = - k1k2dxy / Nx # after normalization, dxy could become d
    K.C[2, 1] = k2
    K.C[1, 2] = k1k2dxy * dxy / Nxy
    K.C[2, 2] = - k1k2dxy / Ny
    return K
end

################################################################################
# structure that holds the potentially heterogeneous types of covariances between
# a kernel value and its derivatives
# mutable only because of value_value?
mutable struct DerivativeKernelElement{T, VV, VG, GV, GG, VH, HV, GH, HG, HH, GI, HI} <: AbstractMatrix{T}
    value_value::VV # value-value covariance (IDEA: if this is MMAtrix, could make structure non-mutable)
    value_gradient::VG # value-gradient covariance
    gradient_value::GV # gradient-value covariance
    gradient_gradient::GG # ...
    value_hessian::VH
    hessian_value::HV
    gradient_hessian::GH
    hessian_gradient::HG
    hessian_hessian::HH
    gradient_indices::GI
    hessian_indices::HI
end

function DerivativeKernelElement(vv, vg, gv, gg, vh = nothing, hv = nothing, gh = nothing, hg = nothing, hh = nothing)
    t = (vv, vg, gv, gg, vh, hv, gh, hg, hh)
    s = something(t...) # get the first input that's not nothing
    T = eltype(s)
    gi = gradient_indices(vv, gg)
    hi = hessian_indices(vv, gg, hh)
    DerivativeKernelElement{T, typeof.(t)..., typeof(gi), typeof(hi)}(t..., gi, hi)
end

function gradient_indices(value_value, gradient_gradient)
    i = 0
    if !isnothing(value_value)
        i += 1
    end
    if !isnothing(gradient_gradient)
        d = size(gradient_gradient, 1)
        return (i + 1) : (i + d)
    else
        return nothing
    end
end

function hessian_indices(value_value, gradient_gradient, hessian_hessian)
    i = 0
    if !isnothing(value_value)
        i += 1
    end
    if !isnothing(gradient_gradient)
        d = size(gradient_gradient, 1)
        i += d
    end
    if !isnothing(hessian_hessian)
        dd = size(hessian_hessian, 1)
        return (i + 1) : (i + dd)
    else
        return nothing
    end
end

Base.eltype(A::DerivativeKernelElement{T}) where {T} = T
function Base.size(A::DerivativeKernelElement)
    n = 0
    if !isnothing(A.value_value)
        n += 1
    end
    if !isnothing(A.gradient_gradient)
        n += size(A.gradient_gradient, 1)
    end
    if !isnothing(A.hessian_hessian)
        n += size(A.hessian_hessian, 1)
    end
    return (n, n)
end
Base.size(A::DerivativeKernelElement, i::Int) = i <= 2 ? size(A)[i] : 1

# TODO: efficient indexing
function Base.getindex(A::DerivativeKernelElement, i::Int, j::Int)
    # Matrix(A)[i, j]
    @boundscheck begin
        n = size(A, 1)
        (1 <= i <= n && 1 <= j <= n) || throw(BoundsError("attempt to access $(size(A)) DerivativeKernelElement at index ($i, $j)"))
    end
    gi = A.gradient_indices
    hi = A.hessian_indices
    if !isnothing(A.value_value)
        (i == 1 && j == 1) && return A.value_value
        if i == 1
            (!isnothing(gi) && j in gi) && return A.value_gradient[findfirst(==(j), gi)]
            (!isnothing(hi) && j in hi) && return A.value_hessian[findfirst(==(j), hi)]
        end
        if j == 1
            (!isnothing(gi) && i in gi) && return A.gradient_value[findfirst(==(i), gi)]
            (!isnothing(hi) && i in hi) && return A.hessian_value[findfirst(==(i), hi)]
        end
    end
    if !isnothing(gi)
        (i in gi && j in gi) && return A.gradient_gradient[findfirst(==(i), gi), findfirst(==(j), gi)]
        if !isnothing(hi)
            (i in gi && j in hi) && return A.gradient_hessian[findfirst(==(i), gi), findfirst(==(j), hi)]
            (i in hi && j in gi) && return A.hessian_gradient[findfirst(==(i), hi), findfirst(==(j), gi)]
        end
    end
    if !isnothing(hi)
        return A.hessian_hessian[findfirst(==(i), hi), findfirst(==(j), hi)]
    end
    error("this should never happen")
end

function LinearAlgebra.mul!(y::AbstractVector, A::DerivativeKernelElement, x::AbstractVector, α::Number = 1, β::Number = 0)
    y .*= β
    gi = A.gradient_indices
    hi = A.hessian_indices
    if !isnothing(A.value_value)
        y[1] += α * A.value_value * x[1]
    end
    if !isnothing(A.value_gradient)
        xg, yg = @views x[gi], y[gi]
        y[1] += α * dot(A.value_gradient, xg)
        @. yg += α * A.gradient_value * x[1]
    end
    if !isnothing(A.gradient_gradient)
        xg, yg = @views x[gi], y[gi]
        mul!(yg, A.gradient_gradient, xg, α, 1)
    end
    if !isnothing(A.value_hessian)
        xh, yh = @views x[hi], y[hi]
        y[1] += α * dot(A.value_hessian, xh)
        @. yh += α * A.hessian_value * x[1]
    end
    if !isnothing(A.gradient_hessian)
        xg, yg = @views x[gi], y[gi]
        xh, yh = @views x[hi], y[hi]
        mul!(yg, A.gradient_hessian, xh, α, 1)
        mul!(yh, A.hessian_gradient, xg, α, 1)
    end
    if !isnothing(A.hessian_hessian)
        xh, yh = @views x[hi], y[hi]
        mul!(yh, A.hessian_hessian, xh, α, 1)
    end
    return y
end

function Base.Matrix(A::DerivativeKernelElement)
    n = size(A, 1)
    M = zeros(eltype(A), n, n)
    gradient_index = 0
    hessian_index = 0
    if !isnothing(A.value_value)
        M[1, 1] = A.value_value
        gradient_index += 1 # if we have a value variance, gradient indices are 2:2+d
        hessian_index += 1
    end
    if !isnothing(A.value_gradient)
        d = length(A.value_gradient)
        @. M[2:1+d, 1] = A.gradient_value
        @. M[1, 2:1+d] = A.value_gradient
    end
    if !isnothing(A.gradient_gradient)
        d = size(A.gradient_gradient, 1)
        gi = gradient_index+1:gradient_index+d
        M[gi, gi] .= Matrix(A.gradient_gradient)
        hessian_index += d + 1 # if we have a gradient covariances, hessian indices are (1 + hasvalue? + d + 1)
    end
    if !isnothing(A.value_hessian)
        dd = length(A.value_hessian)
        hi = hessian_index+1:hessian_index+dd
        @. M[1, hi] = A.value_hessian
        @. M[hi, 1] = A.hessian_value
    end
    if !isnothing(A.gradient_hessian)
        d, dd = size(A.gradient_hessian)
        gi = gradient_index+1:gradient_index+d
        hi = hessian_index+1:hessian_index+dd
        M[gi, hi] .= Matrix(A.gradient_hessian)
        M[hi, gi] .= Matrix(A.hessian_gradient)
    end
    if !isnothing(A.hessian_hessian)
        dd = size(A.hessian_hessian, 1)
        hi = hessian_index+1:hessian_index+dd
        M[hi, hi] .= Matrix(A.hessian_hessian)
    end
    return M
end

########################## using new kernel element ############################
################################################################################
# [f, ∂f] ∼ GP([μ, ∂μ], dK) # value + gradient kernel
# IDEA: For efficiency, maybe create ValueGradientKernelElement like in hessian.jl
# currently, this is an order of magnitude slower than GradientKernel
struct ValueGradientKernel{T, K, IT<:InputTrait} <: MultiKernel{T}
    k::K
    input_trait::IT
    function ValueGradientKernel{T, K}(k, it = input_trait(k)) where {T, K}
        new{T, K, typeof(it)}(k, it)
    end
end
ValueGradientKernel(k::AbstractKernel{T}) where {T} = ValueGradientKernel{T, typeof(k)}(k)
ValueGradientKernel(k) = ValueGradientKernel{Float64, typeof(k)}(k) # use fieldtype here?

input_trait(G::ValueGradientKernel) = G.input_trait

function elsize(G::Gramian{<:AbstractMatrix, <:ValueGradientKernel}, i::Int)
    i ≤ 2 ? length(G.x[1]) + 1 : 1
end

# computes covariances of the function and its derivative.
function (G::ValueGradientKernel)(x, y)
    value_gradient_kernel(G.k, x, y, input_trait(G.k))
end

function value_gradient_kernel(k, x, y, IT::InputTrait = input_trait(G.k))
    d = length(x)
    kxy = k(x, y)
    value_value = kxy
    value_gradient = zeros(typeof(kxy), d)
    gradient_value = zeros(typeof(kxy), d)
    gradient_gradient = gradient_kernel(k, x, y, IT)
    K = DerivativeKernelElement(value_value, value_gradient, gradient_value, gradient_gradient)
    value_gradient_kernel!(K, k, x, y, IT)
end

# IDEA: for Sum and Product kernels, if input_trait is not passed, could default to Generic
# No, should keep track of input type in kernel
function value_gradient_kernel!(K::DerivativeKernelElement, k, x, y)
    value_gradient_kernel!(K, k, x, y, input_trait(k))
end
function value_gradient_kernel!(K::DerivativeKernelElement, k, x, y, IT::InputTrait)
    K.value_value = k(x, y)
    value_gradient_covariance!(K.gradient_value, K.value_gradient, k, x, y, IT)
    K.gradient_gradient = gradient_kernel!(K.gradient_gradient, k, x, y, IT)
    return K
end

# NOTE: since value_gradient_covariance! was added, specializations of value_gradient_kernel!
# only yield marginal performance improvements (<2x)
# could be removed to reduce LOC
function value_gradient_kernel!(K::DerivativeKernelElement, k, x, y, T::DotProductInput)
    xy = dot(x, y)
    k0, k1 = value_derivative(k, xy)
    K.value_value = k0
    @. K.value_gradient = k1*x # ForwardDiff.gradient(z->k(x, z), y)'
    @. K.gradient_value = k1*y # ForwardDiff.gradient(z->k(z, y), x)
    K.gradient_gradient = gradient_kernel!(K.gradient_gradient, k, x, y, T)
    return K
end

function value_gradient_kernel!(K::DerivativeKernelElement, k, x, y, T::IsotropicInput)
    r = K.gradient_value
    @. r = x - y
    d² = sum(abs2, r)
    k0, k1 = value_derivative(k, d²)
    K.value_value = k0
    @. K.value_gradient = r
    @. K.value_gradient *= -2k1 # ForwardDiff.gradient(z->k(x, z), y)'
    @. K.gradient_value *= 2k1 # ForwardDiff.gradient(z->k(z, y), x)
    K.gradient_gradient = gradient_kernel!(K.gradient_gradient, k, x, y, T)
    return K
end

# need this to work efficiently with BlockFactorizations, see blockmul! in gramian
# reuses temporary storage in Gij, speeds up MVM by more than an order of magnitude
function evaluate_block!(Gij::DerivativeKernelElement, k::ValueGradientKernel, x, y, IT::InputTrait = input_trait(k))
    value_gradient_kernel!(Gij, k.k, x, y, IT)
end

################################################################################
# computes covariance between value and gradient, needed for value-gradient kernel
# technically, complexity would be O(d) even with generic implementation
# in practice, at least an order of magitude can be gained by specialized implementations
function value_gradient_covariance!(gx, gy, k, x, y, ::GenericInput)
    # GradientConfig() # for generic version, this could be pre-computed for efficiency gains
    ForwardDiff.gradient!(gx, z->k(z, y), x) # these are bottlenecks
    ForwardDiff.gradient!(gy, z->k(x, z), y)
    return gx, gy
end

function value_gradient_covariance!(gx, gy, k, x, y, ::GenericInput, α::Real, β::Real)
    tx, ty = ForwardDiff.gradient(z->k(z, y), x), ForwardDiff.gradient(z->k(x, z), y)
    @. gx = α * tx + β * gx
    @. gy = α * ty + β * gy
    return gx, gy
end

function value_gradient_covariance!(gx, gy, k::Sum, x, y, ::GenericInput, α::Real = 1, β::Real = 0)
    @. gx *= β
    @. gy *= β
    for h in k.args # input_trait(h) should not be called if h is a composite kernel for higher effiency
        value_gradient_covariance!(gx, gy, h, x, y, input_trait(h), α, 1)
    end
    return gx, gy
end

# this is more tricky to do without additional allocations
# could do it with one more vector
function value_gradient_covariance!(gx, gy, k::Product, x, y, ::GenericInput, α::Real = 1, β::Real = 0)
    @. gx *= β
    @. gy *= β
    prod_k_xy = k(x, y)
    if !iszero(prod_k_xy)
        for i in eachindex(k.args)
            ki = k.args[i]
            αi = α * prod_k_xy / ki(x, y)
            value_gradient_covariance!(gx, gy, ki, x, y, input_trait(ki), αi, 1)
        end
    else # GradientKernel of Product requires less efficient O(dr²) special case if prod_k_xy is zero, for now, throw error
        throw(DomainError("value_gradient_covariance! of Product not supported for products that are zero"))
    end
    return gx, gy
end

function value_gradient_covariance!(gx, gy, k, x, y, ::IsotropicInput, α::Real = 1, β::Real = 0)
    r = difference(x, y)
    r² = sum(abs2, r)
    k1 = ForwardDiff.derivative(k, r²) # IDEA: this computation could be pooled with the gradient computation
    @. gx = α * 2k1 * r + β * gx
    @. gy = -α * 2k1 * r + β * gy
    return gx, gy
end

function value_gradient_covariance!(gx, gy, k, x, y, ::DotProductInput, α::Real = 1, β::Real = 0)
    xy = dot(x, y)
    k1 = ForwardDiff.derivative(k, xy) # IDEA: this computation could be pooled with the gradient computation
    @. gx = α * k1 * y + β * gx
    @. gy = α * k1 * x + β * gy
    return gx, gy
end

function value_gradient_covariance!(gx, gy, k, x, y, ::StationaryLinearFunctionalInput, α::Real = 1, β::Real = 0)
    cr = dot(k.c, difference(x, y))
    k1 = ForwardDiff.derivative(k, cr) # IDEA: this computation could be pooled with the gradient computation
    @. gx = α * k1 * k.c + β * gx
    @. gy = -α * k1 * k.c + β * gy
    return gx, gy
end

################################################################################
# for 1d inputs
struct DerivativeKernel{T, K} <: AbstractKernel{T}
    k::K
end
DerivativeKernel(k) = DerivativeKernel{Float64, typeof(k)}(k)
function (G::DerivativeKernel)(x::Real, y::Real)
    g(y) = ForwardDiff.derivative(z->G.k(z, y), x)
    ForwardDiff.derivative(g, y)
end

################################################################################
struct ValueDerivativeKernel{T, K} <: MultiKernel{T}
    k::K
end
ValueDerivativeKernel(k) = ValueDerivativeKernel{Float64, typeof(k)}(k)
function (k::ValueDerivativeKernel)(x::Real, y::Real)
    x = float(x)
    x, y = promote(x, y)
    value_derivative_kernel!(zeros(eltype(x), (2, 2)), k, x, y)
end
function value_derivative_kernel!(K::AbstractMatrix, g::ValueDerivativeKernel, x::Real, y::Real)
    function f(k, x::Real, y::Real)
        r = DiffResults.DiffResult(zero(y), zero(y)) # this could take pre-allocated temporary storage
        r = ForwardDiff.derivative!(r, z->k(z, y), x)
        vcat(r.value, r.derivs[1])
    end
    value = @view K[:, 1] # value of helper, i.e. original value and gradient stacked
    derivs = @view K[:, 2] # jacobian of helper (higher order terms)
    result = DiffResults.DiffResult(value, derivs)
    ForwardDiff.jacobian!(result, z->f(g.k, x, z[1]), [y])
    return K
end

############################# Helpers ##########################################
# computes value, first, and second derivatives of f at x
# this works even with nested differentiation
function value_derivative_laplacian(f, x::Real)
    fx = f(x)
    fx, derivative_laplacian(f, x, typeof(fx))...
end
# TODO: overwrite for MaternP
function derivative_laplacian(f, x::Real, T::DataType = typeof(x))
    g(x) = ForwardDiff.derivative(f, x)
    value_derivative(g, x, T)
end

# NOTE: if nested differentation through value_derivative is happening,
# need to pass correct output T = f(x)
function value_derivative(f, x::Real, T::DataType = typeof(x))
    r = DiffResults.DiffResult(zero(T), zero(T)) # this could take pre-allocated temporary storage
    r = ForwardDiff.derivative!(r, f, x)
    r.value, r.derivs[1]
end


text:
# Julia wrapper for header: /usr/local/include/libavutil/opt.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

export
    av_find_opt,
    av_set_string3,
    av_set_double,
    av_set_q,
    av_set_int,
    av_get_double,
    av_get_q,
    av_get_int,
    av_get_string,
    av_next_option,
    av_opt_show2,
    av_opt_set_defaults,
    av_opt_set_defaults2,
    av_set_options_string,
    av_opt_free,
    av_opt_flag_is_set,
    av_opt_set_dict,
    av_opt_eval_flags,
    av_opt_eval_int,
    av_opt_eval_int64,
    av_opt_eval_float,
    av_opt_eval_double,
    av_opt_eval_q,
    av_opt_find,
    av_opt_find2,
    av_opt_next,
    av_opt_child_next,
    av_opt_child_class_next,
    av_opt_set,
    av_opt_set_int,
    av_opt_set_double,
    av_opt_set_q,
    av_opt_get,
    av_opt_get_int,
    av_opt_get_double,
    av_opt_get_q

function av_find_opt(obj,name,unit,mask::Integer,flags::Integer)
    ccall((:av_find_opt,libavutil),Ptr{AVOption},(Ptr{Void},Ptr{Uint8},Ptr{Uint8},Cint,Cint),obj,name,unit,mask,flags)
end

function av_set_string3(obj,name,val,alloc::Integer,o_out)
    ccall((:av_set_string3,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Ptr{Uint8},Cint,Ptr{Ptr{AVOption}}),obj,name,val,alloc,o_out)
end

function av_set_double(obj,name,n::Cdouble)
    ccall((:av_set_double,libavutil),Ptr{AVOption},(Ptr{Void},Ptr{Uint8},Cdouble),obj,name,n)
end

function av_set_q(obj,name,n::AVRational)
    ccall((:av_set_q,libavutil),Ptr{AVOption},(Ptr{Void},Ptr{Uint8},AVRational),obj,name,n)
end

function av_set_int(obj,name,n::Int64)
    ccall((:av_set_int,libavutil),Ptr{AVOption},(Ptr{Void},Ptr{Uint8},Int64),obj,name,n)
end

function av_get_double(obj,name,o_out)
    ccall((:av_get_double,libavutil),Cdouble,(Ptr{Void},Ptr{Uint8},Ptr{Ptr{AVOption}}),obj,name,o_out)
end

function av_get_q(obj,name,o_out)
    ccall((:av_get_q,libavutil),AVRational,(Ptr{Void},Ptr{Uint8},Ptr{Ptr{AVOption}}),obj,name,o_out)
end

function av_get_int(obj,name,o_out)
    ccall((:av_get_int,libavutil),Int64,(Ptr{Void},Ptr{Uint8},Ptr{Ptr{AVOption}}),obj,name,o_out)
end

function av_get_string(obj,name,o_out,buf,buf_len::Integer)
    ccall((:av_get_string,libavutil),Ptr{Uint8},(Ptr{Void},Ptr{Uint8},Ptr{Ptr{AVOption}},Ptr{Uint8},Cint),obj,name,o_out,buf,buf_len)
end

function av_next_option(obj,last)
    ccall((:av_next_option,libavutil),Ptr{AVOption},(Ptr{Void},Ptr{AVOption}),obj,last)
end

function av_opt_show2(obj,av_log_obj,req_flags::Integer,rej_flags::Integer)
    ccall((:av_opt_show2,libavutil),Cint,(Ptr{Void},Ptr{Void},Cint,Cint),obj,av_log_obj,req_flags,rej_flags)
end

function av_opt_set_defaults(s)
    ccall((:av_opt_set_defaults,libavutil),Void,(Ptr{Void},),s)
end

function av_opt_set_defaults2(s,mask::Integer,flags::Integer)
    ccall((:av_opt_set_defaults2,libavutil),Void,(Ptr{Void},Cint,Cint),s,mask,flags)
end

function av_set_options_string(ctx,opts,key_val_sep,pairs_sep)
    ccall((:av_set_options_string,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Ptr{Uint8},Ptr{Uint8}),ctx,opts,key_val_sep,pairs_sep)
end

function av_opt_free(obj)
    ccall((:av_opt_free,libavutil),Void,(Ptr{Void},),obj)
end

function av_opt_flag_is_set(obj,field_name,flag_name)
    ccall((:av_opt_flag_is_set,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Ptr{Uint8}),obj,field_name,flag_name)
end

function av_opt_set_dict(obj,options)
    ccall((:av_opt_set_dict,libavutil),Cint,(Ptr{Void},Ptr{Ptr{AVDictionary}}),obj,options)
end

function av_opt_eval_flags(obj,o,val,flags_out)
    ccall((:av_opt_eval_flags,libavutil),Cint,(Ptr{Void},Ptr{AVOption},Ptr{Uint8},Ptr{Cint}),obj,o,val,flags_out)
end

function av_opt_eval_int(obj,o,val,int_out)
    ccall((:av_opt_eval_int,libavutil),Cint,(Ptr{Void},Ptr{AVOption},Ptr{Uint8},Ptr{Cint}),obj,o,val,int_out)
end

function av_opt_eval_int64(obj,o,val,int64_out)
    ccall((:av_opt_eval_int64,libavutil),Cint,(Ptr{Void},Ptr{AVOption},Ptr{Uint8},Ptr{Int64}),obj,o,val,int64_out)
end

function av_opt_eval_float(obj,o,val,float_out)
    ccall((:av_opt_eval_float,libavutil),Cint,(Ptr{Void},Ptr{AVOption},Ptr{Uint8},Ptr{Cfloat}),obj,o,val,float_out)
end

function av_opt_eval_double(obj,o,val,double_out)
    ccall((:av_opt_eval_double,libavutil),Cint,(Ptr{Void},Ptr{AVOption},Ptr{Uint8},Ptr{Cdouble}),obj,o,val,double_out)
end

function av_opt_eval_q(obj,o,val,q_out)
    ccall((:av_opt_eval_q,libavutil),Cint,(Ptr{Void},Ptr{AVOption},Ptr{Uint8},Ptr{AVRational}),obj,o,val,q_out)
end

function av_opt_find(obj,name,unit,opt_flags::Integer,search_flags::Integer)
    ccall((:av_opt_find,libavutil),Ptr{AVOption},(Ptr{Void},Ptr{Uint8},Ptr{Uint8},Cint,Cint),obj,name,unit,opt_flags,search_flags)
end

function av_opt_find2(obj,name,unit,opt_flags::Integer,search_flags::Integer,target_obj)
    ccall((:av_opt_find2,libavutil),Ptr{AVOption},(Ptr{Void},Ptr{Uint8},Ptr{Uint8},Cint,Cint,Ptr{Ptr{Void}}),obj,name,unit,opt_flags,search_flags,target_obj)
end

function av_opt_next(obj,prev)
    ccall((:av_opt_next,libavutil),Ptr{AVOption},(Ptr{Void},Ptr{AVOption}),obj,prev)
end

function av_opt_child_next(obj,prev)
    ccall((:av_opt_child_next,libavutil),Ptr{Void},(Ptr{Void},Ptr{Void}),obj,prev)
end

function av_opt_child_class_next(parent,prev)
    ccall((:av_opt_child_class_next,libavutil),Ptr{AVClass},(Ptr{AVClass},Ptr{AVClass}),parent,prev)
end

function av_opt_set(obj,name,val,search_flags::Integer)
    ccall((:av_opt_set,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Ptr{Uint8},Cint),obj,name,val,search_flags)
end

function av_opt_set_int(obj,name,val::Int64,search_flags::Integer)
    ccall((:av_opt_set_int,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Int64,Cint),obj,name,val,search_flags)
end

function av_opt_set_double(obj,name,val::Cdouble,search_flags::Integer)
    ccall((:av_opt_set_double,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Cdouble,Cint),obj,name,val,search_flags)
end

function av_opt_set_q(obj,name,val::AVRational,search_flags::Integer)
    ccall((:av_opt_set_q,libavutil),Cint,(Ptr{Void},Ptr{Uint8},AVRational,Cint),obj,name,val,search_flags)
end

function av_opt_get(obj,name,search_flags::Integer,out_val)
    ccall((:av_opt_get,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Cint,Ptr{Ptr{Uint8}}),obj,name,search_flags,out_val)
end

function av_opt_get_int(obj,name,search_flags::Integer,out_val)
    ccall((:av_opt_get_int,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Cint,Ptr{Int64}),obj,name,search_flags,out_val)
end

function av_opt_get_double(obj,name,search_flags::Integer,out_val)
    ccall((:av_opt_get_double,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Cint,Ptr{Cdouble}),obj,name,search_flags,out_val)
end

function av_opt_get_q(obj,name,search_flags::Integer,out_val)
    ccall((:av_opt_get_q,libavutil),Cint,(Ptr{Void},Ptr{Uint8},Cint,Ptr{AVRational}),obj,name,search_flags,out_val)
end


text:
struct Pingreq <: Packet
    header::UInt8
end
Pingreq() = Pingreq(PINGREQ)
Base.show(io::IO, x::Pingreq) = print(io, "pingreq")

struct Pingresp <: Packet
    header::UInt8
end
Pingresp() = Pingresp(PINGRESP)
Base.show(io::IO, x::Pingresp) = print(io, "pingresp")

max_issues_repo_name: kivaari/jlMQTT

max_forks_repo_name: kivaari/jlMQTT


text:
using LinearAlgebra
using LazySets

@testset "lazysets" begin
    pt2 = Singleton([1.0,1.0])
    pt3 = Singleton([1.0,1.0,1.0])
    vec2 = SVector{2}([1.0, 0.0])
    vec3 = SVector{3}([1.0, 0.0, 0.0])

    @testset "Ball1" begin
        b1 = Ball1(zeros(2), 1.0)
        ret = minimum_distance(b1, pt2, vec2)
        @test ret ≈ 0.7071067811865476 rtol=0.001
        b1 = Ball1(zeros(3), 1.0)
        ret = minimum_distance(b1, pt3, vec3)
        @test ret ≈ 1.1547005383792515 rtol=0.001
    end

    @testset "Ball2" begin
        b2 = Ball2(zeros(2), 1.0)
        ret = minimum_distance(b2, pt2, vec2)
        @test ret ≈ 0.41421356237309515 rtol=0.001
        b2 = Ball2(zeros(3), 1.0)
        ret = minimum_distance(b2, pt3, vec3)
        @test ret ≈ 0.7320508075688773 rtol=0.001
    end

    @testset "BallInf" begin
        bi = BallInf(zeros(2), 1.0)
        ret = minimum_distance(bi, pt2, vec2)
        @test ret ≈ 0.0 rtol=0.001
        bi = BallInf(zeros(3), 1.0)
        ret = minimum_distance(bi, pt3, vec3)
        @test ret ≈ 0.0 rtol=0.001
    end

    @testset "Ballp" begin
        bp = Ballp(1.5, zeros(2), 1.0)
        ret = minimum_distance(bp, pt2, vec2)
        @test ret ≈ 0.5233148442327559 rtol=0.001
        bp = Ballp(1.5, zeros(3), 1.0)
        ret = minimum_distance(bp, pt3, vec3)
        @test ret ≈ 0.8993676299132733 rtol=0.001
        bp = Ballp(2.5, zeros(2), 1.0)
        ret = minimum_distance(bp, pt2, vec2)
        @test ret ≈ 0.3424400998368018 rtol=0.001
        bp = Ballp(2.5, zeros(3), 1.0)
        ret = minimum_distance(bp, pt3, vec3)
        @test ret ≈ 0.6159276335349732 rtol=0.001
    end

    #  Not testing anymore because rngs can be different
    #  on different versions of Base and StableRNGs is not
    #  fully supported by LazySets
    #
    #  @testset "Ellipsoid" begin
    #      e = rand(Ellipsoid, dim=2, rng=StableRNG(1))
    #      ret = minimum_distance(e, pt2, vec2)
    #      @test ret ≈ 0.5015857643287766  rtol=0.001
    #      e = rand(Ellipsoid, dim=3, seed=8)
    #      ret = minimum_distance(e, pt3, vec3)
    #      @test ret ≈ 0.17395550015075892 rtol=0.001
    #  end
    #
    #  @testset "HPolygon" begin
    #      hp = rand(HPolygon, seed=2)
    #      ret = minimum_distance(hp, pt2, vec2)
    #      @test ret ≈ 1.099023620263928 rtol=0.001
    #      hp = rand(HPolygonOpt, seed=3)
    #      ret = minimum_distance(hp, pt2, vec2)
    #      @test ret ≈ 0.19296447682913367 rtol=0.001
    #  end
    #
    #  @testset "Hyperrectangle" begin
    #      hr = rand(Hyperrectangle, dim=2, seed=3)
    #      ret = minimum_distance(hr, pt2, vec2)
    #      @test ret ≈ 2.546482967114228 rtol=0.001
    #      hr = rand(Hyperrectangle, dim=3, seed=4)
    #      ret = minimum_distance(hr, pt3, vec3)
    #      @test ret ≈ 2.4101143198871897 rtol=0.001
    #  end
    #
    #  @testset "LineSegment" begin
    #      ls = rand(LineSegment, dim=2, seed=5)
    #      ret = minimum_distance(ls, pt2, vec2)
    #      @test ret ≈ 0.44758226978293086 rtol=0.001
    #  end
    #
    #  @testset "VPolytope" begin
    #      vp = rand(VPolygon, dim=2, seed=10)
    #      ret = minimum_distance(vp, pt2, vec2)
    #      @test ret ≈ 1.2323399249078948 rtol=0.001
    #      vp = rand(VPolytope, dim=2, seed=12)
    #      ret = minimum_distance(vp, pt2, vec2)
    #      @test ret ≈ 0.5871018052878227 rtol=0.001
    #      vp = rand(VPolytope, dim=3, seed=14)
    #      ret = minimum_distance(vp, pt3, vec3)
    #      @test ret ≈ 0.9811157850634189 rtol=0.001
    #  end
    #
    #  @testset "Zonotope" begin
    #      z = rand(Zonotope, dim=2, seed=16)
    #      ret = minimum_distance(z, pt2, vec2)
    #      @test ret ≈ 0.9086568206474619 rtol=0.001
    #      z = rand(Zonotope, dim=3, seed=17)
    #      ret = minimum_distance(z, pt3, vec3)
    #      @test ret ≈ 0.1823104257339578 rtol=0.001
    #  end
end


text:
#
#----------------------------------------------------------------------------#
abstract type AbstractBanachSpace end # banach space
#----------------------------------------------------------------------------#
export nnBanachSpace
#--------------------------------------#
struct nnBanachSpace <: AbstractBanachSpace
    nprt::Integer
    ndim::Integer
    NN
end
#--------------------------------------#
function nnBanachSpace(nprt::Integer,ndim::Integer;
                       depth=2,
                       act=tanh, # cos
                       init=Flux.glorot_uniform)

    W = Array{Float32}((0:ndim-1) * ones(1,1)) .* pi
    b = zeros(Float32,ndim)
    chain = Chain(Dense(W,b,act),
                  Chain([Dense(init(ndim,ndim),init(ndim),act)
                                           for i in 1:(depth-1)]...)
                 )
    if(nprt == 1)
        NN = chain
    else
        NN = Parallel(vcat,(deepcopy(chain) for i=1:nprt)...)
    end

    return nnBanachSpace(nprt,ndim,NN)
end
#--------------------------------------#
Flux.@functor nnBanachSpace
DiffEqFlux.initial_params(bsp::nnBanachSpace) = Flux.destructure(bsp)[1]
#--------------------------------------#
function (bsp::nnBanachSpace)(x)

    @unpack ndim,nprt,NN = bsp

    N = size(x)[end]

    ψ = NN(x)
    ψ = reshape(ψ,ndim,nprt,N)

    return ψ
end
#----------------------------------------------------------------------------#
export Cosines
#--------------------------------------#
struct Cosines <: AbstractBanachSpace
    nprt::Integer
    ndim::Integer
    α::AbstractArray # [ndim,nprt]
    β::AbstractArray # [ndim,nprt]
end
#--------------------------------------#
function Cosines(nprt::Integer,ndim::Integer)

    α = Array{Float32}((0:ndim-1)*ones(1,nprt))
    β = Array{Float32}(zeros(ndim,nprt))

    return Cosines(nprt,ndim,α,β)
end
#--------------------------------------#
Flux.@functor Cosines
DiffEqFlux.initial_params(bsp::Cosines) = Flux.destructure(bsp)[1]
#--------------------------------------#
function (bsp::Cosines)(x)
    @unpack α, β = bsp

    siz1 = Zygote.ignore() do
        dimA = length(size(α))
        siz0 = size(x)[2:end]
        return ((1 for i in 1:dimA)...,siz0...)
    end

    x̄ = reshape(x,siz1)

    # x -> (α x + β) -> cos(x)
    #
    x̄ = @. α * x̄ + β
    ψ = @. cos(pi*x̄)

    return ψ
end
#----------------------------------------------------------------------------#
export TensorProdCosines
#--------------------------------------#
struct TensorProdCosines <: AbstractBanachSpace
    nprt::Integer
    ndim::Integer
    nx::Integer
    ny::Integer
    αx::AbstractArray # [ndim,nprt]
    βx::AbstractArray
    αy::AbstractArray
    βy::AbstractArray
end
#--------------------------------------#
function TensorProdCosines(nprt::Integer,nx::Integer,ny::Integer)

    @assert nx == ny

    ndim = nx*ny

    αx = Array{Float32}((0:nx-1)*ones(1,nprt))
    βx = Array{Float32}(zeros(nx,nprt))

    αy = Array{Float32}((0:ny-1)*ones(1,nprt))
    βy = Array{Float32}(zeros(ny,nprt))

    αx = reshape(αx,nx,1,nprt)
    βx = reshape(βx,nx,1,nprt)

    αy = reshape(αy,1,ny,nprt)
    βy = reshape(βy,1,ny,nprt)

    return TensorProdCosines(nprt,ndim,nx,ny,αx,βx,αy,βy)
end
#--------------------------------------#
Flux.@functor TensorProdCosines
DiffEqFlux.initial_params(bsp::TensorProdCosines) = Flux.destructure(bsp)[1]
#--------------------------------------#
function (bsp::TensorProdCosines)(x,y)

    @assert size(x) == size(y)

    @unpack αx, βx, αy, βy = bsp
    @unpack nprt, ndim = bsp

    N = size(x)[end]

    siz1 = Zygote.ignore() do
        dimA = length(size(αx))
        siz0 = size(x)[2:end]
        return ((1 for i in 1:dimA)...,siz0...)
    end

    x̄ = reshape(x,siz1)
    ȳ = reshape(y,siz1)

    x̄ = @. αx * x̄ + βx
    ȳ = @. αy * ȳ + βy

    ψ = @. cos(pi*x̄) * cos(pi*ȳ)
    ψ = reshape(ψ,ndim,nprt,N)

    return ψ
end
#----------------------------------------------------------------------------#
export GL
#--------------------------------------#
struct GL <: AbstractBanachSpace
    nprt::Integer
    ndim::Integer
    ẑ ::AbstractVector
    xg::AbstractArray # element grid [nprt]
    xm::AbstractArray # computational mesh [ndim,nprt]
end
#--------------------------------------#
function GL(ndim::Integer,pou::AbstractPOU) # [nprt+1]

    @unpack nprt,xe,x0,x1 = pou
    xg = addEndPts(xe,x0,x1)

    ẑ,_ = FastGaussQuadrature.gausslegendre(ndim) # ∈ (-1,1)
    z   = @. 0.5 * (ẑ + 1)       # [ndim] ∈ (0,1)

    hg = diff(xg)                # [nprt]
    xm = z * hg' .+ xg[1:end-1]' # [ndim,nprt]

    return GL(nprt,ndim,ẑ,xg,xm)
end
#--------------------------------------#
Flux.@functor GL ()
DiffEqFlux.initial_params(bsp::GL) = Flux.destructure(bsp)[1]
#--------------------------------------#
function (bsp::GL)(x) # [1,N]

    @unpack ndim,nprt,xm = bsp

    u = Zygote.ignore() do
        N = size(x)[end]
        u = zeros(Float32,ndim,nprt,N)  # [ndim,nprt,N]
        for i=1:nprt
            J = SEM.interpMat(x',xm[:,i]) # [N,ndim]
            u[:,i,:] = J'
        end
        return u
    end

    return u
end
#----------------------------------------------------------------------------#


text:
#################### BED byte to RAW math ####################

"""
    maskbyte(bedbyte::UInt8, num_quarters::Integer, qoffset::Integer=0)

Reads `byte` in BED format and returns `byte`, after masking quarters
such that `(qoffset + 1):(num_quarters + qoffset)` are preserved and
other quarters are "zeroed" (with respect to RAW format). Assumes RAW
format!!

"""
function maskbyte(bedbyte::UInt8, num_quarters::Integer, qoffset::Integer=0)
    @inbounds return (bedbyte | Consts.partialbytemasks[num_quarters, qoffset + 1])
end

"""
    hasNAs(bedbyte::UInt8)

Returns `true` if `bedbyte` has any missing values, and `false`
otherwise.

"""
function hasNAs(bedbyte::UInt8)
    @inbounds return Consts.hasNAmap[bedbyte + 1]
end

"""
    zeroNAs(bedbyte::UInt8)

Reads `bedbyte` in BED format and returns `bedbyte` with any missing
values set to zero. Assumes RAW format!!

"""
function zeroNAs(bedbyte::UInt8)
    @inbounds return Consts.natozeromap[bedbyte + 1]
end

"""
    NAsup(bedbyte::UInt8)

Returns `bedbyte` in BED with all non-missing quarters set to (RAW
format) `0b00`, and missing quarters set to `0b01`. Roughly the
complement of zeroNAs.

"""
function NAsup(bedbyte::UInt8)
    @inbounds return Consts.na_sup_map[bedbyte + 1]
end

"""
    countNAs(bedbyte::UInt8)

Returns the number of missing values in `bedbyte`.

"""
function countNAs(bedbyte::UInt8)
    @inbounds return Consts.nacountmap[bedbyte + 1]
end

function bytedot(b1::UInt8, b2::UInt8)
    @inbounds return Consts.bytebytemulttable[b1 + 1, b2 + 1]
end

"""
    bytedot(byteL::UInt8, byteR::UInt8, num_quarters::Integer=4, qoffset::Integer=0)

Return RAW format dot product of two bytes in BED format over quarters
in `(qoffset + 1):(qoffset + num_quarters)`. Missing values are
treated as zero.

"""
function bytedot(byteL::UInt8, byteR::UInt8, num_quarters::Integer, qoffset::Integer=0)
    @boundscheck qoffset + num_quarters <= 4

    bytedot(maskbyte(byteL, num_quarters, qoffset), byteR)
end

function bytedot{T}(byte::UInt8, v::AbstractArray{T}, voffset::Integer=0)
    @boundscheck checkbounds(v, voffset + 4)

    dotsum = zero(T)
    quarters = breakbyte(zeroNAs(byte))

    for j in 1:4
        @inbounds dotsum += quarters[j]*v[j + voffset]
    end

    dotsum
end

"""
    bytedot(byte::UInt8, v::AbstractArray, voffset::Integer=0, num_quarters::Integer=4, qoffset::Integer=0)

Return RAW format dot product of `byte[(qoffset + 1):(num_quarters +
qoffset)]` in BED format with `v[(voffset + 1):(voffset +
num_quarters)]`. Missing values (in `byte`) are treated as zero.

"""
function bytedot{T}(byte::UInt8, v::AbstractArray{T}, voffset::Integer, num_quarters::Integer, qoffset::Integer=0)
    @boundscheck begin
        num_quarters + qoffset <= 4 || error("too many quarters")
        checkbounds(v, voffset + num_quarters)
    end

    quarters = breakbyte(zeroNAs(byte))
    dotsum = zero(T)

    for j in 1:num_quarters
        @inbounds dotsum += v[j + voffset]*quarters[j + qoffset]
    end

    dotsum
end

"""
    byteNAsum(bedbyte::UInt8, v::AbstractArray, voffset::Integer=0, num_quarters::Integer=4, qoffset::Integer=0)

Return RAW format dot product of the NA support of
`byte[(qoffset + 1):(num_quarters + qoffset)]` in BED format with
`v[(voffset + 1):(voffset + num_quarters)]`. Missing values (in `bedbyte`) are
set to unity and all other entries to zero. That is to say, this is the sum of
elements of `v` for which corresponding quarters of `bedbyte` are missing.

"""
byteNAsum(bedbyte::UInt8, v::AbstractArray, voffset::Integer=0) = bytedot(NAsup(bedbyte), v, voffset)
function byteNAsum(bedbyte::UInt8, v::AbstractArray, voffset::Integer, num_quarters::Integer, qoffset::Integer=0)
    @boundscheck begin
        num_quarters + qoffset <= 4 || error("too many quarters")
        checkbounds(v, voffset + num_quarters)
    end

    @inbounds return bytedot(NAsup(bedbyte), v, voffset, num_quarters, qoffset)
end

function byteNAsum2{T}(bedbyte::UInt8, v::AbstractArray{T}, voffset::Integer=0)
    @boundscheck checkbounds(v, voffset + 4)

    dotsum = zero(T)
    quarters = breakbyte(NAsup(bedbyte))

    for j in 1:4
        @inbounds dotsum += quarters[j]*abs2(v[j + voffset])
    end

    dotsum
end

function byteNAsum2{T}(bedbyte::UInt8, v::AbstractArray{T}, voffset::Integer, num_quarters::Integer, qoffset::Integer=0)
    @boundscheck begin
        num_quarters + qoffset <= 4 || error("too many quarters")
        checkbounds(v, voffset + num_quarters)
    end

    quarters = breakbyte(NAsup(bedbyte))
    dotsum = zero(T)

    for j in 1:num_quarters
        @inbounds dotsum += quarters[j + qoffset]*abs2(v[j + voffset])
    end

    dotsum
end

function byteabs2(b::UInt8)
    @inbounds return Consts.bytebytemulttable[b + 1, b + 1]
end

bytesum(byte::UInt8) = bytedot(byte, onebyte)

"""
    bytesum(byte::UInt8, num_quarters::Integer=4, qoffset::Integer=0)

Returns sum of non-missing values in RAW format in `byte[(qoffset +
1):(num_quarters + qoffset)]`.

"""
function bytesum(byte::UInt8, num_quarters::Integer, qoffset::Integer=0)
    @boundscheck num_quarters + qoffset <= 4

    @inbounds return bytedot(byte, onebyte, num_quarters, qoffset)
end

function bytedist(bedbyte::UInt8)
    @inbounds return Consts.distributionmap[bedbyte + 1]
end

"""
    bytedist(bedbyte::UInt8, num_quarters::Integer=4, qoffset::Integer=0)

Returns number of (RAW format) `0x00`s, `0x01`s, `0x02`s, and missing
values in `bedbyte`'s `(1 + qoffset):(num_quarters + qoffset)`
quarters.

!!! warning
    The excluded quarters are included in the `0x00` count. Thus if
    `num_quarters < 4`, one likely wants to use
    `bytedist(bedbyte)[1] - 4 + num_quarters` instead of `bytedist(bedbyte)[1]`.

"""
function bytedist(bedbyte::UInt8, num_quarters::Integer, qoffset::Integer=0)
    @boundscheck num_quarters + qoffset <= 4
    @inbounds dist = Consts.distributionmap[maskbyte(bedbyte, num_quarters, qoffset) + 1]
    return dist
end

"""
    quarterfuncvector{T}(func::Function, ::Type{T}=Int, func_na=0x0)

Convert `func` with return type `T` into a length-4 `Tuple`
corresponding to how `func` acts on the RAW representation of the BED
quarters `(0b00, 0b01, 0b10, 0b11)` converted to type `T`. `func_na`
should be the desired behavior on missing entries, defaults to
`zero(T)`. This fully specifies the behavior of `func` on a
`BEDMatrix{T, Q}`.

"""
quarterfuncvector{T}(func::Function, ::Type{T}=Int, func_na=0x00) = ntuple(q -> q < 4 ? func(convert(T, q - 1)) : convert(T, func_na), 4)

#################### SubArray Interface ####################

Column{T, K, B} = SubArray{T, 1, K, Tuple{Base.Slice{Base.OneTo{Int}}, Int}, B}
BEDColumn{T, K <: BEDMatrix} = Column{T, K, false}

ColumnUnitRange{T, K, R <: Union{AbstractUnitRange, Base.Slice{Base.OneTo{Int}}}, B} = SubArray{T, 1, K, Tuple{R, Int}, B}
BEDColumnUnitRange{T, K <: BEDMatrix, R <: Union{AbstractUnitRange, Base.Slice{Base.OneTo{Int}}}} = ColumnUnitRange{T, K, R, false}

SubColumn{T, K, J, B} = SubArray{T, 1, K, Tuple{J, Int}, B}
BEDSubColumn{T, K <: BEDMatrix, J, B} = SubArray{T, 1, K, Tuple{J, Int}, B}

# Strictly two dimensional view of BEDMatrix
BEDSubMatrix{T, K <: BEDMatrix, R} = SubArray{T, 2, K, R, false}

"""
    hasNAs(B)

Returns `true` if there are missing values in `B`.

"""
function hasNAs(B::BEDMatrix)
    for col in indices(B, 2)
        column_hasNAs(B, col) && return true
    end

    return false
end

function hasNAs(B::BEDSubMatrix)
    for col in indices(B, 2)
        hasNAs(view(B, :, col)) && return true
    end

    return false
end

hasNAs(BC::BEDColumn) = column_hasNAs(parent(BC), parentindexes(BC)[2])

function hasNAs(BC::BEDSubColumn)
    rows, col = parentindexes(BC)
    column_hasNAs(parent(BC), col, rows)
end

"""
    countNAs(B)

Returns the number of missing values in `B`. `B` may be a `BEDMatrix`
or a single column `SubArray` of a `BEDMatrix`.

"""
function countNAs(B::BEDMatrix)
    total = 0

    # parallizable
    for col in indices(B, 2)
        total += column_countNAs(B, col)
    end
    total
end

countNAs(BC::BEDColumn) = column_countNAs(parent(BC), parentindexes(BC)[2])

function countNAs(BC::BEDSubColumn)
    B = parent(BC)
    rows, col = parentindexes(BC)

    @inbounds return column_dist(B, col, rows)[4]
end

Base.sum(BC::BEDColumn) = column_sum(parent(BC), parentindexes(BC)[2])
function Base.sum(BC::BEDSubColumn)
    rows, col = parentindexes(BC)
    column_sum(parent(BC), col, rows)
end

"""
    sum(func::Function, BC::BEDSubColumn; skipna=true)

Applies `func` element-wise to `BC` and returns the sum. If one sets
`skipna=false`, then `func(NArep(parent(BC)))` is used for missing
values; otherwise missing values are not included.

"""
function Base.sum(func::Function, BC::BEDSubColumn; skipna=true)
    B = parent(BC)
    rows, col = parentindexes(BC)

    if skipna
        return column_sum(func, B, col, rows)
    end

    func_na = func(NArep(B))
    column_sum(func, func_na, B, col, rows)
end

# RAW format is positive:
Base.sumabs(BC::BEDSubColumn) = sum(BC)

Base.sumabs2(BC::BEDColumn) = column_sumabs2(parent(BC), parentindexes(BC)[2])
function Base.sumabs2(BC::BEDSubColumn)
    rows, col = parentindexes(BC)
    column_sumabs2(parent(BC), col, rows)
end

"""
    BEDdot(BCL::BEDColumn, BCR::BEDColumn)

Computes the dot product of two BEDMatrix columns, skipping missing
values.

"""
function BEDdot(BCL::BEDColumn, BCR::BEDColumn)
    BL = parent(BCL)
    BR = parent(BCR)
    if BL === BR
        return column_dot(BL, parentindexes(BCL)[2], parentindexes(BCR)[2])
    end

    column_dot(BL, parentindexes(BCL)[2], BR, parentindexes(BCR)[2])
end

"""
    BEDdot(BC::BEDColumn, v::Array)
    BEDdot(v::Array, BC::BEDColumn)

Computes the dot product of a BEDMatrix column `BC` with `v`, skipping
missing values.

"""
BEDdot(BC::BEDColumn, v::AbstractArray) = column_dot(parent(BC), parentindexes(BC)[2], v)
BEDdot(v::AbstractArray, BC::BEDColumn) = BEDdot(BC, v)

# ? Base.mapreduce
# Base.norm

#################### Column tools ####################

# ? column_any
# ? column_all

# Roughly equivalent to StatsBase.countmap
function column_dist(B::BEDMatrix, col::Integer, row::Integer)
    @boundscheck checkbounds(B, row, col)
    byterow, snpind = rowtobytequarter(row)
    @inbounds q = breakbyte(B.X[byterow, col])[snpind]

    if q === 0x00
        return (1, 0, 0, 0)
    elseif q === 0x01
        return (0, 1, 0, 0)
    elseif q === 0x02
        return (0, 0, 1, 0)
    end

    return (0, 0, 0, 1)
end

function column_dist(B::BEDMatrix, col::Integer, indices::AbstractVector{Bool})
    zero_count = 0
    one_count = 0
    two_count = 0
    na_count = 0
    X = B.X
    @inbounds for (row, b) in enumerate(indices)
        if b
            byterow, snpind = rowtobytequarter(row)
            q = breakbyte(X[byterow, col])[snpind]
            if q === 0x00
                zero_count += 1
            elseif q === 0x01
                one_count += 1
            elseif q === 0x02
                two_count += 1
            else
                na_count += 1
            end
        end
    end

    return (zero_count, one_count, two_count, na_count)
end

function column_dist{T <: Integer}(B::BEDMatrix, col::Integer, inds::AbstractVector{T})
    zero_count = 0
    one_count = 0
    two_count = 0
    na_count = 0
    X = B.X
    @inbounds for row in inds
        byterow, snpind = rowtobytequarter(row)
        q = breakbyte(X[byterow, col])[snpind]
        if q === 0x00
            zero_count += 1
        elseif q === 0x01
            one_count += 1
        elseif q === 0x02
            two_count += 1
        else
            na_count += 1
        end
    end

    return (zero_count, one_count, two_count, na_count)
end

"""
    column_dist(B::BEDMatrix, col::Integer, rows=(:))

Returns a length-4 `Tuple` of the number of `0`s, `1`s, `2`s, and
`NArep(B)`s in `B[rows, col]`. `rows` may be a `UnitRange`, `Vector`
of indices, or a logical index (`Vector{Bool}`).

"""
function column_dist(B::BEDMatrix, col::Integer, rrange::AbstractUnitRange)
    bytestart, quarterstart = rowtobytequarter(first(rrange))
    bytestop, quarterstop = rowtobytequarter(last(rrange))

    unsafe_column_dist(B.X, col, bytestart, quarterstart, bytestop, quarterstop)
end

column_dist(B::BEDMatrix, col::Integer, ::Colon) = column_dist(B, col)
column_dist(B::BEDMatrix, col::Integer) = unsafe_column_dist(B.X, col, 1, 1, B._byteheight, B._lastrowSNPheight)

function unsafe_column_dist(X::Matrix{UInt8}, col::Integer, bytestart::Integer, quarterstart::Integer, bytestop::Integer, quarterstop::Integer)
    zero_count = 0
    one_count = 0
    two_count = 0
    na_count = 0

    @inbounds begin
        if quarterstart > 1
            num_quarters = ifelse(bytestop == bytestart, quarterstop - quarterstart + 1, 5 - quarterstart)

            bdist = bytedist(X[bytestart, col], num_quarters, quarterstart - 1)
            zero_count += (bdist[1] - 4 + num_quarters)  # correction for zeroed quarters
            one_count += bdist[2]
            two_count += bdist[3]
            na_count += bdist[4]

            bytestart += 1
        end

        if bytestart <= bytestop
            if quarterstop < 4
                bdist = bytedist(X[bytestop, col], quarterstop)
                zero_count += (bdist[1] - 4 + quarterstop)  # correction for zeroed quarters
                one_count += bdist[2]
                two_count += bdist[3]
                na_count += bdist[4]

                bytestop -= 1
            end

            @simd for x in bytestart:bytestop
                bdist = bytedist(X[x, col])
                zero_count += bdist[1]
                one_count += bdist[2]
                two_count += bdist[3]
                na_count += bdist[4]
            end
        end
    end

    return zero_count, one_count, two_count, na_count
end

"""
    column_sum{T, S}(B::BEDMatrix{T, S}, col::Integer, rows=(:))

Returns the sum of non-missing entries in `B[rows, col]`, optimized
for large `UnitRange`s of rows.

"""
function column_sum{T, S}(B::BEDMatrix{T, S}, col::Integer, rows=(:))
    counts = column_dist(B, col, rows)
    @inbounds return convert(promote_type(T, Int), counts[2] + 2*counts[3])
end

"""
    column_sumabs2{T, S}(B::BEDMatrix{T, S}, col::Integer, rows=(:))

Returns the sum of the squares of non-missing entries in `B[rows, col]`.

"""
function column_sumabs2{T, S}(B::BEDMatrix{T, S}, col::Integer, rows=(:))
    counts = column_dist(B, col, rows)
    @inbounds return convert(promote_type(T, Int), counts[2] + 4*counts[3])
end

"""
    column_norm(B::BEDMatrix, col::Integer, rows=(:), p=2)

Computes the `p`-norm of `B[rows, col]`, where the `p`-norm of a
vector `v` is given by ``\LaTeX``

```math
\\|v\\|_p = \\big(\sum_j |v_j|^p\\big)^\\frac{1}{p}
```

Missing values are ignored; use with care.

"""
function column_norm(B::BEDMatrix, col::Integer, rows=(:), p=2)
    if p == 2
        return sqrt(column_sumabs2(B, col, rows))
    end

    if p == 3
        return cbrt(column_sum((0, 1, 8, 0), B, col, rows))
    end

    @inbounds return (column_sum((0, 1, 2^p, 0), B, col, rows))^(1.0/p)
end

"""
    column_sum(func::Function, B::BEDMatrix, col::Integer, rows=(:))

Apply `func` to `B[:, col]`, and return the sum. `NArep(B)` values are
skipped.

"""
function column_sum{T, S}(f::Function, B::BEDMatrix{T, S}, col::Integer, rows=(:))
    column_sum(quarterfuncvector(f, T), B, col, rows)
end

"""
    column_sum(func::Function, func_na, B::BEDMatrix, col::Integer, rows)

Apply `func` to `B[rows, col]` and return the sum. `func_na` is
treated as the value of `func(NArep(B))`. (If omitted it is `0`.)

"""
function column_sum{T, S}(f::Function, func_na, B::BEDMatrix{T, S}, col::Integer, rows=(:))
    column_sum(quarterfuncvector(f, T, func_na), B, col, rows)
end

"""
    column_sum(func_qvec::Tuple, B::BEDMatrix, col::Integer, rows=(:))

`func_qvec` is a length-4 `Tuple` representation of a function `f`:
`func_qvec = (f(0), f(1), f(2), f(NArep(B)))`.  The result of applying
`f` to `B[rows, col]` and computing the sum is returned.

"""
function column_sum(func_qvec::Tuple, B::BEDMatrix, col::Integer, rows=(:))
    counts = column_dist(B, col, rows)
    @inbounds return tuple_dot(func_qvec, counts)
end

"""
    tuple_dot(t1::Tuple, t2::Tuple)

Compute scalar dot product of two `Tuple`s, treating them like `Vector`s.

"""
function tuple_dot(t1::Tuple, t2::Tuple)
    @boundscheck length(t1) == length(t2) || error("Tuples must have identical length")

    acc = 0
    for j in 1:length(t1)
        @inbounds acc += t1[j]*t2[j]
    end
    acc
end

"""
    column_hasNAs(B::BEDMatrix, col::Integer, rows=(:))

Returns `true` if `B[rows, col]` has missing values.

"""
function column_hasNAs(B::BEDMatrix, col::Integer, row::Integer)
    @inbounds return B[row, col] === NArep(B)
end

function column_hasNAs(B::BEDMatrix, col::Integer, rrange::AbstractUnitRange)
    X = B.X
    nbytes = B._byteheight
    na = NArep(B)
    rowstart, rowend = first(rrange), last(rrange)

    @inbounds begin
        firstquarters, bytestart, qstart = getquarterblock(B, rowstart, col)
        lastquarters, bytestop, qstop = getquarterblock(B, rowend, col)

        if qstart > 1
            stop = bytestart == bytestop ? qstop : 4
            for j in qstart:stop
                firstquarters[j] === na && return true
            end
            bytestart += 1
        end

        if bytestart <= bytestop
            if qstop < 4
                for j in 1:qstop
                    lastquarters[j] === na && return true
                end
                bytestop -= 1
            end

            for r in bytestart:bytestop
                hasNAs(X[r, col]) && return true
            end
        end
    end

    return false
end

column_hasNAs(B::BEDMatrix, col::Integer, ::Colon) = column_hasNAs(B, col)

function column_hasNAs(B::BEDMatrix, col::Integer)
    X = B.X
    nbytes = B._byteheight

    for x in 1:nbytes
        @inbounds hasNAs(X[x, col]) && return true
    end

    return false
end

# Fall-back method
function column_hasNAs(B::BEDMatrix, col::Integer, rows)
    na = NArep(B)

    @inbounds for row in rows
        B[row, col] === na && return true
    end

    return false
end

"""
    column_countNAs(B::BEDMatrix{T, S}, col::Integer, rows=(:))

Returns number of NAs in `B[:, col]`.

"""
function column_countNAs(B::BEDMatrix, col::Integer)
    nacount = 0

    X = B.X
    nbytes = B._byteheight
    @inbounds begin
        for x in 1:(nbytes - 1)
            nacount += countNAs(X[x, col])
        end
        nacount += countNAs(maskbyte(X[nbytes, col], B._lastrowSNPheight))
    end

    nacount
end

column_countNAs(B::BEDMatrix, col::Integer, ::Colon) = column_countNAs(B, col)

# Fall-back method
function column_countNAs(B::BEDMatrix, col::Integer, rows)
    @inbounds return column_dist(B, col, rows)[4]
end

"""
    column_dot{T, S}(B::BEDMatrix{T, S}, col::Integer, v::AbstractVector)

Computes the scalar dot product between `B[:, col]` and `v`, with
missing values in `B` treated as zero.

"""
function column_dot{T, S}(B::BEDMatrix{T, S}, col::Integer, v::AbstractArray)
    @boundscheck B.n == length(v) || throw(DimensionMismatch("v has incorrect length"))

    dotsum = zero(promote_type(T, eltype(v), Int))  # Avoid overflow when T == UInt8

    X = B.X
    nbytes = B._byteheight

    for x in 1:(nbytes - 1)
        @inbounds dotsum += bytedot(X[x, col], v, 4*(x-1))
    end
    @inbounds dotsum += bytedot(X[nbytes, col], v, 4*(nbytes-1), B._lastrowSNPheight)

    dotsum
end

"""
    column_dot{T, S}(B::BEDMatrix{T, S}, colL::Integer, colR::Integer)

Computes the scalar dot product between `B[:, colL]` and `B[:, colR]`,
treating missing values as zero.

"""
function column_dot{T, S}(B::BEDMatrix{T, S}, colL::Integer, colR::Integer)
    @boundscheck begin
        checkbounds(B, :, colL)
        checkbounds(B, :, colR)
    end

    dotsum = zero(promote_type(T, Int))
    X = B.X
    nbytes = B._byteheight

    for x in 1:(nbytes - 1)
        @inbounds dotsum += bytedot(X[x, colL], X[x, colR])
    end
    @inbounds dotsum += bytedot(X[nbytes, colL], X[nbytes, colR], B._lastrowSNPheight)

    dotsum
end

"""
    column_dot{T, S, U, R}(BL::BEDMatrix{T, S}, colL::Integer, BR::BEDMatric{U, R}, colR::Integer)

Computes the scalar dot product between `BL[:, colL]` and
`BR[:, colR]`, treating missing values as zero.

"""
function column_dot{T, S, U, R}(BL::BEDMatrix{T, S}, colL::Integer, BR::BEDMatrix{U, R}, colR::Integer)
    @boundscheck begin
        checkbounds(BL, :, colL)
        checkbounds(BR, :, colR)
        BL.n == BR.n || throw(DimensionMismatch("BEDMatrices must have same height"))
    end

    dotsum = zero(promote_type(T, U, Int))
    XL = BL.X
    XR = BR.X
    nbytes = BL._byteheight

    for x in 1:(nbytes - 1)
        @inbounds dotsum += bytedot(XL[x, colL], XR[x, colR])
    end
    @inbounds dotsum += bytedot(XL[nbytes, colL], XR[nbytes, colR], BL._lastrowSNPheight)

    dotsum
end

function column_NAsup_dot{T, S}(B::BEDMatrix{T, S}, col::Integer, v::AbstractArray)
    @boundscheck B.n == length(v) || throw(DimensionMismatch("v has incorrect length"))

    X = B.X
    dotsum = zero(promote_type(T, eltype(v), Int))  # Avoid overflow when T == UInt8
    nbytes = B._byteheight

    @simd for x in 1:(nbytes - 1)
        @inbounds dotsum += byteNAsum(X[x, col], v, 4*(x-1))
    end
    @inbounds dotsum += byteNAsum(X[nbytes, col], v, 4*(nbytes - 1), B._lastrowSNPheight)

    dotsum
end

"""
    column_dist_dot{T, S}(B::BEDMatrix{T, S}, col::Integer, v::AbstractArray, rows=(:))

Find everything a body needs to do univariate regression in a single
pass over `B[:, col]`. Returns
`(zero_count, one_count, two_count, na_count, dotsum, nasum, na2sum)`.

"""
# this is broken!!
function column_dist_dot(B::BEDMatrix, col::Integer, v::AbstractArray, rrange::AbstractUnitRange)
    bytestart, quarterstart = rowtobytequarter(first(rrange))
    bytestop, quarterstop = rowtobytequarter(last(rrange))

    unsafe_column_dist_dot(B.X, col, bytestart, quarterstart, bytestop, quarterstop, v)
end

column_dist_dot(B::BEDMatrix, col::Integer, v::AbstractArray, ::Colon) = column_dist_dot(B, col, v)
column_dist_dot(B::BEDMatrix, col::Integer, v::AbstractArray) = unsafe_column_dist_dot(B.X, col, 1, 1, B._byteheight, B._lastrowSNPheight, v)

function column_dist_dot{T, S, R <: AbstractUnitRange}(B::BEDMatrix{T, S}, col::Integer, v::AbstractArray, ranges::Vector{R})
    nasum = na2sum = dotsum = zero(promote_type(T, eltype(v), Int))
    zero_count = one_count = two_count = na_count = 0

    for rrange in ranges
        counttup = column_dist_dot(B, col, v, rrange)
        zero_count += counttup[1]
        one_count += counttup[2]
        two_count += counttup[3]
        na_count += counttup[4]
        dotsum += counttup[5]
        nasum += counttup[6]
        na2sum += counttup[7]
    end

    return (zero_count, one_count, two_count, na_count, dotsum, nasum, na2sum)
end

function column_dist_dot{T, S, R <: Integer}(B::BEDMatrix{T, S}, col::Integer, v::AbstractArray, inds::Vector{R})
    nasum = na2sum = dotsum = zero(promote_type(T, eltype(v), Int))

    zero_count = 0
    one_count = 0
    two_count = 0
    na_count = 0

    @inbounds for row in inds
        x = B[row, col]
        if x === NArep(B)
            na_count += 1
            nasum += v[row]
            na2sum += abs2(v[row])
        else
            if x == 0
                zero_count += 1
            elseif x == 1
                one_count += 1
            elseif x == 2
                two_count += 1
            end
            dotsum += x*v[row]
        end
    end

    return (zero_count, one_count, two_count, na_count, dotsum, nasum, na2sum)
end

function unsafe_column_dist_dot(X::Matrix{UInt8}, col::Integer, bytestart::Integer, quarterstart::Integer, bytestop::Integer, quarterstop::Integer, v::AbstractArray)
    zero_count = 0
    one_count = 0
    two_count = 0
    na_count = 0
    nasum = na2sum = dotsum = zero(promote_type(eltype(v), Int))

    num_quarters = 0
    @inbounds begin
        # First byte
        if quarterstart > 1
            num_quarters = ifelse(bytestop == bytestart, quarterstop - quarterstart + 1, 5 - quarterstart)

            bdist = bytedist(X[bytestart, col], num_quarters, quarterstart - 1)
            zero_count += (bdist[1] - 4 + num_quarters)  # correction for zeroed quarters
            one_count += bdist[2]
            two_count += bdist[3]
            na_count += bdist[4]

            dotsum += bytedot(X[bytestart, col], v, 0, num_quarters, quarterstart - 1)
            nasum += byteNAsum(X[bytestart, col], v, 0, num_quarters, quarterstart - 1)
            na2sum += byteNAsum2(X[bytestart, col], v, 0, num_quarters, quarterstart - 1)

            bytestart += 1
        end

        if bytestart <= bytestop
            # Last byte
            if quarterstop < 4
                bdist = bytedist(X[bytestop, col], quarterstop)
                zero_count += (bdist[1] - 4 + quarterstop)  # correction for zeroed quarters
                one_count += bdist[2]
                two_count += bdist[3]
                na_count += bdist[4]

                voffset = num_quarters + 4*(bytestop - bytestart)
                dotsum += bytedot(X[bytestop, col], v, voffset, quarterstop)
                nasum += byteNAsum(X[bytestop, col], v, voffset, quarterstop)
                na2sum += byteNAsum2(X[bytestop, col], v, voffset, quarterstop)

                bytestop -= 1
            end

            # Main course
            @simd for x in bytestart:bytestop
                byte = X[x, col]
                bdist = bytedist(byte)
                zero_count += bdist[1]
                one_count += bdist[2]
                two_count += bdist[3]
                na_count += bdist[4]

                voffset = 4*(x-1) + num_quarters
                dotsum += bytedot(byte, v, voffset)
                nasum += byteNAsum(byte, v, voffset)
                na2sum += byteNAsum2(byte, v, voffset)
            end
        end
    end

    return zero_count, one_count, two_count, na_count, dotsum, nasum, na2sum
end

"""
    column_mean(B::BEDMatrix, col::Integer, rows=(:))

Returns the mean of `B[rows, col]`, skipping missing values.

"""
function column_mean(B::BEDMatrix, col::Integer, rows=(:))
    counts = column_dist(B, col, rows)

    @inbounds return (counts[2] + 2.0*counts[3])/(B.n - counts[4])
end

# name matches `StatsBase.mean_and_var`, but interface does not
"""
    column_mean_and_var(B::BEDMatrix, col::Integer, rows=(:))

Computes the mean and (Bessel-corrected) sample variance of
`B[rows, col]`, skipping missing values.

"""
function column_mean_and_var(B::BEDMatrix, col::Integer, rows=(:))
    counts = column_dist(B, col, rows)

    @inbounds μ = (counts[2] + 2.0*counts[3])/(B.n - counts[4])
    @inbounds σ2 = (counts[1]*abs2(μ) + counts[2]*abs2(1.0 - μ) + counts[3]*abs2(2.0 - μ))/(B.n - counts[4] - 1)

    return μ, σ2
end

# name matches with `StatsBase.mean_and_std`, but interface does not
"""
    column_mean_and_std(B::BEDMatrix, col::Integer, rows=(:))

Computes the mean and (Bessel-corrected) sample standard deviation of
`B[rows, col]`, skipping missing values.

"""
function column_mean_and_std(B::BEDMatrix, col::Integer, rows=(:))
    μ, σ2 = column_mean_and_var(B, col, rows)
    return μ, sqrt(σ2)
end


text:
export hindmarsh_rose

"""
Equations of motion for the Hindmarsh-Rose system.
"""
function eom_hindmarsh_rose(u, p, t)
    a, b, c, d, r, s, xᵣ, I = (p...,)
    x, y, z = (u...,)

	ϕ = -a*x^3 + b*x^2
	ψ = c - d*x^2
    dx = y + ϕ - z + I
	dy = ψ - y
	dz = r*(s*(x - xᵣ) - z)
    return SVector{3}(dx, dy, dz)
end

"""
	hindmarsh_rose(u₀, p)

Initialise a Hindmarsh-Rose system, which is a model of neuronal
spiking.

```math
\\begin{aligned}
\\dfrac{dx}{dt} &= y + \\phi(x) - z + I
\\dfrac{dy}{dt} &= \\psi(x) - y
\\dfrac{dz}{dt} &= r[s(x - x_R) - z],
\\end{aligned}
```
where

```math
\\begin{aligned}
\\phi(x) &= -ax^3+bx^2
\\psi(x) &= c - dx^2
\\end{aligned}
```

If parameters other than the defaults are to be used, they must be
provided as a vector `[a, b, c, d, r, s, xᵣ, I]`.
"""
function hindmarsh_rose(u₀, p)
    ContinuousDynamicalSystem(eom_hindmarsh_rose, u₀, p)
end
hindmarsh_rose(;u₀ = rand(3), a = 1, b = 3, c = 1, d = 5, r = 1e-3, s = 4, xᵣ = - 8/5, I = -8) =
    hindmarsh_rose(u₀, [a, b, c, d, r, s, xᵣ, I])


text:
module MaxHelpingHandSidewaysLava

using ..Ahorn, Maple

@mapdef Entity "MaxHelpingHand/SidewaysLava" SidewaysLava(x::Integer, y::Integer, intro::Bool=false, lavaMode::String="LeftToRight", speedMultiplier::Number=1.0)

const lavaModes = String["LeftToRight", "RightToLeft", "Sandwich"]

const placements = Ahorn.PlacementDict(
    "Sideways Lava (max480's Helping Hand)" => Ahorn.EntityPlacement(
        SidewaysLava
    )
)

Ahorn.editingOptions(entity::SidewaysLava) = Dict{String, Any}(
    "lavaMode" => lavaModes
)

function Ahorn.selection(entity::SidewaysLava)
    x, y = Ahorn.position(entity)

    return Ahorn.Rectangle(x - 12, y - 12, 24, 24)
end

function Ahorn.render(ctx::Ahorn.Cairo.CairoContext, entity::SidewaysLava, room::Maple.Room)
    Ahorn.Cairo.save(ctx)

    lavaMode = get(entity.data, "lavaMode", "LeftToRight")

    Ahorn.rotate(ctx, (lavaMode == "RightToLeft" ? -1 : 1) * pi / 2)
    Ahorn.drawImage(ctx, (lavaMode == "Sandwich" ? Ahorn.Assets.lavaSanwitch : Ahorn.Assets.risingLava), -12, -12)

    Ahorn.Cairo.restore(ctx);
end

end

max_stars_repo_name: Kitty-Cats/MaxHelpingHand

max_issues_repo_name: Kitty-Cats/MaxHelpingHand

max_forks_repo_name: Kitty-Cats/MaxHelpingHand


text:
function compute_log_return_table(data_table::DataFrame,  date_key::Symbol, data_key::Symbol)::DataFrame

    # get the size of the table -
    (number_of_rows, number_of_cols) = size(data_table)

    # initialize -
    return_array = Array{Float64,1}()
    date_array = Array{Date,1}()

    # compute -
    for row_index = 2:number_of_rows

        # setup -
        today_index = row_index
        yesterday_index = row_index - 1

        # compute the return -
        close_price_today = data_table[today_index,data_key]
        close_price_yesterday = data_table[yesterday_index,data_key]
        value = log(close_price_today/close_price_yesterday)
        push!(return_array, value)

        # grab te date -
        today_timestamp = data_table[today_index,date_key]
        push!(date_array, today_timestamp)
    end

    # build data frame and return -
    return_table = DataFrame(timestamp=date_array, log_return=return_array)
    return return_table
end

max_stars_repo_name: Pooksoft/BinomialEquityAssetPriceModel

max_issues_repo_name: Pooksoft/BinomialEquityAssetPriceModel

max_forks_repo_name: Pooksoft/BinomialEquityAssetPriceModel


text:
@testset "Generic Group" begin
    @testset "QuotientGroup" begin
        @test Hecke.quotient_indx(1,1) == [(1,1)]
        @test Hecke.quotient_indx(2,1) == [(1,1), (2,1)]
        @test Hecke.quotient_indx(3,1) == [(1,1), (3,1)]
        @test Hecke.quotient_indx(4,1) == [(1,1), (2,1), (4,1)]
        @test Hecke.quotient_indx(4,2) == [(1,1), (2,1), (2,1), (2,1), (4,2)]
        @test Hecke.quotient_indx(5,1) == [(1,1), (5,1)]
        @test Hecke.quotient_indx(6,1) == [(1,1), (2,1), (6,1)]
        @test Hecke.quotient_indx(6,2) == [(1,1), (2,1), (3,1), (6,2)]
        @test Hecke.quotient_indx(7,1) == [(1,1), (7,1)]
        @test Hecke.quotient_indx(8,1) == [(1,1), (2,1), (4,1), (8,1)]
        @test Hecke.quotient_indx(8,2) == [(1,1), (2,1), (2,1), (2,1), (4,1), (4,1), (4,2), (8,2)]
        @test Hecke.quotient_indx(8,3) == [(1,1), (2,1), (2,1), (2,1), (4,2), (8,3)]
        @test Hecke.quotient_indx(8,4) == [(1,1), (2,1), (2,1), (2,1), (4,2), (8,4)]
        @test Hecke.quotient_indx(8,5) == [(1,1), (2,1), (2,1), (2,1), (2,1), (2,1), (2,1), (2,1), (4,2), (4,2), (4,2), (4,2), (4,2), (4,2), (4,2), (8,5)]
        @test Hecke.quotient_indx(9,1) == [(1,1), (3,1), (9,1)]
        @test Hecke.quotient_indx(9,2) == [(1,1), (3,1), (3,1), (3,1), (3,1), (9,2)]
        @test Hecke.quotient_indx(10,1) == [(1,1), (2,1), (10,1)]
        @test Hecke.quotient_indx(10,2) == [(1,1), (2,1), (5,1), (10,2)]
        @test Hecke.quotient_indx(11,1) == [(1,1), (11,1)]
        @test Hecke.quotient_indx(12,1) == [(1,1), (2,1), (4,1), (6,1), (12,1)]
        @test Hecke.quotient_indx(22,1) == [(1,1), (2,1), (22,1)]
        @test Hecke.quotient_indx(32,15) == [(1,1), (2,1), (2,1), (2,1), (4,1), (4,1), (4,2), (8,2), (8,3), (8,4), (16,4), (32,15)]
        @test Hecke.quotient_indx(42,1) == [(1,1), (2,1), (3,1), (6,2), (42,1)]
        @test Hecke.quotient_indx(52,3) == [(1,1), (2,1), (4,1), (52,3)]
        @test Hecke.quotient_indx(60,5) == [(1,1), (60,5)]
    end

    @testset "derived series" begin
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(1,1))] == [(1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(8,3))] == [(8,3), (2,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(18,1))] == [(18,1), (9,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(24,5))] == [(24,5), (3,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(24,6))] == [(24,6), (6,2), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(24,10))] == [(24,10), (2,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(27,4))] == [(27,4), (3,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(36,10))] == [(36,10), (9,2), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(38,2))] == [(38,2), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(40,13))] == [(40,13), (5,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(43,1))] == [(43,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(46,1))] == [(46,1), (23,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(48,37))] == [(48,37), (6,2), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(51,1))] == [(51,1), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(54,14))] == [(54,14), (27,5), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(56,13))] == [(56,13), (1,1)]
        @test [tuple(Hecke.find_small_group(i[1])[1]...) for i in Hecke.derived_series(Hecke.small_group(59,1))] == [(59,1), (1,1)]
    end

    @testset "psylow" begin
        G,AtoG,GtoA = Hecke.generic_group([1, -1, im, -im], *)
        @test order(psylow_subgroup(G,2)[1]) == 4
        @test order(psylow_subgroup(G,3)[1]) == 1
        @test_throws ErrorException psylow_subgroup(G,10)
    end

    @testset "GrpGenToGrpAb" begin
        G,AtoG,GtoA = Hecke.generic_group([1, -1, im, -im], *)
        GrpAb, GtoGrpAb, GrpAbtoG = gen_2_ab(G)
        @test GrpAb.snf == ZZ.([4])
        @test order(GtoGrpAb(G[1])) == 1
        @test order(GtoGrpAb(G[2])) == 2
        @test order(GtoGrpAb(G[3])) == 4
        @test order(GtoGrpAb(G[4])) == 4

        @test gen_2_ab(small_group(4,2))[1].snf == ZZ.([2, 2])
        @test gen_2_ab(small_group(8,2))[1].snf == ZZ.([2, 4])
        @test gen_2_ab(small_group(58,2))[1].snf == ZZ.([58])
        @test gen_2_ab(small_group(56,13))[1].snf == ZZ.([2, 2, 14])
        @test gen_2_ab(small_group(60,13))[1].snf == ZZ.([2, 30])
        @test gen_2_ab(small_group(54,15))[1].snf == ZZ.([3, 3, 6])

    end
end


text:
#Signal Processing Functions

#Consts
AF_CONV_EXPAND = icxx"AF_CONV_EXPAND;"

# Fourier Transforms
import Base: fft, ifft
fft{T}(a::AFVector{T}) = AFArray{Complex{T}}(af_fft(a))
fft{T}(a::AFMatrix{T}) = AFArray{Complex{T}}(af_fft2(a))
fft{T}(a::AFAbstractArray{T,3}) = AFArray{Complex{T}}(af_fft3(a))

fft{T<:Complex}(a::AFVector{T}) = AFArray{T}(af_fft(a))
fft{T<:Complex}(a::AFMatrix{T}) = AFArray{T}(af_fft2(a))
fft{T<:Complex}(a::AFAbstractArray{T,3}) = AFArray{T}(af_fft3(a))

ifft{T}(a::AFVector{T}) = AFArray{Complex{T}}(af_ifft(a))
ifft{T}(a::AFMatrix{T}) = AFArray{Complex{T}}(af_ifft2(a))
ifft{T}(a::AFAbstractArray{T,3}) = AFArray{Complex{T}}(af_ifft3(a))

ifft{T<:Complex}(a::AFVector{T}) = AFArray{T}(af_ifft(a))
ifft{T<:Complex}(a::AFMatrix{T}) = AFArray{T}(af_ifft2(a))
ifft{T<:Complex}(a::AFAbstractArray{T,3}) = AFArray{T}(af_ifft3(a))

#FIR Filter
fir{T}(a::AFAbstractArray{T}, x::AFAbstractArray{T}) = AFArray{T}(af_fir(a,x))

#IIR Filter
iir{T}(ff::AFAbstractArray{T}, fb::AFAbstractArray{T}, a::AFAbstractArray{T}) = AFArray{T}(af_iir(ff, fb, a))

#Convolutions

import Base: conv, conv2

conv{T}(a::AFAbstractArray{T}, b::AFAbstractArray{T}) = AFArray{T}(af_fftConvolve(a, b, AF_CONV_EXPAND))
conv2{T}(a::AFAbstractArray{T}, b::AFAbstractArray{T}) = AFArray{T}(af_fftConvolve2(a,b, AF_CONV_EXPAND))