### A Pluto.jl notebook ###
# v0.12.21

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

# ╔═╡ f7a6d7c3-37b9-437d-8b8e-853665ddbae3
filter!(LOAD_PATH) do path
	path != "@v#.#"
end;

# ╔═╡ 65780f00-ed6b-11ea-1ecf-8b35523a7ac0
begin
	import Pkg
	Pkg.activate(mktempdir())
	Pkg.add([
			Pkg.PackageSpec(name="Images", version="0.22.4"),
			Pkg.PackageSpec(name="ImageMagick", version="0.7"),
			Pkg.PackageSpec(name="PlutoUI", version="0.7"),
			Pkg.PackageSpec(name="HypertextLiteral", version="0.5"),
			Pkg.PackageSpec(name="OffsetArrays"),
			])

	using Images
	using PlutoUI
	using HypertextLiteral
	using OffsetArrays
end

# ╔═╡ 83eb9ca0-ed68-11ea-0bc5-99a09c68f867
md"_homework 2, version 1_"

# ╔═╡ ac8ff080-ed61-11ea-3650-d9df06123e1f
md"""

# **Homework 2** - _convolutions_
`18.S191`, Spring 2021

`Due Date`: **Friday Mar 5, 2021 at 11:59pm EST**

This notebook contains _built-in, live answer checks_! In some exercises you will see a coloured box, which runs a test case on your code, and provides feedback based on the result. Simply edit the code, run it, and the check runs again.

_For MIT students:_ there will also be some additional (secret) test cases that will be run as part of the grading process, and we will look at your notebook and write comments.

Feel free to ask questions!
"""

# ╔═╡ 911ccbce-ed68-11ea-3606-0384e7580d7c
# edit the code below to set your name and kerberos ID (i.e. email without @mit.edu)

student = (name = "Vidhi Katkoria", kerberos_id = "vidhi")

# press the ▶ button in the bottom right of this cell to run your edits
# or use Shift+Enter

# you might need to wait until all other cells in this notebook have completed running.
# scroll down the page to see what's up

# ╔═╡ 8ef13896-ed68-11ea-160b-3550eeabbd7d
md"""

Submission by: **_$(student.name)_** ($(student.kerberos_id)@mit.edu)
"""

# ╔═╡ 5f95e01a-ee0a-11ea-030c-9dba276aba92
md"_Let's create a package environment:_"

# ╔═╡ e08781fa-ed61-11ea-13ae-91a49b5eb74a
md"""

## **Exercise 1** - _Convolutions in 1D_

As we have seen in the lectures, we can produce cool effects using the mathematical technique of **convolutions**. We input one image $M$ and get a new image $M'$ back.

Conceptually we think of $M$ as a matrix. In practice, in Julia it will be a `Matrix` of color objects, and we may need to take that into account. Ideally, however, we should write a **generic** function that will work for any type of data contained in the matrix.

A convolution works on a small **window** of an image, i.e. a region centered around a given point $(i, j)$. We will suppose that the window is a square region with odd side length $2\ell + 1$, running from $-\ell, \ldots, 0, \ldots, \ell$.

The result of the convolution over a given window, centred at the point $(i, j)$ is a *single number*; this number is the value that we will use for $M'_{i, j}$.
(Note that neighbouring windows overlap.)

To get started, in Exercise 1 we'll restrict ourselves to convolutions in 1D.
So a window is just a 1D region from $-\ell$ to $\ell$.

"""

# ╔═╡ 80108d80-ee09-11ea-0368-31546eb0d3cc
md"""
#### Exercise 1.1

Let's create a vector `v` of random numbers of length `n=100`.
"""

# ╔═╡ 7fcd6230-ee09-11ea-314f-a542d00d582e
n = 100

# ╔═╡ 7fdb34dc-ee09-11ea-366b-ffe10d1aa845
v = rand(n)

# ╔═╡ 7fe9153e-ee09-11ea-15b3-6f24fcc20734
md"_Feel free to experiment with different values!_

Let's use the function `colored_line` to view this 1D number array as a 1D image.
"

# ╔═╡ ff70782e-e8d2-4281-9b24-d45c925f55e2
begin
	colored_line(x::Vector) = hcat(Gray.(Float64.(x)))'
	colored_line(x::Any) = nothing
end

# ╔═╡ 01070e28-ee0f-11ea-1928-a7919d452bdd
colored_line(v)

# ╔═╡ 7522f81e-ee1c-11ea-35af-a17eb257ff1a
md"👉 Try changing `n` and `v` around. Notice that you can run the cell `v = rand(n)` again to regenerate new random values."

# ╔═╡ 801d90c0-ee09-11ea-28d6-61b806de26dc
md"""
#### Exercise 1.2
We need to decide how to handle the **boundary conditions**, i.e. what happens if we try to access a position in the vector `v` beyond `1:n`.  The simplest solution is to assume that $v_{i}$ is 0 outside the original vector; however, this may lead to strange boundary effects.

A better solution is to use the *closest* value that is inside the vector. Effectively we are extending the vector and copying the extreme values into the extended positions. (Indeed, this is one way we could implement this; these extra positions are called **ghost cells**.)

👉 Write a function `extend(v, i)` that checks whether the position $i$ is inside `1:n`. If so, return the $(HTML("<br>")) ``i``th component of `v`; otherwise, return the nearest end value.
"""

# ╔═╡ 802bec56-ee09-11ea-043e-51cf1db02a34
function extend(v::AbstractVector, i)
	println(i)
	if i<1
		return v[1]
	end
	if i>length(v)
		return v[length(v)]
	end
	return v[i]
end

# ╔═╡ b7f3994c-ee1b-11ea-211a-d144db8eafc2
md"_Some test cases:_"

# ╔═╡ 3492b164-7065-48e8-978b-6c96b965d376
example_vector = [0.8, 0.2, 0.1, 0.7, 0.6, 0.4]

# ╔═╡ 02123165-2a0a-49a8-b7a9-458955523511
colored_line(example_vector)

# ╔═╡ 806e5766-ee0f-11ea-1efc-d753cd83d086
md"- Extended with 0:"

# ╔═╡ 38da843a-ee0f-11ea-01df-bfa8b1317d36
colored_line([0, 0, example_vector..., 0, 0])

# ╔═╡ 9bde9f92-ee0f-11ea-27f8-ffef5fce2b3c
md"- Extended with your `extend` function:"

# ╔═╡ 431ba330-0f72-416a-92e9-55f51ff3bcd1
md"""
#### Exercise 1.3
👉 Write (or copy) the `mean` function from Homework 1, which takes a vector and returns the mean.

"""

# ╔═╡ 5fdc5d0d-a52c-476e-b3b5-3b6364b706e4
function mean(v)
	return sum(v)/length(v)
end

# ╔═╡ e84c9cc2-e6e1-46f1-bf4e-9605da5e6f4a
md"""

👉 Write a function `box_blur(v, l)` that blurs a vector `v` with a window of length `l` by averaging the elements within a window from $-\ell$ to $\ell$. This is called a **box blur**. Use your function `extend` to handle the boundaries correctly.

Return a vector of the same size as `v`.
"""

# ╔═╡ ca91bf1a-80ed-11eb-3be0-cd60c9adae9d
colored_line(example_vector)

# ╔═╡ 809f5330-ee09-11ea-0e5b-415044b6ac1f
md"""
#### Exercise 1.4
👉 Apply the box blur to your vector `v`. Show the original and the new vector by creating two cells that call `colored_line`. Make the parameter $\ell$ interactive, and call it `l_box` instead of `l` to avoid a naming conflict.
"""

# ╔═╡ e555a7e6-f11a-43ac-8218-6d832f0ce251
@bind l_box Slider(1:20)

# ╔═╡ 302f0842-453f-47bd-a74c-7942d8c96485
colored_line(v)

# ╔═╡ 80ab64f4-ee09-11ea-29b4-498112ed0799
md"""
#### Exercise 1.5
The box blur is a simple example of a **convolution**, i.e. a linear function of a window around each point, given by

$$v'_{i} = \sum_{m}  \, v_{i - m} \, k_{m},$$

where $k$ is a vector called a **kernel**.

Again, we need to take care about what happens if $v_{i -m }$ falls off the end of the vector.

👉 Write a function `convolve(v, k)` that performs this convolution. You need to think of the vector $k$ as being *centred* on the position $i$. So $m$ in the above formula runs between $-\ell$ and $\ell$, where $2\ell + 1$ is the length of the vector $k$.

   You will either need to do the necessary manipulation of indices by hand, or use the `OffsetArrays.jl` package.
"""

# ╔═╡ cf73f9f8-ee12-11ea-39ae-0107e9107ef5
md"_Edit the cell above, or create a new cell with your own test cases!_"

# ╔═╡ fa463b71-5aa4-44a3-a67b-6b0776236243
md"""
#### Exercise 1.6

👉 Define a function `box_blur_kernel(l)` which returns a _kernel_ (i.e. a vector) which, when used as the kernel in `convolve`, will reproduce a box blur of length `l`.
"""

# ╔═╡ 8a7d3cfd-6f19-43f0-ae16-d5a236f148e7
function box_blur_kernel(l)

	kernel_size = 2*l+1
	box_blur_kernel = fill(1/kernel_size, kernel_size)
	return box_blur_kernel
end

# ╔═╡ a34d1ad8-3776-4bc4-93e5-72cfffc54f15
@bind box_kernel_l Slider(1:5)

# ╔═╡ 971a801d-9c46-417a-ad31-1144894fb4e1
box_blur_kernel_test = box_blur_kernel(box_kernel_l)

# ╔═╡ 5f13b1a5-8c7d-47c9-b96a-a09faf38fe5e
md"""
Let's apply your kernel to our test vector `v` (first cell), and compare the result to our previous box blur function (second cell). The two should be identical.
"""

# ╔═╡ 03f91a22-1c3e-4c42-9d78-1ee36851a120
md"""
#### Exercise 1.7
👉 Write a function `gaussian_kernel`.

The definition of a Gaussian in 1D is

$$G(x) = \frac{1}{\sqrt{2\pi \sigma^2}} \exp \left( \frac{-x^2}{2\sigma^2} \right),$$

or as a Julia function:
"""

# ╔═╡ 48530f0d-49b4-4aec-8109-d69f1ef7f0ee
md"""
Write a function `gauss` that takes `σ` as a keyword argument and implements this function.
"""

# ╔═╡ beb62fda-38a6-4528-a176-cfb726f4b5bd
gauss(x::Real; σ=1) = 1 / sqrt(2π*σ^2) * exp(-x^2 / (2 * σ^2))

# ╔═╡ f0d55cec-2e81-4cbb-b166-2cf4f2a0f43f
md"""
We need to **sample** (i.e. evaluate) this at each pixel in an interval of length $2n+1$,
and then **normalize** so that the sum of the resulting kernel is 1.
"""

# ╔═╡ f8bd22b8-ee14-11ea-04aa-ab16fd01826e
md"""
You can edit the cell above to test your kernel function!

Let's try applying it in a convolution.
"""

# ╔═╡ 2a9dd06a-ee13-11ea-3f84-67bb309c77a8
@bind gaussian_kernel_size_1D Slider(0:6)

# ╔═╡ 8f70cbfc-80f6-11eb-078c-337a02409c6b
x = [0.343, 0.23121, 1.0, 1.0]

# ╔═╡ 6f74c790-80f6-11eb-3200-717715b41e3d
colored_line(x)

# ╔═╡ ce24e486-df27-4780-bc57-d3bf7bee83bb
function create_bar()
	x = zeros(100)
	x[41:60] .= 1
	x
end

# ╔═╡ b01858b6-edf3-11ea-0826-938d33c19a43
md"""

## **Exercise 2** - _Convolutions in 2D_

Now let's move to 2D images. The convolution is then given by a **kernel matrix** $K$:

$$M'_{i, j} = \sum_{k, l}  \, M_{i- k, j - l} \, K_{k, l},$$

where the sum is over the possible values of $k$ and $l$ in the window. Again we think of the window as being *centered* at $(i, j)$.

A common notation for this operation is $\star$:

```math
M' = M \star K
```
"""

# ╔═╡ 7c1bc062-ee15-11ea-30b1-1b1e76520f13
md"""
#### Exercise 2.1
👉 Write a new method for `extend` that takes a matrix `M` and indices `i` and `j`, and returns the closest element of the matrix.
"""

# ╔═╡ 7c2ec6c6-ee15-11ea-2d7d-0d9401a5e5d1
function extend(M::AbstractMatrix, i, j)

	i = max(i, 1)
	i = min(i, size(M,1))

	j = max(j, 1)
	j = min(j, size(M,2))

	return M[i,j]
end

# ╔═╡ d87175b6-7db0-11eb-08a4-456a34fa137e
extend([1],0)

# ╔═╡ 803905b2-ee09-11ea-2d52-e77ff79693b0
extend([5,6,7], 1)

# ╔═╡ 80479d98-ee09-11ea-169e-d166eef65874
extend([5,6,7], -8)

# ╔═╡ 805691ce-ee09-11ea-053d-6d2e299ee123
extend([5,6,7], 10)

# ╔═╡ 45c4da9a-ee0f-11ea-2c5b-1f6704559137
if extend(v,1) === missing
	missing
else
	a = [extend(example_vector, i) for i in -1:length(example_vector)+2]
	colored_line(a)
end

# ╔═╡ 807e5662-ee09-11ea-3005-21fdcc36b023
function box_blur(v::AbstractArray, l)
	blurred_image = [mean(extend(v, j) for j in i-l:i+l) for i in 1:length(v)]
	return blurred_image
end

# ╔═╡ 4f08ebe8-b781-4a32-a218-5ecd8338561d
colored_line(box_blur(example_vector, 2))

# ╔═╡ 808deca8-ee09-11ea-0ee3-1586fa1ce282
let
	try
		test_v = rand(n)
		original = copy(test_v)
		box_blur(test_v, 5)
		if test_v != original
			md"""
			!!! danger "Oopsie!"
			    It looks like your function _modifies_ `v`. Can you write it without doing so? Maybe you can use `copy`.
			"""
		end
	catch
	end
end

# ╔═╡ 7d80a1ea-a0a9-41b2-9cfe-a334717ab2f4
colored_line(box_blur(v,l_box))

# ╔═╡ bbe1a562-8d97-4112-a88a-c45c260f574d
let
	result = box_blur(v, box_kernel_l)
	colored_line(result)
end

# ╔═╡ 28e20950-ee0c-11ea-0e0a-b5f2e570b56e
function convolve(v::AbstractVector, k)
	l = (length(k)-1)÷2
	offset_k = OffsetArray(k,-l:l)
	convolved_vector = [sum(offset_k[j]*extend(v, i+j) for j in -l:l) for i in 1:length(v)]
	return convolved_vector
end

# ╔═╡ 9afc4dca-ee16-11ea-354f-1d827aaa61d2
md"_Let's test it!_"

# ╔═╡ cf6b05e2-ee16-11ea-3317-8919565cb56e
small_image = Gray.(rand(5,5))

# ╔═╡ e3616062-ee27-11ea-04a9-b9ec60842a64
md"- Extended with `0`:"

# ╔═╡ e5b6cd34-ee27-11ea-0d60-bd4796540b18
[get(small_image, (i, j), Gray(0)) for (i,j) in Iterators.product(-1:7,-1:7)]

# ╔═╡ b4e98589-f221-4922-b11e-364d72d0788e

# ╔═╡ d06ea762-ee27-11ea-2e9c-1bcff86a3fe0
md"- Extended with your `extend` function:"

# ╔═╡ e1dc0622-ee16-11ea-274a-3b6ec9e15ab5
[extend(small_image, i, j) for (i,j) in Iterators.product(-1:7,-1:7)]

# ╔═╡ 4bbea325-35f8-4a51-bd66-153aba4aed96
md"""
### Extending Philip
"""

# ╔═╡ c4f5a867-74ba-4106-91d4-195f6ae644d0
url = "https://user-images.githubusercontent.com/6933510/107239146-dcc3fd00-6a28-11eb-8c7b-41aaf6618935.png"

# ╔═╡ c825ebe2-511b-43ba-afdf-6226dbac48d2
philip_filename = download(url) # download to a local file. The filename is returned

# ╔═╡ 2701ab0c-b91d-47fe-b36b-4e0036ecd4aa
philip = load(philip_filename);

# ╔═╡ 84a48984-9adb-40ab-a1f1-1ab7b76c9a19
philip_head = philip[470:800, 140:410];

# ╔═╡ 3cd535e4-ee26-11ea-2482-fb4ad43dda19
[
	extend(philip_head, i, j) for
		i in -50:size(philip_head,1)+51,
		j in -50:size(philip_head,2)+51
]

# ╔═╡ 7c41f0ca-ee15-11ea-05fb-d97a836659af
md"""
#### Exercise 2.2
👉 Implement a new method `convolve(M, K)` that applies a convolution to a 2D array `M`, using a 2D kernel `K`. Use your new method `extend` from the last exercise.
"""

# ╔═╡ 8b96e0bc-ee15-11ea-11cd-cfecea7075a0
function convolve(M::AbstractMatrix, K::AbstractMatrix)
	rows, cols = size(K)
	r = (rows-1)÷2
	c = (cols-1)÷2
	offset_kernel = OffsetArray(K, -r:r, -c:c)
	convolved_matrix = [sum([extend(M, i+m, j+n) * offset_kernel[i,j] for i in -r:r, j in -c:c]) for m in 1:size(M,1), n in 1:size(M,2)]

	return convolved_matrix
end

# ╔═╡ 93284f92-ee12-11ea-0342-833b1a30625c
test_convolution = let
	v = [1, 10, 100, 1000, 10000]
	k = [1, 1, 0]
	convolve(v, k)
end

# ╔═╡ 5eea882c-ee13-11ea-0d56-af81ecd30a4a
colored_line(test_convolution)

# ╔═╡ 338b1c3f-f071-4f80-86c0-a82c17349828
let
	result = convolve(v, box_blur_kernel_test)
	colored_line(result)
end

# ╔═╡ 5a5135c6-ee1e-11ea-05dc-eb0c683c2ce5
md"_Let's test it out! 🎃_"

# ╔═╡ 577c6daa-ee1e-11ea-1275-b7abc7a27d73
test_image_with_border = [get(small_image, (i, j), Gray(0)) for (i,j) in Iterators.product(-1:7,-1:7)]

# ╔═╡ 275a99c8-ee1e-11ea-0a76-93e3618c9588
K_test = [
	0   0  0
	1/2 0  1/2
	0   0  0
]

# ╔═╡ 42dfa206-ee1e-11ea-1fcd-21671042064c
new_image = convolve(test_image_with_border, K_test)

# ╔═╡ 2ac28130-810c-11eb-2cb1-117579142578
typeof(new_image)

# ╔═╡ 1d3b604a-810c-11eb-27fe-538cba860476
typeof(test_image_with_border)

# ╔═╡ 9ce383c8-810b-11eb-0c4c-fd184b9f6e3a
typeof(philip_head)

# ╔═╡ 4739eae4-810a-11eb-07f8-0b47904d9efe
convolve(philip_head, K_test)

# ╔═╡ 6e53c2e6-ee1e-11ea-21bd-c9c05381be07
md"_Edit_ `K_test` _to create your own test case!_"

# ╔═╡ e7f8b41a-ee25-11ea-287a-e75d33fbd98b
[philip_head convolve(philip_head, K_test)]

# ╔═╡ 8a335044-ee19-11ea-0255-b9391246d231
md"""
---

You can create all sorts of effects by choosing the kernel in a smart way. Today, we will implement two special kernels, to produce a **Gaussian blur** and a **Sobel edge detection** filter.

Make sure that you have watched the lecture about convolutions!
"""

# ╔═╡ 79eb0775-3582-446b-996a-0b64301394d0
md"""
#### Exercise 2.3
The 2D Gaussian kernel will be defined using

$$G(x,y)=\frac{1}{2\pi \sigma^2}\exp\left(\frac{-(x^2+y^2)}{2\sigma^2}\right)$$

How can you express this mathematically using the 1D Gaussian function that we defined before?
"""

# ╔═╡ f4d9fd6f-0f1b-4dec-ae68-e61550cee790
gauss(x, y; σ=1) = 2π*σ^2 * gauss(x; σ=σ) * gauss(y; σ=σ)

# ╔═╡ 1c8b4658-ee0c-11ea-2ede-9b9ed7d3125e
function gaussian_kernel_1D(n; σ = 1)
	gaussian_kernel = [gauss(x,σ) for x in -n:n]
	return gaussian_kernel./sum(gaussian_kernel)
end

# ╔═╡ a6149507-d5ba-45c1-896a-3487070d36ec
colored_line(gaussian_kernel_1D(4; σ=5))

# ╔═╡ 38eb92f6-ee13-11ea-14d7-a503ac04302e
test_gauss_1D_a = let
	k = gaussian_kernel_1D(gaussian_kernel_size_1D)

	if k !== missing
		convolve(x, k)
	end
end

# ╔═╡ b424e2aa-ee14-11ea-33fa-35491e0b9c9d
colored_line(test_gauss_1D_a)

# ╔═╡ 24c21c7c-ee14-11ea-1512-677980db1288
test_gauss_1D_b = let
	v = create_bar()
	k = gaussian_kernel_1D(gaussian_kernel_size_1D)

	if k !== missing
		convolve(v, k)
	end
end

# ╔═╡ bc1c20a4-ee14-11ea-3525-63c9fa78f089
colored_line(test_gauss_1D_b)

# ╔═╡ 7c50ea80-ee15-11ea-328f-6b4e4ff20b7e
md"""
👉 Write a function that applies a **Gaussian blur** to an image. Use your previous functions, and add cells to write helper functions as needed!
"""

# ╔═╡ 48ebad62-8107-11eb-34e9-09890a16e6ea
function gaussian_kernel_2D(n, σ)
	gaussian_kernel = [gauss(x, y, σ=σ) for x in -n:n, y in -n:n]
	return gaussian_kernel./sum(gaussian_kernel)
end

# ╔═╡ aad67fd0-ee15-11ea-00d4-274ec3cda3a3
function with_gaussian_blur(image; σ=3, l=5)
	return convolve(image, gaussian_kernel_2D(l, σ))
end

# ╔═╡ 8ae59674-ee18-11ea-3815-f50713d0fa08
md"_Let's make it interactive. 💫_"

# ╔═╡ 96146b16-79ea-401f-b8ba-e05663a18bd8
@bind face_σ Slider(0.1:0.1:10; show_value=true)

# ╔═╡ 2cc745ce-e145-4428-af3b-926fba271b67
@bind face_l Slider(0:20; show_value=true)

# ╔═╡ a75701c4-ee18-11ea-2863-d3042e71a68b
gaussian_blurred_image = with_gaussian_blur(philip_head; σ=face_σ, l=face_l)

# ╔═╡ 9cb8c8a0-810f-11eb-375b-7b4de6432ebd
[philip_head gaussian_blurred_image]

# ╔═╡ 364a0a42-8110-11eb-3907-e9cfc000a302
typeof(gaussian_blurred_image)

# ╔═╡ d5ffc6ab-156b-4d43-ac3d-1947d0176e7f
md"""
When you set `face_σ` to a low number (e.g. `2.0`), what effect does `face_l` have? And vice versa?
"""

# ╔═╡ 7c6642a6-ee15-11ea-0526-a1aac4286cdd
md"""
#### Exercise 2.4
👉 Create a **Sobel edge detection filter**.

Here, we will need to create two filters that separately detect edges in the horizontal and vertical directions, given by the following kernels:

```math
G_x = \begin{bmatrix}
1 & 0 & -1 \\
2 & 0 & -2 \\
1 & 0 & -1 \\
\end{bmatrix};
\qquad
G_y = \begin{bmatrix}
1 & 2 & 1 \\
0 & 0 & 0 \\
-1 & -2 & -1 \\
\end{bmatrix}
```

We can think of these filterrs as derivatives in the $x$ and $y$ directions, as we discussed in lectures.

Then we combine them by finding the magnitude of the **gradient** (in the sense of multivariate calculus) by defining

$$G_\text{total} = \sqrt{G_x^2 + G_y^2},$$

where each operation applies *element-wise* on the matrices.

Use your previous functions, and add cells to write helper functions as needed!
"""

# ╔═╡ 9767e240-8110-11eb-3869-cbfed5e5918d
function horizontal_edge_detect_kernel()
	return [-0.125  -0.25  -0.125;
  0.0     0.0    0.0;
  0.125   0.25   0.125]
end

# ╔═╡ fe8eb9d0-8110-11eb-0fe6-d91ef91f892e
function vertical_edge_detect_kernel()
	return [1 2 1 ; 0 0 0 ; -1 -2 -1]
end

# ╔═╡ 9eeb876c-ee15-11ea-1794-d3ea79f47b75
function with_sobel_edge_detect(image)
	kernel_x = horizontal_edge_detect_kernel()
	kernel_y = vertical_edge_detect_kernel()

	x_edges = convolve(image, kernel_x)
	y_edges = convolve(image,kernel_y)

	y_edges

	edges_sqrt = sqrt.(x_edges.^2 + y_edges.^2)
	edges = x_edges + y_edges
	[edges_sqrt edges]
	# typeof(edges)
	# return edges
end

# ╔═╡ 525e8402-818c-11eb-1f8b-05cd978cffb3
Gray.(with_sobel_edge_detect(Gray.(philip_head)))

# ╔═╡ 260dc208-8113-11eb-05c9-33108c830d89
[Gray.(philip_head) Gray.(with_sobel_edge_detect(Gray.(philip_head)))]

# ╔═╡ d2063f66-818b-11eb-3762-e734ed57af8f
sobel_hor = Kernel.sobel()[1]

# ╔═╡ 10b2e046-818c-11eb-0f8f-af52231e7fca
convolve(philip_head, sobel_hor)

# ╔═╡ f864d0b4-818b-11eb-2c48-7b334976ae22
typeof(horizontal_edge_detect_kernel())

# ╔═╡ 8ffe16ce-ee20-11ea-18bd-15640f94b839
if student.kerberos_id === "jazz"
	md"""
!!! danger "Oops!"
    **Before you submit**, remember to fill in your name and kerberos ID at the top of this notebook!
	"""
end

# ╔═╡ 2d9f3ae4-9e4c-49ce-aab0-5f87aba85c3e
md"## Function library

Just some helper functions used in the notebook."

# ╔═╡ 5516c800-edee-11ea-12cf-3f8c082ef0ef
hint(text) = Markdown.MD(Markdown.Admonition("hint", "Hint", [text]))

# ╔═╡ ea435e58-ee11-11ea-3785-01af8dd72360
hint(md"Have a look at the lecture notes to see examples of adding interactivity with a slider. You can read the Interactivity and the PlutoUI sample notebooks to learn more, you can find them in Pluto's main menu. _(Right click the Pluto logo in the top left -> Open in new tab)_.")

# ╔═╡ 32a07f1d-93cd-4bf3-bac1-91afa6bb88a6
md"""
You can use the `÷` operator (you type `\div<TAB>` to get it with autocomplete) to do _integer division_. For example:

```julia
8 / 6 ≈ 1.3333333 # a floating point number!

8 // 6 == 4 // 3  # a fraction!

8 ÷ 6 == 1        # an integer!
```
""" |> hint

# ╔═╡ 649df270-ee24-11ea-397e-79c4355e38db
hint(md"`num_rows, num_columns = size(M)`")

# ╔═╡ 0cabed84-ee1e-11ea-11c1-7d8a4b4ad1af
hint(md"`num_rows, num_columns = size(K)`")

# ╔═╡ 9def5f32-ee15-11ea-1f74-f7e6690f2efa
hint(md"Can we just copy the 1D code? What is different in 2D?")

# ╔═╡ 57360a7a-edee-11ea-0c28-91463ece500d
almost(text) = Markdown.MD(Markdown.Admonition("warning", "Almost there!", [text]))

# ╔═╡ dcb8324c-edee-11ea-17ff-375ff5078f43
still_missing(text=md"Replace `missing` with your answer.") = Markdown.MD(Markdown.Admonition("warning", "Here we go!", [text]))

# ╔═╡ 58af703c-edee-11ea-2963-f52e78fc2412
keep_working(text=md"The answer is not quite right.") = Markdown.MD(Markdown.Admonition("danger", "Keep working on it!", [text]))

# ╔═╡ f3d00a9a-edf3-11ea-07b3-1db5c6d0b3cf
yays = [md"Great!", md"Yay ❤", md"Great! 🎉", md"Well done!", md"Keep it up!", md"Good job!", md"Awesome!", md"You got the right answer!", md"Let's move on to the next exercise."]

# ╔═╡ 5aa9dfb2-edee-11ea-3754-c368fb40637c
correct(text=rand(yays)) = Markdown.MD(Markdown.Admonition("correct", "Got it!", [text]))

# ╔═╡ f0c3e99d-9eb9-459e-917a-c2338af6683c
let
	result = gaussian_kernel_1D(5)

	if ismissing(result)
		still_missing()
	elseif isnothing(result)
		keep_working(md"Did you forget to write `return`?")
	elseif !(result isa AbstractVector)
		keep_working(md"The returned object is not a `Vector`.")
	elseif size(result) != (11,)
		hint(md"The returned vector has the wrong dimensions.")
	elseif !(sum(result) ≈ 1.0)
		keep_working(md"You need to _normalize_ the result.")
	elseif gaussian_kernel_1D(3; σ=1) == gaussian_kernel_1D(3; σ=2)
		keep_working(md"Use the keyword argument `σ` in your function.")
	else
		correct()
	end
end

# ╔═╡ 74d44e22-edee-11ea-09a0-69aa0aba3281
not_defined(variable_name) = Markdown.MD(Markdown.Admonition("danger", "Oopsie!", [md"Make sure that you define a variable called **$(Markdown.Code(string(variable_name)))**"]))

# ╔═╡ bcf98dfc-ee1b-11ea-21d0-c14439500971
if !@isdefined(extend)
	not_defined(:extend)
else
	let
		result = extend([6,7],-10)

		if ismissing(result)
			still_missing()
		elseif isnothing(result)
			keep_working(md"Did you forget to write `return`?")
		elseif result != 6 || extend([6,7],10) != 7
			keep_working()
		else
			correct()
		end
	end
end

# ╔═╡ 7ffd14f8-ee1d-11ea-0343-b54fb0333aea
if !@isdefined(convolve)
	not_defined(:convolve)
else
	let
		x = [1, 10, 100]
		result = convolve(x, [0, 1, 1])
		shouldbe = [11, 110, 200]
		shouldbe2 = [2, 11, 110]

		if ismissing(result)
			still_missing()
		elseif isnothing(result)
			keep_working(md"Did you forget to write `return`?")
		elseif !(result isa AbstractVector)
			keep_working(md"The returned object is not a `Vector`.")
		elseif size(result) != size(x)
			keep_working(md"The returned vector has the wrong dimensions.")
		elseif result != shouldbe && result != shouldbe2
			keep_working()
		else
			correct()
		end
	end
end

# ╔═╡ d93fa3f6-c361-4dfd-a2ea-f38e682bcd6a
if !@isdefined(box_blur_kernel)
	not_defined(:box_blur_kernel)
else
	let
		result = box_blur_kernel(2)

		if ismissing(result)
			still_missing()
		elseif isnothing(result)
			keep_working(md"Did you forget to write `return`?")
		elseif !(result isa AbstractVector)
			keep_working(md"The returned object is not a `Vector`.")
		elseif size(result) != (5,)
			hint(md"The returned vector has the wrong dimensions.")
		else

			x = [1, 10, 100]
			result1 = box_blur(x, 2)
			result2 = convolve(x, result)

			if result1 ≈ result2
				correct()
			else
				keep_working()
			end
		end
	end
end

# ╔═╡ efd1ceb4-ee1c-11ea-350e-f7e3ea059024
if !@isdefined(extend)
	not_defined(:extend)
else
	let
		input = [42 37; 1 0]
		result = extend(input, -2, -2)

		if ismissing(result)
			still_missing()
		elseif isnothing(result)
			keep_working(md"Did you forget to write `return`?")
		elseif result != 42 || extend(input, -1, 3) != 37
			keep_working()
		else
			correct()
		end
	end
end

# ╔═╡ 115ded8c-ee0a-11ea-3493-89487315feb7
bigbreak = html"<br><br><br><br><br>";

# ╔═╡ 54056a02-ee0a-11ea-101f-47feb6623bec
bigbreak

# ╔═╡ a3067222-a83a-47b8-91c3-24ad78dd65c5
bigbreak

# ╔═╡ 27847dc4-ee0a-11ea-0651-ebbbb3cfd58c
bigbreak

# ╔═╡ 0001f782-ee0e-11ea-1fb4-2b5ef3d241e2
bigbreak

# ╔═╡ 5842895a-ee10-11ea-119d-81e4c4c8c53b
bigbreak

# ╔═╡ dfb7c6be-ee0d-11ea-194e-9758857f7b20
function camera_input(;max_size=200, default_url="https://i.imgur.com/SUmi94P.png")
"""
<span class="pl-image waiting-for-permission">
<style>

	.pl-image.popped-out {
		position: fixed;
		top: 0;
		right: 0;
		z-index: 5;
	}

	.pl-image #video-container {
		width: 250px;
	}

	.pl-image video {
		border-radius: 1rem 1rem 0 0;
	}
	.pl-image.waiting-for-permission #video-container {
		display: none;
	}
	.pl-image #prompt {
		display: none;
	}
	.pl-image.waiting-for-permission #prompt {
		width: 250px;
		height: 200px;
		display: grid;
		place-items: center;
		font-family: monospace;
		font-weight: bold;
		text-decoration: underline;
		cursor: pointer;
		border: 5px dashed rgba(0,0,0,.5);
	}

	.pl-image video {
		display: block;
	}
	.pl-image .bar {
		width: inherit;
		display: flex;
		z-index: 6;
	}
	.pl-image .bar#top {
		position: absolute;
		flex-direction: column;
	}

	.pl-image .bar#bottom {
		background: black;
		border-radius: 0 0 1rem 1rem;
	}
	.pl-image .bar button {
		flex: 0 0 auto;
		background: rgba(255,255,255,.8);
		border: none;
		width: 2rem;
		height: 2rem;
		border-radius: 100%;
		cursor: pointer;
		z-index: 7;
	}
	.pl-image .bar button#shutter {
		width: 3rem;
		height: 3rem;
		margin: -1.5rem auto .2rem auto;
	}

	.pl-image video.takepicture {
		animation: pictureflash 200ms linear;
	}

	@keyframes pictureflash {
		0% {
			filter: grayscale(1.0) contrast(2.0);
		}

		100% {
			filter: grayscale(0.0) contrast(1.0);
		}
	}
</style>

	<div id="video-container">
		<div id="top" class="bar">
			<button id="stop" title="Stop video">✖</button>
			<button id="pop-out" title="Pop out/pop in">⏏</button>
		</div>
		<video playsinline autoplay></video>
		<div id="bottom" class="bar">
		<button id="shutter" title="Click to take a picture">📷</button>
		</div>
	</div>

	<div id="prompt">
		<span>
		Enable webcam
		</span>
	</div>

<script>
	// based on https://github.com/fonsp/printi-static (by the same author)

	const span = currentScript.parentElement
	const video = span.querySelector("video")
	const popout = span.querySelector("button#pop-out")
	const stop = span.querySelector("button#stop")
	const shutter = span.querySelector("button#shutter")
	const prompt = span.querySelector(".pl-image #prompt")

	const maxsize = $(max_size)

	const send_source = (source, src_width, src_height) => {
		const scale = Math.min(1.0, maxsize / src_width, maxsize / src_height)

		const width = Math.floor(src_width * scale)
		const height = Math.floor(src_height * scale)

		const canvas = html`<canvas width=\${width} height=\${height}>`
		const ctx = canvas.getContext("2d")
		ctx.drawImage(source, 0, 0, width, height)

		span.value = {
			width: width,
			height: height,
			data: ctx.getImageData(0, 0, width, height).data,
		}
		span.dispatchEvent(new CustomEvent("input"))
	}

	const clear_camera = () => {
		window.stream.getTracks().forEach(s => s.stop());
		video.srcObject = null;

		span.classList.add("waiting-for-permission");
	}

	prompt.onclick = () => {
		navigator.mediaDevices.getUserMedia({
			audio: false,
			video: {
				facingMode: "environment",
			},
		}).then(function(stream) {

			stream.onend = console.log

			window.stream = stream
			video.srcObject = stream
			window.cameraConnected = true
			video.controls = false
			video.play()
			video.controls = false

			span.classList.remove("waiting-for-permission");

		}).catch(function(error) {
			console.log(error)
		});
	}
	stop.onclick = () => {
		clear_camera()
	}
	popout.onclick = () => {
		span.classList.toggle("popped-out")
	}

	shutter.onclick = () => {
		const cl = video.classList
		cl.remove("takepicture")
		void video.offsetHeight
		cl.add("takepicture")
		video.play()
		video.controls = false
		console.log(video)
		send_source(video, video.videoWidth, video.videoHeight)
	}

	document.addEventListener("visibilitychange", () => {
		if (document.visibilityState != "visible") {
			clear_camera()
		}
	})

	// Set a default image

	const img = html`<img crossOrigin="anonymous">`

	img.onload = () => {
	console.log("helloo")
		send_source(img, img.width, img.height)
	}
	img.src = "$(default_url)"
	console.log(img)
</script>
</span>
""" |> HTML
end

# ╔═╡ 94c0798e-ee18-11ea-3212-1533753eabb6
@bind gauss_raw_camera_data camera_input(;max_size=100)

# ╔═╡ e7b91064-8112-11eb-06c1-cd57fd57133f
gauss_raw_camera_data

# ╔═╡ 1a0324de-ee19-11ea-1d4d-db37f4136ad3
@bind sobel_raw_camera_data camera_input(;max_size=200)

# ╔═╡ e15ad330-ee0d-11ea-25b6-1b1b3f3d7888

function process_raw_camera_data(raw_camera_data)
	# the raw image data is a long byte array, we need to transform it into something
	# more "Julian" - something with more _structure_.

	# The encoding of the raw byte stream is:
	# every 4 bytes is a single pixel
	# every pixel has 4 values: Red, Green, Blue, Alpha
	# (we ignore alpha for this notebook)

	# So to get the red values for each pixel, we take every 4th value, starting at
	# the 1st:
	reds_flat = UInt8.(raw_camera_data["data"][1:4:end])
	greens_flat = UInt8.(raw_camera_data["data"][2:4:end])
	blues_flat = UInt8.(raw_camera_data["data"][3:4:end])

	# but these are still 1-dimensional arrays, nicknamed 'flat' arrays
	# We will 'reshape' this into 2D arrays:

	width = raw_camera_data["width"]
	height = raw_camera_data["height"]

	# shuffle and flip to get it in the right shape
	reds = reshape(reds_flat, (width, height))' / 255.0
	greens = reshape(greens_flat, (width, height))' / 255.0
	blues = reshape(blues_flat, (width, height))' / 255.0

	# we have our 2D array for each color
	# Let's create a single 2D array, where each value contains the R, G and B value of
	# that pixel

	RGB.(reds, greens, blues)
end

# ╔═╡ f461f5f2-ee18-11ea-3d03-95f57f9bf09e
gauss_camera_image = process_raw_camera_data(gauss_raw_camera_data);

# ╔═╡ 1ff6b5cc-ee19-11ea-2ca8-7f00c204f587
sobel_camera_image = Gray.(process_raw_camera_data(sobel_raw_camera_data));

# ╔═╡ 1bf94c00-ee19-11ea-0e3c-e12bc68d8e28
[sobel_camera_image Gray.(with_sobel_edge_detect(Gray.(sobel_camera_image)))]

# ╔═╡ Cell order:
# ╟─83eb9ca0-ed68-11ea-0bc5-99a09c68f867
# ╟─8ef13896-ed68-11ea-160b-3550eeabbd7d
# ╟─ac8ff080-ed61-11ea-3650-d9df06123e1f
# ╠═911ccbce-ed68-11ea-3606-0384e7580d7c
# ╟─5f95e01a-ee0a-11ea-030c-9dba276aba92
# ╠═65780f00-ed6b-11ea-1ecf-8b35523a7ac0
# ╟─f7a6d7c3-37b9-437d-8b8e-853665ddbae3
# ╟─54056a02-ee0a-11ea-101f-47feb6623bec
# ╟─e08781fa-ed61-11ea-13ae-91a49b5eb74a
# ╟─a3067222-a83a-47b8-91c3-24ad78dd65c5
# ╟─80108d80-ee09-11ea-0368-31546eb0d3cc
# ╠═7fcd6230-ee09-11ea-314f-a542d00d582e
# ╠═7fdb34dc-ee09-11ea-366b-ffe10d1aa845
# ╟─7fe9153e-ee09-11ea-15b3-6f24fcc20734
# ╠═01070e28-ee0f-11ea-1928-a7919d452bdd
# ╠═ff70782e-e8d2-4281-9b24-d45c925f55e2
# ╟─7522f81e-ee1c-11ea-35af-a17eb257ff1a
# ╟─801d90c0-ee09-11ea-28d6-61b806de26dc
# ╠═802bec56-ee09-11ea-043e-51cf1db02a34
# ╠═d87175b6-7db0-11eb-08a4-456a34fa137e
# ╟─b7f3994c-ee1b-11ea-211a-d144db8eafc2
# ╠═803905b2-ee09-11ea-2d52-e77ff79693b0
# ╠═80479d98-ee09-11ea-169e-d166eef65874
# ╠═805691ce-ee09-11ea-053d-6d2e299ee123
# ╟─bcf98dfc-ee1b-11ea-21d0-c14439500971
# ╠═3492b164-7065-48e8-978b-6c96b965d376
# ╠═02123165-2a0a-49a8-b7a9-458955523511
# ╟─806e5766-ee0f-11ea-1efc-d753cd83d086
# ╠═38da843a-ee0f-11ea-01df-bfa8b1317d36
# ╟─9bde9f92-ee0f-11ea-27f8-ffef5fce2b3c
# ╠═45c4da9a-ee0f-11ea-2c5b-1f6704559137
# ╟─431ba330-0f72-416a-92e9-55f51ff3bcd1
# ╠═5fdc5d0d-a52c-476e-b3b5-3b6364b706e4
# ╟─e84c9cc2-e6e1-46f1-bf4e-9605da5e6f4a
# ╠═807e5662-ee09-11ea-3005-21fdcc36b023
# ╠═ca91bf1a-80ed-11eb-3be0-cd60c9adae9d
# ╠═4f08ebe8-b781-4a32-a218-5ecd8338561d
# ╟─808deca8-ee09-11ea-0ee3-1586fa1ce282
# ╟─809f5330-ee09-11ea-0e5b-415044b6ac1f
# ╟─e555a7e6-f11a-43ac-8218-6d832f0ce251
# ╠═302f0842-453f-47bd-a74c-7942d8c96485
# ╠═7d80a1ea-a0a9-41b2-9cfe-a334717ab2f4
# ╟─ea435e58-ee11-11ea-3785-01af8dd72360
# ╟─80ab64f4-ee09-11ea-29b4-498112ed0799
# ╠═28e20950-ee0c-11ea-0e0a-b5f2e570b56e
# ╟─32a07f1d-93cd-4bf3-bac1-91afa6bb88a6
# ╟─5eea882c-ee13-11ea-0d56-af81ecd30a4a
# ╠═93284f92-ee12-11ea-0342-833b1a30625c
# ╟─cf73f9f8-ee12-11ea-39ae-0107e9107ef5
# ╟─7ffd14f8-ee1d-11ea-0343-b54fb0333aea
# ╟─fa463b71-5aa4-44a3-a67b-6b0776236243
# ╠═8a7d3cfd-6f19-43f0-ae16-d5a236f148e7
# ╟─a34d1ad8-3776-4bc4-93e5-72cfffc54f15
# ╠═971a801d-9c46-417a-ad31-1144894fb4e1
# ╟─5f13b1a5-8c7d-47c9-b96a-a09faf38fe5e
# ╠═338b1c3f-f071-4f80-86c0-a82c17349828
# ╠═bbe1a562-8d97-4112-a88a-c45c260f574d
# ╟─d93fa3f6-c361-4dfd-a2ea-f38e682bcd6a
# ╟─03f91a22-1c3e-4c42-9d78-1ee36851a120
# ╟─48530f0d-49b4-4aec-8109-d69f1ef7f0ee
# ╠═beb62fda-38a6-4528-a176-cfb726f4b5bd
# ╟─f0d55cec-2e81-4cbb-b166-2cf4f2a0f43f
# ╠═1c8b4658-ee0c-11ea-2ede-9b9ed7d3125e
# ╟─f0c3e99d-9eb9-459e-917a-c2338af6683c
# ╠═a6149507-d5ba-45c1-896a-3487070d36ec
# ╟─f8bd22b8-ee14-11ea-04aa-ab16fd01826e
# ╠═2a9dd06a-ee13-11ea-3f84-67bb309c77a8
# ╠═8f70cbfc-80f6-11eb-078c-337a02409c6b
# ╠═6f74c790-80f6-11eb-3200-717715b41e3d
# ╠═b424e2aa-ee14-11ea-33fa-35491e0b9c9d
# ╠═38eb92f6-ee13-11ea-14d7-a503ac04302e
# ╠═bc1c20a4-ee14-11ea-3525-63c9fa78f089
# ╠═24c21c7c-ee14-11ea-1512-677980db1288
# ╟─ce24e486-df27-4780-bc57-d3bf7bee83bb
# ╟─27847dc4-ee0a-11ea-0651-ebbbb3cfd58c
# ╟─b01858b6-edf3-11ea-0826-938d33c19a43
# ╟─7c1bc062-ee15-11ea-30b1-1b1e76520f13
# ╠═7c2ec6c6-ee15-11ea-2d7d-0d9401a5e5d1
# ╟─649df270-ee24-11ea-397e-79c4355e38db
# ╟─9afc4dca-ee16-11ea-354f-1d827aaa61d2
# ╠═cf6b05e2-ee16-11ea-3317-8919565cb56e
# ╟─e3616062-ee27-11ea-04a9-b9ec60842a64
# ╟─e5b6cd34-ee27-11ea-0d60-bd4796540b18
# ╟─b4e98589-f221-4922-b11e-364d72d0788e
# ╟─d06ea762-ee27-11ea-2e9c-1bcff86a3fe0
# ╟─e1dc0622-ee16-11ea-274a-3b6ec9e15ab5
# ╟─efd1ceb4-ee1c-11ea-350e-f7e3ea059024
# ╟─4bbea325-35f8-4a51-bd66-153aba4aed96
# ╠═c4f5a867-74ba-4106-91d4-195f6ae644d0
# ╠═c825ebe2-511b-43ba-afdf-6226dbac48d2
# ╠═2701ab0c-b91d-47fe-b36b-4e0036ecd4aa
# ╠═84a48984-9adb-40ab-a1f1-1ab7b76c9a19
# ╠═3cd535e4-ee26-11ea-2482-fb4ad43dda19
# ╟─7c41f0ca-ee15-11ea-05fb-d97a836659af
# ╠═8b96e0bc-ee15-11ea-11cd-cfecea7075a0
# ╟─0cabed84-ee1e-11ea-11c1-7d8a4b4ad1af
# ╟─5a5135c6-ee1e-11ea-05dc-eb0c683c2ce5
# ╟─577c6daa-ee1e-11ea-1275-b7abc7a27d73
# ╠═275a99c8-ee1e-11ea-0a76-93e3618c9588
# ╠═42dfa206-ee1e-11ea-1fcd-21671042064c
# ╠═2ac28130-810c-11eb-2cb1-117579142578
# ╠═1d3b604a-810c-11eb-27fe-538cba860476
# ╠═9ce383c8-810b-11eb-0c4c-fd184b9f6e3a
# ╠═4739eae4-810a-11eb-07f8-0b47904d9efe
# ╟─6e53c2e6-ee1e-11ea-21bd-c9c05381be07
# ╠═e7f8b41a-ee25-11ea-287a-e75d33fbd98b
# ╟─8a335044-ee19-11ea-0255-b9391246d231
# ╟─79eb0775-3582-446b-996a-0b64301394d0
# ╠═f4d9fd6f-0f1b-4dec-ae68-e61550cee790
# ╟─7c50ea80-ee15-11ea-328f-6b4e4ff20b7e
# ╠═48ebad62-8107-11eb-34e9-09890a16e6ea
# ╠═aad67fd0-ee15-11ea-00d4-274ec3cda3a3
# ╟─9def5f32-ee15-11ea-1f74-f7e6690f2efa
# ╟─8ae59674-ee18-11ea-3815-f50713d0fa08
# ╠═94c0798e-ee18-11ea-3212-1533753eabb6
# ╠═e7b91064-8112-11eb-06c1-cd57fd57133f
# ╠═a75701c4-ee18-11ea-2863-d3042e71a68b
# ╠═9cb8c8a0-810f-11eb-375b-7b4de6432ebd
# ╠═364a0a42-8110-11eb-3907-e9cfc000a302
# ╠═96146b16-79ea-401f-b8ba-e05663a18bd8
# ╠═2cc745ce-e145-4428-af3b-926fba271b67
# ╟─d5ffc6ab-156b-4d43-ac3d-1947d0176e7f
# ╟─f461f5f2-ee18-11ea-3d03-95f57f9bf09e
# ╟─7c6642a6-ee15-11ea-0526-a1aac4286cdd
# ╠═9767e240-8110-11eb-3869-cbfed5e5918d
# ╠═fe8eb9d0-8110-11eb-0fe6-d91ef91f892e
# ╠═9eeb876c-ee15-11ea-1794-d3ea79f47b75
# ╠═525e8402-818c-11eb-1f8b-05cd978cffb3
# ╠═1a0324de-ee19-11ea-1d4d-db37f4136ad3
# ╠═1bf94c00-ee19-11ea-0e3c-e12bc68d8e28
# ╠═260dc208-8113-11eb-05c9-33108c830d89
# ╠═d2063f66-818b-11eb-3762-e734ed57af8f
# ╠═10b2e046-818c-11eb-0f8f-af52231e7fca
# ╠═f864d0b4-818b-11eb-2c48-7b334976ae22
# ╟─1ff6b5cc-ee19-11ea-2ca8-7f00c204f587
# ╟─0001f782-ee0e-11ea-1fb4-2b5ef3d241e2
# ╟─8ffe16ce-ee20-11ea-18bd-15640f94b839
# ╟─5842895a-ee10-11ea-119d-81e4c4c8c53b
# ╟─2d9f3ae4-9e4c-49ce-aab0-5f87aba85c3e
# ╟─5516c800-edee-11ea-12cf-3f8c082ef0ef
# ╟─57360a7a-edee-11ea-0c28-91463ece500d
# ╟─dcb8324c-edee-11ea-17ff-375ff5078f43
# ╟─58af703c-edee-11ea-2963-f52e78fc2412
# ╟─f3d00a9a-edf3-11ea-07b3-1db5c6d0b3cf
# ╟─5aa9dfb2-edee-11ea-3754-c368fb40637c
# ╟─74d44e22-edee-11ea-09a0-69aa0aba3281
# ╟─115ded8c-ee0a-11ea-3493-89487315feb7
# ╟─dfb7c6be-ee0d-11ea-194e-9758857f7b20
# ╟─e15ad330-ee0d-11ea-25b6-1b1b3f3d7888


using Flux.Data.MNIST;

function loadMNIST()
    global MNISTimgs = MNIST.images();
    global MNISTlabels = MNIST.labels();
end

struct MNISTImg
    dbindex::Int;#The index of this image in the MNIST data base
    label::Int;#The label of the image
    data::FFNetwork.TestDatum;#the input and output of the image.
end

output(img::MNISTImg) = output(img.data);
input(img::MNISTImg) = input(img.data);
outlength(img::MNISTImg) = outlength(img.data);
inlength(img::MNISTImg) = inlength(img.data);
data(img::MNISTImg) = img.data;

MNIST_img_length = 28*28;
Labels_length = 10;

labels_tuple = NTuple{10,Vector{Float64}}(
    ([1.,0.,0.,0.,0., 0.,0.,0.,0.,0.],
    [0.,1.,0.,0.,0., 0.,0.,0.,0.,0.],
    [0.,0.,1.,0.,0., 0.,0.,0.,0.,0.],
    [0.,0.,0.,1.,0., 0.,0.,0.,0.,0.],
    [0.,0.,0.,0.,1., 0.,0.,0.,0.,0.],
    [0.,0.,0.,0.,0., 1.,0.,0.,0.,0.],
    [0.,0.,0.,0.,0., 0.,1.,0.,0.,0.],
    [0.,0.,0.,0.,0., 0.,0.,1.,0.,0.],
    [0.,0.,0.,0.,0., 0.,0.,0.,1.,0.],
    [0.,0.,0.,0.,0., 0.,0.,0.,0.,1.]));

vectorize_label(l::Int)::Vector{Float64} = labels_tuple[l+1];

function get_MNIST_img_data(imgindex::Int)::MNISTImg
    imglabel = MNISTlabels[imgindex];
    imgvalues = reshape(Float64.(MNISTimgs[imgindex]),MNIST_img_length)
    MNISTImg(imgindex ,imglabel, TestDatum(imgvalues,vectorize_label(imglabel)));
end
function get_MNIST_img_data()::Vector{MNISTImg}
    return [MNISTImg(imgindex ,MNISTlabels[imgindex], TestDatum(reshape(Float64.(MNISTimgs[imgindex]),MNIST_img_length),vectorize_label(MNISTlabels[imgindex]))) for imgindex in 1:length(MNISTimgs)];
end

function get_MNIST_train_and_test_data()::NTuple{2,Vector{MNISTImg}}
    MNISTLength = length(MNISTlabels);
    total = shuffle!([get_MNIST_img_data(i) for i in 1:MNISTLength]);
    return (total[1:MNISTLength - 10000], total[MNISTLength - 9999:MNISTLength])
end

function train_network_with_MNIST!(net::Network;
    epoch = 10, eta = 3.0, mini_batch_size = 20, verbose = false)
    @assert outlength(net) == Labels_length;
    @assert inlength(net) == MNIST_img_length;
    loadMNIST();
    traindata, testdata = get_MNIST_train_and_test_data();
    train!(net, data.(traindata) , data.(testdata) , epoch, mini_batch_size, eta; verbose = verbose);
end

# (net::Network, train_data::Vector{TestDatum},
#         test_data::Vector{TestDatum}, epoch::Int,
#         mini_batch_size::Int, eta::Float64; verbose = false)


using ITensors,
      ITensors.NDTensors,
      ITensorsGPU,
      LinearAlgebra, # For tr()
      Combinatorics, # For permutations()
      CUDA,
      Test

@testset "cuITensor $T Contractions" for T ∈ (Float64,ComplexF64)
  mi,mj,mk,ml,ma = 2,3,4,5,6,7
  i = Index(mi,"i")
  j = Index(mj,"j")
  k = Index(mk,"k")
  l = Index(ml,"l")
  a = Index(ma,"a")
  @testset "Test contract cuITensors" begin
      Aij = cuITensor(randomITensor(T,i,j))
      Aji = cuITensor(randomITensor(T,j,i))
      Bij = cuITensor(randomITensor(T,i,j))
      Aik = cuITensor(randomITensor(T,i,k))
      Ajk = cuITensor(randomITensor(T,j,k))
      Ajl = cuITensor(randomITensor(T,j,l))
      Akl = cuITensor(randomITensor(T,k,l))
      Dv = rand(T, mi)
      D = itensor(ITensors.tensor(NDTensors.Diag(CuVector(Dv)), IndexSet(i, i')))
      Ev = rand(T, mi)
      E = itensor(ITensors.tensor(NDTensors.Diag(CuVector(Ev)), IndexSet(i, i'')))
    @testset "Test contract cuITensors (Matrix*Diag -> Matrix)" begin
      C = Aij*D
      @test collect(CuArray(C))≈collect(CuMatrix(Aij, j, i))*diagm(0=>Dv)
    end
    @testset "Test contract cuDiagITensors (Diag*Diag -> Diag)" begin
      C  = E*D
      cC = CuArray(C)
      @test collect(cC) ≈ diagm(0=>Ev)*diagm(0=>Dv)
    end
    @testset "Test contract cuDiagITensors (UniformDiag*Diag -> Diag)" begin
      scal = itensor(ITensors.tensor(NDTensors.Diag(2.0), IndexSet(i, i'')))
      C = scal*D
      @test collect(CuArray(C))≈2.0 .* diagm(0=>Dv)
      C = D*scal
      @test collect(CuArray(C))≈2.0 .* diagm(0=>Dv)
    end
    @testset "Test contract cuITensors (Matrix*UniformDiag -> Matrix)" begin
      scal = itensor(ITensors.tensor(NDTensors.Diag(T(2.0)), IndexSet(i, i')))
      C    = scal*Aij
      @test cpu(C)≈2.0*cpu(replaceind(Aij, i, i')) atol=1e-4
      C    = Aij*scal
      @test_broken cpu(C)≈2.0*cpu(replaceind(permute(Aij, j, i), i, i')) atol=1e-4
    end
  end # End contraction testset
end

@testset "cuITensor $T1, $T2 Contractions" for T1 ∈ (Float64,ComplexF64), T2 ∈ (Float64, ComplexF64)
  mi,mj,mk,ml,ma = 2,3,4,5,6,7
  i = Index(mi,"i")
  j = Index(mj,"j")
  k = Index(mk,"k")
  l = Index(ml,"l")
  a = Index(ma,"a")
  @testset "Test contract cuITensors" begin
      Aij = cuITensor(randomITensor(T1,i,j))
      Aji = cuITensor(randomITensor(T1,j,i))
      Bij = cuITensor(randomITensor(T1,i,j))
      Dv = rand(T2, mi)
      D = itensor(ITensors.tensor(NDTensors.Diag(CuVector(Dv)), IndexSet(i, i')))
      Ev = rand(T2, mi)
      E = itensor(ITensors.tensor(NDTensors.Diag(CuVector(Ev)), IndexSet(i, i'')))
    @testset "Test contract cuITensors (Matrix*Diag -> Matrix)" begin
      C = Aij*D
      cC = CuArray(C)
      @test collect(cC)≈collect(CuMatrix(Aij, j, i))*diagm(0=>Dv)
    end
    @testset "Test contract cuDiagITensors (Diag*Diag -> Diag)" begin
      C = E*D
      cC = CuArray(C)
      @test collect(cC)≈diagm(0=>Ev)*diagm(0=>Dv)
    end
    @testset "Test contract cuDiagITensors (UniformDiag*Diag -> Diag)" begin
      scal = itensor(ITensors.tensor(NDTensors.Diag(T2(2.0)), IndexSet(i, i'')))
      C = scal*D
      cC = CuArray(C)
      @test collect(cC)≈2.0 .* diagm(0=>Dv)
      C = D*scal
      cC = CuArray(C)
      @test collect(cC)≈2.0 .* diagm(0=>Dv)
    end
    @testset "Test contract cuITensors (Matrix*UniformDiag -> Matrix)" begin
      scal = itensor(ITensors.tensor(NDTensors.Diag(T2(2.0)), IndexSet(i, i')))
      C    = scal*Aij
      cC = CuArray(C)
      @test collect(cC)≈array(2.0*cpu(replaceind(Aij, i, i'))) atol=1e-4
      C    = Aij*scal
      cC = CuArray(C)
      @test_broken collect(cC)≈array(2.0*cpu(replaceind(permute(Aij, j, i), i, i'))) atol=1e-4
    end
  end # End contraction testset
end


#=
Created on Saturday 3 August 2019
Last update: Thursday 8 August 2019

@author: Michiel Stock
michielfmstock@gmail.com

Some methods to simulate Kronecker graphs.

Leskovec, J., Chakrabarti, D., Kleinberg, J., Faloutsos, C., & Ghahramani,
Z. (2008). Kronecker graphs: an approach to modeling networks.
Journal of Machine Learning Research, 11, 985–1042.
Retrieved from https://cs.stanford.edu/~jure/pubs/kronecker-jmlr10.pdf
=#

using SparseArrays: spzeros
using StatsBase: sample, Weights

"""
    isprob(A::AbstractArray)

Test if a matrix can be interpeted as a probability matrix, i.e., all elements
are between 0 and 1.
"""
isprob(A::AbstractArray) = all(0 .<= A .<= 1)

"""
    isprob(K::AbstractKroneckerProduct)

Test if a Kronecker product can be interpeted as a probability matrix,
i.e., all elements are between 0 and 1.
"""
function isprob(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return isprob(A) && isprob(B)
end

"""
    naivesample(P::AbstractKroneckerProduct)

Sample a Kronecker graph from a probabilistic Kronecker product P using the
naive method. This method has a time complexity in the size of the Kronecker
product (but is still light in memory use). Consider using `fastsample`.
"""
function naivesample(P::AbstractKroneckerProduct)
    @assert isprob(P) throw(DomainError(
                            "All values of K should be between 0 and 1"))
    G = spzeros(Bool, size(P)...)
    for I in CartesianIndices(P)
        if P[I] > rand()  # QUESTION: is this the most efficient way?
            @inbounds G[I] = true
        end
    end
    return G
end

"""
    sampleindices(A::AbstractMatrix, s::Int)

Samples the indices from an `AbstractMatrix`. Probability of sampling indices is
proportional to the size of the corresponding value. Does not do any checks on A.
"""
sampleindices(A::AbstractMatrix, s::Int) = Tuple.(sample(CartesianIndices(A),
                                                    Weights(vec(A), sum(A)), s,
                                                    replace=true))

"""
sampleindices(K::AbstractKroneckerProduct, s::Int)

Samples the indices from an `AbstractKroneckerProduct`. Probability of
sampling indices is proportional to the size of the corresponding value.
Does not do any checks on A.
"""
function sampleindices(K::AbstractKroneckerProduct, s::Int)
    A, B = getmatrices(K)
    p, q = size(B)
    indicesA = sampleindices(A, s)
    indicesB = sampleindices(B, s)
    indices = similar(indicesA)
    for (o, (Ia, Ib)) in enumerate(zip(indicesA, indicesB))
        (i, j), (k, l) = Ia, Ib
        @inbounds indices[o] = ((i - 1) * p + k, (j - 1) * q + l)
    end
    return indices
end

"""
    fastsample(P::AbstractKroneckerProduct)

Uses the heuristic sampling from Leskovec et al. (2008) to sample a large
Kronecker graph.
"""
function fastsample(P::AbstractKroneckerProduct)
    @assert isprob(P) throw(DomainError(
                            "All values of K should be between 0 and 1"))
    G = spzeros(Bool, size(P)...)
    n = Int(round(sum(P)))  # expected number of edges
    for (i, j) in sampleindices(P, n)
        @inbounds G[i,j] = true
    end
    return G
end


using AugmentedGaussianProcesses
using Distributions
using Distances
using Random
ν = 3.0

testfunction(x) = (1 + x / ν)^(-0.5 * (ν + 1))
testfunction(x) = sech(sqrt(x / 2))
# testfunction(x) = (1+sqrt(3*x))*exp(-sqrt(3*x))
# testfunction(x) = (1+sqrt(5*x)+5/3*x)*exp(-sqrt(5*x))
testdist = AGP.LaplaceTransformDistribution(
    testfunction, 0.0, AGP.BromwichInverseLaplace(19, 1, 10, 5)
)
@btime rand($testdist)
@btime rand(Gamma((ν + 1) / 2, 1 / ν))
rand(testdist)
# @btime AGP.draw(AugmentedGaussianProcesses.PolyaGammaDist(),1,0)
evaluate(KLDivergence(), rand(testdist, 10000), rand(Gamma((ν + 1) / 2, 1 / ν), 10000))
# laptrans(testdist,n=1)[1]
p = display(histogram(rand(testdist, 10000); bins=range(0, 3; length=100), normalize=true))
histogram!(
    rand(Gamma((ν + 1) / 2, 1 / ν), 10000);
    alpha=0.5,
    bins=range(0, 10; length=100),
    normalize=true,
)
# histogram!([AGP.draw(AugmentedGaussianProcesses.PolyaGammaDist(),1,0) for i in 1:10000],alpha=0.5,bins=range(0,10,length=100),normalize=true)


"""
    plsnipals(X, Y, weights = ones(size(X, 1)); nlv)
Partial Least Squares Regression (PLSR) with the NIPALS algorithm
(e.g. Tenenhaus 1998, Wold 2002).
* `X` : X-data.
* `Y` : Y-data.
* `weights` : Weights of the observations.
* `nlv` : Nb. latent variables (LVs) to consider.

In the function, the usual PLS2 NIPALS iterations are replaced by a
direct computation of the PLS weights (`w`) by SVD decomposition of matrix X'Y
(Hoskuldsson 1988 p.213).

Vector `weights` (row-weighting) is internally normalized to sum to 1.
See the help of `plskern` for details.

`X` and `Y` are internally centered.

## References

Hoskuldsson, A., 1988. PLS regression methods. Journal of Chemometrics 2, 211-228.
https://doi.org/10.1002/cem.1180020306

Tenenhaus, M., 1998. La régression PLS: thÃ©orie et pratique. Editions Technip,
Paris, France.

Wold, S., Sjostrom, M., Eriksson, l., 2001. PLS-regression: a basic tool
for chemometrics. Chem. Int. Lab. Syst., 58, 109-130.
"""
function plsnipals(X, Y, weights = ones(size(X, 1)); nlv)
    plsnipals!(copy(X), copy(Y), weights; nlv = nlv)
end

function plsnipals!(X, Y, weights = ones(size(X, 1)); nlv)
    X = ensure_mat(X)
    Y = ensure_mat(Y)
    n = size(X, 1)
    p = size(X, 2)
    q = size(Y, 2)
    nlv = min(nlv, n, p)
    weights = mweights(weights)
    D = Diagonal(weights)
    xmeans = colmeans(X, weights)
    ymeans = colmeans(Y, weights)
    center!(X, xmeans)
    center!(Y, ymeans)
    # Pre-allocation
    XtY = similar(X, p, q)
    T = similar(X, n, nlv)
    P = similar(X, p, nlv)
    W = copy(P)
    C = similar(X, q, nlv)
    TT = similar(X, nlv)
    t   = similar(X, n)
    dt  = similar(X, n)
    zp  = similar(X, p)
    w   = similar(X, p)
    c   = similar(X, q)
    # End
    @inbounds for a = 1:nlv
        XtY .= X' * (D * Y)
        if q == 1
            w .= vec(XtY)
            w ./= sqrt(dot(w, w))
        else
            w .= svd!(XtY).U[:, 1]
        end
        mul!(t, X, w)
        dt .= weights .* t
        tt = dot(t, dt)
        mul!(c, Y', dt)
        c ./= tt
        mul!(zp, X', dt)
        zp ./= tt
        X .-= t * zp'
        Y .-= t * c'
        P[:, a] .= zp
        T[:, a] .= t
        W[:, a] .= w
        C[:, a] .= c
        TT[a] = tt
     end
     R = W * inv(P' * W)
     Plsr(T, P, R, W, C, TT, xmeans, ymeans, weights, nothing)
end


"""
A partition divides a region into several subregions. The main algorithm of
every partition is deciding in which subregion a given point x lies, implemented
by the `partition_index` function.

The intent of a partition is to be combined with a function set on each subregion,
effectively creating a piecewise function.

The pieces of a partition have an index ranging from 1 to length(partition).
"""
abstract type Partition end

"""
A PartitionPiece represents a piece of a partition - it essentially combines
a partition and an index into that partition. It is constructed by indexing
a partition.
"""
struct PartitionPiece{P <: Partition,I}
    part    ::  P
    idx     ::  I
end

partition(piece::PartitionPiece) = piece.part
index(piece::PartitionPiece) = piece.idx

left(piece::PartitionPiece) = left(partition(piece), index(piece))
right(piece::PartitionPiece) = right(partition(piece), index(piece))

Base.in(x, piece::PartitionPiece) = in_partition(partition(piece), index(piece), x)

function checkbounds(part::Partition, i::Int)
    1 <= i <= length(part) || throw(BoundsError())
end

function getindex(part::Partition, i)
    checkbounds(part, i)
    PartitionPiece(part, i)
end

eachindex(part::Partition) = 1:length(part)

"""
A PiecewiseInterval is a subdivision of an interval into a number of subintervals.
The subintervals are determined by a set of breakpoints ``t_i``, ``i=0,...,n``. The main
interval is ``[t_0,t_n]``. The subintervals always contain the right endpoint, i.e.
the first subinterval is ``[t_0,t_1]``, the second is the halfopen interval
``(t_1,t_2]`` and so on until the n-th interval ``(t_{n-1},t_n]``.
"""
struct PiecewiseInterval{T} <: Partition
    points  ::  Array{T,1}
end

# n partitions of equal size on the interval [a,b]
PiecewiseInterval(a, b, n::Int) = PiecewiseInterval(collect(linspace(a,b,n+1)))

length(part::PiecewiseInterval) = length(part.points)-1

support(part::PiecewiseInterval) = Interval(part.points[1],part.points[end])

# Not sure this is a good idea, but let's respect the convention that the
    # interval does not contain the left endpoint, unless it is the first interval.
support(part::PiecewiseInterval, i::Int) = i==1 ? Interval(part.points[1],part.points[2]) : Interval{:open,:closed}(part.points[i],part.points[i+1])
left(part::PiecewiseInterval) = part.points[1]

left(part::PiecewiseInterval, i::Int) = part.points[i]

right(part::PiecewiseInterval) = part.points[end]

right(part::PiecewiseInterval, i::Int) = part.points[i+1]

"Return the index of the region in which the given point x lies."
function partition_index(part::PiecewiseInterval, x)
    @assert in(x,support(part))

    idx = 1
    while x > part.points[idx+1]
        idx += 1
    end
    idx
end

in_partition(part::PiecewiseInterval, i, x) = in(x,support(part,i))

# TODO: applying partition_index repeatedly, say for all points in a grid, could
# be done more efficiently if we can exploit ordering of the points in the grid.
# Perhaps partition_index applied to an iterable set of points could return an
# iterable set of indices?

split_interval(part::PiecewiseInterval, i::Int, x) = PiecewiseInterval([part.points[1:i]...,x,part.points[i+1:end]...])

split_interval(part::PiecewiseInterval, x) = split(part, partition_index(part, x), x)


global t = 0

function braitenberg(n_::Ptr{Csize_t}, in_::Ptr{Cdouble},
                     m_::Ptr{Csize_t}, out_::Ptr{Cdouble},
                     l_::Ptr{Csize_t}, str_::Ptr{UInt8})

  n::Int64 = unsafe_load(n_)
  m::Int64 = unsafe_load(m_)
  l::Int64 = unsafe_load(l_)

  input::Vector{Float64}  = pointer_to_array(in_,  (n,))
  output::Vector{Float64} = pointer_to_array(out_, (m,))
  message::Vector{UInt8}  = pointer_to_array(str_, (l,))

  # START your code here

  global t

  t = t + 1
  status = 0

  if t % 1000 == 0
    status = 1
  end

  if t == 10000
    status = 2
  end

  o = [-tanh(1/3*(input[4] + input[5] + input[6])),
       -tanh(1/3*(input[1] + input[2] + input[3]))]

  str = string(input[1], " ", input[2], " ", input[3], "\n", input[4], " ", input[5], " ", input[6])

  # END your code here

  unsafe_copy!(pointer(output),  pointer(o),   m)
  unsafe_copy!(pointer(message), pointer(str), l)

  return convert(Cint, status)
end

function test_function()
  println("test function called")
end

cfunction(braitenberg, Cint, (Ptr{Csize_t}, Ptr{Cdouble}, Ptr{Csize_t}, Ptr{Cdouble}, Ptr{Csize_t}, Ptr{UInt8}))


using UnbalancedOptimalTransport
using UnbalancedOptimalTransport: KL, TV, RG, Balanced
using Test
using SCS, ECOS
using StaticArrays

include("Convex_formulation.jl")

const TEST_DIVERGENCES = (KL(), KL(2.0), KL(100.0), TV(4.0), RG(0.5, 1.5), Balanced())

function rand_measure(
    n,
    d;
    scale = 100,
    static = false,
    prob_type = Float64,
    set_type = Float64,
)
    if static
        make_vector =
            () -> d == 1 ? Scalar{set_type}(randn()) : SVector{d,set_type}(randn(d))
    else
        make_vector = () -> d == 1 ? [set_type(randn())] : set_type.(randn(d))
    end
    DiscreteMeasure(rand(prob_type, n), [scale * make_vector() for _ = 1:n])
end

function balanced_measures(n1, n2, d; kwargs...)
    a = rand_measure(n1, d; kwargs...)
    b = rand_measure(n2, d; kwargs...)
    b = DiscreteMeasure(sum(a.density) * b.density / sum(b.density), b.set)
    return a, b
end

@testset "UnbalancedOptimalTransport.jl" begin

    # We check the implementation of the Sinkhorn algorithm against its definition as an optimization problem.
    @testset "regularized OT! with balancing via $D" for D in TEST_DIVERGENCES
        for ϵ ∈ (0.1, 1.0)
            # We choose very small test problems, since solving the convex problem is relatively slow
            # Ensure the problem is feasible
            if D isa Balanced
                a, b = balanced_measures(4, 3, 1; scale = 100.0)
            else
                a = rand_measure(4, 1; scale = 100.0)
                b = rand_measure(3, 1; scale = 100.0)
            end

            if D isa RG
                sc = rand() + 0.5
                b = DiscreteMeasure(sc * sum(a.density) * b.density / sum(b.density), b.set)
            end

            # ECOS natively supports the exponential cone we need for the KL divergence
            if D isa TV
                solver = () -> SCS.Optimizer(verbose = 0, eps = 1e-6, max_iters = 20_000)
            else
                solver = () -> ECOS.Optimizer(verbose = 0, eps = 1e-6)
            end

            result = OT_convex(D, a, b, ϵ; solver = solver, verbose = true)

            @test OT!(D, a, b, ϵ) ≈ result.optimal_value rtol = 1e-3 atol = 1e-3

            coupling_sinkhorn = optimal_coupling!(D, a, b, ϵ)
            @test result.optimal_coupling ≈ coupling_sinkhorn rtol = 1e-3 atol = 1e-3

        end
    end

    @testset "Optimal coupling" begin
        for ϵ ∈ (0.1, 1.0), D in TEST_DIVERGENCES
            if D isa Balanced
                a, b = balanced_measures(4, 3, 1; scale = 100.0)
            else
                a = rand_measure(4, 1; scale = 100.0)
                b = rand_measure(3, 1; scale = 100.0)
            end

            if D isa RG
                sc = rand() + 0.5
                b = DiscreteMeasure(sc * sum(a.density) * b.density / sum(b.density), b.set)
            end

            aa = deepcopy(a)
            bb = deepcopy(b)
            OT_val = OT!(D, aa, bb, ϵ)

            # Try both with `dual_potentials_populated` true and false
            π2 = optimal_coupling!(D, aa, bb, ϵ; dual_potentials_populated = true)
            π = optimal_coupling!(D, a, b, ϵ)
            @test π ≈ π2

            obj = objective(π, D, a, b, ϵ)
            @test OT_val ≈ obj rtol = 1e-3
        end
    end

    @testset "Allocations" begin
        a = rand_measure(2, 2; static = true)
        b = rand_measure(2, 2; static = true)
        C = cost_matrix(a, b)
        OT!(KL(), C, a, b) # compile

        a = rand_measure(100, 2; static = true)
        b = rand_measure(80, 2; static = true)
        C = cost_matrix(a, b)
        @test @allocated(OT!(KL(), C, a, b)) <= 500

        a = rand_measure(1000, 2; static = true)
        b = rand_measure(800, 2; static = true)
        C = cost_matrix(a, b)
        @test @allocated(OT!(KL(), C, a, b)) <= 500
    end

    @testset "Prop. 12: Optimized KL-Sinkhorn divergence method" begin
        for ρ ∈ (0.5, 2.0),
            ϵ ∈ (0.1, 1.0, 10.0),
            C ∈ ((x, y) -> norm(x - y, 3.0), (x, y) -> norm(x - y))

            a = rand_measure(100, 2; static = true)
            b = rand_measure(80, 2; static = true)
            @test sinkhorn_divergence!(KL(ρ), C, a, b, ϵ; tol = 1e-6) ≈
                  UnbalancedOptimalTransport.generic_sinkhorn_divergence!(
                KL(ρ),
                C,
                a,
                b,
                ϵ;
                tol = 1e-6,
            ) rtol = 1e-4

            # Test default argument for `C`
            @test sinkhorn_divergence!(KL(ρ), a, b, ϵ; tol = 1e-6) ≈
                  UnbalancedOptimalTransport.generic_sinkhorn_divergence!(
                KL(ρ),
                (x, y) -> norm(x - y),
                a,
                b,
                ϵ;
                tol = 1e-6,
            ) rtol = 1e-4
        end
    end

    @testset "BigFloats" begin
        ϵ = 1.0
        ϵ_big = big(ϵ)
        for D in TEST_DIVERGENCES
            if D isa Balanced
                a, b = balanced_measures(8, 9, 1; static = true, scale = 2.0)
            else
                a = rand_measure(8, 1; static = true, scale = 2.0)
                b = rand_measure(9, 1; static = true, scale = 2.0)
            end
            a_set_big = [Scalar{BigFloat}(x.data) for x in a.set]
            a_big = DiscreteMeasure(big.(a.density), a_set_big)
            b_set_big = [Scalar{BigFloat}(x.data) for x in b.set]
            b_big = DiscreteMeasure(big.(b.density), b_set_big)
            if D isa Balanced
                b_big = DiscreteMeasure(
                    sum(a_big.density) * b_big.density / sum(b_big.density),
                    b_big.set,
                )
            end

            sd = @inferred(sinkhorn_divergence!(D, a, b, ϵ; max_iters = 10^6))
            sd_big =
                @inferred(sinkhorn_divergence!(D, a_big, b_big, ϵ_big; max_iters = 10^6))
            @test sd_big isa BigFloat
            @test sd ≈ sd_big
        end
    end

    @testset "Lemma 11: equality of dual potentials" begin
        for D in (KL(1.0), KL(2.0)), ϵ in (0.1, 1.0, 10.0)
            a = rand_measure(30, 3; static = true)
            b = DiscreteMeasure(copy(a.density), copy(a.set))
            unbalanced_sinkhorn!(D, a, b, ϵ; tol = 1e-7, max_iters = 10^6)
            @test a.dual_potential ≈ b.dual_potential rtol = 1e-4
        end
    end

    @testset "Theorem 4: positive definiteness and (separate) convexity of Sinkhorn divergence" begin
        for D in (KL(1.0), KL(2.0)), ϵ in (0.1, 1.0, 10.0)
            a_set = rand_measure(30, 3; static = true).set
            a_1 = DiscreteMeasure(rand(30), copy(a_set))
            a_2 = DiscreteMeasure(rand(30), copy(a_set))

            b_set = rand_measure(25, 3; static = true).set
            b_1 = DiscreteMeasure(rand(25), copy(b_set))
            b_2 = DiscreteMeasure(rand(25), copy(b_set))

            λ = rand()
            a = DiscreteMeasure(λ .* a_1.density .+ (1 - λ) .* a_2.density, copy(a_set))
            b = DiscreteMeasure(λ .* b_1.density .+ (1 - λ) .* b_2.density, copy(b_set))

            sd_b1 = sinkhorn_divergence!(D, a, b_1, ϵ; tol = 1e-5)
            sd_b2 = sinkhorn_divergence!(D, a, b_2, ϵ; tol = 1e-5)
            sd_a1 = sinkhorn_divergence!(D, a_1, b, ϵ; tol = 1e-5)
            sd_a2 = sinkhorn_divergence!(D, a_2, b, ϵ; tol = 1e-5)
            sd = sinkhorn_divergence!(D, a, b, ϵ; tol = 1e-5)
            @test all(x -> x > 0, (sd_b1, sd_b2, sd_a1, sd_a2, sd)) # positive definite
            @test λ * sd_a1 + (1 - λ) * sd_a2 >= sd # convexity in a
            @test λ * sd_b1 + (1 - λ) * sd_b2 >= sd # convexity in b

            @test sinkhorn_divergence!(D, a, a, ϵ; tol = 1e-5) ≈ 0 atol = 1e-4
            @test sinkhorn_divergence!(D, b, b, ϵ; tol = 1e-5) ≈ 0 atol = 1e-4
        end
    end

    @testset "eltype" begin
        for T in (Float32, Float64, BigFloat)
            a = rand_measure(4, 2; prob_type = T)
            @test eltype(a) == T
        end
    end

    @testset "Warnings and errors" begin
        a, b = balanced_measures(4, 3, 2)
        a = DiscreteMeasure(a.density * 2, a.set) # unbalanced
        @test_logs (:warn, r"for `D==Balanced\(\)`") unbalanced_sinkhorn!(
            Balanced(),
            a,
            b;
            tol = 10.0,
        )

        a = rand_measure(100, 2; static = true)
        b = rand_measure(90, 2; static = true)
        @test_logs (:warn, r"Maximum iterations") unbalanced_sinkhorn!(
            TV(),
            a,
            b,
            max_iters = 1,
        )

        @test_throws ArgumentError DiscreteMeasure(rand(5), rand(4), rand(5))
        @test_throws ArgumentError DiscreteMeasure(rand(5), rand(5), rand(4))
        @test_throws ArgumentError DiscreteMeasure(rand(4), rand(5), rand(5))
        @test_throws ArgumentError sinkhorn_divergence!(
            KL(1),
            randn(100, 90),
            a,
            b,
            ϵ = 1e-1,
        )
    end
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:4] <= 1)

                     @variable(m, -8.186501529651974 <= q <= 9.622483160840265)

                     add_NL_constraint(m, :(softplus(-0.9342689023247241 + -0.06643065035242302*softplus(-0.6967397920955185 + 0.4612305646899797*softplus(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4])) + 0.5917923558177942*softplus(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4]))) + 0.13022424784864484*softplus(-0.8092681495291991 + -0.6577545463201151*softplus(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4])) + 0.816846389391447*softplus(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4])))) + softplus(0.9502905692637253 + 0.2731932353707758*softplus(-0.6967397920955185 + 0.4612305646899797*softplus(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4])) + 0.5917923558177942*softplus(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4]))) + 0.7525923835832526*softplus(-0.8092681495291991 + -0.6577545463201151*softplus(-0.7651411836742423 + 0.736339264438953*$(x[1]) + 0.06323249947984699*$(x[2]) + 0.6326865446578034*$(x[3]) + 0.8935006364653741*$(x[4])) + 0.816846389391447*softplus(0.8959549878207418 + 0.6412415470262633*$(x[1]) + 0.6383904188579423*$(x[2]) + 0.9706157764666923*$(x[3]) + 0.8108132764672078*$(x[4])))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -6.103907844236108 <= q <= 8.378675476959124)

                     add_NL_constraint(m, :(sigmoid(-0.8297607126794757 + -0.5064858067335005*sigmoid(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2])) + -0.8619658571743893*sigmoid(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2])) + 0.8415580445654567*sigmoid(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))) + sigmoid(-0.0809505831099564 + 0.6185835236357913*sigmoid(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2])) + -0.5288114683363174*sigmoid(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2])) + -0.42070009030061817*sigmoid(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))) + sigmoid(0.16011867292674786 + -0.14898607305663658*sigmoid(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2])) + -0.8220978033005912*sigmoid(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2])) + -0.8518875556726564*sigmoid(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


"""
    eⱽ(M, r, a, θ)
Modified from Cunningham et al. (1975) eq. (A2a):
```math
e^\\nu = \\sqrt{\\frac{\\Delta \\Sigma}{A}}.
```
"""
eⱽ(M, r, a, θ) = √(Σ(r, a, θ) * Δ(M, r, a) / A(M, r, a, θ))

"""
    eᶲ(M, r, a, θ)
Modified from Cunningham et al. (1975) eq. (A2b):
```math
e^\\Phi = \\sin \\theta \\sqrt{\\frac{A}{\\Sigma}}.
```
"""
eᶲ(M, r, a, θ) = sin(θ) * √(A(M, r, a, θ) / Σ(r, a, θ))

"""
    ω(M, r, a, θ)
From Cunningham et al. (1975) eq. (A2c):
```math
\\omega = \\frac{2 a M r}{A}.
```
"""
ω(M, r, a, θ) = 2 * a * M * r / A(M, r, a, θ)

"""
    Ωₑ(M, r, a)
Coordinate angular velocity of an accreting gas.
Taken from Cunningham et al. (1975) eq. (A7b):
```math
\\Omega_e = \\frac{\\sqrt{M}}{a \\sqrt{M} + r_e^{3/2}}.
```
# Notes
Fanton et al. (1997) use
```math
\\Omega_e = \\frac{\\sqrt{M}}{a \\sqrt{M} \\pm r_e^{3/2}},
```
where the sign is dependent on co- or contra-rotation. This function may be extended in the future to support this definition.
"""
Ωₑ(M, r, a) = √M / (r^1.5 + a * √M)

"""
    Vₑ(M, r, a, θ)
Velocity of an accreting gas in a locally non-rotating reference frame (see Bardeen et al. 1973).
Taken from Cunningham et al. (1975) eq. (A7b):
```math
V_e = (\\Omega_e - \\omega) e^{\\Phi - \\nu}.
```
"""
Vₑ(M, r, a, θ) = (Ωₑ(M, r, a) - ω(M, r, a, θ)) * eᶲ(M, r, a, θ) / eⱽ(M, r, a, θ)

"""
    Lₑ(M, rms, a)
Angular momentum of an accreting gas within ``r_ms``.
Taken from Cunningham et al. (1975) eq. (A11b):
```math
L_e = \\sqrt{M} \\frac{
        r_{\\text{ms}}^2 - 2 a \\sqrt{M r_{\\text{ms}}} + a^2
    }{
        r_{\\text{ms}}^{3/2} - 2 M \\sqrt{r_{\\text{ms}}} + a \\sqrt{M}
    }.
```
"""
Lₑ(M, rms, a) = √M * (rms^2 - 2 * a * √(M * rms) + a^2) / (rms^1.5 - 2 * M * √rms + a * √M)

"""
    H(M, rms, r, a)
Taken from Cunningham et al. (1975) eq. (A12e):
```math
H = \\frac{2 M r_e - a \\lambda_e}{\\Delta},
```
where we distinguing ``r_e`` as the position of the accreting gas.
"""
H(M, rms, r, a) = (2 * M * r - a * Lₑ(M, rms, a)) / Δ(M, r, a)

"""
    γₑ(M, rms)
Taken from Cunningham et al. (1975) eq. (A11c):
```math
\\gamma_e = \\sqrt{1 - \\frac{
        2M
    }{
        3 r_{\\text{ms}}
    }}.
```
"""
γₑ(M, rms) = √(1 - (2 * M) / (3 * rms))

"""
    uʳ(M, rms, r)
Taken from Cunningham et al. (1975) eq. (A12b):
```math
u^r = - \\sqrt{\\frac{
        2M
    }{
        3 r_{\\text{ms}}
    }} \\left(
        \\frac{ r_{\\text{ms}} }{r_e} - 1
    \\right)^{3/2}.
```
"""
uʳ(M, rms, r) = -√((2 * M) / (3 * rms)) * (rms / r - 1)^1.5

"""
    uᶲ(M, rms, r, a)
Taken from Cunningham et al. (1975) eq. (A12c):
```math
u^\\phi = \\frac{\\gamma_e}{r_e^2} \\left(
        L_e + aH
    \\right).
```
"""
uᶲ(M, rms, r, a) = γₑ(M, rms) / r^2 * (Lₑ(M, rms, a) + a * H(M, rms, r, a))

"""
    uᵗ(M, rms, r, a)
Taken from Cunningham et al. (1975) eq. (A12b):
```math
u^t = \\gamma_e \\left(
        1 + \\frac{2 M (1 + H)}{r_e}
    \\right).
```
"""
uᵗ(M, rms, r, a) = γₑ(M, rms) * (1 + 2 * M * (1 + H(M, rms, r, a)) / r)

"""
Experimental API:
"""
function regular_pdotu_inv(L, M, r, a, θ)
    (eⱽ(M, r, a, θ) * √(1 - Vₑ(M, r, a, θ)^2)) / (1 - L * Ωₑ(M, r, a))
end

@inline function regular_pdotu_inv(m::AbstractMetricParams{T}, u, v) where {T}
    metric = GeodesicBase.metric(m, u)
    # reverse signs of the velocity vector
    # since we're integrating backwards
    p = @inbounds @SVector [-v[1], 0, 0, -v[4]]

    # TODO: this only works for Kerr
    disc_norm = (
        AccretionFormulae.eⱽ(m.M, u[2], m.a, u[3]) *
        √(1 - AccretionFormulae.Vₑ(m.M, u[2], m.a, u[3])^2)
    )

    u_disc =
        @SVector [1 / disc_norm, 0, 0, AccretionFormulae.Ωₑ(m.M, u[2], m.a) / disc_norm]

    # use Tullio to do the einsum
    @tullio g := metric[i, j] * (u_disc[i]) * p[j]
    1 / g
end

function plunging_p_dot_u(E, a, M, L, Q, rms, r, sign_r)
    inv(
        uᵗ(M, rms, r, a) - uᶲ(M, rms, r, a) * L -
        sign_r * uʳ(M, rms, r) * Σδr_δλ(E, L, M, Q, r, a) / Δ(M, r, a),
    )
end

function plunging_p_dot_u(m::AbstractMetricParams{T}, u, v, rms) where {T}
    metric = GeodesicBase.metric(m, u)

    # reverse signs of the velocity vector
    # since we're integrating backwards
    p = @inbounds @SVector [-v[1], v[2], 0, -v[4]]

    u_disc = @inbounds @SVector [
        uᵗ(m.M, rms, u[2], m.a),
        uʳ(m.M, rms, u[2]),
        0,
        uᶲ(m.M, rms, u[2], m.a),
    ]

    @tullio g := metric[i, j] * u_disc[i] * p[j]
    1 / g
end

@inline function redshift_function(m::CarterMethodBL{T}, u, p, λ) where {T}
    isco = rms(m.M, m.a)
    if u[2] > isco
        @inbounds regular_pdotu_inv(p.L, m.M, u[2], m.a, u[3])
    else
        # change sign if we're after the sign flip
        # TODO: this isn't robust to multiple sign changes
        # TODO: i feel like there should be a better way than checking this with two conditions
        #       used to have λ > p.changes[1] to make sure we're ahead of the time flip (but when wouldn't we be???)
        #       now p.changes[1] > 0.0 to make sure there was a time flip at all
        sign_r = (p.changes[1] > 0.0 ? 1 : -1) * p.r
        @inbounds plunging_p_dot_u(m.E, m.a, m.M, p.L, p.Q, isco, u[2], sign_r)
    end
end

@inline function redshift_function(m::AbstractMetricParams{T}, u, v) where {T}
    isco = rms(m.M, m.a)
    if u[2] > isco
        regular_pdotu_inv(m, u, v)
    else
        plunging_p_dot_u(m, u, v, isco)
    end
end

# point functions exports

function _redshift_guard(
    m::CarterMethodBL{T},
    gp::CarterGeodesicPoint{T},
    max_time,
) where {T}
    redshift_function(m, gp.u, gp.p, gp.t)
end
function _redshift_guard(m::AbstractMetricParams{T}, gp, max_time) where {T}
    redshift_function(m, gp.u, gp.v)
end

const redshift = PointFunction(_redshift_guard)


######################################################################
# isogenies.jl: Isogenies between elliptic curves
######################################################################

#Only odd-degree isogenies are implemented; create a special class for degree 2 isogenies

export Isogeny

export degree, rational_fractions, areisomorphic, compose, squarefree_kernel, isomorphism_coefficients, subgroup, dual, isomorphism, find_isomorphism, multiplication_isogeny

######################################################################
# Type definitions
######################################################################

"""
Concrete type for isogenies between Weierstrass or Montgomery elliptic curves.
Only odd degree isogenies are currently implemented.
The kernel polynomial of an l-isogeny is a monic, squarefree polynomial of degree (l-1)/2.
In the rational fractions, the usual factor y in the ordinate is dropped in order to use only univariate polynomials.
"""
#TODO: make this a Nemo.Map
mutable struct Isogeny{T}
	domain::EllipticCurve{T}
	degree::Integer
	kernel::PolyElem{T}
	u::T
	r::T
	s::T
	t::T
	image::EllipticCurve{T}
	x::Union{Frac{U}, Nothing} where U <: PolyElem{T}
	y::Union{Frac{U}, Nothing} where U <: PolyElem{T}
end

base_ring(phi::Isogeny) = base_ring(domain(phi))

domain(phi::Isogeny) = phi.domain

image(phi::Isogeny) = phi.image

kernel(phi::Isogeny) = phi.kernel

degree(phi::Isogeny) = phi.degree

"""
Get the u, r, s, t coefficients of the isomophism coming after the normalized part of the isogeny
"""
isomorphism_coefficients(phi::Isogeny) = (phi.u, phi.r, phi.s, phi.t)

"""
Get the rational fractions defining the isogeny. They are computed at the first call.
"""
function rational_fractions(phi::Isogeny)
	try
		phix = get(phi.rationalx)
		phiy = get(phi.rationaly)
		return (phix, phiy)
	catch
		fx, fy = _compute_rational_fractions(phi)
		phi.x = fx
		phi.y = fy
		return fx, fy
	end
end

"""
Decide if two isogenies differ by postcomposition by an isomorphism.
"""
areisomorphic(phi::Isogeny, psi::Isogeny) = domain(phi) == domain(psi) && degree(phi) == degree(psi) && kernel(phi) == kernel(psi)

"""
Decide if two isogenies are the same.
"""
==(phi::Isogeny, psi::Isogeny) =  areisomorphic(phi, psi) && isomorphism_coefficients(phi) == isomorphism_coefficients(psi)

"""
Shows a description of an isogeny between elliptic curves.
"""
function show(io::IO, phi::Isogeny)
	print("Isogeny of degree ")
	show(phi.degree)
	print(" between ")
	show(io, phi.domain)
	print(" and ")
	show(io, phi.image)
	print(" with kernel ")
	show(io, phi.kernel)
end

"""
Decide if an isogeny has valid formulas.
"""
function isvalid(phi::Isogeny)
	E1 = domain(phi)
	E2 = image(phi)
	a11, a21, a31, a41, a61 = a_invariants(E1) #we have a1 = a3 = 0
	a12, a22, a32, a42, a62 = a_invariants(E2) #we have a1 = a3 = 0
	phix, phiy = rational_fractions(phi)
	X = gen(base_ring(parent(phix)))
	lhs = phiy^2 * (X^3 + a21 * X^2 + a41 * X + a61)
	rhs = (phix^3 + a22 * phix^2 + a42 * phix + a62)
	test = lhs - rhs
	return test==0
end

######################################################################
# Rational fractions and evaluation
######################################################################

"""
Get the evaluation of an isogeny on a point.
"""
function (phi::Isogeny)(P::EllipticPoint)
	Q = normalized(P)
	x1 = Q.X
	y1 = Q.Y
	phix, phiy = rational_fractions(phi)
	try
		x2 = phix(x1)
		y2 = y1 * phiy(x1)
		return EllipticPoint(x2, y2, one(base_ring(Q)), image(phi))
	catch
		return infinity(image(phi))
	end
end

#DOES NOT WORK if u or t != 0...
"""
Compute the rational functions defining an isogeny
"""
function _compute_rational_fractions(phi::Isogeny)
	phixn, phiyn = _normalized_rational_fractions(phi)
	u, r, s, t = isomorphism_coefficients(phi)
	return (u^2 * phixn + r, u^3 * phiyn + s * u^2 * phixn + t)
end

function _normalized_rational_fractions(phi::Isogeny)
	l = degree(phi)
	E = domain(phi)
	K = base_ring(E)
	b2, b4, b6, b8 = b_invariants(E)
	a1, a2, a3, a4, a6 = a_invariants(E)

	if l==1 #phi is an isomorphism
		A, x = PolynomialRing(K, "x")
		Ffield = FractionField(A)
		return (Ffield(x), Ffield(1))

	else #remember l is odd

		n = div(l - 1, 2)
		kpoly = kernel(phi)
		s1 = - coeff(kpoly, n-1)
		kpolysquare = kpoly^2
		dkpoly = derivative(kpoly)
		d2kpoly = derivative(dkpoly)

		FX = parent(kpoly)
		psi2 = FX(divisionpolynomial(E, -1))
		X = gen(FX)
		N = psi2 * (dkpoly^2 - d2kpoly * kpoly) - (6 * X^2 + b2 * X + b4) * dkpoly * kpoly + (l * X - 2 * s1) * kpolysquare

		if K(2) != 0
			dN = derivative(N)
			omega = dN * kpoly - 2 * N * dkpoly    #+ (1 // K(2)) * ((a1 * N + a3 * kpolysquare) * kpoly)(x): we force to have a1 = a3 = 0
			return (N // kpolysquare , omega // (kpoly * kpolysquare))

		else #K has characteristic 2
			throw(NotImplementedError("Rational fractions not yet implemented in characteristic 2"))
		end
	end
end

######################################################################
# Composition
######################################################################

function compose(phi1::Isogeny, phi2::Isogeny)
	E1 = domain(phi1)
	E2 = image(phi1)
	E3 = image(phi2)
	@assert E2 == domain(phi2)

	phi1x, phi1y = rational_fractions(phi1)
	phi2x, phi2y = rational_fractions(phi2)

	resx = phi2x(phi1x)
	resy = phi1y * phi2y(phi1x)

	return Isogeny(E1, E3, resx, resy)
end

######################################################################
# Defining isogenies with various inputs
######################################################################

"""
Build an isogeny given its domain, kernel, degree, and image
"""
function Isogeny(E::AbstractWeierstrass{T}, poly::PolyElem{T}, Eprime::AbstractWeierstrass{T}) where T
	phi1 = Isogeny(E, poly)
	postisomorphism = find_isomorphism(image(phi1), Eprime)
	phi = compose(phi1, postisomorphism)
	return phi
end

function Isogeny(E1::AbstractWeierstrass{T}, E2::AbstractWeierstrass{T}, phix::Frac{U}, phiy::Frac{U}) where U<:PolyElem{T} where T
	K = sqrt(denominator(phix))
	F = base_ring(K)
	normalizedphix, normalizedphiy = rational_fractions(Isogeny(E1, K))
	usquare = phix//normalizedphix
	ucube = phiy//normalizedphiy

	#convert to base field
	num = numerator(usquare)
	den = denominator(usquare)
	@assert den == 1
	@assert degree(num) == 0
	usquare = coeff(num, 0)
	num = numerator(ucube)
	den = denominator(ucube)
	@assert den == 1
	@assert degree(num) == 0
	ucube = coeff(num, 0)

	u = ucube // usquare
	testphiy = u^3 * normalizedphiy
	if testphiy == - phiy
		u = -u
	end
	@assert degree(numerator(phix)) == 2 * degree(K) + 1
	res = Isogeny(E1, 2 * degree(K) + 1, K, u, F(0), F(0), F(0), E2, phix, phiy)
	@assert isvalid(res)
	return res
end

"""
Build an isogeny given its domain and its kernel polynomial.
"""
function Isogeny(E::AbstractWeierstrass{T}, poly::PolyElem{T}) where T

	CurveType = curvetype(E)

	K = base_ring(E)
	a1, a2, a3, a4, a6 = a_invariants(E)
	b2, b4, b6, b8 = b_invariants(E)
	n = Nemo.degree(poly)

	coeff(poly, n) == 1 || throw(ArgumentError("kernel polynomial must be monic"))

	s1, s2, s3 = zero(K), zero(K), zero(K)

	n > 0 && (s1 = - coeff(poly, n - 1))
	n > 1 && (s2 = coeff(poly, n - 2))
	n > 2 && (s3 = - coeff(poly, n - 3))
	t = 6 * (s1^2 - 2 * s2) + b2 * s1 + n * b4
	w = 10 * (s1^3 - 3 * s1 * s2 + 3 * s3) + 2 * b2 * (s1^2 - 2 * s2) + 3 * b4 * s1 + n * b6
	E1 = CurveType(a1, a2, a3, a4 - 5 * t, a6 - b2 * t - 7 * w)

	return Isogeny(E, 2 * n + 1, poly, K(1), K(0), K(0), K(0), E1, nothing, nothing)
end

######################################################################
# BMSS algorithm
######################################################################

include("bmss.jl")

######################################################################
# Isogenies between Weierstrass curves
######################################################################

"""
Get the polynomial associated to an l-torsion rational point, l being an odd prime.

The input is not checked.
"""
function subgroup(Q::EllipticPoint{T}, l::Nemo.Integer) where T<:Nemo.FieldElem
	normalize!(Q)
	K = base_ring(Q)
	R, x = PolynomialRing(K, "x")
	poly = Nemo.one(R)
	point = Q
	for k = 1 : div(l-1, 2)
		poly *= (x - point.X)
		point += Q
	end
	return poly
end

"""
Build an isogeny given an elliptic curve in short Weierstrass form, the j-invariant of the targetted elliptic curve, and the degree, assuming it exists.
Returns a normalized isogeny.
If the keyword argument USE_ATKIN is set to false (default), we use classical modular equations. If USE_ATKIN is set to true, we use Atkin modular equations.
"""

function Isogeny(E::ShortWeierstrass{T}, degree::Nemo.Integer, jprime::T; USE_ATKIN=false) where T
	K = base_ring(E)
	j = j_invariant(E)
	KX, X = PolynomialRing(K, "X")
	KXY, Y = PolynomialRing(KX, "Y")
	if USE_ATKIN
		print("using Atkin polys")
		A_ell_j = ClassPolynomials.AtkinModularPolynomial(degree, X, j)
		A_ell_jprime = ClassPolynomials.AtkinModularPolynomial(degree, X, jprime)
		poly = gcd(A_ell_j, A_ell_jprime)
		test, f = any_root(poly)
		if !test
			throw(ArgumentError("The curves are not linked by a rational isogeny of this degree"))
		end
		return Isogeny(E, degree, jprime, f)
	else
		Phi_l = ClassPolynomials.ClassicalModularPolynomial(degree, X, Y)
		#polynomial in Y over polynomials in X
		l = K(degree)

		derx = derivative(Phi_l)(j)(jprime)
		dery = derivative(Phi_l)(jprime)(j) #works because Phi_l is symmetric

		J = - K(18) // l * (E.b // E.a) * (derx // dery) * j
		jj = jprime * (jprime - 1728)

		aprime = - l^4 * J^2 // (48 * jj)
		bprime = - l^6 * J^3 // (864 * jprime * jj)
		Eprime = ShortWeierstrass(aprime, bprime)
		poly = kernelpoly(E, Eprime, degree)
		return Isogeny(E, degree, poly, K(1), K(0), K(0), K(0), Eprime, nothing, nothing)
	end
end

"""
Another element in the base field may be given as argument. In this case we use the Atkin modular A_\ell equations instead of the classical ones,
and f must satisfy A_\ell(j, f) == A_\ell(jprime, f) == 0. The curve must be given in Short Weierstrass form.
"""
function Isogeny(E::ShortWeierstrass{T}, degree::Nemo.Integer, jprime::T, f::T) where T
	j = j_invariant(E)
	_, _, _, a, b = a_invariants(E)
	K = base_ring(E)
	@assert ClassPolynomials.AtkinModularPolynomial(degree, f, j) == 0
	@assert ClassPolynomials.AtkinModularPolynomial(degree, f, jprime) == 0 #sanity checks
	KX, X = PolynomialRing(K, "X")
	dF = derivative(ClassPolynomials.AtkinModularPolynomial(degree, X, j))(f)
	dF *= f
	dF2 = derivative(ClassPolynomials.AtkinModularPolynomial(degree, X, jprime))(f)
	dF2 *= f
	poly = derivative(ClassPolynomials.AtkinModularPolynomial(degree, f, X))
	dJ = j * poly(j)
	dJ2 = degree * jprime * poly(jprime)
	jj = jprime // (jprime - 1728)
	lambda = (dF2 * dJ * b) // (dJ2 * dF * a)
	aprime = - K(27)//K(4) * degree^4 * lambda^2 * jj
    bprime = - K(27)//K(4) * degree^6 * lambda^3 * jj
	Eprime = ShortWeierstrass(aprime, bprime)
	poly = kernelpoly(E, Eprime, degree)

	return Isogeny(E, degree, poly, K(1), K(0), K(0), K(0), Eprime, nothing, nothing)
end

"""
Build an isogeny given an odd integer l and a rational torsion point of this order. The input is not checked.
"""

function Isogeny(E::AbstractWeierstrass{T}, Q::EllipticPoint{T}, l::Nemo.Integer) where T<:Nemo.FieldElem
	poly = subgroup(Q, l)
	return Isogeny(E, poly)
end

"""
Build an isogeny of given degree between two curves, assuming a normalized separable isogeny of this degree exists.
"""

function Isogeny(E1::AbstractWeierstrass{T}, degree::Nemo.Integer, E2::AbstractWeierstrass{T}) where T<:Nemo.FieldElem
	poly = kernelpoly(E1, E2, degree)
	return Isogeny(E1, poly)
end

"""
Compute the scalar multiplication by m on a weierstrass curve as a normalized isogeny.
"""
function multiplication_isogeny(E::AbstractWeierstrass, m::Nemo.Integer)
	@assert isodd(m)
	K = base_ring(E)
	psim = divisionpolynomial(E, m)
	n = degree(psim)
	phi1 = Isogeny(E, 1 // coeff(psim, n) * psim)
	E2 = image(phi1)
	postisomorphism = isomorphism(E2, 1 // K(m))
	@assert image(postisomorphism) == E
	return compose(phi1, postisomorphism)
end

######################################################################
# Dual isogeny
######################################################################

function dual(phi::Isogeny)
	K = base_ring(phi)
	E1 = domain(phi)
	E2 = image(phi)
	d = degree(phi)
	E1_0 = image(isomorphism(E1, K(d)))
	phi_hat = Isogeny(E2, d, E1_0) #use BMSS
	#de-normalize !
	postisomorphism = isomorphism(image(phi_hat), 1//K(d))
	phi_hat = compose(phi_hat, postisomorphism)
	@assert isvalid(phi_hat)
	@assert image(phi_hat) == E1
	return phi_hat
end

function isomorphism(E::AbstractWeierstrass{T}, u::T) where T
	a1, a2, a3, a4, a6 = a_invariants(E) #we have a1 = a3 = 0
	a2prime = a2 * u^2
	a4prime = a4 * u^4
	a6prime = a6 * u^6
	K = base_ring(E)
	R, X = PolynomialRing(base_ring(E), "X")
	phi = Isogeny(E, 1, R(1), u, K(0), K(0), K(0), EllipticCurve(a2prime, a4prime, a6prime), nothing, nothing)
	return phi
end

function find_isomorphism(E1::AbstractWeierstrass{T}, E2::AbstractWeierstrass{T}) where T
	a11, a21, a31, a41, a61 = a_invariants(E1)
	a12, a22, a32, a42, a62 = a_invariants(E2)
	u = sqrt(a22 // a21)
	@assert a42 = a41 * u^4
	@assert a62 = a61 * u^6
	return isomophism(E1, u)
end

######################################################################
# Operations
######################################################################

function +(phi::Isogeny, psi::Isogeny)
	E = domain(phi)
	@assert E == domain(psi)
	Eprime = image(phi)
	@assert Eprime == image(psi)
	phix, phiy = rational_fractions(phi)
	psix, psiy = rational_fractions(phi)
	throw(NotImplementedError("Use the addition laws on the image curve"))
	#return Isogeny(E, Eprime, phix + psix, phiy + psiy)
end

function -(phi::Isogeny)
	pix, phiy = rational_fractions(phi)
	return Isogeny(domain(phi), image(phi), phix, -phiy)
end

function -(phi::Isogeny, psi::Isogeny)
	return phi + (-psi)
end


@with_kw mutable struct BBParameters{T <: AbstractFloat} <: AbstractPolicyParameters
    γ₀::T = 1.
end

mutable struct BB <: AbstractStep
    params::BBParameters{Float64}
    xprev::Vector{Float64}
    gprev::Vector{Float64}

    function (::Type{BB})(; kwargs...)
        return new(BBParameters(; kwargs...))
    end
end

function initialize!(bb::BB,x₀::Vector{Float64})
    bb.xprev = zeros(length(x₀))
    bb.gprev = zeros(length(x₀))
end

function stepsize(bb::BB,klocal::Integer,kglobal::Integer,fval::Real,x::AbstractVector,g::AbstractVector)
    if kglobal == 1
        bb.xprev .= x
        bb.gprev .= g
        return bb.params.γ₀
    end
    s = x - bb.xprev
    y = g - bb.gprev
    η = norm(s,2)^2/(s⋅y)
    bb.xprev .= x
    bb.gprev .= g
    return η
end


"""
    This package contains a number of kernels frequently used in SPH, as well as their functions and derivatives an 2D and 3D.
"""

module SPHKernels

    export  kernel_value_1D,     𝒲₁,
            kernel_value_2D,     𝒲₂,
            kernel_value_3D,     𝒲₃,
            kernel_deriv_1D,    ∇𝒲₁,
            kernel_deriv_2D,    ∇𝒲₂,
            kernel_deriv_3D,    ∇𝒲₃,
            bias_correction_1D, δρ₁,
            bias_correction_2D, δρ₂,
            bias_correction_3D, δρ₃,
            SPHKernel,
            Cubic,
            Quintic,
            WendlandC2,
            WendlandC4,
            WendlandC6,
            WendlandC8

    """
        SPHKernel

    Supertype for all SPH kernels.
    """
    abstract type SPHKernel end

    include("bsplines/Cubic.jl")
    include("bsplines/Quintic.jl")
    include("wendland/C2.jl")
    include("wendland/C4.jl")
    include("wendland/C6.jl")
    include("wendland/C8.jl")

    # multiple dispatch for nicer look

    """
        𝒲₁( kernel::SPHKernel, u::Real, h_inv::Real)

    Evaluate 1D spline at position ``u = \\frac{x}{h}``.
    """
    𝒲₁( kernel::SPHKernel, u::Real, h_inv::Real) = kernel_value_1D(kernel, u, h_inv)

    """
        ∇𝒲₁( kernel::SPHKernel, u::Real, h_inv::Real)

    Evaluate 1D derivative at position ``u = \\frac{x}{h}``.
    """
    ∇𝒲₁(kernel::SPHKernel, u::Real, h_inv::Real) = kernel_deriv_1D(kernel, u, h_inv)

    """
        𝒲₂( kernel::SPHKernel, u::Real, h_inv::Real)

    Evaluate 2D spline at position ``u = \\frac{x}{h}``.
    """
    𝒲₂( kernel::SPHKernel, u::Real, h_inv::Real) = kernel_value_2D(kernel, u, h_inv)

    """
        ∇𝒲₂( kernel::SPHKernel, u::Real, h_inv::Real)

    Evaluate 1D derivative at position ``u = \\frac{x}{h}``.
    """
    ∇𝒲₂(kernel::SPHKernel, u::Real, h_inv::Real) = kernel_deriv_2D(kernel, u, h_inv)

    """
        𝒲₃( kernel::SPHKernel, u::Real, h_inv::Real)

    Evaluate 3D spline at position ``u = \\frac{x}{h}``.
    """
    𝒲₃( kernel::SPHKernel, u::Real, h_inv::Real) = kernel_value_3D(kernel, u, h_inv)

    """
        ∇𝒲₃( kernel::SPHKernel, u::Real, h_inv::Real)

    Evaluate 1D derivative at position ``u = \\frac{x}{h}``.
    """
    ∇𝒲₃(kernel::SPHKernel, u::Real, h_inv::Real) = kernel_deriv_3D(kernel, u, h_inv)

    """
        δρ₁(kernel::SPHKernel, density::Real, m::Real, h_inv::Real)

    Corrects the 1D density estimate for the kernel bias. See Dehnen&Aly 2012, eq. 18+19.
    """
    δρ₁(kernel::SPHKernel, density::Real, m::Real, h_inv::Real) = bias_correction_1D(kernel, density, m, h_inv)

    """
        δρ₂(kernel::SPHKernel, density::Real, m::Real, h_inv::Real)

    Corrects the 2D density estimate for the kernel bias. See Dehnen&Aly 2012, eq. 18+19.
    """
    δρ₂(kernel::SPHKernel, density::Real, m::Real, h_inv::Real) = bias_correction_2D(kernel, density, m, h_inv)

    """
        δρ₃(kernel::SPHKernel, density::Real, m::Real, h_inv::Real)

    Corrects the 3D density estimate for the kernel bias. See Dehnen&Aly 2012, eq. 18+19.
    """
    δρ₃(kernel::SPHKernel, density::Real, m::Real, h_inv::Real) = bias_correction_3D(kernel, density, m, h_inv)

end # module


using MRIReco

##### simple example ####

N = 256
I = shepp_logan(N)
I = circularShutterFreq!(I,1)
cmap = 1im*quadraticFieldmap(N,N,125*2pi)

# simulation parameters
params = Dict{Symbol, Any}()
params[:simulation] = "fast"
params[:trajName] = "Spiral"
params[:numProfiles] = 1
params[:numSamplingPerProfile] = N*N
params[:windings] = 128
params[:AQ] = 3.0e-2
params[:correctionMap] = cmap[:,:,1]

# do simulation
acqData = simulation(I, params)

# reco parameters
params = Dict{Symbol, Any}()
params[:reco] = "direct"
params[:reconSize] = (N,N)
Ireco = reconstruction(acqData, params)

# export image
filename = joinpath(dirname(pathof(MRIReco)),"../docs/src/assets/fieldmapReco1.png")
exportImage(filename, abs.(Ireco[:,:,1,1,1]) )

# now with fieldmap correction
params[:reco] = "standard"
params[:regularization] = "L2"
params[:iterations] = 3
params[:solver] = "admm"
params[:correctionMap] = cmap
params[:alpha] = 1.75
params[:m] = 4.0
params[:K] = 28

Ireco = reconstruction(acqData, params)

# export images
filename = joinpath(dirname(pathof(MRIReco)),"../docs/src/assets/fieldmapReco2.png")
exportImage(filename, abs.(Ireco[:,:,1,1,1]) )

# export fieldmap image
cmap_ = imag(cmap[:,:,1]) .+ 125*2pi
filename = joinpath(dirname(pathof(MRIReco)),"../docs/src/assets/fieldmap.png")
exportImage(filename, abs.(cmap_) )


using Test; using Laplacians; using MAT; using LinearAlgebra; using BenchmarkTools
include("../src/CombinatorialMultigrid.jl");

using Main.CombinatorialMultigrid

## load example matrix

file = matopen("../example/X.mat"); X = read(file, "X"); close(file)
LX = lap(X);
b1 = rand(Float64, size(X, 1));
b1 = b1 .- sum(b1)/length(b1);

@info "###  Running CMG ###"
t = @elapsed (pfunc, h) = cmg_preconditioner_lap(LX);
@info "Time Required to build CMG Solver: $(t) seconds"
t = @elapsed x = pfunc(b1);
@info "Time Required to find x: $(t) seconds"

@info "###  Running PCG with CMG    ### "
## solve with pcg and cmg preconditioner
f1 = pcgSolver(LX,pfunc);
t = @elapsed x = f1(b1, maxits=40, tol=1e-6,verbose=true);
@info "Time Required to solve system: $(t) seconds"

## solve with approxchol_lap from laplacians
@info "###  Running approxchol_lap    ###"
t = @elapsed solver = approxchol_lap(X; tol=1e-6, maxits=1000, maxtime=Inf, verbose=true, pcgIts=Int[], params=ApproxCholParams());
@info "Time Required to build Lap Solver: $(t) seconds"
t = @elapsed x = solver(b1);
@info "Time Required to find x: $(t) seconds"


function alexnet()
  layers = Chain(Conv((11, 11), 3=>64, stride=(4, 4), relu, pad=(2, 2)),
                 MaxPool((3, 3), stride=(2, 2)),
                 Conv((5, 5), 64=>192, relu, pad=(2, 2)),
                 MaxPool((3, 3), stride=(2, 2)),
                 Conv((3, 3), 192=>384, relu, pad=(1, 1)),
                 Conv((3, 3), 384=>256, relu, pad=(1, 1)),
                 Conv((3, 3), 256=>256, relu, pad=(1, 1)),
                 MaxPool((3, 3), stride=(2, 2)),
                 AdaptiveMeanPool((6,6)),
                 flatten,
                 Dropout(0.5),
                 Dense(256 * 6 * 6, 4096, relu),
                 Dropout(0.5),
                 Dense(4096, 4096, relu),
                 Dense(4096, 1000))
  Flux.testmode!(layers, false)
  return layers
end


struct RateParameter
end

@with_kw mutable struct Rate
    param::RateParameter = RateParameter()
    N::SNNInt = 100
    x::Vector{SNNFloat} = 0.5randn(N)
    r::Vector{SNNFloat} = tanh.(x)
    g::Vector{SNNFloat} = zeros(N)
    I::Vector{SNNFloat} = zeros(N)
    records::Dict = Dict()
end

function integrate!(p::Rate, param::RateParameter, dt::SNNFloat)
    @unpack N, x, r, g, I = p
    @inbounds for i = 1:N
        x[i] += dt * (-x[i] + g[i] + I[i])
        r[i] = tanh(x[i]) #max(0, x[i])
    end
end


include("Density.jl");
include("BiasPotential.jl");

function calc_delta_potential_of_mean_force(umbrella_center, data, sigma, kbt, input)
    delta_pmf = -kbt * log(calc_density(input, data, sigma))
    delta_pmf += kbt * log(calc_density(umbrella_center, data, sigma))
    delta_pmf -= calc_bias_potential(umbrella_center, input) - calc_bias_potential(umbrella_center, umbrella_center)
    return delta_pmf
end

function calc_potential_of_mean_force(umbrella_centers, data, weight, sigma)
    pmf = []
    for i = 1:length(data)
        tmp = 0.0
        for j = 1:length(umbrella_centers)
              tmp += weight[j] * get_gaussian(data[i], umbrella_centers[j], sigma)
        end
        push!(pmf, tmp)
    end

    return pmf
end


import MDP

p1 = 0.1; q1 = 0.3;
p2 = 0.1; q2 = 0.1;

const P1 = [1 - p1 p1; q1 1 - q1];
const P2 = [1 - p2 p2; q2 1 - q2];

const size1 = size(P1,1)
const size2 = size(P2,1)

const M = 100

# Reachable states
const pi1 = [ P1^n for n = 1:M ]
const pi2 = [ P2^n for n = 1:M ]

# Distortion
const d1 = 1 - eye(size1);
const d2 = 1 - eye(size2);

const Davg1 = [ minimum(pi1[n]*d1,1) for n = 1:M ]
const Davg2 = [ minimum(pi2[n]*d1,1) for n = 1:M ]

function bellmanUpdate(v; discount=1.0)
    # Assume that v is size1 x size2 x size1 x M x size2 x M
    v_next = zeros(Float64, size(v))
    g_next = zeros(Int,     size(v))

    # v_next(s1, s2, z1, k1, z2, k2)

    W1 = zeros(Float64, (size1, size2, M))
    W2 = zeros(Float64, (size2, size1, M))

    # We go for a dumb solution to see how fast it runs.
    # If speed becomes an issue, refactor this to reduce the number of loops

    # Calculate W1
    for k2 = 1:M, z2 = 1:size2, s1 = 1:size1
        k2_next = min(k2+1,M)
        # Note that the cost to go is just a quadratic form
        # cost_to_go = 0 ;
        # for s2_next = 1:size2
        #     for s1_next = 1:size1
        #         cost_to_go += ( pi1[1][s1,s1_next] * pi2[k2_next][z2,s2_next]
        #                         * v(s1_next, s2_next, s1, 1, z2, k2_next)
        #                         )
        #     end
        # end
        cost_to_go = pi1[1][s1,:] * v[:,:, s1, 1, z2, k2_next] * pi2[k2_next][z2,:]'
        # In Julia 1x1 matrix is not a scalar. So we need to index cost_to_go
        W1[s1,z2,k2] = Davg2[k2][z2] + discount * cost_to_go[1,1]
    end

    # Calculate W2
    for k1 = 1:M, z1 = 1:size1, s2 = 1:size2
        k1_next = min(k1+1,M)

        cost_to_go = pi1[k1_next][z1,:] * v[:,:, z1, k1_next, s2, 1] * pi2[1][s2,:]'
        W2[s2,z1,k1] = Davg1[k1][z1] + discount * cost_to_go[1,1]
    end

    # Calculate v_next and g_next
    for k2 = 1:M, z2 = 1:size2, k1 = 1:M, z1 = 1:size1, s2 = 1:size2, s1 = 1:size1
        # Choose the option will lower cost
        w1  = W1[s1, z2, k2]
        w2  = W2[s2, z1, k1]

        index = (s1, s2, z1, k1, z2, k2)
        (v_next[index...], g_next[index...]) = (w1 < w2)? (w1, 1) : (w2, 2)
    end

    return (v_next, g_next)

end

model = MDP.DynamicModel(bellmanUpdate; objective=:Min)
v_initial = zeros(size1, size2, size1, M, size2, M)

@time (v,g) = MDP.valueIteration(model, v_initial; discount=0.9)

println(v[1])

# z1, z2 = 1, 1
# s1, s2 = 1, 1
#
# policy = reshape(g[s1,s2,z1,:,z2,:], (M,M))


#=
    CODE FOR SIMULATION AND CONTINUATION OF
    SWIFT-HOHENBERG with Non-local RAMAN response:

        ∂u/∂t = η + μu - u³ + β ∂²u/∂τ² - 4/3 ∂⁴u/∂τ⁴ + ∫χ(τ-τ')u(τ')dτ'

    Author: Martin Bataille-Gonzalez.

    -------------------------------
    Where,
        χ(τ) = √3/2 a fᵣ * exp(-τ₀τ/τ₂) * sin(τ₀τ/τ₁)
    and,
        a = τ₀(τ₁² + τ₂²) / (τ₁τ₂²)

    Parameters will be,
        τ₀ = 14 fs
        fᵣ = 0.18
        β = -1.5
        y = 0.12 = η (?)
        σ = -0.35 ⟹ μ = σ - √3/2fᵣ
    From, Parra-Rivas et al (2021) PRA.
        fᵣ = 0.18
        τ₁ = 12.2 fs
        τ₂ = 32 fs
=#
using DifferentialEquations

using Plots

function χ(τ::Number, p)
    return sqrt(3) / 2 * p[:a] * p[:fᵣ] * exp(-p[:τ₀]/p[:τ₂] * τ) * sin.(p[:τ₀]/p[:τ₁] * τ)
end

function χ(τ::AbstractArray, p)
    [χ(τᵢ, p) for τᵢ=τ]
end

function spectral_derivative(u, Δx, order)
    N = length(u)

    k = fftfreq(N) * N
    L = (N-1) * Δx

    return ifft( (2π / L * im) ^ order * k.^ order .* fft(u))
end

∫χu(u) = ifft(ft_coupling * fft(u))

∂ₓₓ(u, Δx) = spectral_derivative(u, Δx, 2)
∂ₓₓₓₓ(u, Δx) = spectral_derivative(u, Δx, 4)

# Parameters
L = 300
Δx = 0.1
τ = 0:Δx:L
τ₀ = 14
fᵣ = 0.18
β = -1
η = 0.0
σ = -0.35
μ = σ - sqrt(3)/2 * fᵣ
τ₁ = 12.2
τ₂ = 32
a = τ₀ * (τ₁ ^ 2 + τ₂ ^ 2) / (τ₁ * τ₂ ^ 2)

p = (τ₀ = τ₀, fᵣ = fᵣ, β = β, η = η, μ = μ, τ₁ = τ₁, τ₂ = τ₂, a=a, L=L, Δx = Δx)

coupling = χ(τ, p)

plot(τ, coupling)

using FFTW

FFTW.set_provider!("mkl")
ft_coupling = fft(coupling)

coupling[200]
decay = τ₂ / τ₀

function SHRaman(u, p, t)
    return p[:η] .+ p[:μ] * u .- u .^ 3 .+ p[:β] * ∂ₓₓ(u, p[:Δx])
            .- 4 / 3 * ∂ₓₓₓₓ(u, p[:Δx]) .+ ∫χu(u)
end

u = im * zeros(length(τ))

tspan = (0.0, 1.0)
prob = ODEProblem(SHRaman, u, tspan, p)
sol = solve(prob, RK4())

plot(τ, real.(sol(1.0)))

u = SHRaman(u, p, 0)


using Test, JLD2, UnicodePlots
using LinearAlgebra, ForwardDiff, Arpack
using SparseArrays, SparseDiffTools, Symbolics

# load mesh
meshfile = "$(@__DIR__)/data/squaremesh.jld2"
JLD2.@load(meshfile, p, t, jacK, jacM, ps, ts, jacKs, jacMs)
p, t, jacK, jacM =  ps, ts, jacKs, jacMs

np, nt = size(p, 1), size(t, 1)
ε = 1e-6
x, dp = p, 2 * rand(size(p)...) .- 1
xp, xm = p .+ ε * dp, p .- ε * dp
u, v = 2 * rand(np) .- 1, 2 * rand(np) .- 1

# finite differences
f = x -> FemAdjoint.costKproduv(x, t, u, v)
fp, fm = f(xp), f(xm)
g =  FemAdjoint.∇costKproduv(x, t, u, v, jacK = jacK, jacM = jacM)
println("")
println(fp)
println(fm)
println(" ")
println(dot(g, dp))
println((fp - fm) / (2 * ε))
println("")
@testset "test eval gradient dot function for costKproduv" begin
    @test abs(dot(g, dp) - (fp - fm) / (2 * ε)) < 1e-4
end