### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ ab976d70-f59d-4dcd-ac89-771606a74108
begin
	using PlutoUI
	using PlutoTest
	using PlutoTeachingTools
	#eval(Meta.parse(code_for_check_type_funcs))
end

# ╔═╡ 5b826e17-ad51-43d0-9756-89a9fd0ea834
md"""
# Astro 528, Lab 1, Exercise 1
"""

# ╔═╡ acb10cb1-d03d-47c5-90dc-7216e5110ed2
md"""
## Access the lab
Before we delve into the meat of the course, we'll practice accessing the assignment via your private git repository ("repo"), cloning the repo, starting Julia, and launching the Pluto notebook server.
Once you open the Pluto notebook, it will begin executing the entire notebook.  (Depending on your system and how complex the notebook is, sometimes this can take a little time.)  You'll read through this notebook, make some small changes, see how that affects the results, save your changes to your local machine, commit your changes to your local git repository, and then push those to your repository on github.
For many of you, this is likely your first time using all of these tools, and it probably sounds like a lot.  The good news is that before long all of this will become second nature and you'll wonder how you got along without such tools.
"""

# ╔═╡ c15ce452-494f-407c-9c08-02d547929e6a
md"""
## Assigning Variables
In Julia, you assign variables values using the `=` sign.
Assigning integer or floating point numbers uses a natural notation.
Once a cell is executed, Pluto shows the output of each cell above the cell with the relevant code.
For example,"""

# ╔═╡ 25d5497c-3a67-4495-9f37-0bc8f2a64e68
current_year = 2021

# ╔═╡ 774b2e4a-6cd7-4c9a-920d-b4f38354d42d
md"""Strings are created by enclosing text between double quotes.
For example,"""

# ╔═╡ 5917df75-0dab-4778-917e-91bbfddbe03a
begin
	course_number = "Astro 528"
	semester = "Fall 2021"
	instructor = "Eric Ford"
end

# ╔═╡ 0d15ba33-f2c4-46d6-9b08-9181d40596cd
md"""Notice that, when working in a Pluto notebook, cells that assign multiple variables need to be wraped in a `begin...end` block.  (This is not necessary in general (e.g., julia scripts, Jupyter notebooks, or module code.)
Note that only the output from the last line is displayed.  The other variables were still assigned and can be accessed in other cells."""

# ╔═╡ dc0bf251-b3a6-42f7-84f6-796ac3761b2d
course_number

# ╔═╡ dbfceab2-994d-45a1-97ff-d978827bd129
tip(md"""Single quotes are used to indicate a single character (as opposed to a string of length 1).
[Triple quoted string literals](https://docs.julialang.org/en/v1/manual/strings/#Triple-Quoted-String-Literals) are useful for quotes spanning multiple lines (and when you want to include `"` symbols).
""")

# ╔═╡ e2bb48e5-a2b2-42d7-a8e5-5c7532d62e93
md"Assigning integer or floating point numbers uses a natural notation."

# ╔═╡ cda49166-b9fb-4e68-8b3e-b97845c7bf3b
instructor_year = current_year - 1999

# ╔═╡ 0c5d5b14-6c9a-4ac9-8126-0d6e1edba593
favorite_number = 3.1415926535897

# ╔═╡ a3450384-ccfc-4bda-af52-24ed8f429cf8
md"""Now, it's your turn to create variables named `student_name`, `student dept` and `student_year` with values based on yourself.  You can modify the code in the cells below."""

# ╔═╡ 33350a21-7971-406e-b316-6539eca72bf4
begin  # INSERT YOUR RESPONSES BELOW
	student_name = missing
	student_dept = missing
	student_year = missing
end;

# ╔═╡ f15643a4-39d5-4492-98d3-9db4f4cb0d3b
md"""
Run the cell with your code and make sure there are no error messages.
Sometimes, I may provide a few quick checks in the Pluto notebook itself.  For example,
"""

# ╔═╡ 1bc2edaf-a0a7-4c4e-96c2-25872dd30e2e
if !@isdefined(student_year)
	not_defined(:student_year)
else
	let
		if ismissing(student_name)
			still_missing()
		elseif isnothing(student_name)
			keep_working("Could there be a typeo?")
		elseif !(student_name isa String)
			keep_working(md"The value should be a String.")
		else
			correct(md"You set `student_name` to $student_name.")
		end
	end
end

# ╔═╡ c05d62b5-6063-41ca-9823-2499b2d4fbc4
if !@isdefined(student_dept)
	not_defined(:student_dept)
else
	let
		if ismissing(student_dept)
			still_missing()
		elseif isnothing(student_dept)
			keep_working("Could there be a typeo?")
		elseif !(student_dept isa String)
			keep_working(md"The value of `student_year` should be a String.")
		elseif occursin("Astro", student_dept)
			correct(md"Welcome back.")
		else
			correct(md"We're glad to have you join us.")
		end
	end
end

# ╔═╡ 8b666639-e3d0-4535-803d-40ddc4360a9c
if !@isdefined(student_year)
	not_defined(:student_year)
else
	let
		if ismissing(student_year)
			still_missing()
		elseif isnothing(student_year)
			keep_working("Could there be a typeo?")
		elseif !(student_year isa Integer)
			keep_working(md"The value should be an integer.")
		elseif ! (1<= student_year)
			keep_working(md"I'm expecting `student_year` to be one or larger.")
		elseif ! (student_year <= 10)
			keep_working(md"Are you sure?  Most grad students finish in less than 10 years.")
		else
			correct(md"I think that $student_year seems like a reasonable value for `student_year`.")
		end
	end
end

# ╔═╡ ec4ec6e7-9548-4bd6-970a-48fefa9861d7
md"""
## Pass a Test
Once we start doing real calculations, it will be important to test our work.  Therefore, you and I will both write tests that your code should pass. For example, in this case we can test that variables with the correct names exist, they're not empty, that they have the expected type, that their value falls within a reasonable range, etc.
"""

# ╔═╡ ac9e6787-5c42-44d6-907e-e7033e7ab4bc
tip(md"Before using a test, we need to tell Julia to import a module that will provide the `@test` macro.  To simplify the notebook, I've put that code at the bottom of the notebook, so it doesn't distract from our main task.")

# ╔═╡ 087b40c4-7d59-4f9a-b584-65c2aa9f9578
md"
The first few tests below should pass.  Later, I've written a couple of tests with bugs in them for you to fix.  Look at the results of the cells below, find the tests that doesn't pass, and modify the code so that it passes.
"

# ╔═╡ eefa92a1-1d4c-4536-abd3-1c1af98e146f
md"One of the first things to check is that variables have been defined.  If not, then later tests will fail and we might spend time trying to dive deeper into the problem than we need to."

# ╔═╡ 78f81544-fead-4649-b306-48fb1463d397
@test @isdefined(student_name)

# ╔═╡ 3d224633-0eb8-407c-b44c-0d9ec927dfc6
@test @isdefined(student_dept)

# ╔═╡ 2fe3f4e8-5288-4f61-9346-7a455c2abc85
@test @isdefined(student_year)

# ╔═╡ b49789b2-49c0-49e3-b81f-db04e60f9d23
md"Next, we'll test that the types of the variables are what we expect.  If not, then it's unlikely that the values stored will be what we expect."

# ╔═╡ 34f767c9-f95f-4968-9423-ad87bb456775
@test typeof(student_name) == String

# ╔═╡ d6417bdc-9e36-4d2b-8b3e-f5ddafd82437
@test typeof(student_dept) == String

# ╔═╡ 89ac58ea-b782-4dcc-802f-d6d063e7b855
@test typeof(student_year) <: Number

# ╔═╡ 2479fb6a-2e51-4e6f-8155-d47e9e3ae27d
@test typeof(student_year) <: Integer

# ╔═╡ 6d8e0985-3613-49bc-8050-91e2ee73de52
md"We often see the `==` operator when testing for equality of two numbers.  Here the `==` operator is testing for the equivalence of two types.
Alternatively, we can use the `<:` operator to ask if one type is a 'sub-type' of another 'abstract type'."

# ╔═╡ 7f031249-4413-4cd6-8ba5-b3369a1216a0
md"Now, we're ready to test that the values of the variables are consistent with what we expect."

# ╔═╡ 8ff88c72-7982-4025-9057-384a48186dc9
@test !ismissing(student_name) && length(student_name) > 3

# ╔═╡ c408ddb6-af5d-4b6a-ae81-9d338fd15f06
@test !ismissing(student_name) && occursin(' ',student_name)

# ╔═╡ 939a5680-c387-454d-a7f0-16e0d51254d8
md"In the above cell, we test that the `occursin` function retures true.  There should be a slider widget.  Try dragging it left or right.  This can be helpful when figuring out *why* a complicated test fails."

# ╔═╡ d732ab18-65b3-417a-8f8d-5f4502baa777
@test !ismissing(student_dept) && length(student_dept) >= 3

# ╔═╡ bc491991-1c10-43bb-8afe-3ed08c91d3d7
@test !ismissing(student_year) && 1 <= student_year <= 10

# ╔═╡ a328c8c6-6268-4b1f-ae1c-84ee295285c8
@test student_name == instructor

# ╔═╡ 852d7d9d-fc12-42f2-863f-af0b647979df
md"**TODO:** There are (at least) two bugs in the above tests.  Find them and fix them."

# ╔═╡ 6584e095-465d-4b3d-a383-9ac569e307cc
md"In the future, we'll separate out the tests for each exercise into a separate file (located in the `test` directory of the lab repository), so that they don't distract from the code we're writing."

# ╔═╡ 68341392-aa5a-41b7-a76a-76db1c6357f5
md"""Once you're happy with the results, it's time to save your notebook.  Look for the button in the upper right or `Ctrl`+'S'.  If it's not there, then your notebook has been autosaved.

Then read the [instructions for submitting assignments](https://psuastro528.github.io/tips/submitting) for how to commit your changes to your local git repository and push them to github.

When you're done you can look into the second notebook (ex2) in this repository to begin the next exercise."""

# ╔═╡ 51a1bf7c-6a1d-4d31-a735-de3414a97ab6
md"""## Helper code"""

# ╔═╡ 4d860e75-0631-4ef6-853c-818bafe080f4
md"""Students can ignore the code below was quietly used above."""

# ╔═╡ b9f89dc4-a8e8-4afe-9da6-9880cfcf2a37
md"""
The above cell both loads the code from packages and brings the functions (and macros) that it "exports" "into scope", meaing that we can "just call them" without having to include the package name as a prefix.
"""

# ╔═╡ 45efd82d-63e1-4016-838e-d0fbb21db99f
md"The PlutoUI package provides several cool features, but for now we just use it to add the Table of Contents in the upper right corent."

# ╔═╡ 60348e2d-9943-486c-9061-3c5c3bf88883
md"The PlutoTest package provides an alternative `@test` macro for use in Pluto notebooks."

# ╔═╡ ceff74da-4cd2-4091-af1a-8ecde8e13700
md"The PlutoTeachingTools package has several small functions that we use for things like tips."

# ╔═╡ 98ecc733-15f3-4263-821a-483dadfc8ab8
Foldable("Pro Tip!",tip(md"Normally, we'd use the `Test` module for the `@test` macro.  Julia has a large set of modules and packages, that range from very basic functionality to complex science codes.  The quality also varries widely.  Several modules (like `Test`) are included in Julia standard library, so they're already installed for us.

However, inside Pluto, it can be helpful to instead import `PlutoTest`, since it displays the results particularly nicely.  (It's an external package and it's still experimental, so if things break in the future, then we can revert to just using `Test`."))

# ╔═╡ 73a3ce7e-9967-4d11-8ab7-0e0409190975
TableOfContents()

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoTeachingTools = "661c6b06-c737-4d37-b85c-46df65de6f69"
PlutoTest = "cb4044da-4d16-4ffa-a6a3-8cad7f73ebdc"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
PlutoTeachingTools = "~0.1.0"
PlutoTest = "~0.1.0"
PlutoUI = "~0.7.9"
"""

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

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

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

[[HypertextLiteral]]
git-tree-sha1 = "1e3ccdc7a6f7b577623028e0095479f4727d8ec1"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.8.0"

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

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

[[LaTeXStrings]]
git-tree-sha1 = "c7f1c695e06c01b95a67f0cd1d34994f3e7db104"
uuid = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
version = "1.2.1"

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

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

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

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

[[PlutoTeachingTools]]
deps = ["LaTeXStrings", "Markdown", "PlutoUI", "Random"]
git-tree-sha1 = "265980831960aabe7e1f5ae47c898a8459588ee7"
uuid = "661c6b06-c737-4d37-b85c-46df65de6f69"
version = "0.1.3"

[[PlutoTest]]
deps = ["HypertextLiteral", "InteractiveUtils", "Markdown", "Test"]
git-tree-sha1 = "3479836b31a31c29a7bac1f09d95f9c843ce1ade"
uuid = "cb4044da-4d16-4ffa-a6a3-8cad7f73ebdc"
version = "0.1.0"

[[PlutoUI]]
deps = ["Base64", "Dates", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "Suppressor"]
git-tree-sha1 = "44e225d5837e2a2345e69a1d1e01ac2443ff9fcb"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.9"

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

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

[[Reexport]]
git-tree-sha1 = "5f6c21241f0f655da3952fd60aa18477cf96c220"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.1.0"

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

[[Suppressor]]
git-tree-sha1 = "a819d77f31f83e5792a76081eee1ea6342ab8787"
uuid = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
version = "0.2.0"

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

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

# ╔═╡ Cell order:
# ╟─5b826e17-ad51-43d0-9756-89a9fd0ea834
# ╟─acb10cb1-d03d-47c5-90dc-7216e5110ed2
# ╟─c15ce452-494f-407c-9c08-02d547929e6a
# ╠═25d5497c-3a67-4495-9f37-0bc8f2a64e68
# ╟─774b2e4a-6cd7-4c9a-920d-b4f38354d42d
# ╠═5917df75-0dab-4778-917e-91bbfddbe03a
# ╟─0d15ba33-f2c4-46d6-9b08-9181d40596cd
# ╠═dc0bf251-b3a6-42f7-84f6-796ac3761b2d
# ╠═dbfceab2-994d-45a1-97ff-d978827bd129
# ╟─e2bb48e5-a2b2-42d7-a8e5-5c7532d62e93
# ╠═cda49166-b9fb-4e68-8b3e-b97845c7bf3b
# ╠═0c5d5b14-6c9a-4ac9-8126-0d6e1edba593
# ╟─a3450384-ccfc-4bda-af52-24ed8f429cf8
# ╠═33350a21-7971-406e-b316-6539eca72bf4
# ╟─f15643a4-39d5-4492-98d3-9db4f4cb0d3b
# ╟─1bc2edaf-a0a7-4c4e-96c2-25872dd30e2e
# ╟─c05d62b5-6063-41ca-9823-2499b2d4fbc4
# ╟─8b666639-e3d0-4535-803d-40ddc4360a9c
# ╟─ec4ec6e7-9548-4bd6-970a-48fefa9861d7
# ╟─ac9e6787-5c42-44d6-907e-e7033e7ab4bc
# ╟─087b40c4-7d59-4f9a-b584-65c2aa9f9578
# ╟─eefa92a1-1d4c-4536-abd3-1c1af98e146f
# ╠═78f81544-fead-4649-b306-48fb1463d397
# ╠═3d224633-0eb8-407c-b44c-0d9ec927dfc6
# ╠═2fe3f4e8-5288-4f61-9346-7a455c2abc85
# ╟─b49789b2-49c0-49e3-b81f-db04e60f9d23
# ╠═34f767c9-f95f-4968-9423-ad87bb456775
# ╠═d6417bdc-9e36-4d2b-8b3e-f5ddafd82437
# ╠═89ac58ea-b782-4dcc-802f-d6d063e7b855
# ╠═2479fb6a-2e51-4e6f-8155-d47e9e3ae27d
# ╟─6d8e0985-3613-49bc-8050-91e2ee73de52
# ╟─7f031249-4413-4cd6-8ba5-b3369a1216a0
# ╠═8ff88c72-7982-4025-9057-384a48186dc9
# ╠═c408ddb6-af5d-4b6a-ae81-9d338fd15f06
# ╟─939a5680-c387-454d-a7f0-16e0d51254d8
# ╠═d732ab18-65b3-417a-8f8d-5f4502baa777
# ╠═bc491991-1c10-43bb-8afe-3ed08c91d3d7
# ╠═a328c8c6-6268-4b1f-ae1c-84ee295285c8
# ╟─852d7d9d-fc12-42f2-863f-af0b647979df
# ╟─6584e095-465d-4b3d-a383-9ac569e307cc
# ╟─68341392-aa5a-41b7-a76a-76db1c6357f5
# ╟─51a1bf7c-6a1d-4d31-a735-de3414a97ab6
# ╟─4d860e75-0631-4ef6-853c-818bafe080f4
# ╠═ab976d70-f59d-4dcd-ac89-771606a74108
# ╟─b9f89dc4-a8e8-4afe-9da6-9880cfcf2a37
# ╠═45efd82d-63e1-4016-838e-d0fbb21db99f
# ╟─60348e2d-9943-486c-9061-3c5c3bf88883
# ╟─ceff74da-4cd2-4091-af1a-8ecde8e13700
# ╟─98ecc733-15f3-4263-821a-483dadfc8ab8
# ╠═73a3ce7e-9967-4d11-8ab7-0e0409190975
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


# This code is part of QuantumCircuits.
#
# (C) Copyright Rafał Pracht 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

module QBase

export QuantumCircuit, QuantumDevice, QuantumGate, QuantumObject,
       add!, tomatrix, setparameters!, simplify, measure!,
       standardGateError, bindparameters!, decompose

"Quantum object :)"
abstract type QuantumObject end

"Quantum Gate"
abstract type QuantumGate <: QuantumObject end

"Quantum Circuit object"
abstract type QuantumCircuit <: QuantumObject end

"Quantum davice, it may be simulator or real hardware."
abstract type QuantumDevice end

"The quantum gates"
add!(qc::QuantumCircuit, ::QuantumGate) = error("Please implement add! function used to adding gate to circuit for ", typeof(qc))

"Get the matrix version of the quantum object."
tomatrix(qc::QuantumObject) = error("Please implement tomatrix function for ", typeof(qc))
tomatrix(qc::QuantumObject, ::Any) = error("Please implement tomatrix with params function for ", typeof(qc))

"Method return the error of standard gate"
standardGateError(qc::QuantumObject) = error("Please implement standardGateError function for ", typeof(qc))
standardGateError(qc::QuantumObject, ::Any) = error("Please implement standardGateError with params function for ", typeof(qc))

"Set the parameters to the quantum object"
setparameters!(qc::QuantumObject, ::Any) = error("Please implement setparameters! function for ", typeof(qc))

"Bind the parameters to the quantum object"
bindparameters!(qc::QuantumObject) = error("Please implement bindparameters! function for ", typeof(qc))

"Simplify the quantum object"
simplify(qc::QuantumObject) = error("Please implement simplify function for ", typeof(qc))

decompose(qc::QuantumObject) = error("Please implement decompose function for ", typeof(qc))

Base.inv(qc::QuantumObject) = error("Please implement inv function for ", typeof(qc))

measure!(qc::QuantumObject) = error("Please implement measure! function for ", typeof(qc))

end  # module QBase


module TrackerCompat

using ..Tracker: Tracker,
                 TrackedReal,
                 TrackedVector,
                 TrackedMatrix,
                 TrackedArray,
                 TrackedVecOrMat,
                 @grad,
                 track,
                 data,
                 param

import ..Bijectors
using ..Bijectors: Log, SimplexBijector, ADBijector,
    TrackerAD, Inverse, Stacked, Exp

import ChainRulesCore
import LogExpFunctions

using Compat: eachcol
using LinearAlgebra
using Distributions: LocationScale

Bijectors.maporbroadcast(f, x::TrackedArray...) = f.(x...)
function Bijectors.maporbroadcast(
    f,
    x1::TrackedArray{T, N},
    x::AbstractArray{<:TrackedReal}...,
) where {T, N}
    return f.(convert(Array{TrackedReal{T}, N}, x1), x...)
end

Bijectors._eps(::Type{<:TrackedReal{T}}) where {T} = Bijectors._eps(T)
function Base.minimum(d::LocationScale{<:TrackedReal})
    m = minimum(d.ρ)
    if isfinite(m)
        return d.μ + d.σ * m
    else
        return m
    end
end
function Base.maximum(d::LocationScale{<:TrackedReal})
    m = maximum(d.ρ)
    if isfinite(m)
        return d.μ + d.σ * m
    else
        return m
    end
end

# AD implementations
function Bijectors.jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::Real
)
    return data(Tracker.gradient(b, x)[1])
end
function Bijectors.jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::AbstractVector{<:Real}
)
    # We extract `data` so that we don't return a `Tracked` type
    return data(Tracker.jacobian(b, x))
end

# implementations for Shift bijector
function Bijectors._logabsdetjac_shift(a::TrackedReal, x::Real, ::Val{0})
    return tracker_shift_logabsdetjac(a, x, Val(0))
end
function Bijectors._logabsdetjac_shift(a::TrackedReal, x::AbstractVector{<:Real}, ::Val{0})
    return tracker_shift_logabsdetjac(a, x, Val(0))
end
function Bijectors._logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractVector{<:Real},
    ::Val{1}
)
    return tracker_shift_logabsdetjac(a, x, Val(1))
end
function Bijectors._logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractMatrix{<:Real},
    ::Val{1}
)
    return tracker_shift_logabsdetjac(a, x, Val(1))
end
function tracker_shift_logabsdetjac(a, x, ::Val{N}) where {N}
    return param(Bijectors._logabsdetjac_shift(data(a), data(x), Val(N)))
end

# Log bijector

@grad function Bijectors.logabsdetjac(b::Log{1}, x::AbstractVector)
    return -sum(log, data(x)), Δ -> (nothing, -Δ ./ data(x))
end
@grad function Bijectors.logabsdetjac(b::Log{1}, x::AbstractMatrix)
    return -vec(sum(log, data(x); dims = 1)), Δ -> (nothing, .- Δ' ./ data(x))
end
@grad function Bijectors.logabsdetjac(b::Log{2}, x::AbstractMatrix)
    return -sum(log, data(x)), Δ -> (nothing, -Δ ./ data(x))
end

# implementations for Scale bijector
# Adjoints for 0-dim and 1-dim `Scale` using `Real`
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::Real, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::Real, ::Val{0})
    return Bijectors._logabsdetjac_scale(data(a), data(x), Val(0)), Δ -> (inv(data(a)) .* Δ, nothing, nothing)
end
# Need to treat `AbstractVector` and `AbstractMatrix` separately due to ambiguity errors
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::AbstractVector, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{0})
    da = data(a)
    J = fill(inv.(da), length(x))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::AbstractMatrix, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{0})
    da = data(a)
    J = fill(size(x, 1) / da, size(x, 2))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
# adjoints for 1-dim and 2-dim `Scale` using `AbstractVector`
function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(1))
end
@grad function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    # ∂ᵢ (∑ⱼ log|aⱼ|) = ∑ⱼ δᵢⱼ ∂ᵢ log|aⱼ|
    #                 = ∂ᵢ log |aᵢ|
    #                 = (1 / aᵢ) ∂ᵢ aᵢ
    #                 = (1 / aᵢ)
    da = data(a)
    J = inv.(da)
    return Bijectors._logabsdetjac_scale(da, data(x), Val(1)), Δ -> (J .* Δ, nothing, nothing)
end
function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(1))
end
@grad function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    da = data(a)
    Jᵀ = repeat(inv.(da), 1, size(x, 2))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(1)), Δ -> (Jᵀ * Δ, nothing, nothing)
end
# TODO: implement analytical gradient for scaling a vector using a matrix
# function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     track(_logabsdetjac_scale, a, data(x), Val{1})
# end
# @grad function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     throw
# end
# implementations for Stacked bijector
function Bijectors.logabsdetjac(b::Stacked, x::TrackedMatrix{<:Real})
    return map(eachcol(x)) do c
        Bijectors.logabsdetjac(b, c)
    end
end
# TODO: implement custom adjoint since we can exploit block-diagonal nature of `Stacked`
function (sb::Stacked)(x::TrackedMatrix{<:Real})
    return Bijectors.eachcolmaphcat(sb, x)
end
# Simplex adjoints
function Bijectors._simplex_bijector(X::TrackedVecOrMat, b::SimplexBijector{1})
    return track(Bijectors._simplex_bijector, X, b)
end
function Bijectors._simplex_inv_bijector(Y::TrackedVecOrMat, b::SimplexBijector{1})
    return track(Bijectors._simplex_inv_bijector, Y, b)
end
@grad function Bijectors._simplex_bijector(X::AbstractVector, b::SimplexBijector{1})
    Xd = data(X)
    return Bijectors._simplex_bijector(Xd, b), Δ -> (Bijectors.simplex_link_jacobian(Xd)' * Δ, nothing)
end
@grad function Bijectors._simplex_inv_bijector(Y::AbstractVector, b::SimplexBijector{1})
    Yd = data(Y)
    return Bijectors._simplex_inv_bijector(Yd, b), Δ -> (Bijectors.simplex_invlink_jacobian(Yd)' * Δ, nothing)
end

@grad function Bijectors._simplex_bijector(X::AbstractMatrix, b::SimplexBijector{1})
    Xd = data(X)
    return Bijectors._simplex_bijector(Xd, b), Δ -> begin
        Bijectors.maphcat(eachcol(Xd), eachcol(Δ)) do c1, c2
            Bijectors.simplex_link_jacobian(c1)' * c2
        end, nothing
    end
end
@grad function Bijectors._simplex_inv_bijector(Y::AbstractMatrix, b::SimplexBijector{1})
    Yd = data(Y)
    return Bijectors._simplex_inv_bijector(Yd, b), Δ -> begin
        Bijectors.maphcat(eachcol(Yd), eachcol(Δ)) do c1, c2
            Bijectors.simplex_invlink_jacobian(c1)' * c2
        end, nothing
    end
end

Bijectors.replace_diag(::typeof(log), X::TrackedMatrix) = track(Bijectors.replace_diag, log, X)
@grad function Bijectors.replace_diag(::typeof(log), X)
    Xd = data(X)
    f(i, j) = i == j ? log(Xd[i, j]) : Xd[i, j]
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = i == j ? ∇[i, j]/Xd[i, j] : ∇[i, j]
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

Bijectors.replace_diag(::typeof(exp), X::TrackedMatrix) = track(Bijectors.replace_diag, exp, X)
@grad function Bijectors.replace_diag(::typeof(exp), X)
    Xd = data(X)
    f(i, j) = ifelse(i == j, exp(Xd[i, j]), Xd[i, j])
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = ifelse(i == j, ∇[i, j]*exp(Xd[i, j]), ∇[i, j])
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

Bijectors.logabsdetjac(b::SimplexBijector{1}, x::TrackedVecOrMat) = track(Bijectors.logabsdetjac, b, x)
@grad function Bijectors.logabsdetjac(b::SimplexBijector{1}, x::AbstractVector)
    xd = data(x)
    return Bijectors.logabsdetjac(b, xd), Δ -> begin
        (nothing, Bijectors.simplex_logabsdetjac_gradient(xd) * Δ)
    end
end
@grad function Bijectors.logabsdetjac(b::SimplexBijector{1}, x::AbstractMatrix)
    xd = data(x)
    return Bijectors.logabsdetjac(b, xd), Δ -> begin
        (nothing, Bijectors.maphcat(eachcol(xd), Δ) do c, g
            Bijectors.simplex_logabsdetjac_gradient(c) * g
        end)
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractVector)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:α_, :β, :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedVecOrMat)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedVector)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
]
    @eval begin
        function Bijectors._radial_transform($(header...))
            α = LogExpFunctions.log1pexp(α_)            # from A.2
            β_hat = -α + LogExpFunctions.log1pexp(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = vectorof(typeof(β_hat))
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = typeof(z_0)
            else
                T = TV = typeof(z)
            end
            Tr = promote_type(eltype(z), eltype(z_0))
            r::Tr = norm((z .- z_0)::TV)
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractMatrix)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:α_, :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
]
    @eval begin
        function Bijectors._radial_transform($(header...))
            α = LogExpFunctions.log1pexp(α_)            # from A.2
            β_hat = -α + LogExpFunctions.log1pexp(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = matrixof(TV)
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = matrixof(TV)
            else
                T = typeof(z)
                TV = vectorof(T)
            end
            r::TV = eachcolnorm(z .- z_0)
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end
eachcolnorm(X) = map(norm, eachcol(X))
eachcolnorm(X::TrackedMatrix) = track(eachcolnorm, X)
@grad function eachcolnorm(X)
    Xd = data(X)
    y = map(norm, eachcol(Xd))
    y, Δ -> begin
        (Xd .* (Δ ./ y)',)
    end
end

function matrixof(::Type{<:TrackedArray{T,1,Vector{T}}}) where {T<:Real}
    return TrackedArray{T,2,Matrix{T}}
end
function matrixof(::Type{TrackedReal{T}}) where {T<:Real}
    return TrackedArray{T,2,Matrix{T}}
end
function vectorof(::Type{<:TrackedArray{T,2,Matrix{T}}}) where {T<:Real}
    return TrackedArray{T,1,Vector{T}}
end
function vectorof(::Type{TrackedReal{T}}) where {T<:Real}
    return TrackedArray{T,1,Vector{T}}
end

(b::Exp{0})(x::TrackedVector) = exp.(x)::vectorof(float(eltype(x)))
(b::Exp{1})(x::TrackedVector) = exp.(x)::vectorof(float(eltype(x)))
(b::Exp{1})(x::TrackedMatrix) = exp.(x)::matrixof(float(eltype(x)))
(b::Exp{2})(x::TrackedMatrix) = exp.(x)::matrixof(float(eltype(x)))

(b::Log{0})(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
(b::Log{1})(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
(b::Log{1})(x::TrackedMatrix) = log.(x)::matrixof(float(eltype(x)))
(b::Log{2})(x::TrackedMatrix) = log.(x)::matrixof(float(eltype(x)))

Bijectors.logabsdetjac(b::Log{0}, x::TrackedVector) = .-log.(x)::vectorof(float(eltype(x)))
Bijectors.logabsdetjac(b::Log{1}, x::TrackedMatrix) = - vec(sum(log.(x); dims = 1))

Bijectors.getpd(X::TrackedMatrix) = track(Bijectors.getpd, X)
@grad function Bijectors.getpd(X::AbstractMatrix)
    Xd = data(X)
    return Bijectors.LowerTriangular(Xd) * Bijectors.LowerTriangular(Xd)', Δ -> begin
        Xl = Bijectors.LowerTriangular(Xd)
        return (Bijectors.LowerTriangular(Δ' * Xl + Δ * Xl),)
    end
end

Bijectors.lower(A::TrackedMatrix) = track(Bijectors.lower, A)
@grad function Bijectors.lower(A::AbstractMatrix)
    Ad = data(A)
    return Bijectors.lower(Ad), Δ -> (Bijectors.lower(Δ),)
end

Bijectors._inv_link_chol_lkj(y::TrackedMatrix) = track(Bijectors._inv_link_chol_lkj, y)
@grad function Bijectors._inv_link_chol_lkj(y_tracked)
    y = data(y_tracked)

    K = LinearAlgebra.checksquare(y)

    w = similar(y)

    z_mat = similar(y) # cache for adjoint
    tmp_mat = similar(y)

    @inbounds for j in 1:K
        w[1, j] = 1
        for i in 2:j
            z = tanh(y[i-1, j])
            tmp = w[i-1, j]

            z_mat[i, j] = z
            tmp_mat[i, j] = tmp

            w[i-1, j] = z * tmp
            w[i, j] = tmp * sqrt(1 - z^2)
        end
        for i in (j+1):K
            w[i, j] = 0
        end
    end

    function pullback_inv_link_chol_lkj(Δw)
        LinearAlgebra.checksquare(Δw)

        Δy = zero(y)

        @inbounds for j in 1:K
            Δtmp = Δw[j,j]
            for i in j:-1:2
                Δz = Δw[i-1, j] * tmp_mat[i, j] - Δtmp * tmp_mat[i, j] / sqrt(1 - z_mat[i, j]^2) * z_mat[i, j]
                Δy[i-1, j] = Δz / cosh(y[i-1, j])^2
                Δtmp = Δw[i-1, j] * z_mat[i, j] + Δtmp * sqrt(1 - z_mat[i, j]^2)
            end
        end

        return (Δy,)
    end

    return w, pullback_inv_link_chol_lkj
end

Bijectors._link_chol_lkj(w::TrackedMatrix) = track(Bijectors._link_chol_lkj, w)
@grad function Bijectors._link_chol_lkj(w_tracked)
    w = data(w_tracked)

    K = LinearAlgebra.checksquare(w)

    z = similar(w)

    @inbounds z[1, 1] = 0

    tmp_mat = similar(w) # cache for pullback.

    @inbounds for j=2:K
        z[1, j] = atanh(w[1, j])
        tmp = sqrt(1 - w[1, j]^2)
        tmp_mat[1, j] = tmp
        for i in 2:(j - 1)
            p = w[i, j] / tmp
            tmp *= sqrt(1 - p^2)
            tmp_mat[i, j] = tmp
            z[i, j] = atanh(p)
        end
        z[j, j] = 0
    end

    function pullback_link_chol_lkj(Δz)
        LinearAlgebra.checksquare(Δz)

        Δw = similar(w)

        @inbounds Δw[1,1] = zero(eltype(Δz))

        @inbounds for j=2:K
            Δw[j, j] = 0
            Δtmp = zero(eltype(Δz)) # Δtmp_mat[j-1,j]
            for i in (j-1):-1:2
                p = w[i, j] / tmp_mat[i-1, j]
                ftmp = sqrt(1 - p^2)
                d_ftmp_p = -p / ftmp
                d_p_tmp = -w[i,j] / tmp_mat[i-1, j]^2

                Δp = Δz[i,j] / (1-p^2) + Δtmp * tmp_mat[i-1, j] * d_ftmp_p
                Δw[i, j] = Δp / tmp_mat[i-1, j]
                Δtmp = Δp * d_p_tmp + Δtmp * ftmp # update to "previous" Δtmp
            end
            Δw[1, j] = Δz[1, j] / (1-w[1,j]^2) - Δtmp / sqrt(1 - w[1,j]^2) * w[1,j]
        end

        return (Δw,)
    end

    return z, pullback_link_chol_lkj
end

function Bijectors.find_alpha(wt_y::T, wt_u_hat::T, b::T) where {T<:TrackedReal}
    return track(Bijectors.find_alpha, wt_y, wt_u_hat, b)
end
@grad function Bijectors.find_alpha(wt_y::TrackedReal, wt_u_hat::TrackedReal, b::TrackedReal)
    α = Bijectors.find_alpha(data(wt_y), data(wt_u_hat), data(b))

    ∂wt_y = inv(1 + wt_u_hat * sech(α + b)^2)
    ∂wt_u_hat = - tanh(α + b) * ∂wt_y
    ∂b = ∂wt_y - 1
    find_alpha_pullback(Δ::Real) = (Δ * ∂wt_y, Δ * ∂wt_u_hat, Δ * ∂b)

    return α, find_alpha_pullback
end

# `OrderedBijector`
Bijectors._transform_ordered(y::Union{TrackedVector,TrackedMatrix}) = track(Bijectors._transform_ordered, y)
@grad function Bijectors._transform_ordered(y::AbstractVecOrMat)
    x, dx = ChainRulesCore.rrule(Bijectors._transform_ordered, data(y))
    return x, (wrap_chainrules_output ∘ Base.tail ∘ dx)
end

function Bijectors._transform_inverse_ordered(x::Union{TrackedVector,TrackedMatrix})
    return track(Bijectors._transform_inverse_ordered, x)
end
@grad function Bijectors._transform_inverse_ordered(x::AbstractVecOrMat)
    y, dy = ChainRulesCore.rrule(Bijectors._transform_inverse_ordered, data(x))
    return y, (wrap_chainrules_output ∘ Base.tail ∘ dy)
end

# NOTE: Probably doesn't work in complete generality.
wrap_chainrules_output(x) = x
wrap_chainrules_output(x::ChainRulesCore.AbstractZero) = nothing
wrap_chainrules_output(x::Tuple) = map(wrap_chainrules_output, x)

end


module modelutils

using ..Flux   # must be Flux@0.90
using ..BSON
using ..ArgCheck, ..Dates

export save, load!

if Flux.has_cuarrays()
    using CuArrays
    # note I'm using cu(zeros(...)), cu(randn(...)), CuArrays.randn(...) in particular occasionally
    # (apparently non-deterministically) kept bugging out on the GPU side :/.
    randn_repar(σ::CuArray, n, d, stochastic=true) = !stochastic ? cu(zeros(Float32, n, d)) : σ .* cu(randn(Float32, n,d))
    randn_repar(σ::TrackedArray{Float32, N, CuArray{Float32,N}}, n, d, stochastic=true) where N =
        !stochastic ? cu(zeros(Float32, n, d)) : σ .* cu(randn(Float32, n,d))
end

# Add a third channel to a 2-channel image Tensor
chan3cat(x::AbstractArray{T,4}) where T = cat(x, zero(x)[:,:,1:1,:], dims=3)
chan3cat(x::AbstractArray{T,3}) where T = cat(x, zero(x)[:,:,1:1], dims=3)

# get_final_layer_dim from an arbitrary Chain of layers.
# This was created for pretty printing (Base.show) where n_y was unavailable in MT-LDS.
get_final_layer_dim(l) = nothing
get_final_layer_dim(l::Int) = l
get_final_layer_dim(l::Chain) = get_final_layer_dim(l.layers[end])
get_final_layer_dim(l::Dense) = length(l.b)

get_first_layer_dim(l) = nothing
get_first_layer_dim(l::Int) = l
get_first_layer_dim(l::Chain) = get_first_layer_dim(l.layers[1])
get_first_layer_dim(l::Dense) = size(l.W, 2)

"""
    LookupTable(d, mu_sd, logstd_sd, lkp)

A lookup table (Wrapped Dict) for use instead of an encoder for use in standard variational
inference. The lookup table maps training data (arrays) to mean/logstd Flux.params for
a Gaussian variational posterior. Where a queried array does not have an entry, LookupTable
creates one using standard Gaussian variates with `mu_sd` standard deviation for the mean
and `logstd_sd` standard deviation for the logstd.
"""
struct LookupTable
    d::Int
    mu_sd::Float32
    logstd_sd::Float32
    lkp::Dict
end
LookupTable(d, mu_sd=0.001f0, logstd_sd=-1f0) = LookupTable(d, mu_sd, logstd_sd, Dict())
function Flux.mapleaves(f, s::LookupTable)
    for (k,v) in s
        s[k] = f(v)
    end
end

(lkp::LookupTable)(x::Flux.TrackedArray) = error("Unable to accept tracked arrays in LookupTable.")
function (lkp::LookupTable)(x::AbstractVecOrMat{T}) where T
    # seems inefficient to *always* allocate the default value, but `get` doesn't support
    # lazy evaluation (so it seems), and below is *MUCH* faster than using try/catch which
    # adds O(100μs) to what is otherwise a O(1μs) operation.
    def_value = (Flux.param(randn(T, lkp.d, 1)*lkp.mu_sd), Flux.param(ones(T, lkp.d, 1)*lkp.logstd_sd))
    mu_sd = get!(lkp.lkp, vec(x), def_value)
    return (mu_sd[1], σ.(mu_sd[2]))   # pass 'logstd' (!) through sigmoid
end

function (lkp::LookupTable)(x::AbstractArray{T, 3}) where T
    mu_sds = [lkp(x[:,:,i]) for i in 1:size(x, 3)]
    return hcat([mu for (mu, sd) in mu_sds]...), hcat([sd for (mu, sd) in mu_sds]...)
end

# play nicely with Flux's params.
Flux.children(m::Iterators.Flatten) = m
Flux.children(m::Flux.Params) = m        # params is currently not idempotent: this fixes it.
                                         # Not sure if unintended consequences though?
Flux.children(lkp::LookupTable) = () #Iterators.Flatten(lkp.lkp.vals[findall(lkp.lkp.slots .== 1)])
Flux.params(lkp::LookupTable) = Flux.params(Flux.children(lkp))

"""
    PartialSplit(ixs1, ixs2, split_on_d, fn1, fn2)

Splits an input in two on dimension ``split_on_d`` according to the two UnitRange vectors
``ix1`` and ``ix2``, and passes the result into functions ``fn1`` and ``fn2`` respectively.
"""
struct PartialSplit{V,W}
    ixs1::UnitRange
    ixs2::UnitRange
    split_on_d::Int
    fn1::V
    fn2::W
end

function (ps::PartialSplit)(x::AbstractArray)
    x1 = selectdim(x, ps.split_on_d, ps.ixs1)
    x2 = selectdim(x, ps.split_on_d, ps.ixs2)
    fx1 = ps.fn1(x1)
    fx2 = ps.fn2(x2)
    return vcat(fx1, fx2)
end

Flux.@treelike PartialSplit    # new Flux calls this '@functor'. May need to replace.

function Base.show(io::IO, l::PartialSplit)
    l1 = l.fn1 == identity ? "" : format("ixs: {:s} => {:s}", string(l.ixs1),
        string(typeof(l.fn1)))
    l2 = l.fn2 == identity ? "" : format("ixs: {:s} => {:s}", string(l.ixs2),
        string(typeof(l.fn2)))
    out = length(l1) == 0 ? (length(l2) == 0 ? "" : l2) : (length(l2) == 0 ? l1 : l1 * "\n" * l2)
    print(io, "PartialSplit(\n" * out * "\n)")
end

"""
    MultiDense(Dense(in_1, out_1), Dense(in_2, out_2))

2× Dense layer for convenience for variational (μ/σ) like application. Allows
one to keep both layers together, and can be called like:

    m(x)  # = (m.Dense1(x), m.Dense2(x))

for an instance `m`; the result is a Tuple.
"""
struct MultiDense{F1, F2, S, T}
    Dense1::Dense{F1,S,T}
    Dense2::Dense{F2,S,T}
end
Flux.@treelike MultiDense    # new Flux calls this '@functor'. May need to replace.
(a::MultiDense)(x::AbstractArray) = (a.Dense1(x), a.Dense2(x))

function Base.show(io::IO, l::MultiDense)
  fs = [d.σ == identity ? "" : ", " * string(d.σ) for d in [l.Dense1, l.Dense2]]
  print(io, "MultiDense((", size(l.Dense1.W, 2), ", ", size(l.Dense1.W, 1), fs[1], ")")
  print(io, ", (", size(l.Dense2.W, 2), ", ", size(l.Dense2.W, 1), fs[2], "))")
end

"""
    mlp(units::Tuple; activation::Function, final::Function)

A convenience wrapper for defining a fully connected feedforward (MLP)
component. `units` is a tuple which specifies the number of units in each layer,
with the first and last components taken to be the inputs / outputs of the MLP;
all intermediate units are hidden layers.

Example:
    mlp((20,128,128,40)) === Chain(Dense(20,128,relu), Dense(128,128,relu), Dense(128, 40))

The activation functions can be customized via `activation=...` (for hidden layers)
and `final=...` for the final layer.
"""
function mlp(units::Tuple; activation::Function=relu, final::Function=identity)
    last = length(units)-1
    layers = [Dense(i,j, l==last ? final : activation) for (l,(i,j)) in
            enumerate(zip(units[1:end-1], units[2:end]))]
    return last == 1 ? layers[1] : Chain(layers...)
end
mlp(x...; activation::Function=relu, final::Function=identity) = mlp(x, activation=activation, final=final)

"""
    BRNNenc(d_in, d_state)

Bidirectional RNN using an LSTM in each direction with `d_in` dimensional inputs
and `d_state` dimension for _each_. This is an encoder variant which only uses
the final state (of each LSTM), concatenated. Therefore, its usage is as
follows for `xs`, a Vector where each element corresponds to the input at time
``t``:

    m(xs)::Vector

for an instance `m`, where the output vector is the final state concatenation.
"""
struct BRNNenc{V}
    forward::V
    backward::V
end

BRNNenc(in::Integer, out::Integer) = BRNNenc(LSTM(in, out), LSTM(in, out))

function (m::BRNNenc)(xs::AbstractVector)
    m.forward.(xs)
    m.backward.(reverse(xs))
    return vcat(m.forward.state[1], m.backward.state[1])
end

Flux.@treelike BRNNenc

function Base.show(io::IO, l::BRNNenc)
  print(io, "Bidirectional RNN. (Forward / Backward):")
  show(l.forward); print(", "); show(l.backward)
end

################################################################################
##  Generate samples from posterior (using encoder) via reparameterization    ##
################################################################################

randn_repar(σ::AbstractArray, n, d, stochastic=true) = !stochastic ? zeros(Float32, n, d) : σ .* randn(Float32, n,d)

function posterior_sample(enc, dec, input, T_max, stochastic=true)
    Flux.reset!(enc)
    for tt = 1:T_max; enc(input[:,tt,:]); end
    μ_, σ_ = dec(enc.state)
    n, d = size(μ_)
    smp = μ_ + randn_repar(σ_, n, d, stochastic)
    return smp, μ_, σ_
end

function posterior_sample(enc::LookupTable, dec, input, T_max, stochastic=true)
    μ_, σ_ = enc(input)
    n, d = size(μ_)
    smp = μ_ + randn_repar(σ_, n, d, stochastic)
    return smp, μ_, σ_
end

################################################################################
##  Save and Load utils for parameters of arbitrary models.                   ##
##                                                                            ##
##  These utils can be used for *any* models as they don't save the model def ##
##  itself, but uses the fact that the parameter 'vector' can be extracted    ##
##  easily using Flux's treelike structure and similarly loaded easily.       ##
################################################################################

"""
    save(fname::String, ps::Flux.Params, opt::Flux.ADAM; timestamp::Bool=false, force::Bool=false)
    save(fname::String, ps::Flux.Params; timestamp::Bool=false, force::Bool=false)

Save a parameter vector extracted via `Flux.params(models...)` to `fname`. in
BSON format. Since the state of the optimizer is critical for continuing
optimization (should it be required), the optimizer can optionally be given too.
At the moment this is limited to Adam, but this is easily extended.

Two options:
   (i) `timestamp = true` will insert a timestamp just before the extension.
   (ii) `force=true` will overwrite an existing file of the same name. The
        default behaviour is to throw an error instead.
"""
function save(fname::String, ps::Flux.Params, opt::Flux.ADAM; timestamp::Bool=false, force::Bool=false)
    fname = _get_and_check_fname(fname, timestamp, force)
    BSON.bson(fname, ps=[cpu(Tracker.data(p)) for p in ps], etabeta=(opt.eta, opt.beta),
                opt_state=[mapleaves(cpu, opt.state[p]) for p in ps])
end

function save(fname::String, ps::Flux.Params; timestamp::Bool=false, force::Bool=false)
    fname = _get_and_check_fname(fname, timestamp, force)
    BSON.bson(fname, ps=[cpu(Tracker.data(p)) for p in ps])
end

function _get_and_check_fname(fname::String, timestamp::Bool, force::Bool)
    fn = splitext(fname)
    @argcheck fn[2] in [".bson", ""]
    timestamp && (fname = fn[1] * Dates.format(Dates.now(), "yyyymmddTHH:MM:SS") * fn[2])
    !force && (@assert !isfile(fname) "fname already exists. To overwrite, use `force=true`.")
    return fname
end

"""
    load!(ps::Flux.Params, fname::String)
    load!(ps::Flux.Params, opt::Flux.ADAM, fname::String)

The inverse operation to save; loads a parameter vector from a BSON file created
per the `save` function. This operation will overwrite the parameter values in
the `ps` parameter 'vector' and optionally also the optimizer state (and η, β)
if specified as an argument.

The function takes an _existing_ parameter vector in order to link it to some
pre-specified models; this allows the file format to be agnostic to the original
model.
"""
function load!(ps::Flux.Params, fname::String)
    tf = Flux.has_cuarrays() && Tracker.data(first(ps)) isa CuArray ? gpu : identity
    _load_pars!(ps, BSON.load(fname)[:ps], tf)
end

function load!(ps::Flux.Params, opt::Flux.ADAM, fname::String)
    tf = Flux.has_cuarrays() && Tracker.data(first(ps)) isa CuArray ? cu : identity
    f=BSON.load(fname)
    _load_pars!(ps, f[:ps], tf)
    opt.eta, opt.beta = f[:etabeta]
    for (p, p_saved) in zip(ps, f[:opt_state])
        opt.state[p] = tf(p_saved)
    end
end

function _load_pars!(ps_to::Flux.Params, ps_from::Vector, tf::Function)
    try
        for (p, p_saved) in zip(ps_to, ps_from)
            p.data .= tf(p_saved)
        end
    catch e
        # Debugging dimension mismatch:
        if e isa DimensionMismatch
            @warn "Specified file has parameters of different dimensions:"; flush(stderr)
            for (p, p_saved) in zip(ps_to, ps_from)
                txt = ("(cur/new): ", size(p.data), " <-- ", size(p_saved) , "\n")
                if size(p.data) == size(p_saved)
                    print(txt...)
                else
                    printstyled(IOContext(stdout, :color => true), txt...; bold=true)
                end
            end
        end
        rethrow(e)
    end
end

end


#
# Execution machinery
#

function price_input_vector(program::AbstractPricer, scenario::Scenario) :: Vector{Float64}
    return price_input_vector(program.input_riskfactors, scenario, program.currency)
end

# A NonCashRiskFactor has no Currency information, therefore we don't need to apply `exch`
function parse_risk_factor_value(scenario::Scenario, risk_factor::NonCashRiskFactor, ::TO_CURRENCY) :: Float64 where {TO_CURRENCY<:Currencies.Currency}
    return scenario[risk_factor]
end

function parse_risk_factor_value(scenario::Scenario, risk_factor::SpotCurrency{FROM_CURRENCY}, to::TO_CURRENCY) :: Float64 where {FROM_CURRENCY<:Currencies.Currency, TO_CURRENCY<:Currencies.Currency}
    exch(scenario, FROM_CURRENCY(), to)
end

function parse_risk_factor_value(scenario::Scenario, risk_factor::CashRiskFactor, to::Currencies.Currency) :: Float64
    exch(scenario, scenario[risk_factor], to)
end

function price_input_vector(
        input_symbols::OptimizingIR.LookupTable{RiskFactor},
        scenario::Scenario,
        functional_currency::Currencies.Currency) :: Vector{Float64}

    if isempty(input_symbols)
        # there are no risk factors
        return Vector{Float64}()
    end

    input_vector = zeros(Float64, length(input_symbols))

    for risk_factor in input_symbols
        if !isa(risk_factor, SpotCurrency)
            # For currencies, `exch` checks wether currency prices are available.
            # For any other risk factor, the Scenario must provide values for Risk Factors.
            # See `parse_risk_factor_value`.
            @assert haskey(scenario, risk_factor) "Scenario has no value for Risk Factor $risk_factor."
        end

        @inbounds input_vector[ OptimizingIR.indexof(input_symbols, risk_factor) ] = parse_risk_factor_value(scenario, risk_factor, functional_currency)
    end

    return input_vector
end

# converts `program.pricing_function`
# to something that accepts a vector as input
# so that ForwardDiff can do its thing
# Regarding `invokelatest`, see https://felipenoris.github.io/OptimizingIR.jl/dev/#World-Age-Problem-1
function price_closure(program::Pricer) :: Function
    if isa(program.pricing_function, Function)
        x -> Base.invokelatest(program.pricing_function, x)
    else
        x -> program.pricing_function(x)
    end
end

# Regarding `invokelatest`, see https://felipenoris.github.io/OptimizingIR.jl/dev/#World-Age-Problem-1
function price_closure(program::CashflowPricer) :: Function
    # the CashflowPricer returns the contract price along with
    # the prices for each cashflow.

    if isempty(program.output_index_to_cashflow_type)
        # there are no events for this contract
        if isa(program.pricing_function, Function)
            return x -> Base.invokelatest(program.pricing_function, x)
        else
            return x -> program.pricing_function(x)
        end
    else
        if isa(program.pricing_function, Function)
            return x -> Base.invokelatest(program.pricing_function, x)[program.price_output_index]
        else
            return x -> program.pricing_function(x)[program.price_output_index]
        end
    end
end

function price(program::AbstractPricer, scenario::Scenario) :: Number
    x = price_input_vector(program, scenario)
    f = price_closure(program)
    return f(x)
end

function eachcashflow(p::CashflowPricer, scenario::Scenario) :: Vector{Cashflow}

    if isempty(p.output_index_to_cashflow_type)
        return Vector{Cashflow}()
    else

        cashflows = Vector{Cashflow}(undef, length(p.output_index_to_cashflow_type))

        x = price_input_vector(p, scenario)
        result = p.pricing_function(x)

        i = 1
        for (index, cftype) in p.output_index_to_cashflow_type
            val = result[index]
            cf = Cashflow(cftype.event, cftype.maturity, val, p.currency)
            cashflows[i] = cf
            i += 1
        end
    end

    return cashflows
end

function exposures(::DeltaNormalExposuresMethod, program::AbstractPricer, scenario::Scenario) :: ExposureResult
    result = ExposureResult()
    first_order_exposures!(result, program, scenario)
    return result
end

function exposures(::DeltaGammaApproxExposuresMethod, program::AbstractPricer, scenario::Scenario) :: ExposureResult
    result = ExposureResult()
    first_order_exposures!(result, program, scenario)
    second_order_exposures!(result, program, scenario)
    return result
end

function first_order_exposures!(result::ExposureResult, program::AbstractPricer, scenario::Scenario)

    x = price_input_vector(program, scenario)

    if isempty(x)
        return result
    end

    f_grad = args -> ForwardDiff.gradient(price_closure(program), args)
    exposures_vector = x .* f_grad(x)

    # sanity check
    _risk_factors = riskfactors(program)
    @assert length(exposures_vector) == length(_risk_factors) "Size of exposures vector ($(length(exposures_vector))) should equal the size of inputs vector ($(length(_risk_factors)))."

    for i in 1:length(exposures_vector)
        # consider as zero exposure values below eps(1_000_000_000_000.0)
        if abs(exposures_vector[i]) > 0.0001220703125
            @inbounds result[_risk_factors[i]] = exposures_vector[i]
        end
    end

    nothing
end

"""
Returns second order exposures approximation as:

``\\frac{1}{2} \\frac{ \\partial^2 P }{ \\partial S^2 } S^2``

where `P` is the pricing function,
and `S` is a Risk Factor.
"""
function second_order_exposures!(result::ExposureResult, program::AbstractPricer, scenario::Scenario)

    x = price_input_vector(program, scenario)

    if isempty(x)
        return result
    end

    f_hessian = args -> ForwardDiff.hessian(price_closure(program), args)
    exposures_vector = 0.5 * (LinearAlgebra.diag(f_hessian(x)) .* (x.^2))

    # sanity check
    _risk_factors = riskfactors(program)
    @assert length(exposures_vector) == length(_risk_factors) "Size of exposures vector ($(length(exposures_vector))) should equal the size of inputs vector ($(length(_risk_factors)))."

    for i in 1:length(exposures_vector)
        # consider as zero exposure values below eps(1_000_000_000_000.0)
        if abs(exposures_vector[i]) > 0.0001220703125
            @inbounds result[SecondOrderRiskFactor(_risk_factors[i])] = exposures_vector[i]
        end
    end

    nothing
end


function PSI.construct_device!(
    optimization_container::PSI.OptimizationContainer,
    sys::PSY.System,
    model::PSI.DeviceModel{St, BookKeepingwReservationOutage},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = PSI.get_available_components(St, sys)

    if !PSI.validate_available_devices(St, devices)
        return
    end

    # Variables
    PSI.add_variables!(
        optimization_container,
        PSI.ActivePowerInVariable,
        devices,
        BookKeepingwReservationOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.ActivePowerOutVariable,
        devices,
        BookKeepingwReservationOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.ReactivePowerVariable,
        devices,
        BookKeepingwReservationOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.EnergyVariable,
        devices,
        BookKeepingwReservationOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.ReserveVariable,
        devices,
        BookKeepingwReservationOutage(),
    )

    # Initial Conditions
    PSI.initial_conditions!(
        optimization_container,
        devices,
        BookKeepingwReservationOutage(),
    )

    # Constraints
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerOutVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerInVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ReactivePowerVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    outage_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.feedforward!(optimization_container, devices, model, PSI.get_feedforward(model))

    # Energy Balanace limits
    PSI.energy_balance_constraint!(
        optimization_container,
        devices,
        model.formulation,
        S,
        PSI.get_feedforward(model),
    )

    return
end

function PSI.construct_device!(
    optimization_container::PSI.OptimizationContainer,
    sys::PSY.System,
    model::PSI.DeviceModel{St, BookKeepingwReservationOutage},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = PSI.get_available_components(St, sys)

    if !PSI.validate_available_devices(St, devices)
        return
    end

    # Variables
    PSI.add_variables!(
        optimization_container,
        PSI.ActivePowerInVariable,
        devices,
        BookKeepingwReservationOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.ActivePowerOutVariable,
        devices,
        BookKeepingwReservationOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.EnergyVariable,
        devices,
        BookKeepingwReservationOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.ReserveVariable,
        devices,
        BookKeepingwReservationOutage(),
    )

    # Initial Conditions
    PSI.initial_conditions!(
        optimization_container,
        devices,
        BookKeepingwReservationOutage(),
    )

    # Constraints
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerOutVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerInVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    outage_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.feedforward!(optimization_container, devices, model, PSI.get_feedforward(model))

    # Energy Balanace limits
    PSI.energy_balance_constraint!(
        optimization_container,
        devices,
        model.formulation,
        S,
        PSI.get_feedforward(model),
    )

    return
end

function PSI.construct_device!(
    optimization_container::PSI.OptimizationContainer,
    sys::PSY.System,
    model::PSI.DeviceModel{St, EndOfPeriodEnergyTargetOutage},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractPowerModel}
    devices = PSI.get_available_components(St, sys)

    if !PSI.validate_available_devices(St, devices)
        return
    end

    # Variables
    PSI.add_variables!(
        optimization_container,
        PSI.ActivePowerInVariable,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.ActivePowerOutVariable,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.ReactivePowerVariable,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.EnergyVariable,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )

    # Initial Conditions
    PSI.initial_conditions!(
        optimization_container,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )

    # Constraints
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerOutVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerInVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ReactivePowerVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    outage_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.feedforward!(optimization_container, devices, model, PSI.get_feedforward(model))

    # Energy Balanace limits
    PSI.energy_balance_constraint!(
        optimization_container,
        devices,
        model.formulation,
        S,
        PSI.get_feedforward(model),
    )

    # Cost Function
    PSI.cost_function!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )

    return
end

function PSI.construct_device!(
    optimization_container::PSI.OptimizationContainer,
    sys::PSY.System,
    model::PSI.DeviceModel{St, EndOfPeriodEnergyTargetOutage},
    ::Type{S},
) where {St <: PSY.Storage, S <: PM.AbstractActivePowerModel}
    devices = PSI.get_available_components(St, sys)

    if !PSI.validate_available_devices(St, devices)
        return
    end

    # Variables
    PSI.add_variables!(
        optimization_container,
        PSI.ActivePowerInVariable,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.ActivePowerOutVariable,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )
    PSI.add_variables!(
        optimization_container,
        PSI.EnergyVariable,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )

    # Initial Conditions
    PSI.initial_conditions!(
        optimization_container,
        devices,
        EndOfPeriodEnergyTargetOutage(),
    )

    # Constraints
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerOutVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.add_constraints!(
        optimization_container,
        PSI.RangeConstraint,
        PSI.ActivePowerInVariable,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.energy_capacity_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    outage_constraints!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )
    PSI.feedforward!(optimization_container, devices, model, PSI.get_feedforward(model))

    # Energy Balanace limits
    PSI.energy_balance_constraint!(
        optimization_container,
        devices,
        model.formulation,
        S,
        PSI.get_feedforward(model),
    )

    # Cost Function
    PSI.cost_function!(
        optimization_container,
        devices,
        model,
        S,
        PSI.get_feedforward(model),
    )

    return
end


using CUDA
using Adapt
using DataStructures

struct DummyApp{A,B}
    number_time_steps::Int32
    number_inputs::Int64
    number_outputs::Int64
    number_checkboxes::Int64
    number_gui_elements::Int64
    gui_elements_rectangles::A
    checkboxes_checked::B
end

function DummyApp(configuration::OrderedDict, number_individuals::Int)

    number_checkboxes = 8
    number_gui_elements = number_checkboxes + 1

    # Initialize Cuda Array for positions of GUI elements
    gui_elements_rectangles = CUDA.fill(0.0, (number_gui_elements, 4))

    checkboxes_size = 20
    checkboxes_grid_size = 50
    checkboxes_border = 30

    # Place all 8 checkboxes in 4 colums and 2 rows
    n = 1
    for i = 1:4
        for j = 1:2

            x = checkboxes_grid_size * i + checkboxes_border
            y = checkboxes_grid_size * j + checkboxes_border

            gui_elements_rectangles[n, :] = [x, y, checkboxes_size, checkboxes_size]
            n += 1
        end
    end

    # Place button (below checkboxes)
    gui_elements_rectangles[n, :] = [200, 200, 150, 50]

    DummyApp(
        convert(Int32, configuration["number_time_steps"]),
        number_checkboxes,
        2,
        number_checkboxes,
        number_gui_elements,
        gui_elements_rectangles,
        CUDA.fill(false, (number_checkboxes, number_individuals)),
    )
end

Adapt.@adapt_structure DummyApp

function kernel_eval_fitness(individuals, rewards, environment_seeds, number_rounds, brains, environments)

    threadID = threadIdx().x
    blockID = blockIdx().x

    offset_shared_memory = 0

    reward = @cuDynamicSharedMem(Float32, environments.number_gui_elements, offset_shared_memory)
    offset_shared_memory += sizeof(reward)

    sync_threads()

    observation = @cuDynamicSharedMem(Float32, environments.number_inputs, offset_shared_memory)
    offset_shared_memory += sizeof(observation)

    sync_threads()

    action = @cuDynamicSharedMem(Float32, environments.number_outputs, offset_shared_memory)
    offset_shared_memory += sizeof(action)

    sync_threads()

    click_position = @cuDynamicSharedMem(Int32, environments.number_outputs, offset_shared_memory)
    offset_shared_memory += sizeof(click_position)

    sync_threads()

    gui_elements_states = @cuDynamicSharedMem(Int32, environments.number_gui_elements, offset_shared_memory)
    offset_shared_memory += sizeof(gui_elements_states)

    sync_threads()

    # Initialize brains
    initialize(brains, individuals)

    sync_threads()

    # Initialize states of gui elements
    if threadID <= environments.number_gui_elements
        gui_elements_states[threadID] = 0
    end

    sync_threads()

    # Initialize rewards
    if threadID <= environments.number_gui_elements
        reward[threadID] = 0.0
    end

    sync_threads()

    # Iterate over given number of time steps
    for time_step = 1:environments.number_time_steps

        # Set observations
        if threadID <= environments.number_checkboxes
            observation[threadID] = gui_elements_states[threadID]
        end

        sync_threads()

        # Brain step
        step(brains, observation, action, offset_shared_memory)

        # Scale actions to click positions
        if threadID <= 2
            click_position[threadID] = trunc(0.5 * (action[threadID] + 1.0) * 400.0)
            #click_position[threadID] = rand(1:400)
        end

        sync_threads()

        #if blockID == 1 && threadID == 1
        #    @cuprintln("tx=", threadID,"   action_x=", action[1], "  action_y=", action[2])
        #end

        # Process mouse click
        process_click(environments, click_position, gui_elements_states, reward)

        sync_threads()

    end

    # Transfer rewards to output vector
    if threadID == 1
        rewards[blockID] = 0.0

        for i in 1:environments.number_gui_elements
            rewards[blockID] += reward[i]
        end
    end

    return
end

function get_memory_requirements(environments::DummyApp)
    return sizeof(Float32) * environments.number_gui_elements +                             # Reward
           sizeof(Float32) * (environments.number_inputs + environments.number_outputs) +   # Observation + Action
           sizeof(Int32) * environments.number_gui_elements                                 # States of gui elements

end

function get_number_inputs(environments::DummyApp)
    return environments.number_inputs
end

function get_number_outputs(environments::DummyApp)
    return environments.number_outputs
end

function process_click(environments, point, gui_elements_states, reward)

    threadID = threadIdx().x

    if threadID <= environments.number_gui_elements
        rect_x = environments.gui_elements_rectangles[threadID, 1]
        rect_y = environments.gui_elements_rectangles[threadID, 2]
        width = environments.gui_elements_rectangles[threadID, 3]
        height = environments.gui_elements_rectangles[threadID, 4]

        if is_point_in_rect(point[1], point[2], rect_x, rect_y, width, height)

            # Toggle checkboxes
            if threadID <= environments.number_checkboxes
                # Bitwise xor with 1 does toggling: https://stackoverflow.com/questions/11604409/how-to-toggle-a-boolean
                #gui_elements_states[threadID] ⊻= 1

                if gui_elements_states[threadID] == 0
                    reward[threadID] += 1.0
                    gui_elements_states[threadID] = 1
                end
            end
        end
    end

    sync_threads()

    #if blockID == 1 && threadID == 1
    #    @cuprintln("tx=", threadID,"   x=", point[1], "  y=", point[2])
    #end

    return
end

function is_point_in_rect(point_x, point_y, rect_x, rect_y, width, height)

    return rect_x <= point_x <= (rect_x + width) && rect_y <= point_y <= (rect_y + height)
end


#
# List of elements with properties
#
struct Element
 atomic_number::Int
 element::String
 pdb_name::String
 name::String
 mass::Float64
end
const elements = [
             Element( 0      ,  "X"  ,  "X"   ,   "NotFound"    ,  0.00000   )
             Element( 1      ,  "H"  ,  "H"   ,   "Hydrogen"    ,  1.00797   )
             Element( 2      ,  "He" ,  "HEL" ,   "Helium"      ,  4.00260   )
             Element( 3      ,  "Li" ,  "LI"  ,   "Lithium"     ,  6.941     )
             Element( 4      ,  "Be" ,  "BE"  ,   "Beryllium"   ,  9.01218   )
             Element( 5      ,  "B"  ,  "B"   ,   "Boron"       ,  10.81     )
             Element( 6      ,  "C"  ,  "C"   ,   "Carbon"      ,  12.011    )
             Element( 7      ,  "N"  ,  "N"   ,   "Nitrogen"    ,  14.0067   )
             Element( 8      ,  "O"  ,  "O"   ,   "Oxygen"      ,  15.9994   )
             Element( 9      ,  "F"  ,  "F"   ,   "Fluorine"    ,  18.998403 )
             Element( 10     ,  "Ne" ,  "NEO" ,   "Neon"        ,  20.179    )
             Element( 11     ,  "Na" ,  "NA"  ,   "Sodium"      ,  22.98977  )
             Element( 11     ,  "Na" ,  "SOD" ,   "Sodium"      ,  22.98977  )
             Element( 12     ,  "Mg" ,  "MG"  ,   "Magnesium"   ,  24.305    )
             Element( 13     ,  "Al" ,  "AL"  ,   "Aluminum"    ,  26.98154  )
             Element( 14     ,  "Si" ,  "SI"  ,   "Silicon"     ,  28.0855   )
             Element( 15     ,  "P"  ,  "P"   ,   "Phosphorus"  ,  30.97376  )
             Element( 16     ,  "S"  ,  "S"   ,   "Sulfur"      ,  32.06     )
             Element( 17     ,  "Cl" ,  "CL"  ,   "Chlorine"    ,  35.453    )
             Element( 17     ,  "Cl" ,  "CLA" ,   "Chlorine"    ,  35.453    )
             Element( 18     ,  "Ar" ,  "AR"  ,   "Argon"       ,  39.948    )
             Element( 19     ,  "K"  ,  "K"   ,   "Potassium"   ,  39.0983   )
             Element( 19     ,  "K"  ,  "POT" ,   "Potassium"   ,  39.0983   )
             Element( 20     ,  "Ca" ,  "CAL" ,   "Calcium"     ,  40.08     )
             Element( 21     ,  "Sc" ,  "SC"  ,   "Scandium"    ,  44.9559   )
             Element( 22     ,  "Ti" ,  "TI"  ,   "Titanium"    ,  47.90     )
             Element( 23     ,  "V"  ,  "V"   ,   "Vanadium"    ,  50.9415   )
             Element( 24     ,  "Cr" ,  "CR"  ,   "Chromium"    ,  51.996    )
             Element( 25     ,  "Mn" ,  "MN"  ,   "Manganese"   ,  54.9380   )
             Element( 26     ,  "Fe" ,  "FE"  ,   "Iron"        ,  55.847    )
             Element( 27     ,  "Co" ,  "CO"  ,   "Cobalt"      ,  58.9332   )
             Element( 28     ,  "Ni" ,  "NI"  ,   "Nickel"      ,  58.70     )
             Element( 29     ,  "Cu" ,  "CU"  ,   "Copper"      ,  63.546    )
             Element( 30     ,  "Zn" ,  "ZN"  ,   "Zinc"        ,  65.38     )
             Element( 31     ,  "Ga" ,  "GA"  ,   "Gallium"     ,  69.72     )
             Element( 32     ,  "Ge" ,  "GE"  ,   "Germanium"   ,  72.59     )
             Element( 33     ,  "As" ,  "AS"  ,   "Arsenic"     ,  74.9216   )
             Element( 34     ,  "Se" ,  "SE"  ,   "Selenium"    ,  78.96     )
             Element( 35     ,  "Br" ,  "BR"  ,   "Bromine"     ,  79.904    )
             Element( 36     ,  "Kr" ,  "KR"  ,   "Krypton"     ,  83.80     )
             Element( 37     ,  "Rb" ,  "RB"  ,   "Rubidium"    ,  85.4678   )
             Element( 38     ,  "Sr" ,  "SR"  ,   "Strontium"   ,  87.62     )
             Element( 39     ,  "Y"  ,  "Y"   ,   "Yttrium"     ,  88.9059   )
             Element( 40     ,  "Zr" ,  "ZR"  ,   "Zirconium"   ,  91.22     )
             Element( 41     ,  "Nb" ,  "NB"  ,   "Niobium"     ,  92.9064   )
             Element( 42     ,  "Mo" ,  "MO"  ,   "Molybdenum"  ,  95.94     )
             Element( 43     ,  "Tc" ,  "TC"  ,   "Technetium"  ,  98.       )
             Element( 44     ,  "Ru" ,  "RU"  ,   "Ruthenium"   ,  101.07    )
             Element( 45     ,  "Rh" ,  "RH"  ,   "Rhodium"     ,  102.9055  )
             Element( 46     ,  "Pd" ,  "PD"  ,   "Palladium"   ,  106.4     )
             Element( 47     ,  "Ag" ,  "AG"  ,   "Silver"      ,  107.868   )
             Element( 48     ,  "Cd" ,  "CAD" ,   "Cadmium"     ,  112.41    )
             Element( 49     ,  "In" ,  "IN"  ,   "Indium"      ,  114.82    )
             Element( 50     ,  "Sn" ,  "SN"  ,   "Tin"         ,  118.69    )
             Element( 51     ,  "Sb" ,  "SB"  ,   "Antimony"    ,  121.75    )
             Element( 52     ,  "Te" ,  "TE"  ,   "Tellurium"   ,  127.60    )
             Element( 53     ,  "I"  ,  "I"   ,   "Iodine"      ,  126.9045  )
             Element( 54     ,  "Xe" ,  "XE"  ,   "Xenon"       ,  131.30    )
             Element( 55     ,  "Cs" ,  "CES" ,   "Cesium"      ,  132.9054  )
             Element( 56     ,  "Ba" ,  "BA"  ,   "Barium"      ,  137.33    )
             Element( 57     ,  "La" ,  "LA"  ,   "Lanthanum"   ,  138.9055  )
             Element( 58     ,  "Ce" ,  "CER" ,   "Cerium"      ,  140.12    )
             Element( 59     ,  "Pr" ,  "PR"  ,   "Praseodymium",  140.9077  )
             Element( 60     ,  "Nd" ,  "ND"  ,   "Neodymium"   ,  144.24    )
             Element( 61     ,  "Pm" ,  "PM"  ,   "Promethium"  ,  145       )
             Element( 62     ,  "Sm" ,  "SM"  ,   "Samarium"    ,  150.4     )
             Element( 63     ,  "Eu" ,  "EU"  ,   "Europium"    ,  151.96    )
             Element( 64     ,  "Gd" ,  "GD"  ,   "Gadolinium"  ,  157.25    )
             Element( 65     ,  "Tb" ,  "TB"  ,   "Terbium"     ,  158.9254  )
             Element( 66     ,  "Dy" ,  "DY"  ,   "Dysprosium"  ,  162.50    )
             Element( 67     ,  "Ho" ,  "HLM" ,   "Holmium"     ,  164.9304  )
             Element( 68     ,  "Er" ,  "ER"  ,   "Erbium"      ,  167.26    )
             Element( 69     ,  "Tm" ,  "TM"  ,   "Thulium"     ,  168.9342  )
             Element( 70     ,  "Yb" ,  "YB"  ,   "Ytterbium"   ,  173.04    )
             Element( 71     ,  "Lu" ,  "LU"  ,   "Lutetium"    ,  174.967   )
             Element( 72     ,  "Hf" ,  "HAF" ,   "Hafnium"     ,  178.49    )
             Element( 73     ,  "Ta" ,  "TA"  ,   "Tantalum"    ,  180.9479  )
             Element( 74     ,  "W"  ,  "W"   ,   "Tungsten"    ,  183.85    )
             Element( 75     ,  "Re" ,  "RE"  ,   "Rhenium"     ,  186.207   )
             Element( 76     ,  "Os" ,  "OS"  ,   "Osmium"      ,  190.2     )
             Element( 77     ,  "Ir" ,  "IR"  ,   "Iridium"     ,  192.22    )
             Element( 78     ,  "Pt" ,  "PT"  ,   "Platinum"    ,  195.09    )
             Element( 79     ,  "Au" ,  "AU"  ,   "Gold"        ,  196.9665  )
             Element( 80     ,  "Hg" ,  "MER" ,   "Mercury"     ,  200.59    )
             Element( 81     ,  "Tl" ,  "TL"  ,   "Thallium"    ,  204.37    )
             Element( 82     ,  "Pb" ,  "PB"  ,   "Lead"        ,  207.2     )
             Element( 83     ,  "Bi" ,  "BI"  ,   "Bismuth"     ,  208.9804  )
             Element( 84     ,  "Po" ,  "PO"  ,   "Polonium"    ,  209       )
             Element( 85     ,  "At" ,  "AT"  ,   "Astatine"    ,  210       )
             Element( 86     ,  "Rn" ,  "RN"  ,   "Radon"       ,  222       )
             Element( 87     ,  "Fr" ,  "FR"  ,   "Francium"    ,  223       )
             Element( 88     ,  "Ra" ,  "RA"  ,   "Radium"      ,  226.0254  )
             Element( 89     ,  "Ac" ,  "AC"  ,   "Actinium"    ,  227.0278  )
             Element( 90     ,  "Th" ,  "TH"  ,   "Thorium"     ,  232.0381  )
             Element( 91     ,  "Pa" ,  "PA"  ,   "Protactinium",  231.0359  )
             Element( 92     ,  "U"  ,  "U"   ,   "Uranium"     ,  238.029   )
           ]

#
# Retrive index of element in elements list from name. Returns 1 (element "X" of list) if not found
#
function element_index(name::String)
  # Try to find if there is any exact match
  i = findfirst(el -> (el.element == name ||
                       el.name == name ||
                       el.pdb_name == name ), elements)
  (i != nothing) && return i
  #
  # Try to find by PDB name, note that NT2 and 2NT2 must match N, for example
  #
  # First, check if the first char is a number
  i0 = try
    parse(Int,name[1:1])
    2
  catch
    1
  end
  # Now check if the start of this name matches some PDB code
  for (iel,el) in pairs(elements)
    lpdb = length(el.pdb_name)
    # If the name is shorter than the PDB name, no match
    (length(name) - i0 + 1) < lpdb && continue
    if name[i0:i0+lpdb-1] == el.pdb_name
      return iel
    end
  end
  return 1
end


#! format: off

abstract type AbstractAGCFormulation <: AbstractServiceFormulation end
struct PIDSmoothACE <: AbstractAGCFormulation end

"""
Steady State deviation of the frequency
"""
function add_variables!(optimization_container::OptimizationContainer, ::Type{SteadyStateFrequencyDeviation})
    variable_name = make_variable_name(SteadyStateFrequencyDeviation)
    time_steps = model_time_steps(optimization_container)
    variable = add_var_container!(optimization_container, variable_name, time_steps)
    for t in time_steps
        variable[t] = JuMP.@variable(optimization_container.JuMPmodel, base_name = "ΔF_{$(t)}")
    end
end

get_variable_sign(_, ::Type{<:PSY.Area}, ::AbstractAGCFormulation) = NaN
########################## ActivePowerVariable, Area ###########################

get_variable_binary(::ActivePowerVariable, ::Type{<:PSY.Area}, ::AbstractAGCFormulation) = false

########################## SmoothACE, AggregationTopology ###########################

get_variable_binary(::SmoothACE, ::Type{<:PSY.AggregationTopology}, ::AbstractAGCFormulation) = false

########################## DeltaActivePowerUpVariable, Area ###########################

get_variable_binary(::DeltaActivePowerUpVariable, ::Type{<:PSY.Area}, ::AbstractAGCFormulation) = false
get_variable_lower_bound(::DeltaActivePowerUpVariable, ::PSY.Area, ::AbstractAGCFormulation) = 0.0

########################## DeltaActivePowerDownVariable, Area ###########################

get_variable_binary(::DeltaActivePowerDownVariable, ::Type{<:PSY.Area}, ::AbstractAGCFormulation) = false
get_variable_lower_bound(::DeltaActivePowerDownVariable, ::PSY.Area, ::AbstractAGCFormulation) = 0.0

########################## AdditionalDeltaPowerUpVariable, Area ###########################

# Commented out since not in use. These functions will substitute balancing_auxiliary_variables!
# """
# This function add the upwards scheduled regulation variables for power generation output to the model
# """
# function AddVariableSpec(
#     ::Type{AdditionalDeltaActivePowerUpVariable},
#     ::Type{U},
#     optimization_container::OptimizationContainer,
# ) where {U <: PSY.Area}
#     return AddVariableSpec(;
#         variable_name = make_variable_name(AdditionalDeltaActivePowerUpVariable, U),
#         binary = false,
#         lb_value_func = x -> 0.0,
#         expression_name = :emergency_up,
#     )
# end
#
# """
# This function add the variables for power generation output to the model
# """
# function AddVariableSpec(
#     ::Type{AdditionalDeltaActivePowerDownVariable},
#     ::Type{U},
#     optimization_container::OptimizationContainer,
# ) where {U <: PSY.Area}
#     return AddVariableSpec(;
#         variable_name = make_variable_name(AdditionalDeltaActivePowerDownVariable, U),
#         binary = false,
#         lb_value_func = x -> 0.0,
#         expression_name = :emergency_dn,
#     )
# end

########################## AreaMismatchVariable, Area ###########################

make_variable_name(::Type{AreaMismatchVariable}, _) = make_variable_name(AreaMismatchVariable)
get_variable_binary(::AreaMismatchVariable, ::Type{<:PSY.Area}, ::AbstractAGCFormulation) = false

########################## LiftVariable, Area ###########################

make_variable_name(::Type{LiftVariable}, _) = make_variable_name(LiftVariable)
get_variable_binary(::LiftVariable, ::Type{<:PSY.Area}, ::AbstractAGCFormulation) = false
get_variable_lower_bound(::LiftVariable, ::PSY.Area, ::AbstractAGCFormulation) = 0.0

#! format: off

########################## Initial Condition ###########################
function area_control_initial_condition!(
    optimization_container::OptimizationContainer,
    services::Vector{PSY.AGC},
    ::D,
) where {D <: AbstractAGCFormulation}
    key = ICKey(AreaControlError, PSY.AGC)
    _make_initial_conditions!(
        optimization_container,
        services,
        D(),
        nothing,
        key,
        _make_initial_condition_area_control,
        _get_variable_initial_value,
        # Doesn't require Cache
    )

    return
end

function _get_variable_initial_value(
    d::PSY.Component,
    key::ICKey,
    ::AbstractAGCFormulation,
    ::Nothing,
)
    return _get_ace_error(d, key)
end

########################## , ###########################

function balancing_auxiliary_variables!(optimization_container, sys)
    area_names = [PSY.get_name(a) for a in PSY.get_components(PSY.Area, sys)]
    time_steps = model_time_steps(optimization_container)
    R_up_emergency = JuMPVariableArray(undef, area_names, time_steps)
    R_dn_emergency = JuMPVariableArray(undef, area_names, time_steps)
    assign_variable!(
        optimization_container,
        make_variable_name(AdditionalDeltaActivePowerUpVariable, PSY.Area),
        R_up_emergency,
    )
    assign_variable!(
        optimization_container,
        make_variable_name(AdditionalDeltaActivePowerDownVariable, PSY.Area),
        R_dn_emergency,
    )
    emergency_up =
        add_expression_container!(optimization_container, :emergency_up, area_names, time_steps)
    emergency_dn =
        add_expression_container!(optimization_container, :emergency_dn, area_names, time_steps)
    for t in time_steps, a in area_names
        R_up_emergency[a, t] = JuMP.@variable(
            optimization_container.JuMPmodel,
            base_name = "Re_up_{$(a),$(t)}",
            lower_bound = 0.0
        )
        emergency_up[a, t] = R_up_emergency[a, t] + 0.0
        R_dn_emergency[a, t] = JuMP.@variable(
            optimization_container.JuMPmodel,
            base_name = "Re_dn_{$(a),$(t)}",
            lower_bound = 0.0
        )
        emergency_dn[a, t] = R_dn_emergency[a, t] + 0.0
    end

    return
end

function absolute_value_lift(optimization_container::OptimizationContainer, areas)
    time_steps = model_time_steps(optimization_container)
    area_names = [PSY.get_name(a) for a in areas]
    container_lb = JuMPConstraintArray(undef, area_names, time_steps)
    assign_constraint!(optimization_container, "absolute_value_lb", container_lb)
    container_ub = JuMPConstraintArray(undef, area_names, time_steps)
    assign_constraint!(optimization_container, "absolute_value_ub", container_ub)
    mismatch = get_variable(optimization_container, :area_mismatch)
    z = get_variable(optimization_container, :lift)

    for t in time_steps, a in area_names
        container_lb[a, t] =
            JuMP.@constraint(optimization_container.JuMPmodel, mismatch[a, t] <= z[a, t])
        container_ub[a, t] =
            JuMP.@constraint(optimization_container.JuMPmodel, -1 * mismatch[a, t] <= z[a, t])
    end

    JuMP.add_to_expression!(
        optimization_container.cost_function,
        sum(z[a, t] for t in time_steps, a in area_names) * SERVICES_SLACK_COST,
    )
    return
end

"""
Expression for the power deviation given deviation in the frequency. This expression allows updating the response of the frequency depending on commitment decisions
"""
function frequency_response_constraint!(optimization_container::OptimizationContainer, sys::PSY.System)
    time_steps = model_time_steps(optimization_container)
    services = PSY.get_components(PSY.AGC, sys)
    area_names = [PSY.get_name(PSY.get_area(s)) for s in services]
    frequency_response = 0.0
    for area in PSY.get_components(PSY.Area, sys)
        frequency_response += PSY.get_load_response(area)
    end
    for g in PSY.get_components(PSY.RegulationDevice, sys, x -> PSY.get_available(x))
        d = PSY.get_droop(g)
        response = 1 / d
        frequency_response += response
    end

    @assert frequency_response >= 0.0
    # This value is the one updated later in simulation based on the UC result
    inv_frequency_reponse = 1 / frequency_response
    area_balance = get_variable(optimization_container, ActivePowerVariable, PSY.Area)
    frequency = get_variable(optimization_container, "Δf")
    R_up = get_variable(optimization_container, DeltaActivePowerUpVariable, PSY.Area)
    R_dn = get_variable(optimization_container, DeltaActivePowerDownVariable, PSY.Area)
    R_up_emergency =
        get_variable(optimization_container, AdditionalDeltaActivePowerUpVariable, PSY.Area)
    R_dn_emergency =
        get_variable(optimization_container, AdditionalDeltaActivePowerUpVariable, PSY.Area)

    container = JuMPConstraintArray(undef, time_steps)
    assign_constraint!(optimization_container, "frequency_response", container)

    for s in services, t in time_steps
        system_balance = sum(area_balance.data[:, t])
        total_reg = JuMP.AffExpr(0.0)
        for a in area_names
            JuMP.add_to_expression!(total_reg, R_up[a, t])
            JuMP.add_to_expression!(total_reg, -1 * R_dn[a, t])
            JuMP.add_to_expression!(total_reg, R_up_emergency[a, t])
            JuMP.add_to_expression!(total_reg, -1 * R_dn_emergency[a, t])
        end
        container[t] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            frequency[t] == -inv_frequency_reponse * (system_balance + total_reg)
        )
    end
    return
end

function smooth_ace_pid!(optimization_container::OptimizationContainer, services::Vector{PSY.AGC})
    time_steps = model_time_steps(optimization_container)
    area_names = [PSY.get_name(PSY.get_area(s)) for s in services]
    RAW_ACE = add_expression_container!(optimization_container, :RAW_ACE, area_names, time_steps)
    SACE = get_variable(optimization_container, SmoothACE, PSY.Area)
    SACE_pid = JuMPConstraintArray(undef, area_names, time_steps)
    assign_constraint!(optimization_container, "SACE_pid", SACE_pid)

    Δf = get_variable(optimization_container, make_variable_name("Δf"))

    for (ix, service) in enumerate(services)
        kp = PSY.get_K_p(service)
        ki = PSY.get_K_i(service)
        kd = PSY.get_K_d(service)
        B = PSY.get_bias(service)
        Δt = convert(Dates.Second, optimization_container.resolution).value
        a = PSY.get_name(PSY.get_area(service))
        for t in time_steps
            # Todo: Add initial Frequency Deviation
            RAW_ACE[a, t] = -10 * B * Δf[t] + 0.0
            SACE[a, t] =
                JuMP.@variable(optimization_container.JuMPmodel, base_name = "SACE_{$(a),$(t)}")
            if t == 1
                SACE_ini =
                    get_initial_conditions(optimization_container, ICKey(AreaControlError, PSY.AGC))[ix]
                sace_exp = SACE_ini.value + kp * ((1 + Δt / (kp / ki)) * (RAW_ACE[a, t]))
                SACE_pid[a, t] =
                    JuMP.@constraint(optimization_container.JuMPmodel, SACE[a, t] == sace_exp)
                continue
            end
            SACE_pid[a, t] = JuMP.@constraint(
                optimization_container.JuMPmodel,
                SACE[a, t] ==
                SACE[a, t - 1] +
                kp * (
                    (1 + Δt / (kp / ki) + (kd / kp) / Δt) * (RAW_ACE[a, t]) +
                    (-1 - 2 * (kd / kp) / Δt) * (RAW_ACE[a, t - 1])
                )
            )
        end
    end
    return
end

function aux_constraints!(optimization_container::OptimizationContainer, sys::PSY.System)
    time_steps = model_time_steps(optimization_container)
    area_names = [PSY.get_name(a) for a in PSY.get_components(PSY.Area, sys)]
    aux_equation = JuMPConstraintArray(undef, area_names, time_steps)
    assign_constraint!(optimization_container, "balance_aux", aux_equation)
    area_mismatch = get_variable(optimization_container, :area_mismatch)
    SACE = get_variable(optimization_container, SmoothACE, PSY.Area)
    R_up = get_variable(optimization_container, DeltaActivePowerUpVariable, PSY.Area)
    R_dn = get_variable(optimization_container, DeltaActivePowerDownVariable, PSY.Area)
    R_up_emergency =
        get_variable(optimization_container, AdditionalDeltaActivePowerUpVariable, PSY.Area)
    R_dn_emergency =
        get_variable(optimization_container, AdditionalDeltaActivePowerUpVariable, PSY.Area)

    for t in time_steps, a in area_names
        aux_equation[a, t] = JuMP.@constraint(
            optimization_container.JuMPmodel,
            -1 * SACE[a, t] ==
            (R_up[a, t] - R_dn[a, t]) +
            (R_up_emergency[a, t] - R_dn_emergency[a, t]) +
            area_mismatch[a, t]
        )
    end
    return
end


import HTTP

export method, GET, route, page, probability, query

# Request type

method(m::AbstractString, app...) = branch(req -> req[:method] == m, app...)
method(ms, app...) = branch(req -> req[:method] in ms, app...)

GET(app...) = method("GET", app...)

# Path routing

splitpath(p::AbstractString) = split(p, "/", keepempty=false)
splitpath(p) = p

function matchpath(target, path)
  length(target) > length(path) && return
  params = d()
  for i = 1:length(target)
    if startswith(target[i], ":")
      params[Symbol(target[i][2:end])] = path[i]
    else
      target[i] == path[i] || return
    end
  end
  return params
end

function matchpath!(target, req)
  ps = matchpath(target, req[:path])
  ps == nothing && return false
  merge!(params!(req), ps)
  splice!(req[:path], 1:length(target))
  return true
end

route(p::Vector, app...) = branch(req -> matchpath!(p, req), app...)
route(p::AbstractString, app...) = route(splitpath(p), app...)
route(app...) = route([], app...)
route(app::Base.Callable, p) = route(p, app)
route(app1::Base.Callable, app2::Base.Callable) = route([], app1, app2)

page(p::Vector, app...) = branch(req -> length(p) == length(req[:path]) && matchpath!(p, req), app...)
page(p::AbstractString, app...) = page(splitpath(p), app...)
page(app...) = page([], app...)
page(app::Base.Callable, p) = page(p, app)
page(app1::Base.Callable, app2::Base.Callable) = page([], app1, app2)

# Query routing

function matchquery(q, req)
  qdict = HTTP.URIs.queryparams(req[:query])
  length(q) != length(qdict) && return false
  for (key, value) in q
    if haskey(qdict, key) && (value == "" || value == qdict[key])
      continue
    else
      return false
    end
  end
  return true
end

query(q::Dict{<:AbstractString, <:AbstractString}, app...) =
  branch(req -> matchquery(q, req), app...)

# Misc

probability(x, app...) = branch(_->rand()<x, app...)
# Old typo
@deprecate probabilty(x, app...) probability(x, app...)


~n=n|>last


# Similar to JuMP/test/solvers.jl

function try_import(name::Symbol)
    try
        @eval import $name
        return true
    catch e
        return false
    end
end

# Load an available solver
lp_solver = nothing
glp = try_import(:GLPKMathProgInterface)
glp && (lp_solver = GLPKMathProgInterface.GLPKSolverLP())
if lp_solver === nothing
    cbc = try_import(:Cbc)
    if cbc; import Clp; end
    cbc && (lp_solver = Clp.ClpSolver())
end
if lp_solver === nothing
    grb = try_import(:Gurobi) # Gurobi creates BigFloat when the input is BigInt and then cannot handle it
    grb && (lp_solver = Gurobi.GurobiSolver(OutputFlag=0))
end
if lp_solver === nothing
    cpx = try_import(:CPLEX)
end
if lp_solver === nothing
    xpr = try_import(:Xpress)
    xpr && (lp_solver = Xpress.XpressSolver(OUTPUTLOG=0))
end
if lp_solver === nothing
    mos = try_import(:Mosek)
    mos && (lp_solver = Mosek.MosekSolver(LOG=0))
end
if lp_solver === nothing
    ipt = try_import(:Ipopt)
    ipt && (lp_solver = Ipopt.IpoptSolver(print_level=0))
end
if lp_solver === nothing
    eco = try_import(:ECOS)
    eco && (lp_solver = ECOS.ECOSSolver(verbose=false))
end
if lp_solver === nothing
    scs = try_import(:SCS)
    scs && (lp_solver = SCS.SCSSolver(eps=1e-6,verbose=0))
    @assert lp_solver !== nothing
end