import Base.length
import JSON
export GRN, reset!, set_input!, get_output, step!, distance

struct GRN
    nin::Int64
    nout::Int64
    nreg::Int64
    ids::Array{Float64}
    enh::Array{Float64}
    inh::Array{Float64}
    beta::Float64
    delta::Float64
    weights::Array{Float64}
    cons::Array{Float64}
end

function get_weights(ids::Array{Float64}, enh::Array{Float64}, inh::Array{Float64},
                     beta::Float64)
    glen = length(ids)
    weights = zeros(length(ids), length(ids))
    for i in eachindex(ids)
        for j in eachindex(ids)
            # influence of i on j (based on excitatory and inhibitory proteins of i)
            weights[i, j] = (exp(-beta * abs(enh[i] - ids[j])) -
                             exp(-beta * abs(inh[i] - ids[j])))
        end
    end
    weights
end

function GRN(nin::Int64, nout::Int64, nreg::Int64, config::Config)
    grn_size = nin + nout + nreg
    ids = rand(grn_size)
    enh = rand(grn_size)
    inh = rand(grn_size)
    beta = rand() * (config.beta_max - config.beta_min) + config.beta_min
    delta = rand() * (config.delta_max - config.delta_min) + config.delta_min
    weights = get_weights(ids, enh, inh, beta)
    cons = ones(grn_size) ./ grn_size
    GRN(nin, nout, nreg, ids, enh, inh, beta, delta, weights, cons)
end

function GRN(nin::Int64, nout::Int64, nreg::Int64, ids::Array{Float64},
             enh::Array{Float64}, inh::Array{Float64},
             beta::Float64, delta::Float64)
    grn_size = nin + nout + nreg
    weights = get_weights(ids, enh, inh, beta)
    cons = ones(grn_size) ./ grn_size
    GRN(nin, nout, nreg, ids, enh, inh, beta, delta, weights, cons)
end

function GRN(grn::GRN)
    grn_size = grn.nin + grn.nout + grn.nreg
    cons = ones(grn_size) ./ grn_size
    weights = get_weights(grn.ids, grn.enh, grn.inh, grn.beta)
    GRN(grn.nin, grn.nout, grn.nreg, copy(grn.ids), copy(grn.enh),
        copy(grn.inh), grn.beta, grn.delta, weights, cons)
end

function GRN(json_repr::String)
    d = JSON.parse(json_repr)
    ids = Array{Float64}(d["ids"])
    enh = Array{Float64}(d["enh"])
    inh = Array{Float64}(d["inh"])
    GRN(d["nin"], d["nout"], d["nreg"], ids, enh, inh, d["beta"], d["delta"])
end

function length(grn::GRN)
    length(grn.ids)
end

function reset!(grn::GRN)
    grn.cons .= 1./length(grn.ids)
end

function set_input!(grn::GRN, input::Array{Float64})
    @assert length(input) == grn.nin
    # @assert all(input .>= 0.0)
    # @assert all(input .<= 1.0)
    for i in eachindex(input)
        grn.cons[i] = input[i]
    end
end

function get_output(grn::GRN)
    grn.cons[grn.nin+(1:grn.nout)]
end

function step!(grn::GRN)
    reg = deepcopy(grn.cons)
    reg[grn.nin+(1:grn.nout)] .= 0.0
    cons = max.(0.0, grn.delta / length(grn.ids) * (reg' * grn.weights))
    sumcons = sum(cons[(grn.nin+1):end])
    if sumcons > 0
        cons = cons ./ sumcons
    end
    for i in (grn.nin+1):length(grn.ids)
        grn.cons[i] = cons[i]
    end
end

function step!(grn::GRN, nsteps::Int64)
    for i in 1:nsteps
        step!(grn)
    end
end

function step!(grn::GRN, inputs::Array{Float64})
    set_input!(grn, inputs)
    step!(grn)
    get_output(grn)
end

function distance(g1::Array{Float64}, g2::Array{Float64})
    # TODO: align regulatory
    dist = 1.0
    if (length(g1) > 0) || (length(g2) > 0)
        if length(g1) > length(g2)
            dist = sum(abs.(g1 - [g2; zeros(length(g1)-length(g2))]))
        elseif length(g1) < length(g2)
            dist = sum(abs.(g2 - [g1; zeros(length(g2)-length(g1))]))
        else
            dist = sum(abs.(g1 .- g2))
        end
    end
    dist
end

function distance(g1::GRN, g2::GRN, config::Config)
    (distance(g1.ids, g2.ids) * config.dist_coef_ids +
     distance(g1.enh, g2.enh) * config.dist_coef_enh +
     distance(g1.inh, g2.inh) * config.dist_coef_inh +
     (abs(g1.beta - g2.beta) / (config.beta_max - config.beta_min)) +
     (abs(g1.delta - g2.delta) / (config.delta_max - config.delta_min)))
end


type GMR{M<:MvNormal}

    # μ₁₋₂ = μ₁ + Σ₁₂ * Σ₂₂⁻¹ * (x₂ - μ₂) = A*x₂ + b
    vec_A::Vector{Matrix{Float64}} # [n_components [ntargets×nindicators]]
    vec_b::Vector{Vector{Float64}} # [n_components [ntargets]]

    # pdf(p), all pre-computed. Used to compute βⱼ(p)
    mixture_Obs::MixtureModel{Multivariate,Continuous,M} # p(obs), all pre-computed, should never be edited

    # pdf(a|p), means μⱼ_ₚ and weights βⱼ(p) are functions of p and must be updated every time, covariance is constant
    mixture_Act_given_Obs::MixtureModel{Multivariate,Continuous,M}
end
function GMR{M<:MvNormal}(mix::MixtureModel{Multivariate,Continuous,M}, n_targets::Int=2)

    #=
    Construct a Gaussian Mixture Regressor using a Gaussian mixture over both
    the features and the actions, and the features are at the end of the input
    =#

    weights = probs(mix.prior) # [n_components]
    n_components = length(weights)
    n_indicators = length(mix) - n_targets

    vec_A = Array(Matrix{Float64}, n_components) # μ₁₋₂ = μ₁ + Σ₁₂ * Σ₂₂⁻¹ * (x₂ - μ₂) = A*x₂ + b
    vec_b = Array(Vector{Float64}, n_components)
    vec_G = Array(MvNormal, n_components)
    vec_H = Array(MvNormal, n_components)

    for i = 1 : n_components

        μ = mix.components[i].μ
        μₐ = μ[1:n_targets]
        μₚ = μ[n_targets+1:end]

        Σ = full(mix.components[i].Σ)
        Σₐₐ = Σ[1:n_targets,1:n_targets]
        Σₐₚ = Σ[1:n_targets,n_targets+1:end]
        Σₚₚ = nearestSPD(Σ[n_targets+1:end,n_targets+1:end])
        iΣₚₚ = inv(Σₚₚ)

        A = Σₐₚ * iΣₚₚ
        vec_A[i] = A
        vec_b[i] = vec(μₐ - A*μₚ)
        C = nearestSPD(Σₐₐ - Σₐₚ * iΣₚₚ * ((Σₐₚ)'))

        vec_G[i] = MvNormal(Array(Float64, n_targets), C) # p(action|obs), mean and weighting must be updated with each observation, cov is pre-computed
        vec_H[i] = MvNormal(μₚ, Σₚₚ) # p(obs), all pre-computed, should never be edited
    end

    mixture_Act_given_Obs = MixtureModel(vec_G) # p(action|obs), mean and weighting must be updated with each observation, cov is pre-computed
    mixture_Obs = MixtureModel(vec_H, weights) # p(obs), all pre-computed, should never be edited
    GMR(vec_A, vec_b, mixture_Obs, mixture_Act_given_Obs)
end

function Base.print(model::GMR)
    println("GMR:")
    for (i, mat) in enumerate(model.vec_A)
        println(i)
        print("\t[")
        for j in 1:size(mat,2)
            @printf("  %10.6f", mat[1,j])
        end
        @printf("] + [  %10.6f]\n", model.vec_b[i][1])
        print("\t[")
        for j in 1:size(mat,2)
            @printf("  %10.6f", mat[2,j])
        end
        @printf("] + [  %10.6f]\n", model.vec_b[i][2])
    end
    println("\tmixture_Obs: ")
    println("\t\tprior: ", model.mixture_Obs.prior)
end

n_targets(gmr::GMR) = size(gmr.vec_A[1], 1)
n_features(gmr::GMR) = size(gmr.vec_A[1], 2)
n_components(gmr::GMR) = length(gmr.vec_A)
function nsuffstats(gmr::GMR)
    dimA = length(gmr.vec_A[1])
    n_components(gmr) * (2*dimA + 2 # bias
                                  + 3 # covariance in mixture_Act_given_Obs
                                  + div(dimA*dimA,2)) # covariance for mixture_Obs
end
function nearestSPD(A::Matrix{Float64})

    # see http://www.mathworks.com/matlabcentral/fileexchange/42885-nearestspd

    # output:
    #  α, β ≥ 0.0 such that
    #         α ≤ δ₂(A) ≤ β ≤ α + 2 max(fα, tol)
    #  and a PSD matrix X such that |A - X|₂ = β

    n = size(A, 1)
    @assert(n == size(A, 2)) # ensure it is square

    I = eye(n)

    # symmetrize A into B
    B = (A+A')./2

    # Compute the symmetric polar factor of B. Call it H.
    # Clearly H is itself SPD.
    U, σ, V = svd(B)
    H = V*diagm(σ)*V'

    # get Ahat in the above formula
    Ahat = (B+H)/2

    # ensure symmetry
    Ahat = (Ahat + Ahat')/2;

    # test that Ahat is in fact PD. if it is not so, then tweak it just a bit.
    worked = false
    iteration_count = 0
    while !worked && iteration_count < 100

        iteration_count += 1

        try
            chol(Ahat)
            worked = true
        catch
            # do nothing
        end

        if !worked
            # Ahat failed the chol test. It must have been just a hair off,
            # due to floating point trash, so it is simplest now just to
            # tweak by adding a tiny multiple of an identity matrix.

            min_eig = minimum(eigvals(Ahat))
            Ahat = Ahat + (-min_eig*iteration_count.^2 + eps(Float32))*I
        end
    end

    Ahat
end

@compat function (gmr::GMR)(features::Vector{Float64})

    mixture_Act_given_Obs = gmr.mixture_Act_given_Obs
    mixture_Obs = gmr.mixture_Obs
    nc = n_components(gmr)

    for j in 1 : nc

        # compute the β value, unweighted
        # βⱼ(f) ∝ wⱼ Nⱼ(μₚ, Σₚ)
        wⱼ = mixture_Obs.prior.p[j]
        Nⱼ = mixture_Obs.components[j]
        mixture_Act_given_Obs.prior.p[j] = wⱼ * pdf(Nⱼ, features)

        # compute the conditional mean
        #  μₐ_ₚ = A⋅f + b
        A = gmr.vec_A[j]
        b = gmr.vec_b[j]
        copy!(mixture_Act_given_Obs.components[j].μ, A*features + b)
    end

    # normalize the β values
    sum_β = sum(mixture_Act_given_Obs.prior.p)

    if sum_β > 0.0 && !isnan(sum_β) && !isinf(sum_β)
        for i in 1 : nc
            mixture_Act_given_Obs.prior.p[i] /= sum_β
        end
    else
        fill!(mixture_Act_given_Obs.prior.p, 1/nc) # set all to equal weight
        for i in 1 : nc
            fill!(mixture_Act_given_Obs.components[i].μ, 0.0) # set mean to zero
        end
    end

    mixture_Act_given_Obs
end


function ant(puzzle::Dict{Complex,Int}, at::Complex, bursts::Int, update::Function)::Int
	d::Complex = -1
	count::Int = 0
	for _ in 1:bursts
		status = get(puzzle, at, 0)
		if status == 0
			d *= im
		elseif status == 2
			d *= -im
		elseif status == 3
			d *= im * im
		end
		puzzle[at] = status = update(status)
		count += status == 2
		at += d
	end
	count
end

part1(status::Int)::Int = status == 0 ? 2 : 0
part2(status::Int)::Int = (status + 1) % 4

function main()
	lines = readlines()
	w = length(lines)
	h = length(lines[1])
	grid = Dict{Complex,Int}()
	for (i, line) in enumerate(lines)
		for (j, c) in enumerate(line)
			c == '#' && setindex!(grid, 2, i + j*im)
		end
	end

	mid = complex(1+div(w, 2), 1+div(h, 2))
	println(ant(copy(grid), mid, 10000, part1))
	println(ant(grid, mid, 10000000, part2))
end

main()


## estimate MAP by first estimating MAP for each chain
MAP(x::T, data) where T <: Fit = MAP([MAP(ch, data) for ch in x.chains], data)

MAP(x::Chain, data) = x.theta[last(findmax(sum(x.loglik_p, dims=2)[:,1]))]

function MAP(theta::Vector{T}, data) where T <: AbstractParam
    M = length(theta)
    ll = -Inf
    tmax = 1

    for t in 1:M
        new_ll = loglik(theta[t], data)

        if new_ll > ll
            ll = new_ll
            tmax = t
        end
    end

    theta[tmax]
end

weibull_loghaz(time, shape, scale) = log(scale) .+ log(shape) .+ (shape-1) .* log.(time)
weibull_cumhaz(time, shape, scale) = scale .* time.^shape
##inv_weibull_cumhaz(p, shape, scale) = (p ./ scale).^(1/shape)

function loglik_p_surv(T, E, Xe, shape::Tx, scale::Tx, delta::Vector{Tx}) where Tx <: Real
    if (shape .<= 0.0) | (scale .<= 0.0)
        -Inf
    else
        Xed = Xe*delta
        sum(E.*(weibull_loghaz(T .- 1, shape, scale) .+ Xed) .- weibull_cumhaz(T .- 1, shape, scale).*exp.(Xed))
    end
end

function loglik_p_surv(T, E, Xe, haz_basis, cumhaz_basis, w, delta)
    Xed = Xe*delta
    ew = exp.(w)
    sum(E.*(log.(haz_basis*ew) .+ Xed) .- cumhaz_basis*ew.*exp.(Xed))
end

##invquad(L, x) = dot(x, L' \ (L \ x))
##mvn_logpdf(x, mu, L) = -length(mu)/2 * log(2*pi) - log(det(L)) - invquad(L, x-mu)/2

function loglik_p_long(b, beta, mu, Sigma, sigma, Y, times, n, Xl, Z)
    target = zero(eltype(mu))
    d = Normal(0, sigma)

    for j in 1:n
        target += logpdf(d, Y[j] - dot(Z[j,:], b) - dot(Xl[j,:], beta))
        ##target += logpdf(Normal(zero(eltype(sigma)), sigma), Y[j] - dot(Z[j,:], u) - dot(Xl[j,:], beta))
        ## loglikelihood(Normal(0, sigma), Y[1:n] .- Z[1:n,:]*u - Xl[1:n,:]*beta)
    end

    target + logpdf(MvNormal(mu, Sigma), b) ##mvn_logpdf(b, zero(b), Sigma)
end

function loglik_p_long(beta, mu, Sigma, sigma, Y, times, n, Xl, Z)
    m = Z*mu + Xl*beta
    C = Array(Symmetric(Z*Sigma*Z' + sigma^2*diagm(0 => ones(Float64, n))))
    issuccess(cholesky(C; check=false)) || return -Inf
    logpdf(MvNormal(m, C), Y)
end

function loglik_p(theta::Param, data)
    @unpack T, E, Xe, Xl, Xp, Z, Y, times, n = data
    @unpack xi, cind, delta, bl_shape, bl_scale, b, beta, mu, S, L, sigma = theta

    cp = class_prob(xi, Xp)

    N = length(T)
    G = size(beta)[1]

    loglik = zeros(Float64, N)

    Sigma = [build_cov_matrix(S[g,:], L[g,:,:]) for g in 1:G]

    for i in 1:N
        g = cind[i]
        lll_vec = loglik_p_long(b[i,:], beta[g,:], mu[g,:], Sigma[g], sigma, Y[i,1:n[i]], times[i,1:n[i]], n[i], Xl[i,1:n[i],:], Z[i,1:n[i],:])
        lle_vec = loglik_p_surv(T[i], E[i], Xe[i,:]', bl_shape[g], bl_scale[g], delta[g,:])
        loglik[i] = lll_vec + lle_vec + log(cp[i,g])
    end

    loglik
end

function loglik_p(theta::ParamSpline, data)
    @unpack T, E, Xe, Xl, Xp, Z, Y, times, n, haz_basis, cumhaz_basis = data
    @unpack xi, cind, delta, w, b, beta, mu, S, L, sigma = theta

    cp = class_prob(xi, Xp)

    N = length(T)
    G = size(beta)[1]

    ll = zeros(Float64, N)

    Sigma = [build_cov_matrix(S[g,:], L[g,:,:]) for g in 1:G]

    for i in 1:N
        g = cind[i]
        lll_vec = loglik_p_long(b[i,:], beta[g,:], mu[g,:], Sigma[g], sigma, Y[i,1:n[i]], times[i,1:n[i]], n[i], Xl[i,1:n[i],:], Z[i,1:n[i],:])
        lle_vec = loglik_p_surv(T[i], E[i], Xe[i,:]', haz_basis[i,:]', cumhaz_basis[i,:]', w[g,:], delta[g,:])
        ll[i] = lll_vec + lle_vec + log(cp[i,g])
    end
    ll
end

loglik(theta, data) = sum(loglik_p(theta, data))

loglik_array(x::T, data::Data) where T <: AbstractFit = loglik_array(x.chains, data)

function loglik_array(x::AbstractVector{T}, data::Data) where T <: Chain
    M, K, N = length(x[1].theta), length(x), length(data.T)

    ##lp = SharedArray{Float64,3}((M, K, N))
    lp = Array{Float64,3}(undef, M, K, N) ##SharedArray{Float64,3}((M, K, N))

    ##@sync @distributed
    for m in 1:M
        for k in 1:K
            lp[m,k,:] = loglik_p(x[k].theta[m], data)
        end
    end
    lp
end


include("euler/euler.jl")
using .Primes: prime_numbers
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute()::Int
    for prime ∈ reverse(prime_numbers(7_654_321))
        str_prime = string(prime)
        if sort(parse.(Int, ∪(str_prime))) == collect(1:length(str_prime))
            return prime
        end
    end
end

compute()

@benchmark compute()


using Nemo, LinearAlgebra, DelimitedFiles

# using Revise
using SymbolicCRN

M = [63 90 36 -45 -18
    -7 -10 -4 5 2
    28 40 16 -20 -8
    -56 -80 -32 40 16
    7 10 4 -5 -2]

K = [2 5 0 0 0 0
    0 0 1 0 0 1
    0 0 0 5 1 0
    0 7 2 4 0 0
    7 0 0 0 2 5]

R, vrs = PolynomialRing(Nemo.ZZ, vcat(["k$i" for i in 1:5], ["x$i" for i in 1:4]));

@test all(integermultiple(raysof(SymbolicCRN.cone_positiveorthant(5))) .== LinearAlgebra.I(5))
@test integermultiple(raysof(cone_positivenullspace([0 0; 0 0]))) == [1 0; 0 1]
@test raysof(cone_positivenullspace([1 -1; 2 -2])) == transpose([1 1])
Knew = raysof(cone_positivenullspace(M))
@test all(M * Knew .== 0)
@test all(map(>=(0), Knew))
@test K == integermultiple(Knew)
p = prod([vrs[i] - i for i in 1:7])
@test integermultiple(raysof(outernormalcone(p, 5))) == [
    0 1 0 0 0 0 0
    1 0 0 0 0 0 0
    0 0 1 0 0 0 0
    0 0 0 1 0 0 0
    0 0 0 0 0 0 -1
    0 0 0 0 1 0 0
    0 0 0 0 0 1 0
    0 0 0 0 0 0 0
    0 0 0 0 0 0 0
]

# TODO;
# verticesof(polytope) = permutedims(Rational.(polymake_dehomog(polytope.VERTICES)))


# functions to decompose a proper polygon into triangles
export triangulate

"""
`alt_mod(k,n)` is `mod(k,n)` unless the result is zero,
in which case we return `n`.
"""
function alt_mod(k::Int, n::Int)
    r = mod(k,n)
    return r==0 ? n : r
end

"""
`check_diagonal(X::HPolygon, k::Int)` checks to see if the diagonal
between vertices `k-1` and `k+1` (a) does not intersect the boundary of `X`
and (b) goes through the interior of `X`.
"""
function check_diagonal(X::HPolygon, k::Int)::Bool
    n = npoints(X)

    if n <= 3
        return true
    end

    kprev = alt_mod(k-1,n)
    knext = alt_mod(k+1,n)

    a = X.plist[kprev]
    b = X.plist[k]      # this is the current vertex (index k)
    c = X.plist[knext]

    dg = a+c  # this is the diagonal we're testing

    # check if ac contains a vertex (other than a or c)
    for i=1:n
        if i != kprev && i != knext
            v = X.plist[i]
            if in(v,dg)
                return false
            end
        end
    end

    # now check that ac doesn't intersect any of the sides of the polygon
    f(x) = alt_mod(x,n)
    exclude = [ f(k+j) for j=-2:1 ]
    for i=1:n
        if in(i,exclude)
            continue
        end
        S = X.plist[f(i)] +X.plist[f(i+1)]
        if meet_check(dg,S)
            return false
        end

    end

    # check that the diagonal is interior to X
    m = midpoint(dg)
    if !in(m,X)
        return false
    end

    return true
end

"""
`find_ear_diagonal(X::HPolygon)` returns an index `k` such that the
diagonal from `k-1` to `k+1` passes `check_diagonal`.
Returns `0` if no ear diagonal is found.
"""
function find_ear_diagonal(X::HPolygon)::Int
    for k=1:npoints(X)
        if check_diagonal(X,k)
            return k
        end
    end
    return 0
end

"""
`triangulate(X::HPolygon)` returns a list of triangles that
triangulate `X`. The polygon should have at least three sides, and
have no bad angles, and not self-intersect.
"""
function triangulate(X::HPolygon)::Array{HTriangle,1}
    n = npoints(X)

    if n < 3
        return HTriangle[]

    end

    if n == 3
        T = HTriangle(X)
        return [T]
    end

    k = find_ear_diagonal(X)

    if k==0
        @error "The polygon cannot be triangulated. Is it simple and proper?"
    end

    kprev = alt_mod(k-1,n)
    knext = alt_mod(k+1,n)
    a = X.plist[kprev]
    b = X.plist[k]
    c = X.plist[knext]
    T = a+b+c

    qlist = [ HPoint(X.plist[j]) for j=1:n if j!=k ]
    Y = HPolygon(qlist)
    result = triangulate(Y)
    pushfirst!(result, T)
    return result
end

"""
`area(X::HPolygon)` returns the area of the polygon, but is only
reliable if `X` does not self-intersect. See `is_simple`.
"""
function area(X::HPolygon)::Float64
    if npoints(X) < 3
        return 0.0
    end
    TT = triangulate(X)
    return sum(area(T) for T in TT)
end

# DEBUG STUFF

function draw_diagonal(X::HPolygon, k::Int)
    n = npoints(X)
    kprev = alt_mod(k-1,n)
    knext = alt_mod(k+1,n)
    a = X.plist[kprev]
    b = X.plist[knext]
    S = a+b
    set_line_style(S,:dot)
    draw(S)
end

function visualize(P::HPolygon)
    n = npoints(P)
    set_thickness(P,2)
    plot()
    draw(P)
    f(x) = alt_mod(x,n)
    for i=1:n
        a = P.plist[f(i-1)]
        ii = alt_mod(i+1,n)
        b = P.plist[ii]
        S = a+b
        if check_diagonal(P,i)
            set_color(S,:green)
        else
            continue
        end
        if in(midpoint(S),P)
            draw(S)
        end
    end

    for pt in P.plist
        draw(pt)
    end
    finish()
end


function logprior(
  rparams::RiskParameters{S},
  rpriors::RiskPriors{S}) where {
  S <: DiseaseStateSequence}
  lprior = 0.0
  for i in 1:length(rpriors.sparks)
    lprior == -Inf && break
    lprior += Distributions.logpdf(rpriors.sparks[i], rparams.sparks[i])
  end
  for i in 1:length(rpriors.susceptibility)
    lprior == -Inf && break
    lprior += Distributions.logpdf(rpriors.susceptibility[i], rparams.susceptibility[i])
  end
  for i in 1:length(rpriors.infectivity)
    lprior == -Inf && break
    lprior += Distributions.logpdf(rpriors.infectivity[i], rparams.infectivity[i])
  end
  for i in 1:length(rpriors.transmissibility)
    lprior == -Inf && break
    lprior += Distributions.logpdf(rpriors.transmissibility[i], rparams.transmissibility[i])
  end
  if State_E ∈ S
    for i in 1:length(rpriors.latency)
      lprior == -Inf && break
      lprior += Distributions.logpdf(rpriors.latency[i], rparams.latency[i])
    end
  end
  if State_R ∈ S
    for i in 1:length(rpriors.removal)
      lprior == -Inf && break
      lprior += Distributions.logpdf(rpriors.removal[i], rparams.removal[i])
    end
  end
  return lprior
end

function logprior(
  sm::NucleicAcidSubstitutionModel,
  sm_priors::Vector{UnivariateDistribution})
  lprior = 0.0
  for i in zip([getproperty(sm, θ) for θ in [propertynames(sm)...]], sm_priors)
    lprior += Distributions.logpdf(i[2], i[1])
    lprior == -Inf && break
  end
  return lprior
end


function decompose(ln::ReferenceLine,x::Float64)
    @assert x ∈ ln
    a,b = vertices(ln)
    return LagrangeLine(x,a),LagrangeLine(x,b)
end

function decompose(tri::ReferenceTriangle,x)
    @assert x ∈ tri
    a,b,c = vertices(tri)
    return LagrangeTriangle(a,x,b), LagrangeTriangle(b,x,c), LagrangeTriangle(c,x,a)
end

function decompose(sq::ReferenceSquare,x)
    @assert x ∈ sq
    a,b,c,d = vertices(sq)
    return LagrangeTriangle(a,x,b), LagrangeTriangle(b,x,c), LagrangeTriangle(c,x,d), LagrangeTriangle(d,x,a)
end

function decompose(sq::ReferenceSquare,x,target::ReferenceSquare)
    @assert x ∈ sq
    a,b,c,d = vertices(sq)
    return  LagrangeRectangle(x,(0,x[2]),(0,0),(x[1],0)),
            LagrangeRectangle(x,(x[1],0),(1,0),(1,x[2])),
            LagrangeRectangle(x,(1,x[2]),(1,1),(x[1],1)),
            LagrangeRectangle(x,(x[1],1),(0,1),(0,x[2]))
end


using DRIPs
using Test
using LinearAlgebra

@testset "MW2009.jl" begin
	## Parameters
	ρ  = 0.95;
	σq = 0.01;
	σz = 11.8*σq;
	κ  = 3;
	ξ  = 1;
	α  = 1 - 0.15;

	## Primitives of drip
	L  = 21; # length of trunction
	A  = [zeros(1,L);[Matrix(I,L-1,L-1);zeros(1,L-1)]']; # MW truncate the state space with linear irfs of length 20
	Qq = zeros(L,1); Qq[1]=σq;
	Qz = zeros(L,1); Qz[1]=σz;
	H  = zeros(L,1); H[1:21] = Array(1:-1/20:0);

	## Function for solving the fixed point of the aggregate shock given ω
	# Note that for a certain range of ω the fixed point is not unique:
	# there is a fixed point with zero capacity and one with positive capacity.
	# In this range, we would like to find the fixed point in which the firms are
	# acqurining a positive amount of information.

	function agg_drip(ω,A,Qq,        #primitives of drip except for H because H is endogenous
	                  α,             #strategic complementarity
	                  H;             #state space rep. of q
	                  β     = 1,     #optional: discount factor, MW's parameterization implies β = 1
	                  H0    = H,     #optional: initial guess for HΔ (H is the true solution when α=0)
	                  maxit = 10000, #optional: max number of iterations for GE code
	                  tol   = 1e-4,  #optional: tolerance for iterations
	                  w     = 1)     #optional: update weight for RI
	    # set primitives
	    errmin= 1;
	    err   = 1;
	    iter  = 0;
	    L     = length(H);
	    while (err > tol) & (iter < maxit)
	            if iter == 0
	                global agg  = Drip(ω,β,A,Qq,H0;w = w);
	            else
	                global agg  = Drip(ω,β,A,Qq,H0;Ω0 = agg.ss.Ω , Σ0 = agg.ss.Σ_1,w = w);
	            end

	            XFUN(jj) = ((I-agg.ss.K*agg.ss.Y')*agg.A)^jj * (agg.ss.K*agg.ss.Y') * (agg.A')^jj
	            X = DRIPs.infinitesum(XFUN; maxit=200, start = 0);  #E[x⃗]=X×x⃗

	            XpFUN(jj) = α^jj * X^(jj)
	            Xp = DRIPs.infinitesum(XpFUN; maxit=200, start = 0);

	            H1 = (1-α)*Xp'*H;

	            err= 0.5*norm(H1-H0,2)/norm(H0)+0.5*err;
	            if DRIPs.capacity(agg) < 1e-2 # perturb the initial guess if solution is the zero capacity one
	                H0 = H0+rand(L).*(H-H0);
	            else # store the solution if it has positive capacity
	                H0 = H1;
	                if err < errmin
	                    global aggmin = agg;
	                    errmin = err;
	                end
	            end
	            iter += 1;
	    end
	    return(aggmin,min(err,errmin))
	end;
	println("Solving Mackowiack and Wiederholt (2009) ...")
	## Solve for κ = 3
	agg, err = agg_drip(2.4σq^2,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95)
	@test agg.ss.err < 1e-4
	agg, err = agg_drip(2.82σq^2,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95)
	@test agg.ss.err < 1e-4
	agg, err = agg_drip(0.01*σq^2,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95)
	@test agg.ss.err < 1e-4
	agg, err = agg_drip(10*σq^2,A,Qq,α,H; H0 = rand(L), maxit = 500, w = 0.95)
	@test agg.ss.err < 1e-4
	idi  = Drip(2.4σq^2,1,A,Qz,H,w = 0.9);
	@test idi.ss.err < 1e-4
end


using Plots
U(r,eps,sig) = 4*eps*((sig/r)^12 - (sig/r)^6)
# r vector between 0.01 and 15.
r = collect(0.4:0.01:2.5)
# compute U for every r
eps = 2. ; sig = 0.5 ;
u = [ U(x,eps,sig) for x in r ]
plot(r,u,label="",xlabel="r",ylabel="U(r)",
     ylim=[-2.5,10],framestyle=:box,linewidth=2,
     size=(400,300))
savefig("./cutoff1.pdf")


import AoC2019
using Test
task_commands = AoC2019.read_lines("input22.txt")
test_commands1 = """deal with increment 7
deal into new stack
deal into new stack
Result: 0 3 6 9 2 5 8 1 4 7"""
test_commands2 = """deal into new stack
cut -2
deal with increment 7
cut 8
cut -4
deal with increment 7
cut 3
deal with increment 9
deal with increment 3
cut -1"""

function do_commands_helper(commands; num_cards, track_card, do_reverse=false, iterations=1)::Array{Int128,1}
  current_num::Int128 = track_card
  operations = Int128[1,0]
  if do_reverse
    commands = reverse(commands)
  end
  for command in commands
    println(command)
    tokens = split(command)
    if tokens[1] == "deal"
      if tokens[2] == "into"
        current_num = (num_cards - 1) - current_num
        #current_num = mod(-current_num, num_cards)
        operations[2] = (num_cards -1) - operations[2]
        operations[1] *= -1
      elseif tokens[2] == "with"
        local N = parse(Int, tokens[4])
        if ! do_reverse
          current_num = mod(current_num * N, num_cards)
          operations .*= N
        else
          # find inverse to N mod num_cards
          N_inv = invmod(N, num_cards)
          current_num = mod(current_num * N_inv, num_cards)
          operations .*= N_inv
        end
      end
    elseif tokens[1] == "cut"
      local N = parse(Int, tokens[2])
      if ! do_reverse
        current_num = mod(current_num - N, num_cards)
        operations[2] -= N
      else
        current_num = mod(current_num + N, num_cards)
        operations[2] += N
      end
    end
    operations = mod.(operations, num_cards)
    println("operations: $operations")
    println("Current_num: $current_num")
    current_num_alt = mod(track_card * operations[1] + operations[2], num_cards)
    println("current_num_alt: $current_num_alt")
  end
  #println(operations)
  return operations
end

function fast_exp(kernel_operation::Array{Int128, 1}; iterations::Int64, modulo::Int64)::Array{Int128, 1}
  binary = bitstring(iterations)
  combined_operation::Array{Int128, 1} = kernel_operation
  new_operation::Array{Int128, 1} = Int128[1,0]
  for entry in reverse(digits(iterations, base=2))[2:end]
    new_operation = copy(combined_operation)
    new_operation .*= combined_operation[1]
    new_operation[2] += combined_operation[2]
    new_operation = mod.(new_operation, modulo)
    if entry == 1
      new_operation .*= kernel_operation[1]
      new_operation[2] += kernel_operation[2]
    end
    new_operation = mod.(new_operation, modulo)
    combined_operation = new_operation
  end
  return map(x -> Int128(x), combined_operation)
end

@test fast_exp(Int128[2,1], iterations=Int64(2), modulo=Int64(50)) == [4,3]
@test fast_exp(Int128[2,1], iterations=Int64(3), modulo=Int64(50)) == [8,7]
@test fast_exp(Int128[2,1], iterations=Int64(4), modulo=Int64(50)) == [16,15]
@test fast_exp(Int128[2,2], iterations=Int64(2), modulo=Int64(50)) == [4, 6]
@test fast_exp(Int128[2,2], iterations=Int64(4), modulo=Int64(50)) == [16, 30]
@test fast_exp(Int128[2,2], iterations=Int64(3), modulo=Int64(50)) == [8,14]

function do_commands(commands; num_cards, track_card, do_reverse=false, iterations=1)::Int128
  operations::Array{Int128, 1} = do_commands_helper(commands, num_cards=num_cards, track_card=track_card, do_reverse=do_reverse, iterations=iterations)
  operations = fast_exp(operations, iterations=iterations, modulo=num_cards)
  println(operations)
  return mod(track_card * operations[1] + operations[2], num_cards)
end

task_result = do_commands(task_commands, num_cards=10007, track_card=2019, do_reverse=false)
@test task_result == 7860

@test do_commands(split(test_commands1, '\n'), num_cards=10, track_card=7) == 9
@test do_commands(split(test_commands1, '\n'), num_cards=10, track_card=9, do_reverse=true) == 7

test_result1 = do_commands(split(test_commands2, '\n'), num_cards=10007, track_card=2019, do_reverse=false)
test_result2 = do_commands(split(test_commands2, '\n'), num_cards=10007, track_card=test_result1, do_reverse=false)
test_result2_alt = do_commands(split(test_commands2, '\n'), num_cards=10007, track_card=2019, do_reverse=false, iterations=2)
@test test_result2 == test_result2_alt

test_result3 = do_commands(split(test_commands2, '\n'), num_cards=10007, track_card=test_result2, do_reverse=false)
test_result3_alt = do_commands(split(test_commands2, '\n'), num_cards=10007, track_card=2019, do_reverse=false, iterations=3)
@test test_result3 == test_result3_alt

test_result1_back = do_commands(split(test_commands2, '\n'), num_cards=10007, track_card=test_result2, do_reverse=true)
test_result0_back = do_commands(split(test_commands2, '\n'), num_cards=10007, track_card=test_result1_back, do_reverse=true)
@test test_result0_back == 2019

iterations= 101741582076661
num_cards = 119315717514047

task_operations_1 = do_commands(task_commands, num_cards=119315717514047, track_card=2020, do_reverse=true, iterations=1)
task_operations_2 = do_commands(task_commands, num_cards=119315717514047, track_card=task_operations_1, do_reverse=true, iterations=1)
task_operations_2_alt = do_commands(task_commands, num_cards=119315717514047, track_card=2020, do_reverse=true, iterations=2)
@test task_operations_2 == task_operations_2_alt

task_operations_3 = do_commands(task_commands, num_cards=119315717514047, track_card=task_operations_2, do_reverse=true, iterations=1)
task_operations_3_alt = do_commands(task_commands, num_cards=119315717514047, track_card=2020, do_reverse=true, iterations=3)
@test task_operations_3 == task_operations_3_alt
task_operations_4 = do_commands(task_commands, num_cards=119315717514047, track_card=task_operations_3, do_reverse=true, iterations=1)
task_operations_4_alt = do_commands(task_commands, num_cards=119315717514047, track_card=2020, do_reverse=true, iterations=4)
@test task_operations_4 == task_operations_4_alt
task_operations_7 = do_commands(task_commands, num_cards=119315717514047, track_card=task_operations_4, do_reverse=true, iterations=3)
task_operations_7_alt = do_commands(task_commands, num_cards=119315717514047, track_card=2020, do_reverse=true, iterations=7)
@test task_operations_7 == task_operations_7_alt
task_operations_0_back = do_commands(task_commands, num_cards=119315717514047, track_card=task_operations_7, do_reverse=false, iterations=7)
@test task_operations_0_back == 2020

task_operations = do_commands(task_commands, num_cards=119315717514047, track_card=2020, do_reverse=true, iterations=iterations)
#fast_exp(task_operations, UInt64(iterations), UInt64(num_cards))
println(UInt64(task_operations))


function example_domain_H(parent = RealField(precision(BigFloat)))
    angles =
        fmpq.([
            3 // 2,
            1 // 2,
            1 // 2,
            1 // 2,
            1 // 2,
            3 // 2,
            3 // 2,
            1 // 2,
            1 // 2,
            1 // 2,
            1 // 2,
            3 // 2,
        ])
    vertices = [
        parent.(v) for v in [
            (0, 0),
            (0, 1),
            (-1, 1),
            (-1, -2),
            (0, -2),
            (0, -1),
            (1, -1),
            (1, -2),
            (2, -2),
            (2, 1),
            (1, 1),
            (1, 0),
        ]
    ]

    domain = Polygon(angles, vertices, parent)

    u1 = StandaloneVertexEigenfunction(vertex(domain, 1), fmpq(1 // 2), fmpq(3 // 2))
    u2 = StandaloneVertexEigenfunction(vertex(domain, 6), fmpq(0), fmpq(3 // 2))
    u3 = StandaloneVertexEigenfunction(vertex(domain, 7), fmpq(-1 // 2), fmpq(3 // 2))
    u4 = StandaloneVertexEigenfunction(vertex(domain, 12), fmpq(1), fmpq(3 // 2))

    u = CombinedEigenfunction(
        domain,
        [u1, u2, u3, u4],
        us_to_boundary = [
            setdiff(boundaries(domain), (12, 1)),
            setdiff(boundaries(domain), (5, 6)),
            setdiff(boundaries(domain), (6, 7)),
            setdiff(boundaries(domain), (11, 12)),
        ],
    )

    return domain, u
end

function example_domain_ngon(
    n,
    parent = RealField(precision(BigFloat));
    lightning = false,
    linked = false,
)
    θ = fmpq((n - 2) // n)
    angles = fill(θ, n)
    vertices = [(cos(θ), sin(θ)) for θ in (2parent(π) / n) .* (0:n-1)]
    domain = Polygon(angles, vertices, parent)

    if !lightning
        if linked
            us = [
                LinkedEigenfunction(
                    [
                        StandaloneVertexEigenfunction(
                            vertex(domain, i),
                            i * (1 - θ) + θ * 1 // 2,
                            θ,
                        ) for i in boundaries(domain)
                    ],
                    excluded_boundaries = [
                        BitSet([mod1(i - 1, n), i]) for i in boundaries(domain)
                    ],
                ),
            ]
            us_to_boundary = fill(BitSet([1]), length(us))
        else
            us = [
                StandaloneVertexEigenfunction(
                    vertex(domain, i),
                    i * (1 - θ) + θ * 1 // 2,
                    θ,
                ) for i in boundaries(domain)
            ]
            us_to_boundary = [
                setdiff(boundaries(domain), (i, mod1(i - 1, length(boundaries(domain))))) for i in eachindex(us)
            ]
        end
    else
        if linked
            us = [
                LinkedEigenfunction(
                    [
                        StandaloneLightningEigenfunction(
                            vertex(domain, i),
                            i * (1 - θ) + θ * 1 // 2,
                            θ,
                        ) for i in boundaries(domain)
                    ],
                    ones(n),
                ),
                StandaloneInteriorEigenfunction(domain, stride = n),
            ]

            us_to_boundary = fill(BitSet([1]), length(us))
        else
            us = [
                [
                    StandaloneLightningEigenfunction(
                        vertex(domain, i),
                        i * (1 - θ) + θ * 1 // 2,
                        θ,
                    ) for i in boundaries(domain)
                ]
                StandaloneInteriorEigenfunction(domain)
            ]

            us_to_boundary = fill(boundaries(domain), length(us))
        end
    end

    u = CombinedEigenfunction(domain, us, us_to_boundary = us_to_boundary)

    return domain, u
end

function example_domain_triangle_in_triangle(parent = RealField(precision(BigFloat)))
    domain1 = Triangle(fmpq(1 // 3), fmpq(1 // 3), parent)
    domain2 = TransformedDomain(
        Triangle(fmpq(1 // 3), fmpq(1 // 3), parent),
        fmpq(0),
        parent(0.25),
        SVector(cospi(fmpq(1 // 6), parent), sinpi(fmpq(1 // 6), parent)) / 4,
    )

    domain = IntersectedDomain(domain1, domain2)

    u1 = StandaloneLightningEigenfunction(domain1, 1)

    u2 = LinkedEigenfunction([StandaloneLightningEigenfunction(domain1, i) for i = 2:3])

    u3 = StandaloneLightningEigenfunction(domain2, 1, l = parent(0.15), outside = true)

    u4 = LinkedEigenfunction([
        StandaloneLightningEigenfunction(domain2, i, l = parent(0.15), outside = true)
        for i = 2:3
    ])

    u5 = StandaloneInteriorEigenfunction(domain)

    us = [u1, u2, u3, u4, u5]
    #us = vcat([u1], u2.us, [u3], u4.us, [u5])
    orders = [1, 1, 1, 1, 1]
    us_to_boundary = fill(1:6, length(us))

    u = CombinedEigenfunction(domain, us, orders, us_to_boundary = us_to_boundary)

    return domain, u
end

function example_domain_ngon_in_ngon(
    n1,
    n2,
    parent = RealField(precision(BigFloat));
    linked = true,
    even = true,
    T = arb,
)
    θ1 = fmpq((n1 - 2) // n1)
    angles1 = fill(θ1, n1)
    vertices1 = [(cos(θ), sin(θ)) for θ in (2parent(π) / n1) .* (0:n1-1)]
    domain1 = Polygon(angles1, vertices1, parent)

    θ2 = fmpq((n2 - 2) // n2)
    angles2 = fill(θ2, n2)
    vertices2 = [(cos(θ), sin(θ)) for θ in (2parent(π) / n2) .* (0:n2-1)]
    domain2 = TransformedDomain(
        Polygon(angles2, vertices2, parent),
        fmpq(0),
        parent(0.5),
        SVector(parent(0), parent(0)),
    )

    domain = IntersectedDomain(domain1, domain2)

    if linked
        us = [
            LinkedEigenfunction([
                StandaloneLightningEigenfunction{T,arb}(domain1, i; even) for
                i in boundaries(domain1)
            ]),
            LinkedEigenfunction([
                StandaloneLightningEigenfunction{T,arb}(
                    domain2,
                    i,
                    outside = true,
                    l = parent(0.4);
                    even,
                ) for i in boundaries(domain2)
            ]),
            StandaloneInteriorEigenfunction(domain, stride = 4; even),
        ]

        orders = ifelse(even, [2, 2, 1], [3, 3, 2])
        even_boundaries = ifelse(even, [1, 5], [])

        u = CombinedEigenfunction(
            domain,
            us,
            orders,
            us_to_boundary = fill(BitSet([1, 5]), length(us));
            even_boundaries,
        )
    else
        us = vcat(
            [
                StandaloneLightningEigenfunction{T,arb}(domain1, i) for
                i in boundaries(domain1)
            ],
            [
                StandaloneLightningEigenfunction{T,arb}(
                    domain2,
                    i,
                    outside = true,
                    l = parent(0.4),
                ) for i in boundaries(domain2)
            ],
            [StandaloneInteriorEigenfunction(domain)],
        )

        u = CombinedEigenfunction(domain, us, [3, 3, 3, 3, 3, 3, 3, 3, 2])
    end

    return domain, u
end

"""
    example_domain_goal_v1(N, d, h, parent; ...)

Return the domain given by a hexagon with 6 triangles cut out as well
as a corresponding eigenfunction.

The placement of the holes is determined by `(N, d, h)`.

The keyword `symmetry_class` determines which symmetries are assumed
by the eigenfunction. `symmetry_class = 1` assumes it's six-fold
symmetric and even, this is satisfied by the first and second
eigenvalue. `symmetry_class = 2` assumes it's even w.r.t. the x-axis
and odd w.r.t. the y-axis, `symmetry_class = 3` does the opposite,
this is used by the third and fourth eigenvalue.

Computing the first few eigenvalues in Matlab using a finite element
method gives (with `(N, d, h) = (22, 9, 6)`)
```
   31.5943
   75.8995
   75.9633
   75.9633
```
though it's likely not accurate to more than 1 or 2 decimals. If `(N,
d, h) = (27, 11, 6)` then they are
```
   31.0432
   63.2104
   63.7259
   63.7259
   68.2629
   68.2629
```
For the second eigenvalue the value `63.2083359862688427940` might be
slightly more accurate.
"""
function example_domain_goal_v1(
    N::Integer = 27,
    d::Integer = 11,
    h::Integer = 6,
    parent = RealField(precision(BigFloat));
    T = arb,
    lightning = false,
    inner_expansion = true,
    outer_expansion = false,
    symmetry_class = 1,
    even = true,
    reversed = true,
)
    ### Compute the domain ###
    # The main domain is a hexagon
    n = 6
    θ = fmpq((n - 2) // n)
    angles = fill(θ, n)
    vertices = [(cospi(θ, parent), sinpi(θ, parent)) for θ in fmpq(2 // n) .* (0:n-1)]
    exterior = Polygon(angles, vertices, parent)

    # The interior domains are triangles
    @assert mod(h, 3) == 0
    points = [
        SVector(parent(d // N), -parent(h // N) / sqrt(parent(3))),
        SVector(parent((d + h) // N), parent(0)),
        SVector(parent(d // N), parent(h // N) / sqrt(parent(3))),
    ]

    interior_angles = [1 // 3, 1 // 3, 1 // 3]

    interiors = [
        TransformedDomain(Polygon(interior_angles, points, parent), i // 3, 1, [0, 0])
        for i = 0:5
    ]

    domain = IntersectedDomain(exterior, interiors)

    ### Compute the eigenfunction ###
    # Expansions from the vertices of the hexagon
    if symmetry_class == 1
        if lightning
            outer_eigenfunction =
                i -> StandaloneLightningEigenfunction(
                    vertex(exterior, i),
                    fmpq(mod(Rational(1 - θ // 2 + (i - 1) * (1 - θ)), 2)),
                    θ;
                    even,
                )
            outer_excluded_boundary = i -> BitSet()
            order1 = ifelse(even, 2, 3)

        else
            outer_eigenfunction =
                i -> StandaloneVertexEigenfunction(
                    vertex(exterior, i),
                    i * (1 - θ) + θ * 1 // 2,
                    θ,
                    stride = ifelse(even, 2, 1),
                )
            outer_excluded_boundary =
                i -> BitSet([mod1(i - 1, length(boundaries(exterior))), i])
            order1 = ifelse(even, 1, 2)
        end

        u1 = LinkedEigenfunction(
            [outer_eigenfunction(i) for i in boundaries(exterior)],
            excluded_boundaries = [
                outer_excluded_boundary(i) for i in boundaries(exterior)
            ],
        )

        # Expansions from the outer tip of the triangles
        u2 = LinkedEigenfunction([
            StandaloneLightningEigenfunction{T,fmpq}(
                vertex(d, 2),
                7 // 6 + d.rotation,
                2 - d.original.angles[2],
                l = parent(h // 2N),
                even = even,
            ) for d in interiors
        ])
        if even && reversed
            order2 = 3 * 2
        elseif even
            order2 = 2
        else
            order2 = 3
        end

        # Expansions from the inner tips of the triangles
        u3 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, i),
                    ifelse(i == 1, 1 // 2, -1 // 6) + d.rotation,
                    2 - d.original.angles[i],
                    l = parent(h // 2N),
                    even = even && !reversed,
                    reversed = reversed && i == 3,
                ) for d in interiors, i in [1, 3]
            ][:],
        )
        if even && reversed
            order3 = 3 * 3
        elseif even
            order3 = 2
        else
            order3 = 3
        end

        # Expansion from the center
        u4 = StandaloneInteriorEigenfunction(SVector(domain.parent(0), domain.parent(0)), stride = 6; even)
        order4 = ifelse(even, 1, 2)

        # Interior expansions closer to outer boundaries
        u5 = LinkedEigenfunction([
            StandaloneInteriorEigenfunction(
                parent(3 // 4) .*
                SVector(cospi(fmpq(i // 3), parent), sinpi(fmpq(i // 3), parent)),
                fmpq(i // 3),
                stride = 6,
                even = true,
            ) for i = 0:5
        ])
        order5 = ifelse(even, 1, 2)

        us = AbstractPlanarEigenfunction[u1, u2, u3]
        orders = Int[order1, order2, order3]

        if inner_expansion
            push!(us, u4)
            push!(orders, order4)
        end
        if outer_expansion
            push!(us, u5)
            push!(orders, order5)
        end

        us_to_boundary = fill(BitSet([1, 8, 9]), length(us))
        even_boundaries = ifelse(even, Int[1, 9], Int[])
    elseif symmetry_class == 2
        ####################################################################
        is_even = i -> i == 1 || i == 4
        is_reversed = i -> i == 3 || i == 6

        if lightning
            outer_eigenfunction =
                i -> StandaloneLightningEigenfunction(
                    vertex(exterior, i),
                    fmpq(mod(Rational(1 - θ // 2 + (i - 1) * (1 - θ)), 2)),
                    θ;
                    even = is_even(i),
                    reversed = is_reversed(i),
                )
            outer_excluded_boundary = i -> BitSet()
            order11 = 2
            order12 = 3
        else
            outer_eigenfunction =
                i -> StandaloneVertexEigenfunction(
                    vertex(exterior, i),
                    i * (1 - θ) + θ * 1 // 2,
                    θ,
                    stride = ifelse(is_even(i), 2, 1),
                    reversed = is_reversed(i),
                )
            outer_excluded_boundary =
                i -> BitSet([mod1(i - 1, length(boundaries(exterior))), i])
            order11 = 1
            order12 = 1
        end

        u11 = LinkedEigenfunction(
            [outer_eigenfunction(i) for i in (1, 4)],
            [parent(1), parent(-1)],
            excluded_boundaries = [outer_excluded_boundary(i) for i in (1, 4)],
        )

        u12 = LinkedEigenfunction(
            [outer_eigenfunction(i) for i in (2, 3, 5, 6)],
            [parent(1), parent(-1), parent(-1), parent(1)],
            excluded_boundaries = [outer_excluded_boundary(i) for i in (2, 3, 5, 6)],
        )

        # Expansions from the outer tip of the triangles
        u21 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, 2),
                    7 // 6 + d.rotation,
                    2 - d.original.angles[2],
                    l = parent(h // 2N),
                    even = is_even(i),
                ) for (i, d) in ((1, interiors[1]), (4, interiors[4]))
            ],
            [parent(1), parent(-1)],
        )
        order21 = 2

        u22 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, 2),
                    7 // 6 + d.rotation,
                    2 - d.original.angles[2],
                    l = parent(h // 2N),
                    even = is_even(i),
                    reversed = is_reversed(i),
                ) for (i, d) in (
                    (2, interiors[2]),
                    (3, interiors[3]),
                    (5, interiors[5]),
                    (6, interiors[6]),
                )
            ],
            [parent(1), parent(-1), parent(-1), parent(1)],
        )
        order22 = 3

        # Expansions from the inner tips of the triangles
        u31 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, i),
                    ifelse(i == 1, 1 // 2, -1 // 6) + d.rotation,
                    2 - d.original.angles[i],
                    l = parent(h // 2N),
                    reversed = i == 3,
                ) for (j, d) in ((1, interiors[1]), (4, interiors[4])), i in [1, 3]
            ][:],
            [parent(1), parent(-1), parent(1), parent(-1)],
        )
        order31 = 3

        u32 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, i),
                    ifelse(i == 1, 1 // 2, -1 // 6) + d.rotation,
                    2 - d.original.angles[i],
                    l = parent(h // 2N),
                    reversed = i == 3,
                ) for (j, d, i) in (
                    (2, interiors[2], 1),
                    (3, interiors[3], 3),
                    (5, interiors[5], 1),
                    (6, interiors[6], 3),
                )
            ],
            [parent(1), parent(-1), parent(-1), parent(1)],
        )
        order32 = 3

        u33 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, i),
                    ifelse(i == 1, 1 // 2, -1 // 6) + d.rotation,
                    2 - d.original.angles[i],
                    l = parent(h // 2N),
                    reversed = i == 3,
                ) for (j, d, i) in (
                    (2, interiors[2], 3),
                    (3, interiors[3], 1),
                    (5, interiors[5], 3),
                    (6, interiors[6], 1),
                )
            ],
            [parent(1), parent(-1), parent(-1), parent(1)],
        )
        order33 = 3

        # Expansion from the center
        u4 = StandaloneInteriorEigenfunction(domain, stride = 2, offset = 1; even)
        order4 = 1

        us = AbstractPlanarEigenfunction[u11, u12, u21, u22, u31, u32, u33, u4]
        orders =
            Int[order11, order12, 3order21, 3order22, 3order31, 3order32, 3order33, order4]

        us_to_boundary = fill(BitSet([1, 2, 8, 9, 10, 11, 12]), length(us))
        even_boundaries = Int[2, 9]
    elseif symmetry_class == 3
        ####################################################################
        is_odd = i -> i == 1 || i == 4
        is_reversed = i -> i == 3 || i == 4 || i == 6

        if lightning
            outer_eigenfunction =
                i -> StandaloneLightningEigenfunction(
                    vertex(exterior, i),
                    fmpq(mod(Rational(1 - θ // 2 + (i - 1) * (1 - θ)), 2)),
                    θ;
                    odd = is_odd(i),
                    reversed = is_reversed(i),
                )
            outer_excluded_boundary = i -> BitSet()
            order11 = 1
            order12 = 3
        else
            outer_eigenfunction =
                i -> StandaloneVertexEigenfunction(
                    vertex(exterior, i),
                    i * (1 - θ) + θ * 1 // 2,
                    θ,
                    stride = ifelse(is_odd(i), 2, 1),
                    offset = ifelse(is_odd(i), 1, 0),
                    reversed = is_reversed(i),
                )
            outer_excluded_boundary =
                i -> BitSet([mod1(i - 1, length(boundaries(exterior))), i])
            order11 = 1
            order12 = 1
        end

        u11 = LinkedEigenfunction(
            [outer_eigenfunction(i) for i in (1, 4)],
            [parent(1), parent(1)],
            excluded_boundaries = [outer_excluded_boundary(i) for i in (1, 4)],
        )

        u12 = LinkedEigenfunction(
            [outer_eigenfunction(i) for i in (2, 3, 5, 6)],
            [parent(1), parent(1), parent(-1), parent(-1)],
            excluded_boundaries = [outer_excluded_boundary(i) for i in (2, 3, 5, 6)],
        )

        # Expansions from the outer tip of the triangles
        u21 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, 2),
                    7 // 6 + d.rotation,
                    2 - d.original.angles[2],
                    l = parent(h // 2N),
                    odd = is_odd(i),
                    reversed = i == 4,
                ) for (i, d) in ((1, interiors[1]), (4, interiors[4]))
            ],
            [parent(1), parent(1)],
        )
        order21 = 1

        u22 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, 2),
                    7 // 6 + d.rotation,
                    2 - d.original.angles[2],
                    l = parent(h // 2N),
                    odd = is_odd(i),
                    reversed = is_reversed(i),
                ) for (i, d) in (
                    (2, interiors[2]),
                    (3, interiors[3]),
                    (5, interiors[5]),
                    (6, interiors[6]),
                )
            ],
            [parent(1), parent(1), parent(-1), parent(-1)],
        )
        order22 = 3

        # Expansions from the inner tips of the triangles
        u31 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, i),
                    ifelse(i == 1, 1 // 2, -1 // 6) + d.rotation,
                    2 - d.original.angles[i],
                    l = parent(h // 2N),
                    reversed = i == 3,
                ) for (j, d) in ((1, interiors[1]), (4, interiors[4])), i in [1, 3]
            ][:],
            [parent(1), parent(-1), parent(-1), parent(1)],
        )
        order31 = 3

        u32 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, i),
                    ifelse(i == 1, 1 // 2, -1 // 6) + d.rotation,
                    2 - d.original.angles[i],
                    l = parent(h // 2N),
                    reversed = i == 3,
                ) for (j, d, i) in (
                    (2, interiors[2], 1),
                    (3, interiors[3], 3),
                    (5, interiors[5], 1),
                    (6, interiors[6], 3),
                )
            ],
            [parent(1), parent(1), parent(-1), parent(-1)],
        )
        order32 = 3

        u33 = LinkedEigenfunction(
            [
                StandaloneLightningEigenfunction{T,fmpq}(
                    vertex(d, i),
                    ifelse(i == 1, 1 // 2, -1 // 6) + d.rotation,
                    2 - d.original.angles[i],
                    l = parent(h // 2N),
                    reversed = i == 3,
                ) for (j, d, i) in (
                    (2, interiors[2], 3),
                    (3, interiors[3], 1),
                    (5, interiors[5], 3),
                    (6, interiors[6], 1),
                )
            ],
            [parent(1), parent(1), parent(-1), parent(-1)],
        )
        order33 = 3

        # Expansion from the center
        u4 = StandaloneInteriorEigenfunction(domain, stride = 2, odd = true)
        order4 = 1

        us = AbstractPlanarEigenfunction[u11, u12, u21, u22, u31, u32, u33, u4]

        orders =
            Int[order11, order12, 3order21, 3order22, 3order31, 3order32, 3order33, order4]

        us_to_boundary = fill(BitSet([1, 2, 8, 9, 10, 11, 12]), length(us))
        even_boundaries = Int[2, 9]
    else
        throw(ArgumentError("symmetry_class should be 1 or 2, got $symmetry_class"))
    end

    u = CombinedEigenfunction(domain, us, orders; us_to_boundary, even_boundaries)

    return domain, u
end


include("C:\\Users\\darien.shannon\\Documents\\Code\\Julia\\JuliaSE\\src\\2DBeam.jl")

COORD = [Node(0, 0)
         Node(6, 0)
         Node(10, 0)
         Node(20, 0)]

MSUP = [Support(1, 1, 1)
        Support(2, 0, 0)
        Support(3, 1, 0)
        Support(4, 1, 0)]

EM = Dict( 1 => Material(28e6))

CP = Dict( 1 => Section(5.8e9/(1000^4)),
           2 => Section(1.5*5.8e9/(1000^4)))

MPRP = [Element(1, 2, 1, 2)
        Element(2, 3, 1, 1)
        Element(3, 4, 1, 1)]

memberLoads = [(1, DLoad(30, 0, 0, 6))
               (3, CLoad(150, 5))]

nodeLoads = [(2, NLoad((-200, 0)))
             (3, NLoad((0, -90)))]

forces, reactions = run_analysis(COORD, MSUP, EM, CP, MPRP, memberLoads, nodeLoads)
println(forces)
println(reactions)


module ThinWallBCTest

using Gridap
using Test
using LineSearches: BackTracking

using GridapMHD: conductive_thin_wall
using GridapMHD: analytical_shercliff_u
using GridapMHD: analytical_shercliff_j
using GridapMHD: compute_u_j_errors

# Problem setting
ρ = 1.0
ν = 1.0
σ = 1.0
wall_thickness = 0.10
wall_conductivity = 1.0
c_w = wall_conductivity * wall_thickness/σ

U0 = 10.0
B0 = 10.0
L = 1.0

Re = U0 * L / ν
Ha = B0 * L * sqrt(σ/(ρ*ν))
K = Ha / (1-0.825*Ha^(-1/2)-Ha^(-1))
∂p∂z = -Re * K / L^3

f_u = VectorValue(0.0,0.0, -∂p∂z) * L/U0^2

xh, trian, dΩ = conductive_thin_wall(nx=5,ny=5,c_w=c_w,Re=Re,Ha=Ha,α=10.0,f_u=f_u)

uh, ph, jh, φh = xh
divj = (∇⋅jh)

# Scale unknowns
uh = uh * U0
jh = jh * σ * B0 * U0

# Analyical solutions
side_wall_semilength = 1.0
hartmann_wall_semilength = 1.0
number_fourier_sumands = 20
K = Ha / (1-0.825*Ha^(-1/2)-Ha^(-1))
∂p∂z = -Re * K / L^3
u0(x) = analytical_shercliff_u(side_wall_semilength, hartmann_wall_semilength,
  wall_thickness, wall_conductivity, σ, ν*ρ, ∂p∂z, Ha, number_fourier_sumands, x)

j0(x) = analytical_shercliff_j(side_wall_semilength, hartmann_wall_semilength,
  wall_thickness, wall_conductivity, σ, ν*ρ, ∂p∂z, Ha, number_fourier_sumands, x)

writevtk(trian, "results.vtu",
  cellfields=["uh"=>uh, "ph"=>ph, "jh"=>jh, "φh"=>φh, "u"=>u0, "j"=>j0])

eu_l2, ej_l2 = compute_u_j_errors(uh, jh, u0, j0, dΩ)
e_divj = sqrt(sum(∫(divj*divj)*dΩ))

@test eu_l2 < 0.3
@test ej_l2 < 2
@test e_divj < 1e-12

end #module


#####
##### 3D modes on pyramids
#####

"""
    basis(elem::Pyr,N,r,s,t,tol=1e-12)

Computes orthonormal semi-nodal basis on the biunit pyramid element.

Warning: nodal derivative matrices may contain errors for nodes at t = 1.
A way to avoid this is to use weak differentiation matrices computed using
quadrature rules with only interior nodes.
"""
function basis(elem::Pyr, N, r, s, t, tol = 1e3*eps())

    # convert to abc
    a = @. 2*(r+1)/(1-t)-1
    b = @. 2*(s+1)/(1-t)-1
    c = t

    # fix top point
    ids = @. abs(1 - t) < tol
    t[ids] .= 1 - tol # WARNING: hacky (note factor of (1-c) = (1-t) in "pc" below)
    a[ids] .= -1
    b[ids] .= -1

    # change of vars from a to b
    dadr = @. 2 / (1-t)
    dbds = @. 2 / (1-t)
    dadt = @. (2 * r + 2) / (1 - t)^2
    dbdt = @. (2 * s + 2) / (1 - t)^2
    dcdt = 1

    Np = (N + 1) * (N + 2) * (2 * N + 3) ÷ 6
    V, Vr, Vs, Vt = ntuple(x->zeros(length(r), Np), 4)

    ind = 1
    for k = 0:N
        # make nodal quad basis
        bq, aq, wab = quad_nodes(Quad(), k)
        VDM, _ = basis(Quad(), k, aq, bq)
        VDMab, Va, Vb = basis(Quad(), k, a, b)
        Vab, DVa, DVb = map(A->A / VDM, (VDMab, Va, Vb))

        CNk = (N+2) / (2^(2*k+2) * (2*k+3))
        pNk = jacobiP(c, 2*k+3, 0, N-k)
        pc = @. ((1-c) / 2)^k * pNk
        dpNk = grad_jacobiP(c, 2*k+3, 0, N-k)
        dpc = @. ((1-c)/2)^k * dpNk - (k/2^k) * ((1 - c)^(k - 1)) * pNk

        for ij = 1:(k+1)^2
            scale = 1 / sqrt(CNk * wab[ij]) # normalization
            @. V[:,ind]  = scale * Vab[:,ij] * pc
            @. Vr[:,ind] = scale * (DVa[:,ij] * dadr) * pc
            @. Vs[:,ind] = scale * (DVb[:,ij] * dbds) * pc
            @. Vt[:,ind] = scale * ((DVa[:,ij] * dadt + DVb[:,ij] * dbdt) * pc + Vab[:,ij] * dpc * dcdt)
            ind += 1
        end
    end
    return V, Vr, Vs, Vt
end

"""
    abctorst(elem::Pyr,a,b,c)

Converts from Stroud coordinates (a,b,c) on [-1,1]^3 to reference element
coordinates (r,s,t).
"""
function abctorst(elem::Pyr,a,b,c)
    r = @. .5 * (1+a) * (1-c) - 1
    s = @. .5 * (1+b) * (1-c) - 1
    t = @. c
    return r, s, t
end

"""
    nodes(elem::Pyr,N)

Computes interpolation nodes of degree N. Edge nodes coincide with (N+1)-point Lobatto
points. Triangular face nodes coincide with Tri.nodes(N), quadrilateral face nodes
coincide with tensor product (N+1)-point Lobatto points.
"""
function nodes(elem::Pyr, N)

    if N == 1
        return equi_nodes(Pyr(), N)
    end

    # append quad face nodes
    r1D_equi = equi_nodes(Line(), N)
    rst_equi = interp_1D_to_edges(Pyr(), r1D_equi)
    quad_face_pts_equi = (vec.(meshgrid(r1D_equi[2:end-1]))..., -ones((N - 1) * (N - 1)))
    append!.(rst_equi, quad_face_pts_equi)

    # append quad face nodes
    r1D = nodes(Line(), N)
    rst_lobatto = interp_1D_to_edges(Pyr(), r1D)
    quad_face_pts = (vec.(meshgrid(r1D[2:end-1]))..., -ones((N - 1) * (N - 1)))
    append!.(rst_lobatto, quad_face_pts)

    # append quad face "bubble" basis functions
    function edge_face_pyr_basis(N, rst)
        V_edge = edge_basis(Pyr(), N, rst...)
        # WARNING: assumes first 4 vertices of Pyr = quad vertices
        # also assumes first 5 cols of V_edge = vertex functions for Pyr
        V_quad_face = zeros(length(first(rst)), (N - 1) * (N - 1))
        V_quad = vandermonde(Quad(), N-2, rst[1], rst[2])
        V_quad_bubble = @. V_edge[:,1] * V_edge[:,2] * V_edge[:,3] * V_edge[:,4]
        for i = 1:size(V_quad_face, 2)
            @. V_quad_face[:,i] = V_quad_bubble * V_quad[:,i]
        end
        return hcat(V_edge, V_quad_face)
    end

    # same interp problem as before
    V_edge_face = edge_face_pyr_basis(N, rst_equi)
    c = (x->V_edge_face \ x).(rst_lobatto) # should be lsq solve
    return (x->edge_face_pyr_basis(N, equi_nodes(elem, N)) * x).(c)
end

function equi_nodes(elem::Pyr, N)
    # Np = (N+1)*(N+2)*(2*N+3)/6
    a, b, c = ntuple(x->Float64[],3)
    c1D = LinRange(-1, 1, N+1)
    for k = N:-1:0
        if k==0
            r1D = [.5]
        else
            r1D = LinRange(-1, 1, k + 1)
        end
        ak, bk = vec.(meshgrid(r1D, r1D))
        ck = c1D[N + 1 - k]*one.(ak)
        append!.((a, b, c),(ak, bk, ck))
    end
    return abctorst(elem, a, b, c)
end

quad_nodes(elem::Pyr, N) = stroud_quad_nodes(elem, N)

function stroud_quad_nodes(elem::Pyr, N)

    a1D, w1D = gauss_quad(0, 0, N)
    c1D, wc1D = gauss_quad(2, 0, N)

    a, c, b = vec.(meshgrid(a1D, c1D, a1D))
    wa, wc, wb = vec.(meshgrid(w1D, wc1D, w1D))
    w = @. wa * wb * wc
    w = (8/3) * w ./ sum(w) # scale by b*h/3 = volume of pyr. b = 4, h = 2 for biunit

    return abctorst(elem, a, b, c)..., w
end


# ----------
# MHD Analysis Method, providing MHD related quantities function
# ----------

# Scale Decomposition FUnction
function ScaleDecomposition(B1::Array,B2::Array,B3::Array;kf=[1,5],Lx = 2π)
    k1,k2 = minimum(kf),maximum(kf);
    nx,ny,nz = size(B1);
    T    = Float32;
    grid = ThreeDGrid(nx, Lx, T = T);

    B1h = zeros(ComplexF32,(div(nx,2)+1,ny,nz));
    B2h = zeros(ComplexF32,(div(nx,2)+1,ny,nz));
    B3h = zeros(ComplexF32,(div(nx,2)+1,ny,nz));
    Bxhf = copy(B1h);
    Byhf = copy(B2h);
    Bzhf = copy(B3h);
    mul!(B1h, grid.rfftplan, B1);
    mul!(B2h, grid.rfftplan, B2);
    mul!(B3h, grid.rfftplan, B3);

    for i in 1:div(nz,2)+1, j in 1:ny, k in 1:nx
       x,y,z = grid.kr[i],grid.l[j],grid.m[k];
       rr    = sqrt(x^2+y^2+z^2);
       if (( rr >=  k1) && (rr <= k2))
           Bxhf[i,j,k] = B1h[i,j,k];
           Byhf[i,j,k] = B2h[i,j,k];
           Bzhf[i,j,k] = B3h[i,j,k];
       end
    end
    cB1,cB2,cB3 = zeros(T,size(B1)),zeros(T,size(B1)),zeros(T,size(B1));
    ldiv!(cB1, grid.rfftplan,Bxhf);
    ldiv!(cB2, grid.rfftplan,Byhf);
    ldiv!(cB3, grid.rfftplan,Bzhf);
    return cB1,cB2,cB3;
end

#Kenitic Helicity
function h_k(iv,jv,kv;L=2π)
# V ⋅ ( ∇ × V )
	dlx,dly,dlz  = (L/size(iv)[1]),(L/size(iv)[2]),(L/size(iv)[3]);
	dV = dlx*dly*dlz;
	cV1,cV2,cV3 = Curl(iv,jv,kv;Lx=L);
	h_k = @. (cV1*iv + cV2*jv + cV3*kv)*dV
	return h_k
end

#Megnetic Helicity
function h_m(ib,jb,kb)
# A ⋅ B
   Ax,Ay,Az = VectorPotential(ib,jb,kb);
   return Ax.*ib + Ay.*jb + Az.*kb;
end

#Vector Potential
function VectorPotential(B1,B2,B3;L=2π)
#=
    funtion of computing B = ∇ × A using the fourier method
     fft(∇×Vector) -> im * k × A
      | i j k  |
      | x y z  |  =  (y*A3 - z*A2) i - (x*A3 - z*A1) j + (x*A2 - y*A1) k
      |A1 A2 A3|

    Note: We are using the Coulomb gauge ∇ ⋅ A = 0
    Using the relations J = ∇ × (∇ × B) and ∇ ⋅ A = 0 in peroideric condition,
    we will arrive Jₖ = - k² Aₖ.

    Aₖ = (k × B)ᵢ/ k²
=#
    nx,ny,nz = size(B1);
    T    = Float32;
    grid = ThreeDGrid(nx, L, T = T);

    B1h = zeros(ComplexF32,(div(nx,2)+1,ny,nz));
    B2h = zeros(ComplexF32,(div(nx,2)+1,ny,nz));
    B3h = zeros(ComplexF32,(div(nx,2)+1,ny,nz));
    Axh = copy(B1h);
    Ayh = copy(B2h);
    Azh = copy(B3h);
    mul!(B1h, grid.rfftplan, B1);
    mul!(B2h, grid.rfftplan, B2);
    mul!(B3h, grid.rfftplan, B3);

    for i in 1:div(nz,2)+1, j in 1:ny, k in 1:nx
       x,y,z = grid.kr[i],grid.l[j],grid.m[k];
       k²    = x^2 + y^2 + z^2;
       Axh[i,j,k] = im*(y*B3h[i,j,k] - z*B2h[i,j,k])/k²;
       Ayh[i,j,k] = im*(z*B1h[i,j,k] - x*B3h[i,j,k])/k²;
       Azh[i,j,k] = im*(x*B2h[i,j,k] - y*B1h[i,j,k])/k²;
       Axh[i,j,k] = ifelse(k² == 0, 0,  Axh[i,j,k]);
       Ayh[i,j,k] = ifelse(k² == 0, 0,  Ayh[i,j,k]);
       Azh[i,j,k] = ifelse(k² == 0, 0,  Azh[i,j,k]);
    end

    A1,A2,A3 = zeros(T,size(B1)),zeros(T,size(B1)),zeros(T,size(B1));
    ldiv!(A1, grid.rfftplan, deepcopy(Axh));
    ldiv!(A2, grid.rfftplan, deepcopy(Ayh));
    ldiv!(A3, grid.rfftplan, deepcopy(Azh));
    return A1,A2,A3;

end

#Checking for anagular momentum using center point as a reference (r = 0 at the center)
function getL(iv,jv,kv;L=2π)
    nx,ny,nz = size(iv);
    T    = Float32;
    grid = ThreeDGrid(nx, L, T = T);
    Lᵢ,Lⱼ,Lₖ  = getL(iv,jv,kv,grid)
    return Lᵢ,Lⱼ,Lₖ
end

function getL(iv,jv,kv,grid)
    nx,ny,nz = size(iv);
    T    = Float32;
    Lᵢ,Lⱼ,Lₖ = zeros(nx,ny,nz),zeros(nx,ny,nz),zeros(nx,ny,nz);
    for k in 1:nz, j in 1:ny, i in 1:nx
       x,y,z = grid.x[i],grid.y[j],grid.z[k];
       r = sqrt(x^2 + y^2);
       # L = r × p => iv,jv,kv
       # |  i  j  k |
       # |  x  y  z | = (y*kv - z*jv) i - (x*kv - z*iv) j + (x*jv - y*iv) k;
       # | iv jv kv |
       Lᵢ[i,j,k] =   (y*kv[i,j,k] - z*jv[i,j,k]);
       Lⱼ[i,j,k] =  -(x*kv[i,j,k] - z*iv[i,j,k]);
       Lₖ[i,j,k] =   (x*jv[i,j,k] - y*iv[i,j,k]);
    end
    return Lᵢ,Lⱼ,Lₖ
end


using Logging, Random
using JuMP, Gurobi
using DecisionProgramming

Random.seed!(42)

const dᴾ = 1    # Decision node: range for number of patents
const cᵀ = 2    # Chance node:   technical competitiveness
const dᴬ = 3    # Decision node: range for number of applications
const cᴹ = 4    # Chance node:   market share
const DP_states = ["0-3 patents", "3-6 patents", "6-9 patents"]
const CT_states = ["low", "medium", "high"]
const DA_states = ["0-5 applications", "5-10 applications", "10-15 applications"]
const CM_states = ["low", "medium", "high"]

S = States([
    (length(DP_states), [dᴾ]),
    (length(CT_states), [cᵀ]),
    (length(DA_states), [dᴬ]),
    (length(CM_states), [cᴹ]),
])
C = Vector{ChanceNode}()
D = Vector{DecisionNode}()
V = Vector{ValueNode}()
X = Vector{Probabilities}()
Y = Vector{Consequences}()

I_DP = Vector{Node}()
push!(D, DecisionNode(dᴾ, I_DP))

I_CT = [dᴾ]
X_CT = zeros(S[dᴾ], S[cᵀ])
X_CT[1, :] = [1/2, 1/3, 1/6]
X_CT[2, :] = [1/3, 1/3, 1/3]
X_CT[3, :] = [1/6, 1/3, 1/2]
push!(C, ChanceNode(cᵀ, I_CT))
push!(X, Probabilities(cᵀ, X_CT))

I_DA = [dᴾ, cᵀ]
push!(D, DecisionNode(dᴬ, I_DA))

I_CM = [cᵀ, dᴬ]
X_CM = zeros(S[cᵀ], S[dᴬ], S[cᴹ])
X_CM[1, 1, :] = [2/3, 1/4, 1/12]
X_CM[1, 2, :] = [1/2, 1/3, 1/6]
X_CM[1, 3, :] = [1/3, 1/3, 1/3]
X_CM[2, 1, :] = [1/2, 1/3, 1/6]
X_CM[2, 2, :] = [1/3, 1/3, 1/3]
X_CM[2, 3, :] = [1/6, 1/3, 1/2]
X_CM[3, 1, :] = [1/3, 1/3, 1/3]
X_CM[3, 2, :] = [1/6, 1/3, 1/2]
X_CM[3, 3, :] = [1/12, 1/4, 2/3]
push!(C, ChanceNode(cᴹ, I_CM))
push!(X, Probabilities(cᴹ, X_CM))

# Dummy value node
push!(V, ValueNode(5, [cᴹ]))
push!(Y, Consequences(5, zeros(S[cᴹ])))

@info("Validate influence diagram.")
validate_influence_diagram(S, C, D, V)
sort!.((C, D, V, X, Y), by = x -> x.j)

@info("Creating path probability.")
P = DefaultPathProbability(C, X)

@info("Defining DecisionModel")
model = Model()
z = DecisionVariables(model, S, D)

@info("Creating problem specific constraints and expressions")

function variables(model::Model, dims::AbstractVector{Int}; binary::Bool=false)
    v = Array{VariableRef}(undef, dims...)
    for i in eachindex(v)
        v[i] = @variable(model, binary=binary)
    end
    return v
end

n_T = 5                     # number of technology projects
n_A = 5                     # number of application projects
I_t = rand(n_T)*0.1         # costs of technology projects
O_t = rand(1:3,n_T)         # number of patents for each tech project
I_a = rand(n_T)*2           # costs of application projects
O_a = rand(2:4,n_T)         # number of applications for each appl. project

V_A = rand(S[cᴹ], n_A).+0.5 # Value of an application
V_A[1, :] .+= -0.5          # Low market share: less value
V_A[3, :] .+= 0.5           # High market share: more value

x_T = variables(model, [S[dᴾ]...,n_T]; binary=true)
x_A = variables(model, [S[dᴾ]...,S[cᵀ]...,S[dᴬ]..., n_A]; binary=true)

M = 20                      # a large constant
ε = 0.5*minimum([O_t O_a])  # a helper variable, allows using ≤ instead of < in constraints (28b) and (29b)

q_P = [0, 3, 6, 9]          # limits of the technology intervals
q_A = [0, 5, 10, 15]        # limits of the application intervals
z_dP = z.z[1]
z_dA = z.z[2]

@constraint(model, [i=1:3],
    sum(x_T[i,t] for t in 1:n_T) <= z_dP[i]*n_T)            #(25)
@constraint(model, [i=1:3, j=1:3, k=1:3],
    sum(x_A[i,j,k,a] for a in 1:n_A) <= z_dP[i]*n_A)        #(26)
@constraint(model, [i=1:3, j=1:3, k=1:3],
    sum(x_A[i,j,k,a] for a in 1:n_A) <= z_dA[i,j,k]*n_A)    #(27)

#(28a)
@constraint(model, [i=1:3],
    q_P[i] - (1 - z_dP[i])*M <= sum(x_T[i,t]*O_t[t] for t in 1:n_T))
#(28b)
@constraint(model, [i=1:3],
    sum(x_T[i,t]*O_t[t] for t in 1:n_T) <= q_P[i+1] + (1 - z_dP[i])*M - ε)
#(29a)
@constraint(model, [i=1:3, j=1:3, k=1:3],
    q_A[k] - (1 - z_dA[i,j,k])*M <= sum(x_A[i,j,k,a]*O_a[a] for a in 1:n_A))
#(29b)
@constraint(model, [i=1:3, j=1:3, k=1:3],
    sum(x_A[i,j,k,a]*O_a[a] for a in 1:n_A) <= q_A[k+1] + (1 - z_dA[i,j,k])*M - ε)

@constraint(model, [i=1:3, j=1:3, k=1:3], x_A[i,j,k,1] <= x_T[i,1])
@constraint(model, [i=1:3, j=1:3, k=1:3], x_A[i,j,k,2] <= x_T[i,1])
@constraint(model, [i=1:3, j=1:3, k=1:3], x_A[i,j,k,2] <= x_T[i,2])

@info("Creating model objective.")
patent_investment_cost = @expression(model, [i=1:S[1]], sum(x_T[i, t] * I_t[t] for t in 1:n_T))
application_investment_cost = @expression(model, [i=1:S[1], j=1:S[2], k=1:S[3]], sum(x_A[i, j, k, a] * I_a[a] for a in 1:n_A))
application_value = @expression(model, [i=1:S[1], j=1:S[2], k=1:S[3], l=1:S[4]], sum(x_A[i, j, k, a] * V_A[l, a] for a in 1:n_A))
@objective(model, Max, sum( sum( P((i,j,k,l)) * (application_value[i,j,k,l] - application_investment_cost[i,j,k]) for j in 1:S[2], k in 1:S[3], l in 1:S[4] ) - patent_investment_cost[i] for i in 1:S[1] ))

@info("Starting the optimization process.")
optimizer = optimizer_with_attributes(
    () -> Gurobi.Optimizer(Gurobi.Env()),
    "IntFeasTol"      => 1e-9,
    "LazyConstraints" => 1,
)
set_optimizer(model, optimizer)
optimize!(model)

@info("Extracting results.")
Z = DecisionStrategy(z)

@info("Printing decision strategy:")
print_decision_strategy(S, Z)

@info("Extracting path utilities")
struct PathUtility <: AbstractPathUtility
    data::Array{AffExpr}
end
Base.getindex(U::PathUtility, i::Int) = getindex(U.data, i)
Base.getindex(U::PathUtility, I::Vararg{Int,N}) where N = getindex(U.data, I...)
(U::PathUtility)(s::Path) = value.(U[s...])

path_utility = [@expression(model,
    sum(x_A[s[1:3]..., a] * (V_A[s[4], a] - I_a[a]) for a in 1:n_A) -
    sum(x_T[s[1], t] * I_t[t] for t in 1:n_T)) for s in paths(S)]
U = PathUtility(path_utility)

@info("Computing utility distribution.")
udist = UtilityDistribution(S, P, U, Z)

@info("Printing utility distribution.")
print_utility_distribution(udist)

@info("Printing statistics")
print_statistics(udist)


# # optimization spaces
# abstract type Space end

# # Euclidean
# struct Euclidean <: Space
#     n::Int
# end

function conjugate(q)
	s = q[1]
	v = q[2:4]

	return [s; -v]
end

function L_multiply(q)
	s = q[1]
	v = q[2:4]

	SMatrix{4,4}([s -transpose(v);
	              v s * I + skew(v)])
end

function R_multiply(q)
	s = q[1]
	v = q[2:4]

	SMatrix{4,4}([s -transpose(v);
	              v s * I - skew(v)])
end

function multiply(q1, q2)
	L_multiply(q1) * q2
end

# eq. 14 http://roboticexplorationlab.org/papers/planning_with_attitude.pdf
function attitude_jacobian(q)
	s = q[1]
	v = q[2:4]

	[-transpose(v);
	 s * I + skew(v)]
end

function G_func(q)
	quat = q[4:7]
	[1.0 0.0 0.0 0.0 0.0 0.0;
     0.0 1.0 0.0 0.0 0.0 0.0;
	 0.0 0.0 1.0 0.0 0.0 0.0;
     zeros(4, 3) attitude_jacobian(quat)]
end

function static_quaternion(q::UnitQuaternion)
	SVector{4}([q.w, q.x, q.y, q.z])
end

function quaternion_offset(q_parent, q_child)
	L_multiply(conjugate(q_parent)) * q_child
end

function get_quaternion(q)
	return view(q, 4:7)
end

# eq. 16 http://roboticexplorationlab.org/papers/maximal_coordinate_dynamics.pdf
function ω_finite_difference(q1, q2, h)
	2.0 * multiply(conjugate(q1), (q2 - q1) ./ h)[2:4]
end

# https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation#Quaternion-derived_rotation_matrix
function quaternion_rotation_matrix(q)
	r, i, j, k  = q

	r11 = 1.0 - 2.0 * (j^2.0 + k^2.0)
	r12 = 2.0 * (i * j - k * r)
	r13 = 2.0 * (i * k + j * r)

	r21 = 2.0 * (i * j + k * r)
	r22 = 1.0 - 2.0 * (i^2.0 + k^2.0)
	r23 = 2.0 * (j * k - i * r)

	r31 = 2.0 * (i * k - j * r)
	r32 = 2.0 * (j * k + i * r)
	r33 = 1.0 - 2.0 * (i^2.0 + j^2.0)

	SMatrix{3,3}([r11 r12 r13;
	              r21 r22 r23;
				  r31 r32 r33])
end

# Cayley map
function cayley_map(ϕ)
	1.0 / sqrt(1.0 + norm(ϕ)^2.0) * [1.0; ϕ]
end

# slerp
function slerp(q0::UnitQuaternion, q1::UnitQuaternion, t::T) where T
	Rotations.params(q0 * exp(t * log(inv(q0) * q1)))
end

# Square root of quaternion: https://www.johndcook.com/blog/2021/01/06/quaternion-square-roots/
function sqrt_quat(q; ϵ=1e-16)
	r = norm(q)
	# Angle
	theta = acos(q[1]/r)
	# Axis
	u = q[2:4]
	u ./= norm(u) + ϵ
	# Half axis-angle rotation
	x = [cos(theta/2);  sin(theta/2)*u]
	# Sqrt on the norm of the quaternion
	x .*= r^0.5 # useless for unit quaternion
	return x
end

function midpoint(q0, q1)
	# Small delta rotation
	qψ = R_multiply(q0)' * q1
	# Divide delta in two
	sqψ = sqrt_quat(qψ)
	# Apply small rotation
	qmid = L_multiply(sqψ) * q0
	return qmid
end

struct RnQuaternion <: Space
	n::Int # size of the optimization space (typically = num_var - 1)
	mapping # Jacobian mapping i.e. a function of z
	r_idx # euclidean part in z
	Δr_idx # euclidean part in Δ
	quat_idx # quaternion part in z
	Δquat_idx # quaternion part in Δ
end

function rn_quaternion_space(dim, mapping, r_idx, Δr_idx, quat_idx, Δquat_idx)
	RnQuaternion(dim, mapping, r_idx, Δr_idx, quat_idx, Δquat_idx)
end

function candidate_point!(z̄::Vector{T}, s::RnQuaternion, z::Vector{T}, Δ::Vector{T}, α::T) where T
	z̄[s.r_idx] .= z[s.r_idx] - α .* Δ[s.Δr_idx]
	for i = 1:length(s.quat_idx)
		z̄[s.quat_idx[i]] .= L_multiply(z[s.quat_idx[i]]) * cayley_map(-1.0 * α .* Δ[s.Δquat_idx[i]])
	end
	return nothing
end

function mapping!(δz, s::RnQuaternion, δzs, z)
	δz .= s.mapping(z) * δzs
end


import CommutativeRings: pseudo_ideal

@testset "basics" begin
    R = ZZ{Int}
    RX = R[:x]
    RYZ = R[:y,:z]
    @test iszero(Ideal(0))
    @test isone(Ideal(-1))
    @test Ideal(15, 21) == Ideal(3)
    @test Ideal([3, 5]) == R
    @test R == Ideal(-1)
    @test Ideal(3, ZZ(15)) == Ideal(3)

    x, = generators(RX)
    y, z = generators(RYZ)
    @test Ideal(x, x^2-x) == Ideal(x)
    @test Ideal(y + z, y^2 - z^2) == Ideal(y + z)

    @test pseudo_ideal(R, 15) == ZZ(15)
    @test pseudo_ideal(R, ZZ(Int8(127))) == 127
    @test pseudo_ideal(R, [15, 21]) == Ideal(3)
    i = Ideal(ZZ(3))
    @test pseudo_ideal(R, i) === i
    j = Ideal(ZZ(Int16(3)))
    @test i == j
end

@testset "operations" begin
    P = QQ{Int}[:x,:y,:z]
    x, y, z = generators(P)
    @test Ideal(x, y) isa Ideal

    A = [x*z^2 - 3y*z; x^2-2y; x*y - 5z]
    B = [x*y]
    ida = Ideal(A)
    idb = Ideal(B)
    idint = intersect(ida, idb)
    idsum = ida + idb

    @test ida isa Ideal{P}
    @test Ideal(ida.base) == ida
    @test zero(Ideal{P}) == zero(ida)
    @test zero(Ideal{P}) == Ideal(P(0))
    @test one(Ideal{P}) == Ideal(P(1))
    @test isone(Ideal(QQ(2)))
    @test issubset(idint, ida)
    @test issubset(idint, idb)
    @test issubset(ida, idsum)
    @test issubset(idb, idsum)
    @test A[1] * y + A[2] in ida
    @test 2x*y*z^2 - 9x*y^2 in idint
    @test B[1] * x + A[1] * y in idsum

    i0 = zero(ida)
    i1 = one(ida)
    @test i0 ⊆ ida ⊆ i1
    @test i0 + ida == ida
    @test i1 + ida == i1
    @test i1 ∩ ida == ida
    @test i0 ∩ ida == i0
    @test zero(ida) != idint^2 ⊆ ida * idb ⊆ idint
    @test one(ida) != idsum ⊆ one(ida)

    @test isone(idb^0)
    @test idb^1 === idb
    @test idb^2 == idb * idb
    @test idb^3 == idb * idb * idb
    @test idb^4 == idb * idb * idb * idb

    @test idb + x == x + idb
    @test idb * x == x * idb
    @test rem((x + y)^2, idb) == x^2 + y^2
end


using Polytopes
using Colors
using GraphPlot
using LightGraphs

p = Polytope([[1, 2], [2, 3], [3, 1], [2, 3], [3, 4], [4, 2]], [[1, 2, 3], [4, 5, 6]])
g = p.graph
is = Polytopes.start_indices(p.connectivity)
colors = [colorant"red", colorant"turquoise", colorant"green"]
gplot(g, nodelabel=vertices(g), nodefillc=[colors[something(findfirst(x -> i ≤ x, is), 0) + 1] for i ∈ vertices(g)])


doc"""
    Levy(μ, σ)

The *Lévy distribution* with location `μ` and scale `σ` has probability density function

$f(x; \mu, \sigma) = \sqrt{\frac{\sigma}{2 \pi (x - \mu)^3}}
\exp \left( - \frac{\sigma}{2 (x - \mu)} \right), \quad x > \mu$

```julia
Levy()         # Levy distribution with zero location and unit scale, i.e. Levy(0, 1)
Levy(u)        # Levy distribution with location u and unit scale, i.e. Levy(u, 1)
Levy(u, c)     # Levy distribution with location u ans scale c

params(d)      # Get the parameters, i.e. (u, c)
location(d)    # Get the location parameter, i.e. u
```

External links

* [Lévy distribution on Wikipedia](http://en.wikipedia.org/wiki/Lévy_distribution)
"""
immutable Levy{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    σ::T

    Levy(μ::T, σ::T) = (@check_args(Levy, σ > zero(σ)); new(μ, σ))
end

Levy{T<:Real}(μ::T, σ::T) = Levy{T}(μ, σ)
Levy(μ::Real, σ::Real) = Levy(promote(μ, σ)...)
Levy(μ::Integer, σ::Integer) = Levy(Float64(μ), Float64(σ))
Levy(μ::Real) = Levy(μ, 1.0)
Levy() = Levy(0.0, 1.0)

@distr_support Levy d.μ Inf

#### Conversions

convert{T <: Real, S <: Real}(::Type{Levy{T}}, μ::S, σ::S) = Levy(T(μ), T(σ))
convert{T <: Real, S <: Real}(::Type{Levy{T}}, d::Levy{S}) = Levy(T(d.μ), T(d.σ))

#### Parameters

location(d::Levy) = d.μ
params(d::Levy) = (d.μ, d.σ)
@inline partype{T<:Real}(d::Levy{T}) = T

#### Statistics

mean{T<:Real}(d::Levy{T}) = T(Inf)
var{T<:Real}(d::Levy{T}) = T(Inf)
skewness{T<:Real}(d::Levy{T}) = T(NaN)
kurtosis{T<:Real}(d::Levy{T}) = T(NaN)

mode(d::Levy) = d.σ / 3 + d.μ

entropy(d::Levy) = (1 - 3digamma(1) + log(16 * d.σ^2 * π)) / 2

median(d::Levy) = d.μ + d.σ / (2 * T(erfcinv(0.5))^2)

#### Evaluation

function pdf(d::Levy, x::Real)
    μ, σ = params(d)
    z = x - μ
    (sqrt(σ) / sqrt2π) * exp((-σ) / (2z)) / z^(3//2)
end

function logpdf(d::Levy, x::Real)
    μ, σ = params(d)
    z = x - μ
    (log(σ) - log2π - σ / z - 3log(z))/2
end

cdf(d::Levy, x::Real) = erfc(sqrt(d.σ / (2(x - d.μ))))
ccdf(d::Levy, x::Real) = erf(sqrt(d.σ / (2(x - d.μ))))

quantile(d::Levy, p::Real) = d.μ + d.σ / (2*erfcinv(p)^2)
cquantile(d::Levy, p::Real) = d.μ + d.σ / (2*erfinv(p)^2)

mgf{T<:Real}(d::Levy{T}, t::Real) = t == zero(t) ? one(T) : T(NaN)

function cf(d::Levy, t::Real)
    μ, σ = params(d)
    exp(im * μ * t - sqrt(-2im * σ * t))
end

#### Sampling

rand(d::Levy) = d.μ + d.σ / randn()^2


function ContJitter(n, d, num)

    base_obn = convert(Array{Float32},range(d[1],stop=(n[1]-1)*d[1],length = num)) # spacing = 250m on average
    avg_spc = (n[1]-2)*d[1]/(num-1)
    OBN_loc = zeros(Float32,num)

    for i = 1:num
        OBN_loc[i] = max(min(base_obn[i] + avg_spc*rand(Float32) - avg_spc/2, (n[1]-1)*d[1]),0f0)
    end

    return OBN_loc
end

function Jitt_function(N, gamma, xi)
    # 2020.12.09 Yijun
    # jittered subsampling function for simple test
    # N: the full size
    # gamma: undersampling factor
    # xi: control the maximum gap between the neighbor subsampling indexes
    # gamma=4 --> 75 % subsampling
    #rng(rseed);
    test1 = rand(1,Int64(round(N/gamma))).*(xi-1) .- (xi-1)/2 ;
    test2 = ones(1, Int64(round(N/gamma)));
    test3 = cumsum(test2; dims=2);
    n = Int64.(round.((1-gamma)/2 .+ gamma .* test3 .+ test1));
    for i = 1:length(n)
        if n[i] > N
            n[i] = N
        end
    end
    return n
end