### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ be57f61e-a9f1-40dc-a692-676c14f03aa4
begin
	using PlutoUI
	using PrettyTables
end

# ╔═╡ f7c935c8-f577-4e38-b242-7e262bd446e1
md"""
##### Example: Gas Phase Reaction with a non-constant number of mols
The reversible _gas phase_ chemical reaction (R1):

$$\begin{equation}
  A_{(g)}+B_{(g)}\leftrightharpoons{C_{(g)}}
\end{equation}$$

is carried out at T = 298.15 K, and an _unknown_ pressure P (bar). The volume of the reactor V = 20L is constant. At these conditions, the teaching team measured the equilibrium extent of reaction $\epsilon^{\star}$ = 0.90 mol.
Let A = 1, B = 2 and C = 3. We start with $n^{\circ}_{1}$ = 2.0 mol, $n^{\circ}_{2}$ = 1.0 mol and
$n^{\circ}_{3}$ = 0.0 mol.

__Assumptions__:
* (i) P$^{\circ}$ = 1 bar;
* (ii) R = 8.314 $\times$10$^{-2}$ L bar mol$^{-1}$ K$^{-1}$ and R = 8.314 J mol$^{-1}$ K$^{-1}$;
* (iii) the vapor phase is _ideal_;
* (iv) the ideal gas law is valid at these reaction conditions

__Compute__:
  * (a) Compute all state (mol) values at the reaction conditions
  * (b) Starting from the general equilibrium expression for a single vapor phase reaction, estimate a value for $K$ for R1 at T = 298.15 K and system pressure.
  * (c) Estimate the value of the Gibbs energy of reaction $\Delta{G}^{\circ}$ (kJ/mol) from the equilibrium constant $K$. Is the reaction: favorable or not favorable?
"""

# ╔═╡ 861c1cf4-ce56-43ec-bca1-6974d72240f5
md"""
##### (a) Compute the composition table
To compute the entries in the composition table, we use the definition of the extent of reaction, and the reaction stoichiometry, to calculate the final number of mols (given the initial number of mols):

$$n_{i} = n_{i}^{\circ}+\sigma_{i1}\epsilon_{1}$$

In this problem, we have a single reaction, and at equilibrium $\epsilon_{1}=\epsilon^{\star}=0.90$ mol. To compute mol fractions (which appear in the equilibrium constant expressions), we divide each species mol number by the total i.e., $y_{i}$ = $n_{i}/n_{T}$.
"""

# ╔═╡ 5e423631-becf-4351-99f2-10da34eb5675
# Setup initial composition table -
begin

	# how many species do we have?
	number_of_species = 3
	number_of_state_cols = 5

	# setup the state table -> species x items
	initial_state_table = Array{Union{Number,Nothing, String},2}(undef, number_of_species, number_of_state_cols)

	# put stuff into the table from the problem -
	initial_state_table[1,1] = "A" 		# species
	initial_state_table[1,2] = -1 		# stoichiometric coeff -
	initial_state_table[1,3] = 2.0 		# initial number of mols
	initial_state_table[1,4] = nothing 	# final number of mols
	initial_state_table[1,5] = nothing 	# mol fraction

	initial_state_table[2,1] = "B" 		# species
	initial_state_table[2,2] = -1 		# stoichiometric coeff -
	initial_state_table[2,3] = 1.0 		# initial number of mols
	initial_state_table[2,4] = nothing 	# final number of mols
	initial_state_table[2,5] = nothing 	# mol fraction

	initial_state_table[3,1] = "C" 		# species
	initial_state_table[3,2] = 1 		# stoichiometric coeff -
	initial_state_table[3,3] = 0.0 		# initial number of mols
	initial_state_table[3,4] = nothing 	# final number of mols
	initial_state_table[3,5] = nothing 	# mol fraction

	with_terminal() do

		# header row -
		state_table_header_row = (["species","σ","initial n","final n","yᵢ"],["","dimensionless","mol","mol","dimensionless"]);

		# write the table -
		pretty_table(initial_state_table; header=state_table_header_row)
	end
end

# ╔═╡ 4c722cf0-97d5-4ac2-9533-ae216ce89d19
# compute the final number of mols -
begin
	# setup as MV calculation (faster) -
	n_initial = [2.0 ; 1.0 ; 0.0];
	S = [-1 ; -1 ; 1];
	ϵ = 0.90

	# compute final mol -
	n_final = n_initial + S*ϵ
end

# ╔═╡ 67b1906f-0c64-4cca-9d26-4383eeb370c0
# compute the mol frac -
begin
	n_total = sum(n_final)
	y_array = (1/n_total)*n_final
end

# ╔═╡ bf499485-c101-452d-85f8-be4d6f7a34aa
begin

	final_state_table = copy(initial_state_table)

	# put stuff into the table from the problem -
	final_state_table[1,1] = "A" 			# species
	final_state_table[1,2] = -1 			# stoichiometric coeff -
	final_state_table[1,3] = 2.0 			# initial number of mols
	final_state_table[1,4] = n_final[1] 	# final number of mols
	final_state_table[1,5] = y_array[1] 	# mol fraction

	final_state_table[2,1] = "B" 			# species
	final_state_table[2,2] = -1 			# stoichiometric coeff -
	final_state_table[2,3] = 1.0 			# initial number of mols
	final_state_table[2,4] = n_final[2] 	# final number of mols
	final_state_table[2,5] = y_array[2] 	# mol fraction

	final_state_table[3,1] = "C" 			# species
	final_state_table[3,2] = 1 				# stoichiometric coeff -
	final_state_table[3,3] = 0.0 			# initial number of mols
	final_state_table[3,4] = n_final[3] 	# final number of mols
	final_state_table[3,5] = y_array[3] 	# mol fraction

	with_terminal() do

		# header row -
		state_table_header_row = (["species","σ","initial n","final n","yᵢ"],["","dimensionless","mol","mol","dimensionless"]);

		# write the table -
		pretty_table(final_state_table; header=state_table_header_row)
	end
end

# ╔═╡ e305af93-5ef0-413c-8c7d-68737cc1ffc4
md"""
##### (b) Estimate the equilibrium constant K
The general expression for the equilibrium constant for a gas phase reaction is:

$$\begin{equation}
  \left(\frac{P}{P^{\circ}}\right)^{-\sigma_{1}}K = \prod_{i=1}^{\mathcal{M}}\left(y_{i}\hat{\phi}^{v}_{i}\right)^{\sigma_{i1}}
\end{equation}$$

However, we know a few things from the problem. First, the gas phase is ideal so the fugacity coefficient $\hat{\phi}^{v}_{i}\rightarrow{1},\forall{i}$. Next, we know that the overall stoichiometric coefficient $\sigma_{1}$ = -1, and
$P^{\circ}$ = 1 bar. Finally, while we do not know the pressure, we do know that the gas can treated as an _ideal_ gas, meaning we can use the ideal gas law to describe the pressure i.e.,

$$P = \frac{n_{T}RT}{V}$$

Putting these ideas together, the general expression for the equilibrium gas phase equilibrium constant becomes:

$$K = \frac{1}{P}\prod_{i=1}^{\mathcal{M}}\left(y_{i}\right)^{\sigma_{i1}}$$

or after further simplification:

$$K = \frac{1}{P}\left(\frac{y_{3}}{y_{1}y_{2}}\right)$$

"""

# ╔═╡ 9183cb07-8e42-4440-aab4-cde1e101af05
md"""
__First step__: Compute the unkown pressure P (units: bar)
"""

# ╔═╡ 07db3757-926f-4391-aeae-be7240998f09
# compute the Pressure -
begin

	# use the IGL to estimate the equilibrium pressure -
	V = 20.0 					# units: L
	T = 298.15 					# units: K
	R_bar = (8.314)*10^(-2)			# units: L bar mol$^{-1}$ K$^{-1}$
	P = (n_total*R_bar*T)/(V)		# units: bar

	# display the pressure -
	with_terminal() do
		println("Pressure = $(P) bar")
	end
end

# ╔═╡ f51e4de0-1c39-45ca-933f-1826abcbf1f3
md"""
__Compute__: the equilibrium constant K given P and composition
"""

# ╔═╡ 3421b268-a3d6-48b5-819e-d9cbd14ee687
# compute the K -
begin

	# get the mol frac's from the final state table -
	y₁ = final_state_table[1,5]
	y₂ = final_state_table[2,5]
	y₃ = final_state_table[3,5]

	# compute the K -
	K = (1/P)*(y₃/(y₁*y₂))

	with_terminal() do
		# print -
		println("K = $(K)")
	end
end

# ╔═╡ dd52da8d-106b-4385-8576-6fc030ce90d1
md"""
##### (c) Estimate the ΔG$^{\circ}$ for the reaction
At equilibrium we know that the constant K and the ΔG$^{\circ}$ of reaction are related by:

$$\ln K = -\frac{\Delta{G}^{\circ}}{RT}$$

which we can use to solve for ΔG$^{\circ}$ given that we know the equilibrium constant $K$:

$$\Delta{G}^{\circ} = -RT\ln K$$

"""

# ╔═╡ 22ce61af-1a3d-4214-8b98-dc463e3d61b8
with_terminal() do
	R = 8.314*(1/1000) # units: kJ/mol-K
	ΔG = -R*T*log(K)
	println("ΔG = $(ΔG) kJ/mol")
end

# ╔═╡ 7bfe56b8-3639-11ec-1f61-9d52526c146d
html"""
<style>
main {
    max-width: 1200px;
    width: 90%;
    margin: auto;
    font-family: "Roboto, monospace";
}

a {
    color: blue;
    text-decoration: none;
}

.H1 {
    padding: 0px 30px;
}
</style>"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"

[compat]
PlutoUI = "~0.7.16"
PrettyTables = "~1.2.3"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[DataAPI]]
git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.9.0"

[[DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[HypertextLiteral]]
git-tree-sha1 = "5efcf53d798efede8fee5b2c8b09284be359bf24"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.2"

[[IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "f19e978f81eca5fd7620650d7dbea58f825802ee"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.1.0"

[[PlutoUI]]
deps = ["Base64", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "4c8a7d080daca18545c56f1cac28710c362478f3"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.16"

[[PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "d940010be611ee9d67064fe559edbb305f8cc0eb"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.2.3"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "TableTraits", "Test"]
git-tree-sha1 = "fed34d0e71b91734bf0a7e10eb1bb05296ddbcd0"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.6.0"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
"""

# ╔═╡ Cell order:
# ╟─f7c935c8-f577-4e38-b242-7e262bd446e1
# ╟─861c1cf4-ce56-43ec-bca1-6974d72240f5
# ╟─5e423631-becf-4351-99f2-10da34eb5675
# ╠═4c722cf0-97d5-4ac2-9533-ae216ce89d19
# ╠═67b1906f-0c64-4cca-9d26-4383eeb370c0
# ╟─bf499485-c101-452d-85f8-be4d6f7a34aa
# ╟─e305af93-5ef0-413c-8c7d-68737cc1ffc4
# ╟─9183cb07-8e42-4440-aab4-cde1e101af05
# ╠═07db3757-926f-4391-aeae-be7240998f09
# ╟─f51e4de0-1c39-45ca-933f-1826abcbf1f3
# ╠═3421b268-a3d6-48b5-819e-d9cbd14ee687
# ╟─dd52da8d-106b-4385-8576-6fc030ce90d1
# ╠═22ce61af-1a3d-4214-8b98-dc463e3d61b8
# ╟─be57f61e-a9f1-40dc-a692-676c14f03aa4
# ╟─7bfe56b8-3639-11ec-1f61-9d52526c146d
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


@testset "pdf of quadrature" begin

    gaussian(x, μ, σ) = exp(-0.5((x-μ)/σ)^2) / (σ*sqrt(2π))

    vacuum_state = VacuumState(rep=StateMatrix)
    v_μ = QSB.π̂ₓ_μ([0], vacuum_state)[1]
    v_σ = real(sqrt.(QSB.π̂ₓ²_μ([0], vacuum_state) .- v_μ.^2))[1]

    v(x) = gaussian(x, v_μ, v_σ)

    @test all(
        q_pdf(
            VacuumState(rep=StateMatrix, dim=100),
            LinRange(0, 2π, 2),
            LinRange(-10, 10, 100),
        )[i, :] ≈ v.(LinRange(-10, 10, 100))
        for i in 1:2
    )
end


using LinearAlgebra

#if VERSION < v"0.7-"
#  import Base: A_mul_B!
#  mul!(x,B,y) = A_mul_B!(x,B,y)
#end

@testset "Point-Field Routines" begin

  @testset "Point creation" begin
    @test_throws AssertionError VectorData([1,2,3],[1,2])

    @test_throws AssertionError TensorData([1,2,3],[1,2],[2,3],[4,5])

    f = ScalarData(10)
    f .= rand(10)

    ft = TensorData(f)
    ft[25] = 4
    @test ft.dvdx[5] == 4

  end

  @testset "Point operations" begin

    f = ScalarData(10)
    f .= rand(10)

    g = -f
    @test g[4] == -f[4]

    g = ScalarData(f)
    g .= rand(10)

    h = f + g

    @test h[8] == f[8]+g[8]

    h = 2.0*g

    @test h[4] == 2.0*g[4]

    Y = VectorData(4)
    Y .= rand(length(Y))
    X = 2*Y

    @test X.u[1] == 2*Y.u[1] && X.v[1] == 2*Y.v[1]

    X = Y/2

    @test X.u[3] == Y.u[3]/2 && X.v[3] == Y.v[3]/2

    fill!(Y,0.0)

    X = Y + (1,2)
    @test X.v[4] == 2.0

    X2 = (1,2) + Y
    @test X2 == X

    X = Y - (1,2)
    @test X.u[4] == -1.0

    X = (1,2) - Y
    @test X.u[4] == 1.0

    Z = 2.0 × X
    @test typeof(Z) <: VectorData
    @test Z.u[1] == -4.0

    X = TensorData(Y)
    fill!(X,1)
    Z = X + X
    @test Z.dvdx[4] == X.dvdx[4] + X.dvdx[4]

    Z = (1,2) ⋅ X
    @test typeof(Z) <: VectorData
    @test Z.u[3] == 3.0

  end

  @testset "Hadamard products" begin

    fcs = ScalarData(5,dtype=ComplexF64)
    fill!(fcs,2im)
    fcs .= rand(5)*im

    a = 2.0
    frs = TensorData(fcs,dtype=Float64)
    fill!(frs,a)

    out = similar(frs,element_type=ComplexF64)

    product!(out,frs,fcs)

    @test out[3] == frs[3]*fcs[3]

    out = frs ∘ fcs

    @test out[2] == frs[2]*fcs[2]

  end

  n = 10
  x = 0.5 .+ 0.2*rand(n)
  y = 0.5 .+ 0.2*rand(n)
  X = VectorData(x,y)

  nx = 12; ny = 12

  dx = 0.1
  H = Regularize(X,dx)
  H̃ = Regularize(X,dx,filter=true)

  @testset "Regularize vector to primal edges" begin

    f = VectorData(X)
    f.u .= rand(n)
    f.v .= rand(n)

    q = Edges(Primal,(nx,ny))

    H(q,f)
    @test sum(f.u) ≈ sum(q.u)*dx*dx
    @test sum(f.v) ≈ sum(q.v)*dx*dx

  end

  @testset "Regularize vector to dual edges" begin

  f = VectorData(X)
  f.u .= rand(n)
  f.v .= rand(n)

  p = Edges(Dual,(nx,ny))

  H(p,f)
  @test sum(f.u) ≈ sum(p.u)*dx*dx
  @test sum(f.v) ≈ sum(p.v)*dx*dx

  end

  @testset "Regularize vector to dual and primal nodes" begin

  f = VectorData(X)
  f.u .= rand(n)
  f.v .= rand(n)

  w = NodePair(Dual,(nx,ny))

  H(w,f)
  @test sum(f.u) ≈ sum(w.u)*dx*dx
  @test sum(f.v) ≈ sum(w.v)*dx*dx

  w = NodePair(Primal,(nx,ny))

  H(w,f)
  @test sum(f.u) ≈ sum(w.u)*dx*dx
  @test sum(f.v) ≈ sum(w.v)*dx*dx

  end

  @testset "Regularize tensor to edge gradient" begin

  f = TensorData(X)
  f.dudx .= rand(n)
  f.dudy .= rand(n)
  f.dvdx .= rand(n)
  f.dvdy .= rand(n)

  gradq = EdgeGradient(Dual,(nx,ny))

  H(gradq,f)

  @test sum(f.dudx) ≈ sum(gradq.dudx)*dx*dx
  @test sum(f.dudy) ≈ sum(gradq.dudy)*dx*dx
  @test sum(f.dvdx) ≈ sum(gradq.dvdx)*dx*dx
  @test sum(f.dvdy) ≈ sum(gradq.dvdy)*dx*dx

  gradq = EdgeGradient(Primal,(nx,ny))

  H(gradq,f)

  @test sum(f.dudx) ≈ sum(gradq.dudx)*dx*dx
  @test sum(f.dudy) ≈ sum(gradq.dudy)*dx*dx
  @test sum(f.dvdx) ≈ sum(gradq.dvdx)*dx*dx
  @test sum(f.dvdy) ≈ sum(gradq.dvdy)*dx*dx

  end

  @testset "Regularize scalar to primal nodes" begin

  f = ScalarData(X)
  f .= rand(n)

  w = Nodes(Primal,(nx,ny))

  H(w,f)
  @test sum(f) ≈ sum(w)*dx*dx

  end

  @testset "Regularize scalar to dual nodes" begin

  f = ScalarData(X)
  f .= rand(n)

  w = Nodes(Dual,(nx,ny))

  H(w,f)
  @test sum(f) ≈ sum(w)*dx*dx

  end

  @testset "Matrix representation" begin

  f = ScalarData(X)
  w = Nodes(Dual,(nx,ny))
  Hmat = RegularizationMatrix(H,f,w)
  Emat = InterpolationMatrix(H,w,f)

  f .= rand(n)

  w2 = Nodes(Dual,(nx,ny))
  mul!(w,Hmat,f)
  H(w2,f)
  @test w ≈ w2

  @test_throws MethodError mul!(f,Hmat,w)

  f2 = ScalarData(f)
  mul!(f,Emat,w)
  H(f2,w)
  @test f ≈ f2

  w .= rand(nx,ny)
  Ẽmat = InterpolationMatrix(H̃,w,f)

  mul!(f,Ẽmat,w)
  H̃(f2,w)
  @test f ≈ f2

  w = Nodes(Primal,(nx,ny))
  Hmat = RegularizationMatrix(H,f,w)
  Emat = InterpolationMatrix(H,w,f)
  Ẽmat = InterpolationMatrix(H̃,w,f)

  w2 = Nodes(Primal,(nx,ny))
  mul!(w,Hmat,f)
  H(w2,f)
  @test w ≈ w2

  w .= rand(Float64,size(w))
  f2 = ScalarData(f)
  mul!(f,Emat,w)
  H(f2,w)
  @test f ≈ f2

  f2 = ScalarData(f)
  mul!(f,Ẽmat,w)
  H̃(f2,w)
  @test f ≈ f2

  f = VectorData(X)
  f.u .= rand(n)
  f.v .= rand(n)

  p = Edges(Dual,(nx,ny))
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)

  p2 = Edges(Dual,(nx,ny))
  mul!(p,Hmat,f)
  H(p2,f)
  @test p.u ≈ p2.u && p.v ≈ p2.v

  p.u .= rand(Float64,size(p.u))
  p.v .= rand(Float64,size(p.v))
  f2 = VectorData(f)
  mul!(f,Emat,p)
  H(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  p = NodePair(Dual,(nx,ny))
  p2 = deepcopy(p)
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)

  mul!(p,Hmat,f)
  H(p2,f)
  @test p.u ≈ p2.u && p.v ≈ p2.v

  p.u .= rand(Float64,size(p.u))
  p.v .= rand(Float64,size(p.v))
  mul!(f,Emat,p)
  H(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  p = NodePair(Primal,(nx,ny))
  p2 = deepcopy(p)
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)
  mul!(p,Hmat,f)
  H(p2,f)
  @test p.u ≈ p2.u && p.v ≈ p2.v

  p.u .= rand(Float64,size(p.u))
  p.v .= rand(Float64,size(p.v))
  mul!(f,Emat,p)
  H(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.u ≈ f2.u && f.v ≈ f2.v

  f = TensorData(X)
  f.dudx .= rand(n)
  f.dudy .= rand(n)
  f.dvdx .= rand(n)
  f.dvdy .= rand(n)
  f2 = TensorData(f)

  p = EdgeGradient(Dual,(nx,ny))
  p2 = deepcopy(p)
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)

  mul!(p,Hmat,f)
  H(p2,f)
  @test p.dudx ≈ p2.dudx && p.dudy ≈ p2.dudy && p.dvdx ≈ p2.dvdx && p.dvdy ≈ p2.dvdy

  p.dudx .= rand(Float64,size(p.dudx))
  p.dudy .= rand(Float64,size(p.dudy))
  p.dvdx .= rand(Float64,size(p.dvdx))
  p.dvdy .= rand(Float64,size(p.dvdy))
  mul!(f,Emat,p)
  H(f2,p)
  @test f.dudx ≈ f2.dudx && f.dudy ≈ f2.dudy && f.dvdx ≈ f2.dvdx && f.dvdy ≈ f2.dvdy

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.dudx ≈ f2.dudx && f.dudy ≈ f2.dudy && f.dvdx ≈ f2.dvdx && f.dvdy ≈ f2.dvdy

  p = EdgeGradient(Primal,(nx,ny))
  p2 = deepcopy(p)
  Hmat = RegularizationMatrix(H,f,p)
  Emat = InterpolationMatrix(H,p,f)
  Ẽmat = InterpolationMatrix(H̃,p,f)

  mul!(p,Hmat,f)
  H(p2,f)
  @test p.dudx ≈ p2.dudx && p.dudy ≈ p2.dudy && p.dvdx ≈ p2.dvdx && p.dvdy ≈ p2.dvdy

  p.dudx .= rand(Float64,size(p.dudx))
  p.dudy .= rand(Float64,size(p.dudy))
  p.dvdx .= rand(Float64,size(p.dvdx))
  p.dvdy .= rand(Float64,size(p.dvdy))
  mul!(f,Emat,p)
  H(f2,p)
  @test f.dudx ≈ f2.dudx && f.dudy ≈ f2.dudy && f.dvdx ≈ f2.dvdx && f.dvdy ≈ f2.dvdy

  mul!(f,Ẽmat,p)
  H̃(f2,p)
  @test f.dudx ≈ f2.dudx && f.dudy ≈ f2.dudy && f.dvdx ≈ f2.dvdx && f.dvdy ≈ f2.dvdy

  end

### COMPLEX ROUTINES

@testset "Complex point data" begin

  f = ScalarData(10)
  f .= rand(10)

  fc = ScalarData(f,dtype=ComplexF64)
  ftc = TensorData(f,dtype=ComplexF64)

  a = 1.0 + 2.0im
  fc[5] = a

  @test fc[5] == a

  fvc = VectorData(f,dtype=ComplexF64)
  b = (-2.0+0im,5.0im)
  gvc = fvc + b
  @test gvc.u[5] == b[1]
  @test gvc.v[5] == b[2]

  gvc *= 2im
  @test gvc.u[5] == 2im*b[1]
  @test gvc.v[5] == 2im*b[2]

end

end


using Test
using SumOfSquares
using DynamicPolynomials

# Inspired from https://github.com/JuliaOpt/SumOfSquares.jl/issues/79
function concave_then_convex_cubic_test(optimizer, config::MOIT.TestConfig,
                                        MCT::Type{<:MOI.AbstractVectorSet})
    atol = config.atol
    rtol = config.rtol

    model = _model(optimizer)

    @polyvar x
    @variable(model, p, Poly(monomials(x, 0:3)))
    cone = ConvexPolyInnerCone{MCT}()
    @constraint(model,  p in cone, domain = (@set x >= 0))
    @constraint(model, -p in cone, domain = (@set x <= 0))
    @constraint(model, p(x =>  2) ==  8)
    @constraint(model, p(x => -2) == -8)
    @constraint(model, p(x =>  1) ==  1)
    @constraint(model, p(x => -1) == -1)

    optimize!(model)

    @test termination_status(model) == MOI.OPTIMAL

    @test primal_status(model) == MOI.FEASIBLE_POINT
    @test value(p) ≈ x^3 atol=atol rtol=rtol
end

function sos_concave_then_convex_cubic_test(optimizer, config)
    concave_then_convex_cubic_test(optimizer, config,
                                   MOI.PositiveSemidefiniteConeTriangle)
end
function sdsos_concave_then_convex_cubic_test(optimizer, config)
    concave_then_convex_cubic_test(optimizer, config,
                                   SumOfSquares.ScaledDiagonallyDominantConeTriangle)
end
function dsos_concave_then_convex_cubic_test(optimizer, config)
    concave_then_convex_cubic_test(optimizer, config,
                                   SumOfSquares.DiagonallyDominantConeTriangle)
end


abstract type BasicModel end
abstract type ModelResults end

### Mean Forecast Model
mutable struct MeanForecast <: BasicModel
    y::Vector{AbstractFloat}
    h::Integer
    level::Union{Nothing, Vector{AbstractFloat}}
    function MeanForecast(y = []; h = 5, level = [80, 95]) # Constructor
        if length(y) <= 0
            error("The input array is empty")
        end
        if h <= 0
            error("The h-value must be positive")
        end

        cleanPredInterval_(level)
        return new(y, h, level)
    end
end

function predict(model::MeanForecast)
    y = model.y
    h = model.h
    level = model.level
    data_size = length(y)

    # Compute mean, concatenate it onto the original data
    data_mean = mean(y)
    mean_vector = ones(h + data_size) .* data_mean
    fittedvalues = mean_vector
    residuals = y - fittedvalues[1:data_size] # compute residuals

    ### Coverage Probability Estimate, will calculate for 80 and 95 for now.
    c_80 = 1.28
    c_95 = 1.96

    K = 0 # number of parameters estimated in the forecasting method

    lower = zeros(length(level), h)
    upper = lower

    # Compute Prediction Intervals
    if h != 0
        # Calculate one-step prediction intervals
        one_step = sqrt(1 / (data_size - K) * (transpose(residuals) * residuals))
        σ = one_step
    end
    if h > 1
        # Calculate multi-step prediction intervals
        multi_step = fill(one_step * sqrt(1 + 1/data_size), h)
        multi_step[1] = one_step
        σ = multi_step
    end

    # Get the last h values
    yhat = fittedvalues[data_size + 1 : data_size + h]

    # Calculate prediction intervals, store in upper / lower matrices
    lower[1, :] = yhat - c_80 * σ
    lower[2, :] = yhat - c_95 * σ

    upper[1, :] = yhat + c_80 * σ
    upper[2, :] = yhat + c_95 * σ

    results = BasicMethodResults(model, fittedvalues, lower, upper)
    return results
end

function fit(model::MeanForecast)
    return model
end

### Naive Forecast Model
mutable struct NaiveForecast <: BasicModel
    y::Vector{AbstractFloat}
    h::Integer
    level::Union{Nothing, Vector{AbstractFloat}}
    function NaiveForecast(y = []; h = 5, level = [80, 95]) # Constructor
        if length(y) <= 0
            error("The input array is empty")
        end
        if h <= 0
            error("The h-value must be positive")
        end
        cleanPredInterval_(level)
        return new(y, h, level)
    end
end

function predict(model::NaiveForecast)
    y = model.y
    h = model.h
    level = model.level
    data_size = length(y)

    # Forecasts are equal to last value of observations
    lastValue = last(y)
    fitted_chunk = fill(lastValue, (h))
    fittedvalues = vcat(y, fitted_chunk)

    # Compute Prediction Intervals
    lower = zeros(length(level), h)
    upper = lower
    for i in 1:length(level)
        dist = Normal(0.0, 1.0)
        qq = quantile(dist, 0.5 * (1 + level[i]/100))
        lower[i, :] = fitted_chunk - qq * fitted_chunk
        upper[i, :] = fitted_chunk + qq * fitted_chunk
    end

    results = BasicMethodResults(model, fittedvalues, lower, upper)
    return results
end

function fit(model::NaiveForecast)
    return model
end

### Seasonal Naive Model
mutable struct SeasonalNaiveForecast <: BasicModel
    y::Vector{AbstractFloat}
    h::Integer
    m::Integer
    level::Union{Nothing, Vector{AbstractFloat}}
    function SeasonalNaiveForecast(y = []; h = 5, m = 12, level = [80, 95]) # Constructor
        if length(y) <= 0
            error("The input array is empty")
        end
        if h <= 0
            error("The h-value must be positive")
        end
        if m <= 0
            error("The m-value must be positive")
        end
        cleanPredInterval_(level)
        return new(y, h, m, level)
    end
end

function fit(model::SeasonalNaiveForecast)
    y = model.y
    h = model.h
    m = model.m
    level = model.level
    data_size = length(y)

    # Forecasts are equal to last value of observations
    k = trunc(Integer, (h-1)/m)
    fitted_chunk = y[data_size + 1 - m * (k+1) : data_size + h - m * (k+1)]
    fittedvalues = vcat(y, fitted_chunk)

    # Compute Prediction Intervals
    lower = zeros(length(level), h)
    upper = lower
    for i in 1:length(level)
        dist = Normal(0.0, 1.0)
        qq = quantile(dist, 0.5 * (1 + level[i]/100))
        lower[i, :] = fitted_chunk - qq * fitted_chunk
        upper[i, :] = fitted_chunk + qq * fitted_chunk
    end

    results = BasicMethodResults(model, fittedvalues, lower, upper)
    return results
end

# Ensures Prediction Intervals are numerical and on the interval (0, 100)
function cleanPredInterval_(level)
    for i in level
        if !isa(i, Number)
            print(isa(i, Number))
            error("Ensure that all level inputs are numeric")
        elseif i < 0 || i > 100
            error("Ensure that all confidence interval inputs are in the interval [0, 100]")
        end
    end
end

# Results Object - Common across all "Basic Method" models
mutable struct BasicMethodResults <: ModelResults # Stores results from fit(MeanForecast)
    model::BasicModel
    fittedvalues::Vector{Float64} # Vector of fitted values
    lower::Matrix{Float64} # Lower bound of prediction intervals
    upper::Matrix{Float64} # Upper bound of prediction intervals
    function BasicMethodResults(model, fittedvalues, lower, upper)
        return new(model, fittedvalues, lower, upper)
    end
end


# Apply kinetic operator to wave function in reciprocal space

function op_K( pw::PWGrid, psi::Array{ComplexF64,2}, ik::Int64 )
    #
    out = zeros(ComplexF64,size(psi))
    Ncol = size(psi)[2]

    Ngwk = pw.gkvec.Ngw[ik]
    idx_gkk = pw.gkvec.idx_gk[ik]
    G = pw.gvec.G[:,idx_gkk]

    for is = 1:Ncol
        for ig = 1:Ngwk
            Gk = G[:,ig] + pw.gkvec.kpts[:,ik]
            out[ig,is] = psi[ig,is]*( Gk[1]^2 + Gk[2]^2 + Gk[3]^2 )
        end
    end
    return 0.5*out # two minus signs
end


module LincombTests

using Test
using Gridap.Arrays
using Gridap.Fields
using Gridap.Fields: MockField, MockBasis
using FillArrays
using Gridap.TensorValues

np = 4
p = Point(1,2)
x = fill(p,np)
v = 2.0
d = 2
ndof = 8
wi = 3.0
w = fill(wi,ndof)
f = MockBasis{d}(v,ndof)
g = lincomb(f,w)
fx = evaluate(f,x)
∇fx = evaluate(∇(f),x)
gx = fx*w
∇gx = ∇fx*w
test_field(g,x,gx,grad=∇gx)

l = 10
af = Fill(f,l)
ax = fill(x,l)
aw = fill(w,l)
ag = lincomb(af,aw)
agx = fill(gx,l)
a∇gx = fill(∇gx,l)
test_array_of_fields(ag,ax,agx,grad=a∇gx)

l = 0
af = Fill(f,l)
ax = fill(x,l)
aw = fill(w,l)
ag = lincomb(af,aw)
agx = fill(gx,l)
a∇gx = fill(∇gx,l)
test_array_of_fields(ag,ax,agx,grad=a∇gx)

end #module


# test cj

A = [[4.0,1.0] [1.0,3.0]]
        b = [1.0,2.0]
        x = [2.0,1.0]
        maxiter = 100
        eps = float(10^-6.0)
        x = Solvers.cg(A, b, x, maxiter, eps)
        xtrue = [1/11,7/11]
        println("Error: ")
        println(xtrue-x)


"""
    Multiply(input=inputDimension, output=outputDimension; winit=xavier, o...)

Matrix multiplication layer with a matrix.
    (m::Multiply)(x) = m.w * x

Input:  InputDimension by BatchSize `x` matrix where Batchsize ∈ {1, 2,.... N}

Output: Projection of a vector `x` from `inputDimension` to `outputDimension`

# Keywords
* `input=inputDimension`   : input matrix (vector) first dimension
* `output=outputDimension` : output matrix (vector) first dimension
* `winit=xaiver`           : distribution for weight initialization
"""
struct Multiply <: Layer
    w # weight
end

function Multiply(;input::Int, output::Int, winit=xavier, o...)
    w = Parameter(output, input; init=winit, o...)
    return Multiply(w)
end

(m::Multiply)(x) = m.w * x
(m::Multiply)(x::Array{T}) where T<:Integer = m.w[:, x]

"""
    Embed(input=inputDimension, output=outputDimension; winit=xavier, o...)

Lookup layer to encode sparse one-hot vector into dense continous vector.
    (m::Embed)(x) = m.w[:, x]

Input:  An integer array containing 1 locations of one-hot vectors

Output: Input-indexed columns of layer's matrix (m.w)

# Keywords
* `input=inputDimension`   : One-hot vector's dimension
* `output=outputDimension` : Dense vector's dimension
* `winit=xaiver`           : Distribution for weight initialization

"""
Embed = Multiply

"""
    BatchMul(;input=inputDimension, output=outputDimension, winit=xavier, o...)
Matrix tensor multiplication layer with a 2D matrix.

Input: N-dimensional tensor containing instances in the first dimension

Output: Multiplication of all instances in the raws input. If `flatte=true` output is reduced to 2D matrix, otherwise input size is preserved except first dimension

# Keywords
* `input=inputDimension`   : Input tensor's  1st dimension
* `output=outputDimension` : Output tensor's 1st dimension
* `winit=xaiver`           : Distribution for weight initialization
* `flatten=true`           : Reduces the output tensor into 2D matrix

"""
struct BatchMul <: Layer
    w # weight
end

function BatchMul(;input::Int, output::Int, winit=xavier, o...)
    w = Parameter(output, input; init=winit, o...)
    return BatchMul(w)
end

function (m::BatchMul)(x; flatten=false)
    xdims  = size(x)
    x_flat = reshape(x, xdims[1], prod(xdims[2:end]))
    yout = m.w * x_flat
    if flatten
        return yout
    end
    return reshape(yout, size(yout, 1), xdims[2:end]...)
end

"""
    Linear(;input=inputDimension, output=outputDimension, winit=xavier, binit=zeros, o...)
Linear layer applying usual linear transformation `w * x .+ b`.

Input:  inputDimension by BatchSize `x` matrix where Batchsize ∈ {1, 2,.... N}

Output: outputDimension by BatchSize where Batchsize ∈ {1, 2,.... N}

# Keywords
* `input=inputDimension`   : Input tensor's  1st dimension
* `output=outputDimension` : Output tensor's 1st dimension
* `winit=xaiver`           : Distribution for weight initialization
* `binit=zeros`            : Distribution for bias initialization

"""
struct Linear <: Layer
    m::Multiply
    b # bias
end

function Linear(;input::Int, output::Int, winit=xavier, binit=zeros, o...)
    w = Multiply(;input=input, output=output, o...)
    b = Parameter(output; init=binit, o...)
    return Linear(w, b)
end

(l::Linear)(x) = l.m(x) .+ l.b

"""
    Dense(;input=inputDimension, output=outputDimension, winit=xavier, binit=zeros, activation=relu, o...)
Dense layer applying non-linear function, e.g. relu, on top of linear layer. If activation is not provided, Linear transformation applied

Input:  inputDimension by BatchSize `x` matrix where Batchsize ∈ {1, 2,.... N}

Output: outputDimension by Batchsize where Batchsize ∈ {1, 2,.... N}

# Keywords
* `input=inputDimension`   : Input tensor's  1st dimension
* `output=outputDimension` : Output tensor's 1st dimension
* `winit=xaiver`           : Distribution for weight initialization
* `binit=zeros`            : Distribution for bias initialization
* `activation=relu`        : Activation function applied after linear transformation, if nothing Linear layer initialized
"""
struct Dense <: Layer
    l::Linear
    activate # activation
end

function Dense(;input::Int, output::Int, activation=nothing, winit=xavier, binit=zeros, o...)
    l = Linear(;input=input, output=output, winit=winit, binit=binit, o...)
    if activation == nothing
        return l
    end
    activate = activation
    return Dense(l, activate)
end

(d::Dense)(x) = d.activate.(d.l(x))

# TODO: Add dropout layer


using VectorAutoregressions
using Base.Test

#-----------Set base-path------------------------------------------
path = Pkg.dir("VectorAutoregressions")

#--------------------Test VAR model---------------------------------------------------------
# comparison against http://www-personal.umich.edu/~lkilian/figure9_1_chol.zip
y = readdlm(joinpath(path,"test","cholvar_test_data.csv"),',')

V = VAR(y, 4, true)

@test V.p == 4
@test V.Σ ≈ [
        312.524637669663 0.773643992117187 0.919341486304653
        0.773643992117187 0.0515304724687740 0.0148802276681216
        0.919341486304653 0.0148802276681216 0.557021891337960
        ]

@test V.β[:,1:4] ≈ [
        -0.680676727884458 -0.00637960178073601 0.936519611007488 1.75104737240770
        -0.0152242907833718 0.00177981748889471 0.599086054531876 0.0281955980964510
        0.436329751995733 -0.00236721376013309 -0.218737236527547 0.332361646200618
        ]

@test V.ϵ[:,1:4] ≈ [
        70.3934221013155 -0.443580106581005 4.21164200500382 1.22299717886091
        0.189207494365181 0.287694213161629 0.692414550347607 0.405418261036420
        -1.39379479232130 0.277340573620367 -0.573151724739961 -0.395501452903569
        ]

mIRFa = IRFs_a(V,4,true)

@test [mIRFa.IRF[1,1:4]'; mIRFa.IRF[4,1:4]'; mIRFa.IRF[7,1:4]'] ≈ [
        17.6783663744607 0.0192642533976305 -2.85156060358692 0.594909295963393
        0.0437621879606954 0.0591478592307352 0.113331057827868 0.109932348126692
        0.0520037579735191 -0.0341368375885218 -0.111284843056732 -0.0491881671489620
        ]

mIRFb = IRFs_b(V,4,10,true)

@test [mIRFb.IRF[1,1:4]'; mIRFb.IRF[4,1:4]'; mIRFb.IRF[7,1:4]'] ≈ [
        17.6783663744607 0.0192642533976305 -2.85156060358692 0.594909295963393
        0.0437621879606954 0.0591478592307352 0.113331057827868 0.109932348126692
        0.0520037579735191 -0.0341368375885218 -0.111284843056732 -0.0491881671489620
        ]

#--------------------Test BVAR model---------------------------------------------------------
# comparison against http://cremfi.econ.qmul.ac.uk/outgoing/bvar.zip
include(joinpath(path,"src","bvar.jl"))
y = readdlm(joinpath(path,"test","bvar_data.csv"), ',')
y = y[:,1:3]
prior = Hyperparameter()
mForecast = fit_bvar(y,prior)

@test isapprox([0.8510 1.4081 2.2570 2.3415 2.4622 2.5835 2.6867 2.5790 2.5897 2.5767],median(mForecast[:,:,1],1); atol = 1)
@test isapprox([1.9614 2.2587 1.8328 1.8745 2.0870 2.2014 2.3303 2.5225 2.5453 2.59387],median(mForecast[:,:,2],1); atol = 1)
@test isapprox([-0.3827 -0.2272 -0.1532 -0.0735 0.0784 0.1620 0.3764 0.5247 0.7264 0.8861],median(mForecast[:,:,3],1); atol = 1)

#--------------------Test Local Projection IRFs---------------------------------------------------------
# comparison against Kilian and Kim (2011) generated data from VAR(12)

#-----------Load data----------------------------------------------
y      = readcsv(joinpath(path,"test","lp_data.csv"))
irfl   = readcsv(joinpath(path,"test","lp_test_lp_chol_irf.csv"))
stdl   = readcsv(joinpath(path,"test","lp_test_lp_std.csv"))
stdv   = readcsv(joinpath(path,"test","lp_test_var_stdv.csv"))
COVsig = readcsv(joinpath(path,"test","lp_test_var_covsig.csv"))
rfirfl = readcsv(joinpath(path,"test","lp_test_lp_rf_irf.csv"))
mlag   = readcsv(joinpath(path,"test","lp_test_lp_laglength.csv"))

#-----------Hyperparameter-----------------------------------------
const p = 12    # lag length
const H = 24    # horizon
const intercept = true

#-----------Reduced form local projection IRFs-------------------
RF_IRFs = IRFs_localprojection(y, p, H)
mRFIRFs,CI = RF_IRFs.IRF, RF_IRFs.CI

#-----------Structural local projection IRFs-------------------
# Using a VAR(pbar) as auxiliary model for cholesky identification
V = VAR(y,p,intercept)
A0inv = V.Σ |> λ -> cholfact(λ)[:L] |> full
mStd,mCov_Σ = irf_ci_asymptotic(V, H, V.inter)

IRF = IRFs_localprojection(y, p, H, A0inv, mCov_Σ)
mIRF,CI = IRF.IRF, IRF.CI
CIl,CIh = CI.CIl, CI.CIh

@test isapprox(mRFIRFs,rfirfl)
@test isapprox(A0inv,[0.721465    0.0          0.0        0.0
               -0.0842337   2.68627      0.0        0.0
                2.61134     0.724746    26.8394     0.0
                0.0886017  -0.00940385   0.0693017  0.486568],
                atol = 0.0001)
@test isapprox(V.Σ, [0.520512   -0.0607717    1.88399   0.063923
              -0.0607717   7.22317      1.7269   -0.0327246
               1.88399     1.7269     727.695     2.08457
               0.063923   -0.0327246    2.08457   0.249489],
               atol = 0.0001)
@test isapprox(mStd, stdv, atol = 0.1)
@test isapprox(mCov_Σ, COVsig, atol = 0.1)
@test isapprox(mIRF, irfl)
@test isapprox(CIl,(irfl - 1.96*stdl),atol = 0.1)
@test isapprox(CIh,(irfl + 1.96*stdl),atol = 0.1)

#-----------Test LP with different lag-length at different horizon-------
const vP = p*ones(Int64,H) # vector of lag-length

#-----------Reduced form local projection IRFs-------------------
RF_IRFs = IRFs_localprojection(y, vP, H)
mRFIRFs,CI = RF_IRFs.IRF, RF_IRFs.CI

#-----------Structural local projection IRFs-------------------
# Using a VAR(pbar) as auxiliary model for cholesky identification
V = VAR(y,p,intercept)
A0inv = V.Σ |> λ -> cholfact(λ)[:L] |> full
mStd,mCov_Σ = irf_ci_asymptotic(V, H, V.inter)

IRF = IRFs_localprojection(y, vP, H, A0inv, mCov_Σ)
mIRF,CI = IRF.IRF, IRF.CI
CIl,CIh = CI.CIl, CI.CIh

@test isapprox(mRFIRFs,rfirfl)
@test isapprox(mIRF, irfl)
@test isapprox(CIl,(irfl - 1.96*stdl),atol = 0.1)
@test isapprox(CIh,(irfl + 1.96*stdl),atol = 0.1)

#-----------Test LP lag-length selecion procedure----------------
const pbar = 12      # max order of lag to test

#-----------Select lag-length with AIC, BIC, AICC, HQC-----------
mVARlag = (var_lagorder(y,pbar,ic) for ic in ["aic","bic","aicc","hqc"]) |> λ -> hcat(collect(λ)...)
mLPlag  = (lp_lagorder(y,pbar,H,ic) for ic in ["aic","bic","aicc","hqc"]) |> λ -> hcat(collect(λ)...)

@test mLPlag == mlag
@test mVARlag == [12 2 12 3]


# dependencies
using LFAToolkit
using LinearAlgebra
using CSV
using DataFrames

# setup
dimension = 1
numbercomponents = 1
mesh = []
if dimension == 1
    mesh = Mesh1D(1.0)
elseif dimension == 2
    mesh = Mesh2D(1.0, 1.0)
elseif dimension == 3
    mesh = Mesh3D(1.0, 1.0, 1.0)
end
smoothingfactors = DataFrame()

# test range
for P = 1:4
    println("p = ", 2^P)
    # setup
    # -- bases
    p = 2^P

    # -- diffusion operator
    diffusion = GalleryOperator("diffusion", p + 1, p + 1, mesh)

    # -- Jacobi smoother
    jacobi = Jacobi(diffusion)

    # compute smoothing factor
    # -- setup
    numbersteps = 100
    maxeigenvalue = 0
    θ_min = -π / 2
    θ_min_high = π / 2
    θ_max = 3π / 2
    θ_step = 2π / numbersteps
    θ_range = θ_min:θ_step:(θ_max-θ_step)
    ω_range = 0.50:0.01:1.05
    num_ω = max(size(ω_range)...)

    # -- compute
    maxeigenvalue = zeros(num_ω)
    θ_maxeigenvalue = -1 * ones(num_ω, dimension)
    # -- 1D --
    if dimension == 1
        for i = 1:numbersteps
            θ = [θ_range[i]]
            if abs(θ[1]) > π / 128 && θ[1] > θ_min_high
                A = computesymbols(jacobi, [1.0], θ)
                for w = 1:num_ω
                    eigenvalues = [abs(val) for val in eigvals(I - ω_range[w] * (I - A))]
                    currentmaxeigenvalue = max(eigenvalues...)
                    if (currentmaxeigenvalue > maxeigenvalue[w])
                        maxeigenvalue[w] = currentmaxeigenvalue
                        θ_maxeigenvalue[w, :] = θ / π
                    end
                end
            end
        end
        # -- 2D --
    elseif dimension == 2
        for i = 1:numbersteps, j = 1:numbersteps
            θ = [θ_range[i], θ_range[j]]
            if sqrt(abs(θ[1])^2 + abs(θ[2])^2) > π / 128 &&
               (θ[1] > θ_min_high || θ[2] > θ_min_high)
                A = computesymbols(jacobi, [1.0], θ)
                for w = 1:num_ω
                    eigenvalues = [abs(val) for val in eigvals(I - ω_range[w] * (I - A))]
                    currentmaxeigenvalue = max(eigenvalues...)
                    if (currentmaxeigenvalue > maxeigenvalue[w])
                        maxeigenvalue[w] = currentmaxeigenvalue
                        θ_maxeigenvalue[w, :] = θ / π
                    end
                end
            end
        end
        # -- 3D --
    elseif dimension == 3
        for i = 1:numbersteps, j = 1:numbersteps, k = 1:numbersteps
            θ = [θ_range[i], θ_range[j], θ_range[k]]
            if sqrt(abs(θ[1])^2 + abs(θ[2])^2 + abs(θ[3])^2) > π / 128 &&
               (θ[1] > θ_min_high || θ[2] > θ_min_high || θ[3] > θ_min_high)
                A = computesymbols(jacobi, [1.0], θ)
                for w = 1:num_ω
                    eigenvalues = [abs(val) for val in eigvals(I - ω_range[w] * (I - A))]
                    currentmaxeigenvalue = max(eigenvalues...)
                    if (currentmaxeigenvalue > maxeigenvalue[w])
                        maxeigenvalue[w] = currentmaxeigenvalue
                        θ_maxeigenvalue[w, :] = θ / π
                    end
                end
            end
        end
    end
    (_, ω_index) = findmin(maxeigenvalue)
    append!(
        smoothingfactors,
        DataFrame(
            p = p,
            ω = ω_range[ω_index],
            θ = θ_maxeigenvalue[ω_index, :],
            ρ = maxeigenvalue[ω_index],
        ),
    )
end

CSV.write("jacobi-smoothing-factor.csv", smoothingfactors)


#=

Code for Sparsification.

Right now, just implements Spielman-Srivastava

=#

"""
    as = sparsify(a; ep=0.5)

Apply Spielman-Srivastava sparsification: sampling by effective resistances.
`ep` should be less than 1.
"""
function sparsify(a; ep=0.3, matrixConcConst=4.0, JLfac=4.0)

  if ep > 1
    @warn "Calling sparsify with ep > 1 can produce a disconnected graph."
  end

  f = approxchol_lap(a,tol=1e-2);

  n = size(a,1)
  k = round(Int, JLfac*log(n)) # number of dims for JL

  U = wtedEdgeVertexMat(a)
  m = size(U,1)
  R = randn(Float64, m,k)
  UR = U'*R;

  V = zeros(n,k)
  for i in 1:k
    V[:,i] = f(UR[:,i])
  end

  (ai,aj,av) = findnz(triu(a))
  prs = zeros(size(av))
  for h in 1:length(av)
      i = ai[h]
      j = aj[h]
      prs[h] = min(1,av[h]* (norm(V[i,:]-V[j,:])^2/k) * matrixConcConst *log(n)/ep^2)
  end

  ind = rand(Float64,size(prs)) .< prs

  as = sparse(ai[ind],aj[ind],av[ind]./prs[ind],n,n)
  as = as + as'

  return as

end


let
    n = 50
    p = 3
    t = (1:n)/n
    Ut,Vt = egrss.generators(t,p);

    K = Symmetric(tril(Ut'*Vt) + triu(Vt'*Ut,1))
    @test ( isapprox(egrss.full(Ut,Vt),K) )
    @test ( isapprox(egrss.full(Ut,Vt,1e-3),K+Diagonal(1e-3*ones(n))) )
    @test ( isapprox(egrss.full(Ut,Vt,1e-3*ones(n)),K+Diagonal(1e-3*ones(n))) )
end


using Flux, DiffEqFlux, OrdinaryDiffEq
using Zygote

u0 = Float32[2.; 0.]
datasize = 30
tspan = (0.0f0,1.5f0)

function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
t = range(tspan[1],tspan[2],length=datasize)
prob = ODEProblem(trueODEfunc,u0,tspan)
ode_data = Array(solve(prob,Tsit5(),saveat=t))

dudt = Chain(x -> x.^3,
             Dense(2,50,tanh),
             Dense(50,2))
#= ps = Flux.params(dudt) =#

p = DiffEqFlux.destructure(dudt)
ps = Flux.params(p)
dudt_(u,p,t) = DiffEqFlux.restructure(dudt,p)(u)
prob = ODEProblem(dudt_,u0,tspan,p)

diffeq_adjoint(p,u0,prob,Tsit5())

forward(u->diffeq_adjoint(p,u,prob,Tsit5()),u0)

n_ode = x->neural_ode(dudt,x,tspan,Tsit5(),saveat=t,reltol=1e-7,abstol=1e-9)

pred = n_ode(u0) # Get the prediction using the correct initial condition

function predict_n_ode()
  n_ode(u0)
end
loss_n_ode() = sum(abs2,ode_data .- predict_n_ode())

data = Iterators.repeated((), 10)
opt = ADAM(0.1)
cb = function () #callback function to observe training
  display(loss_n_ode())
end

# Display the ODE with the initial parameter values.
cb()

Flux.train!(loss_n_ode, ps, data, opt, cb = cb)


module QuantumAlgebra

export scal,param,a,adag,f,fdag,OpSumAnalytic,ExpVal,Corr
export σx,σy,σz,σp,σm,comm,latex
export Avac,vacA,vacExpVal
export CorrOrExp,ascorr
export @Pr_str, @Pc_str, ∑
#export σ,preftuple,exptuple,optuple,prodtuples,sumtuples

using Printf

# we will want to overload these operators and functions for our custom types
import Base: ==, ≈, *, +, -, isless, length, adjoint, print, zero, one

# define for σx, σy, σz
@enum SpatialIndex x=1 y=2 z=3
SpatialIndex(a::SpatialIndex) = a

# we sometimes want to sort integers and symbols together, with integers coming first
sortsentinel(x::Integer) = (1,x)
sortsentinel(x::Symbol) = (2,x)

const OpIndex = Union{Int,Symbol}
const OpIndices = NTuple{N,OpIndex} where N

abstract type Operator end
abstract type Scalar <: Operator; end
"Represent a scalar value (i.e., a number)"
struct scal{T<:Number} <: Scalar; v::T; end
"`param(g,(:i,:j),'n')`: represent a scalar named parameter ``g_{i,j}``. state can be purely real (`'r'`), not conjugated (`'n'`), or conjugated (`'c'`)"
struct param{T<:OpIndices} <: Scalar
    name::Symbol
    state::Char
    inds::T
    function param(name,state::Char,inds::OpIndex...)
        state in ('n','r','c') || error("state has to be one of n,r,c")
        new{typeof(inds)}(name,state,inds)
    end
    param(name,inds::OpIndex...) = param(name,'n',inds...)
    param(name,state::Char,inds::Tuple) = param(name,state,inds...)
    param(name,inds::Tuple) = param(name,'n',inds...)
end
struct δ <: Scalar
    inds::Tuple{OpIndex,OpIndex}
    δ(inds) = δ(inds...)
    function δ(iA::OpIndex,iB::OpIndex)
        # sort indices
        if iA == iB
            scal(1)
        elseif iA isa Integer && iB isa Integer
            # e.g., δ_1,3 = 0 (integers are not symbolic!)
            scal(0)
        elseif sortsentinel(iB) < sortsentinel(iA)
            new((iB,iA))
        else
            new((iA,iB))
        end
    end
end
function δ(Ainds::OpIndices,Binds::OpIndices)::Operator
    if length(Ainds) == length(Binds)
        prod(collect(Operator,δ.(Ainds,Binds)))
    else
        scal(0)
    end
end

abstract type BaseOperator <: Operator; end
for (op,desc,sym) in (
    (:a,   "bosonic annihilation","a"),
    (:adag,"bosonic creation","a^†"),
    (:f,   "fermionic annihilation","f"),
    (:fdag,"fermionic creation","f^†"),
    (:σminus,"TLS annihilation","σ^-"),
    (:σplus,"TLS creation","σ^+"))
    @eval begin
        "`$($op)(inds)`: represent $($desc) operator ``$($sym)_{inds}``"
        struct $op{T<:OpIndices} <: BaseOperator
            inds::T
            $op(inds::OpIndex...) = new{typeof(inds)}(inds)
            $op(inds::Tuple) = $op(inds...)
        end
    end
end

"`σ(a,n)`: represent Pauli matrix ``σ_{a,n}`` for two-level system (TLS) ``n``, where ``a ∈ \\{x,y,z\\}`` or ``\\{1,2,3\\}`` is the type of Pauli matrix."
struct σ{T<:OpIndices} <: BaseOperator
    a::SpatialIndex
    inds::T
    σ(a,inds::OpIndex...) = new{typeof(inds)}(SpatialIndex(a),inds)
    σ(a,inds::Tuple) = σ(a,inds...)
end

zero(::Type{<:Operator}) = scal(0)
zero(::Operator) = scal(0)
one(::Type{<:Operator}) = scal(1)
one(::Operator) = scal(1)

using_σpm = false
use_σpm(t::Bool=true) = eval(:( using_σpm = $t; nothing ))
use_σxyz() = use_σpm(false)

"`σp(n)`: construct ``σ^+_n = \\frac12 σ_{x,n} + \\frac{i}{2} σ_{y,n}``"
σp(n...) = using_σpm ? σplus(n...) : scal(1//2)*σ(x,n...) + scal(1im//2)*σ(y,n...)
"`σm(n)`: construct ``σ^-_n = \\frac12 σ_{x,n} - \\frac{i}{2} σ_{y,n}``"
σm(n...) = using_σpm ? σminus(n...) : scal(1//2)*σ(x,n...) - scal(1im//2)*σ(y,n...)
"`σx(n)`: construct ``σ_{x,n}``"
σx(n...) = using_σpm ? σminus(n...) + σplus(n...) : σ(x,n...)
"`σy(n)`: construct ``σ_{y,n}``"
σy(n...) = using_σpm ? scal(1im) * (σminus(n...) - σplus(n...)) : σ(y,n...)
"`σz(n)`: construct ``σ_{z,n}``"
σz(n...) = using_σpm ? scal(2)*σplus(n...)*σminus(n...) - scal(1) : σ(z,n...)

struct OpProd <: Operator; A::Operator; B::Operator; end
struct OpSum  <: Operator; A::Operator; B::Operator; end

"`OpSumAnalytic(i::Symbol,A::Operator)` or `∑(i,A)`: represent ``\\sum_{i} A``, with all possible values of ``i`` assumed to be included"
struct OpSumAnalytic <: Operator
    ind::Symbol
    A::Operator
    OpSumAnalytic(ind::Symbol,A::scal) = A.v == 0 ? A : new(ind,A)
    OpSumAnalytic(ind::Symbol,A::Operator) = begin
        if A isa OpProd && A.A isa scal
            A.A*OpSumAnalytic(ind,A.B)
        elseif A isa OpSumAnalytic && A.ind < ind
            # reorder nested sums by indices
            OpSumAnalytic(A.ind,OpSumAnalytic(ind,A.A))
        else
            for t in preftuple(A)
                # if there is a δ in the expression and it has the sum index, the sum disappears
                if t isa δ && ind in t.inds
                    irepl = ind==t.inds[1] ? t.inds[2] : t.inds[1]
                    return replace_index(A,ind,irepl)
                end
            end
            return new(ind,A)
        end
    end
    OpSumAnalytic(ind::Symbol,A::OpSum) = OpSumAnalytic(ind,A.A) + OpSumAnalytic(ind,A.B)
end

const ∑ = OpSumAnalytic

"`ExpVal(A::Operator)`: represent expectation value ``⟨A⟩``"
struct ExpVal <: Scalar
    A::Operator
    ExpVal(A::Scalar) = A
    ExpVal(A::OpSum) = ExpVal(A.A) + ExpVal(A.B)
    ExpVal(A::OpProd) = A.A isa Scalar ? A.A*ExpVal(A.B) : new(A)
    ExpVal(A::OpSumAnalytic) = OpSumAnalytic(A.ind,ExpVal(A.A))
    ExpVal(A::Operator) = new(A)
end

"""
    Corr(A::Operator)

Represents correlations of operator ``A``. ``A`` should be a product for this to
make sense, in which case ``⟨AB⟩_c = ⟨AB⟩ - ⟨A⟩⟨B⟩``, with corresponding
extensions for products of more operators.

See also: [`ascorr`](@ref)
"""
struct Corr <: Scalar
    A::Operator
    Corr(A::Scalar) = A
    Corr(A::OpSum) = Corr(A.A) + Corr(A.B)
    Corr(A::OpProd) = A.A isa Scalar ? A.A*Corr(A.B) : new(A)
    Corr(A::Operator) = new(A)
end

==(A::scal,B::scal) = A.v == B.v
==(A::OpProd,B::OpProd) = A.A == B.A && A.B == B.B
==(A::OpSum, B::OpSum)  = A.A == B.A && A.B == B.B
==(A::OpSumAnalytic, B::OpSumAnalytic) = A.A == B.A && A.ind == B.ind
==(A::ExpVal, B::ExpVal) = A.A == B.A
==(A::Corr, B::Corr) = A.A == B.A
==(A::param, B::param) = (A.name,A.inds,A.state) == (B.name,B.inds,B.state)

≈(A::Operator,B::Operator) = A == B
≈(A::scal,B::scal) = A.v ≈ B.v
≈(A::OpProd,B::OpProd) = A.A ≈ B.A && A.B ≈ B.B
≈(A::OpSum, B::OpSum)  = A.A ≈ B.A && A.B ≈ B.B
≈(A::OpSumAnalytic, B::OpSumAnalytic) = A.A ≈ B.A && A.ind == B.ind
≈(A::ExpVal, B::ExpVal) = A.A ≈ B.A
≈(A::Corr, B::Corr) = A.A ≈ B.A

prodtuple(A::Operator) = (A,)
prodtuple(A::OpProd) = (prodtuple(A.A)...,prodtuple(A.B)...)

# make a tuple from a product, containing only either prefactor types, expectation value types, or operators
for (name,types) in [(:pref,(scal,param,δ)),(:exp,(ExpVal,Corr)),(:op,(adag,a,f,fdag,σ,σplus,σminus))]
    name = Symbol(name,:tuple)
    types = Union{types...}
    @eval $name(A::$types) = (A,)
    @eval $name(A::Operator) = ()
    @eval $name(A::OpSum) = error("cannot get $($name) for OpSum!")
    @eval $name(A::OpSumAnalytic) = ($name(A.A)...,)
    @eval $name(A::OpProd) = ($name(A.A)...,$name(A.B)...)
end
prodtuples(A::Operator) = (preftuple(A), exptuple(A), optuple(A))

sumtuple(A::Operator) = (A,)
sumtuple(A::OpSum) = (sumtuple(A.A)...,sumtuple(A.B)...)

OpOrder = (scal,δ,param,ExpVal,Corr,adag,a,fdag,f,σplus,σminus,σ,OpProd,OpSumAnalytic,OpSum)
for (ii,op1) in enumerate(OpOrder)
    for op2 in OpOrder[ii+1:end]
        @eval isless(::$op1,::$op2) = true
        @eval isless(::$op2,::$op1) = false
    end
end
isless(A::scal,B::scal) = (real(A.v),imag(A.v)) < (real(B.v),imag(B.v))

# do not order parameters by whether they are purely real, or conjugated or not
isless(A::param,B::param) = (A.name,sortsentinel.(A.inds)) < (B.name,sortsentinel.(B.inds))
isless(A::σ,B::σ) = (sortsentinel.(A.inds),A.a) < (sortsentinel.(B.inds),B.a)
for op in (δ,a,adag,f,fdag,σminus,σplus)
    @eval isless(A::$op,B::$op) = sortsentinel.(A.inds) < sortsentinel.(B.inds)
end
for op in (ExpVal,Corr,OpSumAnalytic)
    @eval isless(A::$op,B::$op) = A.A < B.A
end
function isless(A::OpProd,B::OpProd)
    # order operator products first by number of operators (also within expectation values),
    # then by operators, expectation values and correlations, then by reversed prefactors (to order params, not scalars). Again use tuples to write this easily
    Aprefs,Aexps,Aops = prodtuples(A)
    Bprefs,Bexps,Bops = prodtuples(B)
    (length(A),Aops,Aexps,reverse(Aprefs)) < (length(B),Bops,Bexps,reverse(Bprefs))
end

# prefactors do not count for length calculation
for op in [scal,δ,param]
    @eval length(::$op) = 0
end
for op in [adag,a,fdag,f,σ,σminus,σplus]
    @eval length(::$op) = 1
end
for op in [ExpVal,Corr,OpSumAnalytic]
    @eval length(A::$op) = length(A.A)
end
length(A::OpProd) = length(A.A) + length(A.B)

*(A::Operator) = A # needed to make prod((A,)) work
*(A::OpProd,B::Operator) = A.A*(A.B*B)
*(A::OpSum,B::Operator)      = A.A*B + A.B*B
*(A::OpSum,B::OpSum)         = A.A*B + A.B*B # resolve ambiguity
*(A::OpSum,B::OpSumAnalytic) = A.A*B + A.B*B # resolve ambiguity
*(A::Operator,     B::OpSum) = A*B.A + A*B.B
*(A::OpProd,       B::OpSum) = A*B.A + A*B.B # resolve ambiguity
*(A::OpSumAnalytic,B::OpSum) = A*B.A + A*B.B # resolve ambiguity

# allow addition, substraction, and multiplication with a number x by promoting it to scal(x) operator
*(x::Number,A::Operator) = scal(x)*A
*(A::Operator,x::Number) = scal(x)*A
+(x::Number,A::Operator) = scal(x)+A
+(A::Operator,x::Number) = scal(x)+A
-(x::Number,A::Operator) = scal(x)-A
-(A::Operator,x::Number) = A-scal(x)

function *(A::Operator,B::Operator)::Operator
    if A isa scal && A.v==0
        A
    elseif A isa scal && A.v==1
        B
    elseif A isa scal && B isa scal
        scal(A.v*B.v)
    elseif A isa scal && B isa OpProd && B.A isa scal
        scal(A.v*B.A.v)*B.B
    elseif B isa OpProd && (A*B.A) != OpProd(A,B.A)
        # if A*B.A is not ordered as we want, evaluate A*B.A first
        (A*B.A)*B.B
    elseif A isa δ && A.inds[2] in indextuple(B)
        # if second index of δ_iA,iB shows up on RHS, replace by iA
        # (relies on indices in δ being sorted)
        A * replace_index(B,A.inds[2],A.inds[1])
    elseif A isa σ && B isa σ && A.inds == B.inds
        # σa σb = δab + i ϵabc σc = δab + 1/2 [σa,σb]
        A.a==B.a ? scal(1) : scal(1//2)*comm(A,B)
    elseif A isa Union{σplus,σminus,f,fdag} && typeof(B)==typeof(A) && A.inds == B.inds
        scal(0)
    elseif A>B
        B*A + comm(A,B)
    else
        OpProd(A,B)
    end
end

# these are the only overrides for minus, everything else falls back to +
-(A::Operator) = scal(-1)*A
-(A::Operator,B::Operator) = A + scal(-1)*B

separate_prefac(A::Operator) = (1,A)
separate_prefac(A::OpProd) = A.A isa scal ? (A.A.v,A.B) : (1,A)
function combinable(A::Operator,B::Operator)
    sA, oA = separate_prefac(A)
    sB, oB = separate_prefac(B)
    if oA==oB
        true, scal(sA+sB)*oA
    else
        false, scal(0)
    end
end

+(A::OpSum,B::Operator) = A.A + (A.B + B)
function +(A::Operator,B::Operator)::Operator
    A isa scal && A.v == 0 && return B
    A isa scal && B isa scal && return scal(A.v+B.v)
    A isa scal && B isa OpSum && B.A isa scal && return scal(A.v+B.A.v) + B.B

    # ignore scalars in ordering
    sA, oA = separate_prefac(A)
    if B isa OpSum
        sB, oB = separate_prefac(B.A)
        oA>oB && return B.A + (A + B.B)
    else
        sB, oB = separate_prefac(B)
        oA>oB && return B + A
    end

    flag, S = combinable(A,B)
    flag && return S

    if B isa OpSum
        flag, S = combinable(A,B.A)
        flag && return S + B.B
    end
    return OpSum(A,B)
end

adjoint(A::scal) = scal(conj(A.v))
adjoint(A::param) = A.state=='r' ? A : param(A.name,A.state=='n' ? 'c' : 'n',A.inds)
adjoint(A::a) = adag(A.inds)
adjoint(A::adag) = a(A.inds)
adjoint(A::f) = fdag(A.inds)
adjoint(A::fdag) = f(A.inds)
adjoint(A::σminus) = σplus(A.inds)
adjoint(A::σplus) = σminus(A.inds)
adjoint(A::σ) = A
adjoint(A::OpSum) = A.A' + A.B'
adjoint(A::OpProd) = A.B' * A.A'
adjoint(A::ExpVal) = ExpVal(A.A')
adjoint(A::Corr) = Corr(A.A')
adjoint(A::OpSumAnalytic) = OpSumAnalytic(A.ind,A.A')

comm(A::Scalar,B::Operator) = scal(0)
comm(A::Scalar,B::OpProd)   = scal(0)
comm(A::Scalar,B::OpSum)    = scal(0)
comm(A::Scalar,B::OpSumAnalytic) = scal(0)
comm(A::Operator,B::Scalar) = scal(0)
comm(A::OpSum,   B::Scalar) = scal(0)
comm(A::OpProd,  B::Scalar) = scal(0)
comm(A::OpSumAnalytic,B::Scalar) = scal(0)
comm(A::Scalar,B::Scalar)   = scal(0)
comm(A::OpProd,B::Operator) = comm(A.A,B)*A.B + A.A*comm(A.B,B)
comm(A::OpProd,B::OpSumAnalytic) = comm(A.A,B)*A.B + A.A*comm(A.B,B)
comm(A::OpProd,B::OpProd)   = comm(A.A,B)*A.B + A.A*comm(A.B,B)
comm(A::Operator,B::OpProd) = comm(A,B.A)*B.B + B.A*comm(A,B.B)
comm(A::OpSumAnalytic, B::OpProd) = comm(A,B.A)*B.B + B.A*comm(A,B.B)
comm(A::OpSum,B::Operator)  = comm(A.A,B) + comm(A.B,B)
comm(A::OpSum,B::OpSum)     = comm(A.A,B) + comm(A.B,B)
comm(A::OpSum,B::OpProd)    = comm(A.A,B) + comm(A.B,B)
comm(A::Operator,B::OpSum)  = comm(A,B.A) + comm(A,B.B)
comm(A::OpProd,  B::OpSum)  = comm(A,B.A) + comm(A,B.B)
comm(A::OpSumAnalytic, B::OpSum) = comm(A,B.A) + comm(A,B.B)
# different types of operators commute
commgroups = (Union{a,adag},Union{f,fdag},Union{σ,σminus,σplus})
for (ii,op) in enumerate(commgroups)
    for op2 in commgroups[ii+1:end]
        @eval comm(A::$op,B::$op2) = scal(0)
        @eval comm(A::$op2,B::$op) = scal(0)
    end
end
# these operators commute with themselves
for op in (a,adag,σplus,σminus)
    @eval comm(A::$op,B::$op) = scal(0)
end
comm(A::a,B::adag) = δ(A.inds,B.inds)
comm(A::adag,B::a) = -comm(B,A)
comm(A::σplus,B::σminus) = δ(A.inds,B.inds)*σz(A.inds)
comm(A::σminus,B::σplus) = -comm(B,A)

# {f_i, fdag_j} = f_i fdag_j + fdag_j f_i = δ_{i,j}
# f_i fdag_j = δ_{i,j} - fdag_j f_i
# [f_i, fdag_j] = f_i fdag_j - fdag_j f_i = δ_{i,j} - fdag_j f_i - fdag_j f_i = δ_{i,j} - 2 fdag_j f_i
comm(A::f,B::fdag) = δ(A.inds,B.inds) - 2*B*A
comm(A::fdag,B::f) = -comm(B,A)
# {f_i,f_j} = f_i f_j + f_j f_i = 0
# we assume that if we need the commutator, it's because of exchanging order
# [f_i,f_j] = f_i f_j - f_j f_i = -f_j f_i - f_j f_i
comm(A::f,B::f) = -2*B*A
comm(A::fdag,B::fdag) = -2*B*A

# levicivita_lut[a,b] contains the Levi-Cevita symbol ϵ_abc
# for c=6-a-b, i.e, when a,b,c is a permutation of 1,2,3
const levicivita_lut = [0 1 -1; -1 0 1; 1 -1 0]

function comm(A::σ,B::σ)
    if A.a == B.a
        scal(0)
    else
        a = Int(A.a)
        b = Int(B.a)
        # a+b+c == 6 (since a,b,c is a permutation of 1,2,3)
        c = 6 - a - b
        s = levicivita_lut[a,b]
        δ(A.inds,B.inds)*scal(2im*s)*σ(c,A.inds)
    end
end

replace_index(A::scal,iold,inew) = A
replace_index(A::param,iold,inew) = param(A.name,A.state,(n-> n==iold ? inew : n).(A.inds))
replace_index(A::ExpVal,iold,inew) = ExpVal(replace_index(A.A,iold,inew))
replace_index(A::Corr,iold,inew) = Corr(replace_index(A.A,iold,inew))
for op in (δ,a,adag,f,fdag,σminus,σplus)
    @eval replace_index(A::$op,iold,inew) = $op((n-> n==iold ? inew : n).(A.inds))
end
replace_index(A::σ,iold,inew) = σ(A.a,(n-> n==iold ? inew : n).(A.inds))
replace_index(A::OpProd,iold,inew) = replace_index(A.A,iold,inew)*replace_index(A.B,iold,inew)
replace_index(A::OpSum,iold,inew) = replace_index(A.A,iold,inew) + replace_index(A.B,iold,inew)
replace_index(A::OpSumAnalytic,iold,inew) = begin
    (A.ind==iold || A.ind==inew) && error("replace_index in OpSumAnalytic cannot have iold ($iold) or inew ($inew) be the same as the sum index ($(A.ind))!")
    tmp = replace_index(A.A,iold,inew)
    # we have to be careful here - when replacing an index inside the expression, there might be reorderings that
    # have be done with the same index explicitly to make sure that the commutator shows up
    # so take out the term with the same index explicitly, and redo that one term
    OpSumAnalytic(A.ind,tmp) - replace_index(tmp,A.ind,inew) + replace_index(replace_index(A.A,A.ind,iold),iold,inew)
end
replace_index(A::Operator,reps) = begin
    Anew = A
    for (iold,inew) in reps
        Anew = replace_index(Anew,iold,inew)
    end
    Anew
end

function exchange_inds(A::Operator,i1,i2)
    ss = gensym()
    A = replace_index(A,i1,ss)
    A = replace_index(A,i2,i1)
    A = replace_index(A,ss,i2)
    A
end

distribute_indices!(inds,A::scal) = A
distribute_indices!(inds,A::param) = param(A.name,A.state,(popfirst!(inds) for _ in A.inds)...)
distribute_indices!(inds,A::ExpVal) = ExpVal(distribute_indices!(inds,A.A))
distribute_indices!(inds,A::Corr) = Corr(distribute_indices!(inds,A.A))
for op in (a,adag,f,fdag,σminus,σplus)
    @eval distribute_indices!(inds,A::$op) = $op((popfirst!(inds) for _ in A.inds)...)
end
distribute_indices!(inds,A::σ) = σ(A.a,(popfirst!(inds) for _ in A.inds)...)
distribute_indices!(inds,A::OpProd) = distribute_indices!(inds,A.A)*distribute_indices!(inds,A.B)
distribute_indices!(inds,A::OpSum) = distribute_indices!(inds,A.A) + distribute_indices!(inds,A.B)
# on purpose do not define this for OpSumAnalytic or δ

"sum indices have no semantic meaning, so rename them in case they happen to occur in the other expression"
function ensure_compatible_sumind(S::OpSumAnalytic,A::Operator)
    Ainds = indextuple(A)
    if S.ind in Ainds
        oldinds = Set{OpIndex}((Ainds...,indextuple(S)...))
        m = match(r"(.*)_([0-9]+)",string(S.ind))
        indstem, ii = (m === nothing) ? (string(S.ind), 1) : (m.captures[1], 1+parse(Int,m.captures[2]))
        while (newind = Symbol(indstem,:_,ii)) in oldinds
            ii += 1
        end
        OpSumAnalytic(newind,replace_index(S.A,S.ind,newind))
    else
        S
    end
end

# when multiplying (or commuting) with an operator with an index, take into account that the term in the sum with equal index has to be treated specially
*(A::Union{param,ExpVal,Corr,BaseOperator},B::OpSumAnalytic) = (B = ensure_compatible_sumind(B,A); OpSumAnalytic(B.ind,A*B.A))
# no need to check indices here since we just dispatch to another routine
*(A::OpSumAnalytic,B::OpProd) = (A*B.A)*B.B
*(A::OpSumAnalytic,B::Operator) = (A = ensure_compatible_sumind(A,B); OpSumAnalytic(A.ind,A.A*B))

comm(A::Union{param,ExpVal,Corr,BaseOperator},B::OpSumAnalytic) = (B = ensure_compatible_sumind(B,A); OpSumAnalytic(B.ind,comm(A,B.A)))
comm(A::OpSumAnalytic,B::OpSumAnalytic) = (B = ensure_compatible_sumind(B,A); OpSumAnalytic(B.ind,comm(A,B.A)))
comm(A::OpSumAnalytic,B::Operator) = -comm(B,A)

print(io::IO,A::a) = print(io,"a($(A.inds...))")
print(io::IO,A::adag) = print(io,"a†($(A.inds...))")
print(io::IO,A::f) = print(io,"f($(A.inds...))")
print(io::IO,A::fdag) = print(io,"f†($(A.inds...))")
print(io::IO,A::σ) = print(io,"σ$(A.a)($(A.inds...))")
print(io::IO,A::σminus) = print(io,"σ-($(A.inds...))")
print(io::IO,A::σplus) = print(io,"σ+($(A.inds...))")
print(io::IO,A::scal) = print(io,"$(A.v)")
print(io::IO,A::param) = print(io,string(A.name, A.state=='c' ? "'" : "",length(A.inds)==0 ? "" : "($(A.inds...))"))
print(io::IO,A::OpProd) = print(io,"$(A.A) $(A.B)")
print(io::IO,A::OpSum) = print(io,"$(A.A) + $(A.B)")
print(io::IO,A::ExpVal) = print(io,"⟨$(A.A)⟩")
print(io::IO,A::Corr) = print(io,"⟨$(A.A)⟩c")
print(io::IO,A::OpSumAnalytic) = print(io,"∑_$(A.ind) $(A.A)")

mystring(x::Number) = @sprintf "%g" x
mystring(x::Rational) = denominator(x)==1 ? "$(numerator(x))" : "\\frac{$(numerator(x))}{$(denominator(x))}"
mystring(x::Complex) = @sprintf "(%g%+gi)" real(x) imag(x)
mystring(x::Complex{Rational{T}}) where T = @sprintf "\\left(%s%s%si\\right)" mystring(real(x)) (imag(x)>=0 ? "+" : "-") mystring(abs(imag(x)))

Base.show(io::IO, ::MIME"text/latex", A::Operator) = print(io,"\$",latex(A),"\$")
latex(A::σ) = string("\\sigma_{$(A.a)",length(A.inds)>0 ? ",$(A.inds...)}" : "}")
latexindstr(inds::OpIndices) = length(inds)==0 ? "" : "_{$(inds...)}"
latex(A::δ) = "δ" * latexindstr(A.inds)
latex(A::a) = "a" * latexindstr(A.inds)
latex(A::adag) = "a$(latexindstr(A.inds))^\\dagger"
latex(A::f) = "f" * latexindstr(A.inds)
latex(A::fdag) = "f$(latexindstr(A.inds))^\\dagger"
latex(A::σminus) = "\\sigma^-" * latexindstr(A.inds)
latex(A::σplus) = "\\sigma^+" * latexindstr(A.inds)
latex(A::scal) = imag(A.v)==0 ? mystring(real(A.v)) : (real(A.v)==0 ? mystring(imag(A.v))*"i" : mystring(A.v))
latex(A::param)  = string(A.name, latexindstr(A.inds), A.state=='c' ? "^*" : "")
latex(A::OpProd) = string(latex(A.A)," ",latex(A.B))
latex(A::OpSum) = string(latex(A.A)," + ",latex(A.B))
latex(A::ExpVal) = "\\langle $(latex(A.A)) \\rangle"
latex(A::Corr) = "\\langle $(latex(A.A)) \\rangle_{c}"
latex(A::OpSumAnalytic) = string("\\sum_{$(A.ind)}",latex(A.A))

indextuple(A::scal)::OpIndices = ()
indextuple(A::Union{param,δ,a,adag,f,fdag,σ,σminus,σplus})::OpIndices = A.inds
indextuple(A::Union{OpProd,OpSum})::OpIndices = (indextuple(A.A)...,indextuple(A.B)...)
indextuple(A::Union{ExpVal,Corr})::OpIndices = indextuple(A.A)
indextuple(A::OpSumAnalytic)::OpIndices = (A.ind,indextuple(A.A)...)
indexset(A) = Set{OpIndex}(indextuple(A))
sumindextuple(A::Operator)::OpIndices = ()
sumindextuple(A::Union{OpProd,OpSum})::OpIndices = (sumindextuple(A.A)...,sumindextuple(A.B)...)
sumindextuple(A::Union{ExpVal,Corr})::OpIndices = sumindextuple(A.A)
sumindextuple(A::OpSumAnalytic)::OpIndices = (A.ind,sumindextuple(A.A)...)
sumindexset(A) = Set{OpIndex}(sumindextuple(A))

"`extindices(A::Operator)` return externally visible indices of an expression"
extindices(A::Operator) = [ind for ind in indextuple(A) if !in(ind,sumindextuple(A))]

"`symmetric_index_nums(A::Operator)` return sequence of numbers of exchange-symmetric indices"
function symmetric_index_nums(A::Operator)
    inds = extindices(A)
    Nsyms = [1]
    for ii=2:length(inds)
        if A == exchange_inds(A,inds[ii-1],inds[ii])
            Nsyms[end] += 1
        else
            push!(Nsyms,1)
        end
    end
    Nsyms
end

"""
    ascorr(expr::Operator)

Take an expression `expr=A B C + D E...` and write its expectation value in
terms of single-body expectation values ``⟨A⟩, ⟨B⟩, \\ldots``, and many-body
correlations ``⟨AB⟩_c, ⟨ABC⟩_c``, etc. Currently, up to fourth-order
correlations (i.e., products of four operators) are supported.

E.g., `ascorr(adag(:n)*a(:n))` returns ``⟨a^\\dagger_n a_n⟩_c + ⟨a^\\dagger_n⟩
⟨a_n⟩`` (which is equal to ``⟨a^\\dagger_n a_n⟩``), while
`ascorr(adag(:n)*a(:m)*a(:n))` returns ``\\langle a_{n}^\\dagger a_{m} a_{n}
\\rangle_{c} + \\langle a_{n}^\\dagger \\rangle \\langle a_{m} \\rangle \\langle
a_{n} \\rangle + \\langle a_{n}^\\dagger \\rangle \\langle a_{m} a_{n}
\\rangle_{c} + \\langle a_{m} \\rangle \\langle a_{n}^\\dagger a_{n}
\\rangle_{c} + \\langle a_{n} \\rangle \\langle a_{n}^\\dagger a_{m}
\\rangle_{c}``.

See also: [`ExpVal`](@ref), [`Corr`](@ref)"""
function ascorr end

ascorr(A::Scalar) = A
for op in (adag,a,f,fdag,σ,σminus,σplus)
    @eval ascorr(A::$op) = ExpVal(A)
end
ascorr(A::OpSum) = ascorr(A.A) + ascorr(A.B)
ascorr(A::OpSumAnalytic) = begin
    # first calculate the correlation for the term in the sum with the "bare" indices, which means that the sum index
    # is assumed to be distinct from the indices of the expressions
    # then, calculate the correction expression for the term where the sum index is the same as any of the non-summed indices,
    # and add the correction between that and the "incorrect" term where index identity is not taken into account
    tmp = ascorr(A.A)
    res = OpSumAnalytic(A.ind,tmp)
    for ii in setdiff(indexset(A),sumindexset(A))
        res += ascorr(replace_index(A.A,A.ind,ii)) - replace_index(tmp,A.ind,ii)
    end
    res
end
function ascorr(A::OpProd)::Operator
    A.A isa OpSumAnalytic && error("should not occurr!")
    A.A isa scal && A.B isa OpSumAnalytic && return A.A*ascorr(A.B)
    preftup, exptup, optup = prodtuples(A)
    pref = prod((scal(1),preftup...,exptup...))
    if length(optup)==0
        pref
    elseif length(optup)==1
        pref*ExpVal(optup[1])
    elseif length(optup)==2
        A,B = optup
        pref*(Corr(A*B) + ExpVal(A)*ExpVal(B))
    elseif length(optup)==3
        A,B,C = optup
        pref*(Corr(A*B*C) + ExpVal(A)*Corr(B*C) + ExpVal(B)*Corr(A*C) + ExpVal(C)*Corr(A*B) + ExpVal(A)*ExpVal(B)*ExpVal(C))
    elseif length(optup)==4
        A,B,C,D = optup
        pref*(Corr(A*B*C*D) + ExpVal(A)*ExpVal(B)*ExpVal(C)*ExpVal(D)
            + ExpVal(A)*Corr(B*C*D) + ExpVal(B)*Corr(A*C*D) + ExpVal(C)*Corr(A*B*D) + ExpVal(D)*Corr(A*B*C)
            + ExpVal(A)*ExpVal(B)*Corr(C*D) + Corr(A*B)*ExpVal(C)*ExpVal(D) + Corr(A*B)*Corr(C*D)
            + ExpVal(A)*ExpVal(C)*Corr(B*D) + Corr(A*C)*ExpVal(B)*ExpVal(D) + Corr(A*C)*Corr(B*D)
            + ExpVal(A)*ExpVal(D)*Corr(B*C) + Corr(A*D)*ExpVal(B)*ExpVal(C) + Corr(A*D)*Corr(B*C))
    else
        error("ERROR: Only correlations up to fourth order are implemented for now")
    end
end
CorrOrExp(A::Operator) = length(A)==1 ? ExpVal(A) : Corr(A)

Avac(A::a) = scal(0)
Avac(A::adag) = A
Avac(A::f) = scal(0)
Avac(A::fdag) = A
Avac(A::σminus) = scal(0)
Avac(A::σplus) = A
# vacuum is an eigenstate of σz
Avac(A::σ) = (A.a == z) ? scal(-1) : A
Avac(A::OpSum) = Avac(A.A) + Avac(A.B)
Avac(A::OpSumAnalytic) = OpSumAnalytic(A.ind,Avac(A.A))
Avac(A::OpProd) = begin
    Bv = Avac(A.B)
    Bv != A.B && return Avac(A.A*Bv)
    # if A.A and A.B commute, try what happens with Avac(A.B*A.A)
    if comm(A.A,A.B) == scal(0)
        Av = Avac(A.A)
        Av != A.A && return Avac(A.B*Av)
    end
    # did not find a way to simplify
    A
end
Avac(A::Scalar) = A

vacA(A::a) = A
vacA(A::adag) = scal(0)
vacA(A::f) = A
vacA(A::fdag) = scal(0)
vacA(A::σminus) = A
vacA(A::σplus) = scal(0)
# vacuum is an eigenstate of σz
vacA(A::σ) = (A.a == z) ? scal(-1) : A
vacA(A::OpSum) = vacA(A.A) + vacA(A.B)
vacA(A::OpSumAnalytic) = OpSumAnalytic(A.ind,vacA(A.A))
vacA(A::OpProd) = begin
    if A.A isa Scalar
        A.A*vacA(A.B)
    else
        # for applying ⟨0|AB, see if ⟨0|A changes A, and if so,
        # keep going with the new operator
        vA = vacA(A.A)
        vA==A.A ?  A : vacA(vA * A.B)
    end
end
vacA(A::Scalar) = A

"""
    Avac(A::Operator), vacA(A::Operator)

Simplify operator by assuming it is applied to the vacuum from the left or
right, respectively. To be precise, `Avac(A)` returns ``A'`` such that ``A'|0⟩ =
A|0⟩``, while `vacA(A)` does the same for ``⟨0|A``."""
Avac, vacA

"""
    vacExpVal(A::Operator,S::Operator=scal(1))

Calculate the vacuum expectation value ``⟨0|S^\\dagger A S|0⟩``, i.e., the
expectation value ``⟨ψ|A|ψ⟩`` for the state defined by ``|ψ⟩= S|0⟩```.
"""
function vacExpVal(A::Operator,stateop::Operator=scal(1))
    # simplify down as much as possible by applying vacuum from left and right
    vsAsv = vacA(Avac(stateop' * A * stateop))
    # only operators that should survive here as operators are σs
    _vacExpVal(vsAsv)
end
_vacExpVal(A::Scalar) = A
# the terms that survive until here have at most a single σ of any particle
# so it does not matter if we are inside a product
_vacExpVal(A::σ) = A.a == z ? scal(-1) : scal(0)
_vacExpVal(A::OpSum) = _vacExpVal(A.A) + _vacExpVal(A.B)
_vacExpVal(A::OpSumAnalytic) = OpSumAnalytic(A.ind,_vacExpVal(A.A))
# we know that the operators here commute (all a and a† have disappeared, and at most a single σ remaining for each particle)
_vacExpVal(A::OpProd) = _vacExpVal(A.A) * _vacExpVal(A.B)

function parse_paramstr(s)
    s = split(s,"_")
    par = Symbol(s[1])
    @assert length(s) <= 2
    inds = length(s)==1 ? () : Meta.parse.(split(s[2],","))
    par,inds
end

macro Pc_str(s)
    par, inds = parse_paramstr(s)
    param(par,'n',inds...)
end
macro Pr_str(s)
    par, inds = parse_paramstr(s)
    param(par,'r',inds...)
end

include("precompile.jl")

end # module


@testset "Lines on a Quintic surface in 3-space" begin
    @testset "Lines on a quintic surface in 3-space" begin
        sys, q₀, gamma = fano_quintic()
        @time res =
            solve(sys; gamma = gamma, target_parameters = q₀, start_system = :total_degree)
        @test nsolutions(res) == 2875

        @time poly_res = solve(sys; target_parameters = q₀, start_system = :polyhedral)
        @test nsolutions(poly_res) == 2875
    end

end


function bt02()
#
# Brugnano and Trigiante
#
    p = poly([ones(10);-1;-1;im;-im;2])
    z = [1.0 10; -1 2; im 1; -im 1; 2 1];

    p, PolyZeros(z)
end


mass(U, q, w, v, x, t, μ, η) = [1.0, 1.0, 1.0]
source(U, q, w, v, x, t, μ, η) = [0.0, 0.0, 0.0]
ubou(U, q, w, v, x, t, μ, η, û, n, τ) = [0.0, 0.0, 0.0]

function flux(U, q, w, v, x, t, μ, η)
    h, hu, hv = U

    h⁻¹ = 1/h

    uv = hu * h⁻¹ # u
    vv = hv * h⁻¹ # v

    g = μ[1] # gravitational acceleration

    p = g * h^2 / 2 # g * h^2 / 2

    f = [hu, hu * uv + p, hv * uv, hv, hu * vv, hv * vv + p]

    f = reshape(f, (3, 2))

    return f
end

function fbou(U, q, w, v, X, t, μ, η, Û, n, τ)
    f = flux(U, q, w, v, X, t, μ, η)
    f̂ = f[:, 1] * n[1] + f[:, 2] * n[2] + τ[1] * (U - Û)
    return f̂
end

function initu(X, μ, η)
    ϵ = 0.1 # Perturbation magnitude
    l = 0.5 # Gaussian width
    k = 0.5 # Sinusoidal wavenumber

    x = X[1]
    y = X[2]

    # The Bickley jet
    U = 1 / cosh(y)^2

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

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

    Uᵢ₁ = 1.0        # h
    Uᵢ₂ = U + ϵ * u′ # h * u
    Uᵢ₃ = ϵ * v′     # h * v

    return [Uᵢ₁, Uᵢ₂, Uᵢ₃]
end


function getSubParams(Hparam::RegularMesh, M::RegularMesh,i::Array{Int64},NumCells::Array{Int64},overlap::Array{Int64})
	subMesh   = getSubMeshOfCell(NumCells,overlap,i,Hparam);
	return subMesh;
end

function getLap(M::RegularMesh)
	G = getNodalGradientMatrix(M);
	Ar = G'*G;
	Ar = Ar + 1e-5*opnorm(Ar,1)*sparse(1.0I,size(Ar,2),size(Ar,2));
	return Ar;
end

function getDirichletMassNodalMesh(DDparam::DomainDecompositionParam,problem_param::RegularMesh,i::Array{Int64})
	d = getDirichletMassNodal(DDparam.numDomains,DDparam.overlap,i,DDparam.Mesh.n);
	d.*=(0.1*4.0)/prod(DDparam.Mesh.h);
	return d;
end


"""
Misplaced and miscellaneous functions
"""
module Utils

using SpecialFunctions: gamma
using StatsFuns: gammacdf
using HypergeometricFunctions: _₂F₁

"""
Wrapper around hypergeometric function 2F1
"""
function hyp2f1(a, b, c, z)
    return _₂F₁(a, b, c, z)
end

"""
Wrapper around regularized lower incomplete gamma function from Rmath
"""
function gamma_inc(a, x)
    return gammacdf(a, 1, x)
end

"""
Redefined incomplete beta function in terms of hypergeometric function 2F1.
"""
function B_inc(a, b, z)
    F(z) = hyp2f1(a, 1.0 - b, a + 1.0, z)
    return @. a ^ -1.0  * z ^ a * F(z)
end

"""
Natural logarithm of a gaussian with mean μ and standard deviation σ.
"""
function log_gauss(x, mu, sigma)
    return @. -0.5 * (log(2π * sigma^2) + ((x - mu) / sigma)^2)
end

end


using CompScienceMeshes
using BEAST
using Base.Test

fn = joinpath(dirname(@__FILE__),"assets","rect1.in")

mesh = readmesh(fn)
edgs = skeleton(mesh,1)

charts = [chart(edgs,edg) for edg in cells(edgs)]
ctrs = [cartesian(center(cht)) for cht in charts]

ND = BEAST.nedelec(mesh, edgs)

@test numfunctions(ND) == numcells(edgs)
@test length(ND.fns[1]) == 1
@test length(ND.fns[2]) == 1
@test length(ND.fns[3]) == 2
@test length(ND.fns[4]) == 1
@test length(ND.fns[5]) == 1

# center of the diagonal
ctr = center(charts[3])
face_charts = [chart(mesh,fce) for fce in cells(mesh)]
nbd1 = neighborhood(face_charts[1], carttobary(face_charts[1], ctr))
nbd2 = neighborhood(face_charts[2], carttobary(face_charts[2], ctr))
t = tangents(ctr,1)
ut = t / norm(t)

ndlocal = refspace(ND)
fn = ND.fns[3]
v1 = dot(fn[1].coeff*ndlocal(nbd1)[1][1],ut)
v2 = dot(fn[2].coeff*ndlocal(nbd2)[2][1],ut)
@test v1 ≈ v2 ≈ -1/√2 ≈ -1/volume(charts[3])


using TransformUtils
using Base.Test

print("[TEST] constructors SO3, Quaternion, AngleAxis... ")
q = Quaternion(0)
ss = so3(0)
R = SO3(0)
A = AngleAxis(0)
T = SE3(0)
SE3(T.t, q)
SE3(T.t, R)
SE3(T.t, A)
println("[SUCCESS]")

print("[TEST] comparison functions for SO3, Quaternion, AngleAxis... ")
@test compare(Quaternion(0),Quaternion(0))
@test !compare(Quaternion(0),Quaternion(0,[1;0;0]))
@test compare(SO3(0),SO3(0))
IrotSO3 = SO3(0)
for i in 1:1000
  r = 0.1*randn(3)
  if norm(r) > 1e-10
    @test !compare(IrotSO3, convert(SO3,so3( r )) )
  end
end
@test compare(AngleAxis(0), AngleAxis(0))
println("[SUCCESS]")

print("[TEST] trivial case quaterion -> SO3 -> Euler -> quaternion... ")
q = Quaternion(0)
R = convert(SO3, q)
E = convert(Euler, R)
@test compare(q,convert(Quaternion, E))
# @show SO3(0) * Quaternion(0) * so3(0) * AngleAxis(0)
@test compare(SO3(0), SO3(0) * Quaternion(0) * so3(0) * AngleAxis(0))
println("[SUCCESS]")

println("[TEST] convert functions ")
warn("Need better coverage on convert function tests")
# print("[SUCSSESS]")

print("[TEST] compare SO3 and quaternion rotations... ")
q = Quaternion(0)
R = convert(SO3,q)
A = convert(AngleAxis, q)

for dAA in
    [AngleAxis(pi/4, [1;0;0]);
    so3(0.1*randn(3));
    so3(0.1*randn(3));
    so3(0.1*randn(3));
    convert(Quaternion, Euler(0.1*randn(3)));
    convert(SO3,Euler(0.1*randn(3)))  ]
  @show dAA
  q1 = q*dAA
  R1 = R*dAA
  AA1 = A*dAA
  @test compare( q1, convert(Quaternion, R1) )
  @test compare( R1, convert(SO3, q1) )
  @test compare( q1, convert(Quaternion,AA1) )
  @test compare( R1, convert(SO3,AA1) )
end
println("[SUCCESS]")

print("[TEST] basic SE3 mechanics... ")
# SE3 tests
a = SE3([0;0;0],SO3(0))
b = SE3([1;0;0],SO3(0))
@test compare(a.R,b.R)
c = a*b
@test compare(b,c)

ap = SE3(a.t, a.R*so3(0.1*randn(3)))
@test !compare(ap,a)

back = ap*b*SE3(zeros(3),transpose(ap.R))
@test compare(SO3(0),back.R)
println("[SUCCESS]")


"""
    coefficient_diff(β::AbstractVector, method=:concat)

Returns the coefficient difference statistic W[j] = |β[j]| - |β[j + p]|
from a univariate (single response) regression.

# Inputs
+ `β`: Vector of regression coefficients
+ `method`: Either `:concatenated` (default) if all knockoffs are concatenated
    at the end (e.g. [XX̃]) or `:interleaved` if each variable is immediately
    followed by its knockoff (e.g. [x₁x̃₁x₂x̃₂...])
"""
function coefficient_diff(β::AbstractVector, method::Symbol=:concatenated)
    iseven(length(β)) || error("length of β should be even but was odd.")
    p = length(β) >> 1
    W = Vector{eltype(β)}(undef, p)
    if method == :concatenated
        for j in 1:p
            W[j] = abs(β[j]) - abs(β[j + p])
        end
    elseif method == :interleaved
        for j in 1:p
            W[j] = abs(β[2j - 1]) - abs(β[2j])
        end
    else
        error("method should be :concatenated or :interleaved but got $method")
    end
    return W
end

"""
    coefficient_diff(β::AbstractVector, original::AbstractVector{Int}, knockoff::AbstractVector{Int})

Returns the coefficient difference statistic W[j] = |β[j]| - |β[j + p]|
from a univariate (single response) regression, where the `j`th variable is stored
in position `original[j]` of `β`, and its knockoff is stored in position `knockoff[j]`

# Inputs
+ `β`: Vector of regression coefficients
+ `original`: The index of original variables in `β`
+ `knockoff`: The index of knockoff variables in `β`
"""
function coefficient_diff(β::AbstractVector, original::AbstractVector{Int}, knockoff::AbstractVector{Int})
    return abs.(β[original]) - abs.(β[knockoff])
end

"""
    coefficient_diff(β::AbstractVector, original::AbstractVector{Int}, knockoff::AbstractVector{Int})

Returns the coefficient difference statistic for grouped variables
W[G] = sum_{j in G} |β[j]| - sum_{j in G} |β[j + p]|.

# Inputs
+ `β`: Vector of regression coefficients
+ `groups`: Vector storing group membership. `groups[i]` is the group of `β[i]`
+ `original`: The index of original variables in `β`
+ `knockoff`: The index of knockoff variables in `β`
"""
function coefficient_diff(β::AbstractVector, groups::AbstractVector{Int},
    original::AbstractVector{Int}, knockoff::AbstractVector{Int})
    unique_groups = unique(groups)
    β_groups = zeros(length(unique_groups))
    # find which variables are Knockoffs
    knockoff_idx = falses(length(β))
    knockoff_idx[knockoff] .= true
    # loop over each variable
    for i in 1:length(β)
        g = groups[i]
        if knockoff_idx[i]
            β_groups[g] -= abs(β[i])
        else
            β_groups[g] += abs(β[i])
        end
    end
    return β_groups
end

"""
    coefficient_diff(B::AbstractMatrix)

Returns the coefficient difference statistic
W[j] = |B[j, 1]| - |B[j + p, 1]| + ... +  |B[j, r]| - |B[j + p, r]|
from a multivariate (multiple response) regression.

# Inputs
+ `β`: Matrix of regression coefficients. Each column is a vector of β.
+ `method`: Either `:concatenated` if all knockoffs are concatenated at the end
    (e.g. [XX̃]) or `:interleaved` if each variable is immediately followed by
    its knockoff (e.g. [x₁x̃₁x₂x̃₂...])
"""
function coefficient_diff(B::AbstractMatrix, method::Symbol=:concatenated)
    iseven(size(B, 1)) || error("Number of covariates in B should be even but was odd.")
    p = size(B, 1) >> 1
    r = size(B, 2)
    T = eltype(B)
    W = Vector{T}(undef, p)
    if method == :concatenated
        for j in 1:p
            Wj = zero(T)
            for i in 1:r
                Wj += abs(B[j, i]) - abs(B[j + p, i])
            end
            W[j] = Wj
        end
    elseif method == :interleaved
        for j in 1:p
            Wj = zero(T)
            for i in 1:r
                Wj += abs(B[2j - 1, i]) - abs(B[2j, i])
            end
            W[j] = Wj
        end
    else
        error("method should be :concatenated or :interleaved but got $method")
    end
    return W
end

# """
#     coefficient_diff(B::AbstractMatrix, original::AbstractVector{Int}, knockoff::AbstractVector{Int})

# Returns the coefficient difference statistic W[j] = |β[j]| - |β[j + p]|
# from a multivariate (multiple response) regression, where the `j`th variable is stored
# in position `original[j]` of `β`, and its knockoff is stored in position `knockoff[j]`

# # Inputs
# + `β`: Vector of regression coefficients
# + `original`: The index of original variables in `β`
# + `knockoff`: The index of knockoff variables in `β`
# """
# function coefficient_diff(B::AbstractMatrix, original::AbstractVector{Int}, knockoff::AbstractVector{Int})
#     p = size(B, 1) >> 1
#     length(original) == length(knockoff) == p || error("Number of variables in " *
#         "B should be twice the length of original and knockoff.")
#     W = Vector{eltype(B)}(undef, p)
#     for j in 1:p
#         Wj = zero(T)
#         for i in 1:r
#             Wj += abs(B[original[j], i]) - abs(B[knockoff[j], i])
#         end
#         W[j] = Wj
#     end
#     return W
# end

"""
    extract_beta(β̂_knockoff::AbstractVector, fdr::Number, method::Symbol=:concatenated)

Given estimated β of original variables and their knockoffs, compute β for the
original design matrix that controls the FDR.
"""
function extract_beta(β̂_knockoff::AbstractVector{T}, fdr::Number,
    method::Symbol=:concatenated
    ) where T <: AbstractFloat
    # first handle errors
    0 ≤ fdr ≤ 1 || error("Target FDR should be between 0 and 1 but got $fdr")
    p = length(β̂_knockoff) >> 1
    # length(original) == length(knockoff) == p || error("Length of " *
    #     "β should be twice of original and knockoff.")
    # find set of selected predictors
    W = coefficient_diff(β̂_knockoff, method)
    τ = threshold(W, fdr)
    detected = findall(W .≥ τ)
    # construct original β
    β = zeros(T, p)
    if method == :concatenated
        for i in detected
            β[i] = β̂_knockoff[i]
        end
    elseif method == :interleaved
        for i in detected
            β[i] = β̂_knockoff[2i - 1]
        end
    else
        error("method should be :concatenated or :interleaved but got $method")
    end
    return β
end

function extract_beta(β̂_knockoff::AbstractVector{T}, fdr::Number,
    original::AbstractVector{Int}, knockoff::AbstractVector{Int},
    method::Symbol=:knockoff, combine::Bool=false) where T <: AbstractFloat
    # first handle errors
    p = length(β̂_knockoff) >> 1
    0 ≤ fdr ≤ 1 || error("Target FDR should be between 0 and 1 but got $fdr")
    # length(original) == length(knockoff) == p || error("Length of " *
    #     "β should be twice of original and knockoff.")
    # find set of selected predictors
    W = coefficient_diff(β̂_knockoff, original, knockoff)
    τ = threshold(W, fdr, method)
    detected = findall(W .≥ τ)
    # construct final β
    β = zeros(T, p)
    if combine
        for i in detected
            β[i] = β̂_knockoff[original[i]] + β̂_knockoff[knockoff[i]]
        end
    else
        for i in detected
            β[i] = β̂_knockoff[original[i]]
        end
    end
    return β
end

function extract_beta(β̂_knockoff::AbstractVector{T}, fdr::Number, groups::Vector{Int},
    original::AbstractVector{Int}, knockoff::AbstractVector{Int}, method=:knockoff
    ) where T <: AbstractFloat
    # first handle errors
    p = length(β̂_knockoff) >> 1
    0 ≤ fdr ≤ 1 || error("Target FDR should be between 0 and 1 but got $fdr")
    # length(β̂_knockoff) == length(groups) ||
    #     error("β̂_knockoff should have same length as groups")
    # length(original) == length(knockoff) == p ||
    #     error("Length of β should be twice of original and knockoff.")
    # find set of selected predictors
    W = coefficient_diff(β̂_knockoff, groups, original, knockoff)
    τ = threshold(W, fdr, method)
    detected_groups = findall(W .≥ τ)
    # construct the full β
    β = zeros(T, length(β̂_knockoff))
    for g in detected_groups
        group_idx = findall(x -> x == g, groups)
        β[group_idx] .= @view(β̂_knockoff[group_idx])
    end
    return β[original]
end


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

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

# ╔═╡ b1a8e446-76c6-11eb-2474-874d4bc0a9cc
using Plots

# ╔═╡ e338d842-7719-11eb-35c6-efdf74875751
using PlutoUI

# ╔═╡ e80655b8-76c5-11eb-1438-85fd5054780b
next_logistic(x, growth_rate=2.6) = growth_rate .* x .* (1 .- x)

# ╔═╡ 54a5bd48-76c6-11eb-18fa-65a57e2be14f
begin
  x = 0:.001:1
  y = next_logistic.(x);
end

# ╔═╡ cf72f62e-76c6-11eb-3389-d343b8217f7b
plot(x, y, dpi=200, legend=false)

# ╔═╡ ddb1a95a-7730-11eb-0c51-19fc37d551eb
md"""
#### Negative feedback loop
In the sense, once an iteration get a high $x_n$, the next iteration $x_{n+1}$ will be mapped to a small value.

"""

# ╔═╡ 41ab4ace-76c8-11eb-1779-eb7d95ed3842
#@bind n Slider(1:100, default=1, show_value=true)

# ╔═╡ d3ae7692-7724-11eb-30d6-091451792d46
md"""
`n` $(@bind n Slider(1:1000, default=1, show_value=true))
$(html"
")
`x1` $(@bind x1 Slider(0:0.01:1, default=.4, show_value=true))
$(html"
")
`r` $(@bind r Slider(0:0.1:999_999_999, default=2.6, show_value=true))
"""

# ╔═╡ bfd9a498-76c7-11eb-07c7-ad7d1144d4d7
begin
  xn = x1
  for i in 1:n
    xn = next_logistic(xn, r)
  end
  xn
end

# ╔═╡ aebdfe54-7731-11eb-3ad4-e1a139ada52f
begin
  length = 50
  xs = [0. for _ in 1:length]
  xs[1] = x1
  for i in 2:length
	xs[i] = next_logistic(xs[i-1], r)
  end
  xs
end

# ╔═╡ 0e9528c8-7734-11eb-0e4a-bf65284ad4aa
md"""
**tweak** the intial population `x1` to observe that the convergent value changes
- not with `x1`
- but with `r`
"""

# ╔═╡ 79ddeee2-7733-11eb-2f04-0132bae963ff
begin
  good_dpi = 300
  plot(xs, dpi=good_dpi)
  scatter!(xs, dpi=good_dpi, ylims=(0,1),
		title="r = $r,  x1 = $(x1)\nxs[$(length)] = $(xs[end])", legend=false)
end

# ╔═╡ Cell order:
# ╠═e80655b8-76c5-11eb-1438-85fd5054780b
# ╠═b1a8e446-76c6-11eb-2474-874d4bc0a9cc
# ╠═54a5bd48-76c6-11eb-18fa-65a57e2be14f
# ╠═cf72f62e-76c6-11eb-3389-d343b8217f7b
# ╟─ddb1a95a-7730-11eb-0c51-19fc37d551eb
# ╠═e338d842-7719-11eb-35c6-efdf74875751
# ╠═41ab4ace-76c8-11eb-1779-eb7d95ed3842
# ╟─d3ae7692-7724-11eb-30d6-091451792d46
# ╠═bfd9a498-76c7-11eb-07c7-ad7d1144d4d7
# ╠═aebdfe54-7731-11eb-3ad4-e1a139ada52f
# ╟─0e9528c8-7734-11eb-0e4a-bf65284ad4aa
# ╟─79ddeee2-7733-11eb-2f04-0132bae963ff


# using Pkg
# Pkg.add("NLopt")
# Pkg.add("RecipesBase")

using LikelihoodProfiler, Test

include("./cases_func.jl")
methods_list = [:CICO_ONE_PASS, :LIN_EXTRAPOL, :QUADR_EXTRAPOL]

println("Starting tests for get_interval")
@testset "get_interval" begin include("test_get_interval.jl") end

println("Starting tests for get_endpoint")
@testset "get_endpoint CICO_ONE_PASS parameters" begin include("test_get_endpoint_parameters.jl") end
@testset "get_endpoint CICO_ONE_PASS function" begin include("test_get_endpoint_scan_func.jl") end
@testset "get_endpoint errors and warnings" begin include("test_get_endpoint_errors.jl") end

println("Starting tests for all all methods of get_right_endpoint")
@testset "get_right_endpoint methods" begin include("test_methods.jl") end
@testset "get_right_endpoint CICO_ONE_PASS" begin include("test_cico_one_pass.jl") end

println("Starting tests for profile")
@testset "profile" begin include("test_profile.jl") end

println("Starting tests for different fitting alg")
@testset "local_alg" begin include("test_algorithms.jl") end

println("Starting tests for Plot @recipe")
@testset "PLOT_INTERVAL" begin include("test_plots.jl") end

println("Starting tests for loss error")
@testset "LOSS_ERROR" begin include("test_loss_error.jl") end

println("Starting tests for bands")
@testset "Bands" begin include("test_bands.jl") end

println("Starting tests for get_optimal")
@testset "get_optimal(...)" begin include("test_get_optimal.jl") end

# experimental tests

@testset "testing derivative-free algorithms" begin include("test_deriv_free_algs.jl") end
@testset "gradient-based algorithms" begin include("test_grad_algs.jl") end


using Base.Test

include("triangle.jl")

@testset "check equilateral triangles" begin
    @testset "true if all sides are equal" begin
        @test is_equilateral([2, 2, 2]) == true
        @test is_equilateral([0.5, 0.5, 0.5]) == true
    end

    @testset "false if any side is unequal" begin
        @test is_equilateral([2, 3, 2]) == false
    end

    @testset "false if no sides are equal" begin
        @test is_equilateral([5, 4, 6]) == false
    end

    @testset "false if invalid triangle" begin
        @test is_equilateral([0, 0, 0]) == false
    end
end

@testset "check isosceles triangles" begin
    @testset "true if at least 2 sides are equal" begin
        @test is_isosceles([3, 4, 4]) == true
        @test is_isosceles([4, 3, 4]) == true
        @test is_isosceles([4, 4, 3]) == true
        @test is_isosceles([4, 4, 4]) == true
        @test is_isosceles([0.4, 0.5, 0.5]) == true
        @test is_isosceles([0.5, 0.4, 0.5]) == true
        @test is_isosceles([0.5, 0.5, 0.4]) == true
    end

    @testset "false if no sides are equal" begin
        @test is_isosceles([2, 3, 4]) == false
    end

    @testset "false if invalid triangle" begin
        @test is_isosceles([1, 1, 3]) == false
    end
end

@testset "check scalene triangles" begin
    @testset "true if no sides are equal" begin
        @test is_scalene([5, 4, 6]) == true
        @test is_scalene([0.5, 0.4, 0.6]) == true
    end

    @testset "false if at least 2 sides are equal" begin
        @test is_scalene([3, 4, 4]) == false
        @test is_scalene([4, 3, 4]) == false
        @test is_scalene([4, 4, 3]) == false
        @test is_scalene([4, 4, 4]) == false
        @test is_scalene([0.4, 0.5, 0.5]) == false
        @test is_scalene([0.5, 0.4, 0.5]) == false
        @test is_scalene([0.5, 0.5, 0.4]) == false
    end

    @testset "false if invalid triangle" begin
        @test is_scalene([7, 3, 2]) == false
    end
end


using Test
using OMEinsum
using OMEinsum: subindex

@testset "EinCode" begin
    code = EinCode(((1,2), (2,3)), (1,3))
    @test code isa EinCode
    @test OMEinsum.getixs(code) == ((1,2), (2,3))
    @test OMEinsum.getiy(code) == (1,3)
end

@testset "indexer" begin
    si = EinIndexer{()}(())
    @test subindex(si, (1,2,3)) == 1
    si = EinIndexer{(3,2)}((7,6))
    @test OMEinsum.getlocs(si) == (3,2)
    a = randn(7,6)
    @test a[subindex(si, (4,5,2))] == a[2,5]
end

@testset "EinArray" begin
    locs_xs = (EinIndexer{(1,2)}((8,8)), EinIndexer{(2,3)}((8,8)))
    ixs = ((1,2), (2,3))
    iy = (1,3)
    x1 = randn(8, 8)
    x2 = randn(8, 8)
    arr = einarray(EinCode(ixs, iy), (x1, x2), OMEinsum.get_size_dict(ixs, (x1, x2)))
    @test size(arr) == (8,8,8)
    # inner first, then outer
    @test arr[CartesianIndex((3,4,2))] == x1[4,3]*x2[3,2]
    @test arr[3,4,2] == x1[4,3]*x2[3,2]
end


module TestSecondDerivative
    using Base.Test
    import FiniteDiff

    negative_sin(x) = -sin(x)
    negative_cos(x) = -cos(x)
    square(x) = x * x
    two(x) = 2.0

    function check_error(y_true, y_approximate, err = 1 // 10_000)
        # TODO: Decide how to handle very small inputs.
        if y_true > eps(y_true)
            @test abs(y_true - y_approximate) / abs(y_true) < err
        end
    end

    function run_tests(n::Integer)
        output = Array(Float64, 1)

        funcs = (
            (sin, negative_sin),
            (cos, negative_cos),
            (square, two),
            (exp, exp),
        )

        for (f, f′′) in funcs
            for itr in 1:n
                x = 1.0 + itr

                y = FiniteDiff.second_derivative(f, x)
                check_error(f′′(x), y)

                FiniteDiff.second_derivative!(output, f, x)
                check_error(f′′(x), output[1])

                tmp = FiniteDiff.second_derivative(f, mutates=false)
                y = tmp(x)
                check_error(f′′(x), y)

                tmp! = FiniteDiff.second_derivative(f, mutates=true)
                tmp!(output, x)
                check_error(f′′(x), output[1])
            end
        end
    end

    @testset "second_derivative tests" begin
        run_tests(100)
    end
end


using Winston

x = linspace(0, 3pi, 100)
c = cos(x)
s = sin(x)

p = FramedPlot(
        title="title!",
        xlabel="\\Sigma x^2_i",
        ylabel="\\Theta_i")

add(p, FillBetween(x, c, x, s))
add(p, Curve(x, c, color="red"))
add(p, Curve(x, s, color="blue"))

file(p, "example1.png")


using Plots; gr()

function draw_neuron(x, y, r; c=:blue)

    θs = 0:0.1:2pi
    xs = x .+ r.*cos.(θs)
    ys = y .+ r.*sin.(θs)

    plot!(xs, ys, seriestype=:shape, c=c, alpha=0.5, aspect_ratio=1, leg=false)

end

#neuron_coords(x, N, spacing) = range(-(N - 1)/2 * spacing, spacing, N)

"""
Vertical position of neuron in layer i, position j, with a total of N neurons
"""
neuron_coords(j, N, y_spacing) = (-(N - 1)/2 + j) * y_spacing

function draw_neurons(x, N, spacing, r; c=:blue)

    ys = neuron_coords(x, N, spacing)

    draw_neuron.(x, ys, r; c=c)

end

function draw_layer(x, spacing, N1, N2, r)

    plot!(framestyle=:none, grid=:none)

    first_x = x
    second_x = x + 1

    first = neuron_coords(x + 1, N1, spacing)
    second = neuron_coords(x, N2, spacing)

    draw_neurons(x, N1, 1, r; c=:blue)
    draw_neurons(x+1, N2, 1, r; c=:red)

    for i in 1:N1
        for j in 1:N2

            vec = [second_x - first_x, second[j] - first[i]]
            normalize!(vec)

            start = [first_x, first[i]] + 1.2*r*vec
            finish = [second_x, second[j]] - 1.2*r*vec

            plot!([start[1], finish[1]], [start[2], finish[2]], c=:black, alpha=0.5)
        end
    end

end

#draw_layer(1, 1, 3, 4, 0.2)

function draw_link(x1, y1, x2, y2, r)
    vec = [x2 - x1, y2 - y1]
    normalize!(vec)

    start = [x1, y1] + 1.2 * r * vec
    finish = [x2, y2] - 1.2 * r * vec

    plot!([start[1], finish[1]], [start[2], finish[2]], c=:black, alpha=0.5)
end

"""
Takes a vector of neurons per layer
"""
function draw_network(neurons_per_layer)

    x_spacing = 1
    y_spacing = 1
    r = 0.2

    num_layers = length(neurons_per_layer)

    plot(framestyle=:none, grid=:none)

    # draw input links
    N1 = neurons_per_layer[1]

    for j in 1:N1
        y = neuron_coords(j, N1, y_spacing)
        draw_link(0, y, 1, y, r)
    end

    # draw neurons
    for layer in 1:length(neurons_per_layer)
        N = neurons_per_layer[layer]

        if layer == 1
            c = :green
        elseif layer == num_layers
            c = :red
        else
            c = :blue
        end

        for j in 1:N

            draw_neuron(layer, neuron_coords(j, N, y_spacing), r, c=c)

        end
    end

    # draw links
    for layer in 1:length(neurons_per_layer)-1
        N1 = neurons_per_layer[layer]
        N2 = neurons_per_layer[layer+1]

        for j1 in 1:N1
            for j2 in 1:N2

                draw_link(layer,  neuron_coords(j1, N1, y_spacing), layer+1,
                neuron_coords(j2, N2, y_spacing), r)
            end

        end
    end

    # draw output links
    N_last = neurons_per_layer[end]

    for j in 1:N_last
        y = neuron_coords(j, N_last, y_spacing)
        draw_link(num_layers, y, num_layers+1, y, r)
    end

    plot!()

end

draw_network([3, 2, 2])