### A Pluto.jl notebook ###
# v0.14.0

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

# ╔═╡ 54e6bd88-a6a9-11eb-3380-49b79430f6cf
html"""

Section 3.2

 ODEs and parameterized types

Lecture Video

"""

# ╔═╡ 5997bb40-b508-4647-8b12-a7cc5152e550
begin
    import Pkg
    Pkg.activate(mktempdir())

    Pkg.add([
        Pkg.PackageSpec(name="DifferentialEquations", version="6"),
        Pkg.PackageSpec(name="Plots", version="1"),
		Pkg.PackageSpec(name="PlutoUI", version="0.7"),
        Pkg.PackageSpec(name="MappedArrays", version="0.4"),

    ])

    using DifferentialEquations, Plots, PlutoUI, MappedArrays
end

# ╔═╡ 38c74e63-13e9-49a9-8eae-4a898197647a
TableOfContents(aside=true)

# ╔═╡ 7a1d476b-1302-4b98-b9fa-01b990979985
md"""
# Julia concepts
- Using libraries
- Libraries are built around types and objects
- Plot recipes: Defining how to plot our objects
- Making objects into callable functions
- Parameterized types
"""

# ╔═╡ 3004336a-0250-483f-be00-60db592c9fc8
md"""
# ODEs and parameterized types
"""

# ╔═╡ a1cfd00b-825f-494e-91a9-192815f06ad6
md"""
In the previous lecture we started to look at Ordinary Differential Equations (ODEs). These are equations relating a function to its derivatives, telling us how to move at each moment in time.

For example, exponential decay is modelled by

$$\dot{u} = -p \, u,$$

in other words

$$\dot{u}(t) = - p \, u(t)$$

for all real $t \ge 0$,

with initial condition

$$u(0) = u_0$$

where $u(t)$ is the number of light bulbs still functioning at time $t$, the number of radioactive nuclei that have not yet decayed, or the number of people who are still infectious and have not yet recovered. The parameter $p$ is the decay rate.

We saw that **time stepping** methods provide a class of numerical methods for solving such initial-value problems, of which the Euler method is the simplest example.

Unfortunately the Euler method is *not* a good choice, since it is not accurate enough. In numerical analysis classes (like [18.330](https://github.com/mitmath/18330)) various better methods are derived.

Fortunately those methods are available in a Julia **library** or package, `DifferentialEquations.jl`.

"""

# ╔═╡ b2d2b8f0-c81d-437a-83cf-726b5caa77ab
md"""
# The `DifferentialEquations.jl` library
"""

# ╔═╡ d27ec502-ea0c-4d15-a223-7cbf6b784d1c
md"""
[DifferentialEquations.jl](https://diffeq.sciml.ai/v2.0/) (often abbreviated as DiffEq) is a Julia library for solving many types of differential equations. It has become the best (most complete and fastest) general library for **solving** ordinary differential equations (and many other types) in any programming environment. It is also a very good example of the way that Julia libraries are often structured.

We will see some basic usage examples and how Julia **types** are used both to control the library and to store the input and output to the **solvers**.
"""

# ╔═╡ 6f87059a-dbdc-46dd-933b-7bb36980868e
md"""
## Basic usage of DifferentialEquations.jl
"""

# ╔═╡ 557cedbc-6e44-40a8-b0ab-37ae11b63f55
md"""
Let's try to solve the above equation that we derived in the previous lecture, $\dot{u} = -p \, u$. (We are using $u$ as the name of the dependent variable for consistency with the `DifferentialEquations.jl` documentation and internal structure.)

We write a general ODE in 1 variable as

$\dot{u} = f_p(t, u).$

The right-hand side of the ODE is given by a function $f$, which may depend on a **parameter** $p$, i.e. a constant which does not change for any given run of the ODE.
$f$ may also depend explicitly on time.

In our example we do indeed have a parameter $p$ that we want to change, but $f$ does not have an explicit time dependence, so we write
"""

# ╔═╡ a6ae905b-db28-4e75-b567-54840dcda9a5
f(u, p, t) = -p * u

# ╔═╡ f16644d0-e2d2-44ac-988c-aa5b178844f5
md"""
Note that you always need the three input arguments `u`, `p` and `t`, in that order, even if you don't use them.
"""

# ╔═╡ 1ea2acfe-6ba1-44d3-9a24-a904a5fac0c4
md"""
Now we define our initial condition $u_0$ and the time span to integrate over, given as a tuple (initial\_time, final\_time):
"""

# ╔═╡ 264a27cd-9ace-48cc-9556-7eb798298f17
begin
	u0 = 100.0

	time_span = (0.0, 10.0)
end

# ╔═╡ 70ebe76d-6107-492d-a3d2-84c1dd7f59fb
md"""
To set up the **problem** instance we use a type `ODEProblem` that is defined in the `DifferentialEquations.jl` package, into which we pass all the information necessary to define the problem. The parameters must go in the following order:
"""

# ╔═╡ 6e7c8e9e-bef0-4bf8-b1b1-d50c82aa203e
problem = ODEProblem(f, u0, time_span, p)

# ╔═╡ b5379dc2-d97f-47ed-8737-35e3fe59285c
md"""
(For more advanced use there are also some additional, optional, keyword arguments.)
Note that the displayed output does not currently include information about the function, nor the parameters.
"""

# ╔═╡ 0ed8e007-2137-43e0-93b8-6f5dfb9496f5
md"""
To solve the ODE we call the `solve` function:
"""

# ╔═╡ ff91515a-89a4-4423-a186-5572c712493d
solution = solve(problem)

# ╔═╡ 64b26404-1ca0-4aa4-bea0-3c9075b08298
md"""
What happened here? A suitable solver (i.e. an algorithm to calculate the solution) was chosen *automatically*, and it chose certain moments in time at which to output information about the (approximate, but very accurate) solution.

In this particular case it chose to output data at only eight points in time between $t=0$ and $t=10$.

Let's try to plot the `solution` object:
"""

# ╔═╡ b07de2ce-640b-42d6-8b60-39fcbcd116e7
plot(solution, size=(500, 300), label="solution")

# ╔═╡ 0b1855de-a8c9-492e-a249-3238e41fe84c
md"""
### Plot recipes
"""

# ╔═╡ 0190bdb0-97bf-4c35-b3c8-e97c1940b702
md"""
Two surprising things happen. Firstly, there is no reason to expect this to have worked: `solution` is some kind of Julia object, but somehow `Plots.jl` knows how to plot it. This is because `DifferentialEquations.jl` defines a **plot recipe**. This specifies a way to turn solution objects into plots! Any package can do this relatively easily.
"""

# ╔═╡ 555a6328-8aff-4f9f-87f5-51a67450002d
md"""
The second surprise is that the output looks like a smooth curve, rather than just 8 points. Let's see those points on top of the curve. We can extract the relevant data from the `solution` object:
"""

# ╔═╡ d65372df-e6ca-4165-8d9b-b21cc5c9f796
scatter!(solution.t, solution.u, label="discrete output")

# ╔═╡ b274c61b-4fab-4f4a-99df-5d39e0f56aa1
md"""
We see that the package in fact gives not only the value at those points, but it is in fact also capable of calculating an (approximate) solution at *any* intermediate point, using **interpolation**. In fact, we can access this by treating `solution` as if it were a function:
"""

# ╔═╡ e3d9e0be-6cf4-4ae1-8c7f-b444a153a9f5
begin
	tt = 3.5
	solution(tt)
end

# ╔═╡ 180bcdc1-5f51-4b32-8b9c-5000605cdf32
scatter!([tt], [solution(tt)], label="t = $(tt)", ms=5, m=:square)

# ╔═╡ 13b38f88-2ead-46b3-bc96-eae2ea10204d
md"""
For this particular ODE we know the analytical solution. Let's compare them as we vary the parameter $p$:
"""

# ╔═╡ 475ef0e5-7f3a-4bbb-b513-8be29abfd3f9
md"""
p = $(@bind p Slider(0.0:0.1:2.0, show_value=true))
"""

# ╔═╡ 0ed6d65c-707d-4666-b203-2ad0ea822687
let

	problem = ODEProblem(f, u0, time_span, p)
	solution = solve(problem)

	plot(solution,
		linewidth=3, xlabel="t", yaxis="x(t)", label="numerical", size=(500, 300))

	plot!(t -> u0 * exp(-p*t), lw=3, ls=:dash, label="exact")

	ylims!(0, 100)
	title!("p = $p")
end

# ╔═╡ cf9a293c-aa0f-40ce-960b-d4c3b82b8346
md"""
We see that the numerical and exact solutions are (to the eye) indistinguishable, and that the package is fast enough to calculate the solution basically in real time.
"""

# ╔═╡ efcfdc1a-5a80-4780-ba56-43c5e2d6ea36
md"""
## Systems of ODEs
"""

# ╔═╡ 5d7a81d7-8080-4390-924a-72c34c0a5e23
md"""
Now let's try to solve the SIR equations:

$$\begin{aligned}
\dot{s} &= -\beta s i \\
\dot{i} &= +\beta s i - \gamma i\\
\dot{r} &= + \gamma i
\end{aligned}$$
"""

# ╔═╡ 6c607126-dd0b-438d-8724-c1ef5e79c71d
md"""
We need to think of the system in a vector form,

$$\dot{\mathbf{x}} = \mathbf{f}_\mathbf{p}(t, \mathbf{x}),$$

where:
- the vector $\mathbf{x}$ is a vector of all the variables, $\mathbf{x} := (x_1, x_2, x_3) := (s, i, r)$
- the function $\mathbf{f}$ is a vector-valued function $\mathbb{R}^3 \to \mathbb{R}^3$, where $f_k$ gives the right-hand side of the equation for $\dot{x_k}$; and
- the vector $\mathbf{p} = (\beta, \gamma)$ is a vector of the parameters.

"""

# ╔═╡ d16e06f4-55c4-4c0e-bde5-c0ce4b36b043
x0 = [0.99, 0.01, 0.0]

# ╔═╡ aa0e4260-6268-4eff-a662-f3fd55e7d450
function SIR(x, p, t)

	s, i, r = x    # unpack the vectors into scalar values
	β, γ = p

	# build a new vector to return:

	return [-β * s * i,
			+β * s * i - γ * i,
					   + γ * i]

end

# ╔═╡ d0f40681-73df-4cd3-bbd5-3edb8193153e
params = [β, γ]

# ╔═╡ fd122a25-5655-410b-aa7f-34936fc97b53
SIR_problem = ODEProblem(SIR, x0, (0.0, 50.0), params)

# ╔═╡ 8ac7da6b-46e0-470b-91b3-1a61c226fa4a
sol = solve(SIR_problem)

# ╔═╡ 7aa45efe-865d-4e0e-9c71-0c032c72c40d
md"""
Now we see that the solverr has recognised that everything is a vector, and it returns a vector at each time stamp.

Again we can plot:
"""

# ╔═╡ bca8112f-cb61-43dc-ae87-383915c8a89b
gr()

# ╔═╡ a766a141-5d7b-499f-9e18-48bf926ee7ea
plot(sol)

# ╔═╡ 76cbc37d-54c8-4626-8bfe-58b63a602c38
md"""
β = $(@bind β Slider(-0.5:0.01:2.0, default=1.0, show_value=true))

γ = $(@bind γ Slider(-0.5:0.01:2.0, default=0.1, show_value=true))
"""

# ╔═╡ c2765282-bcc7-4110-9822-10557326461e
md"""
It knows to plot each variable separately.

We can instead plot combinations of variables in *phase space* or *state space*:
"""

# ╔═╡ 7671506d-d4b7-4792-9571-e003097235e1
gr()

# ╔═╡ 203e90b0-4d0c-4999-8513-d4eb43f53aac
plot(sol, vars=(1, 2), xlabel="s", ylabel="i", arrow=true, xlims=(-0.1, 1.0), size=(500, 300))

# ╔═╡ 326890ae-0675-4779-b5e8-9b3e0412a52b
md"""
And even in 3D:
"""

# ╔═╡ d7025d1f-c3b8-461a-94dc-2414fbdfd373
plotly()

# ╔═╡ 7f5c41f9-96b9-40d6-87f8-3e81c372b48e
plot(sol, vars=(1, 2, 3), xlabel="s", ylabel="i", zlabel="r")

# ╔═╡ d8f15e9b-aef1-4795-8522-85ea3564e351
md"""
The [ModelingToolkit.jl](https://mtk.sciml.ai/stable/tutorials/ode_modeling/) library  provides ways to make creating ODE models more intuitive, using symbolic equation objects.
"""

# ╔═╡ e723400e-3f8f-47d6-8944-28355a54c698
md"""
## Making objects into callable functions
"""

# ╔═╡ 15fdb430-4157-4401-85b9-665dde795a22
md"""
One of the interesting things we have seen is the way that the `solution` object can be **called** using the normal syntax for calling functions:
"""

# ╔═╡ a410e98d-dd7f-4f55-8812-8b0339a56c86
solution(3.5)

# ╔═╡ 239f39b2-53c6-4da9-ac25-b5e7d19deb5f
md"""
Not all types allow this. For example, if we define a matrix
"""

# ╔═╡ 6630cbc8-1b2f-4ea9-bcd6-e4f4fcb2e844
A = rand(2, 2)

# ╔═╡ 5c1c7016-69aa-4009-8dc8-0c68c925ee66
md"""
we cannot usually call it with a vector:
"""

# ╔═╡ 83719ef1-e8d6-44d0-98ab-a851b1082fa5
A(rand(2))

# ╔═╡ e4d36db5-352b-4fb2-9169-da5c6a3458b7
md"""
If we wanted to do so, we could, however, using the following syntax:
"""

# ╔═╡ 167015ff-1caf-42a8-85f1-070da3e5d5c5
# (A::Matrix)(x) = A * x

# ╔═╡ f8bcd82d-089c-49a2-99d0-b17a5ac5c509
md"""
Note that this makes objects (instances) of the type callable, rather than the type itself, which we can already call as a **constructor** to create objects of that type:
"""

# ╔═╡ cd17e738-089b-4934-af63-d27706c5d1c7
Matrix([1 2; 3 4])

# ╔═╡ fc525f77-ad62-43cb-92e4-0fc1354be99c
md"""
For example, if we wanted to write our own type to store the output of the Euler method, we could do something like
"""

# ╔═╡ ec90842a-a20c-4847-a803-e070d343b98d
begin
	struct SimpleEulerOutput
		times::Vector{Float64}
		values::Vector{Float64}
	end

	function (soln::SimpleEulerOutput)(t)
		# use interpolation to calculate the value at time t
		# using  `soln.times` and `soln.values`
	end
end

# ╔═╡ 9fad6f01-a3b7-4dbb-a8f5-29f2ea1f545a
md"""
The definition of the function makes it callable.

This says that "the way to call an object called `soln`, of type `EulerOutput`, is the following:". Note that inside the function you will need too access the `times` and `values` from *inside* the `soln` object, using `soln.times` and `soln.values`.

Note that at the time of writing, Pluto requires that you put this call definition in the same cell as the type definition.
"""

# ╔═╡ 658d8d15-5b3e-4d35-88b1-4e2fa0f3b17a
md"""
## Digging into objects
"""

# ╔═╡ aafeeb55-943f-4375-9cc0-e0dce659d161
md"""
A key feature of the `DifferentialEquations.jl` library, and of most other libraries in Julia, is the way they define and use new **types**.

As we have seen previously, we can define new types using `struct` which store data, and then define functions which act on those types.
"""

# ╔═╡ 3c45ad21-d18b-41fa-a0fd-64d8988621f3
md"""
Let's start by looking at the `ODEProblem` type:
"""

# ╔═╡ 96e871a8-48b9-459b-a1a0-94fc1deb941e
problem

# ╔═╡ 2622801c-4e37-42f1-ab77-939d461b8d20
md"""
To see what fields, or attributes, the object contains we can use
"""

# ╔═╡ 43e7b270-6452-4a4c-ba6f-c286a337a423
fieldnames(typeof(problem))

# ╔═╡ dbdb9998-ca3b-46fb-87cf-5b7c6e387ab9
md"""
In Pluto and other interactive environments we can get the same information using `problem.<TAB>`, i.e. by typing the TAB key after `problem.`
"""

# ╔═╡ b8fbff51-72fe-4432-9f6f-06f4628320d4
md"""
As usual we can extract the data contained in the object as usual using `.`:
"""

# ╔═╡ ef677abe-c96d-41cc-9b4a-192ff29cdb86
problem.u0

# ╔═╡ 78933b33-8a3c-492e-8a58-19ad7187b45e
problem.tspan

# ╔═╡ 9ef4aa78-2ecd-4d53-87ae-f0909bb46915
md"""
To see everything contained in the object, we can use `Dump` in Pluto, or `dump` if we are not using Pluto:
"""

# ╔═╡ 798bc482-7556-4904-8946-0d9898ce2d33
Dump(problem)

# ╔═╡ 32e9f6e5-8df6-4d40-8092-efaf2bce28d3
md"""
This seems (much) more complicated than you might expect. The library is clearly doing more than just storing the data that we provide in a type! The main thing that has happened is that the function we passed in has been processed by wrapping it into yet another type, `ODEFunction`. This contains information such as how to calculate the derivative of the function that can be provided for more advanced usage of the package.
"""

# ╔═╡ 6dfc2699-fe5e-4b55-8437-0866f715170a
md"""
Similarly we can look inside the solution object:
"""

# ╔═╡ 17d6021b-0d21-4d04-a8ae-64743cc01404
fieldnames(typeof(solution))

# ╔═╡ f4c3b174-57ac-461b-a816-869460d8896d
Dump(solution)

# ╔═╡ 500973f5-faa7-4a31-b812-d0b20dbb9d82
md"""
The solution is even more complicated, containing not only the data that was calculated, but also all information about which algorithms were used to solve the problem and tables of coefficients for interpolation.
"""

# ╔═╡ 4a575627-32c2-4842-a540-52253f79ff89
md"""
## Parameterized types
"""

# ╔═╡ bea223ef-f961-457c-8962-d0ed396f4717
md"""
Let's look now at the *type* of the DiffEq objects:
"""

# ╔═╡ 92bd82dc-eaab-4869-9d95-cb4647092352
typeof(problem)

# ╔═╡ d518828a-1972-4bfa-b96f-0ebbf86b6193
md"""
We see that indeed `problem` is an object of type `ODEProblem`, but that it also has several **type parameters**, which are listed inside the curly braces (`{` and `}`):
"""

# ╔═╡ c3570557-55bd-4189-a332-aa2e54693665
[typeof(problem).parameters...]

# ╔═╡ 1f684b1b-8e43-4e0e-8a47-1f2474cc5e4f
md"""
For example, we see that the second type parameter is of type `Tuple{Float64, Float64}`, which is the type of the variable `time_span`:
"""

# ╔═╡ 1bf84cfe-9ddc-4644-b7ea-8ac43b01d477
typeof(time_span)

# ╔═╡ d2a1f97e-2e3e-4ae9-8287-2de62065d708
md"""
One of the reasons for using parameterized types like this is for efficiency: Julia is fastest when it knows the exact types of every variable, since then it can generate efficient machine code.
"""

# ╔═╡ a71aaf0b-8ce9-4242-b31f-8f05730ce2e7
md"""
# Making parameterized types
"""

# ╔═╡ 09ddc00b-2361-44c5-b194-7b4b6d24bef4
md"""
Let's see how and why we would make a parameterized type.

Let's think again about the `EulerOutput` type that we started to write above.
What happens if, as in the `DifferentialEquations.jl` library, we would like to handle both *scalar* ODEs (one single real dependent variable), and *vector* ODEs?

We can no longer specify that the output type is a `Vector{Float64}`, so we would need to leave it unspecified:
"""

# ╔═╡ c7c16e89-3ccf-4d70-9cdb-d5ae71a5ea49
struct SimpleEulerOutput2
	times
	values
end

# ╔═╡ 1872c30d-aece-44a5-b3b6-36e82c072609
md"""
However, it turns out that we then lose efficiency (speed).

What we would need is two different types: one in which we specify the type of `times` as `Vector{Float64}`, and the other as `Vector{Vector{Float64}}`.

Julia provides a mechanism to define *both at once*: this is what parameterized types do!
"""

# ╔═╡ 5d06734e-ffbd-421e-820e-609e373ccc7b
md"""
Let's define our own `MyODEProblem` type as a simple example. We specify type parameters for each variable:
"""

# ╔═╡ 62d55814-c863-4d4d-8582-908721d37326
struct MyODEProblem{T, U}
	t0::T
	tfinal::T

	u0::U
end

# ╔═╡ 470b7808-5ff6-4cda-82cc-2ac250d51909
md"""
Note that the starting and final times have been specified to require the *same* type `T`.
"""

# ╔═╡ d392de63-137e-43fb-86ca-4f344769bc71
md"""
Now if we create objects of this type, we see that Julia matches the types we pass in to the correct type parameters:
"""

# ╔═╡ e07830be-3aab-4410-a830-786a30eb162a
newprob = MyODEProblem(3, 4, 100.0)

# ╔═╡ 6d809ff0-20c5-4910-aca0-612496955812
newprob.t0

# ╔═╡ 816a64ed-1f88-43a8-9df7-8070372e84ea
newprob.tfinal

# ╔═╡ cdd6a0dd-e188-4d37-ab6d-12e7fecd3e90
newprob.u0

# ╔═╡ c29d4796-9e52-4234-8325-5bf31099c79d
md"""
If we change the types of the inputs, the type parameters change:
"""

# ╔═╡ 7e5b0073-1bc7-41c7-ae74-c9c1bdf6c6ea
MyODEProblem(3.0, 4.0, 100.0)

# ╔═╡ 67e36f7c-0cb9-4ead-82ad-4c47639802a8
md"""
If we mix types for the first two arguments then we get an error, since Julia can't match to a single type `T`:
"""

# ╔═╡ cdc25238-ed48-42e5-9db4-2367577b215d
MyODEProblem(3.0, 5, 100.0)

# ╔═╡ f66b6078-bab5-48ef-9241-8d5fb0dc78ea
md"""
We can use vectors:
"""

# ╔═╡ 73016af9-5efe-4af2-9e99-ce1a7de6fa96
MyODEProblem(3, 4, [17.0, 18.0])

# ╔═╡ Cell order:
# ╟─54e6bd88-a6a9-11eb-3380-49b79430f6cf
# ╠═5997bb40-b508-4647-8b12-a7cc5152e550
# ╠═38c74e63-13e9-49a9-8eae-4a898197647a
# ╟─7a1d476b-1302-4b98-b9fa-01b990979985
# ╟─3004336a-0250-483f-be00-60db592c9fc8
# ╟─a1cfd00b-825f-494e-91a9-192815f06ad6
# ╟─b2d2b8f0-c81d-437a-83cf-726b5caa77ab
# ╟─d27ec502-ea0c-4d15-a223-7cbf6b784d1c
# ╟─6f87059a-dbdc-46dd-933b-7bb36980868e
# ╟─557cedbc-6e44-40a8-b0ab-37ae11b63f55
# ╠═a6ae905b-db28-4e75-b567-54840dcda9a5
# ╟─f16644d0-e2d2-44ac-988c-aa5b178844f5
# ╟─1ea2acfe-6ba1-44d3-9a24-a904a5fac0c4
# ╠═264a27cd-9ace-48cc-9556-7eb798298f17
# ╟─70ebe76d-6107-492d-a3d2-84c1dd7f59fb
# ╠═6e7c8e9e-bef0-4bf8-b1b1-d50c82aa203e
# ╟─b5379dc2-d97f-47ed-8737-35e3fe59285c
# ╟─0ed8e007-2137-43e0-93b8-6f5dfb9496f5
# ╠═ff91515a-89a4-4423-a186-5572c712493d
# ╟─64b26404-1ca0-4aa4-bea0-3c9075b08298
# ╠═b07de2ce-640b-42d6-8b60-39fcbcd116e7
# ╟─0b1855de-a8c9-492e-a249-3238e41fe84c
# ╟─0190bdb0-97bf-4c35-b3c8-e97c1940b702
# ╟─555a6328-8aff-4f9f-87f5-51a67450002d
# ╠═d65372df-e6ca-4165-8d9b-b21cc5c9f796
# ╟─b274c61b-4fab-4f4a-99df-5d39e0f56aa1
# ╠═e3d9e0be-6cf4-4ae1-8c7f-b444a153a9f5
# ╠═180bcdc1-5f51-4b32-8b9c-5000605cdf32
# ╟─13b38f88-2ead-46b3-bc96-eae2ea10204d
# ╟─475ef0e5-7f3a-4bbb-b513-8be29abfd3f9
# ╠═0ed6d65c-707d-4666-b203-2ad0ea822687
# ╟─cf9a293c-aa0f-40ce-960b-d4c3b82b8346
# ╟─efcfdc1a-5a80-4780-ba56-43c5e2d6ea36
# ╟─5d7a81d7-8080-4390-924a-72c34c0a5e23
# ╟─6c607126-dd0b-438d-8724-c1ef5e79c71d
# ╠═d16e06f4-55c4-4c0e-bde5-c0ce4b36b043
# ╠═aa0e4260-6268-4eff-a662-f3fd55e7d450
# ╠═d0f40681-73df-4cd3-bbd5-3edb8193153e
# ╠═fd122a25-5655-410b-aa7f-34936fc97b53
# ╠═8ac7da6b-46e0-470b-91b3-1a61c226fa4a
# ╟─7aa45efe-865d-4e0e-9c71-0c032c72c40d
# ╠═bca8112f-cb61-43dc-ae87-383915c8a89b
# ╠═a766a141-5d7b-499f-9e18-48bf926ee7ea
# ╟─76cbc37d-54c8-4626-8bfe-58b63a602c38
# ╟─c2765282-bcc7-4110-9822-10557326461e
# ╠═7671506d-d4b7-4792-9571-e003097235e1
# ╠═203e90b0-4d0c-4999-8513-d4eb43f53aac
# ╟─326890ae-0675-4779-b5e8-9b3e0412a52b
# ╠═d7025d1f-c3b8-461a-94dc-2414fbdfd373
# ╠═7f5c41f9-96b9-40d6-87f8-3e81c372b48e
# ╟─d8f15e9b-aef1-4795-8522-85ea3564e351
# ╟─e723400e-3f8f-47d6-8944-28355a54c698
# ╟─15fdb430-4157-4401-85b9-665dde795a22
# ╠═a410e98d-dd7f-4f55-8812-8b0339a56c86
# ╟─239f39b2-53c6-4da9-ac25-b5e7d19deb5f
# ╠═6630cbc8-1b2f-4ea9-bcd6-e4f4fcb2e844
# ╟─5c1c7016-69aa-4009-8dc8-0c68c925ee66
# ╠═83719ef1-e8d6-44d0-98ab-a851b1082fa5
# ╟─e4d36db5-352b-4fb2-9169-da5c6a3458b7
# ╠═167015ff-1caf-42a8-85f1-070da3e5d5c5
# ╟─f8bcd82d-089c-49a2-99d0-b17a5ac5c509
# ╠═cd17e738-089b-4934-af63-d27706c5d1c7
# ╟─fc525f77-ad62-43cb-92e4-0fc1354be99c
# ╠═ec90842a-a20c-4847-a803-e070d343b98d
# ╟─9fad6f01-a3b7-4dbb-a8f5-29f2ea1f545a
# ╟─658d8d15-5b3e-4d35-88b1-4e2fa0f3b17a
# ╟─aafeeb55-943f-4375-9cc0-e0dce659d161
# ╟─3c45ad21-d18b-41fa-a0fd-64d8988621f3
# ╠═96e871a8-48b9-459b-a1a0-94fc1deb941e
# ╟─2622801c-4e37-42f1-ab77-939d461b8d20
# ╠═43e7b270-6452-4a4c-ba6f-c286a337a423
# ╟─dbdb9998-ca3b-46fb-87cf-5b7c6e387ab9
# ╟─b8fbff51-72fe-4432-9f6f-06f4628320d4
# ╠═ef677abe-c96d-41cc-9b4a-192ff29cdb86
# ╠═78933b33-8a3c-492e-8a58-19ad7187b45e
# ╟─9ef4aa78-2ecd-4d53-87ae-f0909bb46915
# ╠═798bc482-7556-4904-8946-0d9898ce2d33
# ╟─32e9f6e5-8df6-4d40-8092-efaf2bce28d3
# ╟─6dfc2699-fe5e-4b55-8437-0866f715170a
# ╠═17d6021b-0d21-4d04-a8ae-64743cc01404
# ╠═f4c3b174-57ac-461b-a816-869460d8896d
# ╟─500973f5-faa7-4a31-b812-d0b20dbb9d82
# ╟─4a575627-32c2-4842-a540-52253f79ff89
# ╟─bea223ef-f961-457c-8962-d0ed396f4717
# ╠═92bd82dc-eaab-4869-9d95-cb4647092352
# ╟─d518828a-1972-4bfa-b96f-0ebbf86b6193
# ╠═c3570557-55bd-4189-a332-aa2e54693665
# ╟─1f684b1b-8e43-4e0e-8a47-1f2474cc5e4f
# ╠═1bf84cfe-9ddc-4644-b7ea-8ac43b01d477
# ╟─d2a1f97e-2e3e-4ae9-8287-2de62065d708
# ╟─a71aaf0b-8ce9-4242-b31f-8f05730ce2e7
# ╟─09ddc00b-2361-44c5-b194-7b4b6d24bef4
# ╠═c7c16e89-3ccf-4d70-9cdb-d5ae71a5ea49
# ╟─1872c30d-aece-44a5-b3b6-36e82c072609
# ╟─5d06734e-ffbd-421e-820e-609e373ccc7b
# ╠═62d55814-c863-4d4d-8582-908721d37326
# ╟─470b7808-5ff6-4cda-82cc-2ac250d51909
# ╟─d392de63-137e-43fb-86ca-4f344769bc71
# ╠═e07830be-3aab-4410-a830-786a30eb162a
# ╠═6d809ff0-20c5-4910-aca0-612496955812
# ╠═816a64ed-1f88-43a8-9df7-8070372e84ea
# ╠═cdd6a0dd-e188-4d37-ab6d-12e7fecd3e90
# ╟─c29d4796-9e52-4234-8325-5bf31099c79d
# ╠═7e5b0073-1bc7-41c7-ae74-c9c1bdf6c6ea
# ╟─67e36f7c-0cb9-4ead-82ad-4c47639802a8
# ╠═cdc25238-ed48-42e5-9db4-2367577b215d
# ╟─f66b6078-bab5-48ef-9241-8d5fb0dc78ea
# ╠═73016af9-5efe-4af2-9e99-ce1a7de6fa96


include("graph.jl")
include("opinion.jl")

function row_sum_max(M::Matrix{WEIGHT})
    (n, m) = size(M)
    acc = fill(0.0, n)
    for i = 1:n
        for j = 1:n
            acc[i] += abs(M[i, j])
        end
    end
    return maximum(acc)
end

function main()
    dataset = "cloister"
    filepath = "./data/raw/$(dataset).edgelist"
    @show filepath
    (W, U, n) = read_edgelist(filepath)
    m = 2
    normalize_edgelist!(W)
    normalize_edgelist!(U)
    y = generate_opinions(n, m)
    α, β, γ = generate_factors(n)
    M = (β .* W) - (γ .* U)
    Q = get_fundamental_matrix(β, W, γ, U)
    z = get_equlibrium(α, Q, y)
    z_next = get_next_opinion(α,
                              β,
                              γ,
                              W, U, y,z)
    M = Matrix(β .* W - γ .* U)

    y_trans = y'
    z_trans = z'

    p = scatter(y_trans[1,:], y_trans[2,:], fmt = :png)
    savefig("before_equlilbrium_$(dataset).png")
    p = scatter(z_trans[1,:], z_trans[2,:], fmt = :png)
    savefig("after_equlilbrium_$(dataset).png")

    return y, Q, z
end

y, Q, z = main()

df = DataFrame(X = Float64.(y[:,1]), Y = Float64.(y[:, 2]))

df1 = DataFrame(X = Float64.(z[:,1]), Y = Float64.(z[:, 2]))


"""
Calculate Dynamic Programming using Bayesian approach
    dp_bayes(q1, q1L, q2L, times, cut)
    :param q1: vector (N)
    :param q1L: vector (q1L)
    :param q2L: vector (q2L)
    :param times: MCMC parameter number of simulations
    :param cut: MCMC parameter number of cuts

    :return MatchIn2: warping function
    :return NDist: minimal distance
    :return q2LL: re-parameterized q2L
"""
function dp_bayes(q1, q1L, q2L, times, cut)

    colnum = length(q1L);
    rownum = div(colnum,times);
    q2LLlen = (colnum-1)*times+1;
    q2LL = zeros(q2LLlen);
    tempspan1 = collect(1:q2LLlen)-1;
    temp_span2 = float(tempspan1);
    timesf = float(times);
    q2LL_time = (temp_span2*(1/timesf))+1;
    q2L_time1 = collect(1:colnum);
    q2L_time2 = float(q2L_time1);
    q2LL_i = interpolate((collect(1:length(q2L)),), q2L, Gridded(Linear()))
    q2LL = q2LL_i[q2LL_time];
    ID = zeros(rownum+1, colnum+1);
    S = zeros(ID);
    jend = 0; start = 0; end1 = 0; k = 0; index =0;
    interx = zeros(Int64, times-1);
    intery2 = zeros(times-1);
    intery = zeros(Int64, times-1);
    span1 = (1:(times-1));
    span2 = float(span1);
    q1x = Array(Float64, times-1);
    q2y = Array(Float64, times-1);

    for i in 1:(rownum-1)
        if (i*cut+1) < colnum
            jend = i*cut+1;
        else
            jend = colnum;
        end

        for j in i:jend-1
            start = max(i-1,j-cut);
            end1 = min(j-1, cut*(i-1));
            n = start-1+collect(1:(end1-start+1));
            k = length(n);
            interx = times*(i-1)+span1;
            Energy = Array(Float64, k);

            for m in 1:k
                jf = j;
                divisor = (jf-n[m])/times;
                intery2 = divisor.*span2+n[m];
                for l in 1:(times-1)
                    intery[l] = round(times*intery2[l]);
                    q1x[l] = q1L[interx[l]+1];
                    q2y[l] = q2LL[intery[l]+1];
                end
                Energy[m] = S[i,n[m]+1]+(q1[i]-sqrt(divisor)*q2L[n[m]+1])^2 +
                    (norm(q1x-sqrt(divisor)*q2y))^2;
            end
            min_val, index = findmin(Energy);
            loc = n[index];
            S[i+1,j+1] = min_val;
            ID[i+1,j+1] = loc;
        end
    end

    i = rownum;
    j = colnum;
    start = max(i-1,j-cut);
    end1 = min(j-1,cut*(i-1));
    n = start-1+collect(1:(end1-start+1));
    k = length(n);
    interx = times*(i-1)+span1;
    span2 = float(span1);
    Energy = Array(Float64,k);

    for m in 1:k
        jf = j;
        divisor = (jf-n[m])/times;
        intery2 = divisor.*span2+n[m];
        for l in 1:times-1
            intery[l] = round(times*intery2[l]);
            q1x[l] = q1L[interx[l]+1];
            if intery[l]+1 < q2LLlen
                q2y[l] = q2LL[intery[l]+1];
            else
                q2y[l] = 0.0;
            end
        end
        Energy[m] = S[i,n[m]+1]+(q1[i]-sqrt(divisor)*q2L[n[m]+1])^2 +
                    (norm(q1x-sqrt(divisor)*q2y))^2;
    end

    min_val, index = findmin(Energy);
    loc = n[index];
    S[i+1,j+1] = min_val;
    ID[i+1,j+1] = loc;

    path = Array(Int64,rownum);
    count = round(Integer,ID[i+1,j+1]);
    oldcount = 0;
    path[i] = count;

    while count>1
        i -= 1;
        oldcount = count;
        count = round(Integer, ID[i+1,oldcount+1]);
        path[i] = count;
    end

    MatchIn2 = path+1;
    NDist = S[rownum+1,colnum+1]/colnum;

    return MatchIn2, NDist, q2LL
end

"""
Calculate mean function using Bayesian Dynamic Programming
    DP_mean(f, times=5)
    :param f: array (M,N) of N functions
    :param times: MCMC parameter number of times to split

    Returns Dict containing
    :return distmfamily: distance matrix
    :return match_matrix: match_matrix
    :return position: position
    :return mu_5: mean function
    :return rtmatrix: rtmatrix
    :return sumdist: sum distance
    :return qt_fitted: qt fitted matrix
    :return esimator: estimator
    :return estimator2: estimator2
    :return regcurve: registered curves
"""
function DP_mean(f, times=5)
    cut = 5*times;
    iter = 20;
    timet = LinRange(0,1,size(f,1));
    group = collect(1:size(f, 2));
    m = length(timet) - 1;
    n = size(f, 2);
    qt_matrix = zeros(m, n);

    for j = 1:n
        qt_matrix[:, j] = sign(diff(f[:,j])).*sqrt(abs(diff(f[:,j]))./diff(timet));
        rescale = sqrt(m/sum(qt_matrix[:, j].^2));
        qt_matrix[:, j] = rescale.*qt_matrix[:, j];
    end
    row = collect(1:times:m);
    qt_fitted_matrix = zeros(m, n);

    search = Array(Float64, n);
    meanq = mean(qt_matrix, 2);
    for j = 1:n
        search[j] = Enorm(vec(qt_matrix[:, j] - meanq));
    end
    position = argmin(search);

    mu_5 = qt_matrix[:, position];
    mu_curve = zeros(iter, m+1);
    dist_matrix = zeros(iter+1, n);
    for j = 1:n
        dist_matrix[1, j] = (Enorm(vec(mu_5 - qt_matrix[:, j])))^2/m;
    end
    rtmatrix = zeros(m+1, n);
    match_matrix = zeros(length(row)+1, n);

    i = 1;
    diffdist = 1000.;
    while ((i<iter) & (diffdist > 0.001))
        for j = 1:n
            MatchIn2, NDist, q2LL = dp_bayes(mu_5[row], mu_5, qt_matrix[:, j],
                                             times, cut);
            match = [MatchIn2; m+1];
            f_i = interpolate(([row; m+1],), float(match), Gridded(Linear()))
            idy = round(Integer, f_i[1:m]);
            idy[idy .> m] = m;
            scale1 = sqrt(diff(match)*(1/times));
            scalevec = kron(scale1, ones(times));
            tmpidy = copy(idy) - 1;
            extended_idy = tmpidy[1:m]*times+1;
            q_fitted = scalevec.*q2LL[extended_idy];
            qt_fitted_matrix[:, j] = q_fitted;
            dist_matrix[i+1, j] = NDist;
            rtmatrix[:, j] = [idy; m+1];
            match_matrix[:, j] = match;
        end
        diffdist = abs(sum(dist_matrix[i+1, :]) - sum(dist_matrix[i, :]));
        mu_5 = mean(qt_fitted_matrix, 2);
        rescale = sqrt(m/sum(mu_5.^2));
        mu_5 *= rescale;
        i += 1;
    end

    estimator2 = copy(mu_5);
    invidy, revscalevec = findkarcherinv(match_matrix, times);
    invidy = invidy[1:m];
    invidy[invidy .>= m] = m;
    f_i = interpolate((collect(1:m),), vec(mu_5), Gridded(Linear()))
    mu_5 = revscalevec.*f_i[invidy];
    rescale = sqrt(m/sum(mu_5.^2));
    estimator = rescale.*copy(mu_5);
    reg_curve = zeros(m+1, n);
    for j = 1:n
        f_i = interpolate((collect(0:m),), f[:, j], Gridded(Linear()))
        tmp = f_i[LinRange(0, m, times*(m+1)-1)];
        reg_curve[:, j] = tmp[round(Integer, (rtmatrix[:,j] - 1)*times+1)];
    end
    crossmean = mean(reg_curve, 2);
    sumdist = sum(reg_curve, 2);

    out = Dict("distfamily" => dist_matrix, "match_matrix" => match_matrix,
               "position" => position, "mu_5" => mu_5,
               "rtmatrix" => rtmatrix, "sumdist" => sumdist,
               "qt_fitted" => qt_fitted_matrix, "estimator" => estimator,
               "estimator2" => estimator2, "regcurve" => reg_curve);
    return out
end


export dmrg

function dmrg(H::MPO,
              psi0::MPS,
              sweeps::Sweeps;
              kwargs...)::Tuple{Float64,MPS}

  which_factorization::String = get(kwargs,:which_factorization,"automatic")

  psi = copy(psi0)
  N = length(psi)

  PH = ProjMPO(H)
  position!(PH,psi0,1)
  energy = 0.0

  quiet::Bool = get(kwargs,:quiet,false)
  for sw=1:nsweep(sweeps)
    sw_time = @elapsed begin

    @debug begin
      reset!(timer)
      println("\n\n~~~ This is sweep $sw ~~~\n\n")
    end

    for (b,ha) in sweepnext(N)
      position!(PH,psi,b)

      phi = psi[b]*psi[b+1]

      energy,phi = davidson(PH,phi;kwargs...)

      dir = ha==1 ? "fromleft" : "fromright"
      replaceBond!(psi,b,phi;
                   maxdim=maxdim(sweeps,sw),
                   mindim=mindim(sweeps,sw),
                   cutoff=cutoff(sweeps,sw),
                   dir=dir,
                   which_factorization=which_factorization)
    end
    end
    if !quiet
      @printf("After sweep %d energy=%.12f maxLinkDim=%d time=%.3f\n",sw,energy,maxLinkDim(psi),sw_time)
    end
  end
  return (energy,psi)
end

@doc """
dmrg(H::MPO,psi0::MPS,sweeps::Sweeps;kwargs...)::Tuple{Float64,MPS}

Optimize a matrix product state (MPS) to be the eigenvector
of the Hermitian matrix product operator (MPO) H with minimal
eigenvalue using the density matrix renormalization group
(DMRG) algorithm.

Inputs:
* `H::MPO` - a Hermitian MPO
* `psi0::MPS` - MPS used to initialize the optimization
* `sweeps::Sweeps` - `Sweeps` object used to control the algorithm

Returns:
* `energy::Float64` - eigenvalue of the optimized MPS
* `psi::MPS` - optimized MPS
""" dmrg


using Ripserer
using SparseArrays
using Test
using TupleTools

using Ripserer: adjacency_matrix, births

include("interfacetest.jl")

@testset "Custom filtration 1" begin
    flt = Custom(
        [
            (1,) => 0,
            (4,) => 0,
            (1, 3) => 2,
            (1, 4) => 3,
            (3, 4) => 6.0,
            (1, 2, 3) => 7,
            (1, 2, 4) => 8,
            (1, 3, 4) => 9,
            (1, 2, 3, 4) => 9,
            (3,) => 10_000,
        ];
        threshold=8,
    )

    @test flt isa Custom{Int,Float64}
    @test dim(flt) == 3
    @test sort(flt[0]; by=index) ==
        [Simplex{0}(1, 0.0), Simplex{0}(2, 7.0), Simplex{0}(3, 2.0), Simplex{0}(4, 0.0)]
    @test isempty(flt[10])
    @test simplex(flt, Val(2), (1, 2, 4)) === Simplex{2}((4, 2, 1), 8.0)
    @test simplex(flt, Val(2), (4, 3, 1)) === nothing

    @test all(eachindex(adjacency_matrix(flt))) do cart_ind
        i, j = TupleTools.sort(Tuple(cart_ind); rev=true)
        if i ≠ j
            idx = index((i, j))
            return adjacency_matrix(flt)[i, j] == haskey(flt.dicts[2], idx)
        else
            return !adjacency_matrix(flt)[i, j]
        end
    end

    d0, d1, d2 = ripserer(flt; dim_max=2)
    @test d0 == [(0, 3), (0, Inf)]
    @test d1 == [(6, Inf)]
    @test d2 == []
end

@testset "Custom filtration vs Rips" begin
    dist = [
        0 1 2 1
        1 0 1 2
        2 1 0 1
        1 2 1 0
    ]

    cf = Custom(
        [
            [(i,) => 0 for i in 1:4]
            [(i, j) => dist[i, j] for i in 1:4 for j in (i + 1):4]
            [
                (i, j, k) => max(dist[i, j], dist[i, k], dist[j, k]) for i in 1:4 for
                j in (i + 1):4 for k in (j + 1):4
            ]
        ],
    )
    @test simplex(cf, Val(3), (4, 3, 2, 1)) == nothing
    @test births(cf) == zeros(4)

    @test ripserer(dist) == ripserer(cf)
end

@testset "Custom filtration vs sparse Rips" begin
    spdist = sparse(
        [
            0 1 0 0 0 1
            1 0 1 0 0 0
            0 1 0 1 0 0
            0 0 1 0 1 0
            0 0 0 1 0 1
            1 0 0 0 1 0
        ]
    )
    cf = Custom([
        (1,) => 0,
        (2,) => 0,
        (3,) => 0,
        (4,) => 0,
        (5,) => 0,
        (6,) => 0,
        (1, 2) => 1,
        (2, 3) => 1,
        (3, 4) => 1,
        (4, 5) => 1,
        (5, 6) => 1,
        (6, 1) => 1,
    ])
    sprips = Rips(spdist)

    @test adjacency_matrix(cf) == Bool.(spdist)
    @test threshold(cf) == 1
    @test ripserer(cf) == ripserer(sprips)
end

@testset "Overflow" begin
    big_simplex = Tuple(Ripserer._vertices(Int128(typemax(Int64) ÷ 2), Val(6)))
    @test_throws OverflowError Custom{Int}([big_simplex => 1])
    @test begin
        Custom{Int128}([big_simplex => 1])
        true
    end
end

@testset "Interface" begin
    data = [
        (1,) => 0,
        (4,) => 0,
        (1, 3) => 2,
        (1, 4) => 3,
        (3, 4) => 6.0,
        (1, 2, 3) => 7,
        (1, 2, 4) => 8,
        (1, 3, 4) => 9,
        (1, 2, 3, 4) => 9,
        (3,) => 10_000,
    ]

    test_filtration(Custom, data; flt_kwargs=(; threshold=9), dim_max=2)
end


# World of testing
#Testing the forward-mode of the Poisson solver module Poisson.jl using combinations of random eigenfunctions as a source
#Developed by:
#Niels Grobbe, Massachusetts Institute of Technology, USA.
#In collaboration with: Aimé Fournier & Laurent Demanet, Massachusetts Institute of Technology, USA.
#Date: October, 2017
#Contact: ngrobbe@gmail.com

using GeoPhyInv
using Polynomials
using Test
using StatsBase

#include("core.jl")
#include("code_original.jl")

############################################
# for manual random or uniform sigma testing
# sigma=ones(rand(16:64),rand(16:64)); # random size
# # sigma=ones(40,40);
############################################

# function input random source alterations and sigma
#ntst = 5;#  how many tests
ntst=32; # how many tests: larger number, more more gridpoints (finer grid)
wavs = [rand(1:8,2,5); 8*randn(1,5)]; # pre-set random wavenumbers --> how many complex source terms we want to add..
t = ["Computed potential" "True potential" "interior" "boundary"];
m=["or","vk"];

err=zeros(ntst,4);
for i = 1 : ntst
    sigma=ones(floor(Int,2^(4+2.5*(i - 1)/(ntst - 3))),
	       floor(Int,2^(4+2.5*(i - 1)/(ntst - 3))));

    global fields # to make it a global variable

# Get model size (also in module, but here local for creating the input for the module)
    nz=size(sigma,1);
    nx=size(sigma,2);
    dz= 1.0*inv(nz - 1);
    dx= 1.0*inv(nx - 1);

    z=(0 : nz - 1)*dz;
    x=(0:nx-1)*dx;

# For random tests
    if size(wavs, 1) == 3
   	 l = wavs[1 : 2,:];
   	 j = wavs[3,:]'; # why 3? (take 3rd row which is amplitudes..) # --> make row vector for c definition below
    else
# for just a single test (if random tests not specified)
    	l = [rand(1:ceil(nz/3),1,3) ;rand(1:ceil(nx/3),1,3)];  #% random wavevector ... #% for eigenfunction source
    	j = 8*randn(1, size(l, 2));
    end

    l = pi*l;
# preallocate source and reference electrical potential
    src = zeros(nz,nx);
    psi_ref = src;

    for k = 1 : size(l,2)
        c = j[k].*cos.(l[1,k].*z)*cos.(l[2,k].*x'); # result is size nz*nx
        src = src + c; # add eigenfunction to source (to make source more complex)
        psi_ref = psi_ref - c./(l[:,k]'*l[:,k]);
    end

    psi_ref = psi_ref./sigma;
    println("Solving Poisson equation, testrun number:",i,"\t", size(sigma))

    global pa
    pa=GeoPhyInv.Poisson.Param(nx,nz, sigma)
    psi=reshape(GeoPhyInv.Poisson.applyinvA(src,pa),nz,nx)

   #psi,Ac = Poisson.solve(src,sigma,1)

    # subtract mean for DC
    psi_ref=psi_ref.-StatsBase.mean(psi_ref);
    psi=psi.-StatsBase.mean(psi);

    fields=(psi,psi_ref); #make tuple

#println("testing...")
# Normalized RSME errors for interior points only, and interior+boundary points
    err[i,1] = size(fields[1],1);
    err[i,2] = size(fields[1],2);
    err[i,3] = sqrt(sum(sum((fields[1][2 : end - 1, 2 : end - 1] - # diff recovered-true only for inner points without boundaries
                fields[2][2 : end - 1, 2 : end - 1]).^2))/sum(sum(fields[2][2 : end - 1, 2 : end - 1].^2)));
    err[i, 4] = sqrt(sum(sum([(fields[1][1,1 : end - 1]-
                fields[2][1,1 : end - 1])' fields[1][1 : end - 1,end]'-
                fields[2][1 : end - 1,end]' (fields[1][end,end : -1 : 2] -
                fields[2][end,end : -1 : 2])' fields[1][end : -1: 2,1]' - fields[2][end : -1: 2,1]'].^2))/
    sum(sum([fields[2][1,1 : end - 1]' fields[2][1 : end - 1,end]' fields[2][end,end : -1 : 2]' fields[2][end : -1: 2,1]'].^2)));

# test to study error decay for interior only, and interior+boundary points
 if i > 2
        x = minimum(err[:,1 : 2], dims=2); # error over smallest dimension
      for k = 3:4
           global p,chk
           p = polyfit(log.(x[1:i]), log.(err[1:i,k]),1);  # fit polynomial of first degree using least squares fit
           println(p)
           #chk=round(p[1],4)
	   chk=p[1]
       end
   end

end # end i-loop

# Test passed?

# first pass interior points, then also include boundary points
@test chk<=-0.8
@test (err[end,3]/err[1,3]) < 0.2
@test (err[end,4]/err[1,4]) < 0.2

#if(chk<=-0.8 && err[end,3]/err[1,3] < 0.2 && err[end,4]/err[1,4] < 0.2 ) i
#println("tests successfully passed!")
#println("log-log slope <=-0.8? slope=",chk)
#println("Normalized RMSE interior < 0.2? err[end,3]/err[1,3]=",err[end,3]/err[1,3])
#println("Normalized RMSE int.+boundary < 0.2? err[end,4]/err[1,4]=",err[end,4]/err[1,4])
#
#println("log-log slope 'chk' <=-0.8? slope=",chk)
#println("iteration=",collect(1:length(err[:,3])));
#println("Normalized RMSE interior < 0.5? err[:,3]=",err[:,3])
#println("Normalized RMSE int.+boundary points < 0.5? err[:,4]=",err[:,4])
#error("ERROR: RMSE and/or log-log slope too large: try another run, or try increasing amount of gridpoints via ntst")


#=
% Diagnose the mean or ideal age.
=#
using Revise
using TMI, PyPlot, PyCall

TMIversion = "modern_90x45x33_GH10_GH12"
TMIversion = "modern_180x90x33_GH11_GH12"

A, Alu, γ, TMIfile, L, B = config_from_nc(TMIversion)

a = meanage(TMIversion,Alu,γ)

# plot a section at 330 east longitude (i.e., 30 west)
#lon_section = 330 # only works if exact
#lon_section = -150
lon_section = -151
lims = 0:100:2000
tlabel = string(lon_section)*"°E, Mean Age [yr]"
sectionplot(a,lon_section,lims,titlelabel = tlabel)


#=
## Generate `PowerDynamics.jl`-Node with `BlockSystems`

We want to model the
[`VSIVoltagePT1`](https://juliaenergy.github.io/PowerDynamics.jl/dev/node_types/#PowerDynamics.VSIVoltagePT1)
with the help of `BlockSystems`.

We start by defining some general stuff...
=#
using BlockSystems
using ModelingToolkit

@parameters t
D = Differential(t)
nothing #hide

#=
### Defining the common Blocks
We need some common parts such as filters or integrators. The actual Blocks will
be defined by using these 'blueprints'. Stuff like this should go to a library of
components at some point.

=#
# #### low pass filter
@parameters τ input(t)
@variables filtered(t)

lpf = IOBlock([D(filtered) ~ 1/τ * (- filtered + input)],
              [input], [filtered])

# #### voltage source

@parameters ω(t) v(t) τ
@variables u_i(t) u_r(t) A(t)

## explicit algebraic equation for A will be reduced at connect
voltage_source = IOBlock([A ~ 1/τ * (v/√(u_i^2 + u_r^2) - 1),
                          D(u_r) ~ -ω * u_i + A*u_r,
                          D(u_i) ~  ω * u_r + A*u_i],
                         [ω, v], [u_i, u_r])

# #### Droop control

@parameters K u_ref x_ref x(t)
@variables u(t)

droop_control = IOBlock([
    u ~ - K * (x - x_ref) + u_ref # output is the droop voltage v
    ], [x], [u])

#=
After defining the blueprints we can create the Blocks of the system
based on them:
=#

p_filter = IOBlock(lpf, name = :p_filter)
q_filter = IOBlock(lpf, name = :q_filter)
p_droop = IOBlock(droop_control, name = :p_droop)
q_droop = IOBlock(droop_control, name = :q_droop)
v_source = IOBlock(voltage_source, name = :v_source)

nothing #hide
#=
We can put the blocks together to form this system:

```
     +----------+  +-------------+
 P --| p_filter |--|   p_droop   |   +----------+
     |    τ     |  | u_ref, xref |---|          |
     +----------+  +-------------+  ω| v_source |-- u_r
                                     |    τ     |
     +----------+  +-------------+  v|          |-- u_i
 Q --| q_filter |--|   q_droop   |---|          |
     |    τ     |  | u_ref, xref |   +----------+
     +----------+  +-------------+

```
=#
gfi = IOSystem([p_filter.filtered => p_droop.x,
                q_filter.filtered => q_droop.x,
                p_droop.u => v_source.ω,
                q_droop.u => v_source.v],
               [p_filter, q_filter, p_droop, q_droop, v_source],
               name = :GridForming,
               namespace_map = [p_filter.input => :P_in,
                                q_filter.input => :Q_in,
                                p_filter.filtered => :p_filtered,
                                q_filter.filtered => :q_filtered,
                                ## parameter names which match VSIVoltagePT1
                                v_source.τ => :τ_v, # time constant voltage control delay
                                p_filter.τ => :τ_P, # time constant active power measurement
                                q_filter.τ => :τ_Q, # time constant reactive power measurement
                                p_droop.K  => :K_P, # droop constant frequency droop
                                q_droop.K  => :K_Q, # droop constant voltage droop
                                q_droop.u_ref => :V_r, # reference/ desired voltage
                                p_droop.u_ref => :ω_r, # reference/ desired frequency
                                p_droop.x_ref => :P, # active (real) power infeed
                                q_droop.x_ref => :Q], # reactive (imag) power infeed                .
               outputs = [v_source.u_i, v_source.u_r])

#=
There is still one problem: our connected system has the inputs `P` and `Q`. In order to construct
IONodes the IOSystems have to be systems which convert some complex current to complex voltage:
```
       +---------+
i_r -->|   ???   |--> u_r
i_i -->|         |--> u_i
       +---------+
```

We can achieve this by defining a nother block which converts `(i, u) ↦ (P, Q)`

```
               +----------+  +-------------+
     +-----+   | p_filter |--|   p_droop   |   +----------+
i_r--|     |-P-|   p_τ    |  |   P, ω_r    |---|          |
i_i--| pow |   +----------+  +-------------+  ω| v_source |---+--u_r
     |     |                                   |    τ     |   |
 +---|     |   +----------+  +-------------+  v|          |-+-|--u_i
 | +-|     |-Q-| q_filter |--|   q_droop   |---|          | | |
 | | +-----+   |   q_τ    |  |   Q, v_r    |   +----------+ | |
 | |           +----------+  +-------------+                | |
 | +--------------------------------------------------------+ |
 +-----------------------------------------------------------+

```
=#
@parameters u_i(t) u_r(t) i_i(t) i_r(t)
@variables P_in(t) Q_in(t)
pow = IOBlock([P_in ~ u_r*i_r + u_i*i_i,
               Q_in ~ u_i*i_r - u_r*i_i],
              [u_i, u_r, i_i, i_r], [P_in, Q_in], name=:pow)

gfi2 = IOSystem([gfi.u_i => pow.u_i,
                 gfi.u_r => pow.u_r,
                 pow.P_in => gfi.P_in,
                 pow.Q_in => gfi.Q_in],
                [gfi, pow], outputs=[gfi.u_i, gfi.u_r], name=:gfi_i_to_u)
nothing #hide

# and the system can be reduced to a single IOBlock
connected = connect_system(gfi2)
nothing #hide

#=
The following part should live inside `PowerDynamics.jl` at some point. It defines an
`IONode` which will be constructed from an `IOBlock` as a new `AbstractNode` type.
=#
using PowerDynamics
using NetworkDynamics: ODEVertex
struct IONode{T} <: AbstractNode
    block::IOBlock
    generated::T
    parameters::Vector{Float64}
end

function IONode(blk::IOBlock, parameters::Dict)
    ## BlockSpec: blk must by of type (i_r, i_i) ↦ (u_r, u_i)
    spec = BlockSpec([:i_r, :i_i], [:u_r, :u_i])
    @assert spec(blk) "Block has to follow PowerDynamics i/o conventions!"
    ## TODO check parameters
    gen = generate_io_function(blk,
                               f_states=[blk.u_r, blk.u_i],
                               f_inputs=[blk.i_r, blk.i_i],
                               f_params=keys(parameters), warn=false);
    IONode(blk, gen, collect(values(parameters)))
end

import PowerDynamics.construct_vertex
function construct_vertex(ion::IONode)
    gen = ion.generated
    function rhs!(dx, x, e_s, e_d, _p, t)
        i = total_current(e_s, e_d)
        gen.f_ip(dx, x, (real(i), imag(i)), ion.parameters, t)
    end
    ODEVertex(f = rhs!, dim = length(gen.states), mass_matrix = gen.massm, sym = Symbol.(gen.states))
end

import PowerDynamics.symbolsof
PowerDynamics.symbolsof(ionode::IONode) = Symbol.(ionode.generated.states)
PowerDynamics.dimension(ionode::IONode) = length(ionode.generated.states)

#=
Let's try it out:
The parameters have to be provided as a Dict. This dict should contain all internal parameters of the block.
The list of parameters can be retrieved from the block:
=#
@info "connected system:" connected.iparams

#=
The parameter dict can be defined in several ways, I think the first option is the most convenient
because it is a pure `Symbol` which does not depend on any variables in the scope. The namespace separator
can be typed as `\_+<tab>`.
```
:τ_P => 1.0
:gfi_i_to_u₊τ_P => 1.0
connected.τ_P => 1.0
```
=#
para = Dict(
    :τ_v => 0.005,    # time constant voltage control delay
    :τ_P => 0.5,      # time constant active power measurement
    :τ_Q => 0.5,      # time constant reactive power measurement
    :K_P => 0.396,     # droop constant frequency droop
    :K_Q => 0.198,     # droop constant voltage droop
    :V_r => 1.0,   # reference/ desired voltage
    :P   => 0.303, # active (real) power infeed
    :Q   => 0.126, # reactive (imag) power infeed                .
    :ω_r => 0.0)   # refrence/ desired frequency
nothing #hide

# Now let's test whether this works in PD...
using PowerDynamics: SlackAlgebraic, PiModelLine,PowerGrid,find_operationpoint
using OrderedCollections: OrderedDict

buses=OrderedDict(
    "bus1"=> IONode(connected, para),
    "bus2"=> SlackAlgebraic(U=1))

branches=OrderedDict(
    "branch1"=> PiModelLine(from= "bus1", to = "bus2",y=4.999131600798035-1im*15.263086523179553, y_shunt_km=0.0528/2, y_shunt_mk=0.0528/2))

powergrid = PowerGrid(buses,branches)
operationpoint = find_operationpoint(powergrid)
timespan= (0.0,0.1)
nothing #hide

# simulating a voltage perturbation at node 1
fault1 = ChangeInitialConditions(node="bus1", var=:u_r, f=Inc(0.2))
solution1 = simulate(fault1, powergrid, operationpoint, timespan)
using Plots
plot(solution1.dqsol)

#=
## Compare against `VSIVoltagePT1`
We  want to compare the newly defined node against the `VSIVoltagePT1` node for random data.
Both nodes have slightly different states:
 - `symbolsof(node_bs) == [:u_r, :u_i, :p_filtered, :q_filtered]`
 - `symbolsof(node_pd) == [:u_r, :u_i, :ω, :q_m]`
We need adapt the initial state for the different third variable:
```math
\omega = \omega_r - K_P\cdot(p_{\mathrm{filtered}} - P)\\
\dot\omega = - K_P \cdot \dot{p}_{\mathrm{filtered}}
```
=#
using Test
@testset "Compare IONode vs. PD Node" begin
    for i in 1:100
        para = Dict(
            :τ_v => rand(), # time constant voltage control delay
            :τ_P => rand(), # time constant active power measurement
            :τ_Q => rand(), # time constant reactive power measurement
            :K_P => rand(), # droop constant frequency droop
            :K_Q => rand(), # droop constant voltage droop
            :V_r => rand(), # reference/ desired voltage
            :P   => rand(), # active (real) power infeed
            :Q   => rand(), # reactive (imag) power infeed                .
            :ω_r => 0.0)    # refrence/ desired frequency

        ## create IONode
        node_bs = IONode(connected, para)
        f_bs = construct_vertex(node_bs).f

        ## create PDNode
        ## the PD node does not have the explicit ω_r parameter
        para_pd = delete!(copy(para), :ω_r)
        nt = NamedTuple{Tuple(keys(para_pd))}(values(para_pd))
        node_pd = VSIVoltagePT1(; nt...)
        f_pd = construct_vertex(node_pd).f

        ## create fake "edge data", 4 incoming, 4 outgooing with 4 states each
        es = [randn(4) for i in 1:4]
        ed = [randn(4) for i in 1:4]

        ## select random time
        t = rand()

        ## chose random initial state and account for initial ω in PD node
        x_bs = randn(4)
        x_pd = copy(x_bs)
        x_pd[3] = - para[:K_P] * (x_bs[3] - para[:P])

        ## create arrays for the results
        dx_bs = similar(x_bs)
        dx_pd = similar(x_pd)

        ## call both functions
        f_bs(dx_bs, x_bs, es, ed, nothing, t)
        f_pd(dx_pd, x_pd, es, ed, nothing, t)

        ## compare results
        ## we have to correct variable 3 of bs implementation to match dω
        @test dx_bs[1] ≈ dx_pd[1]                # u_r
        @test dx_bs[2] ≈ dx_pd[2]                # u_i
        @test - para[:K_P] * dx_bs[3] ≈ dx_pd[3] # ω
        @test dx_bs[4] ≈ dx_pd[4]                # q_filtered
    end
end
nothing #hide

#=
it works 🎉

## Benchmark
We can also run a quick benchmark of both node functions:
=#
para = Dict(:τ_v=>rand(),:τ_P=>rand(), :τ_Q=>rand(),
            :K_P=>rand(), :K_Q=>rand(), :V_r=>rand(),
            :P=>rand(), :Q=>rand(), :ω_r=>0.0)

node_bs = IONode(connected, para)
f_bs = construct_vertex(node_bs).f
para_pd = delete!(copy(para), :ω_r)
nt = NamedTuple{Tuple(keys(para_pd))}(values(para_pd))
node_pd = VSIVoltagePT1(; nt...)
f_pd = construct_vertex(node_pd).f

es = [randn(4) for i in 1:4]
ed = [randn(4) for i in 1:4]
t = rand()

## chose random initial state and account for initial ω in PD node
x_bs = randn(4)
x_pd = copy(x_bs)
x_pd[3] = - para[:K_P] * (x_bs[3] - para[:P])
dx = similar(x_bs)

using BenchmarkTools
@btime $f_bs($dx, $x_bs, $es, $ed, nothing, $t)
@btime $f_pd($dx, $x_pd, $es, $ed, nothing, $t)
# it seems like the `IONode` is even a bit faster.


## types for  diffusionmap calculation

struct Diffusionmap
    data::Matrix
    kernel::AbstractKernel
    laplace_type::AbstractLaplacian
    threshold::Int64
end

Diffusionmap(data; kernel::AbstractKernel=InverseDistanceKernel(), laplace_type::AbstractLaplacian=RowNormalizedLaplacian(), threshold::Int64=size(data,1)) = Diffusionmap(data, kernel, laplace_type, threshold)
Diffusionmap(data, func::Function; laplace_type::AbstractLaplacian=RowNormalizedLaplacian(), threshold::Int64=size(data,1)) = Diffusionmap(data, CustomKernel(func), laplace_type, threshold)

function Base.show(io::IO, ::MIME"text/plain", dp::Diffusionmap)
    println(io, "Diffusionmap")
    println(io, "Kernel = $(typeof(dp.kernel))")
    println(io, "Laplace type = $(dp.laplace_type)")
    println(io, "Threshold = $(dp.threshold)")
end


################################################################################
# this implements the right pseudoinverse
# is defined if A has linearly independent columns
# ⁻¹, ⁺ syntax
struct PseudoInverse{T, M} <: Factorization{T}
    parent::M
end
PseudoInverse(A) = PseudoInverse{eltype(A), typeof(A)}(A)

Base.size(P::PseudoInverse) = size(P.parent')
Base.size(P::PseudoInverse, k::Integer) = size(P.parent', k::Integer)
function Base.AbstractMatrix(P::PseudoInverse)
	A = P.parent
    A isa Number ? fill(inv(A), 1, 1) : (A \ I(size(A, 1)))
end
function Base.Matrix(Inv::PseudoInverse)
	M = AbstractMatrix(Inv)
	M isa Matrix ? M : Matrix(M) # since it could be e.g. a Diagonal
end
# Base.Matrix(P::PseudoInverse) = AbstractMatrix(P)
LinearAlgebra.factorize(P::PseudoInverse) = P # same reasoning as for Inverse

# smart constructor
# calls regular inverse if matrix is square
function pseudoinverse end
const pinverse = pseudoinverse
function pseudoinverse(A::AbstractMatOrFac)
    size(A, 1) == size(A, 2) ? inverse(A) : PseudoInverse(A)
end
pseudoinverse(A::Union{Number, UniformScaling}) = inv(A)
pseudoinverse(P::PseudoInverse) = P.parent
pseudoinverse(P::Inverse) = inverse(P) # in case P.parent is scalar


using EarthMoversDistance: emd!, libemd_emd, emd, emd_flows

W1 = [0.4, 0.2, 0.2, 0.1, 0.1]
W2 = [0.6, 0.2, 0.2]
Cost1 = ones(Float64, 5, 3)

@testset "Calling the C function emd directly" begin
    flows = zeros(Float64, 5, 3)
    distance = libemd_emd(length(W1), W1, length(W2), W2, Cost1, flows)
    @test distance == 1.0
    @test isapprox(flows[1, 1], 0.4)
    @test isapprox(flows[2, 1], 0.2)
    @test isapprox(flows[3, 2], 0.2)
    @test isapprox(flows[4, 3], 0.1)
    @test isapprox(flows[5, 3], 0.1)
    @test Cost1 == ones(Float64, 5, 3)
end

@testset "emd" begin
    distance = emd(W1, W2, Cost1)
    @test distance == 1.0
    @test Cost1 == ones(Float64, 5, 3)
end

@testset "emd_flows" begin
    distance, flows = emd_flows(W1, W2, Cost1)
    @test distance == 1.0
    @test isapprox(flows[1, 1], 0.4)
    @test isapprox(flows[2, 1], 0.2)
    @test isapprox(flows[3, 2], 0.2)
    @test isapprox(flows[4, 3], 0.1)
    @test isapprox(flows[5, 3], 0.1)
    @test Cost1 == ones(Float64, 5, 3)
end

@testset "python-emd example1" begin
    # This corresponds to the example in the file: https://github.com/pdinges/python-emd/blob/master/example1.py

    # features1 = [Feature(100, 40, 22), Feature(211, 20, 2),
    #             Feature(32, 190, 150), Feature(2, 100, 100)]
    # weights1  = [0.4, 0.3, 0.2, 0.1]
    features1 = [[100, 40, 22], [211, 20, 2], [32, 190, 150], [2, 100, 100]]
    weights1  = [0.4, 0.3, 0.2, 0.1]

    #features2 = [Feature(0, 0, 0), Feature(50, 100, 80), Feature(255, 255, 255)]
    #weights2  = [0.5, 0.3, 0.2]
    features2 = [[0, 0, 0], [50, 100, 80], [255, 255, 255]]
    weights2  = [0.5, 0.3, 0.2]

    # The costs are the euclidean distances between all pairs
    cost = zeros(Float64, length(features1), length(features2))
    for i in 1:length(features1)
        for j in 1:length(features2)
            cost[i, j] = sqrt(sum((features1[i] .- features2[j]).^2))
        end
    end

    distance = emd(weights1, weights2, cost)
    @show distance
    @test isapprox(distance, 160.542; atol = 1e-3)
end


## Min-cut algorithms using max-flow subroutine

# Min-cut from bfs on classical max-flow
function mincut(g::Network, request = :EdmondsKarp)
  if request == :EdmondsKarp
    val, flows = edmondsKarp(g)
  end
#  println(flows)
  cut = bfs(g,:MinCut,flows)
  return val, cut
end

# α_k of a cut
function partSum(g::Network,cutedges::Vector{Int},k::Int)
  Σ = Vector{Float64}()
  for i in cutedges
    push!(Σ,g.links[i].cap)
  end
  sort!(Σ)
  return sum(Σ[1:end-k])
end

function bestMinkCut(g::Network,cuts::Vector{Vector{Int}},k::Int)
  function aux(cutedges)
    partSum(g,cutedges,k)
  end
  return minimum(map(aux,cuts))
end


#ENO coefficients for uniform mesh
#Based on: Shu C. Essentially Non-Oscillatory and Weighted Essentially
#Non-Oscillatory Schemes for Hyperbolic Conservation Laws
function unif_crj(k::Int)
  if k == 1
      return([1;1])
  end
  crj = zeros(k+1,k)
  for i = 1:(k+1)
    r = i-2
    for j = 1:k
      psum = 0.0
      for m = j:k
        numer = 0.0
        for l = 0:k
          if (l!=m)
            numer_prod = 1.0
            for q = 0:k
              if (q != m && q != l)
                numer_prod = numer_prod*(r-q+1)
              end
            end
            numer = numer + numer_prod
          end
        end
        denom = 1.0
        for l = 0:k
          if (l!=m);denom = denom * (m-l);end
        end
        psum = psum + numer/denom
      end
      crj[i,j] = psum
    end
  end
  crj
end

# Reconstruction types
struct ENO_Reconstruction{kType} <: AbstractReconstruction
    order::kType
    crj::Matrix{Float64}
end
ENO_Reconstruction(order::Int) = ENO_Reconstruction(order, unif_crj(Int((order + 1)/2)))

struct WENO_Reconstruction{kType,eType} <: AbstractReconstruction
    order::kType
    crj::Matrix{Float64}
    ɛ::eType
end
WENO_Reconstruction(order::Int; ɛ = 1e-6) = WENO_Reconstruction(order, unif_crj(Int((order + 1)/2)), ɛ)

struct MWENO_Reconstruction{kType,eType} <: AbstractReconstruction
    order::kType
    crj::Matrix{Float64}
    ɛ::eType
end
MWENO_Reconstruction(order::Int; ɛ = 1e-6) = MWENO_Reconstruction(order, unif_crj(Int((order + 1)/2)), ɛ)

#Eno reconstruction for uniform mesh
function reconstruct(vloc::AbstractVector, dx::Real, eno_rec::ENO_Reconstruction)
  order = eno_rec.order
  k = Int((order + 1)/2)
  crj = eno_rec.crj
  vdiffs = Vector{typeof(vloc)}(undef,0)
  vl = zero(eltype(vloc))
  vr = zero(eltype(vloc))
  N = size(vloc,1)
  if (N != 2*k-1)
    throw("dimension of vloc is not consistent with order $k ENO")
  end
  # Calculation of divided differences
  push!(vdiffs,copy(vloc))
  for i in 2:k
    push!(vdiffs,(vdiffs[i-1][2:end]-vdiffs[i-1][1:end-1])/dx)
  end
  #Calculation of Stencil
  r = 0
  if k < 2;return(vloc[1],vloc[1]);end
  for m = 2:k
    if (abs(vdiffs[m][k-r-1])<abs(vdiffs[m][k-r]))
      r = r + 1
    end
  end
  for j = 0:(k-1)
    vl = vl + vloc[k-r+j]*crj[r+1,j+1]
    vr = vr + vloc[k-r+j]*crj[r+2,j+1]
  end
  return(vl,vr)
end

#Eno reconstruction for unestructured mesh
function reconstruct(vloc::AbstractVector, xloc::AbstractVector, eno_rec::ENO_Reconstruction)
  order = eno_rec.order
  k = Int((order + 1)/2)
  crj = eno_rec.crj
  vdiffs = Vector{typeof(vloc)}()
  vl = zero(eltype(vloc))
  vr = zero(eltype(vloc))
  N = size(vloc,1)
  if (N != 2*k-1)
    throw("dimension of vloc is not consistent with order $k ENO")
  end
  # Calculation of divided differences
  push!(vdiffs,copy(vloc))
  for i in 2:k
    push!(vdiffs,(vdiffs[i-1][2:end]-vdiffs[i-1][1:end-1])./(xloc[i:N]-xloc[1:N-i+1]))
  end
  #Calculation of Stencil
  r = 0
  if k < 2;return(vloc[1],vloc[1]);end
  for m = 2:k
    if (abs(vdiffs[m][k-r-1])<abs(vdiffs[m][k-r]))
      r = r + 1
    end
  end
  for j = 0:(k-1)
    vl = vl + vloc[k-r+j]*crj[r+1,j+1]
    vr = vr + vloc[k-r+j]*crj[r+2,j+1]
  end
  return(vl,vr)
end

#WEno reconstruction for uniform mesh
#Order available: 1, 3, 5
# Jiang, Shu, Efficient Implementation of Weighted ENO Schemes
@inline function get_βk(vloc::Vector, k)
  βk = zeros(k); dr = zeros(k);
  if (k==2)
      dr = [2//3,1//3]
      βk = [(vloc[3]-vloc[2])^2, (vloc[2]-vloc[1])^2]
  elseif (k==3)
      dr = [3//10, 3//5, 1//10]
      βk = [13//12*(vloc[3]-2*vloc[4]+vloc[5])^2 + 1//4*(3*vloc[3]-4*vloc[4]+vloc[5])^2,
      13//12*(vloc[2]-2*vloc[3]+vloc[4])^2 + 1//4*(vloc[2]-vloc[4])^2,
      13//12*(vloc[1]-2*vloc[2]+vloc[3])^2 + 1//4*(3*vloc[3]-4*vloc[2]+vloc[1])^2]
  else
    throw("WENO reconstruction of order $order is not implemented yet!")
  end
  return βk, dr
end

function reconstruct(vloc::AbstractVector, rec_scheme::WENO_Reconstruction)
  order = rec_scheme.order
  k = Int((order + 1)/2)
  crj = rec_scheme.crj
  ɛ = rech_scheme.ɛ
  vl = zero(eltype(vloc))
  vr = zero(eltype(vloc))
  N = size(vloc,1)
  if (N != order)
    throw("dimension of vloc is not consistent with order $order WENO")
  end

  #Special case k = 1
  if (k == 1)
    vl = vloc[1]; vr = vloc[1]
    return vl, vr
  end

  # Apply WENO procedure
  αl = zeros(k); αr = zeros(k);
  ωl = zeros(k); ωr = zeros(k);
  βk = zeros(k); dr = zeros(k);

  # Compute k values of xl and xr based on different stencils
  ulr = zeros(k); urr = zeros(k);
  for r=0:(k-1)
      for i=0:k-1
          urr[r+1] = urr[r+1] + crj[r+2,i+1]*vloc[k-r+i];
          ulr[r+1] = ulr[r+1] + crj[r+1,i+1]*vloc[k-r+i];
      end
  end

  # Set up WENO coefficients for different orders - 2k-1
  βk, dr = get_βk(vloc, k)

  # Compute α parameters
  for r=1:k
      αr[r] = dr[r]/(ɛ+βk[r])^2;
      αl[r] = dr[k+1-r]/(ɛ+βk[r])^2;
  end

  # Compute wENO weights parameters
  for r=1:k
      ωl[r] = αl[r]/sum(αl);
      ωr[r] = αr[r]/sum(αr);
  end

  # Compute cell interface values
  for r=1:k
      vl = vl + ωl[r]*ulr[r];
      vr = vr + ωr[r]*urr[r];
  end
  return(vl,vr)
end

function reconstruct(vmloc::AbstractVector, vploc::AbstractVector,rec_scheme::WENO_Reconstruction)
  order = rec_scheme.order
  k = Int((order + 1)/2)
  crj = rec_scheme.crj
  ɛ = rec_scheme.ɛ
  vl = zero(eltype(vmloc))
  vr = zero(eltype(vploc))
  N = size(vmloc,1)
  if (N != order)
    throw("dimension of vloc is not consistent with order $order WENO")
  end
  #Special case k = 1
  if (k == 1)
    vl = vmloc[1]; vr = vploc[1]
    return vl, vr
  end

  # Apply WENO procedure
  αl = zeros(k); αr = zeros(k);
  ωl = zeros(k); ωr = zeros(k);
  βk = zeros(k); dr = zeros(k);

  # Compute k values of xl and xr based on different stencils
  ulr = zeros(k); urr = zeros(k);
  for r=0:(k-1)
      for i=0:k-1
          urr[r+1] = urr[r+1] + crj[r+2,i+1]*vploc[k-r+i];
          ulr[r+1] = ulr[r+1] + crj[r+1,i+1]*vmloc[k-r+i];
      end
  end

  # Set up WENO coefficients for different orders - 2k-1
  βrk, dr = get_βk(vploc, k)
  βlk, dr = get_βk(vmloc, k)

  # Compute α parameters
  for r=1:k
      αr[r] = dr[r]/(ɛ+βrk[r])^2;
      αl[r] = dr[k+1-r]/(ɛ+βlk[r])^2;
  end

  # Compute wENO weights parameters
  for r=1:k
      ωl[r] = αl[r]/sum(αl);
      ωr[r] = αr[r]/sum(αr);
  end

  # Compute cell interface values
  for r=1:k
      vl = vl + ωl[r]*ulr[r];
      vr = vr + ωr[r]*urr[r];
  end
  return(vl,vr)
end

#Mapped WEno reconstruction for uniform mesh
#Reference:
# A. Henrick, T. Aslam, J. Powers, Mapped weighted essentially non-oscillatory
# schemes: Achiving optimal order near critical points

#Mapping function
@inline function gk(ω::Vector, dr::Vector)
  g = zero(ω)
  for i in 1:size(ω,1)
    g[i] = ω[i]*(dr[i]+dr[i]^2-3*dr[i]*ω[i]+ω[i]^2)/(dr[i]^2+ω[i]*(1-2*dr[i]))
  end
  g
end

function reconstruct(vloc::AbstractVector, rec_scheme::MWENO_Reconstruction)
  order = rec_scheme.order
  k = Int((order + 1)/2)
  crj = rec_scheme.crj
  ɛ = rech_scheme.ɛ
  vl = zero(eltype(vloc))
  vr = zero(eltype(vloc))
  N = size(vloc,1)
  if (N != order)
    throw("dimension of vloc is not consistent with order $order WENO")
  end
  #Special case k = 1
  if (k == 1)
    vl = vloc[1]; vr = vloc[1]
    return vl, vr
  end

  # Apply WENO procedure
  αl = zeros(k); αr = zeros(k);
  ωl = zeros(k); ωr = zeros(k);
  αml = zeros(k); αmr = zeros(k);
  ωml = zeros(k); ωmr = zeros(k);
  βk = zeros(k); dr = zeros(k);

  # Compute k values of xl and xr based on different stencils
  ulr = zeros(k); urr = zeros(k);
  for r=0:(k-1)
      for i=0:k-1
          urr[r+1] = urr[r+1] + crj[r+2,i+1]*vloc[k-r+i];
          ulr[r+1] = ulr[r+1] + crj[r+1,i+1]*vloc[k-r+i];
      end
  end

  # Set up WENO coefficients for different orders - 2k-1
  if (k==2)
      dr = [2//3,1//3]
      βk = [(vloc[3]-vloc[2])^2, (vloc[2]-vloc[1])^2]
  elseif (k==3)
      dr = [3//10, 3//5, 1//10]
      βk = [13//12*(vloc[3]-2*vloc[4]+vloc[5])^2 + 1//4*(3*vloc[3]-4*vloc[4]+vloc[5])^2,
      13//12*(vloc[2]-2*vloc[3]+vloc[4])^2 + 1//4*(vloc[2]-vloc[4])^2,
      13//12*(vloc[1]-2*vloc[2]+vloc[3])^2 + 1//4*(3*vloc[3]-4*vloc[2]+vloc[1])^2]
  else
    throw("WENO reconstruction of order $order is not implemented yet!")
  end

  # Compute α parameters
  for r=1:k
      αr[r] = dr[r]/(ɛ+βk[r])^2;
      αl[r] = dr[k+1-r]/(ɛ+βk[r])^2;
  end

  # Compute wENO weights parameters
  for r=1:k
      ωl[r] = αl[r]/sum(αl);
      ωr[r] = αr[r]/sum(αr);
  end

  # Compute α mapped parameters
  αmr = gk(ωr)
  αml = gk(ωl)

  # Compute mapped wENO weights parameters
  for r=1:k
      ωml[r] = αml[r]/sum(αml);
      ωmr[r] = αmr[r]/sum(αmr);
  end

  # Compute cell interface values
  for r=1:k
      vl = vl + ωml[r]*ulr[r];
      vr = vr + ωmr[r]*urr[r];
  end
  return(vl,vr)
end

function reconstruct(vmloc::AbstractVector, vploc::AbstractVector,rec_scheme::MWENO_Reconstruction)
  order = rec_scheme.order
  k = Int((order + 1)/2)
  crj = rec_scheme.crj
  ɛ = rec_scheme.ɛ
  vl = zero(eltype(vmloc))
  vr = zero(eltype(vploc))
  N = size(vmloc,1)
  if (N != order)
    throw("dimension of vloc is not consistent with order $order WENO")
  end
  #Special case k = 1
  if (k == 1)
    vl = vmloc[1]; vr = vploc[1]
    return vl, vr
  end

  # Apply WENO procedure
  αl = zeros(k); αr = zeros(k);
  ωl = zeros(k); ωr = zeros(k);
  αml = zeros(k); αmr = zeros(k);
  ωml = zeros(k); ωmr = zeros(k);
  βk = zeros(k); dr = zeros(k);

  # Compute k values of xl and xr based on different stencils
  ulr = zeros(k); urr = zeros(k);
  for r=0:(k-1)
      for i=0:k-1
          urr[r+1] = urr[r+1] + crj[r+2,i+1]*vploc[k-r+i];
          ulr[r+1] = ulr[r+1] + crj[r+1,i+1]*vmloc[k-r+i];
      end
  end

  # Set up WENO coefficients for different orders - 2k-1
  βrk, dr = get_βk(vploc, k)
  βlk, dr = get_βk(vmloc, k)

  # Compute α parameters
  for r=1:k
      αr[r] = dr[r]/(ɛ+βrk[r])^2;
      αl[r] = dr[k+1-r]/(ɛ+βlk[r])^2;
  end

  # Compute wENO weights parameters
  for r=1:k
      ωl[r] = αl[r]/sum(αl);
      ωr[r] = αr[r]/sum(αr);
  end

  # Compute α mapped parameters
  drl = reverse(dr)
  αmr = gk(ωr,dr)
  αml = gk(ωl,drl)

  # Compute mapped wENO weights parameters
  for r=1:k
      ωml[r] = αml[r]/sum(αml);
      ωmr[r] = αmr[r]/sum(αmr);
  end

  # Compute cell interface values
  for r=1:k
      vl = vl + ωml[r]*ulr[r];
      vr = vr + ωmr[r]*urr[r];
  end
  return(vl,vr)
end


using PairwiseListMatrices
using DataFrames
using NamedArrays

using Test
using DelimitedFiles
using LinearAlgebra
using Statistics

@testset "PairwiseListMatrices" begin

    list = [1,-2,3]
    mat_diag_triu =  [ 1 -2
                        0 3 ]
    mat_diag_tril =  [ 1 0
                       -2 3 ]
    mat_triu = [ 0 1 -2
                  0 0 3
                  0 0 0 ]
    mat_tril = [ 0 0 0
                  1 0 0
                 -2 3 0 ]
    mat_diag_sym = [ 1 -2
                     -2 3 ]
    mat_sym = [ 0 1 -2
                 1 0 3
                -2 3 0 ]

    plmd_triu = UpperTriangular(PairwiseListMatrix(list, true))
    plmd_tril = LowerTriangular(PairwiseListMatrix(list, true))
    plm_triu  = UpperTriangular(PairwiseListMatrix(list))
    plm_tril  = LowerTriangular(PairwiseListMatrix(list))
    plmd_sym  = PairwiseListMatrix(list, true)
    plm_sym   = PairwiseListMatrix(list)

    @testset "Creation from list" begin

        @test plmd_triu == mat_diag_triu
        @test plmd_tril == mat_diag_tril
        @test plm_triu  == mat_triu
        @test plm_tril  == mat_tril
        @test plmd_sym  == mat_diag_sym
        @test plm_sym   == mat_sym
    end

    @testset "Matrix" begin

        @test Matrix(plmd_triu) == mat_diag_triu
        @test Matrix(plmd_tril) == mat_diag_tril
        @test Matrix(plm_triu ) == mat_triu
        @test Matrix(plm_tril ) == mat_tril
        @test Matrix(plmd_sym ) == mat_diag_sym
        @test Matrix(plm_sym  ) == mat_sym
    end

    @testset "Getters" begin

        @test getlist(plmd_sym) == plmd_sym.list
        @test getlist(plm_sym)  ==  plm_sym.list

        @test getdiag(plmd_sym) == plmd_sym.diag
        @test getdiag(plm_sym)  ==  plm_sym.diag
    end

    @testset "Length" begin

        @test lengthlist(2, true)  == lengthlist(plmd_sym)
        @test lengthlist(3, false) == lengthlist(plm_sym)

        @test length(plmd_sym) == 4
        @test length(plm_sym) == 9
    end

    @testset "Linear indexing" begin

        for i in 1:length(plm_sym)
            @test plm_sym[i] == mat_sym[i]
            plm_sym[i] = 100
            @test plm_sym[i] == 100
            plm_sym[i] = mat_sym[i]
        end
        for i in 1:length(plmd_sym)
            @test plmd_sym[i] == mat_diag_sym[i]
            plmd_sym[i] = 100
            @test plmd_sym[i] == 100
            plmd_sym[i] = mat_diag_sym[i]
        end

    end

    @testset "Convert" begin

        @test mat_diag_sym ==
                convert(PairwiseListMatrix{Int, true, Vector{Int}},  mat_diag_sym)
        @test mat_diag_sym ==
                convert(PairwiseListMatrix{Int, false, Vector{Int}}, mat_diag_sym)

        @test convert(Matrix{Float64}, mat_sym) ==
                convert(PairwiseListMatrix{Float64, true, Vector{Float64}},  mat_sym)
        @test convert(Matrix{Float64}, mat_sym) ==
                convert(PairwiseListMatrix{Float64, false, Vector{Float64}}, mat_sym)

        @test convert(Matrix{Int}, plm_sym)  == mat_sym
        @test convert(Matrix{Int}, plmd_sym) == mat_diag_sym

        @test convert(PairwiseListMatrix{Int, false, Vector{Int}}, plmd_sym) == plmd_sym
        @test convert(PairwiseListMatrix{Int, true, Vector{Int}}, plm_sym)   == plm_sym
    end

    @testset "Unary operations" begin

        for f in [ -, abs, x -> sqrt(abs(x))]
            @test broadcast(f, plmd_triu) == broadcast(f, mat_diag_triu)
            @test broadcast(f, plmd_tril) == broadcast(f, mat_diag_tril)
            @test broadcast(f, plm_triu ) == broadcast(f, mat_triu)
            @test broadcast(f, plm_tril ) == broadcast(f, mat_tril)
            @test broadcast(f, plmd_sym ) == broadcast(f, mat_diag_sym)
            @test broadcast(f, plm_sym  ) == broadcast(f, mat_sym)
        end
    end

    @testset "Binary operations" begin

        for f in [ -, +, * ]
            @test broadcast(f, plmd_triu, plmd_triu) == broadcast(f, mat_diag_triu, mat_diag_triu)
            @test broadcast(f, plmd_tril, plmd_tril) == broadcast(f, mat_diag_tril, mat_diag_tril)
            @test broadcast(f, plm_triu , plm_triu ) == broadcast(f, mat_triu,      mat_triu     )
            @test broadcast(f, plm_tril , plm_tril ) == broadcast(f, mat_tril,      mat_tril     )
            @test broadcast(f, plmd_sym , plmd_sym ) == broadcast(f, mat_diag_sym,  mat_diag_sym )
            @test broadcast(f, plm_sym  , plm_sym  ) == broadcast(f, mat_sym,       mat_sym      )
        end

        @test plmd_sym ./ plmd_sym == mat_diag_sym ./ mat_diag_sym

        # https://github.com/JuliaLang/julia/issues/19615
        for f in [ *, / ]
            @test broadcast(f, plmd_triu, 2) == broadcast(f, mat_diag_triu, 2)
            @test broadcast(f, plmd_tril, 2) == broadcast(f, mat_diag_tril, 2)
            @test broadcast(f, plm_triu , 2) == broadcast(f, mat_triu,      2)
            @test broadcast(f, plm_tril , 2) == broadcast(f, mat_tril,      2)
        end

        for f in [ -, +, *, / ]
            @test broadcast(f, plmd_sym , 2) == broadcast(f, mat_diag_sym,  2)
            @test broadcast(f, plm_sym  , 2) == broadcast(f, mat_sym,       2)
        end

        @testset "./ with Float64 & Int" begin

            plm_t = PairwiseListMatrix([.5, .4, .3], true)
            plm_f = PairwiseListMatrix([1., 1., 1.], false)
            fill!(plm_f, 1.0)
            mat_t = Matrix(plm_t)
            mat_f = Matrix(plm_f)

            for value in  (4, 4.0)
                result =  plm_t ./ value
                @test isa(result, PairwiseListMatrix{Float64,true,Vector{Float64}})
                @test result == (mat_t ./ value)

                result =  value ./ plm_t
                @test isa(result, PairwiseListMatrix{Float64,true,Vector{Float64}})
                @test result == (value ./ mat_t)

                result =  plm_f ./ value
                @test isa(result, PairwiseListMatrix{Float64,false,Vector{Float64}})
                @test result == (mat_f ./ value)

                result =  value ./ plm_f
                @test isa(result, PairwiseListMatrix{Float64,false,Vector{Float64}})
                @test result == (value ./ mat_f)
            end
        end

        # plm = PairwiseListMatrix(rand(500500), true)
        # @btime broadcast(Base.:/::typeof(/), $plm, 10.0);
        # @btime broadcast(Base.:/::typeof(/), $plm.list, 10.0);
        # @btime 0.5 .* $plm ./ 10.0;
        # @btime 0.5 .* $plm.list ./ 10.0;

    end

    @testset "Inplace Broadcast" begin

        for (p,m) in [(plmd_sym, mat_diag_sym), (plm_sym, mat_sym)]
            cp = deepcopy(p)
            cm = deepcopy(m)
            cp .+= 1.0
            cm .+= 1.0
            @test cp == cm
            cp .+= cp
            cm .+= cm
            @test cp == cm
        end
    end

    @testset "Transpose" begin

        @test transpose(plmd_triu) == plmd_tril == transpose(mat_diag_triu)
        @test transpose(plmd_tril) == plmd_triu == transpose(mat_diag_tril)
        @test transpose(plm_triu) == plm_tril == transpose(mat_triu)
        @test transpose(plm_tril) == plm_triu == transpose(mat_tril)

        @test adjoint(plmd_triu) == plmd_tril == adjoint(mat_diag_triu)
        @test adjoint(plmd_tril) == plmd_triu == adjoint(mat_diag_tril)
        @test adjoint(plm_triu) == plm_tril == adjoint(mat_triu)
        @test adjoint(plm_tril) == plm_triu == adjoint(mat_tril)

        @test transpose(plmd_sym) == plmd_sym == transpose(mat_diag_sym)
        @test transpose(plm_sym) == plm_sym == transpose(mat_sym)

        @test adjoint(plmd_sym) == plmd_sym == adjoint(mat_diag_sym)
        @test adjoint(plm_sym) == plm_sym == adjoint(mat_sym)

        @test transpose!(plm_sym) == plm_sym
        @test adjoint!(plm_sym) == plm_sym
    end

    @testset "Linear algebra" begin

        @test plmd_triu * plmd_triu == mat_diag_triu * mat_diag_triu
        @test plmd_tril * plmd_tril == mat_diag_tril * mat_diag_tril
        @test plm_triu  * plm_triu  == mat_triu      * mat_triu
        @test plm_tril  * plm_tril  == mat_tril      * mat_tril
        @test plmd_sym  * plmd_sym  == mat_diag_sym  * mat_diag_sym
        @test plm_sym   * plm_sym   == mat_sym       * mat_sym
        @test Symmetric(plm_sym) * Symmetric(plm_sym) == Symmetric(mat_sym) * Symmetric(mat_sym)

        @test plmd_triu / plmd_triu == mat_diag_triu / mat_diag_triu
        @test plmd_tril / plmd_tril == mat_diag_tril / mat_diag_tril
        @test plmd_sym  / plmd_sym  == mat_diag_sym  / mat_diag_sym
        @test plm_sym   / plm_sym   == mat_sym       / mat_sym

        for p in [:U, :S, :Vt]
            @test getproperty(svd(plmd_triu), p) ≈ getproperty(svd(mat_diag_triu), p)
            @test getproperty(svd(plmd_tril), p) ≈ getproperty(svd(mat_diag_tril), p)
            @test getproperty(svd(plm_triu ), p) ≈ getproperty(svd(mat_triu), p)
            @test getproperty(svd(plm_tril ), p) ≈ getproperty(svd(mat_tril), p)
            @test getproperty(svd(plmd_sym ), p) ≈ getproperty(svd(mat_diag_sym), p)
            @test getproperty(svd(plm_sym  ), p) ≈ getproperty(svd(mat_sym), p)
        end
    end

    @testset "Stats" begin

        for f in [ mean, std ]
            @test f(plmd_triu) == f(mat_diag_triu)
            @test f(plmd_tril) == f(mat_diag_tril)
            @test f(plm_triu ) == f(mat_triu)
            @test f(plm_tril ) == f(mat_tril)
            @test f(plmd_sym ) == f(mat_diag_sym)
            @test f(plm_sym  ) == f(mat_sym)

            @test f(plmd_sym, dims=1) == f(mat_diag_sym, dims=1)
            @test f(plm_sym , dims=1) == f(mat_sym, dims=1)
            @test f(plmd_sym, dims=2) == f(mat_diag_sym, dims=2)
            @test f(plm_sym , dims=2) == f(mat_sym, dims=2)
        end

        @test cor(plmd_sym) ≈ cor(mat_diag_sym)
        @test cor(plm_sym ) ≈ cor(mat_sym)

        # broadcast used in cor(...)
        @test plmd_sym .- mean(plmd_sym,dims=1) == mat_diag_sym .- mean(mat_diag_sym,dims=1)
        @test plmd_sym .- mean(plmd_sym,dims=2) == mat_diag_sym .- mean(mat_diag_sym,dims=2)
    end

    @testset "triu" begin

        @test triu(plmd_sym ) == triu(mat_diag_sym)
        @test triu(plm_sym  ) == triu(mat_sym)

        @test triu(plmd_sym,1) == triu(mat_diag_sym, 1)
        @test triu(plm_sym, 1) == triu(mat_sym, 1)

        @test_throws ErrorException triu!(plm_sym)
        @test_throws ErrorException triu!(plm_sym, 1)
    end

    # + and - definitions are ambiguous with:
    # +(A::AbstractArray{Bool,N<:Any}, x::Bool)
    #
    # @testset "Boolean binary op" begin
    #
    #     list = PairwiseListMatrix([true, true, true])
    #
    #     @test list - true == [ -1  0  0
    #                                  0 -1  0
    #                                  0  0 -1 ]
    #
    #     @test list + true == [ 1  2  2
    #                                 2  1  2
    #                                 2  2  1 ]
    # end

    @testset "Mean without diagonal" begin

        diag_false = PairwiseListMatrix([10,20,30], false)
        diag_true  = PairwiseListMatrix([0,10,20,0,30,0], true)

        for list in [diag_false, diag_true]
            @test vec( mean_nodiag(list, dims=1) ) == [15., 20., 25.]
            @test vec( mean_nodiag(list, dims=2) ) == [15., 20., 25.]
            @test mean_nodiag(list) == 20.
        end
    end
end

@testset "Named PairwiseListMatrix" begin

    list =  [1,2,3]
    labels= ["a", "b", "c"]
    labels_diag = ["A", "B"]

    plm_diag = PairwiseListMatrix(list, true)
    plm = PairwiseListMatrix(list)
    nplm_diag = setlabels(plm_diag, labels_diag)
    nplm = setlabels(plm, labels)

    @testset "set and get labels" begin

        @test isa(nplm_diag, NamedArray)
        @test isa(nplm, NamedArray)
        @test isa(nplm_diag.array, PairwiseListMatrix)
        @test isa(nplm.array, PairwiseListMatrix)

        @test_throws AssertionError setlabels!(nplm_diag, labels)
        @test_throws AssertionError setlabels!(nplm, labels_diag)

        @test getlabels(nplm_diag) == labels_diag
        @test getlabels(nplm) == labels

        setlabels!(nplm, ["A","B","C"])
        @test getlabels(nplm) == ["A","B","C"]
        setlabels!(nplm, ["a","b","c"])
    end

    @testset "get/setindex" begin

        for i in 1:2
            for j in 2:3
                @test nplm[labels[i], labels[j]] == nplm[labels[j], labels[i]]
            end
        end

        @test nplm_diag["A", "B"] == nplm_diag["B", "A"]

        nplm_diag["A", "B"] = 20
        @test nplm_diag["B", "A"] == 20
        nplm_diag["B", "A"] = 2
        @test nplm_diag["A", "B"] == 2
    end

    @testset "join" begin

        left = setlabels(PairwiseListMatrix(collect(1.:15.), true),
            String["a","b","c","d","e"])
        right = setlabels(PairwiseListMatrix(collect(1.:10.), false),
            String["a","e","i","o","u"])

        a, b = join(left, right)
        @test (a, b) == join(left, right, kind=:inner)
        @test getlabels(a) == getlabels(b)
        @test getlabels(a) == ["a", "e"]
        @test size(a) == (2,2)
        @test size(b) == (2,2)

        a, b = join(left, right, kind=:left)
        @test getlabels(a) == getlabels(b)
        @test getlabels(a) == getlabels(left)
        @test a == left
        @test b != right

        a, b = join(left, right, kind=:right)
        @test getlabels(a) == getlabels(b)
        @test getlabels(b) == getlabels(right)
        @test a != left
        @test b == right

        a, b = join(left, right, kind=:outer)
        @test getlabels(a) == getlabels(b)
        @test getlabels(a) == ["a", "b", "c", "d", "e", "i", "o", "u"]
        @test size(a) == (8,8)
        @test size(b) == (8,8)
        @test a != left
        @test b != right
    end

    @testset "Named and not named arrays" begin

        @testset "Symmetric" begin

            @test issymmetric(plm)
            @test issymmetric(nplm)
            @test issymmetric(plm_diag)
            @test issymmetric(nplm_diag)
        end

        @testset "Diagonal" begin

            @test hasdiagonal(plm_diag)
            @test hasdiagonal(nplm_diag)
            @test !hasdiagonal(plm)
            @test !hasdiagonal(nplm)

            @test diagonal(plm_diag) == [1,3]
            @test diagonal(nplm_diag) == [1,3]
            @test diagonal(plm) == [0,0,0]
            @test diagonal(nplm) == [0,0,0]
        end

        @testset "eltype" begin

            @test eltype(plm) == Int
            @test eltype(nplm) == Int
            @test eltype(plm_diag) == Int
            @test eltype(nplm_diag) == Int
        end

        @testset "delegated functions" begin

            for F in (getlist, getdiag, Matrix, lengthlist, sum_nodiag, mean_nodiag, diagonal)
                    @test F(nplm) == F(plm)
                    @test F(nplm_diag) == F(plm_diag)
            end
        end

        @testset "map and broadcast" begin

            mat = Matrix(nplm)
            mat_diag = Matrix(nplm_diag)

            @test map(sqrt,plm) == map(sqrt,mat)
            @test map(sqrt,nplm) == map(sqrt,mat)
            @test map(sqrt,plm_diag) == map(sqrt,mat_diag)
            @test map(sqrt,nplm_diag) == map(sqrt,mat_diag)

            @test broadcast(sqrt,plm) == broadcast(sqrt,mat)
            @test broadcast(sqrt,nplm) == broadcast(sqrt,mat)
            @test broadcast(sqrt,plm_diag) == broadcast(sqrt,mat_diag)
            @test broadcast(sqrt,nplm_diag) == broadcast(sqrt,mat_diag)
        end
    end
end

@testset "Vector{PairwiseListMatrix}" begin

    list_samples = [ PairwiseListMatrix(rand(10), true) for i in 1:100 ]
    list_diag_samples = [ PairwiseListMatrix(rand(6), false) for i in 1:100 ]
    full_samples = Matrix{Float64}[ Matrix(mat) for mat in list_samples ]
    full_diag_samples = Matrix{Float64}[ Matrix(mat) for mat in list_diag_samples ]

    for f in [sum, mean]
        @test f(list_samples) == f(full_samples)
        @test f(list_diag_samples) == f(full_diag_samples)
    end

    @test sum(list_samples[1:1]) == list_samples[1]
    @test sum(list_diag_samples[1:1]) == list_diag_samples[1]

    @test_throws ErrorException sum(PairwiseListMatrix{Int, true, Array{Int,1}}[])
    @test_throws ErrorException sum(PairwiseListMatrix{Int, true, Array{Int,1}}[
                                        PairwiseListMatrix([1, 1, 1], true),
                                        PairwiseListMatrix([1, 1, 1, 1, 1, 1], true) ])
    @test_throws ErrorException std(PairwiseListMatrix{Int, true, Array{Int,1}}[
                                        PairwiseListMatrix([1, 1, 1], true) ])

    full_samples = reshape(hcat(full_samples...), (4,4,100))
    full_diag_samples = reshape(hcat(full_diag_samples...), (4,4,100))
    # @test std(list_samples) ≈ std(full_samples,3)[:,:,1] # It uses n - 1
    @test std(list_samples) ≈ sqrt.(mean(abs2.(full_samples .- mean(full_samples, dims=3)), dims=3))
    # @test std(list_diag_samples) ≈ std(full_diag_samples,3)[:,:,1] # It uses n - 1
    @test std(list_diag_samples) ≈ sqrt.(mean(abs2.(full_diag_samples .- mean(full_diag_samples, dims=3)), dims=3))

    @testset "Z score" begin

        list = PairwiseListMatrix{Float64, false, Vector{Float64}}[
            PairwiseListMatrix([1., 1., 1.]),
            PairwiseListMatrix([3., 3., 3.]) ]
        mat = PairwiseListMatrix([2., 2., 2.])

        @test std(list) == PairwiseListMatrix([1., 1., 1.])
        @test mean(list) == mat
        @test PairwiseListMatrices.zscore(list, mat) == PairwiseListMatrix([0., 0., 0.])

        @test_throws MethodError PairwiseListMatrices.zscore(list, PairwiseListMatrix([2.,2.,2.], true))
        @test_throws ErrorException PairwiseListMatrices.zscore(list, PairwiseListMatrix([2.,2.,2.,2.,2.,2.]))

        list = PairwiseListMatrix{Float64, true, Vector{Float64}}[
            PairwiseListMatrix([1., 1., 1.], true),
            PairwiseListMatrix([3., 3., 3.], true) ]
        mat = PairwiseListMatrix([2., 2., 2.], true)

        @test std(list) == PairwiseListMatrix([1., 1., 1.], true)
        @test mean(list) == mat
        @test PairwiseListMatrices.zscore(list, mat) == PairwiseListMatrix([0., 0., 0.], true)

        @test_throws MethodError PairwiseListMatrices.zscore(list,PairwiseListMatrix([2.], false))
        @test_throws ErrorException PairwiseListMatrices.zscore(list,PairwiseListMatrix([2.,2.,2.,2.,2.,2.],true))
    end
end

@testset "Apply" begin

    PLM = PairwiseListMatrix([1,2,3], false)

    apply2list(PLM) do list, k
        Main.@test(list[k] == k)
    end

    list_values = [1,2,3,4,5,6]
    PLMtrue  = PairwiseListMatrix(list_values, true)
    PLMfalse = PairwiseListMatrix(list_values, false)
    full_t   = Matrix(PLMtrue)
    full_f   = Matrix(PLMfalse)

    apply2list(PLMtrue) do list, k
        @test(list[k] == list_values[k])
    end
    apply2list(PLMfalse) do list, k
        @test(list[k] == list_values[k])
    end

    @iteratelist PLMtrue  Main.@test(list[k] == list_values[k])
    @iteratelist PLMfalse Main.@test(list[k] == list_values[k])

    apply2diag(PLMfalse) do diag, k
        @test(diag[k] == 0)
    end

    apply2upper(PLMtrue, use_diag=true) do list, k, i, j
        list[k] = list_values[k]
    end
    apply2upper(PLMfalse, use_diag=false) do list, k, i, j
        list[k] = list_values[k]
    end

    apply2upper(PLMtrue, use_diag=true) do list, k, i, j
        list[k] = full_t[i, j]
    end
    apply2upper(PLMtrue, use_diag=false) do list, k, i, j
        list[k] = full_t[i, j]
    end

    apply2upper(PLMfalse, use_diag=true) do list, k, i, j
        list[k] = full_f[i, j]
    end
    apply2upper(PLMfalse, use_diag=false) do list, k, i, j
        list[k] = full_f[i, j]
    end

    @iterateupper PLMtrue  true  list[k] = list_values[k]
    @iterateupper PLMfalse false list[k] = list_values[k]

    @iterateupper PLMtrue  true  list[k] = full_t[i,j]
    @iterateupper PLMtrue  false list[k] = full_t[i,j]

    @iterateupper PLMtrue  true  list[k] = full_f[i,j]
    @iterateupper PLMfalse false list[k] = full_f[i,j]
end

@testset "IO" begin

    @testset "to/from table" begin

        table = [ "A" "B" 10
                  "A" "C" 20
                  "B" "C" 30 ]

        list = setlabels(PairwiseListMatrix([10,20,30]), ["A", "B", "C"])
        list_diag = setlabels(PairwiseListMatrix([0,10,20,0,30,0], true), ["A", "B", "C"])

        @test to_table(list, diagonal=false) == table
        @test from_table(table, false, diagonalvalue=0) == list

        @test to_table(list) == [ "A" "A" 0
                                  "A" "B" 10
                                  "A" "C" 20
                                  "B" "B" 0
                                  "B" "C" 30
                                  "C" "C" 0]

        @test to_table(list) == to_table(list_diag)
        @test to_table(list, diagonal=false) == to_table(list_diag, diagonal=false)

        @testset "README" begin
            # using Any type instead of Int to allow @test
            plm  = PairwiseListMatrix(Any[1,2,3], false)
            nplm  = setlabels(plm, ["a","b","c"])
            table = to_table(nplm)
            new_nplm = from_table(table, true)
            @test new_nplm == plm
        end
    end

    @testset "DataFrames" begin

        values   = [1,2,3,4,5,6]
        PLMtrue  = PairwiseListMatrix(values, true)
        PLMfalse = PairwiseListMatrix(values, false)

        df = DataFrame(to_dict(PLMtrue))
        @test PLMtrue == from_table(df, true)

        df = DataFrame(to_dict(PLMtrue, diagonal=false))
        @test triu(PLMtrue,1) == triu(from_table(df, false).array,1)

        df = DataFrame(to_dict(PLMfalse))
        @test PLMfalse == from_table(df, true)

        df = DataFrame(to_dict(PLMfalse, diagonal=false))
        @test PLMfalse == from_table(df, false)
    end

    @testset "Write" begin

        values   = [1,2,3,4,5,6]
        PLMtrue  = PairwiseListMatrix(values, true)
        PLMfalse = PairwiseListMatrix(values, false)
        filename = joinpath(tempdir(), "pairwiselistmatrices_test.temp")

        try
          writedlm(filename, PLMtrue, diagonal=true)
          @test map(string,readdlm(filename,Int)) == map(string,to_table(PLMtrue,diagonal=true))
        finally
          rm(filename)
        end

        try
          writedlm(filename,PLMtrue,diagonal=false)
          @test map(string,readdlm(filename,Int)) == map(string,to_table(PLMtrue,diagonal=false))
        finally
          rm(filename)
        end

        try
          writedlm(filename,PLMfalse,diagonal=true)
          @test map(string,readdlm(filename,Int)) == map(string,to_table(PLMfalse,diagonal=true))
        finally
          rm(filename)
        end

        try
          writedlm(filename,PLMfalse,diagonal=false)
          @test map(string,readdlm(filename,Int)) == map(string,to_table(PLMfalse,diagonal=false))
        finally
          rm(filename)
        end

    end
end


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# ╔═╡ ecf99396-5e16-11ec-1735-c5307d7bf94a
data = [(rand(), rand(), rand()) for _ in 1:200]

# ╔═╡ 026cc2c9-059b-4d59-a95d-e0d884629b84
begin
	using Ripserer
	ripserer(data)
end

# ╔═╡ fca029d7-51f9-4ae2-96eb-d0aa5f086304
begin
	using Plots
	result=ripserer(data,dim_max=2)
	plot(plot(result),barcode(result))
end

# ╔═╡ Cell order:
# ╠═ecf99396-5e16-11ec-1735-c5307d7bf94a
# ╠═026cc2c9-059b-4d59-a95d-e0d884629b84
# ╠═fca029d7-51f9-4ae2-96eb-d0aa5f086304


module TestArithmetic

using Test
using CategoricalDistributions
using StableRNGs
rng = StableRNG(123)

L = ["yes", "no"]
d1 = UnivariateFinite(L, rand(rng, 2), pool=missing)
d2 = UnivariateFinite(L, rand(rng, 2), pool=missing)

@testset "arithmetic" begin

    # addition and subtraction:
    for op in [:+, :-]
        quote
            s = $op(d1, d2 )
            @test $op(pdf.(d1, L), pdf.(d2, L)) ≈ pdf.(s, L)
        end |> eval
    end

    # negative:
    d_neg = -d1
    @test pdf.(d_neg, L) == -pdf.(d1, L)

    # multiplication by scalar:
    d3 = d1*42
    @test pdf.(d3, L) ≈ pdf.(d1, L)*42
    d3 = 42*d1
    @test pdf.(d3, L) ≈ pdf.(d1, L)*42

    # division by scalar:
    d3 = d1/42
    @test pdf.(d3, L) ≈ pdf.(d1, L)/42

    # "probabilities" that aren't:
    d = UnivariateFinite(L, randn(rng, 2) + im*randn(2), pool=missing)
    @test pdf.(3*d -(4*d)/2, L) ≈ pdf.(d, L)
end

p = [0.1, 0.9]
P = vcat(fill(p', 10^5)...);
slow = fill(UnivariateFinite(L, p, pool=missing), 10^5);
fast = UnivariateFinite(L, P, pool=missing);
# @assert pdf(slow, L) == pdf(fast, L)

@testset "performant arithmetic for UnivariateFiniteArray" begin
    @test pdf(slow + slow, L) == pdf(fast + fast, L)
    t_slow = @elapsed @eval slow + slow
    t_fast = @elapsed @eval fast + fast
    @test t_slow/t_fast > 10

    @test pdf(slow - slow, L) == pdf(fast - fast, L)
    t_slow = @elapsed @eval slow - slow
    t_fast = @elapsed @eval fast - fast
    @test t_slow/t_fast > 10

    @test pdf(42*slow, L) == pdf(42*fast, L)
    @test pdf(slow*42, L) == pdf(fast*42, L)
    t_slow = @elapsed @eval 42*slow
    t_fast = @elapsed @eval 42*fast
    @test t_slow/t_fast > 10
    t_slow = @elapsed @eval slow*42
    t_fast = @elapsed @eval fast*42
    @test t_slow/t_fast > 10

    @test pdf(slow/42, L) == pdf(fast/42, L)
    t_slow = @elapsed @eval slow/42
    t_fast = @elapsed @eval fast/42
    @test t_slow/t_fast > 10
end

end # module

true


using ImageQualityIndexes
using Test, ReferenceTests, TestImages
using Statistics

include("testutils.jl")

@testset "ImageQualityIndexes" begin

    include("psnr.jl")
    include("ssim.jl")
    include("colorfulness.jl")

end

nothing


"""
	Cart-pole
"""

struct Cartpole{I, T} <: Model{I, T}
	n::Int
    m::Int
	d::Int

    mc     # mass of the cart in kg
    mp     # mass of the pole (point mass at the end) in kg
    l      # length of the pole in m
    g      # gravity m/s^2
end

function f(model::Cartpole, x, u, w)
    H = @SMatrix [model.mc + model.mp model.mp * model.l * cos(x[2]);
				  model.mp * model.l * cos(x[2]) model.mp * model.l^2.0]
    C = @SMatrix [0.0 -1.0 * model.mp * x[4] * model.l * sin(x[2]);
	 			  0.0 0.0]
    G = @SVector [0.0,
				  model.mp * model.g * model.l * sin(x[2])]
    B = @SVector [1.0,
				  0.0]
    qdd = SVector{2}(-H \ (C * view(x, 3:4) + G - B * u[1]))

    return @SVector [x[3],
					 x[4],
					 qdd[1],
					 qdd[2]]
end

n, m, d = 4, 1, 4
model = Cartpole{Midpoint, FixedTime}(n, m, d, 1.0, 0.2, 0.5, 9.81)

"""
	Cart-pole with Coulomb friction
"""
struct CartpoleFriction{I, T} <: Model{I, T}
	n::Int
    m::Int
	d::Int

    mc     # mass of the cart in kg
    mp     # mass of the pole (point mass at the end) in kg
    l      # length of the pole in m
    g      # gravity m/s^2
	μ      # coefficient of friction
end

function f(model::CartpoleFriction, x, u, w)
    H = @SMatrix [model.mc + model.mp model.mp * model.l * cos(x[2]);
				  model.mp * model.l * cos(x[2]) model.mp * model.l^2.0]
    C = @SMatrix [0.0 -1.0 * model.mp * x[4] * model.l * sin(x[2]);
	 			  0.0 0.0]
    G = @SVector [0.0,
				  model.mp * model.g * model.l * sin(x[2])]
    B = @SVector [1.0,
				  0.0]
    qdd = SVector{2}(-H \ (C * view(x, 3:4) + G - B * (u[1] + u[2] - u[3])))

    return @SVector [x[3],
					 x[4],
					 qdd[1],
					 qdd[2]]
end

n, m, d = 4, 7, 4
model_friction = CartpoleFriction{Midpoint, FixedTime}(n, m, d, 1.0, 0.2, 0.5, 9.81, 0.1)

function visualize!(vis, model, q;
       Δt = 0.1,
	   color = RGBA(1,0,0,1.0))

	default_background!(vis)

	l2 = Cylinder(Point3f0(-model.l * 2.0, 0.0, 0.0),
		Point3f0(model.l * 2.0, 0.0, 0.0),
		convert(Float32, 0.025))

	setobject!(vis["slider"], l2, MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

    l1 = Cylinder(Point3f0(0.0, 0.0, 0.0),
		Point3f0(0.0, 0.0, model.l),
		convert(Float32, 0.025))

    setobject!(vis["arm"], l1,
		MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

    setobject!(vis["base"], HyperSphere(Point3f0(0.0),
        convert(Float32, 0.1)),
        MeshPhongMaterial(color = color))

    setobject!(vis["ee"], HyperSphere(Point3f0(0.0),
        convert(Float32, 0.05)),
        MeshPhongMaterial(color = color))

	anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

	for t = 1:length(q)
	    MeshCat.atframe(anim,t) do
			x = q[t]
			px = x[1] + model.l * sin(x[2])
			pz = -model.l * cos(x[2])
            settransform!(vis["arm"], cable_transform([x[1]; 0;0], [px; 0.0; pz]))
            settransform!(vis["base"], Translation([x[1]; 0.0; 0.0]))
            settransform!(vis["ee"], Translation([px; 0.0; pz]))
	    end
	end

	settransform!(vis["/Cameras/default"],
		compose(Translation(0.0, 0.0, -1.0), LinearMap(RotZ(- pi / 2))))

	MeshCat.setanimation!(vis,anim)
end


function assemble_cell_body_force!(
    systemrhs,
    rhsfunc,
    basis,
    cellquads,
    mesh,
    cellsign,
    cellid,
)

    dim = dimension(mesh)
    detjac = determinant_jacobian(mesh)
    cellmap = cell_map(mesh, cellsign, cellid)
    quad = cellquads[cellsign, cellid]
    rhs = linear_form(rhsfunc, basis, quad, cellmap, dim, detjac)
    nodeids = nodal_connectivity(mesh, cellsign, cellid)
    assemble_cell_rhs!(systemrhs, nodeids, dim, rhs)
end

function assemble_body_force!(systemrhs, rhsfunc, basis, cellquads, mesh)

    ncells = number_of_cells(mesh)
    for cellid = 1:ncells
        cellsign = cell_sign(mesh, cellid)
        check_cellsign(cellsign)

        if cellsign == +1 || cellsign == 0
            assemble_cell_body_force!(
                systemrhs,
                rhsfunc,
                basis,
                cellquads,
                mesh,
                +1,
                cellid,
            )
        end
        if cellsign == -1 || cellsign == 0
            assemble_cell_body_force!(
                systemrhs,
                rhsfunc,
                basis,
                cellquads,
                mesh,
                -1,
                cellid,
            )
        end
    end
end


let S = ZZ{BigInt}, P = S[:x], x = P([0, 1])

@testset "construction of quotient ring" begin
    ideal = x^2 + 1
    @test P / ideal <: Quotient{P}
    @test P / (x^3+1) <: Quotient{P}
    Q = P / ideal
    Qp = P / (x^3+1)

    @test basetype(Q) == P
    @test depth(Q) == 3
    @test P / ideal != nothing
    p = 4x + 5
    @test Qp(p) == p
    @test Qp(1) == 1
    @test Qp(ZZ(1)) == 1
    q = Q(p)
    @test q + q == Q(8x + 10)
    @test q - q == 0
    @test q * q == p * p
    @test 3q == Q(12x + 15)
    pp = (4im + 5)^10
    @test q^10 == Q(P([pp.re, pp.im]))

    @test !iszero(q)
    @test iszero(zero(Q))
    @test !isone(q)
    @test isone(one(Q))
    @test isunit(Q(x))
    @test !isunit(Q(x+1))
    @test inv(Q(x)) == -Q(x)
    @test Q(x) * 5 == Q(5x)

    @test hash(q) != 0
    @test sprint(show, q) isa String

    @test value(q + 3 * ideal) == p
end

@testset "quotient of polynomial over GF" begin
    G = GF(2, 2)
    GX = G[:x]
    x = monom(GX)
    p = x^3 + 1
    Q = GX / p
    @test order(Q) == order(G)^3
end

@testset "show" begin
    P = ZZ{Int}[:x]
    x = monom(P)
    Q = P / (x^2 + 17)
    @test sprint(show, Q(1)) == "1"
    @test sprint(show, Q(x)) == "\u23b7-17"
    @test sprint(show, Q(-2x + 10)) == "-2*\u23b7-17 + 10"
    @test sprint(show, -Q(2x + 10)) == "-2*\u23b7-17 - 10"
    R = P / (x^3 + 1)
    @test sprint(show, R(x^2 + x + 1)) == "x^2 + x + 1 mod(x^3 + 1)"
end

end