### A Pluto.jl notebook ###
# v0.19.0

using Markdown
using InteractiveUtils

# ╔═╡ 9404760f-4015-46b1-8d2b-1bec68024dbb
using Colors, Manopt, Manifolds, PlutoUI

# ╔═╡ 0fabb176-c3e8-11ec-1304-f3d3346be36f
md"""
# Illustration of Jacobi Fields

This tutorial illustrates the usage of Jacobi Fields within
`Manopt.jl`.
For this tutorial you should be familiar with the basic terminology on a
manifold like the exponential and logarithmic map as well as
[shortest geodesic](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.shortest_geodesic-Tuple{Manifold,Any,Any})s.

We first initialize the packages we need
"""

# ╔═╡ cde64039-7fd7-4897-bc31-2e45b7ef18cd
md"and we define some colors from [Paul Tol](https://personal.sron.nl/~pault/)"

# ╔═╡ c8590a06-4825-4e4d-b3fa-7a82a8982063
begin
    black = RGBA{Float64}(colorant"#000000")
    TolVibrantMagenta = RGBA{Float64}(colorant"#EE3377")
    TolVibrantOrange = RGBA{Float64}(colorant"#EE7733")
    TolVibrantCyan = RGBA{Float64}(colorant"#33BBEE")
    TolVibrantTeal = RGBA{Float64}(colorant"#009988")
end;

# ╔═╡ 2028d1ab-3dcb-4f17-a02f-f9ef8c2dfacb
md"And setup our graphics paths"

# ╔═╡ 4e1fc0af-2837-40cc-b6ff-3d8d6d641b6b
begin
    localpath = join(splitpath(@__FILE__)[1:(end - 1)], "/") # files folder
    image_prefix = localpath * "/jacobi_fields"
    @info image_prefix
    render_asy = false # on CI or when you do not have asymptote, this should be false
end;

# ╔═╡ 17803b9a-5b91-4d16-ae52-06a83653dea2
md"""
Assume we have two points on the equator of the [Sphere](https://juliamanifolds.github.io/Manifolds.jl/stable/manifolds/sphere.html) $\mathcal M = \mathbb S^2$
"""

# ╔═╡ bab2db18-0596-4716-a6ce-8bfaf2b8ca24
M = Sphere(2)

# ╔═╡ 9e0e4098-d6c4-49dd-9564-961c73be88d5
p, q = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]

# ╔═╡ ee9a8c21-431e-4448-b758-01e552b1df8c
md"""
their connecting [shortest geodesic](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.shortest_geodesic-Tuple{Manifold,Any,Any}) (sampled at `100` points)
"""

# ╔═╡ 8a0d1039-5cfe-4e4d-9a97-89c3b5705b36
geodesic_curve = shortest_geodesic(M, p, q, [0:0.1:1.0...]);

# ╔═╡ 5e0d7686-7129-41f6-87d7-7298ddb131a5
md"Is just a curve along the equator"

# ╔═╡ b75c32b0-e322-42e5-a8fa-f74893655e18
render_asy && asymptote_export_S2_signals(
    image_prefix * "/jacobi_geodesic.asy";
    curves=[geodesic_curve],
    points=[[p, q]],
    colors=Dict(:curves => [black], :points => [TolVibrantOrange]),
    dot_size=3.5,
    line_width=0.75,
    camera_position=(1.0, 1.0, 0.5),
);

# ╔═╡ 7e3d4b24-ad82-4146-93c4-3bd4011ab151
render_asy && render_asymptote(image_prefix * "/jacobi_geodesic.asy"; render=2);

# ╔═╡ 34622df6-2979-4f05-8be9-224a2ec413f9
PlutoUI.LocalResource(image_prefix * "/jacobi_geodesic.png")

# ╔═╡ f0118459-5343-45dd-aa62-494795019044
md"""
where ``p`` is on the left and ``q`` on the right.
We know solve the following task:

Given a direction $X ∈ T_p\mathcal M$, for example
"""

# ╔═╡ a7b379ac-293e-454e-8ede-6f81bb7caa6e
X = [0.0, 0.4, 0.5]

# ╔═╡ f338d7d0-eef5-4264-af90-33e5bc6ecc62
md"""
we move the start point ``x`` into, how does any point on the geodesic move?

Or mathematically: Compute ``D_p g(t; p,q)`` for some fixed ``t∈[0,1]``
and a given direction ``X_p``.
Of course two cases are quite easy: For ``t=0`` we are in ``x`` and how ``x`` “moves”
is already known, so ``D_x g(0;p,q) = X``. On the other side, for ``t=1``,
``g(1; p,q) = q`` which is fixed, so ``D_p g(1; p,q)`` is the zero tangent vector
(in ``T_q\mathcal M``).

For all other cases we employ a [`jacobi_field`](https://manoptjl.org/stable/functions/Jacobi_fields.html), which is a (tangent)
vector field along the [shortest geodesic](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.shortest_geodesic-Tuple{Manifold,Any,Any}) given as follows: The _geodesic variation_
``\Gamma_{g,X}(s,t)`` is defined for some ``\varepsilon > 0`` as
````math
\Gamma_{g,X}(s,t):=\exp{\gamma_{p,X}(s)}[t\log_{g(s;p,X)}p],\qquad s∈(-\varepsilon,\varepsilon),\ t∈[0,1].
````
Intuitively we make a small step ``s`` into direction ``ξ`` using the geodesic
``g(⋅; p,X)`` and from ``r=g(s; p,X)`` we follow (in ``t``) the geodesic
``g(⋅; r,q)``. The corresponding Jacobi field ``J_{g,X}``
along ``g(⋅; p,q)`` is given by

````math
J_{g,X}(t):=\frac{D}{\partial s}\Gamma_{g,X}(s,t)\Bigl\rvert_{s=0}
````

which is an ODE and we know the boundary conditions ``J_{g,X}(0)=X`` and
``J_{g,X}(t) = 0``. In symmetric spaces we can compute the solution, since the
system of ODEs decouples, see for example [^doCarmo1992],
Chapter 4.2. Within `Manopt.jl` this is implemented as
`jacobi_field(M,p,q,t,X[,β])`, where the optional parameter (function)
``β`` specifies, which Jacobi field we want to evaluate and the one used here is
the default.

We can hence evaluate that on the points on the geodesic at
"""

# ╔═╡ b457c7c8-7f5c-4b92-a064-0f36b188f7a1
T = [0:0.1:1.0...];

# ╔═╡ c2ba61e5-7f74-4cb7-b68d-9057ab896859
md"namely"

# ╔═╡ 3e044a29-dff6-4174-9adf-f1228672b91d
r = shortest_geodesic(M, p, q, T);

# ╔═╡ c17abc34-6e5d-41d9-b428-b984bc8f7cd9
md"The tangent vector now moves as a differential along the geodesic as"

# ╔═╡ d7308356-100f-4766-9afe-588960b925dd
W = jacobi_field.(Ref(M), Ref(p), Ref(q), T, Ref(X));

# ╔═╡ fab5b16b-e05e-4a8a-b001-236422d9e245
md"""Which can also be called using `differential_geodesic_startpoint`.
We can add to the image above by creating extended tangent vectors
the include their base points"""

# ╔═╡ 0713167c-80c6-4907-b1e3-e104d598a804
V = [Tuple([a, b]) for (a, b) in zip(r, W)];

# ╔═╡ e26bb3fb-5f7e-470a-a2a1-7be37eb1bcb9
md"to add the vectors as one further set to the Asymptote export."

# ╔═╡ aa621823-57e3-4f53-8b65-1afbb80943bc
render_asy && asymptote_export_S2_signals(
    image_prefix * "/jacobi_geodesic_diff_start.asy";
    curves=[geodesic_curve],
    points=[[p, q], r],
    tangent_vectors=[V],
    colors=Dict(
        :curves => [black],
        :points => [TolVibrantOrange, TolVibrantCyan],
        :tvectors => [TolVibrantCyan],
    ),
    dot_sizes=[3.5, 2.0],
    line_width=0.75,
    camera_position=(1.0, 1.0, 0.5),
);

# ╔═╡ d02cff4d-8e88-4979-b536-6f367af8ab32
render_asy && render_asymptote(image_prefix * "/jacobi_geodesic_diff_start.asy"; render=2);

# ╔═╡ a17f81a9-066f-435c-b5ba-ae4809ac835b
PlutoUI.LocalResource(image_prefix * "/jacobi_geodesic_diff_start.png")

# ╔═╡ b2ac06aa-4b9f-4869-8b4c-c6f30661b8fe
md"""
The interpretation is as follows: If an infinitesimal change of the start point in direction ``X`` would happen, the infinitesimal change of a point along the line would change as indicated.

Note that each new vector is a tangent vector to its point (up to a small numerical tolerance), so the blue vectors are not just “shifted and scaled versions” of ``X``.
"""

# ╔═╡ fa6e15b8-2201-4095-8ebd-ddcd749fe11d
all([is_vector(M, a[1], a[2]; atol=1e-15) for a in V])

# ╔═╡ 99526567-14aa-4526-9a8d-48e2a3ba15a6
md"If we further move the end point, too, we can derive that Differential in direction"

# ╔═╡ 84090442-ea5f-4a1e-9c66-b0e828463853
begin
    Y = [0.2, 0.0, -0.5]
    W2 = differential_geodesic_endpoint.(Ref(M), Ref(p), Ref(q), T, Ref(Y))
    V2 = [Tuple([a, b]) for (a, b) in zip(r, W2)]
end;

# ╔═╡ bf3205fa-5348-42f9-8c1d-16e3fbcbc06c
md"and we can combine both vector fields we obtained"

# ╔═╡ c319b655-b3e5-4190-bf3c-d5af55da42ba
V3 = [Tuple([a, b]) for (a, b) in zip(r, W2 + W)];

# ╔═╡ 63b0d8ae-9a1b-46ff-978c-753f1cbdefe9
render_asy && asymptote_export_S2_signals(
    image_prefix * "/jacobi_geodesic_complete.asy";
    curves=[geodesic_curve],
    points=[[p, q], r],
    tangent_vectors=[V, V2, V3],
    colors=Dict(
        :curves => [black],
        :points => [TolVibrantOrange, TolVibrantCyan],
        :tvectors => [TolVibrantCyan, TolVibrantMagenta, TolVibrantTeal],
    ),
    dot_sizes=[3.5, 2.0],
    line_width=0.75,
    camera_position=(1.0, 1.0, 0.0),
);

# ╔═╡ 18278235-a049-4bd5-91dd-bddd31bcb8fb
render_asy && render_asymptote(image_prefix * "/jacobi_geodesic_complete.asy"; render=2);

# ╔═╡ cd1fa853-16ea-4975-b91b-5e6c1943523d
PlutoUI.LocalResource(image_prefix * "/jacobi_geodesic_complete.png")

# ╔═╡ ee803867-bc00-43b3-85c5-c4471370b6a8
md"""
Here the first vector field is still in blue, the second is in magenta, and their combined effect is in teal. Sure as a differential this does not make much sense, maybe as an infinitesimal movement of both start and end point cmobined.
"""

# ╔═╡ 83d85bac-85c7-462b-bf1c-097e27a3ca9a
md"""
## Literature

[^doCarmo1992]:
    > do Carmo, Manfredo __Riemannian Geometry__,  Birkhäuser Boston, 1992, ISBN: 0-8176-3490-8.
"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
Colors = "~0.12.8"
Manifolds = "~0.7.8"
Manopt = "~0.3.23"
PlutoUI = "~0.7.38"
"""

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

julia_version = "1.7.2"
manifest_format = "2.0"

[[deps.AbstractFFTs]]
deps = ["ChainRulesCore", "LinearAlgebra"]
git-tree-sha1 = "6f1d9bc1c08f9f4a8fa92e3ea3cb50153a1b40d4"
uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
version = "1.1.0"

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

[[deps.Adapt]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "af92965fb30777147966f58acb05da51c5616b5f"
uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
version = "3.3.3"

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

[[deps.ArnoldiMethod]]
deps = ["LinearAlgebra", "Random", "StaticArrays"]
git-tree-sha1 = "62e51b39331de8911e4a7ff6f5aaf38a5f4cc0ae"
uuid = "ec485272-7323-5ecc-a04f-4719b315124d"
version = "0.2.0"

[[deps.ArrayInterface]]
deps = ["Compat", "IfElse", "LinearAlgebra", "Requires", "SparseArrays", "Static"]
git-tree-sha1 = "c933ce606f6535a7c7b98e1d86d5d1014f730596"
uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
version = "5.0.7"

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

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

[[deps.ChainRulesCore]]
deps = ["Compat", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "9950387274246d08af38f6eef8cb5480862a435f"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.14.0"

[[deps.ChangesOfVariables]]
deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
git-tree-sha1 = "bf98fa45a0a4cee295de98d4c1462be26345b9a1"
uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
version = "0.1.2"

[[deps.ColorSchemes]]
deps = ["ColorTypes", "Colors", "FixedPointNumbers", "Random"]
git-tree-sha1 = "12fc73e5e0af68ad3137b886e3f7c1eacfca2640"
uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
version = "3.17.1"

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

[[deps.Colors]]
deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
git-tree-sha1 = "417b0ed7b8b838aa6ca0a87aadf1bb9eb111ce40"
uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
version = "0.12.8"

[[deps.Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "b153278a25dd42c65abbf4e62344f9d22e59191b"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.43.0"

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

[[deps.CovarianceEstimation]]
deps = ["LinearAlgebra", "Statistics", "StatsBase"]
git-tree-sha1 = "a3e070133acab996660d31dcf479ea42849e368f"
uuid = "587fd27a-f159-11e8-2dae-1979310e6154"
version = "0.2.7"

[[deps.DataAPI]]
git-tree-sha1 = "fb5f5316dd3fd4c5e7c30a24d50643b73e37cd40"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.10.0"

[[deps.DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "3daef5523dd2e769dad2365274f760ff5f282c7d"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.11"

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

[[deps.DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[deps.DensityInterface]]
deps = ["InverseFunctions", "Test"]
git-tree-sha1 = "80c3e8639e3353e5d2912fb3a1916b8455e2494b"
uuid = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
version = "0.4.0"

[[deps.Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[deps.Distributions]]
deps = ["ChainRulesCore", "DensityInterface", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SparseArrays", "SpecialFunctions", "Statistics", "StatsBase", "StatsFuns", "Test"]
git-tree-sha1 = "5a4168170ede913a2cd679e53c2123cb4b889795"
uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
version = "0.25.53"

[[deps.DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "b19534d1895d702889b219c382a6e18010797f0b"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.6"

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

[[deps.Einsum]]
deps = ["Compat"]
git-tree-sha1 = "4a6b3eee0161c89700b6c1949feae8b851da5494"
uuid = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
version = "0.4.1"

[[deps.FillArrays]]
deps = ["LinearAlgebra", "Random", "SparseArrays", "Statistics"]
git-tree-sha1 = "246621d23d1f43e3b9c368bf3b72b2331a27c286"
uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
version = "0.13.2"

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

[[deps.Graphs]]
deps = ["ArnoldiMethod", "Compat", "DataStructures", "Distributed", "Inflate", "LinearAlgebra", "Random", "SharedArrays", "SimpleTraits", "SparseArrays", "Statistics"]
git-tree-sha1 = "57c021de207e234108a6f1454003120a1bf350c4"
uuid = "86223c79-3864-5bf0-83f7-82e725a168b6"
version = "1.6.0"

[[deps.HybridArrays]]
deps = ["LinearAlgebra", "Requires", "StaticArrays"]
git-tree-sha1 = "eb6b23460f5544c5d09efae0818b86736cefcd3d"
uuid = "1baab800-613f-4b0a-84e4-9cd3431bfbb9"
version = "0.4.10"

[[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.IfElse]]
git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1"
uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
version = "0.1.1"

[[deps.Inflate]]
git-tree-sha1 = "f5fc07d4e706b84f72d54eedcc1c13d92fb0871c"
uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9"
version = "0.1.2"

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

[[deps.InverseFunctions]]
deps = ["Test"]
git-tree-sha1 = "91b5dcf362c5add98049e6c29ee756910b03051d"
uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
version = "0.1.3"

[[deps.IrrationalConstants]]
git-tree-sha1 = "7fd44fd4ff43fc60815f8e764c0f352b83c49151"
uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
version = "0.1.1"

[[deps.JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "abc9885a7ca2052a736a600f7fa66209f96506e1"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.4.1"

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

[[deps.Kronecker]]
deps = ["LinearAlgebra", "NamedDims", "SparseArrays", "StatsBase"]
git-tree-sha1 = "a51f46415c844dee694cb8b20a3fcbe6dba342c2"
uuid = "2c470bb0-bcc8-11e8-3dad-c9649493f05e"
version = "0.5.1"

[[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.LogExpFunctions]]
deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "a970d55c2ad8084ca317a4658ba6ce99b7523571"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.12"

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

[[deps.MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "3d3e902b31198a27340d0bf00d6ac452866021cf"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.9"

[[deps.Manifolds]]
deps = ["Colors", "Distributions", "Einsum", "Graphs", "HybridArrays", "Kronecker", "LinearAlgebra", "ManifoldsBase", "Markdown", "Random", "RecipesBase", "RecursiveArrayTools", "Requires", "SimpleWeightedGraphs", "SpecialFunctions", "StaticArrays", "Statistics", "StatsBase"]
git-tree-sha1 = "cc5e34ecd061bf9bedd62a8b8a28abe96c3b7c6e"
uuid = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
version = "0.7.8"

[[deps.ManifoldsBase]]
deps = ["LinearAlgebra", "Markdown"]
git-tree-sha1 = "b8518cbf0f06f1375c95479759b266bd75d89ad9"
uuid = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb"
version = "0.12.13"

[[deps.Manopt]]
deps = ["ColorSchemes", "ColorTypes", "Colors", "DataStructures", "Dates", "LinearAlgebra", "ManifoldsBase", "Markdown", "Printf", "Random", "Requires", "SparseArrays", "StaticArrays", "Statistics", "Test"]
git-tree-sha1 = "a8a292a422b0a6924f46d4e022fc8263721120c5"
uuid = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
version = "0.3.23"

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

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

[[deps.Missings]]
deps = ["DataAPI"]
git-tree-sha1 = "bf210ce90b6c9eed32d25dbcae1ebc565df2687f"
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
version = "1.0.2"

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

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

[[deps.NamedDims]]
deps = ["AbstractFFTs", "ChainRulesCore", "CovarianceEstimation", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
git-tree-sha1 = "0856b62716585eb90cc1dada226ac9eab5f69aa5"
uuid = "356022a1-0364-5f58-8944-0da4b18d706f"
version = "0.2.47"

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

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

[[deps.OpenLibm_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "05823500-19ac-5b8b-9628-191a04bc5112"

[[deps.OpenSpecFun_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.5+0"

[[deps.OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[deps.PDMats]]
deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
git-tree-sha1 = "e8185b83b9fc56eb6456200e873ce598ebc7f262"
uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
version = "0.11.7"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "3b429f37de37f1fc603cc1de4a799dc7fbe4c0b6"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.3.0"

[[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 = "670e559e5c8e191ded66fa9ea89c97f10376bb4c"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.38"

[[deps.Preferences]]
deps = ["TOML"]
git-tree-sha1 = "47e5f437cc0e7ef2ce8406ce1e7e24d44915f88d"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.3.0"

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

[[deps.QuadGK]]
deps = ["DataStructures", "LinearAlgebra"]
git-tree-sha1 = "78aadffb3efd2155af139781b8a8df1ef279ea39"
uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
version = "2.4.2"

[[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.RecipesBase]]
git-tree-sha1 = "6bf3f380ff52ce0832ddd3a2a7b9538ed1bcca7d"
uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
version = "1.2.1"

[[deps.RecursiveArrayTools]]
deps = ["Adapt", "ArrayInterface", "ChainRulesCore", "DocStringExtensions", "FillArrays", "LinearAlgebra", "RecipesBase", "Requires", "StaticArrays", "Statistics", "ZygoteRules"]
git-tree-sha1 = "bfe14f127f3e7def02a6c2b1940b39d0dabaa3ef"
uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
version = "2.26.3"

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

[[deps.Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.3.0"

[[deps.Rmath]]
deps = ["Random", "Rmath_jll"]
git-tree-sha1 = "bf3188feca147ce108c76ad82c2792c57abe7b1f"
uuid = "79098fc4-a85e-5d69-aa6a-4863f24498fa"
version = "0.7.0"

[[deps.Rmath_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "68db32dff12bb6127bac73c209881191bf0efbb7"
uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f"
version = "0.3.0+0"

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

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

[[deps.SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[deps.SimpleTraits]]
deps = ["InteractiveUtils", "MacroTools"]
git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231"
uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
version = "0.9.4"

[[deps.SimpleWeightedGraphs]]
deps = ["Graphs", "LinearAlgebra", "Markdown", "SparseArrays", "Test"]
git-tree-sha1 = "a6f404cc44d3d3b28c793ec0eb59af709d827e4e"
uuid = "47aef6b3-ad0c-573a-a1e2-d07658019622"
version = "1.2.1"

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

[[deps.SortingAlgorithms]]
deps = ["DataStructures"]
git-tree-sha1 = "b3363d7460f7d098ca0912c69b082f75625d7508"
uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
version = "1.0.1"

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

[[deps.SpecialFunctions]]
deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
git-tree-sha1 = "5ba658aeecaaf96923dce0da9e703bd1fe7666f9"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "2.1.4"

[[deps.Static]]
deps = ["IfElse"]
git-tree-sha1 = "87e9954dfa33fd145694e42337bdd3d5b07021a6"
uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
version = "0.6.0"

[[deps.StaticArrays]]
deps = ["LinearAlgebra", "Random", "Statistics"]
git-tree-sha1 = "cd56bf18ed715e8b09f06ef8c6b781e6cdc49911"
uuid = "90137ffa-7385-5640-81b9-e52037218182"
version = "1.4.4"

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

[[deps.StatsAPI]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "c82aaa13b44ea00134f8c9c89819477bd3986ecd"
uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
version = "1.3.0"

[[deps.StatsBase]]
deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]
git-tree-sha1 = "8977b17906b0a1cc74ab2e3a05faa16cf08a8291"
uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
version = "0.33.16"

[[deps.StatsFuns]]
deps = ["ChainRulesCore", "InverseFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"]
git-tree-sha1 = "5950925ff997ed6fb3e985dcce8eb1ba42a0bbe7"
uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
version = "0.9.18"

[[deps.SuiteSparse]]
deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"]
uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"

[[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.ZygoteRules]]
deps = ["MacroTools"]
git-tree-sha1 = "8c1a8e4dfacb1fd631745552c8db35d0deb09ea0"
uuid = "700de1a5-db45-46bc-99cf-38207098b444"
version = "0.2.2"

[[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:
# ╟─0fabb176-c3e8-11ec-1304-f3d3346be36f
# ╠═9404760f-4015-46b1-8d2b-1bec68024dbb
# ╟─cde64039-7fd7-4897-bc31-2e45b7ef18cd
# ╠═c8590a06-4825-4e4d-b3fa-7a82a8982063
# ╟─2028d1ab-3dcb-4f17-a02f-f9ef8c2dfacb
# ╠═4e1fc0af-2837-40cc-b6ff-3d8d6d641b6b
# ╟─17803b9a-5b91-4d16-ae52-06a83653dea2
# ╠═bab2db18-0596-4716-a6ce-8bfaf2b8ca24
# ╠═9e0e4098-d6c4-49dd-9564-961c73be88d5
# ╟─ee9a8c21-431e-4448-b758-01e552b1df8c
# ╠═8a0d1039-5cfe-4e4d-9a97-89c3b5705b36
# ╟─5e0d7686-7129-41f6-87d7-7298ddb131a5
# ╠═b75c32b0-e322-42e5-a8fa-f74893655e18
# ╠═7e3d4b24-ad82-4146-93c4-3bd4011ab151
# ╠═34622df6-2979-4f05-8be9-224a2ec413f9
# ╟─f0118459-5343-45dd-aa62-494795019044
# ╟─a7b379ac-293e-454e-8ede-6f81bb7caa6e
# ╟─f338d7d0-eef5-4264-af90-33e5bc6ecc62
# ╠═b457c7c8-7f5c-4b92-a064-0f36b188f7a1
# ╟─c2ba61e5-7f74-4cb7-b68d-9057ab896859
# ╠═3e044a29-dff6-4174-9adf-f1228672b91d
# ╟─c17abc34-6e5d-41d9-b428-b984bc8f7cd9
# ╠═d7308356-100f-4766-9afe-588960b925dd
# ╟─fab5b16b-e05e-4a8a-b001-236422d9e245
# ╠═0713167c-80c6-4907-b1e3-e104d598a804
# ╟─e26bb3fb-5f7e-470a-a2a1-7be37eb1bcb9
# ╠═aa621823-57e3-4f53-8b65-1afbb80943bc
# ╠═d02cff4d-8e88-4979-b536-6f367af8ab32
# ╟─a17f81a9-066f-435c-b5ba-ae4809ac835b
# ╟─b2ac06aa-4b9f-4869-8b4c-c6f30661b8fe
# ╠═fa6e15b8-2201-4095-8ebd-ddcd749fe11d
# ╟─99526567-14aa-4526-9a8d-48e2a3ba15a6
# ╠═84090442-ea5f-4a1e-9c66-b0e828463853
# ╠═bf3205fa-5348-42f9-8c1d-16e3fbcbc06c
# ╠═c319b655-b3e5-4190-bf3c-d5af55da42ba
# ╠═63b0d8ae-9a1b-46ff-978c-753f1cbdefe9
# ╠═18278235-a049-4bd5-91dd-bddd31bcb8fb
# ╠═cd1fa853-16ea-4975-b91b-5e6c1943523d
# ╟─ee803867-bc00-43b3-85c5-c4471370b6a8
# ╟─83d85bac-85c7-462b-bf1c-097e27a3ca9a
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


"""
    reliabilityplot(probabilities, frequencies; deviation=true, kwargs...)
    reliabilityplot(probabilities, frequencies, low_high; deviation=true, kwargs...)
    reliabilityplot(probabilities, frequencies, low, high; deviation=true, kwargs...)

Plot a reliability diagram of the observed `frequencies` versus the predicted
`probabilities`, optionally with consistency bars ranging from `low` to `high` for
each probability.

If `deviation` is `true` (default), the difference `frequencies - probabilities` is
plotted versus `probabilities`. This can be helpful for inspecting how closely the
observed frequencies match the predicted probabilities.

# Examples

```jldoctest
julia> probabilities = sort(rand(10));

julia> frequencies = rand(10);

julia> reliabilityplot(probabilities, frequencies);

julia> # plot frequencies instead of deviations
       reliabilityplot(probabilities, frequencies; deviation=false);

julia> # use different colors for markers
       reliabilityplot(probabilities, frequencies; markercolor=:red);

julia> # random consistency bars
       low = max.(0, probabilities .- rand.());

julia> high = min.(1, probabilities .+ rand.());

julia> # plot pre-computed consistency bars
       reliabilityplot(probabilities, frequencies, low, high);

julia> # alternatively consistency bars can be specified as vector of tuples
       reliabilityplot(probabilities, frequencies, map(tuple, low, high));
```

# References

Bröcker, J. and Smith, L.A. (2007). *Increasing the reliability of reliability diagrams*.
Weather and forecasting, 22(3), pp. 651-661.
"""
RecipesBase.@userplot struct ReliabilityPlot{X<:AbstractVector,Y<:AbstractVector,T}
    x::X
    y::Y
    low_high::T
end

"""
    reliabilityplot(
        probabilities::AbstractVector{<:Real},
        outcomes::AbstractVector{Bool};
        binning=EqualMass(),
        consistencybars=ConsistencyBars(),
        kwargs...,
    )

Plot a reliability diagram of the observed frequencies versus the mean probabilities
in different clusters together with a set of consistency bars.

Optionally, the binning algorithm that is used to form the clusters and the consistency
bars can be configured with the `binning` and `consistencybars` keyword arguments. If
`consistencybars === nothing`, then no consistency bars are computed.

# Examples

```jldoctest
julia> probabilities = rand(10);

julia> outcomes = rand(Bool, 10);

julia> # default binning and consistency bars
       reliabilityplot(probabilities, outcomes);

julia> # custom options: without consistency bars
       reliabilityplot(probabilities, outcomes; consistencybars=nothing);
```

See also: [`EqualMass`](@ref), [`EqualSize`](@ref), [`ConsistencyBars`](@ref)
"""
reliabilityplot(::AbstractVector{<:Real}, ::AbstractVector{Bool}; kwargs...)

RecipesBase.@recipe function f(plot::ReliabilityPlot)
    # extract x and y values
    x = plot.x
    y = plot.y

    # compute mean probabilities and frequencies and consistency bars
    probs, freqs_nodeviation, low_high = if y isa AbstractVector{Bool}
        binning = get(plotattributes, :binning, EqualMass())
        consistencybars = get(plotattributes, :consistencybars, ConsistencyBars())
        means_and_bars(x, y; binning=binning, consistencybars=consistencybars)
    else
        x, y, plot.low_high
    end
    deviation::Bool = get(plotattributes, :deviation, true)
    freqs = deviation ? freqs_nodeviation - probs : freqs_nodeviation - zero(probs)

    # default attributes
    seriestype --> :line
    markershape --> :circle

    # add consistency bars, taking into account the `deviation` keyword argument
    if low_high !== nothing
        low = map(low_high, freqs_nodeviation) do (low, _), freq
            return freq - low
        end
        high = map(low_high, freqs_nodeviation) do (_, high), freq
            return high - freq
        end
        yerror := (low, high)
    end

    return probs, freqs
end


module test_julia_axes

using Test

A = [1 2; 3 4; 5 6; 7 8;]
#=
4×2 Array{Int64,2}:
 1  2
 3  4
 5  6
 7  8
=#

@test size(A) == (4, 2)
@test axes(A, 1) == Base.OneTo(4)
@test axes(A, 2) == Base.OneTo(2)

end # module test_julia_axes


gate(h, n) = (1:h) .+ h*(n-1)
gate(x::AbstractVector, h, n) = @view x[gate(h,n)]
gate(x::AbstractMatrix, h, n) = x[gate(h,n),:]

# Stateful recurrence

"""
    Recur(cell)

`Recur` takes a recurrent cell and makes it stateful, managing the hidden state
in the background. `cell` should be a model of the form:

    h, y = cell(h, x...)

For example, here's a recurrent network that keeps a running total of its inputs:

```julia
accum(h, x) = (h + x, x)
rnn = Flux.Recur(accum, 0)
rnn(2)      # 2
rnn(3)      # 3
rnn.state   # 5
rnn.(1:10)  # apply to a sequence
rnn.state   # 60
```

Folding over a 3d Array of dimensions `(features, batch, time)` is also supported:

```julia
accum(h, x) = (h .+ x, x)
rnn = Flux.Recur(accum, zeros(Int, 1, 1))
rnn([2])                    # 2
rnn([3])                    # 3
rnn.state                   # 5
rnn(reshape(1:10, 1, 1, :)) # apply to a sequence of (features, batch, time)
rnn.state                   # 60
```

"""
mutable struct Recur{T,S}
  cell::T
  state::S
end

function (m::Recur)(x)
  m.state, y = m.cell(m.state, x)
  return y
end

@functor Recur
trainable(a::Recur) = (a.cell,)

Base.show(io::IO, m::Recur) = print(io, "Recur(", m.cell, ")")

"""
    reset!(rnn)

Reset the hidden state of a recurrent layer back to its original value.

Assuming you have a `Recur` layer `rnn`, this is roughly equivalent to:
```julia
rnn.state = hidden(rnn.cell)
```
"""
reset!(m::Recur) = (m.state = m.cell.state0)
reset!(m) = foreach(reset!, functor(m)[1])

# TODO remove in v0.13
function Base.getproperty(m::Recur, sym::Symbol)
  if sym === :init
    Zygote.ignore() do
      @warn "Recur field :init has been deprecated. To access initial state weights, use m::Recur.cell.state0 instead."
    end
    return getfield(m.cell, :state0)
  else
    return getfield(m, sym)
  end
end

flip(f, xs) = reverse(f.(reverse(xs)))

function (m::Recur)(x::AbstractArray{T, 3}) where T
  h = [m(view(x, :, :, i)) for i in 1:size(x, 3)]
  sze = size(h[1])
  reshape(reduce(hcat, h), sze[1], sze[2], length(h))
end

# Vanilla RNN

struct RNNCell{F,A,V,S}
  σ::F
  Wi::A
  Wh::A
  b::V
  state0::S
end

RNNCell(in::Integer, out::Integer, σ=tanh; init=Flux.glorot_uniform, initb=zeros32, init_state=zeros32) =
  RNNCell(σ, init(out, in), init(out, out), initb(out), init_state(out,1))

function (m::RNNCell{F,A,V,<:AbstractMatrix{T}})(h, x::Union{AbstractVecOrMat{T},OneHotArray}) where {F,A,V,T}
  σ, Wi, Wh, b = m.σ, m.Wi, m.Wh, m.b
  h = σ.(Wi*x .+ Wh*h .+ b)
  sz = size(x)
  return h, reshape(h, :, sz[2:end]...)
end

@functor RNNCell

function Base.show(io::IO, l::RNNCell)
  print(io, "RNNCell(", size(l.Wi, 2), ", ", size(l.Wi, 1))
  l.σ == identity || print(io, ", ", l.σ)
  print(io, ")")
end

"""
    RNN(in::Integer, out::Integer, σ = tanh)

The most basic recurrent layer; essentially acts as a `Dense` layer, but with the
output fed back into the input each time step.

The parameters `in` and `out` describe the size of the feature vectors passed as input and as output. That is, it accepts a vector of length `in` or a batch of vectors represented as a `in x B` matrix and outputs a vector of length `out` or a batch of vectors of size `out x B`.

This constructor is syntactic sugar for `Recur(RNNCell(a...))`, and so RNNs are stateful. Note that the state shape can change depending on the inputs, and so it is good to `reset!` the model between inference calls if the batch size changes. See the examples below.

# Examples
```jldoctest
julia> r = RNN(3, 5)
Recur(
  RNNCell(3, 5, tanh),                  # 50 parameters
)         # Total: 4 trainable arrays, 50 parameters,
          # plus 1 non-trainable, 5 parameters, summarysize 432 bytes.

julia> r(rand(Float32, 3)) |> size
(5,)

julia> Flux.reset!(r);

julia> r(rand(Float32, 3, 10)) |> size # batch size of 10
(5, 10)
```

!!! warning "Batch size changes"

    Failing to call `reset!` when the input batch size changes can lead to unexpected behavior. See the following example:

    ```julia
    julia> r = RNN(3, 5)
    Recur(
      RNNCell(3, 5, tanh),                  # 50 parameters
    )         # Total: 4 trainable arrays, 50 parameters,
              # plus 1 non-trainable, 5 parameters, summarysize 432 bytes.

    julia> r.state |> size
    (5, 1)

    julia> r(rand(Float32, 3)) |> size
    (5,)

    julia> r.state |> size
    (5, 1)

    julia> r(rand(Float32, 3, 10)) |> size # batch size of 10
    (5, 10)

    julia> r.state |> size # state shape has changed
    (5, 10)

    julia> r(rand(Float32, 3)) |> size # erroneously outputs a length 5*10 = 50 vector.
    (50,)
    ```
"""
RNN(a...; ka...) = Recur(RNNCell(a...; ka...))
Recur(m::RNNCell) = Recur(m, m.state0)

# TODO remove in v0.13
function Base.getproperty(m::RNNCell, sym::Symbol)
  if sym === :h
    Zygote.ignore() do
      @warn "RNNCell field :h has been deprecated. Use m::RNNCell.state0 instead."
    end
    return getfield(m, :state0)
  else
    return getfield(m, sym)
  end
end

# LSTM

struct LSTMCell{A,V,S}
  Wi::A
  Wh::A
  b::V
  state0::S
end

function LSTMCell(in::Integer, out::Integer;
                  init = glorot_uniform,
                  initb = zeros32,
                  init_state = zeros32)
  cell = LSTMCell(init(out * 4, in), init(out * 4, out), initb(out * 4), (init_state(out,1), init_state(out,1)))
  cell.b[gate(out, 2)] .= 1
  return cell
end

function (m::LSTMCell{A,V,<:NTuple{2,AbstractMatrix{T}}})((h, c), x::Union{AbstractVecOrMat{T},OneHotArray}) where {A,V,T}
  b, o = m.b, size(h, 1)
  g = m.Wi*x .+ m.Wh*h .+ b
  input = σ.(gate(g, o, 1))
  forget = σ.(gate(g, o, 2))
  cell = tanh.(gate(g, o, 3))
  output = σ.(gate(g, o, 4))
  c = forget .* c .+ input .* cell
  h′ = output .* tanh.(c)
  sz = size(x)
  return (h′, c), reshape(h′, :, sz[2:end]...)
end

@functor LSTMCell

Base.show(io::IO, l::LSTMCell) =
  print(io, "LSTMCell(", size(l.Wi, 2), ", ", size(l.Wi, 1)÷4, ")")

"""
    LSTM(in::Integer, out::Integer)

[Long Short Term Memory](https://www.researchgate.net/publication/13853244_Long_Short-term_Memory)
recurrent layer. Behaves like an RNN but generally exhibits a longer memory span over sequences.

The parameters `in` and `out` describe the size of the feature vectors passed as input and as output. That is, it accepts a vector of length `in` or a batch of vectors represented as a `in x B` matrix and outputs a vector of length `out` or a batch of vectors of size `out x B`.

This constructor is syntactic sugar for `Recur(LSTMCell(a...))`, and so LSTMs are stateful. Note that the state shape can change depending on the inputs, and so it is good to `reset!` the model between inference calls if the batch size changes. See the examples below.

See [this article](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)
for a good overview of the internals.

# Examples
```jldoctest
julia> l = LSTM(3, 5)
Recur(
  LSTMCell(3, 5),                       # 190 parameters
)         # Total: 5 trainable arrays, 190 parameters,
          # plus 2 non-trainable, 10 parameters, summarysize 1.062 KiB.

julia> l(rand(Float32, 3)) |> size
(5,)

julia> Flux.reset!(l);

julia> l(rand(Float32, 3, 10)) |> size # batch size of 10
(5, 10)
```

!!! warning "Batch size changes"
    Failing to call `reset!` when the input batch size changes can lead to unexpected behavior. See the example in [`RNN`](@ref).
"""
LSTM(a...; ka...) = Recur(LSTMCell(a...; ka...))
Recur(m::LSTMCell) = Recur(m, m.state0)

# TODO remove in v0.13
function Base.getproperty(m::LSTMCell, sym::Symbol)
  if sym === :h
    Zygote.ignore() do
      @warn "LSTMCell field :h has been deprecated. Use m::LSTMCell.state0[1] instead."
    end
    return getfield(m, :state0)[1]
  elseif sym === :c
    Zygote.ignore() do
      @warn "LSTMCell field :c has been deprecated. Use m::LSTMCell.state0[2] instead."
    end
    return getfield(m, :state0)[2]
  else
    return getfield(m, sym)
  end
end

# GRU

function _gru_output(Wi, Wh, b, x, h)
  o = size(h, 1)
  gx, gh = Wi*x, Wh*h
  r = σ.(gate(gx, o, 1) .+ gate(gh, o, 1) .+ gate(b, o, 1))
  z = σ.(gate(gx, o, 2) .+ gate(gh, o, 2) .+ gate(b, o, 2))

  return gx, gh, r, z
end

struct GRUCell{A,V,S}
  Wi::A
  Wh::A
  b::V
  state0::S
end

GRUCell(in, out; init = glorot_uniform, initb = zeros32, init_state = zeros32) =
  GRUCell(init(out * 3, in), init(out * 3, out), initb(out * 3), init_state(out,1))

function (m::GRUCell{A,V,<:AbstractMatrix{T}})(h, x::Union{AbstractVecOrMat{T},OneHotArray}) where {A,V,T}
  b, o = m.b, size(h, 1)
  gx, gh, r, z = _gru_output(m.Wi, m.Wh, b, x, h)
  h̃ = tanh.(gate(gx, o, 3) .+ r .* gate(gh, o, 3) .+ gate(b, o, 3))
  h′ = (1 .- z) .* h̃ .+ z .* h
  sz = size(x)
  return h′, reshape(h′, :, sz[2:end]...)
end

@functor GRUCell

Base.show(io::IO, l::GRUCell) =
  print(io, "GRUCell(", size(l.Wi, 2), ", ", size(l.Wi, 1)÷3, ")")

"""
    GRU(in::Integer, out::Integer)

[Gated Recurrent Unit](https://arxiv.org/abs/1406.1078v1) layer. Behaves like an
RNN but generally exhibits a longer memory span over sequences. This implements
the variant proposed in v1 of the referenced paper.

The parameters `in` and `out` describe the size of the feature vectors passed as input and as output. That is, it accepts a vector of length `in` or a batch of vectors represented as a `in x B` matrix and outputs a vector of length `out` or a batch of vectors of size `out x B`.

This constructor is syntactic sugar for `Recur(GRUCell(a...))`, and so GRUs are stateful. Note that the state shape can change depending on the inputs, and so it is good to `reset!` the model between inference calls if the batch size changes. See the examples below.

See [this article](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)
for a good overview of the internals.

# Examples
```jldoctest
julia> g = GRU(3, 5)
Recur(
  GRUCell(3, 5),                        # 140 parameters
)         # Total: 4 trainable arrays, 140 parameters,
          # plus 1 non-trainable, 5 parameters, summarysize 792 bytes.

julia> g(rand(Float32, 3)) |> size
(5,)

julia> Flux.reset!(g);

julia> g(rand(Float32, 3, 10)) |> size # batch size of 10
(5, 10)
```

!!! warning "Batch size changes"
    Failing to call `reset!` when the input batch size changes can lead to unexpected behavior. See the example in [`RNN`](@ref).
"""
GRU(a...; ka...) = Recur(GRUCell(a...; ka...))
Recur(m::GRUCell) = Recur(m, m.state0)

# TODO remove in v0.13
function Base.getproperty(m::GRUCell, sym::Symbol)
  if sym === :h
    Zygote.ignore() do
      @warn "GRUCell field :h has been deprecated. Use m::GRUCell.state0 instead."
    end
    return getfield(m, :state0)
  else
    return getfield(m, sym)
  end
end

# GRU v3

struct GRUv3Cell{A,V,S}
  Wi::A
  Wh::A
  b::V
  Wh_h̃::A
  state0::S
end

GRUv3Cell(in, out; init = glorot_uniform, initb = zeros32, init_state = zeros32) =
  GRUv3Cell(init(out * 3, in), init(out * 2, out), initb(out * 3),
            init(out, out), init_state(out,1))

function (m::GRUv3Cell{A,V,<:AbstractMatrix{T}})(h, x::Union{AbstractVecOrMat{T},OneHotArray}) where {A,V,T}
  b, o = m.b, size(h, 1)
  gx, gh, r, z = _gru_output(m.Wi, m.Wh, b, x, h)
  h̃ = tanh.(gate(gx, o, 3) .+ (m.Wh_h̃ * (r .* h)) .+ gate(b, o, 3))
  h′ = (1 .- z) .* h̃ .+ z .* h
  sz = size(x)
  return h′, reshape(h′, :, sz[2:end]...)
end

@functor GRUv3Cell

Base.show(io::IO, l::GRUv3Cell) =
  print(io, "GRUv3Cell(", size(l.Wi, 2), ", ", size(l.Wi, 1)÷3, ")")

"""
    GRUv3(in::Integer, out::Integer)

[Gated Recurrent Unit](https://arxiv.org/abs/1406.1078v3) layer. Behaves like an
RNN but generally exhibits a longer memory span over sequences. This implements
the variant proposed in v3 of the referenced paper.

The parameters `in` and `out` describe the size of the feature vectors passed as input and as output. That is, it accepts a vector of length `in` or a batch of vectors represented as a `in x B` matrix and outputs a vector of length `out` or a batch of vectors of size `out x B`.

This constructor is syntactic sugar for `Recur(GRUv3Cell(a...))`, and so GRUv3s are stateful. Note that the state shape can change depending on the inputs, and so it is good to `reset!` the model between inference calls if the batch size changes. See the examples below.

See [this article](https://colah.github.io/posts/2015-08-Understanding-LSTMs/)
for a good overview of the internals.

# Examples
```jldoctest
julia> g = GRUv3(3, 5)
Recur(
  GRUv3Cell(3, 5),                      # 140 parameters
)         # Total: 5 trainable arrays, 140 parameters,
          # plus 1 non-trainable, 5 parameters, summarysize 848 bytes.

julia> g(rand(Float32, 3)) |> size
(5,)

julia> Flux.reset!(g);

julia> g(rand(Float32, 3, 10)) |> size # batch size of 10
(5, 10)
```

!!! warning "Batch size changes"
    Failing to call `reset!` when the input batch size changes can lead to unexpected behavior. See the example in [`RNN`](@ref).
"""
GRUv3(a...; ka...) = Recur(GRUv3Cell(a...; ka...))
Recur(m::GRUv3Cell) = Recur(m, m.state0)

@adjoint function Broadcast.broadcasted(f::Recur, args...)
  Zygote.∇map(__context__, f, args...)
end

"""
    Bidirectional{A,B}

A wrapper layer that allows the use of [bidirectional](https://ieeexplore.ieee.org/document/650093) layers. It contains two parts that are Flux layers: `forward` and `backward` where
the forward layer weights are concatenated with the reversed order of the backward layer weights.

It is intended to be used with recurrent layers such as `LSTM`, `GRU` or `RNN` to benefit from the sequential information that recurrent
layers have, but it will not raise an error if used with a different layer such as `Dense`, as long as the layer is compatible with the concatenation function `vcat`.

For flexibility, it is possible to use the contructor `Bidirectional(rnn, in::Int, out::Int, a...; ka...)` by passing the input and output dimensions
together with the desired recurrent layers (one of `LSTM`, `GRU`, or `RNN`). Check the examples below for more details.

# Examples

1. Using the flexible constructor to create a bidirectional LSTM layer (BiLSTM):
```jldoctest
julia> BLSTM = Bidirectional(LSTM, 3, 5)
Bidirectional(
  Recur(
    LSTMCell(3, 5),                     # 190 parameters
  ),
  Recur(
    LSTMCell(3, 5),                     # 190 parameters
  ),
)         # Total: 10 trainable arrays, 380 parameters,
          # plus 4 non-trainable, 20 parameters, summarysize 2.141 KiB.
```
2. Checking the dimension after running the bidirectional layer on an input vector. Shows that the
dimension of the output is twice the dimension of the input:
```jldoctest
julia> Bidirectional(LSTM, 3, 5)(rand(Float32, 3)) |> size
(10,)
```
3. It is possible to use the bidirectional layer inside the `Chain` container:
```
julia> model = Chain(Embedding(10000, 200), Bidirectional(LSTM, 200, 128), Dense(256, 5), softmax)
Chain(
  Embedding(10000, 200),                # 2_000_000 parameters
  Bidirectional(
    Recur(
      LSTMCell(200, 128),               # 168_704 parameters
    ),
    Recur(
      LSTMCell(200, 128),               # 168_704 parameters
    ),
  ),
  Dense(256, 5),                        # 1_285 parameters
  NNlib.softmax,
)         # Total: 13 trainable arrays, 2_338_693 parameters,
          # plus 4 non-trainable, 512 parameters, summarysize 8.922 MiB.
```
4. It is also possible to use the default constructor
```jldoctest
julia> BiLSTM = Bidirectional(GRU(3, 5), GRU(3, 5))
Bidirectional(
  Recur(
    GRUCell(3, 5),                      # 140 parameters
  ),
  Recur(
    GRUCell(3, 5),                      # 140 parameters
  ),
)         # Total: 8 trainable arrays, 280 parameters,
          # plus 2 non-trainable, 10 parameters, summarysize 1.562 KiB.
```
5. And use other parameters available on the recurrent layers
```jldoctest
julia> BiLSTM = Bidirectional(RNN(3, 5, tanh; init=glorot_normal), LSTM(3, 5; initb=zeros32, init_state=zeros32))
Bidirectional(
  Recur(
    RNNCell(3, 5, tanh),                # 50 parameters
  ),
  Recur(
    LSTMCell(3, 5),                     # 190 parameters
  ),
)         # Total: 9 trainable arrays, 240 parameters,
          # plus 3 non-trainable, 15 parameters, summarysize 1.500 KiB.
```
"""
struct Bidirectional{A,B}
  forward::A
  backward::B
end

# Constructor that creates a bidirectional with the same layer for forward and backward
# Needs to have `in` explicitly declared to avoid conflicts with the default construtor
Bidirectional(rnn, in::Integer, args...; ka...) = Bidirectional(rnn(in, args...; ka...), rnn(in, args...; ka...))

# Concatenate the forward and reversed backward weights
function (m::Bidirectional)(x::Union{AbstractVecOrMat{T},OneHotArray}) where {T}
  return vcat(m.forward(x), reverse(m.backward(reverse(x; dims=1)); dims=1))
end

@functor Bidirectional

function Base.show(io::IO, b::Bidirectional)
  print(io, "Bidirectional(")
  show(io, b.forward.cell)
  print(io, ", ")
  show(io, b.backward.cell)
  print(io, ")")
end;


# by lazarusA # HIDE
# using GLMakie # HIDE
using CairoMakie
CairoMakie.activate!() # HIDE
let
    x = -2:0.005:2
    y = -2:0.005:2
    f(z) = (z^3 - 3)/z
    fvals = [f(u+1im*v) for u in x, v in y]
    fvalues = abs.(fvals)
    fargs = angle.(fvals)
    fig = Figure(resolution = (960,450))
    ax1 = Axis(fig, aspect = 1)
    ax2 = Axis(fig, aspect = 1)
    cmap = :diverging_rainbow_bgymr_45_85_c67_n256
    pltobj1 = contour!(ax1, x, y, fargs, levels = 30, colormap = cmap)
    pltobj2 = heatmap!(ax2, x, y, fargs, colorrange = (-π,π), colormap = cmap)
        contour!(ax2, x, y, fargs,levels = 30,color = :white,linewidth = 0.85)
    cbar = Colorbar(fig, pltobj2,
            ticks = ([-π,-π/2,0,π/2,π],["-π", "-π/2", "0", "π/2", "π"]))
    limits!(ax1, -2,2,-2,2)
    fig[1,1] = ax1
    fig[1,2] = ax2
    fig[1,3] = cbar
    fig
    save(joinpath(@__DIR__, "output", "ContourComplexF.png"), fig, px_per_unit = 2.0) # HIDE
end
using Pkg # HIDE
Pkg.status(["CairoMakie"]) # HIDE


using Unitful, FiniteElementDiffEq

### Setup
dx = (1//2^(3))u"m"
mesh = notime_squaremesh([0u"m" 1u"m" 0u"m" 1u"m"],dx,:dirichlet) #Fails at meshgrid because unit ranges not supported

f = (x) -> sin(2π.*map((y)->y.val,x[:,1])).*cos(2π.*map((y)->y.val,x[:,2]))
prob = PoissonProblem(f,mesh,D=1u"m^2/s")

#=
sol = solve(mesh,prob)
=#

#Define a parabolic problem
T = 1u"s"
dt = (1//2^(7))u"s"
mesh = parabolic_squaremesh([0u"m" 1u"m" 0u"m" 1u"m"],dx,dt,T,:dirichlet)
f = (u,x,t)  -> (float(ones(size(x,1))))u"N" - .5u
u0_func = (x) -> map((x)->(x)u"N",zeros(size(x,1)))
u0 = u0_func(mesh.node)
prob = HeatProblem(u0,f,mesh)
#=
println("Euler")
sol = solve(prob,alg=:Euler)
=#

true


# ### m13.3t.jl

using Pkg, DrWatson

@quickactivate "StatisticalRethinkingTuring"
using Turing
using StatisticalRethinking
Turing.setprogress!(false)

df = CSV.read(sr_datadir("UCBadmit.csv"), DataFrame);

dept_map = Dict(key => idx for (idx, key) in enumerate(unique(df.dept)))
df.male = ifelse.(df.gender .== "male", 1, 0)
df.dept_id = [dept_map[de] for de in df.dept]

@model ppl13_3(applications, dept_id, male, admit) = begin
    Rho ~ LKJ(2, 2.)
    sigma_dept ~ filldist(truncated(Cauchy(0, 2), 0, Inf), 2)
    bm ~ Normal(0, 1)
    a ~ Normal(0, 1)

    dist_mu = [a, bm]
    dist_Sigma = sigma_dept .* Rho .* sigma_dept'
    dist_Sigma = (dist_Sigma' + dist_Sigma) / 2

    a_bm_dept ~ filldist(MvNormal(dist_mu, dist_Sigma), 6)

    a_dept, bm_dept = a_bm_dept[1, :], a_bm_dept[2, :]
    logit_p = a_dept[dept_id] + bm_dept[dept_id] .* male

    admit .~ BinomialLogit.(applications, logit_p)
end

m13_3t = ppl13_3(df.applications, df.dept_id, df.male, df.admit)
nchains = 4; sampler = NUTS(0.65); nsamples=2000
nchains = 1
chns13_3t = mapreduce(c -> sample(m13_3t, sampler, nsamples), chainscat, 1:nchains)

m_13_3_rethinking = """
Inference for Stan model: f0d86ec689cbf7921aab4fc0f55616d2.
    4 chains, each with iter=5000; warmup=1000; thin=1;
    post-warmup draws per chain=4000, total post-warmup draws=16000.

                      mean se_mean   sd     2.5%      25%      50%      75%
    bm_dept[1]       -0.79    0.00 0.27    -1.32    -0.97    -0.79    -0.61
    bm_dept[2]       -0.21    0.00 0.33    -0.87    -0.42    -0.20     0.00
    bm_dept[3]        0.08    0.00 0.14    -0.18    -0.01     0.08     0.18
    bm_dept[4]       -0.09    0.00 0.14    -0.37    -0.19    -0.09     0.00
    bm_dept[5]        0.12    0.00 0.19    -0.24    -0.01     0.12     0.25
    bm_dept[6]       -0.12    0.00 0.27    -0.67    -0.29    -0.12     0.05
    a_dept[1]         1.30    0.00 0.25     0.82     1.13     1.30     1.47
    a_dept[2]         0.74    0.00 0.33     0.08     0.52     0.73     0.95
    a_dept[3]        -0.65    0.00 0.09    -0.82    -0.71    -0.65    -0.59
    a_dept[4]        -0.62    0.00 0.10    -0.83    -0.69    -0.62    -0.55
    a_dept[5]        -1.13    0.00 0.11    -1.36    -1.21    -1.13    -1.05
    a_dept[6]        -2.60    0.00 0.20    -3.01    -2.73    -2.60    -2.46
    a                -0.49    0.01 0.73    -1.96    -0.91    -0.49    -0.07
    bm               -0.16    0.00 0.24    -0.65    -0.29    -0.16    -0.02
    sigma_dept[1]     1.67    0.01 0.63     0.88     1.25     1.53     1.94
    sigma_dept[2]     0.50    0.00 0.26     0.16     0.33     0.45     0.61
    Rho[1,1]          1.00     NaN 0.00     1.00     1.00     1.00     1.00
    Rho[1,2]         -0.31    0.00 0.35    -0.87    -0.59    -0.35    -0.08
    Rho[2,1]         -0.31    0.00 0.35    -0.87    -0.59    -0.35    -0.08
    Rho[2,2]          1.00    0.00 0.00     1.00     1.00     1.00     1.00
    lp__          -2593.92    0.04 3.17 -2601.03 -2595.85 -2593.57 -2591.65
                     97.5% n_eff Rhat
    bm_dept[1]       -0.27  6918    1
    bm_dept[2]        0.46  9133    1
    bm_dept[3]        0.36 13016    1
    bm_dept[4]        0.19 12258    1
    bm_dept[5]        0.49 12953    1
    bm_dept[6]        0.40 12011    1
    a_dept[1]         1.81  6993    1
    a_dept[2]         1.39  9335    1
    a_dept[3]        -0.48 13489    1
    a_dept[4]        -0.42 12595    1
    a_dept[5]        -0.91 14659    1
    a_dept[6]        -2.23 13168    1
    a                 0.96  9541    1
    bm                0.31  8946    1
    sigma_dept[1]     3.27  9058    1
    sigma_dept[2]     1.15  6626    1
    Rho[1,1]          1.00   NaN  NaN
    Rho[1,2]          0.44 11358    1
    Rho[2,1]          0.44 11358    1
    Rho[2,2]          1.00 16069    1
    lp__          -2588.70  5310    1
"""

# End of m13.3.jl


using SubmodularMaximization
using Distributions
using Statistics

table = SubmodularMaximization.NormalLookup(increment=0.0001, max=4.0)

maxerror = maximum(0.0:0.00001:5) do x
  trueval = pdf(Normal(0.0, 1.0), -x)
  lookupval = SubmodularMaximization.evaluate(table, error= -x, stddev=1.0)

  abs(trueval - lookupval)
end
println("Max error: ", maxerror)

let sigma = 3.0, error = 5.0
  @show trueval = pdf(Normal(0.0, sigma), error)
  @show lookupval = SubmodularMaximization.evaluate(table, error=error, stddev=sigma)
  @show abs(trueval - lookupval)
end


using NLsolve: OnceDifferentiable, NewtonCache, newton_
using LineSearches

struct NLsolveNewton{T, FT, DT, CT, ST, LT} <: AbstractNewtonSolver{T}
    x::Vector{T}
    f::Vector{T}
    J::Matrix{T}

    cache::CT

    F!::FT
    DF::DT

    line_search::ST
    linear_solver::LT

    params::NonlinearSolverParameters{T}
    status::NonlinearSolverStatus{T}

    function NLsolveNewton(x::AbstractVector{T}, f::AbstractVector{T}, J::AbstractMatrix{T},
                    F!::FT, DF::DT, cache::CT, line_search::ST, linear_solver::LT) where {T,FT,DT,CT,ST,LT}

        nls_params = NonlinearSolverParameters(T)
        nls_status = NonlinearSolverStatus{T}(length(x))

        new{T,FT,DT,CT,ST,LT}(x, f, J, cache, F!, DF, line_search, linear_solver, nls_params, nls_status)
    end
end

function NLsolveNewton(x::AbstractVector{T}, F!::Function; J!::Union{Function,Nothing}=nothing) where {T}
    linear_solver = getLinearSolver(x)

    f  = linear_solver.b
    df = linear_solver.A
    f! = (f, x) -> F!(x, f)

    if J! == nothing
        if get_config(:jacobian_autodiff)
            DF = OnceDifferentiable(f!, x, f, df; autodiff=:forward, inplace=true)
        else
            DF = OnceDifferentiable(f!, x, f, df; autodiff=:finite, inplace=true)
        end
    else
        j! = (df, x) -> J!(x, df)
        DF = OnceDifferentiable(f!, j!, x, f, df; inplace=true)
    end

    NLsolveNewton(x, f, df, F!, DF, NewtonCache(DF), LineSearches.Static(), linear_solver)
end

function linsolve!(s::NLsolveNewton, x, A, b)
    copyto!(s.linear_solver.A, A)
    copyto!(s.linear_solver.b, b)
    factorize!(s.linear_solver)
    solve!(s.linear_solver)
    copyto!(x, s.linear_solver.b)
end

function solve!(s::NLsolveNewton; n::Int=0)
    local nmax::Int = n > 0 ? nmax = n : s.params.nmax

    res=newton_(s.DF, s.x, s.params.stol, s.params.atol, nmax, false, false, false,
                s.line_search, (x, A, b) -> linsolve!(s, x, A, b), s.cache)

    s.x .= res.zero
    s.status.i  = res.iterations
    s.status.rₐ = res.residual_norm

    nothing
end


# perform point-to-point ICP with provided correspondence and distance score functions

squared_deviation(P, Q, metric) = sum(abs2, colwise(metric, P, Q))
squared_deviation(P, Q, metric, matches) = squared_deviation(matched_points(P, Q, matches)..., metric)

rmsd(P, Q, metric, n::Int) = sqrt(squared_deviation(P,Q,metric)/n)
rmsd(P::AbstractMatrix, Q::AbstractMatrix, metric) = rmsd(P, Q, metric, size(P,2))
rmsd(P, Q, metric) = rmsd(P, Q, metric, length(P))

rmsd(P, Q, metric, matches) = rmsd(matched_points(P, Q, matches)..., metric)

function iterate_kabsch(P, Q; iterations=1000, correspondence=closest_points, kwargs...)
    # initial correspondences
    matches = correspondence(P, Q; kwargs...)
    tform = identity

    # iterate until convergence
    it = 0
    while it < iterations
        it += 1
        matchedP, matchedQ = matched_points(P, Q, matches)
        tform = kabsch(matchedP, matchedQ)

        prevmatches = matches
        matches = correspondence(tform(P), Q; kwargs...)
        if matches == prevmatches
            break
        end
    end
    return matches
end


function numdiff(f, x0; ϵ=1e-6)
    f0 = f(x0)
    p = size(f0,1)
    q = size(x0,1)
    J = zeros(p,q)
    for j=1:q
        ei = zeros(q)
        ei[j] = 1.0
        J[:,j] = (f(x0+ϵ*ei) - f0)/ϵ
    end
    return J
end

function get_ss_derivatives(model; method=:symbolic)
    m, s, x, p = model.calibration[:exogenous, :states, :controls, :parameters]
    if method==:symbolic
        # it is still incorrect
        g_diff = transition(model,Val{(1,2,3,4)}, m, s, x, m, p)
        f_diff = arbitrage(model, Val{(1,2,3,4,5,6)}, m, s, x, m, s, x, p)
    else
        g_diff = [numdiff(f,u) for (f,u) in [
            (u->transition(model, u, s, x, m, p),m),
            (u->transition(model, m, u, x, m, p),s),
            (u->transition(model, m, s, u, m, p),x),
            (u->transition(model, m, s, x, u, p),m)
        ]]
        f_diff = [numdiff(f,u) for (f,u) in [
            (u->arbitrage(model, u, s, x, m, s, x, p), m),
            (u->arbitrage(model, m, u, x, m, s, x, p), s),
            (u->arbitrage(model, m, s, u, m, s, x, p), x),
            (u->arbitrage(model, m, s, x, u, s, x, p), m),
            (u->arbitrage(model, m, s, x, m, u, x, p), s),
            (u->arbitrage(model, m, s, x, m, s, u, p), x),
        ]]
    end
    g_diff, f_diff
end

perturb(p::IIDExogenous) = (zeros(0), zeros(0, 0))
perturb(p::VAR1) = (p.mu, p.R)

mutable struct PerturbationResult
    dr::BiTaylorExpansion
    generalized_eigenvalues::Union{Vector, Nothing}
    stable::Union{Bool, Nothing}     # biggest e.v. lam of solution is < 1
    determined::Union{Bool, Nothing} # next eigenvalue is > lam + epsilon (MOD solution well defined)
    unique::Union{Bool, Nothing}     # next eigenvalue is > 1
end

function Base.show(io::IO, pbr::PerturbationResult)
    @printf io "Perturbation Results\n"
    @printf io " * Decision Rule type: %s\n" string(typeof(pbr.dr))
    # @printf io " * Blanchard-Kahn: %s\n" blanchard_kahn(pbr)
    if !(typeof(pbr.stable)<:Nothing)
        @printf io " * stable < %s\n" pbr.stable
    end
    if !(typeof(pbr.determined)<:Nothing)
        @printf io " * determined < %s\n" pbr.determined
    end
end

blanchard_kahn(fos::PerturbationResult) = fos.stable && fos.unique

function perturb_first_order(g_s, g_x, f_s, f_x, f_S, f_X)

    eigtol = 1.0+1e-6

    ns = size(g_s, 2)
    nx = size(g_x, 2)

    nv = ns + nx

    A = [Matrix(1.0I,ns,ns) zeros(ns, nx);
         -f_S     -f_X]

    B = [g_s g_x;
         f_s f_x]

    # do orderd QZ decomposition
    gs = schur(A, B)

    genvals = (abs.(gs.α) ./ abs.(gs.β))
    sort!(genvals, rev=true)
    n_keep = ns # number of eigenvalues to keep
    diff = genvals[n_keep+1] - genvals[n_keep]
    eigtol = genvals[n_keep] + diff/2

    select = (abs.(gs.α) .> eigtol*abs.(gs.β))

    ordschur!(gs, select)
    S, T, Q, Z = gs.S, gs.T, gs.Q, gs.Z
    diag_S = diag(S)
    diag_T = diag(T)
    eigval = abs.(diag_S./diag_T)

    # now obtain first order solution
    Z11 = Z[1:ns, 1:ns]
    Z12 = Z[1:ns, ns+1:end]
    Z21 = Z[ns+1:end, 1:ns]
    Z22 = Z[ns+1:end, ns+1:end]
    S11 = S[1:ns, 1:ns]
    T11 = T[1:ns, 1:ns]

    # first order solution
    C = (Z11'\Z21')'
    return C, genvals
end

function get_gf_derivatives(model::AbstractModel)

    g_diff, f_diff = get_ss_derivatives(model)
    _f_m, _f_s, _f_x, _f_M, _f_S, _f_X = f_diff
    _g_m, _g_s, _g_x, _g_M = g_diff

    (M, R) = perturb(model.exogenous)

    f_x = _f_x
    f_X = _f_X
    _m, _s, x, p = model.calibration[:exogenous, :states, :controls, :parameters]

    if size(R, 1)>0
        f_s = [_f_m _f_s]
        f_S = [_f_M _f_S]
        g_s = [R zeros(size(R, 1), size(_g_s, 2)); _g_m _g_s]
        g_x = [zeros(size(_g_m, 1), size(_g_x, 2)); _g_x]
        s = cat(_m, _s; dims=1)
    else
        f_s = _f_s
        f_S = _f_S
        g_s = _g_s
        g_x = _g_x
        s = _s
    end
    return g_s, g_x, f_s, f_x, f_S, f_X
end

"""
TBD
"""
function perturb(model::Model; method=:qz)

    g_s, g_x, f_s, f_x, f_S, f_X = get_gf_derivatives(model)
    nx = size(g_x, 2)

    (M, R) = perturb(model.exogenous)
    _m, _s, x, p = model.calibration[:exogenous, :states, :controls, :parameters]

    if size(R, 1)>0
        s = cat(_m, _s; dims=1)

    else
        s = _s
    end

    n_s = size(g_s, 1)
    if method==:qz
        tol = 1e-6 # minimum distance betweel lam_n and lam_{n+1}
        C, genvals = perturb_first_order(g_s, g_x, f_s, f_x, f_S, f_X)
        sort!(genvals)
        stable = (genvals[n_s]<1)
        deter = (genvals[n_s+1]-genvals[n_s]>tol)
        uni = (genvals[n_s+1]>1)
    elseif method==:linear_time_iteration
        C, nit = linear_time_iteration(g_s, g_x, f_s, f_x, f_S, f_X)
        genvals = nothing
        stable = nothing
        deter = nit>0
        uni = nothing
    end

    if size(R, 1)>0
        C_exo = C[:, 1:length(_m)]
        C_endo = C[:, (length(_m)+1):end]
        dr = BiTaylorExpansion{nx}(_m, _s, x, C_exo, C_endo)
    else
        dr = BiTaylorExpansion{nx}(_m, _s, x, zeros(nx, length(_m)), C)
    end

    PerturbationResult(
        dr,
        genvals,
        stable,
        deter,
        uni
    )

end

function linear_time_iteration(model; maxit=1000, tol_ϵ=1e-8, improve=20, verbose=false)
    g_s, g_x, f_s, f_x, f_S, f_X = Dolo.get_gf_derivatives(model)
    return linear_time_iteration(g_s, g_x, f_s, f_x, f_S, f_X; maxit=maxit)
end

function L(P, Q, X, Y, u)
    return P*u*Q
end

function linear_time_iteration(g_s, g_x, f_s, f_x, f_S, f_X; maxit=1000, tol_ϵ=1e-8, improve=20, verbose=false)

    n_x, n_s = size(f_s)

    K = f_s + f_S*g_s

    X_0 = 1 .+rand(n_x, n_s)
    local X
    for it=1:maxit
        Y = X_0
        B = -K - f_X*Y*g_s
        A = f_x + f_S*g_x + f_X*Y*g_x
        X = A\B
        ϵ = maximum(abs.(X-X_0))
        if ϵ<tol_ϵ
            return (X,it)
        end
        if verbose
            println(it, "\t: ", ϵ)
        end
        if (improve>0)
            S = (X-X_0)
            δ = (X-X_0)
            P = -(f_x + f_S*g_x + f_X*Y*g_x)\f_X
            Q = g_s + g_x*X
            for k=1:improve
                S = δ + L(P, Q, X, Y, S)
            end
            X_0 = X_0 + S
        else
            X_0 = X
        end
    end

    return (X,-1)

end

function ttest(g_s, g_x, f_s, f_x, f_S, f_X, X)
    res = f_s + f_x*X + f_S*(g_s+g_x*X) + f_X*X*(g_s+g_x*X)
    return res
end


using MeasureBase
using Distributions

@testset "distributions" begin
    d = Rice(ν=2.0, σ=0.5)
    logdensity(d, 1.0)

    samples = [rand(d) for _ in 1:100_000]
    ms = mean(samples)
    vs = var(samples)
    @test isapprox(ms, mean(d); atol=1e-2)
    @test isapprox(vs, var(d); atol=1e-2)

    dhsnr = Rice(ν=5.0, σ=1e-3)
    μ = sqrt(dhsnr.ν^2 - dhsnr.σ^2)
    @test isapprox(logdensity(dhsnr, 4.9), logpdf(Normal(μ, dhsnr.σ), 4.9); atol=1e-1)

    dvm1 = CPVonMises(μ=0.0, σ=0.1)
    dvm2 = CPVonMises(μ=0.0, κ=1/0.1^2)

    # Take off the normalization piece
    ds = VonMises(0.0, 1/0.1^2)
    norm = -log(ds.I0κx) - log(2π)
    @test logpdf(ds, 0.5) - norm ≈ logdensity(dvm2, 0.5)

    @test logdensity(dvm1, 0.5) ≈ logdensity(dvm2, 0.5)
end


using FactorGraph
using OMEinsum: bpcheck, einsum
using Zygote
using Test

Base.conj(arr::AbstractArray{<:Any, 0}) = conj!(copy(arr))

@testset "einsum bp" begin
    a = randn(ComplexF64, 3,3)
    f2(a) = ein"ij,ik,il->jkl"(a, conj(a), a) |> sum |> real

    @test gradient(f2, a)[1] !== nothing
    @test bpcheck(f2, a)
    @test bpcheck(a->treecontract(((1,2),3), ((1,2), (2,3), (3,1)), (a, a, a), ())[] |> real, a)
end

@testset "optcontract" begin
    a = randn(ComplexF64, 3,3)
    @test optcontract(((1,2), (2,3)), (a, a), (1,3)) ≈ a*a
    @test bpcheck(a->ein"ij,ji->"(a, a) |> sum |> real, a)
    @test bpcheck(a->optcontract(((1,2), (2,3)), (a, a), (1,3)) |> sum |> real, a)
    @test bpcheck(a->optcontract(((1,2), (2,3), (3,4)), (a,a,a), (4,1)) |> sum |> real, a)
end


module Decimals
import Base: +, -, *, /, div, inv,
             ==, <, <=, >, >=,
             signbit, sign, decompose,
             one, zero,
             eps, parse, print

export Decimal, @d_str, strip_trailing_zeros,
       RoundUnnecessary, RoundFloor, RoundCeiling

struct Decimal <: AbstractFloat
    c::BigInt  # coefficient
    q::Integer # exponent
end

PRECISION, ROUNDING = 9, RoundDown

function setprecision!(; precision::Integer=16, rounding::RoundingMode=RoundDown)
    global PRECISION = precision
    global ROUNDING = rounding
    precision, rounding
end

include("decimal.jl")
include("arithmetic.jl")

function decompose(x::Decimal)
    x = strip_trailing_zeros(x)
    x.q >= 0 ? decompose(x.c * 10^x.q) : decompose(x.c // 10^-x.q)
end

end


using PyPlot
using Peacock
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using CUDA

get_k = Peacock.get_k
DiagonalMatrix = Peacock.DiagonalMatrix

# Permittivity
function epf(x,y)
    # equation of a circle with radius 0.2a
    if x^2+y^2 <= 0.2^2
        # dielectric inside the circle
        return 8.9
    else
        # air outside the circle
        return 1
    end
end

# Permeability is unity everywhere
function muf(x,y)
    return 1
end

a1 = [1, 0]  # first lattice vector
a2 = [0, 1]  # second lattice vector
d1 = 0.01  # resolution along first lattice vector
d2 = 0.01  # resolution along second lattice vector
geometry = Geometry(epf, muf, a1, a2, d1, d2)
fourier_space_cutoff = 17
solver = Solver(geometry, fourier_space_cutoff)
G = BrillouinZoneCoordinate(  0,   0, "Γ")
X = BrillouinZoneCoordinate(1/2,   0, "X")
M = BrillouinZoneCoordinate(1/2, 1/2, "M")
ks = [G,X,M,G]
ks = [typeof(x)==BrillouinZoneCoordinate ? get_k(x,solver.basis) : x for x in ks]

polarisation = TE

k = ks[1]
bands = 1:2

basis = solver.basis
epc, muc = solver.epc, solver.muc
Kx = DiagonalMatrix(basis.kxs) + k[1]*I
Ky = DiagonalMatrix(basis.kys) + k[2]*I
Kx_d = CuArray(Kx); Ky_d = CuArray(Ky); epc_d = CuArray(epc); muc_d = CuArray(muc);
if polarisation == TE
    LHS = Kx/epc*Kx + Ky/epc*Ky
    RHS = muc
    LHS_d = CuArray(LHS)
    RHS_d = CuArray(RHS)
    LHS_dd = Kx_d/epc_d*Kx_d + Ky_d/epc_d*Ky_d
    RHS_dd = muc_d
    label = L"H_z"
elseif polarisation == TM
    LHS = Kx/muc*Kx + Ky/muc*Ky
    RHS = epc
    label = L"E_z"
end
# Sometimes the generalised eigenvalue problem solver
# fails near Γ when the crystals are symmetric.
# In these cases, rewrite as a regular eigenvalue problem
@time freqs_squared, modes_data = eigen(RHS \ LHS);
@time freqs_squared, modes_data = eigen(Array(RHS_d \ LHS_d));


#=------------------------------------------------------------------------------
   Numerical Routines that are useful for tensors but are not a part of the
   core functionality. Though one could create a seperate package implementing
   these, it would be more effective to store common functions in this package
   that can be optimized for the COOTensor package specifically.

  Operations Glossary
  -------------------
    * Dynamical
------------------------------------------------------------------------------=#

"""-----------------------------------------------------------------------------
    Dynamical_System_Solver(A,x0,h,tol,m=1,update=0)

  This function takes in a sparse symmetric tensor and computes the tensor
eigenvector by solving the dynamical system formed by contracting the tensor
A into a vector, and then computing the mth largest eigenvector of the matrix
Ax^{k-1}. With this we form the dynamical system
					dxdt = :Lambda_m(Ax^{k+1})
which we solve with a forward Euler scheme with a step size of h, starting at
the point x0. We solve this until the norm of dxdt reaches a specified
tolerance.

Input:
------
* A - (SSSTensor or Array{AbstractFloat,k > 2})

    The symmetric tensor to compute the eigenvector of. Dense tensors have
    limited support, primary support is for constructors to the sparse cases.

* x0 - (Array{AbstractFloat,1})

    The initial starting point for solving the dynamical system.
* h - (Float64)

    The step size for running the forward Euler scheme.
* tol - (Float64)

    The tolerance to solve the dynamical system up to, stops when
    norm(dxdt) < tol.
* m - (Int)

    The eigenvector to find when computing the dynamical system. Default is the
    largest eigenvector of the matrix.
* update - (Int)

    Indicates how many steps the program should update the user by, default
    value is 0.

Output:
-------
* x - (Array{AbstractFloat,1}

    The resulting eigenvector computed by the dynamical system method.
* lambda - (AbstractFloat)

    The resulting eigenvalue computed by the dynamical system method.
-----------------------------------------------------------------------------"""
function Dynamical_System_Solver(A::SSSTensor,x0::Array{N,1},h::Float64,
                                 tol::Float64,m::Int64 = 1, update::Int = 0) where
								 N <: AbstractFloat
                                 #start with strictly positive random vector
                                 #check the monotonic dereasing property
  n = A.cubical_dimension
  @assert m <= n
  @assert length(x0) == n
  k = order(A)
  x = copy(x0)/norm(x0)
  step = 1

  while true
    A_x_k_2 = contract(A,x,k-2)
    _,V,_ = eigs(A_x_k_2,nev=m) # check eigs
    dxdt = sign(V[1,m])*V[:,m] - x
	#x /= norm(x)

	if norm(dxdt) <= tol
      return x, x'*A_x_k_2*x , step
    else
      x += h*dxdt
    end

	if update > 0
	  if step % update == 0
	    z = A_x_k_2*x
		lambda = dot(x,z)
		residual = z - lambda*x
	    @printf("step %5d: norm(dxdt) = %.16f | lambda = %.16f | res norm = %.16f \n",
		       step, norm(dxdt),lambda, norm(residual))
	  end
	end
    step += 1
  end
end

#figure out if it's faster to do an internal check than two overloaded functions

function Dynamical_System_Solver(A::Array{N,k},x0::Array{N,1},h::Float64,
                                 tol::Float64,m::Int64 = 1) where {N <: AbstractFloat,k}
  #k = length(size(A))
  n = size(A)[1]
  @assert m <= n
  @assert length(x0) == n
  x = copy(x0)

  while true
    _,V = eigen(contract(A,x,k-2))
    dxdt = sign(real(V[1,m]))*real(V[:,m]) - x
	x /= norm(x)

	if norm(dxdt) <= tol
      return x
    else
      x += h*dxdt
    end
  end
end

"""-----------------------------------------------------------------------------
    SSHOPM(A,x_0,shift,max_iter,tol)

  This function runs the shifts symmetric higher order power method for a super
symmetric tensor with the passed in shift, up to a tolerance or up until a
maximum iteration.

Input:
------
* A - (SSSTensor):

    An instance of the super symmetric tensor class.
* x_0 - (Array{AbstractFloat,1}):

    An initial vector to start the algorithm with.
* shift - (AbstractFloat)

    The shift for the algorithm, can be predetermined to ensure convergence of
    the method.
* max_iter - (Int)

    The maximum number of iterations to run the routine for, prints a warning if
    the method hasn't converged by then.
* tol - (Float)

    The tolerance in difference between subsequent approximate eigenvalues to
    solve the routine up to.

* display - (optional Int)

    Indicates how many iterations in between printing to the user.

Output:
-------
* z - (Array{AbstractFloat,1})

    The final vector produced by the SSHOPM routine.
* lambda_k - (AbstractFloat)

    The final approximate eigenvalue at the last iteration.
* iterations - (Integer)

    The number of iterations the algorithm ran for.
-----------------------------------------------------------------------------"""
function SSHOPM(A::Ten, x_0::Array{N,1},shift::N,max_iter::Int,tol::N,
                display::Int=0) where
				{N <: AbstractFloat, Ten <: AbstractSSTen}
    @assert A.cubical_dimension == length(x_0)

	z = Array{N,1}(undef,A.cubical_dimension)
    x = x_0/norm(x_0)
    iterations = 0
    lambda_k_1 = Inf

    while true

		contract_k_1!(A,x,z)
        if shift != 0
			z .+= shift*x
			if shift < 0
				z .*= -1
			end
        end

        lambda_k = x'*z

        #normalize
		residual = z - lambda_k*x
        z ./= norm(z)

        iterations += 1

		if display != 0
			if iterations % display == 0
				@printf("step = %-3d -- λ_k:% 0.12f -- |λ_k - λ_{k-1}| :%0.12f -- res_norm:%0.12f\n",
						iterations,lambda_k,abs(lambda_k - lambda_k_1),norm(residual))
			end
		end
        if norm(residual) < tol || iterations >= max_iter
            if iterations >= max_iter
                @warn("maximum iterations reached")
            end
            return z, lambda_k, iterations
        else
            lambda_k_1 = lambda_k
            x = copy(z)
        end
    end
end

"""-----------------------------------------------------------------------------
    SSHOPM(indices,vals,n,x_0,shift,max_iter,tol)

  This function runs a low memory usage shifted symmetric higher order power
method for a super symmetric tensor with the passed in shift, up to a tolerance
or up until a maximum iteration.

TODO: update to new SSSTensor class
Input:
------
* indices - (Array{Int,2}):

    Indices associated with
* x_0 - (Array{AbstractFloat,1}):

    An initial vector to start the algorithm with.
* shift - (AbstractFloat)

    The shift for the algorithm, can be predetermined to ensure convergence of
    the method.
* max_iter - (Int)

    The maximum number of iterations to run the routine for, prints a warning if
    the method hasn't converged by then.
* tol - (Float)

    The tolerance in difference between subsequent approximate eigenvalues to
    solve the routine up to.

Output:
-------
* z - (Array{AbstractFloat,1})

    The final vector produced by the SSHOPM routine.
* lambda_k - (AbstractFloat)

    The final approximate eigenvalue at the last iteration.
* iterations - (Integer)

    The number of iterations the algorithm ran for.
-----------------------------------------------------------------------------"""
function SSHOPM(indices::Array{Int,2},vals::Array{N,1},n::Int, x_0::Array{N,1},
                shift::N,max_iter,tol) where N <: AbstractFloat

    x = x_0/norm(x_0)
    iterations = 0
	z = Array{N,1}(undef,n)
    lambda_k_1 = Inf

    while true

		contract_k_1!(indices,vals,x,z)
        if shift != 0

            z += shift*x

			if shift < 0
				x *= -1
			end
        end

        lambda_k = x'*z

		residual = z - lambda_k*x

        #normalize
        z /= norm(z)
        iterations += 1

        if norm(residual) < tol || iterations >= max_iter
            if iterations >= max_iter
                @warn("maximum iterations reached")
            end
            return z, lambda_k, iterations
        else
            lambda_k_1 = lambda_k
            x = z
        end
    end
end

"""-----------------------------------------------------------------------------
    find_shift_for_convergence(A,use_fro)

  This function takes in a tensor and computes a shift bound to ensure
convergence of the SSHOPM. If the contracted tensor is large, this can be
computed with the frobenius norm.

Input:
------
* A - (SSSTensor):

    An instance of a super symmetric tensor class.
* use_fro - (bool):

    Indicates whether or not to use the frobenius norm.

Output:
-------
* shift_bound - (Float)

    A float indicating the lower bound for which the method is guaranteed to
    converge.
-----------------------------------------------------------------------------"""
function find_shift_for_convergence(A::SSSTensor)
    error("unfinished")
    shift_bound = (order(A) -1)
    if use_fro
        shift_bound *= contract(A)
    end

end

"""-----------------------------------------------------------------------------
    HOSVD(A,k)

  Returns the top k singular vectors and values of the Higher Order SVD.

 Notes:
 ------
 Currently outputs the left singular vectors and values from the flattened
 tensor. More functionality can be added in, but the current use case is for the
 starting vectors for the SSHOPM.

TODO: would be better to create a adjoint matvec operation from some variant of
  the contraction

-----------------------------------------------------------------------------"""
function HOSVD(A::SSSTensor,k::Int)

	(U,S,_),_ = svds(flatten(A),nsv=k)
	return U,S

end


using Plots
plotlyjs();

include("../julia/sskernel.jl");
include("../julia/utils.jl");

x_true = readcsv("../data/spikes.csv")[:];
rs_true = readcsv("../data/firing_rate_true.csv")[:];

ws = linspace(0.01,2.5,100);
rs_kernel = sskernel(x_true, ws=ws);

plot(ts, rs_true, label="True")
plot!(ts, rs_kernel, label="Estimated")

xlabel!("Time (s)")
ylabel!("Firing rate (Hz)")

# C = get_costfunc(x_inhomogeneous)

# plot(ws, C.(ws));
# xlabel!("w");
# ylabel!("C(w)");


# Solve 2D Fisher-KPP equation with Dirichlet boundary conditions
# Alex Tam, 03/11/2021

"Solve Fisher-KPP equation"
function fkpp(D, U, ϕ, par, dx, dy, dt, i)
    # Manually assign value at grid points close to interface
    for gp in D
        if (gp.Ω == true) && (gp.dΩ == true)
            U[gp.xInd, gp.yInd] = par.u_f
        end
    end
    # Create vector from matrix data
    u = build_vector(U, D)
    # Construct and solve ODE problem using DifferentialEquations.jl
    prob = ODEProblem((du, u, p, t) -> fkpp_rhs!(du, u, p, t, D, ϕ, par, dx, dy), u, ((i-1)*dt, i*dt))
    sol = solve(prob, Tsit5(), reltol = 1e-3, abstol = 1e-6, saveat = i*dt)
    # Reshape solution to matrix
    U = build_u_matrix(sol[:,end], par, D)
    return U
end

"Construct right-hand vector for use in DifferentialEquations.jl"
function fkpp_rhs!(du, u, p, t, D, ϕ, par, dx, dy)
    U = build_u_matrix(u, par, D) # Generate matrix
    for i = 1:length(D) # Loop over interior grid points
        if (D[i].Ω == true) && (D[i].dΩ == false) # Grid points inside Ω, away from dΩ
            # Obtain density at stencil points
            u_mid = U[D[i].xInd, D[i].yInd]
            if ϕ[D[i].xInd-1, D[i].yInd] >= 0
                u_left = par.u_f
            else
                u_left = U[D[i].xInd-1, D[i].yInd]
            end
            if ϕ[D[i].xInd+1, D[i].yInd] >= 0
                u_right = par.u_f
            else
                u_right = U[D[i].xInd+1, D[i].yInd]
            end
            if ϕ[D[i].xInd, D[i].yInd-1] >= 0
                u_bottom = par.u_f
            else
                u_bottom = U[D[i].xInd, D[i].yInd-1]
            end
            if ϕ[D[i].xInd, D[i].yInd+1] >= 0
                u_top = par.u_f
            else
                u_top = U[D[i].xInd, D[i].yInd+1]
            end
            # Compute Laplacian and source term
            uxx = 2.0*( u_left/(D[i].θxm*(D[i].θxm+D[i].θxp)) - u_mid/(D[i].θxm*D[i].θxp) + u_right/(D[i].θxp*(D[i].θxm+D[i].θxp)) )/(dx^2)
            uyy = 2.0*( u_bottom/(D[i].θym*(D[i].θym+D[i].θyp)) - u_mid/(D[i].θym*D[i].θyp) + u_top/(D[i].θyp*(D[i].θym+D[i].θyp)) )/(dy^2)
            du[i] = par.D*(uxx + uyy) + par.λ*u_mid*(1-u_mid)
        else # Grid points outside Ω or close to dΩ
            du[i] = 0.0
        end
    end
end


push!(LOAD_PATH, "./models")
push!(LOAD_PATH, ".")
module Vessel_network
# Module that implements methods to facilitate the generation and
# manipulation of blood vessel structures.

using Interpolations
using Roots
using QuadGK

export Tree, addchild, Absolute_Vessel, displacement, evaluation, Uniform_sampler

type TreeNode
    parent::Int
    children::Vector{Int}
    curveData::Array{Any,1}
end
type Tree
    nodes::Vector{TreeNode}
end
Tree(data) = Tree([TreeNode(0, Vector{Int}(), data)])
function addchild(tree::Tree, id::Int, curveData::Array{Any,1})
    1 <= id <= length(tree.nodes) || throw(BoundsError(tree, id))
    push!(tree.nodes, TreeNode(id, Vector{}(), curveData))
    child = length(tree.nodes)
    push!(tree.nodes[id].children, child)
    child
end
children(tree, id) = tree.nodes[id].children
parent(tree,id) = tree.nodes[id].parent
data(tree,id) = tree.nodes[id].curveData

function angle_direction(direction)
    return pi/2 - atan(direction[2],direction[1])
end

function segment_generator(Vessel, id)
    parameters = data(Vessel,id)
    x0 = parameters[1]
    tf = parameters[2]
    angle = parameters[3]
    speed = parameters[4]
    curve_type = parameters[5]
    params = parameters[6]
    if curve_type == "straight"
        f  = t -> [0 ; t]
        df = t -> [0 ; 1]
    elseif curve_type == "ellipse"
        if length(params)!= 3
            error("Incorrect number of input parameters")
        end
        side = params[1]
        x_semiax = params[2]
        y_semiax = params[3]
        g(t) = [x_semiax*(1-cos(t))*side; y_semiax*sin(t)]
        dg(t) = [x_semiax*side*sin(t); y_semiax*cos(t) ]
        # now we need to obtain the natural parametrization
            norm_derivative(t) = sqrt(dg(t)[1]^2 + dg(t)[2]^2)
            arclength(t) = quadgk(norm_derivative, 0, t)
            # invert the arc length
            max_sample = 1
            while arclength(max_sample)[1]-tf < 0
                max_sample = max_sample*2
            end
            t_sample = [t for t = linspace(0,max_sample,1000)]
            arclength_sample = [arclength(t)[1] for t in t_sample]
            inv_interp = LinearInterpolation(arclength_sample, t_sample)
        f = t -> g(inv_interp(t))
        df = t-> dg(inv_interp(t))/sqrt( dg(inv_interp(t))[1]^2 + dg(inv_interp(t))[2]^2 )
    else
        error("Not a valid curve type, find the correct names")
    end
    # Rotation matrix
    rotation_mat = [cos(angle) -sin(angle); sin(angle) cos(angle)]
    # Translation, speed scalation, rotation
    segment(t) = rotation_mat*f(t*speed) + x0
    dsegment(t) = rotation_mat*df(t*speed)
    return [segment , dsegment]
end

function Absolute_Vessel(Vessel,id, x0, angle)
    datta = data(Vessel,id)
    datta[2] = datta[2]-angle
    prepend!(datta,[x0])
    Vessel.nodes[id].curveData = datta
    segt, dsegt = segment_generator(Vessel, id)
    new_angle = angle_direction(dsegt(datta[2]))
    new_position = segt(datta[2])
    for j in children(Vessel,id)
        Absolute_Vessel(Vessel,j, new_position, new_angle)
    end
end

# Function to move along the vessels, from a starting position
# (id, t) and a displacement dt
function displacement(Vessel, id, t, dt)
    # This method is not robust to big displacements !
    tf = data(Vessel,id)[2]
    if 0 <= t+dt && t+dt <= tf
        return (id, t+dt)
    elseif t+dt > tf
        next_vessel = children(Vessel,id)
        if isempty(next_vessel)
            return NaN
        else
        # Choose some children at random, assuming at most 2
            new_id = next_vessel[rand(1:2)]
            return (new_id, t+dt-tf)
        end
    else
        # backward case, t+dt < 0
        new_id = parent(Vessel,id)
        if new_id == 0
            return NaN
        else
            new_tf = data(Vessel,new_id)[2]
            return (new_id, new_tf+t+dt)
        end
    end
end

#function to convert the position (id,t) to a point in R2.
function evaluation(Vessel, id, t)
    segt, _ = segment_generator(Vessel, id)
    return segt(t)
end

# Uniform samples on the network
function Uniform_sampler(Vessel)
    times = []
    for id =1:length(Vessel.nodes)
        append!(times, data(Vessel,id)[2])
    end
    tot_time = sum(times)
    selector = rand()*tot_time
    for i=1:length(times)
        if selector < times[i]
            return (i, selector)
        else
            selector = selector - times[i]
        end
    end
end
end


using Revise,LinearAlgebra, Random, StatsBase, Statistics, Distributions

Random.seed!(1111)

printm(m) = show(IOContext(stdout, :limit=>false), MIME"text/plain"(), m)
function printmln(m)
  show(IOContext(stdout, :limit=>false), MIME"text/plain"(), m)
  println()
end

function testex2(;
  Ξ = ([4,8], [4,-8], [8,4], [-4, -8], [-4, 8], [-8, -4]),
  dims=(N=6, T=10^5, K=2))

  #set up the factor structure and dividend realizations
  Θ0 = reduce(hcat, [rand(Normal(), dims.T) for k in 1:dims.K])
  @assert size(Θ0) ≡ (dims.T, dims.K)
  D̄ = 1.0
  ϵ = [rand(Normal(), dims.T) for n ∈ 1:dims.N]
  Dvec = [D̄ .+ Θ0*Ξ[n] .+ ϵ[n] for n ∈ 1:dims.N]
  D = reduce(hcat, Dvec)
  Σ = cov(D)

  #select a baseline and a bunch of comp portfolios
  w1 = [1/6 for n ∈ 1:6]
  walt = [rand(6) for i ∈ 1:10^5]
  walt ./= sum.(walt)

  varbase = w1'*Σ*w1
  @info "baseline: w1 $w1: $(w1'*Σ*w1)"

  vars = (w-> w'*Σ*w).(walt)
  for (i,w) ∈ enumerate(walt)
    if vars[i] > varbase
      #@info "w$i $w: $(vars[i])"
      nothing
    else
      @warn "w$i $w: $(vars[i])"
    end
  end

  printmln(Σ)
end

testex2()


"""
Sampling from DPP and k-DPP.

Methods:
--------
    rand(dpp::DeterminantalPointProcess, N::Int)
    randmcmc(dpp::DeterminantalPointProcess, N::Int)
    rand(dpp::DeterminantalPointProcess, N::Int, k::Int)
    randmcmc(dpp::DeterminantalPointProcess, N::Int, k::Int)

References:
-----------
    [1] Kulesza, A., and B. Taskar. Determinantal point processes for machine
        learning. arXiv preprint arXiv:1207.6083, 2012.
    [2] Kang, B. Fast determinantal point process sampling with application to
        clustering. NIPS, 2013.
"""

"""Sample a mask for an elementary DPP.
"""
function _sample_mask(Λ::SharedArray{Float64},
                      M::SharedMatrix{Bool},
                      i::Int, seed::Int)
    rng = MersenneTwister(seed)

    for j in 1:length(Λ)
        M[j, i] = (rand(rng) < Λ[j] / (Λ[j] + 1))
    end
end

"""Sample a mask for an elementary k-DPP.
"""
function _sample_k_mask(Λ::SharedArray{Float64},
                        M::SharedMatrix{Bool},
                        E::SharedMatrix{Float64},
                        k::Int, i::Int, seed::Int)
    rng = MersenneTwister(seed)

    j = length(Λ)
    remaining = k

    # iteratively sample a k-mask
    while remaining > 0
        # compute marginal of j given that we choose remaining values from 1:j
        if j == remaining
            marg = 1
        else
            marg = Λ[j] * E[remaining, j] / E[remaining + 1, j + 1];
        end

        # sample marginal
        if rand(rng) <= marg
            M[j, i] = true
            remaining -= 1
        end
        j -= 1
      end
end

"""Exact sampling from an elementary DPP. The algorithm based on [1].
"""
function _sample_from_elementary(V::SharedMatrix,
                                 M::SharedMatrix{Bool},
                                 i::Int, seed::Int)
    rng = MersenneTwister(seed)

    # select the elementary DPP
    V_mask = M[:, i]

    # edge case: empty sample
    if !any(V_mask)
        return Int[]
    end

    # select the kernel of the elementary DPP
    L = V[:, V_mask]

    Y = Int[]
    mask = ones(Bool, size(L, 2))
    prob = Array{Float64}(undef,size(L, 1))

    for i in 1:size(L, 2)
        # compute probabilities
        fill!(prob, 0)
        for c in 1:size(L, 2)
            !mask[c] && continue
            for r in 1:size(L, 1)
                prob[r] += L[r, c].^2
            end
        end
        prob ./= sum(prob)

        # sample a point in the original space
        h = findfirst(rand(rng) .<= cumsum(prob))
        push!(Y, h)

        # select and mask-out an element
        j = get_first_nz_idx(L[h, :], mask)
        mask[j] = false

        if any(mask)
            # Subtract scaled Lj from other columns so that their
            # projections on e_s[i] turns into 0. This operation
            # preserves the rank of L_{-j}.
            for c in 1:size(L, 2)
                !mask[c] && continue
                for r in 1:size(L, 1)
                    L[r, c] -= L[r, j] * L[h, c] / L[h, j]
                end
            end

            # Gram-Schmidt orthogonalization
            L[:, mask] = Matrix(qr(L[:, mask]).Q)
        end
    end

    sort(Y)
end

"""Exact sampling from a DPP [1].
"""
function Base.rand(dpp::DeterminantalPointProcess, N::Int)
    Λ = SharedArray{Float64}(dpp.Lfact.values)
    V = SharedMatrix{Float64}(dpp.Lfact.vectors)
    M = SharedMatrix{Bool}(zeros(Bool, dpp.size, N))

    # step I: sample masks for elementary DPPs
    pmap((i, seed) -> _sample_mask(Λ, M, i, seed),
         1:N, abs.(rand(dpp.rng, Int, N)))

    # step II: iteratively sample from a mixture of elementary DPPs
    pmap((i, seed) -> _sample_from_elementary(V, M, i, seed),
         1:N, abs.(rand(dpp.rng, Int, N)))
end

"""Exact sampling from a k-DPP [1].
"""
function Base.rand(dpp::DeterminantalPointProcess, N::Int, k::Int)
    Λ = SharedArray{Float64}(dpp.Lfact.values)
    V = SharedMatrix{Float64}(dpp.Lfact.vectors)
    M = SharedMatrix{Bool}(zeros(Bool, dpp.size, N))

    # compute elementary symmetric polynomials
    E = SharedMatrix{Float64}(elem_symm_poly(dpp.Lfact.values, k))

    # step I: sample masks for elementary DPPs
    pmap((i, seed) -> _sample_k_mask(Λ, M, E, k, i, seed),
         1:N, abs.(rand(dpp.rng, Int, N)))

    # step II: iteratively sample from a mixture of elementary DPPs
    pmap((i, seed) -> _sample_from_elementary(V, M, i, seed),
         1:N, abs.(rand(dpp.rng, Int, N)))
end

"""Perform one MCMC accept-reject transition for DPP.
"""
function _do_mcmc_step!(dpp::DeterminantalPointProcess, state::MCMCState)
    # propose an element to swap
    u = rand(dpp.rng, 1:dpp.size)
    insert = !state[1][u]

    # attempt to make a transition
    if insert
        p = _comp_accept_prob(dpp, state, u, insert)
        rand(dpp.rng) < p && _update_mcmc_state!(dpp, state, u, insert)
    else  # delete
        new_state = _update_mcmc_state(dpp, state, u, insert)
        p = _comp_accept_prob(dpp, new_state, u, insert)
        if rand(dpp.rng) < p
            state[1] .= new_state[1]
            state[2] .= new_state[2]
        end
    end
end

"""Perform one MCMC accept-reject transition for k-DPP.
"""
function _do_mcmc_k_step!(dpp::DeterminantalPointProcess, state::MCMCState)
    z, L_z_inv = state

    # propose the elements to swap
    u, v = rand(findall(z)), rand(findall(z.==true))

    # copy the state and delete the u element
    new_state = _update_mcmc_state(dpp, state, u, false)

    # attempt to make a transition
    p = _comp_accept_prob(dpp, new_state, u, v)
    if rand(dpp.rng) < p
        # insert the v element into the new state
        _update_mcmc_state!(dpp, new_state, v, true)
        state[1] .= new_state[1]
        state[2] .= new_state[2]
    end
end

function _comp_accept_prob(dpp::DeterminantalPointProcess, state::MCMCState,
                           u::Int, insert::Bool)
    """Compute accept probability to insert / delete u from the state.
    """
    z, L_z_inv = state

    d_u = dpp.L[u, u]
    if any(z)
        b_u = dpp.L[z, u]
        d_u -= dot(b_u, L_z_inv[z, z] * b_u)
    end

    insert ? min(1.0, d_u) : min(1.0, 1.0 / d_u)
end

function _comp_accept_prob(dpp::DeterminantalPointProcess, state::MCMCState,
                           u::Int, v::Int)
    """Compute accept probability to swap u and v.
    """
    z, L_z_inv = state

    d_u, d_v = dpp.L[u, u], dpp.L[v, v]
    if any(z)
        b_u, b_v = dpp.L[z, u], dpp.L[z, v]
        d_u -= dot(b_u, L_z_inv[z, z] * b_u)
        d_v -= dot(b_v, L_z_inv[z, z] * b_v)
    end

    min(1.0, d_v / d_u)
end

function _update_mcmc_state!(dpp::DeterminantalPointProcess, state::MCMCState,
                             u::Int, insert::Bool)
    """Update the state after u is inserted / deleted.
    """
    z, L_z_inv = state

    if insert
        d_u = dpp.L[u, u]
        if any(z)
            b_u = dpp.L[z, u]
            x_u = L_z_inv[z, z] * b_u
            d_u -= dot(b_u, x_u)

            L_z_inv[z, z] += (x_u * x_u') / d_u
            L_z_inv[z, u] = L_z_inv[u, z] = - x_u / d_u
        end

        L_z_inv[u, u] = 1. / d_u
        z[u] = true
    else  # delete
        z[u] = false

        e = L_z_inv[z, u]
        f = L_z_inv[u, u]

        L_z_inv[z, z] -= (e * e') / f
    end
end

function _update_mcmc_state(dpp::DeterminantalPointProcess, state::MCMCState,
                            u::Int, insert::Bool)
    """Update the state after u is inserted / deleted.
    """
    new_state = deepcopy(state)
    _update_mcmc_state!(dpp, new_state, u, insert)
    new_state
end

function randmcmc(dpp::DeterminantalPointProcess, N::Int;
                  init_state=nothing,
                  return_final_state::Bool=false,
                  mix_eps::Float64=1e-1,
                  mixing_steps::Int=ceil(Int, dpp.size*log(dpp.size/mix_eps)),
                  steps_between_samples::Int=mixing_steps)
    """MCMC sampling from a DPP [2].

    TODO: Add support for running MCMC in parallel, similar as rand.
          Make sure parallelization produces unbiased and consistent samples.
    """
    # initialize the Markov chain
    state = init_state
    if state == nothing
        L_z_inv = Array{Float64}(undef,size(dpp.L))
        z = bitrand(dpp.rng, dpp.size)  # TODO: improve initialization (?)
        if any(z)
            L_z_inv[z, z] = pinv(dpp.L[z, z])
        end
        state = (z, L_z_inv)
    end

    # sanity check
    @assert typeof(state) == MCMCState

    # mix the Markov chain
    for t in 1:mixing_steps
        _do_mcmc_step!(dpp, state)
    end

    Y = []
    for i in 1:N
        push!(Y, findall(state[1]))
        for t in 1:steps_between_samples
            _do_mcmc_step!(dpp, state)
        end
    end

    return_final_state ? (Y, state) : Y
end

function randmcmc(dpp::DeterminantalPointProcess, N::Int, k::Int;
                  init_state=nothing,
                  return_final_state::Bool=false,
                  mix_eps::Float64=1e-1,
                  mixing_steps::Int=ceil(Int, k*log(k / mix_eps)),
                  steps_between_samples::Int=mixing_steps)
    """MCMC sampling from a k-DPP [2].

    TODO: Add support for running MCMC in parallel, similar as rand.
          Make sure parallelization produces unbiased and consistent samples.
    """
    # initialize the Markov chain
    state = init_state
    if state == nothing
        L_z_inv = Array{Float64}(undef,size(dpp.L))
        z = falses(dpp.size)  # TODO: improve initialization (?)
        z[1:k] .= true
        if any(z)
            L_z_inv[z, z] = pinv(dpp.L[z, z])
        end
        state = (z, L_z_inv)
    end

    # sanity check
    @assert typeof(state) == MCMCState
    @assert sum(state[1]) == k

    # mix the Markov chain
    for t in 1:mixing_steps
        _do_mcmc_k_step!(dpp, state)
    end

    Y = []
    for i in 1:N
        push!(Y, findall(state[1]))
        for t in 1:steps_between_samples
            _do_mcmc_k_step!(dpp, state)
        end
    end

    return_final_state ? (Y, state) : Y
end


using Revise

using IndividualRewrite

using Catlab.Present, Catlab.CategoricalAlgebra
using Plots, Random

Random.seed!(5);

# Define structure of a simulation state
@present ThSIR(FreeSchema) begin
    (S,I,R)::Ob
end
@acset_type SIR(ThSIR)

# Helper function for plotting SIR over time.
function plot_states(states::Vector{SIR})
    ss, is, rs = [[nparts(s, x) for s in states] for x in [:S,:I,:R]]
    x = 1:length(states)
    plot(x,ss; label="S")
    plot!(x, is; label="I")
    plot!(x, rs; label="R")
end

# ACSets for rewrite rules
I = @acset SIR begin I=1 end
R = @acset SIR begin R=1 end
I2 = @acset SIR begin I=2 end
SI = @acset SIR begin S=1; I=1 end

# Rewrite rule morphisms
L_infect =  homomorphism(I, SI)
R_infect = homomorphism(I, I2)
L_recov = homomorphism(SIR(), I)
R_recov = homomorphism(SIR(), R)

# Rewrite rules
rules = Dict(pairs((
  infect=Rule(L_infect,R_infect, 2, 1),
  recover=Rule(L_recov, R_recov, 1, 2)
))...)

# Initial state
init = @acset SIR begin S=10; I=3; R=1 end

# Run the simulation and visualize
states = sim(init, rules, 50; verbose=true)
plot_states(states)

using Debugger
Debugger.@enter sim(init, rules, 50; verbose=true)