"""
Project density on to a staggerred grid using simple interpolation
"""
function get_rhovxI(rhoI::AbstractArray{Float64})
	rhovxI = zero(rhoI);
	get_rhovxI!(rhovxI, rhoI)
	return rhovxI
end
function get_rhovxI!(rhovxI, rhoI::AbstractArray{Float64})
	for ix = 1:size(rhoI, 2)-1
		for iz = 1:size(rhoI,1)
			rhovxI[iz, ix] = 0.5e0 *(rhoI[iz,ix+1] + rhoI[iz,ix])
		end
	end
	# last point, circular boundary conditions
	nx=size(rhoI, 2)
	for iz = 1:size(rhoI,1)
		rhovxI[iz, nx] = 0.5e0 *(rhoI[iz,1] + rhoI[iz,nx])
	end
	return nothing
end # get_rhovxI

"""
Project density on to a staggerred grid using simple interpolation
"""
function get_rhovzI(rhoI::AbstractArray{Float64})
	rhovzI = zero(rhoI);
	get_rhovzI!(rhovzI, rhoI)
	return rhovzI
end
function get_rhovzI!(rhovzI, rhoI::AbstractArray{Float64})
	nz=size(rhoI, 1)
	for ix = 1: size(rhoI, 2)
		for iz = 1: size(rhoI, 1)-1
			rhovzI[iz, ix] =  0.5e0 * (rhoI[iz+1,ix] + rhoI[iz,ix])
		end
		# last point, circular boundary conditions
		rhovzI[nz, ix] =0.5e0 * (rhoI[1,ix] + rhoI[nz,ix])
	end
	return nothing
end # get_rhovzI

function grad_modrr_sprayrr!(grad_modrr_stack,grad_modrhovxI_stack,grad_modrhovzI_stack)
	@simd for i in eachindex(grad_modrr_stack)
		@inbounds grad_modrr_stack[i] = (grad_modrhovxI_stack[i] + grad_modrhovzI_stack[i]) * (0.5)
	end
end

function grad_modrr_sprayrhovxI!(grad_modrr_stack,grad_modrhovxI_stack)
	for ix=1:size(grad_modrr_stack,2)-1
		for iz=1:size(grad_modrr_stack,1)
			@inbounds grad_modrr_stack[iz,ix+1] +=  0.5e0 * grad_modrhovxI_stack[iz,ix]
		end
	end
	nx=size(grad_modrr_stack,2)
	# circular boundary
	for iz=1:size(grad_modrr_stack,1)
		@inbounds grad_modrr_stack[iz,1] +=  0.5e0 * grad_modrhovxI_stack[iz,nx]
	end
end
function grad_modrr_sprayrhovzI!(grad_modrr_stack,grad_modrhovzI_stack)
	nz=size(grad_modrr_stack,1)
	for ix=1:size(grad_modrr_stack,2)
		for iz=1:size(grad_modrr_stack,1)-1
			@inbounds grad_modrr_stack[iz+1,ix] +=  0.5e0 * grad_modrhovzI_stack[iz,ix]
		end
		# circular boundary
		grad_modrr_stack[1,ix] +=  0.5e0 * grad_modrhovzI_stack[nz,ix]
	end
end


function getcombinations(patterns::Vector{Interval}, weights::Vector{Interval})
	N = 2length(patterns)
	inners = 2 .^ collect(0:N-1)
	outers = reverse(inners)
	variables = [patterns; weights]
	combs = Float64[]
	for (inner, outer, var) in zip(inners, outers, variables)
		endpoints = [var.left, var.right]
		c = repeat(endpoints, inner=inner, outer=outer)
		if Base.isempty(combs)
			combs = c
		else
			combs = hcat(combs, c)
		end
	end
	return combs
end

function fwa(levelendpoints::Matrix{Float64})
	n = Int(size(levelendpoints)[2] / 2)
	averages = []
	for endpoints in eachrow(levelendpoints)
		numerator = 0
		denominator = 0
		for i = 1:n
			j = i + n
			x = endpoints[i]
			w = endpoints[j]
			numerator += x * w
			denominator += w
		end
		avg = denominator == 0 ? nothing : numerator / denominator
		if !isnothing(avg)
			push!(averages, avg)
		end
	end
	return Interval(minimum(averages), maximum(averages))
end

function fuzzy_weighted_average(X⃗::FuzzyVector, W⃗::FuzzyVector)
	#FIXME: more than 2-d patterns
	# pattern_vectors = [X.grades for X in X⃗.numbers]
	# weight_vectors = [W.grades for W in W⃗.numbers]

	xx = []
	ww = []
	A₁ = X⃗.numbers[1]
	A₂ = X⃗.numbers[2]
	W₁ = W⃗.numbers[1]
	W₂ = W⃗.numbers[2]
	for i = 1:length(A₁)
		push!(xx, [A₁.grades[i], A₂.grades[i]])
		push!(ww, [W₁.grades[i], W₂.grades[i]])
	end

	combinations = FuzzySets.getcombinations.(xx, ww)
	grades = FuzzySets.fwa.(combinations)
	return FuzzySets.FuzzyNumber(A₁.levels, grades)
end


using LeastSquaresOptim, Base.Test

# simple factor model
# only problem with "real" optimization
# nice example because J'J is not invertible
# but cholfact in sparse handles this case
function factor()
    name = "factor"
    function f!(x, fvec)
        fvec[1] = 3.0 - x[1] * x[4]
        fvec[2] = 2.0 - x[1] * x[5]
        fvec[3] = 5.0 - x[1] * x[6]

        fvec[4] = 4.5 - x[2] * x[4]
        fvec[5] = 3.2 - x[2] * x[5]
        fvec[6] = 2.0 - x[2] * x[6]

        fvec[7] = 5.0 - x[3] * x[4]
        fvec[8] = 1.3 - x[3] * x[5]
        fvec[9] = 1.5 - x[3] * x[6]
    end

    function g!(x, J)
        fill!(J, 0.0)
        J[1, 1] = -x[4]
        J[1, 4] = -x[1]
        J[2, 1] = -x[5]
        J[2, 5] = -x[1]
        J[3, 1] = -x[6]
        J[3, 6] = -x[1]

        J[4, 2] = -x[4]
        J[4, 4] = -x[2]
        J[5, 2] = -x[5]
        J[5, 5] = -x[2]
        J[6, 2] = -x[6]
        J[6, 6] = -x[2]

        J[7, 3] = -x[4]
        J[7, 4] = -x[3]
        J[8, 3] = -x[5]
        J[8, 5] = -x[3]
        J[9, 3] = -x[6]
        J[9, 6] = -x[3]
    end

    function g!(x, J::SparseMatrixCSC)
        Jvals = nonzeros(J)
        i = 0
        i += 1
        Jvals[i] = -x[4]
        i += 1
        Jvals[i] = -x[1]
        i += 1
        Jvals[i] = -x[5]
        i += 1
        Jvals[i] = -x[1]
        i += 1
        Jvals[i] = -x[6]
        i += 1
        Jvals[i] = -x[1]
        i += 1
        Jvals[i] = -x[4]
        i += 1
        Jvals[i] = -x[2]
        i += 1
        Jvals[i] = -x[5]
        i += 1
        Jvals[i] = -x[2]
        i += 1
        Jvals[i] = -x[6]
        i += 1
        Jvals[i] = -x[2]
        i += 1
        Jvals[i] = -x[4]
        i += 1
        Jvals[i] = -x[3]
        i += 1
        Jvals[i] = -x[5]
        i += 1
        Jvals[i] = -x[3]
        i += 1
        Jvals[i] = -x[6]
        i += 1
        Jvals[i] = -x[3]
    end
    x = ones(6)
    return name, f!, g!, x
end

iter = 0
for matrix in (:dense, :sparse)
    for (optimizer, optimizer_abbr) in ((LeastSquaresOptim.Dogleg(), :dl), (LeastSquaresOptim.LevenbergMarquardt(), :lm))
        factorization = matrix == :dense ? LeastSquaresOptim.QR() : LeastSquaresOptim.Cholesky()
        for (solver, solver_abbr) in ((factorization, :fact), (LeastSquaresOptim.LSMR(), :iter))
            iter += 1
            if matrix == :sparse && optimizer == LeastSquaresOptim.LevenbergMarquardt()
                continue
            else
                name, f!, g!, x = factor()
                fcur = ones(9)
                J = ones(9, 6)
                g!(x, J)
                if matrix == :sparse
                    J = sparse(J)
                end
                nls = LeastSquaresProblem(x = x, y = fcur, f! = f!, J = J, g! = g!)
                r = optimize!(nls, optimizer, solver)
                if iter == 1
                    show(r)
                end
                @printf("%-6s %4s %2s %30s %5d %5d   %5d   %10e\n", matrix, solver_abbr, optimizer_abbr, name, r.iterations, r.f_calls, r.g_calls, r.ssr)
                @test r.ssr <= 12
                @test r.converged
                end
            end
    end
end


using transD_GP.GP, Random, PyPlot, Statistics, LinearAlgebra, PositiveFactorizations
##
Random.seed!(12)
f = 1.0/0.5
t = 0:(100*f)^-1:2/f
ntrain = 1
ntries = 4
λ = [0.1]
δ = 0.01
δtry = 5δ
p = 2
demean = true
findmin = true
switchflag = findmin ? -1. : 1.
knownvalue = NaN
acqfun = GP.EI()

y = 100000*switchflag*(2t+ sin.(2*pi*f*t) + sinc.(2*f*(t.-middle(t))))
ynoisy = y + δ*randn(length(t))

tperm = randperm(length(t))
trainidx = tperm[1:ntrain]
ytrain = ynoisy[trainidx]
ttrain = t[trainidx]

ytest, σ2, σ_prior  = GP.GPfit(GP.SqEuclidean(), ytrain', ttrain', t', λ.^2, δtry, p=p, demean=demean)
AF = GP.getAF(acqfun, ytrain, vec(ytest), diag(σ2), findmin=findmin, knownvalue=knownvalue)
f, ax = plt.subplots(2, ntries+1, figsize=(12,5), sharex="col", sharey="row")
ax[1].plot(t,y, "-k")
ax[1].plot(t,ytest)
ax[1].plot(ttrain,ytrain, ".m",markersize=20)
ax[1].plot(t,ytest+2sqrt.(diag(σ2)),"--r",alpha=0.2)
ax[1].plot(t,ytest-2sqrt.(diag(σ2)),"--b",alpha=0.2)
ax[1].grid()
ax[2].plot(t, AF)
ax[2].grid()
ax[1].set_title("step 0")
##
for i = 1:ntries
    global AF
    nextpos = argmax(AF)
    push!(ytrain, ynoisy[nextpos])
    push!(ttrain, t[nextpos])
    ytest, σ2, σ_prior  = GP.GPfit(GP.SqEuclidean(), ytrain', ttrain', t', λ.^2, δtry, p=p, demean=demean)
    AF = GP.getAF(acqfun, ytrain, vec(ytest), diag(σ2), findmin=findmin, knownvalue=knownvalue)
    ax[2i+1].plot(t,y, "-k")
    ax[2i+1].plot(ttrain,ytrain, ".m",markersize=20)
    ax[2i+1].plot(t,ytest)
    ax[2i+1].plot(t,ytest+2sqrt.(diag(σ2)),"--r",alpha=0.2)
    ax[2i+1].plot(t,ytest-2sqrt.(diag(σ2)),"--b",alpha=0.2)
    ax[2i+1].grid()
    ax[2(i+1)].plot(t, AF)
    ax[2(i+1)].grid()
    ax[2i+1].set_title("step $i")
end
# savefig("fig.png",dpi=300)

##
figure()
s1=subplot(211, aspect="auto")
imshow(σ_prior, extent=[t[1],t[end],t[end],t[1]])
s2=subplot(212, sharex=s1, sharey=s1)
imshow(σ2, extent=[t[1],t[end],t[end],t[1]])

figure()
plot(t,y, "-k",label="true")
plot(ttrain,ytrain, "+m", label="train",markersize=10)
nsamp = length(t)
nreal=50
realizations = randn(nreal,nsamp)*cholesky(Positive, var(ytest)*σ2 + 1e-12*Matrix(I,nsamp,nsamp)).U
realizations = ytest .+ realizations'
plot(t,realizations,alpha=0.2)


export armijo, wolfe, armijo_wolfe, shamanskii_stop, goldstein

"""
armijo: check if a step size is admissible according to the Armijo criterion.

Armijo criterion: f(x + θd) - f(x) < τ₀ θ ∇f(x+θd)d

Note: ht, h₀ and g₀ are required in the LSAtT
"""
function armijo(h      :: Any, #LineModel
                h_at_t :: LSAtT;
                τ₀     :: Float64 = 0.01,
                kwargs...)

    if (h_at_t.ht == nothing) || (h_at_t.h₀ == nothing) || (h_at_t.g₀ == nothing)
     #printstyled("Warning: Nothing entries in the State.\n", color = :red)
     return throw(error("Nothing entries in the State. ht, h₀ and g₀ are mandatory."))
    else

     fact = -0.8
     Eps = 1e-10
     hgoal = h_at_t.ht - h_at_t.h₀ - h_at_t.g₀ * h_at_t.x * τ₀
     # Armijo = (h_at_t.ht <= hgoal)# || ((h_at_t.ht <= h_at_t.h₀ + Eps * abs(h_at_t.h₀)) & (h_at_t.gt <= fact * h_at_t.g₀))
	 # Armijo_HZ =
	 # positive = h_at_t.x > 0.0   # positive step
     return max(hgoal, 0.0)
    end
end

"""
wolfe: check if a step size is admissible according to the Wolfe criterion.

Strong Wolfe criterion: |∇f(x+θd)| < τ₁||∇f(x)||.

Note: gt and g₀ are required in the LSAtT
"""
function wolfe(h      :: Any, #LineModel,
               h_at_t :: LSAtT;
               τ₁     :: Float64 = 0.99,
               kwargs...)

    if (h_at_t.g₀ == nothing) || (h_at_t.gt == nothing)
     return throw(error("Nothing entries in the State."))
    else

     #wolfe = (τ₁ .* h_at_t.g₀) - (abs(h_at_t.gt))
     wolfe = abs(h_at_t.gt) - τ₁ * abs(h_at_t.g₀)
	 #positive = h_at_t.x > 0.0   # positive step
     return max(wolfe, 0.0)
    end
end

"""
armijo_wolfe: check if a step size is admissible according to the Armijo and Wolfe criteria.

Note: ht, h₀, gt and g₀ are required in the LSAtT
"""
function armijo_wolfe(h      :: Any, #LineModel,
                      h_at_t :: LSAtT;
                      τ₀     :: Float64 = 0.01,
                      τ₁     :: Float64 = 0.99,
                      kwargs...)

   if (h_at_t.ht == nothing) || (h_at_t.gt == nothing) || (h_at_t.h₀ == nothing) || (h_at_t.g₀ == nothing)
    #printstyled("Warning: Nothing entries in the State.\n", color = :red)
    return throw(error("Nothing entries in the State. ht, h₀, gt and g₀ are mandatory."))
   else
    wolfe  = abs(h_at_t.gt) - τ₁ * abs(h_at_t.g₀)
    armijo = h_at_t.ht - h_at_t.h₀ - h_at_t.g₀ * h_at_t.x * τ₀
    return max(armijo, wolfe, 0.0)
   end
end

"""
shamanskii_stop: check if a step size is admissible according to the "Shamanskii" criteria.
This criteria was proposed in:
Lampariello, F., & Sciandrone, M. (2001).
Global convergence technique for the Newton method with periodic Hessian evaluation.
Journal of optimization theory and applications, 111(2), 341-358.

Note: * h.d accessible (specific LineModel)
      * ht, h₀ are required in the LSAtT
"""
function shamanskii_stop(h      :: Any, #LineModel,
                         h_at_t :: LSAtT;
                         γ      :: Float64 = 1.0e-09,
                         kwargs...)

    admissible = h_at_t.ht - h_at_t.h₀ - γ * (h_at_t.x)^3 * norm(h.d)^3
    return max(admissible, 0.0)
end

"""
goldstein: check if a step size is admissible according to the Goldstein criteria.

Note: ht, h₀ and g₀ are required in the LSAtT
"""
function goldstein(h      :: Any, #LineModel,
                   h_at_t :: LSAtT;
                   τ₀     :: Float64 = 0.0001,
                   τ₁     :: Float64 = 0.9999,
                   kwargs...)

 	goldstein = max(h_at_t.h₀ + h_at_t.x * (1 - τ₀) * h_at_t.g₀ - h_at_t.ht,
                    h_at_t.ht - (h_at_t.h₀ + h_at_t.x *  τ₀ * h_at_t.g₀))
 	# positive = h_at_t.x > 0.0   # positive step
 	return max(goldstein, 0.0) #&& positive
end


immutable Dirichlet <: ContinuousMultivariateDistribution
    alpha::Vector{Float64}
    alpha0::Float64

    function Dirichlet{T <: Real}(alpha::Vector{T})
        alpha0::Float64 = 0.0
        for i in 1:length(alpha)
            if alpha[i] < 0.0
                error("Dirichlet: elements of alpha must be non-negative")
            else
                alpha0 += alpha[i]
            end
        end
        new(float64(alpha), alpha0)
    end
end

Dirichlet(d::Integer, alpha::Real) = Dirichlet(fill(alpha, d))

Dirichlet(dim::Integer) = Dirichlet(ones(dim))

function entropy(d::Dirichlet)
    k = length(d.alpha)
    en = lmnB(d)
    en += (d.alpha0 - k) * digamma(d.alpha0)
    for j in 1:k
        en -= (d.alpha[j] - 1.0) * digamma(d.alpha[j])
    end
    return en
end

function insupport{T <: Real}(d::Dirichlet, x::Vector{T})
    n = length(x)
    if length(d.alpha) != n
        return false
    end
    s = 0.0
    for i in 1:n
        if x[i] < 0.0
            return false
        end
        s += x[i]
    end
    if abs(s - 1.0) > 1e-8
        return false
    end
    return true
end

mean(d::Dirichlet) = d.alpha ./ d.alpha0

function modes(d::Dirichlet)
    k = length(d.alpha)
    x = Array(Float64, k)
    s = d.alpha0 - k
    for i in 1:k
        if d.alpha[i] <= 1.0
            error("modes only defined when alpha[i] > 1 for all i")
        end
        x[i] = (d.alpha[i] - 1.0) / s
    end
    return [x]
end

pdf{T <: Real}(d::Dirichlet, x::Vector{T}) = exp(logpdf(d, x))

function logpdf{T <: Real}(d::Dirichlet, x::Vector{T})
    return dot((d.alpha - 1.0), log(x)) - lmnB(d)
end

function logpdf!{T <: Real}(r::AbstractArray, d::Dirichlet, x::Matrix{T})
    if size(x, 1) != length(d.alpha)
        throw(ArgumentError("Inconsistent argument dimensions."))
    end

    n = size(x, 2)
    if length(r) != n
        throw(ArgumentError("Inconsistent argument dimensions."))
    end
    b::Float64 = lmnB(d)
    At_mul_B(r, log(x), d.alpha - 1.0)
    for i in 1:n
        r[i] -= b
    end
    return
end

function rand!(d::Dirichlet, x::Vector)
    s = 0.0
    n = length(x)
    for i in 1:n
        tmp = rand(Gamma(d.alpha[i]))
        x[i] = tmp
        s += tmp
    end
    for i in 1:n
        x[i] /= s
    end
    return x
end

function rand(d::Dirichlet)
    x = Array(Float64, length(d.alpha))
    return rand!(d, x)
end

function rand!(d::Dirichlet, X::Matrix)
    m, n = size(X)
    for j in 1:n
        for i in 1:m
            X[i, j] = rand(Gamma(d.alpha[i]))
        end
    end
    for j in 1:n
        isum = 0.0
        for i in 1:m
            isum += X[i, j]
        end
        for i in 1:m
            X[i, j] /= isum
        end
    end
    return X
end

function var(d::Dirichlet)
    n = length(d.alpha)
    tmp = d.alpha0^2 * (d.alpha0 + 1.0)
    S = Array(Float64, n, n)
    for j in 1:n
        for i in 1:n
            if i == j
                S[i, j] = d.alpha[i] * (d.alpha0 - d.alpha[i]) / tmp
            else
                S[i, j] = -d.alpha[i] * d.alpha[j] / tmp
            end
        end
    end
    return S
end

# Log multinomial beta
function lmnB(d::Dirichlet)
    s = 0.0
    for i in 1:length(d.alpha)
        s += lgamma(d.alpha[i])
    end
    return s - lgamma(d.alpha0)
end

mnB(d::Dirichlet) = exp(lmnB(d))


randuniform(dims::Base.Dims) = randuniform(Float64, dims)
randuniform(::Type{T}, dims::Base.Dims) where {T<:Number} = rand(T, dims)

randnormal(dims::Base.Dims) = randnormal(Float64, dims)
randnormal(::Type{T}, dims::Base.Dims) where {T<:Number} = randn(T, dims)

randisometry(dims::Base.Dims{2}) = randisometry(Float64, dims)
randisometry(::Type{T}, dims::Base.Dims{2}) where {T<:Number} = _leftorth!(randnormal(T, dims), QRpos())[1]


# Recreate plot with CMF vs. orbital period:

include("../../../src/regress.jl")
using Statistics
using PyPlot
using DelimitedFiles

period = [1.510826, 2.421937,4.049219, 6.101013, 9.207540,12.352446,18.772866]

dir = "../../../data/POSTERIOR_CMF/"
fname = ["T1b_IRON_FRAC.out","T1c_IRON_FRAC.out","T1d_IRON_FRAC.out","T1e_IRON_FRAC.out","T1f_IRON_FRAC.out","T1g_IRON_FRAC.out","T1h_IRON_FRAC.out"]

np = 7; nsamp = 10000

cmf_norm = zeros(np,nsamp)
cmf_mean = zeros(np)
cmf_median = zeros(np)
cmf_sig = zeros(np)
cmf_sig1 = zeros(np)
cmf_sig2 = zeros(np)
for ip=1:np
  data = readdlm(string(dir,fname[ip]))
  cmf_norm[ip,:] .= data[:,2]
  cmf_mean[ip] = mean(cmf_norm[ip,:])
  cmf_median[ip] = median(cmf_norm[ip,:])
  cmf_sig[ip] = std(cmf_norm[ip,:])
  cmf_sig1[ip] = cmf_median[ip]-sort(cmf_norm[ip,:])[1587]
  cmf_sig2[ip] = sort(cmf_norm[ip,:])[8413]-cmf_median[ip]
end

# Now, carry out regressions for each of these:
coeff_samp = zeros(2,nsamp)
fn = zeros(2,np)
fn[1,:] .= 1.0
fn[2,:] .= period ./6
sig = ones(np)
cmf_avg = 0.0
for isamp=1:nsamp
  global cmf_avg += mean(cmf_norm[:,isamp])
  coeff,cov = regress(fn,cmf_norm[:,isamp],sig)
  coeff_samp[:,isamp] = coeff
end

cmf_avg /= nsamp

# Now, compute standard deviation:
cmf_var = 0.0
for isamp=1:nsamp
  global cmf_var += (mean(cmf_norm[:,isamp]) - cmf_avg)^2
end
cmf_rms = sqrt(cmf_var/nsamp)

# Now, draw 1-sigma and 2-sigma confidence intervals:
period_trial = collect(0.0:0.1:20.0)
conf = zeros(5,length(period_trial))
for (i,p) in enumerate(period_trial)
  # Predict samples at this period
  csamp = coeff_samp[1,:] .+ coeff_samp[2,:] .* p /6
  # Now, sort these and pick out 68.3% & 95% confidence limits:
  csamp = sort(csamp)
  conf[:,i] = [csamp[250],csamp[1415],csamp[5000],csamp[8415],csamp[9750]]
end
# Now, plot median, 1-sigma and 2-sigma confidence intervals:

clf()
#plot(period_trial,conf[1,:],color="C1",alpha=0.5,linestyle="--")
errorbar(period,cmf_mean,yerr=cmf_sig,fmt="o")
#errorbar(period,cmf_median,yerr=[cmf_sig1';cmf_sig2'],fmt="o")
#plot(period_trial,conf[5,:],color="C1",alpha=0.5,linestyle="--")
plot(period_trial,conf[3,:],color="C1",linewidth=2)
plot(period_trial,conf[2,:],color="C1",linewidth=2,linestyle=":")
plot(period_trial,conf[4,:],color="C1",linewidth=2,linestyle=":")
#plot(period_trial,cmf_avg  .+ 0.0 .* period_trial,color="C0",linewidth=2)
fill_between([0,20],[cmf_avg-cmf_rms,cmf_avg-cmf_rms],[cmf_avg+cmf_rms,cmf_avg+cmf_rms],alpha=0.2,color="C0")
xlabel("Orbital Period [d]")
ylabel("Iron mass fraction (%)")
# Overplot solar system bodies:
text([16],[28],"Mars")
plot([0,20],[30.0,30.0],"--",color="k")
text([1],[8],"Moon")
plot([0,20],[10.0,10.0],"--",color="k")
text([18],[33],"Earth")
plot([0,20],[32.0,32.0],"--",color="k")
#text([17],[31.5],"Venus")
#plot([0,20],[31.,31.],"--",color="k")
axis([0,20,0,40])
text([15],[20],"Chondrites (L/CI)")
#plot([0,20],[22.0,22.0],"--",color="k")
plot([0,20],[18.7,18.7],"--",color="k")
text([1],[20.6],"b")
text([2.7],[22.6],"c")
text([3.5],[13],"d")
text([6.3],[20],"e")
text([8.7],[15],"f")
text([11.8],[12],"g")
text([18.1],[11.8],"h")

tight_layout()
savefig("../Norm_cmf_vs_period.pdf",bbox_inches = "tight")

println("Coefficients of fit: ",mean(coeff_samp[1,:]),"+-",std(coeff_samp[1,:])," ",mean(coeff_samp[2,:]),"+-",std(coeff_samp[2,:]))


import IterTools: groupby, imap
using Zygote: @nograd

struct BSplineBasis <: Basis1D
    knots::Vector{Float64}
    order::Int
    deriv::Int

    function BSplineBasis(knots, order, deriv=0; extend=true)
        if deriv >= order throw(ArgumentError("Differentiation order not supported")) end

        for (kn, kp) in zip(knots[2:end], knots[1:end-1])
            if kn < kp throw(ArgumentError("Knot vector must be nondecreasing")) end
        end

        if extend
            knots = [fill(knots[1], order-1); knots; fill(knots[end], order-1)]
        else
            for d in 1:order-1
                if knots[1+d] != knots[1] || knots[end-d] != knots[end]
                    throw(ArgumentError("Expected $order repeated knots on either end"))
                end
            end
        end

        new(knots, order, deriv)
    end

    BSplineBasis(lft::Real, rgt::Real, elements::Int, order::Int) =
        BSplineBasis(range(lft, stop=rgt, length=elements+1), order)
end

const BSpline = BasisFunction1D{BSplineBasis}

Base.length(b::BSplineBasis) = length(b.knots) - b.order
Base.size(b::BSplineBasis) = (length(b),)
Base.getindex(b::BSplineBasis, i) = BSpline(b, i, b.deriv)
nderivs(b::BSplineBasis) = b.order - 1

domain(b::BSplineBasis) = Interval(b.knots[1], b.knots[end])
domain(b::BSpline) = Interval(b.basis.knots[b.index], b.basis.knots[b.index+b.basis.order])

degree(b::BSplineBasis) = b.order - b.deriv - 1
degree(b::BSpline) = b.basis.order - b.deriv - 1

deriv(b::BSplineBasis) = BSplineBasis(b.knots, b.order, b.deriv+1; extend=false)
deriv(b::BSplineBasis, order) = BSplineBasis(b.knots, b.order, b.deriv+order; extend=false)

function supported(b::BSplineBasis, pt::T) where {T<:Real}
    kidx = b.order - 1 + searchsorted(b.knots[b.order:end], pt).stop
    stop = b.knots[kidx] == b.knots[end] ? kidx - b.order : kidx
    stop - b.order + 1 : stop
end

function supported(b::BSplineBasis, pts::Vector{T}) where {T<:Real}
    (min, max) = extrema(pts)
    @assert(min in domain(b) && max in domain(b), "Domain error")

    idxs = zeros(Int, length(pts))

    if !issorted(pts)
        for (i, pt) in enumerate(pts)
            idxs[i] = supported(b, pt).stop
        end
    else
        kidx = b.order
        for (i, pt) in enumerate(pts)
            kidx += searchsorted(b.knots[kidx:end], pt).stop - 1
            idxs[i] = b.knots[kidx] == b.knots[end] ? kidx - b.order : kidx
        end
    end

    imap(groupby(i -> i[2], enumerate(idxs))) do i
        (pts[i[1][1]:i[end][1]], i[1][2] - b.order + 1 : i[1][2])
    end
end

@nograd function evaluate_raw(b::BSplineBasis, pts::Vector{T}, deriv::Int, rng::UnitRange{Int}) where {T<:Real}
    # Basis values of order 1 (piecewise constants)
    bvals = zeros(Float64, (b.order, length(pts)))
    bvals[end,:] .= 1.0

    p = b.order
    bi = rng.start + p

    # Order increment
    for k in 0:p-deriv-2
        bvals[p-k:end, :] ./= b.knots[bi:bi+k] - b.knots[bi-k-1:bi-1]

        for (i, kp, kn) in zip(p-k-1:p-1, b.knots[bi-k-2:bi-2], b.knots[bi:bi+k])
            bvals[i,:] .*= (pts .- kp)
            bvals[i,:] .+= bvals[i+1,:] .* (kn .- pts)
        end
        bvals[end,:] .*= (pts .- b.knots[bi-1])
    end

    # Differentiation
    for k = p-deriv-1:p-2
        bvals[p-k:end, :] ./= b.knots[bi:bi+k] - b.knots[bi-k-1:bi-1]

        bvals[1:end-1,:] = -diff(bvals, dims=1)
        bvals .*= k + 1
    end

    bvals
end


using Test, SpecialFunctions, Statistics
import TruncatedNormal as TN

@test isnan(TN.tnvar(1, 0))
@test isnan(TN.tnvar(+Inf, 0))
@test isnan(TN.tnvar(0, -Inf))
@test isnan(TN.tnvar(NaN, NaN))
@test isnan(TN.tnvar(0, NaN))
@test isnan(TN.tnvar(NaN, 0))

@test TN.tnvar(-Inf, +Inf) == 1
@test TN.tnvar(+Inf, +Inf) == 0
@test TN.tnvar(-Inf, -Inf) == 0
@test TN.tnvar(0, +Inf) ≈ 1 - 2/π
@test TN.tnvar(-Inf, 0) ≈ 1 - 2/π

for x = -10:10
    @test TN.tnvar(x, +Inf) ≈ 1 + √(2/π) * x / erfcx(+x / √2) - 2/π / erfcx(+x / √2)^2
    @test TN.tnvar(-Inf, x) ≈ 1 - √(2/π) * x / erfcx(-x / √2) - 2/π / erfcx(-x / √2)^2
end

@test TN.tnvar(50, 70) ≈ 0.0003990431868038995479099272265360593305365

@test TN.tnvar(-Inf, Inf, 0, 1) == 1

for x = -10:10
    @test TN.tnvar(x, +Inf, 0, 1) ≈ TN.tnvar(x, +Inf)
    @test TN.tnvar(-Inf, x, 0, 1) ≈ TN.tnvar(-Inf, x)
end

@test TN.tnvar(50, 70, -2, 3) ≈ 0.029373438107168350377591231295634273607812172191712
@test TN.tnvar(-100.0, 0.0, 0.0, 2.0986317998643735) ≈ 1.6004193412141677189841357987638847137391508803335
@test TN.tnvar(0.0, 0.9, 0.0, 0.07132755843183151) ≈ 0.0018487407287725028827020557707636415445504260892486
@test TN.tnvar(-100.0, 100.0, 0.0, 17.185261847875548) ≈ 295.333163899557735486302841237124507431445
@test TN.tnvar(-100.0, 0.5, 0.0, 0.47383322897860064) ≈ 0.145041095812679283837328561547251019229612
@test TN.tnvar(-100.0, 100.0, 0.0, 17.185261847875548) ≈ 295.333163899557735486302841237124507431445
@test TN.tnvar(-10000, 10000, 0, 1) ≈ 1

# https://github.com/JuliaStats/Distributions.jl/issues/827
@test TN.tnvar(999000, 1e6) ≥ 0
@test TN.tnvar(0, 1000, 1000000, 1) ≥ 0

@test_broken TN.tnvar(1e6, Inf) ≈ 9.99999999994000000000050000000e-13
@test_broken TN.tnvar(999000, 1e6) ≈ 1.00200300399898194688784897455e-12
@test_broken TN.tnvar(0, 1000, 1000000, 1) ≈ 1.00200300399898194688784897455909681215272616880064850605775930845141e-12


using Test
using ArnoldiMethod: reflector!, restore_arnoldi!, Reflector, Arnoldi, reinitialize!, iterate_arnoldi!
using LinearAlgebra
using Random

@testset "Reflector $T" for T in (Float64, ComplexF64)
    n = 20
    x = rand(T, n)
    z = copy(x)
    τ = reflector!(z, n)
    z[n] = 1
    y = x .- τ * dot(z, x) .* z

    @inferred reflector!(copy(x), n)
    @test norm(y[1:n-1]) ≤ 10eps()
    @test abs(real(y[n])) ≈ norm(x)
    @test abs(imag(y[n])) ≤ eps()
    @test 1 ≤ real(τ) ≤ 2
    @test abs(τ - 1) ≤ 1
end

@testset "Trivial reflector $T" for T in (Float64, ComplexF64)
    x = [T(0), T(0), T(5)]
    z = copy(x)
    τ = reflector!(z, 3)
    @test iszero(τ)
end

@testset "rmul! $T" for T in (Float64, ComplexF64)
    A = rand(T, 4, 4)
    B = copy(A)

    # Implicit representation
    G = Reflector{T}(4)
    rand!(G.vec)
    τ = reflector!(G, 4)

    # Matrix representation
    z = [G.vec[1:3]; 1]
    H = I - τ * z * z'

    # Apply
    rmul!(A, G, 1, 4)

    @test A ≈ B * H'
end

@testset "lmul! $T" for T in (Float64, ComplexF64)
    A = rand(T, 4, 4)
    B = copy(A)

    # Implicit representation
    G = Reflector{T}(4)
    rand!(G.vec)
    τ = reflector!(G, 4)

    # Matrix representation
    z = [G.vec[1:3]; 1]
    H = I - τ * z * z'

    # Apply
    lmul!(G, A, 1, 4)

    @test A ≈ H * B
end

@testset "Restore Hessenberg matrix $T" for T in (Float64, ComplexF64)
    n, k = 10, 6
    Q = qr(randn(T, k, k)).Q * Matrix(1.0I, k, k)

    A = rand(T, n, n)
    arn = Arnoldi{T}(n, k)
    reinitialize!(arn)
    iterate_arnoldi!(A, arn, 1:k)

    # Do an orthogonal change of basis.

    H = copy(arn.H)
    H[1:k,1:k] .= Q' * H[1:k,1:k] * Q

    # Restore the relation
    restore_arnoldi!(H, 1, k, Q, Reflector{T}(k))

    # Only now do the change of basis in V with accumulated Q.
    W = [arn.V[:, 1:k] * Q arn.V[:, k+1]]
    @test norm(A * W[:, 1:k] - W[:, 1:k+1] * H) < 1e-14
end


using DiffEqBase: bisection
using Test

# basic tests
# result of bisection should never be zero,
# should have sign same as left limit
# and nexfloat in tdir should be zero or same sign of right limit
r = bisection(sin, (3.0, 4.0), 1.0)
@test !iszero(sin(r))
@test sin(r) > 0
@test sin(nextfloat(r)) <= 0

r = bisection(sin, (4.0, 3.0), -1.0)
@test !iszero(sin(r))
@test sin(r) < 0
@test sin(prevfloat(r)) >= 0

# this example has a big zero interval, makes it go inside the inner loop
for i in 1:10
	f = (x) -> (10.0^(-i))*(x^3)
	r = bisection(f, (-1.0, 1.0), 1.0)
	@test !iszero(f(r))
	@test f(r) < 0
	@test f(nextfloat(r)) >= 0

	r = bisection(f, (1.0, -1.0), -1.0)
	@test !iszero(f(r))
	@test f(r) > 0
	@test f(prevfloat(r)) <= 0
end


function BlockMeans(X::Array{Float64}; alpha::Float64 = 1/2, block_size::Int = -1)
  # estimate the asymptotic variance for a stationary time series using block means
  # see v.d. Vaart, Time Series, Section 5.1
  # alpha determines block size through l = [n^alpha]
  # alpha is overruled by block_size
  if ndims(X) == 1
    n = length(X)
    X = reshape(X,(1,n))
  end
  (dimension,n) = size(X)
  if block_size < 0
    block_size = Int(floor(n^alpha))
  end
  Z = Array{Float64,2}(undef,dimension,n-block_size+1) # block means
  for i = 1:(n-block_size+1)
    Z[:,i] = mean(X[:,i:(i+block_size-1)], dims = 2)
  end
  mu = mean(X,dims=2)
  v = vec(block_size/(n-block_size+1) * sum((Z-repeat(mu,1, n-block_size+1)).^2, dims=2))
  #F(x) = 1/(n-l+1) * sum(sqrt(l)*(Z-mu).<=x)
  # return (v,F)
end

function BatchMeans(X::Matrix{Float64}; alpha::Float64 = 1/2, n_batches::Int = -1, true_mean::Vector{Float64} = Vector{Float64}(undef,0))
  # estimate the asymptotic variance for a stationary time series using batch means

  # if ndims(X) == 1
  #   n = length(X)
  #   X = reshape(X,(1,n))
  # end
  (dimension,n) = size(X)
  if (n_batches <= 0)
    batch_size = Int(floor(n^alpha))
    n_batches = Int(floor(n / n^alpha))
  else
    batch_size = Int(floor(n/n_batches))
  end
  if n_batches > n/2
    return Inf
  end
  Z = Array{Float64,2}(undef,dimension,n_batches) # batch means
  unused = n - batch_size * n_batches
  for i=1:n_batches
    Z[:,i] = mean(X[:,(unused + (i-1)*batch_size + 1):(unused + i * batch_size)], dims=2)
  end
  if length(true_mean) == 0
    vec(var(Z, dims=2)*batch_size)
  else
    sum((Z - repeat(true_mean, 1, n_batches)).^2, dims = 2)/n_batches * batch_size
  end
end

function ESS(X::AbstractArray{Float64};alpha::Float64=1/2, true_variance::Vector{Float64} = Vector{Float64}(undef,0), true_mean::Vector{Float64} = Vector{Float64}(undef,0), n_batches::Int = -1)

  # if ndims(X) == 1
  #   n = length(X)
  #   X = reshape(X,(1,n))
  # end
  asvar = vec(BatchMeans(X,alpha=alpha,n_batches=n_batches,true_mean=true_mean))
  if length(true_variance) == 0
    variance = vec(var(X,dims=2))
  else
    variance = true_variance
  end
  return size(X)[2] * variance ./ asvar
end

function SquaredJumpingDistance(X::AbstractArray{Float64})

  if ndims(X) == 1
    n_samples = length(X)
    differences = X[2:end] - X[1:end-1]
  else
    n_samples = size(X)[2]
    differences = X[:,2:end] - X[:,1:end-1]
  end
  return sumabs2(differences)/(n_samples-1)

end

function PooledVariance(AsVar::Vector{Float64}, n_samples::Vector{Int}, SampleVariance::Vector{Float64})

  return sum(n_samples .* SampleVariance ./ AsVar)/sum(n_samples ./AsVar)

end


"""
    OlsSureFn(Y,X,m=0)

LS of Y on X; for one n dependent variables, same regressors

# Usage
(b,res,Yhat,Covb,R2) = OlsSureFn(Y,X,m)

# Input
- `Y::Array`:     Txn, the n dependent variables
- `X::Array`:     Txk matrix of regressors (including deterministic ones)
- `m::Int`:       scalar, bandwidth in Newey-West

# Output
- `b::Array`:     n*kx1, regression coefficients
- `u::Array`:     Txn, residuals Y - Yhat
- `Yhat::Array`:  Txn, fitted values X*b
- `V::Array`:     matrix, covariance matrix of vec(b)
- `R2::Number`:  n vector, R2 value

"""
function OlsSureFn(Y,X,m=0)

    (T,n) = (size(Y,1),size(Y,2))
    k     = size(X,2)

    b     = X\Y
    Yhat  = X*b
    u     = Y - Yhat

    g     = zeros(T,n*k)
    for i = 1:n
      vv      = (1+(i-1)*k):(i*k)   #1:k,(1+k):2k,...
      g[:,vv] = X.*u[:,i]           #moment conditions for Y[:,i] regression
    end

    S0    = NWFn(g,m)            #Newey-West covariance matrix
    Sxx   = X'X/T
    Sxx_1 = kron(I(n),inv(Sxx))
    V     = Sxx_1 * S0 * Sxx_1/T

    R2   = 1 .- var(u,dims=1)./var(Y,dims=1)

    return b, u, Yhat, V, R2

end


function DiffusionDefinition.diffusion_message(::Val{:Lorenz})
    println("* * * * * * *")
    println("Lorenz system")
    println("* * * * * * *")
    println("Description")
    println("------------")
    println("A three-dimensional elliptic diffusion")
    println("dXₜ = θ₁(Yₜ-Xₜ)dt + σdW¹ₜ,")
    println("dYₜ = [Xₜ(θ₂ - Zₜ)-Yₜ]dt + σdW²ₜ,")
    println("dZₜ = (XₜYₜ - θ₃Zₜ)dt + σdW³ₜ.")
    println("The model has four parameters:")
    println("\tθ1, θ2, θ3, σ")
    println("\nConstructors")
    println("------------")
    println("You may initialize the process either by passing the parameters")
    println("in the order above, as in:")
    println("θ₁, θ₂, θ₃, σ = ...")
    println("P = Lorenz(θ₁, θ₂, θ₃, σ)")
    println("or by keyword arguments")
    println("p = (θ1=..., θ2=..., θ3=..., σ=...)")
    println("P = Lorenz(;p...)")
    println()
end

@diffusion_process Lorenz{T} begin
    :dimensions
    process --> 3
    wiener --> 3

    :parameters
    _ --> (3, T)
    σ --> T

    :additional
    constdiff --> true
    sparsediff --> true
end

function DiffusionDefinition.b(t, x, P::Lorenz)
    @SVector [
        P.p1*(x[2]-x[1]),
        P.p2*x[1] - x[2] - x[1]*x[3],
        x[1]*x[2] - P.p3*x[3]
    ]
end

DiffusionDefinition.σ(t, x, P::Lorenz) = SDiagonal(P.σ, P.σ, P.σ)


"""
    RBEstimator <: Estimator

Unbiased estimator for response rate `p` see also:

> Kunzmann K, Kieser M. Point estimation and p‐values in phase II adaptive two‐stage designs with a binary endpoint. Statistics in medicine. 2017 Mar 15;36(6):971-84.
"""
struct RBEstimator{TD<:Design} <: Estimator

  design::TD

  RBEstimator{TD}(design::TD) where {TD<:Design} = new(design)

end # RBEstimator

"""
    RBEstimator(design::TD) where {TD<:Design}

Create unbiased estimator for response rate `p` see also:

> Kunzmann K, Kieser M. Point estimation and p‐values in phase II adaptive two‐stage designs with a binary endpoint. Statistics in medicine. 2017 Mar 15;36(6):971-84.

"""
RBEstimator(design::TD) where {TD<:Design} = RBEstimator{TD}(design)

Base.show(io::IO, estimator::RBEstimator) = print("RBEstimator")

function estimate(estimator::RBEstimator, x1::T, x2::T) where {T<:Integer}

  checkx1x2(x1, x2, design(estimator))
  n1          = interimsamplesize(design(estimator))
  nominator   = BigInt(0.0) # we need big ints, as binomials get HUGE!
  denominator = BigInt(0.0)
  x = x1 + x2
  n = samplesize(design(estimator), x1)
  for xx1 in 0:(min(x, n1))
    if samplesize(design(estimator), xx1) == n
      nominator   += BigInt(binomial(n1 - 1, xx1 - 1))*BigInt(binomial(n - n1, x - xx1))
      denominator += BigInt(binomial(n1, xx1))*BigInt(binomial(n - n1, x - xx1))
    end
  end
  return convert(Float64, nominator/denominator)

end


using Regression
using Base.Test

srand(54321)  # ensure repeatable experiment
              # avoid CI issues

## auxiliary functions

_relerr(x::Array, x0::Array) = sumabs2(x - x0) / sumabs2(x0)
_corrdist(x::Array, x0::Array) = 1.0 - dot(vec(x), vec(x0)) / (vecnorm(x) * vecnorm(x0))

_classify(U::Matrix) = vec(mapslices(indmax, U, 1))

function _errrate(pred::PredictionModel, X::Matrix, θ::Array, θ0::Array)
    u0 = _classify(predict(pred, θ0, X))
    u  = _classify(predict(pred, θ, X))
    countnz(u .!= u0) / length(u0)
end

const _solvers = [GD(), BFGS(), LBFGS(6)]

function verify_solver(title, loss::Loss, data,
                       θg::Array, vcond::Function, thres::Real)
    println("    $title")

    X = data[1]
    dx = size(X, 1)
    dθ = size(θg, ndims(θg))
    b = dθ == dx ? 0.0 :
        dθ == dx + 1 ? 1.0 :
        error("Unmatched dimensions.")

    pb = ndims(θg) == 1 ? UnivariateRegression(loss, data...; bias=b) :
                          MultivariateRegression(loss, data...; bias=b)

    for solv in _solvers
        println("      - with solver $(solv)")
        n = size(X, ndims(X))
        ret = Regression.solve(pb;
                               reg=SqrL2Reg(1.0e-4),
                               solver=solv,
                               options=Regression.Options(grtol=n * 1.0e-6))

        @test isa(ret, Regression.Solution)
        θe = ret.sol
        @test vcond(θe, θg) < thres
    end
end

## Univariate loss

# data
d = 5
n = 2000
w = randn(d)
wa = [w; randn()]

X = randn(d, n)
u = X'w + 0.01 * randn(n)
ua = u .+ wa[d+1]

# test cases

verify_solver("LinearPred + SqrLoss",
    SqrLoss(), (X, u), w, _relerr, 1.0e-4)

verify_solver("AffinePred + SqrLoss",
    SqrLoss(), (X, ua), wa, _relerr, 1.0e-4)

verify_solver("LinearPred + LogisticLoss",
    SqrLoss(), (X, u), w, _corrdist, 1.0e-3)

verify_solver("AffinePred + LogisticLoss",
    SqrLoss(), (X, ua), wa, _corrdist, 1.0e-3)

## Multivariate loss

# data
d = 5
k = 3
W = randn(k, d)
Wa = [W randn(k)]

X = randn(d, n)
U = W * X + 0.01 * randn(k, n)
Ua = U .+ Wa[:,d+1]

# test cases

verify_solver("MvLinearPred + SumSqrLoss",
    SumSqrLoss(), (X, U, k), W, _relerr, 1.0e-4)

verify_solver("MvAffinePred + SumSqrLoss",
    SumSqrLoss(), (X, Ua, k), Wa, _relerr, 1.0e-4)

y = _classify(U)
verify_solver("MvLinearPred + MultiLogisticLoss",
    MultiLogisticLoss(), (X, y, k), W,
        (θ, θ0) -> _errrate(MvLinearPred(d, k), X, θ, θ0), 0.02)

ya = _classify(Ua)
verify_solver("MvAffinePred + MultiLogisticLoss",
    MultiLogisticLoss(), (X, ya, k), Wa,
        (θ, θ0) -> _errrate(MvAffinePred(d, k), X, θ, θ0), 0.02)


#=
This code is loosely based on a routine originally copyright Federal Reserve Bank of Atlanta
and written by Iskander Karibzhanov.
=#

"""
```
kalman_filter(data, lead, a, F, b, H, var, z0, vz0, Ny0; allout=false)
kalman_filter(data, lead, a, F, b, H, var, Ny0=0; allout=false)
```

Inputs
------

- `data`: a `Ny x T` matrix containing data `y(1), ... , y(T)`.
- `lead`: the number of steps to forecast after the end of the data.
- `a`: an `Nz x 1` vector for a time-invariant input vector in the transition equation.
- `F`: an `Nz x Nz` matrix for a time-invariant transition matrix in the transition
  equation.
- `b`: an `Ny x 1` vector for a time-invariant input vector in the measurement equation.
- `H`: an `Ny x Nz` matrix for a time-invariant measurement matrix in the measurement
  equation.
- `var`: an `Ny + Nz` x `Ny + Nz` matrix for a time-invariant variance matrix for the
  error in the transition equation and the error in the measurement equation, that is,
  `[η(t)', ϵ(t)']'`.
- `z0`: an optional `Nz x 1` initial state vector.
- `vz0`: an optional `Nz x Nz` covariance matrix of an initial state vector.
- `Ny0`: an optional scalar indicating the number of periods of presample (i.e. the number
  of periods which we don't add to the likelihood)
- `allout`: an optional keyword argument indicating whether we want optional output
  variables returned as well

Outputs
-------

- `logl`: value of the average log likelihood function of the SSM under assumption that
  observation noise ϵ(t) is normally distributed
- `pred`: a `Nz` x `T+lead` matrix containing one-step predicted state vectors.
- `vpred`: a `Nz` x `Nz` x `T+lead` matrix containing mean square errors of predicted
  state vectors.
- `filt`: an optional `Nz` x `T` matrix containing filtered state vectors.
- `vfilt`: an optional `Nz` x `Nz` x `T` matrix containing mean square errors of filtered
  state vectors.

Notes
-----
The state space model is defined as follows:
```
z(t+1) = a+F*z(t)+η(t)     (state or transition equation)
y(t) = b+H*z(t)+ϵ(t)       (observation or measurement equation)
```

When `z0` and `Vz0` are omitted, the initial state vector and its covariance matrix of the
time invariant Kalman filters are computed under the stationarity condition:
```
z0 = (I-F)\a
vz0 = (I-kron(F,F))\(V(:),Nz,Nz)
```
where `F` and `V` are the time invariant transition matrix and the covariance matrix of
transition equation noise, and `vec(V)` is an `Nz^2` x `1` column vector that is constructed
by stacking the `Nz` columns of `V`.  Note that all eigenvalues of `F` are inside the unit
circle when the state space model is stationary.  When the preceding formula cannot be
applied, the initial state vector estimate is set to `a` and its covariance matrix is given
by `1E6I`.  Optionally, you can specify initial values.
"""
function kalman_filter{S<:AbstractFloat}(data::Matrix{S},
                                      lead::Int64,
                                      a::Matrix{S},
                                      F::Matrix{S},
                                      b::Matrix{S},
                                      H::Matrix{S},
                                      var::Matrix{S},
                                      z0::Matrix{S},
                                      vz0::Matrix{S},
                                      Ny0::Int = 0;
                                      allout::Bool = false)
    T = size(data, 2)
    Nz = size(a, 1)
    Ny = size(b, 1)

    z = z0
    P = vz0

    # Check input matrix dimensions
    @assert size(data, 1) == Ny
    @assert size(a, 2) == 1
    @assert size(F) == (Nz, Nz)
    @assert size(b, 2) == 1
    @assert size(H) == (Ny, Nz)
    @assert size(var) == (Ny + Nz, Ny + Nz)
    @assert size(z) == (Nz, 1)
    @assert size(P) == (Nz, Nz)

    # V(t) and R(t) are variances of η(t) and ϵ(t), respectively, and G(t) is a covariance
    # of η(t) and ϵ(t)
    # In dsgelh :
    # --- V is same as QQ
    # --- R is same as EE
    # --- G is same as VV = QQ*MM
    V = var[1:Nz, 1:Nz]
    R = var[(Nz+1):end, (Nz+1):end]
    G = var[1:Nz, (Nz+1):end]

    if allout
        pred          = zeros(S, Nz, T)
        vpred         = zeros(S, Nz, Nz, T)
        yprederror    = NaN*zeros(S, Ny, T)
        ystdprederror = NaN*zeros(S, Ny, T)
        filt          = zeros(Nz, T)
        vfilt         = zeros(Nz, Nz, T)
    end

    L = zero(S)

    for t = 1:T
        # If an element of the vector y(t) is missing (NaN) for the observation t, the
        #   corresponding row is ditched from the measurement equation.
        nonmissing = !isnan(data[:, t])
        data_t = data[nonmissing, t]       # data_t = Y_T = [y1, y2, ..., yT] is matrix of observable data time-series
        H_t = H[nonmissing, :]             # H_t = DD is matrix mapping states to observables
        G_t = G[:, nonmissing]             # G_t = Cov(η_t, ϵ_t)
        R_t = R[nonmissing, nonmissing]    # R_t = Var(ϵ_t)
        Ny_t = length(data_t)              # Ny_t = T is length of time
        b_t = b[nonmissing, :]             # b_t = DD

        ## forecasting
        z = a + F*z                        # z_{t|t-1} = a + F(Θ)*z_{t-1|t-1}
        P = F*P*F' + V                     # P_{t|t-1} = F(Θ)*P_{t-1|t-1}*F(Θ)' + F(Θ)*Var(η_t)*F(Θ)'
        dy = data_t - H_t*z - b_t          # dy = y_t - H(Θ)*z_{t|t-1} - DD is prediction error or innovation
        HG = H_t*G_t                       # HG is ZZ*Cov(η_t, ϵ_t)
        D = H_t*P*H_t' + HG + HG' + R_t    # D = ZZ*P_{t+t-1}*ZZ' + HG + HG' + R_t
        D = (D+D')/2

        if allout
            pred[:, t]                   = z
            vpred[:, :, t]               = P
            yprederror[nonmissing, t]    = dy
            ystdprederror[nonmissing, t] = dy./sqrt(diag(D))
        end

        ddy = D\dy
        # We evaluate the log likelihood function by adding values of L at every iteration
        #   step (for each t = 1,2,...T)
        if t > Ny0
            L += -log(det(D))/2 - first(dy'*ddy/2) - Ny_t*log(2*pi)/2
        end

        ## updating
        PHG = P*H_t' + G_t
        z = z + PHG*ddy                    # z_{t|t} = z_{t|t-1} + P_{t|t-1}*H(Θ)' + ...
        P = P - PHG/D*PHG'                 # P_{t|t} = P_{t|t-1} - PHG*(1/D)*PHG

        if allout
            PH = P*H_t'
            filt[:, t]     = z
            vfilt[:, :, t] = P
        end
    end

    zend = z
    Pend = P

    if allout && lead > 1
        for t = (T+2):(T+lead)
            z = F*z + a
            P = F*P*F' + V
            pred[:, t]     = z
            vpred[:, :, t] = P
        end
    end

    if allout
        rmse = sqrt(mean((yprederror.^2)', 1))
        rmsd = sqrt(mean((ystdprederror.^2)', 1))
        return Kalman(L, zend, Pend, pred, vpred, yprederror, ystdprederror, rmse, rmsd, filt, vfilt)
    else
        return Kalman(L, zend, Pend)
    end

end

function kalman_filter{S<:AbstractFloat}(data::Matrix{S},
                                      lead::Int64,
                                      a::Matrix{S},
                                      F::Matrix{S},
                                      b::Matrix{S},
                                      H::Matrix{S},
                                      var::Matrix{S},
                                      Ny0::Int = 0;
                                      allout::Bool = false)
    Nz = size(a, 1)
    V = var[1:Nz, 1:Nz]

    e, _ = eig(F)
    if countnz(e*e' - eye(Nz)) == Nz^2
        z0 = (eye(Nz) - F)\a
        vz0 = reshape((eye(Nz^2)-kron(F,F))\V, Nz, Nz)
    else
        z0 = a
        vz0 = eye(Nz)*1e6
    end

    return kalman_filter(data, lead, a, F, b, H, var, z0, vz0, Ny0; allout=allout)
end

immutable Kalman{S<:AbstractFloat}
    L::S
    zend::Matrix{S}
    Pend::Matrix{S}
    pred::Matrix{S}
    vpred::Array{S,3}
    yprederror::Matrix{S}
    ystdprederror::Matrix{S}
    rmse::Matrix{S}
    rmsd::Matrix{S}
    filt::Matrix{S}
    vfilt::Array{S,3}
end
function Kalman{S<:AbstractFloat}(L::S,
                                  zend::Matrix{S},
                                  Pend::Matrix{S},
                                  pred::Matrix{S}          = Matrix{S}(),
                                  vpred::Array{S,3}        = Array{S}(0,0,0),
                                  yprederror::Matrix{S}    = Matrix{S}(),
                                  ystdprederror::Matrix{S} = Matrix{S}(),
                                  rmse::Matrix{S}          = Matrix{S}(),
                                  rmsd::Matrix{S}          = Matrix{S}(),
                                  filt::Matrix{S}          = Matrix{S}(),
                                  vfilt::Array{S,3}        = Array{S}(0,0,0))
    return Kalman{S}(L,zend,Pend,pred,vpred,yprederror,ystdprederror,rmse,rmsd,filt,vfilt)
end
function Base.getindex(K::Kalman, d::Symbol)
    if d in (:L, :zend, :Pend, :pred, :vpred, :yprederror, :ystdprederror, :rmse, :rmsd,
             :filt, :vfilt)
        return getfield(K, d)
    else
        throw(KeyError(d))
    end
end


using Complementarity
using Test

@info("-------[Testing Complementarity/NLsolve]------------------------------------------")

@testset "mcp test 0 with NLsolve example" begin

    lb = [0., 0., 0., 0.]
    ub = [Inf, Inf, Inf, Inf]
    x0 = [1.25, 0., 0., 0.5]

    m = nothing
    m = MCPModel()

    @variable(m, x[1:4] >= 0)
    @mapping(m, F1, 3*x[1]^2+2*x[1]*x[2]+2*x[2]^2+x[3]+3*x[4]-6)
    @mapping(m, F2, 2*x[1]^2+x[1]+x[2]^2+3*x[3]+2*x[4]-2)
    @mapping(m, F3, 3*x[1]^2+x[1]*x[2]+2*x[2]^2+2*x[3]+3*x[4]-1)
    @mapping(m, F4, x[1]^2+3*x[2]^2+2*x[3]+3*x[4]-3)
    @complementarity(m, [F1 F2 F3 F4], x)

    set_start_value.(x, x0)

    status = solveMCP(m, solver=:NLsolve)
    @show status

    z = result_value.(x)
    # Fz = [result_value(F1) result_value(F2) result_value(F3) result_value(F4)]

    @show z
    # @show Fz

    @test isapprox(z[1], 1.22474, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], -1.378e-19, atol=1e-4)
    @test isapprox(z[4], 0.5, atol=1e-4)

    # @test isapprox(Fz[1], -1.26298e-9, atol=1e-4)
    # @test isapprox(Fz[2], 3.22474, atol=1e-4)
    # @test isapprox(Fz[3], 5.0, atol=1e-4)
    # @test isapprox(Fz[4], 3.62723e-11, atol=1e-4)

    set_start_value.(x, x0)
    status = solveMCP(m, solver=:PATH)
    @show status

    z = result_value.(x)
    # Fz = [result_value(F1) result_value(F2) result_value(F3) result_value(F4)]

    @show z
    # @show Fz

    @test isapprox(z[1], 1.22474, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], -1.378e-19, atol=1e-4)
    @test isapprox(z[4], 0.5, atol=1e-4)

    # @test isapprox(Fz[1], -1.26298e-9, atol=1e-4)
    # @test isapprox(Fz[2], 3.22474, atol=1e-4)
    # @test isapprox(Fz[3], 5.0, atol=1e-4)
    # @test isapprox(Fz[4], 3.62723e-11, atol=1e-4)

    set_start_value.(x, x0)
    status = solveMCP(m, solver=:NLsolve)
    @show status

    z = result_value.(x)
    # Fz = [result_value(F1) result_value(F2) result_value(F3) result_value(F4)]

    @show z
    # @show Fz

    @test isapprox(z[1], 1.22474, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], -1.378e-19, atol=1e-4)
    @test isapprox(z[4], 0.5, atol=1e-4)

    # @test isapprox(Fz[1], -1.26298e-9, atol=1e-4)
    # @test isapprox(Fz[2], 3.22474, atol=1e-4)
    # @test isapprox(Fz[3], 5.0, atol=1e-4)
    # @test isapprox(Fz[4], 3.62723e-11, atol=1e-4)

end

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

@testset "mcp test 1 with NLsolve" begin
    m = nothing
    m = MCPModel()

    M = [0  0 -1 -1 ;
         0  0  1 -2 ;
         1 -1  2 -2 ;
         1  2 -2  4 ]

    q = [2; 2; -2; -6]

    lb = zeros(4)
    ub = Inf*ones(4)

    items = 1:4

    # @variable(m, lb[i] <= x[i in items] <= ub[i])
    @variable(m, x[i in items] >= 0)
    @mapping(m, F[i in items], sum(M[i,j]*x[j] for j in items) + q[i])
    @complementarity(m, F, x)

    status = solveMCP(m, solver=:NLsolve)
    @show status

    z = result_value.(x)
    # Fz = result_value(F)

    @show z
    # @show Fz

    @test isapprox(z[1], 2.8, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], 0.8, atol=1e-4)
    @test isapprox(z[4], 1.2, atol=1e-4)
end

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

@testset "mcp test 2 with NLsolve" begin

    m = nothing
    m = MCPModel()

    M = [0  0 -1 -1 ;
         0  0  1 -2 ;
         1 -1  2 -2 ;
         1  2 -2  4 ]

    q = [2; 2; -2; -6]

    lb = zeros(4)
    ub = Inf*ones(4)

    items = 1:4

    @variable(m, lb[i] <= x[i in items] <= ub[i])
    # @variable(m, x[i in items] >= 0)
    @mapping(m, F[i in items], sum(M[i,j]*x[j] for j in items) + q[i])
    @complementarity(m, F, x)

    status = solveMCP(m, solver=:NLsolve)
    @show status

    z = result_value.(x)
    # Fz = result_value(F) # currently produces an error

    @show z
    # @show Fz

    @test isapprox(z[1], 2.8, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], 0.8, atol=1e-4)
    @test isapprox(z[4], 1.2, atol=1e-4)
end

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

@testset "mcp test 3 with NLsolve" begin

    m = nothing
    m = MCPModel()

    M = [0  0 -1 -1 ;
         0  0  1 -2 ;
         1 -1  2 -2 ;
         1  2 -2  4 ]

    q = [2; 2; -2; -6]

    lb = zeros(4)
    ub = Inf*ones(4)

    @variable(m, lb[i] <= myvariablename[i in 1:4] <= ub[i])
    @mapping(m, myconst[i=1:4], sum(M[i,j]*myvariablename[j] for j in 1:4) + q[i])
    @complementarity(m, myconst, myvariablename)

    status = solveMCP(m, solver=:NLsolve)
    @show status

    z = result_value.(myvariablename)
    # Fz = result_value(myconst)

    @show z
    # @show Fz

    @test isapprox(z[1], 2.8, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], 0.8, atol=1e-4)
    @test isapprox(z[4], 1.2, atol=1e-4)

end

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

@testset "mcp test 4 with NLsolve" begin

    m = nothing
    m = MCPModel()

    lb = zeros(4)
    ub = Inf*ones(4)
    items = 1:4
    @variable(m, lb[i] <= x[i in items] <= ub[i])

    @mapping(m, F1, 3*x[1]^2 + 2*x[1]*x[2] + 2*x[2]^2 + x[3] + 3*x[4] -6)
    @mapping(m, F2, 2*x[1]^2 + x[1] + x[2]^2 + 3*x[3] + 2*x[4] -2)
    @mapping(m, F3, 3*x[1]^2 + x[1]*x[2] + 2*x[2]^2 + 2*x[3] + 3*x[4] -1)
    @mapping(m, F4, x[1]^2 + 3*x[2]^2 + 2*x[3] + 3*x[4] - 3)

    @complementarity(m, F1, x[1])
    @complementarity(m, F2, x[2])
    @complementarity(m, F3, x[3])
    @complementarity(m, F4, x[4])

    set_start_value(x[1], 1.25)
    set_start_value(x[2], 0.)
    set_start_value(x[3], 0.)
    set_start_value(x[4], 0.5)

    status = solveMCP(m, solver=:NLsolve)
    @show status

    z = result_value.(x)
    # Fz = [result_value(F1), result_value(F2), result_value(F3), result_value(F4)]

    @show z
    # @show Fz

    @test isapprox(z[1], 1.22474, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], 0.0, atol=1e-4)
    @test isapprox(z[4], 0.5, atol=1e-4)

    set_start_value(x[1], 1.19)
    set_start_value(x[2], 0.2)
    set_start_value(x[3], 0.1)
    set_start_value(x[4], 0.49)

    status = solveMCP(m, solver=:PATH)
    @show status

    z = result_value.(x)
    # Fz = [result_value(F1), result_value(F2), result_value(F3), result_value(F4)]

    @show z
    # @show Fz

    @test isapprox(z[1], 1.22474, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], 0.0, atol=1e-4)
    @test isapprox(z[4], 0.5, atol=1e-4)

    set_start_value(x[1], 1.40)
    set_start_value(x[2], 0.1)
    set_start_value(x[3], 0.2)
    set_start_value(x[4], 0.43)

    status = solveMCP(m, solver=:NLsolve)
    @show status

    z = result_value.(x)
    # Fz = [result_value(F1), result_value(F2), result_value(F3), result_value(F4)]

    @show z
    # @show Fz

    @test isapprox(z[1], 1.22474, atol=1e-4)
    @test isapprox(z[2], 0.0, atol=1e-4)
    @test isapprox(z[3], 0.0, atol=1e-4)
    @test isapprox(z[4], 0.5, atol=1e-4)
end


###############
## VALUATION ##
###############

"""
    Valuation(x::PVector; affine = true)
    Valuation(x::AbstractVector)

A data structure for approximating the valuation of a solution path ``x(s)``.
It does this by computing an approximation function [`ν`](@ref) and approximating
the first and second derivatives ``ν̇`` and ``ν̈``.
If `affine = true` then the computed valuation is the valuation pulled back
to the affine space.
The valuation is estimated continously along a solution path. For this it is assumed that
the path is tracked in a **logarithmic time scale**.
"""
struct Valuation
    # Estimate
    ν::Vector{Float64}
    ν̇::Vector{Float64}
    ν̈::Vector{Float64}
    s::Base.RefValue{Float64}
    affine::Bool
    # Data for computing the ν̇ and ν̈ by finite differences
    ν_data::NTuple{2,Vector{Float64}}
    s_data::Base.RefValue{NTuple{2,Float64}}
end

function Valuation(n::Integer, affine::Bool = true)
    v = zeros(n)
    ν̇ = zeros(n)
    ν̈ = zeros(n)
    ν_data = (zeros(n), zeros(n))
    s_data = (NaN, NaN)
    Valuation(v, ν̇, ν̈, Ref(NaN), affine, ν_data, Ref(s_data))
end
function Valuation(x::ProjectiveVectors.PVector; affine::Bool = true)
    if affine
        Valuation(length(x) - length(ProjectiveVectors.dims(x)), affine)
    else
        Valuation(length(x), affine)
    end
end
Valuation(x::Vector; affine::Bool = true) = Valuation(length(x))

function Base.show(io::IO, val::Valuation)
    println(io, typeof(val), ":")
    println(io, " • s → ", round.(val.s[]; sigdigits = 4))
    for name in [:ν, :ν̇, :ν̈]
        println(io, " • ", name, " → ", round.(getfield(val, name); sigdigits = 4))
    end
end
Base.show(io::IO, ::MIME"application/prs.juno.inline", v::Valuation) = v

"""
    init!(val::Valuation)

Initialize the valuation estimator to start from scratch.
"""
function init!(val::Valuation)
    val.ν .= NaN
    val.ν̇ .= NaN
    val.ν̈ .= NaN
    val.s[] = NaN
    val.s_data[] = (NaN, NaN)
    val.ν_data[1] .= NaN
    val.ν_data[2] .= NaN
    val
end

#############################
## Valuation approximation ##
#############################

# First to the coordinate wise functions

"""
    ν(z, ż)

Computes the function ``ν(z) = -(xẋ + yẏ) / |z|²`` where `x,y = reim(z)`.
This is equivalent to computing ``d/ds log|x(s)|``.
"""
function ν(z::Complex, ż::Complex)
    x, y = reim(z)
    ẋ, ẏ = reim(ż)
    -(x * ẋ + y * ẏ) / abs2(z)
end

"""
    ν_ν̇_ν̈(z::Complex, ż::Complex, z̈::Complex, z³::Complex)

Computes the function [`ν`](@ref) and the derivaties ``ν̇``, ``ν̈`` by using the analytic
derivatives.
"""
function ν_ν̇_ν̈(z::Complex, ż::Complex, z̈::Complex, z³::Complex)
    x, y = reim(z)
    ẋ, ẏ = reim(ż)
    ẍ, ÿ = reim(z̈)
    x³, y³ = reim(z³)

    μ = x * ẋ + y * ẏ
    # The following is just applying product rule properly
    μ̇ = x * ẍ + ẋ^2 + y * ÿ + ẏ^2
    μ̈ = x * x³ + 3 * ẋ * ẍ + y * y³ + 3 * ẏ * ÿ

    z² = abs2(z)
    ν = -μ / z²
    # ν̇ = -μ̇ / z² + 2μ²/(z²)² and  μ²/(z²)² = ν^2
    ν̇ = 2 * ν^2 - μ̇ / z²
    ν̈ = 4 * ν * ν̇ - μ̈ / z² + 2 * μ̇ * μ / z²^2
    ν, ν̇, ν̈
end

"""
    ν̇_ν̈(z::Complex, ż::Complex, s::Float64, ν₂, s₂, ν₁, s₁)

Computes the function [`ν`](@ref) and the derivaties ``ν̇``, ``ν̈`` by using a finite
difference scheme. This uses the value `ν₂` of ``ν`` at `s₂` and `ν₁` of ``ν`` at `s₁` with
``s > s₂ > s₁``.
Since we have a non-uniform grid, we need a more elaborate difference scheme.
The implementation follows the formulas derived in [^BS05].

[^BS05]: Bowen, M. K., and Ronald Smith. "Derivative formulae and errors for non-uniformly
  spaced points." Proceedings of the Royal Society A: Mathematical, Physical and Engineering
  Sciences 461.2059 (2005): 1975-1997.
"""
function ν̇_ν̈(ν::Float64, s::Float64, ν₂::Float64, s₂::Float64, ν₁::Float64, s₁::Float64)
    Δ₁, Δ₂, Δ₁₂ = s - s₁, s - s₂, s₁ - s₂
    ν̇ = (Δ₂ * ν₁) / (Δ₁₂ * Δ₁) - ((Δ₁₂ + Δ₂) * ν₂) / (Δ₁₂ * Δ₂) - (Δ₁₂ * ν) / (Δ₁ * Δ₂)
    ν̈ = -2ν₁ / (Δ₁₂ * Δ₁) + 2ν₂ / (Δ₁₂ * Δ₂) + 2ν / (Δ₁ * Δ₂)
    ν̇, ν̈
end

"""
    update!(val::Valuation, x, ẋ, s,
            predictor::AbstractPredictorCache = NullPredictorCache())

Compute new approximations of the valuation of the path ``x(s)`` from `x` and `ẋ` at `s`.
This queries the given predictor cache `predictor` for the second and third derivative
of the path ``x(s)``. If this information is available it computes the derivatives
of ``ν`` analytically otherwise a finite differene scheme is used.
"""
function update!(
    val::Valuation,
    z::AbstractVector,
    ż,
    s,
    predictor::AbstractPredictorCache = NullPredictorCache(),
)
    @unpack ν̇, ν̈, ν_data, s_data = val

    ν₂, ν₁ = ν_data
    s₂, s₁ = s_data[]

    if val.affine && isa(z, PVector)
        k = 1
        for (rⱼ, j) in ProjectiveVectors.dimension_indices_homvars(z)
            νⱼ = ν(z[j], ż[j])
            for i in rⱼ
                ν_k = ν(z[i], ż[i]) - νⱼ
                if !isnan(s₁)
                    ν̇_k, ν̈_k = ν̇_ν̈(ν_k, s, ν₂[k], s₂, ν₁[k], s₁)
                    ν̇[k], ν̈[k] = ν̇_k, ν̈_k
                end
                # update ν and shift the previous ones
                val.ν[k], ν₂[k], ν₁[k] = ν_k, val.ν[k], ν₂[k]
                k += 1
            end
        end
    else
        for i = 1:length(val.ν)
            # try to compute the values using higher derivatives from the predictor
            z̈ = second_derivative(predictor)
            z³ = third_derivative(predictor)
            if !isnothing(z̈) && !isnothing(z³)
                νᵢ, ν̇ᵢ, ν̈ᵢ = ν_ν̇_ν̈(z[i], ż[i], z̈[i], z³[i])
                if !isnan(s₁)
                    # The estimates can be sensitive to numerical errors, so we still compare
                    # against the numerical derivatives since this is cheap anyway
                    ν̇ᵢ′, ν̈ᵢ′ = ν̇_ν̈(νᵢ, s, ν₂[i], s₂, ν₁[i], s₁)
                    ν̇[i] = abs(ν̇ᵢ) < abs(ν̇ᵢ′) ? ν̇ᵢ : ν̇ᵢ′
                    ν̈[i] = abs(ν̈ᵢ) < abs(ν̈ᵢ′) ? ν̈ᵢ : ν̈ᵢ′
                else
                    ν̇[i], ν̈[i] = ν̇ᵢ, ν̈ᵢ
                end
            else
                νᵢ = ν(z[i], ż[i])
                if !isnan(s₁)
                    ν̇ᵢ, ν̈ᵢ = ν̇_ν̈(νᵢ, s, ν₂[i], s₂, ν₁[i], s₁)
                    ν̇[i], ν̈[i] = ν̇ᵢ, ν̈ᵢ
                end
            end
            # update ν and shift the previous ones
            val.ν[i], ν₂[i], ν₁[i] = νᵢ, val.ν[i], ν₂[i]
        end
    end

    # Shift the s
    s_data[] = (val.s[], s₂)
    val.s[] = s
    val
end

"""
    enum ValuationVerdict

The possible states the [`judge`](@ref) function returns:

* `VAL_INDECISIVE`: The estimates are not trustworthy enough.
* `VAL_FINITE`: The valuation indicates that the path is finite.
* `VAL_AT_INFINITY`: The valuation indicates that the path is diverging.
"""
@enum ValuationVerdict begin
    VAL_INDECISIVE
    VAL_FINITE
    VAL_AT_INFINITY
end

"""
    judge(val::Valuation; tol = 1e-3, tol_at_infinity = 1e-4)::ValuationVerdict

Judge the current valuation to determine whether a path is diverging. Returns
a [`ValuationVerdict`](@ref).
A path is diverging if one entry of the valuation is negative. To assure that
the estimate is trust worthy we require that the first and second derivative
of [`ν``](@ref) is smaller than `tol_at_infinity`.
The `tol` is used to estimate whether a finite valuation is trustworthy.
"""
function judge(val::Valuation; tol::Float64 = 1e-3, tol_at_infinity::Float64 = 1e-4)
    @unpack ν, ν̇, ν̈ = val
    finite = true
    indecisive = false
    for (i, νᵢ) in enumerate(ν)
        if νᵢ < -max(0.01, tol)
            finite = false
        end

        if νᵢ < -0.05 && abs(ν̇[i]) < tol_at_infinity && abs(ν̈[i]) < tol_at_infinity
            return VAL_AT_INFINITY
        end

        if abs(ν̇[i]) > tol
            indecisive = true
        end
    end
    finite && !indecisive && return VAL_FINITE

    VAL_INDECISIVE
end


###############################################################
# Part of BDisposal.jl package
###############################################################

"""
    efficiencyScores(gI,gO,bO,bI;retToScale,prodStructure,dirGI,dirBI,dirGO,dirBO,startθ,startμ,startλ)

Compute efficiency indicators and convexity test results considering the joined production of good and bad outputs.

Given a set of measures of inputs, "good" ("desiderable") and "bad" ("undesiderable") outputs for different decision making units, compute their
distance to the production frontier, i.e. their degree of efficiency.

## Parameters:
- Positional
  *  `gI`: "Good" inputs (3D matrix by DMUs, input items and periods)
  *  `gO`: "Good" outputs (3D matrix by DMUs, input items and periods)
  *  `bO`: "Bad" outputs (3D matrix by DMUs, input items and periods)
  *  `bI`: "Bad" inputs (optional 3D matrix by DMUs, input items and periods) [default: empty array]
- Keyword
  *  `retToScale`: Wheter the returns to scale should be assumed "constant" or "variable" (default). Non-convex distance and test is computed
      only under the "variable" assumption, except for the `(0,0,1,-1)` distance under multiplicative production function
  *  `prodStructure`: Wheter the production structure should be assumed "additive" (default) or "multiplicative"
  *  `dirGI`,`dirBI`,`dirGO`,`dirBO`: The directions toward where to measure the efficiency (see notes) [default: `(0,0,1,0)`]
  *  `startθ`,`startμ`,`startλ`: Initial values in the convex optimisation problem [default: `(0,0,1.1)`]

## Returns:
- A named tuple with the following items:
  * `λs`: The efficiency indicators coherent with the convexity test (2D matrix by DMUs and periods)
  * `λs_convex`: The efficiency indicators assuming a convex production frontier (2D matrix by DMUs and periods)
  * `λs_nonconvex`:  The efficiency indicators assuming a non-convex production frontier (2D matrix by DMUs and periods)
  * `nonConvTest`: The result of the non-convexity test (2D matrix of boolean by DMUs and periods). Note that "true" here refers to non-convex.
  * `nonConvTest_value`: The value of the non-convexity test (2D matrix by DMUs and periods)

## Description of the function

`efficiencyScores()` defines environmental efficiency indicators for a set
of Decision Making Units (ie., observations). The nature of the DMUs (eg. firms,
countries etc.) depends on the aim of the environmental efficiency analysis. All the
DMUs are characterized by a set of inputs and outputs separated into undesirable (eg.,
polluting) and desirable (eg., no polluting) ones. Note that, inputs separation into
desirable (eg., non emissions-causing) and undesirable (eg., emissions-causing) compo-
nents is optional for the BDisposal package. The BDisposal DEA modelling defines the
smallest environmental production process that contains all DMUs informations (in-
puts/outputs) and satisfies the B-disposal axiomatic pattern.
The BDisposal algorithms for efficiency assessment compute the distance to the best
environmental production procedures; ie., to the efficient production frontier. These
efficiency indices inherit the structure of additive and multiplicative distance functions.
Notice that the BDisposal package computes general shape of additive and multiplicative distance functions. It follows that,
the Bdisposal package introduces a flexible framework for environmental efficiency as-
sessment, with the possibility to analyse different directions for the distance functions.
In addition, the BDisposal package allows to test the (economists’) properties of con-
vexity and disposability.

## Interpretation of the results

### Efficiency indicators

When the additive shape of the distance function is greater than 0 then, the produc-
tion unit doesn’t operate efficiently. If the additive efficiency score is equal to 0 then,
the DMU is efficient. It follows that, the DMU employs one of the best production
technology to transform inputs into outputs; ie., the producer is a benchmark for the
other observations. A similar reasoning holds for the multiplicative distance functions.
If the multiplicative efficiency index is greater than 1 then, the DMU doesn’t perform
efficiently. When the multiplicative efficiency score is equal to 1 then, the production
unit is a benchmark for the other observations.

### Convexity test

When the multiplicative test of convexity is greater than 1
(respectively, equal to 1) then, the production process frontier is non convex (respec-
tively, convex). If the multiplicative test of disposability is greater than 1 (respectively,
equal to 1) then, the frontier of the production set displays B-disposability; ie., there
exists costs for undesirable component decrease. A similar reasoning applies for the
additive framework. The informations provided by these tests define the shape of the
production process boundaries.

## Notes:
* Directions toward where to measure the efficiency distance can be tuned using the `dirGI`,`dirBI`,`dirGO` and `dirBO` parameters.
* The following directions are supported:
  - multiplicative pr. struct.: (-1,0,0,0), (0,-1,0,0), (0,0,1,0), (0,0,0,-1) or (0,0,1,-1)
  - addittive pr. struct.: (1,0,0,0), (0,1,0,0), (0,0,1,0) or (0,0,0,1)
  Other directions can be used to compute the direction under the convex frontier assumption, but the convexity test is not performed and,
  for the multiplicative case, no guarantee is given on solving the underliying (non-linear) problem (different initial value
  `startθ`,`startμ`,`startλ` can be attempted)

"""
function efficiencyScores(gI::Array{Float64,3},gO::Array{Float64,3},bO::Array{Float64,3},bI::Array{Float64,3}=Array{Float64}(undef, (size(gI,1),0,size(gI,3)));
                          retToScale="variable",prodStructure="additive",
                          dirGI=0,dirBI=0,dirGO=1,dirBO=0,
                          startθ=0,startμ=0,startλ=1.1)

    # Data processing
    nDMUs, nPer, nGI, nBI, nGO, nBO = size(gI,1), size(gI,3), size(gI,2), size(bI,2), size(gO,2), size(bO,2)
    λs_convex    = Array{Union{Missing,Float64,String}}(undef,nDMUs,nPer) # Output matrix hosting the results (nDMUs,nPer)
    λs_nonconvex = Array{Union{Missing,Float64,String}}(undef,nDMUs,nPer) # Output matrix hosting the results (nDMUs,nPer)
    directions   = (dirGI,dirBI,dirGO,dirBO)

    # Tests if doing also convex analysis or not
    doNonConvexAnalysis = true

    if retToScale == "constant" && directions != (0,0,1,-1)
        @warn "Non-convex analysis is supported only with variable returns to scale or with direcitons (0,0,1,-1) and multiplicative production structure."
        doNonConvexAnalysis = false
    end

    if  prodStructure == "multiplicative"
        if !( directions in [(-1,0,0,0),(0,-1,0,0),(0,0,1,0),(0,0,0,-1),(0,0,1,-1)])
            @warn "Non-convex analysis not supported with these directions. For the multiplicative case,
            non-convex analysis is supported only with directions (-1,0,0,0),(0,-1,0,0),(0,0,1,0),(0,0,0,-1) or (0,0,1,-1)"
             doNonConvexAnalysis = false
        end
    else
        if !( directions in [(1,0,0,0),(0,1,0,0),(0,0,1,0),(0,0,0,1)])
            @warn "Non-convex analysis not supported with these directions. For the addittive case,
            non-convex analysis is supported only with directions (1,0,0,0),(0,1,0,0),(0,0,1,0) or (0,0,0,1)"
             doNonConvexAnalysis = false
        end
    end
    # temp todo:
    doNonConvexAnalysis = true

    # Loping over the periods
    for t in 1:nPer
        gIₜ = gI[:,:,t]; gOₜ = gO[:,:,t]; bOₜ = bO[:,:,t]
        #bIₜ = nBI > 0 ? bI[:,:,t] : Array{Float64}(undef, 0,0)
        bIₜ = bI[:,:,t]

        #println("Period: $t")
        # Solving for each dmu
        for z in 1:nDMUs
            gIₜ₀ = gIₜ[z,:]; gOₜ₀ = gOₜ[z,:]; bOₜ₀ = bOₜ[z,:]
            #bIₜ₀ = nBI > 0 ? bIₜ[z,:] : Float64[]
            bIₜ₀ = bIₜ[z,:]

        #    println("DMU: $z")

            # CONVEX analysis....
            linearModelPossibleDirections = [(-1,0,0,0),(0,-1,0,0),(0,0,1,0),(0,0,0,-1)]
            forceLinearModel = (dirGI,dirBI,dirGO,dirBO) in linearModelPossibleDirections ?  true : false

            λs_convex[z,t] = problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ;
                                     retToScale=retToScale,prodStructure=prodStructure,
                                     convexAssumption=true,
                                     directions=(dirGI,dirBI,dirGO,dirBO),
                                     startValues=(startθ,startμ,startλ),
                                     forceLinearModel=forceLinearModel)

            # Non-convex analysis (when implemented)...
            λs_nonconvex[z,t] = doNonConvexAnalysis ? problem(gIₜ₀,bIₜ₀,gOₜ₀,bOₜ₀,gIₜ,bIₜ,gOₜ,bOₜ;
                                        retToScale=retToScale,prodStructure=prodStructure,
                                        directions=(dirGI,dirBI,dirGO,dirBO),
                                        convexAssumption=false) : missing
        end # end of each dmu
    end # endof each period

#    println("Done all DMU problems")
    nonConvTest_value = λs_nonconvex ./ λs_convex
    nonConvTest       = nonConvTest_value .< (1.0 - eps())
    λs                = doNonConvexAnalysis ? min.(λs_nonconvex,λs_convex) : λs_convex

    return λs, λs_convex, λs_nonconvex, nonConvTest, nonConvTest_value
end


using Images, FileIO, Colors, FixedPointNumbers, MLDatasets, Statistics
# using WAV

CANVAS_SIZE = 1000

function triangular_step(current_pos)
    random_number = rand()
    if random_number < 1/3
        output = ([0, 0] + current_pos)/2
    elseif random_number < 2/3
        output = ([1, 0] + current_pos)/2
    else
        output = ([0.5, (3^0.5)/2] + current_pos)/2
    end
    output
end

function triangular_chaos(num_points)
    output = zeros(num_points, 2)
    output[1,:] = [0.5, 0.5]
    for i in 2:num_points
        output[i,:] = triangular_step(output[i-1,:])
    end
    output
end

function square_step(current_pos)
    random_number = rand()
    if random_number < 1/4
        output = ([0, 0] + current_pos)/2
    elseif random_number < 1/2
        output = ([1, 0] + current_pos)/2
    elseif random_number < 3/4
        output = ([0, 1] + current_pos)/2
    else
        output = ([1, 1] + current_pos)/2
    end
    output
end

function square_chaos(num_points)
    output = zeros(num_points, 2)
    output[1,:] = [0.5, 0.5]
    for i in 2:num_points
        output[i,:] = square_step(output[i-1,:])
    end
    output
end

function circular_step(current_pos, step)
    theta = 2*pi*step
    output = ([(1+cos(theta))/2, (1+sin(theta))/2] + current_pos)/2
    output
end

function circular_chaos(input)
    num_points = length(input) + 1
    output = zeros(num_points, 2)
    output[1,:] = [0.5, 0.5]
    for i in 2:num_points
        output[i,:] = circular_step(output[i-1,:], input[i-1])
    end
    output
end

function square_driven_step(current_pos, step)
    if step < 1/4
        output = ([0, 0] + current_pos)/2
    elseif step < 1/2
        output = ([1, 0] + current_pos)/2
    elseif step < 3/4
        output = ([0, 1] + current_pos)/2
    else
        output = ([1, 1] + current_pos)/2
    end
    output
end

function square_driven_chaos(input)
    num_points = length(input) + 1
    output = zeros(num_points, 2)
    output[1,:] = [0.5, 0.5]
    for i in 2:num_points
        output[i,:] = square_driven_step(output[i-1,:], input[i-1])
    end
    output
end

function plot_data(data, canvas_size, alpha)
    num_points = size(data)[1]
    canvas = zeros(canvas_size, canvas_size)
    @. data = (data * canvas_size) + 1
    data = floor.(Int, data)
    @. data = canvas_size - data + 1
    for i in 1:num_points
        canvas[data[i,2], data[i,1]] += alpha
    end
    canvas
end

function pixel_color_transform(pixel)
    output = zeros(3)
    #white to blue colormap
    if pixel <= 0
        output = RGB(1.,1.,1.)
    elseif pixel <=0.5
        output = RGB(1-2pixel,1-2pixel,1)
    else
        output = RGB(0.,0.,1.)
    end
    output
end

function render_img(data, canvas_size, alpha)
    data = plot_data(data, canvas_size, alpha)
    x,y = size(data)
    canvas = zeros(RGB,x,y)
    for i in 1:x
        for j in 1:y
            canvas[i,j] = pixel_color_transform(data[i,j])
        end
    end
    canvas
end

function sierpinski_main(num_points, canvas_size, alpha)
    data = triangular_chaos(num_points)
    @. data[data==1] = 0.99999
    img = render_img(data, canvas_size, alpha)
    save("triangle.jpg", img);
end

function square_main(num_points, canvas_size, alpha)
    data = square_chaos(num_points)
    @. data[data==1] = 0.99999
    img = render_img(data, canvas_size, alpha)
    save("square.jpg", img);
end

function circular_main(num_points, canvas_size, alpha)
    data = circular_chaos(rand(num_points))
    @. data[data==1] = 0.99999
    img = render_img(data, canvas_size, alpha)
    save("circular.jpg", img);
end

function munge_dataset(images)
    if length(size(images))==4
        image_data = mean(images, dims=3)
    else
        image_data = images
    end
    num_samples = size(image_data)[end]
    image_data = reshape(image_data, :, num_samples)
    image_data = reshape(image_data, :)
    image_data = image_data + (rand(Float64, size(image_data))/100000)
    quartiles = quantile(image_data, [0.25,0.5,0.75,1.0])
    output = ones(size(image_data))
    @. output[image_data < quartiles[3]] = 0.74
    @. output[image_data < quartiles[2]] = 0.49
    @. output[image_data < quartiles[1]] = 0.24
    output
end

function cifar_10_square(canvas_size, alpha)
    train_x, _ = CIFAR10.traindata()
    chaos_in = munge_dataset(train_x)
    data = square_driven_chaos(chaos_in)
    @. data[data==1] = 0.99999
    img = render_img(data, canvas_size, alpha)
    save(string("cifar_10.jpg"), img)
end

function mnist_square(canvas_size, alpha)
    train_x, _ = MNIST.traindata()
    chaos_in = munge_dataset(train_x)
    data = square_driven_chaos(chaos_in)
    @. data[data==1] = 0.99999
    img = render_img(data, canvas_size, alpha)
    save(string("mnist.jpg"), img)
end

function fashion_mnist_square(canvas_size, alpha)
    train_x, _ = FashionMNIST.traindata()
    chaos_in = munge_dataset(train_x)
    data = square_driven_chaos(chaos_in)
    @. data[data==1] = 0.99999
    img = render_img(data, canvas_size, alpha)
    save(string("fashion_mnist.jpg"), img)
end

function cifar_100_square(canvas_size, alpha)
    train_x, _ = CIFAR100.traindata()
    chaos_in = munge_dataset(train_x)
    data = square_driven_chaos(chaos_in)
    @. data[data==1] = 0.99999
    img = render_img(data, canvas_size, alpha)
    save(string("cifar_100.jpg"), img)
end

# function import_wav(filepath)
    # import_data = wavread(filepath)[1]
    # import_data = mean(import_data, dims=2)
    # import_data = reshape(import_data, :)
    # import_data= import_data + (rand(Float64, size(import_data))/10000000)
    # quartiles = quantile(import_data, [0.25,0.5,0.75,1.0])
    # output = ones(size(import_data))
    # @. output[import_data < quartiles[3]] = 0.74
    # @. output[import_data < quartiles[2]] = 0.49
    # @. output[import_data < quartiles[1]] = 0.24
    # output
# end

# function wav_chaos(filename, canvas_size, alpha)
    # chaos_in = import_wav(string(filename,".wav"))
    # data = square_driven_chaos(chaos_in)
    # @. data[data==1] = 0.99999
    # img = render_img(data, canvas_size, alpha)
    # save(string(filename,"_chaos.jpg"), img)
# end

sierpinski_main(100000000, 500, 0.001)
square_main(100000000, CANVAS_SIZE, 0.001)
cifar_10_square(CANVAS_SIZE,0.1)
mnist_square(CANVAS_SIZE,0.01)
# cifar_100_square(CANVAS_SIZE,0.1)
fashion_mnist_square(CANVAS_SIZE,0.1)


function readinput(filename)
    groups = map(x -> eachmatch(r"(e|se|sw|w|nw|ne)", x), readlines(filename))
end

function sumtuple(a, b, dims)
    return ntuple(i -> a[i] + b[i], dims)
end

function blacktiles(lines)
    blacks = Set()
    for line in lines
        point = (0,0,0)
        for d in line
            point = sumtuple(point, hd[d[1]], 3)
        end
        if point in blacks
            pop!(blacks, point)
        else
            push!(blacks, point)
        end
    end
    return blacks
end

function changetiles(blacks)
    newblacks = Set()
    check = Set()
    for b in blacks
        push!(check,b)
        for d in values(hd)
            push!(check,sumtuple(b,d,3))
        end
    end
    for c in check
        n = 0
        for d in values(hd)
            if sumtuple(c,d,3) in blacks
                n += 1
            end
        end
        if c in blacks && (n == 1 || n == 2)
            push!(newblacks,c)
        end
        if !(c in blacks) && n == 2
            push!(newblacks,c)
        end
    end
    return newblacks
end

function looptiles(blacks, cycles)
    for i in 1:cycles
        blacks = changetiles(blacks)
    end
    return blacks
end

hd = Dict("e"=>(1,-1,0),"se"=>(0,-1,1),"sw"=>(-1,0,1),"w"=>(-1,1,0),"nw"=>(0,1,-1),"ne"=>(1,0,-1))

lines = readinput("d24.in")
blacks1 = blacktiles(lines)
blacks2 = looptiles(blacks1, 100)

println("P1: ", length(blacks1))
println("P2: ", length(blacks2))


using BigNums: BITS, normalize!, set_one!, set_zero!, bitneg!

### Properties ###

###################
### Operational ###
###################

begin
    function checkOne()
        a = one(ArbUInt)
        a.data == [one(ArbDigit)]
    end

    function checkZero()
        a = zero(ArbUInt)
        isempty(a.data)
    end

    function isnormalized(a::ArbUInt)
        findlast(iszero, a.data) != lastindex(a.data)
    end

    function constructNormalized(v::Vector{ArbDigit})
        isnormalized(ArbUInt(v))
    end

    function constructNormalized(v::Union{ArbDigit, DoubleArbDigit})
        isnormalized(ArbUInt(v))
    end

    function setOne(a)
        set_one!(deepcopy(a)) == one(a)
    end

    function setZero(a)
        set_zero!(deepcopy(a)) == zero(a)
    end
end

################
### Ordering ###
################

begin
    function orderTransitive(a,b,c)
        if a < b && b < c
            return a < c
        elseif a < c && c < b
            return a < b
        elseif b < c && c < a
            return b < a
        elseif b < a && a < c
            return b < c
        elseif c < a && a < b
            return c < b
        elseif c < b && b < a
            return c < a
        else
            return a == b || b == c || a == c
        end
    end

    function orderReversal(a,b)
        if a < b
            return b > a
        elseif b < a
            return a > b
        else
            return a == b
        end
    end

    function trichotomy(a,b)
        x1 = a < b
        x2 = a == b
        x3 = a > b
        if x1
            return x2 == x3 == false
        elseif x2
            return x1 == x3 == false
        elseif x3
            return x1 == x2 == false
        end
    end

    function orderPreservedAddition((small,big),constant)
        (small + constant) < (big + constant)
    end
end

##################
### Arithmetic ###
##################

begin ## Addition ##
    function commutative_add(a::ArbUInt, b::ArbUInt)
        c = a + b
        d = b + a
        c == d
    end

    function associative_add(a::ArbUInt, b::ArbUInt, c::ArbUInt)
        d = a + (b + c)
        e = (a + b) + c
        d == e
    end

    function identity_add(a::ArbUInt)
        a == (a + zero(ArbUInt)) && a == (zero(ArbUInt) + a)
    end
end

begin ## Multiplication ##
    function long_karatsuba_eq(a::AbstractVector{ArbDigit}, b::AbstractVector{ArbDigit})
        len = length(a) + length(b)
        long_buf = zeros(ArbDigit, len)
        kara_buf = zeros(ArbDigit, len)

        short, long = length(a) < length(b) ? (a,b) : (b,a)

        BigNums.long_mul!(long_buf, long, short)
        BigNums.karatsuba!(kara_buf, long, short)

        all(Base.splat(==), zip(kara_buf, long_buf))
    end

    function commutative_mul(a::ArbUInt, b::ArbUInt)
        c = a * b
        d = b * a
        c == d
    end

    function associative_mul(a::ArbUInt, b::ArbUInt, c::ArbUInt)
        d = a * (b * c)
        e = (a * b) * c
        d == e
    end

    function distributive_mul(a::ArbUInt, b::ArbUInt, c::ArbUInt)
        d = a * (b + c)
        e = (a*b) + (a*c)
        d == e
    end

    function identity_mul(a::ArbUInt)
        a == a * one(ArbUInt) && a == one(ArbUInt) * a
    end

    function zero_mul(a::ArbUInt)
        iszero(a * zero(ArbUInt)) && iszero(zero(ArbUInt) * a)
    end
end

###############
### Bitwise ###
###############

commutative(op) = (a,b) -> commutative(a,b,op)
function commutative(a, b, op)
    op(a,b) == op(b,a)
end

begin ## Shifts ##
    function leftShiftTrailLeadCorrect(a, sh)
        b = a << sh
        iszero(a) && return trailing_zeros(a) == trailing_zeros(b) == 0

        trail_correct = trailing_zeros(a) + sh == trailing_zeros(b)

        sh %= BITS # => sh < BITS (the part we care about)

        lead_a = leading_zeros(a)
        lead_b = leading_zeros(b)
        lead_correct = begin
            if iszero(sh) # did we just shift the top most word up a number of digits?
                lead_a == lead_b
            elseif sh > lead_a # ok, only a part word overall, but did we overflow?
                lead_b == BITS - (sh - lead_a)
            elseif sh < lead_a # what about without overflow?
                lead_b == lead_a - sh
            else # and what if we shift exactly up to the boundary?
                iszero(lead_b)
            end
        end

        return trail_correct && lead_correct
    end

    function leftShiftNoDataLoss(a, sh)
        b = a << sh
        count_ones(b) == count_ones(a)
    end

    function additionOfShifts(a, (sh1, sh2)::Tuple{T,T}) where T <: Unsigned
        T1 = widen(T)
        ((a << sh1) << sh2) == (a << (T1(sh1) + T1(sh2)))
    end

    function shiftAntiInverse(a, sh)
        ((a << sh) >> sh) == a
    end

    function shiftValueSubtraction(a, (sh1,sh2)::Tuple{T,T}) where T <: Unsigned
        # property assumes sh1 <= sh2
        b = (a << sh2) >> sh1
        c = a << (sh2 - sh1)
        b == c
    end

    function rightShiftSame(a::ArbUInt, sh)
        (a >>> sh) == (a >> sh)
    end

    function shiftIdentity(a)
        (a >> 0x0) == a && (a << 0x0) == a
    end
end

begin ## Negation ##
    function inverse(a)
        a == ~(~a)
    end
end

begin ## And ##
    function zeroing_and(a)
        (a & zero(a)) == zero(a) && (zero(a) & a) == zero(a)
    end

    function identity_and(a)
        (a & a) == a
    end
end

begin ## Or ##
    function identity_or(a, b)
       (a | a) == a
    end
end

begin ## Xor ##
    function zeroing_xor(a)
        (a ⊻ a) == zero(a)
    end

    function identity_xor(a)
        (a ⊻ zero(a)) == a
    end
end

begin ## Nor ##
    function chaining_nor(a,b)
        nor(a,b) == normalize!(~(a | b))
    end
end

begin ## Nand ##
    function chaining_nand(a,b)
        nand(a,b) == normalize!(~(a & b))
    end
end

begin ## Combined Properties ##
    function lossyShift(a, sh)
        b = (a >> sh) << sh
        # mask out the parts we shifted out in b
        # FIXME: This generation of the mask is broken for when we have more than one word
        mask = ArbUInt(unsigned(-1) % ArbDigit) << sh
        # mask = (one(ArbUInt) << sh) - 0x1 # this is proper, but we don't have subtraction yet
        pop!(mask.data)
        bitneg!(mask)
        (b ⊻ (a & mask)) == a
    end
end

begin #= Stringification & Parsing =#
    function parsestring(str, base)
        string(parse(ArbUInt, str); base=base, pad=length(str)) == str
    end

    function stringparse(a, base)
        parse(ArbUInt, string(a; base=base)) == a
    end

    function throwNonHex(str)
        try
            parse(ArbUInt, str; base=16)
        catch ex
            ex isa ArgumentError && return true
            rethrow(ex)
        end
        return false
    end
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    AnisotropicBall(radii, angles; convention=TaitBryanExtr)

An anisotropic ball with `radii` and `angles` according to the rotation `convention`.

- For 2D balls (i.e. ellipses), there are two radii and one rotation angle.
- For 3D balls (i.e. ellipsoids), there are three radii and three rotation angles.

The list of available conventions can be found with `subtypes(RotationConvention)`.
"""
struct AnisotropicBall{V,A,C,M} <: MetricBall
  # input fields
  radii::V
  angles::A
  convention::C

  # state fields
  metric::M
end

function AnisotropicBall(radii::V, angles::A; convention::C=TaitBryanExtr) where {V,A,C}
  Dim, nangles = length(radii), length(angles)
  valid = (Dim == 3 && nangles == 3) || (Dim == 2 && nangles == 1)
  @assert valid "invalid number of radii/angles"

  # invert radii if necessary
  invert = mainaxis(convention) == :Y
  ranges = invert ? [radii[i] for i in reverse(1:Dim,1,2)] : radii

  # scaling matrix
  Λ = Diagonal(SVector{Dim}(one(eltype(ranges))./ranges.^2))

  # rotation matrix
  R = rotmat(angles, convention)

  # ellipsoid metric
  metric = Mahalanobis(Symmetric(R'*Λ*R))

  AnisotropicBall{V,A,C,typeof(metric)}(radii, angles, convention, metric)
end

metric(ball::AnisotropicBall) = ball.metric

radii(ball::AnisotropicBall) = ball.radii

function Base.show(io::IO, ball::AnisotropicBall)
  r = ball.radii
  a = ball.angles
  c = ball.convention
  T = eltype(r)
  print(io, "AnisotropicBall{$T}($r, $a, $c)")
end


# Laplacian

struct Laplacian{T,D} <: LazyQuasiMatrix{T}
    axis::Inclusion{T,D}
end

Laplacian{T}(axis::Inclusion{<:Any,D}) where {T,D} = Laplacian{T,D}(axis)
# Laplacian{T}(domain) where T = Laplacian{T}(Inclusion(domain))
# Laplacian(axis) = Laplacian{eltype(axis)}(axis)

axes(D::Laplacian) = (D.axis, D.axis)
==(a::Laplacian, b::Laplacian) = a.axis == b.axis
copy(D::Laplacian) = Laplacian(copy(D.axis))

@simplify function *(Δ::Laplacian, P::AbstractSphericalHarmonic)
     # Spherical harmonics are the eigenfunctions of the Laplace operator on the unit sphere
     P * Diagonal(mortar(Fill.((-(0:∞)-(0:∞).^2), 1:2:∞)))
end

# Negative fractional Laplacian (-Δ)^α or equiv. abs(Δ)^α

struct AbsLaplacianPower{T,D,A} <: LazyQuasiMatrix{T}
    axis::Inclusion{T,D}
    α::A
end

AbsLaplacianPower{T}(axis::Inclusion{<:Any,D},α) where {T,D} = AbsLaplacianPower{T,D,typeof(α)}(axis,α)

axes(D:: AbsLaplacianPower) = (D.axis, D.axis)
==(a:: AbsLaplacianPower, b:: AbsLaplacianPower) = a.axis == b.axis && a.α == b.α
copy(D:: AbsLaplacianPower) = AbsLaplacianPower(copy(D.axis), D.α)

abs(Δ::Laplacian) = AbsLaplacianPower(axes(Δ,1),1)
-(Δ::Laplacian) = abs(Δ)

^(D::AbsLaplacianPower, k) = AbsLaplacianPower(D.axis, D.α*k)


# This file is part of the Julia package ModularForms.jl
#
# Copyright (c) 2018-2019: Lynn Engelberts and Alexandru Ghitza.

#This file contains the functions delta_poly and delta_qexp which return the
#q-expansion of delta (the weight 12 level 1 cusp form) as a polynomial or
#powerseries (respectively). Moreover, this file contains the function
#delta_k_qexp, which returns the q-expansion of the normalised generator for
#(one of the six) one-dimensional spaces of cusp forms of weight k and level 1.

"""
    delta_poly(prec=10, var="q", K=ZZ)

Return the q-expansion of the normalized cusp form of weight 12 to
precision `prec` as a polynomial over `K` in the variable `var`.

# Arguments
- `prec::Integer=10`: precision of the output
- `var::String="q"`: variable name
- `K=ZZ`: base ring of the output

# Examples still missing
"""
function delta_poly(prec::Int=10, var::String="q", K=ZZ)

   R, q = PolynomialRing(ZZ, var)
   delta = q*eta_qexp(24, prec-1, q)
   RR, q = PolynomialRing(K, var)

   return RR(delta)
end

"""
    delta_qexp(prec=10, var="q", K=ZZ)

Return the q-expansion of the normalized cusp form of weight 12 to
precision `prec` as a power series over `K` in the variable `var`.

# Arguments
- `prec::Integer=10`: precision of the output
- `var::String="q"`: variable name
- `K=ZZ`: base ring of the output

# Examples still missing
"""
function delta_qexp(prec::Int=10, var::String="q", K=ZZ)

   delta = delta_poly(prec, var, K)
   power_series = poly_to_power_series(delta, prec)

   return power_series
end

"""
    delta_k_qexp(k, prec=10, var="q")

Return the q-expansion of the unique normalized eigenform of weight
`k` and level 1 for k in {12, 16, 18, 20, 22, 26} as a power series
to precision `prec` in `var` over ZZ.

These eigenforms are the normalized generators for the six
one-dimensional spaces of cusp forms of level 1.

# Arguments
- `k::Integer`: weight of the unique normalized eigenform
- `prec::Integer=10`: precision of the output
- `var::String="q"`: variable name

# Examples still missing
"""
function delta_k_qexp(k::Int, prec::Int=10, var::String="q")

   if k == 12
      return delta_qexp(prec, var, ZZ)
   elseif k in [16, 18, 20, 22, 26]
      return delta_qexp(prec, var) * eisenstein_series_qexp(k-12, prec, ZZ, var, "integral")
   else
      error("k must be one of 12, 16, 18, 20, 22, 26")
   end

end


using ConstrainedRootSolvers
using PkgUtility
using Test

@testset "ConstrainedRootSolvers --- find_peak" begin
    # segmented function with solution 1
    # segmented function with solution 1
    # segmented function with solution of upper end
    # segmented function with solution 1
    # segmented function with solution of lower bound
    # segmented function with solution 1
    # surface function
    @inline function _f0(x::FT) where {FT<:AbstractFloat}
        return x>1 ? FT(1) : x
    end
    @inline function _f1(x::FT) where {FT<:AbstractFloat}
        return x>1 ? 2-x : x
    end
    @inline function _f2(x::FT) where {FT<:AbstractFloat}
        return x>1 ? (1+x)/2 : x
    end
    @inline function _f3(x::FT) where {FT<:AbstractFloat}
        return x>1 ? 2-x : FT(0)
    end
    @inline function _f4(x::FT) where {FT<:AbstractFloat}
        return x>1 ? 2-x : FT(1)
    end
    @inline function _f5(x::FT) where {FT<:AbstractFloat}
        return x>1 ? FT(0.5) : x
    end
    @inline function _surf_func(x::Array{FT,1}) where {FT<:AbstractFloat}
        return -(x[1]-2)^2 - (x[2]-3)^2
    end
    lr = [1, 1, 5, 1, 0, 1];
    lf = [_f0, _f1, _f2, _f3, _f4, _f5];

    # testing the BisectionMethod and ReduceStepMethod
    for FT in [Float32, Float64]
        ms_bis = BisectionMethod{FT}(x_min=0, x_max=5);
        ms_rst = ReduceStepMethod{FT}(x_min=0, x_max=5, x_ini=1.1, Δ_ini=1);
        st = SolutionTolerance{FT}(1e-4, 50);
        for ms in [ms_bis, ms_rst]
            for i in 1:6
                sol = find_peak(lf[i], ms, st);
                @test !isnan(sol);
                @test typeof(sol) == FT;
                @test sol ≈ lr[i] atol=0.01;
            end
        end
    end

    # testing the NelderMeadMethod and ReduceStepMethodND
    for FT in [Float32, Float64]
        ms_nm1 = NelderMeadMethod{FT}(
                        N = 2,
                        x_inis = FT[rand(FT)+1, rand(FT)+2, 0]);
        ms_nm2 = NelderMeadMethod{FT}(
                        N = 2,
                        x_inis = FT[0.1, 0.1, 0]);
        ms_rst = ReduceStepMethodND{FT}(
                        x_mins = FT[0,0],
                        x_maxs = FT[5,5],
                        x_inis = [1+rand(FT), 2+rand(FT)],
                        Δ_inis = FT[1,1]);
        rt = ResidualTolerance{FT}(1e-5, 50);
        st = SolutionToleranceND{FT}(FT[1e-3, 1e-3], 50);
        for sol in [ find_peak(_surf_func, ms_nm1, rt),
                     find_peak(_surf_func, ms_nm1, st),
                     find_peak(_surf_func, ms_nm2, rt),
                     find_peak(_surf_func, ms_nm2, st),
                     find_peak(_surf_func, ms_rst, st) ]
            @test all( .!isnan.(sol) );
            @test eltype(sol) == FT;
            @test sol[1] ≈ 2 atol=0.01;
            @test sol[2] ≈ 3 atol=0.01;
        end
    end
end

@testset "ConstrainedRootSolvers --- find_zero" begin
    # A S-shape function, e.g. Weibull function
    # A R-shape function
    @inline function _r_func(x::FT) where {FT<:AbstractFloat}
        return exp( -1 * (x/2) ^ FT(0.8) ) - FT(0.5)
    end
    @inline function _s_func(x::FT) where {FT<:AbstractFloat}
        return exp( -1 * (x/2) ^ 5 ) - FT(0.5)
    end

    for FT in [Float32, Float64]
        rt = ResidualTolerance{FT}(1e-3, 50);
        st = SolutionTolerance{FT}(1e-3, 50);

        # Bisection method
        # NewtonBisection method
        # NewtonRaphson method
        # ReduceStep method
        ms_bis = BisectionMethod{FT}(x_min=0, x_max=5);
        ms_nbs = NewtonBisectionMethod{FT}(x_min=0, x_max=5);
        ms_ntr = NewtonRaphsonMethod{FT}(x_ini=2);
        ms_rst = ReduceStepMethod{FT}(x_min=0, x_max=5, x_ini=1.1, Δ_ini=1);

        for ms in [ms_bis, ms_nbs, ms_ntr, ms_rst]
            for tol in [rt, st]
                for f in [_r_func, _s_func]
                    sol = find_zero(f, ms, tol; stepping=true);
                    @test NaN_test(sol);
                    @test FT_test(sol, FT);
                end
            end
        end
    end
end


using Test, BenchmarkTools

input = 9798

xs = 1:300
ys = (1:300)'

rack_id = xs .+ 10 # Find the fuel cell's rack ID, which is its X coordinate plus 10.
# Begin with a power level of the rack ID times the Y coordinate.
# Increase the power level by the value of the grid serial number (your puzzle input).
# Set the power level to itself multiplied by the rack ID.
# Keep only the hundreds digit of the power level (so 12345 becomes 3; numbers with no hundreds digit become 0).
# Subtract 5 from the power level.

function star1(input, sizes=1:300)
    xs = (1:300) .+ 10
    ys = (1:300)'

    power = ((((xs .* ys .+ input) .* xs) .÷ 100) .% 10) .- 5

    maxpower = typemin(Int)
    maxxysz = (0,0,0)
    for sz in sizes
        for x in 1:(300-sz+1)
            for y in 1:(300-sz+1)
                pow = sum(power[xx,yy] for xx in x:(x+sz-1), yy in y:(y+sz-1))
                maxpower = max(maxpower, pow)
                maxxysz = maxpower == pow ? (x,y,sz) : maxxysz
            end
        end
    end
    return maxxysz, maxpower
end

star1(18, 3)

@test star1(18) == ((33,45), 29)
@test star1(42) == ((21,61), 30)

xy, maxp = star1(input)


@testset "Basic" begin
  @testset "Gramian" begin
    x_nu, x_de = [rand(2) for i=1:100], [rand(2) for i=1:200]
    G = DensityRatioEstimation.gaussian_gramian(x_nu, x_de, σ=1.0)
    @test size(G) == (length(x_nu), length(x_de))
    @test all(G .> 0)

    G = DensityRatioEstimation.gaussian_gramian(x_nu, x_nu, σ=2.0)
    @test issymmetric(G)
    @test all(G .> 0)

    # features can be any indexable
    x_nu = [(a=1.,b=2.),(a=3.,b=4.)]
    x_de = [(a=1.,b=2.),(a=3.,b=4.),(a=5.,b=6.)]
    G = DensityRatioEstimation.gaussian_gramian(x_nu, x_de)
    @test size(G) == (2, 3)
    @test all(G .> 0)
  end

  for (d_nu, d_de) in [pair₁, pair₂]
    Random.seed!(123)
    x_nu, x_de = rand(d_nu, 100), rand(d_de, 200)
    @testset "$(typeof(dre).name) -- $optlib" for (dre, optlib) in [(KMM(),   JuMPLib),
                                                                    (KLIEP(), OptimLib),
                                                                    (KLIEP(), ConvexLib),
                                                                    (LSIF(),  OptimLib),
                                                                    (LSIF(),  JuMPLib)]

      r = densratio(x_nu, x_de, dre, optlib=optlib)

      # density ratios must be positive
      @test all(r .> 0)
    end
  end
end