using StatsBase
using Statistics
using Distributions
using Random

# N.B. See Mimi/WIP/load_empirical_dist.jl for helper functions.

struct EmpiricalDistribution{T}
    values::Vector{T}
    weights::ProbabilityWeights
    dist::Distribution

    # Create an empirical distribution from a vector of values and an optional
    # vector of probabilities for each value. If not provided, the values are
    # assumed to be equally likely.
    # N.B. This doesn't copy the values vector, so caller must, if required
    function EmpiricalDistribution(values::Vector{T}, probs::Union{Nothing, Vector{Float64}}=nothing) where T
        n = length(values)
        if probs === nothing
            probs = Vector{Float64}(n)
            probs[:] = 1/n
        elseif length(probs) != n
            error("Vectors of values and probabilities must be equal lengths")
        end

        # If the probs don't exactly sum to 1, but are close, tweak
        # the final value so the sum is 1.
        total = sum(probs)
        if 0 < (1 - total) < 1E-5
            probs[end] += 1 - total
        end

        weights = ProbabilityWeights(probs)

        return new{T}(values, weights, Categorical(probs))
    end
end

#
# Delegate a few functions that we require in our application.
# No need to be exhaustive here.
#
function Statistics.mean(d::EmpiricalDistribution)
    return mean(d.values, d.weights)
end

function Statistics.std(d::EmpiricalDistribution)
    return std(d.values, d.weights, corrected=true)
end

function Statistics.var(d::EmpiricalDistribution)
    return var(d.values, d.weights, corrected=true)
end

function Statistics.quantile(d::EmpiricalDistribution, args...)
    indices = quantile.(d.dist, args...)
    return d.values[indices]
end

function Statistics.rand(d::EmpiricalDistribution, args::Vararg{Integer,N}) where {N}
    indices = rand(d.dist, args...)
    return d.values[indices]
end

function Random.rand!(d::EmpiricalDistribution, args::Vararg{Integer,N}) where {N}
    indices = rand!(d.dist, args...)
    return d.values[indices]
end


"""
    BoltzmannNetworks

Implements neural sampling in second-order Boltzmann distributions as presented in

> Lars Buesing, Johannes Bill, Bernhard Nessler and Wolfgang Maass (2011).
> Neural Dynamics as Sampling: A Model for Stochastic Computation in Recurrent Networks of Spiking Neurons.
> PLoS Computational Biology 7(11), e1002211.

The goal is to approximate a distribution ``p(\\mathbf{z})`` over the binary variables
``\\mathbf{z} = z_1, z_2, \\ldots, z_k`` by a neural network ``\\mathcal{N}`` of neurons
``v_1, v_2, \\ldots, v_k``. The spiking activity of ``\\mathcal{N}`` is interpreted as a
Markov chain implementing Markov Chain Monte Carlo. Specifically, a spike of neuron ``v_i``
sets the variable ``z_i`` to `true` for ``\\tau`` timesteps / milliseconds / samples. The
passing of this period is kept track of via auxiliary variables ``\\zeta_1, \\zeta_2, \\ldots, \\zeta_k``.

Networks are represented by the type `BoltzmannNetwork`. This module implements the
following methods from the paper:

- sampling with absolute refractory period in discrete time via the function `sample_absolute!`
- sampling with relative refractory period in discrete time via the function `sample_relative!`
- Gibbs sampling (not using neural dynamics) as a baseline method

"""
module BoltzmannNetworks

using Random: shuffle
using ..Utils: σ
import ..Utils: spikes2chain
import Base: show, length

export BoltzmannNetwork
export randomnetwork, sample_absolute!, sample_relative!, sample_clamp!

###############################################################################
# Network type                                                                #
###############################################################################
"""
    BoltzmannNetwork

Represents a network of spiking neurons that performs MCMC sampling from the distribution

```math
p(\\mathbf{z}) = \\frac{1}{Z} \\exp \\left( \\sum_{i,j} \\frac{1}{2} W_{ij} z_i z_j + \\sum_i b_i z_i \\right)
```

where `Z` is the partition function, `W` are weights and `b` are biases.

# Fields
- `ζ::Vector{Int}` - vector of auxiliary variables `ζ`
- `z::Vector{Bool}` - vector of variables `z`
- `u::Vector{Float64}` - vector of neural membrane potentials
- `W::Matrix{Float64}` - weights relating the variables (symmetric, zeros on main diagonal)
- `b::Vector{Float64}` - biases governing neural excitability
- `τ::Int` - time constant encoding for how many steps `z` is `true` after a spike
- `g::G` - refractoriness function for sampling with relative refr. period
- `f::F` - activation function for sampling with relative refr. period
"""
struct BoltzmannNetwork{G, F}
    ζ::Vector{Int}
    z::Vector{Bool}
    u::Vector{Float64}
    W::Matrix{Float64}
    b::Vector{Float64}
    τ::Int
    g::G
    f::F

    # constructor to enforce constraints on W
    function BoltzmannNetwork(ζ, z, u, W, b, τ, g, f)
        for i in 1:size(W, 1)
            @assert W[i, i] == 0.0 "Main diagonal of weight matrix has to be all zeros."
            for j in i:size(W, 2)
                @assert W[i, j] == W[j, i] "Weight matrix has to be symmetric."
            end
        end
        return new{typeof(g), typeof(f)}(ζ, z, u, W, b, τ, g, f)
    end
end

"Generate a random `BoltzmannNetwork` of `n` neurons."
function randomnetwork(n::Int; τ=20, g=g) # τ as in Buesing et al. (2011)
    ζ = zeros(Int, n)
    z = zeros(Bool, n)
    u = zeros(Float64, n)
    W = randn(n, n) * 0.3 # scaling as in Buesing et al. (2011), Fig. 3, 5, 6, 7
    for i in 1:n, j in i:n
        if i==j
            W[i, j] = 0. # enforce that main axis is all zeros
        else
            W[i, j] = W[j, i] # enforce that weights are symmetric
        end
    end
    b = randn(n) * 0.5 .- 1.5 # scaling as in Buesing et al. (2011), Fig. 3, 5, 6, 7

    f = gettransferfun(g, τ=τ)
    return BoltzmannNetwork(ζ, z, u, W, b, τ, g, f)
end

###############################################################################
# Sampling with absolute refractory period.                                   #
###############################################################################
"""
    step_absolute!(net::BoltzmannNetwork)

Implements a single step of neural MCMC sampling with absolute refractory period
in discrete time.

# Details

Alters the state of the membrane potentials `net.u`, the variables `net.z` and
the auxiliary variables `net.ζ`. The neurons are updated sequentially (as
explained in the paper, this is necessary for correct sampling). For each
neuron, the membrane potential is calculated in accordance with the
**neural computability condition**:

```math
u_i = b_i + sum_j W_ij * z_i * z_j
```

If the neuron is refractory (``\\zeta \\ge 1``), ζ is decreased deterministically.
Otherwise the neurons spikes with probability ``\\sigma(u - \\log(\\tau))``. A spike
sets `z` to `true` and `ζ` to `τ`.
"""
function step_absolute!(net::BoltzmannNetwork)
    spikes = zeros(Bool, length(net.z))
    # update neurons
    # (in order, since transition operators have to be applied sequentially)
    for i in 1:length(spikes)
        # update membrane potential
        net.u[i] = net.b[i] + net.W[i, :]' * net.z
        # generate spike
        if (net.ζ[i] ≤ 1) && (rand() < σ(net.u[i] - log(net.τ)))
            spikes[i] = true
            net.ζ[i] = net.τ
            net.z[i] = 1
        else
            net.ζ[i] = max(net.ζ[i] - 1)
            if net.ζ[i] == 0
                net.z[i] = 0
            end
        end
    end
    return(spikes)
end

"""
    sample_absolute!(net, steps, burnin, dt)

Performs neural sampling with absolute refractory period in discrete time in
`BoltzmannNet` `net` for `steps` steps after a burnin time of `burnin` steps
with step size `dt`.

# Returns
- `ts::Vector{Float64}` - vector of spike times
- `inds::Vector{Int}` - vector of neuron indices of neurons that spiked
"""
function sample_absolute!(net, steps, burnin, dt)
    for _ in 1:burnin
        step_absolute!(net)
    end
    ts = zeros(0)
    inds = zeros(Int, 0)
    for t in 1:steps
        spikes = step_absolute!(net)
        for spike in findall(spikes)
            push!(ts, t * dt)
            push!(inds, spike)
        end
    end
    return ts, inds
end

###############################################################################
# Sampling with relative refractory period                                    #
###############################################################################
#
# In sampling with a relative refractory period, the neural transfer function
# is slightly more complicated. The refactoriness of a neuron is given by the
# function g(ζ): [0, τ] -> [0, 1]
#
# Depending on the exact function g chosen, the neural transfer function f
# has to fulfill some specific conditions. In general, a suitable f exists but
# cannot be found analytically. Therefore, it has to be approximated numerically.
#
# First, we implement the function g given in the paper.
"""Refractory function g(ζ).
Quantifies how excitable a neuron is based on variable ζ,
which encodes the time since last spike.

Assumes that ζ is in range [0, 1]"""
function g(ζ)
    x = 1 - ζ + sin(2π * ζ) / (2π) # according to Table 1 in Buesing et al. (2011)
    return min(1, max(0, x))
end

"Alternative refractory function"
function g2(ζ) # assumes that ζ is in range [0, 1]
    x = 4 * (1 - ζ) + sin(8π * ζ) / (2π) # according to Table 1 in Buesing et al. (2011)
    return min(1, max(0, x))
end

"Alternative refractory function"
function g3(ζ) # assumes that ζ is in range [0, 1]
    x = 1 - 2ζ + sin(4π * ζ) / (2π) # according to Table 1 in Buesing et al. (2011)
    return min(1, max(0, x))
end

# Next, we need to approximate f for a given g.
"""
    gettransferfun(g; τ=20)

Approximates the neural transfer function `f(u)` for a given function `g(ζ)`.

# Arguments
- `g::Function` - refractoriness; should accept values in range [0, 1].
"""
function gettransferfun(g; τ=20)
    # the following function calculates the error for a given u and f(u)
    ferror(u, fu) = abs(exp(u) - fu *
        sum(prod(1 - g(ζ/τ) * fu for ζ in η+1:τ) for η in 1:τ-1) /
        prod(1 - g(ζ/τ) * fu for ζ in 1:τ))
    # we assume that membrane potentials will stay in the range [-10, 10]
    us = -10:0.01:10
    fs = zeros(length(us))
    # we know that the output of f must lie in [0, 1]
    vals = 0:0.001:1
    # now, for each possible input value u, find the output value that minimizes the error
    for (i, u) in enumerate(us)
        fs[i] = vals[argmin(ferror.(u, vals))]
    end
    # now we can generate the function f:
    # for each input value, find the nearest known input value from us and
    # return the corresponding f(u)
    function f(u; us=us, fs=fs)
        i = argmin(abs.(u .- us))
        return fs[i]
    end
    return f
end

# Given the correct transfer function, we can implement the sampling
"""
    step_relative!(net::BoltzmannNetwork)

Implements a single step of neural MCMC sampling with relative refractory period
in discrete time.

# Details

Alters the state of the membrane potentials `net.u`, the variables `net.z` and
the auxiliary variables `net.ζ`. The neurons are updated sequentially (as
explained in the paper, this is necessary for correct sampling). For each
neuron, the membrane potential is calculated in accordance with the
**neural computability condition**:

```math
u_i = b_i + sum_j W_ij * z_i * z_j
```

The neural refactoriness is given by `g(ζ)`. At each point, the neuron fires with
probability `g(ζ) * f(u)`. A spike sets `z` to `true` and `ζ` to `τ`. If the neuron
does not spike, `ζ` decreases and `z` is set to `false` once `ζ` becomes `0`.
"""
function step_relative!(net::BoltzmannNetwork)
    spikes = zeros(Bool, length(net.z))
    # update neurons
    # (in order, since transition operators have to be applied sequentially)
    for i in 1:length(spikes)
        # update membrane potential
        net.u[i] = net.b[i] + net.W[i, :]' * net.z
        # generate spike with probability g(ζ) * f(u), otherwise decay ζ
        if (rand() < net.g(net.ζ[i] / net.τ) * net.f(net.u[i]))
            spikes[i] = true
            net.ζ[i] = net.τ
            net.z[i] = 1
        else
            net.ζ[i] = max(net.ζ[i] - 1)
            if net.ζ[i] == 0
                net.z[i] = 0
            end
        end
    end
    return(spikes)
end

"""
    sample_relative!(net, steps, burnin, dt)

Performs neural sampling with relative refractory period in discrete time in
`BoltzmannNet` `net` for `steps` steps after a burnin time of `burnin` steps
with step size `dt`.

# Returns
- `ts::Vector{Float64}` - vector of spike times
- `inds::Vector{Int}` - vector of neuron indices of neurons that spiked
"""
function sample_relative!(net, steps, burnin, dt)
    # run the Markov chain for burn-in period
    for _ in 1:burnin
        step_relative!(net)
    end

    # sample spikes
    ts = zeros(0)
    inds = zeros(Int, 0)
    for t in 1:steps
        spikes = step_relative!(net)
        for spike in findall(spikes)
            push!(ts, t * dt)
            push!(inds, spike)
        end
    end
    return ts, inds
end

###############################################################################
# Sampling with clamped values (inferring posterior)                          #
###############################################################################
"""
    step_clamp!(net::BoltzmannNetwork, clamp_on, clamp_off)

Performs one step of neural sampling with relative refractory period in `net`
while clamping the neurons at indices `clamp_on` to `true` and the neurons at
indices `clamp_off` to `false`.
"""
function step_clamp!(net::BoltzmannNetwork, clamp_on, clamp_off)
    spikes = zeros(Bool, length(net.z))
    # update neurons
    # (in order, since transition operators have to be applied sequentially)
    for i in shuffle(1:length(spikes))
        if i in clamp_on
            net.z[i] = true
            continue
        end
        if i in clamp_off
            net.z[i] = false
            continue
        end
        # update membrane potential
        net.u[i] = net.b[i] + net.W[i, :]' * net.z
        # generate spike with probability g(ζ) * f(u), otherwise decay ζ
        if (rand() < net.g(net.ζ[i] / net.τ) * net.f(net.u[i]))
            spikes[i] = true
            net.ζ[i] = net.τ
            net.z[i] = 1
        else
            net.ζ[i] = max(net.ζ[i] - 1)
            if net.ζ[i] == 0
                net.z[i] = 0
            end
        end
    end
    return(spikes)
end

"""
    sample_clamp!(net, steps, burnin, dt, clamp_on, clamp_off)

Performs neural sampling with relative refractory period in discrete time in
`BoltzmannNet` `net` for `steps` steps after a burnin time of `burnin` steps
with step size `dt` while clamping the neurons at indices `clamp_on` to `true`
and the neurons at indices `clamp_off` to `false`.

# Returns
- `ts::Vector{Float64}` - vector of spike times
- `inds::Vector{Int}` - vector of neuron indices of neurons that spiked
"""
function sample_clamp!(net, steps, burnin, dt, clamp_on, clamp_off)
    # run the Markov chain for burn-in period
    for _ in 1:burnin
        step_clamp!(net, clamp_on, clamp_off)
    end

    # sample spikes
    ts = zeros(0)
    inds = zeros(Int, 0)
    for t in 1:steps
        spikes = step_clamp!(net, clamp_on, clamp_off)
        for spike in findall(spikes)
            push!(ts, t * dt)
            push!(inds, spike)
        end
    end
    return ts, inds
end

###############################################################################
# Conversion of spike sequence into Markov chain                              #
###############################################################################
function spikes2chain(ts, inds, net, dt; start=0, stop=nothing)
    if stop == nothing
        stop = maximum(ts) + net.τ * dt
    end
    timesteps = round(Int, (stop - start) / dt)
    chain = zeros(Bool, length(net), timesteps)
    for (t, i) in zip(ts, inds)
        # each spike sets the corresponding variable to true for τ timesteps
        tᵢₙₜ = max(1, round(Int, (t - start) / dt))
        chain[i, tᵢₙₜ:min(timesteps, tᵢₙₜ+net.τ)] .= true
    end
    return chain
end

###############################################################################
# Overloads for pretty printing etc.                                          #
###############################################################################
function Base.show(io::IO, net::BoltzmannNetwork)
    print(io, "BoltzmannNetwork of $(length(net.z)) neurons.")
end

function Base.show(io::IO, ::MIME"text/plain", net::BoltzmannNetwork)
    print(io,
        "BoltzmannNetwork of $(length(net.z)) neurons.\n",
        "\t z = $(net.z) \n",
        "\t u = $(net.u) \n",
        "\t W = $(net.W) \n",
        "\t b = $(net.b) \n",
        "\t τ = $(net.τ)"
    )
end

function Base.length(net::BoltzmannNetwork)
    return length(net.z)
end

end # module BoltzmannNetworks


"""

The structure containing the set of neaerest-neighbor gates when
applied on the initial configuration `initconf` generates the
approximate MPS representation for a gaussian fermionic model. `lx` is
the size of the system. The left acting site of each gate is given in
the `positions` vector and each gate can be constructed using the
`θs` or possibly `ϕs`.

"""
struct GaussianMPS
    lx :: Int
    initconf :: Vector{Int}
    xs :: Vector{Int}
    θs :: Vector{Float64}

    function GaussianMPS(lx::Int,
                         initconf::Vector{Int},
                         xs::Vector{Int},
                         θs::Vector{Float64})
        @assert length(initconf) == lx
        @assert length(xs) == length(θs)

        new(lx, initconf, xs, θs)
    end
end

"""
    corrmat2gmps(corrmat, threshold, verbose)

For a given two-point correlation matrix `corrmat` called lambda
perform the approach explained in arxiv:1504.07701 and generate the
local (nearest-neighbor) gates.

"""

function corrmat2gmps(corrmat::Matrix{T},
                      threshold::Float64=1.e-8;
                      alternate::Bool=false) where {T}
    T <: Complex && error("correlation matrix is complex!")
    lx = size(corrmat)[1]

    #initconf = Vector{Int}(lx)
    initconf = zeros(Int, lx)
    positions = Int[]
    θs = Float64[]

    expected_next_evalue = 1
    for site=1:lx
        block_end = site
        evalue = corrmat[site, site]
        corr_block = [evalue]
        vals = corr_block

        if alternate
            delta = abs(expeted_next_evalue - evalue)
        else
            delta = min(abs(evalue), abs(1-evalue))
        end

        while (delta > threshold && block_end < lx)
            block_end += 1

            corr_block = corrmat[site:block_end, site:block_end]
            vals = eigvals(corr_block)

            if alternate
                delta = (expected_next_evalue == 0) ? abs.(minimum(vals)) : abs.(1-maximum(vals))
            else
                delta = min(abs.(minimum(vals)), abs.(1-maximum(vals)))
            end
        end

        if block_end > site
            block_size = block_end - site + 1
            block_range = site:block_end

            if alternate
                v_index = (expected_next_evalue == 0) ? indmin(abs.(vals)) : argmin(abs.(1-vals))
            else
                v_index = argmax(abs.(vals .- 0.5))
            end
            evalue = vals[v_index]
            v = eigvecs(corr_block)[:,v_index]

            # Find the set of (block_size-1) unitary gate that
            # diagonalize the current correlation block
            ugate_block = Matrix{T}(I, block_size, block_size)

            for p = block_size-1:-1:1 # p here stands for "pivot"

                θ = atan(v[p + 1]/ v[p])
                ugate = Matrix{Float64}(1.0I, block_size, block_size)
                ugate[p:p+1, p:p+1] =
                    [
                        cos(θ) -sin(θ);
                        sin(θ)  cos(θ)
                    ]
                push!(positions, site+p-1)
                push!(θs, θ)
                ## QQQ? acting on the vectors from left! Explain, why
                ## is this a better choice?
                v = transpose(transpose(v) * ugate)

                ugate_block = ugate_block * ugate
            end

            ugate_extended = Matrix{T}(I, lx, lx)
            ugate_extended[block_range, block_range] = ugate_block
            corrmat= Symmetric(ugate_extended' * corrmat * ugate_extended)
            # display(corrmat)
        end
        initconf[site] = round(evalue)
        expected_next_evalue = (round(evalue) == 1) ? 0 : 1
    end
    # display(corrmat)
    GaussianMPS(lx, initconf, positions, θs)
end


using CairoMakie, MLBase, Random
Random.seed!(123)
# fake data
gt = rand(1:10,1000)
pred = rand(1:10,1000)
data = confusmat(10, gt, pred)
ticks = string.(sort(unique(gt)))
# plot
fig = Figure(resolution = (600, 400))
ax = Axis(fig[1,1]; aspect = 1, ylabel = "Ground truth",
    xlabel = "Prediction")
hmap = heatmap!(ax, data, colormap = :magma)
for i in 1:10, j in 1:10
    txtcolor = data[i,j] < 10 ? :white : :black
    text!(ax, "$(round(data[i,j], digits = 2))", position = (i,j),
        color = txtcolor, align = (:center, :center), textsize = 12)
end
cbar = Colorbar(fig[1,2], hmap, label = "Counts", width = 25,
    halign = :left, ticksize = 25, tickalign = 1,
    height = Relative(3.55/4))
ax.xticks = (1:10, ticks)
ax.yticks = (1:10, ticks)
save("confusionMatrix.png", fig, px_per_unit = 2)
fig


@with_kw struct BoxNode <: Node
	lb :: Vector{Float64}
	ub :: Vector{Float64}
	subdivision_level :: Int64 = 0
	parent :: Union{BoxNode, Nothing} = nothing
	sub_boxes :: Vector{BoxNode} = []

    # Test point data
	database :: D where D<:SomeDataBase = DataBase()
    db_indices :: Vector{Int} = [];

end

# required Tree methods (see "Trees.jl")
children( bn :: BoxNode ) = bn.sub_boxes;
parent( bn :: BoxNode ) = bn.parent;
depth( bn :: BoxNode ) = bn.subdivision_level;

# utilities
lb( bn :: BoxNode ) = bn.lb; # use methods for `lb` & `ub` …
ub( bn :: BoxNode ) = bn.ub; # … in case we would rather want to store center & radius
n_vars( bn :: BoxNode ) = length( bn.lb );
width( bn :: BoxNode ) = ub(bn) .- lb(bn); # maybe raname?
center(bn :: BoxNode ) = lb(bn) .+ .5 .* width( bn );
edge_length( bn :: BoxNode, dim :: Int) = ub(bn)[dim] - lb(bn)[dim];

@doc """
    contains(box, point)

Return true if `point` lies in `box`.
"""
function contains( bn :: BoxNode, point :: Vector{R} where R<:Real)
    all( lb(bn) .<= point ) && all( point .<= ub(bn) )
end

function vertices( bn :: BoxNode )
    LB = lb(bn);
    UB = ub(bn);

    other_ind = combinations( eachindex(UB) )
    return [ [LB,];
        [ let v = copy(LB); v[ind] .= UB[ind]; v end for ind in other_ind ]
    ]
end

function edges( bn :: BoxNode )
    verts = vertices( bn );
    return combinations( verts, 2 )
end

@doc "Retrieve sites indexed by `db_indices` from referenced DataBase `db`."
function sites( bn :: BoxNode )
    return sites( bn.database, bn.db_indices);
end

@doc "Retrieve evaluation results indexed by `db_indices` from referenced DataBase `db`."
function results( bn :: BoxNode; include_nothing = true )
    return results( bn.database, bn.db_indices; include_nothing);
end

"Evaluate all sites referenced by field `db_indices`."
function eval!( bn :: BoxNode, f :: F where F<:Function; reevaluate = true)
    for db_id in bn.db_indices
        eval!( bn.database, f, db_id; reevaluate );
    end
end

@doc """
    init_box_tree( lb, ub, db = nothing )

Initialize the root node of a box Tree with variable boundaries
`lb` and `ub`.
If `db` is a database then it is referenced in the node.
"""
function init_box_tree(
        lb :: Vector{R} where R<:Real, ub :: Vector{R} where R<:Real,
        db :: Union{Nothing, DataBase} = nothing
    )
    box_tree = BoxNode(;
        lb = lb,
		ub = ub,
        database = isnothing(db) ? DataBase() : db,
    );
	return box_tree
end

@doc """
    add_sub_box!( parent, sub_lb, sub_ub )

Add a new box node with variable boundaries `sub_lb` and `sub_ub`
to parent.
Modifies `sub_boxes` field of `parent` and sets the parent reference
of the new box correctly.
"""
function add_sub_box!(
        parent :: BoxNode, sub_lb :: Vector{Float64},
        sub_ub :: Vector{Float64}
    )
	# check if boundaries of sub-box lay within parent box
	if all( lb(parent) .<= sub_lb ) && all( sub_ub .<= ub(parent) )
		push!( parent.sub_boxes,
			BoxNode(;
				lb = sub_lb,
				ub = sub_ub,
				subdivision_level = parent.subdivision_level + 1,
				parent = parent,
				sub_boxes = BoxNode[],
                database = parent.database,
			)
		);
	else
		@error("Variable boundaries of subbox must not exceed parent boundaries.")
	end
end

# recursive subdivison function
# usually not called directly by the user
"Divide each of `bn`s edges by divisor or the entries of divisor
and add sub boxes as children to the tree."
function subdivide!(
        bn :: BoxNode, divisor :: Union{Int, Vector{Int}},
        LB :: Vector{Float64}, UB :: Vector{Float64};
    )

    # recursion starts with `LB` and `UB` empty
	var_level = length(LB) + 1; # recursion level

    # divisor for this level
    sub_divisor = isa( divisor, Int ) ? divisor : divisor[ var_level ]

    width = edge_length( bn, var_level );
	width_fraction = width / sub_divisor
	for s = 1 : sub_divisor
		sub_lb = [LB; lb( bn )[ var_level ] + (s-1) * width_fraction ]
		sub_ub = [UB; lb( bn )[ var_level ] + s * width_fraction ]
		if var_level == n_vars( bn )
			add_sub_box!( bn, sub_lb, sub_ub )
		else
			subdivide!( bn, divisor, sub_lb, sub_ub )
		end
	end
end

function subdivide!(
        bn :: BoxNode;
        method = :cycle
    )
    if method == :cycle
        divisor = fill(1, n_vars(bn) );
        divisor[ ( depth(bn) % n_vars(bn) + 1 ) ] = 2;
    else
        divisor = 2;
    end
    subdivide!( bn, divisor, Float64[], Float64[] );
end

@doc """
    subdivide!(parent, divisor=2)

Recursively subdivide node `parent`.
`parent` is divided by `divisor` along each box dimension.
"""
function subdivide!( parent :: BoxNode, divisor :: Union{Vector{Int}, Int} = 2)
	lb = Float64[];
	ub = Float64[];
	subdivide!( parent, divisor, lb, ub )
end

function generate_sites!( bn :: BoxNode; method ::Symbol = :edges_center, kwargs...  )
    generate_sites!(bn , Val(method); kwargs...)
end

function generate_sites!( bn :: BoxNode, ::Val{:edges_center}; num_edge_points = -1)
    if num_edge_points < 0
        num_edge_points = n_vars(bn)
    end

    # reset old database ids
    empty!(bn.db_indices);

    # get points for unit HyperCube (cached for future reference)
    unit_points = unit_sites_edges_center( n_vars(bn); num_edge_points = num_edge_points)

    scale_fn(p) = p .* width(bn) .+ lb(bn);
    new_points = scale_fn.(unit_points);

    db_indices = push_sites!( bn.database, new_points );
    push!(bn.db_indices,db_indices...);
    nothing;
end

@memoize function unit_sites_edges_center(n_vars ::Int; num_edge_points ::Int = -1)
    if num_edge_points < 0
        num_edge_points = n_vars
    end

    bn = init_box_tree( zeros(n_vars), ones(n_vars) );

    ℓ = collect( 1 : num_edge_points )
    steps = ( 2 .* ℓ .- 1 ) ./ (2*num_edge_points);

    # N * num_edges + 1
    box_edges = edges(bn)
    num_edges = length(edges(bn))
    num_new_points = num_edge_points * num_edges + 1;

    # pre-allocate site array
    new_points = Vector{Vector{Float64}}(undef, num_new_points);

    new_points[end] = center(bn);

    counter = 1;
    for e ∈ box_edges
        v1, v2 = e;
        for λ ∈ steps
            new_points[counter] = λ .* v1 .+ (1-λ) .* v2;
            counter += 1;
        end
    end

    return new_points
end


# # Estimate Binomial draw probabilility

using DynamicHMCModels, MCMCChains

Random.seed!(1356779)

# Define a structure to hold the data.

Base.@kwdef struct BernoulliProblem
    "Total number of draws in the data."
    n::Int
    "Number of draws ' == 1' "
    obs::Vector{Int}
end;

# Write a function to return properly dimensioned transformation.

make_transformation(model::BernoulliProblem) =
    as((p = as𝕀, ))

# Add data

model = BernoulliProblem(; n = 9, obs = rand(Binomial(9, 2/3), 3))

# Make the type callable with the parameters *as a single argument*.

function (model::BernoulliProblem)(θ)
    @unpack n, obs = model        # extract the data
    @unpack p = θ
    loglikelihood(Binomial(n, p), obs)
end

# Use a flat priors (the default, omitted) for α

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);

# Sample 4 chains

a3d = Array{Float64, 3}(undef, 1000, 1, 4);
for j in 1:4
  global results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000;
    #initialization = (q = (p = 0.2,)),
    #reporter = NoProgressReport()
    )
  global posterior = P.transformation.(results.chain)

  for i in 1:1000
    a3d[i, 1, j] = values(posterior[i].p)
  end
end

# Create MCMCChains object

parameter_names = ["p"]
sections =   Dict(
  :parameters => parameter_names,
)
chns = create_mcmcchains(a3d, parameter_names, sections, start=1)
show(chns)

println()
DynamicHMC.Diagnostics.EBFMI(results.tree_statistics) |> display

println()
DynamicHMC.Diagnostics.summarize_tree_statistics(results.tree_statistics)


@testset "Katz" begin
    g5 = DiGraph(4)
    add_edge!(g5, 1, 2); add_edge!(g5, 2, 3); add_edge!(g5, 1, 3); add_edge!(g5, 3, 4)
    for g in testdigraphs(g5)
        z = @inferred(katz_centrality(g, 0.4))
        @test round.(z, 2) == [0.32, 0.44, 0.62, 0.56]
    end
end


function fscore(golds::Vector{UnitRange{Int}}, preds::Vector{UnitRange{Int}})
    g = Set(golds)
    p = Set(preds)
    inter = intersect(g, p)
    correct = length(inter)

    prec = correct / length(p)
    recall = correct / length(g)
    fvalue = 2*recall*prec/(recall+prec)
    println("precision: $(prec)")
    println("recall: $(recall)")
    println("f-value: $(fvalue)")
end

function accuracy(golds::Vector{Int}, preds::Vector{Int})
    @assert length(golds) == length(preds)
    correct = 0
    total = 0
    for i = 1:length(golds)
        golds[i] == preds[i] && (correct += 1)
        total += 1
    end
    correct / total
end


# our iterator type
type OrderedSubsets{T}
    members::Vector{T}
    valueFunction
    pq::Collections.PriorityQueue
    started::Bool
end

"Enumerate all the subsets ordered by the provided value function."
function orderedsubsets{T}(members::Vector{T}, f, order=:descending)
    orderby = x->-f([x])
    pqOrder = Base.Order.Reverse
    if order == :ascending
        orderby = x->f([x])
        pqOrder = Base.Order.Forward
    end
    OrderedSubsets(
        sort(members, by=orderby),
        f,
        Collections.PriorityQueue(Vector{Int64}, Float64, pqOrder),
        false
    )
end

Base.eltype(it::OrderedSubsets) = Vector{eltype(it.members)}
function Base.length(it::OrderedSubsets)
    length(it.members) < 40 ? 2^length(it.members) : Inf
end

Base.start{T}(it::OrderedSubsets{T}) = nothing

function Base.next{T}(it::OrderedSubsets{T}, state)
    if !it.started
        it.started = true
        Collections.enqueue!(it.pq, [1], it.valueFunction([it.members[1]]))
        return T[],nothing
    end

    # get the next best subset
    nextSubset = Collections.dequeue!(it.pq)

    # move this subset to look for the next "last member"
    if nextSubset[end] < length(it.members)
        nextSubsetInc = copy(nextSubset)
        nextSubsetInc[end] += 1
        Collections.enqueue!(it.pq, nextSubsetInc, it.valueFunction(T[it.members[i] for i in nextSubsetInc]))
    end

    # add the grown version of this subset
    if nextSubset[end] < length(it.members)
        nextSubsetGrown = [nextSubset; nextSubset[end]+1]
        Collections.enqueue!(it.pq, nextSubsetGrown, it.valueFunction(T[it.members[i] for i in nextSubsetGrown]))
    end

    T[it.members[i] for i in nextSubset],nothing
end

Base.done(it::OrderedSubsets, state) = isempty(it.pq) && it.started

# our iterator type
type ESKernelSubsets
    itr
    complement::Bool
    variances::Vector{Float64}
    lastSubset
    M
    numDone
    numTotal
end

function eskernelSubsetWeight(x, M)
    s = length(x)
    if s == M || s == 0
        return 1e12
    end
    (M-1)/(binomial(M,s)*s*(M-s))
end

"Enumerate all the subsets in descending order by the provided value function."
function eskernelsubsets(members, variances)
    M = length(members)

    # ensure all the variances are unique
    variances = variances .+ collect(M:-1:1)/1e5

    function subsetValue(x)

        s = length(x)
        if s == M || s == 0
            return 1e12
        elseif s > M/2
            return -1
        elseif s == 1
            return 1e10 + minimum(variances[x])
        end

        w = (M-1)/(binomial(M,s)*s*(M-s))
        minimum(variances[x])*w
    end

    ESKernelSubsets(
        orderedsubsets(members, subsetValue),
        false,
        variances,
        Int64[],
        M,
        0,
        M < 40 ? 2^M : Inf
    )
end

Base.eltype(it::ESKernelSubsets) = Vector{Float32}
function Base.length(it::ESKernelSubsets)
    it.numTotal
end

function Base.start(it::ESKernelSubsets)
    start(it.itr)
end

function Base.next(it::ESKernelSubsets, state)
    if !it.complement
        subsetElements,tmp = next(it.itr, nothing)

        it.lastSubset = subsetElements
        it.complement = true
        out = zeros(Float32, it.M)
        out[subsetElements] = 1
    else
        out = ones(Float32, it.M)
        out[it.lastSubset] = 0
        it.complement = false
    end
    it.numDone += 1

    return (out,eskernelSubsetWeight(find(out), it.M)),nothing
end

Base.done(it::ESKernelSubsets, state) = it.numDone >= it.numTotal


# TODO: Add HigherPrecision.jl tests here?
# TODO: Test Interior Point algorithm as well
@testset "Arbitrary Precision" begin
    prob = MVP.UnconstrainedProblems.examples["Rosenbrock"]

    f = MVP.objective(prob)
    g! = MVP.gradient(prob)
    h! = MVP.hessian(prob)

    @testset "BigFloat" begin
        x0 = big.(prob.initial_x)
        res = optimize(f, x0)
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) < 1e-8
        @test Optim.minimizer(res) ≈ [1.0, 1.0] atol=1e-4 rtol=0

        res = optimize(f, g!, x0)
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) < 1e-16
        @test Optim.minimizer(res) ≈ [1.0, 1.0] atol=1e-10 rtol=0

        res = optimize(f, g!, h!, x0)
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) < 1e-16
        @test Optim.minimizer(res) ≈ [1.0, 1.0] atol=1e-10 rtol=0

        lower = big.([-Inf, -Inf])
        upper = big.([0.5, 1.5])
        res = optimize(f, g!, lower, upper, x0, Fminbox(), Optim.Options(outer_g_abstol=sqrt(eps(big(1.0))), g_abstol=sqrt(eps(big(1.0)))))
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) ≈ 0.25 atol=1e-10 rtol=0
        @test Optim.minimizer(res) ≈ [0.5, 0.25] atol=1e-10 rtol=0

        res = optimize(f, lower, upper, x0, Fminbox(), Optim.Options(outer_g_abstol=sqrt(eps(big(1.0))), g_abstol=sqrt(eps(big(1.0)))))
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) ≈ 0.25 atol=1e-10 rtol=0
        @test Optim.minimizer(res) ≈ [0.5, 0.25] atol=1e-10 rtol=0
    end
end

@testset "Float32 doesn't work in Fminbox" begin
    f(x::AbstractVector{T}) where {T} = (T(1.0) - x[1])^2 + T(100.0) * (x[2] - x[1]^2)^2

    function g!(storage::AbstractVector{T}, x::AbstractVector{T}) where {T}
        storage[1] = T(-2.0) * (T(1.0) - x[1]) - T(400.0) * (x[2] - x[1]^2) * x[1]
        storage[2] = T(200.0) * (x[2] - x[1]^2)
    end

    for T = (Float32, Float64, BigFloat)
        lower = T[1.25, -2.1]
        upper = T[Inf, Inf]
        initial_x = T[2.0, 2.0]

        od = OnceDifferentiable(f, g!, initial_x)

        results = optimize(od, lower, upper, initial_x, Fminbox(GradientDescent()))
    end
end


struct Scale{T, N} <: Bijector{N}
    a::T
end

Base.:(==)(b1::Scale{<:Any, N}, b2::Scale{<:Any, N}) where {N} = b1.a == b2.a

function Scale(a::Union{Real,AbstractArray}; dim::Val{D} = Val(ndims(a))) where D
    return Scale{typeof(a), D}(a)
end

# field is a numerical parameter
function Functors.functor(::Type{<:Scale{<:Any,N}}, x) where N
    function reconstruct_scale(xs)
        return Scale{typeof(xs.a),N}(xs.a)
    end
    return (a = x.a,), reconstruct_scale
end

up1(b::Scale{T, N}) where {N, T} = Scale{T, N + 1}(a)

(b::Scale)(x) = b.a .* x
(b::Scale{<:AbstractMatrix, 1})(x::AbstractVecOrMat) = b.a * x
(b::Scale{<:AbstractMatrix, 2})(x::AbstractMatrix) = b.a * x
(ib::Inverse{<:Scale})(y) = Scale(inv(ib.orig.a))(y)
(ib::Inverse{<:Scale{<:AbstractVector}})(y) = Scale(inv.(ib.orig.a))(y)
(ib::Inverse{<:Scale{<:AbstractMatrix, 1}})(y::AbstractVecOrMat) = ib.orig.a \ y
(ib::Inverse{<:Scale{<:AbstractMatrix, 2}})(y::AbstractMatrix) = ib.orig.a \ y

# We're going to implement custom adjoint for this
logabsdetjac(b::Scale{<:Any, N}, x) where {N} = _logabsdetjac_scale(b.a, x, Val(N))

_logabsdetjac_scale(a::Real, x::Real, ::Val{0}) = log(abs(a))
_logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{0}) = fill(log(abs(a)), length(x))
_logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{1}) = log(abs(a)) * length(x)
_logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{1}) = fill(log(abs(a)) * size(x, 1), size(x, 2))
_logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{2}) = log(abs(a)) * length(x)
_logabsdetjac_scale(a::Real, x::AbstractArray{<:AbstractMatrix}, ::Val{2}) = map(x) do x
    _logabsdetjac_scale(a, x, Val(2))
end
_logabsdetjac_scale(a::AbstractVector, x::AbstractVector, ::Val{1}) = sum(x -> log(abs(x)), a)
_logabsdetjac_scale(a::AbstractVector, x::AbstractMatrix, ::Val{1}) = fill(sum(x -> log(abs(x)), a), size(x, 2))
_logabsdetjac_scale(a::AbstractVector, x::AbstractMatrix, ::Val{2}) = sum(x -> log(abs(x)), a)
_logabsdetjac_scale(a::AbstractVector, x::AbstractArray{<:AbstractMatrix}, ::Val{2}) = map(x) do x
    _logabsdetjac_scale(a, x, Val(2))
end
_logabsdetjac_scale(a::AbstractMatrix, x::AbstractVector, ::Val{1}) = logabsdet(a)[1]
_logabsdetjac_scale(a::AbstractMatrix, x::AbstractMatrix{T}, ::Val{1}) where {T} = logabsdet(a)[1] * ones(T, size(x, 2))
_logabsdetjac_scale(a::AbstractMatrix, x::AbstractMatrix, ::Val{2}) = logabsdet(a)[1]
function _logabsdetjac_scale(
    a::AbstractMatrix,
    x::AbstractArray{<:AbstractMatrix},
    ::Val{2},
)
    map(x) do x
        _logabsdetjac_scale(a, x, Val(2))
    end
end


mutable struct Cluster
	points::Vector{Tuple{Int64,Int64}}
	boundary::Rect
	r::Int64
	x_bound::Int64
	y_bound::Int64
end

Cluster(r::Int64,x_bound::Int64,y_bound::Int64) = Cluster([], Rect([ [0,0] [0,0] ]),r,x_bound,y_bound)

function add_point_to_cluster(cluster::Cluster,p::Tuple{Int64,Int64})
	if p in cluster.points
		return false
	end

	push!(cluster.points,p)
	update_rect(cluster::Cluster,p::Tuple{Int64,Int64})
	return true
end

function update_rect(cluster::Cluster,p::Tuple{Int64,Int64})
	if length(cluster.points) == 1
		cluster.boundary = Rect( bounded_diff(p[1],cluster.r,0),
							  	 bounded_diff(p[2],cluster.r,0),
							 	 bounded_sum(p[1],cluster.r,cluster.x_bound),
							  	 bounded_sum(p[2],cluster.r,cluster.y_bound))
		return
	end

	x1,x2 = x(cluster.boundary)
	y1,y2 = y(cluster.boundary)

	if length(cluster.points) == 2
		cluster.boundary = Rect(x1,y1,
								 bounded_sum(p[1],cluster.r,cluster.x_bound),
								 bounded_sum(p[2],cluster.r,cluster.y_bound))
		return
	end

	x1 = min(x1, bounded_diff(p[1],cluster.r,0))
	x2 = max(x2, bounded_sum(p[1],cluster.r,cluster.x_bound))
	y1 = min(y1, bounded_diff(p[2],cluster.r,0))
	y2 = max(y2, bounded_sum(p[2],cluster.r,cluster.y_bound))

	cluster.boundary = Rect(x1,y1,x2,y2)
end

function is_inside_cluster(cluster::Cluster,p::Tuple{Int64,Int64})
	return inside(p,cluster.boundary)
end


mutable struct GLMCopulaVCObs{T <: BlasReal, D, Link}
    # data
    y::Vector{T}
    X::Matrix{T}
    V::Vector{Matrix{T}}
    n::Int
    p::Int          # number of mean parameters in linear regression
    m::Int          # number of variance components
    # working arrays
    ∇β::Vector{T}   # gradient wrt β
    ∇resβ::Matrix{T}# residual gradient matrix d/dβ_p res_ij (each observation has a gradient of residual is px1)
    ∇τ::Vector{T}   # gradient wrt τ
    ∇θ::Vector{T}   # gradient wrt θ
    Hβ::Matrix{T}   # Hessian wrt β
    Hθ::Matrix{T}   # Hessian wrt variance components θ
    Hτ::Matrix{T}   # Hessian wrt τ
    res::Vector{T}  # residual vector res_i
    t::Vector{T}    # t[k] = tr(V_i[k]) / 2
    q::Vector{T}    # q[k] = res_i' * V_i[k] * res_i / 2
    xtx::Matrix{T}  # Xi'Xi
    storage_n::Vector{T}
    m1::Vector{T}
    m2::Vector{T}
    storage_p1::Vector{T}
    storage_p2::Vector{T}
    storage_np::Matrix{T}
    storage_pp::Matrix{T}
    added_term_numerator::Matrix{T}
    added_term2::Matrix{T}
    η::Vector{T}    # η = Xβ systematic component
    μ::Vector{T}    # μ(β) = ginv(Xβ) # inverse link of the systematic component
    varμ::Vector{T} # v(μ_i) # variance as a function of the mean
    dμ::Vector{T}   # derivative of μ
    d::D            # distribution()
    link::Link      # link function ()
    wt::Vector{T}   # weights wt for GLM.jl
    w1::Vector{T}   # working weights in the gradient = dμ/v(μ)
    w2::Vector{T}   # working weights in the information matrix = dμ^2/v(μ)
end

function GLMCopulaVCObs(
    y::Vector{T},
    X::Matrix{T},
    V::Vector{Matrix{T}},
    d::D,
    link::Link) where {T <: BlasReal, D, Link}
    n, p, m = size(X, 1), size(X, 2), length(V)
    @assert length(y) == n "length(y) should be equal to size(X, 1)"
    # working arrays
    ∇β  = Vector{T}(undef, p)
    ∇resβ  = Matrix{T}(undef, n, p)
    ∇τ  = Vector{T}(undef, 1)
    ∇θ  = Vector{T}(undef, m)
    Hβ  = Matrix{T}(undef, p, p)
    Hθ  = Matrix{T}(undef, m, m)
    Hτ  = Matrix{T}(undef, 1, 1)
    res = Vector{T}(undef, n)
    t   = [tr(V[k])/2 for k in 1:m]
    q   = Vector{T}(undef, m)
    xtx = transpose(X) * X
    storage_n = Vector{T}(undef, n)
    m1        = Vector{T}(undef, m)
    m2        = Vector{T}(undef, m)
    storage_p1 = Vector{T}(undef, p)
    storage_p2 = Vector{T}(undef, p)
    storage_np = Matrix{T}(undef, n, p)
    storage_pp = Matrix{T}(undef, p, p)
    added_term_numerator = Matrix{T}(undef, n, p)
    added_term2 = Matrix{T}(undef, p, p)
    η = Vector{T}(undef, n)
    μ = Vector{T}(undef, n)
    varμ = Vector{T}(undef, n)
    dμ = Vector{T}(undef, n)
    wt = Vector{T}(undef, n)
    fill!(wt, one(T))
    w1 = Vector{T}(undef, n)
    w2 = Vector{T}(undef, n)
    # constructor
    GLMCopulaVCObs{T, D, Link}(y, X, V, n, p, m, ∇β, ∇resβ, ∇τ, ∇θ, Hβ, Hθ,
        Hτ, res, t, q, xtx, storage_n, m1, m2, storage_p1, storage_p2, storage_np,
        storage_pp, added_term_numerator, added_term2, η, μ, varμ, dμ, d, link, wt, w1, w2)
end

"""
GLMCopulaVCModel
GLMCopulaVCModel(gcs)
Gaussian copula variance component model, which contains a vector of
`GLMCopulaVCObs` as data, model parameters, and working arrays.
"""
struct GLMCopulaVCModel{T <: BlasReal, D, Link} <: MathProgBase.AbstractNLPEvaluator
    # data
    data::Vector{GLMCopulaVCObs{T, D, Link}}
    Ytotal::T
    ntotal::Int     # total number of singleton observations
    p::Int          # number of mean parameters in linear regression
    m::Int          # number of variance components
    # parameters
    β::Vector{T}    # p-vector of mean regression coefficients
    τ::Vector{T}    # inverse of linear regression variance parameter
    θ::Vector{T}    # m-vector: [θ12, ..., θm2]
    # working arrays
    ∇β::Vector{T}   # gradient from all observations
    ∇τ::Vector{T}
    ∇θ::Vector{T}
    XtX::Matrix{T}  # X'X = sum_i Xi'Xi
    Hβ::Matrix{T}    # Hessian from all observations
    Hθ::Matrix{T}
    Hτ::Matrix{T}
    Ainv::Matrix{T}
    Aevec::Matrix{T}
    M::Matrix{T}
    vcov::Matrix{T}
    ψ::Vector{T}
    TR::Matrix{T}         # n-by-m matrix with tik = tr(Vi[k]) / 2
    QF::Matrix{T}         # n-by-m matrix with qik = res_i' Vi[k] res_i
    storage_n::Vector{T}
    storage_m::Vector{T}
    storage_θ::Vector{T}
    d::Vector{D}
    link::Vector{Link}
    penalized::Bool
end

function GLMCopulaVCModel(gcs::Vector{GLMCopulaVCObs{T, D, Link}}; penalized::Bool = false) where {T <: BlasReal, D<:Union{Poisson, Bernoulli}, Link}
    n, p, m = length(gcs), size(gcs[1].X, 2), length(gcs[1].V)
    β       = Vector{T}(undef, p)
    τ       = [1.0]
    θ       = Vector{T}(undef, m)
    ∇β      = Vector{T}(undef, p)
    ∇τ      = Vector{T}(undef, 1)
    ∇θ      = Vector{T}(undef, m)
    XtX     = zeros(T, p, p) # sum_i xi'xi
    Hβ      = Matrix{T}(undef, p, p)
    Hθ      = Matrix{T}(undef, m, m)
    Hτ      = Matrix{T}(undef, 1, 1)
    Ainv    = zeros(T, p + m, p + m)
    Aevec   = zeros(T, p + m, p + m)
    M       = zeros(T, p + m, p + m)
    vcov    = zeros(T, p + m, p + m)
    ψ       = Vector{T}(undef, p + m)
    TR      = Matrix{T}(undef, n, m) # collect trace terms
    Ytotal  = 0.0
    ntotal  = 0.0
    d       = Vector{D}(undef, n)
    link    = Vector{Link}(undef, n)
    for i in eachindex(gcs)
        ntotal  += length(gcs[i].y)
        Ytotal  += sum(gcs[i].y)
        BLAS.axpy!(one(T), gcs[i].xtx, XtX)
        TR[i, :] = gcs[i].t
        d[i] = gcs[i].d
        link[i] = gcs[i].link
    end
    QF        = Matrix{T}(undef, n, m)
    storage_n = Vector{T}(undef, n)
    storage_m = Vector{T}(undef, m)
    storage_θ = Vector{T}(undef, m)
    GLMCopulaVCModel{T, D, Link}(gcs, Ytotal, ntotal, p, m, β, τ, θ,
        ∇β, ∇τ, ∇θ, XtX, Hβ, Hθ, Hτ, Ainv, Aevec, M, vcov, ψ, TR, QF,
        storage_n, storage_m, storage_θ, d, link, penalized)
end

"""
    loglikelihood!(gc, β, τ, θ)
Calculates the loglikelihood of observing `y` given mean `μ`, for the Poisson and Bernoulli base distribution using the GLM.jl package.
"""
function loglikelihood!(
    gc::Union{GLMCopulaVCObs, Poisson_Bernoulli_VCObs},
    β::Vector{T},
    θ::Vector{T},
    needgrad::Bool = false,
    needhess::Bool = false;
    penalized::Bool = false
    ) where {T <: BlasReal}
    # n, p, m = size(gc.X, 1), size(gc.X, 2), length(gc.V)
    needgrad = needgrad || needhess
    if needgrad
        fill!(gc.∇β, 0)
        fill!(gc.∇θ, 0)
    end
    needhess && fill!(gc.Hβ, 0)
    update_res!(gc, β)
    # @show gc.res
    standardize_res!(gc)
    # @show gc.res
    fill!(gc.∇resβ, 0.0) # fill gradient of residual vector with 0
    std_res_differential!(gc) # this will compute ∇resβ

    # evaluate copula loglikelihood
    @inbounds for k in 1:gc.m
        mul!(gc.storage_n, gc.V[k], gc.res) # storage_n = V[k] * res
        if needgrad
            BLAS.gemv!('T', θ[k], gc.∇resβ, gc.storage_n, 1.0, gc.∇β) # stores ∇resβ*Γ*res (standardized residual)
        end
        gc.q[k] = dot(gc.res, gc.storage_n) / 2
    end
    # loglikelihood
    logl = QuasiCopula.component_loglikelihood(gc)
    tsum = dot(θ, gc.t)
    logl += -log(1 + tsum)
    qsum  = dot(θ, gc.q)
    logl += log(1 + qsum)
    # add L2 ridge penalty
    if penalized
        logl -= 0.5 * dot(θ, θ)
    end
    if needgrad
        inv1pq = inv(1 + qsum)
        inv1pt = inv(1 + tsum) #
        # gc.∇θ .= inv1pq * gc.q .- inv1pt * gc.t
        gc.m1 .= gc.q
        gc.m1 .*= inv1pq
        gc.m2 .= gc.t
        gc.m2 .*= inv1pt
        gc.∇θ .= gc.m1 .- gc.m2
        if penalized
            gc.∇θ .-= θ
        end
        if needhess
            BLAS.syrk!('L', 'N', -abs2(inv1pq), gc.∇β, 1.0, gc.Hβ) # only lower triangular
            # does adding this term to the approximation of the hessian violate negative semidefinite properties?
            fill!(gc.added_term_numerator, 0.0) # fill gradient with 0
            fill!(gc.added_term2, 0.0) # fill hessian with 0
            @inbounds for k in 1:gc.m
                mul!(gc.added_term_numerator, gc.V[k], gc.∇resβ) # storage_n = V[k] * res
                BLAS.gemm!('T', 'N', θ[k], gc.∇resβ, gc.added_term_numerator, one(T), gc.added_term2)
            end
            gc.added_term2 .*= inv1pq
            gc.Hβ .+= gc.added_term2
            gc.Hβ .+= QuasiCopula.glm_hessian(gc)

            # hessian for vc
            fill!(gc.Hθ, 0.0)
            BLAS.syr!('U', one(T), gc.m2, gc.Hθ)
            BLAS.syr!('U', -one(T), gc.m1, gc.Hθ)
            copytri!(gc.Hθ, 'U')
            # gc.Hθ .= gc.m2 * transpose(gc.m2) - gc.m1 * transpose(gc.m1)
        end
        gc.storage_p2 .= gc.∇β .* inv1pq
        # @show gc.res
        gc.res .= gc.y .- gc.μ
        gc.∇β .= QuasiCopula.glm_gradient(gc)
        # @show gc.res
        gc.∇β .+= gc.storage_p2
    end
    logl
end

function loglikelihood!(
    gcm::Union{GLMCopulaVCModel, Poisson_Bernoulli_VCModel},
    needgrad::Bool = false,
    needhess::Bool = false
    )
    logl = 0.0
    if needgrad
        fill!(gcm.∇β, 0)
        fill!(gcm.∇θ, 0)
    end
    if needhess
        fill!(gcm.Hβ, 0)
        fill!(gcm.Hθ, 0)
    end
    logl = zeros(Threads.nthreads())
    Threads.@threads for i in eachindex(gcm.data)
        @inbounds logl[Threads.threadid()] += loglikelihood!(gcm.data[i], gcm.β,
         gcm.θ, needgrad, needhess; penalized = gcm.penalized)
     end
     @inbounds for i in eachindex(gcm.data)
         if needgrad
             gcm.∇β .+= gcm.data[i].∇β
             gcm.∇θ .+= gcm.data[i].∇θ
         end
         if needhess
             gcm.Hβ .+= gcm.data[i].Hβ
             gcm.Hθ .+= gcm.data[i].Hθ
         end
     end
    sum(logl)
end


using FactCheck

#push!(LOAD_PATH, "/Users/dsanders/Dropbox/papers/billiards/")

using BilliardModels

#using Vector2d

facts("Particle tests") do

    p = Particle([1, 2.], [3., 4.5])
    @fact p.x => Vector2D{Float64}(1,2)
    @fact p.v => Vector2D{Float64}(3.0,4.5)

end

facts("Disc tests") do
    d = Disc([1, 2], 3)
    @fact d.radius => 3.0

    d = Disc([0, 0], 1)
    p = Particle([-2, 0], [1, 0])

    @fact collision_time(p, d) => 1

end

facts("Plane tests") do
    o = Plane([0,0], [1, 1])
    p = Particle([-1, 0], [1, 0])

    @fact collision_time(p, o) => 1
end

facts("Billiard table tests") do
    table = Sinai_billiard(0.1)
    p = Particle([0.2, 0.], [1., 0.])

    next_collision = calculate_next_collision(p, table, nothing)
    @fact next_collision[2] => Vector2D(-1., 0.)  # new velocity after reflecting on vertical boundary
    @fact next_collision[3] => 0.3  # col]\

    p = Particle([0.2, 0.], [-1., 0.])
    next_collision = calculate_next_collision(p, table, nothing)
    @fact isa(next_collision[4], Disc) => true

end

facts("Finite plane") do
    f = FinitePlane([0., 0.], [1.,1.])
    p = Particle([0.5, 0.], [0., 1.])

    @fact collision_time(p, f) => 0.5

    p = Particle([1.5, 0.], [0., 1.])
    @fact collision_time(p, f) => -Inf
end


# bdf_utils
@inline function U!(k, U)
  @inbounds for r = 1:k
    U[1,r] = -r
    for j = 2:k
      U[j,r] = U[j-1,r] * ((j-1) - r)/j
    end
  end
  nothing
end

function R!(k, ρ, cache)
  @unpack R = cache
  @inbounds for r = 1:k
    R[1,r] = -r * ρ
    for j = 2:k
      R[j,r] = R[j-1,r] * ((j-1) - r * ρ)/j
    end
  end
  nothing
end

# This functions takes help of D2 array to create backward differences array D
# Ith row of D2 keeps Ith order backward differences (∇ⁱyₙ)
function backward_diff!(cache::OrdinaryDiffEqMutableCache, D, D2, k, flag=true)
  flag && copyto!(D[1], D2[1,1])
  for i = 2:k
    for j = 1:(k-i+1)
      @.. D2[i,j] = D2[i-1,j] - D2[i-1,j+1]
    end
    flag && copyto!(D[i], D2[i,1])
  end
end

function backward_diff!(cache::OrdinaryDiffEqConstantCache, D, D2, k, flag=true)
  flag && (D[1] = D2[1,1])
  for i = 2:k
    for j = 1:(k-i+1)
      D2[i,j] = D2[i-1,j] - D2[i-1,j+1]
    end
    flag && (D[i] = D2[i,1])
  end
end

# this function updates backward difference array D when stepsize gets change
# Formula -> D = D * (R * U)
# and it is taken from the paper -
# Implementation of an Adaptive BDF2 Formula and Comparison with the MATLAB Ode15s paper
# E. Alberdi Celaya, J. J. Anza Aguirrezabala, and P. Chatzipantelidis
function reinterpolate_history!(cache::OrdinaryDiffEqMutableCache, D, R, k)
  @unpack tmp = cache
  fill!(tmp,zero(eltype(D[1])))
  for j = 1:k
    for k = 1:k
      @.. tmp += D[k] * R[k,j]
    end
    D[j] .= tmp
    fill!(tmp, zero(eltype(tmp)))
  end
end

function reinterpolate_history!(cache::OrdinaryDiffEqConstantCache, D, R, k)
  tmp = zero(D[1])
  for j = 1:k
    for k = 1:k
      tmp += D[k] * R[k,j]
    end
    D[j] = tmp
  end
end

global const γₖ = @SVector[sum(1//j for j in 1:k) for k in 1:6]

# this stepsize and order controller is taken from
# Implementation of an Adaptive BDF2 Formula and Comparison with the MATLAB Ode15s paper
# E. Alberdi Celaya, J. J. Anza Aguirrezabala, and P. Chatzipantelidis
function stepsize_and_order!(cache, est, estₖ₋₁, estₖ₊₁, h, k)
  zₛ = 1.2
  zᵤ = 0.1
  Fᵤ = 10

  expo = 1/(k+1)
  z = zₛ * ((est)^expo)
  F = inv(z)

  hₖ₋₁ = 0.0
  hₖ₊₁ = 0.0

  if z <= zₛ
    # step is successful
    # precalculations
    if z <= zᵤ
      hₖ = Fᵤ * h
    elseif zᵤ < z <= zₛ
      hₖ = F * h
    end

    if k > 1
      expo = 1/k
      zₖ₋₁ = 1.3 * ((estₖ₋₁)^expo)
      Fₖ₋₁ = inv(zₖ₋₁)
      if zₖ₋₁ <= 0.1
        hₖ₋₁ = 10 * h
      elseif 0.1 < zₖ₋₁ <= 1.3
        hₖ₋₁ = Fₖ₋₁ * h
      end
    end

    expo = 1/(k+2)
    zₖ₊₁ = 1.4 * ((estₖ₊₁)^expo)
    Fₖ₊₁ = inv(zₖ₊₁)

    if zₖ₊₁<= 0.1
      hₖ₊₁ = 10 * h
    elseif 0.1 < zₖ₊₁ <= 1.4
      hₖ₊₁ = Fₖ₊₁ * h
    end
    # adp order and step conditions
    if hₖ₋₁ > hₖ
      hₙ = hₖ₋₁
      kₙ = max(k-1,1)
    else
      hₙ = hₖ
      kₙ = k
    end
    if hₖ₊₁ > hₙ
      hₙ = hₖ₊₁
      kₙ = min(k+1,5)
    end
    if hₙ < h
      hₙ = h
      kₙ = k
    end
    cache.h = hₙ
    cache.order = kₙ
    return true
  else
    # step is not successful
    if cache.c >= 1  # postfail
      cache.h = h/2
      cache.order = k
      return
    end
    if 1.2 < z <= 10
      hₖ = F * h
    elseif z > 10
      hₖ = 0.1 * h
    end
    hₙ = hₖ
    kₙ = k
    if k > 1
      expo = 1/k
      zₖ₋₁ = 1.3 * ((estₖ₋₁)^expo)
      Fₖ₋₁ = inv(zₖ₋₁)
      if 1.3 < zₖ₋₁ <= 10
        hₖ₋₁ = Fₖ₋₁ * h
      elseif zₖ₋₁ > 10
        hₖ₋₁ = 0.1 * h
      end

      if hₖ₋₁ > hₖ
        hₙ = min(h,hₖ₋₁)
        kₙ = max(k-1,1)
      end
    end
    cache.h = hₙ
    cache.order = kₙ
    return false
  end
end


using StaticTrafficAssignment
using DataFrames

nnodes, ftnode, linkdf, trips, geometry, bestsolution = readtntpdata("examples/data/SiouxFalls/")

nzones = size(trips)[1]
zonedf = DataFrame(:id => 1:nzones, :issource => ones(Bool, nzones), :issink => ones(Bool, nzones), :thruallowed => (1:nzones) .>= ftnode)

network = SimpleNetwork(nnodes, linkdf, zonedf; idkey=:id, issrckey=:issource, issnkkey=:issink, thrukey=:thruallowed,  upnkey=:init_node, dwnkey=:term_node)

uecostfn = CostFunctionUE(network, BPR)
socostfn = CostFunctionSO(network, BPR)

@time aonsoln = allornothing(network, trips, uecostfn);

@time uesoln_msa = msa(network, trips, uecostfn);
@time sosoln_msa = msa(network, trips, socostfn);

@time uesoln_fw = frankwolfe(network, trips, uecostfn);
@time sosoln_fw = frankwolfe(network, trips, socostfn);

@time uesoln_cfw = conjugatefrankwolfe(network, trips, uecostfn);
@time sosoln_cfw = conjugatefrankwolfe(network, trips, socostfn);

@time uesoln_B = algorithmB(network, trips, uecostfn; λ=0.1);
@time sosoln_B = algorithmB(network, trips, socostfn; λ=0.05);


module Transforms

import Base:   ==
import ACE: read_dict, write_dict, transform, transform_d, inv_transform
abstract type DistanceTransform end

export PolyTransform, IdTransform, MorseTransform, AgnesiTransform

poly_trans(p, r0, r) = @fastmath(((1+r0)/(1+r))^p)

poly_trans_d(p, r0, r) = @fastmath((-p/(1+r0)) * ((1+r0)/(1+r))^(p+1))

poly_trans_inv(p, r0, x) = ( (1+r0)/(x^(1/p)) - 1 )

# TODO: generalise the distance transform to allow
#       ((a + r0) / (a + r) )^p

@doc raw"""
Implements the distance transform
```math
   x(r) = \Big(\frac{1 + r_0}{1 + r}\Big)^p
```

Constructor:
```
PolyTransform(p, r0)
```
"""
struct PolyTransform{TP, T} <: DistanceTransform
   p::TP
   r0::T
end

write_dict(T::PolyTransform) =
   Dict("__id__" => "ACE_PolyTransform", "p" => T.p, "r0" => T.r0)

PolyTransform(D::Dict) = PolyTransform(D["p"], D["r0"])

read_dict(::Val{:ACE_PolyTransform}, D::Dict) = PolyTransform(D)

transform(t::PolyTransform, r::Number) = poly_trans(t.p, t.r0, r)

transform_d(t::PolyTransform, r::Number) = poly_trans_d(t.p, t.r0, r)

inv_transform(t::PolyTransform, x::Number) = poly_trans_inv(t.p, t.r0, x)

(t::PolyTransform)(x) = transform(t, x)

"""
`IdTransform`: Implements the distance transform `z -> z`;
Primarily used for the z-coordinate for the EnvPairPots

Constructor: `IdTransform()`
"""
struct IdTransform <: DistanceTransform
end

write_dict(T::IdTransform) =  Dict("__id__" => "ACE_IdTransform")
IdTransform(D::Dict) = IdTransform()
read_dict(::Val{:ACE_IdTransform}, D::Dict) = IdTransform(D)
transform(t::IdTransform, z::Number) = z
transform_d(t::IdTransform, r::Number) = one(r)
inv_transform(t::IdTransform, x::Number) = x

read_dict(::Val{:SHIPs_IdTransform}, D::Dict) =
   read_dict(Val{:ACE_IdTransform}(), D)

@doc raw"""
Implements the distance transform
```math
   x(r) = \exp( - \lambda (r/r_0) )
```

Constructor:
```
MorseTransform(lambda, r0)
```
"""
struct MorseTransform{T} <: DistanceTransform
   lambda::T
   r0::T
end

write_dict(T::MorseTransform) =
   Dict("__id__" => "ACE_MorseTransform", "lambda" => T.p, "r0" => T.r0)
MorseTransform(D::Dict) = MorseTransform(D["lambda"], D["r0"])
read_dict(::Val{:ACE_MorseTransform}, D::Dict) = MorseTransform(D)
transform(t::MorseTransform, r::Number) = exp(- t.lambda * (r/t.r0))
transform_d(t::MorseTransform, r::Number) = (-t.lambda/t.r0) * exp(- t.lambda * (r/t.r0))
inv_transform(t::MorseTransform, x::Number) = - t.r0/t.lambda * log(x)
(t::MorseTransform)(x) = transform(t, x)

@doc raw"""
Implements the distance transform
```math
   x(r) = \frac{1}{1 + a (r/r_0)^p}
```
with default $a = (p-1)/(p+1)$. That default is chosen such that
$|x'(r)|$ is maximised at $r = r_0$. Default for $p$ is $p = 2$. Any value
$p > 1$ is permitted.

Constructor:
```
AgnesiTransform(r0, [p, [, a]])
```
"""
struct AgnesiTransform{T, TP} <: DistanceTransform
   r0::T
   p::TP
   a::T
   c::T
end

AgnesiTransform(r0, p=2, a=(p-1)/(p+1)) = (@assert p > 1;
                             AgnesiTransform(r0, p, a, - a * p / r0))

write_dict(T::AgnesiTransform) =
Dict("__id__" => "ACE_AgnesiTransform", "r0" => T.r0, "p" => T.p, "a" => T.a)
AgnesiTransform(D::Dict) = AgnesiTransform(D["r0"], D["p"], D["a"])
read_dict(::Val{:ACE_AgnesiTransform}, D::Dict) = AgnesiTransform(D)
transform(t::AgnesiTransform, r::Number) = @fastmath 1 / (1 + t.a * (r/t.r0)^t.p)
transform_d(t::AgnesiTransform, r::Number) = (s1 = (r/t.r0); s2 = s1^(t.p-1);
                                         @fastmath t.c * s2 / (1+t.a * s2*s1)^2)
inv_transform(t::AgnesiTransform, x::Number) = t.r0 * ( (1/x-1)/t.a )^(1/t.p)
(t::AgnesiTransform)(x) = transform(t, x)

end


# functions to generate distance fields for stroke types
function circle{T}(uv::Vec{2,T})
    T(1)-norm(uv)
end

function dots(gap, resolution)
    X = linspace(0, gap, resolution[1])
    Y = linspace(0, 1, resolution[2])
    Float16[
        circle(Vec(x-0.5,y-0.5))
    for x=X, y=Y]
end

dist(a, b) = abs(a-b)
mindist(x, a, b) = min(dist(a, x), dist(b, x))
function gappy(x, ps)
    n = length(ps)
    x <= first(ps) && return first(ps) - x
    for j=1:(n-1)
        p0 = ps[j]
        p1 = ps[min(j+1, n)]
        if p0 <= x && p1 >= x
            return mindist(x, p0, p1) * (isodd(j) ? 1 : -1)
        end
    end
    return last(ps) - x
end
function ticks(points, resolution)
    Float16[gappy(x, points) for x=linspace(first(points),last(points), resolution)]
end

points = [0., 0.01, 0.02, 0.03, 0.09, 0.6, 0.8, 0.9]

text = Texture(ticks(points, 100), x_repeat=:repeat)


using Unitful
using UnitfulAtomic
using BesselFunctions
using DifferentialEquations
using LegendrePolynomials
using QuadGK
using Dierckx

const a0 = u"a0_au"
const Eh = u"Eh_au"
const me = u"me_au"
const ħ = u"ħ_au"

kFromE(E,m=1me) = uconvert(a0^-1, sqrt(2*m*E) / ħ)
EFromk(k,m=1me) = uconvert(Eh, ħ^2*k^2/2m )

"""
    DerivFunc(u,u′,r, E,k,potential,l,m=1mₑ)

This is the differential equation derivative function. This function should work
in Unitful quantities.
"""
function DerivFunc(u,u′,r, E,k,potential,l,m=1me)
    # Expressing this as the plane wave expansion with the Y_l0 at the end,
    # gives rise to the sqrt() prefactors.
    jl = sqrt(4pi)*sqrt(2l+1) * 1/k * ricj(l, k*r)

    prefac = 2m/ħ^2
    U = -prefac*potential(r)

    u′′ = @. (l*(l+1)/r^2)*u - ((prefac*E + U)*u + U*jl)

    return u′,u′′
end

"""
    SolveOutwards(rspan, E,k,potential,l ; init_derivs=[0.0,1.0])

Calculates the wavefunction for the off-shell T-matrix problem. Integrates
outwards using `DifferentialEquations` with `rspan`. The ODE derivative is
created from the `potential` given (which should handle `Unitful` quantities).

The incoming wave has a wavenumber of `k` and the energy that the problem is
calculated on is `E`.

TODO: Put the ODE here.

Note: the inhomogeneous part is not included in the solution. i.e. this is
solving for \tilde{f}. Also not that there are two solutions calculated, which
need to be matched to the appropriate boundary conditions using e.g. `ElminateIngoing`.

Returns (r,u,u′)
"""
function SolveOutwards(rspan, E,k,potential,l ; init_derivs=[0.0,1.0])
    # Going to deal with units in the deriv funcs
    # Note: wavefunctions are unitless in the picture with \phi_k = (2\pi)^-3/2 exp(ikr)
    # But u = r*R so there is units for u.

    prob_func = function (f,p,r)
        r = r*a0
        u = f[1,:] .* a0
        u′ = f[2,:]
        du,du′ = DerivFunc(u, u′,r, E, k, potential,l)
        [ustrip.(NoUnits, du)' ; ustrip.(1/a0, du′)']
    end

    # Not sure if I will need to make this a more carefully chosen derivative.
    # At least the value being zero will be necessary.
    u0 = ComplexF64[[0. 0.] ; permutedims(init_derivs)]

    # Needs to include the jl part that is in f, as the BC is for f not tildef.
    jl = sqrt(4pi)*sqrt(2l+1) * 1/k * ricj(l, k*rspan[1])
    u0 .-= [ustrip(a0, jl) ; 0.]

    prob = ODEProblem(prob_func, u0, ustrip.(a0, rspan))
    sol = solve(prob, Tsit5(), abstol=1e-10, reltol=1e-12)

    this_soln = reduce((x,y) -> cat(x,y,dims=3), sol.u)
    r = sol.t .* a0

    u = this_soln[1,:,:] .* a0
    u′ = this_soln[2,:,:]
    return r, u, u′
end

"""
Old function, hasn't been tested since included in the module.
"""
function SolveInwards(r_list, E,k,potential,l, A)

    prob_func = function (f,p,r)
        r = r*a0
        u = f[1] .* a0
        u′ = f[2]
        du,du′ = DerivFunc(u, u′,r, E, k, potential,l)
        [ustrip.(NoUnits, du) ; ustrip.(1/a0, du′)]
    end

    rspan = extrema(r_list)
    rmax = rspan[2]

    κ = kFromE(E)
    u0 = A*[rich₊(l, κ*rmax) ; ustrip(1/a0, drich₊dr(l, κ, rmax))]

    prob = ODEProblem(prob_func, u0, ustrip.(a0, reverse(rspan)))
    sol = solve(prob, Tsit5(), abstol=1e-10, reltol=1e-12)#, callback=DiscreteCallback(RenormCond, RenormAffect!))

    # this_r = filter(x -> rspan[1] <= x <= sol.t[end]*a0, r_list)
    this_r = r_list
    this_soln = sol.(ustrip.(a0,this_r))
    this_soln = reduce(hcat, this_soln)

    r = this_r
    u = this_soln[1,:] .* a0
    u′ = this_soln[2,:]
    # coeffs = MatchSolutionUnitless(u, u′, sol.t[end], ustrip(1/a0,k), l)

    # conv = TransformToPureIngoingOutgoing(coeffs)

    return r, u, u′
end

function MatchSolution(u, u′, r, k, l)
    u = ustrip.(a0,u)
    u′ = u′
    r = ustrip(a0,r)
    k = ustrip(1/a0,k)
    MatchSolutionUnitless(u,u′,r,k,l)
end
function MatchSolution(f, r, k, l)
    r = ustrip(a0,r)
    k = ustrip(1/a0,k)
    MatchSolutionUnitless(f, r, k, l)
end
function MatchSolutionUnitless(u, u′, r, k, l)
    f = [u' ; u′']
    MatchSolutionUnitless(f, r, k, l)
end
function MatchSolutionUnitless(f, r, κ, l)
    M = [rich₊(l,κ*r)    rich₋(l,κ*r) ;
         drich₊dr(l,κ,r) drich₋dr(l,κ,r)]
    coeffs = M\f
end

"""
TransformToPureIngoingOutgoing(C)

Given a coeff matrix `C`, relating two solutions obtained from `SolveOutwards`
to the ingoing/outgoing states, combine the two solutions to satisfy the
boundary conditions.

Note: because of the inhomogeneity, this must satisfy u = A·u1 + (1-A)·u2
instead of the usual freedom of u = A·u1 + B·u2.
"""
function TransformToPureIngoingOutgoing(C)
    x1 = -C[2,2] / (C[2,1] - C[2,2])
    x2 = -C[1,2] / (C[1,1] - C[1,2])

    M = [[x1;(1-x1)] [x2;(1-x2)]]
    return M
end

"""
EliminateIngoing(u, u′, r, k, l)

Matches to the boundary conditions for \tilde{u} using the output from `SolveOutwards`.
"""
function EliminateIngoing(u, u′, r, k, l)
    coeffs = MatchSolution(u[:,end], u′[:,end], r[end], k, l)
    M = TransformToPureIngoingOutgoing(coeffs)
    M = M[:,1:1]
    M = permutedims(M)

    return dropdims(M*u, dims=1),
           dropdims(M*u′, dims=1)
end

"""
    AddBesselPart(r, tildeu, tildeu′, k, l)

Includes the inhomogeneity that was missing from the ODE solution. i.e. converts
from \tilde{u} to u.
"""
function AddBesselPart(r, tildeu, tildeu′, k, l)
    u = sqrt(4pi)*sqrt(2l+1) * 1/k * ricj.(l,k*r) + tildeu
    u′ = sqrt(4pi)*sqrt(2l+1) * 1/k * dricjdr.(l,k,r) + tildeu′

    return u,u′
end

"""
Tmat_l(r,u,potential,p,costh,l)

Once a complete solution for u has been obtained, this function calculates the
`T` matrix element. The `u` going into here must be the full matched solution
and not \tilde{u}.

Note: this was before I considered T_l. This term is instead one of the Legendre
contributions of T(theta) for a specific `costh` and `l`.
"""
function Tmat_l(r,u,potential,p,costh,l)
    # One l component only. costh = p⋅z (or p⋅k more generally)
    V = r -> ustrip(Eh,potential(r*a0))
    p2 = ustrip(1/a0,p)
    r2 = ustrip.(a0,r)

    interp_r = Spline1D(r2,real.(ustrip.(a0,u)),k=1)
    interp_i = Spline1D(r2,imag.(ustrip.(a0,u)),k=1)
    interp = r -> interp_r(r) + im*interp_i(r)
    out, = quadgk(r -> interp(r)*V(r)*ricj(l,p2*r), extrema(r2)...)
    out *= a0*Eh*a0
    # These terms come from the plane wave prefactors for p.
    out *= sqrt(4pi)*sqrt(2l+1) * 1/p * Pl(costh,l)
    # Adjusting for my lack of (2π)^-3/2 in the wavefunctions
    out /= (2π)^3
end

fFromT(T,m=me) = uconvert(a0, -(2pi)^2*m/ħ^2 * T)

"""
TCSandOptTheorem(faa, k)

Compare the total cross section with the optical theorem result, assuming that
`faa` is the forward scattering part (i.e. `faa` = f(θ=0)) and assuming a low
enough energy that the scattering is s-wave.

`faa` can be obtained from `fFromT`.
"""
function TCSandOptTheorem(faa, k)
    DCS = abs2(faa)
    TCS = 4pi * DCS

    opt_theorem = 4pi/k * imag(faa)

    a = sqrt(TCS/4pi)
    return (DCS=DCS, TCS=TCS, opt=opt_theorem, a=a)
end

# |faa|^2 =a^2
# -(2pi)^2*m/\hbar^2 * T = a
# hbr^2 a / (2pi)^2 / m


# Particle in a box
using Plots
using LinearAlgebra: eigen
using HelFEM: HelFEM

a, b = -33.3, 7.1
pb = HelFEM.PolynomialBasis(:lip, 10)
B = HelFEM.FEMBasis(pb, range(a, b, length=21))
DD = HelFEM.radial_integral(B, q->1.0, lderivative=true, rderivative=true)
S = HelFEM.radial_integral(B, q->1.0)
e = eigen(DD, S)
# Compare energies with analytical energies (En = (πn/L)^2)
N=length(B)
es = (1:N) .^ 2 .* (π / (b - a))^2
e.values .- es
# Plot the wavefunctions -- they should be sinewaves:
xs = range(a, b, length=1001)
Bri = B(xs)
ns = 1:10
plot(xs, Bri*e.vectors[:,ns], alpha=(ns .^ -0.6)', legend=false)


using StructJuMP

m = StructuredModel()

@variable(m, xh0[ASSETS] >= 0)
@variable(m, xb0[ASSETS] >= 0)
@constraint(m, (1+Ct)*sum(val[j]*xb0[j] for j=ASSETS) <= BUDGET)

for n1 in NODES
    if PARENT[n1] == 0
        sub = StructuredModel(parent=m)
        @variable(sub, xh1[ASSETS] >= 0)
        @variable(sub, xb1[ASSETS] >= 0)
        @variable(sub, xs1[ASSETS] >= 0)
        for j in ASSETS
            @constraint(sub, xh1[j] == (1+Ret[n1,j])*xh0[j]+xb1[j]-xs1[j])
        end
        @constraint(sub, (1-Ct)*sum(val[j]*xs1[j] for j=ASSETS) == Liability1 + (1+Ct)*sum(val[j]*xb1[j]))
    end

    for n2 in NODES
        if PARENT[n2] == n1
            sub = StructuredModel(parent=m)
            @variable(sub, xh2[ASSETS] >= 0)
            @variable(sub, xb2[ASSETS] >= 0)
            @variable(sub, xs2[ASSETS] >= 0)
            for j in ASSETS
                @constraint(sub, xh2[j] == (1+Ret2[n2,j])*xh1[j] + xb2[j] - xs2[j])
            end
            @constraint(sub, (1-Ct)*sum(val[j]*xs2[j] for j=ASSETS) == Liability2 + (1+Ct)*sum(val[j]*xb2[j]))
            @variable(sub, wealth)
            @constraint(sub, wealth == Prob[n1]*Prob[n2]*sum(val[j]*xh2[j] for j=ASSETS))
            @objective(sub, mu - Rho*((wealth^2-mu^2))*Prob[n1]*Prob[n2])
        end
    end
end

@constraint(m, sum(Prob[n1]*Prob[n2]*wealth[(n1,n2)] for n1=NODES n2=NODES if PARENT[n1]==0 && PARENT[n2]==n1))

# things of note:
# Block(parent::Union(Model,Block), index) constructor
# should do name mangling so that variables get associated with their block
# add function block(parent, child1, child2, ...) that returns reference to subproblem. This should then be passable as index to subproblem variables


using RadialBasisFunctionModels
using MLJBase

r = RadialBasisFunctionModels.RBFInterpolator(; kernel_name = :multiquadric, kernel_args = [1,])
X,y = @load_boston

R = machine(r, X, y)
fit!(R)

predict(R, X)


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ 7cb1f900-375e-11eb-3cc9-cf8a8fb4fc04
file = readlines("day_3.txt")

# ╔═╡ 8f19f0e8-375e-11eb-38fa-b79ed15de216
md"Getting all tree positions on the original map"

# ╔═╡ a6db569a-375e-11eb-3943-051a1aefc042
world = Dict()

# ╔═╡ 7ed9ada8-375f-11eb-1194-93353bb23be0
for (index,line) in enumerate(file)
	trees = eachmatch(r"#",line)

	trees_pos = map(x-> x.offset, trees)

	for pos in trees_pos
		world[(index-1,pos-1)] = true
	end
end

# ╔═╡ 21106f36-375f-11eb-387e-7b11364cc055
world

# ╔═╡ 895f12cc-375f-11eb-3b70-a5c0accd57bc
md"extending the map to all positions"

# ╔═╡ 9818c068-375f-11eb-12a6-2b99d75fa222
begin
	tam_column = size(file)[1]
	tam_row = length(file[1])  #the \n ocupy one space

	has_tree(x::Int64,y::Int64) = haskey(world, ( x,y % tam_row))
end

# ╔═╡ d2c55762-3761-11eb-2e27-4b019a678050
tam_column

# ╔═╡ 5b953f4e-3762-11eb-1e0a-fbd3f82403d4
tam_row

# ╔═╡ ea80ee60-375d-11eb-2c3f-fd90c04e37c3
toboggan_pos = [(x,3*x) for x in 1:tam_column-1]

# ╔═╡ 6e1d3084-3769-11eb-1523-9d7fc6c7b1e3
[has_tree(x,y) for (x,y) in toboggan_pos]

# ╔═╡ 72a220ca-3763-11eb-3622-d33c420c6fb5
sum([has_tree(x,y) for (x,y) in toboggan_pos])

# ╔═╡ 24f92922-3768-11eb-018d-b582fb9346ad
md"# Question 2"

# ╔═╡ 3e9f32cc-3768-11eb-0561-7fac93e67081
slope(dx,dy) = sum([has_tree(x,Int64(dy/dx*x)) for x in 0:dx:tam_column-1])

# ╔═╡ 58b6ee84-3768-11eb-0cb0-5f322f2c4cc8
slope(1,3)

# ╔═╡ 7366bd04-3768-11eb-2366-a19e90bb92ca
slope(1,1)*slope(1,3)*slope(1,5)*slope(1,7)*slope(2,1)

# ╔═╡ 99bc5a36-3768-11eb-23e0-2394b9706da1
[slope(1,1),slope(1,3),slope(1,5),slope(1,7),slope(2,1)]

# ╔═╡ Cell order:
# ╠═7cb1f900-375e-11eb-3cc9-cf8a8fb4fc04
# ╟─8f19f0e8-375e-11eb-38fa-b79ed15de216
# ╠═a6db569a-375e-11eb-3943-051a1aefc042
# ╠═7ed9ada8-375f-11eb-1194-93353bb23be0
# ╠═21106f36-375f-11eb-387e-7b11364cc055
# ╠═895f12cc-375f-11eb-3b70-a5c0accd57bc
# ╠═9818c068-375f-11eb-12a6-2b99d75fa222
# ╠═d2c55762-3761-11eb-2e27-4b019a678050
# ╠═5b953f4e-3762-11eb-1e0a-fbd3f82403d4
# ╠═ea80ee60-375d-11eb-2c3f-fd90c04e37c3
# ╠═6e1d3084-3769-11eb-1523-9d7fc6c7b1e3
# ╠═72a220ca-3763-11eb-3622-d33c420c6fb5
# ╠═24f92922-3768-11eb-018d-b582fb9346ad
# ╠═3e9f32cc-3768-11eb-0561-7fac93e67081
# ╠═58b6ee84-3768-11eb-0cb0-5f322f2c4cc8
# ╠═7366bd04-3768-11eb-2366-a19e90bb92ca
# ╠═99bc5a36-3768-11eb-23e0-2394b9706da1


test("Add path to clipper") do
    path = Vector{IntPoint}()

    push!(path, IntPoint(0, 0))
    push!(path, IntPoint(0, 1))

    c = Clip()

    @test add_path!(c, path, PolyTypeSubject, true) == false

    push!(path, IntPoint(1, 1))

    @test add_path!(c, path, PolyTypeSubject, true) == true
end

test("Union") do
  path1 = Vector{IntPoint}()
  push!(path1, IntPoint(0, 0))
  push!(path1, IntPoint(0, 1))
  push!(path1, IntPoint(1, 1))
  push!(path1, IntPoint(1, 0))

  path2 = Vector{IntPoint}()
  push!(path2, IntPoint(1, 0))
  push!(path2, IntPoint(1, 1))
  push!(path2, IntPoint(2, 1))
  push!(path2, IntPoint(2, 0))

  c = Clip()
  add_path!(c, path1, PolyTypeSubject, true)
  add_path!(c, path2, PolyTypeSubject, true)

  result, polys = execute(c, ClipTypeUnion, PolyFillTypeEvenOdd, PolyFillTypeEvenOdd)

  @test result == true
  @test string(polys) == "Array{Clipper.IntPoint,1}[Clipper.IntPoint[[0,0],[2,0],[2,1],[0,1]]]"
end

test("Difference") do
  path1 = Vector{IntPoint}()
  push!(path1, IntPoint(0, 0))
  push!(path1, IntPoint(0, 10))
  push!(path1, IntPoint(10, 10))
  push!(path1, IntPoint(10, 0))

  path2 = Vector{IntPoint}()
  push!(path2, IntPoint(4, 0))
  push!(path2, IntPoint(4, 10))
  push!(path2, IntPoint(6, 10))
  push!(path2, IntPoint(6, 0))

  c = Clip()
  add_path!(c, path1, PolyTypeSubject, true)
  add_path!(c, path2, PolyTypeClip, true)

  result, polys = execute(c, ClipTypeDifference, PolyFillTypeEvenOdd, PolyFillTypeEvenOdd)

  @test result == true
  @test string(polys) == "Array{Clipper.IntPoint,1}[Clipper.IntPoint[[10,10],[6,10],[6,0],[10,0]],Clipper.IntPoint[[0,10],[0,0],[4,0],[4,10]]]"
end

test("GetBounds") do
  path1 = Vector{IntPoint}()
  push!(path1, IntPoint(0, 0))
  push!(path1, IntPoint(0, 1))
  push!(path1, IntPoint(1, 1))
  push!(path1, IntPoint(1, 0))

  path2 = Vector{IntPoint}()
  push!(path2, IntPoint(1, 0))
  push!(path2, IntPoint(1, 1))
  push!(path2, IntPoint(2, 1))
  push!(path2, IntPoint(2, 0))

  c = Clip()
  add_path!(c, path1, PolyTypeSubject, true)
  add_path!(c, path2, PolyTypeSubject, true)

  rect = get_bounds(c)

  @test rect.left == 0
  @test rect.top == 0
  @test rect.right == 2
  @test rect.bottom == 1
end

test("Clear") do
  path1 = Vector{IntPoint}()
  push!(path1, IntPoint(0, 0))
  push!(path1, IntPoint(0, 10))
  push!(path1, IntPoint(10, 10))
  push!(path1, IntPoint(10, 0))

  path2 = Vector{IntPoint}()
  push!(path2, IntPoint(4, 0))
  push!(path2, IntPoint(4, 10))
  push!(path2, IntPoint(6, 10))
  push!(path2, IntPoint(6, 0))

  c = Clip()
  add_path!(c, path1, PolyTypeSubject, true)
  add_path!(c, path2, PolyTypeSubject, true)

  Clipper.clear!(c)

  rect = get_bounds(c)

  @test rect.left == 0
  @test rect.top == 0
  @test rect.right == 0
  @test rect.bottom == 0
end

test("AddPaths") do
  path1 = Vector{IntPoint}()
  push!(path1, IntPoint(0, 0))
  push!(path1, IntPoint(0, 1))
  push!(path1, IntPoint(1, 1))
  push!(path1, IntPoint(1, 0))

  path2 = Vector{IntPoint}()
  push!(path2, IntPoint(1, 0))
  push!(path2, IntPoint(1, 1))
  push!(path2, IntPoint(2, 1))
  push!(path2, IntPoint(2, 0))

  paths = Vector{IntPoint}[path1, path2]

  c = Clip()
  add_paths!(c, paths, PolyTypeSubject, true)

  result, polys = execute(c, ClipTypeUnion, PolyFillTypeEvenOdd, PolyFillTypeEvenOdd)

  @test result == true
  @test string(polys) == "Array{Clipper.IntPoint,1}[Clipper.IntPoint[[0,0],[2,0],[2,1],[0,1]]]"
end

test("Orientation") do
  path1 = Vector{IntPoint}()
  push!(path1, IntPoint(0, 0))
  push!(path1, IntPoint(0, 1))
  push!(path1, IntPoint(1, 1))
  push!(path1, IntPoint(1, 0))

  @test orientation(path1) == false

  path1 = Vector{IntPoint}()
  push!(path1, IntPoint(1, 0))
  push!(path1, IntPoint(1, 1))
  push!(path1, IntPoint(0, 1))
  push!(path1, IntPoint(0, 0))

  @test orientation(path1) == true
end

test("Area") do
  path1 = Vector{IntPoint}()
  push!(path1, IntPoint(0, 0))
  push!(path1, IntPoint(0, 1))
  push!(path1, IntPoint(1, 1))

  @test area(path1) == -0.5
end


# Generate plots based on data from dev/demo_unsteady.jl

# Speed up loading
module ModifiedStokesSolver
include("../src/types.jl")
end

using FileIO
using HDF5
using PyPlot
using Interpolations

include("../src/plottools.jl")
include("../src/streakplot.jl")

function load_geo(datafile)
    info("Loading ", datafile)
    D = load(datafile)
    X = D["X"]
    Y = D["Y"]
    dcurve = D["dcurve"]
    interior = D["interior"]
    return X, Y, dcurve, interior
end

function make_vortex_plot(;Re=100, vorticity=true, stepnum=15000,
                          streaklines=[], streamlines=false,
                          plotcolorbar=false)
    basedir = "data/vortex_Re$(Re)"
    geoname = "$basedir/geo.jld"
    dataname = i -> @sprintf("%s/vstep_%06d.jld", basedir, i)
    datafile = dataname(stepnum)
    X, Y, dcurve, interior = load_geo(geoname)
    vorticity, U1, U2 = [], [], []
    t = 0.0;
    h5open(datafile, "r") do io
        U1 = read(io, "U1")
        U2 = read(io, "U2")
        dU2dx = read(io, "dU2dx")
        dU1dy = read(io, "dU1dy")
        t = read(io, "ti")
        vorticity = dU2dx - dU1dy
    end
    @. vorticity[interior==false] = NaN

    h5open(datafile, "r") do io
        dU2dx = read(io, "dU2dx")
        dU1dy = read(io, "dU1dy")
        t = read(io, "ti")
        vorticity = dU2dx - dU1dy
    end
    @. vorticity[interior==false] = NaN
    clf()
    s = 1
    plot_boundary(dcurve)
    zoom_to_boundary(dcurve)
    axis("off")
    pcolormesh(X, Y, vorticity, cmap="coolwarm", shading="gouraud")
    clim(-3, 3)
    if plotcolorbar
        colorbar(label="Vorticity", shrink=0.8, pad=0.01)
    end
    if streamlines
        # Start points at inflow and in wake
        ystart = 0.1:0.5:3
        ystart = vcat(-reverse(ystart), ystart)
        xstart = zeros(length(ystart)) .- 4.9
        push!(xstart, -2.5)
        push!(ystart, 0.03)
        push!(xstart, -2.5)
        push!(ystart, -0.03)

        streamplot(y=Y[1,:],
                   x=X[:,1],
                   u=U1',
                   v=U2',
                   density=40,
                   start_points=hcat( xstart, ystart-0.1),
                   #integration_direction="both",
                   maxlength=4,
                   color="black"

                   )
    end

    if !isempty(streaklines)
        starty = 0.1:0.5:3
        starty = vcat(reverse(-starty), starty) .- 0.1
        startx = ones(size(starty)) .* (-5)
        streakplot(X=X, Y=Y, interior=interior,
                   startx=startx,
                   starty=starty,
                   dataname=dataname,
                   ilist=streaklines)
    end

    zoom_to_boundary(dcurve)
    tight_layout(0)

    #pngname = "../docs/paper/figs/vortex_Re$(Re).png"
    #savefig(pngname, dpi=200)
    #run(`mogrify -trim $pngname`)
end

close("all")
figsize = (5.5, 3)

# figure(1, figsize=figsize)
# make_vortex_plot(Re=100, stepnum=6250)

# figure(2, figsize=figsize)
# make_vortex_plot(Re=200, stepnum=14500, plotcolorbar=true)

#figure(3, figsize=figsize)
#sn = 14600
#make_vortex_plot(Re=50, stepnum=sn, plotcolorbar=true, streaklines=sn-4000:1:sn)

figure(4, figsize=figsize)
make_vortex_plot(Re=25, stepnum=10, streamlines=true)


using CSV
using LinearAlgebra: dot, I
using CairoMakie
using naumer_Dimensionality_2022

const Σ = [1.0  0.1  0.1;
           0.1  1.0  0.1;
           0.1  0.1  1.0]

const v = [1.0/sqrt(3), 1.0/sqrt(3), 1.0/sqrt(3)]
const σ² = 1

f(x) = abs2(dot(v, Σ * x)) / (dot(x, Σ * x) + σ² * dot(x, x))

# Helper function to compute extra coordinate
function getZ(x,y)
    if abs2(x) + abs2(y) > 1 - 10*eps()
        return NaN
    end
    return sqrt(1 - abs2(x) - abs2(y))
end

function rewardPositive(x,y)
    return f([x,y,getZ(x,y)])
end

function rewardNegative(x,y)
    return f([x,y,-getZ(x,y)])
end

state = zeros(3,100)
state[3,1] = 1

for i in 2:100
    state[:,i] .= actiongradientstep_1DApprox(state[:,i-1], v, Σ, σ², 0.1)
end

rewards = zeros(100)
for i in 1:100
    rewards[i] = f(state[:,i])
end

# Sample Circle
x = -1.1:0.01:1.1
y = -1.1:0.01:1.1

noto_sans = "./resources/NotoSans-Regular.ttf"
noto_sans_bold = "./resources/NotoSans-Bold.ttf"

tickfontsize    = 26
labelfontsize   = 28
basewidth = 2

fig = Figure(font=noto_sans, resolution=(1000,500), figure_padding=40)

ax1 = Axis( fig[1,1],
            xticklabelsize=tickfontsize,
            yticklabelsize=tickfontsize,
            yticklabelpad=2,
            xlabel="x", ylabel = "y",
            xlabelsize=labelfontsize,
            ylabelsize=labelfontsize,
            titlesize=labelfontsize,
            ratio=AxisAspect(1))

hm = heatmap!(ax1,x,y,rewardPositive)
lines!(ax1, state[1,:], state[2,:], color=:black, linewidth=3)

ax2 = Axis( fig[1,3],
            xticklabelsize=tickfontsize,
            yticklabelsize=tickfontsize,
            yticklabelpad=2,
            xlabel="Iteration", ylabel = "Objective",
            xlabelsize=labelfontsize,
            ylabelsize=labelfontsize,
            titlesize=labelfontsize,
            ratio=AxisAspect(1),
            xminorticksvisible = true,
            xminorgridvisible  = true,
            xminorticks = IntervalsBetween(10),
            yminorticksvisible = true,
            yminorgridvisible  = true,
            yminorticks = IntervalsBetween(5))

lines!(ax2,1:100, rewards, color=:black, width=2)
xlims!(ax2,(0,100))

Colorbar(fig[1, 2], hm, ticklabelsize=tickfontsize)

colgap!(fig.layout,2,80)

Label(fig[1,1,TopLeft()], "A", font=noto_sans_bold, textsize = 36, halign=:left, valign=:bottom, padding=(25,0,10,0))
Label(fig[1,3,TopLeft()], "B", font=noto_sans_bold, textsize = 36, halign=:left, valign=:bottom, padding=(25,0,10,0))

save("out/OptDescent.pdf",fig)


module TestDiscreteRv

using QuantEcon
using Base.Test
using FactCheck
using DataStructures

# set up
srand(42)
n = 10
x = rand(n)
x ./= sum(x)
drv = DiscreteRV(x)

facts("Testing discrete_rv.jl") do
    # test Q sums to 1
    @fact drv.Q[end] --> roughly(1.0)

    # test lln
    draws = draw(drv, 100000)
    c = counter(draws)
    counts = Array(Float64, n)
    for i=1:n
        counts[i] = c[i]
    end
    counts ./= sum(counts)

    @fact Base.maxabs(counts - drv.q) --> roughly(0.0; atol=1e-2)

end  # facts
end  # module


# # Continuous SIMP example
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`csimp.ipynb`](@__NBVIEWER_ROOT_URL__/examples/csimp.ipynb)
#-
# ## Commented Program
#
# What follows is a program spliced with comments.
#md # The full program, without comments, can be found in the next [section](@ref csimp-plain-program).

using TopOpt

# ### Define the problem
E = 1.0 # Young’s modulus
v = 0.3 # Poisson’s ratio
f = 1.0 # downward force

problems = Any[
            PointLoadCantilever(Val{:Linear}, (60, 20, 20), (1.0, 1.0, 1.0), E, v, f),
            PointLoadCantilever(Val{:Linear}, (160, 40), (1.0, 1.0), E, v, f),
            HalfMBB(Val{:Linear}, (60, 20), (1.0, 1.0), E, v, f),
            LBeam(Val{:Linear}, Float64, force = f),
            TieBeam(Val{:Quadratic}, Float64)
            ]
problem_names = [
    "3d cantilever beam",
    "cantilever beam",
    "half MBB beam",
    "L-beam",
    "tie-beam",
    ]

i = 2
println(problem_names[i])
problem = problems[i]

# ### Parameter settings
V = 0.5 # volume fraction
xmin = 0.001 # minimum density
rmin = 3.0
steps = 40 # maximum number of penalty steps, delta_p0 = 0.1
reuse = true # adaptive penalty flag

convcriteria = Nonconvex.GenericCriteria()
penalty = TopOpt.PowerPenalty(1.0)
pcont = Continuation(penalty, steps = steps, xmin = xmin, pmax = 5.0)

# NOTE: non-convexity + computational error lead to different solutions that satisfy the KKT tolerance
mma_options = options = Nonconvex.MMAOptions(maxiter=1000)
maxtol = 0.01 # maximum tolerance
mintol = 0.0001 # minimum tolerance
b = log(mintol / maxtol) / steps
a = maxtol / exp(b)
mma_options_gen = TopOpt.MMAOptionsGen(
    steps = steps,
    initial_options = mma_options,
    ftol_gen = ExponentialContinuation(a, b, 0.0, steps + 1, mintol),
)
csimp_options = TopOpt.CSIMPOptions(
    steps = steps,
    options_gen = mma_options_gen,
    p_gen = pcont,
    reuse = reuse,
)

# ### Define a finite element solver
solver = FEASolver(Displacement, Direct, problem, xmin = xmin, penalty = penalty)

# ### Define compliance objective
comp = Compliance(problem, solver)
filter = if problem isa TopOptProblems.TieBeam
    identity
else
    DensityFilter(solver, rmin = rmin)
end
obj = x -> comp(filter(x))

# ### Define volume constraint
volfrac = TopOpt.Volume(problem, solver)
constr = x -> volfrac(filter(x)) - V

# ### Define subproblem optimizer
x0 = fill(V, length(solver.vars))
optimizer = Optimizer(
    obj, constr, x0, Nonconvex.MMA87(),
    options = mma_options, convcriteria = convcriteria,
)

# ### Define continuation SIMP optimizer
simp = SIMP(optimizer, solver, penalty.p)
cont_simp = ContinuationSIMP(simp, steps, csimp_options)

# ### Solve
result = cont_simp(x0)

@show result.convstate
@show optimizer.workspace.iter
@show result.objval

# ### (Optional) Visualize the result using Makie.jl
# Need to run `using Pkg; Pkg.add(Makie)` first
# ```julia
# using TopOpt.TopOptProblems.Visualization: visualize
# fig = visualize(problem; topology = result.topology,
#     problem; topology = result.topology, default_exagg_scale = 0.07,
#     scale_range = 10.0, vector_linewidth = 3, vector_arrowsize = 0.5,
# )
# Makie.display(fig)
# ```

#md # ## [Plain Program](@id csimp-plain-program)
#md #
#md # Below follows a version of the program without any comments.
#md # The file is also available here: [csimp.jl](csimp.jl)
#md #
#md # ```julia
#md # @__CODE__
#md # ```


export ConvertPanelToDiffCS_2

"""
    ConvertPanelToDiffCS_2(formula,data,idsymbol,tsymbol)

Takes a panel data set in long form (one row represents measurements
on a cross-section unit at a given time) and converts it into a
cross-sectional data set with one row representing differences.

It outputs two matrices:

`Y`, an `n` by 2 matrix, where the first column is the depdent variables
   in the first period, and `y2` is the dependent variables in the second.
`X`, an `n` by `K` matrix, with first differences in all explanatory
   variables.
"""
function ConvertPanelToDiffCS_2(formula,data,idsymbol,tsymbol)

    # Takes a data frame with panel data in two time periods, and returns the data
    #   in a form that is ready for this package's optimization procedure.

    # First, sort the data
    sort!(data, cols = [idsymbol, tsymbol])

    # Second, remove the incomplete records.
    dfcount = by(data,idsymbol, df -> length(df[tsymbol]))
    complete_id = dfcount[dfcount[:x1].==2,:]
    data = join(data,complete_id,on = idsymbol, kind = :inner)

    # Third, extract the model matrix ("X"s)
    MF = ModelFrame(formula,data)
    X = ModelMatrix(MF).m[:,2:end]

    # Take differences, exploiting the fact that T=2 and the data was previously balanced.
    n = size(X)[1] #new effective size after dropping the incomplete ones
    dX = X[filter(iseven,1:n),:] - X[filter(isodd,1:n),:]

    # Fourth, extract the response variable.
    y = model_response(MF)
    y1 = y[filter(isodd,1:n)]
    y2 = y[filter(iseven,1:n)]
    ybar = 1.0*((y1 + y2) .== 1)

    return [y1 y2], dX

end