function relative_prime(a, b)
  gcd(a, b) == 1
end

function find_an_e_that_is_relative_prime(p, q)
  p1q1 = (p-1) * (q-1)

  candidate = -1
  found = false

  while !found
    # Select a random candidate and then loop until relative_prime or reached p1q1
    candidate = round( 1 + rand() * p1q1 )

    while candidate < p1q1 && !relative_prime(candidate, p1q1)
      candidate += 1
    end

    if relative_prime(candidate, p1q1)
      found = true
    end
  end

  return candidate
end

function find_a_d(p, q, e)
  p1q1 = (p-1) * (q-1)
  sgd = gcd(e, p1q1)
  d = e + p1q1 / sgd
  return d
end

function rand_big_int(keylength = 4096)
  bigint = BigInt(0)
  bits_left = keylength
  max_bits = 30
  while bits_left > 0
    numbits = min(bits_left, max_bits)
    bigint = (bigint << numbits) + (rand(1:2^numbits) - 1)
    bits_left -= numbits
  end
  bigint
end

type RSAKey
  # Public key = (n, e) and Private key = (n, d)
  n::Integer
  e::Integer
  d::Integer

  RSAKey(p = rand_big_int(4096)) = begin

    q = next_prime(p)
    n = p * q
    e = find_an_e_that_is_relative_prime(p, q)
    d = find_a_d(p, q, e)

    new(n, e, d)
  end
end

type RSAPublicKey
  n::Integer
  e::Integer
  RSAPublicKey(n::Integer, e::Integer) = new(n, e)
end

function public_key(rsa::RSAKey)
  RSAPublicKey(rsa.n, rsa.e)
end

function encrypt(rsa::Union(RSAKey, RSAPublicKey), P::Integer)
  fast_exponentiation(P, rsa.e, rsa.n)
end

function uint8_array_to_integer(arr::Array{Uint8,1})
  integer = BigInt(0)
  for i = 1:length(arr)
    integer = (integer << 8) + arr[i]
  end
  integer
end

function integer_to_uint8_array(x::Integer)
  bytes = Uint8[]

  while x > 0
    b = uint8(x & 0xff)
    push!(bytes, b)
    x = x >>> 8
  end

  return reverse(bytes)
end

# Convert a string to an integer for RSA encryption.
function string_to_integer(s::String)
  bytes = convert(Vector{Uint8}, utf8(s))
  uint8_array_to_integer(bytes)
end

using Codecs

# Convert a non-negative integer i into an octet string and then base64
# encode it so we can pass it around.
function integer_to_base64_string(x::Integer, len = nothing)
  if typeof(len) <: Integer && (x >= 256^len)
    throw("integer is too large")
  end

  if x < 0
    throw("integer is negative")
  end

  bytes = Uint8[]

  while x > 0
    b = uint8(x & 0xff)
    push!(bytes, b)
    x = x >>> 8
  end

  bytes = reverse(bytes)

  if typeof(len) <: Integer && (length(bytes) < len)
    bytes = [zeros(Uint8, len - str), bytes]
  end

  encoded = encode(Base64, bytes)

  return bytestring(encoded)
end

function base64_string_to_integer(s::String)
  decoded = decode(Base64, s)
  uint8_array_to_integer(decoded)
end

function encrypt(rsa::Union(RSAKey, RSAPublicKey), P::String)
  C = fast_exponentiation(string_to_integer(P), rsa.e, rsa.n)
end

function decrypt(rsa::RSAKey, C::String)
  P = fast_exponentiation(string_to_integer(C), rsa.d, rsa.n)
  integer_to_string(P)
end


### A Pluto.jl notebook ###
# v0.17.4

using Markdown
using InteractiveUtils

# ╔═╡ a2a0a0d7-b457-40bd-ae6b-da14cc557031
begin
	using PlutoUI
	bigbreak = HTML("<br>"^4)
	nothing
end

# ╔═╡ 6643612d-d1a6-4fdc-8398-e12682891eb0
using LinearAlgebra

# ╔═╡ 9426ed50-83ae-11ec-34f2-2f8e36fc8d05
md"""
# Composite types aka `structs`

If your project growths it gets more important to keep things modular.
"""

# ╔═╡ 869614f3-f243-4319-9f8f-64c3cc13eb97
md"## 1. Inmutable structs

The definition is rather easy. It needs name and a list of its variables (optionally with types)."

# ╔═╡ 70d2151b-0de1-41c2-a73a-7ee6f01a2d9f
struct Point
	x
	y
end

# ╔═╡ 24eabc0b-496b-4159-8da6-29910b85ea4f
md"There will always be a default constructor:"

# ╔═╡ 996a3423-bd55-49d7-9401-d270b6e109c1
p = Point(1.0, 2.0)

# ╔═╡ 909cc091-b321-4752-9f7a-62d1aa3a3a3d
p.x

# ╔═╡ be82d829-d6c3-4822-8f57-3f0727776817
p.y

# ╔═╡ 3528ea22-89da-4330-9b66-49fdfef87687
# p.x = 10   # does not work!

# ╔═╡ 260566f0-c909-41a2-9ebf-ddb8103635f4

# ╔═╡ 8f62d5bd-764b-49e5-9b00-05d9f90332f9
md"Two typical use cases of inmutable structs:

- Parameters for a solver which do not change during the simulation.
- Implementation of useful syntax, by methods functions for your type.
- Realisation of an interface for your data. (Construction of a mutable struct is often almost zero overhead.)
"

# ╔═╡ 1515c0ab-f7ed-4273-9b64-c810abbaf941
md"### Example (`FastDiagonal`)"

# ╔═╡ 5790a006-640d-4751-a690-b139654049c6
md"Let's implement a matrix like type, which only stores diagonal elements."

# ╔═╡ 62ee2d5b-d757-47db-9019-37622330e9bc
begin
	struct FastDiag{T}
	    data::Vector{T}
		z
	end

	import Base: *, +
	function *(A::FastDiag{T}, x::Vector{T}) where T
	    return A.data .* x
	end

	function +(A::FastDiag{T}, B::FastDiag{T}) where T
	    return FastDiag{T}( A.data + B.data )
	end

	nothing # just to surpress Pluto output
end

# ╔═╡ 8fabcd94-1c71-420c-b2b9-d10a6eca76dc
begin
	Base.@kwdef struct Foo
		x
		y::Float64 = 1.0
	end
end

# ╔═╡ bbe30b74-8e29-4b3f-bf04-e05e67c82dbf
FastDiag([1], 1)

# ╔═╡ dbd7131b-05a9-4d70-8864-d5ca9b2fe9c5
md"It is always good, if your type is `concrete`. That means Julia can predict each variables type in advance."

# ╔═╡ 3c7c5961-014d-47c6-8ffa-cd06198682ce
isconcretetype(FastDiag{Float64})

# ╔═╡ b9a4a4ac-c07f-4f73-aa48-dbfa123a267a

# ╔═╡ cb9e96e1-597a-4b0e-912a-1972c84edd5f
begin
    N = 1000
    A = [i == j ? rand() : 0.0 for i in 1:N, j in 1:N]
	A_diag = collect(diag(A))
    A_fast = FastDiag(diag(A))

    x = rand(N)
end

# ╔═╡ 4bdf13bd-a27c-4fda-9807-eb12b1132ea0

# ╔═╡ 0fbd5d9a-30dc-4e73-9bff-c088768de5ab
md"Let's first make sure that our methods work correctly:"

# ╔═╡ 4a9409d1-d04d-459c-93a2-fd47315a1c27
(A+A) * x == (A_fast+A_fast) * x

# ╔═╡ 1dc516fd-b6de-466c-8cc2-ec56301cebba
md"Now, let's measure performance:"

# ╔═╡ 49044d10-11a2-44e7-9a86-481ddfd595c4
with_terminal() do
	printstyled("Full matrix             :", color = :red)
	@time A * x

	printstyled("Fast digonal type       :", color = :green)
	@time A_fast * x

	printstyled("On the fly construction :", color = :green)
	@time FastDiag(A_diag) * x
end

# ╔═╡ 35345440-550c-4723-9aec-e08142ff0802
md"### Success 🚀

By the way, there is a build in Diagonal type `Diagonal` in the `LinearAlgebra` standard package.

It has basically the same definition as our example!"

# ╔═╡ 7c12b757-8ca5-4d2b-92b0-d204f94e6d0b

# ╔═╡ 292ee470-e286-47ea-8065-92ad2fc824d7
Diagonal(A_diag) * x

# ╔═╡ 55485242-77c4-43ec-996d-c1177ab34ecf
bigbreak

# ╔═╡ ee6fa527-4d43-4b51-baa3-2cb2fa57d81b
md"### Are inmutable structs constant values? NO!"

# ╔═╡ f491e000-1f10-443d-bf18-18a413f8cf94
begin
	data = [1,2,3]
    D = FastDiag(data)

	D.data[2] = 100

	(;data, D)
end

# ╔═╡ f8afb728-5080-41f0-a0c4-9937416493b8
begin
	struct InM
		x
	end

	mutable struct Mut
		y::InM
	end
end

# ╔═╡ c13de59d-508e-4472-a602-17d334337ca4
md"
The variables (here `data`) are inmutable.

But if the value is of an mutable type, than we can mutate the underlying data indeed!"

# ╔═╡ 2dc06114-7f29-4021-943d-e2b14f787e3d
bigbreak

# ╔═╡ 95fbd51f-62a0-4167-9478-452420d19369
md"## 2. Mutable structs"

# ╔═╡ 91185b65-573b-42de-8238-3cd046aa5b98
md"We just have to add the keyword `mutable`. That's it!"

# ╔═╡ fa237ef5-f1df-4826-8062-81e061be64a2
mutable struct SimulationState
	X::Vector{Float64}
	t::Float64
end

# ╔═╡ f01b0192-2a01-4e04-905b-3500c84c52c1
s = SimulationState( [10.0, 20.0], 0.0)

# ╔═╡ 5e55db0d-6245-47e6-ac65-6a37d66758ae
s.t = 10.0

# ╔═╡ 1ab4ba54-b86c-46ec-8852-326223fce26b
s

# ╔═╡ 72fec38d-4c24-4356-97de-71a2914b821f
bigbreak

# ╔═╡ 82209778-7d13-49fd-9812-1d534ad79a49
md"## 3. Parametric types and functions"

# ╔═╡ 94b65eff-9df8-4621-9a92-74ae0a0c8052
md"Warning: This is more an advanced feature. If you want to use it, your will have to study it a bit more!"

# ╔═╡ 5fa76a28-b1d4-4c94-a7b8-5057115caa7a
struct Unpredictable
	x
end

# ╔═╡ aa7189d4-53a5-4c54-bfdb-c8b242be2613
struct Predictable{T}
	x::T
end

# ╔═╡ 7fc21ab5-52d9-4b12-a0c9-780cf2c62c75

# ╔═╡ 2c21a7fe-3223-46d4-af30-296045c905ec
a = Unpredictable(rand([1.0, 2, "3"]))

# ╔═╡ d7a61c32-0ed9-4d04-91e1-d6849057df2d
b = Predictable(rand([1.0, 2, "3"]))

# ╔═╡ fc842dae-373f-4ec5-815e-e25193fce8fe
with_terminal() do
	println("a = ", a.x, "\t     typeof(a) = ", typeof(a))
	println("b = ", b.x, "\t     typeof(b) = ", typeof(b))
end

# ╔═╡ 9c4c826b-f1a8-43f9-a50c-455826223402
bigbreak

# ╔═╡ 8ad53b1a-55c3-4a85-9f06-0ecd01305120
square(s) = s.x^2

# ╔═╡ 4f962f52-99e4-4d2c-a16c-7b187b03d9d6
md"#### Problem: Type instabilities:"

# ╔═╡ e02cf43c-3389-4020-b2cc-57f201cee0aa
md"The output of `@code_warntype` shows whenever some types are unpredictable at compile time:"

# ╔═╡ 0bfee97a-6b2c-42d7-b095-b356f5858bc0
with_terminal() do
	@code_warntype square(a)
end

# ╔═╡ 4120837d-749b-4813-b53f-0ab22bceacee
with_terminal() do
	@code_warntype square(b)
end

# ╔═╡ 08facfc2-4b2b-4055-891b-b8ea6b0b61d0
bigbreak

# ╔═╡ 27352174-d77c-4d61-bcf4-b897e451c5f6
md"### Parametric functions

You can define always a domain of types which you would like to allow.

**Syntax:**
```julia
function fnc(x::T, y::S) where {T, S <: AbstractVector}
   #
end
```
"

# ╔═╡ e355c3b9-470d-447a-8208-cf99dfbceadc
function pseudo_dot(x::Vector{T}, y::Vector{S}) where {T<:Number, S<:Number}
	res = zero(T)
	for i in 1:min(length(x), length(y))
		res += x[i] * y[i]
	end
	return res
end

# ╔═╡ 99f36b6d-ade6-4f37-8dfb-76c967c445d9
pseudo_dot( [1, 2, 3], [1.0, 1im])

# ╔═╡ cebfa1b4-6876-41ac-b219-032f061f44ac
bigbreak

# ╔═╡ 665b754d-be3e-4982-8e03-02d1b5fbbf39
md"## 4. Abstract types

You might want to define a type hiearachy for your types!

We cannot define relations directly between our custom structs. Instead, we use abstract types to carry the hierachy information!

It is up to convention/documentation what an abstract type actually represents.
Famous example `AbstractArray`, see also [Documentation - Interfaces](https://docs.julialang.org/en/v1/manual/interfaces/#Indexing)
"

# ╔═╡ 7bf7c3c4-85cd-4d86-bbb1-933a7d5c8328

# ╔═╡ 40d2b89f-802d-441d-8053-ffd95feb22a0
md"Let's consider this example"

# ╔═╡ 7f67b94e-495d-4605-b0f4-7978a323a798
abstract type AbstractAlgorithm end

# ╔═╡ 9b214cfc-a24e-4e4d-86e3-71451f647268

# ╔═╡ 4f55481e-7884-4dd8-92ab-14a9fedd2840
md"#### Our `AbstractAlgorithm` Interface:

We assume that everything which is more specific than `AbstractAlgorithm`
has at least a function
```julia
minimise(::Function, ::AbstractAlgorithm)
```

Now, how could someone implement an algorithm? Just define the needed functions!
"

# ╔═╡ 685f46ae-e3dd-4514-bba2-393602831d3a

# ╔═╡ b03b9de7-022a-4ba3-924a-92fdbf599453
f(x::AbstractAlgorithm) = x.n_steps

# ╔═╡ 7f40cabc-8ce8-4ecc-b22e-1dc62bd8b648
struct RandomSolver <: AbstractAlgorithm

	domain::Tuple{Float64, Float64}
	n_steps::Int64

end

# ╔═╡ 7e32559f-7148-472a-bff0-2777bd0df755
random(d::Tuple{Float64,Float64}) = d[1] + (d[2]-d[1]) * rand()

# ╔═╡ f6c9269f-56a1-4e23-a7c7-93783c149e84
function minimise(f, r::RandomSolver)
	minimum( f( random(r.domain) ) for i in 1:r.n_steps )
end

# ╔═╡ 75c1c92f-1e03-4ec3-86dc-65cfb2bfe35e

# ╔═╡ d2e013da-c752-4fe8-a020-c67a5e89c397
bigbreak

# ╔═╡ 9705e4c1-1fdf-40aa-be08-b2df34715d11
md"## Exercises"

# ╔═╡ bfb72e16-8e42-430b-921a-cec2818eaf24

# ╔═╡ 5d2e637e-3f40-49d2-9f4b-59b7c1ee6afe
md"""
7.1

Define your own type, which represents a point on a circle.
Implement a add function, so that two points on a circle can be added!

It should allow operations such as
```julia
CirclePoint( π ) + CirclePoint( π ) == CirclePoint( 0.0 )
```

"""

# ╔═╡ ef533b3f-b705-461f-a590-b0a11fa05f52

# ╔═╡ 38599bbf-6e45-49d8-a117-2d1822688bed
bigbreak

# ╔═╡ 16dc8591-a6b3-4256-bbd9-093548352449
md"""
7.2

Implement your own an `AbstractAlgorithm`, call it `GridAlgorithm` which is compatible with the code below!

The algorithm should try to find the minimum of a function by evaluation of equidistant grid points in it's domain.

"""

# ╔═╡ ac3f9a7f-ecdf-486a-a3df-a7bdd3d91012
struct GridAlgorithm

end

# ╔═╡ c577d146-2b43-4820-956e-c708861e8881
# minimise( x -> x^2,   GridAlgorithm( (-10, 10),  10000) )

# ╔═╡ d638bc5e-44b5-401b-a4ce-e47552c28960

# ╔═╡ f2d11973-64b6-47b5-b34a-c4ae9405be0b
md"Imagine someone wrote code to allow solution of multidimensional functions.

*Does this `ProcutAlgorithm` still work with your implementation?*"

# ╔═╡ 97c78e41-21e5-46b0-b06f-c754143cae0c
struct ProductAlgorithm <: AbstractAlgorithm
	algs::Vector{AbstractAlgorithm}
end

# ╔═╡ 4e2e930c-8e2b-4bd3-98f1-ae35055d7b81
function minimise(f::Function, r::ProductAlgorithm)
	N = length(r.algs)
	x0 = zeros(N)

	for i in 1:N
		f_i = x -> f( [( j == i ) ? x : x0[j] for j in 1:N] )

		x_i = minimise( f_i, r.algs[i] )
		x0[i] = x_i
	end

	return x0
end

# ╔═╡ d69d58e1-a9d2-4933-a78e-a5cecc0c0c82
minimise( x -> x^2,   RandomSolver( (-10, 10),  10000) )

# ╔═╡ ced56d45-1a1d-4389-bb9d-5067e8500628
minimise( x -> x[1]^2 + x[2]^2,
			ProductAlgorithm( [RandomSolver( (-10.0, 10),  1000),
								RandomSolver( (-10.0, 10),  1000)] ))

# ╔═╡ 7383fa13-825b-4e65-af4a-3e90533ca342
bigbreak

# ╔═╡ c09a103b-d086-4718-90d5-9cd59d972286
md"""
7.3 Optional: (Object oriented programming)

Let's imagine three datatypes
"""

# ╔═╡ af24c0dd-b37f-420a-83bd-59e2ccddf17c
abstract type AbstractScientist end

# ╔═╡ fe617e26-b9cb-4252-bb14-45c1514e6c0f
struct Scientist <: AbstractScientist
	name
	field
end

# ╔═╡ e4a7b64f-8c2a-4077-aefb-8e8b9f22f65b
struct Mathematician <: AbstractScientist
	name
	field
	specialisation
end

# ╔═╡ a2fdbb49-dc6c-4ca7-affa-1639b62b7e10
md"""
Task: Complete the code below, such that scientists say "Hi, my name is Anna, I work in the chemisty department".

But mathematicians just say, "Hi, I'm Axel. Do you like optimisation."
"""

# ╔═╡ 7fd35958-c367-4002-9cea-482e0d1f495a
begin
	# your solution
	sayhi(x) = "Hello" # define the function appropriately

end

# ╔═╡ c106b584-471d-4ef7-b14d-ea2c82b9385b
with_terminal() do
	s = Scientist("Anna", "chemistry")
	m = Mathematician("Axel", "mathematics", "optimisation")

	println( sayhi(s) )
	println( sayhi(m) )
	nothing
end

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
PlutoUI = "~0.7.32"
"""

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

julia_version = "1.7.1"
manifest_format = "2.0"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "8eaf9f1b4921132a4cff3f36a1d9ba923b14a481"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.4"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

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

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "024fe24d83e4a5bf5fc80501a314ce0d1aa35597"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.0"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

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

[[deps.Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

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

[[deps.HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

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

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

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

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

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

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

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

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

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

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

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "0b5cfbb704034b5b4c1869e36634438a047df065"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.2.1"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

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

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

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

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

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

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

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

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

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

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

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

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─a2a0a0d7-b457-40bd-ae6b-da14cc557031
# ╟─9426ed50-83ae-11ec-34f2-2f8e36fc8d05
# ╟─869614f3-f243-4319-9f8f-64c3cc13eb97
# ╠═70d2151b-0de1-41c2-a73a-7ee6f01a2d9f
# ╟─24eabc0b-496b-4159-8da6-29910b85ea4f
# ╠═996a3423-bd55-49d7-9401-d270b6e109c1
# ╠═909cc091-b321-4752-9f7a-62d1aa3a3a3d
# ╠═be82d829-d6c3-4822-8f57-3f0727776817
# ╠═3528ea22-89da-4330-9b66-49fdfef87687
# ╟─260566f0-c909-41a2-9ebf-ddb8103635f4
# ╟─8f62d5bd-764b-49e5-9b00-05d9f90332f9
# ╟─1515c0ab-f7ed-4273-9b64-c810abbaf941
# ╟─5790a006-640d-4751-a690-b139654049c6
# ╠═62ee2d5b-d757-47db-9019-37622330e9bc
# ╠═8fabcd94-1c71-420c-b2b9-d10a6eca76dc
# ╠═bbe30b74-8e29-4b3f-bf04-e05e67c82dbf
# ╟─dbd7131b-05a9-4d70-8864-d5ca9b2fe9c5
# ╠═3c7c5961-014d-47c6-8ffa-cd06198682ce
# ╟─b9a4a4ac-c07f-4f73-aa48-dbfa123a267a
# ╠═6643612d-d1a6-4fdc-8398-e12682891eb0
# ╠═cb9e96e1-597a-4b0e-912a-1972c84edd5f
# ╠═4bdf13bd-a27c-4fda-9807-eb12b1132ea0
# ╟─0fbd5d9a-30dc-4e73-9bff-c088768de5ab
# ╠═4a9409d1-d04d-459c-93a2-fd47315a1c27
# ╟─1dc516fd-b6de-466c-8cc2-ec56301cebba
# ╠═49044d10-11a2-44e7-9a86-481ddfd595c4
# ╟─35345440-550c-4723-9aec-e08142ff0802
# ╟─7c12b757-8ca5-4d2b-92b0-d204f94e6d0b
# ╠═292ee470-e286-47ea-8065-92ad2fc824d7
# ╟─55485242-77c4-43ec-996d-c1177ab34ecf
# ╟─ee6fa527-4d43-4b51-baa3-2cb2fa57d81b
# ╠═f491e000-1f10-443d-bf18-18a413f8cf94
# ╠═f8afb728-5080-41f0-a0c4-9937416493b8
# ╟─c13de59d-508e-4472-a602-17d334337ca4
# ╟─2dc06114-7f29-4021-943d-e2b14f787e3d
# ╟─95fbd51f-62a0-4167-9478-452420d19369
# ╟─91185b65-573b-42de-8238-3cd046aa5b98
# ╠═fa237ef5-f1df-4826-8062-81e061be64a2
# ╠═f01b0192-2a01-4e04-905b-3500c84c52c1
# ╠═5e55db0d-6245-47e6-ac65-6a37d66758ae
# ╠═1ab4ba54-b86c-46ec-8852-326223fce26b
# ╟─72fec38d-4c24-4356-97de-71a2914b821f
# ╟─82209778-7d13-49fd-9812-1d534ad79a49
# ╟─94b65eff-9df8-4621-9a92-74ae0a0c8052
# ╠═5fa76a28-b1d4-4c94-a7b8-5057115caa7a
# ╠═aa7189d4-53a5-4c54-bfdb-c8b242be2613
# ╟─7fc21ab5-52d9-4b12-a0c9-780cf2c62c75
# ╠═2c21a7fe-3223-46d4-af30-296045c905ec
# ╠═d7a61c32-0ed9-4d04-91e1-d6849057df2d
# ╠═fc842dae-373f-4ec5-815e-e25193fce8fe
# ╟─9c4c826b-f1a8-43f9-a50c-455826223402
# ╠═8ad53b1a-55c3-4a85-9f06-0ecd01305120
# ╟─4f962f52-99e4-4d2c-a16c-7b187b03d9d6
# ╟─e02cf43c-3389-4020-b2cc-57f201cee0aa
# ╠═0bfee97a-6b2c-42d7-b095-b356f5858bc0
# ╠═4120837d-749b-4813-b53f-0ab22bceacee
# ╟─08facfc2-4b2b-4055-891b-b8ea6b0b61d0
# ╟─27352174-d77c-4d61-bcf4-b897e451c5f6
# ╠═e355c3b9-470d-447a-8208-cf99dfbceadc
# ╠═99f36b6d-ade6-4f37-8dfb-76c967c445d9
# ╟─cebfa1b4-6876-41ac-b219-032f061f44ac
# ╟─665b754d-be3e-4982-8e03-02d1b5fbbf39
# ╟─7bf7c3c4-85cd-4d86-bbb1-933a7d5c8328
# ╟─40d2b89f-802d-441d-8053-ffd95feb22a0
# ╠═7f67b94e-495d-4605-b0f4-7978a323a798
# ╟─9b214cfc-a24e-4e4d-86e3-71451f647268
# ╟─4f55481e-7884-4dd8-92ab-14a9fedd2840
# ╟─685f46ae-e3dd-4514-bba2-393602831d3a
# ╠═b03b9de7-022a-4ba3-924a-92fdbf599453
# ╠═7f40cabc-8ce8-4ecc-b22e-1dc62bd8b648
# ╠═7e32559f-7148-472a-bff0-2777bd0df755
# ╠═f6c9269f-56a1-4e23-a7c7-93783c149e84
# ╟─75c1c92f-1e03-4ec3-86dc-65cfb2bfe35e
# ╠═d69d58e1-a9d2-4933-a78e-a5cecc0c0c82
# ╟─d2e013da-c752-4fe8-a020-c67a5e89c397
# ╟─9705e4c1-1fdf-40aa-be08-b2df34715d11
# ╟─bfb72e16-8e42-430b-921a-cec2818eaf24
# ╟─5d2e637e-3f40-49d2-9f4b-59b7c1ee6afe
# ╠═ef533b3f-b705-461f-a590-b0a11fa05f52
# ╟─38599bbf-6e45-49d8-a117-2d1822688bed
# ╟─16dc8591-a6b3-4256-bbd9-093548352449
# ╠═ac3f9a7f-ecdf-486a-a3df-a7bdd3d91012
# ╠═c577d146-2b43-4820-956e-c708861e8881
# ╟─d638bc5e-44b5-401b-a4ce-e47552c28960
# ╟─f2d11973-64b6-47b5-b34a-c4ae9405be0b
# ╠═97c78e41-21e5-46b0-b06f-c754143cae0c
# ╠═4e2e930c-8e2b-4bd3-98f1-ae35055d7b81
# ╠═ced56d45-1a1d-4389-bb9d-5067e8500628
# ╟─7383fa13-825b-4e65-af4a-3e90533ca342
# ╟─c09a103b-d086-4718-90d5-9cd59d972286
# ╠═af24c0dd-b37f-420a-83bd-59e2ccddf17c
# ╠═fe617e26-b9cb-4252-bb14-45c1514e6c0f
# ╠═e4a7b64f-8c2a-4077-aefb-8e8b9f22f65b
# ╟─a2fdbb49-dc6c-4ca7-affa-1639b62b7e10
# ╟─7fd35958-c367-4002-9cea-482e0d1f495a
# ╠═c106b584-471d-4ef7-b14d-ea2c82b9385b
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


function getNbors(M::Matrix{Int}, pos::CartesianIndex{2}) :: Vector{CartesianIndex{2}}
  s = size(M)
  nbors :: Vector{CartesianIndex{2}} = []
  (pos[1] != 1   ) && push!(nbors, pos + CartesianIndex(-1,0))
  (pos[1] != s[1]) && push!(nbors, pos + CartesianIndex(1,0))
  (pos[2] != 1   ) && push!(nbors, pos + CartesianIndex(0,-1))
  (pos[2] != s[2]) && push!(nbors, pos + CartesianIndex(0,1))
  (pos[1] != 1   ) && (pos[2] != 1   ) && push!(nbors, pos + CartesianIndex(-1,-1))
  (pos[1] != 1   ) && (pos[2] != s[2]) && push!(nbors, pos + CartesianIndex(-1, 1))
  (pos[1] != s[1]) && (pos[2] != 1   ) && push!(nbors, pos + CartesianIndex( 1,-1))
  (pos[1] != s[1]) && (pos[2] != s[2]) && push!(nbors, pos + CartesianIndex( 1, 1))
  return nbors
end

# Performs a flash cycle, returns total number of flashes.
function step!(M::Matrix{Int}) :: Int
  flashed :: Matrix{Bool} = zeros(size(M))

  # A) Everyone increases energy by 1
  M .+= 1

  # B) Octopuses flash
  charged = findall((M .> 9) .& .!flashed)
  while !isempty(charged)
    # Flash
    map(pos -> M[getNbors(M, pos)] .+= 1, charged)
    flashed[charged] .= true

    # Reset
    charged = findall((M .> 9) .& .!flashed)
  end

  # C) All octopuses that have flashed go back to 0
  M[flashed] .= 0
  return count(flashed)
end

function readInput() :: Matrix{Int}
  f = open("./inputs/day11.txt")
  data :: Matrix{Int} = zeros(10,10)

  i = 1
  for line in eachline(f)
    data[i,:] = map(x -> parse(Int,x), collect(line))
    i += 1
  end
  return data
end

problem1!(M::Matrix{Int}) = sum(map(i -> step!(M), 1:100))

function problem2!(M::Matrix{Int}) :: Int
  i :: Int = 1
  while step!(M) != 100
    i += 1
  end
  return i
end

println("Problem 1:")
M = readInput()
println(problem1!(M))

println("Problem 2:")
M = readInput()
println(problem2!(M))


@diffusion_process SIRAux{K,R} begin
    :dimensions
    process --> 2
    wiener --> 2

    :parameters
    (α, β, σ1, σ2) --> K

    :auxiliary_info
    t0 --> Float64
    T --> Float64
    vT --> R

    :additional
    linear --> true
    constdiff --> true
    statespace -->BoundedStateSpace(((1,2), (0.0,0.0)), ((1,2), (1.0,1.0)))
end

function B(t, P::SIRAux)
    @SMatrix[
        (P.α*(1.0 - P.vT[1] - P.vT[2]) - P.β)  0.0;
        P.β   0.0
    ]
end

β(t, P::SIRAux) = @SVector [0.0, 0.0]

function σ(t, P::SIRAux)
    @SMatrix [
        (-P.σ1*(1 - P.vT[1] - P.vT[2])*P.vT[1])  -P.σ2*P.vT[1];
        0.0   P.σ2*P.vT[1]
    ]
end


#Following Hamilton(1994)
#TODO: Implement for omega =/= Identity

#data: lines are different dates, columns are different variables!

function loglike(PI,data)
    T = size(data,1)
    n = size(data,2)
    if(size(PI,1) != n)
        @info "Dimensions do not match"
        return 99999999
    end
    bloc1 = -(T*n)/2*log(2*pi)
    aux1 = data[2,:] - PI'*data[1,:]
    sum_like = aux1'*aux1
    for j in 2:(T-1)
        aux1 = data[j+1,:] - PI'*data[j,:]
        sum_like = sum_like + aux1'*aux1
    end
    bloc3 = -1/2*sum_like
end


###################################
#####                         #####
#####     KOLMOGOROV FLOW     #####
#####                         #####
###################################

mutable struct KolmogorovFlow{T} <: NS2DProblem{T}

    # --> Unforced dynamics.
    unforced_dyn::UnforcedProblem{T}

    # --> Steady forcing.
    f::Array{Complex{T}, 2}

end

function KolmogorovFlow(p::SimParams, ω::Array{Complex{T}, 2}, n::Int) where T <: Real

    # --> Extract the parameters of the simulation.
    @unpack Lx, Ly, nx, ny, ν = p

    # --> Unforced part of the dynamics.
    unforced_dyn = UnforcedProblem(p, ω)

    # --> Kolmogorov flow forcing.
    x, y = mesh(p)
    x, y = collect(x), collect(y)

    f = -n .* cos.(n .* y * one.(x)')
    f = fft(f)

    # --> Return the container.
    prob = KolmogorovFlow(unforced_dyn, f)

    return prob
end

function rhs!(dΩ::Array{Complex{T},2}, Ω::Array{Complex{T},2}, p::KolmogorovFlow{T}, t::T) where T <: Real

    # --> Unforced dynamics.
    rhs!(dΩ, Ω, p.unforced_dyn, t)

    # --> Add forcing.
    dΩ .+= p.f

    return
end

function simulate(nsprob::KolmogorovFlow, T ; alg=Tsit5(), kwargs...)

    # --> Type of numbers used.
    N = findparams(nsprob)

    # --> Set the ODE Problem.
    tspan = (zero(N), convert(N, T))
    prob = ODEProblem(rhs!, copy(nsprob.unforced_dyn.ω₀), tspan, nsprob)

    # --> Integrate forward in time.
    sol = solve(prob, alg ; kwargs...)

    return sol
end


"""
    GeoMeanBridge{T, F, G, H}

The `GeometricMeanCone` is `SecondOrderCone` representable; see [1, p. 105].

The reformulation is best described in an example.

Consider the cone of dimension 4:
```math
t \\le \\sqrt[3]{x_1 x_2 x_3}
```
This can be rewritten as ``\\exists x_{21} \\ge 0`` such that:
```math
\\begin{align*}
  t & \\le x_{21},\\\\
  x_{21}^4 & \\le x_1 x_2 x_3 x_{21}.
\\end{align*}
```

Note that we need to create ``x_{21}`` and not use ``t^4`` directly as ``t`` is
allowed to be negative. Now, this is equivalent to:
```math
\\begin{align*}
  t & \\le x_{21}/\\sqrt{4},\\\\
  x_{21}^2 & \\le 2x_{11} x_{12},\\\\
  x_{11}^2 & \\le 2x_1 x_2, & x_{12}^2 & \\le 2x_3(x_{21}/\\sqrt{4}).
\\end{align*}
```

[1] Ben-Tal, Aharon, and Arkadi Nemirovski. *Lectures on modern convex
    optimization: analysis, algorithms, and engineering applications*. Society
    for Industrial and Applied Mathematics, 2001.
"""
struct GeoMeanBridge{T,F,G,H} <: AbstractBridge
    d::Int
    xij::Vector{MOI.VariableIndex}
    t_upper_bound_constraint::MOI.ConstraintIndex{F,MOI.LessThan{T}}
    rsoc_constraints::Vector{MOI.ConstraintIndex{G,MOI.RotatedSecondOrderCone}}
    x_nonnegative_constraint::Union{
        Nothing,
        MOI.ConstraintIndex{H,MOI.Nonnegatives},
    }
    # We use an internal constructor here so that users cannot pass
    # GeoMeanBridge(args..., nothing) can cause a `H not defined` error.
    function GeoMeanBridge{T,F,G,H}(
        d,
        xij,
        t_upper_bound_constraint,
        rsoc,
        x_nonnegative_constraint,
    ) where {T,F,G,H}
        return new{T,F,G,H}(
            d,
            xij,
            t_upper_bound_constraint,
            rsoc,
            x_nonnegative_constraint,
        )
    end
end

_ilog2(n, i = 0) = n <= (1 << i) ? i : _ilog2(n, i + 1)

function bridge_constraint(
    ::Type{GeoMeanBridge{T,F,G,H}},
    model,
    f::MOI.AbstractVectorFunction,
    s::MOI.GeometricMeanCone,
) where {T,F,G,H}
    f_scalars = MOI.Utilities.eachscalar(f)
    t = f_scalars[1]
    if MOI.dimension(s) == 2
        # The case {(t, x) ∈ R²: x₁ >= 0, t <= x₁}
        # The constraint t <= x₁
        t_upper_bound_constraint = MOI.Utilities.normalize_and_add_constraint(
            model,
            MOI.Utilities.operate!(-, T, t, f_scalars[2]),
            MOI.LessThan(zero(T)),
            allow_modify_function = true,
        )
        # The constraint x₁ >= 0
        x_nonnegative_constraint =
            MOI.add_constraint(model, f_scalars[2:2], MOI.Nonnegatives(1))
        return GeoMeanBridge{T,F,G,H}(
            MOI.dimension(s),
            MOI.VariableIndex[],
            t_upper_bound_constraint,
            MOI.ConstraintIndex{G,MOI.RotatedSecondOrderCone}[],
            x_nonnegative_constraint,
        )
    end
    SG = MOIU.scalar_type(G)
    d = MOI.dimension(s)
    n = d - 1
    l = _ilog2(n)
    N = 1 << l
    rsoc_constraints =
        Vector{MOI.ConstraintIndex{G,MOI.RotatedSecondOrderCone}}(undef, N - 1)
    xij = MOI.add_variables(model, N - 1)
    xl1 = xij[1]
    sN = one(T) / sqrt(N)
    _getx(i)::SG = i > n ? sN * xl1 : f_scalars[1+i]

    # With sqrt(2)^l*t - xl1, we should scale both the ConstraintPrimal and
    # ConstraintDual
    t_upper_bound_constraint = MOIU.normalize_and_add_constraint(
        model,
        MOIU.operate!(+, T, t, -sN * xl1),
        MOI.LessThan(zero(T)),
        allow_modify_function = true,
    )
    offset = offset_next = 0
    for i in 1:l
        num_lvars = 1 << (i - 1)
        offset_next = offset + num_lvars
        for j in 1:num_lvars
            if i == l
                a = _getx(2j - 1)
                b = _getx(2j)
            else
                a = convert(SG, xij[offset_next+2j-1])
                b = convert(SG, xij[offset_next+2j])
            end
            c = xij[offset+j]
            rsoc_constraints[offset+j] = MOI.add_constraint(
                model,
                MOIU.operate(vcat, T, a, b, c),
                MOI.RotatedSecondOrderCone(3),
            )
        end
        offset = offset_next
    end
    @assert offset == length(rsoc_constraints)
    return GeoMeanBridge{T,F,G,H}(
        d,
        xij,
        t_upper_bound_constraint,
        rsoc_constraints,
        nothing,
    )
end

function MOI.supports_constraint(
    ::Type{GeoMeanBridge{T}},
    ::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.GeometricMeanCone},
) where {T}
    return true
end

function MOIB.added_constrained_variable_types(::Type{<:GeoMeanBridge})
    return Tuple{Type}[]
end

function MOIB.added_constraint_types(
    ::Type{<:GeoMeanBridge{T,F,G}},
) where {T,F,G}
    return Tuple{Type,Type}[
        (F, MOI.LessThan{T}),
        (G, MOI.RotatedSecondOrderCone),
        (G, MOI.Nonnegatives),
    ]
end

function concrete_bridge_type(
    ::Type{<:GeoMeanBridge{T}},
    H::Type{<:MOI.AbstractVectorFunction},
    ::Type{MOI.GeometricMeanCone},
) where {T}
    S = MOIU.scalar_type(H)
    A = MOIU.promote_operation(*, T, T, MOI.VariableIndex)
    F = MOIU.promote_operation(+, T, S, A)
    G = MOIU.promote_operation(vcat, T, A, S, MOI.VariableIndex)
    return GeoMeanBridge{T,F,G,H}
end

# Attributes, Bridge acting as a model
MOI.get(b::GeoMeanBridge, ::MOI.NumberOfVariables)::Int64 = length(b.xij)

MOI.get(b::GeoMeanBridge, ::MOI.ListOfVariableIndices) = copy(b.xij)

function MOI.get(
    ::GeoMeanBridge{T,F},
    ::MOI.NumberOfConstraints{F,MOI.LessThan{T}},
)::Int64 where {T,F}
    return 1
end

function MOI.get(
    b::GeoMeanBridge{T,F,G},
    ::MOI.NumberOfConstraints{G,MOI.RotatedSecondOrderCone},
)::Int64 where {T,F,G}
    return length(b.rsoc_constraints)
end

function MOI.get(
    b::GeoMeanBridge{T,F,G},
    ::MOI.NumberOfConstraints{G,MOI.Nonnegatives},
)::Int64 where {T,F,G}
    return b.d > 2 ? 0 : 1
end

function MOI.get(
    b::GeoMeanBridge{T,F},
    ::MOI.ListOfConstraintIndices{F,MOI.LessThan{T}},
) where {T,F}
    return [b.t_upper_bound_constraint]
end

function MOI.get(
    b::GeoMeanBridge{T,F,G},
    ::MOI.ListOfConstraintIndices{G,MOI.RotatedSecondOrderCone},
) where {T,F,G}
    return copy(b.rsoc_constraints)
end

function MOI.get(
    b::GeoMeanBridge{T,F,G,H},
    ::MOI.ListOfConstraintIndices{H,MOI.Nonnegatives},
) where {T,F,G,H}
    if b.d > 2
        return MOI.ConstraintIndex{H,MOI.Nonnegatives}[]
    end
    return [b.x_nonnegative_constraint]
end

function MOI.delete(model::MOI.ModelLike, bridge::GeoMeanBridge)
    MOI.delete(model, bridge.xij)
    MOI.delete(model, bridge.t_upper_bound_constraint)
    MOI.delete(model, bridge.rsoc_constraints)
    if bridge.d == 2
        MOI.delete(model, bridge.x_nonnegative_constraint)
    end
    return
end

# Attributes, Bridge acting as a constraint
function MOI.get(::MOI.ModelLike, ::MOI.ConstraintSet, bridge::GeoMeanBridge)
    return MOI.GeometricMeanCone(bridge.d)
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintFunction,
    bridge::GeoMeanBridge{T,F,G,H},
) where {T,F,G,H}
    d = bridge.d
    f_scalars = Vector{MOIU.scalar_type(H)}(undef, bridge.d)
    tub = MOI.get(model, attr, bridge.t_upper_bound_constraint)
    rhs = MOI.constant(
        MOI.get(model, MOI.ConstraintSet(), bridge.t_upper_bound_constraint),
    )
    tub = MOIU.operate(-, T, tub, rhs)
    if d == 2
        x = MOI.get(model, attr, bridge.x_nonnegative_constraint)
        f_scalars[2] = MOIU.eachscalar(x)[1]
        f_scalars[1] = MOIU.operate(+, T, tub, f_scalars[2])
    else
        t = MOIU.remove_variable(tub, bridge.xij[1])
        f_scalars[1] = t
        n = d - 1
        l = _ilog2(n)
        num_lvars = 1 << (l - 1)
        offset = num_lvars - 1
        for j in 1:num_lvars
            if 2j <= bridge.d
                func = MOI.get(model, attr, bridge.rsoc_constraints[offset+j])
                func_scalars = MOIU.eachscalar(func)
                # Numerical issues can arise when a VectorOfVariables function
                # is turned into a VectorAffineFunction, because the RSOC
                # constraints might return something like `0.999999999x` instead
                # of exactly `1.0x`. To counteract this, use `convert_approx`.
                f_scalars[2j] = MOI.Utilities.convert_approx(
                    MOI.Utilities.scalar_type(H),
                    func_scalars[1],
                )
                if 2j + 1 <= bridge.d
                    f_scalars[2j+1] = MOI.Utilities.convert_approx(
                        MOI.Utilities.scalar_type(H),
                        func_scalars[2],
                    )
                end
            end
        end
    end
    return MOIU.vectorize(f_scalars)
end

function _getconstrattr(model, attr, bridge::GeoMeanBridge{T}) where {T}
    output = Vector{T}(undef, bridge.d)
    output[1] = MOI.get(model, attr, bridge.t_upper_bound_constraint)
    if bridge.d == 2
        output[2] = MOI.get(model, attr, bridge.x_nonnegative_constraint)[1]
    else
        N = length(bridge.xij) + 1
        # div(N, 2) gets layer before original problem variables are involved
        offset = div(N, 2) - 1 # 1 + 2 + ... + n/4
        for i in 1:(bridge.d-1)
            j = ((i - 1) >> 1) + 1
            k = i - 2(j - 1)
            output[1+i] =
                MOI.get(model, attr, bridge.rsoc_constraints[offset+j])[k]
        end
    end
    return output
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimal,
    bridge::GeoMeanBridge,
)
    output = _getconstrattr(model, attr, bridge)
    N = length(bridge.xij) + 1
    # the constraint is t - x_l1/sqrt(2^l) ≤ 0, we need to add the value of x_l1
    if bridge.d == 2
        output[1] += MOI.get(model, attr, bridge.x_nonnegative_constraint)[1]
    else
        output[1] +=
            MOI.get(model, MOI.VariablePrimal(), bridge.xij[1]) / sqrt(N)
    end
    return output
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintDual,
    bridge::GeoMeanBridge,
)
    output = _getconstrattr(model, attr, bridge)
    if bridge.d == 2
        # The transformation is:
        # [t]                [1 -1] [t]
        # [x] in GeoMean <=> [0  1] [x] in R- x R+
        # Hence the dual transformation is
        # [ 1 0] [u]                 [u]
        # [-1 1] [v] in GeoMean* <=> [v] in R- x R+
        output[2] -= MOI.get(model, attr, bridge.t_upper_bound_constraint)
    end
    # Otherwise, the transformation is:
    # [t]                                  [t]
    # [x] in GeoMean <=> exists xij s.t. A [x] + B xij in R- x RSOC^(N-1)
    # Hence the dual transformation is
    # A' y in GeoMean* <=>
    # B' y = 0         <=> y in R- x RSOC^(N-1)
    # Here A has at most one nonzero entry by line and column and this entry is 1.
    # Therefore we can A' = inv(A) hence we can use the same `getconstrattr` function
    # than for getting the `ConstraintPrimal`.
    # We don't need to use B' y = 0.
    return output
end


 using DSP, StatsBase

function ContinuousRP(ts1, ts2)

    # center the data
    ts1 = ts1 .- mean(ts1)
    ts2 = ts2 .- mean(ts2)

    # hilbert transform
    h1 = DSP.hilbert(ts1) |> imag
    h2 = DSP.hilbert(ts2) |> imag
    num = h2.*ts1 - ts2.*h1
    denom = ts2.*ts1 + h2.*h1

    # time series of CRP in radians
    radians = atan.(num,denom)

    # convert to degrees
    pRP = rad2deg.(radians)

    # circular stats in radians
    wgh = ones(size(radians))
    wr = wgh'*exp.(1im*radians)
    meanRP = angle(wr) |> rad2deg
    rvRP = abs(wr)/sum(wgh)
    sdRP = sqrt(2*(1-rvRP))

    (meanRP = meanRP, rvRP = rvRP, sdRP = sdRP, pRP = pRP)

    # Dict("meanRP" => meanRP,
    #     "rvRP" => rvRP,
    #     "sdRP" => sdRP,
    #     "pRP" => pRP)

end


#=
These are functions which calculate the marginal likelihood of a segment of a time series
(upto a constant of proprtionality)

P(s,t) = \int \Pr(y_{s+1:t} | \theta) \pi(\theta) d\theta
where theta is the parameter that changes and \pi(\theta) is its prior
=#

# the first function assumes a constant variance and chnging mean where the mean has a normal prior
# centered on 0 and has standard deviation s
function NormalMeanSegment(data::Array{Float64})
    cd = [0,cumsum( data )]
    cd_2 = [0,cumsum( abs2(data) )]
    cost(s::Int64, t::Int64) = cd_2[t+1] - cd_2[s+1] - abs2(cd[t+1] - cd[s+1])/(t-s)
    return cost
end


using LinearAlgebra
using CVChannel
using Test
using DelimitedFiles

"""
In this script we look at the PPT relaxation of the communication value of the
Werner-Holevo channel for multiple copies and its relation to the parameters
λ and the number of copies.
"""

@testset "PPTcv in dimension" begin
    d_list = append!(append!([3],[10:10:100;]),[100:25:750;])
    n_list = [1:4;]
    λ_list = [0,0.25,0.5]
    println("Here we are going to generate data about the PPT relaxation of cv.")
    println("Our interest is largely in how it scales in d.")
    println("We scan over for n ≤ 4, λ in ", λ_list," and d in ", d_list)
    data_table = zeros(length(n_list),length(d_list),length(λ_list))
    λ_ctr = 1
    for λ_val in λ_list
        println("Now evaluating for λ = ", λ_val)
        for n_val in n_list
            println("Now evaluating for n = ", n_val)
            d_ctr = 1
            for d_val in d_list
                A,B,g,a = generalWHLPConstraints(n_val,d_val,λ_val*ones(n_val))
                cv_ppt, opt_var = wernerHolevoCVPPT(n_val,d_val,A,B,g,a)
                data_table[n_val,d_ctr,λ_ctr] = cv_ppt
                d_ctr += 1
            end
        end
        λ_ctr +=1
    end

    header = hcat(["n\\dim:"], d_list')
    data_table_0 = vcat(header,hcat(n_list,data_table[:,:,1]))
    data_table_1 = vcat(header,hcat(n_list,data_table[:,:,2]))
    data_table_2 = vcat(header,hcat(n_list,data_table[:,:,3]))
    println("Here is a susbset of the results.")
    println("Note that each converges to a value as d increases and that n=4 is numerically unstable.")
    println("\nλ=", λ_list[1],":")
    show(stdout, "text/plain", data_table_0[:,1:12])
    println("\nλ=", λ_list[2],":")
    show(stdout, "text/plain", data_table_1[:,1:3:end])
    println("\nλ=", λ_list[3],":")
    show(stdout, "text/plain", data_table_2[:,1:3:end])

    ell = size(data_table_0)[2]
    data_to_save = vcat(data_table_0,zeros(ell)')
    data_to_save = vcat(data_to_save,vcat(data_table_1,zeros(ell)'))
    data_to_save = vcat(data_to_save,data_table_2)
    info_vec = Vector{Union{Nothing,String}}(nothing, size(data_to_save)[1])
    info_vec[1] = "INFO:"
    info_vec[2] = "Generated by multi-copy-werner-holevo-dimension-scaling.jl"
    info_vec[3] = "Tables descend as lambda increases"
    info_vec[4] = "d_list" * string(d_list)
    info_vec[5] = "n_list" * string(n_list)
    info_vec[6] = "lambda_list" * string(λ_list)
    data_to_save = hcat(data_to_save,info_vec)

    print("\nPlease name the file you'd like to write the results to: \n")
    file_name = readline()
    file_to_open = string(file_name,".csv")
    writedlm(file_to_open, data_to_save, ',')

    print("\nWe verify that the values converge for all lambda for")
    print(" n<4, because for n=4, the LP doesn't behave well (see printed data). \n")
    @test isapprox(data_table_0[2:4,38],data_table_0[2:4,39],atol=1e-2)
    @test isapprox(data_table_1[2:4,38],data_table_1[2:4,39],atol=1e-2)
    @test isapprox(data_table_2[2:4,38],data_table_2[2:4,39],atol=1e-2)
end


function nlsolve!(F::OnceDiffed, x, approach::Tuple=(Newton(), Static(1)); maxiter=200, f_abstol=1e-8, f_reltol=1e-12)
    scheme, linesearch = approach
    xp, d, Fx, Jx = copy(x), copy(x), copy(x), x*x'

    Fx, Jx = F(x, Fx, Jx)
    normFx = norm(Fx, 2)
    T = typeof(normFx)
    stoptol = T(f_reltol)*normFx + T(f_abstol)
    if normFx < stoptol
        return x, Fx, 0
    end
    iter = 1
    while iter ≤ maxiter
        # Update the point of evaluation
        d .= -(Jx\Fx)

        merit = MeritObjective(F, Fx, Jx, d)
        # φ = LineObjective!(F, ∇fz, z, x, d, fx, dot(∇fx, d))
        φ = LineObjective(merit, x, d, (normFx^2)/2, -normFx^2)
        # Perform line search along d
        α, f_α, ls_success = find_steplength(linesearch, φ, T(1.0))

        x .= x .+ α.*d
        Fx, Jx = F(x, Fx, Jx)
        normFx = norm(Fx, 2)
        if normFx < stoptol
            break
        end
        iter += 1
    end
    return x, Fx, iter
end


# functions: return_boundary_nodes,
# Matern_3d_Grid, Laplace_3D_grid,
# Parallel_Matern_3DGrid

"""
  return_boundary_nodes(xpoints, ypoints, zpoints)

...
# Arguments

  - `xpoints::Vector{T} where T<:Real`: the vector containing the x coordinate
  - `ypoints::Vector{T} where T<:Real`: the vector containing the y coordinate
  - `zpoints::Vector{T} where T<:Real`: the vector containing the z coordinate
...

...
# Outputs
  - `BoundaryNodes3D::Vector{Int64}`: vector containing the indices of coordinates
  on the boundary of the rectangular 3D volume
...

"""
function return_boundary_nodes(xpoints, ypoints, zpoints)
    BoundaryNodes3D =[]
    xneighbors = []
    yneighbors = []
    zneighbors = []
    counter = 0
    for k = 1:zpoints
        for j = 1:ypoints
            for i = 1:xpoints
                counter=counter+1
                if(k == 1 || k == zpoints || j == 1|| j == ypoints || i == 1 || i == xpoints)
                    BoundaryNodes3D = push!(BoundaryNodes3D, counter)
                    if(k == 1 || k == zpoints)
                        push!(zneighbors, 1)
                    else
                        push!(zneighbors, 2)
                    end
                    if(j == 1 || j == ypoints)
                        push!(yneighbors, 1)
                    else
                        push!(yneighbors, 2)
                    end
                    if(i == 1 || i == xpoints)
                        push!(xneighbors, 1)
                    else
                        push!(xneighbors, 2)
                    end
                end
            end
        end
    end
    return BoundaryNodes3D, xneighbors, yneighbors, zneighbors
end

"""

  Matern3D_Grid(xpoints, ypoints, zpoints, imgg, epsilon, radius, h, k, l, m)

...
# Arguments
  - `xpoints::Vector{T} where T<:Real`: the vector containing the x coordinate
  - `ypoints::Vector{T} where T<:Real`: the vector containing the y coordinate
  - `zpoints::Vector{T} where T<:Real`: the vector containing the z coordinate
  - `imgg`: the matrix containing the image
  - `epsilon`: Matern parameter epsilon
  - `radius::Vector`: the tuple containing the punch radii
  - `h::Float`: grid spacing along the x-axis
  - `k::Float`: grid spacing along the y-axis
  - `l::Float`: grid spacing along the z-axis
  - `m::Int` : Matern parameter

# Outputs
  - tuple containing the restored image and the punched image.
...

"""
function Matern3D_Grid(xpoints, ypoints, zpoints, imgg, epsilon, radius, h, k, l, m)
    A3D = nablasq_3d_grid(length(xpoints), length(ypoints), length(zpoints), h, k, l)
    sizeA = size(A3D, 1)
    for i = 1:sizeA
        A3D[i, i] = A3D[i, i] + epsilon^2
    end
    A3DMatern = A3D
    for i = 1:m - 1
        A3DMatern = A3DMatern * A3D
    end
    discard = punch_holes_nexus(xpoints, ypoints, zpoints, radius)
    punched_image = copy(imgg)
    punched_image[discard] .= 1
    totalsize = prod(size(imgg))
    C = sparse(I, totalsize, totalsize)
    rhs_a = punched_image[:]
    for i in discard
        C[i, i] = 0
        rhs_a[i] = 0
    end
    Id = sparse(I, totalsize, totalsize)
    u = ((C - (Id - C) * A3DMatern)) \ rhs_a
    return (u, punched_image[:])
end

"""

  Laplace3D_Grid(xpoints, ypoints, zpoints, imgg, radius, h, k, l)

...
# Arguments
  - `xpoints::Vector{T} where T<:Real`: the vector containing the x coordinate
  - `ypoints::Vector{T} where T<:Real`: the vector containing the y coordinate
  - `zpoints::Vector{T} where T<:Real`: the vector containing the z coordinate
  - `imgg`: the matrix containing the image
  - `radius::Vector`: the tuple containing the punch radii
  - `h::Float`: grid spacing along the x-axis
  - `k::Float`: grid spacing along the y-axis
  - `l::Float`: grid spacing along the z-axis

# Outputs
  - tuple containing the restored image and the punched image.
...

"""
function Laplace3D_Grid(xpoints, ypoints, zpoints, imgg, radius, h, k, l)
    A3D = nablasq_3d_grid(length(xpoints), length(ypoints), length(zpoints), h, k, l)
    discard = punch_holes_nexus(xpoints, ypoints, zpoints, radius)
    punched_image = copy(imgg)
    punched_image[discard] .= 1
    totalsize = prod(size(imgg))
    C = sparse(I, totalsize, totalsize)
    rhs_a = punched_image[:]
    for i in discard
        C[i,i] = 0
        rhs_a[i] = 0
    end
    Id = sparse(I, totalsize, totalsize)
    u = ((C - (Id - C) * A3D)) \ rhs_a
    return u, punched_image[:]
end

"""

  Parallel_Laplace3D_Grid(xpoints, ypoints, zpoints, imgg, radius, h, k, l,
          xmin, xmax, ymin, ymax, zmin, zmax)

Compute the spherically-punched, Laplace-interpolated 3D data

...
# Arguments
  - `xpoints::Vector{T} where T<:Real`: the vector containing the x coordinate
  - `ypoints::Vector{T} where T<:Real`: the vector containing the y coordinate
  - `zpoints::Vector{T} where T<:Real`: the vector containing the z coordinate
  - `imgg`: the matrix containing the image
  - `radius::Vector`: the tuple containing the punch radii
  - `h::Float`: grid spacing along the x-axis
  - `k::Float`: grid spacing along the y-axis
  - `l::Float`: grid spacing along the z-axis
  - `xmin::Int64`: Vishwas should fill in the next six fields.
  - `xmax::Int64`:
  - `ymin::Int64`:
  - `ymax::Int64`:
  - `zmin::Int64`:
  - `zmax::Int64`:

# Outputs
  - array containing the restored image

# Example

```julia-repl
(xmin, xmax) = (ymin, ymax) = (zmin,zmax) = (0.0, 1.0)
xpoints = ypoints = zpoints = -0.2:0.2:1.2
h = k = l = 0.2
imgg = randn(8,8,8)
radius = 0.2
restored = Parallel_Laplace3D_Grid(xpoints, ypoints, zpoints, imgg, radius,
                                 h, k, l, xmin, xmax, ymin, ymax, zmin, zmax)
```

...

"""
function Parallel_Laplace3D_Grid(xpoints::Union{StepRangeLen{T},Vector{T}},
                                ypoints::Union{StepRangeLen{T},Vector{T}},
                                zpoints::Union{StepRangeLen{T},Vector{T}},
                                imgg::Array{P,3},
                                radius::Union{Q,Tuple{Q,Q,Q}},
                                h::Float64, k::Float64, l::Float64,
                                xmin::R, xmax::R, ymin::R, ymax::R, zmin::R, zmax::R
                                ) where{T<:Number,P<:Number,Q<:Number,R<:Number}
  #
  fun(x,y,z,w) = Int(round((x -y)/z) - w )
  ran(x,y,z,w) = fun(x, y, z, w) .+ (0, 2*w + 1)
  #
  rad = (typeof(radius) <: Tuple) ? radius : (radius, radius, radius)
  radius_x, radius_y, radius_z = rad
  stride_h = Int64(round(radius_x / h))
  stride_k = Int64(round(radius_y / k))
  stride_l = Int64(round(radius_z / l))
  cartesian_product_boxes = [(ran(i, xpoints[1], h, stride_h)...,
                              ran(j, ypoints[1], k, stride_k)...,
                              ran(kk, zpoints[1], l, stride_l)...)
                         for i in xmin:xmax for j in ymin:ymax for kk in zmin:zmax]
  #
  z3d_restored = copy(imgg)
  #
  Threads.@threads for i = 1:length(cartesian_product_boxes)
    i1, i2, j1, j2, k1, k2 = cartesian_product_boxes[i]
    restored_img = Laplace3D_Grid(xpoints[i1 + 1:i2], ypoints[j1 + 1:j2],
                           zpoints[k1 + 1:k2],
                           imgg[i1 + 1:i2, j1 + 1:j2, k1 + 1:k2],
                           radius, h, k, l)[1]
    z3d_restored[i1 + 1:i2, j1 + 1:j2, k1 + 1:k2] = reshape(restored_img,
                           (2 * stride_h + 1, 2 * stride_k + 1, 2 * stride_l + 1))
  end
  #
  return z3d_restored[:]
  #
end

"""

  Parallel_Matern3D_Grid(xpoints, ypoints, zpoints, imgg, epsilon, radius, h, k, l,
          xmin, xmax, ymin, ymax, zmin, zmax, m)

Compute the spherically-punched, Matern-interpolated 3D data

...
# Arguments
  - `xpoints::Vector{T} where T<:Real`: the vector containing the x coordinate
  - `ypoints::Vector{T} where T<:Real`: the vector containing the y coordinate
  - `zpoints::Vector{T} where T<:Real`: the vector containing the z coordinate
  - `imgg`: the matrix containing the image
  - `epsilon`: one of the matern parameters
  - `radius::Vector`: the tuple containing the punch radii
  - `h::Float`: grid spacing along the x-axis
  - `k::Float`: grid spacing along the y-axis
  - `l::Float`: grid spacing along the z-axis
  - `xmin::Int64`: Vishwas should fill in the next six fields.
  - `xmax::Int64`:
  - `ymin::Int64`:
  - `ymax::Int64`:
  - `zmin::Int64`:
  - `zmax::Int64`:
  - `m::Int64`: The matern order parameter

# Outputs
  - array containing the restored image

# Example

```julia-repl
(xmin, xmax) = (ymin, ymax) = (zmin,zmax) = (0.0, 1.0)
xpoints = ypoints = zpoints = -0.2:0.2:1.2
h = k = l = 0.2
imgg = randn(8,8,8)
m = 2
epsilon = 0.0
radius = 0.2
restored = Parallel_Matern3D_Grid(xpoints, ypoints, zpoints, imgg, epsilon, radius,
                                  h, k, l, xmin, xmax, ymin, ymax, zmin, zmax, m)
```

...

"""
function Parallel_Matern3D_Grid(xpoints::Union{StepRangeLen{T},Vector{T}},
                                ypoints::Union{StepRangeLen{T},Vector{T}},
                                zpoints::Union{StepRangeLen{T},Vector{T}},
                                imgg::Array{P,3}, epsilon::Q,
                                radius::Union{Q,Tuple{Q,Q,Q}},
                                h::Float64, k::Float64, l::Float64,
                                xmin::R, xmax::R, ymin::R, ymax::R, zmin::R, zmax::R,
                                m::Int) where{T<:Number,P<:Number,Q<:Number,R<:Number}
  #
  fun(x,y,z,w) = Int(round((x -y)/z) - w )
  ran(x,y,z,w) = fun(x, y, z, w) .+ (0, 2*w + 1)
  #
  rad = (typeof(radius) <: Tuple) ? radius : (radius, radius, radius)
  radius_x, radius_y, radius_z = rad
  stride_h = Int64(round(radius_x / h))
  stride_k = Int64(round(radius_y / k))
  stride_l = Int64(round(radius_z / l))
  cartesian_product_boxes = [(ran(i, xpoints[1], h, stride_h)...,
                              ran(j, ypoints[1], k, stride_k)...,
                              ran(kk, zpoints[1], l, stride_l)...)
                         for i in xmin:xmax for j in ymin:ymax for kk in zmin:zmax]
  #
  z3d_restored = copy(imgg)
  #
  Threads.@threads for i = 1:length(cartesian_product_boxes)
    i1, i2, j1, j2, k1, k2 = cartesian_product_boxes[i]
    restored_img = Matern3D_Grid(xpoints[i1 + 1:i2], ypoints[j1 + 1:j2],
                           zpoints[k1 + 1:k2],
                           imgg[i1 + 1:i2, j1 + 1:j2, k1 + 1:k2],
                           epsilon, radius, h, k, l, m)[1]
    z3d_restored[i1 + 1:i2, j1 + 1:j2, k1 + 1:k2] = reshape(restored_img,
                           (2 * stride_h + 1, 2 * stride_k + 1, 2 * stride_l + 1))
  end
  #
  return z3d_restored[:]
  #
end

# function Parallel_Matern3D_Grid(xpoints, ypoints, zpoints, imgg, epsilon, radius, h, k, l, xmin, xmax, ymin, ymax, zmin, zmax, m)
#   xbegin = xpoints[1];
#   ybegin = ypoints[1];
#   zbegin = zpoints[1];
#   cartesian_product_boxes = [];
#   stride = Int(round(radius/h));
#   z3d_restored = copy(imgg);
#   for i = xmin:xmax+1
#       i1 = Int(round((i-xbegin) /h))-stride;
#       i2 = i1+2*stride+1;
#       for j = ymin:ymax+1
#           j1 = Int(round((j-ybegin)/h))-stride;
#           j2 = j1+2*stride+1;
#           for k = zmin:zmax+1
#               k1 = Int(round((k-ybegin)/h)) - stride;
#               k2 = k1+2*stride+1;
#               append!(cartesian_product_boxes,[(i1,i2,j1,j2,k1,k2)]);
#           end
#       end
#   end

#   Threads.@threads for i = 1:length(cartesian_product_boxes)
#     i1 = cartesian_product_boxes[i][1];
#     i2 = cartesian_product_boxes[i][2];
#     j1 = cartesian_product_boxes[i][3];
#     j2 = cartesian_product_boxes[i][4];
#     k1 = cartesian_product_boxes[i][5];
#     k2 = cartesian_product_boxes[i][6];
#     z3temp = imgg[i1+1:i2,j1+1:j2,k1+1:k2];
#     restored_img, punched_image = Matern3D_Grid(xpoints[i1+1:i2], ypoints[j1+1:j2], zpoints[k1+1:k2], z3temp, epsilon, radius, h, k, l, m);
#     restored_img_reshape = reshape(restored_img, (2*stride+1,2*stride+1,2*stride+1));
#     z3d_restored[i1+1:i2, j1+1:j2, k1+1:k2] = restored_img_reshape;
#   end
#   return z3d_restored[:];
# end


"""
Struct for Kou Process

		kouProcess=KouProcess(σ::num1,λ::num2,p::num3,λ₊::num4,λ₋::num5) where {num1,num2,num3,num4,num5 <: Number}

Where:\n
		σ  =	volatility of the process.
		λ  = 	jumps intensity.
		p  =	prob. of having positive jump.
		λ₊ =	positive jump size.
		λ₋ =	negative jump size.
"""
mutable struct KouProcess{num <: Number, num1 <: Number, num2 <: Number, num3 <: Number, num4 <: Number, abstrUnderlying <: AbstractUnderlying, numtype <: Number} <: FiniteActivityProcess{numtype}
    σ::num
    λ::num1
    p::num2
    λ₊::num3
    λ₋::num4
    underlying::abstrUnderlying
    function KouProcess(σ::num, λ::num1, p::num2, λ₊::num3, λ₋::num4, underlying::abstrUnderlying) where {num <: Number, num1 <: Number, num2 <: Number, num3 <: Number, num4 <: Number, abstrUnderlying <: AbstractUnderlying}
        if σ <= 0
            error("volatility must be positive")
        elseif λ <= 0
            error("jump intensity must be positive")
        elseif λ₊ <= 0
            error("positive λ must be positive")
        elseif λ₋ <= 0
            error("negative λ must be positive")
        elseif !(0 <= p <= 1)
            error("p must be a probability")
        else
            zero_typed = promote_type(num, num1, num2, num3, num4)
            return new{num, num1, num2, num3, num4, abstrUnderlying, zero_typed}(σ, λ, p, λ₊, λ₋, underlying)
        end
    end
end

export KouProcess;

function compute_jump_size(mcProcess::KouProcess, mcBaseData::MonteCarloConfiguration)
    return if rand(mcBaseData.rng) < mcProcess.p
        quantile_exp(mcProcess.λ₊, rand(mcBaseData.rng))
    else
        -quantile_exp(mcProcess.λ₋, rand(mcBaseData.rng))
    end
end
compute_drift(mcProcess::KouProcess) = mcProcess.p / (mcProcess.λ₊ - 1) - (1 - mcProcess.p) / (mcProcess.λ₋ + 1)


"""
    tracenorm(rho)

Trace norm of `rho`.

It is defined as

```math
T(ρ) = Tr\\{\\sqrt{ρ^† ρ}\\}.
```

Depending if `rho` is hermitian either [`tracenorm_h`](@ref) or
[`tracenorm_nh`](@ref) is called.
"""
function tracenorm(rho::DenseOperator{B,B}) where B<:Basis
    ishermitian(rho) ? tracenorm_h(rho) : tracenorm_nh(rho)
end
function tracenorm(rho::T) where T<:AbstractOperator
    throw(ArgumentError("tracenorm not implemented for $(typeof(rho)). Use dense operators instead."))
end

"""
    tracenorm_h(rho)

Trace norm of `rho`.

It uses the identity

```math
T(ρ) = Tr\\{\\sqrt{ρ^† ρ}\\} = \\sum_i |λ_i|
```

where ``λ_i`` are the eigenvalues of `rho`.
"""
function tracenorm_h(rho::DenseOperator{B,B}) where B<:Basis
    s = eigvals(Hermitian(rho.data))
    sum(abs.(s))
end
function tracenorm_h(rho::T) where T<:AbstractOperator
    throw(ArgumentError("tracenorm_h not implemented for $(typeof(rho)). Use dense operators instead."))
end

"""
    tracenorm_nh(rho)

Trace norm of `rho`.

Note that in this case `rho` doesn't have to be represented by a square
matrix (i.e. it can have different left-hand and right-hand bases).

It uses the identity

```math
    T(ρ) = Tr\\{\\sqrt{ρ^† ρ}\\} = \\sum_i σ_i
```

where ``σ_i`` are the singular values of `rho`.
"""
tracenorm_nh(rho::DenseOperator) = sum(svdvals(rho.data))
function tracenorm_nh(rho::T) where T<:AbstractOperator
    throw(ArgumentError("tracenorm_nh not implemented for $(typeof(rho)). Use dense operators instead."))
end

"""
    tracedistance(rho, sigma)

Trace distance between `rho` and `sigma`.

It is defined as

```math
T(ρ,σ) = \\frac{1}{2} Tr\\{\\sqrt{(ρ - σ)^† (ρ - σ)}\\}.
```

It calls [`tracenorm`](@ref) which in turn either uses [`tracenorm_h`](@ref)
or [`tracenorm_nh`](@ref) depending if ``ρ-σ`` is hermitian or not.
"""
tracedistance(rho::T, sigma::T) where T<:DenseOperator = 0.5*tracenorm(rho - sigma)
function tracedistance(rho::AbstractOperator, sigma::AbstractOperator)
    throw(ArgumentError("tracedistance not implemented for $(typeof(rho)) and $(typeof(sigma)). Use dense operators instead."))
end

"""
    tracedistance_h(rho, sigma)

Trace distance between `rho` and `sigma`.

It uses the identity

```math
T(ρ,σ) = \\frac{1}{2} Tr\\{\\sqrt{(ρ - σ)^† (ρ - σ)}\\} = \\frac{1}{2} \\sum_i |λ_i|
```

where ``λ_i`` are the eigenvalues of `rho` - `sigma`.
"""
tracedistance_h(rho::T, sigma::T) where {B<:Basis,T<:DenseOperator{B,B}}= 0.5*tracenorm_h(rho - sigma)
function tracedistance_h(rho::AbstractOperator, sigma::AbstractOperator)
    throw(ArgumentError("tracedistance_h not implemented for $(typeof(rho)) and $(typeof(sigma)). Use dense operators instead."))
end

"""
    tracedistance_nh(rho, sigma)

Trace distance between `rho` and `sigma`.

Note that in this case `rho` and `sigma` don't have to be represented by square
matrices (i.e. they can have different left-hand and right-hand bases).

It uses the identity

```math
    T(ρ,σ) = \\frac{1}{2} Tr\\{\\sqrt{(ρ - σ)^† (ρ - σ)}\\}
         = \\frac{1}{2} \\sum_i σ_i
```

where ``σ_i`` are the singular values of `rho` - `sigma`.
"""
tracedistance_nh(rho::T, sigma::T) where {T<:DenseOperator} = 0.5*tracenorm_nh(rho - sigma)
function tracedistance_nh(rho::AbstractOperator, sigma::AbstractOperator)
    throw(ArgumentError("tracedistance_nh not implemented for $(typeof(rho)) and $(typeof(sigma)). Use dense operators instead."))
end

"""
    entropy_vn(rho)

Von Neumann entropy of a density matrix.

The Von Neumann entropy of a density operator is defined as

```math
S(ρ) = -Tr(ρ \\log(ρ)) = -\\sum_n λ_n\\log(λ_n)
```

where ``λ_n`` are the eigenvalues of the density matrix ``ρ``, ``\\log`` is the
natural logarithm and ``0\\log(0) ≡ 0``.

# Arguments
* `rho`: Density operator of which to calculate Von Neumann entropy.
* `tol=1e-15`: Tolerance for rounding errors in the computed eigenvalues.
"""
function entropy_vn(rho::DenseOperator{B,B}; tol::Float64=1e-15) where B<:Basis
    evals::Vector{ComplexF64} = eigvals(rho.data)
    evals[abs.(evals) .< tol] .= 0.0im
    sum([d == 0.0im ? 0.0im : -d*log(d) for d=evals])
end
entropy_vn(psi::StateVector; kwargs...) = entropy_vn(dm(psi); kwargs...)

"""
    fidelity(rho, sigma)

Fidelity of two density operators.

The fidelity of two density operators ``ρ`` and ``σ`` is defined by

```math
F(ρ, σ) = Tr\\left(\\sqrt{\\sqrt{ρ}σ\\sqrt{ρ}}\\right),
```

where ``\\sqrt{ρ}=\\sum_n\\sqrt{λ_n}|ψ⟩⟨ψ|``.
"""
fidelity(rho::T, sigma::T) where {B<:Basis,T<:DenseOperator{B,B}} = tr(sqrt(sqrt(rho.data)*sigma.data*sqrt(rho.data)))

"""
    ptranspose(rho, index)

Partial transpose of rho with respect to subspace specified by index.
"""
function ptranspose(rho::DenseOperator{B,B}, index::Int=1) where B<:CompositeBasis

    # Define permutation
    N = length(rho.basis_l.bases)
    perm = [1:N;]
    perm[index] = N
    perm[N] = index

    # Permute indexed subsystem to last position
    rho_perm = permutesystems(rho, perm)

    # Transpose corresponding blocks
    m = Int(prod(rho_perm.basis_l.shape[1:N-1]))
    n = rho_perm.basis_l.shape[N]
    for i=1:n, j=1:n
        rho_perm.data[m*(i-1)+1:m*i, m*(j-1)+1:m*j] = permutedims(rho_perm.data[m*(i-1)+1:m*i, m*(j-1)+1:m*j])
    end

    return permutesystems(rho_perm, perm)
end

"""
    PPT(rho, index)

Peres-Horodecki criterion of partial transpose.
"""
PPT(rho::DenseOperator{B,B}, index::Int) where B<:CompositeBasis = all(real.(eigvals(ptranspose(rho, index).data)) .>= 0.0)

"""
    negativity(rho, index)

Negativity of rho with respect to subsystem index.

The negativity of a density matrix ρ is defined as

```math
N(ρ) = \\|ρᵀ\\|,
```
where `ρᵀ` is the partial transpose.
"""
negativity(rho::DenseOperator{B,B}, index::Int) where B<:CompositeBasis = 0.5*(tracenorm(ptranspose(rho, index)) - 1.0)

"""
    logarithmic_negativity(rho, index)

The logarithmic negativity of a density matrix ρ is defined as

```math
N(ρ) = \\log₂\\|ρᵀ\\|,
```
where `ρᵀ` is the partial transpose.
"""
logarithmic_negativity(rho::DenseOperator{B,B}, index::Int) where B<:CompositeBasis = log(2, tracenorm(ptranspose(rho, index)))

"""
    avg_gate_fidelity(x, y)

The average gate fidelity between two superoperators x and y.
"""
function avg_gate_fidelity(x::T, y::T) where T <: Union{PauliTransferMatrix{B, B} where B, SuperOperator{B, B} where B, ChiMatrix{B, B} where B}
    dim = 2 ^ length(x.basis_l)
    return (tr(transpose(x.data) * y.data) + dim) / (dim^2 + dim)
end

"""
    entanglement_entropy(state, partition, [entropy_fun=entropy_vn])

Computes the entanglement entropy of `state` between the list of sites `partition`
and the rest of the system. The state must be defined in a composite basis.

If `state isa AbstractOperator` the operator-space entanglement entropy is
computed, which has the property
```julia
entanglement_entropy(dm(ket)) = 2 * entanglement_entropy(ket)
```

By default the computed entropy is the Von-Neumann entropy, but a different
function can be provided (for example to compute the entanglement-renyi entropy).
"""
function entanglement_entropy(psi::Ket{B}, partition::Vector{Int}, entropy_fun=entropy_vn) where B<:CompositeBasis
    # check that sites are within the range
    @assert all(partition .<= length(psi.basis.bases))

    rho = ptrace(psi, partition)
    return entropy_fun(rho)
end

function entanglement_entropy(rho::DenseOperator{B,B}, partition::Vector{Int}, args...) where {B<:CompositeBasis}
    # check that sites is within the range
    hilb = rho.basis_l
    all(partition .<= length(hilb.bases)) || throw(ArgumentError("Indices in partition must be within the bounds of the composite basis."))
    length(partition) <= length(hilb.bases) || throw(ArgumentError("Partition cannot include the whole system."))

    # build the doubled hilbert space for the vectorised dm, normalized like a Ket.
    b_doubled = hilb^2
    rho_vec = normalize!(Ket(b_doubled, vec(rho.data)))

    return entanglement_entropy(rho_vec,
                                vcat(partition, partition.+length(hilb.bases)),
                                args...)
end

entanglement_entropy(state, partition::Number, args...) =
    entanglement_entropy(state, [partition], args...)
entanglement_entropy(state, partition, args...) =
    entanglement_entropy(state, collect(partition), args...)


using DynamicalSystems
using DifferentialEquations
using PyPlot
using Interpolations
using LinearAlgebra
using Statistics
using DelimitedFiles

function pyragas_rhs(state, history, params, time)
    chaotic_system, feedback_amplitude_vector, feedback_delay = params
    delayed_state = history(params, time - feedback_delay)
    F = feedback_amplitude_vector .* (delayed_state .- state)
    return chaotic_system.f(state, chaotic_system.p, time) .+ F
end

function pyragas_control(chaotic_system, feedback_amplitude_vector, feedback_delay, tmax;
        dt=0.1, system_transit=0.0, dense=false, ode_algorithm=AutoTsit5(Rosenbrock23()),
        solve_kwargs=(abstol=1e-6, reltol=1e-3, maxiters=10^10))
    ode_problem = ODEProblem(chaotic_system, (- feedback_delay - system_transit, 0.0))
    initial_solution = solve(ode_problem, ode_algorithm; dense=true, solve_kwargs...)
    initial_history(p, t) = initial_solution(t)
    initial_state = initial_solution.u[end]
    params = (chaotic_system, feedback_amplitude_vector, feedback_delay)
    problem = DDEProblem(pyragas_rhs, initial_state, initial_history, (0.0, tmax), params; constant_lags=[feedback_delay])
    dde_algorithm = MethodOfSteps(ode_algorithm)
    return solve(problem, dde_algorithm; solve_kwargs, dense=dense, saveat=dense ? [] : dt)
end

const sys = Systems.roessler()
const system_transit = 1000.0
const v = SVector(0.0, 1.0, 0.0)
const default_K = 0.2
const default_tau = 17.5

function plot_control_example()
    K = default_K
    tau = default_tau
    tmax = 500.0
    t_chaotic = trajectory(sys, tmax; Ttr=system_transit)
    feedback_amplitude_vector = K.*v
    t_control = pyragas_control(sys, feedback_amplitude_vector, tau, tmax; system_transit=system_transit)
    dt = t_control.t[2] - t_control.t[1]
    start = ceil(Int, tau/dt)
    transit_end = 1 + ceil(Int, 480.0/dt)
    F = [dot(feedback_amplitude_vector .* (t_control(t - tau) .- t_control(t)), v) for t in t_control.t[start:end]]

    figure()
    subplot(2,1,1)
    plot(t_chaotic[:, 1], t_chaotic[:, 2], lw=0.5, label="chaos")
    plot(t_control[1, transit_end:end], t_control[2, transit_end:end], lw=2.0, label="control")
    legend(framealpha=1.0)
    xlabel("x")
    ylabel("y")
    subplot(2,1,2)
    plot(t_control.t[start:end], F)
    xlabel("t")
    ylabel("F")
    suptitle("Pyragas Control of the Rossler Oscillator")
    tight_layout()
    savefig("pyragas_control_of_roessler_attractor.pdf")
end

const disp_disp_tau_filename = "pyragas_control_displacement_dispersions_tau.dat"
const disp_disp_K_filename = "pyragas_control_displacement_dispersions_K.dat"

function compute_dispersions_tau()
    tmax = 5000.0
    nsamples = 1000
    control_transit = 1000.0
    taus = range(4.0, 20.0, length=500)
    @time displacement_dispersions = map(taus) do tau
        sol = pyragas_control(sys, default_K.*v, tau, tmax; system_transit=system_transit, dense=true)
        mean((sol(t - tau)[2] - sol(t)[2])^2 for t = tau + control_transit:tmax/nsamples:tmax)
    end
    writedlm(disp_disp_tau_filename, [taus displacement_dispersions])
end

const disp_tau = 5.9

function compute_dispersions_K()
    tmax = 5000.0
    nsamples = 1000
    control_transit = 1000.0
    Ks = range(0.0, 1.1, length=500)
    tau = disp_tau
    @time displacement_dispersions = map(Ks) do K
        sol = pyragas_control(sys, K.*v, tau, tmax; system_transit=system_transit, dense=true)
        mean((sol(t - tau)[2] - sol(t)[2])^2 for t = tau + control_transit:tmax/nsamples:tmax)
    end
    writedlm(disp_disp_K_filename, [Ks displacement_dispersions])
end

function plot_dispersions()
    figure()

    data = readdlm(disp_disp_tau_filename)
    taus, displacement_dispersions = data[:, 1], data[:, 2]
    subplot(2,1,1)
    plot(taus, displacement_dispersions, ".")
    xlabel(raw"$\tau$")
    ylabel(raw"$\langle D(t)^2 \rangle$")
    yscale("log")
    ylim(1e-11, 1e2)
    title("K = $default_K")

    data = readdlm(disp_disp_K_filename)
    Ks, displacement_dispersions = data[:, 1], data[:, 2]
    subplot(2,1,2)
    plot(Ks, displacement_dispersions, ".")
    xlabel("K")
    ylabel(raw"$\langle D(t)^2 \rangle$")
    yscale("log")
    # ylim(1e-11, 1e2)
    title("\$\\tau\$ = $disp_tau")

    tight_layout()
    savefig("pyragas_control_displacement_dispersion_plot.pdf")
end

const max_lyap_filename = "pyragas_max_lyap.dat"

function pyragas_map!(next_history_vector, _initial_history_vector, params, _t)
    chaotic_system, feedback_amplitude_vector, feedback_delay, N = params

    initial_history_vector = [SVector{N}(_initial_history_vector[i:i+N-1]) for i = 1:N:length(_initial_history_vector) - N]

    dt = feedback_delay / (length(initial_history_vector) - 1)
    initial_state = initial_history_vector[end]
    interp = LinearInterpolation(range(-feedback_delay, 0.0, length=length(initial_history_vector)), initial_history_vector)
    initial_history(p, t) = interp(t)
    problem = DDEProblem(pyragas_rhs, initial_state, initial_history, (0.0, feedback_delay + dt), params; constant_lags=[feedback_delay])

    dde_algorithm = MethodOfSteps(AutoTsit5(Rosenbrock23()))
    sol = solve(problem, dde_algorithm; saveat=dt)

    for (i, u) in enumerate(@view(sol.u[2:end]))
        copy!(@view(next_history_vector[(i - 1)*N + 1 : i*N]), u)
    end
    return nothing
end

function max_lyapunov(K, tau; N=80)
    dt = tau / (N - 1)
    sol = pyragas_control(sys, K.*v, tau, tau; system_transit=1000.0, dt=dt)
    ds = DiscreteDynamicalSystem(pyragas_map!, vcat(sol.u...), [sys, K.*v, tau, length(sys.u0)])
    return lyapunov(ds, 100; Ttr=round(Int, 500.0/(tau + dt)), show_progress=true)
end

function compute_max_lyapunov_exponent()
    Ks = range(0.0, 1.1, length=10) # 10 .^ range(-2, 2, length=10)
    @time lambda_maxs = map(K -> max_lyapunov(K, disp_tau), Ks)
    writedlm(max_lyap_filename, [Ks lambda_maxs])
end

function plot_max_lyapunov_exponent()
    data = readdlm(max_lyap_filename)
    Ks, lambda_maxs = data[:, 1], data[:, 2]
    figure()
    plot(Ks, lambda_maxs, ".")
    axhline(0.0, color="black")
    xlabel("K")
    ylabel(raw"$\lambda_\mathrm{max}$")
    #xscale("log")
    savefig("pyragas_max_lyapunov_exponent.pdf")
end

function main()
    plot_control_example()
    compute_dispersions_tau()
    compute_dispersions_K()
    plot_dispersions()
    compute_max_lyapunov_exponent()
    plot_max_lyapunov_exponent()
end

function old_max_lyapunov(K, tau; debug=false)
    tmax = 5000.0
    ks = 1:40
    steps = 1000.0:1.0:tmax
    sol = pyragas_control(sys, K.*v, tau, tmax; system_transit=1000.0, dense=true)
    data = Dataset(sol.(steps))
    Es = lyapunov_from_data(data, ks)
    starts, slops = linear_regions(ks, Es; tol=1.0)
    if debug
        for i = 1:length(starts) - 1
            interval = ks[i] : ks[i + 1]
            plot(interval, Es[interval])
        end
    end
    return slops[1]
end


@def Ax_min begin #(U)

	# minimum longitudinal acceleration for given speed
	Ax_min = Array(Float64,(length(U),1))
	for i in eachindex(U)
		Ax_min[i,1] = AXC[5]*U[i]^3 + AXC[6]*U[i]^2 + AXC[7]*U[i] + AXC[8]
	end
	Ax_min
end


#=
Simple 2D ellipse with convex combination initial state.
=#
using Reachability, Plots

function compute(input_options::Pair{Symbol,<:Any}...)
    # =====================
    # Problem specification
    # =====================
    A = sparse([1, 1, 2, 2], [1, 2, 1, 2], [3., -9., 4., -3.])

    # initial set
    X0_1 = Singleton([1.0, 0.0])
    X0_2 = Singleton([1.5, 0.0])
    alpha = 0.8
    Z0 = Singleton(alpha * X0_1.element + (1.-alpha) * X0_2.element)

    # instantiate continuous LTI system
    S1 = ContinuousSystem(A, X0_1)
    S2 = ContinuousSystem(A, X0_2)
    S3 = ContinuousSystem(A, Z0)

    # ===============
    # Problem solving
    # ===============
    # define solver-specific options
    options = Options(
        :δ => 0.01,
        :approx_model => "nobloating",
        :plot_vars => [1, 2]
        )

    t1 = 1.2
    t2 = 1.0
    t3 = 1.09

    o1 = merge(options, Options(:T => t1))
    o2 = merge(options, Options(:T => t2))
    o3 = merge(options, Options(:T => t3))

    # solve
    result1 = solve(S1, o1)
    result2 = solve(S2, o2)
    result3 = solve(S3, o3)

    # ========
    # Plotting
    # ========
    plot(result1,color="red")
    plot!(result2,color="blue")
    plot!(result3,color="green")
    @eval(savefig(@filename_to_png))
end # function

compute()


struct Interval{T<:Real}
    a::T
    b::T
    # Interval(a,b) = a > b ? error("out of order") : new(a,b)
end

function ∈(x::T, interval::Interval) where T<:Real
    if interval.a ≤ x ≤ interval.b
        return true
    else
        return false
    end
end

function ∩(I::Interval, J::Interval)
    min = maximum([I.a, J.a])
    max = minimum([I.b, J.b])

    return Interval(min, max)
end

mutable struct MapperInformations
    data_number
    original_coordinate
    original_subset
    clustered_subset
    filtered_coordinate
    interval_subset
end

function interval_test(Info, Intervals)
    M = length(Info)
    N = length(Intervals)

    X_subset = [[] for i ∈ 1:N]
    for i ∈ 1:M
        for j ∈ 1:N
            if Info[i].filtered_coordinate ∈ Intervals[j]
                Info[i].interval_subset = Intervals[j]
                Info[i].original_subset = push!(Info[i].original_subset, j)
                push!(X_subset[j], Info[i].original_coordinate)
            end
        end
    end

    return X_subset
end

function domain_clustering(Info, X_subset)

    X_subset_clustering = []

    for i ∈ 1:length(X_subset)
        # i=3
        temp = zeros(length(X_subset[i][1]), length(X_subset[i]))
        for j ∈ 1:length(X_subset[i])
            temp[:, j] = X_subset[i][j]
        end
        scatter(temp[1,:], temp[2,:])

        D = pairwise(Euclidean(), temp, temp)

        SL = hclust(D, linkage=:single)
        plot(SL)

        dbsacn_temp = dbscan(reshape(SL.heights, (1,length(SL.heights))), 0.2maximum(SL.heights))
        cutree_temp = cutree(SL, h=SL.heights[dbsacn_temp[1].core_indices[end]])

        X_subset_clustering_temp = [[] for j ∈ 1:maximum(cutree_temp)]

        for n ∈ 1:length(X_subset[i])
            clustering_index = cutree_temp[n]
            data_index = findfirst(==(X_subset[i][n]), original_X)
            push!(X_subset_clustering_temp[clustering_index], data_index)
            Info[data_index].clustered_subset = push!(Info[data_index].clustered_subset, (i, clustering_index))
        end

        push!(X_subset_clustering, X_subset_clustering_temp)

    end
    return X_subset_clustering
end

function coor_mean(X_subset_clustering, Info)
    center_ = [ [] for i ∈ 1:length(X_subset_clustering)]
    for i ∈ 1:length(X_subset_clustering)
        for j ∈ 1:length(X_subset_clustering[i])
            temp = [Info[v].original_coordinate for v ∈ X_subset_clustering[i][j]]
            center_vector = sum(temp)./length(temp)
            push!(center_[i], center_vector)
        end
    end

    return center_
end

function mapper(X, filter, interval, overlab)

    original_X = []
    for i in 1:size(X,2)
        push!(original_X, X[:,i])
    end
    # findfirst(==(Info[500].original_coordinate), original_X)

    Info = [MapperInformations(i, X[:, i], Int64[], Tuple{Int64, Int64}[], :None, :None) for i ∈ 1:length(X[1,:])]

    filtered_X = filter(X)
    filtered_coordinate.(Info, filtered_X)
    Info

    min_X = minimum(filtered_X)
    max_X = maximum(filtered_X)
    filter_range = Interval(min_X, max_X)

    length_ = (max_X - min_X)/interval
    ϵ = overlab*length_

    Intervals = [∩(Interval(min_X +(i-1)*length_ - ϵ, min_X +i*length_ + ϵ), filter_range) for i ∈ 1:interval]

    X_subset = interval_test(Info, Intervals)

    Info

    X_subset_clustering = domain_clustering(Info, X_subset)

    # function myshowall(io, x, limit = false)
    #     println(io, summary(x), ":")
    #     Base.print_matrix(IOContext(io, :limit => limit), x)
    # end
    # myshowall(Base.stdout, Info, false)

    # X_subset_clustering[1]
    # X_subset_clustering[2]
    # X_subset_clustering[3]
    # X_subset_clustering[4]
    # X_subset_clustering[5]

    link = []

    for i ∈ 1:length(X_subset)-1
        for j ∈ 1:length(X_subset_clustering[i])
            for k ∈ 1:length(X_subset_clustering[i+1])
                if length(∩(X_subset_clustering[i][j], X_subset_clustering[i+1][k])) > 0
                    push!(link, [[i,j], [i+1,k]])
                end
            end
        end
    end

    center_ = coor_mean(X_subset_clustering, Info)

    center_number = 0
    for i ∈ 1:length(X_subset)
        for j ∈ 1:length(X_subset_clustering[i])
            center_number +=1
        end
    end

    center_matrix = zeros(2, center_number)
    k_temp = 1
    for i ∈ 1:length(X_subset)
        for j ∈ 1:length(X_subset_clustering[i])
            center_matrix[:,k_temp] = center_[i][j]
            k_temp += 1
        end
    end

    zcolor_temp = filter(cat(Data, center_matrix, dims=2))

    zcolor = [[] for i ∈ 1:length(X_subset)]
    k_temp = 1
    for i ∈ 1:length(X_subset)
        for j ∈ 1:length(X_subset_clustering[i])
            push!(zcolor[i], zcolor_temp[size(Data)[2]+k_temp])
            k_temp +=1
        end
    end

    return link, X_subset_clustering, Info, center_, zcolor
end

function mapper_plot(link, X_subset_clustering, center_, zcolor, xlim, ylim)

    linked_center = zeros(length(link), 4)
    for i ∈ 1:length(link)
        linked_center[i, 1:2] = center_[link[i][1][1]][link[i][1][2]]
        linked_center[i, 3:4] = center_[link[i][2][1]][link[i][2][2]]
    end
    linked_center

    p = plot([linked_center[1,1], linked_center[1,3]], [linked_center[1,2], linked_center[1,4]], color=:black, xlim=xlim, ylim=ylim, label="")
    for i ∈ 2:length(link)
        plot!([linked_center[i,1], linked_center[i,3]], [linked_center[i,2], linked_center[i,4]], color=:black, label="")
    end

    for i ∈ 1:length(X_subset_clustering)
        for j ∈ 1:length(X_subset_clustering[i])
            scatter!([center_[i][j][1]], [center_[i][j][2]], markersize=25*(452.106/181)*sqrt(length(X_subset_clustering[i][j])/500), zcolor=zcolor[i][j], label="", color=:jet, clim=(0,Inf))
            length_temp = length(X_subset_clustering[i][j])
            annotate!([center_[i][j][1]], [center_[i][j][2]], "$length_temp")
        end
    end

    return p
end


#===============================================================================

    Updating rule

    Based on code on Ben Moll's webite:
            www.princton.edu/~moll/HACTproject.htm

===============================================================================#

#Need to figure out best way to transfer inputs in julia
include("Huggett_initial.jl")
include("Huggett_terminal.jl")

# Get values for g, r, A from initial
g0, gg0, r, A = initial_val()

r00 = r
A0=A
AT0 = A'

# Get values for g_st, v_st, r_st from terminal
g_st, v_st, r_st = terminal()

plot(a,v_st)

plot(a, g)
plot!(a,g0, xlims=(-0.15,1))

# Create time and time increment variables
T = 20
N = 100
H =1000
dt = T/N

# Create and initial guess of the interest rate sequence
r0 = r_st*ones(N,1)

# Set up space to save variables
S = zeros(N,1)

v = zeros(H,2,N)

v[:,:,N] = v_st

r_new = r0
r_t =r0
A_t = []

maxit=1000
r_it =zeros(N,maxit)
Sdist = zeros(maxit,1)
dVf, dVb, c = [zeros(H,2) for i in 1:3]
dS_it=[]
SS_it=[]
gg=zeros(N+1,2*H,1)

convergence_criterion = 10^(-5)

# Speed of updating the interest rate
ξ = 20*(exp.(-0.05*(1:N))) .- exp(-0.05*N)

# main loop

for it in 1:maxit
    global r_t = r_new
    r_it[:,it] = r_t
    V = v_st

    for n in N:-1:1
        v[:,:,n]=V
        #forward difference
        dVf[1:H-1,:] = (V[2:H,:]-V[1:H-1,:])/da
        dVf[H,:] = (z .+ r_t[n].*amax).^(-s) #will never be used, but impose state constraint a<=amax just in case
        #backward difference
        dVb[2:H,:] = (V[2:H,:]-V[1:H-1,:])/da
        dVb[1,:] = (z .+ r_t[n].*amin).^(-s) #state constraint boundary condition

        I_concave = dVb .> dVf #indicator whether value function is concave (problems arise if this is not the case)

        # consumption and savings with forward difference
        cf = dVf.^(-1/s)
        ssf = zz + r_t[n].*aa - cf
        # consumption and savings with backward difference
        cb = dVb.^(-1/s)
        ssb = zz + r_t[n].*aa - cb
        # consumption and derivative of value function at steady state
        c0 = zz + r_t[n].*aa
        dV0 = c0.^(-s)

        #= dV_upwind makes a choice of forward or backward differences based on
         the sign of the drift    =#
        If = ssf .> 0 #positive drift → forward difference
        Ib = ssb .< 0 #negative drift → backward difference
        I0 = (1.0 .-If-Ib) #at steady state

        dV_Upwind = dVf.*If + dVb.*Ib + dV0.*I0  #important to include third term
        c = dV_Upwind.^(-1/s)
        u = c.^(1-s)/(1-s)

        # CONSTRUCT MATRIX
        X = - min.(ssb,0)/da
        Y = - max.(ssf,0)/da + min.(ssb,0)/da
        Z = max.(ssf,0)/da

        A1=spdiagm(0 => Y[:,1], -1 => X[2:H,1], 1 =>Z[1:H-1,1])
        A2=spdiagm(0=> Y[:,2], -1 =>X[2:H,2], 1 => Z[1:H-1,2])

        global A = [A1 spzeros(H,H);spzeros(H,H) A2] + Aswitch

        push!(A_t,A)
        B = (1/dt + ρ)*sparse(I, 2*H, 2*H) - A

        u_stacked = [u[:,1];u[:,2]]
        V_stacked = [V[:,1];V[:,2]]

        b = u_stacked + V_stacked/dt
        V_stacked = B\b #SOLVE SYSTEM OF EQUATIONS

        V = [V_stacked[1:H] V_stacked[H+1:2*H]]

        global ss = zz+ r_t[n].*aa -c
    end

    gg[1,:,:] = gg0

    #loop for KFE, solve forward
    global SS = zeros(N,1)
    global dS = zeros(N,1)

    for q in 1:N
        global SS, dS = SS, dS
        AT=A_t[end-(q-1)]' #the indexing is backward from original code.
        gg[q+1,:,:] = (sparse(I,2*H,2*H)-AT*dt)\gg[q,:,:]
        global SS[q] = (gg[q,:,:]'*aa[:].*da)[1]
        global dS[q] = (gg[q+1,:,:]'*aa[:]*da - gg[q,:,:]'*aa[:]*da)[1]
    end

    push!(dS_it,dS)
    push!(SS_it, SS)

    global r_new = r_t - ξ.*dS

    Sdist[it] = findmax(abs.(dS))[1]
    println("Iteration = $(it)")
    println("Convergence Criterion= $(Sdist[it])")

    if Sdist[it] < convergence_criterion
        break
    end
end

plot(Sdist[:])

# Correct graph, previously had some scoping issues with Julia
plot(1:N, SS_it[1], ylabel="Excess Supply", xlabel="Excess Demand", label="First iteration" )
plot!(1:N, SS_it[end], label="Last Iternation")

plot(r_t[:], legend=false, xlabel="Period", ylabel="r")
plot!(1:N, r_st.*ones(N,1), line=:dash)

# More plots

N1 =4
T1=-N1*dt
time = (1:N)'*dt
time1 = T1 .+ (1:N1)'*dt
time2 = [time1'; time']
r_t2 = [r00*ones(N1,1); r_t]

plot(time2, r_t2, xlims=(T1, 10), ylims=(-0.05, 0.035), legend=:false)
plot!(time2, r_st*ones(N1+N,1), line=:dash,
        xlabel="Year", title="Equilibrium Interest Rate")

amax1 = 0.5
gmax=3
n=2

p1 = plot(a, gg[n,1:H,:],xlims=(amin,amax1),
        ylims=(0, gmax), xlabel="Wealth", ylabel="Densities",legend=false)
plot!(a, gg[n,H+1:end,:], title="\$t=0.1\$")

# set up new index
t=2
n=convert(Int64, t/dt)

p2 = plot(a, gg[n,1:H,:],xlims=(amin,amax1),
        ylims=(0, gmax), xlabel="Wealth", ylabel="Densities",legend=false)
plot!(a, gg[n,H+1:end,:], title="\$t=2\$")

t=5
n=convert(Int64, t/dt)

p3 = plot(a, gg[n,1:H,:],xlims=(amin,amax1),
        ylims=(0, gmax), xlabel="Wealth", ylabel="Densities",legend=false)
plot!(a, gg[n,H+1:end,:], title="\$t=5\$")

t=T
n=convert(Int64, t/dt)

p4 = plot(a, gg[n,1:H,:],xlims=(amin,amax1),
        ylims=(0, gmax), xlabel="Wealth", ylabel="Densities",legend=false)
plot!(a, gg[n,H+1:end,:], title="\$t=\\infty\$")

plot(p1,p2,p3,p4,legend=false,layout = (2,2))


#=
Find nodes for tet elements.
=#
include("tet_quadrature_table.jl");

# Build global nodes for a triangular element with refel and vertices v
function tetrahedron_element_nodes(refel, v)
    return  tetrahedron_refel_to_xyz(refel.r[:,1], refel.r[:,2], refel.r[:,3], v);
end

# Set up refel nodal array
function tetrahedron_refel_nodes!(refel)
    (eqx, eqy, eqz) = tetrahedron_equilateral_nodes(refel.N);
    (r, s, t) = tetrahedron_equilateral_to_rst(eqx,eqy,eqz);

    # Adjust things within 1e-15 of the -1 boundaries to be exactly -1 and close to 0 to be 0
    tet_snap_to_bdry!(r, s, t, 1e-15);

    refel.r = zeros(refel.Np,3);
    refel.wr = zeros(refel.Np);

    refel.r[:,1] = r;
    refel.r[:,2] = s;
    refel.r[:,3] = t;

    # quadrature nodes/weights from a table
    xyzw = tetrahedron_quadrature_nodes_weights(refel.N+1);
    refel.g = xyzw[:,1:3];
    refel.wg = xyzw[:,4];

    # face node maps
    tol = 1e-12;
    tf1(x) = abs(x[1] + 1) < tol;
    tf2(x) = abs(x[2] + 1) < tol;
    tf3(x) = abs(x[3] + 1) < tol;
    tf4(x) = abs(x[1] + x[2] + x[3] + 1) < tol;
    refel.face2local = [get_face2local_map(refel.r, tf1),
                        get_face2local_map(refel.r, tf2),
                        get_face2local_map(refel.r, tf3),
                        get_face2local_map(refel.r, tf4)];

end

# Purpose  : Compute (x,y,z) nodes in equilateral tetrahedron for polynomial of order N
function tetrahedron_equilateral_nodes(N)
    # optimal alpha values for N up to 16 (from book)
    alpopt = [0, 0, 0, 0.1002, 1.1332, 1.5608, 1.3413, 1.2577, 1.1603, 1.10153, 0.6080, 0.4523, 0.8856, 0.8717, 0.9655];

    # Set optimized parameter, alpha, depending on order N
    if (N<16)
        alpha = alpopt[N];
    else
        alpha = 1;
    end

    # total number of nodes
    Np = Int64((N+1)*(N+2)*(N+3)/6);

    # Create equidistributed nodes on equilateral tetrahedron
    L1 = zeros(Np);
    L2 = zeros(Np);
    L3 = zeros(Np);
    L4 = zeros(Np);
    sk = 1;
    for n=1:(N+1)
        for m=1:(N+2-n)
            for q=1:(N+3-n-m)
                L1[sk] = (n-1)/N; # (1 + (-1 + (n-1)*2/N))/2
                L2[sk] = (m-1)/N; # (1 + (-1 + (m-1)*2/N))/2
                L3[sk] = 1 - (q+m+n-3)/N; # -(1 + -1 + (q-1)*2/N + -1 + (m-1)*2/N + -1 + (n-1)*2/N)/2
                L4[sk] = (q-1)/N; # (1 + (-1 + (q-1)*2/N))/2
                sk = sk+1;
            end
        end
    end

    # set vertices of tetrahedron
    invsqrt3 = 0.5773502691896258; # 1/sqrt(3)=0.5773502691896258
    invsqrt6 = 0.4082482904638631; # 1/sqrt(6)=0.4082482904638631
    v1 = [-1, -invsqrt3,   -invsqrt6];
    v2 = [ 1, -invsqrt3,   -invsqrt6];
    v3 = [ 0,  2*invsqrt3, -invsqrt6];
    v4 = [ 0,  0,          3*invsqrt6];

    # orthogonal axis tangents on faces 1-4
    t1 = zeros(4,3);
    t2 = zeros(4,3)
    t1[1,:] = v2.-v1;
    t1[2,:] = v2.-v1;
    t1[3,:] = v3.-v2;
    t1[4,:] = v3.-v1;
    t2[1,:] = v3.-0.5 .* (v1.+v2);
    t2[2,:] = v4.-0.5 .* (v1.+v2);
    t2[3,:] = v4.-0.5 .* (v2.+v3);
    t2[4,:] = v4.-0.5 .* (v1.+v3);

    for n=1:4 # normalize tangents
        t1[n,:] = t1[n,:] ./ norm(t1[n,:]);
        t2[n,:] = t2[n,:] ./ norm(t2[n,:]);
    end

    # Warp and blend for each face (accumulated in shiftXYZ)
    XYZ = zeros(Np,3);
    shift = zeros(Np,3);
    for i=1:Np
        for j=1:3
            XYZ[i,j] = L3[i]*v1[j] + L4[i]*v2[j] + L2[i]*v3[j] + L1[i]*v4[j]; # form undeformed coordinates
        end
    end

    for face=1:4
        if face==1     La = L1; Lb = L2; Lc = L3; Ld = L4;
        elseif face==2 La = L2; Lb = L1; Lc = L3; Ld = L4;
        elseif face==3 La = L3; Lb = L1; Lc = L4; Ld = L2;
        elseif face==4 La = L4; Lb = L1; Lc = L3; Ld = L2;
        end

        # compute warp tangential to face
        (warp1, warp2) = tetrahedron_warpfactor(N, Np, alpha, Lb, Lc, Ld);

        blend = Lb.*Lc.*Ld;   # compute volume blending

        denom = (Lb .+ .5 .* La).*(Lc .+ .5 .* La).*(Ld .+ .5 .* La);   # modify linear blend
        tol = 1e-10
        for ids=1:Np
            if denom[ids] > tol
                blend[ids] = (1 .+ (alpha.*La[ids]).^2).*blend[ids]./denom[ids];
            end
        end

        # compute warp & blend
        shift = shift .+ (blend.*warp1) * t1[face,:]' .+ (blend.*warp2) * t2[face,:]';

        # fix face warp
        for ids=1:Np
            numtrue = 0;
            if Lb[ids]>tol
                numtrue += 1;
            end
            if Lc[ids]>tol
                numtrue += 1;
            end
            if Ld[ids]>tol
                numtrue += 1;
            end
            if La[ids]<tol && (numtrue < 3)
                shift[ids,:] = warp1[ids]*t1[face,:] + warp2[ids]*t2[face,:];
            end
        end
    end

    XYZ = XYZ .+ shift;

    return (XYZ[:,1], XYZ[:,2], XYZ[:,3]);
end

# Compute scaled warp function at order N based on rout interpolation nodes
function tetrahedron_warpfactor(N, Np, alpha, L1, L2, L3)
    # Compute LGL and equidistant node distribution
    (LGLr, w) = jacobi_LGL_quad(N);

    # 2) compute blending function at each node for each edge
    blend1 = L2.*L3;
    blend2 = L1.*L3;
    blend3 = L1.*L2;

    # 3) amount of warp for each node, for each edge
    warp1 = zeros(Np);
    warp2 = zeros(Np);
    warp3 = zeros(Np);

    dx = zeros(Np);
    dy = zeros(Np);

    xeq = zeros(N+1);
    for i=1:N+1
        xeq[i] = -1 + 2*(N+1-i)/N;
    end

    for wi=1:Np
        for i=1:N+1
            d1 = -LGLr[i] - xeq[i];
            d2 = -LGLr[i] - xeq[i];
            d3 = -LGLr[i] - xeq[i];
            for j=2:N
                if i!=j
                    d1 = d1.*(L3[wi]-L2[wi]-xeq[j])/(xeq[i]-xeq[j]);
                    d2 = d2.*(L1[wi]-L3[wi]-xeq[j])/(xeq[i]-xeq[j]);
                    d3 = d3.*(L2[wi]-L1[wi]-xeq[j])/(xeq[i]-xeq[j]);
                end
            end

            if i!=1
                d1 = -d1/(xeq[i]-xeq[1]);
                d2 = -d2/(xeq[i]-xeq[1]);
                d3 = -d3/(xeq[i]-xeq[1]);
            end

            if i!=(N+1)
                d1 = d1/(xeq[i]-xeq[N+1]);
                d2 = d2/(xeq[i]-xeq[N+1]);
                d3 = d3/(xeq[i]-xeq[N+1]);
            end

            warp1[wi] = warp1[wi]+d1;
            warp2[wi] = warp2[wi]+d3;
            warp3[wi] = warp3[wi]+d3;
        end

        # 4) combine blend & warp
        warp1[wi] = 4 * blend1[wi] * warp1[wi] * (1 + (alpha * L1[wi])^2);
        warp2[wi] = 4 * blend2[wi] * warp2[wi] * (1 + (alpha * L2[wi])^2);
        warp3[wi] = 4 * blend3[wi] * warp3[wi] * (1 + (alpha * L3[wi])^2);

        # 5) evaluate shift in equilateral triangle
        # dx[wi] = 1*warp1[wi] + cos(2*pi/3)*warp2[wi] + cos(4*pi/3)*warp3[wi];
        # dy[wi] = 0*warp1[wi] + sin(2*pi/3)*warp2[wi] + sin(4*pi/3)*warp3[wi];
        dx[wi] = warp1[wi] - 0.5*warp2[wi] - 0.5*warp3[wi];
        dy[wi] = 0.8660254037844387*warp2[wi] - 0.8660254037844387*warp3[wi];
    end

    return (dx, dy);
end

# From (x,y,z) in equilateral tetrahedron to (r,s,t) coordinates in standard tetrahedron
function tetrahedron_equilateral_to_rst(x,y,z)
    invsqrt3 = 0.5773502691896258; # 1/sqrt(3)=0.5773502691896258
    invsqrt6 = 0.4082482904638631; # 1/sqrt(6)=0.4082482904638631
    v1 = [-1, -invsqrt3,   -invsqrt6];
    v2 = [ 1, -invsqrt3,   -invsqrt6];
    v3 = [ 0,  2*invsqrt3, -invsqrt6];
    v4 = [ 0,  0,          3*invsqrt6];

    # back out right tet nodes
    rhs = zeros(3,length(x));
    rhs[1,:] = x .- (0.5 * (v2[1]+v3[1]+v4[1]-v1[1]));
    rhs[2,:] = y .- (0.5 * (v2[2]+v3[2]+v4[2]-v1[2]));
    rhs[3,:] = z .- (0.5 * (v2[3]+v3[3]+v4[3]-v1[3]));

    A = [0.5 .* (v2 .- v1) 0.5 .* (v3 .- v1) 0.5 .* (v4 .- v1)];
    RST = A \ rhs;

    return (RST[1,:], RST[2,:], RST[3,:]);
end

# # From (r,s,t) coordinates in reference tetrahedron to (x,y,z) in tetrahedron with vertices v
# # v is a 3x4 array [x1 x2 x3 x4; y1 y2 y3 y4; z1 z2 z3 z4]
# function tetrahedron_refel_to_xyz(r, s, t, v)
#     d = v[:,1];
#     A = [v[:,2].-d v[:,3].-d v[:,4].-d];

#     np = length(r);
#     mv = zeros(3,np);
#     for i=1:np
#         tmp = [(r[i]+1)/2, (s[i]+1)/2, (t[i]+1)/2];
#         mv[:,i] = A*tmp + d;
#     end

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

#     return (x, y, z);
# end

# function get_face2local_map(r, compare)
#     n = size(r,1);
#     map = zeros(Int,n);
#     nf = 0;
#     for i=1:n
#         if compare(r[i,:])
#             nf = nf+1;
#             map[nf] = i;
#         end
#     end

#     return map[1:nf];
# end

function tet_snap_to_bdry!(r, s, t, tol)
    n = length(r);
    for i=1:n
        if abs(r[i]+1) < tol
            r[i] = -1;
        end
        if abs(s[i]+1) < tol
            s[i] = -1;
        end
        if abs(t[i]+1) < tol
            t[i] = -1;
        end

        if abs(r[i]) < tol
            r[i] = 0;
        end
        if abs(s[i]) < tol
            s[i] = 0;
        end
        if abs(t[i]) < tol
            t[i] = 0;
        end
    end
end


# D2Q21
# See lectures in LBM

# Same as D2Q17, but vertical and horizontal vectors are extended by one

struct D2Q21{
    T <: Real,
    Abscissaes <: AbstractMatrix{<:Integer},
    Weights <: AbstractVector{<:Real},
} <: Quadrature
    abscissae::Abscissaes
    weights::Weights
    speed_of_sound_squared::T
end
function D2Q21()
    cs = sqrt(2 / 3)
    w_0 = 91 / 324
    w_1 = 1 / 12
    w_2 = 2 / 27
    w_3 = 7 / 360
    w_4 = 1 / 432
    w_5 = 0 / 36
    w_6 = 1 / 1620

    # https://strathprints.strath.ac.uk/29195/1/Zhang_YH_Pure_Accuracy_analysis_of_high_order_lattice_Boltzmann_models_for_rarefied_gas_flows_21_Jan_2011.pdf
    w_0 = 91 / 324
    w_1 = 1 / 12
    w_2 = 2 / 27
    w_3 = 7 / 360
    w_4 = 1 / 432
    w_5 = 1 / 1620
    w_6 = 0
    cs = 3 / 2

    return D2Q21(
        [   #A #B          #C          #D          #E          #F
            0 1 -1 0 0 1 -1 1 -1 2 -2 0 0 2 -2 2 -2 3 -3 0 0 3 -3 -3 3
            0 0 0 1 -1 1 -1 -1 1 0 0 2 -2 2 -2 -2 2 0 0 3 -3 3 -3 3 -3
        ],
        [
            w_0,
            w_1,
            w_1,
            w_1,
            w_1,
            w_2,
            w_2,
            w_2,
            w_2,
            w_3,
            w_3,
            w_3,
            w_3,
            w_4,
            w_4,
            w_4,
            w_4,
            w_5,
            w_5,
            w_5,
            w_5,
            w_6,
            w_6,
            w_6,
            w_6,
        ],
        cs,
    )
end

order(q::D2Q21) = 7
Base.show(io::IO, q::D2Q21) = show(io, "D2Q21")
Base.string(q::D2Q21) = "D2Q21"


using ChaosTools
using Test

@testset "DyCA Tests" begin
    println("Running dyca() tests...")

    @testset "Eigenvalues of Roessler system" begin
        # Initial conditions from `dyca()` paper
        eigenthresold = collect(0.10:0.01:0.995)
        ds = Systems.roessler(rand(3); a = 0.15, b = 0.2, c = 10);
        ts = trajectory(ds, 1000.0, dt = 0.05);
        for thresold in eigenthresold
            eigenvalues,_ = dyca(Matrix(ts),thresold)
            @test sum([round(i) for i in eigenvalues]) == 2.0
        end
    end

    @testset "Eigenvalues of Embeded Roessler system" begin
        # Initial conditions from `dyca()` paper
        eigenthresold = collect(0.90:0.01:0.995)
        ds = Systems.roessler(rand(3); a = 0.15, b = 0.2, c = 10);
        ts = trajectory(ds, 1000.0, dt = 0.05);
        cols = columns(ts);
        for thresold in eigenthresold
            for i in 1:3
                Embedded_system = Matrix(embed(cols[i], 25, 1));
                eigenvals,_ = dyca(Embedded_system,thresold)
                # eigenvalues satisfy `dyca` condtiion and have imaginary part = 0.0
                wanted_eigenvalues = eigenvals[vec(0.999 .< abs.(eigenvals) .< 1.0) .& vec(imag(eigenvals) .== 0)]
                # due to floating point errors more than 2 eigenvalues might satisy the above condition
                # might be better to use `length(psum) in [1,2,3]`
                @test length(wanted_eigenvalues) in 0:4
            end
        end
    end
end