### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 5c3096b0-01b6-11eb-2c01-5ba0aac27b70
using CardioModels, PlutoUI, Plots, Distributions

# ╔═╡ 0c720f5f-a8be-4d8e-a803-461118c3b2be
md"
# DeBoer Model

The following equations manifest the model:

#### Effective pressure:
The systolic pressure is transformed to be in accordance with the sigmoid-shaped activation curve of the baroreflex. The constant $S_0$ defines the working point.

$S_{n}^{\star} = f(S_n)\hspace{5mm}\text{with} \hspace{5mm} f(S_n) = S_0 + 18\text{atan}\left(\frac{S_{n} -S_0 }{18}\right)$

#### Baroreflex on Heart Rate
The transfer function $G_a(z)$ relates the effective systolic pressure to the RR interval. It models fast vagally mediated and slower, sympathetic, baroreflex effects.

$I_{n} = G_a (z) S_{n}^{\star}$

#### Baroreflex on Peripheral Resistance
The peripheral resistance is influenced by the baroreflex mediated through the sympathetic nervous system reflected in the transfer function $G_b(z)$. $T_0$ is an empirical offset constant (around 3000 ms).

$T_{ n } = R_n C = T_{ 0 } + G_b(z)S_{n}^{\star}$

#### Properties of Myocardium & Respiration
Starling's law governs the pulse pressure through the coupling constant $\gamma$. Respiration also enters through the pulse pressure, modeled as a sine.

$P_{ n } = \gamma I_{n}  + A_{\rho}\sin(2\pi f_{\rho} t) \hspace{5mm} \text{with} \hspace{5mm} t = \sum I_n$

#### Windkessel
Diastolic pressure is derived from the two-element Windkessel model.

$D_{ n } = S_{n} \exp \left( - \frac{I_{n}}{ T_{n} } \right)$

#### Time Propagation
The next beat is calculated through the standard relation:

$S_{n+1} = P_{n} + D_{n}$
"

# ╔═╡ 2e8472f6-01b4-11eb-1f49-b35e9ee7c70b
model = DeBoerModel();

# ╔═╡ 373eadc4-01b6-11eb-0ef3-73d0e4518a24
md"""
N:
$(@bind N NumberField(1:1000, default = 100))
"""

# ╔═╡ 232d91d8-01b6-11eb-08d6-3343be95b77f
S, D, I, T, ρ = predict!(model, N, burnIn = 100);

# ╔═╡ b19bf544-b153-41a0-a013-992da6cdd0cd
md"""
Show Spectrum:
$(@bind psd CheckBox())
"""

# ╔═╡ 579f32aa-01b6-11eb-3a79-8f2207aeb03c
begin
	if psd
		psdplot(S)
	else
		s = plot(S, lab = "", ylabel = "mmHg", color = :red)
		plot!(D, lab = "", title = "BP", ylabel = "mmHg", color = :green)
		t = plot(T, lab = "", title = "T", ylabel = "a.u.", color = :orange)
		i = plot(I, lab = "", title = "I", ylabel = "ms", color = :blue)
		r = plot(ρ, lab = "", title = "ρ", ylabel = "a.u.", xlab = "beats", color = :black)
		plot(s,t,i,r, layout = (4,1), dpi = 500)
	end
end

# ╔═╡ ac3a7eac-6bba-11eb-37ec-49854b22053f
md"
# Simulating a phenylephrine injection

"

# ╔═╡ 53b158a2-6bbb-11eb-0fa7-87390c33cb64
begin
	model_phe = DeBoerModel(a0 = 9, A = 1,  noiseI = Normal(0, 25), noiseS = Normal(0, 2))
	# pre-injection
	Npre = 50
	S1,D1,I1,T1,R1 = predict!(model_phe, Npre, burnIn = 100)

	# injection
	duration = 15
	S2,D2,I2,T2,R2 = phenylephrine(model_phe, 1000, duration)

	# post-injection
	Npost = 10 # length of plateau
	S3,D3,I3,T3,R3 = predict!(model_phe, Npost, burnIn = 0)

	# back to default
	S4,D4,I4,T4,R4 = phenylephrine(model_phe, -1000, 10)

	plot([S1; S2; S3; S4], lab = "", title = "BP", ylab = "mmHg", color = :red)
	pBP = plot!([D1; D2; D3; D4], lab = "", color = :green)
	pIBI = plot([I1; I2; I3; I4], lab = "", ylab = "ms", title = "I", color = :blue)
	plot(pBP, pIBI, layout = (2,1))
	vline!([Npre Npre], color = :black, lab = "", dpi = 500)

end

# ╔═╡ b312bc92-6bbf-11eb-131e-c946cf3eb23f
begin
	xbp = 0:30
	yms = 9 .* xbp
	scatter(S2 .- minimum(S2), I2 .- minimum(I2), lab = "", title = "Simulated response")
	plot!(xbp, yms, color = :black, lab = "BRS 9 ms/mmHg", ylab = "\\Delta I [ms]", xlab = "\\Delta S [mmHg]", dpi = 500)

end

# ╔═╡ Cell order:
# ╠═5c3096b0-01b6-11eb-2c01-5ba0aac27b70
# ╟─0c720f5f-a8be-4d8e-a803-461118c3b2be
# ╠═2e8472f6-01b4-11eb-1f49-b35e9ee7c70b
# ╟─373eadc4-01b6-11eb-0ef3-73d0e4518a24
# ╠═232d91d8-01b6-11eb-08d6-3343be95b77f
# ╟─b19bf544-b153-41a0-a013-992da6cdd0cd
# ╟─579f32aa-01b6-11eb-3a79-8f2207aeb03c
# ╟─ac3a7eac-6bba-11eb-37ec-49854b22053f
# ╠═53b158a2-6bbb-11eb-0fa7-87390c33cb64
# ╠═b312bc92-6bbf-11eb-131e-c946cf3eb23f


import Distributions
using LoopVectorization
using Random

export MvNormalFast
"""
    MvNormalFast
This is a specialized MvNormal that assumes a **diagonal** covariance. This is substantially
faster than the usual Distributions.jl version that for some reason has large overhead.

# Notes
Only logpdf is defined so no calling rand!
"""
struct MvNormalFast{T<:AbstractVector, S<:AbstractVector, N} <: Distributions.AbstractMvNormal
    μ::T
    Σ::S
    lnorm::N
    function MvNormalFast(μ::T, Σ::S) where {T,S}
        lnorm = -0.5*sum(log.(Σ)) - length(μ)/2*log(2π)
        return new{T,S,typeof(lnorm)}(μ, Σ, lnorm)
    end
end

Base.length(d::MvNormalFast) = length(d.μ)
function Distributions.sampler(d::MvNormalFast)
    return sampler(Distributions.MvNormal(d.μ, sqrt.(d.Σ)))
end

Base.eltype(::MvNormalFast{T,N}) where {T,N} = eltype(T)

@inline function Distributions._logpdf(d::MvNormalFast, x::AbstractVector{T}) where {T}
    μ,Σ = d.μ, d.Σ
    acc = zero(eltype(x))
    @inbounds @fastmath for i in eachindex(x)
        acc += -(x[i]-μ[i])^2/Σ[i]/2
    end
    return acc + d.lnorm
end

function Distributions._rand!(rng::Random.AbstractRNG, d::MvNormalFast, x::AbstractArray)
    return Distributions._rand!(rng, Distributions.MvNormal(d.μ, sqrt.(d.Σ)), x)
end

# These are some soss stuff to make Dirichlet and other stuff work.
#Soss.xform(d::Distributions.Dirichlet, _data::NamedTuple=NamedTuple()) = Soss.TV.UnitSimplex(length(d.alpha))
#Soss.xform(d::Distributions.LogNormal, _data::NamedTuple=NamedTuple()) = Soss.TV.as_positive_real
#Soss.TV.as(d::Distributions.LogNormal) = Soss.TV.as_positive_real


export ShamanskiiDir

function ShamanskiiDir(nlp            :: AbstractNLPModel,
                       nlp_at_x       :: NLPAtX,
                       x              :: AbstractVector,
                       Nwtdirection   :: Function)

    dₙ = Nwtdirection(nlp_at_x.Hx, nlp_at_x.gx, verbose = false)

    xₚ = x + dₙ

    dS = Nwtdirection(nlp_at_x.Hx, grad(nlp, xₚ), verbose = false)

    return xₚ, dS
end


reset_default_graph(); sess = Session()
function test_flow(sess, flows, dim)
    prior = ADCME.MultivariateNormalDiag(loc=zeros(dim))
    model = NormalizingFlowModel(prior, flows)
    x = rand(3, dim)
    z, _, J = model(x)
    xs, iJ = ADCME.backward(model, z[end])
    init(sess)
    x_, J_, iJ_ = run(sess, [xs[end], J, iJ])
    x, x_, J_, iJ_
end

@testset "LinearFlow" begin
    dim = 3
    flows = [LinearFlow(dim)]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "AffineConstantFlow" begin
    dim = 3
    flows = [AffineConstantFlow(dim)]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "ActNorm" begin
    dim = 3
    flows = [ActNorm(dim)]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "SlowMAF" begin
    dim = 3
    flows = [SlowMAF(dim, false, [x->fc(x, [20,20,2], "fc1"); x->fc(x,[20,20,2],"fc2")])]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "MAF" begin
    dim = 3
    flows = [MAF(dim, false, [20,20,20])]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "IAF" begin
    dim = 3
    flows = [IAF(dim, false, [20,20,20])]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "Invertible1x1Conv" begin
    dim = 3
    flows = [Invertible1x1Conv(dim)]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "AffineHalfFlow" begin
    dim = 3
    flows = [AffineHalfFlow(dim, true)]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "NeuralCouplingFlow" begin
    @test_skip begin
        dim = 6
        K = 8
        n = dim÷2
        flows = [NeuralCouplingFlow(dim, x->fc(x, [20,20,n*(3K-1)], "ncf1"), x->fc(x,[20,20,n*(3K-1)],"ncf2"), K)]
        x1, x2, j1, j2 = test_flow(sess, flows, dim)
        @test x1≈x2
        @test j1≈-j2
    end
end

@testset "Permute" begin
    dim = 6
    flows = [Permute(dim, randperm(6) .- 1)]
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

@testset "composite" begin
    function create_linear_transform(dim)
        permutation = randperm(dim) .- 1
        flow = [Permute(dim, permutation);LinearFlow(dim)]
    end
    dim = 5
    flows = create_linear_transform(dim)
    x1, x2, j1, j2 = test_flow(sess, flows, dim)
    @test x1≈x2
    @test j1≈-j2
end

# @testset "composite 2" begin
#     function create_base_transform(dim, K=8)
#         n1 = dim÷2
#         n2 = dim - n1
#         r1 = x->Resnet1D(n1 * (3K-1), 256, dropout_probability=0.0, use_batch_norm=false)(x)
#         r2 = x->Resnet1D(n2 * (3K-1), 256, dropout_probability=0.0, use_batch_norm=false)(x)
#         NeuralCouplingFlow(dim, r1, r2)
#     end
#     dim = 6
#     flows = [create_base_transform(dim)]
#     x1, x2, j1, j2 = test_flow(sess, flows, dim)
#     @test x1≈x2
#     @test j1≈-j2
# end


module Gabor
using DSP

function pad_zeros(v, n)
    d = max(n - length(v),0)
    pre,post = fld(d,2),cld(d,2)
    [zeros(pre);v;zeros(post)]
end

function gabor{T}(t, f::Vector{T}, σ = 1.0, tol = 1e-5,
                  wf = n -> hanning(n))
    Δt = sqrt(-2σ^2*log(tol))
    Δi = round(Int, 2Δt/(t[2]-t[1]))

    w = wf(2Δi+1)

    nt = length(t)
    G = zeros(Complex{real(T)}, (2Δi+1, nt))
    for i = 1:nt
        sel = max(1,i-Δi):min(nt,i+Δi)
        G[:,i] = fftshift(fft(pad_zeros(exp(-(t[sel]-t[i]).^2/2(σ^2)).*f[sel], 2Δi+1).*w))
    end
    fftshift(collect(fftfreq(2Δi+1,1.0/(t[2]-t[1])))), G/√(2π)
end

export gabor

end


function solve_simplex(input::Input)

    midterm = MidTerm()
    if check_phase_1(input)
        midterm.base =  collect((input.n - input.m + 1):input.n)
        midterm.nbase = collect(1:(input.n-input.m))
        val, i = findmin(input.b)

        aux = midterm.base[i]
        push!(midterm.nbase, aux)
        midterm.base[i] = input.n + 1
        c_mem = copy(input.c)
        input.A = [input.A -ones(input.m)]
        input.c = zeros(input.n + 1)
        input.c[end] = -1
        input.n = input.n + 1

        init_log_simplex1(input)
        while midterm.termination_status == 0 && midterm.iter < input.max_iter
            iterate!(input, midterm)
        end
        update_midterm!(input, midterm)
        if midterm.z < -input.tol
            midterm.termination_status = 3 #infeasible
            output = write_output(input, midterm)
            return output
        end
        midterm.iter = 0
        input.n = input.n - 1
        cache = findfirst(x->x==input.n + 1,midterm.nbase)
        input.c = c_mem
        midterm.termination_status = 0

        if cache !== nothing
            deleteat!(midterm.nbase, cache)
        else
            l = findfirst(x->x==input.n + 1,midterm.base)
            B = view(input.A,:,midterm.base)
            for (k,w) in enumerate(midterm.nbase)
                sol = B \ input.A[:,w]
                if abs.(sol[l]) > input.tol
                    push!(midterm.base, w)
                    deleteat!(midterm.nbase, k)
                    deleteat!(midterm.base, l)
                    break
                end
                if k == length(midterm.nbase)
                    input.A = input.A[1:end .!= l, :]
                    input.b = input.b[1:end .!= l]
                    deleteat!(midterm.base, l)
                    input.m -= 1
                    break
                end
            end
        end
        input.A = input.A[:,1:input.n]
    else
        midterm.base =  collect((input.n - input.m + 1):input.n)
        midterm.nbase = collect(1:(input.n-input.m))
    end

    init_log_simplex2(input)
    while midterm.termination_status == 0 && midterm.iter < input.max_iter
        iterate!(input, midterm)
    end
    update_midterm!(input, midterm)
    output = write_output(input, midterm)
    return output
end

function check_phase_1(input::Input)
    b_min = minimum(input.b)
    return b_min < 0
end

function iterate!(input::Input, midterm::MidTerm)

    midterm.iter += 1
    A = input.A
    b = input.b
    c = input.c
    base = midterm.base
    nbase = midterm.nbase
    tol = input.tol
    B = view(A,:,base)
    N = view(A,:,nbase)
    xB = B \ b
    y = B' \ c[base]
    midterm.red_cost = c[nbase] - N'*y
    val = maximum(midterm.red_cost)

    if val <= tol
        midterm.termination_status = 1
        return #optimal
    end
    midterm.j = findfirst(x->x>tol,midterm.red_cost)

    d = zeros(length(c))
    d_base = B \ N[:,midterm.j]
    d[base] = - d_base
    d[nbase[midterm.j]] = 1
    midterm.d = d
    d_base = max.(d_base, 0)
    r = max.(xB, tol) ./ d_base
    val, midterm.i = findmin(r)

    if val == Inf
        midterm.termination_status = 2
        return #unbounded
    end
    midterm.z = c[base]'xB
    midterm.x = zeros(input.n)
    midterm.x[base] = xB
    iteration_log(input, midterm)
    base[midterm.i], nbase[midterm.j] = nbase[midterm.j], base[midterm.i]
    return #max iteration
end


using Stheno, Plots, Distributions, ColorTypes, Random, FixedPointNumbers
using Stheno: @model, Titsias

rng = MersenneTwister(12345);

###########################  Define and inspect our model  ###########################

# Specify generative model.
@model function gp(σ²)
    f = 1.5 * GP(EQ())
    y = f + GP(Noise(σ²))
    return f, y
end

# Construct generative model.
σ² = 1.0;
f, y = gp(σ²);

# Sample prior at random locations.
Xy = rand(Uniform(-5, 5), 100);
y_ = rand(y(Xy));

# Approximately condition processes. Play around with different lengths of Z.
Z = range(-5, stop=5, length=10);
f′ = f | Titsias(f(Xy)←y_, f(Z), sqrt(σ²));
f′_exact = f | (y(Xy)←y_);

#######################  Do posterior inference give a few samples  #######################

Np, S = 500, 10;
Xp = range(-6; stop=6, length=Np);
f′Xp = rand(rng, f′(Xp), S);

# Get posterior marginals for approximate and exact conditioning.
μ′f, σ′f = marginals(f′(Xp));
μf′_exact, σf′_exact = marginals(f′_exact(Xp));

####################################  Plot results  ####################################

plotly();

posterior_plot = plot(
    legend=:topright,
    legendfont=Plots.Font(
        "sans-serif",
        10,
        :hcenter,
        :vcenter,
        0.0,
        RGB{Normed{UInt8, 8}}(0.0,0.0,0.0)
    ),
    background_color_legend=RGBA(1, 1, 1, 0),
    foreground_color_legend=RGBA(1, 1, 1, 0),
);

plot!(posterior_plot, Xp, μf′_exact;
    linewidth=2.0,
    linecolor=:blue,
    label="");
plot!(posterior_plot, Xp, μf′_exact;
    linewidth=0.0,
    linecolor=:blue,
    linealpha=0.2,
    fillrange=[μf′_exact.- 3  .* σf′_exact],
    fillalpha=0.2,
    fillcolor=:blue,
    label=["exact posterior, f", ""]);
plot!(posterior_plot, Xp, μf′_exact;
    linewidth=0.0,
    linecolor=:blue,
    linealpha=0.2,
    fillrange=[μf′_exact .+ 3 .* σf′_exact],
    fillalpha=0.2,
    fillcolor=:blue,
    label="");

plot!(posterior_plot, Xp, μ′f;
    linewidth=2.0,
    linecolor=:red,
    label="");
plot!(posterior_plot, Xp, μ′f;
    linewidth=0.0,
    linecolor=:red,
    linealpha=0.2,
    fillrange=[μ′f.- 3  .* σ′f],
    fillalpha=0.2,
    fillcolor=:red,
    label=["posterior, f", ""]);
plot!(posterior_plot, Xp, μ′f;
    linewidth=0.0,
    linecolor=:red,
    linealpha=0.2,
    fillrange=[μ′f .+ 3 .* σ′f],
    fillalpha=0.2,
    fillcolor=:red,
    label="");

scatter!(posterior_plot, Xy, y_;
    markercolor=:red,
    markershape=:circle,
    markerstrokewidth=0.0,
    markersize=6,
    markeralpha=0.3,
    label="y_");

display(posterior_plot);


# Checking the sampler tables

module TablesTest  # separate namespace

using Base.Test

include("../src/randtn_zigtables.jl")

# basic properties
@test all(tnzig_yl .<= tnzig_yu)
@test all(diff(tnzig_x) .== tnzig_d)

# ncell index into x
n = length(tnzig_x)
xsimn = 10^6
for i = 1:xsimn
    x = tnzig_xmin + (tnzig_xmax - tnzig_xmin) * rand()
    i = tnzig_i0 + ifloor(x * tnzig_invh)
    # array boundaries
    @test i >= 1
    @test i <= length(tnzig_ncell)
    j = tnzig_ncell[i]
    @test j >= 1
    @test j <= n - 1
    # x[j] <= x < x[j+2] (or x[j+1] for j = n-1)
    @test tnzig_x[j] <= x
    @test tnzig_x[min(n, j+2)] > x
end

end


#
# CR3BP Plots
#

"""
Plot specified lagrange points in the rotating
reference frame of CR3BP system μ.
"""
function lagrangeplot(μ, L=1:5; kwargs...)

    defaults = (; title="Nondimensional Lagrange Points",
                  xlabel="X (DU)", ylabel="Y (DU)",
                  labels=["Body 1" "Body 2" [string("L",i) for i ∈ L]...])
    options  = merge(defaults, kwargs)

    lagrange_points = lagrange(μ)

    fig = scatter([-μ], [0]; markersize=10, markercolor=:lightblue, label="Body 1")
    scatter!(fig, [1-μ], [0]; markersize=6, markercolor=:gray, label="Body 2")

    colors = (:red, :orange, :tan, :cyan, :indigo)
    for point ∈ zip(lagrange_points, 1:length(lagrange_points))
        p, i = point
        scatter!(fig, [p[1]], [p[2]];
                markershape=:x, markercolor=colors[i], label=string("L", i))
    end

    scatter!(fig; options...)
    for i ∈ 1:min(length(fig.series_list), length(options.labels))
        fig.series_list[i].plotattributes[:label] = options.labels[i]
    end

    fig

end


const rng = MersenneTwister(1234)

struct Data
    X::Array{Array{<:AbstractFloat,1},1}
    Y::Array{Array{<:AbstractFloat,1},1}
    # add check here
end

struct DiffableParameters
    σ::AbstractArray # measurement noise per output
    G::AbstractArray# all smoothing kenrel paramters
    u::AbstractArray # u kernel paramters
end

mutable struct GaussianProcess
    base_kernel::Function
    D::Int # number of outputs
    C::Int # terms in Volterra series
    P::Int # Number of paramters in each smoothing kernel

    data::Union{Data,Missing} # data to fit

    dpars::Union{DiffableParameters,Missing} # all the differentiable hyperparamters
end

GaussianProcess(base_kernel, D, C, P ) = GaussianProcess(base_kernel, D, C, P, missing, init_dpars(D, C, P))
GaussianProcess(base_kernel, D, C, P, data ) = GaussianProcess(base_kernel, D, C, P, data, init_dpars(D, C, P))

"""
initialise the differentiable paramters
"""
function init_dpars(D::Int, C::Int, P::Int)::DiffableParameters
    G = 0.1 * ones(Float64, (D, sum(1:C), P))
    DiffableParameters(fill(0.3, D), G, [.01])
end

"""
get K for one output
"""
function fill_sub_K(t::Array{<:AbstractFloat}, tp::Array{<:AbstractFloat}, d::Int, dp::Int, gp::GaussianProcess)
    K  = reduce(hcat,
        map.((tpi ->
            map.(ti ->
                kernel(ti, tpi, d, dp, gp),
            t)),
        tp))
    # println(d, dp, " ", minimum(eigvals(K)))
    # Matrix(Hermitian(K))
    K
end

"""
get K for all outputs
"""
function fill_K(t, tp, gp::GaussianProcess)
    reduce(hcat,
        map.(dpi ->
            reduce(vcat,
                map.(di ->
                    fill_sub_K(t[di], tp[dpi], di, dpi, gp),
                 1:gp.D)),
        1:gp.D))
end

"""
get mean for one output
"""
function fill_sub_μ(t::Array{<:AbstractFloat}, d::Int, gp::GaussianProcess)
	map.(ti -> full_E(ti, d, gp), t)
end

"""
get mean for all outputs
"""
function fill_μ(t, gp::GaussianProcess)
    reduce(vcat, map.(di -> fill_sub_μ(t[di], di, gp), 1:gp.D))
end

function posterior(t, gp::GaussianProcess; jitter=1e-5)
    Σ = Diagonal(vcat([gp.dpars.σ[i]^2 * ones(size(gp.data.X[i])[1]) for i in 1:gp.D]...))

    Koo = fill_K(gp.data.X, gp.data.X, gp) + Σ + jitter * I

    if !ishermitian(Koo)
        ∇ = maximum(Koo' - Koo)
        if ∇ > sqrt(eps())
            print("WARNING, Hermitian check faliure not rounding error! ", ∇)
        end
        Koo = Matrix(Hermitian(Koo))
    end

    Kop = fill_K(gp.data.X, t, gp)
    Kpp = fill_K(t, t, gp)

    μo = fill_μ(gp.data.X, gp)
    μp = fill_μ(t, gp)

    y = vcat(gp.data.Y...)

    Loo = cholesky(Koo).L

    μ_post = μp + Kop' * (Loo' \ (Loo \ (y - μo)))

	Lop = Loo \ Kop
	K_post = Kpp - Lop' * Lop
	K_post = Matrix(Hermitian(K_post))
	(μ_post, K_post)
end

function negloglikelihood(gp::GaussianProcess)

    Σ = Diagonal(vcat([gp.dpars.σ[i]^2 * ones(size(gp.data.X[i])[1]) for i in 1:gp.D]...))
    K = fill_K(gp.data.X, gp.data.X, gp) + Σ + 1e-5 * I

    if !ishermitian(K)
        ∇ = maximum(K' - K)
        if ∇ > sqrt(eps())
            print("WARNING, Hermitian check faliure not rounding error! ", ∇)
        end
        K = Matrix(Hermitian(K))
    end

    μ = fill_μ(gp.data.X, gp)

    y = vcat(gp.data.Y...)
    dist = MvNormal(μ, K)
    -1 * logpdf(dist, y)
end

function negloglikelihood(σ::AbstractArray, G::AbstractArray, u::AbstractArray, gp::GaussianProcess)
    gpn = copy(gp)
    gpn.dpars = DiffableParameters(σ, G, u)
    negloglikelihood(gpn)
end


#=
  https://www.youtube.com/watch?v=e1Ykk_CqKTY&t=458s

  Probabilistic Programming: What It Is and How It Works - Noel Welsh

  We can see either 1, 2, or 3 cats.
  There are 3 different enticements:

  - Milkshake
  - Fish
  - Nothing

  And there are different probabilities how many cats there are given
  an enticement, see below.

  Now: We see 3 cats, what is the probability that it's a milkshake?

  The video got the following (for 3 cats):
   - milkshake: 0.42
   - fish: 0.04
   - nothing: 0.03

  Normalized to percentage (from the video):

  0.42/(0.42 + 0.04 + 0.03) milkshake
      0.85714285714285714286
  0.04/(0.42 + 0.04 + 0.03) fish
      0.081632653061224489796
  0.03/(0.42 + 0.04 + 0.03)  nothing
      0.061224489795918367347

  Here are two models.
  Model2 seems to be the correct one (i.e. corresponds to the video),
  i.e. where the enticements are mutual exclusive.
  In model1 we accept that fish and milkshake can coexist.

  See ~/cplint/cat.pl
      ~/blog/cat.blog
      ~/psi/cat.psi
=#

using Turing # , StatsPlots, DataFrames
include("jl_utils.jl")
# using ReverseDiff, Zygote, Tracker
# Turing.setadbackend(:reversediff)
# Turing.setadbackend(:zygote)
# Turing.setadbackend(:tracker)

# This is not correct
@model function cat1()
    NoCat = 0
    OneCat = 1
    TwoCats = 2
    ThreeCats = 3

    # probability of an enticement in the garden
    milkshake      ~ Bernoulli(0.6)
    fish           ~ Bernoulli(0.1)
    no_enticement  ~ Bernoulli(0.3)

    # Number of cats per enticement
    vs = [NoCat,OneCat,TwoCats,ThreeCats]
    vs_cat ~ DiscreteUniform(NoCat,ThreeCats)
    if !no_enticement && milkshake
        cat ~ Categorical(simplex([0.0,0.1,0.2,0.7]))
    elseif !no_enticement && fish
        cat ~ Categorical(simplex([0.0,0.2,0.3,0.4]))
    elseif no_enticement && !fish && !milkshake
        cat ~ Categorical(simplex([0.0,0.6,0.3,0.1]))
    else
        cat ~ Categorical(simplex([1.0,0.0,0.0,0.0]))
    end
    true ~ Dirac(vs_cat == vs[cat])
    true ~ Dirac(vs_cat == ThreeCats)

end

#=
println("Model 1:")
model = cat1()
num_chns = 4
# chns = sample(model, Prior(), MCMCThreads(), 1000, num_chns)
chns = sample(model, MH(), MCMCThreads(), 10_000, num_chns)
# chns = sample(model, MH(), 1000)
display(chns)
=#

# This is correct (according to the video)
#=
   Summary Statistics
            parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
                Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

            enticement    1.2108    0.5429     0.0014    0.0093   2253.8077    1.0013       70.4381
  enticement_milkshake    0.8532    0.3539     0.0009    0.0061   2150.1791    1.0013       67.1994
       enticement_fish    0.0824    0.2750     0.0007    0.0046   2344.6207    1.0014       73.2763
    enticement_nothing    0.0642    0.2452     0.0006    0.0041   2630.1451    1.0013       82.1997
                  cats    2.9992    0.0403     0.0001    0.0008   2814.4694    1.0013       87.9604

   Distributions of variable enticement
   milkshake  =>  136546  (0.853413)
   fish       =>   13181  (0.082381)
   no_enticement =>   10273  (0.064206)

=#
enticements = []
enticements_dict = Dict(1=>"milkshake",2=>"fish",3=>"no_enticement")
@model cat2() = begin

    milkshake = 1
    fish = 2
    no_enticement = 3

    # Priors
    # probability of an enticement in the garden
    enticement ~ Categorical(simplex([0.6,0.1,0.3]))
    # Note: This seems a little too complex...
    enticement_milkshake ~ Bernoulli(0.5)
    enticement_fish      ~ Bernoulli(0.5)
    enticement_nothing   ~ Bernoulli(0.5)
    cats ~ DiscreteUniform(1,3)
    if enticement == milkshake
        cats ~ Categorical(simplex([0.1,0.2,0.7]))
        enticement_milkshake ~ Bernoulli(1.0)
        enticement_fish      ~ Bernoulli(0.0)
        enticement_nothing   ~ Bernoulli(0.0)
    elseif enticement == fish
        cats ~ Categorical(simplex([0.2,0.4,0.4]))
        enticement_milkshake ~ Bernoulli(0.0)
        enticement_fish      ~ Bernoulli(1.0)
        enticement_nothing   ~ Bernoulli(0.0)
    else
        cats ~ Categorical(simplex([0.6,0.3,0.1]))
        enticement_milkshake ~ Bernoulli(0.0)
        enticement_fish      ~ Bernoulli(0.0)
        enticement_nothing   ~ Bernoulli(1.0)
    end

    # We observe 3 cats
    true ~ Dirac(cats == 3)

    return enticement
end

println("\n\nModel 2:")
model = cat2()
num_chns = 4
# chns = sample(model, Prior(), MCMCThreads(), 1000, num_chns)
chns = sample(model, MH(), MCMCThreads(), 40_000, num_chns)
# chns = sample(model, MH(), 40_000)
# chns = sample(model, PG(20), MCMCThreads(), 1000, num_chns)
# chns = sample(model, SMC(), MCMCThreads(), 1000, num_chns)
# chns = sample(model, IS(), MCMCThreads(), 1000, num_chns)

display(chns)

show_var_dist_pct(chns, :enticement,["milkshake","fish","no_enticement"])


using Gadfly, Distributions

set_default_plot_size(6inch, 3inch)

vals = rand(MultivariateNormal([0.0, 0.0], [1.0 0.5; 0.5 1.0]), 10000)
Gadfly.with_theme(:dark) do
    plot(x=vals[1,:], y=vals[2,:], Geom.hexbin)
end


struct LocallyConstant <: AbstractRegression

end

"""
    Apply the analog method on catalog of historical data
    to generate forecasts.
"""
function predict(LocallyConstant, xf::Array{Float64,2}, weight)

    # weighted mean and covariance
    xf_mean = vec(sum(xf .* weights', dims = 2))
    Exf = xf .- xf_mean

    Symmetric(1.0 ./ (1.0 .- sum(weights .^ 2)) .* (Exf .* weights') * Exf')

end


x = 12
println(typeof(x))
y = convert(UInt8, x)
println(y)
println(typeof(y))

# This will flag an error because Uint8 holds up to 2^8 - 1 = 255
# convert(UInt8, 256)

struct Point{T<:Real}
    x::T
    y::T
end

Base.convert(::Type{Point{T}}, x::Array{T, 1}) where {T<:Real} = Point(x...)
println(convert(Point{Int64}, [1, 2]))

a, b = 1.5, 2
println("$(typeof(a)), $(typeof(b))")
c, d = promote(a, b)
println("$(typeof(c)), $(typeof(d))")
# To add functionality for promote, add definition for promote_rule.
# Eg. promote_rule(::Type{Float64}, ::Type{Int32}) = Float64


import Base: eye, rand, *, inv

# Type definitions

immutable OrthogonalGroup{T, N} <: Group
    α::T
    β::T
    γ::T
end

O3{T} = OrthogonalGroup{T, 3}
O3{T}(α::T, β::T, γ::T) = OrthogonalGroup{T, 3}(α, β, γ)

# Identity element

eye(::Type{O3}) = O3(0.0, 0.0, 0.0)

# Random group element

function rand(::Type{O3})
    α, β, γ = 2π*rand(), π*rand(), 2π*rand()
    O3(α, β, γ)
end

# Defining representation

@inbounds function defrep!{T}(M::Matrix{T}, g::O3{T})
    α, β, γ = g.α, g.β, g.γ
    M[1, 1] = cos(γ)⋅cos(α)⋅cos(β) - sin(γ)⋅sin(α)
    M[2, 1] = sin(α)⋅cos(γ)⋅cos(β) + sin(γ)⋅cos(α)
    M[3, 1] = -sin(β)⋅cos(γ)
    M[1, 2] = -sin(γ)⋅cos(α)⋅cos(β) - cos(γ)⋅sin(α)
    M[2, 2] = -sin(γ)⋅sin(α)⋅cos(β) + cos(γ)⋅cos(α)
    M[3, 2] = sin(β)⋅sin(γ)
    M[1, 3] = sin(β)⋅cos(α)
    M[2, 3] = sin(β)⋅sin(α)
    M[3, 3] = cos(β)

    return M
end

defrep{T}(g::O3{T}) = defrep!(Matrix{T}(3, 3), g)

# Hacky and suboptimal. Optimize this.
# Procedure from
#   http://mathworld.wolfram.com/EulerAngles.html
#   http://mathworld.wolfram.com/NonlinearLeastSquaresFitting.html

function nearest{T}(::Type{O3{T}}, M_defrep::Matrix{T}, tol = T(1e-8))
    b = vec(M_defrep)
    tol_sq = tol^2

    α, β, γ = 2π*rand(T), π*rand(T), 2π*rand(T)
    dα, dβ, dγ = T(Inf), T(Inf), T(Inf)

    while dα^2 + dβ^2 + dγ^2 > tol_sq
        A = [(-cos(γ)*sin(α)*cos(β) - sin(γ)*cos(α))    (-cos(γ)*cos(α)*sin(β))     (-sin(γ)*cos(α)*cos(β) - cos(γ)*sin(α))
             (cos(α)*cos(γ)*cos(β) - sin(γ)*sin(α))     (-sin(α)*cos(γ)*sin(β))     (-sin(α)*sin(γ)*cos(β) + cos(γ)*cos(α))
             (zero(T))                                  (-cos(β)*cos(γ))            (sin(β)*sin(γ))
             (sin(γ)*sin(α)*cos(β) - cos(γ)*cos(α))     (sin(γ)*cos(α)*sin(β))      (-cos(γ)*cos(α)*cos(β) + sin(γ)*sin(α))
             (-sin(γ)*cos(α)*cos(β) - cos(γ)*sin(α))    (sin(γ)*sin(α)*sin(β))      (-cos(γ)*sin(α)*cos(β) - sin(γ)*cos(α))
             (zero(T))                                  (cos(β)*sin(γ))             (sin(β)*cos(γ))
             (-sin(β)*sin(α))                           (cos(β)*cos(α))             (zero(T))
             (sin(β)*cos(α))                            (cos(β)*sin(α))             (zero(T))
             (zero(T))                                  (-sin(β))                   (zero(T))]

        dα, dβ, dγ = A \ b
        α = mod2pi(α + dα)
        β = mod2pi(2β + 2dβ) / 2
        γ = mod2pi(γ + dγ)
    end

    return O3(α, β, γ)
end

nearest{T}(::Type{O3}, M::Matrix{T}) = nearest(O3{T}, M)

# Multiplication
# Hacky.  Find an analytic expression for multiplying Euler-ZYZ angles
function *{T}(g₁::O3{T}, g₂::O3{T})
    nearest(O3, defrep(g₁) * defrep(g₂))
end

# Inversion
# Hacky. Find an analytic expression for inverting Euler-ZYZ angles
function inv{T}(g::O3{T})
    nearest(O3, inv(defrep(g)))
end


# This file defines the EmpiricalOrthogonalFunction types

# abstract EOF type
"""
Abstract type to build subtypes from
"""
abstract type AbstractEmpiricalOrthogonalFunction end

# eof struct for 2d data (space, time)
"""
    EmpiricalOrthogonalFunction(dataset; center=true, ddof=1)

Create an Empirical Orthogonal Function object.
The EOF solution is computed at initialization time. Method
calls are used to retrieve or update computed quantities.
"""
mutable struct EmpiricalOrthogonalFunction <: AbstractEmpiricalOrthogonalFunction
    ""
    dataset::AbstractArray{<:Any,2}

    ""
    center::AbstractArray{<:Any,2}

    valididx::AbstractArray{Int64,1}

    eofs::AbstractArray{<:Any,2}

    eigenvals::AbstractArray{<:Any,1}

    pcs::AbstractArray{<:Any,2}

    EmpiricalOrthogonalFunction(
        dataset::AbstractArray{<:Any,2},
        center::AbstractArray{<:Any,2},
        valididx::AbstractArray{Int64,1},
        eofs::AbstractArray{<:Any,2},
        eigenvals::AbstractArray{<:Any,1},
        pcs::AbstractArray{<:Any,2},
    ) = new(dataset, center, valididx, eofs, eigenvals, pcs)

    """
    """
    function EmpiricalOrthogonalFunction(
        dataset::AbstractArray{<:Any,2};
        center::Bool = true,
        ddof::Int = 1,
    )

        # Store information about the shape/size of the input data.
        shape = size(dataset)
        timelen = shape[1]

        if shape[1] > shape[2]
            @error(
                "time dimension has more records than the space dimensions and an EOF cannot be calculated, time = $(shape[1]), space= $(shape[2])"
            )
            return
        end

        # Remove the time mean of the input data unless explicitly told
        # not to by the "center" argument.
        if center
            spatial_mean = Statistics.mean(dataset, dims = 1)
            centered = @. dataset - spatial_mean
        else
            spatial_mean = zeros(shape[1])
            centered = dataset
        end

        # Find the indices of values that are not missing in one row. All the
        # rows will have missing values in the same places provided the
        # array was centered. If it wasn't then it is possible that some
        # missing values will be missed and the singular value decomposition
        # will produce not a number for everything.
        isnotmissing = findall(x -> x !== missing, spatial_mean[1, :])

        # Compute the singular value decomposition of the design matrix.
        U, S, V = svd(centered[:, isnotmissing])

        # Singular values are the square-root of the eigenvalues of the
        # covariance matrix. Construct the eigenvalues appropriately and
        # normalize by N-ddof where N is the number of observations. This
        # corresponds to the eigenvalues of the normalized covariance matrix.
        norm = timelen - 1
        eigenvals = @. S * S / norm

        eofs = Array{Union{Missing,Float64}}(missing, shape[2], shape[1])
        eofs[isnotmissing, :] = V[:, :]

        # Remove the scaling on the principal component time-series that is
        # implicitily introduced by using SVD instead of eigen-decomposition.
        # The PCs may be re-scaled later if required.
        pcs = similar(U)
        for i = 1:length(S)
            pcs[:, i] = U[:, i] * S[i]
        end

        EmpiricalOrthogonalFunction(
            dataset,
            spatial_mean,
            isnotmissing,
            eofs,
            eigenvals,
            pcs,
        )
    end

    """

    """
    function EmpiricalOrthogonalFunction(
        dataset::AbstractArray{<:Any,3};
        timedim::Int = 3,
        kwargs...,
    )
        geodims = [x for x in 1:3 if x != timedim]
        shape = size(dataset)
        spatial_shape = shape[geodims]
        time_shape = shape[timedim]

        if timedim != 1
            dataset = permutedims(dataset, (timedim, geodims...))
        end

        flattened = reshape(dataset, (time_shape, prod(spatial_shape)))

        EmpiricalOrthogonalFunction(flattened; kwargs...)
    end

end


module Kronecker

export GeneralizedKroneckerProduct, AbstractKroneckerProduct
export AbstractKroneckerSum, KroneckerSum
export KroneckerPower
export issquare, order, issymmetric, isposdef, getmatrices, collect!
export ⊗, kronecker, ⊕, kroneckersum
export CholeskyKronecker

export isprob, naivesample, fastsample, sampleindices

using LinearAlgebra
import LinearAlgebra: mul!, lmul!, rmul!, pinv, ldiv!
import Base: collect, *, getindex, size, eltype, inv, adjoint
using SparseArrays
using LinearAlgebra: checksquare

include("base.jl")
include("kroneckerpowers.jl")
include("kroneckersum.jl")
include("multiply.jl")
include("vectrick.jl")
include("indexedkroncker.jl")
include("eigen.jl")
include("factorization.jl")
include("kroneckergraphs.jl")
include("names.jl")

end # module


using Test

using FastFourierTransform

@testset "FFT Recursive Naive" begin
    for i in 1:5
        n = 2 ^ rand(0:10)
        t = complex(rand(n))
        f_direct = dft(t)
        @test idft(f_direct) ≈ t

        f_naive = fft_recursive_naive(t)
        f_recursive = fft_recursive(t)
        f_iterative = fft_iterative(t)
        @test f_naive ≈ f_direct
        @test f_recursive ≈ f_direct
        @test f_iterative ≈ f_direct

        t_naive = ifft_recursive_naive(f_naive)
        t_recursive = ifft_recursive(f_recursive)
        t_iterative = ifft_iterative(f_iterative)
        @test t_naive ≈ t
        @test t_recursive ≈ t
        @test t_iterative ≈ t

        # Polynomial multiplication
        p1 = rand(n)
        p2 = rand(n)
        x = rand(n)
        prod = @. eval_poly($Ref(p1), x) * eval_poly($Ref(p2), x)
        p = poly_mul(p1, p2)
        prod_fft = eval_poly.(Ref(p), x)
        @test prod ≈ prod_fft
    end
end


using Optim

using SteinDistributions: SteinDistribution, gradlogdensity, logdensity, numdimensions

function getapproxmode(target::SteinDistribution;
                       x0::Array{Float64,1}=zeros(numdimensions(target)))
    function objective(x::Vector)
        -logdensity(target, x)
    end
    function gradient!(x::Vector, gradientres::Vector)
        gradient = gradlogdensity(target, x)
        for ii in 1:length(gradient)
            gradientres[ii] = -gradient[ii]
        end
    end
    res = optimize(objective,
                   gradient!,
                   x0,
                   method = LBFGS())
    res.minimizer
end


module ReactiveMPModelsGMMTest

using Test, InteractiveUtils
using Rocket, ReactiveMP, GraphPPL, Distributions
using BenchmarkTools, Random, Plots, Dates, LinearAlgebra
using StableRNGs

## Model definition
## -------------------------------------------- ##
@model [ default_factorisation = MeanField() ] function univariate_gaussian_mixture_model(n)

    s ~ Beta(1.0, 1.0)

    m1 ~ NormalMeanVariance(-2.0, 1e3)
    w1 ~ GammaShapeRate(0.01, 0.01)

    m2 ~ NormalMeanVariance(2.0, 1e3)
    w2 ~ GammaShapeRate(0.01, 0.01)

    z = randomvar(n)
    y = datavar(Float64, n)

    for i in 1:n
        z[i] ~ Bernoulli(s)
        y[i] ~ NormalMixture(z[i], (m1, m2), (w1, w2))
    end

    return s, m1, w1, m2, w2, z, y
end

@model function multivariate_gaussian_mixture_model(rng, L, nmixtures, n)

    z = randomvar(n)
    m = randomvar(nmixtures)
    w = randomvar(nmixtures)

    basis_v = [ 1.0, 0.0 ]

    for i in 1:nmixtures
        # Assume we now only approximate location of cluters's mean
        approximate_angle_prior = ((2π + + rand(rng)) / nmixtures) * (i - 1)
        approximate_basis_v  = L / 2 * (basis_v .+ rand(rng, 2))
        approximate_rotation = [ cos(approximate_angle_prior) -sin(approximate_angle_prior); sin(approximate_angle_prior) cos(approximate_angle_prior) ]
        mean_mean_prior     = approximate_rotation * approximate_basis_v
        mean_mean_cov       = [ 1e6 0.0; 0.0 1e6 ]

        m[i] ~ MvNormalMeanCovariance(mean_mean_prior, mean_mean_cov)
        w[i] ~ Wishart(3, [ 1e2 0.0; 0.0 1e2 ])
    end

    s ~ Dirichlet(ones(nmixtures))

    y = datavar(Vector{Float64}, n)

    means = tuple(m...)
    precs = tuple(w...)

    for i in 1:n
        z[i] ~ Categorical(s) where { q = MeanField() }
        y[i] ~ NormalMixture(z[i], means, precs) where { q = MeanField() }
    end

    return s, z, m, w, y
end
## -------------------------------------------- ##
## Inference definition
## -------------------------------------------- ##
function inference_univariate(data, n_its)
    n = length(data)
    model, (s, m1, w1, m2, w2, z, y) = univariate_gaussian_mixture_model(n);

    mswitch = keep(Marginal)
    mm1 = keep(Marginal)
    mm2 = keep(Marginal)
    mw1 = keep(Marginal)
    mw2 = keep(Marginal)

    fe = keep(Float64)

    m1sub = subscribe!(getmarginal(m1), mm1)
    m2sub = subscribe!(getmarginal(m2), mm2)
    w1sub = subscribe!(getmarginal(w1), mw1)
    w2sub = subscribe!(getmarginal(w2), mw2)
    switchsub = subscribe!(getmarginal(s), mswitch)
    fesub = subscribe!(score(Float64, BetheFreeEnergy(), model), fe)

    setmarginal!(s, vague(Beta))
    setmarginal!(m1, NormalMeanVariance(-2.0, 1e3))
    setmarginal!(m2, NormalMeanVariance(2.0, 1e3))
    setmarginal!(w1, vague(GammaShapeRate))
    setmarginal!(w2, vague(GammaShapeRate))

    for i in 1:n_its
        update!(y, data)
    end

    unsubscribe!((fesub, switchsub, m1sub, m2sub, w1sub, w2sub))

    return mswitch, mm1, mm2, mw1, mw2, fe
end

function inference_multivariate(rng, L, nmixtures, data, viters)
    n = length(data)

    model, (s, z, m, w, y) = multivariate_gaussian_mixture_model(rng, L, nmixtures, n)

    means_estimates  = keep(Vector{Marginal})
    precs_estimates  = keep(Vector{Marginal})
    switch_estimates = keep(Marginal)
    fe_values        = keep(Float64)

    s_sub = subscribe!(getmarginal(s), switch_estimates)
    m_sub = subscribe!(getmarginals(m), means_estimates)
    p_sub = subscribe!(getmarginals(w), precs_estimates)
    f_sub = subscribe!(score(Float64, BetheFreeEnergy(), model), fe_values)

    setmarginal!(s, vague(Dirichlet, nmixtures))

    basis_v = [ 1.0, 0.0 ]

    for i in 1:nmixtures
        # Assume we now only approximate location of cluters's mean
        approximate_angle_prior = ((2π + + rand(rng)) / nmixtures) * (i - 1)
        approximate_basis_v  = L / 2 * (basis_v .+ rand(rng, 2))
        approximate_rotation = [ cos(approximate_angle_prior) -sin(approximate_angle_prior); sin(approximate_angle_prior) cos(approximate_angle_prior) ]
        mean_mean_prior     = approximate_rotation * approximate_basis_v
        mean_mean_cov       = [ 1e6 0.0; 0.0 1e6 ]

        setmarginal!(m[i], MvNormalMeanCovariance(mean_mean_prior, mean_mean_cov))
        setmarginal!(w[i], Wishart(3, [ 1e2 0.0; 0.0 1e2 ]))
    end

    for i in 1:viters
        update!(y, data)
    end

    unsubscribe!((s_sub, m_sub, p_sub, f_sub))

    return switch_estimates, means_estimates, precs_estimates, fe_values
end

@testset "Gaussian Mixture Model" begin

    @testset "Univariate" begin
        ## -------------------------------------------- ##
        ## Data creation
        ## -------------------------------------------- ##
        n = 150

        rng = StableRNG(12345)

        switch = [ 1/3, 2/3 ]
        z      = rand(rng, Categorical(switch), n)
        y      = Vector{Float64}(undef, n)

        μ1 = -10.0
        μ2 = 10.0
        w  = 1.777

        dists = [
            Normal(μ1, sqrt(inv(w))),
            Normal(μ2, sqrt(inv(w))),
        ]

        for i in 1:n
            y[i] = rand(rng, dists[z[i]])
        end
        ## -------------------------------------------- ##
        ## Inference execution
        mswitch, mm1, mm2, mw1, mw2, fe = inference_univariate(y, 10);
        ## -------------------------------------------- ##
        ## Test inference results
        @test length(mswitch) === 10
        @test length(mm1) === 10
        @test length(mm2) === 10
        @test length(mw1) === 10
        @test length(mw2) === 10
        @test length(fe) === 10 && all(filter(e -> abs(e) > 1e-3, diff(getvalues(fe))) .< 0)
        @test abs(last(fe) - 280.3276104) < 0.01

        ms = mean(last(mswitch))

        @test ((abs(ms - switch[1]) < 0.1) || (abs(ms - switch[2]) < 0.1))

        ems = sort([ last(mm1), last(mm2) ], by = mean)
        rms = sort([ μ1, μ2 ])

        foreach(zip(rms, ems)) do (r, e)
            @test abs(r - mean(e)) < 0.2
        end

        ews = sort([ last(mw1), last(mw2) ], by = mean)
        rws = sort([ w, w ])

        foreach(zip(rws, ews)) do (r, e)
            @test abs(r - mean(e)) < 0.2
        end
        ## -------------------------------------------- ##
        ## Form debug output
        base_output = joinpath(pwd(), "_output", "models")
        mkpath(base_output)
        timestamp        = Dates.format(now(), "dd-mm-yyyy-HH-MM")
        plot_output      = joinpath(base_output, "univariate_gmm_model_plot_$(timestamp)_v$(VERSION).png")
        benchmark_output = joinpath(base_output, "univariate_gmm_model_benchmark_$(timestamp)_v$(VERSION).txt")
        ## -------------------------------------------- ##
        ## Create output plots
        dim(d) = (a) -> map(r -> r[d], a)
        mp = plot(mean.(mm1), ribbon = var.(mm1) .|> sqrt, label = "m1 prediction")
        mp = plot!(mean.(mm2), ribbon = var.(mm2) .|> sqrt, label = "m2 prediction")
        mp = plot!(mp, [ μ1 ], seriestype = :hline, label = "real m1")
        mp = plot!(mp, [ μ2 ], seriestype = :hline, label = "real m2")

        wp = plot(mean.(mw1), ribbon = var.(mw1) .|> sqrt, label = "w1 prediction", legend = :bottomleft, ylim = (-1, 3))
        wp = plot!(wp, [ w ], seriestype = :hline, label = "real w1")
        wp = plot!(wp, mean.(mw2), ribbon = var.(mw2) .|> sqrt, label = "w2 prediction")
        wp = plot!(wp, [ w ], seriestype = :hline, label = "real w2")

        swp = plot(mean.(mswitch), ribbon = var.(mswitch) .|> sqrt, label = "Switch prediction")

        swp = plot!(swp, [ switch[1] ], seriestype = :hline, label = "switch[1]")
        swp = plot!(swp, [ switch[2] ], seriestype = :hline, label = "switch[2]")

        fep = plot(fe[2:end], label = "Free Energy", legend = :bottomleft)

        p = plot(mp, wp, swp, fep, layout = @layout([ a b; c d ]), size = (1000, 700))
        savefig(p, plot_output)
        ## -------------------------------------------- ##
        ## Create output benchmarks
        benchmark = @benchmark inference_univariate($y, 10) seconds=15
        open(benchmark_output, "w") do io
            show(io, MIME("text/plain"), benchmark)
            versioninfo(io)
        end
        ## -------------------------------------------- ##
    end

    @testset "Multivariate" begin
        ## -------------------------------------------- ##
        ## Data creation
        ## -------------------------------------------- ##
        rng = StableRNG(43)

        L         = 50.0
        nmixtures = 3
        n_samples = 500

        probvec = ones(nmixtures)
        probvec = probvec ./ sum(probvec)

        switch = Categorical(probvec)

        gaussians = map(1:nmixtures) do index
            angle      = 2π / nmixtures * (index - 1)
            basis_v    = L * [ 1.0, 0.0 ]
            rotationm  = [ cos(angle) -sin(angle); sin(angle) cos(angle) ]
            mean       = rotationm * basis_v
            covariance = Matrix(Hermitian(rotationm * [ 10.0 0.0; 0.0 20.0 ] * transpose(rotationm)))
            return MvNormal(mean, covariance)
        end

        z = rand(rng, switch, n_samples)
        y = Vector{Vector{Float64}}(undef, n_samples)

        for i in 1:n_samples
            y[i] = rand(rng, gaussians[z[i]])
        end
        ## -------------------------------------------- ##
        ## Inference execution
        s, m, w, fe = inference_multivariate(rng, L, nmixtures, y, 25);
        ## -------------------------------------------- ##
        ## Test inference results
        @test length(s) === 25
        @test length(m) === 25
        @test length(w) === 25
        @test length(fe) === 25
        @test all(filter(e -> abs(e) > 1e-3, diff(getvalues(fe))) .< 0)
        @test abs(last(fe) - 3442.4015524445967) < 0.01

        ems = sort(mean.(last(m)), by = x -> atan(x[2] / x[1]))
        rms = sort(mean.(gaussians), by = x -> atan(x[2] / x[1]))

        foreach(zip(ems, rms)) do (estimated, real)
            @test norm(normalize(estimated) .- normalize(real)) < 0.1
        end
        ## -------------------------------------------- ##
        ## Form debug output
        base_output = joinpath(pwd(), "_output", "models")
        mkpath(base_output)
        timestamp        = Dates.format(now(), "dd-mm-yyyy-HH-MM")
        plot_output      = joinpath(base_output, "multivariate_gmm_model_plot_$(timestamp)_v$(VERSION).png")
        benchmark_output = joinpath(base_output, "multivariate_gmm_model_benchmark_$(timestamp)_v$(VERSION).txt")
        ## -------------------------------------------- ##
        ## Create output plots
        sdim(n) = (a) -> map(d -> d[n], a)

        pe = plot(xlim = (-1.5L, 1.5L), ylim = (-1.5L, 1.5L))

        rp = scatter(y |> sdim(1), y |> sdim(2))

        pe = scatter!(pe, y |> sdim(1), y |> sdim(2))

        e_means = mean.(m[end])
        e_precs = mean.(w[end])

        for (e_m, e_w) in zip(e_means, e_precs)
            gaussian = MvNormal(e_m, Matrix(Hermitian(inv(e_w))))
            pe = contour!(pe, range(-2L, 2L, step = 0.25), range(-2L, 2L, step = 0.25), (x, y) -> pdf(gaussian, [ x, y ]), levels = 7, colorbar = false)
        end

        pfe = plot(fe[2:end], label = "Free Energy")

        p = plot(rp, pe, pfe, size = (600, 600), layout = @layout([ a b; c ]))

        savefig(p, plot_output)
        ## -------------------------------------------- ##
        ## Create output benchmarks
        benchmark = @benchmark inference_multivariate(StableRNG(123), $L, $nmixtures, $y, 25) seconds=15
        open(benchmark_output, "w") do io
            show(io, MIME("text/plain"), benchmark)
            versioninfo(io)
        end
        ## -------------------------------------------- ##
    end

end

end


using Downloads: download
using OrdinaryDiffEq
using Trixi

###############################################################################
# semidiscretization of the shallow water equations with a continuous
# bottom topography function

equations = ShallowWaterEquations2D(gravity_constant=9.812, H0=2.0)

function initial_condition_stone_throw(x, t, equations::ShallowWaterEquations2D)
  # Set up polar coordinates
  inicenter = SVector(0.15, 0.15)
  x_norm = x[1] - inicenter[1]
  y_norm = x[2] - inicenter[2]
  r = sqrt(x_norm^2 + y_norm^2)

  # Calculate primitive variables
  H = equations.H0
  v1 = r < 0.6 ? 1.75 : 0.0
  v2 = r < 0.6 ? -1.75 : 0.0
  # bottom topography taken from Pond.control in [HOHQMesh](https://github.com/trixi-framework/HOHQMesh)
  x1, x2 = x
  b = (  1.5 / exp( 0.5 * ((x1 - 1.0)^2 + (x2 - 1.0)^2) )
       + 0.75 / exp( 0.5 * ((x1 + 1.0)^2 + (x2 + 1.0)^2) ) )

  return prim2cons(SVector(H, v1, v2, b), equations)
end

initial_condition = initial_condition_stone_throw

boundary_condition = Dict( :OuterCircle => boundary_condition_slip_wall )

###############################################################################
# Get the DG approximation space

volume_flux = (flux_wintermeyer_etal, flux_nonconservative_wintermeyer_etal)
solver = DGSEM(polydeg=6, surface_flux=(flux_hll, flux_nonconservative_fjordholm_etal),
               volume_integral=VolumeIntegralFluxDifferencing(volume_flux))

###############################################################################
# Get the unstructured quad mesh from a file (downloads the file if not available locally)

default_mesh_file = joinpath(@__DIR__, "mesh_outer_circle.mesh")
isfile(default_mesh_file) || download("https://gist.githubusercontent.com/andrewwinters5000/9beddd9cd00e2a0a15865129eeb24928/raw/be71e67fa48bc4e1e97f5f6cd77c3ed34c6ba9be/mesh_outer_circle.mesh",
                                      default_mesh_file)
mesh_file = default_mesh_file

mesh = UnstructuredMesh2D(mesh_file)

# Create the semi discretization object
semi = SemidiscretizationHyperbolic(mesh, equations, initial_condition, solver,
                                    boundary_conditions=boundary_condition)

###############################################################################
# ODE solvers, callbacks, etc.

tspan = (0.0, 3.0)
ode = semidiscretize(semi, tspan)

summary_callback = SummaryCallback()

analysis_interval = 100
analysis_callback = AnalysisCallback(semi, interval=analysis_interval, save_analysis=false,
                                    extra_analysis_integrals=(energy_kinetic,
                                                              energy_internal))

alive_callback = AliveCallback(analysis_interval=analysis_interval)

save_solution = SaveSolutionCallback(interval=100,
                                     save_initial_solution=true,
                                     save_final_solution=true)

callbacks = CallbackSet(summary_callback, analysis_callback, alive_callback, save_solution)

###############################################################################
# run the simulation

# use a Runge-Kutta method with automatic (error based) time step size control
sol = solve(ode, RDPK3SpFSAL49(), abstol=1.0e-8, reltol=1.0e-8,
            save_everystep=false, callback=callbacks);
summary_callback() # print the timer summary


# This implements viscous creep laws and routines to compute with them
#
# Note that various simple creep laws are defined in this file;
# more complex ones (such as DislocationCreep) are in separate files
# in the same directory
#
# In case you want to add new creep laws, have a look at how the ones
# here are implemented. Please add tests as well!

export  DislocationCreep,
        SetDislocationCreep

# Dislocation Creep ------------------------------------------------
"""
    DislocationCreep(n = 1.0NoUnits, r = 0.00.0NoUnits, A = 1.5MPa/s, E = 476.0kJ/mol, V = 6e-6m^3/mol, apparatus = "AxialCompression" )

Defines the flow law parameter of a dislocation creep law

The (isotropic) dislocation creep law, as used by experimtalists, is given by
```math
     \\dot{\\gamma} = A \\sigma_\\mathrm{d}^n f_\\mathrm{H2O}^r \\exp(-\\frac{E+PV}{RT})
```

where ``n`` is the power law exponent,
``r`` is the exponent of fugacity dependence,
``A`` is a pre-exponential factor [MPa^(n+r)] (if manually defined, n and r must be either pre-defined or substituted),
``E`` is the activation energy [kJ/mol], ``V`` is the activation volume [m^3/mol]. ``\\dot{\\gamma}`` is the ordinary strain rate [1/s],
and ``\\sigma_\\mathrm{d}`` is the differential stress which are converted into second invariants using the apparatus type that can be
either "AxialCompression", "SimpleShear" or "Invariant".
If the flow law paramters are already given as a function of second invariants, choose apparatus = "Invariant"

```julia-repl
julia> x2      =   DislocationCreep(n=3)
DislocationCreep: n=3, r=0.0, A=1.5 MPa^-3 s^-1, E=476.0 kJ mol^-1, V=6.0e-6 m^3 mol^-1, Apparatus=AxialCompression
```
"""
@with_kw_noshow mutable struct DislocationCreep <: AbstractCreepLaw
    equation::LaTeXString   =   L"\tau_{ij} = 2 \eta  \dot{\varepsilon}_{ij}" # TO BE UPDATED
    n::GeoUnit        = 1.0NoUnits         # power-law exponent
    r::GeoUnit        = 0.0NoUnits         # exponent of water-fugacity dependence
    A::GeoUnit        = 1.5MPa^(-n-r)/s    # pre-exponential factor
    E::GeoUnit        = 476.0kJ/mol        # activation energy
    V::GeoUnit        = 6e-6m^3/mol        # activation volume
    R::GeoUnit        = 8.314J/mol/K       # Universal gas constant
    Apparatus         = "AxialCompression" # type of experimental apparatus, either AxialCompression, SimpleShear or Invariant
    Comment::String   = ""                 # Some remarks you want to add about this creep law implementation
    BibTex_Reference  = ""                 # BibTeX reference
end

# Calculation routines for linear viscous rheologies
# All inputs must be non-dimensionalized (or converted to consitent units) GeoUnits
function ComputeCreepLaw_EpsII(TauII, a::DislocationCreep, p::CreepLawVariables)
    @unpack n           = a
    @unpack r           = a
    @unpack A           = a
    @unpack E           = a
    @unpack V           = a
    @unpack R           = a
    @unpack Apparatus   = a
    @unpack P           = p
    @unpack T           = p
    @unpack f           = p
    if Apparatus == "AxialCompression"
        FT = sqrt(3.0)NoUnits               # relation between differential stress recorded by apparatus and TauII
        FE = 2.0/sqrt(3.0)NoUnits            # relation between gamma recorded by apparatus and EpsII
    elseif Apparatus == "SimpleShear"
        FT = 2.0NoUnits                      # it is assumed that the flow law parameters were derived as a function of differential stress, not the shear stress. Must be modidified if it is not the case
        FE = 2.0NoUnits
    elseif Apparatus == "Invariant"
        FT = 1.0NoUnits
        FE = 1.0NoUnits
    end
    return A.val*(TauII.val*FT)^n.val*f.val^r.val*exp(-(E.val+P.val*V.val)/(R.val*T.val))/FE;
end

# EpsII .= A.*(TauII.*FT).^n.*f.^r.*exp.(-(E.+P.*V)./(R.*T))./FE; Once we have a
# All inputs must be non-dimensionalized (or converted to consitent units) GeoUnits
function ComputeCreepLaw_TauII(EpsII, a::DislocationCreep, p::CreepLawVariables)
    @unpack n           = a
    @unpack r           = a
    @unpack A           = a
    @unpack E           = a
    @unpack V           = a
    @unpack R           = a
    @unpack Apparatus   = a
    @unpack P           = p
    @unpack T           = p
    @unpack f           = p
    if Apparatus == "AxialCompression"
        FT = sqrt(3.0)               # relation between differential stress recorded by apparatus and TauII
        FE = 2.0/sqrt(3.0)           # relation between gamma recorded by apparatus and EpsII
    elseif Apparatus == "SimpleShear"
        FT = 2.0                     # it is assumed that the flow law parameters were derived as a function of differential stress, not the shear stress. Must be modidified if it is not the case
        FE = 2.0
    elseif Apparatus == "Invariant"
        FT = 1.0
        FE = 1.0
    end
    return A.val^(-1/n.val)*(EpsII.val*FE)^(1/n.val)*f.val^(-r.val/n.val)*exp((E.val+P.val*V.val)/(n.val*R.val*T.val))/FT;
end

# Print info
function show(io::IO, g::DislocationCreep)
    print(io, "DislocationCreep: n=$(g.n.val), r=$(g.r.val), A=$(g.A.val), E=$(g.E.val), V=$(g.V.val), Apparatus=$(g.Apparatus)" )
end
#-------------------------------------------------------------------------

# Add pre-defined creep laws
"""
    SetDislocationCreep["Name of Dislocation Creep"]

This is a dictionary with pre-defined creep laws
"""
SetDislocationCreep = Dict([

# Olivine rheology
("Dry Olivine | Hirth & Kohlstedt (2003)",
# after Hirth, G. & Kohlstedt (2003), D. Rheology of the upper mantle and the mantle wedge: A view from the experimentalists.
# Inside the subduction Factory 83?105. Table 1, "dry dislocation" parameters
    DislocationCreep(
        n = 3.5NoUnits,
        A = 1.1e5MPa^(-3.5)/s,
        E = 530kJ/mol,
        V = 15e-6m^3/mol,
        Apparatus =   "AxialCompression",
        r = 0NoUnits,
        Comment = "Still to be verified with the original publication (BK). Values checked, plots are not reproduced (DK).",
        BibTex_Reference = parse_bibtex("""
            @incollection{eiler_rheology_2003,
            address = {Washington, D. C.},
            title = {Rheology of the upper mantle and the mantle wedge: {A} view from the experimentalists},
            volume = {138},
            isbn = {978-0-87590-997-4},
            shorttitle = {Rheology of the upper mantle and the mantle wedge},
            url = {http://www.agu.org/books/gm/v138/138GM06/138GM06.shtml},
            language = {en},
            urldate = {2019-10-09},
            booktitle = {Geophysical {Monograph} {Series}},
            publisher = {American Geophysical Union},
            author = {Hirth, Greg and Kohlstedt, David},
            editor = {Eiler, John},
            year = {2003},
            doi = {10.1029/138GM06},
            pages = {83--105},
            }
        """)
    )
)

# Olivine rheology
("Wet Olivine | Hirth & Kohlstedt (2003)",
    # After Hirth, G. & Kohlstedt (2003), D. Rheology of the upper mantle and the mantle wedge: A view from the experimentalists.
    #   Inside the subduction Factory 83?105. Table 1, "wet dislocation" parameters
    #  Note that this assumes C_OH=1000
    DislocationCreep(
        n = 3.5NoUnits,
        A = 90MPa^(-3.5)/s,
        E = 480kJ/mol,
        V = 11e-6m^3/mol,
        r   = 1.2NoUnits,
        Apparatus =   "AxialCompression",
        Comment = "Still to be verified with the original publication (BK). Values checked, plots are not reproduced (DK).",
        BibTex_Reference = parse_bibtex("""
            @incollection{HirthKohlstedt_OlivineRheology_2003,
            address = {Washington, D. C.},
            title = {Rheology of the upper mantle and the mantle wedge: {A} view from the experimentalists},
            volume = {138},
            isbn = {978-0-87590-997-4},
            shorttitle = {Rheology of the upper mantle and the mantle wedge},
            url = {http://www.agu.org/books/gm/v138/138GM06/138GM06.shtml},
            language = {en},
            urldate = {2019-10-09},
            booktitle = {Geophysical {Monograph} {Series}},
            publisher = {American Geophysical Union},
            author = {Hirth, Greg and Kohlstedt, David},
            editor = {Eiler, John},
            year = {2003},
            doi = {10.1029/138GM06},
            pages = {83--105},
            }
        """);
    )
)

]); # end of setting pre-defined creep laws


#using LinearAlgebra
#using DelimitedFiles
#using Statistics
#using Plots
#using DelimitedFiles
#using Random
#using PyCall
#using SparseArrays

#plt = pyimport("matplotlib.pyplot")

#include("Geometry_twisted.jl")

### for triangular lattice, use pi-flux hopping!!!

function get_valley_operator_twisted_bilayer(R,r,mode,theta,r0,layer_factor=1)  # d: inter-layer distance, lambda: parameter, mode="uniform_hop" or "pi-flux"
    t0=now()
    N=size(R,1)

    I,J,value=ones(Int,N*13),ones(Int,N*13),zeros(Complex,N*13)
    n=1
    for i=1:N
        for j=1:N
            dR=R[i]-(R[j]+r)
            z=dR[3]
            L=norm(dR)

            if abs(z)<0.01 && L<2
                t=get_valley_operator_hopping(R[i],R[j]+r,mode,theta,r0,layer_factor)
                if abs(t)>0
                    I[n]=i
                    J[n]=j
                    value[n]=t
                    n+=1
                end
            end
        end
    end
    P=sparse(I,J,value,N,N)
    P=dropzeros(P)
    t1=now()
    println("time to construct P: ",t1-t0)
    return P
end

function get_valley_operator_hopping(R,R2,mode,theta,r0,layer_factor=1)
    factor=1
    U=rotation_along_z_3D(-(theta))
    if R[3]==1
        R=U*(R-r0)
        R2=U*(R2-r0)
        factor=layer_factor
    end
    x=R[1]-R2[1]
    y=R[2]-R2[2]
    if mode=="uniform_hop"    #valley operator for honeycomb
        hop=0
        if round(x/sqrt(3),digits=3)==1 && round(y,digits=3)==0
            hop=im
        elseif round(x/sqrt(3),digits=3)==-1 && round(y,digits=3)==0
            hop=-im
        elseif round(x*2/sqrt(3),digits=3)==-1 && round(y,digits=3)==3/2
            hop=im
        elseif round(x*2/sqrt(3),digits=3)==1 && round(y,digits=3)==-3/2
            hop=-im
        elseif round(x*2/sqrt(3),digits=3)==-1 && round(y,digits=3)==-3/2
            hop=im
        elseif round(x*2/sqrt(3),digits=3)==1 && round(y,digits=3)==3/2
            hop=-im
        end
        return hop*layer_factor/(3*sqrt(3))
    elseif mode=="pi_flux"
        hop=0
        if round(x,digits=3)==1 && round(y,digits=3)==0
            hop=-im
        elseif round(x,digits=3)==-1 && round(y,digits=3)==0
            hop=im
        elseif round(x,digits=3)==0 && round(y/sqrt(3),digits=3)==-1
            hop=-im
        elseif round(x,digits=3)==0 && round(y/sqrt(3),digits=3)==1
            hop=im
        end
        return hop*factor/4
    else
        println("invalid mode for valley operator")
    end
end

function get_valley_operator_inter_twisted_bilayer(g::geometry_twisted,mode,layer_factor=1)
    R=g.sites
    theta=g.twist_angle
    r0=g.inter_alignment

    P_intra=get_valley_operator_twisted_bilayer(R,[0,0,0],mode,theta,r0,layer_factor)
    P_tx=get_valley_operator_twisted_bilayer(R,g.inter_vector.x,mode,theta,r0,layer_factor)
    P_ty=get_valley_operator_twisted_bilayer(R,g.inter_vector.y,mode,theta,r0,layer_factor)
    P_txy=get_valley_operator_twisted_bilayer(R,g.inter_vector.xy,mode,theta,r0,layer_factor)
    P=get_hamiltonian(P_intra,P_tx,P_ty,P_txy,g)
    return P
end

function get_Pk(P::get_hamiltonian,k::Array{Float64,1})
    P0=exp(im*dot(k,P.geometry.inter_vector.x))*P.tx+exp(im*dot(k,P.geometry.inter_vector.y))*P.ty+exp(im*dot(k,P.geometry.inter_vector.xy))*P.txy
    Pk=P.intra+P0+adjoint(P0)
    return Pk
end

function get_Pk_ribbon(P::get_hamiltonian,k::Array{Float64,1})   # here k represents the phase difference
    @assert k[1]*k[2]==0 "please input 1D momenta"
    k[1]==0 ? P0=exp(im*k[2])*P.ty : P0=exp(im*k[1])*P.tx

    Pk=P.intra+P0+adjoint(P0)
    return Pk
end

function plot_P(P,R)
    plt.figure(figsize=(6,6),dpi=80)
    N=size(R,1)
    for i=1:N
        for j=i:N
            if imag(P[i,j])>0.1
                plt.plot((R[i][1],R[j][1]), (R[i][2],R[j][2]),color="red",linewidth=norm(P[i,j])*1)
            elseif imag(P[i,j])<-0.1
                plt.plot((R[i][1],R[j][1]), (R[i][2],R[j][2]),color="blue",linewidth=norm(P[i,j])*1)
            end
        end
    end
    plt.axis([-5,5,-5,5])
    plt.xlabel("")
    plt.ylabel("")
    plt.show()
end

function get_layer_operator_twisted_bilayer(R)  # d: inter-layer distance, lambda: parameter, mode="uniform_hop" or "pi-flux"
    N=size(R,1)

    L=spzeros(Complex,N,N)
    for i=1:N
        if R[i][3]==0
            L[i,i]=1
        else
            L[i,i]=-1
        end
    end
    return L
end


x = randn(16,5)         # Set of 1D signals
y = randn(16,16,5)      # Set of 2D signals
wt = wavelet(WT.haar)

xw = wpdall(x,wt)
xsw = swpdall(x,wt)
xacw = acwpdall(x,wt)
yw = wpdall(y,wt)
ysw = swpdall(y,wt)
# yacw = acwpdall(y,wt)     # TODO: not available yet

# bb
@test isvalidtree(x[:,1], bestbasistree(xw[:,:,1], BB()))
@test isvalidtree(x[:,1], bestbasistreeall(xw, BB())[:,1])
@test isvalidtree(x[:,1], bestbasistreeall(xw, BB(LogEnergyEntropyCost(), false))[:,1])
@test isvalidtree(x[:,1], bestbasistreeall(xsw, BB(redundant=true))[:,1])
@test isvalidtree(x[:,1], bestbasistreeall(xacw, BB(redundant=true))[:,1])
@test isvalidtree(y[:,:,1], bestbasistree(yw[:,:,:,1], BB()))
@test isvalidtree(y[:,:,1], bestbasistreeall(yw, BB())[:,1])
@test isvalidtree(y[:,:,1], bestbasistreeall(yw, BB(LogEnergyEntropyCost(), false))[:,1])
@test isvalidtree(y[:,:,1], bestbasistreeall(ysw, BB(redundant=true))[:,1])
# @test isvalidtree(y[:,:,1], bestbasistreeall(yacw, BB(redundant=true))[:,1])

# jbb
@test isvalidtree(x[:,1], bestbasistree(xw))
@test isvalidtree(x[:,1], bestbasistree(xw, JBB(NormCost(), false)))
@test isvalidtree(x[:,1], bestbasistree(xsw, JBB(redundant=true)))
@test isvalidtree(x[:,1], bestbasistree(xacw, JBB(redundant=true)))
@test isvalidtree(y[:,:,1], bestbasistree(yw))
@test isvalidtree(y[:,:,1], bestbasistree(yw, JBB(NormCost(), false)))
@test isvalidtree(y[:,:,1], bestbasistree(ysw, JBB(redundant=true)))
# @test isvalidtree(y[:,:,1], bestbasistree(yacw, JBB(redundant=true)))

# lsdb
@test isvalidtree(x[:,1], bestbasistree(xw, LSDB()))
@test isvalidtree(x[:,1], bestbasistree(xsw, LSDB(redundant=true)))
@test isvalidtree(x[:,1], bestbasistree(xacw, LSDB(redundant=true)))
@test isvalidtree(y[:,:,1], bestbasistree(yw, LSDB()))
@test isvalidtree(y[:,:,1], bestbasistree(ysw, LSDB(redundant=true)))
# @test isvalidtree(y[:,:,1], bestbasistree(yacw, LSDB(redundant=true)))

# siwpd_bb
x = randn(16)
xw0 = siwpd(x, wt)
xw4 = siwpd(circshift(x,4), wt)
bt0 = map(node -> sum(node), bestbasistree(xw0, 4, SIBB()))
bt4 = map(node -> sum(node), bestbasistree(xw4, 4, SIBB()))
@test bt0 == bt4

# misc
@test_throws ArgumentError BestBasis.bestbasis_treeselection(randn(15), 8, :fail)
@test_throws AssertionError BestBasis.bestbasis_treeselection(randn(7), 3, :fail) # true n=4


module SymFinel

    using Calculus

    include("CoordExpr.jl")
    include("Interpolator.jl")
    include("Regressor.jl")
    include("Domain.jl")
    include("Finel.jl")
    include("Residuals.jl")

end # module

module SymFinelTests

    using Calculus

    include("CoordExpr.jl")
    include("Interpolator.jl")
    include("Regressor.jl")
    include("Domain.jl")
    include("Finel.jl")
    include("Residuals.jl")

    export get_basefun, get_basefun_derivative, get_linpar_contribution, get_interp_linmat,
        _get_gradfunc, regress, finel_get_regr, get_abs_coords!, from_finel_jac!, get_punctual_resfunc,
            vals_to_polys, polys_to_vars

end # module


using ADCME
using AdFem

include("chip_unstructured_solver.jl")
include("chip_unstructured_geometry.jl")
include("plot_inverse_one_iter.jl")

trialnum = 3

k_mold = 0.014531
k_chip_ref = 2.60475
k_air = 0.64357

# θ0 = [2.60475; 0.49; 1.0]
# θ = Variable([2.60475; 0.49; 1.0])
# paramA = θ[1]
# paramB = θ[2]
# paramC = θ[3]

function k_exact(x, y)
    k_mold + 1000 * k_chip_ref * (x-0.49)^2 / (1 + x^2)
end

function k_nn(x, y, θ)
    constant(@. k_exact(x, y)) +  fc(x, [20,20,20,1], θ)
end

nu = 0.47893  # equal to 1/Re
power_source = 82.46295  #82.46295 = 1.0e6 divide by air rho cp   #0.0619 = 1.0e6 divide by chip die rho cp
buoyance_coef = 299102.83

u_std = 0.001
p_std = 0.000001225
T_infty = 300

NT = 14    # number of iterations for Newton's method, 8 is good for m=400

n = 901  #ae_num([20,20,20,1]) => 861?
θ = Variable(randn(n))
θ = θ ./ norm(θ)

xy = mesh.nodes
xy2 = zeros(mesh.nedge, 2)
for i = 1:mesh.nedge
    xy2[i,:] = (mesh.nodes[mesh.edges[i,1], :] + mesh.nodes[mesh.edges[i,2], :])/2
end
xy = [xy;xy2]

x, y = xy[chip_fem_idx, 1], xy[chip_fem_idx, 2]
k_chip =  k_nn(x, y, θ)
k_chip_exact = constant( @. k_exact(x, y) )

heat_source_fem = zeros(ndof)
heat_source_fem[chip_fem_top_idx] .= power_source
heat_source_gauss = dof_to_gauss_points(heat_source_fem, mesh)
heat_source = compute_fem_source_term1(constant(heat_source_gauss), mesh)

B = constant(compute_interaction_matrix(mesh))
Laplace = constant(compute_fem_laplace_matrix1(nu * constant(ones(ngauss)), mesh))

# apply Dirichlet to velocity and temperature; set left bottom two points to zero to fix rank deficient problem for pressure
bd = [bd; bd .+ ndof;
      fvm_bd .+ 2*ndof;
      bd .+ (2*ndof+nelem)]

# add solid region into boundary condition for u, v, p, i.e. exclude solid when solving Navier Stokes
bd = [bd; solid_fem_idx; solid_fem_idx .+ ndof; solid_fvm_idx .+ 2*ndof]

S0 = constant(zeros(nelem+3*ndof))
S = solve_navier_stokes(S0, NT, k_chip)
S_computed = S[end, :]
S_data = matread("fn1/xz_chip_unstructured_data.mat")["V"]

loss =  mean((S_computed .- S_data)^2)
loss = loss * 1e10
sess = Session(); init(sess)

# ---------------------------------------------------
n = 901
θ0 = zeros(n)
@info run(sess, loss, θ=>θ0)
lineview(sess, θ, loss, θ0, ones(n))
savefig("fn$trialnum/lineview.png")
gradview(sess, θ, loss, zeros(n))
savefig("fn$trialnum/gradview.png")


function solve_disc(model::State_Space_Form{T}, obj::State_Space_Objective{T}, tol::T, maxiters::S) where {T <: AbstractFloat, S <: Int}

  nx    = copy(model.nx)
  ny    = copy(model.ny)
  beta  = copy(obj.beta)
  sigma = copy(model.sigma)

  a = copy(model.a)
  b = copy(model.b)
  c = copy(model.c)

  q = copy(obj.q)
  u = copy(obj.u)
  r = copy(obj.r)

  n = nx+ny

  a11 = a[1:nx,1:nx]
  a12 = a[1:nx,(nx+1):n]
  a21 = a[(nx+1):n,1:nx]
  a22 = a[(nx+1):n,(nx+1):n]

  b1 = b[1:nx,:]
  b2 = b[(nx+1):n,:]

  q11 = q[1:nx,1:nx]
  q12 = q[1:nx,(nx+1):n]
  q21 = q[(nx+1):n,1:nx]
  q22 = q[(nx+1):n,(nx+1):n]

  u1 = u[1:nx,:]
  u2 = u[(nx+1):n,:]

  j = (I-a22)\a21   # I = eye(ny)
  k = (I-a22)\b2    # I = eye(ny)
  f = (k'q22*k+k'u2+u2'k+r)\(k'q21+k'q22*j+u1'+u2'j)
  h = j-k*f
  v = q11+2*q12*h+h'q22*h-2*u1*f-2*h'u2*f+f'r*f

  len = Inf
  iters = 0
  while len > tol && iters < maxiters

    j    = (a22-h*a12)\(h*a11-a21)
    k    = (a22-h*a12)\(h*b1-b2)
    abar = a11+a12*j
    bbar = b1+a12*k
    qbar = q11+j'q21+q12*j+j'q22*j
    ubar = q12*k+j'q22*k+u1+j'u2
    rbar = r+k'q22*k+u2'k+k'u2
    fn   = (rbar+beta*bbar'v*bbar)\(ubar'+beta*bbar'v*abar)
    vn   = beta*(abar-bbar*fn)'v*(abar-bbar*fn)+qbar-fn'ubar'-ubar*fn+fn'rbar*fn
    hn   = j-k*fn

    len = maximum(maximum(abs,vn-v), maximum(abs,vn-v), maximum(abs,fn-f))
    h = hn
    v = vn
    f = fn

    iters += 1

  end

  if iters <= maxiters
    retcode = true
  else
    retcode = false
  end

  f = -f
  p = a11+a12*h+b1*f
  h = [eye(nx); h]
  k = c[1:nx,:]

  soln = State_Space_Soln(p,k,h,f,v,sigma,retcode)

  return soln

end

function solve_disc(model::Generalized_State_Space_Form{T}, obj::State_Space_Objective{T}, tol::T, maxiters::S) where {T <: AbstractFloat, S <: Int}

  nx    = copy(model.nx)
  ny    = copy(model.ny)
  beta  = copy(obj.beta)
  sigma = copy(model.sigma)

  a0 = copy(model.a0)
  a  = copy(model.a)
  b  = copy(model.b)
  c  = copy(model.c)

  q = copy(obj.q)
  u = copy(obj.u)
  r = copy(obj.r)

  n = nx+ny

  a11 = a[1:nx,1:nx]
  a12 = a[1:nx,(nx+1):n]
  a21 = a[(nx+1):n,1:nx]
  a22 = a[(nx+1):n,(nx+1):n]

  b1 = b[1:nx,:]
  b2 = b[(nx+1):n,:]

  q11 = q[1:nx,1:nx]
  q12 = q[1:nx,(nx+1):n]
  q21 = q[(nx+1):n,1:nx]
  q22 = q[(nx+1):n,(nx+1):n]

  u1 = u[1:nx,:]
  u2 = u[(nx+1):n,:]

  j = (a0-a22)\a21
  k = (a0-a22)\b2
  f = (k'q22*k+k'u2+u2'k+r)\(k'q21+k'q22*j+u1'+u2'j)
  h = j-k*f
  v = q11+2*q12*h+h'q22*h-2*u1*f-2*h'u2*f+f'r*f

  len = Inf
  iters = 0
  while len > tol && iters < maxiters

    j    = (a22-a0*h*a12)\(a0*h*a11-a21)
    k    = (a22-a0*h*a12)\(a0*h*b1-b2)
    abar = a11+a12*j
    bbar = b1+a12*k
    qbar = q11+j'q21+q12*j+j'q22*j
    ubar = q12*k+j'q22*k+u1+j'u2
    rbar = r+k'q22*k+u2'k+k'u2
    fn   = (rbar+beta*bbar'v*bbar)\(ubar'+beta*bbar'v*abar)
    vn   = beta*(abar-bbar*fn)'v*(abar-bbar*fn)+qbar-fn'ubar'-ubar*fn+fn'rbar*fn
    hn   = j-k*fn

    len = maximum(maximum(abs,vn-v), maximum(abs,vn-v), maximum(abs,fn-f))
    h = hn
    v = vn
    f = fn

    iters += 1

  end

  if iters <= maxiters
    retcode = true
  else
    retcode = false
  end

  f = -f
  p = a11+a12*h+b1*f
  h = [eye(nx); h]
  k = c[1:nx,:]

  soln = State_Space_Soln(p,k,h,f,v,sigma,retcode)

  return soln

end

function solve_disc(model::Structural_Form{T}, obj::Structural_Objective{T}, tol::T, maxiters::S) where {T <: AbstractFloat, S <: Int}

  a0    = copy(model.a0)
  a1    = copy(model.a1)
  a2    = copy(model.a2)
  a3    = copy(model.a3)
  a5    = copy(model.a5)
  sigma = copy(model.sigma)

  q = copy(obj.q)
  r = copy(obj.r)

  beta = copy(obj.beta)

  h1 = a0\a1
  h2 = a0\a5
  f1 = zeros(size(a3'))
  f2 = zeros(size(a3,2),size(a5,2))
  v = q

  a3_bar = Array(T,size(a3))
  a5_bar = Array(T,size(a5))

  len = Inf
  iters = 0
  while len > tol && iters < maxiters

    a1_bar = (a0-a2*h1)\a1
    a3_bar = (a0-a2*h1)\a3
    a5_bar = (a0-a2*h1)\a5

    vn = q+beta*(a1_bar+a3_bar*f1)'v*(a1_bar+a3_bar*f1)

    f1n = -(r+a3_bar'vn*a3_bar)\(a3_bar'vn*a1_bar)
    h1n = a1_bar+a3_bar*f1

    len = maximum(maximum(abs,vn-v), maximum(abs,h1n-h1), maximum(abs,f1n-f1))
    h1 = h1n
    f1 = f1n
    v  = vn

    iters += 1

  end

  if iters <= maxiters
    retcode = true
  else
    retcode = false
  end

  f2 = -(r+a3_bar'v*a3_bar)\(a3_bar'v*a5_bar)
  h2 = a5_bar+a3_bar*f2

  h = [[h1 zeros(size(a3))];
       [f1 zeros(size(r))]]

  g = [h2; f2]

  soln = Structural_Soln(h,g,v,sigma,retcode)

  return soln

end

function solve_disc(model::Generalized_Structural_Form{T}, obj::Structural_Objective{T}, tol::T, maxiters::S) where {T <: AbstractFloat, S <: Int}

  a0    = copy(model.a0)
  a1    = copy(model.a1)
  a2    = copy(model.a2)
  a3    = copy(model.a3)
  a4    = copy(model.a4)
  a5    = copy(model.a5)
  sigma = copy(model.sigma)

  q = copy(obj.q)
  r = copy(obj.r)

  beta = copy(obj.beta)

  h1 = a0\a1
  f1 = zeros(size(a3'))
  di = Array(T,size(a0))
  v  = q

  len = Inf
  iters = 0
  while len > tol && iters < maxiters

    (v,retcode_doubling) = doubling(beta*h1',h1,q+beta*f1'r*f1,tol,maxiters)
    di  = inv(a0-a2*h1-a4*f1)
    f1n = -(r+a3'di'v*di*a3)\(a3'di'v*di*a1)
    h1  = di*(a1+a3*f1n)

    len = maximum(abs,f1n-f1)
    f1 = f1n

    iters += 1

  end

  if iters <= maxiters
    retcode = true
  else
    retcode = false
  end

  f2 = -(r+a3'di'v*di*a3)\(a3'di'v*di*a5)
  h2 = (a0-a2*h1-a4*f1)\(a5+a3*f2)

  h = [h1 zeros(size(a3));
       f1 zeros(size(r))]

  g = [h2; f2]

  soln = Structural_Soln(h,g,v,sigma,retcode)

end


using BlockArrays
using ControlSystems
using LinearAlgebra
using PredictiveControl
using Test

const MPC = PredictiveControl

# Some utilities for testing
include( "../../src/utilities.jl" )
include( "../testUtils.jl" )

# Create a sample system
A = [1.0 0.0 0.0 0.0;
     1.0 2.0 0.0 0.0;
     0.0 0.0 3.0 0.0;
     0.0 0.0 1.0 4.0]
B = [1.0 0.0;
     0.0 0.0;
     0.0 1.0;
     0.0 0.0]
C = to_matrix( Float64, I, 4 )

sys = StateSpace( A, B, C, 0, 0.1 )

Q = diagm( [1.0, 2.0, 3.0, 4.0] )
R = diagm( [1.0, 2.0] )

N  = 5
nx = ControlSystems.nstates( sys )
nu = ControlSystems.ninputs( sys )

###################################################################
# Test the matrices of an uncontrolled system
###################################################################
clqr = ConstrainedTimeInvariantLQR( sys, N, Q, R, :Q )

# Form a propagation matrix to test
H = MPC.hessian( clqr, PrimalFullCondensing )
J = MPC.linearcoefficients( clqr, PrimalFullCondensing )

@test blocksize( H ) == (N, N)
@test size( H ) == (nu*N, nu*N)

@test blocksize( J ) == (N, 1)
@test size( J ) == (nu*N, nx)

# Make sure the blocks have the proper values for the system
for i = 1:N
#    b = view( mat, Block( i, 1 ) )

#    @test size( b ) == (nx, nx)
#    @test b == sys.A^i
end

###################################################################
# Test the matrices of a controlled system
###################################################################
clqr = ConstrainedTimeInvariantLQR( sys, N, Q, R, :Q, K = :dlqr )

# Form a propagation matrix to test
H = MPC.hessian( clqr, PrimalFullCondensing )
J = MPC.linearcoefficients( clqr, PrimalFullCondensing )

@test blocksize( H ) == (N, N)
@test size( H ) == (nu*N, nu*N)

@test blocksize( J ) == (N, 1)
@test size( J ) == (nu*N, nx)

# Make sure the blocks have the proper values for the system
for i = 1:N
#    b = view( mat, Block( i, 1 ) )

#    @test size( b ) == (nx, nx)
#    @test b == sys.A^i
end


#Data Loading specifics
data_path = "/path/to/data/"
data_prefix = "data_prefix"  #If the data file name is bob.npy, this should be 'bob'

#Model Parameters
iterations = 100
hard_clustering = false  #Soft or hard assignments
initial_clusters = 1
argmax_sample_stop = 5 #Change to hard assignment from soft at iterations - argmax_sample_stop
split_stop  = 5#Stop split/merge moves at  iterations - split_stop

random_seed = nothing #When nothing, a random seed will be used.

max_split_iter = 20
burnout_period = 20
max_clusters = Inf

#Model hyperparams
α = 10.0 #Concetration Parameter
hyper_params = niw_hyperparams(1.0,
    zeros(Float32,2),
    5,
    Matrix{Float32}(I, 2, 2)*1.0)

outlier_mod = 0.05 #Concetration Parameter
outlier_hyper_params = niw_hyperparams(1.0,
    zeros(Float32,2),
    5,
    Matrix{Float32}(I, 2, 2)*1.0)

#Saving specifics:
enable_saving = true
model_save_interval = 1000
save_path = "/path/to/save/dir/"
overwrite_prec = false
save_file_prefix = "checkpoint_"


using AverageShiftedHistograms
using StatPlots
using Plots
using Distributions
plotly()
using JLD2

function Bias(A::Array{<:Real,2},B::Array{<:Real,2})
    return  Bias =    (A.-B)./std(A .-B)
end
Bias(A::SharedArray{<:Real,2},B::SharedArray{<:Real,2}) = BiasRel( convert(Array{Real},A),convert(Array{Real},B))

##

N = 10
GBA = 1
GW = 10
N_est = 900

Tvals= [500 1000 5000 10000] #   [100 250 500 1000]#
useStartVal = false
useStartVal? startValFlag = "T" : startValFlag = "F";# false #false #
folderName = "./data/estimatesTest/"
file_name = folderName * "EstTest_$(N)_$(GW)_$(GBA)_$(Tvals)_$(N_est)_$(useStartVal).jld"
estDataT = @load( file_name,simPar_1,simPar_2,simPar_3,estPar_1,estPar_2,estPar_3,est_times)
#

T  = Tvals[1]
ind = find(x->x==T,Tvals)

plot_tot =Vector{Plots.Plot{Plots.PlotlyBackend}}

rng_space = -3:.1:3
pW1 =0
pB1 =0
pA1 =0
for ind = 1:4
#ind = 1

errA = Bias(simPar_3[ind,:,:],estPar_3[ind,:,:])
errB = Bias(simPar_2[ind,:,:],estPar_2[ind,:,:])
errW =  (simPar_1[ind,:,:].- estPar_1[ind,:,:])./(repmat(std(estPar_1[ind,:,:],1),N_est,1))

err= errW
tmpAsh = ash(err;rng =rng_space )
pW1 = plot(tmpAsh,
    title ="  T= $(Tvals[ind]),
    W   mu =  $(round(mean(err[1,:]),2)) ")
plot!(Normal(0,1),linewidth = 5)

err= errB
tmpAsh = ash(err;rng = rng_space)
pB1 = plot(tmpAsh,
title =    " $(N_est) Est   B   mu =  $(round(mean(err[1,:]),2)) ")
plot!(Normal(0,1),linewidth = 5)

err= errA
tmpAsh = ash(err;rng = rng_space)
pA1 = plot(tmpAsh,
title = startValFlag * "  $(GW) / $(GBA) / $(N) A  mu =  $(round(mean(err[1,:]),2)) ")
plot!(Normal(0,1),linewidth = 5)

exp = "pl_T_$(ind)=  plot(pW1,pB1,pA1,layout=(3,1) )"
eval(parse(exp))
end

exp2 = " " ; for i = 1:ind exp2 *=  "pl_T_$(i) , " end

exp3 = "plot( " * exp2 * " layout = (1,ind),size=(1350,600),legend =:none)"
eval(parse(exp3))
#plot(pl_T_1 , pl_T_2,pl_T_3 , pl_T_4, layout = (1,4),size=(1350,600),legend =:none)

##

plot(errW[1,:])

##


#=
mri_trajectory.jl
based on mri_trajectory.m
Jing Dong
2020-03

using MIRT: image_geom_mri, ndgrid, mri_kspace_spiral, ir_mri_kspace_ga_radial, prompt
=#

export mri_trajectory

using Random: seed!
using Plots; default(markerstrokecolor=:auto, markersize=2)
using Test: @test

"""
    kspace, omega, wi = mri_trajectory( ; ktype, N, fov, arg_wi, kwargs...)

Generate kspace trajectory samples and density compensation functions.

option
* `ktype::Symbol` k-space trajectory type; default `:radial`
* `N::Dims` target image size; default (32,30)
* `fov` field of view in x and y (and z); default [250,250] mm
* `arg_wi` options to pass to `ir_mri_density_comp` - not yet done
* `kwargs` options for the specific trajectory

out
* `kspace [Nk 2|3]` kspace samples in units 1/fov
* `omega [Nk 2|3]` trajectory samples over [-π,π)
* `wi [Nk 1]` (optional) density compensation factors

trajectory types:
`:cartesian`
`:radial`
`:cart_y_2`
`:random`
`:half8`
`:epi_sin`
`:spiral0 :spiral1 :spiral3`
`:rosette3`
`:epi_under`
`:gads` (emulate golden-angle data sharing per winkelmann:07:aor)
"""
function mri_trajectory( ;
		ktype::Symbol = :radial,
		N::Dims = (32,30),
		fov::AbstractVector{<:Number} = 250 * ones(length(N)),
		wi::Union{AbstractArray{<:Real},Nothing} = nothing,
		arg_wi::Any = nothing, # not done
		kwargs...) # optional arguments for specific trajectories

	length(N) ∉ (2,3) && throw("only 2d and 3d done")

	# default density compensation, which works for ordinary Cartesian (only)
	wi = nothing

	# trajectory choices

	if (ktype == :cartesian)
		o1 = ((0:(N[1]-1))/N[1] .- 0.5)*2π
		o2 = ((0:(N[2]-1))/N[2] .- 0.5)*2π
		if (length(N) == 2)
			o1, o2 = ndgrid(o1, o2)
			omega = [o1[:] o2[:]]
		end

		if (length(N) == 3)
			o3 = ((0:(N[3]-1))/N[3] .- 0.5)*2π
			o1, o2, o3 = ndgrid(o1, o2, o3)
			omega = [o1[:] o2[:] o3[:]]
		end

		wi = 1 / prod(fov)
	end

	if (ktype == :epi_under)
		omega, wi = mri_trajectory_epi_under(N, fov ; kwargs...)
	end

	if (ktype == :gads)
		omega, wi = mri_trajectory_gads(N, fov ; kwargs...)
	end

	if (ktype == :radial)
@show N, fov
@show N[1:2], fov[1:2]
		omega, wi = mri_trajectory_radial(N[1:2], fov[1:2] ; kwargs...)

		if (length(N) == 3) # 'stack of radials'
			omega = mri_trajectory_stack(omega, N[3])
			wi = repeat(wi, N[3], 1)
			wi = wi[:] / fov[3]
		end
	end

	if (ktype == :rosette3)
		omega, wi = mri_trajectory_rosette3(N, fov ; kwargs...)
	end

	# echo-planar with sinusoid:
	if (ktype == :epi_sin)
		tmp = ( ; oversample=1) -> oversample
		oversample = tmp( ; kwargs...)
		Npt = round(Int, oversample * prod(N))
		t = (0:(Npt-1))/Npt
		omega = [π*sin.(2π*t*N[2]/2) t*2π .- π]
	end

	# bad spiral:
	if (ktype == :spiral0)
		Nspiral = 400
		omega = LinRange(0, 10*2π, Nspiral)
		omega = π * [cos.(omega) sin.(omega)] .* [omega omega]/maximum(omega)
		wi = abs.(omega[:,1] + 1im * omega[:,2]) # simple |r| weighting
	end

	# crude spiral:
	if (ktype == :spiral1)
		Nspiral = round(Int, prod(N) * π/4)
		omega = LinRange(0, N[1]*2π, Nspiral)
		max_omega = maximum(omega)
		omega = π * [cos.(omega) sin.(omega)].* [omega omega]/max_omega
		wi = abs.(omega[:,1] + 1im * omega[:,2]) # simple |r| weighting
	end

	# 3T spiral:
	if (ktype == :spiral3)
		kspace, omega = mri_kspace_spiral(N = maximum(N[1:2]), fov = maximum(fov[1:2]))

		if length(N) == 3 #stack of these spirals
			omega = mri_trajectory_stack(omega, N[3])
		end

		wi = abs.(omega[:,1] + 1im * omega[:,2]) # simple |r| weighting
		wi = wi / (fov[1]*fov[2]) # approximate scaling
		if length(N) == 3
			length(fov) != 3 && throw("fov not 3D")
			wi = wi / fov[3] # cartesian-style weighting in z
		end
	end

	# random
	if (ktype == :random)
		seed!(0)
		omega = (rand(prod(N), 2) .- 0.5) * 2π
	end

	# 2D half Cartesian + 8 rows
	if (ktype == :half8)
		extra = 8
		o1 = ((0:(N[1]-1))/N[1] .- 0.5)*2π
		o2 = (-N[2]/2:extra)/N[2] * 2π
		o1,o2 = ndgrid(o1, o2)
		omega = [o1[:] o2[:]]
	end

	# 2D FT, undersampled in "y" (phase encode) direction
	if (ktype == :cart_y_2)
		under = 2
		o1 = ((0:(N[1]/1-1))/(N[1]/1) .- 0.5) * 2π
		o2 = ((0:(N[2]/2-1))/(N[2]/under) .- 0.5) * 2π
		o1,o2 = ndgrid(o1, o2)
		omega = [o1[:] o2[:]]
	end

	# convert to physical units
	kspace = zeros(Float32, size(omega))
	for id = 1:length(N)
		dx = fov[id] / N[id]
		kspace[:,id] = omega[:,id] / (2π) / dx
	end

	# wi = ir_mri_density_comp(kspace, arg_wi{:}) # todo

	return kspace, omega, wi
end

# mri_trajectory_stack()
# make 3D "stack of ..." trajectory from a 2D trajectory
function mri_trajectory_stack(omega2, N3::Int)
	o3 = ((0:(N3-1))/N3 .- 0.5)*2π
	o3 = repeat(o3, 1, size(omega2,1)) # [N3,N12]
	o3 = o3'[:] # [N12*N3]
	omega = repeat(omega2, N3) # [N12*N3,2]
	omega = [omega o3] # [N12*N3,3]
	return omega
end

# mri_trajectory_epi_under()
# EPI, with optional under-sampling
function mri_trajectory_epi_under(N, fov ;
		under::Int = 2, # default is every other phase encode
		samp::AbstractVector{<:Bool} = rem.(0:(N[2]-1), under) .== 0,
	)
	nx = N[1]
	ny = N[2]
	omx = (-nx/2:nx/2-1) / nx * 2π # [-π,π) in x
	omy = (-ny/2:ny/2-1) / ny * 2π # [-π,π) in y
	omega = zeros(0,2)
	for iy in findall(samp)
		omega = [omega; [omx omy[iy]*ones(nx)]]
		omx = reverse(omx, dims = 1) # non-flyback
	end
	wi = ones(size(omega,1)) / prod(fov)
	return omega, wi
end

"""
    omega, wi = mri_trajectory_gads(N, fov ; ...)
emulate 2D golden angle radial sampling with data sharing

option:
`Nro` # of samples in each readout/spoke
`shift` shift along read-out due to gradient delays (stress)
`kmax_frac` fractions of maximum krad (0.5) for rings
`under` under-sampling factor for each annulus
"""
function mri_trajectory_gads(N, fov ;
		Nro::Int = maximum(N),
		delta_ro::Real = 1/Nro,
		shift::Real = -0.75,
		kmax_frac::Array = [0.20, 0.35, 0.501],
		nspoke::Array = floor.(Int, π * kmax_frac * Nro),
		under::Array = [1, 1, 0.6],
		start::AbstractArray = (0:2) * π/4
	)

	nspoke = floor.(Int, nspoke .* under)
	if true # make fibonacci for more uniform coverage
		phi = (1 + sqrt(5)) / 2
		n = round.(log.(nspoke * sqrt(5) .+ 0.5) / log(phi))
		nspoke = phi.^n / sqrt(5) .+ 0.5
		nspoke = floor.(Int, nspoke)
	end
	nring = length(nspoke)
	kmax_frac = [0; kmax_frac]
	omega = zeros(0,2)
	for ir=1:nring
		#todo: not sure here
		kspace = ir_mri_kspace_ga_radial(Nspoke = nspoke[ir],
			Nro = Nro, delta_ro = delta_ro, shift = shift, start = start[ir])

		kspace = reshape(kspace, :, 2)
		krad = sqrt.(sum(kspace.^2, dims = 2)[:])
		good = (kmax_frac[ir] .<= krad) .& (krad .< kmax_frac[ir+1])
		kspace = kspace[good, :]
		omega = [omega; 2π*kspace]
	end
	wi = [] # no default DCF
	return omega, wi
end

"""
    mri_trajectory_radial()

option:
* `na_nr` default ensures proper sampling at edge of k-space
* `na` angular spokes; default: na_nr * nr
* `nr` radial samples per spoke
* `ir` default: `0:nr`

todo: generalize to 3D using barger:02:trc
"""
function mri_trajectory_radial(N::Dims, fov ;
		nr::Int = round(Int, maximum(N)/2),
		ir::AbstractArray{Int} = 0:nr,
		na_nr::Real = 2π,
		na::Int = 4*ceil(Int, na_nr * nr/4), # mult of 4
	)

	om = ir/nr * π
	ang = (0:na-1)/na * 2π
	om, ang = ndgrid(om, ang) # [nr+1, na]
	omega = [(om .* cos.(ang))[:] (om .* sin.(ang))[:]]

	# density compensation factors based on "analytical" voronoi
	fov[2] != fov[1] && @warn("fov not square")
	du = 1/fov[1] # assume this radial sample spacing
	wi = π * du^2 / na * 2 * ir # see lauzon:96:eop, joseph:98:sei
	wi = collect(wi)
	wi[ir .== 0] .= π * (du/2)^2 / na

	wi = repeat(wi, na)[:]
	return omega, wi
end

"""
    mri_trajectory_rosette3(N, fov ; ...)
3d rosette, with default parameters from bucholz:08:miw

options:
omax: maximum omega
nt : time samples (65.536 ms for 4 usec dt)
dt : time sample spacing (4 usec)
ti : time samples
"""
function mri_trajectory_rosette3(N, fov ;
		f1::Real = 211,
		f2::Real = 117.13,
		f3::Real = length(N) == 3 ? 73.65 : 0,
		omax::Real = π,
		nshot::Int = 32, # todo: shots vs arms
		nt::Int = length(N) == 3 ? 16385 : 2^9, # time samples (65.536 ms for 4 usec dt)
		dt::Real = 4e-6, # time sample spacing (4 usec)
		ti::AbstractVector = (0:nt-1) * dt, # time samples
	)

	tmp = 2π * ti
	p1 = f1 * tmp
	p2 = f2 * tmp
	p3 = f3 * tmp
	kx = omax * sin.(p1) .* cos.(p2) .* cos.(p3)
	ky = omax * sin.(p1) .* sin.(p2) .* cos.(p3)
	kz = omax * sin.(p1) .* sin.(p3)
	omega = [kx ky kz]
	for is = 1:nshot-1 # n-shot, rotate kx,ky by 2π / N
		ang = is * 2π / nshot
		c = cos(ang)
		s = sin(ang)
		ox = c * kx + s * ky
		oy = -s * kx + c * ky
		omega = [omega; [ox oy kz]]
	end
	wi = omax^3 * abs.(sin.(p1).^2 .* cos.(p1) .* cos.(p3)) # bucholz:08:miw
	if length(N) == 2
		omega = omega[:,1:2]
	end
	return omega, wi
end

"""
    mri_trajectory(:test)
self test
"""
function mri_trajectory(test::Symbol)
	test != :test && throw(DomainError(test, "Not valid"))
	ig = image_geom_mri(nx = 2^5, ny = 2^5-4, fov = 250) # 250 mm FOV
	N = ig.dim

	ktype3 = [:cartesian, :spiral3, :radial] # 3D
	for ktype in ktype3 # 3D tests
		@show ktype
		kspace, omega, wi = mri_trajectory( ; ktype=ktype, N=(N...,4))
		@test size(kspace,2) == size(omega,2) == 3
		@test size(kspace,1) == size(omega,1)
		@test length(wi) ∈ (1, size(kspace,1))
		scatter(omega[:,1]/π, omega[:,2]/π, omega[:,3]/π, label="$ktype")
	end
	prompt()

	ktypes = [:spiral3, :spiral1, :spiral0, :radial, :gads, :rosette3,
		:epi_sin, :epi_under, :cartesian, :half8, :cart_y_2, :random,]
	# arg_wi = [:voronoi]
	plots = Any[]
	for ktype in ktypes
		@show ktype
		args =
			(ktype == :epi_sin) ? Dict(:oversample => 1/1,) :
			(ktype == :radial) ? Dict(:nr => 10, :na_nr => π/2) :
				Dict{Symbol,Nothing}()

		kspace, omega, wi = mri_trajectory( ;
			ktype=ktype, N=N, fov = ig.fovs, args...)
		p = scatter(omega[:,1]/π, omega[:,2]/π, aspect_ratio=1,
			xlim = [-1,1]*1.1, xtick = -1:1, xlabel = "omega1/π",
			ylim = [-1,1]*1.1, ytick = -1:1, ylabel = "omega2/π",
			title = string(ktype), label="")
		push!(plots, p)
	end
	plot(plots...); gui()

	return true
end


@testset "testing TICR, Dirichlet distribution" begin
# previously in PhyloNetworks

df = CSV.read(joinpath(dirname(Base.find_package("PhyloNetworks")),"..","examples","buckyCF.csv"));
d = readTableCF(df);

@testset "ticr! on data frame, tree" begin
net1 = readTopology("((((D:0.4,C:0.4):4.8,((A:0.8,B:0.8):2.2)#H1:2.2::0.7):4.0,(#H1:0::0.3,E:3.0):6.2):2.0,O:11.2);");
setGamma!(net1.edge[9],0.0);
net2 = deepcopy(net1);
# without optimizing branch lengthes
result1 = ticr!(net1,df,false);
@test result1[2] ≈ 1.480872194397731   # z statistic, from R: prop.test(2,15, p=0.05, alternative="greater", correct=F)
@test result1[1] ≈ 0.06932031690660927 # p-value, from R
@test result1[3] == Dict("[0.0, 0.01)"=>2, "[0.01, 0.05)"=>0, "[0.05, 0.1)"=>2, "[0.1, 1.0]"=>11)
@test result1[4][2] ≈ 48.152697007372566418 # pseudo log-lik obtained from R
@test result1[4][1] ≈ 10.576940922426542713 atol=1e-5 # alpha obtained from R
#with optimizing branch lengthes
result2 = ticr!(net2,d,true);
net2_1 = readTopology("(D,C,(((A,B):1.297)#H1:1.412::1.0,((#H1:0.0::0.0,E):6.877,O):1.354):9.352);");
result3 = ticr!(net2_1,d,false);
@test result3[2] ≈ 1.480872194397731   # z statistic, from R, same as above
@test result3[1] ≈ 0.06932031690660927 # p-value, from R
@test result3[3] == Dict("[0.0, 0.01)"=>2, "[0.01, 0.05)"=>0, "[0.05, 0.1)"=>2, "[0.1, 1.0]"=>11)
@test result3[4][2] ≈ 54.241562916699216 # pseudo log-lik obtained from R
@test result3[4][1] ≈ 20.128258663235194 # alpha obtained from R
result3_1 = ticr!(net2_1,d,false; test=:goodness);
@test result3_1[2] ≈ 25.962962962962965463 # chi-squared statistic obtained from R
@test result3_1[1] ≈ 9.7092282251534852702e-06 # p-value obtained from R
@test result3_1[3] == Dict("[0.0, 0.01)"=>2, "[0.01, 0.05)"=>0, "[0.05, 0.1)"=>2, "[0.1, 1.0]"=>11)
@test result3[4][2] ≈ 54.241562916699216 atol=1e-3 # pseudo log-lik obtained from R
@test result3[4][1] ≈ 20.128258663235194 atol=1e-4 # alpha obtained from R
end

@testset "ticr! on data frame, network, p value from maximum observed values, dirichlet distribution" begin
net3 = readTopology("((((D:0.4,C:0.4):4.8,((A:0.8,B:0.8):2.2)#H1:2.2::0.7):4.0,(#H1:0::0.3,E:3.0):6.2):2.0,O:11.2);");
# without optimizing branch lengths
netresult1 = ticr!(net3,d,false);
@test netresult1[2] ≈ -0.8885233166386386  # z stat, from R: prop.test(0,15, p=0.05, alternative="greater", correct=F)
@test netresult1[1] ≈ 0.8128703403598878   # p-value, from R
@test netresult1[3] == Dict("[0.0, 0.01)"=>0, "[0.01, 0.05)"=>0, "[0.05, 0.1)"=>2, "[0.1, 1.0]"=>13)
@test netresult1[4][2] ≈ 68.03708830981597 # pseudo log-lik
@test netresult1[4][1] ≈ 29.34808731515701 atol=1e-5 # alpha
end

@testset "ticr! on data frame, network, minimum pvalue, dirichlet distribution" begin
net3 = readTopology("((((D:0.4,C:0.4):4.8,((A:0.8,B:0.8):2.2)#H1:2.2::0.7):4.0,(#H1:0::0.3,E:3.0):6.2):2.0,O:11.2);");
# without optimizing branch lengths
netresult1 = ticr!(net3,d,false; quartetstat=:minpval);
@test netresult1[2] ≈ 8.589058727506838  # z stat
@test netresult1[1] ≈ 4.384234705965304e-18   # p-value
@test netresult1[3] == Dict("[0.05, 0.1)" => 1, "[0.0, 0.01)"  => 4, "[0.01, 0.05)" => 4, "[0.1, 1.0]"  => 6)
@test netresult1[4][2] ≈ 68.03708830981597 # pseudo log-lik
@test netresult1[4][1] ≈ 29.34808731515701 atol=1e-5 # alpha
end

end


########################
# Generic approximation
########################

"""
The `approximation` function returns an operator that can be used to approximate
a function in the function set. This operator maps a grid to a set of coefficients.
"""
approximation(Φ; discrete = true, options...) =
  discrete ?
    discrete_approximation(Φ; options...) :
    continuous_approximation(Φ; options...)

function discrete_approximation(Φ::Dictionary; options...)
    if isbasis(Φ) && hasinterpolationgrid(Φ)
        interpolation(Φ; options...)
    else
        leastsquares(Φ; options...)
    end
end

continuous_approximation(Φ::Dictionary; options...) = inv(gram(Φ; options...))

# Automatically sample a function if an operator is applied to it
function (*)(op::DictionaryOperator, f::Function)
    op * project(src(op), f)
end

project(Φ::GridBasis, f::Function; options...) = sample(Φ, f)

approximate(Φ::Dictionary, f::Function; options...) =
    Expansion(Φ, approximation(Φ; options...) * f)

"""
The 2-argument approximation exists to allow you to transform any
Dictionary coefficients to any other Dictionary coefficients, including
Discrete Grid Spaces.
If a transform exists, the transform is used.
"""
function approximation(A::Dictionary, B::Dictionary, options...)
    if hastransform(A, B)
        return transform(A, B; options...)
    else
        error("Don't know a transformation from ", A, " to ", B)
    end
end

#################
# Interpolation
#################

function interpolation_matrix(Φ::Dictionary, pts, T=coefficienttype(Φ))
    @assert length(Φ) == length(pts)
    evaluation_matrix(Φ, pts, T)
end

interpolation(Φ::Dictionary, args...; options...) =
    interpolation(operatoreltype(Φ), Φ, args...; options...)

interpolation(::Type{T}, Φ::Dictionary; options...) where {T} =
    interpolation(T, Φ, interpolation_grid(Φ); options...)

interpolation(::Type{T}, Φ::Dictionary, grid::AbstractGrid; options...) where {T} =
    interpolation(T, Φ, GridBasis{T}(grid); options...)

function interpolation(::Type{T}, Φ::Dictionary, gb::GridBasis; options...) where {T}
    if hasinterpolationgrid(Φ) && interpolation_grid(Φ) == grid(gb) && hastransform(Φ, gb)
        transform(T, gb, Φ; options...)
    else
        default_interpolation(Φ, gb; options...)
    end
end

default_interpolation(Φ::Dictionary, gb::GridBasis; options...) =
    default_interpolation(operatoreltype(Φ, gb), Φ, gb; options...)

default_interpolation(::Type{T}, Φ::Dictionary, gb::GridBasis; options...) where {T} =
    QR_solver(evaluation(T, Φ, grid(gb)))

interpolate(Φ::Dictionary, pts, f, T=coefficienttype(s)) = Expansion(interpolation(T, Φ, pts) * f)

########################
# Generic least squares
########################

function leastsquares_matrix(Φ::Dictionary, pts, T=coefficienttype(Φ))
    @assert length(Φ) <= length(pts)
    evaluation_matrix(Φ, pts, T)
end

leastsquares(Φ::Dictionary, args...; options...) =
    leastsquares(operatoreltype(Φ), Φ, args...; options...)

function leastsquares(::Type{T}, Φ::Dictionary; samplingfactor = 2, options...) where {T}
    if hasinterpolationgrid(Φ)
        Φ2 = resize(Φ, samplingfactor*length(Φ))
        ls_grid = interpolation_grid(Φ2)
    else
        ls_grid = EquispacedGrid(samplingfactor*length(Φ), support(Φ))
    end
    leastsquares(T, Φ, ls_grid; options...)
end

leastsquares(::Type{T}, Φ::Dictionary, grid::AbstractGrid; options...) where {T} =
    leastsquares(T, Φ, GridBasis{T}(grid); options...)

function leastsquares(T, Φ::Dictionary, gb::GridBasis; options...)
    if hasinterpolationgrid(Φ)
        larger_Φ = resize(Φ, size(gb))
        if interpolation_grid(larger_Φ) == grid(gb) && hastransform(larger_Φ, gb)
            R = restriction(T, larger_Φ, Φ; options...)
            E = transform(T, gb, larger_Φ; options...)
            R * E
        else
            default_leastsquares(Φ, gb, T; options...)
        end
    else
        default_leastsquares(Φ, gb, T; options...)
    end
end

default_leastsquares(Φ::Dictionary, gb::GridBasis, T=operatoreltype(Φ,gb); options...) =
    QR_solver(ArrayOperator(leastsquares_matrix(Φ, grid(gb), T), Φ, gb))