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

using Markdown
using InteractiveUtils

# ╔═╡ 9f0add20-11db-11ec-31c8-1933ca9b19c4
begin
	using MortalityTables
	using DataFrames
	using IterTools
	using BenchmarkTools
end

# ╔═╡ 1ed5e587-d0eb-4dc2-bedc-00d13f258ecd
begin
	using PlutoUI
	TableOfContents()
end

# ╔═╡ d8fa5c6d-4fd3-40b6-abb3-6349b9fca7b7
md"# MortalityTables.jl Design notes"

# ╔═╡ 4980df5b-d2aa-4200-b517-bffcf4775971
md"## MortalityTables.jl types"

# ╔═╡ 5455ba9e-11ef-49e4-b877-d12dc45a22f8
md"There are two types of MortalityTable subtypes for SOA data: `SelectUltimateTable` and `UltimateTable`, so define a function for each:"

# ╔═╡ 1242d939-581b-4236-b00c-608efe9a4d08
vbt = MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Male Nonsmoker, ANB")

# ╔═╡ 2181f078-f538-48b4-a2a3-d7e6c9678ce7
cso = MortalityTables.table("1941 CSO Basic Table, ANB")

# ╔═╡ fcffeb28-ed5d-44f9-9260-f9b631d6988d
vbt isa MortalityTables.SelectUltimateTable

# ╔═╡ 9a9aa632-f773-49ef-bd85-a64148d1d0a4
cso isa MortalityTables.UltimateTable

# ╔═╡ 95b3a432-ad19-47c8-8cf3-831fab934425
md"""
## Candidate function for getting the "long" version of dataframes.
"""

# ╔═╡ 215585ab-d2c3-47e7-b54f-e2719bdeb230
"""
	long(m::MortalityTable)

Return an array of tuples containing `issue_age`,`attained_age`,`duration`,`select` rate, and `ultimate` rate.

"""
function long(m::MortalityTables.SelectUltimateTable)
	earliest_age = min(firstindex(m.select),firstindex(m.ultimate))
	last_age = max(lastindex(m.select),lastindex(m.ultimate))

	table = map(earliest_age:last_age) do issue_age
		map(issue_age:last_age) do attained_age
			# use `get` to provide default missing value
			ultimate = get(m.ultimate,attained_age,missing)
			if issue_age <= lastindex(m.select)
				select = get(m.select[issue_age],attained_age,missing)
			else
				select = missing
			end
			duration = attained_age - issue_age + 1
			(;issue_age,attained_age, duration, select,ultimate)
		end
	end

	vcat(table...)

end

# ╔═╡ 99a1a535-87de-4fea-bb83-2bcc0267d70f
"""
	long(m::MortalityTable)

Return an array of tuples containing `issue_age`,`attained_age`,`duration`,`select` rate, and `ultimate` rate.

"""
function long(m::MortalityTables.UltimateTable)
	earliest_age = firstindex(m.ultimate)
	last_age = lastindex(m.ultimate)

	table = map(earliest_age:last_age) do issue_age
		map(issue_age:last_age) do attained_age
			# use `get` to provide default missing value
			ultimate = get(m.ultimate,attained_age,missing)
			select = missing
			duration = attained_age - issue_age + 1
			(;issue_age,attained_age, duration, select,ultimate)
		end
	end

	vcat(table...)

end

# ╔═╡ c155fb29-275d-4b56-b125-ca0084c4182e
md"Now this can be called on tabular mortality tables and it will return a vector of named tuples, a very efficient and transferrable structure."

# ╔═╡ 73173e80-d76d-4819-a3e8-dc7138d532ac
long(vbt)

# ╔═╡ 6ead1870-0e5c-44e6-92b7-3e6f90c9e387
md"
**Why use a vector of tuples and not straight to DataFrame?**

In Julia, composability and package interoperability is really easy because packages don't need to know about, or can extend when necessary, other packages' types. See https://www.youtube.com/watch?v=kc9HwsxE1OY for more.

DataFrames can then directly use this result, but an array of tuples is also used by anything that uses the `Tables.jl` interface (e.g. all of these: https://github.com/JuliaData/Tables.jl/blob/main/INTEGRATIONS.md)"

# ╔═╡ 38a9ea42-1c9b-48c0-8d08-3acc00319223
DataFrame(long(vbt))

# ╔═╡ 8f999e94-bfc6-4b08-be35-1da44953dc55
md"## More on the design of MortalityTables.jl"

# ╔═╡ f70c9f2c-d24c-441f-bf24-863c5ba596dd
md"
The design went through many iterations, and it may not be the final design, but here are some of the benefits of the current design:"

# ╔═╡ 9b43e322-a873-421b-8260-d4368d6cbaa3
md"### Size in memory

Here, the long format is 5x the size of the regular type because duplicate values not needed:"

# ╔═╡ bd1a6dcf-0d91-4774-a267-e19abb8dc0e4
Base.summarysize(long(vbt))

# ╔═╡ 3f31a05a-0a5b-45f6-a718-f43744758472
Base.summarysize(vbt)

# ╔═╡ 09c00dd9-7a2a-478c-bb4f-9a1289f4796e
md"### Encourage data-driven design

Probably best explained with an example. In the following example, we calculate the ultimate survivorship for every age in the `cso` table from above. The table itself gives us the right range of ages to do this with. No need to assume a 100 or 121 omega age and handle edge cases when working with tables that don't line up to the usual expectations."

# ╔═╡ 31b42e8e-9361-45d1-b84a-2b8b54153ec4
map(eachindex(cso.ultimate)) do issue_age
	survival(cso.ultimate,issue_age,omega(cso.ultimate))
end

# ╔═╡ d2c825f3-7203-4905-828a-9be0fd42fe41
md"Of course, here we are limited by the quality of the input tables. For example, in the `vbt` table from the SOA loaded above, select rates are missing from under-18s, even though select rates are defined for a life issued under-18."

# ╔═╡ afa8ae7e-5beb-45fd-87cd-33d1f263b692
md"### Speed of calculation

In modeling contexts, efficiency is often a key part of the model. This is probably close to the fastest computation of survival probabilities possible to acheive:"

# ╔═╡ 03fb84cf-8fe9-4622-b7a3-6b8ff49907c4
@benchmark survival($cso.ultimate,55,100)

# ╔═╡ 01a5d74d-c71b-49f2-833f-1a937cbc89e8
md"### Parametric models

With MortalityTables.jl, virtually transparent to the user of the pacakge, you can drop in parametric models in place of tabular-SOA-like tables. Here is a Makeham mortality model and the same syntax is used as with the SOA tables:

"

# ╔═╡ f84c087d-8963-4835-81d9-8e72ab297b13
survival(MortalityTables.Makeham(),25,50) # 25 to 50 year old survival

# ╔═╡ b5dc0d67-f80b-40ae-a65c-7d8f085e84ae
md"### Table metadata

When you display a table, you get to see related metadata, which can be lost if simply parsing into a matrix or dataframe:"

# ╔═╡ 05a54094-1587-43eb-a3f2-d8945803f482
cso

# ╔═╡ 9b3425d6-0c3a-4395-a2c2-af73928aed01
md"### Parital Year assumptions

Built in are partial year assumptions, making it easy to use more realistic assumptions like Uniform death distribution:"

# ╔═╡ 04162e25-9a74-4bc1-9ceb-66655111897e
survival(cso.ultimate,55,55.5,MortalityTables.Uniform())

# ╔═╡ a48efa0b-3c2b-44a9-b808-cf0975fce72d
survival(cso.ultimate,55,55.5,MortalityTables.Constant())

# ╔═╡ 841a1fb5-fc52-494c-a2eb-15930b41ced8
survival(cso.ultimate,55,55.5,MortalityTables.Balducci())

# ╔═╡ fbbc5e91-3e43-43c8-b23e-0ccd00ba0c71
md"## Where MortalityTables.jl falls short"

# ╔═╡ c74af2e0-99eb-4526-a68a-8fdab5717f16
md"
In no particular order:
- it mainly relies on data from mort.SOA.org. Some of the tables have issues like inconsistent naming, missing fields
- would be nice to offer 'sets' of table to users. Early tests failed to do this in an automated way because of issues with the naming consistency of tables mentioned in the first bullet
- (currently) no automatic conversion to a dataframe-compatible, though this is discussed above
- MortalityTables.jl is a misnomer, as the `ORT` of MORT.SOA.org is 'other rate tables' which are offered in MortalityTables.jl
- no built-in support for mortality improvement
- has issues with autodifferentiation because of the discrete and non-continuous nature of the annual rate tables
"

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
MortalityTables = "4780e19d-04b9-53dc-86c2-9e9aa59b5a12"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
BenchmarkTools = "~1.1.4"
DataFrames = "~1.2.2"
IterTools = "~1.3.0"
MortalityTables = "~2.1.4"
PlutoUI = "~0.7.9"
"""

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

[[Adapt]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "84918055d15b3114ede17ac6a7182f68870c16f7"
uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
version = "3.3.1"

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

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

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

[[BenchmarkTools]]
deps = ["JSON", "Logging", "Printf", "Statistics", "UUIDs"]
git-tree-sha1 = "42ac5e523869a84eac9669eaceed9e4aa0e1587b"
uuid = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
version = "1.1.4"

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

[[Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[DataAPI]]
git-tree-sha1 = "bec2532f8adb82005476c141ec23e921fc20971b"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.8.0"

[[DataFrames]]
deps = ["Compat", "DataAPI", "Future", "InvertedIndices", "IteratorInterfaceExtensions", "LinearAlgebra", "Markdown", "Missings", "PooledArrays", "PrettyTables", "Printf", "REPL", "Reexport", "SortingAlgorithms", "Statistics", "TableTraits", "Tables", "Unicode"]
git-tree-sha1 = "d785f42445b63fc86caa08bb9a9351008be9b765"
uuid = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
version = "1.2.2"

[[DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "7d9d316f04214f7efdbb6398d545446e246eff02"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.10"

[[DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

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

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

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

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

[[EzXML]]
deps = ["Printf", "XML2_jll"]
git-tree-sha1 = "0fa3b52a04a4e210aeb1626def9c90df3ae65268"
uuid = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615"
version = "1.1.0"

[[Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[Future]]
deps = ["Random"]
uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"

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

[[InvertedIndices]]
git-tree-sha1 = "bee5f1ef5bf65df56bdd2e40447590b272a5471f"
uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
version = "1.1.0"

[[IterTools]]
git-tree-sha1 = "05110a2ab1fc5f932622ffea2a003221f4782c18"
uuid = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
version = "1.3.0"

[[IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "642a199af8b68253517b80bd3bfd17eb4e84df6e"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.3.0"

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

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

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

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

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

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

[[Libiconv_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "42b62845d70a619f063a7da093d995ec8e15e778"
uuid = "94ce4f54-9a6c-5748-9c1c-f9c7231a4531"
version = "1.16.1+1"

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

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

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

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

[[Missings]]
deps = ["DataAPI"]
git-tree-sha1 = "2ca267b08821e86c5ef4376cffed98a46c2cb205"
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
version = "1.0.1"

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

[[MortalityTables]]
deps = ["OffsetArrays", "Parsers", "Pkg", "QuadGK", "Requires", "UnPack", "XMLDict"]
git-tree-sha1 = "cae9281196e40e24cb5885b8463a6f955ae0f3eb"
uuid = "4780e19d-04b9-53dc-86c2-9e9aa59b5a12"
version = "2.1.4"

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

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

[[OffsetArrays]]
deps = ["Adapt"]
git-tree-sha1 = "c870a0d713b51e4b49be6432eff0e26a4325afee"
uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
version = "1.10.6"

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

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "438d35d2d95ae2c5e8780b330592b6de8494e779"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.0.3"

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

[[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"

[[PooledArrays]]
deps = ["DataAPI", "Future"]
git-tree-sha1 = "a193d6ad9c45ada72c14b731a318bedd3c2f00cf"
uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
version = "1.3.0"

[[Preferences]]
deps = ["TOML"]
git-tree-sha1 = "00cfd92944ca9c760982747e9a1d0d5d86ab1e5a"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.2"

[[PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "0d1245a357cc61c8cd61934c07447aa569ff22e6"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.1.0"

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

[[QuadGK]]
deps = ["DataStructures", "LinearAlgebra"]
git-tree-sha1 = "12fbe86da16df6679be7521dfb39fbc861e1dc7b"
uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
version = "2.4.1"

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

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

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

[[Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "4036a3bd08ac7e968e27c203d45f5fff15020621"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.1.3"

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

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

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

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

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

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

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

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

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

[[TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "TableTraits", "Test"]
git-tree-sha1 = "368d04a820fe069f9080ff1b432147a6203c3c89"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.5.1"

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

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

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

[[UnPack]]
git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b"
uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
version = "1.0.2"

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

[[XML2_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Pkg", "Zlib_jll"]
git-tree-sha1 = "1acf5bdf07aa0907e0a37d3718bb88d4b687b74a"
uuid = "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a"
version = "2.9.12+0"

[[XMLDict]]
deps = ["EzXML", "IterTools", "OrderedCollections"]
git-tree-sha1 = "d9a3faf078210e477b291c79117676fca54da9dd"
uuid = "228000da-037f-5747-90a9-8195ccbf91a5"
version = "0.4.1"

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

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

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

# ╔═╡ Cell order:
# ╠═d8fa5c6d-4fd3-40b6-abb3-6349b9fca7b7
# ╠═9f0add20-11db-11ec-31c8-1933ca9b19c4
# ╟─4980df5b-d2aa-4200-b517-bffcf4775971
# ╟─5455ba9e-11ef-49e4-b877-d12dc45a22f8
# ╠═1242d939-581b-4236-b00c-608efe9a4d08
# ╠═2181f078-f538-48b4-a2a3-d7e6c9678ce7
# ╠═fcffeb28-ed5d-44f9-9260-f9b631d6988d
# ╠═9a9aa632-f773-49ef-bd85-a64148d1d0a4
# ╟─95b3a432-ad19-47c8-8cf3-831fab934425
# ╠═215585ab-d2c3-47e7-b54f-e2719bdeb230
# ╠═99a1a535-87de-4fea-bb83-2bcc0267d70f
# ╟─c155fb29-275d-4b56-b125-ca0084c4182e
# ╠═73173e80-d76d-4819-a3e8-dc7138d532ac
# ╟─6ead1870-0e5c-44e6-92b7-3e6f90c9e387
# ╠═38a9ea42-1c9b-48c0-8d08-3acc00319223
# ╟─8f999e94-bfc6-4b08-be35-1da44953dc55
# ╟─f70c9f2c-d24c-441f-bf24-863c5ba596dd
# ╟─9b43e322-a873-421b-8260-d4368d6cbaa3
# ╠═bd1a6dcf-0d91-4774-a267-e19abb8dc0e4
# ╠═3f31a05a-0a5b-45f6-a718-f43744758472
# ╟─09c00dd9-7a2a-478c-bb4f-9a1289f4796e
# ╠═31b42e8e-9361-45d1-b84a-2b8b54153ec4
# ╟─d2c825f3-7203-4905-828a-9be0fd42fe41
# ╟─afa8ae7e-5beb-45fd-87cd-33d1f263b692
# ╠═03fb84cf-8fe9-4622-b7a3-6b8ff49907c4
# ╟─01a5d74d-c71b-49f2-833f-1a937cbc89e8
# ╠═f84c087d-8963-4835-81d9-8e72ab297b13
# ╟─b5dc0d67-f80b-40ae-a65c-7d8f085e84ae
# ╠═05a54094-1587-43eb-a3f2-d8945803f482
# ╟─9b3425d6-0c3a-4395-a2c2-af73928aed01
# ╠═04162e25-9a74-4bc1-9ceb-66655111897e
# ╠═a48efa0b-3c2b-44a9-b808-cf0975fce72d
# ╠═841a1fb5-fc52-494c-a2eb-15930b41ced8
# ╠═fbbc5e91-3e43-43c8-b23e-0ccd00ba0c71
# ╟─c74af2e0-99eb-4526-a68a-8fdab5717f16
# ╟─1ed5e587-d0eb-4dc2-bedc-00d13f258ecd
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002

max_stars_repo_name: JeannotJeannot/Learn

max_issues_repo_name: JeannotJeannot/Learn

max_forks_repo_name: JeannotJeannot/Learn


text:
# Goals: minimize calls to ordered_states (allocates memory)

# needs pomdp for stateindex in pdf(b, s)
# needs list of ordered_states for rand(b)

"""
    DiscreteBelief

A belief specified by a probability vector.

Normalization of `b` is assumed in some calculations (e.g. pdf), but it is only automatically enforced in `update(...)`, and a warning is given if normalized incorrectly in `DiscreteBelief(pomdp, b)`.

# Constructor
    DiscreteBelief(pomdp, b::Vector{Float64}; check::Bool=true)

# Fields
- `pomdp` : the POMDP problem
- `state_list` : a vector of ordered states
- `b` : the probability vector
"""
struct DiscreteBelief{P<:POMDP, S}
    pomdp::P
    state_list::Vector{S}       # vector of ordered states
    b::Vector{Float64}
end

function DiscreteBelief(pomdp, b::Vector{Float64}; check::Bool=true)
    if check
        if !isapprox(sum(b), 1.0, atol=0.001)
            @warn """
                  b in DiscreteBelief(pomdp, b) does not sum to 1.

                  To suppress this warning use `DiscreteBelief(pomdp, b, check=false)`
                  """
        end
        if !all(0.0 <= p <= 1.0 for p in b)
            @warn """
                  b in DiscreteBelief(pomdp, b) contains entries outside [0,1].

                  To suppress this warning use `DiscreteBelief(pomdp, b, check=false)`
                  """
        end
    end
    return DiscreteBelief(pomdp, ordered_states(pomdp), b)
end

"""
     uniform_belief(pomdp)

Return a DiscreteBelief with equal probability for each state.
"""
function uniform_belief(pomdp)
    state_list = ordered_states(pomdp)
    ns = length(state_list)
    return DiscreteBelief(pomdp, state_list, ones(ns) / ns)
end

pdf(b::DiscreteBelief, s) = b.b[stateindex(b.pomdp, s)]

function Random.rand(rng::Random.AbstractRNG, b::DiscreteBelief)
    i = sample(rng, Weights(b.b))
    return b.state_list[i]
end

function Base.fill!(b::DiscreteBelief, x::Float64)
    fill!(b.b, x)
    return b
end

Base.length(b::DiscreteBelief) = length(b.b)

support(b::DiscreteBelief) = b.state_list

Statistics.mean(b::DiscreteBelief) = sum(b.state_list.*b.b)/sum(b.b)
StatsBase.mode(b::DiscreteBelief) = b.state_list[argmax(b.b)]

==(b1::DiscreteBelief, b2::DiscreteBelief) = b1.state_list == b2.state_list && b1.b == b2.b

Base.hash(b::DiscreteBelief, h::UInt) = hash(b.b, hash(b.state_list, h))

"""
    DiscreteUpdater

An updater type to update discrete belief using the discrete Bayesian filter.

# Constructor
    DiscreteUpdater(pomdp::POMDP)

# Fields
- `pomdp <: POMDP`
"""
mutable struct DiscreteUpdater{P<:POMDP} <: Updater
    pomdp::P
end

uniform_belief(up::DiscreteUpdater) = uniform_belief(up.pomdp)

function initialize_belief(bu::DiscreteUpdater, dist::Any)
    state_list = ordered_states(bu.pomdp)
    ns = length(state_list)
    b = zeros(ns)
    belief = DiscreteBelief(bu.pomdp, state_list, b)
    for s in support(dist)
        sidx = stateindex(bu.pomdp, s)
        belief.b[sidx] = pdf(dist, s)
    end
    return belief
end

function update(bu::DiscreteUpdater, b::DiscreteBelief, a, o)
    pomdp = bu.pomdp
    state_space = b.state_list
    bp = zeros(length(state_space))

    for (si, s) in enumerate(state_space)

        if pdf(b, s) > 0.0
            td = transition(pomdp, s, a)

            for (sp, tp) in weighted_iterator(td)
                spi = stateindex(pomdp, sp)
                op = obs_weight(pomdp, s, a, sp, o) # shortcut for observation probability from POMDPModelTools

                bp[spi] += op * tp * b.b[si]
            end
        end
    end

    bp_sum = sum(bp)

    if bp_sum == 0.0
        error("""
              Failed discrete belief update: new probabilities sum to zero.

              b = $b
              a = $a
              o = $o

              Failed discrete belief update: new probabilities sum to zero.
              """)
    end

    # Normalize
    bp ./= bp_sum

    return DiscreteBelief(pomdp, b.state_list, bp)
end

update(bu::DiscreteUpdater, b::Any, a, o) = update(bu, initialize_belief(bu, b), a, o)

ext: jl


text:
# Submodule that constructs the system of Shearlets in 2D of size rows x cols

######################################################################
## ype of Filterconfigsuration to check the sizes
struct Filterswedgebandlow
    wedge::AbstractArray
    bandpass::AbstractArray
    lowpass::AbstractArray
end

#######################################################################
# Function that generates the whole Shearlet System filters (wedge, bandpass and lowpass) of size rows x cols
"""
...
getwedgebandpasslowpassfilters2D(rows,cols,directionalFilter = filt_gen("directional_shearlet"), scalingFilter = filt_gen("scaling_shearlet"),gpu = 0)
generates the Wedge, Bandpass and LowspassFilter of size rows x cols
...
"""
function getwedgebandpasslowpassfilters2D(rows::Int,cols::Int,shearLevels,directionalFilter = filt_gen("directional_shearlet"),scalingFilter = filt_gen("scaling_shearlet"),waveletFilter = mirror(scalingFilter),
                                                   scalingFilter2 = scalingFilter,gpu = 0)
    FFTW.set_num_threads(Sys.CPU_THREADS)
    # Make shearLevels integer
    shearLevels = map(Int,shearLevels);
    # The number of scales
    NScales = length(shearLevels);
    # Initialize the bandpass and wedge filter
    # Bandpass filters partion the frequency plane into different scales
    if gpu == 1
        bandpass = zeros(Float32,rows,cols,NScales);
        bandpass = AFArray(bandpass+im*zeros(bandpass));
    else
        bandpass = zeros(Float64,rows,cols,NScales);
        bandpass = bandpass+im*zeros(size(bandpass));
    end

    # Wedge filters partition the frequency plane into different directions
    wedge = Array{Any}(undef, 1, round(Int64,maximum((shearLevels) .+ 1)));

    # normalize the directional filter directional filter
    if gpu == 1
        directionalFilter = AFArray(convert(Array{Float32},directionalFilter/sum(abs.(directionalFilter[:]))));
    else
        directionalFilter = directionalFilter/sum(abs.(directionalFilter[:]));
    end

    # Compute 1D high and lowpass filters at different scales

    # filterHigh{NScales} = g, and filterHigh{1} = g_J (analogous with filterLow)
    filterHigh = Array{Any}(undef, 1,NScales);
    filterLow = Array{Any}(undef, 1,NScales);
    filterLow2 = Array{Any}(undef, 1,round(Int64,maximum(shearLevels)+1));

    # Initialize wavelet highpass and lowpass filters
    if gpu == 1
        filterHigh[size(filterHigh,2)] = AFArray(convert(Array{Float32},waveletFilter));
        filterLow[size(filterLow,2)] = AFArray(convert(Array{Float32},scalingFilter));
        filterLow2[size(filterLow2,2)] = AFArray(convert(Array{Float32},scalingFilter2));

        # Lets compute the filters in the other scales
        for j = (size(filterHigh)[2]-1):-1:1
            filterLow[j] = conv2(filterLow[size(filterLow,2)],upsample(filterLow[j+1],1,1,1))
            filterHigh[j] = conv(filterLow[size(filterLow,2)],upsample(filterHigh[j+1],1,1,1))
        end

        # Lets compute the filters in the other scales
        for j=(size(filterLow2)[2]-1):-1:1
            filterLow2[j] = conv(filterLow2[size(filterLow2,2)],upsample(filterLow2[j+1],1,1,1))
        end
        # Construct the bandpassfilter
        # Need to convert first to complex array since
        for j = 1:size(filterHigh,2)
           bandpass[:,:,j] = -fftshift(fft(ifftshift(padarray(filterHigh[j],[rows,cols],1))));
        end
    else
        filterHigh[size(filterHigh,2)] = waveletFilter;
        filterLow[size(filterLow,2)] = scalingFilter;
        filterLow2[size(filterLow2,2)] = scalingFilter2;

        # Lets compute the filters in the other scales
        for j = (size(filterHigh)[2]-1):-1:1
            filterLow[j] = conv(filterLow[size(filterLow,2)],upsample(filterLow[j+1],1,1))
            filterHigh[j] = conv(filterLow[size(filterLow,2)],upsample(filterHigh[j+1],1,1))
        end

        # Lets compute the filters in the other scales
        for j=(size(filterLow2)[2]-1):-1:1
            filterLow2[j] = conv(filterLow2[size(filterLow2,2)],upsample(filterLow2[j+1],1,1))
        end

        # Construct the bandpassfilter
        # Need to convert first to complex array since
        for j = 1:size(filterHigh,2)
           bandpass[:,:,j] = fftshift(fft(ifftshift(padarray(filterHigh[j],[rows,cols]))));
        end
    end

    ## construct wedge filters for achieving directional selectivity.
    # as the entries in the shearLevels array describe the number of differently
    # sheared atoms on a certain scale, a different set of wedge
    # filters has to be constructed for each value in shearLevels.
    if gpu == 1
        for shearLevel = unique(shearLevels)
            #preallocate a total of floor(2^(shearLevel+1)+1) wedge filters, where
            #floor(2^(shearLevel+1)+1) is the number of different directions of
            #shearlet atoms associated with the horizontal (resp. vertical)
            #frequency cones.
            wedge[shearLevel+1] = AFArray(zeros(Complex{Float32},rows,cols,floor(2^(shearLevel+1)+1)));
            #upsample directional filter in y-direction. by upsampling the directional
            #filter in the time domain, we construct repeating wedges in the
            #frequency domain ( compare abs.(fftshift(fft2(ifftshift(directionalFilterUpsampled)))) and
            #abs.(fftshift(fft2(ifftshift(directionalFilter)))) ).
            directionalFilterUpsampled = upsample(directionalFilter,1,2^(shearLevel+1)-1,1);

            #remove high frequencies along the y-direction in the frequency domain.
            #by convolving the upsampled directional filter with a lowpass filter in y-direction, we remove all
            #but the central wedge in the frequency domain.
            wedgeHelp = conv2(directionalFilterUpsampled, filterLow2[size(filterLow2,2)-shearLevel][:,:])
            wedgeHelp = padarray(wedgeHelp,[rows,cols],1);
            #please note that wedgeHelp now corresponds to
            #conv(p_j,h_(J-j*alpha_j/2)') in the language of the paper. to see
            #this, consider the definition of p_j on page 14, the definition of w_j
            #on the same page an the definition of the digital sheralet filter on
            #page 15. furthermore, the g_j part of the 2D wavelet filter w_j is
            #invariant to shearings, hence it suffices to apply the digital shear
            #operator to wedgeHelp.;

            ## application of the digital shear operator (compare equation (22))

            #upsample wedge filter in x-direction. this operation corresponds to
            #the upsampling in equation (21) on page 15.
            wedgeUpsampled = upsample(wedgeHelp,2,2^shearLevel-1,1);
            #convolve wedge filter with lowpass filter, again following equation
            #(21) on page 14.
            lowpassHelp = padarray(filterLow2[size(filterLow2,2)-max(shearLevel-1,0)],size(wedgeUpsampled),1);
            if shearLevel >= 1
                wedgeUpsampled = fftshift(ifft(ifftshift(fftshift(fft(ifftshift(lowpassHelp))).*fftshift(fft(ifftshift(wedgeUpsampled))))));
            end
            lowpassHelpFlip = fliplr(lowpassHelp,1);
            #traverse all directions of the upper part of the left horizontal
            #frequency cone
            for k = -2^shearLevel:2^shearLevel
                #resample wedgeUpsampled as given in equation (22) on page 15.
                wedgeUpsampledSheared = dshear(wedgeUpsampled,k,2,1);
                #convolve again with flipped lowpass filter, as required by equation (22) on
                #page 15.
                if shearLevel >= 1
                   wedgeUpsampledSheared = fftshift(ifft(ifftshift(fftshift(fft(ifftshift(lowpassHelpFlip))).*fftshift(fft(ifftshift(wedgeUpsampledSheared))))));
                end
                #obtain downsampled and renormalized and sheared wedge filter in the
                #frequency domain, according to equation (22) on page 15
                wedge[shearLevel+1][:,:,fix(2^shearLevel)+1-k] = fftshift(fft(ifftshift(2^shearLevel*wedgeUpsampledSheared[:,1:2^shearLevel:(2^shearLevel*cols-1)])));
            end
        end
    else
        for shearLevel = unique(shearLevels)
            #preallocate a total of floor(2^(shearLevel+1)+1) wedge filters, where
            #floor(2^(shearLevel+1)+1) is the number of different directions of
            #shearlet atoms associated with the horizontal (resp. vertical)
            #frequency cones.
            wedge[shearLevel+1] = zeros(rows,cols,floor(2^(shearLevel+1)+1))+zeros(rows,cols,floor(2^(shearLevel+1)+1))*im;

            #upsample directional filter in y-direction. by upsampling the directional
            #filter in the time domain, we construct repeating wedges in the
            #frequency domain ( compare abs.(fftshift(fft2(ifftshift(directionalFilterUpsampled)))) and
            #abs.(fftshift(fft2(ifftshift(directionalFilter)))) ).
            directionalFilterUpsampled = upsample(directionalFilter,1,2^(shearLevel+1)-1);

            #remove high frequencies along the y-direction in the frequency domain.
            #by convolving the upsampled directional filter with a lowpass filter in y-direction, we remove all
            #but the central wedge in the frequency domain.

            wedgeHelp = conv2(directionalFilterUpsampled, filterLow2[size(filterLow2,2)-shearLevel][:,:])
            wedgeHelp = padarray(wedgeHelp,[rows,cols]);

            #please note that wedgeHelp now corresponds to
            #conv(p_j,h_(J-j*alpha_j/2)') in the language of the paper. to see
            #this, consider the definition of p_j on page 14, the definition of w_j
            #on the same page an the definition of the digital sheralet filter on
            #page 15. furthermore, the g_j part of the 2D wavelet filter w_j is
            #invariant to shearings, hence it suffices to apply the digital shear
            #operator to wedgeHelp.;

            ## application of the digital shear operator (compare equation (22))

            #upsample wedge filter in x-direction. this operation corresponds to
            #the upsampling in equation (21) on page 15.
            wedgeUpsampled = upsample(wedgeHelp,2,2^shearLevel-1);

            #convolve wedge filter with lowpass filter, again following equation
            #(21) on page 14.
            lowpassHelp = padarray(filterLow2[size(filterLow2,2)-max(shearLevel-1,0)],size(wedgeUpsampled));
            if shearLevel >= 1
                wedgeUpsampled = fftshift(ifft(ifftshift(fftshift(fft(ifftshift(lowpassHelp))).*fftshift(fft(ifftshift(wedgeUpsampled))))));
            end
            lowpassHelpFlip = fliplr(lowpassHelp);
            #traverse all directions of the upper part of the left horizontal
            #frequency cone
            for k = -2^shearLevel:2^shearLevel
                #resample wedgeUpsampled as given in equation (22) on page 15.
                wedgeUpsampledSheared = dshear(wedgeUpsampled,k,2);
                #convolve again with flipped lowpass filter, as required by equation (22) on
                #page 15.
                if shearLevel >= 1
                   wedgeUpsampledSheared = fftshift(ifft(ifftshift(fftshift(fft(ifftshift(lowpassHelpFlip))).*fftshift(fft(ifftshift(wedgeUpsampledSheared))))));
                end
                #obtain downsampled and renormalized and sheared wedge filter in the
                #frequency domain, according to equation (22) on page 15
                wedge[shearLevel+1][:,:,fix(2^shearLevel)+1-k] = fftshift(fft(ifftshift(2^shearLevel*wedgeUpsampledSheared[:,1:2^shearLevel:(2^shearLevel*cols-1)])));
            end
        end
    end

    ## compute low pass filter of shearlet system
    if gpu == 1
        lowpass = fftshift(fft(ifftshift(padarray(transpose(filterLow[1]')*filterLow[1]',[rows,cols],1))));
    else
        lowpass = fftshift(fft(ifftshift(padarray(transpose(filterLow[1]')*filterLow[1]',[rows,cols]))));
    end

    # Generate the final array
    return Filterswedgebandlow(wedge,bandpass,lowpass)
end #getwedgebandpasslowpassfilters2D

##############################################################
# Create a type for the Preparedfilters
struct Preparedfilters
    size
    shearLevels
    cone1
    cone2
end

#######################################################################
# Function that prepare the filters
"""
...
SlprepareFilters2D(rows, cols, nScales, shearLevels = ceil.((1:nScales)/2),
    directionalFilter = filt_gen("directional_shearlet"),
    scalingFilter = filt_gen("scaling_shearlet"),
    waveletFilter = mirror(scalingFilter),
    scalingFilter2 = scalingFilter, gpu = 0) function that prepare the filters to generate
		 the shearlet system
...
"""
function preparefilters2D(rows, cols, nScales, shearLevels = ceil.((1:nScales)/2),
    directionalFilter = filt_gen("directional_shearlet"),
    scalingFilter = filt_gen("scaling_shearlet"),
    waveletFilter = mirror(scalingFilter),
    scalingFilter2 = scalingFilter, gpu = 0)

    #Make sure the shearLevles are integer
    shearLevels = map(Int,shearLevels)
    # check filter sizes
    filters = checkfiltersizes(rows,cols,shearLevels,directionalFilter,scalingFilter,waveletFilter,scalingFilter2);

    # Save the new filters
    directionalFilter = filters.directionalFilter;
    scalingFilter = filters.scalingFilter;
    waveletFilter = filters.waveletFilter;
    scalingFilter2 = filters.scalingFilter2;

    # Define the cones
    cone1 = getwedgebandpasslowpassfilters2D(rows,cols,shearLevels,directionalFilter,scalingFilter,waveletFilter,scalingFilter2,gpu);
    if rows == cols
    	cone2 = cone1;
    else
    	cone2 = getwedgebandpasslowpassfilters2D(cols,rows,shearLevels,directionalFilter,scalingFilter,waveletFilter,scalingFilter2,gpu);
    end
    return Preparedfilters([rows,cols],shearLevels,cone1,cone2)
end # preparefilters2D

#######################################################################
# Function compute a index set describing a 2D shearlet system
"""
...
getshearletidxs2D(shearLevels, full=0, restriction_type = [],restriction_value = []) Compute a index set describing a 2D shearletsystem, with restriction_type in
["cones","scales","shearings"], where shearLevels is a 1D array spcifying the level of shearing on each scale, and full is a boolean that orders to generate the full shearlet indices.
...
"""
function getshearletidxs2D(shearLevels, full=0, restriction_type = [],restriction_value = [])
		shearLevels = map(Int,shearLevels)
    # Set default values
    shearletIdxs = [];
    includeLowpass = 1;
    scales = 1:length(shearLevels);
    shearings = -2^(maximum(shearLevels)):2^(maximum(shearLevels));
    cones = 1:2;

    for i = 1:length(restriction_type)
        includeLowpass = 0;
        if restriction_type[i] == "scales"
            scales = restriction_value[i];
        elseif restriction_type[i] == "shearings"
            shearings = restriction_value[i];
        elseif restriction_type[i] == "cones"
            cones = restriction_value[i];
        end
    end

    for cone = intersect(1:2,cones)
        for scale = intersect(1:length(shearLevels),scales)
            for shearing = intersect(-2^shearLevels[scale]:2^shearLevels[scale],shearings)
                if convert(Bool,full) || cone == 1 || abs.(shearing)<2^(shearLevels[scale])
                    shearletIdxs=[shearletIdxs;[cone,scale,shearing]];
                end
            end
        end
    end

    shearletIdxs = transpose(reshape(shearletIdxs,3,Int(length(shearletIdxs)/3)))
    if convert(Bool,includeLowpass) || convert(Bool,sum(0 .== scales)) || convert(Bool,sum(0 .== cones))
        shearletIdxs = [shearletIdxs;[0 0 0]];
    end
    return shearletIdxs
end # getshearletidxs2D

####################################################################
# Type of shearletsystem in 2D
struct Shearletsystem2D
    shearlets
    size
    shearLevels
    full
    nShearlets
    shearletIdxs
    dualFrameWeights
    RMS
    isComplex
		gpu
end

#######################################################################
# Function that generates the desired shearlet system
"""
...
 getshearletsystem2D(rows,cols,nScales,shearLevels=ceil.((1:nScales)/2),full = 0,directionalFilter = filt_gen("directional_shearlet"),quadratureMirrorFilter= filt_gen("scaling_shearlet"), gpu = 0) generates the desired shearlet system
...
"""
function getshearletsystem2D(rows,cols,nScales,
                                shearLevels=ceil.((1:nScales)/2),
                                full= 0,
                                directionalFilter = filt_gen("directional_shearlet"),
                                quadratureMirrorFilter= filt_gen("scaling_shearlet"),gpu=0)

    # Set default value generates the desired shearlet systems
    shearLevels = Int.(shearLevels)

    #Generate prepared Filters and indices
    Preparedfilters = preparefilters2D(rows,cols,nScales,shearLevels,directionalFilter,quadratureMirrorFilter,mirror(quadratureMirrorFilter),quadratureMirrorFilter, gpu);
    shearletIdxs = getshearletidxs2D(shearLevels,full);

    # Generate shearlets, RMS(rootmeansquare), dualFrameWeights
    rows = Preparedfilters.size[1];
    cols = Preparedfilters.size[2];
    nShearlets = size(shearletIdxs,1);
    if gpu == 1
        shearlets = AFArray(zeros(Complex{Float32},rows,cols,nShearlets));
    else
        shearlets = zeros(rows,cols,nShearlets)+im .* zeros(rows,cols,nShearlets);
    end
    # Compute shearlets
    for j = 1:nShearlets
        cone = shearletIdxs[j,1];
        scale = shearletIdxs[j,2];
        shearing = shearletIdxs[j,3];

        if cone == 0
            shearlets[:,:,j] = Preparedfilters.cone1.lowpass;
        elseif cone == 1
            #here, the fft of the digital shearlet filters described in
            #equation (23) on page 15 of "ShearLab 3D: Faithful Digital
            #Shearlet Transforms based on Compactly Supported Shearlets" is computed.
            #for more details on the construction of the wedge and bandpass
            #filters, please refer to getwedgebandpasslowpassfilters2D.
            shearlets[:,:,j] = Preparedfilters.cone1.wedge[Preparedfilters.shearLevels[scale]+1][:,:,-shearing+2^Preparedfilters.shearLevels[scale]+1].*conj(Preparedfilters.cone1.bandpass[:,:,scale]);
        else
            shearlets[:,:,j] = permutedims(Preparedfilters.cone2.wedge[Preparedfilters.shearLevels[scale]+1][:,:,shearing+2^Preparedfilters.shearLevels[scale]+1].*conj(Preparedfilters.cone2.bandpass[:,:,scale]),[2,1]);
        end
    end

	RMS =  transpose(sum(reshape(sum(real.(abs.(shearlets)).^2, dims=1), size(shearlets, 2), size(shearlets, 3)), dims=1));
    if gpu == 1
        RMS = (sqrt.(RMS)/convert(Float32,sqrt.(rows*cols)));
        dualFrameWeights = sum(real(abs.(shearlets)).^2,dims=3);
    else
        RMS = (sqrt.(RMS)/sqrt.(rows*cols));
        dualFrameWeights = dropdims(sum(abs.(shearlets).^2,dims=3),dims=3); # need to stream to host before they fix abs of comples AFArray
    end
    #return the system
    return Shearletsystem2D(shearlets, Preparedfilters.size,
    Preparedfilters.shearLevels, full, size(shearletIdxs,1),
    shearletIdxs, dualFrameWeights, RMS, 0, gpu)
end #getshearletsystem2D

# type for individual shearlets2D
struct Shearlets2D
		shearlets
		RMS
		dualFrameWeights
		gpu
end

#######################################################################
# Function that generates the desired shearlets
"""
...
 getshearlets2D(PreparedFilters,shearletIdxs,gpu = 0) generates the 2D shearlets in the frequency domain
...
"""
function getshearlets2D(Preparedfilters, shearletIdxs = getshearletidxs2D(Preparedfilters.shearLevels),gpu = 0)
    # Generate shearlets, RMS(rootmeansquare), dualFrameWeights
    rows = Preparedfilters.size[1];
    cols = Preparedfilters.size[2];
    nShearlets = size(shearletIdxs,1);
		if gpu == 1
    	shearlets = AFArray(zeros(Complex{Float32},rows,cols,nShearlets));
		else
    	shearlets = zeros(rows,cols,nShearlets)+im*zeros(rows,cols,nShearlets);
		end
    # Compute shearlets
    for j = 1:nShearlets
        cone = shearletIdxs[j,1];
        scale = shearletIdxs[j,2];
        shearing = shearletIdxs[j,3];

        if cone == 0
            shearlets[:,:,j] = Preparedfilters.cone1.lowpass;
        elseif cone == 1
            #here, the fft of the digital shearlet filters described in
            #equation (23) on page 15 of "ShearLab 3D: Faithful Digital
            #Shearlet Transforms based on Compactly Supported Shearlets" is computed.
            #for more details on the construction of the wedge and bandpass
            #filters, please refer to getwedgebandpasslowpassfilters2D.
            shearlets[:,:,j] = Preparedfilters.cone1.wedge[Preparedfilters.shearLevels[scale]+1][:,:,-shearing+2^Preparedfilters.shearLevels[scale]+1].*conj(Preparedfilters.cone1.bandpass[:,:,scale]);
        else
            shearlets[:,:,j] = permutedims(Preparedfilters.cone2.wedge[Preparedfilters.shearLevels[scale]+1][:,:,shearing+2^Preparedfilters.shearLevels[scale]+1].*conj(Preparedfilters.cone2.bandpass[:,:,scale]),[2,1]);
        end
    end
		RMS =  transpose(sum(reshape(sum(real(abs.(shearlets)).^2,dims=1),size(shearlets,2),size(shearlets,3)),dims=1));
    if gpu == 1
        RMS = (sqrt.(RMS)/convert(Float32,sqrt.(rows*cols)));
        dualFrameWeights = sum(real(abs.(shearlets)).^2,dims=3);
    else
        RMS = (sqrt.(RMS)/sqrt.(rows*cols));
        dualFrameWeights = dropdims(sum(abs.(shearlets).^2,dims=3),3); # need to stream to host before they fix abs of comples AFArray
    end
		return Shearlets2D(shearlets, RMS, dualFrameWeights,gpu)
end #getshearlets2D


text:
using FermiCG
using Printf
using Test
using LinearAlgebra
using Profile
using HDF5
using Random
using StatProfilerHTML

using PyCall
pydir = joinpath(dirname(pathof(FermiCG)), "python")
pushfirst!(PyVector(pyimport("sys")."path"), pydir)
ENV["PYTHON"] = Sys.which("python")

#@testset "Clusters" begin
#function run()
    atoms = []
    clusters = []
    na = 0
    nb = 0
    init_fspace = []

    function generate_H_ring(n,radius)
        theta = 2*pi/n

        atoms = []
        for i in 0:n-1
            push!(atoms,Atom(i+1,"H",[radius*cos(theta*i), radius*sin(theta*i), 0]))
        end
        return atoms
    end

    #
    # Test basic Tucker stuff
    Random.seed!(2);
    A = rand(4,6,3,3,5)
    tuck = FermiCG.Tucker(A, thresh=20, verbose=1)

    display(size(tuck.core))
    display(size.(tuck.factors))
    B = FermiCG.recompose(tuck)
    println()
    println(FermiCG.dims_large(tuck))
    @test all(FermiCG.dims_small(tuck) .== [4, 1, 3, 3, 1])
    @test all(FermiCG.dims_large(tuck) .== [4, 6, 3, 3, 5])

    A = rand(4,6,3,3,5)
    tuck = FermiCG.Tucker(A, thresh=-1, verbose=1)
    B = FermiCG.recompose(tuck)
    @test isapprox(abs.(A), abs.(B), atol=1e-12)

    A = rand(4,6,3,3,5)*.1
    B = rand(4,6,3,3,5)*.1
    C = A+B

    #tuckA = FermiCG.Tucker(A, thresh=-1, verbose=1, max_number=2)
    #tuckB = FermiCG.Tucker(B, thresh=-1, verbose=1, max_number=2)
    #tuckC = FermiCG.Tucker(C, thresh=-1, verbose=1, max_number=2)
    tuckA = FermiCG.Tucker(A, thresh=-1, verbose=0)
    tuckB = FermiCG.Tucker(B, thresh=-1, verbose=0)
    tuckC = FermiCG.Tucker(C, thresh=-1, verbose=0)

    # test Tucker addition
    test = tuckA + tuckB
    @test isapprox(FermiCG.dot(tuckC,tuckC), FermiCG.dot(test,test), atol=1e-12)

    #
    # Now test basis transformation
    A = rand(4,6,3,3,5)*.1

    trans1 = Dict{Int,Matrix{Float64}}()
    trans1[2] = rand(6,5)
    trans1[4] = rand(3,2)

    trans2 = []
    for i = 1:5
        if haskey(trans1, i)
            push!(trans2, trans1[i])
        else
            push!(trans2, Matrix(1.0I,size(A,i),size(A,i)))
        end
    end

    display((length(A), size(A)))

    A1 = FermiCG.transform_basis(A, trans1)
    display((length(A1), size(A1)))

    A2 = FermiCG.transform_basis(A, trans2)
    display((length(A1), size(A2)))

    #A2 = FermiCG.tucker_recompose(A, trans2)
    #display((length(A2), size(A2)))

    @test isapprox(A1, A2, atol=1e-12)

    if false
        r = 1
        a = 0
        push!(atoms,Atom(1,"H", [0,.0*a,0*r]))
        push!(atoms,Atom(2,"H", [0,.0*a,1*r]))
        push!(atoms,Atom(3,"H", [0,.1*a,2*r]))
        push!(atoms,Atom(4,"H", [0,.1*a,3*r]))
        push!(atoms,Atom(5,"H", [0,.2*a,4*r]))
        push!(atoms,Atom(6,"H", [0,.2*a,5*r]))
        push!(atoms,Atom(7,"H", [0,.3*a,6*r]))
        push!(atoms,Atom(8,"H", [0,.3*a,7*r]))
        push!(atoms,Atom(9,"H", [0,.4*a,8*r]))
        push!(atoms,Atom(10,"H",[0,.4*a,9*r]))
        push!(atoms,Atom(11,"H",[0,.5*a,10*r]))
        push!(atoms,Atom(12,"H",[0,.5*a,11*r]))
        push!(atoms,Atom(13,"H",[0,.6*a,12*r]))
        push!(atoms,Atom(14,"H",[0,.6*a,13*r]))
        push!(atoms,Atom(15,"H",[0,.7*a,14*r]))
        push!(atoms,Atom(16,"H",[0,.7*a,15*r]))
        push!(atoms,Atom(17,"H",[0,.8*a,16*r]))
        push!(atoms,Atom(18,"H",[0,.8*a,17*r]))

        clusters    = [(1:2),(3:4),(5:6),(7:8),(9:10),(11:12)]
        init_fspace = [(1,1),(1,1),(1,1),(1,1),(1,1),(1,1)]
        na = 6
        nb = 6
        clusters    = [(1:2),(3:4),(5:6),(7:8)]
        init_fspace = [(1,1),(1,1),(1,1),(1,1)]
        na = 4
        nb = 4
        clusters    = [(1:6),(7:12)]
        init_fspace = [(3,3),(3,3)]
        na = 6
        nb = 6
        clusters    = [(1:6),(7:12),(13:18)]
        init_fspace = [(3,3),(3,3),(3,3)]
        na = 6
        nb = 6
    elseif false
        push!(atoms,Atom(1,"H",[-1.30,0,0.00]))
        push!(atoms,Atom(2,"H",[-1.30,0,1.00]))
        push!(atoms,Atom(3,"H",[ 0.00,0,0.00]))
        push!(atoms,Atom(4,"H",[ 0.00,0,1.00]))
        push!(atoms,Atom(5,"H",[ 1.33,0,0.00]))
        push!(atoms,Atom(6,"H",[ 1.30,0,1.00]))

        clusters    = [(1:2),(3:4),(5:6)]
        init_fspace = [(1,1),(1,1),(1,1)]
        na = 3
        nb = 3
    elseif true

        rad = 3

        atoms = generate_H_ring(10,rad)
        clusters    = [(1:2),(3:4),(5:6),(7:8),(9:10)]
        init_fspace = [(1,1),(1,1),(1,1),(1,1),(1,1)]
        clusters    = [(1:4),(5:6),(7:8),(9:10)]
        init_fspace = [(2,2),(1,1),(1,1),(1,1)]
        clusters    = [(1:4),(5:8),(9:10)]
        init_fspace = [(2,2),(2,2),(1,1)]
        na = 5
        nb = 5

        atoms = generate_H_ring(8,rad)
        clusters    = [(1:2),(3:4),(5:6),(7:8)]
        init_fspace = [(1,1),(1,1),(1,1),(1,1)]
        clusters    = [(1:4),(5:6),(7:8)]
        init_fspace = [(2,2),(1,1),(1,1)]
        clusters    = [(1:4),(5:8)]
        init_fspace = [(2,2),(2,2)]
        na = 4
        nb = 4

        atoms = generate_H_ring(12,rad)
        clusters    = [(1:2),(3:4),(5:6),(7:8),(9:10),(11:12)]
        init_fspace = [(1,1),(1,1),(1,1),(1,1),(1,1),(1,1)]
        clusters    = [(1:4),(5:8),(9:12)]
        init_fspace = [(2,2),(2,2),(2,2)]
        clusters    = [(1:4),(5:8),(9:10),(11:12)]
        init_fspace = [(2,2),(2,2),(1,1),(1,1)]
        clusters    = [(1:6),(7:12)]
        init_fspace = [(3,3),(3,3)]
        na = 6
        nb = 6

    end

    basis = "6-31g"
    basis = "sto-3g"
    mol     = Molecule(0,1,atoms,basis)

    # get integrals
    mf = FermiCG.pyscf_do_scf(mol)
    nbas = size(mf.mo_coeff)[1]
    ints = FermiCG.pyscf_build_ints(mol,mf.mo_coeff, zeros(nbas,nbas));
    #e_fci, d1_fci, d2_fci = FermiCG.pyscf_fci(ints, na, nb, conv_tol=1e-10,max_cycle=100, nroots=2)

    #run fci with pyscf
    if false
        pyscf = pyimport("pyscf")
        fci = pyimport("pyscf.fci")
        mp = pyimport("pyscf.mp")
        mp2 = mp.MP2(mf)
        cisolver = pyscf.fci.direct_spin1.FCI()
        cisolver.max_cycle = 100
        cisolver.conv_tol = 1e-10
        nelec = na + nb
        norb = size(ints.h1)[1]
        e_fci, ci = cisolver.kernel(ints.h1, ints.h2, norb , nelec, ecore=0, nroots = 1, verbose=100)
        e_fci = min(e_fci...)
        @printf(" FCI Energy: %12.8f\n", e_fci)
    end

    # localize orbitals
    C = mf.mo_coeff
    Cl = FermiCG.localize(mf.mo_coeff,"lowdin",mf)
    FermiCG.pyscf_write_molden(mol,Cl,filename="lowdin.molden")
    S = FermiCG.get_ovlp(mf)
    U =  C' * S * Cl
    println(" Rotate Integrals")
    flush(stdout)
    ints = FermiCG.orbital_rotation(ints,U)
    println(" done.")
    flush(stdout)

    #
    # define clusters
    clusters = [Cluster(i,collect(clusters[i])) for i = 1:length(clusters)]
    display(clusters)

    rdm1 = zeros(size(ints.h1))
    #e_cmf, U, Da, Db  = FermiCG.cmf_oo(ints, clusters, init_fspace, rdm1,
    #                                   max_iter_oo=40, verbose=0, gconv=1e-6, method="gd", alpha=1e-1)
    #ints = FermiCG.orbital_rotation(ints,U)

    e_cmf, U, Da, Db  = FermiCG.cmf_oo(ints, clusters, init_fspace, rdm1,
                                       max_iter_oo=40, verbose=0, gconv=1e-6, method="bfgs")
    FermiCG.pyscf_write_molden(mol,Cl*U,filename="cmf.molden")
    ints = FermiCG.orbital_rotation(ints,U)

    e_ref = e_cmf - ints.h0

    max_roots = 100
    # build Hamiltonian, cluster_basis and cluster ops
    #display(Da)
    #cluster_bases = FermiCG.compute_cluster_eigenbasis(ints, clusters, verbose=2, max_roots=max_roots)
    cluster_bases = FermiCG.compute_cluster_eigenbasis(ints, clusters, verbose=0, max_roots=max_roots,
                                                       init_fspace=init_fspace, rdm1a=Da, rdm1b=Db)
    clustered_ham = FermiCG.extract_ClusteredTerms(ints, clusters)
    cluster_ops = FermiCG.compute_cluster_ops(cluster_bases, ints);

    p_spaces = Vector{FermiCG.ClusterSubspace}()
    q_spaces = Vector{FermiCG.ClusterSubspace}()

    # define p spaces
    for ci in clusters
        tss = FermiCG.ClusterSubspace(ci)
        tss[init_fspace[ci.idx]] = 1:1
        push!(p_spaces, tss)
    end

    # define q spaces
    for tssp in p_spaces
        tss = FermiCG.get_ortho_compliment(tssp, cluster_bases[tssp.cluster.idx])
        push!(q_spaces, tss)
    end

    println(" ================= Cluster P Spaces ===================")
    display.(p_spaces)
    println(" ================= Cluster Q Spaces ===================")
    display.(q_spaces)

    nroots = 1
    ci_vector = FermiCG.TuckerState(clusters, p_spaces, q_spaces, na, nb)
    ref_vector = deepcopy(ci_vector)

    # for FOI space
    foi_space = FermiCG.define_foi_space(ref_vector, clustered_ham, nbody=2)
    ci_vector = FermiCG.TuckerState(clusters, p_spaces, q_spaces, foi_space)

    #
    # initialize with eye
    FermiCG.set_vector!(ref_vector, Matrix(1.0I, length(ref_vector),nroots))
    FermiCG.add!(ci_vector, ref_vector)

    #FermiCG.print_fock_occupations(ci_vector)

    cts_ref  = FermiCG.CompressedTuckerState(ref_vector, thresh=-1);

    cts = FermiCG.CompressedTuckerState(ci_vector, thresh=1e-5)
    #display(cts_fois)
    if false
        FermiCG.scale!(ci_vector, 1.0/sqrt(FermiCG.dot(ci_vector, ci_vector)[1]))
        println(" Length of CI Vector: ", length(ci_vector))
        @time e_nb2, ci_vector = FermiCG.tucker_ci_solve(ci_vector, cluster_ops, clustered_ham)
        @printf(" E(CI):   Electronic %16.12f Total %16.12f\n", e_nb2[1], e_nb2[1]+ints.h0)
        FermiCG.print_fock_occupations(ci_vector)

        println(" Now compress and resolve")
        cts = FermiCG.CompressedTuckerState(ci_vector, thresh=1e-5)
        #FermiCG.compress!(cts,thresh=1e-5)
        FermiCG.normalize!(cts)
        display(length(cts))
        @time e_cts, cts = FermiCG.tucker_ci_solve(cts, cluster_ops, clustered_ham, tol=1e-5)
        @printf(" E(cCI):  Electronic %16.12f Total %16.12f\n", e_cts[1], e_cts[1]+ints.h0)
        FermiCG.display(cts)
        #FermiCG.print_fock_occupations(cts)
    end

    if true
        @time e_ref, cts_ref = FermiCG.tucker_ci_solve(cts_ref, cluster_ops, clustered_ham, tol=1e-5)

        cts_var = deepcopy(cts_ref)
        e_var = 0.0
        e_pt2 = 0.0
        #display(abs.(cluster_ops[1]["H"][((2,2),(2,2))]) - abs.(cluster_ops[2]["H"][((2,2),(2,2))]))

        @time e_var, v_var = FermiCG.solve_for_compressed_space(cts_ref, cluster_ops, clustered_ham, nbody=4, thresh_var=1e-4, thresh_foi=1e-6, tol_ci=1e-5, tol_tucker=1e-5)
#        for i in 1:10
#            #@profilehtml e_var, e_pt2, cts_var = FermiCG.iterate_pt2!(cts_var, cluster_ops, clustered_ham, nbody=4, thresh=1e-7, tol=1e-5, do_pt=true)
#            @time e_var, e_pt2, cts_var = FermiCG.iterate_pt2!(cts_var, cluster_ops, clustered_ham, nbody=4, thresh=1e-7, tol=1e-5, do_pt=true, method="ci", ratio=1e3)
#            @printf(" E(Ref)      = %12.8f = %12.8f\n", e_ref[1], e_ref[1] + ints.h0 )
#            @printf(" E(PT2) tot  = %12.8f = %12.8f\n", e_ref[1]+e_pt2, e_ref[1]+e_pt2 + ints.h0 )
#            @printf(" E(var) tot  = %12.8f = %12.8f\n", e_var[1], e_var[1] + ints.h0 )
#
#            if abs(e_ref[1] - e_var[1]) < 1e-8
#                println("*Converged")
#                break
#            end
#            cts_ref = cts_var
#            e_ref = e_var
#        end
#        return cts_var
    end

    if false
        e_ref, ref_vector = FermiCG.tucker_ci_solve(ref_vector, cluster_ops, clustered_ham)
        println(" Reference State:" )
        FermiCG.print_fock_occupations(ref_vector)

        @time e_cepa, x_cepa = FermiCG.tucker_cepa_solve!(ref_vector, ci_vector, cluster_ops, clustered_ham)
        @printf(" E(CEPA): Electronic %16.12f Total %16.12f\n", e_cepa, e_cepa+ints.h0)
        println(e_cepa)
        FermiCG.print_fock_occupations(ci_vector)
    end

    function do_work(cts, cluster_ops, clustered_ham; nbody=3, thresh=1e-3, nfoi=3, n_iter=3)
        for iter in 1:n_iter
            for n in 1:nfoi
                cts  = FermiCG.open_sigma(cts, cluster_ops, clustered_ham, nbody=nbody, thresh=thresh)
            end
            FermiCG.normalize!(cts)
            display(length(cts))
            @time e_cts, v_cts = FermiCG.tucker_ci_solve!(cts, cluster_ops, clustered_ham, tol=1e-3)
            @printf(" E(cCI):  Electronic %16.12f Total %16.12f\n", e_cts[1], e_cts[1]+ints.h0)
        end
    end

    if false

        #
        #
        cts  = FermiCG.open_sigma(cts_ref, cluster_ops, clustered_ham, nbody=4, thresh=1e-4)
        FermiCG.normalize!(cts)
        display(length(cts))
        @time FermiCG.hylleraas_compressed_mp2!(cts, cluster_ops, clustered_ham, tol=1e-6, thresh=1e-6)
    end

    if false
        cts = cts_fois

        @time e_ref, v_ref = FermiCG.tucker_ci_solve!(cts_ref, cluster_ops, clustered_ham, tol=1e-12)
        tmp = deepcopy(cts_ref)
        FermiCG.scale!(tmp, -e_ref[1])

        cts_sig  = FermiCG.open_sigma(cts_ref, cluster_ops, clustered_ham, nbody=4, thresh=1e-4)
        @printf(" Should be E0  : %12.8f\n", FermiCG.nonorth_dot(cts_sig, cts_ref))
        FermiCG.nonorth_add!(cts_sig, tmp)
        @printf(" Should be zero: %12.8f\n", FermiCG.nonorth_dot(cts_sig, cts_ref))
        display(cts_sig)
        display(FermiCG.nonorth_dot(cts_sig, cts_ref))

        cts = cts_sig
        FermiCG.normalize!(cts)
        display(length(cts))
        @time e_cts, v_cts = FermiCG.tucker_ci_solve!(cts, cluster_ops, clustered_ham, tol=1e-3)
        @printf(" E(cCI):  Electronic %16.12f Total %16.12f\n", e_cts[1], e_cts[1]+ints.h0)
        display(cts)

        #@time do_work(cts, cluster_ops, clustered_ham, nbody=4, thresh=1e-2, nfoi=3, n_iter=10)
        #display(cts)
    end
    #end
#end

#run()

max_stars_repo_name: nmayhall-vt/FermiCG

max_issues_repo_name: nmayhall-vt/FermiCG

max_forks_repo_name: nmayhall-vt/FermiCG


text:
"""
    MOIError{F <: Function} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `f::F`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOIError(f, dim = 0) = begin
        #= none:5 =#
        new{typeof(f)}(f, dim)
    end`: DESCRIPTION
"""
struct MOIError{F <: Function} <: MOI.AbstractVectorSet
    f::F
    dimension::Int

    MOIError(f, dim = 0) = new{typeof(f)}(f, dim)
end

"""
    MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIError}) = begin

DOCSTRING

# Arguments:
- ``: DESCRIPTION
- ``: DESCRIPTION
- ``: DESCRIPTION
"""
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIError{F}}
) where {F <: Function}
    return true
end

"""
    MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIError)

DOCSTRING

# Arguments:
- `optimizer`: DESCRIPTION
- `vars`: DESCRIPTION
- `set`: DESCRIPTION
"""
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIError{F}) where {F <: Function}
    cidx = constraint!(optimizer, set.f, map(x -> x.value, vars.variables))
    return CI{VOV, MOIError{F}}(cidx)
end

"""
    Base.copy(set::MOIError) = begin

DOCSTRING
"""
Base.copy(set::MOIError) = MOIError(deepcopy(set.f), set.dimension)

"""
    Error{F <: Function} <: JuMP.AbstractVectorSet

The solver will compute a straightforward error function based on the `concept`. To run the solver efficiently, it is possible to provide an *error function* `err` instead of `concept`. `err` must return a nonnegative real number.

```julia
@constraint(model, X in Error(err))
```
"""
struct Error{F <: Function} <: JuMP.AbstractVectorSet
    f::F
end

# @autodoc
JuMP.moi_set(set::Error{F}, dim::Int) where {F <: Function} = MOIError(set.f, dim)

"""
    MOIPredicate{F <: Function} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `f::F`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOIPredicate(f, dim = 0) = begin
        #= none:5 =#
        new{typeof(f)}(f, dim)
    end`: DESCRIPTION
"""
struct MOIPredicate{F <: Function} <: MOI.AbstractVectorSet
    f::F
    dimension::Int

    MOIPredicate(f, dim = 0) = new{typeof(f)}(f, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIPredicate{F}}
) where {F <: Function}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIPredicate{F}) where {F <: Function}
    err = x -> convert(Float64, !set.f(x))
    cidx = constraint!(optimizer, err, map(x -> x.value, vars.variables))
    return CI{VOV, MOIPredicate{F}}(cidx)
end

Base.copy(set::MOIPredicate) = MOIPredicate(deepcopy(set.f), copy(set.dimension))

"""
    Predicate{F <: Function} <: JuMP.AbstractVectorSet

Assuming `X` is a (collection of) variables, `concept` a boolean function over `X`, and that a `model` is defined. In `JuMP` syntax we can create a constraint based on `concept` as follows.

```julia
@constraint(model, X in Predicate(concept))
```
"""
struct Predicate{F <: Function} <: JuMP.AbstractVectorSet
    f::F
end
JuMP.moi_set(set::Predicate, dim::Int) = MOIPredicate(set.f, dim)

"""
    MOIAllDifferent <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIAllDifferent <: MOI.AbstractVectorSet
    dimension::Int

    MOIAllDifferent(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIAllDifferent}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIAllDifferent)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:all_different])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIAllDifferent}(cidx)
end
Base.copy(set::MOIAllDifferent) = MOIAllDifferent(copy(set.dimension))

"""
Global constraint ensuring that all the values of a given configuration are unique.

```julia
@constraint(model, X in AllDifferent())
```
"""
struct AllDifferent <: JuMP.AbstractVectorSet end
JuMP.moi_set(::AllDifferent, dim::Int) = MOIAllDifferent(dim)

"""
    MOIAllEqual <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIAllEqual <: MOI.AbstractVectorSet
    dimension::Int

    MOIAllEqual(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIAllEqual}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIAllEqual)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:all_equal])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIAllEqual}(cidx)
end

Base.copy(set::MOIAllEqual) = MOIAllEqual(copy(set.dimension))

"""
Global constraint ensuring that all the values of `X` are all equal.

```julia
@constraint(model, X in AllEqual())
```
"""
struct AllEqual <: JuMP.AbstractVectorSet end
JuMP.moi_set(::AllEqual, dim::Int) = MOIAllEqual(dim)

"""
    MOIEq <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIEq <: MOI.AbstractVectorSet
    dimension::Int

    MOIEq(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIEq}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIEq)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:eq])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIEq}(cidx)
end

Base.copy(set::MOIEq) = MOIEq(copy(set.dimension))

"""
Equality between two variables.

```julia
@constraint(model, X in Eq())
```
"""
struct Eq <: JuMP.AbstractVectorSet end
JuMP.moi_set(::Eq, dim::Int) = MOIEq(dim)

"""
    MOIAlwaysTrue <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIAlwaysTrue <: MOI.AbstractVectorSet
    dimension::Int

    MOIAlwaysTrue(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIAlwaysTrue}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIAlwaysTrue)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:always_true])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIAlwaysTrue}(cidx)
end

Base.copy(set::MOIAlwaysTrue) = MOIAlwaysTrue(copy(set.dimension))

"""
Always return `true`. Mainly used for testing purpose.

```julia
@constraint(model, X in AlwaysTrue())
```
"""
struct AlwaysTrue <: JuMP.AbstractVectorSet end
JuMP.moi_set(::AlwaysTrue, dim::Int) = MOIAlwaysTrue(dim)

"""
    MOIOrdered <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIOrdered <: MOI.AbstractVectorSet
    dimension::Int

    MOIOrdered(dim = 0) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIOrdered}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIOrdered)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:ordered])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIOrdered}(cidx)
end

Base.copy(set::MOIOrdered) = MOIOrdered(copy(set.dimension))

"""
Global constraint ensuring that all the values of `x` are ordered.

```julia
@constraint(model, X in Ordered())
```
"""
struct Ordered <: JuMP.AbstractVectorSet end
JuMP.moi_set(::Ordered, dim::Int) = MOIOrdered(dim)

"""
    MOIDistDifferent <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOIDistDifferent <: MOI.AbstractVectorSet
    dimension::Int

    MOIDistDifferent(dim = 4) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIDistDifferent}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOIDistDifferent)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:dist_different])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIDistDifferent}(cidx)
end
Base.copy(set::MOIDistDifferent) = MOIDistDifferent(copy(set.dimension))

"""
Local constraint ensuring that, given a vector `X` of size 4, `|X[1] - X[2]| ≠ |X[3] - X[4]|)`.

```julia
@constraint(model, X in DistDifferent())
```
"""
struct DistDifferent <: JuMP.AbstractVectorSet end
JuMP.moi_set(::DistDifferent, dim::Int) = MOIDistDifferent(dim)

"""
    MOIAllEqualParam{T <: Number} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `param::T`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOIAllEqualParam(param, dim = 0) = begin
        #= none:5 =#
        new{typeof(param)}(param, dim)
    end`: DESCRIPTION
"""
struct MOIAllEqualParam{T <: Number} <: MOI.AbstractVectorSet
    param::T
    dimension::Int

    MOIAllEqualParam(param, dim = 0) = new{typeof(param)}(param, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIAllEqualParam{T}}
) where {T <: Number}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIAllEqualParam{T}
) where T <: Number
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=set.param, dom_size=max_dom_size) -> error_f(
        usual_constraints[:all_equal_param])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIAllEqualParam{T}}(cidx)
end

Base.copy(set::MOIAllEqualParam) = MOIAllEqualParam(copy(set.param), copy(set.dimension))

"""
Global constraint ensuring that all the values of `X` are all equal to a given parameter `param`.

```julia
@constraint(model, X in AllEqualParam(param))
```
"""
struct AllEqualParam{T <: Number} <: JuMP.AbstractVectorSet
    param::T
end
JuMP.moi_set(set::AllEqualParam, dim::Int) = MOIAllEqualParam(set.param, dim)

"""
    MOISumEqualParam{T <: Number} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `param::T`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOISumEqualParam(param, dim = 0) = begin
        #= none:5 =#
        new{typeof(param)}(param, dim)
    end`: DESCRIPTION
"""
struct MOISumEqualParam{T <: Number} <: MOI.AbstractVectorSet
    param::T
    dimension::Int

    MOISumEqualParam(param, dim = 0) = new{typeof(param)}(param, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOISumEqualParam{T}}
) where {T <: Number}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOISumEqualParam{T}
) where {T <: Number}
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=set.param, dom_size=max_dom_size) -> error_f(
        usual_constraints[:sum_equal_param])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOISumEqualParam{T}}(cidx)
end

Base.copy(set::MOISumEqualParam) = MOISumEqualParam(copy(set.param),
copy(set.dimension))

"""
Global constraint ensuring that the sum of the values of `X` is equal to a given parameter `param`.

```julia
@constraint(model, X in SumEqualParam(param))
```
"""
struct SumEqualParam{T <: Number} <: JuMP.AbstractVectorSet
    param::T
end
JuMP.moi_set(set::SumEqualParam, dim::Int) = MOISumEqualParam(set.param, dim)

"""
    MOILessThanParam{T <: Number} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `param::T`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOILessThanParam(param, dim = 0) = begin
        #= none:5 =#
        new{typeof(param)}(param, dim)
    end`: DESCRIPTION
"""
struct MOILessThanParam{T <: Number} <: MOI.AbstractVectorSet
    param::T
    dimension::Int

    MOILessThanParam(param, dim = 0) = new{typeof(param)}(param, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOILessThanParam{T}}
) where {T <: Number}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOILessThanParam{T}
) where {T <: Number}
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=set.param, dom_size=max_dom_size) -> error_f(
        usual_constraints[:less_than_param])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOILessThanParam{T}}(cidx)
end

Base.copy(set::MOILessThanParam) = MOILessThanParam(copy(set.param),
copy(set.dimension))

"""
Constraint ensuring that the value of `x` is less than a given parameter `param`.

```julia
@constraint(model, x in LessThanParam(param))
```
"""
struct LessThanParam{T <: Number} <: JuMP.AbstractVectorSet
    param::T
end
JuMP.moi_set(set::LessThanParam, dim::Int) = MOILessThanParam(set.param, dim)

"""
    MOIMinusEqualParam{T <: Number} <: MOI.AbstractVectorSet

DOCSTRING

# Arguments:
- `param::T`: DESCRIPTION
- `dimension::Int`: DESCRIPTION
- `MOIMinusEqualParam(param, dim = 0) = begin
        #= none:5 =#
        new{typeof(param)}(param, dim)
    end`: DESCRIPTION
"""
struct MOIMinusEqualParam{T <: Number} <: MOI.AbstractVectorSet
    param::T
    dimension::Int

    MOIMinusEqualParam(param, dim = 0) = new{typeof(param)}(param, dim)
end
function MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOIMinusEqualParam{T}}
) where {T <: Number}
    return true
end
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, set::MOIMinusEqualParam{T}
) where {T <: Number}
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=set.param, dom_size=max_dom_size) -> error_f(
        usual_constraints[:minus_equal_param])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOIMinusEqualParam{T}}(cidx)
end

Base.copy(set::MOIMinusEqualParam) = MOIMinusEqualParam(copy(set.param),
copy(set.dimension))

"""
Constraint ensuring that the value of `x` is less than a given parameter `param`.

```julia
@constraint(model, x in MinusEqualParam(param))
```
"""
struct MinusEqualParam{T <: Number} <: JuMP.AbstractVectorSet
    param::T
end
JuMP.moi_set(set::MinusEqualParam, dim::Int) = MOIMinusEqualParam(set.param, dim)

"""
    MOISequentialTasks <: MOI.AbstractVectorSet

DOCSTRING
"""
struct MOISequentialTasks <: MOI.AbstractVectorSet
    dimension::Int

    MOISequentialTasks(dim = 4) = new(dim)
end
MOI.supports_constraint(::Optimizer, ::Type{VOV}, ::Type{MOISequentialTasks}) = true
function MOI.add_constraint(optimizer::Optimizer, vars::MOI.VectorOfVariables, ::MOISequentialTasks)
    max_dom_size = max_domains_size(optimizer, map(x -> x.value, vars.variables))
    e = (x; param=nothing, dom_size=max_dom_size) -> error_f(
        usual_constraints[:sequential_tasks])(x; param=param, dom_size=dom_size
    )
    cidx = constraint!(optimizer, e, map(x -> x.value, vars.variables))
    return CI{VOV, MOISequentialTasks}(cidx)
end
Base.copy(set::MOISequentialTasks) = MOISequentialTasks(copy(set.dimension))

"""
Local constraint ensuring that, given a vector `X` of size 4, `|X[1] - X[2]| ≠ |X[3] - X[4]|)`.

```julia
@constraint(model, X in SequentialTasks())
```
"""
struct SequentialTasks <: JuMP.AbstractVectorSet end
JuMP.moi_set(::SequentialTasks, dim::Int) = MOISequentialTasks(dim)


text:
#
# Helpers and definitions for transformer networks:
#

"""
    function positional_encoding_sincos(n_embed, n_seq)

Calculate and return a matrix of size `[n_embed, n_seq]` of
positional encoding values
following the sin and cos style in the paper
*Vaswani, A. et al.; Attention Is All You Need;
31st Conference on Neural Information Processing Systems (NIPS 2017),
Long Beach, CA, USA, 2017.*
"""
function positional_encoding_sincos(n_embed, n_seq)

    angl = [1/(10000^(2*i/n_embed)) for i in 1:n_embed/2]
    angl = angl * permutedims(1:n_seq)
    pos_enc = vcat(sin.(angl), cos.(angl))
    return convert2KnetArray(pos_enc)
end

"""
    struct PositionalEncoding <: Layer

Positional encoding layer. Only *sincos*-style (according to
Vaswani, et al., NIPS 2017) is implemented.

The layer takes an array of any any number of dimensions (>=2), calculates
the Vaswani-2017-style positional encoding and adds the encoding to each plane
of the array.
"""
struct PositionalEncoding <: Layer
    style
    PositionalEncoding(;style=:sincos) = new(style)
end

function (l::PositionalEncoding)(x)
    # only one style implemented yet:
    # if l.style == sincos
        x = x .+ positional_encoding_sincos(size(x)[1], size(x)[2])
        return x
end

"""
    function mk_padding_mask(x; pad=3)

Make a padding mask; i.e. return an Array of type
`KnetArray{Float32}` (or `Array{Float32}`) similar to `x` but with
two additional dimension of size 1 in teh middle (this will represent the
2nd seq_len and the number of heads) in multi-head attention
and the
value `1.0` at each position where `x` is `pad` and `0.0` otherwise.

The function can be used for creating padding masks for attention
mechanisms.
"""
function mk_padding_mask(x; pad=3)

    return reshape(convert2KnetArray(x .== pad), size(x)[1],1,1,size(x)[2])
end

"""
    function mk_peek_ahead_mask(x; dim=1)

Return a matrix of size `[n_seq, n_seq]` filled with 1.0 and the *uppper triangle*
set to 0.0.
Type is `KnetArray{Float32}` in GPU context, `Array{Float32}` otherwise.
The matrix can be used as peek-ahead mask in transformers.

`dim=1` specifies the dimension in which the sequence length is
represented. For un-embedded data this is normally `1`, i.e. the
shape of x is [n_seq, n_mb]. After embedding the shape probably is
[depth, n_seq, n_mb].
"""
function mk_peek_ahead_mask(x; dim=1)

    n_seq = size(x)[dim]
    return convert2KnetArray(1 .- UpperTriangular(ones(n_seq, n_seq)))
end

"""
    function dot_prod_attn(q, k, v; mask=nothing)

Generic scaled dot product attention following the paper of
Vaswani et al., (2017), *Attention Is All You Need*.

### Arguments:
+ `q`: query of size `[depth, n_seq_q, ...]`
+ `k`: key of size `[depth, n_seq_v, ...]`
+ `v`: value of size `[depth, n_seq_v, ...]`
+ `mask`: mask for attention factors may have different shapes but must be
        broadcastable for addition to the scores tensor (which as the same size as
        alpha `[n_seq_v, n_seq_q, ...]`). In transformer context typical masks are one of:
        padding mask of size `[n_seq_v, ...]` or a peek-ahead mask of size `[n_seq_v, n_seq_v]`
        (which is only possible in case of self-attention when all seqencee lengths
        are identical).

`q, k, v` must have matching leading dimensions (i.e. same depth or embedding).
`k` and `v` must have the same sequence length.

### Return values:
+ `c`: context as alpha-weighted sum of values with size [depth, n_seq_v, ...]
+ `alpha`: attention factors of size [n_seq_v, n_seq_q, ...]
"""
function dot_prod_attn(q, k, v; mask=nothing)

    score = bmm(k, q, transA=true) ./ Float32(sqrt(size(k)[1]))  # [s_v x s_k x mb]

    if mask != nothing
        score = score .+ mask * Float32(-1e9)
    end

    α = softmax(score, dims=1)
    c = bmm(v, α)
    return c, α
end


text:
# The code to generate Fig 6B
using PyPlot
using SurNoR_2020
using Statistics
using HypothesisTests

PyPlot.svg(true)
rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
rcParams["svg.fonttype"] = "none"
rcParams["pdf.fonttype"] = 42

# ------------------------------------------------------------------------------
# Training Algorithms
# ------------------------------------------------------------------------------
Agents_SurNoR  = SurNoR_2020.Func_SurNoR_CV_train()
Agents_HybN    = SurNoR_2020.Func_HybN_CV_train()
Agents_Tot_Alg = cat(Agents_SurNoR, Agents_HybN, dims=4)

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# Fig 6B1 and Fig 6B2: Action preferences in state 3 and 7 in B2E1 and B2E2
# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
Block = 2
Sub_Num = size(Agents_Tot_Alg)[1]
Alg_Num = size(Agents_Tot_Alg)[4]

Map_S = [1,8,3,5,2,9,10,4,6,7]

n_visit = 1 # Looking at the 1st time visiting each state
for Epi = 1:2
    for State = [3,10]

        # A matrix to save participant's actions at each state
        Matrix_Data = zeros(Sub_Num,1)
        # A matrix to save action probabilities of different algorithms for each participant at each state
        Matrix_alg_P_act = zeros(Sub_Num,4,Alg_Num)
        Trap_set = [2,6,7]

        for Sub = 1:Sub_Num
            # Reading real data
            Agent = Agents_Tot_Alg[Sub,Block,Epi,1]

            Ind = 1:length(Agent.Input.obs)
            Ind = Ind[Agent.Input.obs .== State]

            if Epi >= 2
                Agent_old = Agents_Tot_Alg[Sub,Block,Epi-1,1]
                Ind_old = 1:length(Agent_old.Input.obs)
                Ind_old = Ind_old[Agent_old.Input.obs .== State]
            else
                Ind_old = [1]
            end

            if (length(Ind) < n_visit) || (length(Ind_old) < 1)
                Matrix_Data[Sub,:] .= NaN
                Matrix_alg_P_act[Sub,:,:] .= NaN
            else
                a1 = Agent.Input.act[Ind[n_visit]]  # Taken action
                Matrix_Data[Sub,1] = a1
                for Alg = 1:Alg_Num
                    Agent = Agents_Tot_Alg[Sub,Block,Epi,Alg]
                    Matrix_alg_P_act[Sub,:,Alg] = Agent.P_Action[:,Ind[n_visit]]'
                end
            end
        end

        # Choosing the ones who have been in the particular state of interest
        Ind = 1:Sub_Num
        Ind = Ind[isnan.(Matrix_Data[Ind,1]).== 0]

        Color = ["#474848","#8723BE","#B159E0"]
        fig = figure(figsize=(7,12)); ax = gca()
        x = 1:4
        y = [mean(Matrix_Data[Ind,1] .== 0),
             mean(Matrix_Data[Ind,1] .== 1),
             mean(Matrix_Data[Ind,1] .== 2),
             mean(Matrix_Data[Ind,1] .== 3)]
        ax.bar(x, y, align="center", color = Color[1])
        for Alg = 1:Alg_Num
            x = (1:4) .+ (Alg * 5)
            y = mean(Matrix_alg_P_act[Ind,:,Alg],dims=1)
            ax.bar(x, y[:], align="center", color = Color[Alg+1])
        end
        Legends = ["Data", "SurNoR", "Hyb+N"]
        ax.legend(Legends)

        x = 1:4
        y = [mean(Matrix_Data[Ind,1] .== 0),
             mean(Matrix_Data[Ind,1] .== 1),
             mean(Matrix_Data[Ind,1] .== 2),
             mean(Matrix_Data[Ind,1] .== 3)]
        dy = [std(Matrix_Data[Ind,1] .== 0),
             std(Matrix_Data[Ind,1] .== 1),
             std(Matrix_Data[Ind,1] .== 2),
             std(Matrix_Data[Ind,1] .== 3)] ./ sqrt(length(Ind))
        ax.errorbar(x,y[:],yerr=dy[:],color="k",linewidth=1,capsize=5,linestyle ="")
        for Alg = 1:Alg_Num
            x = (1:4) .+ (Alg * 5)
            y = mean(Matrix_alg_P_act[Ind,:,Alg],dims=1)
            dy = std(Matrix_alg_P_act[Ind,:,Alg],dims=1) ./ sqrt(length(Ind))
            ax.errorbar(x,y[:],yerr=dy[:],color="k",linewidth=1,capsize=5,linestyle ="")
        end
        ax.set_ylim([0,1])
        ax.set_xlim([0.25,14.75])
        ax.set_yticks([0,0.5,1])
        xticks = Array(1:4)
        for Alg = 1:Alg_Num
            append!(xticks,Array((1:4) .+ (Alg * 5)))
        end
        ax.set_xticks(xticks)
        ax.plot([0.25,14.75],[0.25,0.25],"--k")
        ax.set_xticklabels(repeat(["a1","a2","a3","a4"],3))
        ax.set_title(string("Visit number ", string(n_visit), " of state ", string(Map_S[State]), " in B", string(Block), "E", string(Epi), " - N=", string(length(Ind))))

        # Conducting the t-tests
        for Alg = 1:Alg_Num
            y1 = Matrix_alg_P_act[Ind,1,Alg]
            y2 = Matrix_alg_P_act[Ind,4,Alg]
            println("----------------------------------------")
            println("----------------------------------------")
            @show(Legends[Alg+1])
            @show(Map_S[State])
            @show(Block)
            @show(Epi)
            @show OneSampleTTest(y1,y2)
        end
    end
end


text:
# The multi distribution represents `n` independent draws of the
# distribution given as a parameter
struct Multi{T<:Distribution}
    n::Int
    distribution::T
end

function Distributions.logpdf(multi::Multi, samples)
    sum([logpdf(multi.distribution, sample) for sample in samples])
end

function Base.rand(multi::Multi)
    rand(multi.distribution, multi.n)
end


text:
## p-value adjustment methods ##

type Bonferroni <: PValueAdjustmentMethod
end

adjust(pvals, method::Bonferroni) = bonferroni(pvals)

function bonferroni{T<:AbstractFloat}(pValues::Vector{T})
    validPValues(pValues)
    return min(pValues * length(pValues), 1.)
end

type BenjaminiHochberg <: PValueAdjustmentMethod
end

adjust(pvals, method::BenjaminiHochberg) = benjamini_hochberg(pvals)

function benjamini_hochberg{T<:AbstractFloat}(pValues::Vector{T})
    validPValues(pValues)
    n = length(pValues)
    if n <= 1
        return pValues
    end
    sortedIndexes, originalOrder = reorder(pValues)
    sortedPValues = pValues[sortedIndexes]
    stepup!(sortedPValues, bejamini_hochberg_multiplier, n)
    min(sortedPValues[originalOrder], 1.)
end

bejamini_hochberg_multiplier(i::Int, n::Int) = n/(n-i)

type BenjaminiHochbergAdaptive <: PValueAdjustmentMethod
    pi0estimator::Pi0Estimator
end

## default to BenjaminiHochberg
BenjaminiHochbergAdaptive(π0::AbstractFloat) = BenjaminiHochbergAdaptive(Oracle(π0))

BenjaminiHochbergAdaptive() = BenjaminiHochbergAdaptive(1.0)

function adjust(pvals, method::BenjaminiHochbergAdaptive)
  π0 = estimate_pi0(pvals, method.pi0estimator)
  benjamini_hochberg(pvals, π0)
end

function benjamini_hochberg{T<:AbstractFloat}(pValues::Vector{T}, pi0::T)
    validPValues([pi0])
    benjamini_hochberg(pValues) .* pi0
end

type BenjaminiYekutieli <: PValueAdjustmentMethod
end

adjust(pvals, method::BenjaminiYekutieli) = benjamini_yekutieli(pvals)

function benjamini_yekutieli{T<:AbstractFloat}(pValues::Vector{T})
    validPValues(pValues)
    n = length(pValues)
    if n <= 1
        return pValues
    end
    sortedIndexes, originalOrder = reorder(pValues)
    sortedPValues = pValues[sortedIndexes]
    stepup!(sortedPValues, benjamini_yekutieli_multiplier, n)
    min(sortedPValues[originalOrder], 1.)
end

function benjamini_yekutieli_multiplier(i::Int, n::Int)
    c = sum([1/i for i in 1:n])
    return ((n*c)/(n-i))
end

type BenjaminiLiu <: PValueAdjustmentMethod
end

adjust(pvals, method::BenjaminiLiu) = benjamini_liu(pvals)

function benjamini_liu{T<:AbstractFloat}(pvalues::Vector{T})
    validPValues(pvalues)
    n = length(pvalues)
    if n <= 1
        return pvalues
    end
    sortedIndexes, originalOrder = reorder(pvalues)
    sortedPValues = pvalues[sortedIndexes]
    general_stepdown!(sortedPValues, benjamini_liu_step, n)
    min(sortedPValues[originalOrder], 1.)
end

function benjamini_liu_step{T<:AbstractFloat}(p::T, i::Int, n::Int)
    # a bit more involved because cutoffs at significance α have the form:
    # P_(i) <= 1- [1 - min(1, m/(m-i+1)α)]^{1/(m-i+1)}
    adjp = (1-(1-p)^(n-i+1))*(n-i+1)/n
    if adjp*n/(n-i+1) > 1
        return 0.0
    else
        return adjp
    end
end

type Hochberg <: PValueAdjustmentMethod
end

adjust(pvals, method::Hochberg) = hochberg(pvals)

function hochberg{T<:AbstractFloat}(pValues::Vector{T})
    validPValues(pValues)
    n = length(pValues)
    if n <= 1
        return pValues
    end
    sortedIndexes, originalOrder = reorder(pValues)
    sortedPValues = pValues[sortedIndexes]
    stepup!(sortedPValues, hochberg_multiplier, n)
    min(sortedPValues[originalOrder], 1.)
end

hochberg_multiplier(i::Int, n::Int) = (i+1)

type Holm <: PValueAdjustmentMethod
end

adjust(pvals, method::Holm) = holm(pvals)

function holm{T<:AbstractFloat}(pValues::Vector{T})
    validPValues(pValues)
    n = length(pValues)
    if n <= 1
        return pValues
    end
    sortedIndexes, originalOrder = reorder(pValues)
    sortedPValues = pValues[sortedIndexes]
    stepdown!(sortedPValues, holm_multiplier, n)
    min(sortedPValues[originalOrder], 1.)
end

holm_multiplier(i::Int, n::Int) = (n-i+1)

type Hommel <: PValueAdjustmentMethod
end

adjust(pvals, method::Hommel) = hommel(pvals)

function hommel{T<:AbstractFloat}(pValues::Vector{T})
    validPValues(pValues)
    n = length(pValues)
    if n <= 1
        return pValues
    end
    sortedIndexes, originalOrder = reorder(pValues)
    sortedPValues = pValues[sortedIndexes]
    q = fill(minimum(n .* pValues./[1:n; ]), n)
    pa = fill(q[1], n)
    for j in (n-1):-1:2
        ij = 1:(n-j+1)
        i2 = (n-j+2):n
        q1 = minimum(j .* sortedPValues[i2]./([2:j; ]))
        q[ij] = min(j .* sortedPValues[ij], q1)
        q[i2] = q[n-j+1]
        pa = max(pa, q)
    end
    max(pa, sortedPValues)[originalOrder]
end

type Sidak <: PValueAdjustmentMethod
end

adjust(pvals, method::Sidak) = sidak(pvals)

function sidak{T<:AbstractFloat}(pValues::Vector{T})
    validPValues(pValues)
    return min(1-(1-pValues).^length(pValues), 1.)
end

type ForwardStop <: PValueAdjustmentMethod
end

adjust(pvals, method::ForwardStop) = forwardstop(pvals)

function forwardstop{T<:AbstractFloat}(pvalues::Vector{T})
    validPValues(pvalues)
    n = length(pvalues)
    logsums = - cumsum(log(1-pvalues))
    stepup!(logsums, forwardstop_multiplier, n)
    max(min(logsums, 1.),0.)
end

forwardstop_multiplier(i::Int, n::Int) = 1/(n-i)


text:
module GeneralizedMetropolisHastings

import PDMats
import Distributions
import Sundials

import Base: ==, size, length, eltype, show, display, time, similar, copy!, copy
import Base: mean, median, std, quantile, var, rand, rand!, scale!

import Distributions: pdf,logpdf!,logpdf,location,location!,scale

import Compat: @compat,view

export
    ###Abstract Types
    AbstractPolicyTrait,AbstractPolicy,
    AbstractParameter,
    AbstractData,
    AbstractNoiseModel,
    AbstractSample,
    AbstractProposalDensity,SymmetricDensity,ASymmetricDensity,
    AbstractSampler,AbstractSamplerState,
    AbstractTuner,AbstractTunerState,
    AbstractModel,
    AbstractChain,
    AbstractIndicatorMatrix,
    AbstractJobSegment,AbstractRemoteSegments,
    AbstractRunner,AbstractMHRunner,
    ###Functions
    trait,traitvalue,traittype,
    policy,
    parameters,
    data,dataindex,datavalues,numvalues,numvars,
    distribution,distributions,
    density,
    noise,
    sampler,
    tuner,
    model,evaluate!,applynoise!,measurements,numparas,
    runner,
    run!,
    samples,numsamples,loglikelihood,logprior,logposterior,
    print_gmh_module_loaded

include("util/util.jl")
include("policies/traits.jl")
include("policies/policies.jl")
include("policies/MHRuntimePolicy.jl")
include("parameters/parameters.jl")
include("data/data.jl")
include("data/DataArray.jl")
include("data/DataFunction.jl")
include("distributions/distributions.jl")
include("distributions/Bactrian.jl")
include("proposals/proposals.jl")
include("proposals/DistributionWrapper.jl")
include("noise/noise.jl")
include("noise/NoiseModelGaussian.jl")
include("samples/samples.jl")
include("samples/base.jl")
include("samples/gradient.jl")
include("samples/tensor.jl")
include("samplers/samplers.jl")
include("samplers/MetropolisHastings.jl")
include("samplers/AdaptiveMetropolis.jl")
include("samplers/SmMALA.jl")
include("indicators/indicators.jl")
include("indicators/IndicatorStationary.jl")
include("models/models.jl")
include("models/geometry.jl")
include("models/TargetModel.jl")
include("models/ODEModel.jl")
include("tuners/tunefunctions.jl")
include("tuners/tuners.jl")
include("tuners/MonitorTuner.jl")
include("tuners/ScaleTuner.jl")
include("chains/chains.jl")
include("chains/ChainStandard.jl")
include("chains/ChainGradient.jl")
include("jobsegments/jobsegments.jl")
include("jobsegments/GMHSegment.jl")
include("jobsegments/remotesegments.jl")
include("jobsegments/MHRemoteSegments.jl")
include("runners/runners.jl")
include("runners/MHRunner.jl")
include("runners/SMHRunner.jl")
include("runners/GMHRunner.jl")

function print_gmh_module_loaded()
  println("$module_name(current_module()) module loaded successfully")
end

end # module


text:
module CaretLearners

export CaretLearner,fit!,transform!
export caretrun

using TSML.TSMLTypes
using TSML.Utils
using DataFrames

import TSML.TSMLTypes.fit!   # importing to overload
import TSML.TSMLTypes.transform! # importing to overload

using RCall
R"library(caret)"
R"library(e1071)"
R"library(gam)"
R"library(randomForest)"
R"library(nnet)"
R"library(kernlab)"
R"library(grid)"
R"library(MASS)"
R"library(pls)"
#R"install.packages('xgboost',repos='https://cloud.r-project.org')"
#R"install.packages('pls',repos='https://cloud.r-project.org')"
#R"install.packages('caret', dependencies = c('Depends', 'Suggests'))"
R"library(xgboost)"

mutable struct CaretLearner <: TSLearner
    model
    args

    function CaretLearner(args=Dict())
        #fitControl=:(R"trainControl(method = 'repeatedcv',number = 5,repeats = 5)")
        fitControl="trainControl(method = 'cv',number = 5,repeats = 5)"
        default_args = Dict(
            :output => :class,
            :learner => "rf",
            :fitControl=>fitControl,
            :impl_args => Dict()
        )
        new(nothing,mergedict(default_args,args))
    end
end

function fit!(crt::CaretLearner,x::T,y::Vector) where  {T<:Union{Vector,Matrix}}
    xx = x |> DataFrame
    yy = y |> Vector
    rres = rcall(:train,xx,yy,method=crt.args[:learner],trControl = reval(crt.args[:fitControl]))
    #crt.model = R"$rres$finalModel"
    crt.model = rres
end

function transform!(crt::CaretLearner,x::T) where  {T<:Union{Vector,Matrix}}
    xx = x |> DataFrame
    res = rcall(:predict,crt.model,xx) #in robj
    return rcopy(res) # return extracted robj
end

function caretrun()
    crt = CaretLearner(Dict(:learner=>"rf",:fitControl=>"trainControl(method='cv')"))
    iris=getiris()
    x=iris[:,1:4]  |> Matrix
    y=iris[:,5] |> Vector
    fit!(crt,x,y)
    print(crt.model)
    transform!(crt,x)
end

end


text:
### Model
m = SmetsWoutersOrig()
TTT, RRR, CCC = solve(m)
meas = measurement(m, TTT, RRR, CCC)
sys = compute_system(m)

nothing


text:
using ModelingToolkit, LinearAlgebra, SparseArrays

# Define the constants for the PDE
const α₂ = 1.0
const α₃ = 1.0
const β₁ = 1.0
const β₂ = 1.0
const β₃ = 1.0
const r₁ = 1.0
const r₂ = 1.0
const _DD = 100.0
const γ₁ = 0.1
const γ₂ = 0.1
const γ₃ = 0.1
const N = 8
const X = reshape([i for i in 1:N for j in 1:N],N,N)
const Y = reshape([j for i in 1:N for j in 1:N],N,N)
const α₁ = 1.0.*(X.>=4*N/5)

const Mx = Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1])
const My = copy(Mx)
Mx[2,1] = 2.0
Mx[end-1,end] = 2.0
My[1,2] = 2.0
My[end,end-1] = 2.0

# Define the initial condition as normal arrays
@variables du[1:N,1:N,1:3] u[1:N,1:N,1:3] MyA[1:N,1:N] AMx[1:N,1:N] DA[1:N,1:N]

# Define the discretized PDE as an ODE function
function f(du,u,p,t)
   A = @view  u[:,:,1]
   B = @view  u[:,:,2]
   C = @view  u[:,:,3]
  dA = @view du[:,:,1]
  dB = @view du[:,:,2]
  dC = @view du[:,:,3]
  mul!(MyA,My,A)
  mul!(AMx,A,Mx)
  @. DA = _DD*(MyA + AMx)
  @. dA = DA + α₁ - β₁*A - r₁*A*B + r₂*C
  @. dB = α₂ - β₂*B - r₁*A*B + r₂*C
  @. dC = α₃ - β₃*C + r₁*A*B - r₂*C
end

f(du,u,nothing,0.0)

multithreadedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.MultithreadedForm())[2])
_du = rand(N,N,3)
_u = rand(N,N,3)
multithreadedf(_du,_u)

using Distributed
addprocs(4)
distributedf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DistributedForm())[2])

using Dagger
daggerf = eval(ModelingToolkit.build_function(du,u,parallel=ModelingToolkit.DaggerForm())[2])

jac = sparse(ModelingToolkit.jacobian(vec(du),vec(u),simplify=false))
serialjac = eval(ModelingToolkit.build_function(vec(jac),u)[2])
multithreadedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.MultithreadedForm())[2])
distributedjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.DistributedForm())[2])
daggerjac = eval(ModelingToolkit.build_function(vec(jac),u,parallel=ModelingToolkit.DaggerForm())[2])

MyA = zeros(N,N)
AMx = zeros(N,N)
DA  = zeros(N,N)

f(_du,_u,nothing,0.0)
multithreadedf(_du,_u)
#distributedf(_du,_u)
#daggerf(_du,_u)

#=
using BenchmarkTools
@btime f(_du,_u,nothing,0.0)
@btime multithreadedf(_du,_u)
@btime distributedf(_du,_u)
@btime daggerf(_du,_u)

_jac = similar(jac,Float64)
@btime serialjac(_jac,_u)
@btime multithreadedjac(_jac,_u)
@btime distributedjac(_jac,_u)
@btime daggerjac(_jac,_u)
=#