using Clustering
using LinearAlgebra
using Plots
using Plots.PlotMeasures

mutable struct Posterior_similarity_matrix
    psm::Vector
    names::Vector{String}
    Posterior_similarity_matrix(K, n_obs) = new([Matrix(1.0I, n_obs, n_obs) for k in 1:(K + (K > 1))],
                                                Vector{String}(undef, K + (K > 1)))
end

"""
`generate_psm(outputFile::String, burnin::Int64 = 0, thin::Int64 = 1)`

Generates a posterior similarity matrix from pmdi output
## Input
- `outputFile::S    tring` a string referring to the location on disk of pmdi output
- `burnin::Int64` an integer of the number of initial iterations to discard as burn-in
- `thin::Int64` an integer for the rate at which to thin output, `thin = 2` discards
every second iteration

## Output
A struct of class `Posterior_similarity_matrix` containing:
    - `psm` K n × n matrices where element (i, j) measures how frequently observation
    i and j are co-clustered in dataset k. If `K > 1` an overall consensus matrix
    is also included which is an element-wise average of the K other matrices
    - `names` a vector of the dataset names
"""
function generate_psm(outputFile::String, burnin::Int64 = 0, thin::Int64 = 1)
    outputNames = split(readline(outputFile), ',')
    output = readdlm(outputFile, ',', header = false, skipstart = burnin + 1)
    K = sum(map(outputNames) do str
        occursin(r"MassParameter", str)
    end)

    output = output[1:thin:end, (K + binomial(K, 2) + (K == 1) + 2):end]

    n_obs = size(output, 2) / K
    @assert mod(n_obs, 1) == 0 "Error: Datasets have different number of observations"
    n_obs = Int64(n_obs)
    n_iter = size(output, 1)

    psm = Posterior_similarity_matrix(K, n_obs)
    psm.names[1:K] = unique(map(x -> split(x, '_')[1], outputNames[(K + (K == 1) + binomial(K, 2) + 2):end]))
    if K > 1
        psm.names[K + 1] = "Overall"
    end
    for k = 1:K
        for j = 1:(n_obs - 1)
            for i = (j + 1):n_obs
                psm.psm[k][i, j] = sum(output[:, i + n_obs * (k - 1)] .== output[:, j + n_obs * (k - 1)]) / n_iter
            end
        end
    end
    if K > 1
        for k = 1:K
            psm.psm[K + 1] += psm.psm[k] / K
        end
            psm.psm[K + 1][diagind(psm.psm[K + 1])] .= 1.0
    end

    return psm
end

function generate_psm2(output, K, n_obs)
    output = output[(K + binomial(K, 2) + (K == 1) + 2):end]

    n_iter = size(output, 1)

    psm = Posterior_similarity_matrix(K, n_obs)
    for k = 1:K
        for j = 1:(n_obs - 1)
            for i = (j + 1):n_obs
                psm.psm[k][i, j] = sum(output[i + n_obs * (k - 1)] .== output[j + n_obs * (k - 1)]) / n_iter
            end
        end
    end
    if K > 1
        for k = 1:K
            psm.psm[K + 1] += psm.psm[k] / K
        end
            psm.psm[K + 1][diagind(psm.psm[K + 1])] .= 1.0
    end
    return psm
end

"""
`get_consensus_allocations(psm::Posterior_similarity_matrix, nclust::Int64, orderby::Int64 = 0)`
"""
function get_consensus_allocations(psm::Posterior_similarity_matrix; k::Union{Int64, Nothing}=nothing, h::Union{Float64, Nothing} = nothing,
    linkage = :ward, orderby::Int64 = 0)
    @assert (k != nothing) | (h != nothing) "You must specify either k (number of clusters) or h (height to cut dendrogram)"
    if orderby == 0
        orderby = size(psm.psm, 1)
    end
    hc = hclust(1 .- Symmetric(psm.psm[orderby], :L), linkage = linkage)
    if k != nothing
        return cutree(hc, k = k)
    elseif h != nothing
        return cutree(hc, h = h)
    end

end

"""
`consensus_map(psm::Posterior_similarity_matrix; k::Union{Int64, Nothing} = nothing,
h::Union{Float64, Nothing} = nothing,
orderby::Int64 = 0)`

Generates a consensus map plot for a given posterior similarity matrix
## Input
- `psm` a posterior similarity matrix struct as output from `generate_psm`
- `k` optional - the number of clusters to highlight in the data
- `h` optional - the distance at which to cut the dendrogram. One of `k` or `h` must be specified.
- `orderby` observations in all plots are ordered in a way to separate clusters as well
as possible. `orderby` specifies which dataset should be used to inform this ordering.
`orderby = 0` will let the overall consensus dictate the ordering.
## Output
- K + 1 (or K if K == 1) consensus maps illustrating the clustering output in each
of K datasets.
"""
function consensus_map(psm::Posterior_similarity_matrix;
                       k::Union{Int64, Nothing} = nothing,
                       h::Union{Float64, Nothing} = nothing,
                       orderby::Int64 = 0,
                       linkage = :ward)
    @assert (k != nothing) | (h != nothing) "You must specify either k (number of clusters) or h (height to cut dendrogram)"
    if (orderby == 0) | (orderby == - 1)
        orderby_main = size(psm.psm, 1)
    else orderby_main = orderby
    end
    hc = hclust(1 .- Symmetric(psm.psm[orderby_main], :L), linkage = linkage)
    if h == nothing
        cuts = cutree(hc, k = k)[hc.order]
    elseif k == nothing
        cuts = cutree(hc, h = h)[hc.order]
    end
    nclust = length(unique(cuts))
    ticks = indexin(1:nclust, cuts) .- 0.5
    sort!(ticks)
    append!(ticks, size(psm.psm[1], 1) + 0.5)
    # Surely a better way of doing this
    ticks1 = [[ticks[i], ticks[i+1], NaN][j] for i = 1:nclust for j = 1:3]
    ticks2 = [[ticks[i], ticks[i], NaN][j] for i = 1:nclust for j = 1:3]
    ticks3 = [[ticks[i], ticks[i], NaN][j] for i = 2:(nclust + 1) for j = 1:3]
    K = length(psm.psm)
    if (K > 1) & (orderby != - 1)
        plots = [Plots.heatmap(Symmetric(psm.psm[k], :L)[hc.order, hc.order],
                        ticks = false,
                        yflip = false,
                        title = psm.names[k],
                        legend = false,
                        aspect_ratio = 1,
                        widen = false,
                        c = :viridis,
                        clim = (0, 1)) for k in [K; 1:max(K - 1, 1)]]
        l = @layout [a{0.85w} grid(K - 1, 1)]
    elseif (K > 1) & (orderby == - 1)
        plot_order = [hclust(1 .- Symmetric(psm.psm[k], :L), linkage = :ward).order for k in 1:K]
        plots = [Plots.heatmap(Symmetric(psm.psm[k], :L)[plot_order[k], plot_order[k]],
                        ticks = false,
                        yflip = false,
                        title = psm.names[k],
                        legend = false,
                        aspect_ratio = 1,
                        c = :viridis,
                        clim = (0, 1),
                        widen = false) for k in [K; 1:max(K - 1, 1)]]
        l = @layout [a{0.85w} grid(max(K - 1, 1), 1)]
    elseif K == 1
        plots = [Plots.heatmap(Symmetric(psm.psm[1], :L)[hc.order, hc.order],
                        ticks = false,
                        yflip = false,
                        title = psm.names[1],
                        legend = false,
                        aspect_ratio = 1,
                        widen = false,
                        # titlefont = Plots.font(family = "serif", pointsize = 12),
                        c = :viridis
                        )]
        l = @layout [a;]
    end

    Plots.plot(plots..., layout = l,
               framestyle = :none)
    plot!(ticks1, ticks2, seriestype = :path, c = "#FFFFFF", linestyle = :dash, linewidth = 2.5)
    plot!(ticks1, ticks3, seriestype = :path, c = "#FFFFFF", linestyle = :dash, linewidth = 2.5)
    plot!(ticks2, ticks1, seriestype = :path, c = "#FFFFFF", linestyle = :dash, linewidth = 2.5)
    plot!(ticks3, ticks1, seriestype = :path, c = "#FFFFFF", linestyle = :dash, linewidth = 2.5, widen = false)

end


"""-----------------------------------------------------------------------------
    contracts a tensor represented with indices + values in a symmetric COO
  format. Rows of the indices correspond to sorted indices of a symmetric
  hyper-edge. Values correspond to the edge weights, and v is vector being
  contracted with the
-----------------------------------------------------------------------------"""
function tensor_vector_contraction(A::ThirdOrderSymTensor, x::Array{Float64,1})
    @assert length(x) == A.n

    y = zeros(Float64,A.n)
    edge_count = size(A.indices,1)
	j =1
    @inbounds for (i_1,i_2,i_3) in eachrow(A.indices)

        y[i_1] += 2 * x[i_2]*x[i_3]*A.values[j]
        y[i_2] += 2 * x[i_1]*x[i_3]*A.values[j]
        y[i_3] += 2 * x[i_1]*x[i_2]*A.values[j]

		j +=1
    end

    return y
end

function tensor_vector_contraction(A::UnweightedThirdOrderSymTensor, x::Array{Float64,1})
    @assert length(x) == A.n

    y = zeros(Float64,A.n)

    @inbounds for (i_1,i_2,i_3) in eachrow(A.indices)

        y[i_1] += 2 * x[i_2]*x[i_3]
        y[i_2] += 2 * x[i_1]*x[i_3]
        y[i_3] += 2 * x[i_1]*x[i_2]

    end

    return y
end

"""-----------------------------------------------------------------------------
    contracts a tensor represented with indices + values in a symmetric COO
  format with a vector to return a sparse matrix.
-----------------------------------------------------------------------------"""
function tri_sub_tensor(A::ThirdOrderSymTensor,x::Array{Float64,1})

	index = 0
	nnzs = size(A.indices,1)
	I = zeros(Int64,6*nnzs)
	J = zeros(Int64,6*nnzs)
	V = zeros(Float64,6*nnzs)

    @inbounds for nnz_index = 1:nnzs
		index += 6
		i,j,k = A.indices[nnz_index,:]
		val = A.values[nnz_index]

		I[index-5] = i
		J[index-5] = j
		I[index-4] = j
		J[index-4] = i
		V[index-5] = val*x[k]
		V[index-4] = val*x[k]

		I[index-3] = i
		J[index-3] = k
		I[index-2] = k
		J[index-2] = i
		V[index-3] = val*x[j]
		V[index-2] = val*x[j]

		I[index-1] = j
		J[index-1] = k
		I[index] = k
		J[index] = j
		V[index-1] = val*x[i]
		V[index] = val*x[i]

	end
	return sparse(I,J,V,A.n,A.n)
end

"""-----------------------------------------------------------------------------
    contracts a tensor represented with indices + values in a symmetric COO
  format with a vector to return a sparse matrix.
-----------------------------------------------------------------------------"""
function tri_sub_tensor(A::ThirdOrderSymTensor,x::Array{Float64,1},
    					I::Array{Int64,1},J::Array{Int64,1},V::Array{Float64,1})

	index = 0
	nnzs = size(A.indices,1)
    for nnz_index = 1:nnzs
		index += 6
		i,j,k = A.indices[nnz_index,:]
		val = A.values[nnz_index]

		I[index-5] = i
		J[index-5] = j
		I[index-4] = j
		J[index-4] = i
		V[index-5] = val*x[k]
		V[index-4] = val*x[k]

		I[index-3] = i
		J[index-3] = k
		I[index-2] = k
		J[index-2] = i
		V[index-3] = val*x[j]
		V[index-2] = val*x[j]

		I[index-1] = j
		J[index-1] = k
		I[index] = k
		J[index] = j
		V[index-1] = val*x[i]
		V[index] = val*x[i]

	end
	return sparse(I,J,V,A.n,A.n)
end

"""-----------------------------------------------------------------------------
    contracts a tensor represented with indices + values in a symmetric COO
  format with a vector to return a sparse matrix.
-----------------------------------------------------------------------------"""
function tri_sub_tensor(A::UnweightedThirdOrderSymTensor,x::Array{Float64,1})

	index = 0
	nnzs = 2*sum([length(x) for x in A_U_ten.indices])
	I = zeros(Int64,2*nnzs)
	J = zeros(Int64,2*nnzs)
	V = zeros(Float64,2*nnzs)

	for i in 1:A.n
		for (j,k) in A.indices[i]
			index += 2

			I[index-1] = j
			J[index-1] = k
			I[index] = k
			J[index] = j
			V[index-1] = x[i]
			V[index] = x[i]
		end
	end
	return sparse(I[1:index],J[1:index],V[1:index],A.n,A.n)
end

function kron_contract(A::ThirdOrderSymTensor,B::ThirdOrderSymTensor,
                       X::SparseMatrixCSC{Float64,Int64},rank::Int=minimum(size(X)))

	(U,S,VT),_ = svds(X,nsv = rank)
	println("input rank: ",rank)
	println("output rank:",length(S))

	#singular_indexes = [i for i in 1:length(S) if S[i] > S[1]*eps(Float64)*minimum(size(X))]

	#println("rank in kron contract is",length(singular_indexes))

    kron_contract(A,B,U,VT*diagm(S))
end

function kron_contract(A::ThirdOrderSymTensor,B::ThirdOrderSymTensor,
                       U::Array{Float64,2},V::Array{Float64,2})
    n,d1 = size(U)
    m,d2 = size(V)
    @assert d1 == d2
    @assert A.n == n
    @assert B.n == m

	#preallocate memory used for building sparse matrices
	A_nnzs = size(A.indices,1)
	A_I = zeros(Int64,6*A_nnzs)
	A_J = zeros(Int64,6*A_nnzs)
	A_V = zeros(Float64,6*A_nnzs)

	B_nnzs = size(B.indices,1)
	B_I = zeros(Int64,6*B_nnzs)
	B_J = zeros(Int64,6*B_nnzs)
	B_V = zeros(Float64,6*B_nnzs)

    result = zeros(A.n * B.n)

    @inbounds for i in 1:d1

	#	sub_A_i = tri_sub_tensor(A,U[:,i])
	#	sub_B_i = tri_sub_tensor(B,V[:,i])
        sub_A_i = tri_sub_tensor(A,U[:,i],A_I,A_J,A_V)
        sub_B_i = tri_sub_tensor(B,V[:,i],B_I,B_J,B_V)
        for j in 1:i

            A_update = (sub_A_i*U[:,j])
            B_update = (sub_B_i*V[:,j])

            if i == j
               result += kron(B_update,A_update)
            else
                result += kron(2*B_update,A_update)
            end
        end
    end

    return reshape(result, A.n,B.n)
end

function get_kron_contract_comps(A::Union{ThirdOrderSymTensor,UnweightedThirdOrderSymTensor},
    							 B::Union{ThirdOrderSymTensor,UnweightedThirdOrderSymTensor},
              			         U::Array{Float64,2},V::Array{Float64,2})

    n,d1 = size(U)
    m,d2 = size(V)
    @assert d1 == d2
    @assert A.n == n
    @assert B.n == m

#    result = zeros(A.n * B.n)

    ileave = (i,j) -> i + A.n*(j-1)

    max_rank = Int((d1+1)*d1/2)

	A_comps = zeros(n,max_rank)
	B_comps = zeros(m,max_rank)
	index = 1

	A_nnzs = size(A.indices,1)
	A_I = zeros(Int64,6*A_nnzs)
	A_J = zeros(Int64,6*A_nnzs)
	A_V = zeros(Float64,6*A_nnzs)

	B_nnzs = size(B.indices,1)
	B_I = zeros(Int64,6*B_nnzs)
	B_J = zeros(Int64,6*B_nnzs)
	B_V = zeros(Float64,6*B_nnzs)

    @inbounds for i in 1:d1

        sub_A_i = tri_sub_tensor(A,U[:,i],A_I,A_J,A_V)
        sub_B_i = tri_sub_tensor(B,V[:,i],B_I,B_J,B_V)

        for j in 1:i

			if i == j
				A_comps[:,index] = (sub_A_i*U[:,j])
				B_comps[:,index] = (sub_B_i*V[:,j])
			else

				A_comps[:,index] = sqrt(2)*(sub_A_i*U[:,j])
				B_comps[:,index] = sqrt(2)*(sub_B_i*V[:,j])
			end
			index += 1
        end
    end

    return A_comps, B_comps
end

function implicit_contraction(A::ThirdOrderSymTensor,B::ThirdOrderSymTensor,x::Array{Float64,1})

	error("Code is Broken, must be fixed")
    @assert length(x) == A.n*B.n
    m = A.n
    n = B.n
	X = reshape(x,A.n,B.n)
    Y = similar(X)

    #ileave = (i,j) -> i + m*(j-1)
	i = 1
	j = 1
    @inbounds for (i_1,i_2,i_3) in eachrow(A.indices)

		A_val = A.values[i]
		i += 1
		j = 1
#        for (i_1,i_2,i_3) in permutations(A.indices[i,:])
		for (j_1,j_2,j_3) in eachrow(B.indices)

			#i_1,i_2,i_3
			@inbounds Y[i_1,j_1] += 2*A_val*B.values[j]*X[i_2,j_2]*X[i_3,j_3]
			@inbounds Y[i_2,j_2] += 2*A_val*B.values[j]*X[i_1,j_1]*X[i_3,j_3]
			@inbounds Y[i_3,j_3] += 2*A_val*B.values[j]*X[i_1,j_1]*X[i_2,j_2]

			#i_1,i_3,i_2
			@inbounds Y[i_1,j_1] += 2*A_val*B.values[j]*X[i_3,j_2]*X[i_2,j_3]
			@inbounds Y[i_3,j_2] += 2*A_val*B.values[j]*X[i_1,j_1]*X[i_2,j_3]
			@inbounds Y[i_2,j_3] += 2*A_val*B.values[j]*X[i_1,j_1]*X[i_3,j_2]

			#i_2,i_1,i_3
			@inbounds Y[i_2,j_1] += 2*A_val*B.values[j]*X[i_1,j_2]*X[i_3,j_3]
			@inbounds Y[i_1,j_2] += 2*A_val*B.values[j]*X[i_2,j_1]*X[i_3,j_3]
			@inbounds Y[i_3,j_3] += 2*A_val*B.values[j]*X[i_2,j_1]*X[i_1,j_2]

			#i_2,i_3,i_1
			@inbounds Y[i_2,j_1] += 2*A_val*B.values[j]*X[i_3,j_2]*X[i_1,j_3]
			@inbounds Y[i_3,j_2] += 2*A_val*B.values[j]*X[i_2,j_1]*X[i_1,j_3]
			@inbounds Y[i_1,j_3] += 2*A_val*B.values[j]*X[i_2,j_1]*X[i_3,j_2]

			#i_3,i_2,i_1
			@inbounds Y[i_3,j_1] += 2*A_val*B.values[j]*X[i_2,j_2]*X[i_1,j_3]
			@inbounds Y[i_2,j_2] += 2*A_val*B.values[j]*X[i_3,j_1]*X[i_1,j_3]
			@inbounds Y[i_1,j_3] += 2*A_val*B.values[j]*X[i_3,j_1]*X[i_2,j_2]

			#i_3,i_1,i_2
			@inbounds Y[i_3,j_1] += 2*A_val*B.values[j]*X[i_1,j_2]*X[i_2,j_3]
			@inbounds Y[i_1,j_2] += 2*A_val*B.values[j]*X[i_3,j_1]*X[i_2,j_3]
			@inbounds Y[i_2,j_3] += 2*A_val*B.values[j]*X[i_3,j_1]*X[i_1,j_2]

		end
		j += 1
  #      end
    end

    return Y[:]
end

function impTTVnodesym(A::UnweightedThirdOrderSymTensor,
                       B::UnweightedThirdOrderSymTensor,x::Vector{Float64})

	X = reshape(x,A.n,B.n)
	Y = similar(X)
	Y .= 0
	@inbounds for a = 1:A.n
		for b = 1:B.n
			newval = 0.0
			for (jp,kp) in A.indices[a]
				@simd for pair in B.indices[b]
					j,k = pair
					@inbounds newval += X[jp,j]*X[kp,k]+X[kp,j]*X[jp,k]
				end
			end
			Y[a,b] = 2*newval
		end
	end
	y = Y[:]
	return y

end

function impTTVnodesym(nG::Int,nH::Int,x::Vector{Float64},Gti, Hti)

	X = reshape(x,nG,nH)
	#Xt = copy(X')
	Y = similar(X)
	Y .= 0
	@inbounds for g = 1:nG
		for h = 1:nH
			newval = 0.0
			for (jp,kp) in Gti[g]
				@simd for pair in Hti[h]
					j,k = pair
					@inbounds newval += X[jp,j]*X[kp,k]+X[kp,j]*X[jp,k]
				end
			end
			Y[g,h] = 2*newval
		end
	end
	y = Y[:]
	return y
end

#TODO: move to appropriate file
function SSHOPM(A::ThirdOrderSymTensor,β::Float64, max_iter::Int, tol::Float64,
	            x_0::Array{Float64,1}=ones(A.n);update_user::Int=-1)

	#normalize init vector
	x_k = copy(x_0)
	x_k ./=norm(x_0)

    lambda_k = Inf
    i = 1

    while true

		x_k_1 = tensor_vector_contraction(A,x_k)
		lambda_k_1 = x_k_1'*x_k

		if β != 0.0
			x_k_1 .+= β*x_k
        end

		x_k_1 ./= norm(x_k_1)

		res = abs(lambda_k_1 - lambda_k)
		if update_user != -1 && i % update_user == 0
			println("iteration $(i)    λ: $(lambda_k_1)   |λ_k_1 - λ_k| = $res")
		end

        if res < tol || i >= max_iter
            return x_k_1, lambda_k_1
        else
			lambda_k = copy(lambda_k_1)
			x_k = copy(x_k_1)
            i += 1
        end

    end

end


using Test
using LinearAlgebra
using Rotations
using StaticArrays
using InteractiveUtils: subtypes
using Quaternions
import Unitful

import Random

# Check that there are no ambiguities beyond those present in StaticArrays
ramb = detect_ambiguities(Rotations, Base, Core)
samb = detect_ambiguities(StaticArrays, Base, Core)
@test isempty(setdiff(ramb, samb))

include("util_tests.jl")
include("2d.jl")
include("rotation_tests.jl")
include("derivative_tests.jl")
include("principal_value_tests.jl")
include("unitquat.jl")
include("rodrigues_params.jl")
include("quatmaps.jl")
include("rotation_error.jl")
include("distribution_tests.jl")
include("eigen.jl")
include("nearest_rotation.jl")
include("rotation_generator.jl")
include("logexp.jl")
include("deprecated.jl")

include(joinpath(@__DIR__, "..", "perf", "runbenchmarks.jl"))


using LinearAlgebra

"""

Structure to hold data for Ehrenfest dynamics

"""
struct EhrenfestOps
    No :: Int64
    ρ  :: Matrix{Complex{Float64}}
    p  :: Vector{Float64}
    q  :: Vector{Float64}
    function EhrenfestOps(ρ0, sbm :: SBModel)
        ρ = zeros(ComplexF64,2,2)
        ρ .= ρ0
        p = zeros(Float64,sbm.No)
        q = zeros(Float64,sbm.No)
        new(sbm.No,ρ,p,q)
    end
end

"""

Leapfrog bootstrap of data (Ehrenfest)

"""
function ehbootstrap!(ops :: EhrenfestOps, oldops :: EhrenfestOps,
                      dotops :: EhrenfestOps, dt :: Float64)
    oldops.ρ .= ops.ρ - dt * dotops.ρ
    oldops.p .= ops.p - dt * dotops.p
    oldops.q .= ops.q - dt * dotops.q

    return nothing
end

"""

Leapfrog step forward (Ehrenfest)

"""
function ehforward!(ops :: EhrenfestOps, oldops :: EhrenfestOps,
                    dotops :: EhrenfestOps, dt :: Float64)
    oldops.ρ .= oldops.ρ + 2 * dt * dotops.ρ
    oldops.p .= oldops.p + 2 * dt * dotops.p
    oldops.q .= oldops.q + 2 * dt * dotops.q

    return nothing
end

"""

Leapfrog Ehrenfest calculate EOM RHS

"""
function ehcalcdots!(dotops :: EhrenfestOps, ops :: EhrenfestOps,
                    dt :: Float64, sbm :: SBModel)
    dotops.ρ .= - IOHBAR * comm(H(sbm,ops.q,ops.p),ops.ρ)
    for i in 1:ops.No
        dotops.p[i] = real(tr(F(sbm,ops.q,i)*ops.ρ))
        dotops.q[i] = ops.p[i]
    end

    return nothing
end

"""

Ehrenfest total energy

"""
function ehenergy(ops :: EhrenfestOps, sbm :: SBModel)
    energy =  real(tr(H(sbm,ops.q,ops.p) * ops.ρ))
    for ν in 1:ops.No
        energy += 0.5 * ops.p[ν]^2
    end
    return energy
end

"""

This function integrates the Ehrenfest EOMs starting from the
initial condition provided in ops for nsteps of dt.
The store! callback function is called at every timestep.

"""
function RunEhrenfest!(sbm, ops, nsteps, dt, store!, storage)

    dotops = EhrenfestOps(zm,sbm)
    oldops = EhrenfestOps(zm,sbm)

    # Bootstrap the integrator
    ehcalcdots!(dotops, ops ,dt , sbm)
    ehbootstrap!(ops,oldops,dotops,dt)

    # call results storage function at t = 0
    t = 0.0
    store!(storage,t,ops,sbm)

    # integrate
    for i in 1:nsteps-1
        t = i*dt
        ehcalcdots!(dotops,ops,dt,sbm)
        ehforward!(ops,oldops,dotops,dt)
        (ops,oldops) = (oldops,ops)
        store!(storage,t,ops,sbm)
    end

    return nothing
end


#_____________________________________________________________________
#                           Charlie's code
#

# Here I use `Gnuplot.jl`, but any other package would work...
#=
using Gnuplot
setverbose(false)
@gp "set size ratio -1" "set grid" "set key bottom right" xr=(-1.5, 2.5) :-
=#

# Methods to plot a Polygon
plotlabel(p::Polygon) = "Polygon (" * string(length(p.x)) * " vert.)"
plotlabel(p::RegularPolygon) = "RegularPolygon (" * string(vertices(p)) * " vert., area=" * string(round(area(p) * 100) / 100) * ")"
plotlabel(p::Named) = name(p) * " (" * string(vertices(p)) * " vert., area=" * string(round(area(p) * 100) / 100) * ")"
function plot(p::AbstractPolygon; dt=1, color="black")
    x = coords_x(p); x = [x; x[1]]
    y = coords_y(p); y = [y; y[1]]
    title = plotlabel(p)
    # @gp :- x y "w l tit '$title' dt $dt lw 2 lc rgb '$color'"
end

# Finally, let's have fun with the shapes!
line = Polygon([0., 1.], [0., 1.])
triangle = RegularPolygon(3, 1)
square = Named{RegularPolygon}("Square", 4, 1)

plot(line, color="black")
plot(triangle, color="blue")
plot(square, color="red")

rotate!(triangle, 90)
move!(triangle, 1, 1)
scale!(triangle, 0.32)
scale!(square, sqrt(2))

plot(triangle, color="blue", dt=2)
plot(square, color="red", dt=2)

# Increase number of vertices
p = Named{RegularPolygon}("Heptagon", 7, 1)
plot(p, color="orange", dt=4)
circle = Named{RegularPolygon}("Circle", 1000, 1)
plot(circle, color="dark-green", dt=4)


"""
    emb_sz_rule(n_cat)

Compute an embedding size corresponding to the number of classes for a
categorical variable using the rule of thumb present in python fastai.
(see https://github.com/fastai/fastai/blob/2742fe844573d06e700f869839fb9ec5f3a9bca9/fastai/tabular/model.py#L12)
"""
emb_sz_rule(n_cat) = min(600, round(Int, 1.6 * n_cat^0.56))

"""
    get_emb_sz(cardinalities::AbstractVector, [size_overrides::AbstractVector])

Given a vector of `cardinalities` of each categorical column
(i.e. each element of `cardinalities` is the number of classes in that categorical column),
compute the output embedding size according to [`emb_sz_rule`](#).
Return a vector of tuples where each element is `(in_size, out_size)` for an embedding layer.

## Keyword arguments

- `size_overrides`: A collection of integers (or `nothing` to skip override) where the value present at any index
    will be used to as the output embedding size for that column.
"""
get_emb_sz(cardinalities::AbstractVector{<:Integer}, size_overrides=fill(nothing, length(cardinalities))) =
    map(zip(cardinalities, size_overrides)) do (cardinality, override)
        emb_dim = isnothing(override) ? emb_sz_rule(cardinality + 1) : Int64(override)
        return (cardinality + 1, emb_dim)
    end

sigmoidrange(x, low, high) = @. Flux.sigmoid(x) * (high - low) + low

function tabular_embedding_backbone(embedding_sizes, dropout_rate=0.)
    embedslist = [Flux.Embedding(ni, nf) for (ni, nf) in embedding_sizes]
    emb_drop = iszero(dropout_rate) ? identity : Dropout(dropout_rate)
    Chain(
        x -> tuple(eachrow(x)...),
        Parallel(vcat, embedslist),
        emb_drop
    )
end

tabular_continuous_backbone(n_cont) = BatchNorm(n_cont)

"""
    TabularModel(catbackbone, contbackbone, [finalclassifier]; kwargs...)

Create a tabular model which operates on a tuple of categorical values
(label or one-hot encoded) and continuous values.
The categorical backbones (`catbackbone`) and continuous backbone (`contbackbone`) operate on each element of the input tuple.
The output from these backbones is then passed through a series of linear-batch norm-dropout layers before a `finalclassifier` block.

## Keyword arguments

- `outsize`: The output size of the final classifier block. For single classification tasks,
    this would be the number of classes, and for regression tasks, this would be the
    number of target continuous variables.
- `layersizes`: A vector of sizes for each hidden layer in the sequence of linear layers.
- `dropout_rates`: Dropout probabilities for the linear-batch norm-dropout layers.
    This could either be a single number which would be used for for all the layers,
    or a collection of numbers which are cycled through for each layer.
- `batchnorm`: Set to `false` to skip each batch norm in the linear-batch norm-dropout sequence.
- `activation`: The activation function to use in the classifier layers.
- `linear_first`: Controls if the linear layer comes before or after batch norm and dropout.
"""
function TabularModel(
        catbackbone,
        contbackbone;
        outsize,
        layersizes=(200, 100),
        kwargs...)
    TabularModel(catbackbone, contbackbone, Dense(layersizes[end], outsize); layersizes=layersizes, kwargs...)
end

function TabularModel(
        catbackbone,
        contbackbone,
        finalclassifier;
        layersizes=(200, 100),
        dropout_rates=0.,
        batchnorm=true,
        activation=Flux.relu,
        linear_first=true)

    tabularbackbone = Parallel(vcat, catbackbone, contbackbone)

    classifierin = mapreduce(layer -> size(layer.weight)[1], +, catbackbone[2].layers;
                             init = contbackbone.chs)
    dropout_rates = Iterators.cycle(dropout_rates)
    classifiers = []

    first_ps, dropout_rates = Iterators.peel(dropout_rates)
    push!(classifiers, linbndrop(classifierin, first(layersizes);
                                 use_bn=batchnorm, p=first_ps, lin_first=linear_first, act=activation))

    for (isize, osize, p) in zip(layersizes[1:(end-1)], layersizes[2:end], dropout_rates)
        layer = linbndrop(isize, osize; use_bn=batchnorm, p=p, act=activation, lin_first=linear_first)
        push!(classifiers, layer)
    end

    Chain(
        tabularbackbone,
        classifiers...,
        finalclassifier
    )
end

"""
    TabularModel(n_cont, outsize, [layersizes; kwargs...])

Create a tabular model which operates on a tuple of categorical values
(label or one-hot encoded) and continuous values. The default categorical backbone (`catbackbone`) is
a [`Flux.Parallel`](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.Parallel) set of `Flux.Embedding` layers corresponding to each categorical variable.
The default continuous backbone (`contbackbone`) is a single [`Flux.BatchNorm`](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.BatchNorm).
The output from these backbones is concatenated then passed through a series of linear-batch norm-dropout layers before a `finalclassifier` block.

## Arguments

- `n_cont`: The number of continuous columns.
- `outsize`: The output size of the model.
- `layersizes`: A vector of sizes for each hidden layer in the sequence of linear layers.

## Keyword arguments

- `cardinalities`: A collection of sizes (number of classes) for each categorical column.
- `size_overrides`: An optional argument which corresponds to a collection containing
    embedding sizes to override the value returned by the "rule of thumb" for a particular index
    corresponding to `cardinalities`, or `nothing`.
"""
function TabularModel(
        n_cont::Number,
        outsize::Number,
        layersizes=(200, 100);
        cardinalities,
        size_overrides=fill(nothing, length(cardinalities)))
    embedszs = get_emb_sz(cardinalities, size_overrides)
    catback = tabular_embedding_backbone(embedszs)
    contback = tabular_continuous_backbone(n_cont)

    TabularModel(catback, contback; layersizes=layersizes, outsize=outsize)
end


# Get displacement at a given iteration
function get_fields( fname::String , itt::Int)

    #Get Variables
    fid         = h5open( fname , "r")

    max_iter = read_attribute(fid, "max_iter")

    xmin = read_attribute(fid, "xmin")
    xmax = read_attribute(fid, "xmax")
    Nx   = read_attribute(fid, "xnodes")

    ymin = read_attribute(fid, "ymin")
    ymax = read_attribute(fid, "ymax")
    Ny   = read_attribute(fid, "ynodes")

    #Discretize space
    dx  = (xmax - xmin)/(Nx-1)
    x   = [xmin + (i-1)*dx for i in 1:Nx]

    dy  = (ymax - ymin)/(Ny-1)
    y   = [ymin + (i-1)*dy for i in 1:Ny]

    #Check it time iteration is OK
    if( itt > max_iter)
        println("ERROR: out of bounds (max iteration is" , max_iter,")")
        return
    end

    #Get Displacement
    g1      = open_group(fid, "ψ")

    dname   = string("iteration_" , itt)
    dset    = open_dataset(g1, dname )
    ttt     = read_attribute(dset,"time")

    ψ = reshape(Array(dset) , Nx,Ny )

    #Get Derivative
    g1      = open_group(fid, "dψ")

    dname   = string("iteration_" , itt)
    dset    = open_dataset(g1, dname )

    dψ = reshape(Array(dset),Nx,Ny)

    close(fid)

    #Return all data
    return ttt , x , y , ψ , dψ

end


using BenchmarkTools, Compat, DataFrames, IntervalRootFinding, IntervalArithmetic, StaticArrays

function randVec(n::Int)
    a = randn(n)
    A = Interval.(a)
    mA = MVector{n}(A)
    sA = SVector{n}(A)
    return A, mA, sA
end

function randMat(n::Int)
    a = randn(n, n)
    A = Interval.(a)
    mA = MMatrix{n, n}(A)
    sA = SMatrix{n, n}(A)
    return A, mA, sA
end

function benchmark(max=10)
    df = DataFrame()
    df[:Method] = ["Array", "MArray", "SArray", "Contractor", "ContractorMArray", "ContractorSArray"]
    for n in 1:max
        A, mA, sA = randMat(n)
        b, mb, sb = randVec(n)
        t1 = @belapsed gauss_seidel_interval($A, $b)
        t2 = @belapsed gauss_seidel_interval($mA, $mb)
        t3 = @belapsed gauss_seidel_interval($sA, $sb)
        t4 = @belapsed gauss_seidel_contractor($A, $b)
        t5 = @belapsed gauss_seidel_contractor($mA, $mb)
        t6 = @belapsed gauss_seidel_contractor($sA, $sb)
        df[Symbol("$n")] = [t1, t2, t3, t4, t5, t6]
    end
    a = []
    for i in 1:max
        push!(a, Symbol("$i"))
    end
    df1 = stack(df, a)
    dfnew = unstack(df1, :variable, :Method, :value)
    dfnew = rename(dfnew, :variable => :n)
    println(dfnew)
    dfnew
end

function benchmark_elimination(max=10)
    df = DataFrame()
    df[:Method] = ["Gauss Elimination", "Base.\\"]
    for n in 1:max
        A, mA, sA = randMat(n)
        b, mb, sb = randVec(n)
        t1 = @belapsed gauss_elimination_interval($A, $b)
        t2 = @belapsed gauss_elimination_interval1($A, $b)
        df[Symbol("$n")] = [t1, t2]
    end
    a = []
    for i in 1:max
        push!(a, Symbol("$i"))
    end
    df1 = stack(df, a)
    dfnew = unstack(df1, :variable, :Method, :value)
    dfnew = rename(dfnew, :variable => :n)
    println(dfnew)
    dfnew
end


# - Julia online documentation
# - Collections and Data Structures
# - Iterable Collections
using Base.Test

notin(x, itr) = any(y -> y != x, itr)

@testset "Iterable Collections" begin
@testset "Base.in" begin
# --------------------------------------------- #
c::Vector{Int64} = 1:10
@test in(1, c)
@test notin(11, c)
end
@testset "eltype" begin
c::Vector{Int64} = 1:10
@test eltype(c) == Int64
@test eltype(typeof(c)) == Int64
end
# ============================================= #
@testset "Base.map!" begin
let
a::Float64 = 0.5
x::Vector{Float64} = 1:10
y::Vector{Float64} = 1:10
# map!(function, destination, collection...)
map!((x,y)->a*x+y, y, x, y)
@test y[1] == 1.5
@test y[10] == 15
end
let
a::Float64 = 0.5
x::Vector{Float64} = 1:10
y::Vector{Float64} = 1:10
f(x,y) = a*x + y
map!(f, y, x, y)
@test y[1] == 1.5
@test y[10] == 15
end
end
# ============================================= #
end


#!/usr/bin/env julia6
using ReadGlobal, MyStats, JLD2

function histstdindices(input::AbstractString,nbins::Integer,nstd::Number)
    field = readfield(input)
    med,stdv,_,_ = read_info(input*".info")
    minf = med - nstd*stdv
    maxf = med + nstd*stdv
    ind = hist_indices(field,minf,maxf,nbins)
    bins = Bins(minf,maxf,nbins)
    @save "pdf_std$(nstd)_indices_$(input)_nbins$(nbins).jld2" ind bins
end

function main(f,ns,nstds)
    n = parse(Int,ns)
    nstd = parse(Float64,nstds)
    histstdindices(f,n,nstd)
end

function main(f,ns)
    n = parse(Int,ns)
    histstdindices(f,n,2.)
end

main(f) = histstdindices(f,50,2.)

main(ARGS...)


#-------------------------------------------------------------------------------
# Test suite for 'lovo_utils.jl' file
#-------------------------------------------------------------------------------

@testset "lovo_utils" begin

    @testset "fmin_eval" begin

        f(x) = ( x[1] + x[2] ) ^ 2.0
        g(x) = x[1] - x[2]
        h(x) = 1.0 - x[1] - x[2]
        l(x) = sin( x[1] + x[2] )
        func_list = [f, g, h, l]
        r = 4

        # Test parameters
        y = [0.0, 0.0]
        set = zeros(Bool, r)
        # Expected return
        fmin_return = 0.0
        imin_return = 1
        # Test
        fmin, imin = LOWDER.fmin_eval!(func_list, r, y, set)
        @test( fmin == fmin_return )
        @test( imin == imin_return )
        @test( set == [true, true, false, true] )

        # Test parameters
        y = [0.0, 1.0]
        set = zeros(Bool, r)
        # Expected return
        fmin_return = -1.0
        imin_return = 2
        # Test
        fmin, imin = LOWDER.fmin_eval!(func_list, r, y, set)
        @test( fmin == fmin_return )
        @test( imin == imin_return )
        @test( set == [false, true, false, false] )

        # Test parameters
        y = [1.0, 1.0]
        set = zeros(Bool, r)
        # Expected return
        fmin_return = -1.0
        imin_return = 3
        # Test
        fmin, imin = LOWDER.fmin_eval!(func_list, r, y, set)
        @test( fmin == fmin_return )
        @test( imin == imin_return )
        @test( set == [false, false, true, false] )

        # Test parameters
        y = [0.0, - 1.0]
        set = zeros(Bool, r)
        # Expected return
        fmin_return = sin(- 1.0)
        imin_return = 4
        # Test
        fmin, imin = LOWDER.fmin_eval!(func_list, r, y, set)
        @test( fmin == fmin_return )
        @test( imin == imin_return )
        @test( set == [false, false, false, true] )

    end

    @testset "fi_eval" begin

        f(x) = ( x[1] + x[2] ) ^ 2.0
        g(x) = x[1] - x[2]
        h(x) = 1.0 - x[1] - x[2]
        l(x) = sin( x[1] + x[2] )
        func_list = [f, g, h, l]

        # Test parameters
        index = 1
        y = [0.0, - 1.0]
        # Expected return
        value_return = 1.0
        # Test
        value = LOWDER.fi_eval(func_list, index, y)
        @test(value == value_return)

        # Test parameters
        index = 2
        y = [1.0, 1.0]
        # Expected return
        value_return = 0.0
        # Test
        value = LOWDER.fi_eval(func_list, index, y)
        @test(value == value_return)

        # Test parameters
        index = 3
        y = [0.0, 1.0]
        # Expected return
        value_return = 0.0
        # Test
        value = LOWDER.fi_eval(func_list, index, y)
        @test(value == value_return)

        # Test parameters
        index = 4
        y = [0.0, 0.0]
        # Expected return
        value_return = 0.0
        # Test
        value = LOWDER.fi_eval(func_list, index, y)
        @test(value == value_return)

    end

    @testset "fmin_partial_eval" begin

        f(x) = ( x[1] + x[2] ) ^ 2.0
        g(x) = x[1] - x[2]
        h(x) = 1.0 - x[1] - x[2]
        l(x) = sin( x[1] + x[2] )
        func_list = [f, g, h, l]
        r = 4
        set = zeros(Bool, r)

        # Test parameters
        y = [0.0, 0.0]

        idx = 1
        fi_y = 0.0

        # Test 01
        fmin, imin = LOWDER.fmin_partial_eval!( func_list, r, idx, fi_y, y, set )
        @test( fmin == 0.0 )
        @test( imin == 1 )
        @test( set == [true, true, false, true] )

        # Test parameters
        idx = 2
        fi_y = 0.0

        # Test 02
        fmin, imin = LOWDER.fmin_partial_eval!( func_list, r, idx, fi_y, y, set )
        @test( fmin == 0.0 )
        @test( imin == 2 )
        @test( set == [true, true, false, true] )

        # Test parameters
        idx = 3
        fi_y = 1.0

        # Test 03
        fmin, imin = LOWDER.fmin_partial_eval!( func_list, r, idx, fi_y, y, set )
        @test( fmin == 0.0 )
        @test( imin == 1 )
        @test( set == [true, true, false, true] )

        # Test parameters
        idx = 4
        fi_y = 0.0

        # Test 04
        fmin, imin = LOWDER.fmin_partial_eval!( func_list, r, idx, fi_y, y, set )
        @test( fmin == 0.0 )
        @test( imin == 4 )
        @test( set == [true, true, false, true] )

        # ---

        # Test parameters
        y = [0.0, 1.0]
        idx = 3
        fi_y = 0.0

        # Test
        fmin, imin = LOWDER.fmin_partial_eval!( func_list, r, idx, fi_y, y, set )
        @test( fmin == -1.0 )
        @test( imin == 2 )
        @test( set == [false, true, false, false] )

        # ---

        # Test parameters
        y = [1.0, 1.0]
        idx = 1
        fi_y = 4.0

        # Test
        fmin, imin = LOWDER.fmin_partial_eval!( func_list, r, idx, fi_y, y, set )
        @test( fmin == -1.0 )
        @test( imin == 3 )
        @test( set == [false, false, true, false] )

        # ---

        # Test parameters
        y = [0.0, -1.0]
        idx = 1
        fi_y = 0.0

        # Test
        fmin, imin = LOWDER.fmin_partial_eval!( func_list, r, idx, fi_y, y, set )
        @test( fmin == l(y) )
        @test( imin == 4 )
        @test( set == [false, false, false, true] )

    end

end


using StaticArrays, BenchmarkTools, Unrolled

const i1 = @SVector [1,1,1,1,1,1,2,2,3,4,3,4,2,2,3,4,3,4,2,2,3,4,3,4,5,5,6,7,6,7,5,5,6,7,6,7,8,9,8,9,10,10,8,9,8,9,10,10,5,5,6,7,6,7,8,9,8,9,10,10,]
const i2 = @SVector [3,4,2,2,4,3,3,4,2,2,4,3,1,1,1,1,1,1,4,3,4,3,2,2,2,2,3,4,3,4,1,1,1,1,1,1,3,4,2,2,4,3,3,4,2,2,4,3,7,6,7,6,5,5,6,7,5,5,7,6,]
const i3 = @SVector [5,5,6,7,6,7,5,5,6,7,6,7,8,9,8,9,10,10,8,9,8,9,10,10,6,7,5,5,7,6,8,9,8,9,10,10,5,5,6,7,6,7,9,8,10,10,8,9,8,9,8,9,10,10,9,8,10,10,8,9,]
x = @SVector rand(10)

tdot_A(x) = (dot(x[i1] .* x[i2], x[i3]))
tdot_B(x) = sum(x[i1] .* x[i2] .* x[i3])
# tdot_C(x) = sum(x[i1[n]] .* x[i2[n]] .* x[i3[n]] for n = 1:60)
function tdot_D(x)
    y = zero(eltype(x))
    for n = 1:length(i1)
        @inbounds a = x[i1[n]] * x[i2[n]]
        @inbounds y = muladd(a, x[i3[n]], y)
    end
    return y
end

@btime tdot_A($x)
@btime tdot_B($x)
# @btime tdot_C($x)
@btime tdot_D($x)

# @code_llvm(tdot_A(x))
@code_llvm(tdot_D(x))

@btime ForwardDiff.gradient($tdot_A, $x)
@btime ForwardDiff.gradient($tdot_D, $x)

const j1 = @SVector [1,1,1,2,2,2,3,4,3,4,4,3,4,5,6,7,5,6,7,7,8,9,8,9,10,10,9,5,5,6,7,6,7,8,9,8,9,10,10,]
const j2 = @SVector [3,4,2,1,3,4,2,2,4,3,1,1,1,2,3,4,1,1,1,1,3,4,2,2,4 ,3 ,2,7,6,7,6,5,5,6,7,5,5,7 ,6 ,]
A12_ = rand([0,1], (10,10))
const A12 = SMatrix{10,10}(A12_...)

ddot_A(x) = (dot(x[j1], x[j2]))
ddot_B(x) = (dot(x, A12 * x) )
@inline @unroll function ddot_C(x, j1,j2)
   y = zero(eltype(x))
   @unroll for n = 1:length(x)
     @inbounds y = muladd(x[j1[n]], x[j2[n]], y)
   end
   return y
end

x = @SVector rand(10)

@btime ddot_A($x)
@btime ddot_B($x)
@btime ddot_C($x, $j1, $j2)


#__precompile__()
module NumericalBasis

#   This file is part of NumericalBasis.jl. It is licensed under the MIT license
#   Copyright (C) 2018 Michael Reed

using SyntaxTree, ForceImport, Reduce
export floatset, optimal, roots, polyhorner, polyfactors, polyexpand

function floatset(T::DataType,N;scale=x->x)
    l = scale(eps(T))
    u = scale(prevfloat(T(Inf)))
    return l:(u-l)/(N-1):u
end

polyhorner(x,a) = polyhorner(x,a,1)
polyhorner(x,a::Array{<:Any,1},k) = k==length(a) ? a[k] : Algebra.:+(a[k],Algebra.:*(x,polyhorner(x,a,k+1)))

polyfactors(x,a) = polyfactors(x,a,1)
polyfactors(x,a::Array{<:Any,1},k) = k==length(a) ? Algebra.:-(x,a[k]) : Algebra.:*(Algebra.:-(x,a[k]),polyfactors(x,a,k+1))

polyexpand(x,a) = polyexpand(x,a,length(a))
polyexpand(x,a::Array{<:Any,1},k) = k==1 ? a[k] : Algebra.:+(Algebra.:*(a[k],Algebra.:^(x,k-1)),polyexpand(x,a,k-1))

roots = factor

function optimal(expr)
    h = horner(expr)
    f = factor(h)
    eh = SyntaxTree.exprval(h)[1]
    ef = SyntaxTree.exprval(f)[1]
    eo = SyntaxTree.exprval(expr)[1]
    if eh ≤ ef
        return eh ≤ eo ? h : expr
    else
        return ef ≤ eo ? f : expr
    end
end

__init__() = nothing

end # module


numbers = parse.(Int, split(readline("day4/bingo.txt"), ','))

boards = open("day4/bingo.txt") do io
    readline(io)  # remove first line
    readlines(io)
end

boards = filter(!isempty, boards)
boards = permutedims(hcat(split.(boards)...))
boards = parse.(Int, boards)
boards = [boards[i:i+4, :] for i in 1:5:size(boards, 1)]

function play_bingo(numbers, boards)
    bingos = []
    marked_boards = [zeros(Bool, size(board)) for board in boards]
    is_bingo_board = zeros(Bool, length(boards))
    for number in numbers
        for i in eachindex(boards, marked_boards, is_bingo_board)
            is_bingo_board[i] && continue
            board = boards[i]
            marked_board = marked_boards[i]
            for j in eachindex(board, marked_board)
                if board[j] == number
                    marked_board[j] = true
                end
            end
            is_marked_col = sum(marked_board, dims = 1) .== 5
            is_marked_row = sum(marked_board, dims = 2) .== 5
            is_bingo_board[i] = any(is_marked_col) | any(is_marked_row)
            is_bingo_board[i] && push!(bingos, (number, board, marked_board))
        end
    end
    bingos
end

bingos = play_bingo(numbers, boards)

number, board, marked_board = bingos[1]
p1 = number * sum(marked_board[i] ? 0 : board[i] for i in eachindex(board))

@show p1

number, board, marked_board = bingos[end]
p2 = number * sum(marked_board[i] ? 0 : board[i] for i in eachindex(board))

@show p2


using BenchmarkTools

# Originally from Teo ShaoWei via Zulip

const INT = UInt32
const cups = parse.(INT, collect("389125467"))

@inline function peek(next, at, n::Int, result=similar(next,n))
  for i in 1:n
    result[i] = next[at]
    at = next[at]
  end
  result
end

@inline function peek(next, at, ::Val{3})
    a = next[at]
    b = next[ a ]
    c = next[ b ]
    ( a, b, c )
end

function next1(cups)
  N = length(cups)
  next = similar(cups)
  for i in 1:N
    next[cups[i]] = cups[mod1(i+1,N)]
  end
  next
end
function next2(cups)
  next = similar(cups)
  next[cups] = circshift(cups,-1)
  next
end

const cups_next = next2(cups)

function run(current::INT, next::Vector{INT}, steps::Int=1)
  N = INT(length(next))

  @inbounds for i in 1:steps
    #pickups = peek(next, current, Val(3))
    a = next[ current ]
    b = next[ a ]
    c = next[ b ]

    #dst = mod1(current-1, N)
    #while dst in pickups
    #  dst = mod1(dst-1, N)
    #end
    dst = current - INT(1)
    if dst == INT(0)
        dst = N
    end
    while dst == a || dst == b || dst == c
        dst -= INT(1)
        if dst == INT(0)
            dst = N
        end
    end

    next[current] = next[c]
    next[c] = next[dst]
    next[dst] = a
    current = next[current]
  end

  return next
end

part1() = join(peek(run(cups[1], cups_next, 10), 1, 8))
#part2() = prod(peek(run(vcat(cups, 10:1_000_000), 10_000_000), 1, 2))
function part2()
    #cups2 = vcat(cups, Int32(10):Int32(1_000_000))
    next = vcat(cups_next, INT(11):INT(1_000_000), cups[1])
    next[cups[end]] = 10
    # cups2_next == next2(cups2) == next1(cups2)
    next = run(cups[1], next, 10_000_000)
    prod( peek( next , 1 , 2) )
end
println( part2() )
@btime part2() #   171.104 ms (30 allocations: 3.82 MiB)
# 149245887792


########################### Stopping Criteria ###################################

# could generally be function of f, x, (d), α, t

# compositie stopping critera
# struct Any{T, A<:Tuple{Vararg{StoppingCriterion}}} <: StoppingCriterion{T}
#     args::A
# end
# (T::Any)(x, f, g, t) = any(c(x, f, g, t) for c in T.args)
#
# struct All{T, A<:Tuple{Vararg{StoppingCriterion}}} <: StoppingCriterion{T}
#     args::A
# end
# (T::All)(x, f, g, t) = all(c(x, f, g, t) for c in T.args)

# this is closest to the definition of fixed point finding
# struct XStop{T, V} <: StoppingCriterion{T}
#     ε::T
#     x::V
# end
# (T::XStop)(x::AbstractVector) = norm(difference(x, T.x)) < ε
# update!(T::XStop, x) = (T.x .= x) # update the stopping criterion
#
# # if function value does not decay enough (could be done in step size selection?)
# # should this be mutable since old_y is scalar but has to change (same with TStop)
# mutable struct YStop{T} <: StoppingCriterion{T}
#     δ::T
#     old_y::T
# end

# if we reach maximum iterations
# struct TStop{T} <: StoppingCriterion
#     max::Int
#    i::Int
# end

# if gradient falls below zero (more generally increment?)
# struct GradientStop{T} <: StoppingCriterion{T}
#     ε::T
# end


@recipe function f(grid::AbstractGrid{T,N}, vals::AbstractMatrix) where {T,N}
    seriestype --> :surface
    size --> (500,400)
    xrange = LinRange(endpoints(support(element(grid,1)))..., size(grid,1))
    yrange = LinRange(endpoints(support(element(grid,2)))..., size(grid,2))
    xrange, yrange, vals'
end

@recipe function f(grid::AbstractGrid{<:AbstractVector}, vals::AbstractVector)
    seriestype --> :surface
    legend=false
    size --> (500,400)
    xrange = [x[1] for x in grid]
    yrange = [x[2] for x in grid]
    xrange, yrange, vals
end

# Plot an Nd grid
@recipe function f(grid::AbstractGrid)
    seriestype --> :scatter
    size --> (500,400)
    legend --> false
    broadcast(x->tuple(x...),collect(grid)[1:end])
end

# Plot a 1D grid
@recipe function f(grid::AbstractGrid1d)
    seriestype --> :scatter
    yticks --> []
    ylims --> [-1 1]
    size --> (500,200)
    legend --> false
    collect(grid), zeros(size(grid))
end

@recipe function f(a::Vector{SArray{Tuple{2},T,1,2}}) where T
    A = zeros(T,length(a),2)
    for (l,e) in enumerate(a)
            A[l,:] .= e
    end
    A[:,1],A[:,2]
end


module Objective
export cost_minimization

using JuMP

function cost_minimization(model_dict)
    """
    Minimizes total system monetary cost.
    """
    backend_model= model_dict["backend_model"]
    cost = model_dict["variables"]["cost"]
    loc_techs_cost = model_dict["sets"]["loc_techs_cost"]
    @objective(backend_model, :Min,
        sum(cost["monetary", loc_tech] for loc_tech in loc_techs_cost)
    )
end
end


include("basic_setup.jl")

using DataFrames

function prepare(mean_s)
    t_s = 100
    t_x = 20

    κ = mean_s / t_s
    λ = 1 / t_s
    ρ = 1 / t_x
    μ = 1 / t_x
    mean_x = mean_s * ρ / μ

    # see Tostevin, ten Wolde, eq. 27
    sigma_squared_ss = mean_s
    sigma_squared_sx = ρ * mean_s / (λ + μ)
    sigma_squared_xx = mean_x * (1 + ρ / (λ + μ))

    joint_stationary = MvNormal([mean_s, mean_x], [sigma_squared_ss sigma_squared_sx; sigma_squared_sx sigma_squared_xx])
    signal_stationary = Poisson(mean_s)

    gen = Trajectories.configuration_generator(sn, rn, [κ, λ], [ρ, μ], signal_stationary, joint_stationary)

    (system, initial) = Trajectories.generate_configuration(gen; duration=100.0)
end

function get_results(algorithm, num_evals, system, initial)
    results = zeros(num_evals)
    Threads.@threads for i in 1:num_evals
        r = Trajectories.simulate(algorithm, initial, system)
        results[i] = Trajectories.log_marginal(r)
    end
    DataFrame(log_marginal=results, κ=system.sparams[1], λ=system.sparams[2], estimate=Trajectories.name(algorithm))
end

an = AnnealingEstimate(15, 50, 100)
ti = TIEstimate(1024, 6, 2^14)
direct = DirectMCEstimate(50_000)

function compute_data()
    df = DataFrame()
    for mean_s in [20, 50, 100, 200]
        system, initial = prepare(mean_s)
        df = vcat(df, get_results(an, 128, system, initial))
        df = vcat(df, get_results(ti, 128, system, initial))
        df = vcat(df, get_results(direct, 128, system, initial))
    end
    df
end

df.mean_s

# df = compute_data()
# using CSV
# CSV.write(projectdir("data", "estimator_variance", "estimator_comparison.csv"), df)

df["mean_s"] = df.κ ./ df.λ
df["mean_s_cat"] = string.(df.mean_s)

df2 = transform(groupby(df, [:mean_s]), :log_marginal => (x-> x .- mean(x[x.estimate == "AIS"])) => :deviation)
show(df2)

using Plots
using StatsPlots
using Statistics

@df df2 dotplot(:mean_s_cat, :deviation, group=(:estimate))

@show combine(groupby(df, [:estimate, :mean_s]), :log_marginal => mean, :log_marginal => std)

jitter_ti = randn(size(ti_samples)...)
jitter_an = randn(size(an_samples)...)
jitter_di = randn(size(di_samples)...)
p = scatter(jitter_ti, ti_samples, label="TI", xlim=(-15, 15))
scatter!(jitter_an, an_samples, label="AIS")
scatter!(jitter_di, di_samples, label="Direct MC")
plot!(p, [-3,3], mean(ti_samples) .* ones(2))
plot!(p, [-3,3], mean(an_samples) .* ones(2))
plot!(p, [-3,3], mean(di_samples) .* ones(2))

density([ti_samples, an_samples, di_samples], label=hcat("TI", "AIS", "Direct MC"))


using APCA
using Test
using Random

@testset "APCA.jl" begin
    v= [7.,5.,5.,3.,3.,3.,4.,6.]
    M= 3
    @test runAPCA(v,M) == [6.,6.,3.5,3.5,3.5,3.5,5.,5.]

    Random.seed!(1234)
    v = vcat(rand(10).+3,rand(10).+2,rand(10).-1)
    M = 4
    @test runAPCA(v,M) == [3.1124992225766124, 3.1124992225766124, 3.1124992225766124, 3.1124992225766124,
    3.1124992225766124, 3.1124992225766124, 3.1124992225766124, 3.1124992225766124,
    3.1124992225766124, 3.1124992225766124, 3.1124992225766124, 3.1124992225766124, 3.1124992225766124,
    3.1124992225766124, 3.1124992225766124, 3.1124992225766124, 2.456760310559139, 2.456760310559139, 2.456760310559139,
    2.456760310559139, -0.08728818214843481, -0.08728818214843481, -0.08728818214843481, -0.08728818214843481, -0.7055353573422534,
    -0.7055353573422534, -0.7055353573422534, -0.7055353573422534, -0.7055353573422534, -0.7055353573422534]

end


export
    SpectralConv,
    FourierOperator

struct SpectralConv{P, N, T, S, F}
    weight::T
    in_channel::S
    out_channel::S
    modes::NTuple{N, S}
    σ::F
end

"""
    SpectralConv(
        ch, modes, σ=identity;
        init=c_glorot_uniform, permuted=false, T=ComplexF32
    )

## Arguments

* `ch`: Input and output channel size, e.g. `64=>64`.
* `modes`: The Fourier modes to be preserved.
* `σ`: Activation function.
* `permuted`: Whether the dim is permuted. If `permuted=true`, layer accepts
    data in the order of `(ch, ..., batch)`, otherwise the order is `(..., ch, batch)`.

## Example

```jldoctest
julia> SpectralConv(2=>5, (16, ))
SpectralConv(2 => 5, (16,), σ=identity, permuted=false)

julia> using Flux

julia> SpectralConv(2=>5, (16, ), relu)
SpectralConv(2 => 5, (16,), σ=relu, permuted=false)

julia> SpectralConv(2=>5, (16, ), relu, permuted=true)
SpectralConv(2 => 5, (16,), σ=relu, permuted=true)
```
"""
function SpectralConv(
    ch::Pair{S, S},
    modes::NTuple{N, S},
    σ=identity;
    init=c_glorot_uniform,
    permuted=false,
    T::DataType=ComplexF32
) where {S<:Integer, N}
    in_chs, out_chs = ch
    scale = one(T) / (in_chs * out_chs)
    weights = scale * init(out_chs, in_chs, prod(modes))
    W = typeof(weights)
    F = typeof(σ)

    return SpectralConv{permuted,N,W,S,F}(weights, in_chs, out_chs, modes, σ)
end

Flux.@functor SpectralConv

Base.ndims(::SpectralConv{P,N}) where {P,N} = N

permuted(::SpectralConv{P}) where {P} = P

function Base.show(io::IO, l::SpectralConv{P}) where {P}
    print(io, "SpectralConv($(l.in_channel) => $(l.out_channel), $(l.modes), σ=$(string(l.σ)), permuted=$P)")
end

function spectral_conv(m::SpectralConv, 𝐱::AbstractArray)
    n_dims = ndims(𝐱)

    𝐱_fft = fft(Zygote.hook(real, 𝐱), 1:ndims(m)) # [x, in_chs, batch]
    𝐱_flattened = reshape(view(𝐱_fft, map(d->1:d, m.modes)..., :, :), :, size(𝐱_fft, n_dims-1), size(𝐱_fft, n_dims))
    𝐱_weighted = apply_spectral_pattern(𝐱_flattened, m.weight) # [prod(m.modes), out_chs, batch], only 3-dims
    𝐱_shaped = reshape(𝐱_weighted, m.modes..., size(𝐱_weighted, 2), size(𝐱_weighted, 3))
    𝐱_padded = spectral_pad(𝐱_shaped, (size(𝐱_fft)[1:end-2]..., size(𝐱_weighted, 2), size(𝐱_weighted, 3))) # [x, out_chs, batch] <- [modes, out_chs, batch]
    𝐱_ifft = real(ifft(𝐱_padded, 1:ndims(m))) # [x, out_chs, batch]

    return m.σ.(𝐱_ifft)
end

function (m::SpectralConv{false})(𝐱)
    𝐱ᵀ = permutedims(𝐱, (ntuple(i->i+1, ndims(m))..., 1, ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch]
    𝐱_out = spectral_conv(m, 𝐱ᵀ) # [x, out_chs, batch]
    𝐱_outᵀ = permutedims(𝐱_out, (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch]

    return 𝐱_outᵀ
end

function (m::SpectralConv{true})(𝐱)
    return spectral_conv(m, 𝐱) # [x, out_chs, batch]
end

############
# operator #
############

struct FourierOperator{L, C, F}
    linear::L
    conv::C
    σ::F
end

"""
    FourierOperator(ch, modes, σ=identity; permuted=false)

## Arguments

* `ch`: Input and output channel size for spectral convolution, e.g. `64=>64`.
* `modes`: The Fourier modes to be preserved for spectral convolution.
* `σ`: Activation function.
* `permuted`: Whether the dim is permuted. If `permuted=true`, layer accepts
    data in the order of `(ch, ..., batch)`, otherwise the order is `(..., ch, batch)`.

## Example

```jldoctest
julia> FourierOperator(2=>5, (16, ))
FourierOperator(2 => 5, (16,), σ=identity, permuted=false)

julia> using Flux

julia> FourierOperator(2=>5, (16, ), relu)
FourierOperator(2 => 5, (16,), σ=relu, permuted=false)

julia> FourierOperator(2=>5, (16, ), relu, permuted=true)
FourierOperator(2 => 5, (16,), σ=relu, permuted=true)
```
"""
function FourierOperator(
    ch::Pair{S, S},
    modes::NTuple{N, S},
    σ=identity;
    permuted=false
) where {S<:Integer, N}
    linear = permuted ? Conv(Tuple(ones(Int, length(modes))), ch) : Dense(ch.first, ch.second)
    conv = SpectralConv(ch, modes; permuted=permuted)

    return FourierOperator(linear, conv, σ)
end

Flux.@functor FourierOperator

function Base.show(io::IO, l::FourierOperator)
    print(io, "FourierOperator($(l.conv.in_channel) => $(l.conv.out_channel), $(l.conv.modes), σ=$(string(l.σ)), permuted=$(permuted(l.conv)))")
end

function (m::FourierOperator)(𝐱)
    return m.σ.(m.linear(𝐱) + m.conv(𝐱))
end

#########
# utils #
#########

c_glorot_uniform(dims...) = Flux.glorot_uniform(dims...) + Flux.glorot_uniform(dims...)*im

# [prod(modes), out_chs, batch] <- [prod(modes), in_chs, batch] * [out_chs, in_chs, prod(modes)]
apply_spectral_pattern(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m]

spectral_pad(𝐱::AbstractArray, dims::NTuple) = spectral_pad!(similar(𝐱, dims), 𝐱)

function spectral_pad!(𝐱_padded::AbstractArray, 𝐱::AbstractArray)
    fill!(𝐱_padded, eltype(𝐱)(0)) # zeros(eltype(𝐱), dims)
    𝐱_padded[map(d->1:d, size(𝐱))...] .= 𝐱

    return 𝐱_padded
end

function ChainRulesCore.rrule(::typeof(spectral_pad), 𝐱::AbstractArray, dims::NTuple)
    function spectral_pad_pullback(𝐲̄)
        return NoTangent(), view(𝐲̄, map(d->1:d, size(𝐱))...), NoTangent()
    end

    return spectral_pad(𝐱, dims), spectral_pad_pullback
end


using Test
include("HelmholtzTest.jl");
include("ShiftedLaplacianTest.jl");
include("AttenuationTest.jl");
include("testFictitiousSource2D.jl")
include("testHybridKacmarzHelmholtz.jl")


"""
A StateSpace model with a partioning imposed according to

A  | B1  B2
——— ————————
C1 | D11 D12
C2 | D21 D22

It corresponds to partioned input and output signals
u = [u1 u2]^T
y = [y1 y2]^T

"""
struct PartionedStateSpace{S<:AbstractStateSpace} <: LTISystem
    P::S
    nu1::Int
    ny1::Int
end
# For converting between different S
PartionedStateSpace{S}(partsys::PartionedStateSpace) where {S<:StateSpace} =
    PartionedStateSpace{S}(S(partsys.P), partsys.nu1, partsys.ny1)

function getproperty(sys::PartionedStateSpace, d::Symbol)
    P = getfield(sys, :P)
    nu1 = getfield(sys, :nu1)
    ny1 = getfield(sys, :ny1)

    if d === :Ts
        return P.Ts # Will throw deprecation until removed # DEPRECATED
    elseif d === :P
        return P
    elseif d === :nu1
        return nu1
    elseif d === :ny1
        return ny1
    elseif d === :A
        return P.A
    elseif d === :B1
        return P.B[:, 1:nu1]
    elseif d === :B2
        return P.B[:, nu1+1:end]
    elseif d === :C1
        return P.C[1:ny1, :]
    elseif d === :C2
        return P.C[ny1+1:end, :]
    elseif d === :D11
        return P.D[1:ny1, 1:nu1]
    elseif d === :D12
        return P.D[1:ny1, nu1+1:end]
    elseif d === :D21
        return P.D[ny1+1:end, 1:nu1]
    elseif d === :D22
        return P.D[ny1+1:end, nu1+1:end]
    else
        return getfield(P, d)
    end
end

timeevol(sys::PartionedStateSpace) = timeevol(sys.P)

function +(s1::PartionedStateSpace, s2::PartionedStateSpace)
    timeevol = common_timeevol(s1,s2)

    A = blockdiag(s1.A, s2.A)

    B = [[s1.B1; s2.B1] blockdiag(s1.B2, s2.B2)]

    C = [[s1.C1 s2.C1];
    blockdiag(s1.C2, s2.C2)]

    D = [(s1.D11 + s2.D11) s1.D12 s2.D12;
    [s1.D21; s2.D21] blockdiag(s1.D22, s2.D22)]

    P = StateSpace(A, B, C, D, timeevol) # How to handle discrete?
    PartionedStateSpace(P, s1.nu1 + s2.nu1, s1.ny1 + s2.ny1)
end

"""
    Series connection of partioned StateSpace systems.
"""
function *(s1::PartionedStateSpace, s2::PartionedStateSpace)
    timeevol = common_timeevol(s1,s2)

    A = [s1.A                           s1.B1*s2.C1;
    zeros(size(s2.A,1),size(s1.A,2))      s2.A]

    B = [s1.B1*s2.D11                         s1.B2           s1.B1*s2.D12;
    s2.B1              zeros(size(s2.B2,1),size(s1.B2,2))          s2.B2]

    C = [s1.C1                       s1.D11*s2.C1;
    s1.C2                        s1.D21*s2.C1;
    zeros(size(s2.C2,1),size(s1.C2,2))         s2.C2]

    D = [s1.D11*s2.D11           s1.D12        s1.D11*s2.D12;
    s1.D21*s2.D11           s1.D22        s1.D21*s2.D12;
    s2.D21          zeros(size(s2.D22,1),size(s1.D22,2))          s2.D22        ]

    P = StateSpace(A, B, C, D, timeevol)
    PartionedStateSpace(P, s2.nu1, s1.ny1)
end

# QUESTION: What about algebraic loops and well-posedness?! Perhaps issue warning if P1(∞)*P2(∞) > 1
function feedback(s1::PartionedStateSpace, s2::PartionedStateSpace)
    timeevol = common_timeevol(s1,s2)
    X_11 = (I + s2.D11*s1.D11)\[-s2.D11*s1.C1  -s2.C1]
    X_21 = (I + s1.D11*s2.D11)\[s1.C1  -s1.D11*s2.C1]

    # For the case of two outputs
    #    X_12 = [I   -s2.D11   -s2.D11*s1.D12   -s2.D12]
    #    X_22 = [s1.D11  I     s1.D12          -s1.D11*s2.D12]
    X_12 = (I + s2.D11*s1.D11)\[I      -s2.D11*s1.D12   -s2.D12]
    X_22 = (I + s1.D11*s2.D11)\[s1.D11   s1.D12          -s1.D11*s2.D12]

    A = [s1.B1 * X_11 ; s2.B1 * X_21] + blockdiag(s1.A, s2.A)

    B = [s1.B1 * X_12 ; s2.B1 * X_22]
    tmp = blockdiag(s1.B2, s2.B2)
    B[:, end-size(tmp,2)+1:end] .+= tmp

    C = [s1.D11 * X_11 ;
         s1.D21 * X_11 ;
         s2.D21 * X_21 ] + [s1.C1 zeros(size(s1.C1,1),size(s2.C1,2)); blockdiag(s1.C2, s2.C2)]

    D = [s1.D11 * X_12 ;
        s1.D21 * X_12 ;
        s2.D21 * X_22 ]
    tmp = [s1.D12 zeros(size(s1.D12,1),size(s2.D12,2)); blockdiag(s1.D22, s2.D22)]
    D[:, end-size(tmp,2)+1:end] .+= tmp

    # in case it is desired to consider both outputs
    # C = [s1.D11 * X_11 ;
    #      s2.D11 * X_21 ;
    #      s1.D21 * X_11 ;
    #      s2.D21 * X_21 ] + [blockdiag(s1.C1, s2.C1); blockdiag(s1.C2, s2.C2)]
    #
    # D = [s1.D11 * X_12 ;
    #     s2.D11 * X_22 ;
    #     s1.D21 * X_12 ;
    #     s2.D21 * X_22 ]
    #tmp = [blockdiag(s1.D12, s2.D12); blockdiag(s1.D22, s2.D22)]
    #D[:, end-size(tmp,2)+1:end] .+= tmp

    P = StateSpace(A, B, C, D, timeevol)
    PartionedStateSpace(P, s2.nu1, s1.ny1)
end

""" Concatenate systems vertically with
    same first input u1
    second input [u2_1; u2_2 ...]
    and output  y1 = [y1_1; y1_2, ...]
                y2 = [y2_1; y2_2, ...]
    for u1_i, u2_i, y1_i, y2_i, where i denotes system i
"""
function vcat_1(systems::PartionedStateSpace...)
    # Perform checks
    timeevol = common_timeevol(systems...)

    nu1 = systems[1].nu1
    if !all(s.nu1 == nu1 for s in systems)
        error("All systems must have same first input dimension")
    end

    A = blockdiag([s.A for s in systems]...)

    B1 = reduce(vcat, [s.B1 for s in systems])
    B2 = blockdiag([s.B2 for s in systems]...)

    C1 = blockdiag([s.C1 for s in systems]...)
    C2 = blockdiag([s.C2 for s in systems]...)

    D11 = reduce(vcat, [s.D11 for s in systems])
    D12 = blockdiag([s.D12 for s in systems]...)
    D21 = reduce(vcat, [s.D21 for s in systems])
    D22 = blockdiag([s.D22 for s in systems]...)

    sysnew = StateSpace(A, [B1 B2], [C1; C2], [D11 D12; D21 D22], timeevol)
    return PartionedStateSpace(sysnew, nu1, sum(s -> s.ny1, systems))
end

""" Concatenate systems horizontally with
    same first ouput y1 being the sum
    second output y2 = [y2_1; y2_2, ...]
    and inputs  u1 = [u1_1; u1_2, ...]
    u2 = [u2_1; u2_2, ...]
    for u1_i, u2_i, y1_i, y2_i, where i denotes system i
"""
function hcat_1(systems::PartionedStateSpace...)
    # Perform checks
    timeevol = common_timeevol(systems...)

    ny1 = systems[1].ny1
    if !all(s.ny1 == ny1 for s in systems)
        error("All systems must have same first ouput dimension")
    end

    A = blockdiag([s.A for s in systems]...)

    B1 = blockdiag([s.B1 for s in systems]...)
    B2 = blockdiag([s.B2 for s in systems]...)

    C1 = reduce(hcat, [s.C1 for s in systems])
    C2 = blockdiag([s.C2 for s in systems]...)

    D11 = reduce(hcat, [s.D11 for s in systems])
    D12 = reduce(hcat, [s.D12 for s in systems])
    D21 = blockdiag([s.D21 for s in systems]...)
    D22 = blockdiag([s.D22 for s in systems]...)

    sysnew = StateSpace(A, [B1 B2], [C1; C2], [D11 D12; D21 D22], timeevol)
    return PartionedStateSpace(sysnew, sum(s -> s.nu1, systems), ny1)
end

function Base.convert(::Type{<:PartionedStateSpace}, sys::T) where T<: StateSpace
    PartionedStateSpace(sys,sys.nu,sys.ny)
end


#= convex.jl
Philip Barrett, pobarrett@gmail.com
01jul2016, Chicago

Tests the (two) convex hull routines to guarantee consistency
=#

Z = [ 1 -1; 1 1; -1 1; 0 -1 ; -1 -1 ]
    # Add a redundant point
a = Polygon( pts=Z ) ;
b=grahamScan(a)

c = Point{2, Float64}[Point(Z[i, 1], Z[i, 2]) for i=1:size(Z)[1]]
d = CHull2D.convexhull(c)

Y = [ 1 -1; 1 1; -1 1; 0 -1 + 1e-06 ; -1 -1 ]
e = Polygon( pts=Y ) ;
f = grahamScan(e)

g = Point{2, Float64}[Point(Y[i, 1], Y[i, 2]) for i=1:size(Y)[1]]
h = CHull2D.convexhull(g)


# bootstrapping traits
const avg_rate = 0.05
const max_rate = 1.0

const AVG_HAZARD_RATE = 0.01
const MAX_HAZARD_RATE = 1.0

# Discount bootstrap trait
type Discount <: BootstrapTrait end
type HazardRate <: BootstrapTrait end

# Type aliases
# typealias InterpolatedHazardRateCurve{P, T} InterpolatedDefaultProbabilityCurve{P, HazardRate}

initial_value(::Discount) = 1.0
max_iterations(::Discount) = 100

initial_value(::HazardRate) = AVG_HAZARD_RATE
max_iterations(::HazardRate) = 30

function guess{T <: TermStructure}(::Discount, i::Int, ts::T, valid::Bool)
  if valid
    # return previous iteration value
    return ts.data[i]
  end

  if i == 2
    # first pillar
    return 1.0 / (1.0 + avg_rate * ts.times[2])
  end

  r = -log(ts.data[i - 1]) / ts.times[i - 1]
  return exp(-r * ts.times[i])
end

function min_value_after{T <: TermStructure}(::Discount, i::Int, ts::T, valid::Bool)
  if valid
    return ts.data[end] / 2.0
  end

  dt = ts.times[i] - ts.times[i - 1]
  return ts.data[i - 1] * exp( -max_rate * dt)
end

function max_value_after{T <: TermStructure}(::Discount, i::Int, ts::T, ::Bool)
  dt = ts.times[i] - ts.times[i - 1]
  return ts.data[i - 1] * exp(max_rate * dt)
  #return ts.data[i - 1]
end

function update_guess!{T <: TermStructure}(::Discount, i::Int, ts::T, discount::Float64)
  ts.data[i] = discount
  return ts
end

function guess(::HazardRate, i::Int, ts::TermStructure, valid::Bool)
  if valid
    return ts.data[i]
  end

  if i == 2 # first pillar
    return AVG_HAZARD_RATE
  end

  # extrapolate
  d = ts.dates[i]
  return hazard_rate(ts, d)
end

function min_value_after(::HazardRate, i::Int, ts::TermStructure, valid::Bool)
  if valid
    r = min(ts.data[1], ts.data[end])
    return r / 2.0
  end

  return eps()
end

function max_value_after(::HazardRate, i::Int, ts::TermStructure, valid::Bool)
  if valid
    r = max(ts.data[1], ts.data[end])
    return r * 2.0
  end

  return MAX_HAZARD_RATE
end

function update_guess!(::HazardRate, i::Int, ts::TermStructure, rate::Float64)
  ts.data[i] = rate
  if i == 2
    ts.data[1] = rate # first point handled as well
  end

  return ts
end


# multimodal_gmm_langevin_floor
#
# Considers the stein discrepancy for a 2 mixture GMM with gap between modes
# gap > 0. It compares a bad distribution [which only samples from one mode]
# and the actual iid sequence.

include("experiment_utils.jl")

# this controls the seed used for the approximating the wasserstein metric
@parseintcli seed "s" "seed" 7
# the gap between the two modes
@parseintcli gap "g" "gap" 3
# should we use multiple cores?
@parseintcli numcores "n" "numcores" 1
# the dimension of the space
@parseintcli d "d" "dimension" 1
# the gap between the two modes
@parsestringcli discrepancytype "t" "discrepancytype" "kernel"
# the gap between the two modes
@parsestringcli sampler "q" "sampler" "iid"
# parallelize it
if (numcores > 1)
    if (numcores > CPU_CORES)
        error(@sprintf("Requested %d cores, only %d available.", numcores, CPU_CORES))
    end
    addprocs(numcores - 1)
    @assert (numcores == nprocs())
end

using Distributions: MvNormal

using SteinDistributions:
    SteinGMM,
    gradlogdensity,
    cdf
using SteinDiscrepancy:
    SteinInverseMultiquadricKernel,
    ksd,
    gsd,
    wasserstein1d

# select largest n
maxn = 30000
# Select an optimization problem solver
if discrepancytype == "graph"
#    solver = "clp"
    import Gurobi
    solver = Gurobi.GurobiSolver(Threads=numcores)
end
# the number of components
K = 2
# distribution sampled from
distname = "$K-gmm"
# target distribution
gmmmus = gap .* [-0.5, 0.5]
gmmmus = hcat(gmmmus, zeros(2, d-1))
gmmsigmas = eye(d)
gmmweights = [0.5, 0.5]
target = SteinGMM(gmmmus, gmmsigmas)
# define gradlogp
function gradlogp(x::Array{Float64,1})
    gradlogdensity(target, x)
end
# Sample sizes at which optimization problem will be solved
ns = vcat([5],
          10:10:min(100,maxn),
          200:100:min(1000,maxn),
          2000:1000:min(10000,maxn),
          20000:10000:min(100000,maxn))

if sampler == "iid"
    X = @setseed rand(target, maxn)
    if discrepancytype == "kernel"
        # this is so the g-h plot in paper are in same range
        X = -X
    end
elseif sampler == "unimodal"
    wrongtarget = MvNormal(vec(gmmmus[1,:]), gmmsigmas)
    X = @setseed rand(wrongtarget, maxn)'
end

println("Beginning optimization!")
for i in ns
    # initialize all the variables
    objectivevalue = nothing
    edgetime = nothing
    solvetime = nothing
    points = X[1:i,:]
    g = nothing
    gradg = nothing
    operatorg = nothing
    if discrepancytype == "graph"
        res = gsd(points=X[1:i,:],
                  gradlogdensity=gradlogp,
                  solver=solver)
        objectivevalue = res.objectivevalue
        edgetime = res.edgetime
        solvetime = res.solvetime
        points = res.points
        g = res.g
        gradg = res.gradg
        operatorg = res.operatorg
    elseif discrepancytype == "kernel"
        kernel = SteinInverseMultiquadricKernel()
        res = ksd(points=X[1:i,:],
                  gradlogdensity=gradlogp,
                  kernel=kernel)
        objectivevalue = sqrt(res.discrepancy2)
        solvetime = res.solvetime
    elseif discrepancytype == "wasserstein"
        @assert (d == 1)
        (objectivevalue, error) = wasserstein1d(X[1:i,:],
                                                targetcdf=x::Float64 -> cdf(target, x))
    end

    println("\tn = $(i), objective = $(objectivevalue)")
    # Package and save results
    instance_data = Dict{Any, Any}(
        "seed" => seed,
        "mus" => gmmmus,
        "sigmas" => gmmsigmas,
        "distweights" => gmmweights,
        "sampler" => sampler,
        "gap" => gap,
        "distname" => distname,
        "X" => points,
        "g" => g,
        "gradg" => gradg,
        "operatorg" => operatorg,
        "n" => i,
        "d" => d,
        "objectivevalue" => objectivevalue,
        "discrepancytype" => discrepancytype,
        "edgetime" => edgetime,
        "solvetime" => solvetime,
        "ncores" => nprocs(),
    )

    save_json(
        instance_data;
        dir="multimodal_gmm_langevin_floor",
        distname=distname,
        sampler=sampler,
        discrepancytype=discrepancytype,
        n=i,
        d=d,
        gap=gap,
        seed=seed,
        numcores=nprocs(),
    )
end

println("COMPLETE!")


using ModelingToolkit, OrdinaryDiffEq, DiffEqParamEstim, MinimallyDisruptiveCurves, ForwardDiff, LinearAlgebra, Test

function make_model(input)
    @parameters t
    @parameters k, c, m
    D = Differential(t)
    @variables pos(t) vel(t)

    eqs = [D(pos) ~ vel,
        D(vel) ~ (-1 / m) * (c * vel + k * pos - input(t))
  ]

    ps = [k,c,m] .=>  [2.,1.,4.]
    ics = [pos, vel] .=> [1.,0.]
    od = ODESystem(eqs, t, first.(ics), first.(ps)
                                 , defaults=merge(Dict(first.(ics) .=> last.(ics))
                                 , Dict(first.(ps) .=> last.(ps))), name=:mass_spring
                                 )
    tspan = (0., 100.)
  # prob = ODEProblem(od, ics, tspan, ps)
    return od, ics, tspan, ps
end

od, ics, tspan, ps = make_model(t -> 0.)

"""
take a log transform of the od parameter space in two ways: at the level of the ODESystem and the level of the ODEProblem
"""

# to_fix = ["c2c","c2","c2a","c3c", "c1c", "a2"]
# tstrct_fix = fix_params(last.(ps), get_name_ids(ps, to_fix))

p0 = last.(ps)
tr = logabs_transform(p0)
log_od = transform_ODESystem(od, tr)
@test typeof(log_od) == ODESystem

prob1 = ODEProblem(od, [], tspan, [])

log_od2, log_ics2, log_ps2 = transform_problem(prob1, tr; unames=ModelingToolkit.get_states(od), pnames=ModelingToolkit.get_ps(od))

"""
check if the two manners of transforming the ODE system give the same output
"""

@test repr.(ModelingToolkit.get_ps(log_od)) == repr.(ModelingToolkit.get_ps(log_od2))

log_prob1 = ODEProblem(log_od, [], tspan, [])
log_prob2 = ODEProblem(log_od2, [], tspan, [])

sol1 = solve(log_prob1, Tsit5())
sol2 = solve(log_prob2, Tsit5())

@test sol1[end] == sol2[end]

"""
check if  log transforming the cost function on od gives the same result as an untransformed cost function on log_od
"""
tsteps = tspan[1]:1.:tspan[end]
nom_sol = solve(prob1, Tsit5())
lossf(sol) = sum([sum(abs2, el1 - el2) for (el1, el2) in zip(sol(tsteps).u, nom_sol(tsteps).u) ])

nom_cost = build_loss_objective(prob1, Tsit5(), lossf; mpg_autodiff=true)
log_cost = build_loss_objective(log_prob1, Tsit5(), lossf; mpg_autodiff=true)

@test nom_cost(p0) == log_cost(log.(p0))

tr_cost, newp0 = transform_cost(nom_cost, p0, tr)
@test tr_cost(newp0) == log_cost(log.(p0))
grad_holder = deepcopy(p0)
g2 = deepcopy(grad_holder)
"""
test that summing losses works
"""
ll = sum_losses([nom_cost,nom_cost], p0)
@test ll(p0 .+ 1., grad_holder) == 2nom_cost(p0 .+1, g2)
@test grad_holder == 2g2

"""
test gradients of cost functions are zero at minimum as a proxy for correctness of their gradients
"""

for el in (log_cost, tr_cost)
    el(newp0, grad_holder)
    @test norm(grad_holder) < 1e-3 # 0 gradient at minimum
end

"""
test that mdc curve evolves, and listens to mdc_callbacks
"""
H0 = ForwardDiff.hessian(tr_cost, newp0)
mom = 1.
span = (-1., 1.);

newdp0 = (eigen(H0)).vectors[:, 1]

eprob = MDCProblem(log_cost, newp0, newdp0, mom, span);

cb = [
    Verbose([CurveDistance(0.1:1:10), HamiltonianResidual(2.3:4:10)]),
    ParameterBounds([1,3], [-10.,-10.], [10.,10.])
    ]

@time mdc = evolve(eprob, Tsit5; mdc_callback=cb);

"""
check mdc works with mdc_callback vector of subtype T <: CallbackCallable, strict subtype.
"""

cb = [
    Verbose([CurveDistance(0.1:1:10), HamiltonianResidual(2.3:4:10)])
    ]

@time mdc = evolve(eprob, Tsit5; mdc_callback=cb);

"""
 test MDC works and gives reasonable output
"""

@test log_cost(mdc.sol[1][1:3]) < 1e-3
@test log_cost(mdc.sol[end][1:3]) < 1e-3

"""
 test injection loss works OK
"""
l2 = build_injection_loss(prob1, Tsit5(), tsteps)
@test l2(p0, grad_holder) == 0
@test norm(grad_holder) < 1e-5

"""
test fixing parameters (i.e. a transformation that changes the number of parameters) works ok
"""
to_fix = ["c"]
trf = fix_params(last.(ps), get_name_ids(ps, to_fix))
de = MinimallyDisruptiveCurves.transform_ODESystem(od, trf)
@test length(ModelingToolkit.get_ps(de)) == 2

"""
test jumpstarting works
"""
jprob = jumpstart(eprob, 1e-2, true)
mdcj = evolve(jprob, Tsit5; mdc_callback=cb);
@test log_cost(mdcj.sol[1][1:3]) < 1e-3


export MCVanilla, Vegas, Domain
export integral, ∫

using QuickTypes: @qstruct
using ArgCheck
using LinearAlgebra, Statistics
using StaticArrays
using Base.Threads: @threads
using Setfield: @settable

import Random

abstract type MCAlgorithm end

@settable @qstruct MCVanilla{R}(
        neval::Int64=10^6,
        rng::R=GLOBAL_PAR_RNG,
    ) do
        @argcheck neval > 2
end <: MCAlgorithm

"""
    Domain

Encodes a domain of integration, i.e. a product of finte intervals.
"""
struct Domain{N,T}
    lower::SVector{N,T}
    upper::SVector{N,T}
end

# function cartesian_product(d1::Domain, d2::Domain)
#     Domain(vcat(d1.lower, d2.lower), vcat(d1.upper, d2.upper))
# end

function pointtype(::Type{Domain{N,T}}) where {N,T}
    SVector{N,T}
end

function volume(dom::Domain)
    prod(dom.upper - dom.lower)
end

function Base.ndims(dom::Domain{N,T}) where {N,T}
    N
end

function Domain(interval::NTuple{2,Number})
    (a,b) = float.(promote(interval...))
    Domain(@SVector[a], @SVector[b])
end

function Domain(dom::Domain)
    dom
end

function Domain(lims)
    lower = SVector(map(float ∘ first, lims))
    upper = SVector(map(float ∘ last, lims))
    Domain(lower, upper)
end

function uniform(rng, dom::Domain)
    V = pointtype(typeof(dom))
    Δ = (dom.upper - dom.lower)
    dom.lower .+ rand(rng, V) .* Δ
end

"""

    integral(f, dom [,alg])

Compute the integral of the function `f` over a domain `dom`, via the algorithm `alg`.
"""
function integral(f, dom, alg=Vegas())
    f2, dom2, alg2 = canonicalize(f, dom, alg)
    integral_kernel(f2, dom2, alg2)
end

function canonicalize(f, dom, alg)
    f, Domain(dom), alg
end

function canonicalize(f, I::NTuple{2,Number}, alg)
    f_v = f ∘ first
    dom = Domain(I)
    f_v, dom, alg
end

const ∫ = integral

function integral_kernel(f, dom::Domain, alg::MCVanilla)
    mc_kernel(f, alg.rng, dom, neval=alg.neval)
end
function draw(rng, dom::Domain)
    vol = volume(dom)
    x = uniform(rng, dom)
    (position=x, weight = volume(dom))
end

"""
    mc_kernel(f, rng::AbstractRNG, dom; neval)

Monte Carlo integration of function `f` over `dom`.
`dom` must support the following methods:
* volume(dom): Return the volume of dom
* draw(rng, dom): Return an object with properties `position::SVector`, `weight::Real`.
"""
function mc_kernel(f, rng::AbstractRNG, dom; neval)
    N = neval
    x = uniform(rng, dom)
    y = float(f(x)) * volume(dom)
    sum = y
    sum2 = y.^2
    for _ in 1:(N-1)
        s = draw(rng, dom)
        y = s.weight * f(s.position)
        sum += y
        sum2 += y.^2
    end

    mean = sum/N
    var_f = (sum2/N) - mean .^2
    var_f = max(zero(var_f), var_f)
    var_fmean = var_f / N
    std = sqrt.(var_fmean)
    (value = mean, std=std, neval=N)
end

function mc_kernel(f, p::ParallelRNG, dom; neval)
    rngs = p.rngs
    res1 = mc_kernel(f, rngs[1], dom, neval=2)
    T = typeof(res1)
    nthreads = length(rngs)
    results = Vector{T}(undef, nthreads)
    neval_i = ceil(Int, neval / nthreads)
    @threads for i in 1:nthreads
        res = mc_kernel(f, rngs[i], dom, neval=neval_i)
        results[i] = res
    end
    fuseall(results)
end

function fuseall(results)
    N = sum(res->res.neval, results)
    value = sum(res->res.value * res.neval/N   , results)
    var   = sum(res->(res.std * res.neval/N).^2, results)
    (value=value, std=sqrt.(var), neval=N)
end

# function fuse(res1, res2)
#     var1 = res1.std .^ 2
#     var2 = res2.std .^ 2
#     w1, w2 = fusion_weights(var1, var2)
#     (
#         value = @.(w1 * res1.value + w2*res2.value),
#         std = @.(sqrt(w1^2*var1 + w2^2*var2))
#     )
# end
#
# function fusion_weights(var1::AbstractVector, var2::AbstractVector)
#     pairs = fusion_weights_scalar.(var1, var2)
#     first.(pairs), last.(pairs)
# end
#
# function fusion_weights(var1, var2)
#     fusion_weights_scalar(var1, var2)
# end
#
# function fusion_weights_scalar(var1, var2)
#     z = zero(var1)
#     o = one(var1)
#     if iszero(var1)
#         (o, z)
#     elseif iszero(var2)
#         (z,o)
#     else
#         p1 = o/var1
#         p2 = o/var2
#         p  = p1 + p2
#         (p1/p, p2/p)
#     end
# end


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 7927ff00-1806-11ec-2787-4db3c2bd8448
# imports
begin
	using Unitful, UnitfulAstro
	using UnitfulRecipes
	using Revise
	using LinearAlgebra
	using Pkg
	Pkg.develop(path="/Users/miles/dev/julia/Transits")
	using Transits
	using Plots
	using Markdown
end

# ╔═╡ 87591ebf-f36b-4643-bbf1-a961ed1825ff
md"""
# Equillibrium temperature in Sirius

When considering a planet orbiting Sirius B, it is interesting to consider the irradiation from Sirius A. This is especially significant due to Sirius B's eccentric orbit (``e=0.6``), leading to an overall range of separations from 7.9-32 AU.

The calculation for equilibrium temperature from a single star is
```math
T_\mathrm{eq} = \left[\frac{L(1 - A_B)}{16\sigma\pi a^2} \right]^{1/4}
```

we assume that adding an additional star will linearly add to the incident flux, so we can extend this to two stars with
```math
T_\mathrm{eq} = \sum_i{\left[\frac{L_i(1 - A_B)}{16\sigma\pi a_i^2} \right]^{1/4}}
```

This does not take into effect any reheating or geometric phase effects, which means this is likely an *upper limit*.
"""

# ╔═╡ b4375a6e-c669-4f64-aa9c-72375bf96868
begin
	plx = 0.3789 # arcsec
	dist = inv(plx)u"pc" # parsec
end

# ╔═╡ 0bdbeded-0d6f-426b-80bd-0d20d058baf8
# don't include i, Ω, or ω because we don't care about the
# line-of-sight rotation, only the relative position
# orbital elements from Bond et al. 2017
orbit = KeplerianOrbit(
	period=50.1284u"yr",
	a=7.4957u"AU" / plx, # convert arcsec to AU
	#incl=136.336u"°",
	#Omega=45.400u"°",
	t_0=1994.5715u"yr",
	ecc=0.59142,
	omega=0u"°"
	#omega=149.161u"°"
)

# ╔═╡ 0dd1df5c-7faa-4795-b4cd-0ebb9ba9e0c0
function relative_position(orbit, t)
	# orbit is in xzy because it wasn't rotated
	xzy = Transits.Orbits.relative_position(orbit, t)
	return [xzy[1], xzy[3]]
end

# ╔═╡ 181ea5bd-18f0-47ae-8911-9c823b8b5531
# plot orbit
begin
	t = range(0u"d", orbit.period, length=1000)
	t_today = 2020u"yr" + 260u"d" |> u"d"
	# calculate separations
	pos = mapreduce(ti -> relative_position(orbit, ti), hcat, t)
	plot(eachrow(pos)...,
		c=:black,
		leg=false,
		xlabel="Δx [AU]", ylabel="Δy [AU]")
	scatter!([0], [0], m=:star, ms=10, c=:black)
	pos_today = relative_position(orbit, t_today)
	scatter!(pos_today[1:1], pos_today[2:2], m=:star, ms=5, c=:black)
end

# ╔═╡ 1235c53f-e01f-457d-9057-7587af56a9f8
function eq_temp(orbit, t, dx, dy)
	pos = relative_position(orbit, t)u"AU"
	d = [dx, dy]
	a1 = norm(d)        # dist to Sirius A
	a2 = norm(d .- pos) # dist to Sirius B
	# assume albedo is negligible
	L1 = 25.4u"Lsun"
	L2 = 0.056u"Lsun"
	Ts = (L1 / (16π * Unitful.σ * a1^2))^(1/4) +
	     (L2 / (16π * Unitful.σ * a2^2))^(1/4)
	return Ts |> u"K"
end

# ╔═╡ 7738d09a-b7c5-4a38-b0fe-aafcba7e092a
begin
	dxs = range(-10, 35, length=500)u"AU"
	dys = range(-17, 17, length=500)u"AU"
end

# ╔═╡ 0dfa9c8b-0dd1-4967-99ce-6e729a3506a9
function eq_temp(orbit, t, dist)
	pos = relative_position(orbit, t)u"AU"
	d = norm(pos) # total orbital distance
	a1 = d - dist # distance to Sirius A
	a2 = dist     # distance to Sirius B
	# assume albedo is negligible
	L1 = 25.4u"Lsun"
	L2 = 0.056u"Lsun"
	T1 = (L1 / (16π * Unitful.σ * a1^2))^(1/4)
	T2 = (L2 / (16π * Unitful.σ * a2^2))^(1/4)
	return (T1 + T2, T1, T2) .|> u"K"
end

# ╔═╡ 57a6fbf8-0b5b-4a0e-b60a-e7d367692446
let temps
	temps = eq_temp.(orbit, t_today, dxs', dys)
	heatmap(dxs, dys, log10.(ustrip.(temps)),
		# levels=100,
		xlim=extrema(dxs),
		ylim=extrema(dys),
		c=:acton,
		cbar_title="log10(Teq)",
		title="today"
	)
	plot!(eachrow(pos)..., c=:white, leg=false, xlabel="Δx [AU]", ylabel="Δy [AU]")
	scatter!([0], [0], m=:star, ms=10, c=:white)
	scatter!(pos_today[1:1], m=:star, ms=5, pos_today[2:2], c=:white)
end

# ╔═╡ 56cd4737-3bdc-43da-a1fc-530baca135a3
let temps
	t_per = orbit.t_0 + orbit.t_ref
	temps = eq_temp.(orbit, t_per, dxs', dys)
	heatmap(dxs, dys, log10.(ustrip.(temps)),
		# levels=100,
		xlim=extrema(dxs),
		ylim=extrema(dys),
		c=:acton,
		cbar_title="log10(Teq)",
		title="periapsis"
	)
	plot!(eachrow(pos)..., c=:white, leg=false, xlabel="Δx [AU]", ylabel="Δy [AU]")
	scatter!([0], [0], m=:star, ms=10, c=:white)
	pos_per = relative_position(orbit, t_per)
	scatter!(pos_per[1:1], pos_per[2:2], m=:star, ms=5, c=:white)
end

# ╔═╡ 0523a4b3-7ecc-444e-853f-3adb46b8c04d
let temps
	t_apo = orbit.t_0 + orbit.t_ref + orbit.period/2
	temps = eq_temp.(orbit, t_apo, dxs', dys)
	heatmap(dxs, dys, log10.(ustrip.(temps)),
		# levels=100,
		xlim=extrema(dxs),
		ylim=extrema(dys),
		c=:acton,
		cbar_title="log10(Teq)",
		title="apoapsis"
	)
	plot!(eachrow(pos)..., c=:white, leg=false, xlabel="Δx [AU]", ylabel="Δy [AU]")
	scatter!([0], [0], m=:star, ms=10, c=:white)
	pos_apo = relative_position(orbit, t_apo)
	scatter!(pos_apo[1:1], pos_apo[2:2], m=:star, ms=5, c=:white)
end

# ╔═╡ cb972d70-1652-464a-9b35-215921674735
begin
	dists = range(0, 2, length=100)u"AU"
	plot(dists, map(d -> eq_temp(orbit, t_today, d)[1], dists),
		lab="",
		ylabel="Teq",
		xlabel="separation from Sirius B",
		title="today",
		xlim=extrema(dists)
	)
	plot!(dists, map(d -> eq_temp(orbit, t_today, d)[2], dists),
		ls=:dash,
		lab="Sirius A"
	)
	plot!(dists, map(d -> eq_temp(orbit, t_today, d)[3], dists),
		ls=:dash,
		lab="Sirius B"
	)
end

# ╔═╡ 6c63961a-f6b6-47e7-8dad-854ff1746aeb
begin
	t_per = orbit.t_0 + orbit.t_ref
	plot(dists, map(d -> eq_temp(orbit, t_per, d)[1], dists),
		lab="",
		ylabel="Teq",
		xlabel="separation from Sirius B",
		title="periapsis",
		xlim=extrema(dists)
	)
	plot!(dists, map(d -> eq_temp(orbit, t_per, d)[2], dists),
		ls=:dash,
		lab="Sirius A"
	)
	plot!(dists, map(d -> eq_temp(orbit, t_per, d)[3], dists),
		ls=:dash,
		lab="Sirius B"
	)
end

# ╔═╡ 860fc3e3-08d5-4cd5-8a4a-3b9db25a2b0e
begin
	t_apo = orbit.t_0 + orbit.t_ref + orbit.period/2
	plot(dists, map(d -> eq_temp(orbit, t_apo, d)[1], dists),
		lab="",
		ylabel="Teq",
		xlabel="separation from Sirius B",
		title="apoapsis",
		xlim=extrema(dists)
	)
	plot!(dists, map(d -> eq_temp(orbit, t_apo, d)[2], dists),
		ls=:dash,
		lab="Sirius A"
	)
	plot!(dists, map(d -> eq_temp(orbit, t_apo, d)[3], dists),
		ls=:dash,
		lab="Sirius B"
	)
end

# ╔═╡ d97e3642-f701-46e2-8c59-13148d2bebae
begin
	curves = hcat(
		map(d -> eq_temp(orbit, t_today, d)[1], dists),
		map(d -> eq_temp(orbit, t_per, d)[1], dists),
		map(d -> eq_temp(orbit, t_apo, d)[1], dists),
	)
	plot(dists, curves,
		lab=["today" "periapsis" "apoapsis"],
		ylabel="Teq",
		xlabel="separation from Sirius B",
		xlim=extrema(dists)
	)
end

# ╔═╡ Cell order:
# ╟─87591ebf-f36b-4643-bbf1-a961ed1825ff
# ╠═b4375a6e-c669-4f64-aa9c-72375bf96868
# ╠═0bdbeded-0d6f-426b-80bd-0d20d058baf8
# ╠═0dd1df5c-7faa-4795-b4cd-0ebb9ba9e0c0
# ╟─181ea5bd-18f0-47ae-8911-9c823b8b5531
# ╠═1235c53f-e01f-457d-9057-7587af56a9f8
# ╠═7738d09a-b7c5-4a38-b0fe-aafcba7e092a
# ╟─57a6fbf8-0b5b-4a0e-b60a-e7d367692446
# ╟─56cd4737-3bdc-43da-a1fc-530baca135a3
# ╟─0523a4b3-7ecc-444e-853f-3adb46b8c04d
# ╠═0dfa9c8b-0dd1-4967-99ce-6e729a3506a9
# ╟─cb972d70-1652-464a-9b35-215921674735
# ╟─6c63961a-f6b6-47e7-8dad-854ff1746aeb
# ╟─860fc3e3-08d5-4cd5-8a4a-3b9db25a2b0e
# ╟─d97e3642-f701-46e2-8c59-13148d2bebae
# ╠═7927ff00-1806-11ec-2787-4db3c2bd8448


@safetestset hilbert_decode_zero = "hilbert decode zero" begin
  using BijectiveHilbert
  table2 = [
    15 12 11 10
    14 13  8  9
    1  2  7  6
    0  3  4  5
  ]
  table2 = table2[end:-1:1, :]

  table2ind = CartesianIndices(table2)
  gg = Simple2D(Int)
  X = zeros(Int, 2)
  for z in 0:15
    BijectiveHilbert.decode_hilbert_zero!(gg, X, z)
    found = table2ind[table2 .== z][1]
    @test found[2] - 1 == X[1]
    @test found[1] - 1 == X[2]
  end
end

@safetestset simple2d_own_inverse = "simple2d is its own inverse" begin
  using BijectiveHilbert: Simple2D
  using ..HilbertTestSuite: check_own_inverse
  gg = Simple2D(Int)
  for b in 2:7
    @test check_own_inverse(gg, b, 2)
  end
end

@safetestset simple2d_complete_set = "simple2d is a complete set" begin
  using BijectiveHilbert: Simple2D
  using ..HilbertTestSuite: check_complete_set
  gg = Simple2D(Int)
  for b in 2:7
    @test check_complete_set(gg, b, 2)
  end
end

@safetestset simple2d_type_interactions = "Simple2D type interactions" begin
    using BijectiveHilbert
    using UnitTestDesign
    using Random
    rng = Random.MersenneTwister(9790323)
    for retrial in 1:5
        AxisTypes = shuffle(rng, [Int8, Int, UInt, Int128, UInt8, UInt128])
        IndexTypes = shuffle(rng, [Int8, UInt8, Int, UInt, Int128, UInt128])
        Count= shuffle(rng, [0, 1])
        Dims = shuffle(rng, [2, 3, 4])
        Bits = shuffle(rng, [2, 3, 4, 5])
        test_set = all_pairs(
            AxisTypes, IndexTypes, Count, Dims, Bits;
        )
        for (A, I, C, D, B) in test_set
            gg = Simple2D(I)
            if B * D > log2(typemax(I))
                continue
            end
            # Add this because these values aren't in Hilbert order because
            # It's an asymmetrical space.
            if B * D > log2(typemax(A))
              continue
            end
          last = (one(I) << (B * D)) - one(I) + I(C)
            mid = one(I) << (B * D - 1)
            few = 5
            X = zeros(A, D)
            hlarr = vcat(C:min(mid, few), max(mid + 1, last - few):last)
            for hl in hlarr
                hli = I(hl)
                if C == 0
                    decode_hilbert_zero!(gg, X, hli)
                    hl2 = encode_hilbert_zero(gg, X)
                    if hl2 != hli
                        @show A, I, C, D, B, X
                        @test hl2 == hli
                    end
                    @test typeof(hl2) == typeof(hli)
                else
                    decode_hilbert!(gg, X, hli)
                    hl3 = encode_hilbert(gg, X)
                    @test hl3 == hli
                    @test typeof(hl3) == typeof(hli)
                end
            end
        end
    end
end


ENV["GKSwstype"] = "nul"

using Printf
using Statistics
using Plots

using Oceananigans
using Oceananigans.Operators: ζ₃ᶠᶠᶜ
using Oceananigans.Advection
using Oceananigans.AbstractOperations
using Oceananigans.OutputWriters
using Oceananigans.Grids
using Oceananigans.Fields
using Oceananigans.Utils: prettytime
using Oceananigans.Models.HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceModel, ExplicitFreeSurface
using Oceananigans.OutputReaders: FieldTimeSeries

using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom
using Oceananigans.Advection: WENOVectorInvariantVel, WENOVectorInvariantVort, VectorInvariant, VelocityStencil, VorticityStencil

using CUDA: device!

device!(2)

#####
##### The Bickley jet
#####

Ψ(y) = - tanh(y)
U(y) = sech(y)^2

# A sinusoidal tracer
C(y, L) = sin(2π * y / L)

# Slightly off-center vortical perturbations
ψ̃(x, y, ℓ, k) = exp(-(y + ℓ/10)^2 / 2ℓ^2) * cos(k * x) * cos(k * y)

# Vortical velocity fields (ũ, ṽ) = (-∂_y, +∂_x) ψ̃
ũ(x, y, ℓ, k) = + ψ̃(x, y, ℓ, k) * (k * tan(k * y) + y / ℓ^2)
ṽ(x, y, ℓ, k) = - ψ̃(x, y, ℓ, k) * k * tan(k * x)

"""
    run_bickley_jet(output_time_interval = 2, stop_time = 200, arch = CPU(), Nh = 64, ν = 0,
                    momentum_advection = VectorInvariant())

Run the Bickley jet validation experiment until `stop_time` using `momentum_advection`
scheme or formulation, with horizontal resolution `Nh`, viscosity `ν`, on `arch`itecture.
"""
function run_bickley_jet(; output_time_interval = 2, stop_time = 200, arch = CPU(), Nh = 64, ν = 0,
                           momentum_advection = VectorInvariant())

    Nφ = Int(Nh/2)

    grid = LatitudeLongitudeGrid(arch, size=(Nh, Nφ, 1),
                                 radius = 1,
                                 longitude = (-180, 180), latitude=(-80, 80), z=(0, 1), halo = (4, 4, 4),
                                 precompute_metrics = true)

    @inline toplft(x, y) = (((x > 45) & (x < 135)) & (((y > 15) & (y < 30)) | ((y > 45) & (y < 60))))
    @inline botlft(x, y) = (((x > 45) & (x < 135)) & (((y < -15) & (y > -30)) | ((y < -45) & (y > -60))))

    @inline toprgt(x, y) = (((x < - 45) & (x > - 135)) & (((y > 15) & (y < 30)) | ((y > 45) & (y < 60))))
    @inline botrgt(x, y) = (((x < - 45) & (x > - 135)) & (((y < -15) & (y > -30)) | ((y < -45) & (y > -60))))

    @inline bottom(x, y) = Int(toplft(x, y) | botlft(x, y) | toprgt(x, y) | botrgt(x, y))

    grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -1))

    free_surface = ImplicitFreeSurface(solver_method=:HeptadiagonalIterativeSolver, gravitational_acceleration=1)

    model = HydrostaticFreeSurfaceModel(momentum_advection = momentum_advection,
                                          tracer_advection = WENO5(),
                                                      grid = grid,
                                                   tracers = :c,
                                                   closure = nothing,
                                              free_surface = free_surface,
                                                  coriolis = HydrostaticSphericalCoriolis(rotation_rate=1),
                                                  buoyancy = nothing)

    # ** Initial conditions **
    #
    # u, v: Large-scale jet + vortical perturbations
    #    c: Sinusoid

    # Parameters
    ϵ = 0.1 # perturbation magnitude
    ℓ = 0.5 # Gaussian width
    k = 0.5 # Sinusoidal wavenumber

    dr(x) = deg2rad(x)

    # Total initial conditions
    uᵢ(x, y, z) = U(dr(y)*8) + ϵ * ũ(dr(x)*2, dr(y)*8, ℓ, k)
    vᵢ(x, y, z) = ϵ * ṽ(dr(x)*2, dr(y)*4, ℓ, k)
    cᵢ(x, y, z) = C(dr(y)*8, grid.Ly)

    set!(model, u=uᵢ, v=vᵢ, c=cᵢ)

    wizard = TimeStepWizard(cfl=0.1, max_change=1.1, max_Δt=10.0)

    c = sqrt(model.free_surface.gravitational_acceleration)
    Δt = 0.1 * Array(model.grid.Δxᶜᶠᵃ.parent)[5] / c

    simulation = Simulation(model, Δt=Δt, stop_time=stop_time)

    progress(sim) = @printf("Iter: %d, time: %s, Δt: %s, max|u|: %.3f, max|η|: %.3f \n",
                            iteration(sim), prettytime(sim), prettytime(sim.Δt),
                            maximum(abs, model.velocities.u), maximum(abs, model.free_surface.η))

    simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))

    wizard = TimeStepWizard(cfl=0.1, max_change=1.1, max_Δt=10.0)

    simulation.callbacks[:wizard]   = Callback(wizard, IterationInterval(10))

    # Output: primitive fields + computations
    u, v, w, c = merge(model.velocities, model.tracers)

    ζ_op = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᶜ, grid; computed_dependencies=(u, v))

    ζ = Field(ζ_op)
    outputs = merge(model.velocities, model.tracers, (ζ=ζ, η=model.free_surface.η))

    name = typeof(model.advection.momentum).name.wrapper
    if model.advection.momentum isa WENOVectorInvariantVel
        name = "WENOVectorInvariantVel"
    elseif model.advection.momentum isa WENOVectorInvariantVort
        name = "WENOVectorInvariantVort"
    end

    @show experiment_name = "bickley_jet_Nh_$(Nh)_$(name)"

    simulation.output_writers[:fields] =
        JLD2OutputWriter(model, outputs,
                                schedule =TimeInterval(output_time_interval),
                                filename = experiment_name,
                                overwrite_existing = true)

    @info "Running a simulation of an unstable Bickley jet with $(Nh)² degrees of freedom..."

    run!(simulation)

    return experiment_name
end

"""
    visualize_bickley_jet(experiment_name)

Visualize the Bickley jet data associated with `experiment_name`.
"""
function visualize_bickley_jet(experiment_name)

    @info "Making a fun movie about an unstable Bickley jet..."

    filepath = experiment_name * ".jld2"

    ζ_timeseries = FieldTimeSeries(filepath, "ζ", boundary_conditions=nothing, location=(Face, Face, Center))
    c_timeseries = FieldTimeSeries(filepath, "c", boundary_conditions=nothing, location=(Face, Center, Center))

    grid = c_timeseries.grid

    xζ, yζ, zζ = nodes(ζ_timeseries)
    xc, yc, zc = nodes(c_timeseries)

    anim = @animate for (i, iteration) in enumerate(c_timeseries.times)

        @info "    Plotting frame $i from iteration $iteration..."

        ζ = ζ_timeseries[i]
        c = c_timeseries[i]
        t = ζ_timeseries.times[i]

        ζi = interior(ζ)[:, :, 1]
        ci = interior(c)[:, :, 1]

        kwargs = Dict(
                      :aspectratio => 1,
                      :linewidth => 0,
                      :colorbar => :none,
                      :ticks => nothing,
                      :clims => (-1, 1),
                      :xlims => (-grid.Lx/2, grid.Lx/2),
                      :ylims => (-grid.Ly/2, grid.Ly/2)
                     )

        ζ_plot = Plots.heatmap(xζ, yζ, clamp.(ζi, -1, 1)'; color = :balance, kwargs...)
        c_plot = Plots.heatmap(xc, yc, clamp.(ci, -1, 1)'; color = :thermal, kwargs...)

        ζ_title = @sprintf("ζ at t = %.1f", t)
        c_title = @sprintf("u at t = %.1f", t)

        Plots.plot(ζ_plot, c_plot, title = [ζ_title c_title], size = (4000, 2000))
    end

    mp4(anim, experiment_name * ".mp4", fps = 8)
end

# experiment_name = run_bickley_jet(momentum_advection=VectorInvariant(), Nh=128)
# experiment_name = run_bickley_jet(momentum_advection=WENO5(zweno=true, vector_invariant=true), Nh=128)

advection_schemes = [WENO5(vector_invariant=VelocityStencil()),
                     WENO5(vector_invariant=VorticityStencil()),
                     VectorInvariant()]

for Nx in [64, 128, 256, 512, 1024]
    for advection in advection_schemes
        experiment_name = run_bickley_jet(arch=GPU(), momentum_advection=advection, Nh=Nx)
    end
end


module BayesianModels

using ExpFamilyDistributions

export Model
export BayesianParameter
export AbstractConjugateParameter
export ConjugateParameter
export HierarchicalConjugateParameter
export elbo
export getconjugateparams

"""
    BayesianParameter

Supertype for all Bayesian parameters, i.e. parameter having a prior
and a posterior.
"""
abstract type BayesianParameter end

"""
    AbstractConjugateParameter <: BayesianParameter

Bayesian parameter having the same type of prior and posterior.
"""
abstract type AbstractConjugateParameter <: BayesianParameter end

Base.show(io::IO, p::AbstractConjugateParameter) = print(io, "$(typeof(p))")

"""
    ConjugateParameter <: AbstractConjugateParameter

Standard implementation of a conjugate parameter.
"""
mutable struct ConjugateParameter{T <: ExpFamilyDistribution} <: AbstractConjugateParameter
    prior::T
    posterior::T
    stats::Function

    function ConjugateParameter(prior::ExpFamilyDistribution,
                                post::ExpFamilyDistribution)
        new{typeof(post)}(prior, post)
    end
end

"""
    getconjugateparams(model)

Return the list of the conjugate parameters of an object.
"""
function getconjugateparams(model)
    params = []
    for name in fieldnames(typeof(model))
        field = getfield(model, name)
        if typeof(field) <: ConjugateParameter
            push!(params, field)
        end
    end
    return params
end

"""
    HierarchicalConjugateParameter <: AbstractConjugateParameter
Σ
Conjugate parameter whose prior depends on another parameter.
"""
mutable struct HierarchicalConjugateParameter{T <: ExpFamilyDistribution} <: AbstractConjugateParameter
    getprior::Function
    posterior::T
    stats::Function

    function HierarchicalConjugateParameter(getprior::Function,
                                            post::ExpFamilyDistribution)
        new{typeof(post)}(getprior, post)
    end
end

function Base.getproperty(p::HierarchicalConjugateParameter, sym::Symbol)
    if sym ≡ :prior
        return p.getprior()
    else
        return getfield(p, sym)
    end
end

"""
    Model

Supertype for JuBeer models.
"""
abstract type Model end

# Imitate pytorch printing style of model.
function Base.show(io::IO, model::Model; pad = 0, padshift = 2, prefix = "")
    originalpad = pad
    modelname = "$(prefix)$(typeof(model))("
    println(io, lpad(modelname, length(modelname) + pad))
    pad += padshift
    for name in propertynames(model)
        pval = getproperty(model, name)
        if isa(pval, Model)
            pstr = "($(name)): "
            print(io, lpad(pstr, length(pstr) + pad))
            Base.show(io, pval, pad=pad)
        elseif isa(pval, BayesianParameter)
            pstr = "($(name)): $(pval)"
            println(io, lpad(pstr, length(pstr) + pad))
        elseif isa(pval, AbstractVector{<:Model})
            pstr = "($(name)): ["
            println(io, lpad(pstr, length(pstr) + pad))
            for (i, m) in enumerate(pval)
                Base.show(io, m, pad = pad + padshift, prefix = "($i) ")
            end
            pstr = "]"
            println(io, lpad(pstr, length(pstr) + pad))
        end
    end
    println(io, lpad(")", 1 + originalpad))
    return nothing
end

"""
    elbo(model)

Compute the evidence lower-bound of the model.
"""
function elbo(model::Model, X...)
    KL = sum([kldiv(param.posterior, param.prior)
              for param in getconjugateparams(model)])
    return sum(model(X...)) - KL
end

include("LogisticRegression.jl")
include("ARNormals.jl")

# Helper functions for plotting parameters of the models.
include("Plotting.jl")

end # BayesianModels


addqueen(queens::Array{Vector{Int}}, queen::Vector{Int}) = push!(copy(queens), queen)

hitsany(queen::Vector{Int}, queens::Array{Vector{Int}}) = any(map(x->hits(queen, x), queens))
hits(a::Array{Int}, b::Array{Int}) = any(a .== b) || abs(a-b)[1] == abs(a-b)[2]

function solve(x, y, n, d=Array(Vector{Int}, 0))
  if n == 0
    return d
  end
  for px = 1:x
    for py = 1:y
      if !hitsany([px, py], d)
        s = solve(x, y, n-1, addqueen(d, [px, py]))
        s != nothing && return s
      end
    end
  end
  return nothing
end