@testset "basic building blocks" begin
    mt = UltimateMortality([0.5, 0.5])

    # whole life insurance
    ins = Insurance(
        SingleLife(mortality = mt, issue_age = 0),
        Yields.Constant(0.05)
    )

    @test all(survival(SingleLife(mortality = mt, issue_age = 0)) .== [1.0, 0.5, 0.25])

    @test [timepoints(ins)...] == [1.0, 2.0]
    @test [survival(ins)...] == [1.0, 0.5]
    @test [discount(ins)...] == [1.0 / 1.05, 1 / 1.05^2]
    @test benefit(ins) == 1.0
    @test [probability(ins)...] == [0.5, 0.25]
    @test [cashflows(ins)...] == [0.5, 0.25]
    @test [cashflows(ins)...] == benefit(ins) .* probability(ins)
    @test present_value(ins) ≈ 0.5 / 1.05 + 0.5 * 0.5 / 1.05^2
    @test present_value(ins) ≈ LifeContingencies.A(ins.life, ins.int)

    # basic life contingency tests
    @test survival(LifeContingency(ins), 0, 1) ≈ 0.5
    @test survival(LifeContingency(ins), 0, 2) ≈ 0.25
    @test discount(LifeContingency(ins), 0, 1) ≈ 1 / 1.05
    @test discount(LifeContingency(ins), 0, 2) ≈ 1 / 1.05^2

    # term life insurance
    ins = Insurance(
        SingleLife(mortality = mt, issue_age = 0),
        Yields.Constant(0.05),
        1
    )

    @test [survival(ins)...] == [1.0]
    @test [discount(ins)...] == [1.0 / 1.05]
    @test benefit(ins) == 1.0
    @test [probability(ins)...] == [0.5]
    @test [cashflows(ins)...] == [0.5]
    @test [cashflows(ins)...] == benefit(ins) .* probability(ins)
    @test [timepoints(ins)...] == [1.0]
    @test present_value(ins) ≈ 0.5 / 1.05

    # annuity due
    ins = AnnuityDue(
        SingleLife(mortality = mt, issue_age = 0),
        Yields.Constant(0.05)
    )

    @test [survival(ins)...] == [1.0, 0.5, 0.25]
    @test [discount(ins)...] == [1.0, 1 / 1.05^1, 1 / 1.05^2]
    @test benefit(ins) == 1.0
    @test [timepoints(ins)...] == [0.0, 1.0, 2.0]
    @test [probability(ins)...] == [1.0, 0.5, 0.25]
    @test [cashflows(ins)...] == [1.0, 0.5, 0.25]
    @test [cashflows(ins)...] == benefit(ins) .* probability(ins)
    @test present_value(ins) ≈ 1 + 1 * 0.5 / 1.05 + 1 * 0.25 / 1.05^2

    ins = AnnuityDue(
        SingleLife(mortality = mt, issue_age = 0),
        Yields.Constant(0.05),
        2
    )
    @test present_value(ins) ≈ 1 + 1 * 0.5 / 1.05

    ins = AnnuityImmediate(
        SingleLife(mortality = mt, issue_age = 0),
        Yields.Constant(0.05),
        1
    )

    @test present_value(ins) ≈ 1 * 0.5 * 1 / 1.05

    ins = AnnuityDue(
        SingleLife(mortality = mt, issue_age = 0),
        Yields.Constant(0.05),
        1,
        start_time = 1
    )

    @test timepoints(ins) == [1.0]
    @test [probability(ins)...] == [0.5]
    @test present_value(ins) ≈ 1 * 0.5 * 1 / 1.05

end

@testset "one  year" begin
    ins = LifeContingency(
        SingleLife(
            mortality = UltimateMortality([0.5]),
            issue_age = 0
        ),
        Yields.Constant(0.05)
    )

    @test survival(ins, 1) ≈ 0.5
    @test survival(ins, 0) ≈ 1.0
    @test survival(ins.life, 0, 0.5) ≈ 1 - 0.5 * 0.5
    @test survival(ins, 0.5) ≈ 1 - 0.5 * 0.5

    @test omega(ins) ≈ 1
    @test present_value(AnnuityDue(ins)) ≈ 1 + 1 * 0.5 / 1.05
    @test present_value(AnnuityDue(ins, 1)) ≈ 1
    @test present_value(AnnuityDue(ins, 0)) == 0
    @test present_value(AnnuityImmediate(ins, 1)) ≈ 1 * 0.5 / 1.05
    @test present_value(AnnuityImmediate(ins, 0)) == 0

    @test present_value(Insurance(ins)) ≈ 0.5 / 1.05
    @test present_value(Insurance(ins, 1)) ≈ 0.5 / 1.05
    @test present_value(Insurance(ins, 0)) ≈ 0

    ins_jl = LifeContingency(
        JointLife((ins.life, ins.life), LastSurvivor(), Frasier()),
        Yields.Constant(0.05)
    )

    @test omega(ins_jl) ≈ 1
    @test present_value(AnnuityDue(ins_jl)) ≈ 1 + 1 * 0.75 / 1.05
    @test present_value(AnnuityDue(ins_jl, 1)) ≈ 1
    @test present_value(AnnuityImmediate(ins_jl, 1)) ≈ 1 * 0.75 / 1.05
    @test present_value(AnnuityDue(ins_jl, 2)) ≈ 1 + 1 * 0.75 / 1.05
    @test present_value(AnnuityDue(ins_jl, 2; certain = 2)) ≈ 1 + 1 / 1.05
    @test present_value(AnnuityDue(ins_jl, 0)) == 0

    @test survival(ins_jl, 1) ≈ 0.5 + 0.5 - 0.5 * 0.5
    @test all(survival(JointLife((ins.life, ins.life), LastSurvivor(), Frasier())) .== [1.0, 0.75])
    @test present_value(Insurance(ins_jl)) ≈ 0.25 / 1.05
    @test present_value(Insurance(ins_jl, 1)) ≈ 0.25 / 1.05
    @test present_value(Insurance(ins_jl, 0)) ≈ 0
end

@testset "two year no discount" begin
    ins = LifeContingency(
        SingleLife(
            mortality = UltimateMortality([0.0, 0.0]),
            issue_age = 0
        ),
        Yields.Constant(0.00)
    )

    @test omega(ins) ≈ 2
    @test present_value(AnnuityDue(ins)) ≈ 3
    @test present_value(AnnuityDue(ins; start_time = 1)) ≈ 2
    @test present_value(AnnuityDue(ins, 1)) ≈ 1
    @test present_value(AnnuityDue(ins, 2)) ≈ 2
    @test_throws BoundsError present_value(AnnuityDue(ins, 2; start_time = 2))
    @test present_value(AnnuityDue(ins, 3)) ≈ 3
    @test present_value(AnnuityDue(ins, 0)) == 0
    @test present_value(AnnuityImmediate(ins, 0)) == 0
    @test present_value(AnnuityImmediate(ins, 1)) ≈ 1
    @test present_value(AnnuityImmediate(ins, 1; start_time = 1)) ≈ 1
    @test_throws BoundsError present_value(AnnuityImmediate(ins; start_time = 2))
    @test present_value(AnnuityImmediate(ins, 2)) ≈ 2
    @test present_value(AnnuityImmediate(ins)) ≈ 2
    @test present_value(AnnuityImmediate(ins; certain = 0)) ≈ 2
    @test present_value(AnnuityImmediate(ins; certain = 2)) ≈ 2

    @test present_value(Insurance(ins)) ≈ 0
    @test present_value(Insurance(ins, 1)) ≈ 0
    @test present_value(Insurance(ins, 0)) ≈ 0

    ins_jl = LifeContingency(
        JointLife(
            lives = (ins.life, ins.life),
            contingency = LastSurvivor(),
            joint_assumption = Frasier()),
        Yields.Constant(0.00)
    )

    @test survival(ins_jl, 0) ≈ 1.0
    @test survival(ins_jl, 1) ≈ 1
    @test survival(ins_jl, 2) ≈ 1

end

@testset "two year with interest" begin
    ins = LifeContingency(
        SingleLife(
            mortality = UltimateMortality([0.5, 0.5]),
            issue_age = 0
        ),
        Yields.Constant(0.05)
    )

    @test omega(ins) ≈ 2
    @test present_value(AnnuityDue(ins)) ≈ 1 + 1 * 0.5 * 1 / 1.05 + 1 * 0.25 / 1.05^2
    @test present_value(AnnuityDue(ins, 1)) ≈ 1
    @test present_value(AnnuityDue(ins, 2)) ≈ 1 + 1 * 0.5 * 1 / 1.05
    @test present_value(AnnuityDue(ins, 3)) ≈ 1 + 1 * 0.5 * 1 / 1.05 + 1 * 0.25 / 1.05^2
    @test present_value(AnnuityDue(ins, 0)) ≈ 0
    @test present_value(AnnuityImmediate(ins, 0)) ≈ 0
    @test present_value(AnnuityImmediate(ins, 1)) ≈ 1 * 0.5 * 1 / 1.05

    @test present_value(Insurance(ins)) ≈ 0.5 / 1.05 + 0.5 * 0.5 / 1.05^2
    @test present_value(Insurance(ins, 1)) ≈ 0.5 / 1.05
    @test present_value(Insurance(ins, 0)) ≈ 0

    ins_jl = LifeContingency(
        JointLife(
            lives = (ins.life, ins.life),
            contingency = LastSurvivor(),
            joint_assumption = Frasier()),
        Yields.Constant(0.05)
    )

    @test survival(ins_jl, 0) ≈ 1.0
    @test survival(ins_jl, 1) ≈ 0.5 + 0.5 - 0.5 * 0.5
    @test survival(ins_jl, 2) ≈ (0.25 + 0.25 - 0.25 * 0.25)

end

@testset "two years" begin
    ins = LifeContingency(
        SingleLife(
            mortality = UltimateMortality([0.5, 0.5]),
            issue_age = 0
        ),
        Yields.Constant(0.05)
    )

    @test omega(ins) ≈ 2
    @test present_value(AnnuityDue(ins, 1)) ≈ 1
    @test present_value(AnnuityDue(ins, 2)) ≈ 1 + 1 * 0.5 * 1 / 1.05
    @test present_value(AnnuityDue(ins, 3)) ≈ 1 + 1 * 0.5 * 1 / 1.05 + 1 * 0.25 * 1 / 1.05^2
    @test present_value(AnnuityDue(ins, 0)) == 0
    @test present_value(AnnuityImmediate(ins, 0)) == 0
    @test present_value(AnnuityImmediate(ins, 1)) ≈ 1 * 0.5 * 1 / 1.05

    @test present_value(Insurance(ins)) ≈ 0.5 / 1.05 + 0.5 * 0.5 / 1.05^2
    @test present_value(Insurance(ins, 1)) ≈ 0.5 / 1.05
    @test present_value(Insurance(ins, 0)) ≈ 0

    ins_jl = LifeContingency(
        JointLife(
            lives = (ins.life, ins.life),
            contingency = LastSurvivor(),
            joint_assumption = Frasier()),
        Yields.Constant(0.05)
    )

    @test omega(ins_jl) ≈ 2
    @test present_value(AnnuityDue(ins_jl, 1)) ≈ 1
    @test present_value(AnnuityDue(ins_jl, 2)) ≈ 1 + 1 * 0.75 / 1.05
    @test present_value(AnnuityDue(ins_jl, 3)) ≈ 1 + 1 * 0.75 / 1.05 + 1 * survival(ins_jl, 2) / 1.05^2
    @test present_value(AnnuityDue(ins_jl, 0)) == 0
    @test present_value(AnnuityImmediate(ins_jl, 0)) == 0
    @test present_value(AnnuityImmediate(ins_jl, 1)) ≈ 1 * 0.75 / 1.05

end

# assumes embedded 'testMort' table in combination with MortalityTables
t = UltimateMortality(maleMort)

@testset "demo mortality" begin
    i = Yields.Constant(0.05)
    ins = LifeContingency(SingleLife(mortality = t, issue_age = 0), i)

    @test l(ins, 0) ≈ 1.0
    @test l(ins, 1) ≈ 0.993010000000000
    @test l(ins, 2) ≈ 0.992566124530000

    @test D(ins, 0) ≈ 1.0
    @test D(ins, 1) ≈ 0.945723809523809
    @test D(ins, 2) ≈ 0.900286734267574

    @test N(ins, 0) ≈ 20.113017073119200
    @test N(ins, 1) ≈ 19.113017073119200
    @test N(ins, 2) ≈ 18.167293263595400

    @test C(ins, 0) ≈ 0.006657142857142910
    @test C(ins, 1) ≈ 0.000402608136054441
    @test C(ins, 2) ≈ 0.000258082197156658

    @test M(ins, 0) ≈ 0.0422372822324182
    @test M(ins, 1) ≈ 0.0355801393752753
    @test M(ins, 2) ≈ 0.0351775312392208

    @test present_value(Insurance(ins)) ≈ 0.04223728223
    @test present_value(Insurance(ins, 0)) ≈ 0.0
    @test present_value(Insurance(ins, 1)) ≈ 0.0066571428571429
    @test present_value(AnnuityDue(ins, 0)) ≈ 0.0
    @test present_value(AnnuityDue(ins, 1)) ≈ 1.0
    @test present_value(AnnuityDue(ins, 2)) ≈ 1.0 + survival(ins, 1) / 1.05
    @test present_value(AnnuityImmediate(ins, 0)) ≈ 0.0
    @test present_value(AnnuityImmediate(ins, 1)) ≈ survival(ins, 1) / 1.05

    @test present_value(Insurance(ins, 30)) ≈ 0.0137761089686975

    @test N(ins, 26) ≈ 5.156762988852310
    @test D(ins, 26) ≈ 0.275358702015970
    @test present_value(AnnuityDue(ins, 26)) ≈ 14.9562540842669

end


import Random
using Plots

println("Pool Testing Optimization\nMonte Carlo Simulation")

min_total = 1
total_step = 1
max_total = 100
prob_step = 0.01
iterations = 1000

function run(total, prob, n_groups)
	# Trivial cases
	# No one is infected or testing everybody
	if prob==0 || n_groups==total
		return n_groups
	# Only 1 group (there is only 1 grouping step)
	elseif n_groups==1
		return total+1
	end

	# Calculate infected from probability using stochastic rounding
	infected = prob*total
	if rand(Float32)<(infected-floor(infected))
		infected = Int(ceil(infected))
	else
		infected = Int(floor(infected))
	end

	# No one is infected
	if infected==0
		return n_groups
	end

	people = zeros(Bool, total)
	people[1:infected] .= true
	Random.shuffle!(people)

	n_per_group = total ÷ n_groups
	n_plus_groups = total % n_groups

	#groups = zeros(UInt, n_groups)
	groups = zeros(Bool, n_groups)

	for i in 1:infected
		j = 1+(Random.rand(UInt) % n_groups)
		#groups[j] += 1
		groups[j] = true
	end

	#infected_groups = count(g -> g>0, groups)
	#infected_groups = count(groups)

	tests = n_groups
	tests += count(groups[1:n_plus_groups]) * (n_per_group+1)
	tests += count(groups[n_plus_groups+1:end]) * n_per_group

	#=
	println("Total   \t", total)
	println("Infected\t", infected)
	println("People\t", people)
	println("People per group\t", n_per_group)
	println("# groups with 1 more\t", n_plus_groups)
	println("Infected per group\t", groups)
	println("Number of infected groups\t", infected_groups)
	println("Number of tests\t", tests)
	=#

	return tests
end

n_total = Int(floor((max_total-min_total)/total_step)+1)
n_prob = Int(floor(1/prob_step))+1
results = Array{Float32, 2}(undef, n_total*n_prob, 3)

i = 1
for total in min_total:total_step:max_total
	for prob in 0:prob_step:1
		println("total ", total, "  probability ", prob)

		optimum_n_groups = total
		previous_sum_tests = iterations*total

		for n_groups in 1:total
			sum_tests = 0

			for iteration in 1:iterations
				tests = run(total, prob, n_groups)
				sum_tests += tests
				#println("total ", total, "  prob ", prob, "  n_groups ", n_groups, "  iteration ", iteration, "  tests ", tests)
			end

			if sum_tests<previous_sum_tests
				optimum_n_groups = n_groups
			end
		end

		global i
		results[i, :] = [total, prob, optimum_n_groups]
		i += 1
	end
end

println("Calculated")

matrix = results[:, 3]./results[:, 1]
matrix = reshape(matrix, n_prob, n_total)'

#println("results ", results, '\n')
#println("matrix ", matrix, '\n')

heatmap(min_total:total_step:max_total, 0:prob_step:1, matrix',
		xlabel="Total", ylabel="% Infected",
		title="Optimal number of groups / Total",
		show=true)
println("Plotted")
savefig("./result.png")
println("Saved")

#=
gui()
scatter(results[:, 2], results[:, 3]./results[:, 1], show=true)
x = prob_step:prob_step:1
y = x.^(-0.5)
plot!(x, y)
=#

readline()

println("End")


export compute_fem_stiffness_matrix,
compute_fem_source_term,
fem_impose_Dirichlet_boundary_condition,
eval_f_on_gauss_pts,
eval_f_on_fvm_pts,
eval_f_on_fem_pts,
compute_fvm_mass_matrix,
compute_interaction_matrix,
compute_fvm_source_term,
compute_fvm_mechanics_term,
compute_fvm_tpfa_matrix,
trim_coupled,
compute_elasticity_tangent,
compute_fem_traction_term,
compute_fem_traction_term1,
compute_fem_normal_traction_term,
coupled_impose_pressure,
compute_von_mises_stress_term,
compute_fem_mass_matrix,
eval_f_on_boundary_node,
eval_f_on_boundary_edge,
compute_fem_mass_matrix1,
compute_fem_stiffness_matrix1,
compute_fem_source_term1,
compute_fem_flux_term1,
fem_impose_Dirichlet_boundary_condition1,
fem_impose_coupled_Dirichlet_boundary_condition,
eval_strain_on_gauss_pts,
eval_stress_on_gauss_pts,
compute_fvm_mass_matrix,
compute_strain_energy_term,
compute_plane_strain_matrix,
compute_fem_laplace_matrix1,
compute_fem_laplace_matrix,
eval_grad_on_gauss_pts1,
eval_grad_on_gauss_pts,
compute_plane_stress_matrix,
eval_f_on_boundary_gauss_pts

####################### Mechanics #######################
@doc raw"""
    compute_fem_stiffness_matrix(K::Array{Float64,2}, m::Int64, n::Int64, h::Float64)

Computes the term
```math
\int_{A}\delta \varepsilon :\sigma\mathrm{d}x = \int_A u_AB^TKB\delta u_A\mathrm{d}x
```
where the constitutive relation is given by
```math
\begin{bmatrix}\sigma_{xx}\\\sigma_{yy}\\\sigma_{xy}\end{bmatrix} = K \begin{bmatrix}\varepsilon_{xx}\\\varepsilon_{yy}\\2\varepsilon_{xy}\end{bmatrix}
```
"""
function compute_fem_stiffness_matrix(K::Array{Float64,2}, m::Int64, n::Int64, h::Float64)
    I = Int64[]; J = Int64[]; V = Float64[]
    function add(ii, jj, kk)
        for  i = 1:length(ii)
            for j = 1:length(jj)
                push!(I,ii[i])
                push!(J,jj[j])
                push!(V,kk[i,j])
            end
        end
    end
    Ω = zeros(8,8)
    for i = 1:2
        for j = 1:2
            ξ = pts[i]; η = pts[j]
            B = [
            -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η 0.0 0.0 0.0 0.0
            0.0 0.0 0.0 0.0 -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
            ]
            Ω += B'*K*B*0.25*h^2
        end
    end

    for i = 1:m
        for j = 1:n
            ii = [i;i+1;i;i+1]
            jj = [j;j;j+1;j+1]
            kk = (jj .- 1)*(m+1) + ii
            kk = [kk; kk .+ (m+1)*(n+1)]
            add(kk, kk, Ω)
        end
    end
    return sparse(I, J, V, (m+1)*(n+1)*2, (m+1)*(n+1)*2)
end

@doc raw"""
    compute_fem_stiffness_matrix1(K::Array{Float64,2}, m::Int64, n::Int64, h::Float64)

Computes the term
```math
\int_{A} (K \nabla u) \cdot \nabla \delta u \mathrm{d}x = \int_A u_A B^T K B \delta u_A\mathrm{d}x
```
Returns a $(m+1)\times (n+1)$ matrix
"""
function compute_fem_stiffness_matrix1(K::Array{Float64,2}, m::Int64, n::Int64, h::Float64)
    I = Int64[]; J = Int64[]; V = Float64[]
    @assert size(K,1)==2 && size(K,2)==2
    function add(ii, jj, kk)
        for  i = 1:length(ii)
            for j = 1:length(jj)
                push!(I,ii[i])
                push!(J,jj[j])
                push!(V,kk[i,j])
            end
        end
    end
    Ω = zeros(4,4)
    for i = 1:2
        for j = 1:2
            ξ = pts[i]; η = pts[j]
            B = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
                -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            ]
            Ω += B'*K*B*0.25*h^2
        end
    end

    for i = 1:m
        for j = 1:n
            ii = [i;i+1;i;i+1]
            jj = [j;j;j+1;j+1]
            kk = (jj .- 1)*(m+1) + ii
            add(kk, kk, Ω)
        end
    end
    return sparse(I, J, V, (m+1)*(n+1), (m+1)*(n+1))
end

@doc raw"""
    compute_fem_source_term(f1::Array{Float64}, f2::Array{Float64},
    m::Int64, n::Int64, h::Float64)

Computes the term
```math
\int_\Omega \mathbf{f}\cdot\delta u \mathrm{d}x
```
Returns a $2(m+1)(n+1)$ vector.
"""
function compute_fem_source_term(f1::Array{Float64}, f2::Array{Float64},
    m::Int64, n::Int64, h::Float64)
    rhs = zeros((m+1)*(n+1)*2)
    k = 0
    for i = 1:m
        for j = 1:n
            idx = (j-1)*m + i
            for p = 1:2
                for q = 1:2
                    ξ = pts[p]; η = pts[q]

                    k = (idx-1)*4 + 2*(q-1) + p
                    val1 = f1[k] * h^2 * 0.25
                    val2 = f2[k] * h^2 * 0.25

                    rhs[(j-1)*(m+1) + i] += val1 * (1-ξ)*(1-η)
                    rhs[(j-1)*(m+1) + i+1] += val1 * ξ*(1-η)
                    rhs[j*(m+1) + i] += val1 * (1-ξ)*η
                    rhs[j*(m+1) + i+1] += val1 * ξ  * η

                    rhs[(m+1)*(n+1) + (j-1)*(m+1) + i] += val2* (1-ξ)*(1-η)
                    rhs[(m+1)*(n+1) + (j-1)*(m+1) + i+1] += val2* ξ*(1-η)
                    rhs[(m+1)*(n+1) + j*(m+1) + i] += val2 * (1-ξ)*η
                    rhs[(m+1)*(n+1) + j*(m+1) + i+1] += val2 * ξ  * η

                end
            end

        end
    end
    return rhs
end

@doc raw"""
    compute_fem_source_term1(f::Array{Float64},
    m::Int64, n::Int64, h::Float64)

Computes the term
```math
\int_\Omega f \delta u dx
```
Returns a $(m+1)\times (n+1)$ vector. `f` is a length $4mn$ vector, given by its values on Gauss points.
"""
function compute_fem_source_term1(f::Array{Float64}, m::Int64, n::Int64, h::Float64)
    rhs = zeros((m+1)*(n+1))
    for i = 1:m
        for j = 1:n
            idx = (j-1)*m + i
            for p = 1:2
                for q = 1:2
                    ξ = pts[p]; η = pts[q]
                    k = (idx-1)*4 + 2*(q-1) + p
                    val1 = f[k] * h^2 * 0.25
                    rhs[(j-1)*(m+1) + i] += val1 * (1-ξ)*(1-η)
                    rhs[(j-1)*(m+1) + i+1] += val1 * ξ*(1-η)
                    rhs[j*(m+1) + i] += val1 * (1-ξ)*η
                    rhs[j*(m+1) + i+1] += val1 * ξ  * η
                end
            end

        end
    end
    return rhs
end

@doc raw"""
    fem_impose_Dirichlet_boundary_condition(A::SparseMatrixCSC{Float64,Int64},
    bd::Array{Int64}, m::Int64, n::Int64, h::Float64)

Imposes the Dirichlet boundary conditions on the matrix `A`.

Returns 2 matrix,
```math
\begin{bmatrix}
A_{BB} & A_{BI} \\
A_{IB} & A_{II}
\end{bmatrix} \Rightarrow \begin{bmatrix}
I & 0 \\
0 & A_{II}
\end{bmatrix}, \quad \begin{bmatrix}
0 \\
A_{IB}
\end{bmatrix}
```
"""
function fem_impose_Dirichlet_boundary_condition(A::SparseMatrixCSC{Float64,Int64},
    bd::Array{Int64}, m::Int64, n::Int64, h::Float64)
    M, N = size(A)
    bd = [bd; bd .+ (m+1)*(n+1)]
    Q = A[:, bd]; Q[bd,:] = spzeros(length(bd), length(bd))
    A[bd,:] = spzeros(length(bd), N)
    A[:,bd] = spzeros(M, length(bd))
    A[bd,bd] = spdiagm(0=>ones(length(bd)))
    return A, Q
end

"""
"""
function fem_impose_coupled_Dirichlet_boundary_condition(A::SparseMatrixCSC{Float64,Int64},
    bd::Array{Int64}, m::Int64, n::Int64, h::Float64)
    bd = [bd; bd .+ (m+1)*(n+1)]
    Q = A[:, bd]; Q[bd,:] = spzeros(length(bd), length(bd))
    A[bd,:] = spzeros(length(bd), 2(m+1)*(n+1)+m*n)
    A[:,bd] = spzeros(2(m+1)*(n+1)+m*n, length(bd))
    A[bd,bd] = spdiagm(0=>ones(length(bd)))
    return A, Q
end

@doc raw"""
    fem_impose_Dirichlet_boundary_condition1(A::SparseMatrixCSC{Float64,Int64},
        bd::Array{Int64}, m::Int64, n::Int64, h::Float64)

Imposes the Dirichlet boundary conditions on the matrix `A`
Returns 2 matrix,
```math
\begin{bmatrix}
A_{BB} & A_{BI} \\
A_{IB} & A_{II}
\end{bmatrix} \Rightarrow \begin{bmatrix}
I & 0 \\
0 & A_{II}
\end{bmatrix}, \quad \begin{bmatrix}
0 \\
A_{IB}
\end{bmatrix}
```

`bd` must NOT have duplicates.
"""
function fem_impose_Dirichlet_boundary_condition1(A::SparseMatrixCSC{Float64,Int64},
    bd::Array{Int64}, m::Int64, n::Int64, h::Float64)
    @warn "Consider imposing boundary conditions using an algebraic approach: impose_Dirichlet_boundary_conditions" maxlog=1
    M, N = size(A)
    Q = A[:, bd]; Q[bd,:] = spzeros(length(bd), length(bd))
    A[bd,:] = spzeros(length(bd), N)
    A[:,bd] = spzeros(M, length(bd))
    A[bd,bd] = spdiagm(0=>ones(length(bd)))
    return dropzeros(A),  dropzeros(Q)
end

####################### Interaction #######################
@doc raw"""
    compute_interaction_matrix(m::Int64, n::Int64, h::Float64)

Computes the interaction term
```math
\int_A p \delta \varepsilon_v\mathrm{d}x = \int_A p [1,1,0]B^T\delta u_A\mathrm{d}x
```
Here $\varepsilon_v = \text{tr}\; \varepsilon = \text{div}\; \mathbf{u}$.

The output is a $mn \times 2(m+1)(n+1)$ matrix.
"""
function compute_interaction_matrix(m::Int64, n::Int64, h::Float64)
    I = Int64[]
    J = Int64[]
    V = Float64[]
    function add(k1, k2, v)
        push!(I, k1)
        push!(J, k2)
        push!(V, v)
    end
    B = zeros(4, 3, 8)
    for i = 1:2
        for j = 1:2
            η = pts[i]; ξ = pts[j]
            B[(j-1)*2+i,:,:] = [
            -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η 0.0 0.0 0.0 0.0
            0.0 0.0 0.0 0.0 -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
            ]
        end
    end
    for i = 1:m
        for j = 1:n
            for p = 1:2
                for q = 1:2
                    idx = [(j-1)*(m+1)+i;(j-1)*(m+1)+i+1;j*(m+1)+i;j*(m+1)+i+1]
                    idx = [idx; idx .+ (m+1)*(n+1)]
                    Bk = B[(q-1)*2+p,:,:]
                    Bk = [1. 1. 0.]*Bk # 1 x 8
                    for s = 1:8
                        add((j-1)*m+i, idx[s], Bk[1,s]*0.25*h^2)
                    end
                end
            end
        end
    end
    sparse(I, J, V, m*n, 2(m+1)*(n+1))
end

####################### Fluids #######################
@doc raw"""
    compute_fvm_source_term(f::Array{Float64}, m::Int64, n::Int64, h::Float64)

Computes the source term
```math
\int_{A_i} f\mathrm{d}x
```
Here $f$ has length $4mn$ or $mn$. In the first case, an average value of four quadrature nodal values of $f$ is used per cell.
"""
function compute_fvm_source_term(f::Array{Float64}, m::Int64, n::Int64, h::Float64)
    out = zeros(m*n)
    @assert length(f) == 4*m*n || length(f)==m*n
    if length(f)==4*m*n
        f = (f[1:4:end] + f[2:4:end] + f[3:4:end] + f[4:4:end])/4
    end
    for i = 1:m
        for j = 1:n
            k = (j-1)*m + i
            out[k] = h^2*f[k]
        end
    end
    out
end

@doc raw"""
    compute_fvm_mechanics_term(u::Array{Float64}, m::Int64, n::Int64, h::Float64)

Computes the mechanic interaction term
```math
\int_{A_i} \varepsilon_v\mathrm{d}x
```
Here
```math
\varepsilon_v = \mathrm{tr} \varepsilon = \varepsilon_{xx} + \varepsilon_{yy}
```
Numerically, we have
```math
\varepsilon_v = [1 \ 1 \ 0] B^T \delta u_A
```
"""
function compute_fvm_mechanics_term(u::Array{Float64}, m::Int64, n::Int64, h::Float64)
    out = zeros(m*n)
    B = zeros(4, 3, 8)
    for i = 1:2
        for j = 1:2
            ξ = pts[i]; η = pts[j]
            B[(j-1)*2+i,:,:] = [
            -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η 0.0 0.0 0.0 0.0
            0.0 0.0 0.0 0.0 -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
            ]
        end
    end
    for i = 1:m
        for j = 1:n
            idx = [(j-1)*(m+1)+i;(j-1)*(m+1)+i+1;j*(m+1)+i;j*(m+1)+i+1]
            idx = [idx; idx .+ (m+1)*(n+1)]
            for p = 1:2
                for q = 1:2
                    Bk = B[(q-1)*2+p,:,:]
                    Bk = [1. 1. 0.]*Bk*u[idx]*0.25*h^2
                    out[(j-1)*m+i] += Bk[1,1]
                end
            end
        end
    end
    out
end

"""
    compute_fvm_mechanics_term(u::PyObject, m::Int64, n::Int64, h::Float64)
"""
function compute_fvm_mechanics_term(u::PyObject, m::Int64, n::Int64, h::Float64)
    volumetric_strain_ = load_op_and_grad(libadfem,"volumetric_strain")
    u,m_,n_,h = convert_to_tensor([u,m,n,h], [Float64,Int32,Int32,Float64])
    strain = volumetric_strain_(u,m_,n_,h)
    set_shape(strain, (m*n,))
end

@doc raw"""
    compute_fvm_tpfa_matrix(m::Int64, n::Int64, h::Float64)

Computes the term with two-point flux approximation
```math
\int_{A_i} \Delta p \mathrm{d}x = \sum_{j=1}^{n_{\mathrm{faces}}} (p_j-p_i)
```
![](https://raw.githubusercontent.com/ADCMEMarket/ADCMEImages/master/AdFem/tpfa.png)

!!! warning
    No flow boundary condition is assumed.
"""
function compute_fvm_tpfa_matrix(m::Int64, n::Int64, h::Float64)
    I = Int64[]; J = Int64[]; V = Float64[]
    function add(i, j, v)
        push!(I, i)
        push!(J, j)
        push!(V, v)
    end
    for i = 1:m
        for j = 1:n
            k = (j-1)*m + i
            for (ii,jj) in [(i+1,j),(i-1,j),(i,j+1),(i,j-1)]
                if 1<=ii<=m && 1<=jj<=n
                    kp = (jj-1)*m + ii
                    add(k, kp, 1.)
                    add(k, k, -1.)
                end
            end
        end
    end
    sparse(I, J, V, m*n, m*n)
end

@doc raw"""
    compute_fvm_tpfa_matrix(K::Array{Float64}, m::Int64, n::Int64, h::Float64)

Computes the term with two-point flux approximation with distinct permeability at each cell
```math
\int_{A_i} K_i \Delta p \mathrm{d}x = K_i\sum_{j=1}^{n_{\mathrm{faces}}} (p_j-p_i)
```

Note that $K$ is a length $mn$ vector, representing values per cell.
"""
function compute_fvm_tpfa_matrix(K::Array{Float64}, m::Int64, n::Int64, h::Float64)
    I = Int64[]; J = Int64[]; V = Float64[]
    function add(i, j, v)
        push!(I, i)
        push!(J, j)
        push!(V, v)
    end
    for i = 1:m
        for j = 1:n
            k = (j-1)*m + i
            for (ii,jj) in [(i+1,j),(i-1,j),(i,j+1),(i,j-1)]
                if 1<=ii<=m && 1<=jj<=n
                    kp = (jj-1)*m + ii
                    add(k, kp, K[k])
                    add(k, k, -K[k])
                end
            end
        end
    end
    sparse(I, J, V, m*n, m*n)
end

@doc raw"""
    compute_fvm_tpfa_matrix(K::Array{Float64}, bc::Array{Int64,2}, pval::Array{Float64,1}, m::Int64, n::Int64, h::Float64)

Computes the term with two-point flux approximation with distinct permeability at each cell

```math
\int_{A_i} K_i \Delta p \mathrm{d}x = K_i\sum_{j=1}^{n_{\mathrm{faces}}} (p_j-p_i)
```

Here $K$ is a length $mn$ vector, representing values per cell.

Additionally, Dirichlet boundary conditions are imposed on the boundary edges `bc` (a $N\times 2$ integer matrix),
i.e., the $i$-th edge has value `pval`. The ghost node method is used for imposing the Dirichlet boundary condition.
The other boundaries are no-blow boundaries, i.e., $\frac{\partial T}{\partial n} = 0$.
The function outputs a length $mn$ vector and $mn\times mn$ matrix $M$.

$$\int_{A_i} K_i \Delta p \mathrm{d}x = f_i + M_{i,:}\mathbf{p}$$

Returns both the sparse matrix `A` and the right hand side `rhs`.

!!! info
    `K` can also be missing, in which case `K` is treated as a all-one vector.
"""
function compute_fvm_tpfa_matrix(K::Union{Array{Float64}, Missing}, bc::Array{Int64,2}, pval::Array{Float64,1}, m::Int64, n::Int64, h::Float64)
    bval = Dict{Tuple{Int64, Int64}, Float64}()
    for i = 1:size(bc, 1)
        bval[(minimum(bc[i,:]), maximum(bc[i,:]))] = pval[i]
    end
    I = Int64[]; J = Int64[]; V = Float64[]
    function add(i, j, v)
        push!(I, i)
        push!(J, j)
        push!(V, v)
    end
    K = coalesce(K, ones(m*n))
    @assert length(K) == m*n # in case users input a 4mn vector.
    rhs = zeros(m*n)
    for i = 1:m
        for j = 1:n
            k = (j-1)*m + i
            for (ii,jj) in [(i+1,j),(i-1,j),(i,j+1),(i,j-1)]
                if 1<=ii<=m && 1<=jj<=n
                    kp = (jj-1)*m + ii
                    add(k, kp, K[k])
                    add(k, k, -K[k])
                else  # boundary edges
                    ed = (0, 0)
                    if ii<=0 # left
                        ed = ((j-1)*(m+1)+1, j*(m+1)+1)
                    elseif ii>m # right
                        ed = ((j-1)*(m+1)+m+1, j*(m+1)+m+1)
                    elseif jj<=0 # upper
                        ed = (i, i+1)
                    elseif jj>n # lower
                        ed = (n*(m+1)+i, n*(m+1)+i+1)
                    end
                    if haskey(bval, ed)
                        add(k, k, -2*K[k])
                        rhs[k] += 2*K[k]*bval[ed]
                    end
                end
            end
        end
    end
    sparse(I, J, V, m*n, m*n), rhs
end

@doc raw"""
    compute_fem_traction_term(t::Array{Float64, 2},
    bdedge::Array{Int64,2}, m::Int64, n::Int64, h::Float64)

Computes the traction term
```math
\int_{\Gamma} t(\mathbf{n})\cdot\delta u \mathrm{d}
```

The number of rows of `t` is equal to the number of edges in `bdedge`.
The first component of `t` describes the $x$ direction traction, while the second
component of `t` describes the $y$ direction traction.

Also see [`compute_fem_normal_traction_term`](@ref).

![](https://raw.githubusercontent.com/ADCMEMarket/ADCMEImages/master/AdFem/traction.png)
"""
function compute_fem_traction_term(t::Array{Float64, 2},
    bdedge::Array{Int64,2}, m::Int64, n::Int64, h::Float64)
    @assert size(t,1)==size(bdedge,1) || size(t,2)==2
    rhs = zeros(2*(m+1)*(n+1))
    for k = 1:size(bdedge, 1)
        ii, jj = bdedge[k,:]
        rhs[ii] += t[k,1]*0.5*h
        rhs[jj] += t[k,1]*0.5*h
        rhs[ii+(m+1)*(n+1)] += t[k,2]*0.5*h
        rhs[jj+(m+1)*(n+1)] += t[k,2]*0.5*h
    end
    rhs
end

@doc raw"""
    compute_fem_traction_term1(t::Array{Float64, 2},
    bdedge::Array{Int64,2}, m::Int64, n::Int64, h::Float64)

Computes the traction term
```math
\int_{\Gamma} t(n) \delta u \mathrm{d}
```

The number of rows of `t` is equal to the number of edges in `bdedge`.
The output is a length $(m+1)*(n+1)$ vector.

Also see [`compute_fem_traction_term`](@ref).
"""
function compute_fem_traction_term1(t::Array{Float64, 1},
    bdedge::Array{Int64,2}, m::Int64, n::Int64, h::Float64)
    @assert length(t)==size(bdedge,1)
    rhs = zeros((m+1)*(n+1))
    for k = 1:size(bdedge, 1)
        ii, jj = bdedge[k,:]
        rhs[ii] += t[k,1]*0.5*h
        rhs[jj] += t[k,1]*0.5*h
    end
    rhs
end

@doc raw"""
    compute_fem_flux_term1(t::Array{Float64},
    bdedge::Array{Int64,2}, m::Int64, n::Int64, h::Float64)

Computes the traction term
```math
\int_{\Gamma} q \delta u \mathrm{d}
```
"""
function compute_fem_flux_term1(q::Array{Float64},
    bdedge::Array{Int64,2}, m::Int64, n::Int64, h::Float64)
    @assert length(size(q))==1
    rhs = zeros((m+1)*(n+1))
    for k = 1:size(bdedge, 1)
        ii, jj = bdedge[k,:]
        rhs[ii] += q[k]*0.5*h
        rhs[jj] += q[k]*0.5*h
    end
    rhs
end

@doc raw"""
    compute_fem_normal_traction_term(t::Array{Float64,1}, bdedge::Array{Int64},
    m::Int64, n::Int64, h::Float64)
    compute_fem_normal_traction_term(t::Float64, bdedge::Array{Int64},
    m::Int64, n::Int64, h::Float64)

Computes the normal traction term
```math
\int_{\Gamma} t(\mathbf{n})\cdot\delta u \mathrm{d}
```
Here $t(\mathbf{n})\parallel\mathbf{n}$ points **outward** to the domain and the magnitude is given by `t`.
`bdedge` is a $N\times2$ matrix and each row denotes the indices of two endpoints of the boundary edge.

See [`compute_fem_traction_term`](@ref) for graphical illustration.
"""
function compute_fem_normal_traction_term(t::Array{Float64,1}, bdedge::Array{Int64},
    m::Int64, n::Int64, h::Float64)
    @assert size(t,1)==size(bdedge,1)
    rhs = zeros(2*(m+1)*(n+1))
    for k = 1:size(bdedge,1)
        normal_ = get_edge_normal(bdedge[k,:], m, n, h)
        ii, jj = bdedge[k,:]
        rhs[ii] += t[k]*normal_[1]*0.5*h
        rhs[jj] += t[k]*normal_[1]*0.5*h
        rhs[ii+(m+1)*(n+1)] += t[k]*normal_[2]*0.5*h
        rhs[jj+(m+1)*(n+1)] += t[k]*normal_[2]*0.5*h
    end
    return rhs
end

function compute_fem_normal_traction_term(t::Float64, bdedge::Array{Int64},
    m::Int64, n::Int64, h::Float64)
    t = t * ones(size(bdedge,1))
    return compute_fem_normal_traction_term(t, bdedge, m, n, h)
end

@doc raw"""
    trim_coupled(pd::PoreData, Q::SparseMatrixCSC{Float64,Int64}, L::SparseMatrixCSC{Float64,Int64},
    M::SparseMatrixCSC{Float64,Int64},
    bd::Array{Int64}, Δt::Float64, m::Int64, n::Int64, h::Float64)

Assembles matrices from mechanics and flow and assemble the coupled matrix

$$\begin{bmatrix}
\hat M & -\hat L^T\\
\hat L & \hat Q
\end{bmatrix}$$

`Q` is obtained from [`compute_fvm_tpfa_matrix`](@ref), `M` is obtained from [`compute_fem_stiffness_matrix`](@ref),
and `L` is obtained from [`compute_interaction_matrix`](@ref).
"""
function trim_coupled(pd::PoreData, Q::SparseMatrixCSC{Float64,Int64}, L::SparseMatrixCSC{Float64,Int64},
    M::SparseMatrixCSC{Float64,Int64},
    bd::Array{Int64}, Δt::Float64, m::Int64, n::Int64, h::Float64)

    A22 = 1/pd.M/Δt*h^2*spdiagm(0=>ones(m*n)) -
    pd.kp/pd.Bf/pd.μ*Q
    A21 = pd.b/Δt * L
    A12 = -A21'
    A11 = M

    A = [A11 A12
    A21 A22]
    A[bd,:] = spzeros(length(bd), 2(m+1)*(n+1)+m*n)
    A[:, bd] = spzeros(2(m+1)*(n+1)+m*n, length(bd))
    A[bd, bd] = spdiagm(0=>ones(length(bd)))
    dropzeros(A)
end

"""
    coupled_impose_pressure(A::SparseMatrixCSC{Float64,Int64}, pnode::Array{Int64},
    m::Int64, n::Int64, h::Float64)

Returns a trimmed matrix.
"""
function coupled_impose_pressure(A::SparseMatrixCSC{Float64,Int64}, pnode::Array{Int64},
    m::Int64, n::Int64, h::Float64)
    pnode = pnode.+2(m+1)*(n+1)
    A[pnode, :] = spzeros(length(pnode), 2(m+1)*(n+1)+m*n)
    A[pnode, pnode] = spdiagm(0=>ones(length(pnode)))
    A
end

function get_gauss_points(m, n, h)
    pts_ = []
    for i = 1:m
        for j = 1:n
            for p = 1:2
                for q = 1:2
                    x = pts[p]*h+(i-1)*h
                    y = pts[q]*h+(j-1)*h
                    push!(pts_, [x y])
                end
            end
        end
    end
    vcat(pts_...)
end

@doc raw"""
    compute_plane_strain_matrix(E::Float64, ν::Float64)

Computes the stiffness matrix for 2D plane strain. The matrix is given by

$$\frac{E(1-\nu)}{(1+\nu)(1-2\nu)}\begin{bmatrix}
1 & \frac{\nu}{1-\nu} & \frac{\nu}{1-\nu}\\
\frac{\nu}{1-\nu} & 1 & \frac{\nu}{1-\nu} \\
\frac{\nu}{1-\nu} & \frac{\nu}{1-\nu} & 1
\end{bmatrix}$$
"""
function compute_plane_strain_matrix(E::Float64, ν::Float64)
    E*(1-ν)/(1+ν)/(1-2ν)*[
    1 ν/(1-ν) ν/(1-ν)
    ν/(1-ν) 1 ν/(1-ν)
    ν/(1-ν) ν/(1-ν) 1
    ]
end

@doc raw"""
    compute_plane_stress_matrix(E::Float64, ν::Float64)

Computes the stiffness matrix for 2D plane stress. The matrix is given by

$$\frac{E}{(1+\nu)(1-2\nu)}\begin{bmatrix}
1-\nu & \nu & 0\\
\nu & 1 & 0 \\
0 & 0 & \frac{1-2\nu}{2}
\end{bmatrix}$$
"""
function compute_plane_stress_matrix(E::Float64, ν::Float64)
    E/(1+ν)/(1-2ν)*[
        1-ν ν 0.0
        ν 1-ν 0.0
        0.0 0.0 (1-2ν)/2
    ]
end

@doc raw"""
    compute_plane_strain_matrix(E::Union{PyObject, Array{Float64, 1}}, nu::Union{PyObject, Array{Float64, 1}})

Returns the pointwise plane strain matrix of size $N\times 3 \times 3$. Here $N$ is the length of $E$ or $\nu$.
$N$ can be any number.
"""
function compute_plane_strain_matrix(E::Union{PyObject, Array{Float64, 1}}, nu::Union{PyObject, Array{Float64, 1}})
    mode = 0
    N = length(nu)
    @assert length(E)==N
    plane_strain_and_stress_ = load_op_and_grad(AdFem.libmfem,"plane_strain_and_stress")
    e,nu,mode = convert_to_tensor(Any[E,nu,mode], [Float64,Float64,Int32])
    out = plane_strain_and_stress_(e,nu,mode)
    set_shape(out, (N, 3, 3))
end

@doc raw"""
    compute_plane_stress_matrix(E::Union{PyObject, Array{Float64, 1}}, nu::Union{PyObject, Array{Float64, 1}})

Returns the pointwise plane stress matrix of size $N\times 3 \times 3$. Here $N$ is the length of $E$ or $\nu$.
$N$ can be any number.
"""
function compute_plane_stress_matrix(E::Union{PyObject, Array{Float64, 1}}, nu::Union{PyObject, Array{Float64, 1}})
    mode = 1
    N = length(nu)
    @assert length(E)==N
    plane_strain_and_stress_ = load_op_and_grad(AdFem.libmfem,"plane_strain_and_stress")
    e,nu,mode = convert_to_tensor(Any[E,nu,mode], [Float64,Float64,Int32])
    out = plane_strain_and_stress_(e,nu,mode)
    set_shape(out, (N, 3, 3))
end

"""
    compute_von_mises_stress_term(K::Array{Float64}, u::Array{Float64}, m::Int64, n::Int64, h::Float64)

Compute the [von Mises stress](https://en.wikipedia.org/wiki/Von_Mises_yield_criterion#Multi-axial_(2D_or_3D)_stress) on the Gauss quadrature nodes.
"""
function compute_von_mises_stress_term(K::Array{Float64}, u::Array{Float64}, m::Int64, n::Int64, h::Float64;
    b::Float64 = 1.0)
    I = Int64[]; J = Int64[]; V = Float64[]
    B = zeros(4, 3, 8)
    for i = 1:2
        for j = 1:2
            ξ = pts[i]; η = pts[j]
            B[(j-1)*2+i,:,:] = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η 0.0 0.0 0.0 0.0
                0.0 0.0 0.0 0.0 -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
                -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
            ]
        end
    end

    pval = Float64[]
    for j = 1:n
        for i = 1:m
            idx = [(j-1)*(m+1)+i;(j-1)*(m+1)+i+1;j*(m+1)+i;j*(m+1)+i+1]
            idx = [idx; idx .+ (m+1)*(n+1)]
            uA = u[idx]
            for p = 1:2
                for q = 1:2
                    Bk = B[(q-1)*2+p,:,:]
                    σ = K * Bk * uA # 3 components
                    σ11, σ22, σ12 = σ
                    σv = sqrt(σ11^2 - σ11*σ22 + σ22^2 + 3σ12^2)
                    push!(pval, σv)
                end
            end
        end
    end
    return pval
end

@doc raw"""
    compute_von_mises_stress_term(Se::Array{Float64,2},  m::Int64, n::Int64, h::Float64)

`Se` is a $4mn\times3$ array that stores the stress data at each Gauss point.
"""
function compute_von_mises_stress_term(Se::Array{Float64,2},  m::Int64, n::Int64, h::Float64)
    pval = Float64[]
    S = zeros(4*m*n)
    for i = 1:4*m*n
        σ11, σ22, σ12 = Se[i,:]
        S[i] = sqrt(σ11^2 - σ11*σ22 + σ22^2 + 3σ12^2)
    end
    return (S[1:4:end]+S[2:4:end]+S[3:4:end]+S[4:4:end])/4
end

@doc raw"""
    compute_fem_mass_matrix1(ρ::Array{Float64,1}, m::Int64, n::Int64, h::Float64)

Computes the mass matrix for a scalar value $u$
```math
\int_A \rho u \delta u \mathrm{d} x
```
The output is a $(m+1)(n+1)\times (m+1)(n+1)$ sparse matrix.
"""
function compute_fem_mass_matrix1(ρ::Array{Float64}, m::Int64, n::Int64, h::Float64)
    I = Int64[]; J = Int64[]; V = Float64[]
    function add(ii, jj, vv)
        for i = 1:4
            for j = 1:4
                push!(I,ii[i])
                push!(J,jj[j])
                push!(V,vv[i,j])
            end
        end
    end

    Me = zeros(4, 4, 4)
    for p = 1:2
        for q = 1:2
            ξ = pts[p]; η = pts[q]
            A = zeros(4)
            A[1] = (1-ξ)*(1-η)
            A[2] = ξ*(1-η)
            A[3] = (1-ξ)*η
            A[4] = ξ*η
            Me[(p-1)*2+q,:,:] = A*A'*0.25*h^2
        end
    end
    for i = 1:m
        for j = 1:n
            elem_idx = (j-1)*m + i
            idx = [i+(j-1)*(m+1); i+1+(j-1)*(m+1); i+j*(m+1); i+1+j*(m+1)]
            for p = 1:2
                for q = 1:2
                    k = (elem_idx-1)*4 + 2*(q-1) + p
                    ρ_ = ρ[k]
                    add(idx, idx, ρ_*Me[(p-1)*2+q,:,:])
                end
            end
        end
    end
    sparse(I, J, V, (m+1)*(n+1), (m+1)*(n+1))
end

@doc raw"""
    compute_fem_mass_matrix1(m::Int64, n::Int64, h::Float64)

Computes the mass matrix for a scalar value $u$
```math
\int_A u \delta u \mathrm{d} x
```
The output is a $(m+1)(n+1)\times (m+1)(n+1)$ sparse matrix.
"""
function compute_fem_mass_matrix1(m::Int64, n::Int64, h::Float64)
    ρ = ones(4*m*n)
    compute_fem_mass_matrix1(ρ, m, n, h)
end

@doc raw"""
    compute_fem_mass_matrix(m::Int64, n::Int64, h::Float64)

Computes the finite element mass matrix

```math
\int_{\Omega} u \delta u \mathrm{d}x
```

The matrix size is $2(m+1)(n+1) \times 2(m+1)(n+1)$.
"""
function compute_fem_mass_matrix(m::Int64, n::Int64, h::Float64)
    M = compute_fem_mass_matrix1(m, n, h)
    Z = spzeros((m+1)*(n+1), (m+1)*(n+1))
    [M Z;Z M]
end

@doc raw"""
    eval_f_on_gauss_pts(f::Function, m::Int64, n::Int64, h::Float64; tensor_input::Bool = false)

Evaluates `f` at Gaussian points and return the result as $4mn$ vector `out` (4 Gauss points per element)

If `tensor_input = true`, the function `f` is assumed to map a tensor to a tensor output.

![](https://raw.githubusercontent.com/ADCMEMarket/ADCMEImages/master/AdFem/gauss.png)
"""
function eval_f_on_gauss_pts(f::Function, m::Int64, n::Int64, h::Float64; tensor_input::Bool = false)
    xy = gauss_nodes(m, n, h)
    x, y = xy[:,1], xy[:,2]
    if tensor_input
        return f(constant(x), constant(y))
    end
    return f.(x, y)
end

@doc raw"""
    eval_f_on_fem_pts(f::Function, m::Int64, n::Int64, h::Float64; tensor_input::Bool = false)

Returns $f(x_i, y_i)$ where $(x_i,y_i)$ are FEM nodes.

If `tensor_input = true`, the function `f` is assumed to map a tensor to a tensor output.
"""
function eval_f_on_fem_pts(f::Function, m::Int64, n::Int64, h::Float64; tensor_input::Bool = false)
    xy = fem_nodes(m, n, h)
    x, y = xy[:,1], xy[:,2]
    if tensor_input
        return f(constant(x), constant(y))
    end
    return f.(x, y)
end

@doc raw"""
    eval_f_on_fvm_pts(f::Function, m::Int64, n::Int64, h::Float64; tensor_input::Bool = false)

Returns $f(x_i, y_i)$ where $(x_i,y_i)$ are FVM nodes.

If `tensor_input = true`, the function `f` is assumed to map a tensor to a tensor output.
"""
function eval_f_on_fvm_pts(f::Function, m::Int64, n::Int64, h::Float64; tensor_input::Bool = false)
    xy = fvm_nodes(m, n, h)
    x, y = xy[:,1], xy[:,2]
    if tensor_input
        return f(constant(x), constant(y))
    end
    return f.(x, y)
end

@doc raw"""
    eval_f_on_boundary_node(f::Function, bdnode::Array{Int64}, m::Int64, n::Int64, h::Float64)

Returns a vector of the same length as `bdnode` whose entries corresponding to `bdnode` nodes
are filled with values computed from `f`.

`f` has the following signature
```
f(x::Float64, y::Float64)::Float64
```
"""
function eval_f_on_boundary_node(f::Function, bdnode::Array{Int64}, m::Int64, n::Int64, h::Float64)
    out = zeros(length(bdnode))
    for i = 1:length(bdnode)
        i1, j1 = femidx(bdnode[i], m)
        out[i] = f((i1-1)*h, (j1-1)*h)
    end
    out
end

@doc raw"""
    eval_f_on_boundary_edge(f::Function, bdedge::Array{Int64,2}, m::Int64, n::Int64, h::Float64)

Returns a vector of the same length as `bdedge` whose entries corresponding to `bdedge` nodes
are filled with values computed from `f`.

`f` has the following signature
```
f(x::Float64, y::Float64)::Float64
```
"""
function eval_f_on_boundary_edge(f::Function, bdedge::Array{Int64,2}, m::Int64, n::Int64, h::Float64)
    out = zeros(size(bdedge,1))
    for i = 1:size(bdedge,1)
        i1, j1 = femidx(bdedge[i,1], m)
        i2, j2 = femidx(bdedge[i,2], m)
        x = ((i1-1)*h + (i2-1)*h)/2
        y = ((j1-1)*h + (j2-1)*h)/2
        out[i] = f(x, y)
    end
    out
end

@doc raw"""
    eval_strain_on_gauss_pts(u::Array{Float64}, m::Int64, n::Int64, h::Float64)

Computes the strain on Gauss points.
Returns a $4mn\times3$ matrix, where each row denotes $(\varepsilon_{11}, \varepsilon_{22}, 2\varepsilon_{12})$
at the corresponding Gauss point.
"""
function eval_strain_on_gauss_pts(u::Array{Float64}, m::Int64, n::Int64, h::Float64)
    I = Int64[]
    J = Int64[]
    V = Float64[]
    function add(k1, k2, v)
        push!(I, k1)
        push!(J, k2)
        push!(V, v)
    end
    B = zeros(4, 3, 8)
    for i = 1:2
        for j = 1:2
            ξ = pts[i]; η = pts[j]
            B[(i-1)*2+j,:,:] = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η 0.0 0.0 0.0 0.0
                0.0 0.0 0.0 0.0 -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
                -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
            ]
        end
    end
    strain = zeros(4m*n, 3)
    for i = 1:m
        for j = 1:n
            for p = 1:2
                for q = 1:2
                    idx = [(j-1)*(m+1)+i;(j-1)*(m+1)+i+1;j*(m+1)+i;j*(m+1)+i+1]
                    idx = [idx; idx .+ (m+1)*(n+1)]
                    Bk = B[(p-1)*2+q,:,:] # 3 x 8
                    strain[4*((j-1)*m+i-1)+(p-1)*2+q,:] = Bk * u[idx]
                end
            end
        end
    end
    strain
end

@doc raw"""
    eval_stress_on_gauss_pts(u::Array{Float64}, K::Array{Float64,2}, m::Int64, n::Int64, h::Float64)

Returns the stress on the Gauss points for elasticity.
"""
function eval_stress_on_gauss_pts(u::Array{Float64}, K::Array{Float64,2}, m::Int64, n::Int64, h::Float64)
    strain = eval_strain_on_gauss_pts(u, m, n, h)
    strain * K
end

@doc raw"""
    compute_fvm_mass_matrix(m::Int64, n::Int64, h::Float64)

Returns the FVM mass matrix
```math
\int_{A_i} p_i \mathrm{d}x = h^2 p_i
```
"""
function compute_fvm_mass_matrix(m::Int64, n::Int64, h::Float64)
    return spdiagm(0=>ones(m*n))*h^2
end

@doc raw"""
    compute_strain_energy_term(S::Array{Float64, 2}, m::Int64, n::Int64, h::Float64)

Computes the strain energy
```math
\int_{A} \sigma : \delta \varepsilon \mathrm{d}x
```
where $\sigma$ is provided by `S`, a $4mn \times 3$ matrix.
The values $\sigma_{11}, \sigma_{22}, \sigma_{12}$ are defined on 4 Gauss points per element.
"""
function compute_strain_energy_term(S::Array{Float64, 2}, m::Int64, n::Int64, h::Float64)
    f = zeros(2(m+1)*(n+1))

    B = zeros(4, 3, 8)
    for i = 1:2
        for j = 1:2
            ξ = pts[i]; η = pts[j]
            B[(i-1)*2+j,:,:] = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η 0.0 0.0 0.0 0.0
                0.0 0.0 0.0 0.0 -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
                -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
            ]
        end
    end

    for i = 1:m
        for j = 1:n
            elem = (j-1)*m + i
            σ = S[4(elem-1)+1:4elem, :] # 4×3
            dof = [(j-1)*(m+1)+i;(j-1)*(m+1)+i+1;j*(m+1)+i;j*(m+1)+i+1]
            dof = [dof; dof .+ (m+1)*(n+1)]
            for p = 1:2
                for q = 1:2
                    idx = 2(q-1) + p
                    f[dof] += (σ[idx,:]' * B[idx,:,:])[:]*0.25*h^2 # length 8 vector
                end
            end
        end
    end
    f
end

@doc raw"""
    compute_strain_energy_term1(S::PyObject, m::Int64, n::Int64, h::Float64)

Computes the strain energy
```math
\int_{A} \sigma : \delta \varepsilon \mathrm{d}x
```
where $\sigma$ is provided by `S`, a $4mn \times 2$ matrix.
The values $\sigma_{31}, \sigma_{32}$ are defined on 4 Gauss points per element.
"""
function compute_strain_energy_term1(S::Array{Float64, 2}, m::Int64, n::Int64, h::Float64)
    @assert size(S,2)==2
    f = zeros((m+1)*(n+1))

    B = zeros(4, 2, 4)
    for p = 1:2
        for q = 1:2
            ξ = pts[p]; η = pts[q]
            B[(q-1)*2+p,:,:] = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
                -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            ]
        end
    end

    for i = 1:m
        for j = 1:n
            elem = (j-1)*m + i
            σ = S[4(elem-1)+1:4elem, :] # 4×2
            dof = [(j-1)*(m+1)+i;(j-1)*(m+1)+i+1;j*(m+1)+i;j*(m+1)+i+1]
            for p = 1:2
                for q = 1:2
                    idx = 2(q-1) + p
                    f[dof] += (σ[idx,:]' * B[idx,:,:])[:]*0.25*h^2 # length 8 vector
                end
            end
        end
    end
    f
end

@doc raw"""
    compute_fem_laplace_matrix1(K::Array{Float64, 1}, m::Int64, n::Int64, h::Float64)

Computes the coefficient matrix for

```math
\int_\Omega K \nabla u \cdot \nabla (\delta u) \; dx
```
Here $K\in \mathbf{R}^{2\times 2}$, $u$ is a scalar variable, and `K` is a $4mn \times 2 \times 2$ matrix.

Returns a $(m+1)(n+1)\times (m+1)(n+1)$ sparse matrix.
"""
function compute_fem_laplace_matrix1(K::Array{Float64, 3}, m::Int64, n::Int64, h::Float64)
    @assert size(K,2)==size(K,3)==2

    B = zeros(4, 2, 4)
    for q = 1:2
        for p = 1:2
            ξ = pts[p]; η = pts[q]
            B[(q-1)*2+p,:,:] = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
                -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            ]
        end
    end

    I = Int64[]
    J = Int64[]
    V = Float64[]
    function add(k1, k2, v)
        push!(I, k1)
        push!(J, k2)
        push!(V, v)
    end
    k = 0
    for j = 1:n
        for i = 1:m
            for q = 1:2
                for p = 1:2
                    k += 1
                    B0 = B[(q-1)*2+p,:,:]
                    Ω = B0'*K[k,:,:]*B0*0.25*h^2
                    idx = [(j-1)*(m+1)+i; (j-1)*(m+1)+i+1; j*(m+1)+i; j*(m+1)+i+1]
                    for i_ = 1:4
                        for j_ = 1:4
                            add(idx[i_], idx[j_], Ω[i_, j_])
                        end
                    end
                end
            end
        end
    end

    return sparse(I,J,V,(m+1)*(n+1),(m+1)*(n+1))
end

"""
    compute_fem_laplace_matrix1(K::Array{Float64, 2}, m::Int64, n::Int64, h::Float64)

`K` is duplicated on each Gauss point.
"""
function compute_fem_laplace_matrix1(K::Array{Float64, 2}, m::Int64, n::Int64, h::Float64)
    Knew = zeros(4m*n, 2, 2)
    for i = 1:4*m*n
        Knew[i,:,:] = K
    end
    return compute_fem_laplace_matrix1(Knew, m, n, h)
end

@doc raw"""
    compute_fem_laplace_matrix1(K::Array{Float64, 1}, m::Int64, n::Int64, h::Float64)

Computes the coefficient matrix for

```math
\int_\Omega K\nabla u \cdot \nabla (\delta u) \; dx
```

Here `K` is a vector with length $4mn$ (defined on Gauss points).

Returns a $(m+1)(n+1)\times (m+1)(n+1)$ sparse matrix.
"""
function compute_fem_laplace_matrix1(K::Array{Float64, 1}, m::Int64, n::Int64, h::Float64)
    @assert length(K)==4*m*n
    Knew = zeros(4m*n, 2, 2)
    for i = 1:4*m*n
        Knew[i,:,:] = K[i] * diagm(0=>ones(2))
    end
    return compute_fem_laplace_matrix1(Knew, m, n, h)
end

@doc raw"""
    compute_fem_laplace_matrix1(m::Int64, n::Int64, h::Float64)

Computes the coefficient matrix for

```math
\int_\Omega \nabla u \cdot \nabla (\delta u) \; dx
```

Returns a $(m+1)(n+1)\times (m+1)(n+1)$ sparse matrix.
"""
function compute_fem_laplace_matrix1(m::Int64, n::Int64, h::Float64)
    Knew = zeros(4m*n, 2, 2)
    for i = 1:4*m*n
        Knew[i,:,:] = diagm(0=>ones(2))
    end
    return compute_fem_laplace_matrix1(Knew, m, n, h)
end

@doc raw"""
    compute_fem_laplace_matrix(m::Int64, n::Int64, h::Float64)

Computes the coefficient matrix for

```math
\int_\Omega \nabla \mathbf{u} \cdot \nabla (\delta \mathbf{u}) \; dx
```

Here

$$\mathbf{u}  = \begin{bmatrix} u \\ v \end{bmatrix}$$

and

$$\nabla \mathbf{u} = \begin{bmatrix}u_x & u_y \\ v_x & v_y \end{bmatrix}$$

Returns a $2(m+1)(n+1)\times 2(m+1)(n+1)$ sparse matrix.
"""
function compute_fem_laplace_matrix(m::Int64, n::Int64, h::Float64)
    K = compute_fem_laplace_matrix1(m, n, h)
    Z = spzeros((m+1)*(n+1), (m+1)*(n+1))
    [K Z;Z K]
end

@doc raw"""
    compute_fem_laplace_matrix(K::Array{Float64, 1}, m::Int64, n::Int64, h::Float64)

Computes the coefficient matrix for

```math
\int_\Omega K \nabla \mathbf{u} \cdot \nabla (\delta \mathbf{u}) \; dx
```

Here $K$ is a scalar defined on Gauss points. `K` is a vector of length $4mn$
"""
function compute_fem_laplace_matrix(K::Array{Float64, 1}, m::Int64, n::Int64, h::Float64)
    K1 = compute_fem_laplace_matrix1(K, m, n, h)
    K2 = compute_fem_laplace_matrix1(K, m, n, h)
    Z = spzeros((m+1)*(n+1), (m+1)*(n+1))
    [K1 Z;Z K2]
end

@doc raw"""
    eval_grad_on_gauss_pts1(u::Array{Float64,1}, m::Int64, n::Int64, h::Float64)

Evaluates $\nabla u$ on each Gauss point. Here $u$ is a scalar function.

The input `u` is a vector of length $(m+1)*(n+1)$. The output is a matrix of size $4mn\times 2$.
"""
function eval_grad_on_gauss_pts1(u::Array{Float64,1}, m::Int64, n::Int64, h::Float64)
    @assert length(u) == (m+1)*(n+1)
    ret = zeros(4*m*n, 2)
    B = zeros(4, 2, 4)
    for q = 1:2
        for p = 1:2
            k = (q-1)*2 + p
            ξ = pts[p]; η = pts[q]
            B[k, :, :] = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
                -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            ]
        end
    end
    ret_idx = 0
    for j = 1:n
        for i = 1:m
            idx = [(j-1)*(m+1)+i; (j-1)*(m+1)+i+1; j*(m+1)+i; j*(m+1)+i+1 ]
            uA = u[idx]
            for q = 1:2
                for p = 1:2
                    k = (q-1)*2 + p
                    B0 = B[k, :, :]
                    ret_idx += 1
                    ret[ret_idx, :] = B0 * uA
                end
            end
        end
    end
    return ret
end

@doc raw"""
    eval_grad_on_gauss_pts(u::Array{Float64,1}, m::Int64, n::Int64, h::Float64)

Evaluates $\nabla u$ on each Gauss point. Here $\mathbf{u} = (u, v)$.

$$\texttt{g[i,:,:]} = \begin{bmatrix} u_x & u_y\\ v_x & v_y \end{bmatrix}$$

The input `u` is a vector of length $2(m+1)*(n+1)$. The output is a matrix of size $4mn\times 2 \times 2$.
"""
function eval_grad_on_gauss_pts(u::Array{Float64,1}, m::Int64, n::Int64, h::Float64)
    r1 = eval_grad_on_gauss_pts1(u[1:(m+1)*(n+1)], m, n, h)
    r2 = eval_grad_on_gauss_pts1(u[(m+1)*(n+1)+1:end], m, n, h)
    ret = zeros(4*m*n, 2, 2)
    for i = 1:4*m*n
        ret[i, 1, :] = r1[i,:]
        ret[i, 2, :] = r2[i,:]
    end
    return ret
end

@doc raw"""
    compute_interaction_term(p::Array{Float64, 1}, m::Int64, n::Int64, h::Float64)

Computes the FVM-FEM interaction term

```math
 \begin{bmatrix} \int p \frac{\partial \delta u}{\partial x} dx \\  \int p \frac{\partial \delta v}{\partial y}  dy \end{bmatrix}
```

The input is a vector of length $mn$. The output is a $2(m+1)(n+1)$ vector.
"""
function compute_interaction_term(pres::Array{Float64, 1}, m::Int64, n::Int64, h::Float64)
    @assert length(pres)==m*n
    rhs = zeros((m+1)*(n+1)*2)
    B = zeros(4, 2, 8)
    for i = 1:2
        for j = 1:2
            η = pts[i]; ξ = pts[j]
            B[(j-1)*2+i,:,:] = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η 0.0 0.0 0.0 0.0
                0.0 0.0 0.0 0.0 -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            ]
        end
    end
    for i = 1:m
        for j = 1:n
            pA = pres[(j-1)*m+i]
            for p = 1:2
                for q = 1:2
                    idx1 = [(j-1)*(m+1)+i;(j-1)*(m+1)+i+1;j*(m+1)+i;j*(m+1)+i+1]
                    idx2 = idx1 .+ (m+1)*(n+1)
                    Bk = B[(q-1)*2+p,:,:]
                    rhs[idx1] += B[(q-1)*2+p,1,1:4] * pA * h * h * 0.25
                    rhs[idx2] += B[(q-1)*2+p,2,5:end] * pA * h * h * 0.25
                end
            end
        end
    end
    rhs
end

function compute_fem_advection_matrix1(u0::Array{Float64, 1},v0::Array{Float64, 1},m::Int64,n::Int64,h::Float64)
    @assert length(u0) == 4*m*n
    @assert length(v0) == 4*m*n

    B = zeros(4, 2, 4)
    for q = 1:2
        for p = 1:2
            ξ = pts[p]; η = pts[q]
            B[(q-1)*2+p,:,:] = [
                -1/h*(1-η) 1/h*(1-η) -1/h*η 1/h*η
                -1/h*(1-ξ) -1/h*ξ 1/h*(1-ξ) 1/h*ξ
            ]
        end
    end

    I = Int64[]
    J = Int64[]
    V = Float64[]
    function add(k1, k2, v)
        push!(I, k1)
        push!(J, k2)
        push!(V, v)
    end
    k = 0
    for j = 1:n
        for i = 1:m
            for q = 1:2
                for p = 1:2
                    k += 1
                    B0 = B[(q-1)*2+p,:,:]
                    Ω = B0'*K[k,:,:]*B0*0.25*h^2
                    idx = [(j-1)*(m+1)+i; (j-1)*(m+1)+i+1; j*(m+1)+i; j*(m+1)+i+1]
                    for i_ = 1:4
                        for j_ = 1:4
                            add(idx[i_], idx[j_], Ω[i_, j_])
                        end
                    end
                end
            end
        end
    end
end


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

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

module EulerTransforms

using Nemo, NumberTheory

export ModuleEulerTransforms
export EulerTransform
export V006171, L006171, V107895, L107895, V061256, L061256
export V190905, L190905, V275585, L275585, V290351, L290351
export V052847, L052847

"""
* V006171, L006171, V107895, L107895, V061256, L061256, V190905, L190905, V275585, L275585, V290351, L290351
"""
const ModuleEulerTransforms = ""

"""
Return the Euler transform of f.
"""
function EulerTransform(f::Function)
    function E(n)
        haskey(CacheET, (f, n)) && return CacheET[(f, n)]
        n == 0 && return ZZ(1)
        a(j) = ZZ(sum(d*f(Int(d)) for d in Divisors(j)))
        b = sum(a(j) * E(n-j) for j in 1:n)
        r = div(b, n)
        CacheET[(f, n)] = r
        return r
    end
end

const CacheET = Dict{Tuple{Function, Int}, fmpz}()

"""
Return the number of factorization patterns of polynomials of degree n over integers.
"""
V006171(n) = EulerTransform(τ)(n)

"""
Return a list of length len of the Euler transform of tau.
"""
L006171(len) = [V006171(n) for n in 0:len-1]

"""
Return the Euler transform of sigma.
"""
V061256(n) = EulerTransform(σ)(n)

"""
Return a list of length len of the Euler transform of sigma.
"""
L061256(len) = [V061256(n) for n in 0:len-1]

"""
Return the Euler transform of sigma_2.
"""
V275585(n) = EulerTransform(σ2)(n)

"""
Return a list of length len of the Euler transform of sigma_2.
"""
L275585(len) = [V275585(n) for n in 0:len-1]

"""
Return the Euler transform of the factorial.
"""
V107895(n) = EulerTransform(fac)(n)

"""
Return a list of length len of the Euler transform of the factorial.
"""
L107895(len) = [V107895(n) for n in 0:len-1]

"""
Return the Euler transform of the swinging factorial (A056040).
"""
V190905(n) = EulerTransform(n -> div(fac(n), fac(div(n, 2))^2))(n)

"""
Return a list of length len of the Euler transform of the swinging factorial.
"""
L190905(len) = [V190905(n) for n in 0:len-1]

"""
Return the Euler transform of the Bell numbers.
"""
V290351(n) = EulerTransform(Nemo.bell)(n)

"""
Return a list of length len of the Euler transform of the Bell numbers.
"""
L290351(len) = [V290351(n) for n in 0:len-1]

"""
Return the Euler transform of [0, 1, 2, 3, ...].
"""
V052847(n) = EulerTransform(k -> k-1)(n)

"""
Return a list of length len of the Euler transform of [0, 1, 2, 3, ...].
"""
L052847(len) = [V052847(n) for n in 0:len-1]

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

using Test, SeqTests

function test()

    @testset "EulerTransform" begin

        if is_oeis_installed()

            L = [L107895, L190905, L061256, L275585, L006171, L290351, L052847]
            SeqTest(L, 'L')

            V = [V107895, V190905, V061256, V275585, V006171]
            for v in V SeqTest(v, 'V', 0) end
        end
    end
end

function demo()
end

"""
for n in 0:150 V107895(n) end ::
    0.303605 seconds (1.01 M allocations: 32.276 MiB, 29.90% gc time)
"""
function perf()
    @time (for n in 0:150 V107895(n) end)
end

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

main()

end # module


function createPointStandards(n::Int64)
    points,weights  = computeXYZandWeights(n)
	return points,weights
end

function computeXYZandWeights(_norder::Int64)
        n = _norder
        pointsxyz = zeros(Float64,n*n,3); # Even though we only need the x and y coordinate
                                      # we still store the corresponding z coordinate
                                      # as our quadrature point actually lives on the sphere
        weights = zeros(n*n);

        # Construct Gauss quadrature
        mu,gaussweights = gausslegendre(n)

        # around z axis equidistant
        phi = [(k+0.5)*pi/n for k=0:2*n-1]

        range = 1:Int(n/2) # we only use the upper half of the sphere as quadrature point since we do pseudo three d

        # Transform between (mu,phi) and (x,y,z)
        x = sqrt.(1.0 .- mu[range].^2).*cos.(phi)'
        y = sqrt.(1.0 .- mu[range].^2).*sin.(phi)'
        z =           mu[range]    .*ones(size(phi))'
        weights = 2.0*pi/n*repeat(gaussweights[range],1,2*n)

        # assign
        pointsxyz[:,1] = x[:]
        pointsxyz[:,2] = y[:]
        pointsxyz[:,3] = z[:]

        weights = weights[:]

        return pointsxyz, weights
    end


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

using Markdown
using InteractiveUtils

# ╔═╡ e9fdca3a-2d74-4a42-acda-2b289172a1d4
begin
	import Pkg
	Pkg.activate(tempname())
	Pkg.add("DynamicPolynomials")
	Pkg.add("NLopt")
	Pkg.add("StaticArrays")
	Pkg.add("GLMakie")
	Pkg.add("Makie")
	Pkg.add("Combinatorics")
end

# ╔═╡ 9956f874-e481-4e3e-97fc-150aef8b69fa
begin
	using DynamicPolynomials
	using BenchmarkTools
	import NLopt
	using StaticArrays
	using Makie, GLMakie
	using Combinatorics
end

# ╔═╡ 0bf0c2a6-7487-49a9-8adf-73381ea2e9fc
md"""
We denote by ``Π_n^d`` the space of ``n``-variate Polyoniams of degree at most ``d``.
To construct polynomials we use `DynamicPolynomials.jl`.
"""

# ╔═╡ af1b5146-5b3a-425d-95a7-6d0532bcca2f
md"""
The canonical basis is obtained by calculating the non-negative integral solutions to the euqation
```math
x_1 + … + x_n \le d.
```
"""

# ╔═╡ 6c51dda0-5c69-4f58-8348-d62245b1eaeb
md"""
These solutions can be found using the `Combinatorics` package via `multiexponents(n,d)` (`d` must be successively increased).
"""

# ╔═╡ 802e1cd2-fceb-4ee7-b178-62b05ced5409
function non_negative_ineq_solutions(deg, n_vars)
	collect(Iterators.flatten( ( collect(multiexponents( n_vars, d )) for d = 0 : deg ) ))
end

# ╔═╡ c94bdc4e-2c95-42ae-a7dd-4172da278585
function get_poly_basis( deg, n_vars)
	exponents = non_negative_ineq_solutions(deg, n_vars )
	polys = let
		@polyvar x[1:n_vars]
		[ prod(x.^e) for e in exponents ]
	end
	return polys
end

# ╔═╡ b70e7d30-6856-4011-8a59-4727c690634f
md"""
We are going to use the canonical basis to determine a **poised set** of points.
This does in fact work with any polynomial basis for ``Π_n^d``. \
In the process of doing so, we also modify (a copy of?) the basis so that it becomes the Lagrange basis for the returned point set.
"""

# ╔═╡ 42ac42a6-d11e-4c03-bb18-0276d6a319fa
md"The Larange basis is formed by normalizing and orthogonalizing with respect to the point set:"

# ╔═╡ 688d4cf3-4e4d-45de-a43c-141646f01b00
function orthogonalize_polys( poly_arr, x, i )
	# normalize i-th polynomial with respect to `x`
	p_i = poly_arr[i] / poly_arr[i](x)

	# orthogoalize
	return [ j != i ? poly_arr[j] - poly_arr[j](x) * p_i : p_i for j = eachindex( poly_arr ) ]
end

# ╔═╡ e8136cef-9235-4ad8-bfac-a9e176923de7
md"""
We use Algorithm 6.2 and Algorithm 6.3 from the book
"Introduction to Derivative-Free Optimization"
by Conn et. al. \
Algorithm 6.2 makes the set poised (suited for interpolation) and returns the corresponding Lagrange basis.
Algorithm 6.3 takes the poised set and the Lagrange basis and tries to make it ``Λ``-poised. ``Λ`` must be greater 1 and a smaller value makes the set more suited for good models.
"""

# ╔═╡ f48b8d93-e98c-4f30-960c-203a8077b6b6
"""
    get_poised_set( basis, points; solver = :LN_BOBYQA, max_solver_evals = -1 )

Compute a point set suited for polynomial interpolation.

Input:
* `basis`: A vector of polynomials constituting a basis for the polynomial space.
* `points`: (optional) A set of candidate points to be tried for inclusion into the poised set.
* `solver`: NLopt solver to use. Should be derivative-free.
* `max_solver_evals`: Maximum number of evaluations in each optimization run.

Return:
* `poised_points :: Vector{T}` where `T` is either a `Vector{F}` or an `SVector{n_vars, F}` and `F` is the precision of the points in `points`, but at least `Float32`.
* `lagrange_basis :: Vector{<:AbstractPolynomialLike}`: The Lagrange basis corresponding to `poised_points`.
* `point_indices`: An array indicating which points from `points` are also in `poised_points`. A positive entry corresponds to the index of a poised point in `points`. If a poised point is new, then the entry is `-1`.
"""
function get_poised_set( basis, points :: AbstractArray{T} = Vector{Float32}[];
		solver = :LN_BOBYQA, max_solver_evals = -1 ) where {
		T <: AbstractArray{<:Real}
	}

	p = length(basis)
	@assert p > 0 "`basis` must not be an empty array."

	vars = variables( basis[end] )
	n_vars = length(vars)
	@assert n_vars > 0 "The number of variables must be positive."

	if max_solver_evals < 0
		max_solver_evals = 2000 * n_vars
	end

	F = promote_type( eltype( T ), Float32 )
	P_type = n_vars > 100 ? Vector{Vector{F}} : Vector{SVector{n_vars, F}}
	ZERO_TOL = min(eps(F) * 100, eps(Float16) * 10)

	# indicates which points from points have been accepted
	point_indices = fill(-1, p)
	not_accepted_indices = collect( eachindex( points ) )
	# return array of points that form a poised set
	poised_points = P_type(undef, p)

	new_basis = basis
	for i = 1 : p
		_points = points[not_accepted_indices]

		# find the point that maximizes the i-th polynomial
		# if the polynomial is constant, then the first remaining point is used (j = 1)
		l_max, j = if isempty(_points)
			0.0, 0
		else
			findmax( abs.( [ new_basis[i]( x ) for x in _points ] ) )
		end

		if l_max > ZERO_TOL
			# accept the `j`-th point from `_points`
			poised_points[i] = _points[j]
			## indicate what the actual point index was
			point_indices[i] = not_accepted_indices[j]
			## delete from further consideration
			deleteat!(not_accepted_indices, j)
		else
			# no point was suitable to add to the set
			# trying to find the maximizer for a | l_i(x) |
			opt = NLopt.Opt( solver, n_vars )
			opt.lower_bounds = zeros(F, n_vars )
            opt.upper_bounds = ones(F, n_vars )
            opt.maxeval = max_solver_evals
            opt.xtol_rel = 1e-3
            opt.max_objective = (x,g) -> abs( new_basis[i](x) )

			x₀_tmp = [ rand(F, n_vars) for i = 1 : 50 * n_vars ]
            x₀ = x₀_tmp[argmax( abs.(new_basis[i].(x₀_tmp)) ) ]

			_, ξ, ret = NLopt.optimize(opt, x₀)

			poised_points[i] = ξ
		end

		new_basis = orthogonalize_polys( new_basis, poised_points[i], i )
	end

	return poised_points, new_basis, point_indices
end

# ╔═╡ 85d345a0-2d3f-46ed-8b39-f88f700d545b
"""
    make_set_lambda_poised( basis, points;
        LAMBDA = 1.5, solver = :LN_BOBYQA, max_solver_evals = -1, max_loops = -1, skip_indices = [1,] )

Make the output of `get_poised_set` even better suited for interpolation.

Input:
* `basis`: A vector of polynomials constituting a Lagrange basis for the polynomial space.
* `points`: The vector of points belonging to the Lagrange basis.
* `LAMBDA :: Real > 1`: Determines the quality of the interpolation.
* `solver`: NLopt solver to use. Should be derivative-free.
* `max_solver_evals`: Maximum number of evaluations in each optimization run.
* `max_loops`: Maximum number of loops that try to make the set Λ-poised.
* `skip_indices`: Inidices of points to discard last.

Return:
* `poised_points :: Vector{T}` where `T` is either a `Vector{F}` or an `SVector{n_vars, F}` and `F` is the precision of the points in `points`, but at least `Float32`.
* `lagrange_basis :: Vector{<:AbstractPolynomialLike}`: The Lagrange basis corresponding to `poised_points`.
* `point_indices`: An array indicating which points from `points` are also in `poised_points`. A positive entry corresponds to the index of a poised point in `points`. If a poised point is new, then the entry is `-1`.
"""
function make_set_lambda_poised( basis, points :: AbstractArray{T};
		LAMBDA :: Real = 1.5, solver = :LN_BOBYQA, max_solver_evals = -1,
		max_loops = -1, skip_indices = [1,] ) where {
		T <: AbstractArray{<:Real}
	}

	@assert length(basis) == length(points) "Polynomial array `basis` and point array `points` must have the same length."

	if length(points) > 0
		n_vars = length(points[1])
		@assert n_vars > 0 "The number of variables must be positive."

		F = promote_type( eltype( T ), Float32 )
		P_type = n_vars > 100 ? Vector{Vector{F}} : Vector{SVector{n_vars, F}}

		if max_loops < 0
			max_loops = length(basis) * 100
		end

		if max_solver_evals < 0
			max_solver_evals = 2000 * n_vars
		end

		new_basis = basis
		new_points = P_type(points)
		point_indices = collect(eachindex(new_points))

		for k = 1 : max_loops
			iₖ = -1
			xₖ = points[1]
			for (i, polyᵢ) in enumerate(new_basis)
				opt = NLopt.Opt( solver, n_vars )
				opt.lower_bounds = zeros(F, n_vars)
				opt.upper_bounds = ones(F, n_vars)
				opt.maxeval = max_solver_evals
				opt.xtol_rel = 1e-3
				opt.max_objective = (x,g) -> abs( polyᵢ( x ) )

				x₀_tmp = [ rand(F, n_vars) for i = 1 : 50 * n_vars ]
				x₀ = x₀_tmp[argmax( abs.(new_basis[i].(x₀_tmp)) ) ]

				abs_lᵢ, xᵢ, _ = NLopt.optimize(opt, x₀)

				if abs_lᵢ > LAMBDA
					iₖ = i
					xₖ = xᵢ
					if i ∉ skip_indices
						# i is not prioritized we can brake here
						break
					end#if
				end#if
			end#for

			if iₖ > 0
				# perform a point swap
				new_points[iₖ] = xₖ
				point_indices[iₖ] = -1
				# adapt coefficients of lagrange basis
				new_basis = orthogonalize_polys( new_basis, xₖ, iₖ )
			else
				# we are done, the set is lambda poised
				break
			end#if
		end#for

		return new_points, new_basis, point_indices
	else
		return points, basis, collect(eachindex(points))
	end

end

# ╔═╡ c417a8e4-2564-451a-8105-58ae5f713720
md"For the sake of convenience we combine both functions:"

# ╔═╡ a57e64c4-cd79-485a-b21d-97591784cd1c
function get_lambda_poised_set( basis, points; solver1 = :LN_BOBYQA, solver2 = :LN_BOBYQA, max_solver_evals1 = -1, max_solver_evals2 = -1, LAMBDA = 1.5, max_lambda_loops = -1 )
	lagrange_points, lagrange_basis, lagrange_indices = get_poised_set(
		basis, points; solver = solver1, max_solver_evals = max_solver_evals1 )
	lambda_points, lambda_basis, lambda_indices = make_set_lambda_poised(
		lagrange_basis, lagrange_points; LAMBDA, max_loops = max_lambda_loops,
		solver = solver2, max_solver_evals = max_solver_evals2 )
	combined_indices = [ i < 0 ? i : lagrange_indices[j] for (j,i) in enumerate( lambda_indices ) ]
	return lambda_points, lambda_basis, combined_indices
end

# ╔═╡ ffc8cd79-9d85-4f04-81e1-3ae0193395b9
md"Let's have a look at what the points look like:"

# ╔═╡ f72d72de-119b-4c4c-a0d7-38b8ac76c2d6
begin
	basis = get_poly_basis(2,2)
	custom_points =  [ ones(Float32, 2), ones(Float32,2)]

	#lambda_points, lambda_basis, c_indices = get_lambda_poised_set( basis,custom_points)
	lambda_points, lambda_basis, c_indices = get_poised_set( basis,custom_points)
	c_indices
end

# ╔═╡ 9659c50c-c30f-4e78-a80a-eaee23209a4d
scatter(Tuple.(lambda_points))

# ╔═╡ Cell order:
# ╠═e9fdca3a-2d74-4a42-acda-2b289172a1d4
# ╠═9956f874-e481-4e3e-97fc-150aef8b69fa
# ╟─0bf0c2a6-7487-49a9-8adf-73381ea2e9fc
# ╠═af1b5146-5b3a-425d-95a7-6d0532bcca2f
# ╠═6c51dda0-5c69-4f58-8348-d62245b1eaeb
# ╠═802e1cd2-fceb-4ee7-b178-62b05ced5409
# ╠═c94bdc4e-2c95-42ae-a7dd-4172da278585
# ╠═b70e7d30-6856-4011-8a59-4727c690634f
# ╟─42ac42a6-d11e-4c03-bb18-0276d6a319fa
# ╠═688d4cf3-4e4d-45de-a43c-141646f01b00
# ╠═e8136cef-9235-4ad8-bfac-a9e176923de7
# ╠═f48b8d93-e98c-4f30-960c-203a8077b6b6
# ╠═85d345a0-2d3f-46ed-8b39-f88f700d545b
# ╟─c417a8e4-2564-451a-8105-58ae5f713720
# ╠═a57e64c4-cd79-485a-b21d-97591784cd1c
# ╟─ffc8cd79-9d85-4f04-81e1-3ae0193395b9
# ╠═f72d72de-119b-4c4c-a0d7-38b8ac76c2d6
# ╠═9659c50c-c30f-4e78-a80a-eaee23209a4d


using SortingLab, SortingAlgorithms

a = rand(Float32, 1_000_000)

sortperm_sorttwo(x) = begin
	i = collect(1:length(x))
	SortingLab.sorttwo!(copy(x), i)
	i
end

@time fsortperm(a)
@time sortperm(a)
@time sortperm(a, alg=RadixSort)
@time ass = sortperm_sorttwo(a)

using BenchmarkTools
@benchmark fsortperm(a)
@benchmark sortperm(a)
@benchmark sortperm(a, alg=RadixSort)
@benchmark ass = sortperm_sorttwo(a)

# sortperm_sorttwo seems to win out!
# which is implemented in fsortperm()


using GeometricIntegratorsDiffEq
using Test

using DiffEqProblemLibrary
using GeometricIntegrators

using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: prob_ode_2Dlinear

prob = prob_ode_2Dlinear

sol = solve(prob,GIEuler(),dt=0.1)
sol = solve(prob,GIMidpoint(),dt=0.1)
sol = solve(prob,GIHeun2(),dt=0.1)
sol = solve(prob,GIHeun3(),dt=0.1)
sol = solve(prob,GIRalston2(),dt=0.1)
sol = solve(prob,GIRalston3(),dt=0.1)
sol = solve(prob,GIRunge(),dt=0.1)
sol = solve(prob,GIKutta(),dt=0.1)
sol = solve(prob,GIRK416(),dt=0.1)
sol = solve(prob,GIRK438(),dt=0.1)
sol = solve(prob,GISSPRK3(),dt=0.1)
sol = solve(prob,GICrankNicolson(),dt=0.1)
sol = solve(prob,GIKraaijevangerSpijker(),dt=0.1)
sol = solve(prob,GIQinZhang(),dt=0.1)
sol = solve(prob,GICrouzeix(),dt=0.1)
sol = solve(prob,GIImplicitEuler(),dt=0.1)
sol = solve(prob,GIImplicitMidpoint(),dt=0.1)
sol = solve(prob,GISRK3(),dt=0.1)
sol = solve(prob,GIGLRK(2),dt=0.1)
sol = solve(prob,GIRadauIA(2),dt=0.1)
sol = solve(prob,GIRadauIIA(2),dt=0.1)
sol = solve(prob,GILobattoIIIA(2),dt=0.1)
sol = solve(prob,GILobattoIIIB(2),dt=0.1)
sol = solve(prob,GILobattoIIIC(2),dt=0.1)
sol = solve(prob,GILobattoIIIC̄(2),dt=0.1)
sol = solve(prob,GILobattoIIID(2),dt=0.1)
sol = solve(prob,GILobattoIIIE(2),dt=0.1)
sol = solve(prob,GILobattoIIIF(2),dt=0.1)

u0 = zeros(2)
v0 = ones(2)
f2 = function (dv,v,u,t)
  dv .= -u
end
function (::typeof(f2))(::Type{Val{:analytic}}, y0, p, x)
  u0, v0 = y0
  ArrayPartition(u0*cos(x) + v0*sin(x), -u0*sin(x) + v0*cos(x))
end
prob = SecondOrderODEProblem(f2,u0,v0,(0.0,5.0))

sol = solve(prob,GISymplecticEulerA(),dt=0.1)
sol = solve(prob,GISymplecticEulerB(),dt=0.1)
sol = solve(prob,GILobattoIIIAIIIB(2),dt=0.1)
sol = solve(prob,GILobattoIIIBIIIA(2),dt=0.1)


function printinrange(a, b)
    println(a)
    if a == b
        return
    else
        return printinrange(a + 1, b)
    end
end

# printinrange(2, 10)

function fibonacci(n)
    if n == 0
        return 0
    elseif n == 1
        return 1
    else
        return fibonacci(n - 1) + fibonacci(n - 2)
    end
end

# println(fibonacci(3))

const table = Dict(0=>0, 1=>1)
function fibmemo(n)
    if n in keys(table)
        return table[n]
    else
        res = table[n] = fibmemo(n - 1) + fibmemo(n - 2)
        return res
    end
end

# println(fibmemo(1000))


## Simplex Benchmark
## =================
# This benchmark samples from a small region around corners of a n-dimensional
# Simplex. If the regions are of equal size then we should expect an equal
# number of samples in each region.

"""The vertex coordinates for an ndim simplex
From http://people.sc.fsu.edu/~jburkardt/m_src/simplex_coordinates/"""
function simplex_coordinates(n::Int)
  x = zeros(Float64, n, n+1)
  for i = 1 : n
    x[i,i] = sqrt(1.0 - sum(x[1:i-1,i].^2))
    for j = i+1: n+1
      b = (-1.0/n - (x[1:i-1,i]' * x[1:i-1,j])) / x[i,i]
      x[i,j] = b[1]
    end
  end
  return x
end

"Is point p in box with center region and halfwidth halfwidth"
function point_near_region(p, region, halfwidth)
  conds = Array(Any,length(p))
  for i = 1:length(p)
    conds[i] = (p[i] >= region[i] - halfwidth) & (p[i] <= region[i] + halfwidth)
  end
  (&)(conds...)
end

"Constrain a point to be within a region around the vertices of a simplex"
function simplex(n::Int, box::RandArray, holesize)
  vertexcoords = simplex_coordinates(n)
  conds = RandVar{Bool}[]
  for i = 1:n+1
    push!(conds,point_near_region(box,vertexcoords[:,i],holesize))
  end
  box,(|)(conds...)
end

"Ground truth histrogram"
groundtruth(ndims) = Dict(i::Int => (1/(ndims+1))::Float64 for i = 1:(ndims+1))

"Constrain a point to be within a region around the vertices of a simplex"
simplex(n::Int) = simplex(n,mvuniform(-2,2,n))

"From a set of samples compute discrete distribution over the vertices"
function vertex_distribution(samples,n,holesize)
  vertexcoords = simplex_coordinates(n)
  counts =  DefaultDict(Int, Float64, 0.0)
  for sample in samples
    for i = 1:n+1
      if point_near_region(sample,vertexcoords[:,i], holesize)
        counts[i] += 1
        break
      end
    end
  end
  Dict(counts)
end

## Benchmark
## =========
"""Sample from a prior over R^ndim constrained to cubes of size holesize around
vertices of ndim Simplex"""
immutable Simplex <: Problem
  ndims::Int
  capture::Vector{Symbol}
  nsamples::Int
  holesize::Float64
  prior::RandArray
end

function benchmark(a::Algorithm, b::Simplex)
  print("Starting Sampler")
  Sigma.restart_counter!()
  captures::Vector{Symbol} = vcat(a.capture,b.capture)
  groundtruth = Dict(i => 1/(b.ndims+1) for i = 1:(b.ndims+1))
  model, condition = simplex(b.ndims, b.prior, b.holesize)
  @show b.nsamples

  value, results = capture(()->sample(a,model,condition,b.nsamples), captures)
  @show length(value)
  results
end

## Algorithm-dependent implementations of sample
## ============================================
"Sample with SMT based solver"
function sample(a::SigmaSMT, model, condition, nsamples)
  @show model
  @show condition
  @show nsamples
  rand(model, condition, nsamples; preimage_sampler = a.preimage_sampler,
                                   RandVarType = a.randvartype,
                                   cores = a.ncores,
                                   split = a.split)
end

"Sample with rejection sampling"
function sample(a::SigmaRej, X::AllRandVars, Y::RandVar{Bool}, nsamples::Integer)
  box = Sigma.LazyOmega(Float64)
  for dim in dims(Y)
    box[dim]
  end

  Y_func = Sigma.lambda(Y)
  X_func = Sigma.lambda(X)

  samples = []
  ntries = 10000000
  for i = 1:nsamples
    for i = 1:ntries
      ω = rand(box)
      if Y_func(ω)
        println("Got sample with rejection")
        push!(samples,X_func(ω))
        break
      end
    end
  end

  @show length(samples), nsamples
  length(samples) != nsamples && error("Couldn't draw sample")
  return samples
end

"Sample using breadth first search"
function sample(a::SigmaBFS, X::AllRandVars, Y::RandVar{Bool}, nsamples::Integer)
  Sigma.rand_bfs(X,Y,nsamples)
end

## Runs
## ====
include("runs/kl.jl")
# include("runs/dimensions.jl")


# dynamics functions, including the restricted three-body problem

export r3bp_hci_eom, twobody_eom, propagate_twobody, propagate_r3bp_hci

"""
Derivatives function for the restricted three-body problem in the
Heliocentric Inertial (HCI) reference frame xHCI is the stacked state
vector for the secondary planet + satellite: [xPlanet, xSat]
"""
function r3bp_hci_eom(t,
                      xHCI::Vector{T},
                      xdot,
                      mu2 = muEarth) where {T<:AbstractFloat}

    derivatives = zeros(12)
    derivatives[1:3] = xHCI[4:6]   # Planet's velocity
    derivatives[7:9] = xHCI[10:12] # Satellite velocity

    rPlanet = xHCI[1:3]
    rSat    = xHCI[7:9]
    derivatives[4:6]   = -muSun/(norm(rPlanet)^3) .* rPlanet
    derivatives[10:12] = -muSun/(norm(rSat)^3) .* rSat -
                            muEarth/(norm(rSat - rPlanet)^3) .* (rSat - rPlanet)
    xdot[:]=derivatives
end

"Derivatives function for the two-body problem"
function twobody_eom(t,
                     x::Vector{T},
                     xdot,
                     mu = muEarth) where {T<:AbstractFloat}

    xdot[1:3] = x[4:6]
    xdot[4:6] = -mu/(norm(x[1:3])^3) * x[1:3]
end

"""
Propagate the R3BP in the HCI frame vectors are expected in the HCI
frame

t should be an array of time values
"""
function propagate_r3bp_hci(x0Planet::Vector{T},
                            x0Sat::Vector{T},
                            t) where {T<:AbstractFloat}

    return Sundials.cvode(r3bp_hci_eom,
                   vcat(x0Planet, x0Sat),
                   t,
                   abstol=1e-12,reltol=1e-12)'
end

"Propagate the two-body problem"
function propagate_twobody(x0::Vector{T},
                           t) where {T<:AbstractFloat}

    return Sundials.cvode(twobody_eom,
                          x0,
                          t,
                          abstol=1e-12,reltol=1e-12)'
end


# Benchmark on column-wise distance evaluation

using Distances

macro bench_colwise_dist(repeat, dist, x, y)
    quote
        println("bench ", typeof($dist))

        # warming up
        r1 = evaluate($dist, ($x)[:,1], ($y)[:,1])
        colwise($dist, $x, $y)

        # timing

        t0 = @elapsed for k = 1 : $repeat
            n = size($x, 2)
            r = Array(typeof(r1), n)
            for j = 1 : n
                r[j] = evaluate($dist, ($x)[:,j], ($y)[:,j])
            end
        end
        @printf "    loop:     t = %9.6fs\n" (t0 / $repeat)

        t1 = @elapsed for k = 1 : $repeat
            r = colwise($dist, $x, $y)
        end
        @printf "    colwise:  t = %9.6fs  |  gain = %7.4fx\n" (t1 / $repeat) (t0 / t1)
        println()
    end
end

m = 200
n = 10000

x = rand(m, n)
y = rand(m, n)
w = rand(m)

Q = rand(m, m)
Q = Q' * Q

@bench_colwise_dist 20 SqEuclidean() x y
@bench_colwise_dist 20 Euclidean() x y
@bench_colwise_dist 20 Cityblock() x y
@bench_colwise_dist 20 Chebyshev() x y
@bench_colwise_dist  5 Minkowski(3.0) x y
@bench_colwise_dist 20 Hamming() x y

@bench_colwise_dist 20 CosineDist() x y
@bench_colwise_dist 10 CorrDist() x y
@bench_colwise_dist 20 ChiSqDist() x y
@bench_colwise_dist 10 KLDivergence() x y
@bench_colwise_dist  5 JSDivergence() x y

@bench_colwise_dist 10 BhattacharyyaDist() x y
@bench_colwise_dist 10 HellingerDist() x y

@bench_colwise_dist 20 WeightedSqEuclidean(w) x y
@bench_colwise_dist 20 WeightedEuclidean(w) x y
@bench_colwise_dist 20 WeightedCityblock(w) x y
@bench_colwise_dist  5 WeightedMinkowski(w, 3.0) x y
@bench_colwise_dist 20 WeightedHamming(w) x y

@bench_colwise_dist 10 SqMahalanobis(Q) x y
@bench_colwise_dist 10 Mahalanobis(Q) x y


function polydivisible(x::Int)::Bool
  accu = 0
  for (i,digit) in x |> digits |> reverse |> enumerate
    accu = muladd( accu, 10, digit )
    if accu % i != 0 return false end
  end
  return true
end


module MiscTest

import NaNMath

using Test
using ForwardDiff
using DiffTests
using SparseArrays: sparse

include(joinpath(dirname(@__FILE__), "utils.jl"))

##########################
# Nested Differentiation #
##########################

# README example #
#----------------#

x = rand(5)

f = x -> sum(sin, x) + prod(tan, x) * sum(sqrt, x)
g = x -> ForwardDiff.gradient(f, x)
j = x -> ForwardDiff.jacobian(g, x)

@test isapprox(ForwardDiff.hessian(f, x), j(x))

# higher-order derivatives #
#--------------------------#

function tensor(f, x)
    n = length(x)
    out = ForwardDiff.jacobian(y -> ForwardDiff.hessian(f, y), x)
    return reshape(out, n, n, n)
end

test_tensor_output = reshape([240.0  -400.0     0.0;
                             -400.0     0.0     0.0;
                                0.0     0.0     0.0;
                             -400.0     0.0     0.0;
                                0.0   480.0  -400.0;
                                0.0  -400.0     0.0;
                                0.0     0.0     0.0;
                                0.0  -400.0     0.0;
                                0.0     0.0     0.0], 3, 3, 3)

@test isapprox(tensor(DiffTests.rosenbrock_1, [0.1, 0.2, 0.3]), test_tensor_output)

test_nested_jacobian_output = [-sin(1)  0.0     0.0;
                               -0.0    -0.0    -0.0;
                               -0.0    -0.0    -0.0;
                                0.0     0.0     0.0;
                               -0.0    -sin(2) -0.0;
                               -0.0    -0.0    -0.0;
                                0.0     0.0     0.0;
                               -0.0    -0.0    -0.0;
                               -0.0    -0.0    -sin(3)]

sin_jacobian = x -> ForwardDiff.jacobian(y -> broadcast(sin, y), x)

@test isapprox(ForwardDiff.jacobian(sin_jacobian, [1., 2., 3.]), test_nested_jacobian_output)

# Issue #59 example #
#-------------------#

x = rand(2)

f = x -> sin(x)/3 * cos(x)/2
df = x -> ForwardDiff.derivative(f, x)
testdf = x -> (((cos(x)^2)/3) - (sin(x)^2)/3) / 2
f2 = x -> df(x[1]) * f(x[2])
testf2 = x -> testdf(x[1]) * f(x[2])

@test isapprox(ForwardDiff.gradient(f2, x), ForwardDiff.gradient(testf2, x))

######################################
# Higher-Dimensional Differentiation #
######################################

x = rand(5, 5)

@test isapprox(ForwardDiff.jacobian(inv, x), -kron(inv(x'), inv(x)))

#########################################
# Differentiation with non-Array inputs #
#########################################

x = rand(5,5)

# Sparse
f = x -> sum(sin, x) + prod(tan, x) * sum(sqrt, x)
gfx = ForwardDiff.gradient(f, x)
@test isapprox(gfx, ForwardDiff.gradient(f, sparse(x)))

# Views
jinvx = ForwardDiff.jacobian(inv, x)
@test isapprox(jinvx, ForwardDiff.jacobian(inv, view(x, 1:5, 1:5)))

########################
# Conversion/Promotion #
########################

# target function returns a literal (Issue #71) #
#-----------------------------------------------#

@test ForwardDiff.derivative(x->zero(x), rand()) == ForwardDiff.derivative(x->1.0, rand())
@test ForwardDiff.gradient(x->zero(x[1]), [rand()]) == ForwardDiff.gradient(x->1.0, [rand()])
@test ForwardDiff.hessian(x->zero(x[1]), [rand()]) == ForwardDiff.hessian(x->1.0, [rand()])
@test ForwardDiff.jacobian(x->[zero(x[1])], [rand()]) == ForwardDiff.jacobian(x->[1.0], [rand()])

# arithmetic element-wise functions #
#-----------------------------------#

N = 4
a = fill(1.0, N)
jac0 = reshape(vcat([[fill(0.0, N*(i-1)); a; fill(0.0, N^2-N*i)] for i = 1:N]...), N^2, N)

for op in (:.-, :.+, :./, :.*)
    @eval begin
        f = x -> [$op(x[1], a); $op(x[2], a); $op(x[3], a); $op(x[4], a)]
        jac = ForwardDiff.jacobian(f, a)
        @test isapprox(jac0, jac)
    end
end

# NaNs #
#------#

@test ForwardDiff.partials(NaNMath.pow(ForwardDiff.Dual(-2.0,1.0),ForwardDiff.Dual(2.0,0.0)),1) == -4.0

# Partials{0} #
#-------------#

x, y = rand(3), rand(3)
h = ForwardDiff.hessian(y -> sum(hypot.(x, y)), y)
@test h[1, 1] ≈ (x[1]^2) / (x[1]^2 + y[1]^2)^(3/2)
@test h[2, 2] ≈ (x[2]^2) / (x[2]^2 + y[2]^2)^(3/2)
@test h[3, 3] ≈ (x[3]^2) / (x[3]^2 + y[3]^2)^(3/2)
let i, j
    for i in 1:3, j in 1:3
        i != j && @test h[i, j] ≈ 0.0
    end
end

########
# misc #
########

# issue 267
@noinline f267(z, x) = x[1]
z267 = ([(1, (2), [(3, (4, 5, [1, 2, (3, (4, 5), [5])]), (5))])])
let z = z267,
    g = x -> f267(z, x),
    h = x -> g(x)
    @test ForwardDiff.hessian(h, [1.]) == fill(0.0, 1, 1)
end

# issue #290
@test ForwardDiff.derivative(x -> rem2pi(x, RoundUp), rand()) == 1
@test ForwardDiff.derivative(x -> rem2pi(x, RoundDown), rand()) == 1

end # module


include("utils.jl")

@testset "Hyperbolic Space" begin
    M = Hyperbolic(2)
    @testset "Hyperbolic Basics" begin
        @test repr(M) == "Hyperbolic(2)"
        @test base_manifold(M) == M
        @test typeof(get_embedding(M)) ==
              MetricManifold{Euclidean{Tuple{3},ℝ},MinkowskiMetric}
        @test representation_size(M) == (3,)
        @test isinf(injectivity_radius(M))
        @test isinf(injectivity_radius(M, ExponentialRetraction()))
        @test isinf(injectivity_radius(M, [0.0, 0.0, 1.0]))
        @test isinf(injectivity_radius(M, [0.0, 0.0, 1.0], ExponentialRetraction()))
        @test !is_manifold_point(M, [1.0, 0.0, 0.0, 0.0])
        @test !is_tangent_vector(M, [0.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0])
        @test_throws DomainError is_manifold_point(M, [2.0, 0.0, 0.0], true)
        @test !is_manifold_point(M, [2.0, 0.0, 0.0])
        @test !is_tangent_vector(M, [1.0, 0.0, 0.0], [1.0, 0.0, 0.0])
        @test_throws DomainError is_tangent_vector(
            M,
            [1.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            true,
        )
        @test !is_tangent_vector(M, [0.0, 0.0, 1.0], [1.0, 0.0, 1.0])
        @test_throws DomainError is_tangent_vector(
            M,
            [0.0, 0.0, 1.0],
            [1.0, 0.0, 1.0],
            true,
        )
        @test is_default_metric(M, MinkowskiMetric())
        @test Manifolds.default_metric_dispatch(M, MinkowskiMetric()) === Val{true}()
        @test manifold_dimension(M) == 2
    end
    types = [
        Vector{Float64},
        SizedVector{3,Float64},
    ]
    TEST_FLOAT32 && push!(types, Vector{Float32})

    for T in types
        @testset "Type $T" begin
            pts = [
                convert(T, [0.0, 0.0, 1.0]),
                convert(T, [1.0, 0.0, sqrt(2.0)]),
                convert(T, [0.0, 1.0, sqrt(2.0)]),
            ]
            test_manifold(
                M,
                pts,
                test_project_tangent = true,
                test_musical_isomorphisms = true,
                test_vector_transport = true,
                is_tangent_atol_multiplier = 10.0,
                exp_log_atol_multiplier = 10.0,
                retraction_methods = (ExponentialRetraction(),),
                test_vee_hat = false,
            )
        end
    end
    @testset "Embedding test" begin
        p = [0.0, 0.0, 1.0]
        X = [1.0, 1.0, 0.0]
        @test embed(M, p) == p
        q = similar(p)
        embed!(M, q, p)
        @test q == p
        @test embed(M, p, X) == X
        Y = similar(X)
        embed!(M, Y, p, X)
        @test Y == X
    end
    @testset "Hyperbolic mean test" begin
        pts = [
            [0.0, 0.0, 1.0],
            [1.0, 0.0, sqrt(2.0)],
            [-1.0, 0.0, sqrt(2.0)],
            [0.0, 1.0, sqrt(2.0)],
            [0.0, -1.0, sqrt(2.0)],
        ]
        @test isapprox(M, mean(M, pts), pts[1]; atol = 10^-4)
    end
end


# Derivative of a function

function derivative(f)

   return function(x)
        # pick a small value for h
        h = x == 0 ? sqrt(eps(Float64)) : sqrt(eps(Float64)) * x

        # floating point arithmetic gymnastics
        xph = x + h
        dx = xph - x

        # evaluate f at x + h
        f1 = f(xph)

        # evaluate f at x
        f0 = f(x)

        # divide the difference by h
        return (f1 - f0) / dx
    end

end


# This file is a part of Julia. License is MIT: http://julialang.org/license

# Finite-difference time-domain seismic simulation in 2D using a staggered grid.
# The intent is to test performance of @simd on the inner loop of a 2D loop nest.
#
# If @simd ever supports forward-lexical dependences, then the kernels in updateV
# and updateU can be merged.

# Update velocity component fields Vx and Vy
function updateV( irange, jrange, U, Vx, Vy, A )
    for j in jrange
        @simd for i in irange
            @inbounds begin
                Vx[i,j] += (A[i,j+1]+A[i,j])*(U[i,j+1]-U[i,j])
                Vy[i,j] += (A[i+1,j]+A[i,j])*(U[i+1,j]-U[i,j])
            end
        end
    end
end

# Update pressure field U
function updateU( irange, jrange, U, Vx, Vy, B )
    for j in jrange
        @simd for i in irange
            @inbounds begin
                U[i,j] += B[i,j]*((Vx[i,j]-Vx[i,j-1]) + (Vy[i,j]-Vy[i-1,j]))
            end
        end
    end
end

# Alternate updates for given number of steps
function flog_fdtd( steps, U, Vx, Vy, A, B )
    m,n = size(U)
    for k=1:steps
         updateV(2:m-1,2:n-1,U,Vx,Vy,A)
         updateU(2:m-1,2:n-1,U,Vx,Vy,B)
    end
end


module potentials

export Potential, BoxPotential

type Potential
    A::Function
    B::Function
    E::Function
    basis_function::Function
end

fcos(s, n::Int, i::Int, j::Int) = cospi((i+j)/2)*sinc((i+j+2*n*s)/2)
fsin(s, n::Int, i::Int, j::Int) = sinpi((i+j+n)/2)*sinc((i+j+n*s)/2)

function A(s, n::Int, i::Int, j::Int)
    @assert n>=0
    @assert i>0
    @assert j>0
    # i += 1
    # j += 1
    # n += 1
    k = (i==j ? 0.5 : 0.)
    k + (-1)^n/4 * (fcos(s,n,i,j) + fcos(s,-n,i,j) - fcos(s,n,i,-j) - fcos(s,-n,i,-j))
end

function B(s, n::Int, i::Int, j::Int)
    @assert n>=0
    @assert i>0
    @assert j>0
    # i += 1
    # j += 1
    # n += 1
    return 0.5 * (-fsin(s,n,i,j) + fsin(s,n,-i,j) + fsin(s,n,i,-j) + fsin(s,-n,i,+j))
end

E(E0, i::Int) = E0*i^2

function basis_function(index::Int, x)
    k = pi*index/2
    if (x<-1 || x>1)
        return zero(x)
    else
        return 1./sqrt(2)*sin(k*(x+1))
    end
end

basis_function(index::Int, x::Vector) = map(xi->basis_function(index, xi), x)

const BoxPotential = Potential(A, B, E, basis_function)

end


#Clearing Error Function
function  z(x::Matrix{Int64}, q::Vector{Int64}, p::Vector{Int64})
   # S is the number of students, N is the number of courses
   # x is the SxN allocation matrix
   # q is the Nx1 course capacity vector
   # p is the Nx1 price vector
   N = length(q)
   zer = zeros(N)        #Create empty Nx1 vector of clearing errors
   numstud = transpose(sum(x, 1)) #Create Nx1 vector with the number of students enrolled in each course.
   for i in 1:N        #For every course
       if p[i] > 0     #As long as the price in this course is positive.
           zer[i] = numstud[i] - q[i]
       else
           zer[i] = max(numstud[i] - q[i], 0)
       end
   end
   return zer
end

#TEST
x = [1 0 1;
     1 0 1;
     1 0 1;
     0 0 1;
     1 1 0;
     1 1 1]
q = [4;3;3]
N = 3
p = [30; 0; 50]

z(x, q, p, N)