function polrtest(nm::OrdinalMultinomialModel, Z::AbstractVecOrMat; test=:score)
    if test == :score
        polrtest(OrdinalMultinomialScoreTest(nm, reshape(Z, size(Z, 1), size(Z, 2))))
    elseif test == :LRT
        polrtest(OrdinalMultinomialLrtTest(nm, reshape(Z, size(Z, 1), size(Z, 2))))
    else
        throw(ArgumentError("unrecognized test $test"))
    end
end

polrtest(nm::StatsModels.TableRegressionModel{<:OrdinalMultinomialModel}, Z::AbstractVecOrMat; kwargs...) =
polrtest(nm.model, Z; kwargs...)

###########################
# Score test
###########################

struct OrdinalMultinomialScoreTest{TY<:Integer, T<:BlasReal, TL<:GLM.Link}
    "`q`: number of covariates to test significance"
    q::Int
    "`Z`: n-by-q covariate matrix to test significance"
    Z::Matrix{T}
    "`nm`: fitted null model"
    nm::OrdinalMultinomialModel{TY,T,TL}
    "`scoreγ`: score vector of γ"
    scoreγ::Vector{T}
    "`∂γ∂θβ`: information of γ vs (θ,β)"
    ∂γ∂θβ::Matrix{T}
    "`∂γ∂γ`: information of γ vs γ"
    ∂γ∂γ::Matrix{T}
    "`scratchv1`: working vector, same size as `scoreγ`"
    scratchv1::Vector{T}
    "`scratchm1`: working matrix, same size as `Z`"
    scratchm1::Matrix{T}
    "`scratchm2`: working matrix, same size as `∂γ∂θβ`"
    scratchm2::Matrix{T}
    "`scratchm3`: working matrix, same size as `∂γ∂γ`"
    scratchm3::Matrix{T}
end

function OrdinalMultinomialScoreTest(nm::OrdinalMultinomialModel, Z::Matrix)
    TY, T, TL = eltype(nm.Y), eltype(nm.X), typeof(nm.link)
    q = size(Z, 2)
    scoreγ = zeros(T, q)
    ∂γ∂θβ  = zeros(T, q, nm.npar)
    ∂γ∂γ   = zeros(T, q, q)
    scratchv1 = similar(scoreγ)
    scratchm1 = similar(Z)
    scratchm2 = similar(∂γ∂θβ)
    scratchm3 = similar(∂γ∂γ)
    OrdinalMultinomialScoreTest{TY, T, TL}(q, Z, nm, scoreγ, ∂γ∂θβ, ∂γ∂γ, scratchv1, scratchm1, scratchm2, scratchm3)
end

function polrtest(d::OrdinalMultinomialScoreTest)
    #
    # partition the informatrix matrix as
    # [P  W]
    # [W' Q],
    # test statistics is R^T (Q - W' P^{-1} W)^{-1} * R, where R is the score
    #
    # compute R = d.scoreγ, score of γ
    if isempty(d.nm.wts)
        mul!(d.scoreγ, transpose(d.Z), d.nm.wt∂β)
    else
        d.nm.wtwk .= d.nm.wts .* d.nm.wt∂β
        mul!(d.scoreγ, transpose(d.Z), d.nm.wtwk)
    end
    # compute W' = d.∂γ∂θβ
    if isempty(d.nm.wts)
        @views mul!(d.∂γ∂θβ[:, 1:d.nm.J-1], transpose(d.Z), d.nm.wt∂θ∂β) # wt∂θ∂β is n-by-(J-1)
    else
        copyto!(d.scratchm1, d.Z)
        lmul!(Diagonal(d.nm.wts), d.scratchm1)
        @views mul!(d.∂γ∂θβ[:, 1:d.nm.J-1], transpose(d.scratchm1), d.nm.wt∂θ∂β)
    end
    if isempty(d.nm.wts)
        copyto!(d.nm.wtwk, d.nm.wt∂β∂β)
    else
        d.nm.wtwk .= d.nm.wts .* d.nm.wt∂β∂β
    end
    copyto!(d.nm.scratchm1, d.nm.X)
    lmul!(Diagonal(d.nm.wtwk), d.nm.scratchm1)
    @views mul!(d.∂γ∂θβ[:, d.nm.J:end], transpose(d.Z), d.nm.scratchm1)
    # compute Q = ∂γ∂γ
    copyto!(d.scratchm1, d.Z)
    lmul!(Diagonal(d.nm.wtwk), d.scratchm1)
    mul!(d.∂γ∂γ, transpose(d.Z), d.scratchm1)
    # compute scratchm3 = Q - W' P^{-1} W
    mul!(d.scratchm2, d.∂γ∂θβ, d.nm.vcov)
    mul!(d.scratchm3, d.scratchm2, transpose(d.∂γ∂θβ))
    d.scratchm3 .= d.∂γ∂γ .- d.scratchm3
    # compute ts = R^T (Q - W' P^{-1} W)^{-1} R
    eigfact = eigen!(Symmetric(d.scratchm3))
    mul!(d.scratchv1, transpose(eigfact.vectors), d.scoreγ)
    T = eltype(d.scoreγ)
    atol = 1e-8 # tolerance for determining rank
    rk = 0 # rank
    ts = zero(T)
    for j in 1:d.q
        if eigfact.values[j] > atol
            ts += abs2(d.scratchv1[j]) / eigfact.values[j]
            rk += 1
        end
    end
    pval = ts ≤ 0 ? 1.0 : ccdf(Chisq(rk), ts)
end

###########################
# LRT test
###########################

struct OrdinalMultinomialLrtTest{TY<:Integer, T<:BlasReal, TL<:GLM.Link}
    "`q`: number of covariates to test significance"
    q::Int
    "`Z`: n-by-q covariate matrix to test significance"
    Z::Matrix{T}
    "`nm`: fitted null model"
    nm::OrdinalMultinomialModel{TY,T,TL}
    "`Xaug`: augmented X matrix for full model"
    Xaug::Matrix{T}
end

function OrdinalMultinomialLrtTest(nm::OrdinalMultinomialModel, Z::AbstractVecOrMat)
    TY, T, TL = eltype(nm.Y), eltype(nm.X), typeof(nm.link)
    q = size(Z, 2)
    Xaug = Matrix{T}(undef, nm.n, nm.p + q)
    @views copyto!(Xaug[:, 1:nm.p], nm.X)
    @views copyto!(Xaug[:, nm.p+1:end], Z)
    OrdinalMultinomialLrtTest{TY, T, TL}(q, Z, nm, Xaug)
end

function polrtest(
    d::OrdinalMultinomialLrtTest;
    solver::TSOLVER = NLopt.Optimizer()
    ) where {TSOLVER <: MOI.AbstractOptimizer}
    if typeof(solver) == NLopt.Optimizer
        algo = MOI.get(solver, MOI.RawOptimizerAttribute("algorithm"))
        if algo == :none
            MOI.set(solver, MOI.RawOptimizerAttribute("algorithm"), :LD_SLSQP)
            MOI.set(solver, MOI.RawOptimizerAttribute("max_iter"), 4000)
        end
    end
    am = polr(d.Xaug, d.nm.Y, d.nm.link, solver; wts = d.nm.wts)
    ccdf(Chisq(d.q), deviance(d.nm) - deviance(am))
end


import Base:
    +, show, *

doc"""N-dimensional Volume with lower and upper bounds"""
immutable Vol{N,T}
    bounds::Interval{T}
end

Vol_name(i::Integer) = i == 1 ? "length" : i == 2 ? "area" : i == 3 ? "Volume" : "hyper-Volume"

show{N,T}(io::IO, v::Vol{N,T}) = print(io, "$N-dimensional $(Vol_name(N)): $(v.bounds)")

Vol{N,T}(X::IntervalBox{N,T}) = Vol{N,T}(prod([Interval(x.hi) - Interval(x.lo) for x in X]))

Vol{T}(x::Interval{T}) = Vol(IntervalBox(x))

+{N,T}(v::Vol{N,T}, w::Vol{N,T}) = Vol{N,T}(v.bounds + w.bounds)

*{N1,N2,T}(v::Vol{N1,T}, w::Vol{N2,T}) = Vol{N1+N2,T}(v.bounds * w.bounds)

Vol{N,T}(XX::SubPaving{N,T}) = sum([Vol(X) for X in XX])

function Vol{N,T}(P::Paving{N,T})
    Vin = Vol(P.inner)
    Vout = Vin + Vol(P.boundary)

    return Vol{N,T}(hull(Vin.bounds, Vout.bounds))
end


using Plots
using Random
using VegaLite
using DataFrames
using LinearAlgebra
using TripletEmbeddings

Random.seed!(4)

# We define a random Embedding in ℜ²
d = 2
n = 100
Y = rand(d, n)

triplets = Triplets(Y)
loss = STE(σ = 1/sqrt(2))
X = Embedding(size(Y))

@time violations = fit!(loss, triplets, X; verbose=true, max_iterations=200)
procrustes!(X, Y)

scatter(Y[1,:], Y[2,:]); scatter!(X[1,:], X[2,:])

dfX = DataFrame(embedding = "X", time = 1:n, x = X[1,:], y = X[2,:])
dfY = DataFrame(embedding = "Y", time = 1:n, x = Y[1,:], y = Y[2,:])

vcat(dfX, dfY) |> @vlplot(:point, x = :x, y = :y, color=:embedding, width=600, height=400)


"""
    SeisOMP(in,operator,param,iter,lp,imax)

Robust Orthognal Matching Pursuit algorithm with the coefficient updated by the conjugate gradient pursuit

# Arguments
- `dn`: Noisy data
- `operator`: Any operator passed the dot test
- `param`: parameters for the operator
- `iter`: Max iterations for the matching pursuit
- `lp`: lp-norm space inner product
- `imax`: Max interation for the lp-norm inner product

"""
function SeisOMP(dn,operator,param,iter,lp,imax)
    m=operator(dn,param,-1,lp,imax);
    r=dn;
    n=1;
    ind_r=[];
    ind_c=[];
    v=[];
    d=[];
    dr=zeros(size(dn))
    for i=1:iter
        g=zeros(size(m))
        gn=operator(r,param,-1,lp,imax)
        ind_ri=Tuple(findmax(abs.(gn))[2])[1]
        ind_ci=Tuple(findmax(abs.(gn))[2])[2]
        push!(ind_r,ind_ri)
        push!(ind_c,ind_ci)
        for j=length(ind_r)
            g[ind_r[j],ind_c[j]]=gn[ind_r[j],ind_c[j]]
        end
        if n==1
            d=g;
            v=operator(d,param,1,lp,imax);
        else
            w=operator(g,param,1,lp,imax)
            u=-LP_norm(vec(v),vec(w),lp,imax,1)
            d=g+u*d
            v=w+u*v
        end
        beta=LP_norm(vec(v),vec(r),lp,imax,1)
        dr=dr+beta*v
        r=r-beta*v
        n=n+1
    end
    return dr
end


using PyCall
using PyPlot

function plot2Dclassifier(X,y,model;Xtest=[],ytest=[],biasIncluded=false,k=2)

	if biasIncluded
		f1 = 2
		f2 = 3
	else
		f1 = 1
		f2 = 2
	end

	increment = 100

	# Pick some colors for the classes (bor binary it uses red/blue)
	colours = [(0,1,0)
			(1,0,0)
			(0,0,1)
			(1,1,1)
			(1,0,1)
			(0,1,1)
			(1,1,0)
			(.1,.1,.1)
			(1,.5,0)
			(0,.5,0)
			(.5,.5,.5)
			(.5,.25,0)
			(.5,0,.5)
			(0,.5,1)]

	figure()
	if k == 2
		plot(X[y.==1,f1],X[y.==1,f2],"b+")
		plot(X[y.==2,f1],X[y.==2,f2],"ro")
		if !isempty(Xtest)
			plot(Xtest[ytest.==1,f1],Xtest[ytest.==1,f2],"bx");
			plot(Xtest[ytest.==2,f1],Xtest[ytest.==2,f2],"rs");
		end
	else

		for c in 1:k
			plot(X[y.==c,f1],X[y.==c,f2],"+",color=colours[c])
			if !isempty(Xtest)
				plot(Xtest[ytest.==c,f1],Xtest[ytest.==c,f2],"x",color=colours[c])
			end
		end
	end

	(xmin,xmax) = xlim()
	xDomain = linspace(xmin,xmax,increment)
	(ymin,ymax) = ylim()
	yDomain = linspace(ymin,ymax,increment)

	xValues = repmat(xDomain,1,length(xDomain))
	yValues = repmat(yDomain',length(yDomain),1)

	if biasIncluded
		t = length(xValues)
		z = model.predict([ones(t,1) xValues[:] yValues[:]])
	else
		z = model.predict([xValues[:] yValues[:]])
	end
	@assert(length(z) == length(xValues),"Size of model function's output is wrong");

	zValues = reshape(z,size(xValues))

	@pyimport matplotlib.colors as matcolors
	if k == 2
		if all(zValues[:] == 1)
    		cm = [(0,0,.5)];
		elseif all(zValues[:] == 2)
    		cm = [(.5,0,0)];
		else
    		cm = [(0,0,.5);(.5,0,0)];
		end
		cmap = matcolors.ListedColormap(cm,"A")

		contourf(xValues,yValues,zValues,cmap=cmap)
	else
		cm = []
		nColours = 0
		for c in 1:k
			if any(zValues[:] .== c)
				nColours += 1
				colour = (.5colours[c][1],.5colours[c][2],.5colours[c][3])
				zValues[zValues .== c] = nColours
				push!(cm,colour)
			end
		end
		levels = 0:nColours
		cmap = matcolors.ListedColormap(cm,"A")
		contourf(xValues,yValues,zValues,levels=levels,cmap=cmap)
	end
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.6
#   kernelspec:
#     display_name: Julia 1.0.3
#     language: julia
#     name: julia-1.0
# ---

module NHosts1Vector
# # Model Specifications

# ## Host Specifications
#
# The host may be in one of 4 states:
#
# - susceptible (S_H)
#
# - incubation period (E_H)
#
# - infectious (I_H)
#
# - recovered (R_H)
#
# ## Vector Specifications
#
# The vector may be in one of 3 states:
#
# - susceptible (S_V)
#
# - incubation period (E_V)
#
# - infectious (I_V)

# # Model Parameters
#
# σ_H, σ_V: the incubation rates for hosts & vectors (units: per time)
#
# μ_H, μ_V: the mortality rates for hosts & vectors (units: per time)
#
# λ: the clearance (or recovery) rate for hosts (units: per time)
#
# β: the infection rate (units: per capita per time)

# $$N_H^i = S_H^i + E_H^i + I_H^i + R_H^i, \dot{N_H^i}=0$$
#
# $$N_H = \sum_{i=1}^N N_H^i$$
#
# $$\dot{S_H}^i = \mu_H^i(E_H^i + I_H^i + R_H^i) - \frac{\beta^i}{N_H} S_H^i I_V$$
#
# $$\dot{E_H}^i = \frac{\beta^i}{N_H} S_H^i I_V - (\sigma_H^i + \mu_H^i) E_H^i$$
#
# $$\dot{I_H}^i = \sigma_H^i E_H^i - (\lambda^i + \mu_H^i) I_H^i$$
#
# $$\dot{R_H}^i = \lambda^i I_H^i - \mu_H^i R_H^i$$
#
# $$\dot{S_V} = \mu_V(E_V + I_V) - \frac{\sum_{i=1}^N\beta^i I_H^i}{N_H} S_V$$
#
# $$\dot{E_V} = \frac{\sum_{i=1}^N\beta^i I_H^i}{N_H} S_V - (\sigma_V + \mu_V) E_V$$
#
# $$\dot{I_V} = \sigma_V E_V - \mu_V I_V$$

# This version considers multiple hosts, which introduces some indexing complications. This approach adopts the indexing scheme of:
#
# - all of host $i$ compartments ($S_H$, $E_H$, etc), for $i∈1…N$
# - vector compartments last
# With this approach, we can re-use the solution for one host, with slight modifications:

# +
H_comps = 4
V_comps = 3

# in sub functions, du / u are the particular relevant slices only
function F1H(du,u,p,t,I_V,N_H)
    S_H, E_H, I_H, R_H = u

    # host dynamics
    host_infection = (p.β*S_H*I_V)/N_H
    host_mortality = p.μ_H .* u # include S_H, so easier to remove mortality
    host_births = sum(host_mortality)
    host_progression = p.σ_H*E_H
    recovery = p.λ*I_H

    du[1] = -host_infection + host_births
    du[2] = host_infection - host_progression
    du[3] = host_progression - recovery
    du[4] = recovery
    du[1:end] -= host_mortality
end

# in sub functions, du / u are the particular relevant slices only
function FV(du,u,p,t,sum_β_I_H)
    S_V, E_V, I_V = u
    vec_infection = sum_β_I_H*S_V/p.N_H
    vec_mortality = p.μ_V .* u # include S_V, so easier to remove mortality
    vec_births = sum(vec_mortality)
    vec_progression = p.σ_V*E_V

    du[1] = -vec_infection + vec_births
    du[2] = vec_infection - vec_progression
    du[3] = vec_progression
    du[1:end] -= vec_mortality
end

function F(du,u,p,t)
    uvec = @view(u[(p.nHosts*H_comps+1):end]) # grab the vector compartments
    S_V, E_V, I_V = uvec
    sum_β_I_H = 0.0
    for host in 0:(p.nHosts-1)
        slice = (1:H_comps).+(H_comps*host)
        F1H(@view(du[slice]), @view(u[slice]), p.host[host+1], t, I_V, p.vec.N_H)
        # must use @view here, so that these arrays can be modified in F1H
        sum_β_I_H += p.host[host+1].β * u[slice[3]] # this host's I compartment
    end
    FV(@view(du[(p.nHosts*4+1):end]), uvec, p.vec,t,sum_β_I_H)
end
# -

# First, state initial conditions. This code generates them randomly for convenience, though they could be assigned based on data, desired parameter space, or algorithmically as part of a fitting process:

# +
using Random
nH = 5
Random.seed!(0)

S_Hs = ones(nH) .* 100.0
E_Hs = zeros(nH)
I_Hs = shuffle(vcat(zeros(nH-1),[1.0]))
R_Hs = zeros(nH)
host0 = reshape(hcat(S_Hs,E_Hs,I_Hs,R_Hs)', nH*4, 1)
vec0 = [10000.0, 0.0, 0.0]
u0 = vcat(host0, vec0);
# -

# Now, generate dynamic parameters. Again: this code generates them randomly for convenience, though they could be assigned based on data, desired parameter space, or algorithmically as part of a fitting process:

# +
Random.seed!(1)

μs = 1 ./ (rand(nH) .* 360)
σs = 1 ./ (rand(nH) .* 6)
λs = 1 ./ (rand(nH) .* 28)
βs = rand(nH) ./ 10.0

p = (
  nHosts = nH,
  vec = (μ_V=1/30, σ_V=1/7, N_H = sum(host0)),
  host = [(μ_H=μs[i], σ_H=σs[i], λ=λs[i], β=βs[i]) for i in 1:nH]
  # just building up a random collection of params for demonstration
);

# +
using DifferentialEquations
using IterableTables, DataFrames

tspan = (0.0, 365.0)
prob = ODEProblem(F, u0, tspan, p)
sol = solve(prob,Tsit5(),reltol=1e-8,abstol=1e-8,saveat=collect(range(0,stop = 365,length = 365*10+1)));
# -

df = DataFrame(sol)
namespace = zip(names(df)[2:end],[
    :S_H1,:E_H1,:I_H1,:R_H1,
    :S_H2,:E_H2,:I_H2,:R_H2,
    :S_H3,:E_H3,:I_H3,:R_H3,
    :S_H4,:E_H4,:I_H4,:R_H4,
    :S_H5,:E_H5,:I_H5,:R_H5,
    :S_V,:E_V,:I_V])
namespace = collect(namespace)
rename!(df,namespace);

using Plots
host_1 = plot(df[:timestamp],[df[:,:S_H1],df[:,:E_H1],df[:,:I_H1],df[:,:R_H1]])
host_2 = plot(df[:timestamp],[df[:,:S_H2],df[:,:E_H2],df[:,:I_H2],df[:,:R_H2]])
host_3 = plot(df[:timestamp],[df[:,:S_H3],df[:,:E_H3],df[:,:I_H3],df[:,:R_H3]])
host_4 = plot(df[:timestamp],[df[:,:S_H4],df[:,:E_H4],df[:,:I_H4],df[:,:R_H4]])
plot(host_1,host_2,host_3,host_4,layout=(4,1),label=["S","E","I","R"])

end


# DFS implementation optimized from http://www.cs.nott.ac.uk/~psznza/G5BADS03/graphs2.pdf
# Depth-first visit / traversal

"""
    is_cyclic(g)

Return `true` if graph `g` contains a cycle.

### Implementation Notes
Uses DFS.
"""
function is_cyclic end
@traitfn is_cyclic(g::::(!IsDirected)) = ne(g) > 0
# see https://github.com/mauro3/SimpleTraits.jl/issues/47#issuecomment-327880153 for syntax
@traitfn function is_cyclic(g::AG::IsDirected) where {T, AG<:AbstractGraph{T}}
    vcolor = zeros(UInt8, nv(g))
    for v in vertices(g)
        vcolor[v] != 0 && continue
        S = Vector{T}([v])
        vcolor[v] = 1
        while !isempty(S)
            u = S[end]
            w = 0
            for n in outneighbors(g, u)
                if vcolor[n] == 1
                    return true
                elseif vcolor[n] == 0
                    w = n
                    break
                end
            end
            if w != 0
                vcolor[w] = 1
                push!(S, w)
            else
                vcolor[u] = 2
                pop!(S)
            end
        end
    end
    return false
end

# Topological sort using DFS
"""
    topological_sort_by_dfs(g)

Return a [topological sort](https://en.wikipedia.org/wiki/Topological_sorting) of a directed
graph `g` as a vector of vertices in topological order.
"""
function topological_sort_by_dfs end
# see https://github.com/mauro3/SimpleTraits.jl/issues/47#issuecomment-327880153 for syntax
@traitfn function topological_sort_by_dfs(g::AG::IsDirected) where {T, AG<:AbstractGraph{T}}
    vcolor = zeros(UInt8, nv(g))
    verts = Vector{T}()
    for v in vertices(g)
        vcolor[v] != 0 && continue
        S = Vector{T}([v])
        vcolor[v] = 1
        while !isempty(S)
            u = S[end]
            w = 0
            for n in outneighbors(g, u)
                if vcolor[n] == 1
                    error("The input graph contains at least one loop.") # TODO 0.7 should we use a different error?
                elseif vcolor[n] == 0
                    w = n
                    break
                end
            end
            if w != 0
                vcolor[w] = 1
                push!(S, w)
            else
                vcolor[u] = 2
                push!(verts, u)
                pop!(S)
            end
        end
    end
    return reverse(verts)
end

"""
    dfs_tree(g, s)

Return an ordered vector of vertices representing a directed acyclic graph based on
depth-first traversal of the graph `g` starting with source vertex `s`.
"""
dfs_tree(g::AbstractGraph, s::Integer; dir=:out) = tree(dfs_parents(g, s; dir=dir))

"""
    dfs_parents(g, s[; dir=:out])

Perform a depth-first search of graph `g` starting from vertex `s`.
Return a vector of parent vertices indexed by vertex. If `dir` is specified,
use the corresponding edge direction (`:in` and `:out` are acceptable values).

### Implementation Notes
This version of DFS is iterative.
"""
dfs_parents(g::AbstractGraph, s::Integer; dir=:out) =
(dir == :out) ? _dfs_parents(g, s, outneighbors) : _dfs_parents(g, s, inneighbors)

function _dfs_parents(g::AbstractGraph{T}, s::Integer, neighborfn::Function) where T
    parents = zeros(T, nv(g))

    seen = zeros(Bool, nv(g))
    S = Vector{T}([s])
    seen[s] = true
    parents[s] = s
    while !isempty(S)
        v = S[end]
        u = 0
        for n in neighborfn(g, v)
            if !seen[n]
                u = n
                break
            end
        end
        if u == 0
            pop!(S)
        else
            seen[u] = true
            push!(S, u)
            parents[u] = v
        end
    end
    return parents
end


struct AugmentedLagrangian
    λ               # duals - general constraints
    λl              # duals - primal lower bounds
    λu              # duals - primal upper bounds

    ρ               # penalties - general constraints
    ρl              # penalties - primal lower bounds
    ρu              # penalties - primal upper bounds
    ρ0              # initial penalty
    s               # penalty scaling

    as              # active set for general constraints
    asl             # active set for primal lower bounds
    asu             # active set for primal upper bounds

    c               # memory for constraints
    ∇c              # memory for constraints Jacobian

    xl              # primal lower bounds
    xu              # primal upper bounds
    cl              # memory for xl - x
    cu              # memory for x - xu

    idx_ineq        # inequality constraint indices
    idx_l           # primal lower bound indices
    idx_u           # primal upper bound indices
    idx_soc         # second-order cone constraint indices
end

function augmented_lagrangian(n, m;
        ρ0 = 1.0, s = 10.0,
        idx_ineq = (1:0),
        xl = -Inf * ones(n), xu = Inf * ones(n),
        idx_soc = (1:0))

    bool_l = isfinite.(xl)
    bool_u = isfinite.(xu)
    idx_l = (1:length(xl))[bool_l]
    idx_u = (1:length(xu))[bool_u]
    ml = sum(bool_l)
    mu = sum(bool_u)

    λ = zeros(m)
    λl = zeros(ml)
    λu = zeros(mu)

    ρ = ρ0 * ones(m)
    ρl = ρ0 * ones(ml)
    ρu = ρ0 * ones(mu)

    as = ones(Bool, m)
    asl = ones(Bool, ml)
    asu = ones(Bool, mu)

    c = zeros(m)
    ∇c = spzeros(m, n)

    cl = zeros(ml)
    cu = zeros(mu)

    AugmentedLagrangian(λ, λl, λu,
        ρ, ρl, ρu, ρ0, s,
        as, asl, asu,
        c, ∇c,
        xl, xu, cl, cu,
        idx_ineq, idx_l, idx_u, idx_soc)
end

function update!(al::AugmentedLagrangian)
    dual_update!(al)
    penalty_update!(al)
    active_set_update!(al)
end

function dual_update!(al::AugmentedLagrangian)
    al.λ .+= al.ρ .* al.c
    al.λl .+= al.ρl .* al.cl
    al.λu .+= al.ρu .* al.cu

    inequality_projection!(al)
    #TODO: soc project
end

function inequality_projection!(al::AugmentedLagrangian)
    al.λ[al.idx_ineq] = max.(0.0, view(al.λ, al.idx_ineq))
    al.λl .= max.(0.0, al.λl)
    al.λu .= max.(0.0, al.λu)
end

function soc_projection!(al::AugmentedLagrangian)
    return nothing
end

function penalty_update!(al::AugmentedLagrangian)
    al.ρ .*= al.s
    al.ρl .*= al.s
    al.ρu .*= al.s
end

function active_set_update!(al::AugmentedLagrangian)
    fill!(al.as, 1)
    fill!(al.asl, 1)
    fill!(al.asu, 1)

    al.as[al.idx_ineq] = ((view(al.λ, al.idx_ineq) .> 0.0) .| (view(al.c, al.idx_ineq) .>= 0.0)) # (view(al.c, al.idx_ineq) .> 0.0) #
    al.asl .= ((al.λl .> 0.0) .| (al.cl .>= 0.0)) # (al.cl .> 0.0) #
    al.asu .= ((al.λu .> 0.0) .| (al.cu .>= 0.0)) # (al.cu .> 0.0) #
    #TODO: soc as
end

function bounds!(al::AugmentedLagrangian, x)
    al.cl .= view(al.xl, al.idx_l) - view(x, al.idx_l)
    al.cu .= view(x, al.idx_u) - view(al.xu, al.idx_u)
end

function constraint_violation(al::AugmentedLagrangian, x)
    c!(al.c, x)
    bounds!(al, x)
    al.c[al.idx_ineq] = max.(0.0, al.c[al.idx_ineq])
    return max(norm(al.c, Inf),
               norm(max.(0.0, al.cl), Inf),
               norm(max.(0.0, al.cu), Inf))
end

function reset!(al::AugmentedLagrangian)
    fill!(al.λ, 0.0)
    fill!(al.λl, 0.0)
    fill!(al.λu, 0.0)
    fill!(al.ρ, al.ρ0)
    fill!(al.ρl, al.ρ0)
    fill!(al.ρu, al.ρ0)
    fill!(al.as, 1)
    fill!(al.asl, 1)
    fill!(al.asu, 1)
    fill!(al.c, 0.0)
    fill!(al.cl, 0.0)
    fill!(al.cu, 0.0)
end

# n = 20
# m = 10
# al = augmented_lagrangian(n, m, xl = zeros(n))
# update!(al)
# reset!(al)
#
# al.asl
# al.λl
# al.cl
# ((al.λl .== 0.0) .& (al.cl .< 0.0))


@testset "Monin_Obukhov_length" begin
    datetime, ustar, Tair, pressure, H = tha48[24,Cols(:datetime, :ustar, :Tair, :pressure, :H)]
    MOL24 = @inferred Monin_Obukhov_length(Tair, pressure, ustar, H)
    MOL24 = @inferred Monin_Obukhov_length(Tair, pressure, ustar, H; constants =BigleafConstants())
    @test ≈(MOL24, -104.3, rtol = 1/1000)
    #
    df = copy(tha48)
    @inferred Monin_Obukhov_length!(df)
    @test df.MOL[24] == MOL24
end

@testset "stability_parameter" begin
    df = copy(tha48)
    df24 = df[24,:]
    MOL24 = Monin_Obukhov_length(df24.Tair, df24.pressure, df24.ustar, df24.H)
    z=40;d=15
    zeta = @inferred stability_parameter(z,d,MOL24)
    @test ≈(zeta , -0.240, rtol = 1/100)
    zeta2 = @inferred stability_parameter(z,d,df24.Tair, df24.pressure, df24.ustar, df24.H)
    @test zeta2 == zeta
    #
    @inferred stability_parameter!(df;z,d)
    @test df.zeta[24] == zeta
    zeta_scalar = df.zeta
    #
    # also test zr as a vector
    df = copy(tha48)
    df[!,:zi] .= z
    df.zi[1] = z/2
    zetas = stability_parameter!(df; z=df.zi, d).zeta
    @test zetas[2:24] == zeta_scalar[2:24]
    @test zetas[1] != zeta_scalar[1]
    #
    # specifyngn MOL directly
    df4 = copy(tha48)
    Monin_Obukhov_length!(df4)
    zetas2 = stability_parameter!(df4; z, d, MOL = df4.MOL./2).zeta
    @test all(zetas2[2:end] .== zetas[2:end] .* 2)
end

@testset "stability_correction zeta" begin
    zeta = -2:0.5:0.5
    @inferred stability_correction(first(zeta))
    df2 = DataFrame(stability_correction.(zeta))
    @test all(isapprox.(df2.psi_h, SA[2.431,2.197,1.881,1.386,0,-2.5], rtol=1e-3))
    @test all(isapprox.(df2.psi_m, SA[2.275, 2.061, 1.772, 1.317, 0, -2.5], rtol=1e-3))
    df2 = DataFrame(stability_correction.(zeta; stab_formulation=Businger1971()))
    @test all(isapprox.(df2.psi_h, SA[2.085, 1.862, 1.564, 1.106,0, -3.9], rtol=1e-3))
    @test all(isapprox.(df2.psi_m, SA[2.418, 2.200, 1.904, 1.435,0, -3], rtol=1e-3))
    #
    # https://bitbucket.org/juergenknauer/bigleaf/issues/8/logical-inconsistency-possible-bug-with
    resm = @inferred stability_correction(missing)
    @test keys(resm) == (:psi_h, :psi_m)
    @test all(ismissing.(values(resm)))
    #
    resm = @inferred stability_correction(first(zeta);
        stab_formulation = NoStabilityCorrection())
    @test resm == (psi_h = 0.0, psi_m = 0.0)
end

@testset "stability_correction metvars" begin
    z=40;d=15
    df = DataFrame(Tair=25.0, pressure=100.0, ustar=0.2:0.1:1.0, H=40:20.0:200)
    df1 = df[1,:]
    zeta1 = stability_parameter(z,d,df1.Tair, df1.pressure, df1.ustar, df1.H)
    res1 = @inferred stability_correction(z,d, df1.Tair, df1.pressure, df1.ustar, df1.H)
    @test res1 == stability_correction(zeta1)
end

@testset "stability_correction DataFrame variant" begin
    z=40;d=15
    dfo = DataFrame(Tair=25.0, pressure=100.0, ustar=0.2:0.1:1.0, H=40:20.0:200)
    df = copy(dfo)
    res = stability_correction(df; z, d)
    # for type stability, use columntable(df)
    stability_correction!(df; z, d)
    propertynames(df)[(end-1):end] == SA[:psi_h, :psi_m]
    @test DataFrame(res) == df[!, (end-1):end]
    #
    dfm = allowmissing(dfo)
    dfm.ustar[1] = missing
    res = stability_correction(
        dfm; z, d, stab_formulation = NoStabilityCorrection())
    stability_correction!(
        dfm; z, d, stab_formulation = NoStabilityCorrection())
    propertynames(dfm)[(end-1):end] == SA[:psi_h, :psi_m]
    @test all(iszero.(dfm.psi_h))
    @test all(iszero.(dfm.psi_m))
    @test DataFrame(res) == dfm[!, (end-1):end]
    #
    # test specifying zeta instead of z and d
    df2 = copy(dfo)
    stability_parameter!(df2; z, d) # adds zeta
    stability_correction!(df2, zeta=df2.zeta)
    @test df2.psi_h == df.psi_h
    @test df2.psi_m == df.psi_m
    #
    # test specifying zr as a vector
    df3 = copy(dfo)
    df3[!,:zi] .= z
    df3.zi[1] = z/2
    res = stability_correction(df3; z=df3.zi, d)
    @inferred stability_correction!(df3; z=df3.zi, d)
    @test df3.psi_h[2:end] == df.psi_h[2:end]
    @test df3.psi_h[1] != df.psi_h[1]
    @test DataFrame(res) == df3[!, (end-1):end]
end

@testset "stability_correction DataFrame variant Float32" begin
    z=40;d=15
    dfo = DataFrame(Tair=25.0, pressure=100.0, ustar=0.2:0.1:1.0, H=40:20.0:200)
    float_cols = names(dfo, All())
    dfo = transform(dfo,  float_cols .=> ByRow(passmissing(Float32)) .=> float_cols)

    df = copy(dfo)
    res = stability_correction(df; z, d)
    @test eltype(res.psi_h) == Float32
    # for type stability, use columntable(df)
    stability_correction!(df; z, d)
    propertynames(df)[(end-1):end] == SA[:psi_h, :psi_m]
    @test eltype(df.psi_h) == Float32
    @test DataFrame(res) == df[!, (end-1):end]
    #
    dfm = allowmissing(dfo)
    dfm.ustar[1] = missing
    res = stability_correction(
        dfm; z, d, stab_formulation = NoStabilityCorrection())
    @test eltype(res.psi_h) == Float32
    stability_correction!(
        dfm; z, d, stab_formulation = NoStabilityCorrection())
    propertynames(dfm)[(end-1):end] == SA[:psi_h, :psi_m]
    @test all(iszero.(dfm.psi_h))
    @test all(iszero.(dfm.psi_m))
    #@test eltype(dfm.psi_h) == typeof(z) #Float32
    @test DataFrame(res) == dfm[!, (end-1):end]
    #
    # test specifying zeta instead of z and d
    df2 = copy(dfo)
    stability_parameter!(df2; z, d) # adds zeta
    stability_correction!(df2, zeta=df2.zeta)
    @test eltype(df2.psi_h) == Float32
    @test df2.psi_h == df.psi_h
    @test df2.psi_m == df.psi_m
    #
    # test specifying zr as a vector
    df3 = copy(dfo)
    df3[!,:zi] .= z
    df3.zi[1] = z/2
    df3_before = copy(df3)
    res = stability_correction(df3; z=df3.zi, d)
    @test df3 == df3_before
    @test eltype(res.psi_h) == Float32
    @inferred stability_correction!(df3; z=df3.zi, d)
    @test eltype(df3.psi_h) == Float32
    @test df3.psi_h[2:end] == df.psi_h[2:end]
    @test df3.psi_h[1] != df.psi_h[1]
    @test DataFrame(res) == df3[!, (end-1):end]
end


using DiffPrivacy, Test

@info "Testing one parallel database"
x = [1,2,3,4,5,6,7,8,9]
y = DiffPrivacy.get_one_parallel_database(x, 2)
@test y == [1 3 4 5 6 7 8 9]

@info "Testing multiple parallel databases"
x = [1,2,3,4]
y = DiffPrivacy.get_multiple_parallel_databases(x)
@test y == [[2 3 4],[1 3 4],[1 2 4],[1 2 3]]

@info "Testing sum sensitivity"
x = [1,2,3,4]
y = DiffPrivacy.sum_sensitivity(x)
@test y == 4

@info "Testing mean sensitivity"
x = [3,6,9,12]
y = DiffPrivacy.mean_sensitivity(x)
@test y == 1.5

@info "All tests completed"


@testset "70.climbing-stairs.jl" begin
    @test climbing_stairs(1) == 1
    @test climbing_stairs(2) == 2
    @test climbing_stairs(3) == 3
    @test climbing_stairs(8) == 34
    @test climbing_stairs(45) == 1836311903
end


function _nodepath(p::T, q::T, target::T, path::Vector{T}, keep::Bool) where T <: AbstractNode
    for link in q.links
        link.to == p && continue
        if link.to == target
            keep = true
            push!(path, link.to)
            break
        end
        istip(link.to) && continue
        keep, path = _nodepath(q, link.to, target, path, keep)
        keep && break
    end
    keep && push!(path, q)

    return keep, path
end

"""
    nodepath(p, q)

Find the shortest path from node `p` to node `q`.
"""
function nodepath(p::T, q::T) where T <: AbstractNode
	keep, path = _nodepath(q, q, p, Vector{T}(), false)
	keep = keep ? true : p == q
	keep || throw(InvalidTopology("no path was found between the two nodes."))

	return path
end

function _nodedist(p::AbstractNode, q::AbstractNode, target::AbstractNode, d::Int, keep::Bool)
	for link in q.links
		link.to == p && continue
		if link.to == target
			keep = true
			d += 1
			break
		end
		istip(link.to) && continue
		keep, d = _nodedist(q, link.to, target, d, keep)
		keep && break
	end
	d += keep ? 1 : 0

	return keep, d
end

"""
    nodedist(p, q)

Count the nodes along the shortest path between node `p` and node `q`.
"""
function nodedist(p::AbstractNode, q::AbstractNode)
	keep, d = _nodedist(q, q, p, 0, false)
	keep = keep ? true : p == q
	keep || throw(InvalidTopology("no path was found between the two nodes."))

	return d > 0 ? d - 2 : 0
end

function _brdist(p::AbstractNode, q::AbstractNode, target::AbstractNode, d::Float64, keep::Bool)
	for link in q.links
		link.to == p && continue
		if link.to == target
			keep = true
			d += link.branch.length
			break
		end
		istip(link.to) && continue
		keep, d = _brdist(q, link.to, target, d, keep)
		keep && break
	end
	d += keep ? brlength(p, q) : 0.0

	return keep, d
end

"""
	brdist(p, q)

Get the sum of the lengths of the branches along the shortest path between node `p` and node `q`.
"""
function brdist(p::AbstractNode, target::AbstractNode)
	d = 0.0
	keep = false
	for link in p.links
		if link.to == target
			return link.branch.length
		end
		istip(link.to) && continue
		keep, d = _brdist(p, link.to, target, d, keep)
		keep && break
	end

	keep = keep ? true : p == target
	keep || throw(InvalidTopology("no path was found between the two nodes."))

	return d
end
# brdist(p::AbstractNode, q::AbstractNode) =
# 	_brdist(q, q, p, 0.0, false)[2]


module SobolSearches

export SobolSearch, setup_sobol_search, sobol_search_next!, sobol_search!, sobol_search

using ArgCheck: @argcheck
using DocStringExtensions: SIGNATURES
using Parameters: @unpack
import Sobol

struct SobolSearch{TM,TF,TV,TS} <: AbstractVector{TM}
    minima::Vector{TM}
    minimand::TF
    lower_bounds::TV
    upper_bounds::TV
    sobol_seq::TS
end

function _transform_position(lower_bounds, upper_bounds, s)
    map((l, u, α) -> l*(1-α) + u*α, lower_bounds, upper_bounds, s)
end

function _evaluated_position(minimand, lower_bounds, upper_bounds, s)
    position = _transform_position(lower_bounds, upper_bounds, s)
    (value = minimand(position), position = position)
end

function setup_sobol_search(minimand, lower_bounds, upper_bounds, keep;
                            skip = 0)
    @argcheck length(lower_bounds) == length(upper_bounds)
    @argcheck all(lower_bounds .< upper_bounds)
    @argcheck keep ≥ 1
    sobol_seq = Sobol.SobolSeq(length(lower_bounds))
    skip > 0 && Sobol.skip(sobol_seq, skip)
    minima = [_evaluated_position(minimand, lower_bounds, upper_bounds, Sobol.next!(sobol_seq))
              for _ in 1:keep]
    sort!(minima, by = x -> x.value)
    SobolSearch(minima, minimand, lower_bounds, upper_bounds, sobol_seq)
end

function sobol_search_next!(search::SobolSearch)
    @unpack lower_bounds, upper_bounds, minimand, sobol_seq, minima = search
    candidate = _evaluated_position(minimand, lower_bounds, upper_bounds,
                                    Sobol.next!(sobol_seq))
    len = length(minima)
    if candidate.value < minima[len].value
        n = searchsortedfirst(minima, candidate, by = x -> x.value)
        for i in (len - 1):-1:n
            minima[i + 1] = minima[i]
        end
        minima[n] = candidate
        true
    else
        false
    end
end

function sobol_search!(search::SobolSearch, steps)
    for _ in 1:steps
        sobol_search_next!(search)
    end
    search
end

"""
$(SIGNATURES)

Find the `keep` lowers minima of `miniman` along a Sobol sequence of length `steps` in the
hypercube between `lower_bounds` and `upper_bounds`.

Return a [`SobolSearch`](@ref) object.

```julia
```
"""
function sobol_search(minimand, lower_bounds, upper_bounds, keep, steps; skip = 0)
    search = setup_sobol_search(minimand, lower_bounds, upper_bounds, keep;
                                skip = skip)
    sobol_search!(search, steps)
end

Base.IndexStyle(::Type{<:SobolSearch{TM}}) where {TM} = Base.IndexStyle(Vector{TM})

Base.size(search::SobolSearch) = size(search.minima)

Base.getindex(search::SobolSearch, I) = getindex(search.minima, I)

end # module


@testset "118.pascals-triangle.jl" begin
    @test pascals_triangle(1) == [[1]]
    @test pascals_triangle(2) == [[1], [1, 1]]
    @test pascals_triangle(5) == [[1], [1, 1], [1, 2, 1], [1, 3, 3, 1], [1, 4, 6, 4, 1]]
end


using Distributions
using Test
using StableRNGs

@testset "Categorical" begin

for p in Any[
    [0.5, 0.5],
    [0.5f0, 0.5f0],
    [1//2, 1//2],
    [0.1, 0.3, 0.2, 0.4],
    [0.15, 0.25, 0.6] ]

    d = Categorical(p)
    k = length(p)
    println("    testing $d as Categorical")

    @test isa(d, Categorical)
    @test probs(d) == p
    @test minimum(d) == 1
    @test maximum(d) == k
    @test extrema(d) == (1, k)
    @test ncategories(d) == k

    c = 0.0
    for i = 1:k
        c += p[i]
        @test @inferred(pdf(d, i)) == p[i]
        @test @inferred(pdf(d, float(i))) == p[i]
        @test @inferred(logpdf(d, i)) === log(p[i])
        @test @inferred(logpdf(d, float(i))) === log(p[i])
        @test @inferred(cdf(d, i))  ≈ c
        @test @inferred(cdf(d, i + 0.5)) ≈ c
        @test @inferred(ccdf(d, i)) ≈ 1 - c
        @test @inferred(ccdf(d, i + 0.5)) ≈ 1 - c
        @test @inferred(logcdf(d, i)) ≈ log(c)
        @test @inferred(logcdf(d, i + 0.5)) ≈ log(c)
        @test @inferred(logccdf(d, i)) ≈ log1p(-c)
        @test @inferred(logccdf(d, i + 0.5)) ≈ log1p(-c)
    end

    @test pdf(d, 0) == 0
    @test pdf(d, k+1) == 0
    @test logpdf(d, 0) == -Inf
    @test logpdf(d, k+1) == -Inf
    @test iszero(cdf(d, -Inf))
    @test iszero(cdf(d, 0))
    @test isone(cdf(d, k+1))
    @test isone(cdf(d, Inf))
    @test isnan(cdf(d, NaN))
    @test isone(ccdf(d, -Inf))
    @test isone(ccdf(d, 0))
    @test iszero(ccdf(d, k+1))
    @test iszero(ccdf(d, Inf))
    @test isnan(ccdf(d, NaN))

    @test pdf.(d, support(d)) == p
    @test pdf.(d, 1:k) == p

    # The test utilities are currently only able to handle Float64s
    if partype(d) === Float64
        test_distr(d, 10^6)
    end
end

d = Categorical(4)
println("    testing $d as Categorical")
@test minimum(d) == 1
@test maximum(d) == 4
@test extrema(d) == (1, 4)
@test probs(d) == [0.25, 0.25, 0.25, 0.25]

p = ones(10^6) * 1.0e-6
@test Distributions.isprobvec(p)

@test convert(Categorical{Float64,Vector{Float64}}, d) === d
for x in (d, probs(d))
    d32 = convert(Categorical{Float32,Vector{Float32}}, d)
    @test d32 isa Categorical{Float32,Vector{Float32}}
    @test probs(d32) == map(Float32, probs(d))
end

@testset "test args... constructor" begin
    @test Categorical(0.3, 0.7) == Categorical([0.3, 0.7])
end

@testset "reproducibility across julia versions" begin
    d= Categorical([0.1, 0.2, 0.7])
    rng = StableRNGs.StableRNG(600)
    @test rand(rng, d, 10) == [2, 1, 3, 3, 2, 3, 3, 3, 3, 3]
end

end


"""
    Kasai's algorithm for linear-time construction of LCP array from Suffix Array
"""
function lcparray{T<:Integer}(sa::Vector{Int}, text::Vector{T})
    n = length(sa)
    lcps = similar(sa)
    rank = similar(sa)
    for i = 1:n
        rank[sa[i]] = i
    end

    lcp = 0
    for i = 1:n
        rank[i] > 1 || continue
        j = sa[rank[i]-1]
        while i+lcp <= n && j+lcp <= n && text[i+lcp] == text[j+lcp]
            lcp += 1
        end
        lcps[rank[i]] = lcp
        lcp > 0 && (lcp -= 1)
    end
    lcps[1] = 0
    lcps
end


using .TurbulenceClosures: ▶xy_cfa, ▶xy_fca

#####
##### Functions for non-rotating models
#####

@inline x_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
@inline y_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, ::Nothing, U) where FT = zero(FT)

#####
##### The 'FPlane' approximation. This is equivalent to a model with a constant
##### rotation rate around its vertical axis.
#####

"""
    FPlane{FT} <: AbstractRotation

A parameter object for constant rotation around a vertical axis.
"""
struct FPlane{FT} <: AbstractRotation
    f :: FT
end

"""
    FPlane([FT=Float64;] f=nothing, rotation_rate=nothing, latitude=nothing)

Returns a parameter object for constant rotation at the angular frequency
`f/2`, and therefore with background vorticity `f`, around a vertical axis.
If `f` is not specified, it is calculated from `rotation_rate` and
`latitude` according to the relation `f = 2*rotation_rate*sind(latitude).

Also called `FPlane`, after the "f-plane" approximation for the local effect of
Earth's rotation in a planar coordinate system tangent to the Earth's surface.
"""
function FPlane(FT=Float64; f=nothing, rotation_rate=nothing, latitude=nothing)

    if f == nothing && rotation_rate != nothing && latitude != nothing
        return FPlane{FT}(2rotation_rate*sind(latitude))
    elseif f != nothing && rotation_rate == nothing && latitude == nothing
        return FPlane{FT}(f)
    else
        throw(ArgumentError("Either both keywords rotation_rate and
                             latitude must be specified, *or* only f
                             must be specified."))
    end
end

@inline x_f_cross_U(i, j, k, grid, coriolis::FPlane, U) = - coriolis.f * ▶xy_fca(i, j, k, grid, U.v)
@inline y_f_cross_U(i, j, k, grid, coriolis::FPlane, U) =   coriolis.f * ▶xy_cfa(i, j, k, grid, U.u)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::FPlane, U) where FT = zero(FT)

#####
##### The Beta Plane
#####

"""
    BetaPlane{T} <: AbstractRotation

A parameter object for meridionally increasing Coriolis
parameter (`f = f₀ + βy`).
"""
struct BetaPlane{T} <: AbstractRotation
    f₀ :: T
     β :: T
end

"""
    BetaPlane([T=Float64;] f₀=nothing, β=nothing,
                           rotation_rate=nothing, latitude=nothing, radius=nothing)

A parameter object for meridionally increasing Coriolis parameter (`f = f₀ + βy`).

The user may specify both `f₀` and `β`, or the three parameters
`rotation_rate`, `latitude`, and `radius` that specify the rotation rate and radius
of a planet, and the central latitude at which the `β`-plane approximation is to be made.
"""
function BetaPlane(T=Float64; f₀=nothing, β=nothing,
                              rotation_rate=nothing, latitude=nothing, radius=nothing)

    f_and_β = f₀ != nothing && β != nothing
    planet_parameters = rotation_rate != nothing && latitude != nothing && radius != nothing

    (f_and_β && all(Tuple(p === nothing for p in (rotation_rate, latitude, radius)))) ||
    (planet_parameters && all(Tuple(p === nothing for p in (f₀, β)))) ||
        throw(ArgumentError("Either both keywords f₀ and β must be specified,
                            *or* all of rotation_rate, latitude, and radius."))

    if planet_parameters
        f₀ = 2rotation_rate * sind(latitude)
         β = 2rotation_rate * cosd(latitude) / radius
     end

    return BetaPlane{T}(f₀, β)
end

@inline x_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
    @inbounds - (coriolis.f₀ + coriolis.β * grid.yC[j]) * ▶xy_fca(i, j, k, grid, U.v)
@inline y_f_cross_U(i, j, k, grid, coriolis::BetaPlane, U) =
    @inbounds   (coriolis.f₀ + coriolis.β * grid.yF[j]) * ▶xy_cfa(i, j, k, grid, U.u)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{FT}, coriolis::BetaPlane, U) where FT = zero(FT)


# ---
# title: 1252. Cells with Odd Values in a Matrix
# id: problem1252
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/cells-with-odd-values-in-a-matrix/description/>
# hidden: true
# ---
#
# Given `n` and `m` which are the dimensions of a matrix initialized by zeros
# and given an array `indices` where `indices[i] = [ri, ci]`. For each pair of
# `[ri, ci]` you have to increment all cells in row `ri` and column `ci` by 1.
#
# Return _the number of cells with odd values_ in the matrix after applying the
# increment to all `indices`.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2019/10/30/e1.png)
#
#
#
#     Input: n = 2, m = 3, indices = [[0,1],[1,1]]
#     Output: 6
#     Explanation: Initial matrix = [[0,0,0],[0,0,0]].
#     After applying first increment it becomes [[1,2,1],[0,1,0]].
#     The final matrix will be [[1,3,1],[1,3,1]] which contains 6 odd numbers.
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2019/10/30/e2.png)
#
#
#
#     Input: n = 2, m = 2, indices = [[1,1],[0,0]]
#     Output: 0
#     Explanation: Final matrix = [[2,2],[2,2]]. There is no odd number in the final matrix.
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 50`
#   * `1 <= m <= 50`
#   * `1 <= indices.length <= 100`
#   * `0 <= indices[i][0] < n`
#   * `0 <= indices[i][1] < m`
#
#
## @lc code=start
using LeetCode

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


function convergence_parameters(p::Param , q0::Int)

    #Base number of nodes
    N0 = 2^q0 + 1

    #Create parameter list
    sim_param = []

    for q in 1:3

        #Number of nodes
        Nx = (2^(q-1))*(N0-1) + 1

        #Name of file
        filename    = string( p.fname , "_Nx_", Nx , "_Ny_", Nx , "_tmax_" , p.tmax)
        #Create simulation parameters
        s   = Param( p ,  xnodes = Nx , ynodes = Nx , fname = filename )

        #Add to list
        push!(sim_param , s)
    end

    return sim_param

end

function test_convergence( p::Param , q0::Int )

    #Create parameter list
    sim_param = convergence_parameters(p, q0)

    for x in sim_param
        solve_wave_equation_2D(x)
    end

end

function get_data_costum(p::Param , itt::Int , q0::Int, q::Int)

    #Get data at a given time
    ttt , xxx , yyy , zzz  , _ =  get_fields(p , itt)

    N0 = 2^q0 + 1

    #Obtain common nodes
    x_out   = [ xxx[ j + ( 2^(q-q0) - 1 )*(j-1)  ] for j in 1:N0  ]
    y_out   = [ yyy[ j + ( 2^(q-q0) - 1 )*(j-1)  ] for j in 1:N0  ]

    ψ = reshape(zzz , (p.xnodes , p.ynodes) )
    ψ_out  = zeros(N0,N0)

    for i in 1:N0 , j in 1:N0
        ψ_out[i,j]  = ψ[ i + ( 2^(q-q0) - 1 )*(i-1)  , j + ( 2^(q-q0) - 1 )*(j-1) ]
    end

    return x_out , y_out , ψ_out
end

function explore_convergence(p::Param , q0::Int)

    sim_param = convergence_parameters(p , q0)

    #Max iteration
    max_iter , _ = get_t_max(sim_param[1])

    #Create Figure
    fig = Figure( resolution = (1000 , 950), font  = "Helvetica")

    ax1 = fig[4,1]      = Axis(fig  , width = 600 , height = 600 )
    ax2 = fig[4,2]      = Axis(fig  , width = 200 , height = 600 , xlabelvisible = :false , yaxisposition = :right)
    ax3 = fig[1:3,1]    = Axis(fig , width = 600 , height = 200 , yticksvisible = true , xaxisposition = :top)

    linkyaxes!(ax1, ax2)
    linkxaxes!(ax1, ax3)

    slider1 = Slider( fig[1,2] , range = 1:1:max_iter , height = 200/3 , width = 200)
    slider2 = Slider( fig[2,2] , range = 1:1:sim_param[1].xnodes  , height = 200/3 ,  width = 200)
    slider3 = Slider( fig[3,2] , range = 1:1:sim_param[1].ynodes  , height = 200/3 , width = 200)

    #Data for plotting

    x , y , _ = get_data_costum(sim_param[1] , 1 , q0 , q0 )

    #Get fields at a certain iteration that is controled by slider 1
    ψ1 =  lift(slider1.value) do v
        get_data_costum(sim_param[1] , v[] , q0 ,q0 )[3]
#            get_fields( sim_param[1] , v[] )[4]
        end

    ψ2 =  lift(slider1.value) do v
        get_data_costum(sim_param[2] , v[] , q0 ,q0+1 )[3]
        end

    ψ3 =  lift(slider1.value) do v
        get_data_costum(sim_param[3] , v[] , q0 ,q0+2 )[3]
        end

    #Get slices and differences for convergence analysis and scale them
    scale_factor = 100
    ψ1_x_slice = lift(slider2.value) do index
        scale_factor*(ψ1[][index,:] -ψ2[][index,:])
    end
    ψ2_x_slice = lift(slider2.value) do index
        scale_factor*2*(ψ2[][index,:] - ψ3[][index,:])
    end

    ψ3_x_slice = lift(slider2.value) do index
        scale_factor*4*(ψ2[][index,:] - ψ3[][index,:])
    end

    ψ4_x_slice = lift(slider2.value) do index
        scale_factor*16*(ψ2[][index,:] - ψ3[][index,:])
    end

    scale_factor = 100
    ψ1_y_slice = lift(slider3.value) do index
        scale_factor*(ψ1[][:,index] -ψ2[][:,index])
    end
    ψ2_y_slice = lift(slider3.value) do index
        scale_factor*2*(ψ2[][:,index] - ψ3[][:,index])
    end

    ψ3_y_slice = lift(slider3.value) do index
        scale_factor*4*(ψ2[][:,index] - ψ3[][:,index])
    end

    ψ4_y_slice = lift(slider3.value) do index
        scale_factor*16*(ψ2[][:,index] - ψ3[][:,index])
    end

    #Get vertical and horizontal lines to see cuts on histogram
    x_coord = lift(slider2.value) do index
        [ x[index] for i in 1:(2^q0+1) ]
    end

    y_coord = lift(slider3.value) do index
        [ y[index] for i in 1:(2^q0+1) ]
    end

    #y3 = lift(slider3.value) do index
    #    y[][:,index]
    #end

    #Plot histogram of lowest resolution and slices
    heatmap!(ax1 , x , y ,ψ1 ,colorrange=(-5,5) , colormap = :vik )
    lines!(ax1 , x_coord  , x )
    lines!(ax1 , x  , y_coord )

    #Plot x-slices with convergence analysis
    l1 = lines!(ax2 , ψ1_x_slice  , x,  linewidth  = 5.0  ,color = RGBA{Float32}(0.0f0, 0.8f0 , 0.9f0 , 1.0f0)) # color = ψ1_x_slice, colormap = :vik)
    l2 = lines!(ax2 , ψ2_x_slice  , x , color = :black , linestyle = :solid )
    l3 = lines!(ax2 , ψ3_x_slice  , x , color = :red   , linestyle = :solid)
    l4 = lines!(ax2 , ψ4_x_slice  , x , color = :green , linestyle = :solid)
    axislegend(ax1, [l1, l2 ,l3 ,l4], ["ψ1-ψ2", "2(ψ2-ψ3)" , "4(ψ2-ψ3)" , "16(ψ2-ψ3)"], position = :rb,
    orientation = :horizontal)

    #Plot y slices
    lines!(ax3 , y , ψ1_y_slice  ,  linewidth  = 5.0  ,color = RGBA{Float32}(0.0f0, 0.8f0 , 0.9f0 , 1.0f0)) # color = ψ1_x_slice, colormap = :vik)
    lines!(ax3 , y , ψ2_y_slice  ,  color = :black , linestyle = :solid )
    lines!(ax3 , y , ψ3_y_slice  ,  color = :red   , linestyle = :solid)
    lines!(ax3 , y , ψ4_y_slice  ,  color = :green , linestyle = :solid)

    #Adjust limits of heatmap
    xlims!(ax1 , p.xmin ,p.xmax)
    ylims!(ax1 , p.ymin ,p.ymax)

    #Adjust limits of vertical
    xlims!(ax2 , -1 ,1)
    ylims!(ax3 , -1 ,1)

    #Do heatmap

    display(fig)

end

function see_energy_covergence(p::Param , q0::Int)

    sim_param = convergence_parameters(p , q0)

    #Max iteration
    max_iter , _ = get_t_max(sim_param[2])

    #Create Figure
    fig = Figure( resolution = (1000 , 950), font  = "Helvetica")
    ax = fig[1,1] = Axis(fig , title = "Test")

    ls = []
cols = [:red , :blue , :black]
    for (ix , x) in enumerate(sim_param)
        t , x = get_energy(x)
        push!(ls , lines!(ax , t , x  , color = cols[ix]) )
    end

    axislegend(ax, ls, ["1" , "2" , "3"], position = :rb)

     #Adjust limits of heatmap
     ylims!(ax , 0,10)
    display(fig)

end

function plot_energy(p::Param , savetrue::Bool , name::String)

    #Max iteration
    max_iter , _ = get_t_max(p)

    #Create Figure
    fig = Figure( resolution = (800 , 450), font  = "Helvetica")
    ax = fig[1,1] = Axis(fig)

    t , x = get_energy(p)

    lines!(ax , t , x )

    ax.ylabel = "Energy Density "
    ax.xlabel = "Time"

    ylims!(ax , 0,10)
    display(fig)

    if(savetrue)
        save( name , fig )
    end

end


module Periodic_Cov
using Statistics, LinearAlgebra
using GaussianRandomFields
import GaussianRandomFields: shortname, apply, apply_symmetric!, apply_non_symmetric!

abstract type Distance end

shortname(::Distance) = "custom distance"
function apply(::Distance, x::NTuple{d, Real}, y::NTuple{d, Real}) where d end

struct CustomDistanceCovarianceStructure{T} <: CovarianceStructure{T}
    distance::Distance
    cov::IsotropicCovarianceStructure{T}
end

apply(s::CustomDistanceCovarianceStructure, ρ::Real) = apply(s.cov, ρ)

shortname(s::CustomDistanceCovarianceStructure) = shortname(s.cov) * " on " * shortname(s.distance)

Statistics.std(c::IsotropicCovarianceStructure) = c.σ
Statistics.std(s::CustomDistanceCovarianceStructure) = Statistics.std(s.cov)
Statistics.std(f::AbstractCovarianceFunction) = Statistics.std(f.cov)

function apply_symmetric!(
    C::Matrix,
    s::CustomDistanceCovarianceStructure,
    x::NTuple{d, AbstractVector},
    y::NTuple{d, AbstractVector}
) where d
    xiterator = enumerate(Iterators.product(x...))
    for (j, idy) in enumerate(Iterators.product(y...))
        for (i, idx) in Iterators.take(xiterator, j)
            @inbounds C[i, j] = apply(s.cov, apply(s.distance, idx, idy))
        end
    end
    Symmetric(C, :U)
end

function apply_non_symmetric!(
    C::Matrix,
    s::CustomDistanceCovarianceStructure,
    x::NTuple{d, AbstractVector},
    y::NTuple{d, AbstractVector}
) where d
    xiterator = enumerate(Iterators.product(x...))
    for (j, idy) in enumerate(Iterators.product(y...))
        for (i, idx) in xiterator
            @inbounds C[i, j] = apply(s.cov, apply(s.distance, idx, idy))
        end
    end
    C
end

struct SphericalDistance <: Distance end

shortname(::SphericalDistance) = "spherical distance"

function apply(::SphericalDistance, (λ₁, φ₁)::Tuple{Float64, Float64}, (λ₂, φ₂)::Tuple{Float64, Float64})
    Δλ = λ₂ - λ₁  # longitude difference
    Δφ = φ₂ - φ₁  # latitude difference
    # haversine formula
    a = sin(Δφ/2)^2 + cos(φ₁)*cos(φ₂)*sin(Δλ/2)^2
    # distance
    2 * ( asin( min(√a, one(a)) ))
end

struct EuclideanDistance <: Distance end
shortname(::EuclideanDistance) = "Euclidean distance"
apply(::EuclideanDistance, x::NTuple{d, Real}, y::NTuple{d, Real}) where d = norm(x .- y)
end


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module QueensProblems

export ModuleQueensProblems
export L319284, Queens

"""

For some background see: [Backtracking with profiles](https://wp.me/paipV7-E).

* Queens, L319284
"""
const ModuleQueensProblems = ""

function solve!(profile, level, size, start, cols, diag4, diag1)

    if level > 0
        for i ∈ start:size
            save = cols & (1 << i) + diag1 & (1 << (i + level)) +
                   diag4 & (1 << (32 + i - level))

            if save == 0

                cols = xor(cols, 1 << i)
                diag1 = xor(diag1, 1 << (i + level))
                diag4 = xor(diag4, 1 << (32 + i - level))

                solve!(profile, level - 1, size, 0, cols, diag4, diag1)

                cols = xor(cols, 1 << i)
                diag1 = xor(diag1, 1 << (i + level))
                diag4 = xor(diag4, 1 << (32 + i - level))

                profile[level+1] += 1
            end
        end
    else
        for i ∈ 0:size
            save = cols & (1 << i) + diag1 & (1 << i) + diag4 & (1 << (32 + i))
            save == 0 && (profile[1] += 1)
        end
    end
end

function search(n::Int)
    profile = zeros(Int, n + 1)
    cols = diag4 = diag1 = Int(0)
    solve!(profile, n - 1, n - 1, 0, cols, diag4, diag1)
    return profile
end

"""

Returns the profile of the backtrack tree for the n queens problem (see `A319284`).
"""
function Queens(n::Int)
    n == 0 && return [1]
    profile = search(n)
    profile[n+1] = 1  # add the root
    [profile[n-i+1] for i = 0:n]
end

"""

Returns the profile of the backtrack tree for the n queens problem.
"""
L319284(n) = Queens(n)

#START-TEST-########################################################

using Test

function test()
    @testset "QueensProblem" begin
        levels = [1, 10, 72, 364, 1400, 3916, 7552, 9632, 7828, 4040, 724]
        @test all(L319284(10) .== levels)
    end
end

function demo()

    up_to = 10
    for n ∈ 0:up_to
        print("elapsed: ")
        @time profile = Queens(n)
        println("size:      ", n)
        println("profile:   ", profile)
        println("nodes:     ", sum(profile))
        println("solutions: ", profile[n+1])
        println()
    end
end

"""

L319284(15) :: 18.731622 seconds (4 allocations: 480 bytes)
"""
function perf()
    @time L319284(15)
end

function main()
    test()
    demo()
    #perf()
end

main()

end # module


# Plotting examples
# =================
# ## Plotters
#  All plot functions in ExtendableGrids.jl have a `Plotter` keyword argument
#  which defaults to `nothing`.  This allows to pass a module as plotting backend
#  without creating a dependency. Fully supported are `PyPlot` and `GLMakie`.
#  `WGLMakie` and `CairoMakie` work in principle but in the moment don't deliver
#  all necessary functionality.
#
# Also supported is [`VTKView`](https://github.com/j-fu/VTKView.jl)  which is exeprimental and works only on linux
#
# Generally, grids and P1 FEM functions on grids can be plotted up to now.
#
# ## Grid plots
# Here, we define some sample grids for plotting purposes.
function grid1d(;n=50)
    X=collect(0:1/n:1)
    g=simplexgrid(X)
end

function grid2d(;n=20)
    X=collect(0:1/n:1)
    g=simplexgrid(X,X)
end

function grid3d(;n=15)
    X=collect(0:1/n:1)
    g=simplexgrid(X,X,X)
end
#
# Now, we can use the plot command of ExtendableGrids to plot grids
# Note the kwargs `xplane`, `yplane` and `zplane` which allow to control
# cutplanes which peel off some elements from the grid in 3d and allow to
# explore the inner triangulation.
#
# For Makie and VTKView, the cutplane values can be controlled interactively.
function plotting_grid3d(;Plotter=nothing)
    visualize(grid3d(), Plotter=Plotter, zplane=0.49)
end
# ![](plotting_grid3d.svg)

# ## Function plots
# Let us define some functions
function func1d(;n=50)
    g=grid1d(n=n)
    g,map(x->sinpi(2*x[1]),g)
end

function func2d(;n=20)
    g=grid2d(n=n)
    g,map((x,y)->sinpi(2*x)*sinpi(3.5*y),g)
end

function func3d(;n=15)
    g=grid3d(n=n)
    g, map((x,y,z)->sinpi(2*x)*sinpi(3.5*y)*sinpi(1.5*z),g)
end

#
# Plotting a function then goes as follows:
# `xplane`, `yplane` and `zplane` now define cut planes where
# the function projection is plotted as a heatmap.
# The additional `flevel` keyword argument allows
# to control an isolevel.
#
# For Makie and VTKView, the cutplane values and the flevel can be controlled interactively.
function plotting_func3d(;Plotter=nothing)
    g,f=func3d()
    visualize(g,f, Plotter=Plotter, zplane=0.49,xplane=0.49,flevel=0.25)
end
# ![](plotting_func3d.svg)

# ## Multiscene plots
# We can combine multiple plots into one scene according to
# some layout grid given by the layout parameter.
#
# The ',' key for GLMakie and the '*' key for VTKView allow to
# switch between gallery view (default) and focused view of only
# one subscene.
function plotting_multiscene(;Plotter=nothing)
    p=GridVisualizer(;Plotter=Plotter,layout=(2,3),clear=true,resolution=(800,500))

    # 3D must come first for Makie, see https://github.com/JuliaPlots/Makie.jl/issues/805
    visualize!(p[1,3],grid3d(),zplane=0.49,title="3D grid")
    visualize!(p[2,3],func3d()...,zplane=0.49,flevel=0.5,colormap=:bamako, title="3D grid function")
    visualize!(p[1,1],grid1d(), title="1D grid")
    visualize!(p[2,1],func1d()..., title="1D grid function")
    visualize!(p[1,2],grid2d(),title="2D grid")
    visualize!(p[2,2],func2d()...,colormap=:bamako,title="2D grid function")
    reveal(p)
end
# ![](plotting_multiscene.svg)


export PowerDivergenceTest, ChisqTest, MultinomialLRTest

const Levels{T} = Tuple{UnitRange{T},UnitRange{T}}

function boundproportion(x::T) where T<:Real
    max(min(convert(Float64,x),1.0),0.0)
end

struct PowerDivergenceTest <: HypothesisTest
    lambda::Float64
    theta0::Vector{Float64}
    stat::Float64
    df::Int64
    observed::Matrix{Int64}
    n::Int64
    thetahat::Vector{Float64}

    expected::Matrix{Float64}
    residuals::Matrix{Float64}
    stdresiduals::Matrix{Float64}
end

#test name
function testname(x::PowerDivergenceTest)
    if x.lambda == 1
        return "Pearson's Chi-square Test"
    elseif x.lambda == 0
        return "Multinomial Likelihood Ratio Test"
    elseif x.lambda == -1
        return "Minimum Discrimination Information Test"
    elseif x.lambda == -2
        return "Neyman Modified Chi-square Test"
    elseif x.lambda == -0.5
        return "Freeman-Tukey Test"
    else
        return "Power Divergence Test"
    end
end

# parameter of interest: name, value under h0, point estimate
population_param_of_interest(x::PowerDivergenceTest) = ("Multinomial Probabilities", x.theta0, x.thetahat)
default_tail(test::PowerDivergenceTest) = :right

pvalue(x::PowerDivergenceTest; tail=:right) = pvalue(Chisq(x.df),x.stat; tail=tail)

"""
    confint(test::PowerDivergenceTest, alpha = 0.05; tail = :both, method = :sison_glaz)

Compute a confidence interval with coverage 1-`alpha` for multinomial proportions using
one of the following methods. Possible values for `method` are:

  - `:sison_glaz` (default): Sison-Glaz intervals
  - `:bootstrap`: Bootstrap intervals
  - `:quesenberry_hurst`: Quesenberry-Hurst intervals
  - `:gold`: Gold intervals (asymptotic simultaneous intervals)

# References

  * Agresti, Alan. Categorical Data Analysis, 3rd Edition. Wiley, 2013.
  * Sison, C.P and Glaz, J. Simultaneous confidence intervals and sample size determination
    for multinomial proportions. Journal of the American Statistical Association,
    90:366-369, 1995.
  * Quesensberry, C.P. and Hurst, D.C. Large Sample Simultaneous Confidence Intervals for
    Multinational Proportions. Technometrics, 6:191-195, 1964.
  * Gold, R. Z. Tests Auxiliary to ``χ^2`` Tests in a Markov Chain. Annals of
    Mathematical Statistics, 30:56-74, 1963.
"""
function StatsBase.confint(x::PowerDivergenceTest, alpha::Float64=0.05;
                           tail::Symbol=:both, method::Symbol=:sison_glaz, correct::Bool=true,
                           bootstrap_iters::Int64=10000, GC::Bool=true)
    check_alpha(alpha)

    m  = length(x.thetahat)

    if tail == :left
        i = StatsBase.confint(x, alpha*2,method=method, GC=GC)
        Tuple{Float64,Float64}[(0.0, i[j][2]) for j in 1:m]
    elseif tail == :right
        i = StatsBase.confint(x, alpha*2,method=method, GC=GC)
        Tuple{Float64,Float64}[(i[j][1], 1.0) for j in 1:m]
    elseif tail == :both
        if method == :gold
            ci_gold(x,alpha,correct=correct,GC=GC)
        elseif method == :sison_glaz
            ci_sison_glaz(x,alpha, skew_correct=correct)
        elseif method == :quesenberry_hurst
            ci_quesenberry_hurst(x,alpha,GC=GC)
        elseif method == :bootstrap
            ci_bootstrap(x,alpha,bootstrap_iters)
        else
            throw(ArgumentError("method=$(method) is invalid or not implemented yet"))
        end
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end
end

# Bootstrap
function ci_bootstrap(x::PowerDivergenceTest,alpha::Float64, iters::Int64)
    m = mapslices(x -> quantile(x, [alpha / 2, 1 - alpha / 2]), rand(Multinomial(x.n, convert(Vector{Float64}, x.thetahat)),iters) / x.n, dims=2)
    Tuple{Float64,Float64}[(boundproportion(m[i,1]), boundproportion(m[i,2])) for i in 1:length(x.thetahat)]
end

# Quesenberry, Hurst (1964)
function ci_quesenberry_hurst(x::PowerDivergenceTest, alpha::Float64; GC::Bool=true)
    m  = length(x.thetahat)
    cv = GC ? quantile(Chisq(1), 1 - alpha / m) : quantile(Chisq(m - 1), 1 - alpha)

    u = (cv .+ 2 .* x.observed .+ sqrt.(cv .* (cv .+ 4 .* x.observed .* (x.n .- x.observed) ./ x.n))) ./ (2 .* (x.n .+ cv))
    l = (cv .+ 2 .* x.observed .- sqrt.(cv .* (cv .+ 4 .* x.observed .* (x.n .- x.observed) ./ x.n))) ./ (2 .* (x.n .+ cv))
    Tuple{Float64,Float64}[(boundproportion(l[j]), boundproportion(u[j])) for j in 1:m]
end

# asymptotic simultaneous confidence interval
# Gold (1963)
function ci_gold(x::PowerDivergenceTest, alpha::Float64; correct::Bool=true, GC::Bool=true)
    m  = length(x.thetahat)
    cv = GC ? quantile(Chisq(1), 1 - alpha / 2m) : quantile(Chisq(m - 1), 1 - alpha / 2)

    u = [x.thetahat[j] + sqrt.(cv * x.thetahat[j] * (1 - x.thetahat[j]) / x.n) + ifelse(correct, inv(2x.n), 0) for j in 1:m]
    l = [x.thetahat[j] - sqrt.(cv * x.thetahat[j] * (1 - x.thetahat[j]) / x.n) - ifelse(correct, inv(2x.n), 0) for j in 1:m]
    Tuple{Float64,Float64}[ (boundproportion(l[j]), boundproportion(u[j])) for j in 1:m]
end

# Simultaneous confidence interval
# Sison, Glaz (1995)
# adapted from SAS macro by May, Johnson (2000)
function ci_sison_glaz(x::PowerDivergenceTest, alpha::Float64; skew_correct::Bool=true)
    k = length(x.thetahat)
    probn = inv(pdf(Poisson(x.n), x.n))

    c = 0
    p = 0.0
    p_old = 0.0
    m1, m2, m3, m4, m5 = zeros(k), zeros(k), zeros(k), zeros(k), zeros(k)
    mu = zeros(4)

    for _c in 1:x.n
        #run truncpoi
        for i in 1:k
            lambda = x.observed[i]
            #run moments
            a = lambda + _c
            b = max(lambda - _c, 0)
            poislama = cdf(Poisson(lambda), a)
            poislamb = cdf(Poisson(lambda), b - 1)
            den = b > 0.0 ? poislama-poislamb : poislama

            for r in 1:4
                plar = cdf(Poisson(lambda), a - r)
                plbr = cdf(Poisson(lambda), b - r - 1)

                poisA = ifelse( (a - r) >= 0, poislama - plar, poislama)
                poisB = 0.0
                if  (b - r - 1) >= 0
                    poisB = poislamb - plbr
                end
                if (b - r - 1) < 0 && b - 1 >= 0
                    poisB = poislamb
                end
                if (b - r - 1) < 0 && (b - 1) < 0
                    poisB = 0.0
                end
                mu[r] = lambda^r * (1 - (poisA - poisB) / den)
            end
            # end of moments
            m1[i] = mu[1]
            m2[i] = mu[2] + mu[1] - mu[1]^2
            m3[i] = mu[3] + mu[2] * (3 - 3mu[1]) + (mu[1] - 3mu[1]^2 + 2mu[1]^3)
            m4[i] = mu[4] + mu[3] * (6 - 4mu[1]) + mu[2] * (7 - 12mu[1] + 6mu[1]^2) + mu[1] - 4mu[1]^2 + 6mu[1]^3 - 3mu[1]^4
            m5[i] = den
        end
        for i in 1:k
            m4[i] -= 3m2[i]^2
        end
        s1, s2, s3, s4 = sum(m1), sum(m2), sum(m3), sum(m4)
        z  = (x.n - s1) / sqrt(s2)
        g1 = s3 / s2^(3/2)
        g2 = s4 / s2^2

        poly = 1 + g1 * (z^3 - 3z) / 6 + g2 * (z^4 - 6z^2 + 3) / 24 + g1^2 * (z^6 - 15z^4 + 45z^2 - 15) / 72
        f = poly * exp(-z^2 / 2) / sqrt(2π)
        probx = prod(m5)

        p = probn * probx * f / sqrt(s2)
        # end of truncpoi

        if p > 1 - alpha && p_old < 1 - alpha
            c = _c
            break
        else
            c = x.n
        end
        p_old = p
    end

    delta = (1 - alpha - p_old) / (p - p_old)
    out = zeros(k, 5)
    num = zeros(k, 1)

    c -= 1

    vol1 = 1
    vol2 = 1
    for i in 1:k
        num[i,1] = i
        cn = c / x.n
        onen = 1 / x.n
        out[i,1] = x.thetahat[i]

        out[i,2] = max(x.thetahat[i] - cn, 0)
        out[i,3] = min(x.thetahat[i] + cn + 2delta / x.n, 1)

        out[i,4] = max(x.thetahat[i] - cn - onen, 0)
        out[i,5] = min(x.thetahat[i] + cn + onen, 1)
    end
    if skew_correct
        return Tuple{Float64,Float64}[(boundproportion(out[i,2]), boundproportion(out[i,3])) for i in 1:k]
    else
        return Tuple{Float64,Float64}[(boundproportion(out[i,4]), boundproportion(out[i,5])) for i in 1:k]
    end
end

# power divergence statistic for testing goodness of fit
# Cressie and Read 1984; Read and Cressie 1988
# lambda    =  1: Pearson's Chi-square statistic
# lambda ->  0: Converges to Likelihood Ratio test stat
# lambda -> -1: Converges to minimum discrimination information statistic (Gokhale, Kullback 1978)
# lambda =  -2: Neyman Modified chi-squared statistic (Neyman 1949)
# lambda = -.5: Freeman-Tukey statistic (Freeman, Tukey 1950)

# Under regularity conditions, their asymptotic distributions are all the same (Drost 1989)
# Chi-squared null approximation works best for lambda near 2/3
"""
    PowerDivergenceTest(x[, y]; lambda = 1.0, theta0 = ones(length(x))/length(x))

Perform a Power Divergence test.

If `y` is not given and `x` is a matrix with one row or column, or `x` is a vector, then
a goodness-of-fit test is performed (`x` is treated as a one-dimensional contingency
table). In this case, the hypothesis tested is whether the population probabilities equal
those in `theta0`, or are all equal if `theta0` is not given.

If `x` is a matrix with at least two rows and columns, it is taken as a two-dimensional
contingency table. Otherwise, `x` and `y` must be vectors of the same length. The contingency
table is calculated using the `counts` function from the `StatsBase` package. Then the power
divergence test is conducted under the null hypothesis that the joint distribution of the
cell counts in a 2-dimensional contingency table is the product of the row and column
marginals.

Note that the entries of `x` (and `y` if provided) must be non-negative integers.

The power divergence test is given by
```math
    \\dfrac{2}{λ(λ+1)}\\sum_{i=1}^I \\sum_{j=1}^J n_{ij} \\left[(n_{ij}
    /\\hat{n}_{ij})^λ -1\\right]
```
where ``n_{ij}`` is the cell count in the ``i`` th row and ``j`` th column and ``λ`` is a
real number determining the nature of the test to be performed:

  * ``λ = 1``: equal to Pearson's chi-squared statistic
  * ``λ \\to 0``: converges to the likelihood ratio test statistic
  * ``λ \\to -1``: converges to the minimum discrimination information statistic
    (Gokhale and Kullback, 1978)
  * ``λ = -2``: equals Neyman modified chi-squared (Neyman, 1949)
  * ``λ = -1/2``: equals the Freeman-Tukey statistic (Freeman and Tukey, 1950).

Under regularity conditions, the asymptotic distributions are identical (see Drost et. al.
1989). The ``χ^2`` null approximation works best for ``λ`` near ``2/3``.

Implements: [`pvalue`](@ref), [`confint(::PowerDivergenceTest)`](@ref)

# References

  * Agresti, Alan. Categorical Data Analysis, 3rd Edition. Wiley, 2013.
"""
function PowerDivergenceTest(x::AbstractMatrix{T}; lambda::U=1.0, theta0::Vector{U} = ones(length(x))/length(x)) where {T<:Integer,U<:AbstractFloat}

    nrows, ncols = size(x)
    n = sum(x)

    #validate date
    (any(x .< 0) || any(!isfinite, x)) && throw(ArgumentError("all entries must be nonnegative and finite"))
    n == 0 && throw(ArgumentError("at least one entry must be positive"))
    (!isfinite(nrows) || !isfinite(ncols) || !isfinite(nrows*ncols)) && throw(ArgumentError("invalid number of rows or columns"))

    if nrows > 1 && ncols > 1
        rowsums = sum(x, dims=2)
        colsums = sum(x, dims=1)
        df = (nrows - 1) * (ncols - 1)
        thetahat = x ./ n
        xhat = rowsums * colsums / n
        theta0 = xhat / n
        V = Float64[(colsums[j]/n) * (rowsums[i]/n) * (1 - rowsums[i]/n) * (n - colsums[j]) for i in 1:nrows, j in 1:ncols]
    elseif nrows == 1 || ncols == 1
        df = length(x) - 1
        xhat = reshape(n * theta0, size(x))
        thetahat = x / n
        V = reshape(n .* theta0 .* (1 .- theta0), size(x))

        abs(1 - sum(theta0)) <= sqrt(eps()) || throw(ArgumentError("Probabilities must sum to one"))
    else
        throw(ArgumentError("Number of rows or columns must be at least 1"))
    end

    stat = 0
    if lambda == 0
        for i in 1:length(x)
            stat += x[i] * (log(x[i]) - log(xhat[i]))
        end
        stat *= 2
    elseif lambda == -1
        for i in 1:length(x)
            stat += xhat[i] * (log(xhat[i]) - log(x[i]))
        end
        stat *= 2
    else
        for i in 1:length(x)
            stat += x[i] * ((x[i] / xhat[i])^lambda - 1)
        end
        stat *= 2 / (lambda * (lambda + 1))
    end

    PowerDivergenceTest(lambda, vec(theta0), stat, df, x, n, vec(thetahat), xhat, (x - xhat) ./ sqrt.(xhat), (x - xhat) ./ sqrt.(V))
end

#convenience functions

#PDT
function PowerDivergenceTest(x::AbstractVector{T}, y::AbstractVector{T}, levels::Levels{T}; lambda::U=1.0) where {T<:Integer,U<:AbstractFloat}
    d = counts(x, y, levels)
    PowerDivergenceTest(d, lambda=lambda)
end

function PowerDivergenceTest(x::AbstractVector{T}, y::AbstractVector{T}, k::T; lambda::U=1.0) where {T<:Integer,U<:AbstractFloat}
    d = counts(x, y, k)
    PowerDivergenceTest(d, lambda=lambda)
end

PowerDivergenceTest(x::AbstractVector{T}; lambda::U=1.0, theta0::Vector{U} = ones(length(x))/length(x)) where {T<:Integer,U<:AbstractFloat} =
    PowerDivergenceTest(reshape(x, length(x), 1), lambda=lambda, theta0=theta0)

#ChisqTest
"""
    ChisqTest(x[, y][, theta0 = ones(length(x))/length(x)])

Perform a Pearson chi-squared test (equivalent to a [`PowerDivergenceTest`](@ref)
with ``λ = 1``).

If `y` is not given and `x` is a matrix with one row or column, or `x` is a vector, then
a goodness-of-fit test is performed (`x` is treated as a one-dimensional contingency
table). In this case, the hypothesis tested is whether the population probabilities equal
those in `theta0`, or are all equal if `theta0` is not given.

If `x` is a matrix with at least two rows and columns, it is taken as a two-dimensional
contingency table. Otherwise, `x` and `y` must be vectors of the same length. The contingency
table is calculated using `counts` function from the `StatsBase` package. Then the power
divergence test is conducted under the null hypothesis that the joint distribution of the
cell counts in a 2-dimensional contingency table is the product of the row and column
marginals.

Note that the entries of `x` (and `y` if provided) must be non-negative integers.

Implements: [`pvalue`](@ref), [`confint`](@ref)
"""
function ChisqTest(x::AbstractMatrix{T}) where T<:Integer
    PowerDivergenceTest(x, lambda=1.0)
end

function ChisqTest(x::AbstractVector{T}, y::AbstractVector{T}, levels::Levels{T}) where T<:Integer
    d = counts(x, y, levels)
    PowerDivergenceTest(d, lambda=1.0)
end

function ChisqTest(x::AbstractVector{T}, y::AbstractVector{T}, k::T) where T<:Integer
    d = counts(x, y, k)
    PowerDivergenceTest(d, lambda=1.0)
end

ChisqTest(x::AbstractVector{T}, theta0::Vector{U} = ones(length(x))/length(x)) where {T<:Integer,U<:AbstractFloat} =
    PowerDivergenceTest(reshape(x, length(x), 1), lambda=1.0, theta0=theta0)

#MultinomialLRTest
"""
    MultinomialLRTest(x[, y][, theta0 = ones(length(x))/length(x)])

Perform a multinomial likelihood ratio test (equivalent to a [`PowerDivergenceTest`](@ref)
with ``λ = 0``).

If `y` is not given and `x` is a matrix with one row or column, or `x` is a vector, then
a goodness-of-fit test is performed (`x` is treated as a one-dimensional contingency
table). In this case, the hypothesis tested is whether the population probabilities equal
those in `theta0`, or are all equal if `theta0` is not given.

If `x` is a matrix with at least two rows and columns, it is taken as a two-dimensional
contingency table. Otherwise, `x` and `y` must be vectors of the same length. The contingency
table is calculated using `counts` function from the `StatsBase` package. Then the power
divergence test is conducted under the null hypothesis that the joint distribution of the
cell counts in a 2-dimensional contingency table is the product of the row and column
marginals.

Note that the entries of `x` (and `y` if provided) must be non-negative integers.

Implements: [`pvalue`](@ref), [`confint`](@ref)
"""
function MultinomialLRTest(x::AbstractMatrix{T}) where T<:Integer
    PowerDivergenceTest(x, lambda=0.0)
end

function MultinomialLRTest(x::AbstractVector{T}, y::AbstractVector{T}, levels::Levels{T}) where T<:Integer
    d = counts(x, y, levels)
    PowerDivergenceTest(d, lambda=0.0)
end

function MultinomialLRTest(x::AbstractVector{T}, y::AbstractVector{T}, k::T) where T<:Integer
    d = counts(x, y, k)
    PowerDivergenceTest(d, lambda=0.0)
end

MultinomialLRTest(x::AbstractVector{T}, theta0::Vector{U} = ones(length(x))/length(x)) where {T<:Integer,U<:AbstractFloat} =
    PowerDivergenceTest(reshape(x, length(x), 1), lambda=0.0, theta0=theta0)

function show_params(io::IO, x::PowerDivergenceTest, ident="")
    println(io, ident, "Sample size:        $(x.n)")
    println(io, ident, "statistic:          $(x.stat)")
    println(io, ident, "degrees of freedom: $(x.df)")
    print(io, ident, "residuals:          ")
    show(io, vec(x.residuals))
    println(io)
    print(io, ident, "std. residuals:     ")
    show(io, vec(x.stdresiduals))
    println(io)
end


using Test, SpecialFunctions
import TruncatedNormal as TruncNorm

@test TruncNorm.tnmom1c(0, -Inf, Inf) == 0
@test TruncNorm.tnmom1c(0, -Inf, 0) ≈ TruncNorm.tnmean(-Inf, 0)
@test TruncNorm.tnmom1c(0, 0, +Inf) ≈ TruncNorm.tnmean(0, +Inf)

@test TruncNorm.tnmom1c(0, +Inf, +Inf) == +Inf
@test TruncNorm.tnmom1c(0, -Inf, -Inf) == -Inf
@test TruncNorm.tnmom1c(1, +Inf, +Inf) == +Inf
@test TruncNorm.tnmom1c(1, -Inf, -Inf) == -Inf
@test TruncNorm.tnmom1c(-Inf, +Inf, +Inf) == +Inf
@test TruncNorm.tnmom1c(+Inf, -Inf, -Inf) == -Inf

@test isnan(TruncNorm.tnmom1c(+Inf, +Inf, +Inf))
@test isnan(TruncNorm.tnmom1c(-Inf, -Inf, -Inf))

@test isnan(TruncNorm.tnmom1c(0, 1, 0))
@test isnan(TruncNorm.tnmom1c(0, +Inf, 0))
@test isnan(TruncNorm.tnmom1c(0, 0, -Inf))
@test isnan(TruncNorm.tnmom1c(0, +Inf, -Inf))
@test isnan(TruncNorm.tnmom1c(1, 1, 0))
@test isnan(TruncNorm.tnmom1c(1, +Inf, 0))
@test isnan(TruncNorm.tnmom1c(1, 0, -Inf))

for c = (0, 1), x = (0, NaN), y = (0, NaN)
    if isnan(x) || isnan(y)
        @test isnan(TruncNorm.tnmom1c(c, x, y))
    end
end

for x = exp.(-100:100), c = -100:100
    @test TruncNorm.tnmom1c(c, -x, x) ≈ -c
end

for x = -10:10, c = -10:10
    @test TruncNorm.tnmom1c(c, x, +Inf) ≈ TruncNorm.tnmean(x, +Inf) - c
    @test TruncNorm.tnmom1c(c, -Inf, x) ≈ TruncNorm.tnmean(-Inf, x) - c
end

#= TruncNorm.tnmom1c(c, a, b) = 0 when c == tnmean(a,b) =#
for a = exp.(-10:100), b = exp.(-10:100)
    a ≤ b || continue
    c = TruncNorm.tnmean(a, b)
    @test abs(TruncNorm.tnmom1c(c, a, b)) ≤ max(1e-10, c * 1e-10)
    @test TruncNorm.tnmom1c(a, a, b) ≥ 0
    @test TruncNorm.tnmom1c(b, a, b) ≤ 0
end

@test TruncNorm.tnmom1c(0, -Inf, Inf, 0, 1) ≈ 0

for x = -10:10
    @test TruncNorm.tnmom1c(0, x, +Inf, 0, 1) ≈ TruncNorm.tnmean(x, +Inf)
    @test TruncNorm.tnmom1c(0, -Inf, x, 0, 1) ≈ TruncNorm.tnmean(-Inf, x)
end

@test TruncNorm.tnmom1c(10, 0, 10) ≈ -9.20211543919713464412024961258
@test TruncNorm.tnmom1c(1, 3, 4) ≈ 2.26045428559002115390313889722
@test TruncNorm.tnmom1c(-1, 3, 4) ≈ 4.26045428559002115390313889722
@test TruncNorm.tnmom1c(2, 3, 4) ≈ 1.26045428559002115390313889722
@test TruncNorm.tnmom1c(-2, 3, 4) ≈ 5.26045428559002115390313889722
@test TruncNorm.tnmom1c(0, -1e200, 1e200) ≈ 0
@test TruncNorm.tnmom1c(1, -1e200, 1e200) ≈ -1
@test TruncNorm.tnmom1c(1e200, -1e200, 1e200) ≈ -1e200
@test TruncNorm.tnmom1c(0, -1e6, -999000) ≈ TruncNorm.tnmean(-1e6, -999000)
@test TruncNorm.tnmom1c(0, +1e6, +Inf) ≈ TruncNorm.tnmean(+1e6, +Inf)
@test TruncNorm.tnmom1c(0, -Inf, -1e6) ≈ TruncNorm.tnmean(-Inf, -1e6)

@test TruncNorm.tnmom1c(999000, 999000, 1e6) ≈ 1.00100100099899498898098100910e-6
@test TruncNorm.tnmom1c(999500, 999000, 1e6) ≈ -499.999998998998999001005011019
@test TruncNorm.tnmom1c(1e6, 999000, 1e6) ≈ -999.999998998998999001005011019


using Distributions
using ParticleMDI

@static if VERSION < v"0.7.0-DEV.2005"
    using Base.Test
else
    using Test
end

## Test data types
# Gaussian
n = 1000
test_data = rand(Normal(0, 1), n, 1)
test_cluster = ParticleMDI.GaussianCluster(test_data)
# Basic check it's been initialised
@test test_cluster.n == 0
# Check the logprob
# ParticleMDI.calc_logprob(test_data[1, :], test_cluster)
for obs in test_data
    ParticleMDI.cluster_add!(test_cluster, [obs], [true])
end

@test test_cluster.n == n
@test isapprox(test_cluster.Σ[1], sum(test_data))
@test isapprox(test_cluster.μ[1], test_cluster.Σ[1] / (n + 0.001))
xbar = test_cluster.Σ[1] / n
s2 = sum((test_data .- xbar) .^ 2)
beta = 0.5 + 0.5 * (s2 + (0.001 * n * xbar ^ 2) / (n + 0.001))
@test isapprox(test_cluster.β[1], beta)
@test isapprox(test_cluster.λ[1], ((0.5 + n * 0.5) * (n + 0.001)) / (test_cluster.β[1] * (n + 1.001)))

xcentred = (test_data[end, :][1] - test_cluster.μ[1]) * sqrt(test_cluster.λ[1])
truelogprob = logpdf(TDist(test_cluster.n + 1), xcentred) + 0.5 * log(test_cluster.λ[1])
estlogprob = ParticleMDI.calc_logprob(test_data[end, :], test_cluster, [true])
@test isapprox(truelogprob, estlogprob)

# Categorical
test_data = rand(1:10, 1000, 1)
test_cluster = ParticleMDI.CategoricalCluster(test_data)
@test test_cluster.n == 0

for obs in test_data
    ParticleMDI.cluster_add!(test_cluster, [obs], [true])
end

@test test_cluster.n == 1000
test_cluster.counts
for x in unique(test_data)
    @test sum(test_data .== x) == test_cluster.counts[x, 1]
end

@test isapprox(ParticleMDI.calc_logprob([1], test_cluster, [true]),
      log((sum(test_data .== 1) + 0.5) / (1005)))

# Test normalising constant calculation
for N = 2:20
    for K = 1:5
        c_combn = Matrix{Int64}(undef, N ^ K, K)
        for k in 1:K
            c_combn[:, K - k + 1] = div.(0:(N ^ K - 1), N ^ (K - k)) .% N .+ 1
        end
        γc = rand(Gamma(1.0 / N, 1), N, K)
        Γ = Matrix{Float64}(undef, N ^ K, K)
        for i in 1:(N ^ K)
            for k in 1:K
                Γ[i, k] = γc[c_combn[i], k]
            end
        end
        if K > 1
            Φ = rand(Gamma(1, 5), binomial(K, 2), 1)
        else
            Φ = zeros(1)
        end
        phi_index = zeros(Int64, K- 1, K)
        num = 1
        for i in 1:(K - 1)
            for j in (i + 1):K
                phi_index[i, j] = num
                num += 1
            end
        end
        Z = 0.0
        for i in 1:(N ^K)
            tmp = prod(Γ[i, :])
            if K > 1
                for k1 in 1:(K - 1)
                    for k2 in (k1 + 1):K
                        tmp *=  (1 + (Φ[phi_index[k1, k2]] * (c_combn[i, k1] == c_combn[i, k2])))
                    end
                end
            end
            Z += tmp
        end
        Φ_index = K > 1 ? Matrix{Bool}(undef, N ^ K, Int64(K * (K - 1) / 2)) : fill(1, (N, 1))
        if K > 1
            i = 1
            for k1 in 1:(K - 1)
                for k2 in (k1 + 1):K
                    Φ_index[:, i] = (c_combn[:, k1] .== c_combn[:, k2])
                    i += 1
                end
            end
        end
        @test isapprox(Z, ParticleMDI.update_Z(Φ, Φ_index, log.(Γ)))
    end
end

# Test cluster alignment
# All datasets strongly agree
K = 5
N = 10
s = rand(1:N, 10000, K)
Φ = [10 for k in 1:binomial(K, 2)]
γ = rand(Gamma(1.0 / N, 1), N, K)
# Perfect agreement up to label permutation
for k in 2:K
    shuf = Distributions.shuffle(1:N)
    s[:, k] .= shuf[s[:, 1]]
    γ[:, k] = γ[indexin(1:N, shuf), 1]
end
# No guarantee that agreement occurs at first attempt
# Run over and over and check that if all the cluster values align
# So do all the gammas
for i = 1:10
    ParticleMDI.align_labels!(s, Φ, γ, N, K)
    @test all(s[:, 2:K] .== s[:, 1]) == all(γ[:, 2:K] .== γ[:, 1])
end
ParticleMDI.align_labels!(s, Φ, γ, N, 0)

@test all(s[:, 2:K] .== s[:, 1])
@test all(γ[:, 2:K] .== γ[:, 1])

# Check the particle IDs are doing what they should
# ie particle IDs
data1 = [rand(Normal(2, 1), 50, 16);
        rand(Normal(-2, 1), 50, 16)]
data2 = [rand(Normal(2, 1), 50, 16);
        rand(Normal(-2, 1), 50, 16)]
data3 = [rand(Normal(2, 1), 50, 16);
        rand(Normal(-2, 1), 50, 16)]
dataFiles = [data1, data2, data3]
dataTypes = [ParticleMDI.GaussianCluster, ParticleMDI.GaussianCluster, ParticleMDI.GaussianCluster]
particle, cluster, counts, n_ops = ParticleMDI.__pmdi(dataFiles, dataTypes, 10, 2, 25 / 100, 1, featureSelect = nothing)
@test all([sum([cluster[k][particle[i, j, k]].n for i in 1:10]) for j in 1:2 for k in 1:3] .== 100)
# Check that cluster counts is working correctly
for k in 1:3
    for i in 1:21
    @test counts[i, k] == count(x -> x == i, particle[:, :, k])
    end
end

particle, cluster, counts, n_ops = ParticleMDI.__pmdi(dataFiles, dataTypes, 10, 1024, 25 / 100, 100, featureSelect = nothing)
@test all([sum([cluster[k][particle[i, j, k]].n for i in 1:10]) for j in 1:1024 for k in 1:3] .== 100)
# Check that cluster counts is working correctly
for k in 1:3
    for i in 1:10241
    @test counts[i, k] == count(x -> x == i, particle[:, :, k])
    end
end


module KNeighbors
  using ...Statistics
  using Distances

  function KNearestNeighbors(trainingSet, testSet, training_y,K=1)
    # m = Statistics.meanColumn(trainingSet)
    # s = Statistics.stdColumn(trainingSet)
    # println("m and s calculatd!")
    # trainingSet=DataManipulation.DataTransform.zScale(trainingSet,m,s)
    # testSet=DataManipulation.DataTransform.zScale(testSet,m,s)
    # println("Scaled!")

    predicted_y=zeros(Float64, size(testSet)[1],1)
    dot_prod = zeros(Float64,size(testSet)[1], size(trainingSet)[1])
    for i=1:size(testSet)[1]
      testSetRow=testSet[i,:]
      for j=1:size(trainingSet)[1]
        dot_prod[i,j]=euclidean(testSetRow,trainingSet[j,:])
      end
    end
    #println("we got our matrix!")
    for i=1:size(dot_prod)[1]
      #find the K smallest distances and get their labels. Take a majority. predicted_y[i] = majority
      indices = Statistics.getMinIndices(collect(dot_prod[i,:]),K)
      # println(typeof(indices));
      # println(indices)
      # println(typeof(training_y))
      # println(training_y)
      predicted_y[i]=Statistics.getMajority(training_y[indices];)

    end
    return predicted_y
    #return fn(test_y, predicted_y)

  end

end


-- prints -3
x = 2 * (1 + 3) / 4 - 5
print(x)


function ll_obj(params;xs=X,ys=Y)
    n = size(X)[1]
    k = Integer(size(X)[2]/2) + 1

    #case 1, y2 comes first
    util11 = params[1] .+ xs[:,1].*params[2]  .+ ys[:,2].*params[3]
    util21 = params[4] .+ xs[:,2].*params[5]

    ll_game1 = exp.(util11).^(ys[:,1])./(1 .+ exp.(util11))
    ll_game1 = ll_game1 .* exp.(util21).^(ys[:,2])./(1 .+ exp.(util21))

    #case 2, y1 comes first
    util12 = params[1] .+ xs[:,1].*params[2]
    util22 = params[4] .+ xs[:,2].*params[5] .+ ys[:,1].*params[6]

    ll_game2 = exp.(util12).^(ys[:,1])./(1 .+ exp.(util12))
    ll_game2 = ll_game2 .* exp.(util22).^(ys[:,2])./(1 .+ exp.(util22))

    lambda = exp(params[7])/(1+exp(params[7]))

    #need to adjust for 11?
    util13 = params[1] .+ xs[:,1].*params[2]  .+ ys[:,2].*params[3]
    util23 = params[4] .+ xs[:,2].*params[5] .+ ys[:,1].*params[6]
    ll_game3 = exp.(util13).^(ys[:,1])./(1 .+ exp.(util13))
    ll_game3 = ll_game3 .* exp.(util23).^(ys[:,2])./(1 .+ exp.(util23))

    ll = (1 .- ys[:,1].*ys[:,2]) .* (lambda .* ll_game1 + (1-lambda) .* ll_game2) + ys[:,1].*ys[:,2] .*ll_game3

    return -sum(log.(ll))
end

function ll_obj(params;xs=X,ys=Y)
    n = size(X)[1]
    k = Integer(size(X)[2]/2) + 1

    p = []

    for bound=[[0,0],[1,1],[0,1],[1,0]]
        util1 = params[1] .+ xs[:,1].*params[2]  .+ bound[2].*params[3]
        util2 = params[4] .+ xs[:,2].*params[5]  .+ bound[1].*params[6]
        bound = exp.(util1).^(bound[1])./(1 .+ exp.(util1))
        bound = bound .* exp.(util2).^(bound[2])./(1 .+ exp.(util2))
        append!(p,[bound])
    end

    p00 = (1 .- ys[:,1]).*(1 .- ys[:,2])
    p11 = ys[:,1].*ys[:,2]
    p10 = ys[:,1].* (1 .- ys[:,2])
    p01 = (1 .- ys[:,1]) .* ys[:,2]

    lambda = exp(params[7])/(1+exp(params[7]))

    ll = p00 .* log.(p[1]) .+ p11 .* log.(p[2])
    ll = ll .+ p01 .* log.( lambda .* p[3] .+ (1 .-lambda) .* p[4] )
    ll = ll .+ p10 .* log. ( lambda .* p[3] .+ (1 .-lambda) .* p[4])
    return -sum(ll)
end

function ll_obj(params;xs=X,ys=Y)
    n = size(X)[1]
    k = Integer(size(X)[2]/2) + 1

    p = []

    for bound=[[0,0],[1,1],[0,1],[1,0]]
        util1 = params[1] .+ xs[:,1].*params[2]  .+ bound[2].*params[3]
        util2 = params[4] .+ xs[:,2].*params[5]  .+ bound[1].*params[6]
        bound = exp.(util1).^(bound[1])./(1 .+ exp.(util1))
        bound = bound .* exp.(util2).^(bound[2])./(1 .+ exp.(util2))
        append!(p,[bound])
    end

    p00 = (1 .- ys[:,1]).*(1 .- ys[:,2])
    p11 = ys[:,1].*ys[:,2]
    p10 = ys[:,1].* (1 .- ys[:,2])
    p01 = (1 .- ys[:,1]) .* ys[:,2]

    lambda = exp(params[7])/(1+exp(params[7]))
    #println(1 .- p[1] .- p[2] .- p[4])
    ll = p00 .* log.(p[1]) + p11 .* log.(p[2])
    ll = p01 .* log.(lambda .* p[3] + (1-lambda).*(1 .- p[1] .- p[2] .- p[4] ) )
    ll = p10 .* log.((1-lambda) .* p[4] + lambda.*(1 .- p[1] .- p[2] .- p[3] ) )
    return -sum(ll)
end


using Gadfly, Distributions

function plot2D(G::Array{ConjugatePostDistribution};
    clim::(Float64, Float64, Float64, Float64) = (-10.0, 10.0, -10.0, 10.0), res::Int=50)

    sx = (clim[2]-clim[1])/(res-1)
    sy = (clim[4]-clim[3])/(res-1)
    A = Array(Float64,2,res^2)
    for (i,a) in enumerate(clim[1]:sx:clim[2])
        for (j,b) in enumerate(clim[3]:sy:clim[4])
            A[:,(i-1)*res+j] = [a,b]
        end
    end

    p = zeros(res * res)
    for d in G

        μ = vec(d.sums / d.n)
        if μ != μ
            μ = vec(zeros(size(d.sums)))
        end

        Σ = (d.ssums / d.n) + (μ * μ')

        p += pdf(MvNormal(μ, Σ), A)
    end

    z= reshape(p,res,res)
    return Gadfly.plot(z=z,x=[clim[1]:sx:clim[2]],y=[clim[3]:sy:clim[4]],Gadfly.Geom.contour)
end


# Set the variables
means = [620; 388; 583; 313]
@randomvariable(m, T[i = 1:4], mean = means[i])
@recoursevariable(m, Qc)
@variable(m, x)

# Define the constraints
@constraint(m, 0.67Qc - 2T[2] - x >= -100)
@constraint(m, -250 - T[2] == x)
@constraint(m, 1388.5 + 0.5Qc - 0.75T[1] - T[2] - T[3] <= 0.0)
@constraint(m, 2044 + Qc - 1.5T[1] - 2T[2] - T[3] <= 0.0)
@constraint(m, 2830 + Qc - 1.5T[1] - 2T[2] - T[3] - 2T[4] <= 0.0)
@constraint(m, -Qc + 1.5T[1] + 2T[2] + T[3] + 3T[4] <= 3153)
true


using Distributions
using SpecialFunctions
using JLD2
using FileIO
using Distributed
using LaTeXStrings
using Glob
using QuadGK

#This has the exponent (alpha d)^phi

function norm_gen_gaus(ϕ,α)
    return (2 /α * gamma(1+ 1/ϕ))
end

function integral_KB_del_ph(δ,ϕ,α)
    integral, err = quadgk(x -> exp(-α^(ϕ) *(abs(x))^ϕ)* (-α^(ϕ)* (abs((x)^ϕ) -abs((x-δ)^ϕ) )), - Inf, Inf, rtol=1e-3)
    return integral/norm_gen_gaus(ϕ,α)
end

function takememory_phi(β,v,α,ϕ)
    prob =probmemory_phi(β,v,α,ϕ)
    wahrscheinlikeit = rand()
    result=false
    if wahrscheinlikeit < prob
        result=true
    end
    result
end

function probmemory_phi(β,v,α,ϕ)
    prob = 1- exp(-β *energiememory_phi(v,α,ϕ) )
    prob
end

function energiememory_phi(v,α,ϕ)
    energie =exp(-(α^(ϕ)*(abs(v))^ϕ))/norm_gen_gaus(ϕ,α)
end

function freeenergy_real_phi(β,v,α,ϕ)
    energy = log( exp(β *energiememory_phi(v,α,ϕ)) -1)/β
    energy
end

function dissipation_phi(β,v,α,ϕ)
    diss = integral_KB_del_ph(v,ϕ,α)/β
    diss
end

function netUtility_phi(β,v,α,ϕ)
    unet=freeenergy_real_phi(β,v,α,ϕ)- dissipation_phi(β,v,α,ϕ)
    unet
end

function newv_drift(δ,n)
    vnew = sqrt(n)δ+randn()*δ*0.05*sqrt(n)
    #vnew = sqrt(n)δ+randn()*δ*0.1
    vnew
end

function onestep_dritf_real_phi(β,δ,oldv,α,n,ϕ)

    if (takememory_phi(β,oldv,α,ϕ))
        free=freeenergy_real_phi(β,oldv,α,ϕ)
        diss=dissipation_phi(β,oldv,α,ϕ)
        unet=free -diss
        nextv=newv_drift(δ,n)
        n+=1
        naiveuse=0
    else
        unet = 0
        free=0
        diss=0
        nextv=newv_drift(δ,1)
        naiveuse=1
        n=2
    end

    return (unet,nextv,naiveuse,free,diss,n)
end

function livesycle_drift_real_phi(β,δ,α,meanencounters,reps,ϕ)
    meanunet=0 #initializing for all the rounds for fist cost
    meanfree=0
    meandiss=0
    naiveuses = 0
    for round in 1:reps

        localnetu= 0
        netu = 0
        localfree= 0
        free = 0
        localdiss= 0
        diss = 0

        v=newv_drift(δ,1) #start the cirus at the memory wich is allways at 0
        naiveuses += 1/reps
        n=2
        for encounter in 2:meanencounters  #start at 1 because initialization is allways at 0 in the first round without any prior memory
             naiveuse =0    # Used naive for the first virus

             (netu,v,naiveuse,free,diss,n)=onestep_dritf_real_phi(β,δ,v,α,n,ϕ)
            localnetu+=netu
            localfree+=free
            localdiss+=diss
            naiveuses += naiveuse/reps
        end

        meanunet+= localnetu/reps
        meanfree+= localfree/reps
        meandiss+= localdiss/reps
    end

    return (meanunet,naiveuses,meanfree,meandiss)
end

function findoptalphaoflive_drift_real_phi(β,δ,meanencounters,reps,αsteps,αmax,ϕ)

    start= 0.01
    α = range(start,stop=αmax,length=αsteps)

    data=livesycle_drift_real_phi.(β,δ,α,meanencounters,reps,ϕ)
    usesmaive=zeros(αsteps)
    memuses=zeros(αsteps)
    utility=zeros(αsteps)
    free=zeros(αsteps)
    diss=zeros(αsteps)
    for ii in 1:αsteps
        utility[ii] = data[ii][1]
        usesmaive[ii] = data[ii][2]
        free[ii] = data[ii][3]
        diss[ii] = data[ii][4]
    end
    memuses= meanencounters./usesmaive

    amax= argmax(utility)
    optalpha = α[amax]
    Nusesatmax=memuses[amax]
    argmaxn = argmax(memuses)
    maxNpossible=memuses[argmaxn]
    utiliymax = utility[amax]
    freemax=free[amax]
    dissmax=diss[amax]

    return Nusesatmax,utiliymax,optalpha,maxNpossible,freemax,dissmax
end

function scandeltaliveGeqBparallel_drift_real_phi(β,meanencounters,reps,αsteps,αmax,deltasteps,ϕ)

    δ = range(0.005,stop=1.5/αmax,length= deltasteps)
    memstepsatmax =zeros(deltasteps)
    utilityopt =zeros(deltasteps)
    alphaopt =zeros(deltasteps)
    posnmax =zeros(deltasteps)
    freemax =zeros(deltasteps)
    dissmax =zeros(deltasteps)
    Threads.@threads for i =1:deltasteps
        memstepsatmax[i],utilityopt[i],alphaopt[i],posnmax[i],freemax[i],dissmax[i]=findoptalphaoflive_drift_real_phi(β,δ[i],meanencounters,reps,αsteps,αmax,ϕ)
    end
    return memstepsatmax,utilityopt,alphaopt,posnmax,δ,freemax,dissmax

end

function scanbetaeqGdeltaliveparallel_drift_real_phi(meanencounters,reps,αsteps,αmax,deltasteps,betasteps,ϕ)

    β = range(0.1,length=betasteps,stop=10.0)

    data = scandeltaliveGeqBparallel_drift_real_phi.(β,meanencounters,reps,αsteps,αmax,deltasteps,ϕ)

    memstepsatmax= zeros(betasteps,deltasteps)
    utilityopt= zeros(betasteps,deltasteps)
    alphaopt= zeros(betasteps,deltasteps)
    posnmax= zeros(betasteps,deltasteps)
    freemax =zeros(betasteps,deltasteps)
    dissmax =zeros(betasteps,deltasteps)
    for ii in 1:betasteps
        memstepsatmax[ii,:]=data[ii][1]
        utilityopt[ii,:]=data[ii][2]
        alphaopt[ii,:] = data[ii][3]
        posnmax[ii,:]=data[ii][4]
        freemax[ii,:]=data[ii][6]
        dissmax[ii,:]=data[ii][7]

    end
    δ = data[1][5]

    introduction="This gives the data meory_used_at_the_max  utility_atmax  otimal_alpha posible_maximum_memoryuses delta and beta and free_at_the_max and diss_at_the_max the data is Gamma first and delta second index"
 return memstepsatmax,utilityopt,alphaopt,posnmax,δ,β,freemax,dissmax,introduction

end

####Data Production

function simulation_DM_cluster_drift_real_oneB_phi(meanencounters,reps,asteps,bsteps,dsteps,amax,thisb,ϕ)

    βas = range(0.1,length=bsteps,stop=10.0)
    β=βas[thisb]
    parameters = (meanencounters,reps,asteps,bsteps,dsteps,amax,β,ϕ)

    filename = "Data/Gen_Bell/Singel_beta/meanencounters_$(meanencounters)/reps_$(reps)/Phi_$(ϕ)/DATA_Asteps_$(asteps)-Bsteps_$(bsteps)-Dsteps_$(dsteps)-amax_$(amax)_real-thisb_$(thisb).jld2"
    println("I want to do  $(filename)")
    if isfile(filename)
            println("Data $(filename) allready exists")
    else
        if !isdir("Data")
            mkdir("Data")
        end
        if !isdir("Data/Gen_Bell")
            mkdir("Data/Gen_Bell")
        end
        if !isdir("Data/Gen_Bell/Singel_beta")
            mkdir("Data/Gen_Bell/Singel_beta")
        end

        if !isdir("Data/Gen_Bell/Singel_beta/meanencounters_$(meanencounters)")
            mkdir("Data/Gen_Bell/Singel_beta/meanencounters_$(meanencounters)")
        end
        if !isdir("Data/Gen_Bell/Singel_beta/meanencounters_$(meanencounters)/reps_$(reps)")
            mkdir("Data/Gen_Bell/Singel_beta/meanencounters_$(meanencounters)/reps_$(reps)")
        end
        if !isdir("Data/Gen_Bell/Singel_beta/meanencounters_$(meanencounters)/reps_$(reps)/Phi_$(ϕ)")
            mkdir("Data/Gen_Bell/Singel_beta/meanencounters_$(meanencounters)/reps_$(reps)/Phi_$(ϕ)")
        end

        memstepsatmax,utilityopt,alphaopt,posnmax,δ,freemax,dissmax = scandeltaliveGeqBparallel_drift_real_phi(β,meanencounters,reps,asteps,amax,dsteps,ϕ)

    #filename = "/Users/oskar/Dropbox/PHD/Simulations/b-cell/Simulation/src/julia/data/test.jld"
        save(filename,"meory_used_at_the_min",  memstepsatmax,"max_utility",utilityopt,"otimal_alpha",alphaopt,"posible_maximum_memoryuses",posnmax,"parameters",parameters,"beta",β,"delta",δ,"freemax",freemax,"dissmax",dissmax,"thisb",thisb)
        println("I'm saving the file $(filename)")
    end
end

function addup_simulation_DM_cluster_drift_real_oneB_phi(meanencounters,reps,asteps,bsteps,dsteps,amax,ϕ)

    filename = "Data/Gen_Bell/full/meanencounters_$(meanencounters)/reps_$(reps)/Phi_$(ϕ)/DATA_Asteps_$(asteps)-Bsteps_$(bsteps)-Dsteps_$(dsteps)-amax_$(amax)_real.jld2"
    deltas = range(0.0001,length=dsteps,stop=0.4)

    parameters = (meanencounters,reps,asteps,bsteps,dsteps,amax,ϕ)

    if !isdir("Data/Gen_Bell/full")
        mkdir("Data/Gen_Bell/full")
    end
        if !isdir("Data/Gen_Bell/full/meanencounters_$(meanencounters)")
            mkdir("Data/Gen_Bell/full/meanencounters_$(meanencounters)")
        end
        if !isdir("Data/Gen_Bell/full/meanencounters_$(meanencounters)/reps_$(reps)")
            mkdir("Data/Gen_Bell/full/meanencounters_$(meanencounters)/reps_$(reps)")
        end
        if !isdir("Data/Gen_Bell/full/meanencounters_$(meanencounters)/reps_$(reps)/Phi_$(ϕ)")
            mkdir("Data/Gen_Bell/full/meanencounters_$(meanencounters)/reps_$(reps)/Phi_$(ϕ)")
        end

        βas = range(0.1,length=bsteps,stop=10.0)

    direct= "Data/Gen_Bell/Singel_beta/meanencounters_$(meanencounters)/reps_$(reps)/Phi_$(ϕ)"
    filenames = "DATA_Asteps_$(asteps)-Bsteps_$(bsteps)-Dsteps_$(dsteps)-amax_$(amax)_real-thisb_*.jld2"

    files2 = get_files(direct,filenames)
    l=length(files2)
    if l<bsteps
        println("!!!!  Achtung !!!!!!!!!!")
        println("es gibt nicht alle files ")
        println("!!!!  Achtung !!!!!!!!!!")
        println(" ")
        println("I have $(l)  of $(bsteps)")
        println(" ")
        println("Please do again ")
        println(" meanencounters_$(meanencounters)/reps_$(reps)/Phi_$(ϕ)/DATA_Asteps_$(asteps)-Bsteps_$(bsteps)-Dsteps_$(dsteps)-amax_$(amax)")
        #return -1
    end

    memstepsatmax= zeros(bsteps,dsteps)
    utilityopt= zeros(bsteps,dsteps)
    alphaopt= zeros(bsteps,dsteps)
    posnmax= zeros(bsteps,dsteps)
    freemax =zeros(bsteps,dsteps)
    dissmax =zeros(bsteps,dsteps)
    for ii in 1:length(files2)
        thisb = load(files2[ii],"thisb")
        memstepsatmax[thisb,:]=load(files2[ii],"meory_used_at_the_min")
        utilityopt[thisb,:]=load(files2[ii],"max_utility")
        alphaopt[thisb,:] = load(files2[ii],"otimal_alpha")
        posnmax[thisb,:]=load(files2[ii],"posible_maximum_memoryuses")
        freemax[thisb,:]=load(files2[ii],"freemax")
        dissmax[thisb,:]=load(files2[ii],"dissmax")
    end

    δ = load(files2[1],"delta")

    #filename = "/Users/oskar/Dropbox/PHD/Simulations/b-cell/Simulation/src/julia/data/test.jld"
        save(filename,"meory_used_at_the_min",  memstepsatmax,"max_utility",utilityopt,"otimal_alpha",alphaopt,"posible_maximum_memoryuses",posnmax,"parameters",parameters,"beta",βas,"delta",δ,"freemax",freemax,"dissmax",dissmax,"phi",ϕ)
            println("I'm saving the file $(filename)")

end

function get_files(direct,filenames)
    filenames=glob(filenames,direct)
    return filenames
end