text:
__precompile__()
module DataTableaux

export Small, IntegerSet, DataTableau, row, head, countmap, FrameToTableauScheme

# hoping I can dump this:
# export push!, in, round, show, showall, Float64

# currently commenting out next because of anticipated confilicts with Koala
# export fit!, transform, fit_transform!, inverse_transform

using DataFrames
import CSV

# to be extended:
import StatsBase.countmap
import DataFrames.head
import Base: getindex, show, showall, convert, length
import Base: size, start, next, done, issparse, ndims, copy
import Base: in, push!, Float64, round

## Constants:
const Small=UInt8
const Big=UInt64
const SMALL_MAX = Small(52)
const BIG_MAX = Big(2^(Int(SMALL_MAX) + 1) - 1)
const TWO = Big(2)
const ZERO = Big(0)
const ONE = Small(1)
const srcdir = dirname(@__FILE__)

## Helpers:

tail(n) = "..."*string(n)[end-3:end]

"""Load a reduced version of the public Ames Housing dataset."""
load_reduced_ames() =
    CSV.read(joinpath(srcdir, "data", "reduced_ames.csv"))

function second_less_than(pair1, pair2)
    return pair1[2]<pair2[2]
end

function keys_ordered_by_values{T, S<:Real}(d::Dict{T,S})
    len = length(d)
    items = collect(d) # 1d array containing the (key, value) pairs
    sort!(items, lt=second_less_than)
    return T[pair[1] for pair in items]
end

function mode(v)
    d = countmap(v)
    return keys_ordered_by_values(d)[end]
end

"""
## mutable struct `IntegerSet`

A type of collection for storing subsets of {0, 1, 2, ... 52}. Every
such subset can be stored as an Float64 object. To convert an
IntegerSet object `s` to a floating point number, use Float64(s). To
recover the original object from a float `f`, use `round(IntegerSet,
f)`.

To instantiate an empty collection use, `IntegerSet()`. To add an
element `i::Integer` use `push!(s, i)` which is quickest if `i` is
type `Small=UInt8`. Membership is tested as usual with `in`. One can also
instantiate an object with multiple elements as in the following example:

    julia> 15 in IntegerSet([1, 24, 16])
    false

"""
mutable struct IntegerSet
    coded::Big
end

function IntegerSet(v::AbstractVector{T} where T <: Integer)
    s = IntegerSet()
    for k in v
        push!(s, k) # push! defined below
    end
    return s
end

IntegerSet() = IntegerSet(ZERO)

function Base.in(k::Small, s::IntegerSet)
    large = TWO^k
    return large & s.coded == large
end

function Base.in(k::Integer, s::IntegerSet)
    large = TWO^Small(k)
    return large & s.coded == large
end

function Base.push!(s::IntegerSet, k::Small)
    if k > SMALL_MAX
        throw(Base.error("Cannot push! an integer larger
           than $(Int(SMALL_MAX)) into an IntegerSet object."))
    end
    if !(k in s)
        s.coded = s.coded | TWO^k
    end
    return s
end

function Base.push!(s::IntegerSet, k::Integer)
    if k > SMALL_MAX
        throw(Base.error("Cannot push! an integer larger
           than $(Int(SMALL_MAX)) into an IntegerSet object."))
    end
    push!(s, Small(k))
end

function Base.show(stream::IO, s::IntegerSet)
    for i in 0:62
        if i in s
            print(stream, "$i, ")
        end
    end
end

Base.Float64(s::IntegerSet) = Float64(s.coded)

function Base.round(T::Type{IntegerSet}, f::Float64)
    if f < 0 || f > BIG_MAX
        throw(Base.error("Float64 numbers outside the
           range [0,BIG_MAX] cannot be rounded to IntegerSet values"))
    end
    return IntegerSet(round(Big, f))
end

"""
Scheme for encoding a mapping from a set of strings (not exceeding 52
in number) to a corresponding number of consecutive SmallFloats,
beginning at 0. Here "SmallFloat" refers to the `Float64`
representation of a `Small` integer.

"""
struct StringToSmallFloatScheme
    small_given_string::Dict{String, Small}
    string_given_small::Dict{Small, String}
    mode::String

    function StringToSmallFloatScheme(v::Vector{String}; overide=false)
        small_given_string = Dict{String, Small}()
        string_given_small = Dict{Small, String}()
        vals = sort(collect(Set(v)))
        if length(vals) > SMALL_MAX + 1 && !overide
            throw(Base.error("Trying to construct a StringToSmallFloatScheme with a vector
                             having more than $(SMALL_MAX + 1) values."))
        end
        i = Small(0)
        for c in vals
            small_given_string[c] = i
            string_given_small[i] = c
            i = i + ONE
        end
        return new(small_given_string, string_given_small, mode(v))
    end
end

const VoidScheme=StringToSmallFloatScheme([""])

function show(stream::IO, s::StringToSmallFloatScheme)
    if s == VoidScheme
        print(stream, "VoidScheme")
    else
        print(stream, "StringToSmallFloatScheme@$(tail(hash(s)))")
    end
end

function showall(stream::IO, s::StringToSmallFloatScheme)
    if s == VoidScheme
        show(stream, s)
    else
        show(stream, s)
        str = "\n"
        for k in sort(collect(keys(s.small_given_string)))
            str = string(str, k, " => ", Float64(s.small_given_string[k]), "\n")
        end
        print(stream, str)
    end
end

"""
## `function transform(s::StringToSmallFloatScheme, c::String)`

If the string input `c` does not appear in the dictionary of scheme
`s` then `s.mode` is returned, and a warning issued.

"""
function transform(s::StringToSmallFloatScheme, c::String)
    small = s.mode
    try
        small = s.small_given_string[c]
    catch
        warn("String not in `StringToSmallFloatScheme` encountered in `transform`. "*
             "Using mode.")
        small = s.small_given_string[s.mode]
    end
    return Float64(small)
end

transform(s::StringToSmallFloatScheme, v::Vector{String}) =
    Float64[transform(s, c) for c in v]
inverse_transform(s::StringToSmallFloatScheme, f::Float64) =
    s.string_given_small[round(Small, f)]
inverse_transform(s::StringToSmallFloatScheme, v::Vector{Float64}) =
    String[inverse_transform(s, f) for f in v]

## The `DataTableau` type

struct DataTableauEncoding

    schemes::Vector{StringToSmallFloatScheme}
    is_ordinal::Vector{Bool}
    names::Vector{Symbol}

end

Base.copy(s::DataTableauEncoding) =
    DataTableauEncoding(copy(s.schemes), copy(s.is_ordinal), copy(s.names))

struct DataTableau

    raw::Array{Float64,2}
    names::Vector{Symbol}
    nrows::Int
    ncols::Int
    encoding::DataTableauEncoding

end

copy(X::DataTableau) = DataTableau(copy(X.raw), copy(X.names), X.nrows, X.ncols, copy(X.encoding))

# Before defining `DataTableau` constructors, we need schemes for
# transforming dataframes into datatables:

mutable struct FrameToTableauScheme

    # postfit:
    encoding::DataTableauEncoding

    fitted::Bool

    function FrameToTableauScheme()
        ret = new()
        ret.fitted = false
        return ret
    end

end

function show(stream::IO, s::FrameToTableauScheme)
    print(stream, "FrameToTableauScheme@$(tail(hash(s)))")
end

# Note: we later define `function FrameToTableauScheme(X::DataTableau)`

function fit!(scheme::FrameToTableauScheme, df::AbstractDataFrame)
    ncols = size(df, 2)
    schemes = Array{StringToSmallFloatScheme}(ncols)
    is_ordinal = Array{Bool}(ncols)
    for j in 1:ncols
        col = [df[j]...] # shrink eltype
        column_type = eltype(col)
        if column_type <: AbstractFloat
            is_ordinal[j] = true
            schemes[j] = VoidScheme
        else
            is_ordinal[j] = false
            col = sort([string(s) for s in col])
            schemes[j] = StringToSmallFloatScheme(col)
        end
    end
    scheme.encoding = DataTableauEncoding(schemes, is_ordinal, names(df))
    scheme.fitted = true
    return scheme
end

function FrameToTableauScheme(df::AbstractDataFrame)
    s = FrameToTableauScheme()
    fit!(s, df)
    return s
end

function transform(scheme::FrameToTableauScheme, df::AbstractDataFrame)

    if !scheme.fitted
        warn("Attempting to transform DataFrame using unfitted FrameToTableauScheme. " *
             "Calling `fit_transform!` instead of `transform`.")
        return fit_transform!(scheme, df)
    end

    encoding = scheme.encoding

    encoding.names == names(df) || error("Attempt to transform AbstractDataFrame "*
                                         "object into DataTableau object using "*
                                         " incompatible encoding.")
    nrows, ncols = size(df)
    raw    = Array{Float64}((nrows, ncols))

    for j in 1:ncols
        if encoding.is_ordinal[j]
            for i in 1:nrows
                raw[i,j] = Float64(df[i,j])
            end
        else
            for i in 1:nrows
                raw[i,j] = transform(encoding.schemes[j], string(df[i,j]))
            end
        end
    end

    return DataTableau(raw, encoding.names, nrows, ncols, encoding)

end

function fit_transform!(s::FrameToTableauScheme, df::AbstractDataFrame)
    fit!(s, df)
    return transform(s, df)
end

function convert(T::Type{DataFrame}, dt::DataTableau)
    return DataFrame(columns(dt), dt.names) # columns(dt) defined later
end

function inverse_transform(s::FrameToTableauScheme, df::DataTableau)
    return convert(DataFrame, df)
end

# `DataTableau constructors:`

DataTableau(df::AbstractDataFrame) = transform(FrameToTableauScheme(df), df)

function DataTableau(;columns::Vector=Any[], names::Vector{Symbol}=Symbol[])

    ncols = length(columns)

    if ncols == 0
        throw(Base.error("Error constructing DataTableau object.
                         It must have at least one column."))
    end

    if isempty(names)
        names = Vector{Symbol}(ncols)
        for i in 1:ncols
            names[i] = Symbol(string("x", i))
        end
    end

    if length(names) != ncols
        throw(Base.error("You must supply one column name per column."))
    end

    nrows = length(columns[1])
    if sum([length(v)!=nrows for v in columns]) != 0
        throw(Base.error("Error constructing DataTableau object.
                         All columns must have same length."))
    end
    df = DataFrame(columns, names)
    return DataTableau(df)
end

########################################################
# function DataTableau(column_tuple...)                #
#     n_cols = length(column_tuple)                    #
#     cols = collect(column_tuple)                     #
#     colnames = Vector{Symbol}(n_cols)                #
#     for i in 1:n_cols                                #
#         colnames[i]=Symbol(string("x",i))            #
#     end                                              #
#     return DataTableau(columns=cols, names=colnames) #
# end                                                  #
########################################################

function FrameToTableauScheme(X::DataTableau)
    s = FrameToTableauScheme()
    s.encoding = X.encoding
    s.fitted = true
    return s
end

function columns(dt::DataTableau)
    cols = Any[]
    for j in 1:dt.ncols
        rawcol = Float64[dt.raw[i,j] for i in 1:dt.nrows]
        if dt.encoding.is_ordinal[j]
            push!(cols, rawcol)
        else
            push!(cols, inverse_transform(dt.encoding.schemes[j], rawcol))
        end
    end
    return cols
end

function show(stream::IO, dt::DataTableau; nrows=0)
    if nrows == 0
        nrows = dt.nrows
    end
    ncols = dt.ncols
    types = Array{String}(ncols)
    for j in 1:ncols
        if dt.encoding.is_ordinal[j]
            types[j] = "ord"
        else
            types[j] = "cat"
        end
    end
    header = [Symbol(string(dt.names[j], " ($j,$(types[j]))")) for j in 1:ncols]
    println("(Displaying DataTableau as DataFrame)")
    show(stream, DataFrame(columns(dt), header)[1:nrows,:])
    println()
end

head(dt::DataTableau) =  show(STDOUT, dt; nrows=min(4, dt.nrows))

getindex(dt::DataTableau, i::Int, j::Int) = dt.raw[i,j]
getindex(dt::DataTableau, j::Int) = dt.raw[:,j]

function getindex(dt::DataTableau, col_name::Symbol)
    j = 0
    for k in eachindex(dt.names)
        if dt.names[k] == col_name
            j = k
        end
    end
    if j == 0
        throw(DomainError)
    end
    return dt[j]
end

function getindex(dt::DataTableau, bs::Vector{Symbol})
    index_given_name = Dict{Symbol,Int}()
    for j in eachindex(dt.names)
        index_given_name[dt.names[j]] = j
    end
    b = [index_given_name[sym] for sym in bs]
    raw = dt.raw[:,b]
    col_names = dt.names[b]
    nrows = dt.nrows
    ncols = length(b)
    encoding = DataTableauEncoding(dt.encoding.schemes[b], dt.encoding.is_ordinal[b], bs)
    return DataTableau(raw, col_names, nrows, ncols, encoding)
end

# function getindex(dt::DataTableau, c::Colon, j::Int)
#     raw = dt.raw[:,j:j] # j:j forces two-dimensionality of the array
#     col_names = dt.names[j:j]
#     nrows = dt.nrows
#     ncols = 1
#     encoding = DataTableauEncoding()
#     encoding.schemes = dt.encoding.schemes[j:j]
#     encoding.is_ordinal = dt.encoding.is_ordinal[j:j]
#     return DataTableau(raw, col_names, nrows, ncols, encoding)
# end

getindex(dt::DataTableau, c::Colon, j::Int) = dt.raw[:,j]
getindex(dt::DataTableau, i::Int, c::Colon) = dt.raw[i,:]
row(dt::DataTableau, i::Int) = dt.raw[i,:]

function getindex(dt::DataTableau, a::Vector{Int}, c::Colon)
    raw = dt.raw[a,:]
    col_names = dt.names
    nrows = length(a)
    ncols = dt.ncols
    encoding = dt.encoding
    return DataTableau(raw, col_names, nrows, ncols, encoding)
end

getindex(dt::DataTableau, a::UnitRange{Int}, c::Colon) = dt[collect(a),:]

function getindex(dt::DataTableau, c::Colon, b::Vector{Int})
    raw = dt.raw[:,b]
    col_names = dt.names[b]
    nrows = dt.nrows
    ncols = length(b)
    encoding = DataTableauEncoding(dt.encoding.schemes[b], dt.encoding.is_ordinal[b], col_names)
    return DataTableau(raw, col_names, nrows, ncols, encoding)
end

getindex(dt::DataTableau, c::Colon, b::UnitRange{Int}) = dt[:,collect(b)]

function getindex(dt::DataTableau, a::Vector{Int}, b::Vector{Int})
    raw = dt.raw[a,b]
    col_names = dt.names[b]
    nrows = length(a)
    ncols = length(b)
    encoding = DataTableauEncoding(dt.encoding.schemes[b], dt.encoding.is_ordinal[b], col_names)
    return DataTableau(raw, col_names, nrows, ncols, encoding)
end
getindex(dt::DataTableau, a::UnitRange{Int}, b::UnitRange{Int}) = dt[collect(a),collect(b)]
getindex(dt::DataTableau, a::UnitRange{Int}, b::Vector{Int}) = dt[collect(a),b]
getindex(dt::DataTableau, a::Vector{Int}, b::UnitRange{Int}) = dt[a,collect(b)]

length(dt::DataTableau) = dt.ncols
issparse(dt::DataTableau) = false
ndims(dt::DataTableau) = 2
size(dt::DataTableau) = (dt.nrows, dt.ncols)
size(dt::DataTableau, n::Integer) = (n == 1 ? dt.nrows : dt.ncols)

function size(dt::DataTableau, i::Int)
    if  i == 1
        return dt.nrows
    elseif i == 2
        return dt.ncols
    else
        throw(BoundsError)
    end
end

# Iteration methods:
start(dt::DataTableau) = 1
next(dt::DataTableau, i) = (dt[i], i + 1)
done(dt::DataTableau, i) = (i > dt.ncols)

end # of module


text:
# homogeneous velocity and density model
nz = 101; nx = 201;
vel = 3000 * ones(nz, nx);  # m/s
vel[51:end,:] .= 3500;  # m/s
rho = 2000 * ones(nz, nx);  # kg/m^3

# vertical and horizontal grid size
dz = 10; dx = 10;

# time step size and maximum modelling length
dt = 0.001; tmax = 1.0;  # use second as unit

# receiver location
irx = collect(1 :2 :nx);
irz = 2 * ones(length(irx));

# ==============================================================================
#                   dot-product test for the adjoint wavefield
# ==============================================================================
test_type = [(true , Float32, 1.0e-6 ),
             (true , Float64, 1.0e-12),
             (false, Float32, 1.0e-6 ),
             (false, Float64, 1.0e-12)
            ]
@testset "multi-step forward modelling and its adjoint operator" begin

    for (free_surface, data_format, tol) in test_type

        # generate finite difference stencil
        params = TdParams(rho, vel, free_surface, dz, dx, dt, tmax;
                 data_format=data_format, fd_flag="taylor", order=3, npml=20, apml=900.);

        # allocate snapshots for forward and adjoint operator
        spt1_f = Snapshot(params);
        spt2_f = Snapshot(params);
        spt1_b = Snapshot(params);
        spt2_b = Snapshot(params);

        # initialize spt1_f and spt2_b with random numbers
        N = params.Nz * params.Nx
        spt1_f.vz .= randn(params.data_format, N); spt2_b.vz .= randn(params.data_format, N);
        spt1_f.vx .= randn(params.data_format, N); spt2_b.vx .= randn(params.data_format, N);
        spt1_f.pz .= randn(params.data_format, N); spt2_b.pz .= randn(params.data_format, N);
        spt1_f.px .= randn(params.data_format, N); spt2_b.px .= randn(params.data_format, N);

        # temporary variables
        tmp    = zeros(params.data_format, params.Nz * params.Nx);
        tmp_z1 = zeros(params.data_format, params.Nz);
        tmp_z2 = zeros(params.data_format, params.Nz);
        tmp_x1 = zeros(params.data_format, params.Nx);
        tmp_x2 = zeros(params.data_format, params.Nx);

        # nt-step forward
        nt = 78
        for it = 1 : nt
            one_step_forward!(spt2_f, spt1_f, params, tmp_z1, tmp_z2, tmp_x1, tmp_x2);
            copy_snapshot!(spt1_f, spt2_f);
        end

        # nt-step adjoint
        for it = 1 : nt
            one_step_adjoint!(spt1_b, spt2_b, params, tmp, tmp_z1, tmp_z2, tmp_x1, tmp_x2);
            copy_snapshot!(spt2_b, spt1_b);
        end

        # inner product
        tmp1 = (dot(spt1_f.vz, spt1_b.vz) + dot(spt1_f.vx, spt1_b.vx)
              + dot(spt1_f.pz, spt1_b.pz) + dot(spt1_f.px, spt1_b.px))

        tmp2 = (dot(spt2_f.vz, spt2_b.vz) + dot(spt2_f.vx, spt2_b.vx)
              + dot(spt2_f.pz, spt2_b.pz) + dot(spt2_f.px, spt2_b.px))

        @test (tmp1-tmp2) / tmp1 < tol
    end
end

# ==============================================================================
#           adjoint Pz = Px in computational area
# ==============================================================================
@testset "test pz == px for the adjoint wavefield" begin

    for (free_surface, data_format, tol) in test_type

        # generate finite difference stencil
        params = TdParams(rho, vel, free_surface, dz, dx, dt, tmax;
                 data_format=data_format, fd_flag="taylor", order=3, npml=20, apml=900.);

        # generate band-limited random recordings
        w   = ricker(12.0, params.dt)
        hl  = floor(Int64, length(w)/2)
        rec = Recordings(irz, irx, params);
        idx_o= [1]

        # generate data trace by trace
        for i = 1 : rec.nr
            tmp = conv(randn(params.nt)*1000, w)
            copyto!(rec.p, idx_o[1], tmp, hl+1, params.nt)
            idx_o[1] = idx_o[1] + params.nt
        end

        # save adjoint pz and px as an 3D cube
        path_spt = joinpath(homedir(), "Desktop/snapshot.rsf")
        multi_step_adjoint(rec, params; path_spt=path_spt)

        # # read the adjoint snapshot cube
        (hdr, d) = read_RSdata(path_spt);

        zl = params.ntop+1; zu = zl + params.nz - 1;
        xl = params.npml+1; xu = xl + params.nx - 1;

        @test norm(d[zl:zu,xl:xu,3,:] - d[zl:zu,xl:xu,4,:]) / norm(d[zl:zu,xl:xu,3,:]) < tol
    end
end

ext: jl


text:
"""
Composite type to represent a dyadic kernel density estimator.
"""
Base.@kwdef mutable struct DyadicKernelDensityEstimator

# input parameters
    kernel_name::String
    bandwidth::Float64
    significance_level::Float64
    n_resample::Int
    sdp_solver::String
    evals::Vector{Float64}
    data::UpperTriangular{Float64}
    meta::Dict

    # final products
    n_evals::Int
    evals_min::Float64
    evals_max::Float64
    n_data::Int
    N_data::Int
    data_vec::Vector{Float64}
    fhat::Vector{Float64}
    Sigmahat::Symmetric{Float64}
    Sigmahatplus::Symmetric{Float64}
    eigmin_Sigmahat::Float64
    eigmin_Sigmahatplus::Float64
    sdp_error::Float64
    ucb::Array{Float64, 2}
    pci::Array{Float64, 2}
    bci::Array{Float64, 2}

end

"""
    DyadicKernelDensityEstimator(kernel_name, bandwidth, significance_level,
                                 n_resample, sdp_solver, evals, data, meta)

Construct a dyadic kernel density estimator.

# Arguments
- `kernel_name::String`: which kernel to use.
- `bandwidth::Float64`: the bandwidth for the estimator.
- `significance_level::Float64`: for the confidence band/intervals.
- `n_resample::Int`: the number of resamples used to construct the confidence band/intervals.
- `sdp_solver::String`: semi-definite program solver.
- `evals::Vector{Float64}`: points at which to evaluate the density estimator.
- `data::UpperTriangular{Float64}`: array of dyadic data.
- `meta::Dict`: any extra information to pass to the estimator.
"""
function DyadicKernelDensityEstimator(
    kernel_name::String,
    bandwidth::Float64,
    significance_level::Float64,
    n_resample::Int,
    sdp_solver::String,
    evals::Vector{Float64},
    data::UpperTriangular{Float64},
    meta::Dict)

    @assert bandwidth > 0
    @assert 0 <= significance_level <= 1
    @assert n_resample >= 1

    n_evals = length(evals)
    n_data = size(data, 1)
    N_data = Int(n_data * (n_data - 1) // 2)

    est = DyadicKernelDensityEstimator(

        # input parameters
        kernel_name,
        bandwidth,
        significance_level,
        n_resample,
        sdp_solver,
        evals,
        data,
        meta,

        # final products
        n_evals,
        minimum(evals),
        maximum(evals),
        n_data,
        N_data,
        fill(NaN, (N_data)),
        fill(NaN, (n_evals)),
        Symmetric(fill(NaN, (n_evals, n_evals))),
        Symmetric(fill(NaN, (n_evals, n_evals))),
        NaN,
        NaN,
        NaN,
        fill(NaN, (2, n_evals)),
        fill(NaN, (2, n_evals)),
        fill(NaN, (2, n_evals)),
    )

    index = 1
    for i in 1:n_data
        for j in 1:n_data
            if i < j
                est.data_vec[index] = est.data[i,j]
                index += 1
            end
        end
    end

    return est

end

function estimate_fhat(est::DyadicKernelDensityEstimator)

    for k in 1:est.n_evals

        est.fhat[k] = sum(kernel.(est.data_vec, est.evals[k], est.bandwidth,
                                  est.evals_min, est.evals_max, est.kernel_name))

        est.fhat[k] /= est.N_data
    end

end

function estimate_conditional_expectation(est::DyadicKernelDensityEstimator)

    cond_exp = zeros((est.n_evals, est.n_data))

    for k in 1:est.n_evals

        for i in 1:est.n_data
            for j in 1:est.n_data

                if i < j
                    if abs(est.data[i,j] - est.evals[k]) <= est.bandwidth
                        cond_exp[k,i] += kernel(est.data[i,j], est.evals[k], est.bandwidth,
                                                est.evals_min, est.evals_max, est.kernel_name)
                    end

                elseif j < i
                    if abs(est.data[j,i] - est.evals[k]) <= est.bandwidth
                        cond_exp[k,i] += kernel(est.data[j,i], est.evals[k], est.bandwidth,
                                                est.evals_min, est.evals_max, est.kernel_name)
                    end
                end
            end
        end
    end

    cond_exp /= (est.n_data - 1)

    return cond_exp
end

function estimate_Sigmahat(est::DyadicKernelDensityEstimator)

    n = est.n_data

    Sigmahat = zeros((est.n_evals, est.n_evals))
    term1 = zeros((est.n_evals, est.n_evals))
    term2 = zeros((est.n_evals, est.n_evals))
    term3 = zeros((est.n_evals, est.n_evals))

    cond_exp = estimate_conditional_expectation(est)

    for w1 in 1:est.n_evals
        for w2 in 1:est.n_evals

            # Si Si' term
            for i in 1:n
                term1[w1,w2] += cond_exp[w1,i] * cond_exp[w2,i]
            end
            term1[w1,w2] *= (4 / (n^2))

            # kij kij' term
            w_w1 = est.evals[w1]
            w_w2 = est.evals[w2]

            if abs(w_w1 - w_w2) <= 2 * est.bandwidth
                for i in 1:est.N_data
                    s = est.data_vec[i]
                    term2[w1,w2] +=
                        kernel(s, w_w1, est.bandwidth, est.evals_min,
                               est.evals_max, est.kernel_name) *
                        kernel(s, w_w2, est.bandwidth, est.evals_min,
                               est.evals_max, est.kernel_name)
                end
            end

            term2[w1,w2] *= (4 / (n^2 * (n - 1)^2))

            # fhat fhat' term
            term3[w1,w2] += est.fhat[w1] * est.fhat[w2]
            term3[w1,w2] *= (4 * n - 6) / (n * (n - 1))

            # combine terms
            Sigmahat[w1,w2] = term1[w1,w2] - term2[w1,w2] - term3[w1,w2]

        end
    end

    est.Sigmahat = Symmetric(Sigmahat)
    est.eigmin_Sigmahat = eigmin(est.Sigmahat)

end

function matrix_lipschitz_number(mat::Symmetric{Float64}, v::Vector{Float64})

    @assert size(mat, 1) == length(v)
    steps = diff(mat, dims=1) ./ diff(v)
    return maximum(steps)
end

function estimate_Sigmahatplus(est::DyadicKernelDensityEstimator)

    if est.eigmin_Sigmahat >= 0
        # do nothing if already psd
        est.Sigmahatplus = est.Sigmahat

    else
        # formulate optimization problem
        d = diag(est.Sigmahat)
        sinv = 1 ./ sqrt.(d .+ d')
        C = Semidefinite(est.n_evals)
        objective = maximum(abs(sinv .* (C - est.Sigmahat)))
        problem = minimize(objective)

        # solve optimization problem
        if est.sdp_solver == "mosek"
            solve!(problem, Mosek.Optimizer, silent_solver=true)
        elseif est.sdp_solver == "cosmo"
            solve!(problem, COSMO.Optimizer, silent_solver=true)
        else
            error("unknown sdp_solver")
        end

        # get answer
        Sigmahatplus = Symmetric(evaluate(C))
        mineig_optsol = max(-eigmin(Sigmahatplus), 0)
        est.Sigmahatplus = Symmetric(Sigmahatplus + 2 * mineig_optsol * I)

        # check lipschitz property
        n = est.n_data
        h = est.bandwidth
        L = matrix_lipschitz_number(est.Sigmahatplus, est.evals)
        @assert L <= 4 / (n * h^3)

    end

    est.eigmin_Sigmahatplus = eigmin(est.Sigmahatplus)
    est.sdp_error = maximum(abs.(est.Sigmahatplus - est.Sigmahat))
    @assert est.eigmin_Sigmahatplus >= 0

end

function get_quantile(A::Vector{<:Real}, q::Real)

    sorted_A = sort(A)
    idx = round(Int, length(A) * q)
    quantile = sorted_A[idx]

    return quantile
end

function estimate_uniform_confidence_band(est::DyadicKernelDensityEstimator)

    # make GP distribution and scaling
    gp_mu = zeros((est.n_evals))
    gp = MvNormal(gp_mu, est.Sigmahatplus)
    Sigmahatplus_diag_sqrt = sqrt.(diag(est.Sigmahatplus))

    sup_scaled_gps = zeros((est.n_resample))

    for rep in 1:est.n_resample

        # sample and scale gp
        Z_hat = rand(gp)
        Z_hat_scaled = Z_hat ./ Sigmahatplus_diag_sqrt

        # get sup absolute value
        sup_scaled_gps[rep] = maximum(abs.(Z_hat_scaled))

    end

    # get quantile
    c = get_quantile(sup_scaled_gps, 1-est.significance_level)
    ucb_width = c .* Sigmahatplus_diag_sqrt
    est.ucb[1,:] .= est.fhat .- ucb_width
    est.ucb[2,:] .= est.fhat .+ ucb_width

end

function estimate_pointwise_confidence_intervals(est::DyadicKernelDensityEstimator)

    for k in 1:est.n_evals

        w = est.evals[k]
        variance = est.Sigmahatplus[k,k]
        normal = Normal(0.0, sqrt(variance))
        q = Statistics.quantile(normal, 1-(est.significance_level/2))

        est.pci[1,k] = est.fhat[k] - q
        est.pci[2,k] = est.fhat[k] + q

    end

end

function estimate_bonferroni_confidence_intervals(est::DyadicKernelDensityEstimator)

    bonferroni_significance_level = est.significance_level / est.n_evals

    for k in 1:est.n_evals

        w = est.evals[k]
        variance = est.Sigmahatplus[k,k]
        normal = Normal(0.0, sqrt(variance))
        q = Statistics.quantile(normal, 1-(bonferroni_significance_level/2))

        est.bci[1,k] = est.fhat[k] - q
        est.bci[2,k] = est.fhat[k] + q

    end

end

"""
    estimate_ROT_bandwidth(data::UpperTriangular{Float64},
                           kernel_name::String)

Estimate a rule-of-thumb bandwidth from dyadic data.
"""
function estimate_ROT_bandwidth(data::UpperTriangular{Float64},
                                kernel_name::String)

    # get data_vec
    n_data = size(data, 1)
    N_data = Int(n_data * (n_data - 1) // 2)
    data_vec = Float64[]

    for i in 1:n_data
        for j in 1:n_data
            if i < j
                if -Inf < data[i,j] < Inf
                    push!(data_vec, data[i,j])
                end
            end
        end
    end

    if kernel_name == "epanechnikov_order_2"

        # sigmahat_W
        mean_W = sum(data_vec) / length(data_vec)
        sum_of_squares_W = sum((data_vec .- mean_W).^2)
        sigmahat_W_squared = sum_of_squares_W / (length(data_vec) - 1)
        sigmahat_W = sqrt(sigmahat_W_squared)

        # IQRhat
        upper_quartile = get_quantile(data_vec, 0.75)
        lower_quartile = get_quantile(data_vec, 0.25)
        IQRhat = upper_quartile - lower_quartile

        # n_rate
        n_rate = N_data^(-1/5)

        return 2.435 * min(sigmahat_W, IQRhat / 1.349) * n_rate

    else
        error("unknown kernel_name")
    end

end

"""
    fit(est::DyadicKernelDensityEstimator)

Fit a dyadic kernel density estimator to data.
"""
function fit(est::DyadicKernelDensityEstimator)

    estimate_fhat(est)
    estimate_Sigmahat(est)
    estimate_Sigmahatplus(est)
    estimate_uniform_confidence_band(est)
    estimate_pointwise_confidence_intervals(est)
    estimate_bonferroni_confidence_intervals(est)

end

"""
    display(est::DyadicKernelDensityEstimator)

Display a dyadic kernel density estimator.
"""
function Base.display(est::DyadicKernelDensityEstimator)

    println("DyadicKernelDensityEstimator")
    println("----------------------------")

    println("Kernel: ", est.kernel_name)
    println("Bandwidth: ", est.bandwidth)
    println("Significance level: ", est.significance_level)
    println("Num. GP resamples: ", est.n_resample)
    println("SDP solver: ", est.sdp_solver)
    println("Num. data nodes (n): ", est.n_data)
    println("Num. data points (N): ", est.N_data)
    println("Num. evaluation points: ", est.n_evals)
    println("Min. evaluation point: ", est.evals_min)
    println("Max. evaluation point: ", est.evals_max)
    println("Min. eigval of Sigmahat: ", est.eigmin_Sigmahat)
    println("Min. eigval of Sigmahatplus: ", est.eigmin_Sigmahatplus)
    println("SDP error: ", est.sdp_error)

end


text:
# ---
# title: 565. Array Nesting
# id: problem565
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array
# link: <https://leetcode.com/problems/array-nesting/description/>
# hidden: true
# ---
#
# A zero-indexed array A of length N contains all integers from 0 to N-1. Find
# and return the longest length of set S, where S[i] = {A[i], A[A[i]],
# A[A[A[i]]], ... } subjected to the rule below.
#
# Suppose the first element in S starts with the selection of element A[i] of
# index = i, the next element in S should be A[A[i]], and then A[A[A[i]]]… By
# that analogy, we stop adding right before a duplicate element occurs in S.
#
#
#
# **Example 1:**
#
#
#
#     Input: A = [5,4,0,3,1,6,2]
#     Output: 4
#     Explanation:
#     A[0] = 5, A[1] = 4, A[2] = 0, A[3] = 3, A[4] = 1, A[5] = 6, A[6] = 2.
#
#     One of the longest S[K]:
#     S[0] = {A[0], A[5], A[6], A[2]} = {5, 6, 2, 0}
#
#
#
#
# **Note:**
#
#   1. N is an integer within the range [1, 20,000].
#   2. The elements of A are all distinct.
#   3. Each element of A is an integer within the range [0, N-1].
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
using Base.Test

using DriveShare

trace = get_dataset(:trace)
high_level = get_dataset(:highlevel)


text:
"""
    RemoteData

A placeholder type to describe the location and properties of remote data that
is too large to fit directly in a tag's spot in the IFD. Calling [`load!`](@ref)
on an IFD object replaces all `RemoteData`s with the respective data.

$(FIELDS)
"""
struct RemoteData{O <: Unsigned}
    """Position of this data in the stream"""
    position::O

    """The datatype stored at this location"""
    datatype::DataType

    """The length of the data"""
    count::O
end

Base.position(o::RemoteData) = Int(o.pos)

const filepattern = r"<file (.*)>"
const fdpattern = r"<fd (.*)>"

"""
    extract_filename(io) -> String

Extract the name of the file backing a stream
"""
function extract_filename(io::IOStream)
    filename = match(filepattern, io.name)
    if filename !== nothing
        return filename[1]
    elseif match(fdpattern, io.name) !== nothing
        return ""
    else
        error("Can't extract filename from the given stream")
    end
end
extract_filename(io::Stream) = io.filename

"""
    check_bswap(io::Stream)

Check endianness of TIFF file to see if we need to swap bytes
"""
function need_bswap(io::Stream)
    seekstart(io)
    endianness = String(read(io, 2))
    # check if we need to swap byte order
    bswap = endianness != (ENDIAN_BOM == 0x04030201 ? "II" : "MM")
    return bswap
end

function offsetsize(io::Stream)
    seek(io, 2)
    tiff_version = Array{UInt8}(undef, 2)
    read!(io, tiff_version)
    if tiff_version == [0x2a, 0x00] || tiff_version == [0x00, 0x2a]
        return UInt32
    elseif tiff_version == [0x2b, 0x00] || tiff_version == [0x00, 0x2b]
        offset_bytesize = read(io, UInt16)
        constant = read(io, UInt16)
        return UInt64
    else
        error("Unknown TIFF version")
    end
end

const _mask1_uint128 = (UInt128(0x5555555555555555) << 64) | UInt128(0x5555555555555555)
const _mask2_uint128 = (UInt128(0x3333333333333333) << 64) | UInt128(0x3333333333333333)
const _mask4_uint128 = (UInt128(0x0f0f0f0f0f0f0f0f) << 64) | UInt128(0x0f0f0f0f0f0f0f0f)

"""
    reversebits(x)

Reverse the bit order in each byte of `x`. Julia assumes the first bit of each
byte is packed as a 0x01, while in TIFFs it's actually packed as 0x80. This
function flips the bit packing order.

!!! note
    This function does not flip the byte packing order, a la `bswap`.

This function should be applied to a BitArray after it is read in via
[`reversebits!`](@ref).

Code adapted from https://github.com/JuliaLang/julia/pull/34791

```jldoctest
julia> TiffImages.reversebits(0x0101) == 0x8080
true
```
"""
function reversebits(x::Base.BitInteger)
    z = unsigned(x)
    mask1 = _mask1_uint128 % typeof(z)
    mask2 = _mask2_uint128 % typeof(z)
    mask4 = _mask4_uint128 % typeof(z)
    z = ((z & mask1) << 1) | ((z >> 1) & mask1)
    z = ((z & mask2) << 2) | ((z >> 2) & mask2)
    z = ((z & mask4) << 4) | ((z >> 4) & mask4)
    return z
end

"""
    reversebits!(x)

Flips the bits in every byte backing `x`. Does not change byte order.
"""
function reversebits!(x::BitArray)
    n = length(x.chunks)
    for i in 1:n
        x.chunks[i] = reversebits(x.chunks[i])
    end
end

# using rephorm's table from
# https://github.com/rephorm/TIFF.jl/blob/master/tiff.jl#L30
const tiff_to_julian = Dict(
  0x0001 => UInt8,            # CHAR
  0x0002 => String,           # ASCII
  0x0003 => UInt16,           # SHORT
  0x0004 => UInt32,           # LONG
  0x0005 => Rational{UInt32}, # RATIONAL
  0x0006 => Int8,             # SBYTE
  0x0007 => Any,              # UNDEFINED
  0x0008 => Int16,            # SSHORT
  0x0009 => Int32,            # SLONG
  0x000a => Rational{Int32},  # SRATIONAL
  0x000b => Float32,          # FLOAT
  0x000c => Float64,          # DOUBLE
  0x000d => UInt32,           # IFD
  0x0010 => UInt64,           # LONG8
  0x0011 => Int64,            # SLONG8
  0x0012 => UInt64,           # IFD8
)

const julian_to_tiff = Dict(
    UInt8               => 0x0001,   # CHAR
    String              => 0x0002,   # ASCII
    UInt16              => 0x0003,   # SHORT
    UInt32              => 0x0004,   # LONG
    Rational{UInt32}    => 0x0005,   # RATIONAL
    Int8                => 0x0006,   # SBYTE
    Any                 => 0x0007,   # UNDEFINED
    Int16               => 0x0008,   # SSHORT
    Int32               => 0x0009,   # SLONG
    Rational{Int32}     => 0x000a,   # SRATIONAL
    Float32             => 0x000b,   # FLOAT
    Float64             => 0x000c,   # DOUBLE
    UInt64              => 0x0010,   # LONG8
    Int64               => 0x0011,   # SLONG8
)

# sampleformat, bitspersample => Julian type
const rawtype_mapping = Dict(
    (SAMPLEFORMAT_UINT, 1) => Bool,
    (SAMPLEFORMAT_UINT, 8) => UInt8,
    (SAMPLEFORMAT_UINT, 16) => UInt16,
    (SAMPLEFORMAT_UINT, 32) => UInt32,
    (SAMPLEFORMAT_UINT, 64) => UInt64,
    (SAMPLEFORMAT_INT, 8) => Int8,
    (SAMPLEFORMAT_INT, 16) => Int16,
    (SAMPLEFORMAT_INT, 32) => Int32,
    (SAMPLEFORMAT_INT, 64) => Int64,
    (SAMPLEFORMAT_IEEEFP, 16) => Float16,
    (SAMPLEFORMAT_IEEEFP, 32) => Float32,
    (SAMPLEFORMAT_IEEEFP, 64) => Float64,
)

Base.bswap(c::Colorant{T, N}) where {T, N} = mapc(bswap, c)