text:
module POMDPModelTools

using POMDPs
using Random
using LinearAlgebra
using SparseArrays
using UnicodePlots

import POMDPs: actions, actionindex
import POMDPs: states, stateindex
import POMDPs: observations, obsindex
import POMDPs: sampletype, generate_sr, initialstate, isterminal, discount
import POMDPs: implemented
import Distributions: pdf, mode, mean, support
import Random: rand, rand!
import Statistics: mean
import Base: ==

export
    render
include("visualization.jl")

# info interface
export
    add_infonode,
    action_info,
    solve_info,
    update_info,
    generate_sri,
    generate_sori
include("info.jl")

export
    ordered_states,
    ordered_actions,
    ordered_observations
include("ordered_spaces.jl")

export
    TerminalState,
    terminalstate
include("terminal_state.jl")

export GenerativeBeliefMDP
include("generative_belief_mdp.jl")

export FullyObservablePOMDP
include("fully_observable_pomdp.jl")

export UnderlyingMDP
include("underlying_mdp.jl")

export obs_weight
include("obs_weight.jl")

export
    probability_check,
    obs_prob_consistency_check,
    trans_prob_consistency_check

export
    weighted_iterator
include("distributions/weighted_iteration.jl")

export
    SparseCat
include("distributions/sparse_cat.jl")

export
    BoolDistribution
include("distributions/bool.jl")

export
    Deterministic
include("distributions/deterministic.jl")

export
    Uniform,
    UnsafeUniform
include("distributions/uniform.jl")

# convenient implementations
include("convenient_implementations.jl")

export
    evaluate
include("policy_evaluation.jl")

export
    showdistribution
include("distributions/pretty_printing.jl")

export
    SparseTabularMDP,
    SparseTabularPOMDP,
    transition_matrix,
    reward_vector,
    observation_matrix,
    transition_matrices,
    reward_matrix,
    observation_matrices

include("sparse_tabular.jl")

end # module


text:
#=
# a customized function to convert number to formatted number
function number_to_string(number::Number; digit::Int=6)
    if digit==0
        return @sprintf("%0.0f", number)
    elseif digit==1
        return @sprintf("%0.1f", number)
    elseif digit==2
        return @sprintf("%0.2f", number)
    elseif digit==3
        return @sprintf("%0.3f", number)
    elseif digit==4
        return @sprintf("%0.4f", number)
    elseif digit==5
        return @sprintf("%0.5f", number)
    else
        return @sprintf("%0.6f", number)
    end
end

# there is an option to save the figure
function visualize_struct_tree(tree::Tree, save::Bool=false)
    # clear plot first
    figure("Tree Visualization",figsize=(16,9), dpi=120)
    clf()

    # display table header
    text(-1.00, tree.h+0.05, L"$P_\mathrm{water}$"            , color=:red   , ha="right", va="bottom", fontsize=8)
    text( 1.00, tree.h+0.05, L"$Q_\mathrm{water}$"            , color=:blue  , ha="left" , va="bottom", fontsize=8)
    text( 1.40, tree.h+0.05, L"$\overline{\mathrm{PAR_{sl}}}$", color=:orange, ha="left" , va="bottom", fontsize=8)
    text( 1.60, tree.h+0.05, L"$\mathrm{PAR_{sh}}$"           , color=:orange, ha="left" , va="bottom", fontsize=8)

    # plot the root system
    for i in 1:length(tree.roots.root_list)
        rooti = tree.roots.root_list[i]

        lx = [         0,          0]
        ly = [rooti.z_hi, rooti.z_lo]
        plot(lx, ly, color=:brown)

        lx = [-rooti.r_layer,          0, rooti.r_layer]
        ly = [ rooti.z_lo   , rooti.z_hi, rooti.z_lo   ]
        plot(lx, ly, color=:brown)

        text(-rooti.r_layer, rooti.z_lo, number_to_string(rooti.p_ups , digit=3), color=:red , ha="right", va="top"   , fontsize=8)
        text(-rooti.r_layer, rooti.z_lo, number_to_string(rooti.p_rhiz, digit=3), color=:red , ha="left" , va="bottom", fontsize=8)
        text( rooti.r_layer, rooti.z_lo, number_to_string(rooti.q     , digit=6), color=:blue, ha="left" , va="top"   , fontsize=8)
    end

    # plot the trunk
    lx = [0, 0, 0]
    ly = [0, tree.trunk.z_lo, tree.trunk.z_hi]
    plot(lx, ly, color=:black)
    text(0, tree.trunk.z_lo    , number_to_string(tree.trunk.p_ups, digit=3), color=:red , ha="right", va="bottom", fontsize=8)
    text(0, tree.trunk.z_hi    , number_to_string(tree.trunk.p_dos, digit=3), color=:red , ha="right", va="top"   , fontsize=8)
    text(0, tree.trunk.z_hi*0.5, number_to_string(tree.trunk.q    , digit=6), color=:blue, ha="left" , va="top"   , fontsize=8)

    # plot the canopy stem and leaf system
    for i in 1:length(tree.branch.branch_list)
        branchi = tree.branch.branch_list[i]
        canopyi = tree.canopy.canopy_list[i]

        # plot the branch vertically
        lx = [           0,            0]
        ly = [branchi.z_lo, branchi.z_hi]
        plot(lx, ly, color=:green)

        # plot the branch horizentally
        lx = [-branchi.r_layer,            0, branchi.r_layer]
        ly = [ branchi.z_hi   , branchi.z_lo, branchi.z_hi   ]
        plot(lx, ly, color=:green)

        # label the stem-leaf joint pressure and flow rate through each branch
        text(-branchi.r_layer, branchi.z_hi, number_to_string(branchi.p_dos, digit=3), color=:red , ha="right", va="top", fontsize=8)
        text( branchi.r_layer, branchi.z_hi, number_to_string(branchi.q    , digit=6), color=:blue, ha="left" , va="top", fontsize=8)

        # calculate and display mean PAR for sunlit leaves and shaded leaves
        par_sl = sum(canopyi.par_list[1:end-1] .* canopyi.la_list[1:end-1]) / sum(canopyi.la_list[1:end-1])
        par_sh = canopyi.par_list[end]
        text(1.40, branchi.z_hi, number_to_string(par_sl, digit=1), color=:orange, ha="left", va="top", fontsize=8)
        text(1.60, branchi.z_hi, number_to_string(par_sh, digit=1), color=:orange, ha="left", va="top", fontsize=8)
    end

    # plot the ground
    plot([-1.2,4], [0,0], "k-", color=:black, linewidth=0.2)

    # set the title and label
    title("Tree system")
    xlabel("Distance from the tree base (m)", fontsize=16)
    ylabel("z (m)"                          , fontsize=16)

    # save the figure
    if save
        savefig("zzz_tree.pdf", bbox_inches="tight")
    end
end
=#

ext: jl

max_stars_repo_name: bbuman/Land

max_issues_repo_name: bbuman/Land

max_forks_repo_name: bbuman/Land


text:
using Test
#using BenchmarkTools
using LogicCircuits

function benchmark_pass_down()
    circuit = load_smooth_logical_circuit(zoo_psdd_file("plants.psdd"));
    data = train(dataset(twenty_datasets("plants")));
    upflow = UpFlowΔ(circuit, num_examples(data), Bool,  (max_factors = 0, compact⋀ = false, compact⋁ = false));
    pass_up(upflow, data);
    circuit(data);

    @time begin
        downflow_circuit = DownFlowΔ(upflow,  (max_factors = 0, compact⋀ = false, compact⋁ = false));
        pass_down(downflow_circuit);
    end
    # 547.595 ms (2416251 allocations: 474.71 MiB)

    @time Logical.pass_down2(circuit, data);
    # 379.885 ms (1687709 allocations: 419.05 MiB)

    for (i,n) in enumerate(downflow_circuit)
        if n isa HasDownFlow
            @test downflow(n) == n.origin.origin.data[1]
        else
            @assert n isa DownFlowLeaf "$(typeof(n))"
        end
    end
end


text:
using RedPitayaDAQServer
using PyPlot

rp = RedPitaya("192.168.178.46")

dec = 64
modulus = 4800
base_frequency = 125000000
periods_per_step = 1
samples_per_period = div(modulus, dec)*periods_per_step
periods_per_frame = div(20, periods_per_step) # about 0.5 s frame length

decimation(rp, dec)
samplesPerPeriod(rp, samples_per_period)
periodsPerFrame(rp, periods_per_frame)
slowDACPeriodsPerFrame(rp, periods_per_frame)
numSlowDACChan(rp, 1)
lut = collect(range(0,1,length=periods_per_frame))
setSlowDACLUT(rp, lut)

modeDAC(rp, "STANDARD")
frequencyDAC(rp, 1, 1, base_frequency / modulus)

println(" frequency = $(frequencyDAC(rp,1,1))")
amplitudeDAC(rp, 1, 1, 4000)
phaseDAC(rp, 1, 1, 0.0 ) # Phase has to be given in between 0 and 1
ramWriterMode(rp, "TRIGGERED")
triggerMode(rp, "INTERNAL")

numTrials = 30

signals = zeros(samples_per_period*periods_per_frame, 2, numTrials)

for i=1:numTrials
    @info "measure trial $i"
    masterTrigger(rp, false)
    startADC(rp, 0)
    masterTrigger(rp, true)
    currFr = enableSlowDAC(rp, true, 1, 0.5, 1.0)
    data = readData(rp, currFr, 1)
    signals[:,1,i] .= vec(data[:,1,:,:])
    signals[:,2,i] .= vec(data[:,2,:,:])
    stopADC(rp)
end

figure(1)
clf()
subplot(2,1,1)
for i=1:numTrials
  plot(vec(signals[:,1,i]))
end
subplot(2,1,2)
for i=1:numTrials
  plot(vec(signals[:,2,i]))
end


text:
using FFTW

function compressaudiodct(channels, threshold)
    # Split channels
    left = channels[:,1]
    right = channels[:,2]

    # Perfom transforms
    fleft = dct(left)
    fright = dct(right)

    # Compress
    tleft = threshold * maximum(abs.(fleft))
    ileft = abs.(fleft) .> tleft
    tright = threshold * maximum(abs.(fright))
    iright = abs.(fright) .> tright

    # Generate return values
    fleft = fleft[ileft]
    fright = fright[iright]
    s = sizeof(fleft) + sizeof(fright) + sizeof(ileft) + sizeof(iright) +
        sizeof(length(left))
    return ((fleft, ileft, fright, iright, length(left)), s)
end

function decompressaudiodct(compression)
    # Unpack compression
    cleft, ileft, cright, iright, len = compression

    # Decompress left channel
    fleft = zeros(length(ileft))
    j = 1
    for i = 1:length(fleft)
        if ileft[i]
            fleft[i] = cleft[j]
            j += 1
        end
    end

    # Decompress right channel
    fright = zeros(length(iright))
    j = 1
    for i = 1:length(fright)
        if iright[i]
            fright[i] = cright[j]
            j += 1
        end
    end

    # Perform inverse transforms
    left = idct(fleft)
    right = idct(fright)

    return [left[1:len] right[1:len]]
end


text:
# function_plots.jl

"""
Script for creating either 3D Makie plots, or 2D Plotly plots, of functions from the Functions directory.
"""

using ..PolaronMakie
using Plots
using ProgressBars
using Printf
plotly()

"""
D(u) Function
"""

# p = (:v, :w, :β)
# i = (1.0, 1.0, 5.0)
# pl = ((0.0, 10.0), (0.0, 10.0), (0.0, 10.0))
# al = ((-25, 25), (-25, 25), (-25, 25))
# a = Dict(
#     "ℜ_title" => "Real D(u)",
#     "ℜ_x" => "ℜ(u)",
#     "ℜ_y" => "ℑ(u)",
#     "ℜ_z" => "ℜ[D(u)]",
#     "ℑ_title" => "Imaginary D(u)",
#     "ℑ_x" => "ℜ(u)",
#     "ℑ_y" => "ℑ(u)",
#     "ℑ_z" => "ℑ[D(u)]",
# )

# vf = viewfunction(ℜD, ℑD, p, i, pl, a, axis_limits = al, len = 200)

"""
S(u) Function
"""

# p = (:v, :w, :β, :α)
# i = (2.0, 1.0, 10, 6.0)
# pl = ((0.0, 10.0), (0.0, 10.0), (0.0, 10.0), (0.0, 10.0))
# al = ((-25, 25), (-25, 25), (-25, 25))
# a = Dict(
#     "ℜ_title" => "Real S(u)",
#     "ℜ_x" => "ℜ(u)",
#     "ℜ_y" => "ℑ(u)",
#     "ℜ_z" => "ℜ[S(u)]",
#     "ℑ_title" => "Imaginary S(u)",
#     "ℑ_x" => "ℜ(u)",
#     "ℑ_y" => "ℑ(u)",
#     "ℑ_z" => "ℑ[S(u)]",
# )

# vf = viewfunction(ℜS, ℑS, p, i, pl, a, axis_limits = al, len = 100)

"""
X(u) Function
"""

# p = (:α, :v, :w)
# i = (5, 4.0, 2.1)
# pl = ((0.01, 10), (0.01, 10), (0.01, 10))
# al = ((0.01, 10), (0.01, 10), (-50, 100))
# a = Dict(
#     "ℜ_title" => "Real Χ(u)",
#     "ℜ_x" => "ℜ(u)",
#     "ℜ_y" => "ℑ(u)",
#     "ℜ_z" => "ℜ[X(u)]",
#     "ℑ_title" => "Imaginary Χ(u)",
#     "ℑ_x" => "ℜ(u)",
#     "ℑ_y" => "ℑ(u)",
#     "ℑ_z" => "ℑ[X(u)]",
# )
#
# vf = viewfunction(ℜχ, ℑχ, p, i, pl, a, axis_limits = al, len = 200)

Ω = range(0.0, stop = 10.0, length = 100)
params = (3.0, 7.0, 5.8, 1.6)
ReX = Array{Float64}(undef, length(Ω))
ImX = Array{Float64}(undef, length(Ω))
iter = ProgressBar(range(1, stop = length(Ω), step = 1))
Threads.@threads for i in iter
    ReX[i] = ℜχ(Ω[i], params...)
    ImX[i] = ℑχ(Ω[i], params...)
    set_description(iter, string(@sprintf("Ω: %.2f | ", Ω[i])) * string(@sprintf("ℜχ: %.2f | ", ℜχ(Ω[i], params...))) * string(@sprintf("ℑχ: %.2f |", ℑχ(Ω[i], params...))))
end
ReX_plot = Plots.plot(Ω, ReX, label = "ℜχ")
ImX_plot = Plots.plot!(Ω, ImX, label = "ℑχ")
Plots.xlabel!("Ω")
Plots.ylabel!("χ(Ω)")


text:
# The SteinGaussianUnboundedPowerKernel has base kernel of the form:
#
# k(x, y) = ((1 + gamma*||x||^alpha)*(1 + gamma*||y||^alpha)) *
#   exp(-1/(2*beta)*||x-y||^2)
#
# where ||.|| is the L2 norm. This is gauranteed to enforce
# that k is a universal kernel that vanishes at infinity.

mutable struct SteinGaussianUnboundedPowerKernel <: SteinGaussianWeightedKernel
    # the alpha parameter
    alpha::Float64
    # the beta parameter
    beta::Float64
    # the gamma parameter
    gamma::Float64
end

# have default parameters of alpha = 2.5,  beta = 1.0, gamma = 1.0
SteinGaussianUnboundedPowerKernel() = SteinGaussianUnboundedPowerKernel(2.5, 1.0, 1.0)

SteinGaussianUnboundedPowerKernel(alpha::Float64) = SteinGaussianUnboundedPowerKernel(alpha, 1.0, 1.0)

SteinGaussianUnboundedPowerKernel(alpha::Int) = SteinGaussianUnboundedPowerKernel(float(alpha), 1.0, 1.0)

# The term (1 + gamma*||z||^alpha)
function _Qz(ker::SteinGaussianUnboundedPowerKernel, z::Array{Float64, 1})
    1.0 + ker.gamma * norm(z,2)^ker.alpha
end

# The term ((1 + gamma*||x||^alpha)*(1 + gamma*||y||^alpha))^-1
function Q(ker::SteinGaussianUnboundedPowerKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    _Qz(ker, x) * _Qz(ker, y)
end

function dxj_Q(ker::SteinGaussianUnboundedPowerKernel, x::Array{Float64,1}, y::Array{Float64,1}, j::Int64)
    fy = _Qz(ker, y)
    mx = norm(x,2)^(ker.alpha - 2)
    mx * ker.gamma * ker.alpha * fy * x[j]
end

function dxjyj_Q(ker::SteinGaussianUnboundedPowerKernel, x::Array{Float64,1}, y::Array{Float64,1}, j::Int64)
    mx = norm(x,2)^(ker.alpha - 2)
    my = norm(y,2)^(ker.alpha - 2)

    ker.gamma^2 * ker.alpha^2 * mx * my * x[j] * y[j]
end


text:
function load_distributions_samples()
return quote

using .Distributions: Distribution, Uniform, Normal
import .Distributions

# ======================================================
# Sampling from a uniform distribution on balls/spheres
# ======================================================

"""
    _sample_unit_nsphere_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                 rng::AbstractRNG=GLOBAL_RNG,
                                 seed::Union{Int, Nothing}=nothing) where {N}

Draw samples from a uniform distribution on an ``n``-dimensional unit
sphere using Muller's method.

### Input

- `D`    -- output, vector of points
- `n`    -- dimension of the sphere
- `p`    -- number of random samples
- `rng`  -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed` -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `nsamples` vectors.

### Algorithm

This function implements Muller's method of normalised Gaussians [1] to uniformly
sample over the ``n``-dimensional sphere ``S^n`` (which is the bounding surface
of the ``n``-dimensional unit ball).

Given ``n`` canonical Gaussian random variables ``Z₁, Z₂, …, Z_n``, the
distribution of the vectors

```math
\\dfrac{1}{α}\\left(z₁, z₂, …, z_n\\right)^T,
```
where ``α := \\sqrt{z₁² + z₂² + … + z_n²}``, is uniform over ``S^n``.

[1] Muller, Mervin E. *A note on a method for generating points uniformly on
    n-dimensional spheres.* Communications of the ACM 2.4 (1959): 19-20.
"""
function _sample_unit_nsphere_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                      rng::AbstractRNG=GLOBAL_RNG,
                                      seed::Union{Int, Nothing}=nothing) where {N}
    rng = reseed(rng, seed)
    Zdims = [Normal() for _ in 1:n] # normal distributions, one for each dimension
    v = Vector{N}(undef, n) # sample direction
    @inbounds for j in 1:p
        α = zero(N)
        for i in 1:n
            v[i] = rand(rng, Zdims[i])
            α += v[i]^2
        end
        D[j] = v ./ sqrt(α)
    end
    return D
end

"""
    _sample_unit_nball_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                               rng::AbstractRNG=GLOBAL_RNG,
                               seed::Union{Int, Nothing}=nothing) where {N}

Draw samples from a uniform distribution on an ``n``-dimensional unit ball
using Muller's method.

### Input

- `D`    -- output, vector of points
- `n`    -- dimension of the ball
- `p`    -- number of random samples
- `rng`  -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed` -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `nsamples` vectors.

### Algorithm

This function implements Muller's method of normalised Gaussians [1] to
sample from the interior of the ball.

Given ``n`` Gaussian random variables ``Z₁, Z₂, …, Z_n`` and a uniformly
distributed random variable ``r`` with support in ``[0, 1]``, the distribution
of the vectors

```math
\\dfrac{r^{1/n}}{α} \\left(z₁, z₂, …, z_n\\right)^T,
```
where ``α := \\sqrt{z₁² + z₂² + … + z_n²}``, is uniform over the
``n``-dimensional unit ball.

[1] Muller, Mervin E. *A note on a method for generating points uniformly on
    n-dimensional spheres.* Communications of the ACM 2.4 (1959): 19-20.
"""
function _sample_unit_nball_muller!(D::Vector{Vector{N}}, n::Int, p::Int;
                                    rng::AbstractRNG=GLOBAL_RNG,
                                    seed::Union{Int, Nothing}=nothing) where {N}

    rng = reseed(rng, seed)
    Zdims = [Normal() for _ in 1:n] # normal distributions, one for each dimension
    Zrad = Uniform() # distribution to pick random radius
    one_over_n = one(N)/n
    v = Vector{N}(undef, n) # sample direction
    @inbounds for j in 1:p
        α = zero(N)
        for i in 1:n
            v[i] = rand(rng, Zdims[i])
            α += v[i]^2
        end
        r = rand(rng, Zrad)
        β = r^one_over_n / sqrt(α)
        D[j] = v .* β
    end
    return D
end

# ========================
# Sampling from a LazySet
# ========================

"""
    Sampler

Abstract type for defining new sample methods.

### Notes

All subtypes should implement a `_sample!` method.
"""
abstract type Sampler end

"""
    sample(X::LazySet{N}, num_samples::Int;
           rng::AbstractRNG=GLOBAL_RNG,
           seed::Union{Int, Nothing}=nothing) where {N}

Sampling of an arbitrary `LazySet` `X`.

### Input

- `X`           -- lazyset
- `num_samples` -- number of random samples
- `sampler`     -- Sampler used (default: `RejectionSampler`)
- `rng`         -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`        -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `num_samples` vectors.
If `num_samples` is not passed one sample as a single vector is returned.

### Algorithm

See the documentation of `RejectionSampler`.
"""
function sample(X::LazySet{N}, num_samples::Int;
                sampler=RejectionSampler,
                rng::AbstractRNG=GLOBAL_RNG,
                seed::Union{Int, Nothing}=nothing) where {N<:Real}
    require(:Distributions; fun_name="sample")
    @assert isbounded(X) "this function requires that the set `X` is bounded"*
                            ", but it is not"

    D = Vector{Vector{N}}(undef, num_samples) # preallocate output
    _sample!(D, sampler(X); rng=rng, seed=seed)
    return D
end

# without argument, returns a single element (instead of a singleton)
function sample(X::LazySet{N}; kwargs...) where {N<:Real}
    return sample(X, 1; kwargs...)[1]
end

# =====================
# Rejection Sampling
# =====================

"""
    RejectionSampler{S<:LazySet, D<:Distribution} <: Sampler

Type used for rejection sampling of an arbitrary `LazySet` `X`.

### Fields

- `X`           -- lazyset
- `box_approx`  -- Distribution from which the sample is drawn

### Algorithm

Draw a sample ``x`` from a uniform distribution of a box-overapproximation of the
original set ``X`` in all ``n`` dimensions. The function rejects a drawn sample ``x``
and redraws as long as the sample is not contained in the original set ``X``,
i.e., ``x ∉ X``.
"""
struct RejectionSampler{S<:LazySet, D<:Distribution} <: Sampler
    X::S
    box_approx::Vector{D}
end

function RejectionSampler(X, distribution=Uniform)
    B = box_approximation(X)
    canonical_support = hcat(low(B), high(B))
    dims = size(canonical_support, 1)
    box_approx = [distribution(canonical_support[i,:]...) for i = 1:dims]
    return RejectionSampler(X, box_approx)
end

"""
    _sample!(D::Vector{Vector{N}},
             sampler::RejectionSampler;
             rng::AbstractRNG=GLOBAL_RNG,
             seed::Union{Int, Nothing}=nothing) where {N<:Real}

Sample points using rejection sampling.

### Input

- `D`           -- output, vector of points
- `sampler`     -- Sampler from which the points are sampled
- `rng`         -- (optional, default: `GLOBAL_RNG`) random number generator
- `seed`        -- (optional, default: `nothing`) seed for reseeding

### Output

A vector of `num_samples` vectors.
"""
function _sample!(D::Vector{Vector{N}},
                  sampler::RejectionSampler;
                  rng::AbstractRNG=GLOBAL_RNG,
                  seed::Union{Int, Nothing}=nothing) where {N<:Real}
    rng = reseed(rng, seed)
    @inbounds for i in 1:length(D)
        w = rand.(Ref(rng), sampler.box_approx)
        while w ∉ sampler.X
            w = rand.(Ref(rng), sampler.box_approx)
        end
        D[i] = w
    end
    return D
end

end # quote
end # function load_distributions_samples()


text:
function measure_feature(property::BoundingBox, t::IndexedTable, labels::AbstractArray, N::Int = 0)
    N = N == 0 ? maximum(labels) : N
    init = StepRange(typemax(Int),-1,-typemax(Int))
    coords = [(init, init) for i = 1:N]

    for i in CartesianIndices(labels)
        l = labels[i]
        if  l != 0
            rs, cs = coords[l]
            rs = i[1] < first(rs) ? StepRange(i[1], 1, last(rs)) : rs
            rs = i[1] > last(rs) ? StepRange(first(rs), 1, i[1]) : rs
            cs = i[2] < first(cs) ? StepRange(i[2], 1, last(cs)) : cs
            cs = i[2] > last(cs) ? StepRange(first(cs), 1, i[2]) : cs
            coords[l] = rs, cs
        end
    end
    coords_matrix = reshape(reinterpret(StepRange{Int,Int},coords),(2,N))
    t = pushcol(t, :r => view(coords_matrix,1,:), :c => view(coords_matrix,2,:))
    fill_properties(property, t)

end

function fill_properties(property::BoundingBox, t::IndexedTable)
    t = property.box_area ? compute_box_area(t) : t
end

function compute_box_area(t::IndexedTable)
    @transform t {box_area = length(:r) *  length(:c)}
end


text:
# Pure Julia implementation of decompression of zlib and gzip
# compressed data, as specified by RFC 1950-1952.
#
# Historical note: In-memory decompression was implemented in 2013 in
# a gist, https://gist.github.com/GunnarFarneback/8254567. Streaming
# decompression was added in 2018 when it was also turned into a Julia
# package.

"""
    Inflate

Pure Julia implementation of decompression of the Deflate format and
the Zlib and Gzip wrapper formats.

In-memory decompression is done by the following functions:

| function | decompresses |
| -------- | ------------ |
| `inflate(data::Vector{UInt8})` | Deflate data |
| `inflate_zlib(data::Vector{UInt8})` | Zlib data |
| `inflate_gzip(data::Vector{UInt8})` | Gzip data |

Streaming decompression is done using the following types:

| stream | decompresses |
| ------ | ------------ |
| `InflateStream(stream::IO)` | Deflate stream |
| `InflateZlibStream(stream::IO)` | Zlib stream |
| `InflateGzipStream(stream::IO)` | Gzip stream |
"""
module Inflate

export inflate, inflate_zlib, inflate_gzip,
       InflateStream, InflateZlibStream, InflateGzipStream

# Huffman codes are internally represented by Vector{Vector{Int}},
# where code[k] are a vector of the values with code words of length
# k. Codes are assigned in order from shorter to longer codes and in
# the order listed. E.g.
# [[], [2, 7], [1, 3, 5], [6, 4]]
# would be the code
# 00    - 2
# 01    - 7
# 100   - 1
# 101   - 3
# 110   - 5
# 1110  - 6
# 1111  - 4

const fixed_literal_or_length_table = (Vector{Int})[Int[],
                                                    Int[],
                                                    Int[],
                                                    Int[],
                                                    Int[],
                                                    Int[],
                                                    Int[256:279;],
                                                    Int[0:143; 280:287],
                                                    Int[144:255;]]

const fixed_distance_table = (Vector{Int})[Int[],
                                           Int[],
                                           Int[],
                                           Int[],
                                           Int[0:31;]]

abstract type AbstractInflateData end

mutable struct InflateData <: AbstractInflateData
    bytes::Vector{UInt8}
    current_byte::Int
    bytepos::Int
    bitpos::Int
    literal_or_length_code::Vector{Vector{Int}}
    distance_code::Vector{Vector{Int}}
    update_input_crc::Bool
    crc::UInt32
end

function InflateData(source::Vector{UInt8})
    InflateData(source, 0, 1, 0, fixed_literal_or_length_table,
                fixed_distance_table, false, init_crc())
end

function get_input_byte(data::InflateData)
    byte = data.bytes[data.bytepos]
    data.bytepos += 1
    if data.update_input_crc
        data.crc = update_crc(data.crc, byte)
    end
    return byte
end

# This isn't called when reading Gzip header, so no need to
# consider updating crc.
function get_input_bytes(data::InflateData, n::Int)
    bytes = @view data.bytes[data.bytepos:(data.bytepos + n - 1)]
    data.bytepos += n
    return bytes
end

function getbit(data::AbstractInflateData)
    if data.bitpos == 0
        data.current_byte = Int(get_input_byte(data))
    end
    b = data.current_byte & 1
    data.bitpos += 1
    if data.bitpos == 8
        data.bitpos = 0
    else
        data.current_byte >>= 1
    end
    return b
end

function getbits(data::AbstractInflateData, n::Int)
    b = 0
    for i = 0:(n-1)
        b |= getbit(data) << i
    end
    return b
end

function skip_bits_to_byte_boundary(data::AbstractInflateData)
    data.bitpos = 0
    return
end

# It is the responsibility of the caller to make sure that bitpos is
# at zero, e.g. by calling skip_bits_to_byte_boundary.
function get_aligned_byte(data::AbstractInflateData)
    return get_input_byte(data)
end

function get_value_from_code(data::AbstractInflateData,
                             code::Vector{Vector{Int}})
    v = 0
    for i = 1:length(code)
        v = (v << 1) | getbit(data)
        if v < length(code[i])
            return code[i][1 + v]
        end
        v -= length(code[i])
    end
    error("incomplete code table")
end

function get_literal_or_length(data::AbstractInflateData)
    return get_value_from_code(data, data.literal_or_length_code)
end

const base_length = [11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131, 163, 195, 227]
const extra_length_bits = [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5]

function getlength(data::AbstractInflateData, v::Int)
    if v <= 264
        return v - 254
    elseif v <= 284
        return base_length[v - 264] + getbits(data, extra_length_bits[v - 264])
    else
        return 258
    end
end

function getdist(data::AbstractInflateData)
    b = get_value_from_code(data, data.distance_code)
    if b <= 3
        return b + 1
    else
        extra_bits = fld(b - 2, 2)
        return 1 + ((2 + b % 2) << extra_bits) + getbits(data, extra_bits)
    end
end

function transform_code_lengths_to_code(code_lengths::Vector{Int})
    code = Vector{Int}[]
    for i = 1:length(code_lengths)
        n = code_lengths[i]
        if n > 0
            while n > length(code)
                push!(code, Int[])
            end
            push!(code[n], i - 1)
        end
    end
    return code
end

const order = [16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15]

function read_code_tables(data::AbstractInflateData)
    hlit = getbits(data, 5) + 257
    hdist = getbits(data, 5) + 1
    hclen = getbits(data, 4) + 4
    code_length_code_lengths = zeros(Int, 19)
    for i = 1:hclen
        code_length_code_lengths[1 + order[i]] = getbits(data, 3)
    end
    code_length_code = transform_code_lengths_to_code(code_length_code_lengths)
    code_lengths = zeros(Int, hlit + hdist)
    i = 1
    while i <= hlit + hdist
        c = get_value_from_code(data, code_length_code)
        n = 1
        l = 0
        if c < 16
            l = c
        elseif c == 16
            n = 3 + getbits(data, 2)
            l = code_lengths[i-1]
        elseif c == 17
            n = 3 + getbits(data, 3)
        else
            # A code of length 19 can only yield values between 0 and
            # 18 so we can only get here if c == 18.
            n = 11 + getbits(data, 7)
        end
        code_lengths[i:(i+n-1)] .= l
        i += n
    end
    data.literal_or_length_code = transform_code_lengths_to_code(code_lengths[1:hlit])
    data.distance_code = transform_code_lengths_to_code(code_lengths[(hlit+1):end])
end

function _inflate(data::InflateData)
    out = UInt8[]
    final_block = false
    while !final_block
        final_block = getbits(data, 1) == 1
        compression_mode = getbits(data, 2)
        if compression_mode == 0
            skip_bits_to_byte_boundary(data)
            len = getbits(data, 16)
            nlen = getbits(data, 16)
            if len ⊻ nlen != 0xffff
                error("corrupted data")
            end
            append!(out, get_input_bytes(data, len))
            continue
        elseif compression_mode == 1
            data.literal_or_length_code = fixed_literal_or_length_table
            data.distance_code = fixed_distance_table
        elseif compression_mode == 2
            read_code_tables(data)
        else
            error("invalid block compression mode 3")
        end

        while true
            v = get_literal_or_length(data)
            if v < 256
                push!(out, UInt8(v))
            elseif v == 256
                break
            else
                length = getlength(data, v)
                distance = getdist(data)
                if length <= distance
                    append!(out, @view out[(end - distance + 1):(end - distance + length)])
                else
                    for i = 1:length
                        push!(out, out[end - distance + 1])
                    end
                end
            end
        end
    end

    return out
end

function init_adler()
    return (0, 1)
end

function update_adler(adler::Tuple{Int, Int}, x::UInt8)
    s2, s1 = adler
    s1 += x
    if s1 >= 65521
        s1 -= 65521
    end
    s2 += s1
    if s2 >= 65521
        s2 -= 65521
    end
    return (s2, s1)
end

function finish_adler(adler)
    s2, s1 = adler
    return (UInt32(s2) << 16) | UInt32(s1)
end

@inline function compute_adler_checksum(x::Vector{UInt8})
    adler = init_adler()
    for b in x
        adler = update_adler(adler, b)
    end
    return finish_adler(adler)
end

const crc_table = zeros(UInt32, 256)

function make_crc_table()
    for n = 1:256
        c = UInt32(n - 1)
        for k = 1:8
            if (c & 0x00000001) != 0
                c = 0xedb88320 ⊻ (c >> 1)
            else
                c >>= 1
            end
        end
        crc_table[n] = c
    end
end

function init_crc()
    if crc_table[1] == 0
        make_crc_table()
    end
    return 0xffffffff
end

@inline function update_crc(c::UInt32, x::UInt8)
    @inbounds return crc_table[1 + ((c ⊻ x) & 0xff)] ⊻ (c >> 8)
end

function finish_crc(c)
    return c ⊻ 0xffffffff
end

function crc(x::Vector{UInt8})
    c = init_crc()
    for b in x
        c = update_crc(c, b)
    end
    return finish_crc(c)
end

function read_zero_terminated_data(data::AbstractInflateData)
    s = UInt8[]
    while true
        c = get_aligned_byte(data)
        push!(s, c)
        if c == 0
            break
        end
    end
    return s
end

function read_zlib_header(data::AbstractInflateData)
    CMF = get_aligned_byte(data)
    FLG = get_aligned_byte(data)
    CM = CMF & 0x0f
    CINFO = CMF >> 4
    FLEVEL = FLG >> 6
    FDICT = (FLG >> 5) & 0x01
    if CM != 8
        error("unsupported compression method")
    end
    if CINFO > 7
        error("invalid LZ77 window size")
    end
    if FDICT != 0
        error("preset dictionary not supported")
    end
    if mod((UInt(CMF) << 8) | FLG, 31) != 0
        error("header checksum error")
    end
end

function read_gzip_header(data::AbstractInflateData, headers, compute_crc)
    data.update_input_crc = compute_crc
    ID1 = get_aligned_byte(data)
    ID2 = get_aligned_byte(data)
    if ID1 != 0x1f || ID2 != 0x8b
        error("not gzipped data")
    end
    CM = get_aligned_byte(data)
    if CM != 8
        error("unsupported compression method")
    end
    FLG = get_aligned_byte(data)
    MTIME = getbits(data, 32)
    XFL = get_aligned_byte(data)
    OS = get_aligned_byte(data)

    if headers != nothing
        headers["mtime"] = MTIME
        headers["os"] = OS
    end

    if (FLG & 0x04) != 0   # FLG.FEXTRA
        xlen = getbits(data, 16)
        if headers != nothing
            headers["fextra"] = zeros(UInt8, xlen)
        end

        for i = 1:xlen
            b = get_aligned_byte(data)
            if headers != nothing
                headers["fextra"][i] = b
            end
        end
    end

    if (FLG & 0x08) != 0   # FLG.FNAME
        name = read_zero_terminated_data(data)
        if headers != nothing
            headers["fname"] = String(name[1:end-1])
        end
    end

    if (FLG & 0x10) != 0   # FLG.FCOMMENT
        comment = read_zero_terminated_data(data)
        if headers != nothing
            headers["fcomment"] = String(comment[1:end-1])
        end
    end

    if (FLG & 0xe0) != 0
        error("reserved FLG bit set")
    end

    data.update_input_crc = false
    if (FLG & 0x02) != 0   # FLG.FHCRC
        crc16 = getbits(data, 16)
        if compute_crc
            header_crc = finish_crc(data.crc)
            if crc16 != (header_crc & 0xffff)
                error("corrupted data, header crc check failed")
            end
        end
    end
end

"""
    inflate(source::Vector{UInt8})

Decompress in-memory `source`, in unwrapped deflate format. The
output will also be a `Vector{UInt8}`. For a streaming counterpart,
see `InflateStream`.

Reference: [RFC 1951](https://www.ietf.org/rfc/rfc1951.txt)
"""
inflate(source::Vector{UInt8}) = _inflate(InflateData(source))

"""
    inflate_zlib(source::Vector{UInt8})

Decompress in-memory `source`, in Zlib compressed format. The
output will also be a `Vector{UInt8}`. For a streaming counterpart,
see `InflateZlibStream`.

    inflate_zlib(source::Vector{UInt8}; ignore_checksum = true)

Skip computing Adler checksum for consistency checking.

Reference: [RFC 1950](https://www.ietf.org/rfc/rfc1950.txt)
"""
function inflate_zlib(source::Vector{UInt8}; ignore_checksum = false)
    data = InflateData(source)
    read_zlib_header(data)

    out = _inflate(data)

    skip_bits_to_byte_boundary(data)
    stored_adler = 0
    for i = [24, 16, 8, 0]
        stored_adler |= Int(get_aligned_byte(data)) << i
    end
    if !ignore_checksum && compute_adler_checksum(out) != stored_adler
        error("corrupted data, adler checksum error")
    end

    return out
end

"""
    inflate_gzip(source::Vector{UInt8})

Decompress in-memory `source`, in Gzip compressed format. The
output will also be a `Vector{UInt8}`. For a streaming counterpart,
see `InflateGzipStream`.

    gzip_headers = Dict{String, Any}()
    inflate_gzip(source::Vector{UInt8}; headers = gzip_headers)

Also retrieve gzip headers in the provided `Dict`.

    inflate_gzip(source::Vector{UInt8}; ignore_checksum = true)

Skip computing CRC of the content, as well as the header, for
consistency checking.

Reference: [RFC 1952](https://www.ietf.org/rfc/rfc1952.txt)
"""
function inflate_gzip(source::Vector{UInt8}; headers = nothing,
                      ignore_checksum = false)
    data = InflateData(source)
    read_gzip_header(data, headers, !ignore_checksum)
    out = _inflate(data)

    skip_bits_to_byte_boundary(data)
    crc32 = getbits(data, 32)
    if !ignore_checksum && crc32 != crc(out)
        error("corrupted data, crc check failed")
    end
    isize = getbits(data, 32)
    if isize != length(out)
        error("corrupted data, length check failed")
    end

    return out
end

"""
    inflate_gzip(filename::AbstractString)

Convenience wrapper for reading a gzip compressed text file. The
result is returned as a string.
"""
function inflate_gzip(filename::AbstractString; kwargs...)
    return String(inflate_gzip(read(filename); kwargs...))
end

### Streaming interface. ###

# Max distance for repeated strings is 32768 (RFC 1951). Make it
# twice the size for better performance.
const buffer_size = 65536

mutable struct StreamingInflateData <: AbstractInflateData
    stream::IO
    input_buffer::Vector{UInt8}
    input_buffer_pos::Int
    current_byte::Int
    bitpos::Int
    literal_or_length_code::Vector{Vector{Int}}
    distance_code::Vector{Vector{Int}}
    output_buffer::Vector{UInt8}
    write_pos::Int
    read_pos::Int
    waiting_for_new_block::Bool
    pending_bytes::Int
    distance::Int
    reading_final_block::Bool
    update_input_crc::Bool
    crc::UInt32
end

function StreamingInflateData(stream::IO)
    return StreamingInflateData(stream, UInt8[], 1, 0, 0,
                                fixed_literal_or_length_table,
                                fixed_distance_table,
                                zeros(UInt8, buffer_size), 1, 1,
                                true, 0, -2, false, false, init_crc())
end

function get_input_byte(data::StreamingInflateData)
    if data.input_buffer_pos > length(data.input_buffer)
        data.input_buffer = read(data.stream, 65536)
        data.input_buffer_pos = 1
    end
    byte = data.input_buffer[data.input_buffer_pos]
    data.input_buffer_pos += 1
    if data.update_input_crc
        data.crc = update_crc(data.crc, byte)
    end
    return byte
end

# This isn't called when reading Gzip header, so no need to
# consider updating crc.
function get_input_bytes(data::StreamingInflateData, n)
    if data.input_buffer_pos > length(data.input_buffer)
        data.input_buffer = read(data.stream, 65536)
        data.input_buffer_pos = 1
    end
    n = min(n, length(data.input_buffer) - data.input_buffer_pos + 1)
    bytes = @view data.input_buffer[data.input_buffer_pos:(data.input_buffer_pos + n - 1)]
    data.input_buffer_pos += n
    return bytes
end

abstract type AbstractInflateStream <: IO end

"""
    InflateStream(stream::IO)

Streaming decompression of unwrapped deflate compressed `stream`. For
an in-memory counterpart, see `inflate`.

Reference: [RFC 1951](https://www.ietf.org/rfc/rfc1951.txt)
"""
mutable struct InflateStream <: AbstractInflateStream
    data::StreamingInflateData
end

function InflateStream(stream::IO)
    stream = InflateStream(StreamingInflateData(stream))
    getbyte(stream)
    return stream
end

"""
    InflateZlibStream(stream::IO)

Streaming decompression of Zlib compressed `stream`. For an in-memory
counterpart, see `inflate_zlib`.

    InflateZlibStream(stream::IO; ignore_checksum = true)

Skip computing Adler checksum for consistency checking.

Reference: [RFC 1950](https://www.ietf.org/rfc/rfc1950.txt)
"""
mutable struct InflateZlibStream <: AbstractInflateStream
    data::StreamingInflateData
    adler::Tuple{Int, Int}
    compute_adler::Bool
end

function InflateZlibStream(stream::IO; ignore_checksum = false)
    stream = InflateZlibStream(StreamingInflateData(stream), init_adler(),
                               !ignore_checksum)
    read_zlib_header(stream.data)
    getbyte(stream)
    if eof(stream)
        read_trailer(stream)
    end
    return stream
end

"""
    InflateGzipStream(stream::IO)

Streaming decompression of Gzip compressed `stream`. For an in-memory
counterpart, see `inflate_gzip`.

    gzip_headers = Dict{String, Any}()
    InflateGzipStream(stream::IO; headers = gzip_headers)

Also retrieve gzip headers in the provided `Dict`. The headers are
available directly after the object is constructed.

    InflateGzipStream(stream::IO; ignore_checksum = true)

Skip computing CRC of the content, as well as the header, for
consistency checking.

Reference: [RFC 1952](https://www.ietf.org/rfc/rfc1952.txt)
"""
mutable struct InflateGzipStream <: AbstractInflateStream
    data::StreamingInflateData
    crc::UInt32
    num_bytes::Int
    compute_crc::Bool
end

function InflateGzipStream(stream::IO; headers = nothing,
                           ignore_checksum = false)
    stream = InflateGzipStream(StreamingInflateData(stream), init_crc(), 0,
                               !ignore_checksum)
    read_gzip_header(stream.data, headers, !ignore_checksum)
    getbyte(stream)
    if eof(stream)
        read_trailer(stream)
    end
    return stream
end

read_trailer(stream::InflateStream) = nothing

function read_trailer(stream::InflateZlibStream)
    computed_adler = finish_adler(stream.adler)
    skip_bits_to_byte_boundary(stream.data)
    stored_adler = 0
    for i = [24, 16, 8, 0]
        stored_adler |= Int(get_aligned_byte(stream.data)) << i
    end
    if stream.compute_adler && computed_adler != stored_adler
        error("corrupted data, adler checksum error")
    end
end

function read_trailer(stream::InflateGzipStream)
    crc = finish_crc(stream.crc)
    skip_bits_to_byte_boundary(stream.data)
    crc32 = getbits(stream.data, 32)
    if stream.compute_crc && crc32 != crc
        error("corrupted data, crc check failed")
    end
    isize = getbits(stream.data, 32)
    if isize != stream.num_bytes
        error("corrupted data, length check failed")
    end
end

@inline function read_output_byte(data::StreamingInflateData)
    @inbounds byte = data.output_buffer[data.read_pos]
    data.read_pos += 1
    if data.read_pos > buffer_size
        data.read_pos -= buffer_size
    end
    return byte
end

function read_output_bytes!(data::StreamingInflateData, out, i)
    if data.read_pos < data.write_pos
        m = data.write_pos - data.read_pos
    else
        m = buffer_size - data.read_pos + 1
    end
    if m == 1
        out[i] = read_output_byte(data)
        return 1
    end
    n = length(out) - i + 1
    n = min(m, n)
    copyto!(out, i, data.output_buffer, data.read_pos, n)
    data.read_pos += n
    if data.read_pos > buffer_size
        data.read_pos -= buffer_size
    end
    return n
end

function read_output_byte(stream::AbstractInflateStream)
    byte = read_output_byte(stream.data)
    getbyte(stream)
    return byte
end

function read_output_bytes!(stream::AbstractInflateStream, out, n)
    i = 1
    resized = false
    while i <= n && !eof(stream)
        if length(out) < i
            resize!(out, min(length(out) * 2, n))
            resized = true
        end
        i += read_output_bytes!(stream.data, out, i)
        getbyte(stream)
    end
    i -= 1
    if resized
        resize!(out, i)
    end
    return i
end

function write_to_buffer(data::StreamingInflateData, x::UInt8)
    data.output_buffer[data.write_pos] = x
    data.write_pos += 1
    if data.write_pos > buffer_size
        data.write_pos -= buffer_size
    end
end

function write_to_buffer(data::StreamingInflateData, x::AbstractVector{UInt8})
    n = length(x)
    copyto!(data.output_buffer, data.write_pos, x, 1, n)
    data.write_pos += n
    if data.write_pos > buffer_size
        data.write_pos -= buffer_size
    end
end

function write_to_buffer(stream::InflateStream, x)
    write_to_buffer(stream.data, x)
end

function write_to_buffer(stream::InflateZlibStream, x::UInt8)
    write_to_buffer(stream.data, x)
    if stream.compute_adler
        stream.adler = update_adler(stream.adler, x)
    end
end

function write_to_buffer(stream::InflateZlibStream, x::AbstractVector{UInt8})
    write_to_buffer(stream.data, x)
    if stream.compute_adler
        for y in x
            stream.adler = update_adler(stream.adler, y)
        end
    end
end

function write_to_buffer(stream::InflateGzipStream, x::UInt8)
    write_to_buffer(stream.data, x)
    if stream.compute_crc
        stream.crc = update_crc(stream.crc, x)
    end
    stream.num_bytes += 1
end

function write_to_buffer(stream::InflateGzipStream, x::AbstractVector{UInt8})
    write_to_buffer(stream.data, x)
    if stream.compute_crc
        for y in x
            stream.crc = update_crc(stream.crc, y)
        end
    end
    stream.num_bytes += length(x)
end

function getbyte(stream::AbstractInflateStream)
    if stream.data.write_pos != stream.data.read_pos
        return
    end

    if stream.data.pending_bytes > 0
        if stream.data.read_pos <= stream.data.write_pos
            n = buffer_size - stream.data.write_pos + 1
        else
            n = stream.data.read_pos - stream.data.write_pos - 1
        end
        n = min(stream.data.pending_bytes, n)
        if stream.data.distance < 0
            # Incompressible data.
            bytes = get_input_bytes(stream.data, n)
            stream.data.pending_bytes -= length(bytes)
            write_to_buffer(stream, bytes)
        else
            pos = stream.data.write_pos - stream.data.distance
            if pos <= 0
                n = min(n, 1 - pos)
                pos += buffer_size
            end
            stream.data.pending_bytes -= n
            write_to_buffer(stream, @view stream.data.output_buffer[pos:(pos + n - 1)])
        end
        return
    end

    if stream.data.waiting_for_new_block
        if stream.data.reading_final_block
            return
        end
        stream.data.reading_final_block = getbits(stream.data, 1) == 1
        compression_mode = getbits(stream.data, 2)
        if compression_mode == 0
            skip_bits_to_byte_boundary(stream.data)
            len = getbits(stream.data, 16)
            nlen = getbits(stream.data, 16)
            if len ⊻ nlen != 0xffff
                error("corrupted data")
            end
            stream.data.distance = -1
            stream.data.pending_bytes = len
            getbyte(stream)
            return
        elseif compression_mode == 1
            stream.data.literal_or_length_code = fixed_literal_or_length_table
            stream.data.distance_code = fixed_distance_table
        elseif compression_mode == 2
            read_code_tables(stream.data)
        else
            error("invalid block compression mode 3")
        end
        stream.data.waiting_for_new_block = false
    end

    v = get_literal_or_length(stream.data)
    if v < 256
        write_to_buffer(stream, UInt8(v))
    elseif v == 256
        stream.data.waiting_for_new_block = true
        getbyte(stream)
    else
        stream.data.pending_bytes = getlength(stream.data, v)
        stream.data.distance = getdist(stream.data)
        getbyte(stream)
    end
end

function Base.eof(stream::AbstractInflateStream)
    return stream.data.read_pos == stream.data.write_pos
end

function Base.read(stream::AbstractInflateStream, ::Type{UInt8})
    if eof(stream)
        throw(EOFError())
    end

    byte = read_output_byte(stream)

    if eof(stream)
        read_trailer(stream)
    end

    return byte
end

function Base.readbytes!(stream::AbstractInflateStream,
                         b::AbstractVector{UInt8}, nb = length(b))
    if eof(stream)
        return 0
    end

    n = read_output_bytes!(stream, b, nb)

    if eof(stream)
        read_trailer(stream)
    end

    return n
end

end # module


text:
#####
##### transformations
#####

"""
$(TYPEDEF)

Abstract type used internally for simplifying forwarding. Assumes a slot `ℓ`, and forwards
`dimension` to it.
"""
abstract type LogDensityTransformation <: SamplingLogDensity end

dimension(ℓ::LogDensityTransformation) = dimension(ℓ.ℓ)

####
#### linear transformation
####

struct Linear{L,M,S,T} <: LogDensityTransformation
    ℓ::L
    A::M
    divA::S
    logabsdetA::T
end

"""
$(SIGNATURES)

Internal method for returning an object that makes `A \\ x` fast and stable.
"""
_fastdiv(A::AbstractMatrix) = lu(A)
_fastdiv(A::AbstractTriangular) = A
_fastdiv(A::Diagonal) = A

"""
$(SIGNATURES)

Internal method for logabsdet, work around
https://github.com/JuliaLang/julia/issues/32988
"""
_logabsdet(A::AbstractMatrix) = first(logabsdet(A))
_logabsdet(A::Diagonal) = sum(log ∘ abs, diag(A))

"""
$(SIGNATURES)

Transform a distribution on `x` to `y = Ax`, where `A` is a conformable square matrix.

Since the log Jacobian is constant, it is dropped in the log density.
"""
function linear(A::AbstractMatrix, ℓ)
    K = dimension(ℓ)
    @argcheck checksquare(A) == K
    Linear(ℓ, A, _fastdiv(A), _logabsdet(A))
end

logdensity(ℓ::Linear, x) = logdensity(ℓ.ℓ, ℓ.divA \ x) - ℓ.logabsdetA

function logdensity_and_gradient(ℓ::Linear, x)
    f, ∇f = logdensity_and_gradient(ℓ.ℓ, ℓ.divA \ x)
    f - ℓ.logabsdetA, (ℓ.divA') \ ∇f # note here \ is needed because / is not defined for LU
end

hypercube_transform(ℓ::Linear, x) = ℓ.A * hypercube_transform(ℓ.ℓ, x)

####
#### shift (translation)
####

struct Shift{L, T <: AbstractVector} <: LogDensityTransformation
    ℓ::L
    b::T
end

"""
$(SIGNATURES)

Transform a distribution on `x` to `y = x + b`, where `b` is a conformable vector.
"""
function shift(b::AbstractVector, ℓ)
    @argcheck length(b) == dimension(ℓ)
    Shift(ℓ, b)
end

logdensity(ℓ::Shift, x) = logdensity(ℓ.ℓ, x - ℓ.b)

# The log Jacobian adjustment is zero
logdensity_and_gradient(ℓ::Shift, x) = logdensity_and_gradient(ℓ.ℓ, x - ℓ.b)

hypercube_transform(ℓ::Shift, x) = ℓ.b .+ hypercube_transform(ℓ.ℓ, x)

####
#### elongate
####

struct Elongate{L, T <: Real} <: LogDensityTransformation
    ℓ::L
    k::T
end

"""
$(SIGNATURES)

Transform a distribution on `x` to ``y = (1 + ‖x‖²)ᵏ⋅x`` where ``‖ ‖`` is the Euclidean norm
and `k` is a real number. `k > 0` values make the tails heavier.
"""
function elongate(k::Real, ℓ)
    Elongate(ℓ, k)
end

"""
$(SIGNATURES)

Helper function that solves
```math
y = x ⋅ (1 + x²)ᵏ
```
for `x`, given `y ≥ 0` and `k`, using Newton's method.

Absolute tolerance `atol` should not be set *too* low to avoid occasional numerical cycles
near the root.
"""
function _find_x_norm(y, k; x = zero(y), atol = 16 * eps(y))
    for _ in 1:100
        x2 = abs2(x)
        A = (1 + x2)^k
        f = A * x - y
        abs(f) ≤ atol && return x
        B = 2 * k * x2 * (1 + x2)^(k - 1)
        x = (B * x + y) / (A + B)
    end
    error("internal error: reached maximum number of iterations, y = $(y), k = $(k)")
end

function logdensity_and_gradient(ℓ::Elongate, y)
    @unpack ℓ, k = ℓ
    d = dimension(ℓ)
    ynorm = norm(y, 2)
    xnorm = _find_x_norm(ynorm, k)
    xnorm2 = abs2(xnorm)
    D = (1 + xnorm2)^(-k)       # x = D ⋅ y
    x = D .* y
    ℓx, ∇ℓx = logdensity_and_gradient(ℓ, x)
    ℓy = ℓx - (k*d - 1) * log1p(xnorm2) - log1p((1 + 2*k) * xnorm2)
    A = 1 + xnorm2
    B = 1 + (1 + 2*k) * xnorm2
    L1 = (I - (2 * k / B) .* (x * x')) * ∇ℓx .* D
    L2 = (((k * d - 1)/A + (1 + 2 * k)/B) * 2 * A^(1 - 2*k) / B) .* y
    ℓy, L1 .- L2
end

function hypercube_transform(ℓ::Elongate, z)
    @unpack ℓ, k = ℓ
    x = hypercube_transform(ℓ, z)
    (1 + sum(abs2, x))^k .* x
end


text:
using MLDatasets
using HDF5
using StatsBase
using LinearAlgebra

# load full training set
train_x, train_y = MNIST.traindata()
# load full test set
# test_x,  test_y  = MNIST.testdata()

# u = 40 # from "VISUALIZING DATA USING T-SNE"
u = 50 # from "Accelerating t-SNE using Tree-Based Algorithms"
n = size(train_x, 3)

X = train_x
X = reshape(X, :, n) # vectorize images
labels = train_y
if false
    k = 1024
    subsample = sample(1:n, k, replace = false)
    X = X[:, subsample]
    labels = labels[subsample]
end

# reduce dimensionality using PCA to
U, S, Vt = svd(X)
V = Vt'
d = 50
X = V[1:d, :]
@. X *= S[1:d]

using StochasticNeighborhoodEmbeddings
using StochasticNeighborhoodEmbeddings: OptimizationParameters,
            optimized_symmetrized_neighbor, initialize_embedding!

total_iter = 64
early_iter = 64
max_iter = total_iter - early_iter

beta_1 = 0.9
beta_2 = 0.99
weight_decay = 0
lr = 3e-1
dofast = false

early_beta_1 = 0.5
early_beta_2 = 0.9
early_lr = 3e-1
early_compression = 0
early_exaggeration = 8

params = OptimizationParameters(max_iter = max_iter, beta_1 = beta_1,
        beta_2 = beta_2, lr = lr, weight_decay = weight_decay,
        early_iter = early_iter, early_beta_1 = early_beta_1, early_beta_2 = early_beta_2,
        early_lr = early_lr, early_compression = early_compression,
        early_exaggeration = early_exaggeration, dofast = dofast,
        verbose = true, debias = false)

Y = 1e-2randn(params.d, size(X, 2))
P = optimized_symmetrized_neighbor(X, u, params.dofast)
G = zeros(params.d, size(X, 2)) # gradient of Y
G² = zeros(params.d, size(X, 2)) # second moments of gradient
exaggerations = 8:-1:1
for i in exaggerations # slowly reducing early_exaggeration
    params.early_exaggeration = i
    @time initialize_embedding!(Y, G, G², P, params)
end
params.early_iter = 0
params.max_iter = 1024 - length(exaggerations) * early_iter
params.early_exaggeration = 0
@time tsne!(Y, G, G², P, params) # fine-tuning

f = h5open("MNIST_TSNE_full.h5", "w")
f["Y"] = Y
f["X"] = X
f["labels"] = labels
f["max_iter"] = max_iter
f["lr"] = lr
f["beta_1"] = beta_1
f["beta_2"] = beta_2
f["early_iter"] = early_iter
f["early_lr"] = early_lr
f["early_beta_1"] = early_beta_1
f["early_beta_2"] = early_beta_2
f["early_compression"] = early_compression
f["early_exaggeration"] = early_exaggeration
close(f)

# using Plots
# doplot = true
# if doplot
#     YC = [Y[:, ==(i).(labels)] for i in 1:10]
#     scatter()
#     for (i, yc) in enumerate(YC)
#         scatter!(yc[1, :], yc[2, :], label = "$(i-1)")
#     end
#     gui()
# end


text:
using UKCarbonIntensityData
using UKCarbonIntensityData: _unpack_df, _unpack_values, _parse_datetime, _parse_data
using UKCarbonIntensityData: _convert_mix
using DataFrames
using Dates
using Test
using TimeZones

function _test_df_inequality(a, b)
    @test issetequal(names(a), names(b))
    @test select(a, sort(names(a))) == select(b, sort(names(a)))
end

@testset "UKCarbonIntensityData.jl" begin
    include("fetch.jl")
end


text:
export Shader, set, loadShader, shader, ShaderType, resetShader

function initShaders()
	# basicShapes
	const vshBS = """
	$(get_glsl_version_string())

	in vec3 position;
	in vec4 color;

	out vec4 vColor;

	uniform mat4 MVP;

	void main() {
		vColor = color;
		gl_Position = MVP * vec4(position, 1.0);
	}
	"""

	# const gshBS = """
	# $(get_glsl_version_string())

	# #extension GL_EXT_gpu_shader4 : enable
	# #extension GL_EXT_geometry_shader4 : enable

	# varying in vec3 vertWorldPos[3];
	# varying in vec3 vertWorldNormal[3];
	# varying in vec4 vColor[4];

	# varying out vec3 worldNormal;
	# varying out vec3 worldPos;
	# varying out vec4 OUTColor;

	# uniform vec2 WIN_SCALE;

	# noperspective varying vec3 dist;

	# void main(void)
	# {
	#	vec2 p0 = WIN_SCALE * gl_PositionIn[0].xy/gl_PositionIn[0].w;
	#	vec2 p1 = WIN_SCALE * gl_PositionIn[1].xy/gl_PositionIn[1].w;
	#	vec2 p2 = WIN_SCALE * gl_PositionIn[2].xy/gl_PositionIn[2].w;

	#	vec2 v0 = p2-p1;
	#	vec2 v1 = p2-p0;
	#	vec2 v2 = p1-p0;

	#	float area = abs(v1.x*v2.y - v1.y * v2.x);

	#	dist = vec3(area/length(v0),0,0);
	#	worldPos = vertWorldPos[0];
	#	worldNormal = vertWorldNormal[0];
	#	gl_Position = gl_PositionIn[0];
	#	EmitVertex();

	#	dist = vec3(0,area/length(v1),0);
	#	worldPos = vertWorldPos[1];
	#	worldNormal = vertWorldNormal[1];
	#	gl_Position = gl_PositionIn[1];
	#	EmitVertex();

	#	dist = vec3(0,0,area/length(v2));
	#	worldPos = vertWorldPos[2];
	#	worldNormal = vertWorldNormal[2];
	#	gl_Position = gl_PositionIn[2];
	#	EmitVertex();

	#	EndPrimitive();
	# }
	# """

	const fshBS = """
	$(get_glsl_version_string())

	in vec4 vColor;

	out vec4 outColor;

	void main() {
		outColor = vColor;
	}
	"""

	push!(globjs.vaos, glGenVertexArray()) # basic shape
	glBindVertexArray(globjs.vaos[1])
    push!(globjs.ebos,  glGenBuffer()) # basic shape
	glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, globjs.ebos[1])

    vertexShader = createShader(vshBS, GL_VERTEX_SHADER)
    fragmentShader = createShader(fshBS, GL_FRAGMENT_SHADER)
    state.program = createShaderProgram(vertexShader, fragmentShader)
    shaderBank["basicShapes"] = state.program
    glUseProgram(shaderBank["basicShapes"])

    push!(globjs.posvbos, glGenBuffer()) # basic shape
    glBindBuffer(GL_ARRAY_BUFFER, globjs.posvbos[1])
    positionAttribute = glGetAttribLocation(state.program, "position")
    glEnableVertexAttribArray(positionAttribute)
    glVertexAttribPointer(positionAttribute, 3, GL_FLOAT, false, 3*sizeof(GLfloat), 0)

    push!(globjs.colvbos, glGenBuffer()) # basic shape
    glBindBuffer(GL_ARRAY_BUFFER, globjs.colvbos[1])
    colorAttribute = glGetAttribLocation(state.program, "color")
    glEnableVertexAttribArray(colorAttribute)
    glVertexAttribPointer(colorAttribute, 4, GL_FLOAT, false, 4*sizeof(GLfloat), 0)

	glUniformMatrix4fv(glGetUniformLocation(state.program, "MVP"), 1, false, GLmatState.currMatrix)

	# texturedShapes
	const vshTS = """
	$(get_glsl_version_string())

	in vec3 position;
	in vec4 color;
	in vec2 texcoord;

	out vec4 vColor;
	out vec2 Texcoord;

	uniform mat4 MVP;

	void main() {
		vColor = color;
		Texcoord = texcoord;
		gl_Position = MVP * vec4(position, 1.0);
	}
	"""

	const fshTS = """
	$(get_glsl_version_string())

	in vec4 vColor;
	in vec2 Texcoord;

	out vec4 outColor;

	uniform sampler2D texture;

	void main() {
		outColor = texture(texture, Texcoord) * vColor;
	}
	"""

	push!(globjs.vaos, glGenVertexArray()) # textured shapes
	glBindVertexArray(globjs.vaos[2])
    push!(globjs.ebos,  glGenBuffer()) # textured shapes
	glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, globjs.ebos[2])

	vertexShader = createShader(vshTS, GL_VERTEX_SHADER)
	fragmentShader = createShader(fshTS, GL_FRAGMENT_SHADER)
	shaderBank["texturedShapes"] = createShaderProgram(vertexShader, fragmentShader)
	glUseProgram(shaderBank["texturedShapes"])

	push!(globjs.posvbos, glGenBuffer()) # textured shape
    glBindBuffer(GL_ARRAY_BUFFER, globjs.posvbos[2])
	positionAttribute = glGetAttribLocation(shaderBank["texturedShapes"], "position")
	glEnableVertexAttribArray(positionAttribute)
	glVertexAttribPointer(positionAttribute, 3, GL_FLOAT, false, 3*sizeof(GLfloat), 0)

	push!(globjs.colvbos, glGenBuffer()) # textured shape
    glBindBuffer(GL_ARRAY_BUFFER, globjs.colvbos[2])
	colorAttribute = glGetAttribLocation(shaderBank["texturedShapes"], "color")
	glEnableVertexAttribArray(colorAttribute)
	glVertexAttribPointer(colorAttribute, 4, GL_FLOAT, false, 4*sizeof(GLfloat), 0)

	push!(globjs.texvbos, glGenBuffer()) # textured shape
    glBindBuffer(GL_ARRAY_BUFFER, globjs.texvbos[1])
	texAttrib = glGetAttribLocation(shaderBank["texturedShapes"], "texcoord")
	glEnableVertexAttribArray(texAttrib)
	glVertexAttribPointer(texAttrib, 2, GL_FLOAT, false, 2*sizeof(GLfloat), 0)

	glUniformMatrix4fv(glGetUniformLocation(shaderBank["texturedShapes"], "MVP"), 1, false, GLmatState.currMatrix)
	glUniform1i(glGetUniformLocation(shaderBank["texturedShapes"], "texture"), 2)

	# fontDrawing
	const vshFD = """
	$(get_glsl_version_string())

	in vec2 position;
	in vec2 texcoord;

	out vec2 TexCoord;

	uniform mat4 proj;

	void main()
	{
		gl_Position = proj * vec4(position, 0.0, 1.0);
		TexCoord = texcoord;
	}
	"""

	const fshFD = """
	$(get_glsl_version_string())

	in vec2 TexCoord;

	out vec4 color;

	uniform sampler2D text;
	uniform vec3 textColor;

	void main()
	{
		vec4 sampled = vec4(1.0, 1.0, 1.0, texture(text, TexCoord).r);
		color = vec4(textColor, 1.0) * sampled;
	}
	"""

	push!(globjs.vaos, glGenVertexArray()) # font drawing
	glBindVertexArray(globjs.vaos[3])

	vertexShader = createShader(vshFD, GL_VERTEX_SHADER)
	fragmentShader = createShader(fshFD, GL_FRAGMENT_SHADER)
	shaderBank["fontDrawing"] = createShaderProgram(vertexShader, fragmentShader)
	glUseProgram(shaderBank["fontDrawing"])

	push!(globjs.posvbos, glGenBuffer()) # font drawing
    glBindBuffer(GL_ARRAY_BUFFER, globjs.posvbos[3])
	positionAttribute = glGetAttribLocation(shaderBank["fontDrawing"], "position")
	glEnableVertexAttribArray(positionAttribute)
	glVertexAttribPointer(positionAttribute, 2, GL_FLOAT, false, 2*sizeof(GLfloat), 0)

	push!(globjs.texvbos, glGenBuffer()) # font drawing
    glBindBuffer(GL_ARRAY_BUFFER, globjs.texvbos[2])
	texAttrib = glGetAttribLocation(shaderBank["fontDrawing"], "texcoord")
	glEnableVertexAttribArray(texAttrib)
	glVertexAttribPointer(texAttrib, 2, GL_FLOAT, false, 2*sizeof(GLfloat), 0)

	glUniform3f(glGetUniformLocation(shaderBank["fontDrawing"], "textColor"), GLfloat(state.fillCol[1].r), GLfloat(state.fillCol[1].g), GLfloat(state.fillCol[1].b))

	# by default, we always use texture 1 for fonts
	glUniform1i(glGetUniformLocation(shaderBank["fontDrawing"], "text"), 1)

	# text is rendered with an orthographic projection
	projection = GLfloat[2/state.width 0 0 -1;
						0 2/state.height 0 -1;
						0 0 -1 0;
						0 0 0 1]

	glUniformMatrix4fv(glGetUniformLocation(shaderBank["fontDrawing"], "proj"), 1, false, projection)

	# framebuffer
	const vshDFB = """
	$(get_glsl_version_string())

	in vec2 position;
	in vec2 texcoord;

	out vec2 Texcoord;

	void main() {
		Texcoord = texcoord;
		gl_Position = vec4(position, 0.0, 1.0);
	}
	"""

	const fshDFB = """
	$(get_glsl_version_string())

	in vec2 Texcoord;

	out vec4 outColor;

	uniform sampler2D texFramebuffer;

	void main() {
		outColor = texture(texFramebuffer, Texcoord);
	}
	"""

	push!(globjs.vaos, glGenVertexArray()) # quad with framebuffer texture
	glBindVertexArray(globjs.vaos[4])

	vertexShader = createShader(vshDFB, GL_VERTEX_SHADER)
	fragmentShader = createShader(fshDFB, GL_FRAGMENT_SHADER)
	shaderBank["drawFramebuffer"] = createShaderProgram(vertexShader, fragmentShader)
	glUseProgram(shaderBank["drawFramebuffer"])

	push!(globjs.posvbos, glGenBuffer()) # quad with framebuffer texture
	glBindBuffer(GL_ARRAY_BUFFER, globjs.posvbos[4])
	positionAttribute = glGetAttribLocation(shaderBank["drawFramebuffer"], "position")
	glEnableVertexAttribArray(positionAttribute)
	glVertexAttribPointer(positionAttribute, 2, GL_FLOAT, false, 2*sizeof(GLfloat), 0)

	# fill the quad data on the gpu for later when we render the
	# framebuffer texture to the screen.
	# this only needs to be done once, since we never re-define the drawing
	# framebuffer.
	glBufferData(GL_ARRAY_BUFFER, sizeof(quadPosCoords), quadPosCoords, GL_STATIC_DRAW)

	push!(globjs.texvbos, glGenBuffer()) # quad with framebuffer texture
	glBindBuffer(GL_ARRAY_BUFFER, globjs.texvbos[3])
	texAttrib = glGetAttribLocation(shaderBank["drawFramebuffer"], "texcoord")
	glEnableVertexAttribArray(texAttrib)
	glVertexAttribPointer(texAttrib, 2, GL_FLOAT, false, 2*sizeof(GLfloat), 0)

	# same idea as the buffer allocation just above.
	glBufferData(GL_ARRAY_BUFFER, sizeof(quadTexCoords), quadTexCoords, GL_STATIC_DRAW)

	glUniform1i(glGetUniformLocation(shaderBank["drawFramebuffer"], "texFramebuffer"), 0)
end

@enum ShaderType POINT=1 LINE=2 TRIANGLE=3

type Shader
	shaderName::String
	fragFilename::String
	vertFilename::String
	vaoind::Int
	posind::Int
	colind::Int
	texind::Int
	drawTexture::Bool
end

Shader(sn, ff, vf) = Shader(sn, ff, vf, 0, 0, 0, 0, false)

function shader(whichShader::String, shaderType::ShaderType=TRIANGLE)
	state.program = shaderBank[whichShader]
	glUseProgram(shaderBank[whichShader])

	if whichShader == "basicShapes"
		glBindVertexArray(globjs.vaos[1])
		globjs.posind = 1
		globjs.colind = 1
	elseif whichShader == "texturedShapes"
		glBindVertexArray(globjs.vaos[2])
		globjs.posind = 2
		globjs.colind = 2
		globjs.texind = 1
	elseif whichShader == "fontDrawing"
		glBindVertexArray(globjs.vaos[3])
		globjs.posind = 3
		globjs.texind = 2

		glUniform3f(glGetUniformLocation(shaderBank["fontDrawing"], "textColor"), GLfloat(state.fillCol[1].r), GLfloat(state.fillCol[1].g), GLfloat(state.fillCol[1].b))
	elseif whichShader == "drawFramebuffer"
		glBindVertexArray(globjs.vaos[4])
		globjs.posind = 4
		globjs.texind = 3
	end

	if whichShader == "texturedShapes"
		state.drawTexture = true
	else
		state.drawTexture = false
	end
end

function shader(whichShader::Shader, shaderType::ShaderType=TRIANGLE)
	state.program = shaderBank[whichShader.shaderName]
	glUseProgram(shaderBank[whichShader.shaderName])

	glBindVertexArray(globjs.vaos[whichShader.vaoind])
	globjs.posind = whichShader.posind
	globjs.colind = whichShader.colind
	state.drawTexture = whichShader.drawTexture
	if whichShader.drawTexture
		globjs.texind = whichShader.texind
	end

	state.alternateShader = whichShader
end

function set(shader::Shader, uniformName::String, value)
	ul = glGetUniformLocation(shaderBank[shader.shaderName], uniformName)
	if length(value) == 1
		glUniform1f(ul, value)
	elseif length(value) == 3
		glUniform3fv(ul, GLsizei(1), Ref(value))
	elseif value <: AbstractArray
		glUniformMatrix4fv(ul, GLsizei(1), GL_FALSE, Ref(value))
	end
end

function loadShader(fragFilename)
	ff = open(fragFilename, "r")
	fsh = readall(ff)
	close(ff)

	shader = Shader(randstring(), fragFilename, "")

	vsh = """
	$(get_glsl_version_string())

	in vec3 position;
	in vec4 color;

	out vec4 vColor;

	uniform mat4 MVP;

	void main() {
		vColor = color;
		gl_Position = MVP * vec4(position, 1.0);
	}
	"""

	vertexShader = createShader(vsh, GL_VERTEX_SHADER)
	fragmentShader = createShader(fsh, GL_FRAGMENT_SHADER)
	program = createShaderProgram(vertexShader, fragmentShader)
	shaderBank[shader.shaderName] = program
	glUseProgram(program)

	push!(globjs.vaos, glGenVertexArray())
	glBindVertexArray(globjs.vaos[end])
	shader.vaoind = length(globjs.vaos)
    push!(globjs.ebos, glGenBuffer())
	glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, globjs.ebos[end])

	push!(globjs.posvbos, glGenBuffer())
    glBindBuffer(GL_ARRAY_BUFFER, globjs.posvbos[end])
    shader.posind = length(globjs.posvbos)
    positionAttribute = glGetAttribLocation(program, "position")
    glEnableVertexAttribArray(positionAttribute)
    glVertexAttribPointer(positionAttribute, 3, GL_FLOAT, false, 3*sizeof(GLfloat), 0)

    push!(globjs.colvbos, glGenBuffer())
    glBindBuffer(GL_ARRAY_BUFFER, globjs.colvbos[end])
    shader.colind = length(globjs.colvbos)
    colorAttribute = glGetAttribLocation(program, "color")
    glEnableVertexAttribArray(colorAttribute)
    glVertexAttribPointer(colorAttribute, 4, GL_FLOAT, false, 4*sizeof(GLfloat), 0)

	glUniformMatrix4fv(glGetUniformLocation(program, "MVP"), 1, false, GLmatState.currMatrix)

	glUseProgram(state.program)

	return shader
end

function loadShader(fragFilename, vertFilename)
	ff = open(fragFilename, "r")
	fsh = readall(ff)
	close(ff)

	vf = open(vertFilename, "r")
	vsh = readall(vf)
	close(vf)

	shader = Shader(randstring(), fragFilename, vertFilename)

	vertexShader = createShader(vsh, GL_VERTEX_SHADER)
	fragmentShader = createShader(fsh, GL_FRAGMENT_SHADER)
	program = createShaderProgram(vertexShader, fragmentShader)
	shaderBank[shader.shaderName] = program
	glUseProgram(program)

	push!(globjs.vaos, glGenVertexArray())
	glBindVertexArray(globjs.vaos[end])
	shader.vaoind = length(globjs.vaos)
    push!(globjs.ebos, glGenBuffer())
	glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, globjs.ebos[end])

	push!(globjs.posvbos, glGenBuffer())
    glBindBuffer(GL_ARRAY_BUFFER, globjs.posvbos[end])
    shader.posind = length(globjs.posvbos)
    positionAttribute = glGetAttribLocation(program, "position")
    glEnableVertexAttribArray(positionAttribute)
    glVertexAttribPointer(positionAttribute, 3, GL_FLOAT, false, 3*sizeof(GLfloat), 0)

    push!(globjs.colvbos, glGenBuffer())
    glBindBuffer(GL_ARRAY_BUFFER, globjs.colvbos[end])
    shader.colind = length(globjs.colvbos)
    colorAttribute = glGetAttribLocation(program, "color")
    glEnableVertexAttribArray(colorAttribute)
    glVertexAttribPointer(colorAttribute, 4, GL_FLOAT, false, 4*sizeof(GLfloat), 0)

	return shader
end

function resetShader()
	shader("basicShapes")
	state.alternateShader = nothing
end


text:
@testset "Transition flux plots" begin

excelTransPlot( mpSim, "excelrequest/plotRequest" )

end  # @testset "Transition flux plots"


text:
using ClimateMachine

using ClimateMachine.Atmos
using ClimateMachine.Orientations
using ClimateMachine.ConfigTypes
using ClimateMachine.Diagnostics
using ClimateMachine.GenericCallbacks
using ClimateMachine.ODESolvers
using ClimateMachine.SystemSolvers
using ClimateMachine.Mesh.Filters
using ClimateMachine.Thermodynamics
using ClimateMachine.TemperatureProfiles
using ClimateMachine.TurbulenceClosures
using ClimateMachine.VariableTemplates
using ClimateMachine.NumericalFluxes
using ClimateMachine.VTK

using StaticArrays
using Test
using Printf
using MPI
using ArgParse

using CLIMAParameters
using CLIMAParameters.Atmos.SubgridScale: C_smag
using CLIMAParameters.Planet: R_d, cp_d, cv_d, MSLP, grav
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

# ------------------------ Description ------------------------- #
# 1) Dry Rising Bubble (circular potential temperature perturbation)
# 2) Boundaries - `All Walls` : Impenetrable(FreeSlip())
#                               Laterally periodic
# 3) Domain - 20000m[horizontal] x 10000m[vertical] (2-dimensional)
# 4) Timeend - 1000s
# 5) Mesh Aspect Ratio (Effective resolution) 2:1
# 7) Overrides defaults for
#               `init_on_cpu`
#               `solver_type`
#               `sources`
#               `C_smag`
# 8) Default settings can be found in `src/Driver/Configurations.jl`
# ------------------------ Description ------------------------- #
function init_risingbubble!(problem, bl, state, aux, localgeo, t)
    (x, y, z) = localgeo.coord

    FT = eltype(state)
    param_set = parameter_set(bl)
    R_gas::FT = R_d(param_set)
    c_p::FT = cp_d(param_set)
    c_v::FT = cv_d(param_set)
    γ::FT = c_p / c_v
    p0::FT = MSLP(param_set)
    _grav::FT = grav(param_set)

    xc::FT = 10000
    zc::FT = 2000
    r = sqrt((x - xc)^2 + (z - zc)^2)
    rc::FT = 2000
    θ_ref::FT = 300
    Δθ::FT = 0

    if r <= rc
        Δθ = FT(2) * cospi(0.5 * r / rc)^2
    end

    # Perturbed state:
    θ = θ_ref + Δθ # potential temperature
    π_exner = FT(1) - _grav / (c_p * θ) * z # exner pressure
    ρ = p0 / (R_gas * θ) * (π_exner)^(c_v / R_gas) # density
    q_tot = FT(0)
    ts = PhaseEquil_ρθq(param_set, ρ, θ, q_tot)
    q_pt = PhasePartition(ts)

    ρu = SVector(FT(0), FT(0), FT(0))

    # State (prognostic) variable assignment
    e_kin = FT(0)
    e_pot = gravitational_potential(bl.orientation, aux)
    ρe_tot = ρ * total_energy(e_kin, e_pot, ts)
    state.ρ = ρ
    state.ρu = ρu
    state.energy.ρe = ρe_tot
    state.moisture.ρq_tot = ρ * q_pt.tot
end

function config_risingbubble(FT, N, resolution, xmax, ymax, zmax, fast_method)

    # Choose fast solver
    if fast_method == "LowStorageRungeKutta2N"
        ode_solver = ClimateMachine.MISSolverType(
            splitting_type = ClimateMachine.SlowFastSplitting(),
            fast_model = AtmosAcousticGravityLinearModel,
            mis_method = MIS2,
            fast_method = LSRK54CarpenterKennedy,
            nsubsteps = (50,),
        )
    elseif fast_method == "StrongStabilityPreservingRungeKutta"
        ode_solver = ClimateMachine.MISSolverType(
            splitting_type = ClimateMachine.SlowFastSplitting(),
            fast_model = AtmosAcousticGravityLinearModel,
            mis_method = MIS2,
            fast_method = SSPRK33ShuOsher,
            nsubsteps = (12,),
        )
    elseif fast_method == "MultirateInfinitesimalStep"
        ode_solver = ClimateMachine.MISSolverType(
            splitting_type = ClimateMachine.HEVISplitting(),
            fast_model = AtmosAcousticGravityLinearModel,
            mis_method = MIS2,
            fast_method = (dg, Q, nsubsteps) -> MultirateInfinitesimalStep(
                MISKWRK43,
                dg,
                (dgi, Qi) -> LSRK54CarpenterKennedy(dgi, Qi),
                Q,
                nsubsteps = nsubsteps,
            ),
            nsubsteps = (12, 2),
        )
    elseif fast_method == "MultirateRungeKutta"
        ode_solver = ClimateMachine.MISSolverType(
            splitting_type = ClimateMachine.HEVISplitting(),
            fast_model = AtmosAcousticGravityLinearModel,
            mis_method = MIS2,
            fast_method = (dg, Q, nsubsteps) -> MultirateRungeKutta(
                LSRK144NiegemannDiehlBusch,
                dg,
                Q,
                steps = nsubsteps,
            ),
            nsubsteps = (12, 4),
        )
    elseif fast_method == "AdditiveRungeKutta"
        ode_solver = ClimateMachine.MISSolverType(
            splitting_type = ClimateMachine.HEVISplitting(),
            fast_model = AtmosAcousticGravityLinearModel,
            mis_method = MISRK3,
            fast_method = (dg, Q, dt, nsubsteps) -> AdditiveRungeKutta(
                ARK548L2SA2KennedyCarpenter,
                dg,
                LinearBackwardEulerSolver(ManyColumnLU(), isadjustable = true),
                Q,
                dt = dt,
                nsubsteps = nsubsteps,
            ),
            nsubsteps = (12,),
        )
    else
        error("Invalid --fast_method=$fast_method")
    end

    # Set up the model
    C_smag = FT(0.23)
    ref_state =
        HydrostaticState(DryAdiabaticProfile{FT}(param_set, FT(300), FT(0)))

    physics = AtmosPhysics{FT}(
        param_set;
        turbulence = SmagorinskyLilly{FT}(C_smag),
        ref_state = ref_state,
    )

    model = AtmosModel{FT}(
        AtmosLESConfigType,
        physics;
        source = (Gravity(),),
        init_state_prognostic = init_risingbubble!,
    )

    # Problem configuration
    config = ClimateMachine.AtmosLESConfiguration(
        "DryRisingBubbleMIS",
        N,
        resolution,
        xmax,
        ymax,
        zmax,
        param_set,
        init_risingbubble!,
        solver_type = ode_solver,
        model = model,
    )
    return config
end

function config_diagnostics(driver_config)
    interval = "10000steps"
    dgngrp = setup_atmos_default_diagnostics(
        AtmosLESConfigType(),
        interval,
        driver_config.name,
    )
    return ClimateMachine.DiagnosticsConfiguration([dgngrp])
end

function main()

    rbb_args = ArgParseSettings(autofix_names = true)
    add_arg_group!(rbb_args, "RisingBubbleBryan")
    @add_arg_table! rbb_args begin
        "--fast_method"
        help = "Choice of fast solver for the MIS method"
        metavar = "<name>"
        arg_type = String
        default = "AdditiveRungeKutta"
    end

    cl_args = ClimateMachine.init(parse_clargs = true, custom_clargs = rbb_args)
    fast_method = cl_args["fast_method"]

    # Working precision
    FT = Float64
    # DG polynomial order
    N = 3
    # Domain resolution and size
    Δx = FT(125)
    Δy = FT(125)
    Δz = FT(125)
    resolution = (Δx, Δy, Δz)
    # Domain extents
    xmax = FT(20000)
    ymax = FT(1000)
    zmax = FT(10000)
    # Simulation time
    t0 = FT(0)
    timeend = FT(20)

    # Time-step size (s)
    Δt = FT(0.4)

    driver_config =
        config_risingbubble(FT, N, resolution, xmax, ymax, zmax, fast_method)
    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        init_on_cpu = true,
        ode_dt = Δt,
    )
    dgn_config = config_diagnostics(driver_config)

    # Invoke solver (calls solve! function for time-integrator)
    result = ClimateMachine.invoke!(
        solver_config;
        diagnostics_config = dgn_config,
        check_euclidean_distance = true,
    )

    @test isapprox(result, FT(1); atol = 1.5e-3)
end

main()


text:
using Genie.Router

route("/") do
  serve_static_file("welcome.html")
end

route("/hello") do
  "Welcome to Genie!"
end


text:
using FluxRC, KitBase, Plots, OrdinaryDiffEq
using ProgressMeter: @showprogress

cd(@__DIR__)
ps = KitBase.UnstructPSpace("../assets/linesource.msh")

begin
    # quadrature
    quadratureorder = 8
    points, weights = KitBase.legendre_quadrature(quadratureorder)
    #points, triangulation = KitBase.octa_quadrature(quadratureorder)
    #weights = KitBase.quadrature_weights(points, triangulation)
    nq = size(points, 1)
    vs = KitBase.UnstructVSpace(-1.0, 1.0, nq, points, weights)

    # IC
    s2 = 0.03^2
    init_field(x, y) = max(1e-4, 1.0 / (4.0 * π * s2) * exp(-(x^2 + y^2) / 4.0 / s2))

    # particle
    SigmaS = ones(size(ps.cellid, 1))
    SigmaA = zeros(size(ps.cellid, 1))
    SigmaT = SigmaS + SigmaA

    # time
    tspan = (0.0, 0.3)
    cfl = 0.7
end

N = deg = 2
Np = (N + 1) * (N + 2) ÷ 2
ncell = size(ps.cellid, 1)
nface = size(ps.faceType, 1)

J = rs_jacobi(ps.cellid, ps.points)

spg = global_sp(ps.points, ps.cellid, N)
fpg = global_fp(ps.points, ps.cellid, N)

pl, wl = tri_quadrature(N)

V = vandermonde_matrix(N, pl[:, 1], pl[:, 2])
Vr, Vs = ∂vandermonde_matrix(N, pl[:, 1], pl[:, 2])
∂l = ∂lagrange(V, Vr, Vs)

ϕ = correction_field(N, V)

pf, wf = triface_quadrature(N)
ψf = zeros(3, N + 1, Np)
for i = 1:3
    ψf[i, :, :] .= vandermonde_matrix(N, pf[i, :, 1], pf[i, :, 2])
end

lf = zeros(3, N + 1, Np)
for i = 1:3, j = 1:N+1
    lf[i, j, :] .= V' \ ψf[i, j, :]
end

u = zeros(size(ps.cellid, 1), Np, nq)
for i in axes(u, 1), j in axes(u, 2)
    u[i, j, :] .= init_field(spg[i, j, 1], spg[i, j, 2])
end

cell_normal = zeros(ncell, 3, 2)
for i = 1:ncell
    pids = ps.cellid[i, :]

    cell_normal[i, 1, :] .= unit_normal(ps.points[pids[1], :], ps.points[pids[2], :])
    cell_normal[i, 2, :] .= unit_normal(ps.points[pids[2], :], ps.points[pids[3], :])
    cell_normal[i, 3, :] .= unit_normal(ps.points[pids[3], :], ps.points[pids[1], :])

    p =
        [
            ps.points[pids[1], :] .+ ps.points[pids[2], :],
            ps.points[pids[2], :] .+ ps.points[pids[3], :],
            ps.points[pids[3], :] .+ ps.points[pids[1], :],
        ] / 2

    for j = 1:3
        if dot(cell_normal[i, j, :], p[j][1:2] - ps.cellCenter[i, 1:2]) < 0
            cell_normal[i, j, :] .= -cell_normal[i, j, :]
        end
    end
end

δu = heaviside.(vs.u[:, 1])
δv = heaviside.(vs.u[:, 2])

function dudt!(du, u, p, t)
    du .= 0.0

    deg = p
    ncell = size(u, 1)
    nsp = size(u, 2)
    nq = size(u, 3)

    f = zeros(ncell, nsp, nq, 2)
    for i in axes(f, 1)
        for j in axes(f, 2), k in axes(f, 3)
            fg = vs.u[k, 1] * u[i, j, k]
            gg = vs.u[k, 2] * u[i, j, k]
            f[i, j, k, :] .= inv(J[i]) * [fg, gg]
        end
    end

    u_face = zeros(ncell, 3, deg + 1, nq)
    f_face = zeros(ncell, 3, deg + 1, nq, 2)
    for i = 1:ncell, j = 1:3, k = 1:deg+1, l = 1:nq
        u_face[i, j, k, l] = sum(u[i, :, l] .* lf[j, k, :])
        f_face[i, j, k, l, 1] = sum(f[i, :, l, 1] .* lf[j, k, :])
        f_face[i, j, k, l, 2] = sum(f[i, :, l, 2] .* lf[j, k, :])
    end

    n = [[0.0, -1.0], [1 / √2, 1 / √2], [-1.0, 0.0]]

    fn_face = zeros(ncell, 3, deg + 1, nq)
    for i = 1:ncell, j = 1:3, k = 1:deg+1, l = 1:nq
        fn_face[i, j, k, l] = sum(f_face[i, j, k, l, :] .* n[j])
    end

    fn_interaction = zeros(ncell, 3, deg + 1, nq)
    for i = 1:ncell
        for j = 1:3, k = 1:deg+1
            ni, nj, nk = neighbor_fpidx([i, j, k], ps, fpg)
            if ni > 0
                for l = 1:nq
                    uL = u_face[i, j, k, l]
                    uR = u_face[ni, nj, nk, l]

                    fx = vs.u[l, 1] * (uL * δu[l] + uR * (1.0 - δu[l]))
                    fy = vs.u[l, 2] * (uL * δv[l] + uR * (1.0 - δv[l]))

                    _f = inv(J[i]) * [fx, fy]

                    fn_interaction[i, j, k, l] = _f .* n[j] |> sum
                end
            else
                fn_interaction[i, j, k, :] .= 0.0
            end
        end
    end

    rhs1 = zeros(ncell, nsp, nq)
    for i in axes(rhs1, 1), j in axes(rhs1, 2), k = 1:nq
        rhs1[i, j, nq] =
            -sum(f[i, :, nq, 1] .* ∂l[j, :, 1]) - sum(f[i, :, nq, 2] .* ∂l[j, :, 2])
    end

    rhs2 = zero(rhs1)
    for i = 1:ncell, k = 1:nq
        if ps.cellType[i] != 1
            for j = 1:nsp
                rhs2[i, j, k] =
                    -sum((fn_interaction[i, :, :, k] .- fn_face[i, :, :, k]) .* ϕ[:, :, j])
            end
        end
    end

    #du .= rhs1 .+ rhs2
    for i = 1:ncell
        if ps.cellType[i] == 0
            du[i, :, :] .= rhs1[i, :, :] .+ rhs2[i, :, :]
        end
    end

    return nothing
end

tspan = (0.0, 1.0)
p = N
prob = ODEProblem(dudt!, u, tspan, p)

dt = 0.01
itg = init(prob, Euler(), save_everystep = false, adaptive = false, dt = dt)

@showprogress for iter = 1:10
    step!(itg)
end

sol = zeros(ncell)
for i in eachindex(sol)
    sol[i] = (vs.weights .* itg.u[i, 2, :]) |> sum
end
write_vtk(ps.points, ps.cellid, sol)

du = zero(u)
dudt!(du, u, p, 0.0)


text:
include("../src/PrivateMultiplicativeWeights.jl")

using PrivateMultiplicativeWeights
using Distributions
using Plots
#plotlyjs()
#plotly()
#pyplot()
gr()

d_continuous = Truncated(MixtureModel(Normal, [(.275, .075), (.65, .1)], [0.65, 0.35]), 0, 1)
xs = linspace(d_continuous.lower, d_continuous.upper, 300)
ys = pdf.(d_continuous, xs)
plot(xs, ys);

num_samples = 1000
domain_size = 100
samples = rand(d_continuous, num_samples)

function histogram_from_samples(samples, num_bins)
    @assert 0 <= minimum(samples)
    @assert maximum(samples) <= 1
    h = zeros(num_bins)
    new_samples = samples*num_bins
    for x in round.(Int64, new_samples)
        h[x] += 1
    end
    collect(1:num_bins)/num_bins, h
end

xs, ys = histogram_from_samples(samples, domain_size)
bar(xs, ys, width=1/domain_size);
#gui()

mw = mwem(RangeQueries(domain_size), Histogram(ys/sum(ys), num_samples))
bar(xs, mw.synthetic.weights*num_samples, color="red", alpha=0.3, width=1/domain_size);
#gui()

mw = mwem(RangeQueries(domain_size), Histogram(ys/sum(ys), num_samples), MWParameters(epsilon=100, iterations=80))
bar(xs, mw.synthetic.weights*num_samples, color="red", alpha=0.3, width=1/domain_size);
#gui()

a = bar(xs, ys)#, alpha=1/domain_size))
bar!(xs, mw.synthetic.weights*num_samples, color="red", alpha=0.3, width=1/domain_size)

display(a)

max_stars_repo_name: giladroyz/MWEM-project

max_issues_repo_name: giladroyz/MWEM-project

max_forks_repo_name: giladroyz/MWEM-project


text:
const ntrials = 5
print_output = isempty(ARGS)
codespeed = length(ARGS) > 0 && ARGS[1] == "codespeed"

if codespeed
    using JSON
    using HTTPClient.HTTPC

    # Ensure that we've got the environment variables we want:
    if !haskey(ENV, "JULIA_FLAVOR")
        error( "You must provide the JULIA_FLAVOR environment variable identifying this julia build!" )
    end

    # Setup codespeed data dict for submissions to codespeed's JSON endpoint.  These parameters
    # are constant across all benchmarks, so we'll just let them sit here for now
    csdata = Dict()
    csdata["commitid"] = Base.GIT_VERSION_INFO.commit
    csdata["project"] = "Julia"
    csdata["branch"] = Base.GIT_VERSION_INFO.branch
    csdata["executable"] = ENV["JULIA_FLAVOR"]
    csdata["environment"] = chomp(readall(`hostname`))
    csdata["result_date"] = join( split(Base.GIT_VERSION_INFO.date_string)[1:2], " " )    #Cut the timezone out
end

# Takes in the raw array of values in vals, along with the benchmark name, description, unit and whether less is better
function submit_to_codespeed(vals,name,desc,unit,test_group,lessisbetter=true)
    # Points to the server
    codespeed_host = "julia-codespeed.csail.mit.edu"

    csdata["benchmark"] = name
    csdata["description"] = desc
    csdata["result_value"] = mean(vals)
    csdata["std_dev"] = std(vals)
    csdata["min"] = minimum(vals)
    csdata["max"] = maximum(vals)
    csdata["units"] = unit
    csdata["units_title"] = test_group
    csdata["lessisbetter"] = lessisbetter

    println( "$name: $(mean(vals))" )
    ret = post( "http://$codespeed_host/result/add/json/", {"json" => json([csdata])} )
    if ret.http_code != 200 && ret.http_code != 202
        error("Error submitting $name [HTTP code $(ret.http_code)], dumping headers and text: $(ret.headers)\n$(bytestring(ret.body))\n\n")
        return false
    end
    return true
end

macro output_timings(t,name,desc,group)
    quote
        # If we weren't given anything for the test group, infer off of file path!
        test_group = length($group) == 0 ? basename(dirname(Base.source_path())) : $group[1]
        if codespeed
            submit_to_codespeed( $t, $name, $desc, "seconds", test_group )
        elseif print_output
            @printf "julia,%s,%f,%f,%f,%f\n" $name minimum($t) maximum($t) mean($t) std($t)
        end
        gc()
    end
end

macro timeit(ex,name,desc,group...)
    quote
        t = zeros(ntrials)
        for i=0:ntrials
            e = 1000*(@elapsed $(esc(ex)))
            if i > 0
                # warm up on first iteration
                t[i] = e
            end
        end
        @output_timings t $name $desc $group
    end
end

macro timeit1(ex,name,desc,group...)
    quote
        t = 0.0
        for i=0:1
            t = 1000*(@elapsed $(esc(ex)))
        end
        @output_timings [t] $name $desc $group
    end
end

macro timeit_init(ex,init,name,desc,group...)
    quote
        t = zeros(ntrials)
        for i=0:ntrials
            $(esc(init))
            e = 1000*(@elapsed $(esc(ex)))
            if i > 0
                # warm up on first iteration
                t[i] = e
            end
        end
        @output_timings t $name $desc $group
    end
end

# seed rng for more consistent timings
srand(1776)

max_stars_repo_name: porterjamesj/julia

max_issues_repo_name: porterjamesj/julia

max_forks_repo_name: porterjamesj/julia


text:
using LinearAlgebra # You'll need the inv(...) function to calculate an inverse
import ReverseDiff, ForwardDiff

# n = 100
# beta = [1.0, 2.2]
# X = [ones(n) rand(n)]
# Y = X*beta + randn(n)
# beta_hat = inv(X'X)*X'Y;

# f(β) = (Y - X*β)' * (Y - X*β)

# grad_f = ForwardDiff.gradient(f, beta_hat)
# grad_r = ReverseDiff.gradient(f, beta_hat)

# n_dims = 2
# window = range(0, 3, length=1000)
# grid = [[x,y] for x in window, y in window]

# gradients = map(
#     β_guess -> ForwardDiff.gradient(f, β_guess),
#     grid
# )

# norms = map(
#     g -> sum(g .^ 2),
#     gradients
# )

# grid[findmin(norms)[2]]

# struct Something{A, B}
#     a::A
#     b::B
# end
# Something(1.0)

# zeros(eltype([1.0]), 5)

# Imports
using Distributions, LinearAlgebra
using Optim

# Create the true distribution
n = 1000
means = [1.0, 0.0, 3.0, 4.0]
dist = MvNormal(means, I)
xs = rand(dist, n)

target(m) = -sum(logpdf(MvNormal(m,I), xs))
res = optimize(target, zeros(4))


text:
type KalmanSmoothed{T}
	predicted::Array{T}
	filtered::Array{T}
	smoothed::Array{T}
	error_cov::Array{T}
	model::StateSpaceModel
	y::Array{T}
  u::Array{T}
	loglik::T
end

type KalmanSmoothedMinimal{T}
	y::Array{T}
	smoothed::Array{T}
	error_cov::Array{T}
  u::Array{T}
	model::StateSpaceModel
	loglik::T
end

function show{T}(io::IO, smoothed::KalmanSmoothed{T})
    n = size(smoothed.y, 1)
    dx, dy = smoothed.model.nx, smoothed.model.ny
    println("KalmanSmoothed{$T}")
    println("$n observations, $dx-D process x $dy-D observations")
    println("Negative log-likelihood: $(smoothed.loglik)")
end

function show{T}(io::IO, smoothed::KalmanSmoothedMinimal{T})
    n = size(smoothed.y, 1)
    dx, dy = smoothed.model.nx, smoothed.model.ny
    println("KalmanSmoothedMinimal{$T}")
    println("$n observations, $dx-D process x $dy-D observations")
    println("Negative log-likelihood: $(smoothed.loglik)")
end

# Durbin-Koopman Smoothing
function kalman_smooth(y::Array, model::StateSpaceModel; u::Array=zeros(size(y,1), model.nu))

    y_orig = copy(y)
    y = y'
    u = u'
    n = size(y, 2)

    @assert !any(isnan(u))
    @assert size(y, 1) == model.ny
    @assert size(u, 1) == model.nu
    @assert size(u, 2) == n

    y_notnan = !isnan(y)
    y = y .* y_notnan

    x_pred,  V_pred     = zeros(model.nx, n), zeros(model.nx, model.nx, n)
    x_smooth, V_smooth  = copy(x_pred), copy(V_pred)
    x_pred_t, V_pred_t  = model.x1, model.P1

    innov               = zeros(model.ny, n)
    innov_cov_inv       = zeros(model.ny, model.ny, n)
    K                   = zeros(model.nx, model.ny, n)

    innov_t             = zeros(model.ny)
    innov_cov_t         = zeros(model.ny, model.ny)
    innov_cov_inv_t     = copy(innov_cov_t)
    Kt                  = zeros(model.nx, model.ny)

    At, Ct, Dt, Wt, ut  = model.A(1), model.C(1), model.D(1), model.W(1), zeros(model.nu)

    log_likelihood = n*model.ny*log(2pi)/2
    marginal_likelihood(innov::Vector, innov_cov::Matrix, innov_cov_inv::Matrix) =
         (dot(innov, innov_cov_inv * innov) + logdet(innov_cov))/2

    # Prediction and filtering
    for t = 1:n

        # Predict using last iteration's values
        if t > 1
            x_pred_t = At * x_pred_t + model.B(t-1) * ut + Kt * innov_t
            V_pred_t = At * V_pred_t * (At - Kt * Ct)' + model.V(t)
            V_pred_t = (V_pred_t + V_pred_t')/2
        end #if
        x_pred[:, t]    = x_pred_t
        V_pred[:, :, t] = V_pred_t

        # Assign new values
        At = model.A(t)
        Ct = model.C(t)
        Dt = model.D(t)
        Wt = model.W(t)
        ut = u[:, t]

        # Check for and handle missing observation values
        missing_obs = !all(y_notnan[:, t])
        if missing_obs
            ynnt = y_notnan[:, t]
            I1, I2 = diagm(ynnt), diagm(!ynnt)
            Ct, Dt = I1 * Ct, I1 * Dt
            Wt = I1 * Wt * I1 + I2 * Wt * I2
        end #if

        # Compute nessecary filtering quantities
        innov_t         = y[:, t] - Ct * x_pred_t - Dt * ut
        innov_cov_t     = Ct * V_pred_t * Ct' + Wt
        innov_cov_inv_t = inv(innov_cov_t)
        Kt              = At * V_pred_t * Ct' * innov_cov_inv_t

        innov[:, t]             = innov_t
        innov_cov_inv[:, :, t]  = innov_cov_inv_t
        K[:, :, t]              = Kt
        log_likelihood += marginal_likelihood(innov_t, innov_cov_t, innov_cov_inv_t)

    end #for

    # Smoothing
    r, N = zeros(model.nx), zeros(model.nx, model.nx)

    for t = n:-1:1

        Ct = model.C(t)
        innov_cov_inv_t = innov_cov_inv[:, :, t]
        V_pred_t = V_pred[:, :, t]
        L = model.A(t) - K[:, :, t] * Ct

        r = Ct' * innov_cov_inv_t * innov[:, t] + L' * r
        N = Ct' * innov_cov_inv_t * Ct + L' * N * L

        x_smooth[:, t] = x_pred[:, t] + V_pred_t * r
        V_smooth_t = V_pred_t - V_pred_t * N * V_pred_t
        V_smooth[:, :, t] =  (V_smooth_t + V_smooth_t')/2

    end #for

    return KalmanSmoothedMinimal(y_orig, x_smooth', V_smooth, u', model, log_likelihood)

end #smooth

# Bad performance - will be much better with sparse matrices
function lag1_smooth(y::Array, u::Array, m::StateSpaceModel)

    A_stack(t) = [m.A(t) zeros(m.nx, m.nx); eye(m.nx) zeros(m.nx, m.nx)]
    B_stack(t) = [m.B(t); zeros(m.nx, m.nu)]
    V_stack(t) = [m.V(t) zeros(m.nx, m.nx); zeros(m.nx, 2*m.nx)]
    C_stack(t) = [m.C(t) zeros(m.ny, m.nx)]
    x1_stack   = [m.x1; zeros(m.x1)]
    P1_stack   = [m.P1 zeros(m.nx, m.nx); zeros(m.nx, 2*m.nx)]
    stack_model = StateSpaceModel(A_stack, V_stack, C_stack, m.W,
                                      x1_stack, P1_stack, B=B_stack, D=m.D)

    stack_smoothed = kalman_smooth(y, stack_model, u=u)
    x     = stack_smoothed.smoothed[:, 1:m.nx]'
    P     = stack_smoothed.error_cov[1:m.nx, 1:m.nx, :]
    Plag1 = stack_smoothed.error_cov[1:m.nx, (m.nx+1):end, :]
    return x, P, Plag1, stack_smoothed.loglik

end #function

# Rauch-Tung-Striebel Smoothing
function kalman_smooth(filt::KalmanFiltered)

    n = size(filt.y, 1)
    model = filt.model
    x_pred = filt.predicted'
    x_filt = filt.filtered'
    x_smooth = zeros(x_filt)
    P_pred = filt.pred_error_cov
    P_filt = filt.error_cov
    P_smoov = zeros(P_filt)

    x_smooth[:, n] = x_filt[:, n]
    P_smoov[:, :, n] = P_filt[:, :, n]
    for i = (n-1):-1:1
        J = P_filt[:, :, i] * model.A(i)' * inv(P_pred[:,:,i+1])
        x_smooth[:, i] = x_filt[:, i] + J * (x_smooth[:, i+1] - x_pred[:, i+1])
        P_smoov[:, :, i] = P_filt[:, :, i] + J * (P_smoov[:, :, i+1] - P_pred[:,:,i+1]) * J'
    end

    return KalmanSmoothed(x_pred', x_filt', x_smooth', P_smoov, model, filt.y, filt.u, filt.loglik)
end


text:
Pkg.init()
Pkg.update()

Pkg.add("PyCall")
ENV["PYTHON"]=""
Pkg.build("PyCall")

Pkg.add("PyPlot")
using PyPlot

max_stars_repo_name: florianLepretre/ucbature

max_issues_repo_name: florianLepretre/ucbature

max_forks_repo_name: florianLepretre/ucbature


text:
function inverse_demand(df ; included = ones(size(eachcol(df[:, r"shares"]),1), size(eachcol(df[:, r"shares"]),1)),
    bO = 2 .*ones(size(included,2),1), iv::Integer = 0, marketvars = nothing, constrained::Bool = false)
# Function to estimate inverse demand function as in Compiani (2018)
    # bO should be a (J x 1) matrix where each entry is
    # the order of the polynomials for the demand function for one product.
# include("fullInteraction.jl")
# include("bern.jl")
# include("dbern.jl")
# include("b.jl")
# include("db.jl")

 if size(included,2) > 4
     bigProblem = 1
 else
     bigProblem = 0
 end
J = size(included,2);

# Unpack DataFrame df
s = Matrix(df[:, r"shares"]);
pt = Matrix(df[:, r"prices"]);
xt = Matrix(df[:, r"x"]);
zt = Matrix(df[:, r"demand_instruments"]);

bernO = bO;
order = bO;
T = size(s,1);
IVbernO = bO.+iv;
if size(bO,1) !=J
    if size(bO,2)==J
        bO=bO';
    else
        throw("Error: Matrix of polynomial orders has wrong number of elements \n")
    end
end

# --------------------------------------------
# Check normalization
if maximum(zt)>1 #|| mininum(zt)<0
    throw("Error: Instruments are not normalized to be between 0 and 1 \n")
end

if marketvars != nothing
    if maximum(marketvars)>1 || minimum(zt)<0
        throw("Error: Market-level variables are not normalized to be between 0 and 1 \n")
    end
end

# --------------------------------------------
# Prep for design matrix and constraints
delta = zeros(T,J)
for j = 1:J
    delta[:,j] = -1*pt[:,j]
end

BERN = []
A = []
design = []

for xj = 1:J
    BERN_xj = zeros(T,1)
    if xj>1
        if xj<J
            perm_order = hcat(convert.(Integer, xj:J)', convert.(Integer, 1:xj-1)');
            perm_order = convert.(Integer, perm_order)
        else
            perm_order = convert.(Integer,1:J)
            perm_order = perm_order .- 1
            perm_order[1] = J
        end
    else
        perm_order = convert.(Integer, 1:J)
    end
    s_perm = similar(s)
    for i = 1:J
        s_perm[:,i] = s[:,perm_order[i]];
    end
    # Market shares
    for j = 1:1:J
         if included[xj,perm_order[j]]==1
            BERN_xj = [BERN_xj bern(s_perm[:,j],convert(Integer, bernO[xj,1])) ]
         end
    end

    # Market characteristics
    if marketvars!=nothing
        for j = 1:1:size(marketvars,2)
            BERN_xj = [BERN_xj bern(marketvars[:,j], bernO[xj,1])]
        end
    end

    if bigProblem == 1
        print("Done Constructing Design Matrix \n")
    end
    # Define instruments
    cols = convert.(Bool, included[xj,:]); # only need instruments for included products
    z_xj = zt[:,cols];
    if marketvars!=nothing
        z_xj = [zt[:,cols] marketvars];
    end

    # --------------------------------------------
    # Instruments
    A_xj = zeros(T,1)
    ztemp = z_xj;
    for zj = 1:1:size(ztemp, 2)
        A_xj = [A_xj bern(ztemp[:,zj], convert(Integer,IVbernO[xj,1]) )]
    end
    A_xj = A_xj[:, 2:end]

    numVars = convert(Integer, size(A_xj, 2)./(IVbernO[xj,1]+1) )
    A_xj = fullInteraction(A_xj, numVars, convert(Integer, IVbernO[xj,1]))
    if bigProblem == 1
        print("Done with instruments \n")
    end
    # --------------------------------------------
    # Prep for estimation
    if marketvars==nothing
        numVars = sum(included[xj,:]);
    else
        numVars = sum(included[xj,:]) + size(marketvars,2);
    end
    BERN_xj = BERN_xj[:,2:end]
    design_xj = fullInteraction(BERN_xj, convert(Integer, numVars), convert(Integer,bernO[xj,1]) )
    if bigProblem == 1
        print("Done with estimation prep \n")
    end
    push!(BERN, BERN_xj)
    push!(design, design_xj)
    push!(A, A_xj)
end

# --------------------------------------------
# Estimation

if constrained==false
    sig = []
    designs = []
    for xj = 1:J
        # Unconstrained Estimator
        if bigProblem == 1
            print("UNCONSTRAINED ESTIMATIION BEGINNING \n")
        end
        deltatemp = delta[:,xj]
        Atemp = A[xj]
        designtemp = design[xj]
        sdt = size(designtemp,2)
        if bigProblem == 1
            print("Number of parameters to estimate is: $sdt \n")
        end
        TOL = sqrt(eps(real(float(one(eltype(Atemp))))))
        sigtemp = pinv((designtemp'*Atemp*pinv(Atemp'*Atemp, rtol = TOL)*Atemp'*designtemp), rtol=TOL)*designtemp'*Atemp*pinv(Atemp'*Atemp, rtol=TOL)*Atemp'deltatemp;
        push!(sig, sigtemp)
        push!(designs, designtemp)
        if bigProblem == 1
            print("Estimation for product $xj is done \n")
        end
    end
else
    sig = []
    for xj = 1:1:J
        # Constrained Estimator
        if bigProblem == 1
            print("CONSTRAINED ESTIMATIION BEGINNING \n")
        end
        deltatemp = delta[:,xj]
        Atemp = A[xj]
        designtemp = design[xj]
        ApA = Atemp'*Atemp
        #eval(['ApA = transpose(A', num2str(xj), ')*A',num2str(xj),';']);
        if marketvars!=nothing
            numVars = sum(included(xj,:)) + size(marketvars,2);
        else
            numVars = sum(included[xj,:]);
        end
        #eval(['constMat = makeConstraint(BERN_x', num2str(xj), ', numVars, bernO(xj,1),1,1,1);']);
        numVars = convert(Integer, numVars)
        bernO = convert.(Integer, bernO)
        constMat = makeConstraint(BERN[xj], numVars, bernO[xj,1],1,1,1);
        ineqMat = 0*ones(size(constMat[1],1),1);
        sdt =size(designtemp,2)
        if bigProblem == 1
            print("Number of parameters to estimate is: $sdt \n")
        end
        ApA = pinv(ApA);
        constMat = constMat[1];
        bigA = Atemp*ApA*Atemp';
        #obj(x) = objective_priceIndex(x, designtemp ,deltatemp , Atemp, ApA);
        function obj1(x::Vector)
            xi = (deltatemp - designtemp*x);
            out = xi'*bigA*xi;
            return out
        end
        function con1(x::Vector)
            out = constMat*x
            return out
        end
        g = x -> ForwardDiff.gradient(obj1, x)
        g2 = x-> ForwardDiff.gradient(con1, x)
        function obj(x::Vector, grad::Vector)
            if length(grad) > 0
                grad = g(x);
            end
            xi = (deltatemp - designtemp*x);
            out = xi'*bigA*xi;
            return out
        end
        function constraint(result::Vector, x::Vector, grad::Matrix, constMat)
            result = constMat*x
            if length(grad) > 0
                grad = g2(x);
            end
            return result
        end
        opt = Opt(:LN_COBYLA, size(designtemp,2))
        opt.xtol_rel = 1e-4
        opt.min_objective = obj
        inequality_constraint!(opt, (result,x,g) -> constraint(result, x, g, constMat), dropdims(1e-8.*ones(1,size(designtemp,2)),dims=1) )
        (minf,minx,ret) = NLopt.optimize(opt, dropdims(ones(size(designtemp,2),1),dims = 2) )
        sigtemp = minx
        push!(sig, sigtemp)
        designs = [];
    end
end
inv_sigma = []
for j = 1:1:J
    push!(inv_sigma, sig[j])
end

return inv_sigma, designs
end