type PossibleExtremaReached
    x0::Real
end

type Dithering
    x0::Real
end

"""

SOLVE button from the HP-34C
follows roughly algorithm described http://www.hpl.hp.com/hpjournal/pdfs/IssuePDFs/1979-12.pdf
though some modifications were made.

Goal is to return a value `x` with either:
* `f(x) == 0.0` or
* `f(prevfloat(x)) * f(nextfloat(x)) < 0`.
if a bracket is found that can be done.

Otherwise, values with `f(x) <= ftol` are given status of answers
and others are marked as a possible extrema.

If the algorithm gets stuck, a Dithering error is thrown.

While this is slower than fzero with order 2, 5, 8, or 16 it is
more robust to the initial condition.

"""
function SOLVE{T<:AbstractFloat}(f, x0::T;
               ftol::Real=10*eps(one(x0)),
               xtol::Real=zero(x0),
               xtolrel::Real=zero(x0),
               maxeval::Int=30,
               verbose::Bool=false)

    fx0 = f(x0)
    abs(fx0) <= ftol && return(x0)

    try
        ## get initial guess. Use Steffensen guess if reasonable size
        stepsize = max(1/100, min(abs(fx0), abs(x0/100)))
        if fx0 < 0
            g(x) = -f(x)
            secant_method_no_bracket(g, x0 - stepsize, x0, ftol=ftol, xtol=xtol, xtolrel=xtolrel, maxeval=maxeval, verbose=verbose)
        else
            secant_method_no_bracket(f, x0 - stepsize, x0, ftol=ftol, xtol=xtol, xtolrel=xtolrel, maxeval=maxeval, verbose=verbose)
        end
    catch ex
        if isa(ex, StateConverged)
            return convert(T,ex.x0)
        elseif isa(ex, PossibleExtremaReached)
            guess = convert(T, ex.x0)
            if sign(f(nextfloat(guess))) * sign(f(prevfloat(guess))) <= 0
                return guess
            end
            # heuristic for case where it thinks it is at a minimum
            if abs(f(guess)) < ftol
                return guess
            else
                rethrow(ex)
            end
        else
            rethrow(ex)
        end
    end
    error("huh? Shouldn't get here")
end

"""

Solve f(x) = 0 if we have a bracket [a,b] and starting point a < c < b

Use Algorithm 4.2 of Alefeld, Potra, Shi unless it dithers, in which case the root is found by bisection.

"""
function have_bracket(f, a, b, c=(0.5)*(a+b);
                      xtol::Real=zero(c), xtolrel::Real=zero(c), maxeval=10, verbose=false)
    ## try a42a unless it fails, then go to failsafe
    verbose && println("Have a bracket, switching to bracketing method")
    fa, fb = f(a), f(b)
    fa == 0.0 && return(a)
    fb == 0.0 && return(b)
    isinf(fa) && throw(ConvergenceFailed("f(a) must be finite"))
    isinf(fb) && throw(ConvergenceFailed("f(b) must be finite"))

    if sign(fa) * sign(fb) > 0
        throw(ConvergenceFailed("Interval [$a, $b] is not a bracket"))
    end

    a,b = sort([a,b])
    if !(a <= c <= b)
        throw(ConvergenceFailed("Value $c is not in interval ($a, $b)"))
    end

    try
        x0 = a42a(f, a, b, c, xtol=xtol, maxeval=maxeval, verbose=verbose)
    catch e
        if isa(e, StateConverged)
            return e.x0
        else
            verbose && println("Dithering, switching to bisection method")
            return find_zero(f, a, b, verbose=verbose)
        end
    end
end

## assume f(b) > 0
## steffensen of secant line method, but we don't assume a bracket
## we "bend" large trajectories
function secant_method_no_bracket{T<:AbstractFloat}(f, a::T, b::T;
                                  ftol=10*eps(one(float(a))), xtol::Real=10*eps(one(float(a))),
                                  xtolrel::Real=10*eps(one(float(a))),
                                  maxeval::Int=100, verbose::Bool=false)

    alpha, beta = a, b ## beta keeps smallest values
    falpha, fbeta = f(alpha), f(beta)

    ## checks -- have a bracket
    if sign(falpha) * sign(fbeta) < 0
        a = a42(f, alpha, beta, xtol=xtol, maxeval=maxeval, verbose=verbose)
        throw(StateConverged(a))
    end

    ## keep fbeta smallest
    if falpha < fbeta
        alpha, beta, falpha, fbeta = beta, alpha, fbeta, falpha
    end

    gamma = Inf

    ctr, quad_ctr = 0, 0

    while fbeta > 0
        if ctr > maxeval
            if abs(f(beta)) < ftol
                throw(StateConverged(beta))
            else
                ## try a higher order method? May get us there
                o = thukral8_itr(f, beta, xtol=float(xtol), xtolrel=float(xtolrel), ftol = float(ftol))
                for (x,fx) in o
                    nothing
                end
                xn = o.xn[end]
                if o.state == :converged && !isnan(xn) && !isinf(xn)
                    throw(StateConverged(xn))
                else
                    throw(Dithering(beta))
                end
            end
        else
            (ctr = ctr + 1)
        end
        ## use steffensen or secant line depending on fbeta's size
        if abs(fbeta) < abs(beta - alpha)
            stepsize = fbeta * fbeta / (f(beta + fbeta) - fbeta)  # steffensen
        else
            stepsize =  fbeta * (beta - alpha)  / (fbeta - falpha) # secant line
        end
        gamma = beta - stepsize

        ## bend in some cases. XXX What is best bending scheme? XXX
        if abs(gamma-beta) >= 5e2 * abs(beta - alpha)
            stepsize = cbrt( (beta - alpha) / (fbeta - falpha) * fbeta)
#            gamma = beta + (1/4) * (beta - alpha) / (fbeta - falpha) * fbeta
            gamma = beta - stepsize
        end

        fgamma = f(gamma)

        # try
        #     fgamma = f(gamma) ## what to do with domain error?
        # catch e
        #     if isa(e, DomainError)
        #         error("domain error -- fix me")
        #     end
        # end

        if abs(fgamma) <= ftol
            throw(StateConverged(gamma))
        end
        if fgamma < 0
            x0 = have_bracket(f, gamma, beta, maxeval=20, verbose=verbose)
            throw(StateConverged(x0))
        end

        ## increasing value, try quadratic
        if fgamma >= fbeta
            quad_ctr = quad_ctr + 1
            if quad_ctr > 5*10
                abs(fbeta) <  ftol ? throw(StateConverged(beta)) : throw(Dithering(beta))
            end
            ## secant line stopped decreasing
            if alpha == gamma || gamma == beta
                throw(PossibleExtremaReached(gamma)) ## poorly titled error
            end

            ## quadratic interpolation to get gamma
            denom = (beta - alpha) * (fbeta - fgamma)  - (beta - fgamma) * (fbeta - falpha)
            if denom == 0.0
                throw(PossibleExtremaReached(alpha))
            end
            gamma = beta -  ((beta - alpha)^2 * (fbeta - fgamma) - (beta - gamma)^2 * (fbeta - falpha))/denom/2

            fgamma = f(gamma)

            if fgamma == 0
                throw(StateConverged(gamma))
            end
            if fgamma < 0
                x0 = have_bracket(f, gamma, beta, maxeval=20, verbose=verbose)
                throw(StateConverged(x0))
            end
        else
            quad_ctr = 0        # reset
        end

        if fgamma > fbeta
            alpha, beta, falpha, fbeta = gamma, beta, fgamma, fbeta
        elseif fgamma < fbeta
            alpha, beta, falpha, fbeta = beta, gamma, fbeta, fgamma
        else
            throw(PossibleExtremaReached(beta))
        end
        verbose && println(@sprintf("x_%s:%18.15f\t f(x_%s):%18.15f", ctr, beta, ctr, f(beta)))
    end
    if abs(f(beta)) < ftol
        throw(StateConverged(beta))
    else
        throw(ConvergenceFailed("More than $maxeval iterations before convergence"))
    end
end


using PyPlot
using Gen

import Random
Random.seed!(1)
include("poisson_process.jl")
include("load_data.jl")

# Example from Section 4 of Reversible jump Markov chain Monte Carlo
# computation and Bayesian model determination

#########
# model #
#########

const K = :k
const EVENTS = :events
const UNSORTED_CHANGEPT = :unsorted_changept
const UNSORTED_RATE = :unsorted_rate

function sort_change_pts(unsorted_change_pts::Vector)
    k = length(unsorted_change_pts)
    sorted_to_unsorted = sortperm(unsorted_change_pts)
    unsorted_to_sorted = Vector{Int}(undef, k)
    for i in 1:k
        unsorted_to_sorted[sorted_to_unsorted[i]] = i
    end
    sorted_change_pts = unsorted_change_pts[sorted_to_unsorted]
    return (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted)
end

function test_sort()
    unsorted_change_pts = [0.3, 0.1, 0.2, 0.0, -1.0]
    (change_pts, unsorted_to_sorted, sorted_to_unsorted) = sort_change_pts(unsorted_change_pts)
    @assert change_pts == [-1.0, 0.0, 0.1, 0.2, 0.3]
    @assert sorted_to_unsorted == [5, 4, 2, 3, 1]
    @assert unsorted_to_sorted == [5, 3, 4, 2, 1]
end

test_sort()

@gen function model(T::Float64)

    # prior on number of change points
    k = @trace(poisson(3.), K)

    # prior distribution on change points
    (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted) = sort_change_pts(
        Float64[({(UNSORTED_CHANGEPT, i)} ~ uniform(0, T)) for i in 1:k])

    # k + 1 rate values
    # unsorted_rates[i] for i > 0 is the rate of the segment immediately before
    # the i'th unsorted changepoint. unsorted_rates[k+1] is the rate of the
    # last segment
    alpha = 1.0; beta = 200.0
    unsorted_rates = Float64[@trace(Gen.gamma(alpha, 1. / beta), (UNSORTED_RATE, i)) for i=1:k+1]
    sorted_rates = Vector{Float64}(undef, k+1)
    sorted_rates[1:k] = unsorted_rates[sorted_to_unsorted]
    sorted_rates[k+1] = unsorted_rates[k+1]

    # poisson process
    bounds = vcat([0.], sorted_change_pts, [T])
    @trace(piecewise_poisson_process(bounds, sorted_rates), EVENTS)

    return (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted, sorted_rates)
end

get_t_min(trace) = 0.0
get_t_max(trace) = get_args(trace)[1]
get_num_change_pts(trace) = trace[:k]
get_sorted_change_pts(trace) = get_retval(trace)[1]
get_sorted_rates(trace) = get_retval(trace)[4]

function render(trace; ymax=0.02)
    k = get_num_change_pts(trace)
    bounds = vcat([get_t_min(trace)], get_sorted_change_pts(trace), [get_t_max(trace)])
    rates = get_sorted_rates(trace)
    for i=1:length(rates)
        lower = bounds[i]
        upper = bounds[i+1]
        rate = rates[i]
        plot([lower, upper], [rate, rate], color="black", linewidth=2)
    end
    points = trace[EVENTS]
    scatter(points, -rand(length(points)) * (ymax/5.), color="black", s=5)
    ax = gca()
    xlim = [get_t_min(trace), get_t_max(trace)]
    plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
end

function show_prior_samples()
    figure(figsize=(16,16))
    T = 40000.
    for i=1:16
        println("simulating $i")
        subplot(4, 4, i)
        (trace, ) = generate(model, (T,), EmptyChoiceMap())
        render(trace; ymax=0.015)
    end
    tight_layout(pad=0)
    savefig("prior_samples.pdf")
end

#############
# rate move #
#############

@gen function rate_proposal(trace)

    # pick a random segment whose rate to change
    i = @trace(uniform_discrete(1, trace[K]+1), :i)

    # propose new value for the rate
    #cur_rate = trace[(UNSORTED_RATE, i)]
    #I@trace(uniform_continuous(cur_rate/2., cur_rate*2.), :new_rate)
    new_rate_scaled ~ uniform(0, 1)

    nothing
end

# it is an involution because it:
# - maintains i = fwd_choices[:i] constant
# - swaps choices[(RATE, i)] with fwd_choices[:new_rate]

@transform rate_involution (model_in, aux_in) to (model_out, aux_out) begin
    i = @read(aux_in[:i], :discrete)
    @write(aux_out[:i], i, :discrete)
    new_rate_scaled = @read(aux_in[:new_rate_scaled], :continuous)
    cur_rate = @read(model_in[(UNSORTED_RATE, i)], :continuous)
    lower_bound = cur_rate / 2.0
    upper_bound = cur_rate * 2.0
    new_rate = lower_bound + new_rate_scaled * (upper_bound - lower_bound)
    @write(model_out[(UNSORTED_RATE, i)], new_rate, :continuous)
    prev_rate_scaled = (cur_rate - (new_rate / 2.0)) / (new_rate * 2.0 - new_rate / 2.0)
    @write(aux_out[:new_rate_scaled], prev_rate_scaled, :continuous)
end

is_involution!(rate_involution)

rate_move(trace) = metropolis_hastings(trace, rate_proposal, (), rate_involution, check=false)

#################
# position move #
#################

function get_sorted_idx(trace, unsorted_idx::Int)
    (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted, _) = get_retval(trace)
    return unsorted_to_sorted[unsorted_idx]
end

function get_lower_bound(trace, unsorted_idx::Int)
    if get_sorted_idx(trace, unsorted_idx) == 1
        return get_t_min(trace)
    else
        return sorted_change_pts[sorted_idx - 1]
    end
end

function get_upper_bound(trace, unsorted_idx::Int)
    if get_sorted_idx(trace, unsorted_idx) == get_num_change_pts(trace)
        return get_t_max(trace)
    else
        return sorted_change_pts[sorted_idx + 1]
    end
end

# TODO, note: the density is not nonzero everywhere with respect to the reference measure...
# but it could easily be made so by sampling from uniform(0, 1) the location between lower and upper bound...

@gen function position_proposal(trace)
    k = trace[K]
    @assert k > 0

    # pick a random unsorted changepoint to change
    i = @trace(uniform_discrete(1, k), :i)

    sorted_idx = get_sorted_idx(trace, i)
    sorted_to_unsorted = get_retval(trace)[3]
    #lower = get_lower_bound(trace, i)
    #upper = get_upper_bound(trace, i)
    @trace(uniform_continuous(0, 1), :new_changept_scaled)
    #@trace(uniform_continuous(lower, upper), :new_changept)

    return (sorted_idx, sorted_to_unsorted)
end

# it is an involution because it:
# - maintains i = fwd_choices[:i] constant
# - swaps choices[(CHANGEPT, i)] with fwd_choices[:new_changept]

@transform position_involution (model_in, aux_in) to (model_out, aux_out) begin
    (sorted_idx, sorted_to_unsorted) = @read(aux_in[], :discrete)
    k = @read(model_in[:k], :discrete)
    i = @read(aux_in[:i], :discrete)
    @write(aux_out[:i], i, :discrete)
    new_changept_scaled = @read(aux_in[:new_changept_scaled], :continuous)
    lower_bound = sorted_idx == 1 ? 0.0 : @read(model_in[(UNSORTED_CHANGEPT, sorted_to_unsorted[sorted_idx-1])], :continuous)
    upper_bound = sorted_idx == k ? T : @read(model_in[(UNSORTED_CHANGEPT, sorted_to_unsorted[sorted_idx+1])], :continuous)
    new_changept = lower_bound + new_changept_scaled * (upper_bound - lower_bound)
    @write(model_out[(UNSORTED_CHANGEPT, i)], new_changept, :continuous)

    prev_changept = @read(model_in[(UNSORTED_CHANGEPT, i)], :continuous)
    prev_changept_scaled = (prev_changept - lower_bound) / (upper_bound - lower_bound)
    @write(aux_out[:new_changept_scaled], prev_changept_scaled, :continuous)
end

is_involution!(position_involution)

position_move(trace) = metropolis_hastings(trace, position_proposal, (), position_involution, check=false)

######################
# birth / death move #
######################

const CHOSEN = :chosen
const IS_BIRTH = :is_birth
const NEW_CHANGEPT = :new_changept
const U = :u

@gen function birth_death_proposal(trace)
    #T = get_t_max(trace)
    k = trace[K]
    # if k = 0, then always do a birth move
    # if k > 0, then randomly choose a birth or death move
    isbirth = (k == 0) ? true : @trace(bernoulli(0.5), IS_BIRTH)
    if isbirth
        {NEW_CHANGEPT} ~ uniform_continuous(get_t_min(trace), get_t_max(trace))
        {U} ~ uniform_continuous(0.0, 1.0)
    end
    (_, unsorted_to_sorted, sorted_to_unsorted, _) = get_retval(trace)
    return (unsorted_to_sorted, sorted_to_unsorted) # NOTE: based on testing predicates on real values
end

function new_rates(cur_rate, u, cur_cp, prev_cp, next_cp)
    d_prev = cur_cp - prev_cp
    d_next = next_cp - cur_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_cur_rate = log(cur_rate)
    log_ratio = log(1 - u) - log(u)
    prev_rate = exp(log_cur_rate - (d_next / d_total) * log_ratio)
    next_rate = exp(log_cur_rate + (d_prev / d_total) * log_ratio)
    @assert prev_rate > 0.
    @assert next_rate > 0.
    (prev_rate, next_rate)
end

function new_rates_inverse(prev_rate, next_rate, cur_cp, prev_cp, next_cp)
    d_prev = cur_cp - prev_cp
    d_next = next_cp - cur_cp
    @assert d_prev > 0
    @assert d_next > 0
    d_total = d_prev + d_next
    log_prev_rate = log(prev_rate)
    log_next_rate = log(next_rate)
    cur_rate = exp((d_prev / d_total) * log_prev_rate + (d_next / d_total) * log_next_rate)
    u = prev_rate / (prev_rate + next_rate)
    @assert cur_rate > 0.
    (cur_rate, u)
end

# it is an involution because:
# - it switches back and forth between birth move and death move
# - it maintains fwd_choices[CHOSEN] constant (applying it twice will first insert a new
#   changepoint and then remove that same changepoint)
# - new_rates, curried on cp_new, cp_prev, and cp_next, is the inverse of new_rates_inverse.

@transform birth_death_involution (model_in, aux_in) to (model_out, aux_out) begin
    T = get_args(model_in)[1]
    (unsorted_to_sorted, sorted_to_unsorted) = @read(aux_in[], :discrete)

    # current number of changepoints
    k = @read(model_in[K], :discrete)

    # if k == 0, then we can only do a birth move
    isbirth = (k == 0) || @read(aux_in[IS_BIRTH], :discrete)

    # if we are a birth move, the inverse is a death move
    if k > 1 || isbirth
        @write(aux_out[IS_BIRTH], !isbirth, :discrete)
    end

    if isbirth
        @tcall(birth(k, unsorted_to_sorted, sorted_to_unsorted))
    else
        @tcall(death(k, unsorted_to_sorted, sorted_to_unsorted))
    end
end

@transform birth(k::Int, unsorted_to_sorted, sorted_to_unsorted) (model_in, aux_in) to (model_out, aux_out) begin
    k = @read(model_in[K], :discrete)
    @write(model_out[K], k+1, :discrete)
    cp_new = @read(aux_in[NEW_CHANGEPT], :continuous)

    # set new changepoint
    @copy(aux_in[NEW_CHANGEPT], model_out[(UNSORTED_CHANGEPT, k+1)])

    # re-sort the changepoints, with the new one added
    unsorted_change_pts = [@read(model_in[(UNSORTED_CHANGEPT, i)], :continuous) for i in 1:k]
    push!(unsorted_change_pts, cp_new)
    (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted) = sort_change_pts(unsorted_change_pts)
    sorted_idx = unsorted_to_sorted[k+1]

    # the unsorted index of the next rate choice
    prev_next_unsorted_rate_idx = (sorted_idx == k+1) ? k+1 : sorted_to_unsorted[sorted_idx+1]
    new_next_unsorted_rate_idx = (sorted_idx == k+1) ? k+2 : sorted_to_unsorted[sorted_idx+1]

    # the previous and next changepoints
    cp_prev = (sorted_idx == 1) ? 0.0 : sorted_change_pts[sorted_idx-1]
    cp_next = (sorted_idx == k+1) ? T : sorted_change_pts[sorted_idx+1]

    # compute new rates
    h_cur = @read(model_in[(UNSORTED_RATE, prev_next_unsorted_rate_idx)], :continuous)
    u = @read(aux_in[U], :continuous)
    (h_prev, h_next) = new_rates(h_cur, u, cp_new, cp_prev, cp_next)

    # set new rates
    @write(model_out[(UNSORTED_RATE, k+1)], h_prev, :continuous)
    @write(model_out[(UNSORTED_RATE, new_next_unsorted_rate_idx)], h_next, :continuous)

    # shift the final rate up, only if it is not where we inserted it...
    if sorted_idx != k+1
        @copy(model_in[(UNSORTED_RATE, k+1)], model_out[(UNSORTED_RATE, k+2)])
    end
end

# TODO rename the last rate to 'last_rate' so we don't need to shift it around...

@transform death(k::Int, unsorted_to_sorted, sorted_to_unsorted) (model_in, aux_in) to (model_out, aux_out) begin
    k = @read(model_in[K], :discrete)
    @write(model_out[K], k-1, :discrete)
    cp_deleted = @read(model_in[(UNSORTED_CHANGEPT, k)], :continuous) # NOTE: will not be copied...

    # set proposed changepoint
    @copy(model_in[(UNSORTED_CHANGEPT, k)], aux_out[NEW_CHANGEPT]) # but gets copied here..

    # re-sort the changepoints (we could use the result of the sorting from the model trace, if we wanted)
    unsorted_change_pts = [@read(model_in[(UNSORTED_CHANGEPT, i)], :continuous) for i in 1:k] # all get copied..
    (sorted_change_pts, unsorted_to_sorted, sorted_to_unsorted) = sort_change_pts(unsorted_change_pts)
    sorted_idx = unsorted_to_sorted[k]

    # the unsorted index of the next rate choice
    prev_next_unsorted_rate_idx = (sorted_idx == k) ? k+1 : sorted_to_unsorted[sorted_idx+1]
    new_next_unsorted_rate_idx = (sorted_idx == k) ? k : sorted_to_unsorted[sorted_idx+1]
    @assert prev_next_unsorted_rate_idx != k

    # the previous and next changepoints
    cp_prev = (sorted_idx == 1) ? 0.0 : sorted_change_pts[sorted_idx-1]
    cp_next = (sorted_idx == k) ? T : sorted_change_pts[sorted_idx+1]

    # compute cur rate and u
    h_prev = @read(model_in[(UNSORTED_RATE, k)], :continuous)
    h_next = @read(model_in[(UNSORTED_RATE, prev_next_unsorted_rate_idx)], :continuous)
    (h_cur, u) = new_rates_inverse(h_prev, h_next, cp_deleted, cp_prev, cp_next)
    @write(aux_out[U], u, :continuous)

    # set cur rate
    @write(model_out[(UNSORTED_RATE, new_next_unsorted_rate_idx)], h_cur, :continuous)

    # shift the final rate, only if we are not deleting the last one..
    if sorted_idx != k
        @copy(model_in[(UNSORTED_RATE, k+1)], model_out[(UNSORTED_RATE, k)])
    end
end

is_involution!(birth_death_involution)

@gen function permutation_proposal(trace)
    k = get_num_change_pts(trace)
    to_swap ~ uniform_discrete(1, k-1)
end

@transform permutation_involution (model_in, aux_in) to (model_out, aux_out) begin
    k = @read(model_in[:k], :discrete)
    to_swap_with_k = @read(aux_in[:to_swap], :discrete)
    @copy(aux_in[:to_swap], aux_out[:to_swap])
    @copy(model_in[(UNSORTED_RATE, k)], model_out[(UNSORTED_RATE, to_swap_with_k)])
    @copy(model_in[(UNSORTED_RATE, to_swap_with_k)], model_out[(UNSORTED_RATE, k)])
    @copy(model_in[(UNSORTED_CHANGEPT, k)], model_out[(UNSORTED_CHANGEPT, to_swap_with_k)])
    @copy(model_in[(UNSORTED_CHANGEPT, to_swap_with_k)], model_out[(UNSORTED_CHANGEPT, k)])
end

is_involution!(permutation_involution)

function birth_death_move(trace)
    if trace[:k] > 1
        trace, acc = metropolis_hastings(trace, permutation_proposal, (), permutation_involution, check=false)
        @assert acc
    end
    return metropolis_hastings(trace, birth_death_proposal, (), birth_death_involution, check=false)
end

function mcmc_step(trace)
    (trace, _) = rate_move(trace)
    if trace[K] > 0
        (trace, _) = position_move(trace)
    end
    (trace, _) = birth_death_move(trace)
    trace
end

function simple_mcmc_step(trace)
    (trace, _) = rate_move(trace)
    if trace[K] > 0
        (trace, _) = position_move(trace)
    end
    (trace, _) = metropolis_hastings(trace, k_selection)
    trace
end

function do_mcmc(T, num_steps::Int)
    (trace, _) = generate(model, (T,), observations)
    for iter=1:num_steps
        k = trace[K]
        if iter % 1000 == 0
            println("iter $iter of $num_steps, k: $k")
        end
        trace = mcmc_step(trace)
    end
    trace
end

const k_selection = select(K)

function do_simple_mcmc(T, num_steps::Int)
    (trace, _) = generate(model, (T,), observations)
    for iter=1:num_steps
        k = trace[K]
        if iter % 1000 == 0
            println("iter $iter of $num_steps, k: $k")
        end
        trace = simple_mcmc_step(trace)
    end
    trace
end

########################
# inference experiment #
########################

import Random
Random.seed!(1)

const points = load_data_set()
const T = maximum(points)
const observations = choicemap()
observations[EVENTS] = points

function show_posterior_samples()
    figure(figsize=(16,16))
    for i=1:16
        println("replicate $i")
        subplot(4, 4, i)
        trace = do_mcmc(T, 200)
        render(trace; ymax=0.015)
    end
    tight_layout(pad=0)
    savefig("posterior_samples.pdf")

    figure(figsize=(16,16))
    for i=1:16
        println("replicate $i")
        subplot(4, 4, i)
        trace = do_simple_mcmc(T, 200)
        render(trace; ymax=0.015)
    end
    tight_layout(pad=0)
    savefig("posterior_samples_simple.pdf")
end

function get_rate_vector(trace, test_points)
    k = trace[K]
    cps = get_sorted_change_pts(trace)
    hs = get_sorted_rates(trace)
    rate = Vector{Float64}()
    cur_h_idx = 1
    cur_h = hs[cur_h_idx]
    next_cp_idx = 1
    upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
    for x in test_points
        while x > upper
            next_cp_idx += 1
            upper = (next_cp_idx == k + 1) ? T : cps[next_cp_idx]
            cur_h_idx += 1
            cur_h = hs[cur_h_idx]
        end
        push!(rate, cur_h)
    end
    rate
end

# compute posterior mean rate curve

function plot_posterior_mean_rate()
    test_points = collect(1.0:10.0:T)
    rates = Vector{Vector{Float64}}()
    num_samples = 0
    num_steps = 8000
    for reps=1:20
        (trace, _) = generate(model, (T,), observations)
        for iter=1:num_steps
            if iter % 1000 == 0
                println("iter $iter of $num_steps, k: $(trace[K])")
            end
            trace = mcmc_step(trace)
            if iter > 4000
                num_samples += 1
                rate_vector = get_rate_vector(trace, test_points)
                @assert length(rate_vector) == length(test_points)
                push!(rates, rate_vector)
            end
        end
    end
    posterior_mean_rate = zeros(length(test_points))
    for rate in rates
        posterior_mean_rate += rate / Float64(num_samples)
    end
    ymax = 0.010
    figure()
    plot(test_points, posterior_mean_rate, color="black")
    scatter(points, -rand(length(points)) * (ymax/6.), color="black", s=5)
    ax = gca()
    xlim = [0., T]
    plot(xlim, [0., 0.], "--")
    ax[:set_xlim](xlim)
    ax[:set_ylim](-ymax/5., ymax)
    savefig("posterior_mean_rate.pdf")
end

function plot_trace_plot()
    figure(figsize=(8, 4))

    # reversible jump
    (trace, _) = generate(model, (T,), observations)
    rate1 = Float64[]
    num_clusters_vec = Int[]
    burn_in = 0
    for iter=1:burn_in + 1000
        trace = mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, trace[K])
        end
    end
    subplot(2, 1, 1)
    plot(num_clusters_vec, "b")

    # simple MCMC
    (trace, _) = generate(model, (T,), observations)
    rate1 = Float64[]
    num_clusters_vec = Int[]
    burn_in = 0
    for iter=1:burn_in + 1000
        trace = simple_mcmc_step(trace)
        if iter > burn_in
            push!(num_clusters_vec, trace[K])
        end
    end
    subplot(2, 1, 2)
    plot(num_clusters_vec, "b")

    savefig("trace_plot.pdf")
end

println("showing prior samples...")
show_prior_samples()

println("showing posterior samples...")
show_posterior_samples()

println("estimating posterior mean rate...")
plot_posterior_mean_rate()

println("making trace plot...")
plot_trace_plot()


module SimulatedAnnealing

using Parameters
using Statistics

export
    AnnealingOptimization, AnnealingState,
    OVGCriterion, SSVCriterion,
    ConstantDecrementRule, AVLDecrementRule,
    simulated_annealing

"""
    AnnealingOptimization

Object representing an optimization done using a Simulated Annealing
algorithm.

Implement single stage homogeneous Simulated Annealing as (very well) described
in
Varanelli, 1996, *On the acceleration of simulated annealing*, Chapter 2.

Notation and naming convention is consistent with that work.

A configuration of the system being optimized is represented by the parametric
type `C`. The configurations to optimize can be of any type.

The object is iterable to allow access to the internal state during the
optimization process.

Parameters
==========
energy: callable(::C)::T
    Callable returning the energy of a configuration.
propose_candiate: callable(::C)::Tuple{::C, ::T}
    Callable returning a new configuration together based on an existing one
    togetehr with the energy difference between the two.
stop_criterion: callable(::AnnealingState)::Bool
    Callable returning whether the annealing must stop.
    Default is Otten-van Ginneken adaptive stop criterion.
decrement_rule: callable(::AneealingState)::T
    Callable returning the new temperature to use at the end of a Markov step.
    Default is Aarts and van Laarhoven decrement rule.
initial_temperature: Real
initial_configuration: C
neighborhood_size: integer
    Number of states that can be accessed from a given one.
"""
struct AnnealingOptimization{T, C, EFUNC, PFUNC, SFUNC, TFUNC}
    energy::EFUNC
    propose_candidate::PFUNC
    stop_criterion::SFUNC
    decrement_rule::TFUNC
    initial_temperature::T
    initial_configuration::C
    neighborhood_size::Int
end

function AnnealingOptimization(
        energy,
        propose_candidate,
        initial_configuration,
        initial_temperature::Real,
        neighborhood_size::Integer)

    return AnnealingOptimization(
        energy,
        propose_candidate,
        SSVCriterion(),
        AVLDecrementRule(),
        initial_temperature,
        initial_configuration,
        neighborhood_size)
end

"""
    AnnealingOptimization(energy, propose_candidate, samples, neighborhood_size)

Given samples of the configuration space, determine the initial temperature.
The first sample is used as initial configuration.

Initial temperature is the standard deviation of the energy in the sample,
according to White criterion (see eq. 2.36 in Varanelli).
"""
function AnnealingOptimization(energy, propose_candidate, samples::Vector,
                               neighborhood_size::Integer)

    energies = energy.(samples)
    return AnnealingOptimization(
        energy,
        propose_candidate,
        samples[1],
        std(energies),
        neighborhood_size)
end

"""
    AnnealingState

Struct representing the sate of the Simulated Annealing algorithm at the end
of a Markov chain step.

Fields
======
temperature
current_configuration
current_energy
bsf_configuration
    Best configuration found so far.
bsf_energy
    Energy of the best configuratin found so far.
energies
    Array of all energies encountered during the last Markov chain step.
"""
struct AnnealingState{T, C}
    temperature::T
    current_configuration::C
    current_energy::T
    bsf_configuration::C
    bsf_energy::T
    energies::Vector{T}
end

function initial_state(search::AnnealingOptimization{T}) where T
    config = search.initial_configuration
    E = search.energy(config)
    return AnnealingState(search.initial_temperature,
                          config, E,
                          config, E,
                          zeros(T, 0))
end

include("decrement_rule.jl")
include("stop_criterion.jl")

Base.eltype(::Type{A}) where {T, C, A <: AnnealingOptimization{T, C}} = AnnealingState{T, C}
Base.IteratorSize(::Type{A}) where {A <: AnnealingOptimization} = Base.SizeUnknown()

function Base.iterate(
        search::AnnealingOptimization{T},
        state=initial_state(search)) where T
    length(state.energies) > 0 && search.stop_criterion(state) && return nothing

    # Initialize all internal loop variables
    configuration = state.current_configuration
    E = state.current_energy
    bsf = state.bsf_configuration
    bsf_energy = state.bsf_energy
    energies = zeros(T, search.neighborhood_size)

    # Markov chain at constant temperature
    for k in 1:search.neighborhood_size
        candidate, dE = search.propose_candidate(configuration)

        if dE < 0 || rand() < exp(-dE/state.temperature)
            configuration = candidate
            E += dE

            if E < state.bsf_energy
                bsf = configuration
                bsf_energy = E
            end
        end

        energies[k] = E
    end

    new_state = AnnealingState(search.decrement_rule(state),
                               configuration, E,
                               bsf, bsf_energy,
                               energies)

    return new_state, new_state
end

function simulated_annealing(search::AnnealingOptimization)
    local state = nothing  # Avoid being shadowed inside the loop

    for s in search  # Go to the end of the search
        state = s
    end

    return state.bsf_configuration, state.bsf_energy
end

end


"""
`extrap_prep_dim(::Type{Reflect}, d)`

First, translate x into the domain over [lbound, 2(ubound-lbound)) i.e. into twice the size of the domain.
Next, if x is now in the upper part of this ''double-domain´´, reflect over the middle to obtain a new value x' for which f(x') == f(x), but where x' is inside the domain
"""
function extrap_prep_dim(::Type{Reflect}, d)
    xs_d = Symbol("xs_", d)
    quote
        start = lbound(etp.itp, $d)
        width = ubound(etp.itp, $d) - start

        $xs_d = mod($xs_d - start, 2width) + start
        $xs_d > start + width && (xs_d = start + width - $xs_d)
    end
end

extrap_prep{N,d}(::Type{Reflect}, ::Val{N}, ::Val{d}) = extrap_prep_dim(Reflect, d)
function extrap_prep{N,d}(::Type{Reflect}, ::Val{N}, ::Val{d}, ::Val{:lo})
    xs_d = Symbol("xs_", d)
    :($xs_d < lbound(etp.itp, $d) && $(extrap_prep_dim(Reflect, d)))
end
function extrap_prep{N,d}(::Type{Reflect}, ::Val{N}, ::Val{d}, ::Val{:hi})
    xs_d = Symbol("xs_", d)
    :($xs_d > ubound(etp.itp, $d) && $(extrap_prep_dim(Reflect, d)))
end


function insertionsort!(data::Array{Int64})
	for i in 2 : length(data)
		current_value = data[i]
		for j in i - 1 : -1 : 1
			if current_value < data[j]
				data[j + 1] = data[j]
				if j == 1
					data[j] = current_value
				end
			else
				data[j + 1] = current_value
				break
			end
		end
	end
end

function builtinsort!(data::Array{Int64})
	sort!(data)
end

function quicksort!(data::Array, start=-1, stop=-1, pivot=-1)
	if start == -1
		start = 1
		stop = length(data)
		pivot = 0
		for i in data
			pivot += i
		end
		pivot = pivot / stop
	end

	print(start)
	print(", ")
	print(stop)
	print(", ")
	println(pivot)

	if stop - start <= 2
		if stop - start == 2
			if data[start] > data[start + 1]
				buffer = data[start]
				data[start] = data[start + 1]
				data[start + 1] = buffer
			end
			if data[start + 1] > data[start + 2]
				buffer = data[start + 1]
				data[start + 1] = data[start + 2]
				data[start + 2] = buffer
			end
			if data[start] > data[start + 1]
				buffer = data[start]
				data[start] = data[start + 1]
				data[start + 1] = buffer
			end
			println(data)
			return
		end

		if stop - start == 1
			if data[start] > data[start + 1]
				buffer = data[start]
				data[start] = data[start + 1]
				data[start + 1] = buffer
			end
			println(data)
			return
		end

		if stop - start == 0
			return
		end
	end

	leftPointer = start
	rightPointer = stop

	done = false

	while !done
		print(leftPointer)
		print(", ")
		println(rightPointer)

		left = data[leftPointer] <= pivot
		right = data[rightPointer] > pivot

		if rightPointer - leftPointer <= 0
			done = true
		elseif !left && !right
			buffer = data[leftPointer]
			data[leftPointer] = data[rightPointer]
			data[rightPointer] = buffer
			if rightPointer - leftPointer > 1
				leftPointer += 1
				rightPointer -= 1
			else
				done = true
			end
		elseif !left && right
			rightPointer -= 1
		elseif left && !right
			leftPointer += 1
		elseif rightPointer - leftPointer > 1
			leftPointer += 1
			rightPointer -= 1
		else
			done = true
		end
		println(data)
	end

	pivotPointer = leftPointer

	if leftPointer == rightPointer && data[leftPointer] > pivot
		pivotPointer -= 1
	end

	quicksort!(data, start, pivotPointer, pivot / 2)
	quicksort!(data, pivotPointer + 1, stop, pivot * 3 / 2)
end


function simulate_gtilde(; NC, NT, etaC, etaT, loc, scale, df, skew)
  K = length(loc)

  @assert K == length(etaC) == length(etaT)
  @assert K == length(scale) == length(df) == length(skew)

  mmC = MixtureModel(SkewT.(loc, scale, df, skew), etaC)
  mmT = MixtureModel(SkewT.(loc, scale, df, skew), etaT)

  yC = rand(mmC, NC)
  yT = rand(mmT, NT)

  return (yC=yC, yT=yT, etaC=etaC, etaT=etaT, loc=loc, scale=scale, df=df,
          skew=skew, mmC=mmC, mmT=mmT)
end

function simulate_ftilde(; NC, NT, etaC, etaT, loc, scale, df, skew, gammaC, gammaT)
  QC = sum(gammaC .> rand(NC))
  QT = sum(gammaT .> rand(NT))

  NC_finite = NC - QC
  NT_finite = NT - QT

  pzero = (gammaC=gammaC, gammaT=gammaT, QC=QC, QT=QT, NC=NC, NT=NT)

  gtilde = simulate_gtilde(NC=NC_finite, NT=NT_finite, etaC=etaC, etaT=etaT, loc=loc,
                           scale=scale, df=df, skew=skew)

  return merge(pzero, gtilde)
end


using DiscreteChoice
using Optim
using ForwardDiff
using Random, StatsBase

#################
# Simulate data
#################

# Simulate J=100 case

Random.seed!(42)

N = 50_000
K1 = 1
K2 = 5
L = 4
J = 100

Y = rand(1:J,N)
X1 = randn(N,J,K1)
X2 = randn(J,K2)
D = randn(N,L)

α = randn(K1)
Π = randn(K2, L)
ξ_all = vcat(0, rand(J-1))
ξ = ξ_all[2:end]

Py = DiscreteChoice.y_prob(X1,X2,D, α, Π, ξ_all)
Y = zeros(Int64, N)
for i in eachindex(Y)
    Y[i] = sample(1:J, Weights(Py[i,:]))
end
[mean(Y.==i) for i in 1:J] |> minimum

DiscreteChoice.llk_sample(Y, X1, X2, D, α, Π, ξ)

#######################
# Estimate the model
#######################

model_homo = DCModel(Y,X1,X2,D; hetero_preference=false)
estDCModel!(model_homo; init_guess = vcat(α, vec(Π), ξ))
# estDCModel!(model_homo; init_guess = randn(K1+K2*L+J-1))
updateCoef!(model_homo)

@show model_homo.α .- α
@show model_homo.Π .- Π
@show model_homo.ξ .- ξ


# Structures of Vector calculus operators store all the operators required during computations.

# Operator for calculating gradients of a N-dimensional function

struct GradientOperator{T,N,O}
    ops :: O
end

function Gradient(approximation_order :: Int, dx::Union{NTuple{N,AbstractVector{T}},NTuple{N,T}},
        len::NTuple{N,Int}; coeff_func=1) where {T<:Real, N}

    ops = permutedims([CenteredDifference{n}(1, approximation_order, dx[n], len[n], coeff_func) for n in 1:N])

    GradientOperator{T,N,typeof(ops)}(
        ops
        )
end

# Operator for calculating curls of a given 3-dimensional vector stored as space-tensor

struct CurlOperator{T,O}
    ops :: O
end

function Curl(approximation_order :: Int, dx::Union{NTuple{3,AbstractVector{T}},NTuple{3,T}},
    len::NTuple{3,Int}; coeff_func=1 ) where {T<:Real}

    @assert len[1] == len[2] && len[2] == len[3] "All dimensions must have equal no. of grid points"
    ops = permutedims([CenteredDifference{n}(1, approximation_order, dx[n], len[n], coeff_func) for n in 1:3])

    CurlOperator{T,typeof(ops)}(
        ops
        )

end

# Operator for calculating divergence of a given N-dimensional vector stored as space-tensor

struct DivergenceOperator{T,N,O}
    ops :: O
end

function Divergence(approximation_order :: Int, dx::Union{NTuple{N,AbstractVector{T}},NTuple{N,T}},
    len::NTuple{N,Int}; coeff_func=1) where {T<:Real, N}

    ops = permutedims([CenteredDifference{n}(1, approximation_order, dx[n], len[n], coeff_func) for n in 1:N])

    DivergenceOperator{T,N,typeof(ops)}(
        ops
    )

end

function *(A::GradientOperator{T},M::AbstractArray{T,N}) where {T<:Real,N}

    size_x_temp = [size(M)...].-2

    x_temp = zeros(T,size_x_temp...,N)

    mul!(x_temp,false,A,M)
    return x_temp
end

function *(A::CurlOperator{T},M::AbstractArray{T,4}) where {T<:Real}

    size_x_temp = [size(M)...].-2
    size_x_temp[4] += 2
    x_temp = zeros(T,size_x_temp...)
    mul!(x_temp, A, M, overwrite = false)
    return x_temp
end

function *(A::DivergenceOperator{T},M::AbstractArray{T,N}) where {T<:Real, N}

    size_x_temp = [size(M)[1:end-1]...].-2

    x_temp = zeros(T,size_x_temp...)

    mul!(x_temp,true,A,M)
    return x_temp
end

function *(c::Number, A::GradientOperator{T,N}) where {T,N}
    ops = permutedims([c*A.ops[i] for i in 1:N])
    GradientOperator{T,N,typeof(ops)}(
        ops
        )
end

function *(c::Number, A::CurlOperator{T}) where {T}
    ops = permutedims([c*A.ops[i] for i in 1:3])
    CurlOperator{T,typeof(ops)}(
        ops
        )
end

function *(c::Number, A::DivergenceOperator{T,N}) where {T,N}
    ops = permutedims([c*A.ops[i] for i in 1:N])
    DivergenceOperator{T,N,typeof(ops)}(
        ops
        )
end


using NodesAndModes
using Test
using LinearAlgebra

@testset "1D tests" begin
    tol = 1e2*eps()

    N = 0
    r,w = gauss_lobatto_quad(0,0,N)
    @test sum(w) ≈ 2

    for N = 1:2
        r,w = gauss_quad(0,0,N)
        @test sum(w) ≈ 2
        @test abs(sum(w.*r)) < tol

        V = vandermonde(Line(),N,r)
        @test V'*diagm(w)*V ≈ I

        Vr = grad_vandermonde(Line(),N,r)
        Dr = Vr/V
        @test norm(sum(Dr,dims=2)) < tol
        @test Dr*r ≈ ones(N+1)

        r,w = gauss_lobatto_quad(0,0,N)
        @test sum(w) ≈ 2
        @test abs(sum(w.*r)) < tol

        # check if Lobatto is exact for 2N-2 polynoms
        V = vandermonde(Line(),N-1,r)
        @test V'*diagm(w)*V ≈ I

        Dr = grad_vandermonde(Line(),N,r)/vandermonde(Line(),N,r)
        @test norm(sum(Dr,dims=2)) < tol
        @test Dr*r ≈ ones(N+1)
    end
end

area(elem::Quad) = 4.0
area(elem::Tri) = 2.0

@testset "2D $elem basis tests" for elem in (Tri(), Quad())

    tol = 1e2*eps()

    N = 3
    rq, sq, wq = quad_nodes(elem, N)
    @test sum(wq) ≈ area(elem)

    Vq = vandermonde(elem, N, rq, sq)
    @test Vq' * diagm(wq) * Vq ≈ I

    r, s = nodes(elem, N)
    V = vandermonde(elem, N, r, s)
    Dr, Ds = (A->A / V).(grad_vandermonde(elem, N, r, s))
    @test norm(sum(Dr, dims=2)) + norm(sum(Ds, dims=2)) < tol
    @test norm(Dr * s) + norm(Ds * r) < tol
    @test Dr * r ≈ ones(length(r))
    @test Ds * s ≈ ones(length(s))

    r, s = equi_nodes(elem, N)
    V = vandermonde(elem, N, r, s)
    Dr,Ds = (A->A / V).(grad_vandermonde(elem, N, r, s))
    @test norm(sum(Dr, dims=2)) + norm(sum(Ds, dims=2)) < tol
    @test norm(Dr * s) + norm(Ds * r) < tol
    @test Dr*r ≈ ones(length(r))
    @test Ds*s ≈ ones(length(s))

    # check for Kronecker structure
    if elem == Quad()
        r, s = nodes(elem, N)
        V = vandermonde(elem, N, r, s)
        V1D = vandermonde(Line(), N, nodes(Line(), N))
        @test V ≈ kron(V1D, V1D)
    end
end

volume(elem::Hex) = 8
volume(elem::Wedge) = 4
volume(elem::Pyr) = 8/3
volume(elem::Tet) = 4/3

@testset "3D $elem basis tests" for elem in (Hex(), Wedge(), Pyr(), Tet())

    tol = 5e2 * eps()

    N = 3
    rq, sq, tq, wq = quad_nodes(elem,N)
    @test sum(wq) ≈ volume(elem)
    if elem==Hex()
        @test abs(sum(rq .* wq)) + abs(sum(sq .* wq)) + abs(sum(tq .* wq)) < tol
    end

    Vq = vandermonde(elem, N, rq, sq, tq)
    @test Vq' * diagm(wq) * Vq ≈ I

    if elem != Pyr()
        r, s, t = nodes(elem, N)
        V = vandermonde(elem, N, r, s, t)
        Dr, Ds, Dt = (A->A / V).(grad_vandermonde(elem, N, r, s, t))
        @test norm(sum(Dr, dims=2)) + norm(sum(Ds, dims=2)) + norm(sum(Dt, dims=2)) < tol
        @test norm(Dr * s) + norm(Dr * t) < tol
        @test norm(Ds * r) + norm(Ds * t) < tol
        @test norm(Dt * r) + norm(Dt * s) < tol
        @test Dr * r ≈ ones(length(r))
        @test Ds * s ≈ ones(length(s))
        @test Dt * t ≈ ones(length(t))

        r, s, t = equi_nodes(elem, N)
        V = vandermonde(elem, N, r, s, t)
        Dr, Ds, Dt = (A->A / V).(grad_vandermonde(elem, N, r, s, t))
        @test norm(sum(Dr, dims=2)) + norm(sum(Ds, dims=2)) + norm(sum(Dt, dims=2)) < tol
        @test norm(Dr * s) + norm(Dr * t) + norm(Ds * r) + norm(Ds * t) + norm(Dt * r) + norm(Dt * s) < tol
        @test Dr * r ≈ one.(r)
        @test Ds * s ≈ one.(s)
        @test Dt * t ≈ one.(t)
    elseif elem == Pyr()
        r, s, t = nodes(elem, N)
        rq, sq, tq, wq = quad_nodes(elem, N)
        Vq, Vr, Vs, Vt = (A -> A / vandermonde(elem, N, r, s, t)).(basis(elem, N, rq, sq, tq))
        M = Vq' * diagm(wq) * Vq

        # construct weak differentiation matrices instead of nodal differentiation matrices since
        # the pyramid basis is singular at the node at the tip of the pyramid.
        Dr, Ds, Dt = (A -> M \ (Vq' * diagm(wq) * A)).((Vr, Vs, Vt))
        @test norm(sum(Dr, dims=2)) < tol
        @test norm(sum(Ds, dims=2)) < tol
        @test norm(sum(Dt, dims=2)) < tol
        @test norm(Dr * s) + norm(Dr * t) < tol
        @test norm(Ds * r) + norm(Ds * t) < tol
        @test norm(Dt * r) + norm(Dt * s) < tol
        @test Dr * r ≈ one.(r)
        @test Ds * s ≈ one.(s)
        @test Dt * t ≈ one.(t)
    end
end

@testset "Test for Kronecker structure in the Hex basis matrix" begin
    tol = 5e2*eps()
    N = 3
    VDM_1D = vandermonde(Line(), N, nodes(Line(), N))
    VDM_quad_1D = vandermonde(Line(), N, quad_nodes(Line(), N)[1])
    @test abs(norm(vandermonde(Hex(), N, nodes(Hex(), N)...) - kron(VDM_1D, VDM_1D, VDM_1D))) < tol
    @test abs(norm(vandermonde(Hex(), N, quad_nodes(Hex(), N)[1:3]...) - kron(VDM_quad_1D, VDM_quad_1D, VDM_quad_1D))) < tol
end

# test high order quadrature
@testset "Simplicial Stroud quadrature" begin
    tol = 1e3*eps()

    N = 14
    rq, sq, wq = quad_nodes(Tri(), N)
    @test sum(wq) ≈ 2
    @test sum(rq .* wq) ≈ -2/3
    @test sum(sq .* wq) ≈ -2/3

    N = 8
    rq, sq, tq, wq = quad_nodes(Tet(), N)
    @test sum(wq) ≈ 4/3
    @test sum(rq .* wq) ≈ -2/3
    @test sum(sq .* wq) ≈ -2/3
    @test sum(tq .* wq) ≈ -2/3
end

if VERSION >= v"1.6" # apparently inference got better after 1.5?
    @testset "Inferrability for $elementType" for elementType in (Line(), Tri(), Quad(), Tet(), Pyr(), Wedge(), Hex())
        N = 1
        if elementType == Line()
            @test_nowarn (@inferred gauss_quad(0, 0, N)) == ([-0.5773502691896257, 0.5773502691896257], [0.9999999999999998, 0.9999999999999998])
            @test_nowarn (@inferred gauss_lobatto_quad(0, 0, N)) == ([-1.0, 1.0], [1.0, 1.0])
        end

        @test_nowarn @inferred nodes(elementType, N)
        @test_nowarn @inferred quad_nodes(elementType, N)
        @test_nowarn @inferred equi_nodes(elementType, N)
        if elementType == Line()
            @test_nowarn @inferred basis(elementType, N, nodes(elementType,N))
        else
            @test_nowarn @inferred basis(elementType, N, nodes(elementType,N)...)
        end
    end
end


module continuousEntropyEstimators

using Distances, NearestNeighbors, SpecialFunctions
# using Statistics
# TODO: implement version in 1D, possibly 2D that doesn't rely on trees
# TODO: implement version that runs in parallel
export entropyEstimator, kraskovEstimator, NNestimator, entropy, mutualInformation, conditionalMutualInformation

# because I like math symbols
Γ(x) = gamma(x)
abstract type entropyEstimator end
mutable struct kraskovEstimator<: entropyEstimator end
mutable struct NNestimator<: entropyEstimator end

"""
    entropy(data::Array{T,2},  estimator::kraskovEstimator=kraskovEstimator(); k::Int64=3, logbase::S=2, treeType=Chebyshev()) where {T<:AbstractFloat, S<:Real}
    entropy(data::Array{Int64}, logbase::Float64=2.0)

computes the entropy of the data, which should be a d×N array.

If data is an array of floats, it is treated as a differential entropy calculation, by default done using knn's with k=3. treeType should be a distance from Distances. Some options include

| Distance     | As a p norm |
|--------------|:-----------:|
| Chebyshev()  | ``∞``       |
| Euclidean()  | ``2``       |
| Cityblock()  | ``1``       |
| Minkowski(p) | ``p``       |

If data is an array of ints, it uses the plug-in estimator.

TODO: Calculate the entropy in the method of Gao et al for mixed distributions
"""
function entropy(data::Array{T,2},  estimator::kraskovEstimator=kraskovEstimator(); k::Int64=3, logbase::S=2, treeType=Chebyshev()) where {T<:AbstractFloat, S<:Real}
    d, N = size(data)
    # perturb all the points to avoid problems with deterministic functions
    X = copy(data) + 1e-10*randn(size(data))
    tree = KDTree(X, treeType)
    if treeType==Chebyshev()
        cb = d*log(logbase, 2)
    elseif treeType==Euclidean()
        p=2
    elseif treeType==Cityblock()
        p=1
    elseif typeof(treeType)==Minkowski{Int64} || typeof(treeType)==Minkowski{Float64}
        p = treeType.p
    end
    if treeType!=Chebyshev()
        print(treeType)
        cb = d*log(logbase,2) + d*log(logbase, Γ(1+1/p))- log(logbase, Γ(1+d/p))
    end
    entropyEstimate = digamma(N)-digamma(k) + cb + d*log(logbase, 2) # this last factor is left implicit in eq 20
    for i=1:N
        _, dist = knn(tree, X[:,i], k+1, true)
        # check to see if this point is discrete, in which case Gao et al suggest using a plug-in estimator
        # if dist[end]<1e-10
        entropyEstimate += d/N*log(logbase, dist[end])
    end
    return entropyEstimate
end
continuousEntropyEstimators.entropy
# function entropy(data::Array{Float64,2}, estimator::NNestimator=NNestimator(); logbase=2, treeType=Chebyshev())
#     d, N = size(data)
#     # perturb all the points to avoid problems with deterministic functions
#     X = copy(data) + 1e-10*randn(size(data))
#     tree = KDTree(X, treeType)
#     entropyEstimate = γ/log(logbase) + log(logbase, 2)
#     for i=1:N
#         _, dist = knn(tree, X[:,i], 2, true)
#         entropyEstimate += 1/N*log(logbase, N*dist[end])
#     end
#     return entropyEstimate
# end
# unique([1.5, 1.5, 3, 201])
# data = rand([1,3,2], 10)
# uniqueVal = unique(data)
# data[data.==uniqueVal[1]]
# maximum(data.==uniqueVal[1])
function entropy(data::Array{Int64}, logbase::Float64=2.0)
    pEstimate, _ = pluginEstimator(data)
    pEstimate[pEstimate.==0] = 1
    return -pEstimate'*log.(logbase, pEstimate)
end

"""
    pluginEstimator(data::Array{Int64,(1,2)})
estimate the probability distribution of a discrete variable by using the counts
"""
function pluginEstimator(data::Array{Int64,2})
    Y = [data[:,i] for i=1:size(data, 2)]
    uniqueVal = unique(Y)
    counts = zeros(size(uniqueVal))
    for i = 1:size(counts,1)
        for j = 1:length(Y)
            counts[i] += Y[j]==uniqueVal[i]
        end
    end
    counts = counts/length(data)
    return (counts, uniqueVal)
end

function pluginEstimator(data::Array{Int64,1})
    uniqueVal = unique(data)
    counts = zeros(size(uniqueVal))
    for i = 1:size(counts,1)
        counts[i] = sum(data.==uniqueVal[i])
    end
    counts = counts/length(data)
    return (counts, uniqueVal)
end
# abs(entropy(rand([1,3], 1000000))-1)<1e-5 # uniform variable on two values has unit entropy base 2
"""
    mutualInformation(Xc::Array{Float64,2}, Y::Array{Float64,2}, k=3; logbase=2, treeType=Chebyshev(), version::Int64=2)
    mutualInformation(Xc::Array{Float64,2}, Yc::Array{Int64}; k=3, logbase=2, treeType=Chebyshev())
    mutualInformation(Xc::Array{Int64,2}, Yc::Array{Float64}; k=3, logbase=2, treeType=Chebyshev())

Calculate the mutual information between Xc and Yc using knn nearest neighbor-type methods.

make sure that the data is `d`×`N`, where `d` is the dimension and `N` is the number of examples. `Xc` and `Yc` should match in the second dimension, of course. `version` corresponds to the versions in the Kraskov et. al paper, and is either 1 or 2. Default is chosen for high dimension. It uses neighborhoods restricted to `Xc` (resp `Yc`) to calculate the contribution from that direction. In the case that one of the variables is discrete and want it treated as such, make sure it is integer. It is calculated using ``I(X;Y) = H(X)-H(X|Y)`` If both are discrete, it just uses ``H(X) + H(Y) - H(X,Y)``
NOTE: only Chebyshev is implemented at this time
"""
function mutualInformation(Xc::Array{Float64,2}, Yc::Array{Float64,2}; k=3, logbase=2, treeType=Chebyshev(), version::Int64=2)
    @assert size(Xc,2)==size(Yc,2)
    @assert treeType==Chebyshev()
    dx, N = size(Xc)
    dy, N = size(Yc)
    X = copy(Xc) + 1e-10*randn(dx, N)
    Y = copy(Yc) + 1e-10*randn(dy, N)
    XY = [X; Y]
    tree = KDTree(XY, treeType)
    entropyEstimate = digamma(k) + digamma(N)
    if version==2
        entropyEstimate-=1/k
    end
    xTree = KDTree(X, treeType)
    yTree = KDTree(Y, treeType)
    for i=1:N
        neighborhood, dist = knn(tree, XY[:,i], k+1, true)
        neighbor = neighborhood[end]
        dist = dist[end]
        if version==1
            # count the points within x(y) distance of ε/2
            nx = length(inrange(xTree, X[:,i], dist, false)) # note that we both need to add one for the algorithm and we "accidentally" overcount by one, so there's no change
            ny = length(inrange(yTree, Y[:,i], dist, false)) # same w/y
        elseif version==2
            εx = chebyshev(X[:, neighbor],X[:,i]) # \varepsilon
            nx = length(inrange(xTree, X[:,i], εx, false)) # note that we both need to add one for the algorithm and we "accidentally" overcount by one, so there's no change
            εy = chebyshev(Y[:, neighbor],Y[:,i])
            ny = length(inrange(yTree, Y[:,i], dist, false)) # note that we both need to add one for the algorithm and we "accidentally" overcount by one, so there's no change
        else
            error("``(version)'s not an allowed value of version. Try 1 or 2")
        end
        entropyEstimate -= 1/N*digamma(nx)
    end
    return entropyEstimate
end

function mutualInformation(Xc::Array{Float64,2}, Yc::Array{Int64}; k=3, logbase=2, treeType=Chebyshev())
    @assert size(Xc,2)==size(Yc,2)
    @assert treeType==Chebyshev()
    dx, N = size(Xc)
    dy, N = size(Yc)
    X = copy(Xc) + 1e-10*randn(dx, N)
    Y = copy(Yc)
    # if Xis pure continuous, Y discrete, it's best to use I(X,Y) = H(Y)-H(Y|X)
    # First order of business is getting a plug-in estimate for Y
    pEstimate, yVals = pluginEstimator(Y)
    Nxgy =pEstimate*N # the number of Xs given the value of y in index
    eeXgY = digamma(k) + digamma.(Nxgy) # a vector of entropy estimates given the value of y
    entropyEstimate = entropy(X,k=k, logbase=logbase, treeType=treeType)
    #
    if dy>1
        Yp = [data[:,i] for i=1:size(data, 2)]
        for i = 1:length(pEstimate)
            for j = 1:length(Yp)
                tmp = pEstimate[i]*entropy(X[:, Yp[j]==yVals[i]])
                entropyEstimate-=pEstimate[i]*entropy(X[:, Yp[j]==yVals[i]])
            end
        end
    else
        # Y is one dimensional
        for i = 1:length(pEstimate)
            entropyEstimate-=pEstimate[i]*entropy(X[:, reshape(Y.==yVals[i],(:))], k=k, logbase=logbase, treeType=treeType)
        end
    end
    return abs(entropyEstimate)
end

function mutualInformation(Xc::Array{Int64}, Yc::Array{Int64}, logbase::Float64=2.0)
    entropy(Xc,logbase=logbase) + entropy(Yc,logbase=logbase) - entropy([Xc;Yc],logbase=logbase)
end

# Various kinds of conditional mutual information
"""
    conditionalMutualInformation(Xc::Array{X},Yc::Array{Y}, Zc::Array{Z}, k::Int64=3, logbase::Float64=2.0, treeType=Chebyshev(), version::Int64=2)) where {X<:Real, Y<:Real, Z<:Real}

Calculate the conditional mutual information between `Xc`, `Yc` given `Zc`, ``I(X;Y|Z) = I(X,Z; Y) - I(Z;Y)``.
"""
function conditionalMutualInformation(Xc::Array{Float64},Yc::Array{Float64}, Zc::Array{Float64}, k::Int64=3, logbase::Float64=2.0, treeType=Chebyshev(), version::Int64=2)
    mutualInformation([Xc; Zc], Yc,k=k,logbase=logbase, treeType=treeType, version=version) - mutualInformation(Zc, Yc, k=k, logbase=logbase, treeType=treeType, version=version)
end
function conditionalMutualInformation(Xc::Array{Float64},Yc::Array{Int64}, Zc::Array{Float64}, k::Int64=3, logbase::Float64=2.0, treeType=Chebyshev(), version::Int64=2)
    mutualInformation([Xc; Zc], Yc,k=k,logbase=logbase, treeType=treeType) - mutualInformation(Zc, Yc, k=k, logbase=logbase, treeType=treeType)
end
end # module

# entropy([rand([1,3,2,4,2],5000)'; rand([4,3,2,1,4],5000)'])
# data = [rand([1,3,2],500)'; rand([4,3,2],500)']
# X = [data[:,i] for i=1:size(data, 2)]
# uniqueVal = unique(X)
# uniqueVal[1]
# entropy([2.5*rand([1,3,2],500)'; 2.5*rand([4,3,2],500)'])
# d = 20
# X = randn(d,  1000)
# Y = X+.01*randn(d, 1000)
# @time mutualInformation(X, Y, version=2)
# @time entropy(X)
# @time entropy(Y)
# get_entropy(rand(1000))
# d = 3
# X = randn(d, 100000)
# @time entropy(X, kraskovEstimator(), k=3, treeType = Chebyshev())- 1/2*log2((e*2π)^d)
# @time entropy(X, NNestimator(), treeType = Chebyshev())
# 1/2*log2((e*2π)^d)
# using InformationMeasures
# 1.3.^range(-.001, -15, 5)
# 1.3.^range(3, 42, 5)
# 2.^range(4,14,5)
# 1/2*log(2,(2π)^2)
# Ns = [10, 50, 100, 500, 1000] #, 3000, 5000, 8000, 10000]
# esttimatesk8 = [zeros(x) for x in Ns]
# for (i,estimates) in enumerate(esttimatesk8)
#     print(i)
#     @time for j=1:length(estimates)
#         data = randn(10,Ns[i])
#         estimates[j] = entropy(data,8)
#     end
# end
# # simple binning
#
# using Cumulants
# c = cumulants(data')
# using InformationMeasures
# using Plots
# plot(Ns, [mean.(esttimates), mean.(esttimatesk8)])
# std.(esttimates)
# mean(estimates1)
# std(estimates1)
# entropyEstimate
# @assert N>=d
# data+=1e-10*randn(size(data))
#
#
# test = randn(10,10000)
# get_entropy(test)
# 1/2*log2((2π)^10)
# function sReLU(x)
#     if x>=0
#         return x*exp(-1/(1-x)^2)
#     else
#         return 0
#     end
# end
#
# using Plots
# t=-.5:.00001:10
# plot(t,[sReLU.(t), t])


# Extension operator: joConvolve
# Author : Rajiv Kumar
## helper module
module joConvolve_etc
    using JOLI: jo_convert
    using FFTW
    ## forward 1D convolution
# function Convolve1D_fwd{VT<:Number}(x::AbstractVector{VT}, fKernel::AbstractVector,idx::AbstractVector, k::Integer,m::Integer,cflag::Bool,rdt=DataType)
function Convolve1D_fwd(x::AbstractVector, fKernel::AbstractVector,idx::AbstractVector, k::Integer,m::Integer,cflag::Bool,rdt=DataType)
        #fx = fft([full(x);zeros(VT,k-1,1)[:]]);
        #fx = fft([full(x);zeros(k-1,1)[:]]);
         fx = fft([x;zeros(k-1,1)[:]]);
         y = ifft(fKernel.*fx);
         y = y[idx];
         if ~cflag && isreal(y)
             y = real(y);
         end
        return jo_convert(rdt,y,false)
 end
  ## transpose 1D convolution
 #function Convolve1D_tran{VT<:Number}(x::AbstractVector{VT}, fKernel::AbstractVector,idx::AbstractVector, k::Integer,m::Integer,cflag::Bool,rdt=DataType)
 function Convolve1D_tran(x::AbstractVector, fKernel::AbstractVector,idx::AbstractVector, k::Integer,m::Integer,cflag::Bool,rdt=DataType)
        #z      = zeros(VT,m+k-1,1)[:];
        z      = zeros(m+k-1,1)[:];
        #z[idx] = full(x);
        z[idx] = x;
        y      = ifft(conj(fKernel).*fft(z));
        y      = y[1:m];
         if ~cflag && isreal(y)
             y = real(y);
         end
        return jo_convert(rdt,y,false)
 end
end
using .joConvolve_etc

export joConvolve

#function joConvolve{KDT<:Number}(m::Integer, n::Integer, kernel::AbstractVector{KDT}; DDT::DataType=joFloat,RDT::DataType=DDT)
function joConvolve(m::Integer, n::Integer, kernel::AbstractVector; DDT::DataType=joFloat,RDT::DataType=DDT)
                # convolution loop start
                offset = 1
                k      = length(kernel);
                cflag  = (eltype(kernel)<:Complex); #~isreal(kernel);
                # Shift kernel and add internal padding
                # kernel = [kernel[offset:end];zeros(KDT,m-1,1)[:];kernel[1:offset-1]];
                kernel = [kernel[offset:end];zeros(m-1,1)[:];kernel[1:offset-1]];
                   idx = collect(1:m);

                # Precompute kernel in frequency domain
               #fKernel = fft(full(kernel));
               fKernel = fft(kernel);
               # Convolution
               return joLinearFunctionFwd(length(idx),m,
                                        v1->joConvolve_etc.Convolve1D_fwd(v1,fKernel,idx, k,m,cflag,RDT),
                                        v2->joConvolve_etc.Convolve1D_tran(v2,fKernel,idx, k,m,cflag,DDT),
                                        v2->joConvolve_etc.Convolve1D_tran(v2,fKernel,idx, k,m,cflag,DDT),
                                        v1->joConvolve_etc.Convolve1D_fwd(v1,fKernel,idx, k,m,cflag,RDT),
                                        DDT,RDT;fMVok=false,
                                        name="joConvolve1D")
    return nothing
end


module InterpolationFunctions

	using Dates
	import Statistics
	export interpolatedMax, interpolate, interpolatedSum, addArraysShiftedInterpolated, medianfilter, averageSamples, smooth

	function interpolatedMax(discreteMax, values)
	  inta=values[discreteMax-1]
	  intb=values[discreteMax]
	  intc=values[discreteMax+1]
	  minAC = min(inta, intc)
	  max = discreteMax + (intc - inta)/(inta + intb + intc - 3*minAC) # Tricky: -0.5 if a=b, +0.5 if b=c, +0.0 if a=c, interpol in between
	return max
	end

	function interpolate(x::AbstractFloat, xAxis, yAxis)
	    try
	      if x <= xAxis[1]
		return yAxis[1]
	      elseif x >= xAxis[end]
		return yAxis[end]
	      end
	      indexLow = searchsortedfirst(xAxis,x) - 1
	      fraction = (x - xAxis[indexLow]) / (xAxis[indexLow+1] - xAxis[indexLow])
	      return fraction * yAxis[indexLow+1] + (1-fraction) * yAxis[indexLow]
	    catch
		println("Interpolation Error")
	    end
	end

	function interpolate(x::DateTime, xAxis, yAxis)
	    x=Dates.datetime2unix(x)
	    xAxis = Dates.datetime2unix.(xAxis)
	    try
	      if x <= xAxis[1]
		return yAxis[1]
	      elseif x >= xAxis[end]
		return yAxis[end]
	      end
	      indexLow = searchsortedfirst(xAxis,x) - 1
	      fraction = (x - xAxis[indexLow]) / (xAxis[indexLow+1] - xAxis[indexLow])
	      return fraction * yAxis[indexLow+1] + (1-fraction) * yAxis[indexLow]
	    catch
		println("Interpolation Error")
	    end
	end

	function interpolate(x::AbstractFloat, yAxis)
	    try
	      if x <= 1
		return yAxis[1]
	      elseif x >= length(yAxis)
		return yAxis[end]
	      end
	      indexLow = Int64(floor(x))
	      fraction = (x - indexLow)
	      return fraction * yAxis[indexLow+1] + (1-fraction) * yAxis[indexLow]
	  catch
	      println("Interpolation Error")
	  end
	end

	function interpolate(x::AbstractArray, xAxis, yAxis)
	  y = Array{typeof(yAxis[1])}(undef,length(x))
	  Threads.@threads for i = 1 : length(x)
	    y[i] = InterpolationFunctions.interpolate(x[i],xAxis,yAxis)
	  end
	  return y #convert(Array,y)
	end

	function interpolateMatrix_1D(x::AbstractArray, xAxis, yMatrix)
		yMatrix_intpol = Array{typeof(yMatrix[1])}(undef,length(x), size(yMatrix)[2])
	  	for i in 1:size(yMatrix)[2]
			yMatrix_intpol[:,i] = interpolate(x, xAxis, yMatrix[:,i])
		end
		return yMatrix_intpol
	end

	function interpolate(x::AbstractArray, yAxis)
	  y = Array{typeof(yAxis[1])}(undef,length(x))
	  Threads.@threads for i = 1 : length(x)
	    y[i] = InterpolationFunctions.interpolate(x[i],yAxis)
	  end
	  return y #convert(Array,y)
	end

	function interpolatedSum(startX::AbstractFloat, endX::AbstractFloat, xAxis, yAxis)
	  firstCompleteIndex = searchsortedfirst(xAxis, startX)
	  lastCompleteIndex = searchsortedfirst(xAxis, endX) - 1
	  if (firstCompleteIndex > 0) && (firstCompleteIndex < length(yAxis)) && (lastCompleteIndex > 0) && (lastCompleteIndex < length(yAxis))

	    fractionStart = xAxis[firstCompleteIndex] - startX
	    firstFractionValueContribution = fractionStart * yAxis[firstCompleteIndex-1]

	    fractionEnd = endX - xAxis[lastCompleteIndex]
	    lastFractionValueContribution = fractionEnd * yAxis[lastCompleteIndex + 1]

	    if firstCompleteIndex == lastCompleteIndex
	      return (firstFractionValueContribution + lastFractionValueContribution)
	    end
	    return (sum(yAxis[firstCompleteIndex:lastCompleteIndex] + firstFractionValueContribution + lastFractionValueContribution))
	  end
	  return 0
	end

	function interpolatedSum(startIndexExact::AbstractFloat, endIndexExact::AbstractFloat, yAxis)
	  subIdxStart::Int64 = ceil(startIndexExact)
	  subIdxStartRoundError = subIdxStart - startIndexExact

	  subIdxEnd::Int64 = floor(endIndexExact)
	  subIdxEndRoundError = endIndexExact - subIdxEnd
	  ret = 0
	  if (subIdxStart>1) && (subIdxEnd+1 < length(yAxis))
	    ret = sum(yAxis[subIdxStart:subIdxEnd]) + yAxis[subIdxStart-1]*subIdxStartRoundError + yAxis[subIdxEnd+1]*subIdxEndRoundError
	  end
	  return ret
	end

	function addArraysShiftedInterpolated(destinationArray::Array, sourceArray::Array, indexShift::Number)
	    if ceil(indexShift) == indexShift
		lowContrib = 1
	    else
		lowContrib = ceil(indexShift) - indexShift
	    end
	    highContrib = 1-lowContrib
	    minIdx = Int64(floor(indexShift)+1)
	    maxIdx = Int64(minIdx + length(sourceArray)-1)
	    if minIdx > length(destinationArray)
		return destinationArray
	    end
	    if maxIdx > length(destinationArray)
		maxIdx = length(destinationArray)
	    end
	    destinationArray[minIdx:maxIdx] += lowContrib*sourceArray[1:maxIdx-minIdx+1]
	    if maxIdx > length(destinationArray)-1
		maxIdx = length(destinationArray)-1
	    end
	    destinationArray[minIdx+1:maxIdx+1] += highContrib*sourceArray[1:maxIdx-minIdx+1]
	    return destinationArray
	end

	function medianfilter(v,ws)
	  [median(v[i:(i+ws-1)]) for i=1:(length(v)-ws+1)]
	end

	function nanmean(x::AbstractArray)
		if length(filter(!isnan, x)) > 0
			return Statistics.mean(filter(!isnan,x))
		elseif length(filter(!isnan, x)) == 0
			return NaN
		end
	end

	function nanmean(x::AbstractArray,d)
		return mapslices(nanmean,x;dims = d)
	end

	# Averaging multidimensional array in one dimension
	# adapted from https://julialang.org/blog/2016/02/iteration
	function averageSamples(data, averagePoints; dim=1, ignoreNaNs = false)
	    if averagePoints > 1
		dt=false
		if typeof(data[1]) == DateTime
		    data = Dates.datetime2unix.(data)
		    dt=true
		end
		len = Int64(floor(size(data,dim)/averagePoints))
		sz = [size(data)...]
		sz[[dim...]] .= len
		averaged = Array{eltype(data)}(undef, sz...)
		Rpre = CartesianIndices(size(data)[1:dim-1])
		Rpost = CartesianIndices(size(data)[dim+1:end])
		for Ipost in Rpost
		  for Ipre in Rpre
		    for i=1:len
		      	if ignoreNaNs
				println(data[Ipre, ((i-1)*averagePoints+1):(i*averagePoints), Ipost])
				averaged[Ipre, i, Ipost] = nanmean(data[Ipre, ((i-1)*averagePoints+1):(i*averagePoints), Ipost])
		      	else
		      		averaged[Ipre, i, Ipost] = Statistics.mean(data[Ipre, ((i-1)*averagePoints+1):(i*averagePoints), Ipost])
			end
		    end
		  end
		end
		if dt
		    averaged = Dates.unix2datetime.(averaged)
		end
		return averaged
	    else
		return data
	    end
	end

	function smooth(x, α, dim::Integer=1)
	    s = similar(x)
	    Rpre = CartesianIndices(size(x)[1:dim-1])
	    Rpost = CartesianIndices(size(x)[dim+1:end])
	    _smooth!(s, x, α, Rpre, size(x, dim), Rpost)
	end

	function _smooth!(s, x, α, Rpre, n, Rpost)
	    for Ipost in Rpost
	        # Initialize the first value along the filtered dimension
	        for Ipre in Rpre
	            s[Ipre, 1, Ipost] = x[Ipre, 1, Ipost]
	        end
	        # Handle all other entries
	        for i = 2:n
	            for Ipre in Rpre
	                s[Ipre, i, Ipost] = α*x[Ipre, i, Ipost] + (1-α)*s[Ipre, i-1, Ipost]
	            end
	        end
	    end
	    s
	end

end


using CharFuncPricing, TaylorSeries

@testset "CumulantsJoshi" begin
    κ = 4.0
    θ = 0.25
    σ = 1.0
    ρ = -0.5
    v0 = 0.01
    τ = 0.01
    params = HestonParams{Float64}(v0, κ, θ, ρ, σ)
    cf = DefaultCharFunc{HestonParams{Float64},Taylor1{Complex}}(params)
    t = Taylor1(Float64, 4)
    cft = CharFuncPricing.evaluateLogCharFuncCui(cf, t, τ)
    c1, c2, c4 = computeCumulants(params, τ)
    println(c1, " ", c2 / 2, " ", c4 / (2 * 3 * 4), " ", cft)
    @test isapprox(imag(cft.coeffs[2]), c1, atol = 1e-12)
    @test isapprox(-real(cft.coeffs[3]) * 2, c2, atol = 1e-12)
    @test isapprox(cft.coeffs[5] * 2 * 3 * 4, c4, atol = 1e-16)
end
@testset "CumulantsAlan" begin
    κ = 4.0
    θ = 0.25
    σ = 1.0
    ρ = -0.5
    v0 = 0.04
    τ = 1.0
    params = HestonParams{Float64}(v0, κ, θ, ρ, σ)
    cf = DefaultCharFunc{HestonParams{Float64},Taylor1{Complex}}(params)
    t = Taylor1(Float64, 4)
    cft = CharFuncPricing.evaluateLogCharFuncCui(cf, t, τ)
    c1, c2, c4 = computeCumulants(params, τ)
    println(c1, " ", c2 / 2, " ", c4 / (2 * 3 * 4), " ", cft)
    @test isapprox(imag(cft.coeffs[2]), c1, atol = 1e-12)
    @test isapprox(-real(cft.coeffs[3]) * 2, c2, atol = 1e-12)
    @test isapprox(cft.coeffs[5] * 2 * 3 * 4, c4, atol = 1e-15)
end
@testset "CumulantsNearZeroKappa" begin
    #Low accuracy with small kappa due to powers of kappa in the denominator.
    κ = 0.02
    θ = 0.25
    σ = 1.0
    ρ = -0.5
    v0 = 0.04
    τ = 1.0
    params = HestonParams{Float64}(v0, κ, θ, ρ, σ)
    cf = DefaultCharFunc{HestonParams{Float64},Taylor1{Complex}}(params)
    t = Taylor1(Float64, 4)
    cft = CharFuncPricing.evaluateLogCharFuncCui(cf, t, τ)
    c1, c2, c4 = computeCumulants(params, τ)
    println(c1, " ", c2 / 2, " ", c4 / (2 * 3 * 4), " ", cft)
    @test isapprox(imag(cft.coeffs[2]), c1, atol = 1e-12)
    @test isapprox(-real(cft.coeffs[3]) * 2, c2, atol = 1e-3)
    @test isapprox(cft.coeffs[5] * 2 * 3 * 4, c4, atol = 1e-3)
end


using MittagLefflerFunctions
using Polynomials
using PyPlot
using SpecialFunctions: erfcx

α = 1/2
β = 1.0
N = 5
r = 2N + 1
m = r + 1
n = r

#Eref = MLQuad(α, β, 2N, :hyperbola)
Eref(z) = erfcx(-z)

EP = MLPade(α, β, m, n, :svd)
p, q = Polynomial(EP.p), Polynomial(EP.q)
χ = roots(q)
EPade(z) = p(-z) / q(-z)

EQ = MLQuad(α, β, N, :hyperbola)
w, C, A = EQ.qs.w, EQ.qs.C, EQ.qs.A
wα = w.^α
P = [ conj.(wα[N:-1:1]); wα[0:N] ]
R1 = -A * C .* w.^(α-β)
R = [ conj.(R1[N:-1:1]); R1[0:N] ]

function mlf_quad1(z)
    s = 0.0
    for j = 1:length(P)
        s += R[j] / ( z - P[j] )
    end
    return s
end

figure(1)
plot(real.(-χ), imag.(-χ), "o",
     real.(P), imag.(P), "o")
legend((L"$-\chi_j$", L"$P_j$"))
grid(true)

Nx = 301
Ny = 301
xmin, xmax = -5.0, 3.0
ymin, ymax = -4.0, 4.0
x = range(xmin, xmax, length=Nx)
y = range(ymin, ymax, length=Ny)

z = Complex{Float64}[ complex(x[j], y[i]) for i=1:Ny, j=1:Nx ]

err_quad = Float64[ abs(Eref(z[i,j])-mlf_quad1(z[i,j])) for i = 1:Ny, j=1:Nx ]
err_pade = Float64[ abs(Eref(z[i,j])-EPade(z[i,j])) for i = 1:Ny, j=1:Nx ]

figure(2)
contour(x, y, log10.(err_quad), levels=collect(-9:4))
plot(real.(P), imag.(P), "ro")
colorbar()
xlabel(L"$\Re z$")
ylabel(L"$\Im z$")
grid(true)
savefig("figure11L.pdf")

figure(3)
contour(x, y, log10.(err_pade), levels=collect(-14:4))
plot(real.(-χ), imag.(-χ), "ro")
colorbar()
xlabel(L"$\Re z$")
ylabel(L"$\Im z$")
grid(true)
savefig("figure11R.pdf")


using DifferentialEquations
using LinearAlgebra, DiffEqSensitivity, Optim
using Flux: flatten, params
using DiffEqFlux, Flux
using Plots
using Flux:train!
using NNlib

function dxdt_train(du, u, p, t)
    x, y, z = u
    du[1] = 10.0*(y - x)
    du[2] = x*(28.0-z) - y
    du[3] = x*y - (8/3)*z
end
u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
dt = 0.1
prob = ODEProblem(dxdt_train,u0,tspan)
ts = collect(0.0:0.1:tspan[2])
prob_train = ODEProblem{true}(dxdt_train, u0, tspan, p=nothing)
data_train = Array(solve(prob_train, Tsit5(), saveat=ts))
A = [PolynomialBasis(2), PolynomialBasis(2)]
nn0 = TensorLayer(A, 1)
nn1 = TensorLayer(A, 1)
nn2 = TensorLayer(A, 1)
nn2([0.1,0.1])
f = x -> min(100one(x), x)

function dxdt_pred(du, u, p, t)
    x, y, z = u
    du[1] = 10*(y - x)
    du[2] = x*(p[1]-z) - y
    du[3] = x*y -  (8/3)*z
end

function dxdt_pred1(du, u, p, t)
    x, y, z = u
    du[1] = nn0([x,y],p[1:4])[1]
    du[2] = nn1([x,z],p[5:8])[1] - p[13]*y
    du[3] = nn2([x,y],p[9:12])[1] - p[14]*z
end

α = 27*ones(1)

prob_pred = ODEProblem{true}(dxdt_pred, u0, tspan, p=nothing)
function predict_adjoint(θ)
    x = Array(solve(prob_pred, Tsit5(), p=θ, saveat=ts))
end

function loss_adjoint(θ)
    x = predict_adjoint(θ)
    loss = log.(x) - log.(data_train)
    return loss
end

function cb(θ, l)
    @show l
    return false
end
res1 = DiffEqFlux.sciml_train(loss_adjoint, α, ADAM(0.02), cb=cb, maxiters=500)
res2 = DiffEqFlux.sciml_train(loss_adjoint, res1.u, BFGS(), cb=cb)#, maxiters=10000)
opt = res2.u
data_pred = predict_adjoint(opt)
plot(ts, data_train[1,:], label="X (ODE)")
plot!(ts, data_train[2,:], label="Y (ODE)")
plot!(ts, data_train[3,:], label="Z (ODE)")
plot!(ts, data_pred[1,:], label="X (NN)")
plot!(ts, data_pred[2,:],label="V (NN)")
plot!(ts, data_pred[3,:], label="Z (NN)")

using DiffEqFlux, DifferentialEquations, Plots

ann_node = FastChain(FastDense(2, 32, tanh), FastDense(32,2, tanh))

using Pkg, Dates
 Pkg.gc(collect_delay = Day(0))


const μ = 3.986e14 * 60^2
const r = 42164e3
const mc = 500.0
const n = sqrt(μ / r^3)

#K₁ = [-28.8287 0.1005 -1449.9754 0.0046;
#      -0.087 -33.2562 0.00462 -1451.5013]
const K₁11 = -28.8287; const K₁12 = 0.1005; const K₁13 = -1449.9754; const K₁14 = 0.0046;
const K₁21 = -0.087; const K₁22 = -33.2562; const K₁23 = 0.00462; const K₁24 = -1451.5013;

#K₂ = [-288.0288 0.1312 -9614.9898 0.0;
#      -0.1312 -288.0 0.0 -9614.9883]
const K₂11 = -288.0288; const K₂12 = 0.1312; const K₂13 = -9614.9898; const K₂14 = 0.0;
const K₂21 = -0.1312; const K₂22 = -288.0; const K₂23 = 0.0; const K₂24 = -9614.9883;

# dynamics in the 'approaching' mode
@taylorize function spacecraft_approaching!(t, x, dx)

    rc = sqrt((r + x[1])^2 + x[2]^2)
    uxy1 = (K₁11 * x[1] + K₁12 * x[2]) + (K₁13 * x[3] + K₁14 * x[4])
    uxy2 = (K₁21 * x[1] + K₁22 * x[2]) + (K₁23 * x[3] + K₁24 * x[4])

    # x' = vx
    dx[1] = x[3]

    # y' = vy
    dx[2] = x[4]

    # vx' = n²x + 2n*vy + μ/(r^2) - μ/(rc^3)*(r+x) + ux/mc
    dx[3] = (n^2*x[1] + 2*(n*x[4])) + ((μ/(r^2) - μ/(rc^3)*(r + x[1])) + uxy1/mc)

    # vy' = n²y - 2n*vx - μ/(rc^3)y + uy/mc
    dx[4] = (n^2*x[2] - 2*(n*x[3])) - (μ/(rc^3)*x[2] - uxy2/mc)

    # t' = 1
    dx[5] = 1.0 + 0.0*x[1]

    return dx
end

# dynamics in the 'rendezvous attempt' mode
@taylorize function spacecraft_rendezvous_attempt!(t, x, dx)

    rc = sqrt((r + x[1])^2 + x[2]^2)
    uxy1 = (K₂11 * x[1] + K₂12 * x[2]) + (K₂13 * x[3] + K₂14 * x[4])
    uxy2 = (K₂21 * x[1] + K₂22 * x[2]) + (K₂23 * x[3] + K₂24 * x[4])

    # x' = vx
    dx[1] = x[3]

    # y' = vy
    dx[2] = x[4]

    # vx' = n²x + 2n*vy + μ/(r^2) - μ/(rc^3)*(r+x) + ux/mc
    dx[3] = (n^2*x[1] + 2*(n*x[4])) + ((μ/(r^2) - μ/(rc^3)*(r + x[1])) + uxy1/mc)

    # vy' = n²y - 2n*vx - μ/(rc^3)y + uy/mc
    dx[4] = (n^2*x[2] - 2*(n*x[3])) - (μ/(rc^3)*x[2] - uxy2/mc)

    # t' = 1
    dx[5] = 1.0 + 0.0*x[1]

    return dx
end

# dynamics in the 'aborting' mode
@taylorize function spacecraft_aborting!(t, x, dx)

    rc = sqrt((r + x[1])^2 + x[2]^2)

    # x' = vx
    dx[1] = x[3]

    # y' = vy
    dx[2] = x[4]

    # vx' = n²x + 2n*vy + μ/(r^2) + μ/(rc^3)*(r+x)
    dx[3] = (n^2*x[1] + 2*n*x[4]) + μ/(r^2) + μ/(rc^3)*(r + x[1])

    # vy' = n²y - 2n*vx - μ/(rc^3)y
    dx[4] = (n^2*x[2] - 2*n*x[3]) - μ/(rc^3)*x[2]

    # t' = 1
    dx[5] = 1.0 + 0.0*x[1]

    return dx
end


#= ======================
   LINEAR REGRESSOR (OLS)
   ====================== =#

"""
Standard linear regression model with objective function

``|Xθ - y|₂²/2``

## Parameters

$TYPEDFIELDS

$(example_docstring("LinearRegressor"))
"""
@with_kw_noshow mutable struct LinearRegressor <: MMI.Deterministic
    "whether to fit the intercept or not."
    fit_intercept::Bool    = true
    "type of solver to use (if `nothing` the default is used). The solver is
    Cholesky by default but can be Conjugate-Gradient as well. See `?Analytical`
    for more information."
    solver::Option{Solver} = nothing
end

glr(m::LinearRegressor) = LinearRegression(fit_intercept=m.fit_intercept)

descr(::Type{LinearRegressor}) = "Regression with objective function ``|Xθ - y|₂²/2``."

#= ===============
   RIDGE REGRESSOR
   =============== =#

"""
Ridge regression model with objective function

``|Xθ - y|₂²/2 + n⋅λ|θ|₂²/2``

where ``n`` is the number of samples `size(X, 1)`.
With `scale_penalty_with_samples = false` the objective function is
``|Xθ - y|₂²/2 + λ|θ|₂²/2``.

## Parameters

$TYPEDFIELDS

$(example_docstring("RidgeRegressor"))
"""
@with_kw_noshow mutable struct RidgeRegressor <: MMI.Deterministic
    "strength of the L2 regularisation."
    lambda::Real             = 1.0
    "whether to fit the intercept or not."
    fit_intercept::Bool      = true
    "whether to penalize the intercept."
    penalize_intercept::Bool = false
    "whether to scale the penalty with the number of samples."
    scale_penalty_with_samples::Bool = true
    "type of solver to use (if `nothing` the default is used). The
     solver is Cholesky by default but can be Conjugate-Gradient as
     well. See `?Analytical` for more information."
    solver::Option{Solver}   = nothing
end

glr(m::RidgeRegressor) =
    RidgeRegression(m.lambda,
                    fit_intercept=m.fit_intercept,
                    penalize_intercept=m.penalize_intercept,
                    scale_penalty_with_samples=m.scale_penalty_with_samples)

descr(::Type{RidgeRegressor}) = "Regression with objective function ``|Xθ - y|₂²/2 + λ|θ|₂²/2``."

#= ===============
   LASSO REGRESSOR
   =============== =#

"""
Lasso regression model with objective function

``|Xθ - y|₂²/2 + n⋅λ|θ|₁``

where ``n`` is the number of samples `size(X, 1)`.
With `scale_penalty_with_samples = false` the objective function is
``|Xθ - y|₂²/2 + λ|θ|₁``

## Parameters

$TYPEDFIELDS

$(example_docstring("LassoRegressor"))
"""
@with_kw_noshow mutable struct LassoRegressor <: MMI.Deterministic
    "strength of the L1 regularisation."
    lambda::Real             = 1.0
    "whether to fit the intercept or not."
    fit_intercept::Bool      = true
    "whether to penalize the intercept."
    penalize_intercept::Bool = false
    "whether to scale the penalty with the number of samples."
    scale_penalty_with_samples::Bool = true
    "type of solver to use (if `nothing` the default is used). Either `FISTA` or
    `ISTA` can be used (proximal methods, with/without acceleration)."
    solver::Option{Solver}   = nothing
end

glr(m::LassoRegressor) =
    LassoRegression(m.lambda,
                    fit_intercept=m.fit_intercept,
                    penalize_intercept=m.penalize_intercept,
                    scale_penalty_with_samples=m.scale_penalty_with_samples)

descr(::Type{LassoRegressor}) = "Regression with objective function ``|Xθ - y|₂²/2 + λ|θ|₁``."

#= =====================
   ELASTIC NET REGRESSOR
   ===================== =#

"""
Elastic net regression model with objective function

``|Xθ - y|₂²/2 + n⋅λ|θ|₂²/2 + n⋅γ|θ|₁``

where ``n`` is the number of samples `size(X, 1)`.
With `scale_penalty_with_samples = false` the objective function is
``|Xθ - y|₂²/2 + λ|θ|₂²/2 + γ|θ|₁``

## Parameters

$TYPEDFIELDS

$(example_docstring("ElasticNetRegressor"))
"""
@with_kw_noshow mutable struct ElasticNetRegressor <: MMI.Deterministic
    "strength of the L2 regularisation."
    lambda::Real             = 1.0
    "strength of the L1 regularisation."
    gamma::Real              = 0.0
    "whether to fit the intercept or not."
    fit_intercept::Bool      = true
    "whether to penalize the intercept."
    penalize_intercept::Bool = false
    "whether to scale the penalty with the number of samples."
    scale_penalty_with_samples::Bool = true
    "type of solver to use (if `nothing` the default is used). Either `FISTA` or
    `ISTA` can be used (proximal methods, with/without acceleration)."
    solver::Option{Solver}   = nothing
end

glr(m::ElasticNetRegressor) =
    ElasticNetRegression(m.lambda, m.gamma,
                         fit_intercept=m.fit_intercept,
                         penalize_intercept=m.penalize_intercept,
                         scale_penalty_with_samples=m.scale_penalty_with_samples)

descr(::Type{ElasticNetRegressor}) = "Regression with objective function ``|Xθ - y|₂²/2 + λ|θ|₂²/2 + γ|θ|₁``."

#= ==========================
   ROBUST REGRESSOR (General)
   ========================== =#

"""
Robust regression model with objective function

``∑ρ(Xθ - y) + n⋅λ|θ|₂² + n⋅γ|θ|₁``

where ``ρ`` is a robust loss function (e.g. the Huber function) and
``n`` is the number of samples `size(X, 1)`.
With `scale_penalty_with_samples = false` the objective function is
``∑ρ(Xθ - y) + λ|θ|₂² + γ|θ|₁``.

## Parameters

$TYPEDFIELDS

$(example_docstring("RobustRegressor"))
"""
@with_kw_noshow mutable struct RobustRegressor <: MMI.Deterministic
    "the type of robust loss to use (see `HuberRho`, `TalwarRho`, ...)"
    rho::RobustRho           = HuberRho(0.1)
    "strength of the regulariser if `penalty` is `:l2` or `:l1`.
    Strength of the L2 regulariser if `penalty` is `:en`."
    lambda::Real             = 1.0
    "strength of the L1 regulariser if `penalty` is `:en`."
    gamma::Real              = 0.0
    "the penalty to use, either `:l2`, `:l1`, `:en` (elastic net) or `:none`."
    penalty::SymStr          = :l2
    "whether to fit the intercept or not."
    fit_intercept::Bool      = true
    "whether to penalize the intercept."
    penalize_intercept::Bool = false
    "whether to scale the penalty with the number of samples."
    scale_penalty_with_samples::Bool = true
    "type of solver to use, default if `nothing`."
    solver::Option{Solver}   = nothing
end

glr(m::RobustRegressor) =
    RobustRegression(m.rho, m.lambda, m.gamma;
                     penalty=Symbol(m.penalty),
                     fit_intercept=m.fit_intercept,
                     penalize_intercept=m.penalize_intercept,
                     scale_penalty_with_samples=m.scale_penalty_with_samples)

descr(::Type{RobustRegressor}) = "Robust regression with objective ``∑ρ(Xθ - y) + λ|θ|₂² + γ|θ|₁`` for a given robust `ρ`."

#= ===============
   HUBER REGRESSOR
   =============== =#

"""
Huber Regression is the same as `RobustRegressor` but with the robust loss
set to `HuberRho`.

## Parameters

$TYPEDFIELDS

$(example_docstring("HuberRegressor"))
"""
@with_kw_noshow mutable struct HuberRegressor <: MMI.Deterministic
    "parametrises the `HuberRho` function (radius of the ball within which the loss
is a quadratic loss)"
    delta::Real              = 0.5
    "strength of the regulariser if `penalty` is `:l2` or `:l1`.
    Strength of the L2 regulariser if `penalty` is `:en`."
    lambda::Real             = 1.0
    "strength of the L1 regulariser if `penalty` is `:en`."
    gamma::Real              = 0.0
    "the penalty to use, either `:l2`, `:l1`, `:en` (elastic net) or `:none`."
    penalty::SymStr          = :l2
    "whether to fit the intercept or not."
    fit_intercept::Bool      = true
    "whether to penalize the intercept."
    penalize_intercept::Bool = false
    "whether to scale the penalty with the number of samples."
    scale_penalty_with_samples::Bool = true
    "type of solver to use, default if `nothing`."
    solver::Option{Solver}   = nothing
end

glr(m::HuberRegressor) =
    HuberRegression(m.delta, m.lambda, m.gamma;
                    penalty=Symbol(m.penalty),
                    fit_intercept=m.fit_intercept,
                    penalize_intercept=m.penalize_intercept,
                    scale_penalty_with_samples=m.scale_penalty_with_samples)

descr(::Type{HuberRegressor}) = "Robust regression with objective ``∑ρ(Xθ - y) + λ|θ|₂² + γ|θ|₁`` where `ρ` is the Huber Loss."

#= ==================
   QUANTILE REGRESSOR
   ================== =#

"""
Quantile Regression is the same as `RobustRegressor` but with the robust
loss set to `QuantileRho`.

## Parameters

$TYPEDFIELDS

$(example_docstring("QuantileRegressor"))
"""
@with_kw_noshow mutable struct QuantileRegressor <: MMI.Deterministic
    "parametrises the `QuantileRho` function (indicating the  quantile to use
with default `0.5` for the median regression)"
    delta::Real              = 0.5
    "strength of the regulariser if `penalty` is `:l2` or `:l1`.
    Strength of the L2 regulariser if `penalty` is `:en`."
    lambda::Real             = 1.0
    "strength of the L1 regulariser if `penalty` is `:en`."
    gamma::Real              = 0.0
    "the penalty to use, either `:l2`, `:l1`, `:en` (elastic net) or `:none`."
    penalty::SymStr          = :l2
    "whether to fit the intercept or not."
    fit_intercept::Bool      = true
    "whether to penalize the intercept."
    penalize_intercept::Bool = false
    "whether to scale the penalty with the number of samples."
    scale_penalty_with_samples::Bool = true
    "type of solver to use, default if `nothing`."
    solver::Option{Solver}   = nothing
end

glr(m::QuantileRegressor) =
    QuantileRegression(m.delta, m.lambda, m.gamma;
                       penalty=Symbol(m.penalty),
                       fit_intercept=m.fit_intercept,
                       penalize_intercept=m.penalize_intercept,
                       scale_penalty_with_samples=m.scale_penalty_with_samples)

descr(::Type{QuantileRegressor}) = "Robust regression with objective ``∑ρ(Xθ - y) + λ|θ|₂² + γ|θ|₁`` where `ρ` is the Quantile Loss."

#= ==================================
   LEAST ABSOLUTE DEVIATION REGRESSOR
   ================================== =#

"""
Least Absolute Deviation regression with with objective function

``∑ρ(Xθ - y) + n⋅λ|θ|₂² + n⋅γ|θ|₁``

where ``ρ`` is the absolute loss and
``n`` is the number of samples `size(X, 1)`.
With `scale_penalty_with_samples = false` the objective function is
``∑ρ(Xθ - y) + λ|θ|₂² + γ|θ|₁``

See also `RobustRegressor`.

## Parameters

$TYPEDFIELDS

$(example_docstring("LADRegressor"))
"""
@with_kw_noshow mutable struct LADRegressor <: MMI.Deterministic
    "strength of the regulariser if `penalty` is `:l2` or `:l1`.
    Strength of the L2 regulariser if `penalty` is `:en`."
    lambda::Real             = 1.0
    "strength of the L1 regulariser if `penalty` is `:en`."
    gamma::Real              = 0.0
    "the penalty to use, either `:l2`, `:l1`, `:en` (elastic net) or `:none`."
    penalty::SymStr          = :l2
    "whether to fit the intercept or not."
    fit_intercept::Bool      = true
    "whether to penalize the intercept."
    penalize_intercept::Bool = false
    "whether to scale the penalty with the number of samples."
    scale_penalty_with_samples::Bool = true
    "type of solver to use, default if `nothing`."
    solver::Option{Solver}   = nothing
end

glr(m::LADRegressor) =
    LADRegression(m.lambda, m.gamma;
                  penalty=Symbol(m.penalty),
                  fit_intercept=m.fit_intercept,
                  penalize_intercept=m.penalize_intercept,
                  scale_penalty_with_samples=m.scale_penalty_with_samples)

descr(::Type{LADRegressor}) = "Robust regression with objective ``∑ρ(Xθ - y) + λ|θ|₂² + γ|θ|₁`` where `ρ` is the Absolute Loss."


using StructuresKit

#Test torsion calculation for a simply-supported beam with
#twist fixed warping free ends.
#ax=1.0 in., offset uniform load

#Ix Iy Ixy J Cw
sectionProperties = [(100.0, 100.0, 0.0, 10.0, 10.0)]

#E  ν
materialProperties = [(29500, 0.30)]

#ax ay
loadLocation = [(1.0, 0.0)]

#kx kϕ
springStiffness = [(0.0, 0.0)]

#member information
#L dL SectionProperties MaterialProperties LoadLocation SpringStiffness
memberDefinitions = [(30*12,3,1,1,1,1)]

#end boundary conditions
#type=1 u''=v''=ϕ''=0 (simply supported), type=2 u'=v'=ϕ'=0  (fixed), type=3 u''=v''=ϕ''=u'''=v'''=ϕ'''=0 (free end, e.g., a cantilever)
endBoundaryConditions = [1 1]

#supports
#location where u=v=ϕ=0
supports = [0.0 30*12]

#load  qx   qy
uniformLoad = (0.0, 5.0/1000)

z, u, v, ϕ, beamProperties = Beam.solve(memberDefinitions, sectionProperties, materialProperties, loadLocation, springStiffness, endBoundaryConditions, supports, uniformLoad)

T = InternalForces.torsion(z, beamProperties.dm, ϕ, beamProperties.E, beamProperties.G, beamProperties.J, beamProperties.Cw)

maxT = maximum(T)
maxT_theory = (uniformLoad[2] * beamProperties.ax[1] * z[end])/2

Error= abs((maxT - maxT_theory))/ maxT_theory

#accept 1% error from numerical solution
@test Error <= 0.01


function test_empty(B)
    @test length(B) == 0
    @test iterate(B|42) == nothing
end

unit(n,i) = [i == j for j = 1:n]

matrix_reference(fun,x) = [fj(xi) for xi in x, fj in fun]
matrix_matrix(B,x) = Matrix(B,x)
matrix_iterate(B,x) = invoke(Matrix, Tuple{ApproxTools.AbstractBasis,AbstractVector}, B,x)
matrix_getindex(B,x) = [B[j](xi) for xi in x, j in 1:length(B)]
matrix_evaltransform_1D(B,x) = hcat((ApproxTools.evaltransform(B,x,unit(length(B),i)) for i = 1:length(B))...)
matrix_evaltransform_2D(B,x) = ApproxTools.evaltransform(B,x,Matrix(I,(length(B),length(B))))

function test_values(B,fun,args...)
    ref = matrix_reference(fun,args...)
    @test matrix_matrix(B,args...) ≈ ref
    @test matrix_iterate(B,args...) ≈ ref
    @test matrix_getindex(B,args...) ≈ ref
    @test matrix_evaltransform_1D(B,args...) ≈ ref
    @test matrix_evaltransform_2D(B,args...) ≈ ref
end

function test_inferred(B,x)
    @inferred(Matrix(B,x))
    @inferred(collect(B|x[1]))
    @inferred(B[1](x[1]))
    @inferred(ApproxTools.evaltransform(B,x,unit(length(B),1)))
    @inferred(ApproxTools.evaltransform(B,x,Matrix(I,(length(B),length(B)))))
end

function test_matrix_eval(B,x)
    ref = Matrix(B,x)
    M = Diagonal(collect(x))
    for (i,Mi) in enumerate(B|M)
        @test diag(Mi) ≈ ref[:,i]
        @test diag(B[i](M)) ≈ ref[:,i]
    end
end

function test_matvec_eval(B,x)
    ref = Matrix(B,x)
    M = Diagonal(collect(x))
    v = ones(length(x))
    for (i,vi) in enumerate(B|(M,v))
        @test vi ≈ ref[:,i]
        @test B[i]((M,v)) ≈ ref[:,i]
    end
end

function test_approximation(B, fun)
    n = length(B)
    for i = 1:n
        @test coeffs(approximate(fun[i],B)) ≈ unit(n,i)
    end
end

@testset "Monomials" begin
    fun = [one,identity,x->x^2,x->x^3,x->x^4]
    B = n->Monomials(n)

    test_empty(B(0))
    test_inferred(B(3), TrapezoidalPoints(3))
    @testset for nB = 1:length(fun)
        @testset for nx = 1:5
            test_values(B(nB), fun[1:nB], TrapezoidalPoints(nx))
            test_matrix_eval(B(nB), TrapezoidalPoints(nx))
            test_matvec_eval(B(nB), TrapezoidalPoints(nx))
        end
        test_approximation(B(nB), fun[1:nB])
    end
end

@testset "NegativeMonomials" begin
    fun = [x->x^-k for k = 1:6]
    B = n->NegativeMonomials(n)

    test_empty(B(0))
    test_inferred(B(4), TrapezoidalPoints(4))
    @testset for nB = 1:length(fun)
        @testset for nx = 2:2:6
            test_values(B(nB), fun[1:nB], TrapezoidalPoints(nx))
            test_matrix_eval(B(nB), TrapezoidalPoints(nx))
            test_matvec_eval(B(nB), TrapezoidalPoints(nx))
        end
    end
end

@testset "Chebyshev" begin
    fun = [one, identity, x->2x^2-1, x->4x^3-3x, x->8x^4-8x^2+1]
    B = n->Chebyshev(n)

    test_empty(B(0))
    test_inferred(B(3), TrapezoidalPoints(3))
    @testset for nB = 1:length(fun)
        @testset for nx = 1:5
            test_values(B(nB), fun[1:nB], TrapezoidalPoints(nx))
            test_matrix_eval(B(nB), TrapezoidalPoints(nx))
            test_matvec_eval(B(nB), TrapezoidalPoints(nx))
        end
        test_approximation(B(nB), fun[1:nB])
    end

    test_inferred(B(3), ChebyshevPoints(3))
    @testset for nB = 1:length(fun)
        @testset for nx = 1:5
            test_values(B(nB), fun[1:nB], ChebyshevPoints(nx))
        end
    end
end

@testset "Weighted" begin
    exp_approx = approximate(exp, Monomials(5))
    monomials = [one,identity,x->x^2,x->x^3,x->x^4]
    fun = [x->exp_approx(x)*p(x) for p in monomials]
    B = n->Weighted(Monomials(n),exp_approx)

    test_empty(B(0))
    test_inferred(B(3), TrapezoidalPoints(3))
    @testset for nB = 1:length(fun)
        @testset for nx = 1:5
            test_values(B(nB), fun[1:nB], TrapezoidalPoints(nx))
            test_matrix_eval(B(nB), TrapezoidalPoints(nx))
            test_matvec_eval(B(nB), TrapezoidalPoints(nx))
        end
        test_approximation(B(nB), fun[1:nB])
    end
end

@testset "Mapped" begin
    monomials = [one,identity,x->x^2,x->x^3,x->x^4]
    x2x̂ = x->x+2I
    x̂2x = x̂->x̂-2I
    fun = [x->p(x2x̂(x)) for p in monomials]
    B = n->Mapped(Monomials(n),x2x̂,x̂2x)

    test_empty(B(0))
    test_inferred(B(3), TrapezoidalPoints(3))
    @testset for nB = 1:length(fun)
        @testset for nx = 1:5
            test_values(B(nB), fun[1:nB], TrapezoidalPoints(nx))
            test_matrix_eval(B(nB), TrapezoidalPoints(nx))
            test_matvec_eval(B(nB), TrapezoidalPoints(nx))
        end
        test_approximation(B(nB), fun[1:nB])
    end
end

@testset "Poles" begin
    z = [2,3,4]
    fun = [x->1/(x - z) for z in z]
    B = n->Poles(z[1:n])

    test_empty(B(0))
    test_inferred(B(3), TrapezoidalPoints(3))
    @testset for nB = 1:length(fun)
        @testset for nx = 1:5
            test_values(B(nB), fun[1:nB], TrapezoidalPoints(nx))
            test_matrix_eval(B(nB), TrapezoidalPoints(nx))
            test_matvec_eval(B(nB), TrapezoidalPoints(nx))
        end
    end
end

@testset "Basis" begin
    fun = (exp,sin,cos)
    B = n->Basis(fun[1:n])

    test_empty(B(0))
    @testset for nB = 1:length(fun)
        @testset for nx = 1:5
            test_values(B(nB), fun[1:nB], TrapezoidalPoints(nx))
        end
    end
end

@testset "Combined" begin
    monomials = [one,identity,x->x^2]
    chebyshev = [one, identity, x->2x^2-1]

    test_empty(Combined(Monomials(0),Chebyshev(0)))
    test_inferred(Combined(Monomials(3),Chebyshev(3)), TrapezoidalPoints(3))
    @testset for nx = 1:5
        test_values(Combined(Monomials(3),Chebyshev(3)), [monomials;chebyshev;], TrapezoidalPoints(nx))
        test_values(Combined(Monomials(0),Chebyshev(3)), [          chebyshev;], TrapezoidalPoints(nx))
        test_values(Combined(Monomials(3),Chebyshev(0)), [monomials;          ], TrapezoidalPoints(nx))
    end
end

@testset "Barycentric" begin
    nodepoly = (i,x,y=Float64[]) -> (x̂ -> prod(x̂.-x[[1:i-1;i+1:length(x)]])./prod(x̂.-y))
    fun = (x,y=Float64[])->[x̂->nodepoly(i,x,y)(x̂)/nodepoly(i,x,y)(x[i]) for i = 1:length(x)]

    test_empty(Barycentric(Float64[]))
    test_inferred(Barycentric(TrapezoidalPoints(3)),TrapezoidalPoints(3))
    @testset for Points in (TrapezoidalPoints,ChebyshevPoints)
        @testset for nB = 1:3
            @testset for nx = 1:5
                x = Points(nx)
                test_values(Barycentric(x), fun(x), TrapezoidalPoints(nx))
                @testset for ny = 1:nB-1
                    y = 1 .+ (1:ny)
                    test_values(Barycentric(x,y), fun(x,y), TrapezoidalPoints(nx))
                end
            end
        end
    end
end

@testset "LanczosBasis" begin
    fun = [ x->1/sqrt(2), x->x*sqrt(3/2), x->(3x^2-1)*sqrt(5/8), x->(5x^3-3x)*sqrt(7/8) ]
    B = n->LanczosBasis(n)

    test_empty(B(0))
    test_inferred(B(3), TrapezoidalPoints(3))
    @testset for nB = 1:length(fun)
        @testset for nx = 1:4
            test_values(B(nB), fun[1:nB], TrapezoidalPoints(nx))
            # test_matrix_eval(B(nB), TrapezoidalPoints(nx))
            # test_matvec_eval(B(nB), TrapezoidalPoints(nx))
        end
    end
end


module Delsarte

using AbstractAlgebra, ClusteredLowRankSolver, BasesAndSamples, QuadGK

export delsarte, Nspherical_cap_packing

function delsarte(n, d,costheta, precision=512; all_free = false, kwargs...)
    #set up the polynomial field
    setprecision(precision)
    FF = RealField
    P, (u, ) = PolynomialRing(FF, ["u"])

    #compute both the gegenbauer polynomials and the sos basis
    gbasis = basis_gegenbauer(2d, n, u)
    samples = sample_points_chebyshev(2d,-1,costheta)
    sosbasis, samples = approximatefekete([u^k for k=0:2d], samples)

    #construct the constraint ∑_k a_k P^n_k(u) + SOS + (u+1)(cos(θ)-u)*SOS = - 1
    c = Dict()
    for k=0:2d
        c[Block(k)] = LowRankMatPol([gbasis[k+1]], [[P(1)]])
    end
    c[Block(:A)] = LowRankMatPol([P(1)], [sosbasis[1:d+1]])
    c[Block(:B)] = LowRankMatPol([(u+1)*(costheta-u)], [sosbasis[1:d]])
    constraint = Constraint(-P(1), c, Dict(), samples)

    #Construct the objective 1 + ∑_k a_k
    objective = Objective(1, Dict(Block(k) => hcat([FF(1)]) for k=0:2d), Dict())

    #Construct the SOS problem: minimize the objective s.t. the constraint
    sos = LowRankPolProblem(false, objective, [constraint])

    #Construct the SDP with or without using free variables for the a_k
    if all_free
        sdp = ClusteredLowRankSDP(sos,[k for k=0:2d])
    else
        sdp = ClusteredLowRankSDP(sos)
    end

    #Solve the SDP and return results
    return status, sol,time, errorcode = solvesdp(sdp; kwargs...)
end

function weights(theta, n)
    #approximate the integral up to a certain precision. For odd n, this can be done exactly (for even n, there is a term √(1-y^2) in the integral, so thats more difficult)
    # exact answer uses hypergeometric 2F1 functions. For n odd>=3, these terminate. Not sure about convergence at x=1.
    # exact answer of integral: [x 2F1(1/2, (3 - n)/2, 3/2, x)]_cos(theta)^1
    # Wolframalpha:  at 1 it equals sqrt(pi) Gamma(n/2-1/2) Gamma(n/2) = 1/the factor in front
    if n==3
        return 1//2*(1-cos(theta)) #odd n give an easy integral (a polynomial in the integration variable)
    elseif n==4
        return (theta-cos(theta)*abs(sin(theta)))/pi #converted to the type of theta
    elseif n==5
        return 1//2-3//4*cos(theta)+1//4 * cos(theta)^3
    end
    #for the others we just compute the integral numerically
    integral =
        sqrt(BigFloat(pi))^(-1) * gamma(BigFloat(n) / 2) / gamma(BigFloat(n - 1) / 2) *
        quadgk(y -> (1 - y^2)^((n - 3) / 2), cos(theta), BigFloat(1), rtol = 10^(-40))[1] # -30 goes relatively fast, -50 is slow
    return integral
end#

function Nspherical_cap_packing(n,d,thetas,N = length(thetas),precision=precision(BigFloat);all_free = false,solver_kwargs...)
    #Construct the polynomial ring and the gegenbauer polynomials
    setprecision(precision)
    FF = RealField
    w = weights.(thetas,n)

    P, (u, ) = PolynomialRing(FF, ["u"])
    basis = basis_gegenbauer(2d, n, u)

    constraints = []
    for i=1:N
        for j=1:i
            # sum_k a_k,i,j P_k^n(x) <= - sqrt(w(theta_i)*w(theta_j)) for x in [-1, cos(theta_i+theta_j)]
            # With sum of squares polynomials this equals:
            # sum_k a_k,i,j P_k^n(x) + SOS + (x+1)*(cos(theta_i+theta_j)-x)* SOS = - sqrt(w(theta_i)*w(theta_j))
            c = Dict()
            for k=0:2d
                if i != j
                    # We set both sides of the matrix
                    c[Block(k,i,j)] = LowRankMatPol([1//2*basis[k+1]], [[P(1)]])
                    c[Block(k,j,i)] = LowRankMatPol([1//2*basis[k+1]], [[P(1)]])
                else
                    c[Block(k,i,i)] = LowRankMatPol([basis[k+1]], [[P(1)]])
                end
            end
            c[Block((:SOS1,i,j))] = LowRankMatPol([P(1)], [[u^k for k=0:d]])
            c[Block((:SOS2,i,j))] = LowRankMatPol([(u+1)*(cos(thetas[i]+thetas[j])-u)], [[u^k for k=0:d-1]])

            samples = sample_points_chebyshev(2d,-1,cos(thetas[i]+thetas[j])) # TODO possibly take more samples
            push!(constraints, Constraint(-P(sqrt(w[i]*w[j])), c, Dict(), samples))
        end

        # The objective is M, but constrained by ∑_k a_kii ≦  M - w_i, i.e. ∑_k a_kii + pos.slack_i - M = -w_i
        obj = Dict(Block(k,i,i) => LowRankMatPol([P(1)],[[P(1)]]) for k=0:2d)
        #an easy mistake would be to use the same slack variable for every i by forgetting the i in the tuple (:SOS_obj, i).
        obj[Block((:SOS_obj,i))] = LowRankMatPol([P(1)],[[P(1)]])
        objective_constraint = Constraint(-P(w[i]),obj,Dict(:M=>P(-1)),[[0]])
        push!(constraints,objective_constraint)
    end

    #The objective is M
    objective = Objective(FF(0), Dict(), Dict(:M => FF(1)))

    #Construct the SOS - problem, minimizing the objective
    sos = LowRankPolProblem(false, objective, constraints)

    #construct and solve the SDP, with or without using free variables for the A_k
    if all_free
        sdp = ClusteredLowRankSDP(sos,[k for k=0:2d])
    else
        sdp = ClusteredLowRankSDP(sos)
    end
    status, sol, time, errorcode = solvesdp(sdp;solver_kwargs...)
end

end # end module