### A Pluto.jl notebook ###
# v0.19.5

using Markdown
using InteractiveUtils

# ╔═╡ fbce1e27-45db-465d-83f0-a0a7935abf14
using ActuaryUtilities

# ╔═╡ 2f700532-bc7b-4408-8268-f9f632a75451
using DataFrames

# ╔═╡ a3a5b911-1097-454a-b0ec-463874be92b4
using Setfield

# ╔═╡ ba85bf56-5401-4430-a0e2-eef7e06347a5
md" # Nested Projection Mechanics

A simple example of how one could define a nested projection system. Includes the following examples:
- Outer loop policy projections only
- Outer + Inner loop policy projections with padded cashflows determining reserves and capital
- Outer + Inner loop with a stochastic interest rate for the reserves

"

# ╔═╡ 1f4e28e0-c2f0-4c06-a9f8-d0d71999114a
md" ## Policy Data & Methods"

# ╔═╡ 001d049c-efd2-11ec-0083-cff70f84b072
abstract type Policy end

# ╔═╡ c7dbfd5b-a97e-4628-a2f8-8651f12c60f6
struct Term <: Policy
	inforce::Float64 # count of inforce
	term::Int        # length of benefit period (in months?)
	premium::Float64
	face::Float64
end

# ╔═╡ 50569a1e-5deb-48b6-be72-5be41dcb5b43
md"""
For consistency, use the function signature of:

```
(policy,assumptions,time) -> result
```
"""

# ╔═╡ 280c7ef0-61d5-4d28-87b2-d26855c76b6b
premiums(pol::Term,assumptions,time) = pol.inforce * pol.premium / 12

# ╔═╡ 741777ff-2d40-4393-a0d4-669f3460cc64
qx(pol::Term,assumptions,time)= assumptions.q / 12

# ╔═╡ b8fe0af1-29b7-419a-ad3e-7b120eb5cbe1
deaths(pol::Term,assumptions,time) = pol.inforce * qx(pol,assumptions,time)

# ╔═╡ 35bed858-6b76-4dbb-b49a-505f36b13434
claims(pol::Term,assumptions,time) = deaths(pol,assumptions,time) * pol.face

# ╔═╡ 1f59e516-4dc1-450e-87a4-9ce6ac64029d
md" ## Projection Assumptions"

# ╔═╡ dde19062-8a22-4caa-ae68-23895628dc3e
assumptions = (
	q = 0.012,
	int_reserve = 0.02,
	capital_factor = 0.1, # rate * reserves
)

# ╔═╡ 18e32b95-2d13-4502-8500-48b0c380291c
innerloop_assumption(outer_assump) = @set outer_assump.q *= 1.2

# ╔═╡ 10c9e253-2913-4f85-9989-8b3fffba423e
md"## Projection Logic"

# ╔═╡ ac656984-3bdd-4a33-90be-45042e0080a5
"""
	project(policy,assumptions;start_time=1,additional_processing=res)

The kwarg `additional_processing` defines an intermediate processing step where one can take the current model state and perform additional work, including nested projections. If left out of the arguments, the default for `additional_processing` is `res`, where `res` is (pol,assumpitons, result)->result (ie will just return the model point's results with no additional work being done).

"""
function project(
		pol::Term,
		assumptions;
		start_time=1,
		additional_processing=(pol,assumpitons, result)->result
	)
	A = assumptions
	map(start_time:pol.term) do t
		timestep = t
		premium = premiums(pol,A,t)
		q = qx(pol,A,t)
		death = deaths(pol,A,t)
		claim = claims(pol,A,t)
		net_cf = premium-claim
		inforce = pol.inforce - death
		pol = @set pol.inforce = inforce

		# return a vector of name tuples with the results
		result = (;
			timestep,
			premium,
			death,
			claim,
			net_cf,
			inforce,
			q,
		)

		# apply additional processing function
		additional_processing(pol,A,result)
	end
end

# ╔═╡ 1a903be4-57fb-466c-96e3-f2ad36acc77c
# Function signature: (policy, assumptions, result) -> updated result
function run_inner(policy,assumptions,result)
	A = innerloop_assumption(assumptions)
	p = project(policy,A;start_time=result.timestep+1)

	reserves = -pv(A.int_reserve,[modelpoint.net_cf for modelpoint in p])
	capital = reserves * A.capital_factor
	additional_results = (; reserves, capital)
	return merge(result,additional_results)

end

# ╔═╡ 324be56b-512b-4664-9f5b-943b6d9aeade
# Function signature: (policy, assumptions, result) -> updated result
function run_inner_stochastic(policy,assumptions,result)
	A = innerloop_assumption(assumptions)
	p = project(policy,A;start_time=result.timestep+1)

	# simple stochastic interest rate
	n = 100

	reserves = let
		i = A.int_reserve
		f = pv(i + 0.005 * randn(),[modelpoint.net_cf for modelpoint in p])

		-sum(f for _ in 1:n) / n

	end
	capital = reserves * A.capital_factor
	additional_results = (; reserves, capital)
	return merge(result,additional_results)

end

# ╔═╡ 5c13430a-96c4-447f-89a7-1fe99078feee
md" ## Projections

First, define a sample policy:"

# ╔═╡ be7bbb76-fc0f-40b5-982b-7e46126f9c28
p = Term(1.,120,1300.,100_000.)

# ╔═╡ b5013701-848a-4920-b4f5-360815a9f2e2
md" A projection without any additional processing:"

# ╔═╡ 1d59569f-c56b-45d0-ae57-ead57f68ed7b
project(p,outerloop_assumptions) |> DataFrame

# ╔═╡ 7dd4015f-6b58-4cf2-af3f-30bf3d43d468
md" And an example which uses a PADed inner loop to determine the resserves and capital:"

# ╔═╡ 2e7a438e-435f-4dc1-86c2-947ea8ceb8d4
project(p,assumptions;additional_processing = run_inner) |> DataFrame

# ╔═╡ 9e824c6d-a303-427f-973c-bc1b9658709d
project(p,assumptions;additional_processing = run_inner) |> DataFrame

# ╔═╡ 973a5a58-9b42-4532-be80-21f18fb6791e
project(p,assumptions;additional_processing = run_inner_stochastic) |> DataFrame

# ╔═╡ 98720096-e605-40e9-a32f-d1d5ade04be6
md"## Endnotes

### Further Work

This example is simple, but could be greatly optimized to reduce intermediate variable allocations, refine the timing of cashflows, add additional decrements, handle different types of `Policy`s, abstract some of the projection mechanics into an `Iterable` object, etc.

### Disclaimer
Created as a proof of concept and not indended to be interpreted as a meaningful projection."

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
ActuaryUtilities = "bdd23359-8b1c-4f88-b89b-d11982a786f4"
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"

[compat]
ActuaryUtilities = "~3.4.0"
DataFrames = "~1.3.4"
Setfield = "~1.0.0"
"""

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

julia_version = "1.7.0"
manifest_format = "2.0"

[[deps.ActuaryUtilities]]
deps = ["Dates", "Distributions", "ForwardDiff", "QuadGK", "Roots", "StatsBase", "Yields"]
git-tree-sha1 = "4e3ff9120b210afe7b43ad8967c64d36e28f5a69"
uuid = "bdd23359-8b1c-4f88-b89b-d11982a786f4"
version = "3.4.0"

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

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

[[deps.ArrayInterfaceCore]]
deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
git-tree-sha1 = "d618d3cf75e8ed5064670e939289698ecf426c7f"
uuid = "30b0a656-2188-435a-8636-2ec0e6a096e2"
version = "0.1.12"

[[deps.ArrayLayouts]]
deps = ["FillArrays", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "ce5666a4081cfb66035a1a46e45a34ce889bac26"
uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
version = "0.8.7"

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

[[deps.AxisAlgorithms]]
deps = ["LinearAlgebra", "Random", "SparseArrays", "WoodburyMatrices"]
git-tree-sha1 = "66771c8d21c8ff5e3a93379480a2307ac36863f7"
uuid = "13072b0f-2c55-5437-9ae7-d433b7a33950"
version = "1.0.1"

[[deps.BSplineKit]]
deps = ["ArrayLayouts", "BandedMatrices", "FastGaussQuadrature", "Interpolations", "LazyArrays", "LinearAlgebra", "Random", "Reexport", "SparseArrays", "StaticArrays"]
git-tree-sha1 = "651e79138b30580c5a1586b67c3298be6c7b8afe"
uuid = "093aae92-e908-43d7-9660-e50ee39d5a0a"
version = "0.8.4"

[[deps.BandedMatrices]]
deps = ["ArrayLayouts", "FillArrays", "LinearAlgebra", "Random", "SparseArrays"]
git-tree-sha1 = "019aa88766e2493c59cbd0a9955e1bac683ffbcd"
uuid = "aae01518-5342-5314-be14-df237901396f"
version = "0.16.13"

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

[[deps.Bzip2_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "19a35467a82e236ff51bc17a3a44b69ef35185a2"
uuid = "6e34b625-4abd-537c-b88f-471c36dfa7a0"
version = "1.0.8+0"

[[deps.CEnum]]
git-tree-sha1 = "eb4cb44a499229b3b8426dcfb5dd85333951ff90"
uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"
version = "0.4.2"

[[deps.Calculus]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "f641eb0a4f00c343bbc32346e1217b86f3ce9dad"
uuid = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
version = "0.5.1"

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

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

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

[[deps.ColorVectorSpace]]
deps = ["ColorTypes", "FixedPointNumbers", "LinearAlgebra", "SpecialFunctions", "Statistics", "TensorCore"]
git-tree-sha1 = "d08c20eef1f2cbc6e60fd3612ac4340b89fea322"
uuid = "c3611d14-8923-5661-9e6a-0046d554d3a4"
version = "0.9.9"

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

[[deps.CommonSolve]]
git-tree-sha1 = "332a332c97c7071600984b3c31d9067e1a4e6e25"
uuid = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
version = "0.2.1"

[[deps.CommonSubexpressions]]
deps = ["MacroTools", "Test"]
git-tree-sha1 = "7b8a93dba8af7e3b42fecabf646260105ac373f7"
uuid = "bbf7d656-a473-5ed7-a52c-81e309532950"
version = "0.3.0"

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

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

[[deps.ConstructionBase]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "f74e9d5388b8620b4cee35d4c5a618dd4dc547f4"
uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
version = "1.3.0"

[[deps.Contour]]
deps = ["StaticArrays"]
git-tree-sha1 = "9f02045d934dc030edad45944ea80dbd1f0ebea7"
uuid = "d38c429a-6771-53c6-b99e-75d170b6e991"
version = "0.5.7"

[[deps.Crayons]]
git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.1.1"

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

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

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

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

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

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

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

[[deps.DiffResults]]
deps = ["StaticArrays"]
git-tree-sha1 = "c18e98cba888c6c25d1c3b048e4b3380ca956805"
uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
version = "1.0.3"

[[deps.DiffRules]]
deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"]
git-tree-sha1 = "28d605d9a0ac17118fe2c5e9ce0fbb76c3ceb120"
uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
version = "1.11.0"

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

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

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

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

[[deps.DualNumbers]]
deps = ["Calculus", "NaNMath", "SpecialFunctions"]
git-tree-sha1 = "5837a837389fccf076445fce071c8ddaea35a566"
uuid = "fa6b7ba4-c1ee-5f82-b5fc-ecf0adba8f74"
version = "0.6.8"

[[deps.EarCut_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "3f3a2501fa7236e9b911e0f7a588c657e822bb6d"
uuid = "5ae413db-bbd1-5e63-b57d-d24a61df00f5"
version = "2.2.3+0"

[[deps.FastGaussQuadrature]]
deps = ["LinearAlgebra", "SpecialFunctions", "StaticArrays"]
git-tree-sha1 = "58d83dd5a78a36205bdfddb82b1bb67682e64487"
uuid = "442a2c76-b920-505d-bb47-c5924d526838"
version = "0.4.9"

[[deps.FileIO]]
deps = ["Pkg", "Requires", "UUIDs"]
git-tree-sha1 = "9267e5f50b0e12fdfd5a2455534345c4cf2c7f7a"
uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
version = "1.14.0"

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

[[deps.FiniteDiff]]
deps = ["ArrayInterfaceCore", "LinearAlgebra", "Requires", "SparseArrays", "StaticArrays"]
git-tree-sha1 = "a0700c21266b55bf62c22e75af5668aa7841b500"
uuid = "6a86dc24-6348-571c-b903-95158fe2bd41"
version = "2.12.1"

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

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

[[deps.ForwardDiff]]
deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "LogExpFunctions", "NaNMath", "Preferences", "Printf", "Random", "SpecialFunctions", "StaticArrays"]
git-tree-sha1 = "2f18915445b248731ec5db4e4a17e451020bf21e"
uuid = "f6369f11-7733-5829-9624-2563aa707210"
version = "0.10.30"

[[deps.FreeType]]
deps = ["CEnum", "FreeType2_jll"]
git-tree-sha1 = "cabd77ab6a6fdff49bfd24af2ebe76e6e018a2b4"
uuid = "b38be410-82b0-50bf-ab77-7b57e271db43"
version = "4.0.0"

[[deps.FreeType2_jll]]
deps = ["Artifacts", "Bzip2_jll", "JLLWrappers", "Libdl", "Pkg", "Zlib_jll"]
git-tree-sha1 = "87eb71354d8ec1a96d4a7636bd57a7347dde3ef9"
uuid = "d7e528f0-a631-5988-bf34-fe36492bcfd7"
version = "2.10.4+0"

[[deps.FreeTypeAbstraction]]
deps = ["ColorVectorSpace", "Colors", "FreeType", "GeometryBasics"]
git-tree-sha1 = "b5c7fe9cea653443736d264b85466bad8c574f4a"
uuid = "663a7486-cb36-511b-a19d-713bb74d65c9"
version = "0.9.9"

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

[[deps.GeometryBasics]]
deps = ["EarCut_jll", "IterTools", "LinearAlgebra", "StaticArrays", "StructArrays", "Tables"]
git-tree-sha1 = "83ea630384a13fc4f002b77690bc0afeb4255ac9"
uuid = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
version = "0.4.2"

[[deps.HypergeometricFunctions]]
deps = ["DualNumbers", "LinearAlgebra", "SpecialFunctions", "Test"]
git-tree-sha1 = "cb7099a0109939f16a4d3b572ba8396b1f6c7c31"
uuid = "34004b35-14d8-5ef3-9330-4cdb6864b03a"
version = "0.3.10"

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

[[deps.Interpolations]]
deps = ["AxisAlgorithms", "ChainRulesCore", "LinearAlgebra", "OffsetArrays", "Random", "Ratios", "Requires", "SharedArrays", "SparseArrays", "StaticArrays", "WoodburyMatrices"]
git-tree-sha1 = "b7bc05649af456efc75d178846f47006c2c4c3c7"
uuid = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
version = "0.13.6"

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

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

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

[[deps.IterTools]]
git-tree-sha1 = "fa6287a4469f5e048d763df38279ee729fbd44e5"
uuid = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
version = "1.4.0"

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

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

[[deps.LazyArrays]]
deps = ["ArrayLayouts", "FillArrays", "LinearAlgebra", "MacroTools", "MatrixFactorizations", "SparseArrays", "StaticArrays"]
git-tree-sha1 = "721bebe4d0f8581c18fccf272c62000e22a80a2d"
uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
version = "0.22.10"

[[deps.LazyModules]]
git-tree-sha1 = "a560dd966b386ac9ae60bdd3a3d3a326062d3c3e"
uuid = "8cdb02fc-e678-4876-92c5-9defec4f444e"
version = "0.3.1"

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

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

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

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

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

[[deps.LineSearches]]
deps = ["LinearAlgebra", "NLSolversBase", "NaNMath", "Parameters", "Printf"]
git-tree-sha1 = "f27132e551e959b3667d8c93eae90973225032dd"
uuid = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
version = "7.1.1"

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

[[deps.LogExpFunctions]]
deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "09e4b894ce6a976c354a69041a04748180d43637"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.15"

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

[[deps.LsqFit]]
deps = ["Distributions", "ForwardDiff", "LinearAlgebra", "NLSolversBase", "OptimBase", "Random", "StatsBase"]
git-tree-sha1 = "91aa1442e63a77f101aff01dec5a821a17f43922"
uuid = "2fda8390-95c7-5789-9bda-21331edee243"
version = "0.12.1"

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

[[deps.MarchingCubes]]
deps = ["StaticArrays"]
git-tree-sha1 = "3bf4baa9df7d1367168ebf60ed02b0379ea91099"
uuid = "299715c1-40a9-479a-aaf9-4a633d36f717"
version = "0.1.3"

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

[[deps.MatrixFactorizations]]
deps = ["ArrayLayouts", "LinearAlgebra", "Printf", "Random"]
git-tree-sha1 = "2212d36f97e01347adb1460a6914e20f2feee853"
uuid = "a3b82374-2e81-5b9e-98ce-41277c0e4c87"
version = "0.9.1"

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

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

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

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

[[deps.NLSolversBase]]
deps = ["DiffResults", "Distributed", "FiniteDiff", "ForwardDiff"]
git-tree-sha1 = "50310f934e55e5ca3912fb941dec199b49ca9b68"
uuid = "d41bc354-129a-5804-8e4c-c37616107c6c"
version = "7.8.2"

[[deps.NaNMath]]
git-tree-sha1 = "b086b7ea07f8e38cf122f5016af580881ac914fe"
uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
version = "0.3.7"

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

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

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

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

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

[[deps.Optim]]
deps = ["Compat", "FillArrays", "ForwardDiff", "LineSearches", "LinearAlgebra", "NLSolversBase", "NaNMath", "Parameters", "PositiveFactorizations", "Printf", "SparseArrays", "StatsBase"]
git-tree-sha1 = "7a28efc8e34d5df89fc87343318b0a8add2c4021"
uuid = "429524aa-4258-5aef-a3af-852621145aeb"
version = "1.7.0"

[[deps.OptimBase]]
deps = ["NLSolversBase", "Printf", "Reexport"]
git-tree-sha1 = "9cb1fee807b599b5f803809e85c81b582d2009d6"
uuid = "87e2bd06-a317-5318-96d9-3ecbac512eee"
version = "2.0.2"

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

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

[[deps.Parameters]]
deps = ["OrderedCollections", "UnPack"]
git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe"
uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a"
version = "0.12.3"

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

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

[[deps.PositiveFactorizations]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "17275485f373e6673f7e7f97051f703ed5b15b20"
uuid = "85a6dd25-e78a-55b7-8502-1745935b8125"
version = "0.2.4"

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

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

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

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

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

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

[[deps.Ratios]]
deps = ["Requires"]
git-tree-sha1 = "dc84268fe0e3335a62e315a3a7cf2afa7178a734"
uuid = "c84ed2f1-dad5-54f0-aa8e-dbefe2724439"
version = "0.4.3"

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

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

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

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

[[deps.Roots]]
deps = ["CommonSolve", "Printf"]
git-tree-sha1 = "06ba8114bf7fc4fd1688e2e4d2259d2000535985"
uuid = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
version = "1.2.0"

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

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

[[deps.Setfield]]
deps = ["ConstructionBase", "Future", "MacroTools", "Requires"]
git-tree-sha1 = "77172cadd2fdfa0c84c87e3a01215a4ca7723310"
uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46"
version = "1.0.0"

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

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

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

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

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

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

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

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

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

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

[[deps.StructArrays]]
deps = ["Adapt", "DataAPI", "StaticArrays", "Tables"]
git-tree-sha1 = "9097e2914e179ab1d45330403fae880630acea0b"
uuid = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
version = "0.6.9"

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

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

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

[[deps.Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "OrderedCollections", "TableTraits", "Test"]
git-tree-sha1 = "5ce79ce186cc678bbb5c5681ca3379d1ddae11a1"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.7.0"

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

[[deps.TensorCore]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "1feb45f88d133a655e001435632f019a9a1bcdb6"
uuid = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50"
version = "0.1.1"

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

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

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

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

[[deps.UnicodePlots]]
deps = ["ColorTypes", "Contour", "Crayons", "Dates", "FileIO", "FreeTypeAbstraction", "LazyModules", "LinearAlgebra", "MarchingCubes", "NaNMath", "Printf", "SparseArrays", "StaticArrays", "StatsBase", "Unitful"]
git-tree-sha1 = "05fa79b807e11f1a123ab636abcbc6cf4b6c51c4"
uuid = "b8865327-cd53-5732-bb35-84acbb429228"
version = "2.12.3"

[[deps.Unitful]]
deps = ["ConstructionBase", "Dates", "LinearAlgebra", "Random"]
git-tree-sha1 = "b649200e887a487468b71821e2644382699f1b0f"
uuid = "1986cc42-f94f-5a68-af5c-568840ba703d"
version = "1.11.0"

[[deps.WoodburyMatrices]]
deps = ["LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "de67fa59e33ad156a590055375a30b23c40299d3"
uuid = "efce3f68-66dc-5838-9240-27a6d6f5f9b6"
version = "0.5.5"

[[deps.Yields]]
deps = ["BSplineKit", "ForwardDiff", "LinearAlgebra", "LsqFit", "Optim", "Roots", "UnicodePlots"]
git-tree-sha1 = "e82a40304025d049e51defe7e6f9c29c650f5603"
uuid = "d7e99b2f-e7f3-4d9e-9f01-2338fc023ad3"
version = "3.0.1"

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

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

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

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

# ╔═╡ Cell order:
# ╟─ba85bf56-5401-4430-a0e2-eef7e06347a5
# ╠═fbce1e27-45db-465d-83f0-a0a7935abf14
# ╠═2f700532-bc7b-4408-8268-f9f632a75451
# ╠═1f4e28e0-c2f0-4c06-a9f8-d0d71999114a
# ╠═001d049c-efd2-11ec-0083-cff70f84b072
# ╠═c7dbfd5b-a97e-4628-a2f8-8651f12c60f6
# ╠═a3a5b911-1097-454a-b0ec-463874be92b4
# ╟─50569a1e-5deb-48b6-be72-5be41dcb5b43
# ╠═280c7ef0-61d5-4d28-87b2-d26855c76b6b
# ╠═b8fe0af1-29b7-419a-ad3e-7b120eb5cbe1
# ╠═35bed858-6b76-4dbb-b49a-505f36b13434
# ╠═741777ff-2d40-4393-a0d4-669f3460cc64
# ╠═1f59e516-4dc1-450e-87a4-9ce6ac64029d
# ╠═dde19062-8a22-4caa-ae68-23895628dc3e
# ╠═18e32b95-2d13-4502-8500-48b0c380291c
# ╠═10c9e253-2913-4f85-9989-8b3fffba423e
# ╠═ac656984-3bdd-4a33-90be-45042e0080a5
# ╠═1a903be4-57fb-466c-96e3-f2ad36acc77c
# ╠═324be56b-512b-4664-9f5b-943b6d9aeade
# ╟─5c13430a-96c4-447f-89a7-1fe99078feee
# ╠═be7bbb76-fc0f-40b5-982b-7e46126f9c28
# ╟─b5013701-848a-4920-b4f5-360815a9f2e2
# ╠═1d59569f-c56b-45d0-ae57-ead57f68ed7b
# ╠═7dd4015f-6b58-4cf2-af3f-30bf3d43d468
# ╠═2e7a438e-435f-4dc1-86c2-947ea8ceb8d4
# ╠═9e824c6d-a303-427f-973c-bc1b9658709d
# ╠═973a5a58-9b42-4532-be80-21f18fb6791e
# ╟─98720096-e605-40e9-a32f-d1d5ade04be6
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Plane(p₀, v, w)

A plane coincident with point `p₀`, defined by non-parallel vectors `v` and `w`.
It can be called as p(s, t) with numeric parameters `s` and `t` to cast it at
`p₀ + s*v + t*w`.
"""
struct Plane{Dim,T} <: Primitive{Dim,T}
  pₒ::Point{Dim,T}
  v::Vec{Dim,T}
  w::Vec{Dim,T}
end

Plane(pₒ::Tuple, v::Tuple, w::Tuple) = Plane(Point(pₒ), Vec(v), Vec(w))

function (p::Plane)(s, t)
  p.pₒ + s*p.v + t*p.w
end

paramdim(::Type{<:Plane}) = 2


using ConservationLawsParticles
using RecursiveArrayTools, DifferentialEquations, Plots

# external velocities
Vr(t, x) = 2.
Vl(t, x) = -2.
# interactions
W_attr(t, r) = 5 * log(abs(r) + 1)
W_rep(t, r) = -2 * log(abs(r) + 1)
W(t, r) = 2 * (exp(abs(r)/4) + exp(-2abs(r)))
# mobilities
mobρ(ρ, σ) = max(2 - ρ - 0.5σ, 0)
mobσ(ρ, σ) = max(2 - σ - 0.5ρ, 0)

imodel = IntegratedModel(
    (Vr, Vl),
    ((W, W_rep),
     (W_rep, W)),
    (mobρ, mobσ))

n = 20
x0 = ArrayPartition(
    vcat(range(-2., -1.5, length=n), range(-1., -.5, length=n)),
    vcat(range(.5, 1.5, length=2n)))

tspan = (0., 1.2)

prob = ODEProblem(velocities!, x0, tspan, imodel)

sol = solve(prob, BS5(), reltol=1e-6, abstol=1e-6)

plot(legend=false)
plot!(sol, vars=1:2n, color=:blue)
plot!(sol, vars=2n+1:4n, color=:red)
plot!(title="2-species integrated scheme", xlabel="time", ylabel="position")

savefig("two-species.png")

n = 60
x0 = ArrayPartition(
    vcat(range(-2., -1.5, length=n), range(-1., -.5, length=n)),
    vcat(range(.5, 1.5, length=2n)))

prob = ODEProblem(velocities!, x0, tspan, imodel)

sol = solve(prob, BS5(), reltol=1e-7, abstol=1e-7)

anim = @animate for t in tspan[1]:1/48:tspan[2]
    p = plot(title="2-species integrated scheme", xlabel="position", ylabel="density",
        legend=false, xlims=(-4,4), ylims=(0,1))
    plot_density!(p, sol(t).x[1], color=:blue)
    plot_density!(p, sol(t).x[2], color=:red)
end

gif(anim, "two-species.gif")


module Point2DTools

using StaticArrays

export Point2D, convex, convexperm, centroid, bounds, scalefns

const Point2D = SVector{2}

Base.atan(p::Point2D) = atan(p.y, p.x)
Base.atan(p1::Point2D, p2::Point2D) = atan(p2-p1)

dtan(p::Point2D) = rad2deg(a)

"""
    convex(points)

Return the list of Point2D points which form the convex hull of the supplied points.
"""
convex(points) = points[convexperm(points)]

×(p1::Point2D, p2::Point2D) = p1.x * p2.y - p2.x * p1.y

direction(p1, p2, p3) = ×(p3-p1, p2-p1)

right(p1, p2, p3) = direction(p1, p2, p3) > 0
left(p1, p2, p3) = direction(p1, p2, p3) < 0
collinear(p1, p2, p3) = direction(p1, p2, p3) == 0 # is 0 too large ?
"""
    convexperm(points)

 Return the list of indexes of the Point2D points which form the convex hull of the supplied points.
"""
function convexperm(points)
    ips = sortperm(points)

    Lupper = [ips[1], ips[2]]

    for i in 3:length(ips)
        push!(Lupper, ips[i])
        while length(Lupper) > 2 && !right(points[Lupper[end-2:end]]...)
            Lupper[end-1] = Lupper[end]
            pop!(Lupper)
        end
    end
    Lower = [ips[end], ips[end-1]]
    for i in length(ips)-2:-1:1
        push!(Lower, ips[i])
        while length(Lower) > 2 && !right(points[Lower[end-2:end]]...)
            Lower[end-1] = Lower[end]
            pop!(Lower)
        end
    end
    vcat(Lupper, Lower[2:end])
end

"""
    centroid(points)
Calculate the unweighted centre of the closed polygon defined by `points`.
# Arguments
- `points` a Vector of points, assumed to include the "start" and "end" point twice, so points[1] is ignored
"""
centroid(points) = length(points) < 2 ? points[1] : sum(points[2:end]) / length(points[2:end])

"""
    bounds(points)
Return the [Minimum::Point2D, Maximum::Point2D] (so you can vcat with a list of lists)
# Arguments
- `points` Vector of points with .x and .y properties
# Examples
julia> b1 = bounds([Point2D(-10,0), Point2D(1,2), Point2D(2,1)])
2-element Vector{StaticArrays.SVector{2, Int64}}:
 [-10, 0]
 [2, 2]
julia> b2 = bounds([Point2D(0,0), Point2D(1,2), Point2D(2,11)])
2-element Vector{StaticArrays.SVector{2, Int64}}:
 [0, 0]
 [2, 11]
 julia> bounds(vcat(b1, b2))
 2-element Vector{StaticArrays.SVector{2, Int64}}:
  [-10, 0]
  [2, 11]
"""
function bounds(points)
    minx, miny, maxx, maxy = Inf, Inf, -Inf, -Inf
    for p in points
        if p[1] < minx
            minx = p[1]
        elseif p[1] > maxx
            maxx = p[1]
        end
        if p[2] < miny
            miny = p[2]
        elseif p.y > maxy
            maxy = p[2]
        end
    end
    [Point2D(minx, miny), Point2D(maxx, maxy)]
end

"""
    scalefns(points, width, height)
    scalefns(min_max, width, height)
return functions to apply to fx and fy of the points or bounds to scale to specified `width` and `height` - preserving aspect
# Arguments
- `points` Vector of points with .x and .y properties
"""
scalefns(points, width, height) = scalefns(bounds(points)..., width, height)
function scalefns(minxy, maxxy, width, height)
    scale = min(width, height) / min(maxxy.x - minxy.x, maxxy.y - minxy.y)
    x -> scale * (x - minxy.x), y -> scale * (y - minxy.y)
end

###
end


export ExpConstraint, conic_form!, vexity

### (Primal) exponential cone constraint ExpConstraint(x,y,z) => y exp(x/y) <= z & y>=0
struct ExpConstraint <: Constraint
    head::Symbol
    id_hash::UInt64
    children::Tuple{AbstractExpr, AbstractExpr, AbstractExpr} # (x, y, z)
    size::Tuple{Int, Int}
    dual::ValueOrNothing

    function ExpConstraint(x::AbstractExpr, y::AbstractExpr, z::AbstractExpr)
        @assert(x.size == y.size == z.size,
                "Exponential constraint requires x, y, and z to be of same size")
        # @assert(x.size == (1,1),
        #         "Exponential constraint requires x, y, and z to be scalar for now")
        sz = x.size
        id_hash = hash((x,y,z, :exp))
        return new(:exp, id_hash, (x, y, z), sz, nothing)
    end
end

ExpConstraint(x::AbstractExpr, y, z::AbstractExpr) = ExpConstraint(x, Constant(y), z)
ExpConstraint(x::AbstractExpr, y::AbstractExpr, z) = ExpConstraint(x, y, Constant(z))
ExpConstraint(x, y::AbstractExpr, z::AbstractExpr) = ExpConstraint(Constant(x), y, z)

function vexity(c::ExpConstraint)
    # TODO: check these...
    if vexity(c.x) == ConcaveVexity()
        error("Exponential constraint requires x to be convex")
    end
    if vexity(c.y)!=ConstVexity()
        error("Exponential constraint requires y to be constant")
    end
    if vexity(c.z) == ConvexVexity()
        error("Exponential constraint requires z to be concave")
    end
    return ConvexVexity()
end

function conic_form!(c::ExpConstraint, unique_conic_forms::UniqueConicForms=UniqueConicForms())
    if !has_conic_form(unique_conic_forms, c)
        conic_constrs = ConicConstr[]
        if c.size == (1, 1)
            objectives = Vector{ConicObj}(undef, 3)
            @inbounds for iobj = 1:3
                objectives[iobj] = conic_form!(c.children[iobj], unique_conic_forms)
            end
            push!(conic_constrs, ConicConstr(objectives, :ExpPrimal, [1, 1, 1]))
        else
            for i = 1:c.size[1]
                for j = 1:c.size[2]
                    objectives = Vector{ConicObj}(undef, 3)
                    @inbounds for iobj = 1:3
                        objectives[iobj] = conic_form!(c.children[iobj][i,j], unique_conic_forms)
                    end
                    push!(conic_constrs, ConicConstr(objectives, :ExpPrimal, [1, 1, 1]))
                end
            end
        end
        cache_conic_form!(unique_conic_forms, c, conic_constrs)
    end
    return get_conic_form(unique_conic_forms, c)
end


using PtFEM, Test

data = Dict(
  # Frame(nels, nn, ndim, nst, nip, finite_element(nod, nodof))
  :struc_el => Frame(4, 5, 3, 1, 1, Line(2, 2)),
  :properties => [5.0e5;],
  :x_coords => [0.0, 1.25, 3.5, 4.0, 2.0],
  :y_coords => [0.0, 3.0, 2.0, 1.0, 1.5],
  :z_coords => [0.0, 0.0, 0.0, 0.0, 3.0],
  :g_num => [1 2 3 4;5 5 5 5],
  :support => [
    (1, [0 0 0]),
    (2, [0 0 0]),
    (3, [0 0 0]),
    (4, [0 0 0])
    ],
  :loaded_nodes => [
    (5, [20.0 -20.0 30.0])],
  :fixed_freedoms => [
    (5, 2, -0.0005)],
  :penalty => 1e19
)

@time m, dis_df, fm_df = p42(data)

@test round.(m.displacements[:, 5], digits=7) == [0.0002569, -0.0005, 7.61e-5]


using Plots
pyplot()
using Random

include("hexbinplots_orig.jl")

function shapecoords(hx::HexBinPlots.HexHistogram)
    h,vh = HexBinPlots.make_shapes(hx)
    vmax = maximum(vh)
    xsh = Vector{Float64}()
    ysh = Vector{Float64}()
    for k in eachindex(h)
        append!(xsh,h[k].x)
        push!(xsh,h[k].x[1])
        push!(xsh,NaN)
        append!(ysh,h[k].y)
        push!(ysh,h[k].y[1])
        push!(ysh,NaN)
    end
    return (xsh, ysh, vh)
end

using Interpolations
function control_point(xi, xj, yi, yj, dist_from_mid)
    xmid = 0.5 * (xi+xj)
    ymid = 0.5 * (yi+yj)
    # get the angle of y relative to x
    theta = atan((yj-yi) / (xj-xi)) + 0.5pi
    # dist = sqrt((xj-xi)^2 + (yj-yi)^2)
    # dist_from_mid = curvature_scalar * 0.5dist
    # now we have polar coords, we can compute the position, adding to the midpoint
    (xmid + dist_from_mid * cos(theta),
     ymid + dist_from_mid * sin(theta))
end

function _add_edge_points!(xlines,ylines,xsi,ysi,xdi,ydi,curve_amount,interplen)
    xpt, ypt = control_point(xsi,xdi,ysi,ydi,curve_amount)
    xpts = [xsi, xpt, xdi]
    ypts = [ysi, ypt, ydi]
    t = range(0, stop=1, length=3)
    A = hcat(xpts, ypts)
    itp = scale(interpolate(A, BSpline(Cubic(Natural(OnGrid())))), t, 1:2)
    tfine = range(0, 1, length=interplen)
    for t in tfine
        push!(xlines, itp(t, 1))
    end
    for t in tfine
        push!(ylines, itp(t, 2))
    end
end
# Mutating version from HexBinPlots
function _xy2counts!(counts::Dict{(Tuple{Int, Int}), Int},
                     x::AbstractArray,y::AbstractArray,
                     xsize::Real,ysize::Real,x0,y0)
    @inbounds for i in eachindex(x)
        h = convert(HexagonOffsetOddR,
                    cube_round(x[i] - x0,y[i] - y0,xsize, ysize))
        idx = (h.q, h.r)
        counts[idx] = 1 + get(counts,idx,0)
    end
    counts
end

using Statistics
using Hexagons

# Generate counts for a graph edge plot histogram
# hhparam = (xsize,ysize,x0,y0) are the hex-histogrma params
function graphedgehist(src,dst,x,y,ptscale,hhparams,curvature = 0.05)
    xlines = similar(x,0)
    ylines = similar(y,0)
    sizehint!(xlines,max(length(src)+1000,10000000)) # we'll process them in batches this big
    sizehint!(ylines,max(length(src)+1000,10000000))
    # this is the counts of the edge vectors...
    edgecounts = Dict{(Tuple{Int, Int}), Int}()
    # now, run through each edges...
    for ei in eachindex(src)
        si = src[ei]
        di = dst[ei]
        xsi = x[si]
        ysi = y[si]
        xdi = x[di]
        ydi = y[di]
        curve_amount = curvature*sign((di-si))
        edgelen = sqrt((xdi-xsi)^2 + (ydi-ysi)^2)
        interplen = max(ceil(Int,2*edgelen/ptscale),2)
        if edgelen == 0
            continue
        end
        _add_edge_points!(xlines, ylines, xsi, ysi, xdi, ydi, curve_amount, interplen)
        if length(xlines) >= 9500000
            # compress all to edgecounts
            println("Compressing counts... $(ei) / $(length(src))")
            _xy2counts!(edgecounts, xlines, ylines, hhparams...)
            resize!(xlines,0)
            resize!(ylines,0)
        end
    end
    # handle final update
    return _xy2counts!(edgecounts, xlines, ylines, hhparams...)
end

function hexbingraphplot(src,dst,x,y;nbins=500)
    xbins = nbins
    ybins = nbins
    # compute params for hex histogram bins
    xmin,xmax = extrema(x)
    ymin,ymax = extrema(y)
    xspan, yspan = xmax - xmin, ymax - ymin
    xsize, ysize = xspan / xbins, yspan / ybins
    x0, y0 = xmin - xspan / 2,ymin - yspan / 2
    hhparams = (xsize, ysize, x0, y0) # save them as a tuple
    edgescale = 6
    hsize = min((xmax-xmin)/(xbins*edgescale),(ymax-ymin)/(ybins*edgescale))
    # need to compute hhparams
    # get coords for edges
    counts = graphedgehist(src,dst,x,y,hsize,hhparams)
    # do this for the nodes...
    #counts = HexBinPlots.xy2counts_(x,y,hhparams...)
    xh,yh,vh = HexBinPlots.counts2xy_(counts,hhparams...)
    hedges = HexBinPlots.HexHistogram(xh,yh,vh,xsize,ysize,false)
    return hedges
end

function generate_nice_plot(xy, src, dst, imagefilename;
                            labels=[],
                            mymarkeralpha=0.1,
                            mymarkersize=5,
                            invpermval=false,
                            colorscale=[],
                            skipedges=false)

    if invpermval
        x = xy[:,1]
        y = xy[:,2]
        x = invperm(sortperm(x))
        y = invperm(sortperm(y))
        xy = Float64.(hcat(x,y))
    end

    if !isempty(labels)
	labels = Int.(labels)
	ulabels = unique(labels)
	d = Dict(sort(ulabels) .=> 1:length(ulabels))
	color_palette = distinguishable_colors(length(ulabels))
        color_palette[end] = RGB(1,1,1);
	colorstouse = map(i->color_palette[d[labels[i]]],1:length(labels))
    end

    xd = xy[:,1]
    yd = xy[:,2]

    plot()
    if !skipedges
        hedges = hexbingraphplot(src, dst, xd, yd, nbins=900)
        xsh, ysh, axform = shapecoords(hedges)
        avals = axform/maximum(axform)
        avals = map(x -> x <= 1/4 ? sqrt(x) : ((sqrt(1/4)-1/4) + x), avals )
        avals = 1/1.25.*avals
        plot!(xsh, ysh,
              seriestype=:shape,
              alpha=avals,
              fillcolor=:darkblue,
              legend=false,
              linealpha=0,
              framestyle=:none,
              background=nothing)
    end

    #cscale = colorscale ./ maximum(colorscale)
    #rgbvals = get(ColorSchemes.valentine,cscale)
    #ids = findall(colorscale.!=0)
    ids = collect(length(xd):-1:1)
    scatter!(xd[ids], yd[ids],
             color=colorstouse[ids],
             alpha=mymarkeralpha,
             markersize=mymarkersize,
             legend=false,
             colorbar=false,
             markerstrokewidth=0,
             framestyle=:none,
             background=nothing)

    #plot!(dpi=300,size=(800,800))
    #plot!(dpi=200,size=(600,600))
    savefig(imagefilename)
end

function generate_nice_biplot(xy1, xy2, src, dst, imagefilename;
                              labels=[],
                              mymarkeralpha1=0.1, mymarkeralpha2=0.1,
                              mymarkersize1=5, mymarkersize2=5,
                              incledges=true,
                              inclsecond=true)

    xy = [xy1; xy2]
    xd = xy[:,1]
    yd = xy[:,2]

    plot()
    if incledges
        hedges = hexbingraphplot(src, dst, xd, yd, nbins=900)
        xsh, ysh, axform = shapecoords(hedges)
        avals = axform/maximum(axform)
        avals = map(x -> x <= 1/4 ? sqrt(x) : ((sqrt(1/4)-1/4) + x), avals )
        avals = 1/1.25.*avals
        plot!(xsh, ysh,
              seriestype=:shape,
              alpha=avals,
              fillcolor=:darkblue,
              legend=false,
              linealpha=0,
              framestyle=:none,
              background=nothing)
    end

    if inclsecond
        scatter!(xy2[:, 1], xy2[:, 2],
                 color=RGB(1,1,1),
                 alpha=mymarkeralpha2,
                 markersize=mymarkersize2,
                 legend=false,
                 colorbar=false,
                 framestyle=:none,
                 background=nothing)

    end

    function get_colors()
        if isempty(labels)
            return :red
        end
	labels = Int.(labels)
	ulabels = unique(labels)
	d = Dict(sort(ulabels) .=> 1:length(ulabels))
        # skip first in palette
	color_palette = distinguishable_colors(length(ulabels) + 1)
	return map(i->color_palette[d[labels[i]] + 1], 1:length(labels))
    end
    colorstouse = get_colors()
    scatter!(xy1[:, 1], xy1[:, 2],
             color=colorstouse,
             alpha=mymarkeralpha1,
             markersize=mymarkersize1,
             legend=false,
             colorbar=false,
             framestyle=:none,
             background=nothing)

    plot!(dpi=300,size=(800,800))
    #plot!(dpi=200,size=(400,400))
    savefig(imagefilename)
end


# Image resizing utilities

export anti_aliasing_filter

function anti_aliasing_filter(n::Integer; fact::Integer=1, flat::Bool=false, coeff::Real=1, T::DataType=Float32)
    kmax = T(pi)
    k = kmax*coord_norm(n)
    filt = zeros(T, n)
    if ~flat
        i1 = abs.(k) .<= kmax/2^fact-kmax/2^(fact+1)
        i2 = (abs.(k) .> kmax/2^fact-kmax/2^(fact+1)) .&& (abs.(k) .<= kmax/2^fact+kmax/2^(fact+1))
        filt[i2] .= ((cos.(range(-T(pi), T(pi); length=length(findall(i2)))).+1)/2).^T(coeff)
    else
        i1 = abs.(k) .<= kmax/2^fact
    end
    filt[i1] .= 1
    return filt
end

function anti_aliasing_filter(K::KSpaceCartesianSampling{T}; fact::Integer=1, flat::Bool=false, coeff::Real=1) where {T<:Real}
    nt, nk = size(K)
    perm = (K.phase_encoding..., K.readout)
    filt = reshape(permutedims(anti_aliasing_filter(K.X.n; fact=fact, flat=flat, coeff=coeff), perm), nt, nk)[K.subsampling,:]
    i_pe, i_r = downscale_phase_encode_index(K; fact=fact, readout=true)
    return filt[i_pe, i_r]
end

function anti_aliasing_filter(n::NTuple{3,Integer}; fact::Integer=1, flat::Bool=false, coeff::Real=1, T::DataType=Float32)
    filt_x = reshape(anti_aliasing_filter(n[1]; fact=fact, flat=flat, coeff=coeff, T=T), :, 1, 1)
    filt_y = reshape(anti_aliasing_filter(n[2]; fact=fact, flat=flat, coeff=coeff, T=T), 1, :, 1)
    filt_z = reshape(anti_aliasing_filter(n[3]; fact=fact, flat=flat, coeff=coeff, T=T), 1, 1, :)
    return filt_x.*filt_y.*filt_z
end

anti_aliasing_filter(X::RegularCartesianSpatialSampling{T}; fact::Integer=1, flat::Bool=false, coeff::Real=1) where {T<:Real} = anti_aliasing_filter(X.n; fact=fact, flat=flat, coeff=coeff, T=T)

function downscale(u::AbstractArray{CT,3}, X::RegularCartesianSpatialSampling{T}; fact::Integer=1, flat::Bool=false, coeff::Real=1) where {T<:Real,CT<:RealOrComplex{T}}
    (fact == 0) && (return u)
    x, y, z = coord(X)
    itp = interpolate((vec(x), vec(y), vec(z)), u, Gridded(Linear()))
    u_ = ifft(ifftshift(anti_aliasing_filter(X; fact=fact, flat=flat, coeff=coeff)).*fft(u))
    X_h = downscale(X; fact=fact)
    u_h = similar(u_, size(X_h))
    x_h, y_h, z_h = coord(X_h)
    @inbounds for i = 1:X_h.n[1], j = 1:X_h.n[2], k = 1:X_h.n[3]
        u_h[i, j, k] = itp(x_h[i], y_h[j], z_h[k])
    end
    return u_h
end

function upscale(u::AbstractArray{CT,3}, X::RegularCartesianSpatialSampling{T}; fact::Integer=1) where {T<:Real,CT<:RealOrComplex{T}}
    (fact == 0) && (return u)
    x, y, z = coord(X)
    itp = extrapolate(interpolate((vec(x), vec(y), vec(z)), u, Gridded(Linear())), T(0))
    X_h = upscale(X; fact=fact)
    u_h = similar(u, size(X_h))
    x_h, y_h, z_h = coord(X_h)
    @inbounds for i = 1:X_h.n[1], j = 1:X_h.n[2], k = 1:X_h.n[3]
        u_h[i, j, k] = itp(x_h[i], y_h[j], z_h[k])
    end
    return u_h
end


function digital_root(a::Integer)::Int
    while a >= 10
        a = sum(digits(a))
    end
    a
end

function digital_root(a::Vector{<:Integer})::Int
    a = cumprod(big.(a))
    digital_root(sum(a))
end


# Input data
data = map(x -> parse(Int64, x), split(readline("Day07.input"), ","));

# Part 1
mindistance =
    (F, data) ->
        reduce((acc, v) -> min(acc, v), map(x -> F(x, data), range(extrema(data)...)))

ans = mindistance((x, data) -> sum(map(v -> abs(v - x), data)), data);
println("Part 1: $ans")

# Part 2
ans = mindistance(
    (x, data) -> sum(map(v -> (n = abs(v - x); (n * (n + 1)) >> 1), data)),
    data,
);
println("Part 2: $ans")


function swsetupmatrix(s1, s2)
    setupmatrix = zeros(Int, length(s1)+1, length(s2)+1)
    return setupmatrix
end

function swscorematrix(seq1, seq2; match=1, mismatch=-1, gap=-1)
    scoremat = swsetupmatrix(seq1, seq2)
    for i in 2:size(scoremat, 1) # iterate through row indices
        for j in 2:size(scoremat, 2) # iterate through column indices
            above = nwscore(seq1[i-1], nothing; match, mismatch, gap) + scoremat[i-1,j]
            left = nwscore(nothing, seq2[j-1]; match, mismatch, gap) + scoremat[i,j-1]
            diagonal = nwscore(seq1[i-1], seq2[j-1]; match, mismatch, gap) + scoremat[i-1,j-1]
            scoremat[i,j] = max(above, left, diagonal)
            # or `scoremat[i,j] = max(above, left, diagonal, 0)`

            if scoremat[i,j] < 0
                scoremat[i,j] = 0
            end
        end
    end
    return scoremat
end


using Test
using DECAES
using Logging

function compare_epg(
        work₁::DECAES.AbstractEPGWorkspace{T,ETL},
        work₂::DECAES.AbstractEPGWorkspace{T,ETL},
    ) where {T,ETL}
    α, TE, T2, T1, β = T(163.0), T(11e-3), T(39e-3), T(1.1), T(151.0)
    θ   = DECAES.EPGOptions((; α, TE, T2, T1, β), Val(ETL), T)
    dc₁ = DECAES.EPGdecaycurve!(work₁, θ)
    dc₂ = DECAES.EPGdecaycurve!(work₂, θ)
    if !(dc₁ ≈ dc₂)
        @info "Comparing: $((nameof(typeof(work₁)), nameof(typeof(work₂))))"
        @info "    max error:   $(maximum(abs, dc₁ .- dc₂))"
        @info "    diff vector: $(round.(abs.(dc₁ .- dc₂)'; sigdigits = 4))"
    end
    @test dc₁ ≈ dc₂
end

@testset "EPG algorithms" begin
    # NOTE: Generated function approach is extremely slow to compile for large ETL (around 16)
    epg_algs = DECAES.EPG_Algorithms
    for i in 1:length(epg_algs), j in 1:i-1, T in [Float32, Float64], ETL in [4,5,6,7] # test four ETL >= 4 which are unique mod 4
        algᵢ, algⱼ = epg_algs[i], epg_algs[j]
        compare_epg(algᵢ(T, ETL), algⱼ(T, ETL))
    end
end

nothing


"""Sampler object"""
struct GeneralizedInverseGaussian{T<:Real} <: Distributions.ContinuousUnivariateDistribution
    a::T
    b::T
    p::T
    function GeneralizedInverseGaussian{T}(a::T, b::T, p::T) where {T}
        Distributions.@check_args(GeneralizedInverseGaussian, a > zero(a) && b > zero(b))
        return new{T}(a, b, p)
    end
end

function GeneralizedInverseGaussian(a::T, b::T, p::T) where {T}
    return GeneralizedInverseGaussian{T}(a::T, b::T, p::T)
end

Distributions.params(d::GeneralizedInverseGaussian) = (d.a, d.b, d.p)
@inline Distributions.partype(::GeneralizedInverseGaussian{T}) where {T<:Real} = T

function Distributions.mean(d::GeneralizedInverseGaussian)
    a, b, p = params(d)
    q = sqrt(a * b)
    return (sqrt(b) * besselk(p + 1, q)) / (sqrt(a) * besselk(p, q))
end

function Distributions.var(d::GeneralizedInverseGaussian)
    a, b, p = params(d)
    q = sqrt(a * b)
    r = besselk(p, q)
    return (b / a) * ((besselk(p + 2, q) / r) - (besselk(p + 1, q) / r)^2)
end

function Distributions.mode(d::GeneralizedInverseGaussian)
    return ((d.p - 1) + sqrt((d.p - 1)^2 + d.a * d.b)) / d.a
end

function Distributions.pdf(d::GeneralizedInverseGaussian{T}, x::Real) where {T<:Real}
    if x > 0
        a, b, p = params(d)
        (((a / b)^(p / 2)) / (2 * besselk(p, sqrt(a * b)))) *
        (x^(p - 1)) *
        exp(-(a * x + b / x) / 2)
    else
        zero(T)
    end
end

function Distributions.logpdf(d::GeneralizedInverseGaussian{T}, x::Real) where {T<:Real}
    if x > 0
        a, b, p = params(d)
        (p / 2) * (log(a) - log(b)) - log(2 * besselk(p, sqrt(a * b))) + (p - 1) * log(x) -
        (a * x + b / x) / 2
    else
        -T(Inf)
    end
end

function Distributions.rand(d::GeneralizedInverseGaussian)
    a, b, p = params(d)
    α = sqrt(a / b)
    β = sqrt(a * b)
    λ = abs(p)
    if (λ > 1) || (β > 1)
        x = _rou_shift(λ, β)
    elseif β >= min(0.5, (2 / 3) * sqrt(1 - λ))
        x = _rou(λ, β)
    else
        x = _hormann(λ, β)
    end
    return p >= 0 ? x / α : 1 / (α * x)
end
function _gigqdf(x::Real, λ::Real, β::Real)
    return (x^(λ - 1)) * exp(-β * (x + 1 / x) / 2)
end

function _hormann(λ::Real, β::Real)
    # compute bounding rectangles
    m = β / (1 - λ + sqrt((1 - λ)^2 + β^2))  # mode
    x0 = β / (1 - λ)
    xstar = max(x0, 2 / β)
    # in subdomain (0, x0)
    k1 = _gigqdf(m, λ, β)
    a1 = k1 * x0
    # in subdomain (x0, 2 / β), may be empty
    if x0 < 2 / β
        k2 = exp(-β)
        a2 = λ == 0 ? k2 * log(2 / (β^2)) : k2 * ((2 / β)^λ - x0^λ) / λ
    else
        k2 = 0
        a2 = 0
    end
    # in subdomain (xstar, Inf)
    k3 = xstar^(λ - 1)
    a3 = 2k3 * exp(-xstar * β / 2) / β
    a = a1 + a2 + a3

    # perform rejection sampling
    while true
        u = rand()
        v = a * rand()
        if v <= a1
            # in subdomain (0, x0)
            x = x0 * v / a1
            h = k1
        elseif v <= a1 + a2
            # in subdomain (x0, 2 / β)
            v -= a1
            x = λ == 0 ? β * exp(v * exp(β)) : (x0^λ + v * λ / k2)^(1 / λ)
            h = k2 * x^(λ - 1)
        else
            # in subdomain (xstar, Inf)
            v -= a1 + a2
            x = -2log(exp(-xstar * β / 2) - v * β / (2k3)) / β
            h = k3 * exp(-x * β / 2)
        end
        if u * h <= _gigqdf(x, λ, β)
            return x
        end
    end
end

function _rou(λ::Real, β::Real)
    # compute bounding rectangle
    m = β / (1 - λ + sqrt((1 - λ)^2 + β^2))  # mode
    xpos = (1 + λ + sqrt((1 + λ)^2 + β^2)) / β
    vpos = sqrt(_gigqdf(m, λ, β))
    upos = xpos * sqrt(_gigqdf(xpos, λ, β))

    # perform rejection sampling
    while true
        u = upos * rand()
        v = vpos * rand()
        x = u / v
        if v^2 <= _gigqdf(x, λ, β)
            return x
        end
    end
end

function _rou_shift(λ::Real, β::Real)
    # compute bounding rectangle
    m = (λ - 1 + sqrt((λ - 1)^2 + β^2)) / β  # mode
    a = -2(λ + 1) / β - m
    b = 2(λ - 1) * m / β - 1
    p = b - (a^2) / 3
    q = 2(a^3) / 27 - (a * b) / 3 + m
    ϕ = acos(-(q / 2) * sqrt(-27 / (p^3)))  # Cardano's formula
    r = sqrt(-4p / 3)
    xneg = r * cos(ϕ / 3 + 4π / 3) - a / 3
    xpos = r * cos(ϕ / 3) - a / 3
    vpos = sqrt(_gigqdf(m, λ, β))
    uneg = (xneg - m) * sqrt(_gigqdf(xneg, λ, β))
    upos = (xpos - m) * sqrt(_gigqdf(xpos, λ, β))

    # perform rejection sampling
    while true
        u = (upos - uneg) * rand() + uneg
        v = vpos * rand()
        x = max(u / v + m, 0)
        if v^2 <= _gigqdf(x, λ, β)
            return x
        end
    end
end


using eFEMpart
include("AdvDiff_example_param.jl")

function ad_example(p::AD_Para)
  square = [p.xmin,p.xmax,p.ymin,p.ymax]
  mesh = squareMesh(square,p.Nintervals,p.basisorder)

  varname      = ""
  OperatorType = :AdvDiff2D

  dNodes = Dirichlet(:left,:top)
  nNodes = Neumann(:right,:bottom)

  dBCf   = Dirichlet((x,y) -> (x==p.xmin ? p.dirVal : 0.0))
  ff     = Forcing((x,y) -> 0.0)

  # parameter
  κ = 1/p.Pe
  windX(x,y) =  -y;  windY(x,y) = x
  wx = [windX.(mesh.xy[i].x,mesh.xy[i].y) for i=1:length(mesh.xy)]
  wy = [windY.(mesh.xy[i].x,mesh.xy[i].y) for i=1:length(mesh.xy)]
  param  = AdvDiffParam(wx,wy,κ)

  Nodes = [dNodes]
  bcfun = [dBCf,ff]

  prob = Problem(mesh,Nodes,bcfun,varname,OperatorType)
  sol = solve(prob,mesh,param)

  # save solution
  vtkname = Path(p.solFolder,p.solFile)
  sd = ScalarData(sol.u)
  sn = ScalarNames(p.solName)
  vd = VectorData([wx,wy])
  vn = VectorNames("velocity")
  vtksave(mesh,sd,sn,vd,vn,vtkname)

  return nothing
end

function main()
  p = AD_Para()
  ad_example(p)
end

main()


using RecyclePermanent
using RNGPool
using JLD2
using StaticArrays
using LinearAlgebra
using Random
import MonteCarloMarkovKernels: simulateChainProgress, makeAMKernel

include("../models/bpModel.jl")

@load "data/bpData.jld" ys θ
θ0 = θ

d = 3

@inline function toBPθ(v::SVector{d, Float64})
  return BPθ(v[1], v[2], v[3])
end

t0 = SVector{d, Float64}(θ0.λ, θ0.a, θ0.b)

sigmaProp = 1.0*SMatrix{d, d, Float64}(Matrix{Float64}(I, d, d))

@inline function lglogprior(θ::BPθ)
  if θ.λ < 0 return -Inf end
  if θ.a < 0 return -Inf end
  if θ.b < 0 return -Inf end
  v::Float64 = 0.0
  v -= 0.001*θ.λ
  v -= 0.1*θ.a
  v -= 0.1*θ.b
  return v
end

N = 40 * length(ys)
chainLength = 1000000


@symbol_func function G_LO(cur_reactor::AbstractReactor)
  cur_inner_radius = R_CS(cur_reactor)

  cur_outer_radius = cur_inner_radius + d(cur_reactor)

  cur_G = cur_inner_radius * cur_outer_radius

  cur_G += cur_inner_radius ^ 2

  cur_G += cur_outer_radius ^ 2

  cur_G /= 3

  cur_G
end


"""
 Centermanifold computation and simplest normal form of Hopf bifurcation.
 This code is available to compute centermanifold and normal form of only semisimple Jacobian case (i.e. diagonalizable Jacobian).
 For nonsemisimple case, algorithm should be modified.

 SymEngine package is used.

 For theoretical background of Centermanifold computation see:
    http://www.orcca.on.ca/TechReports/TechReports/2000/TR-00-15.pdf
 for simplest normal form of Hopf bifurcation see:
    https://iopscience.iop.org/article/10.1088/0951-7715/16/1/317
"""
module HopfNormalForm

using SymEngine, OffsetArrays, StaticArrays

export calculate_normal_form, verify, Systems, Flutter_opt

################### Basic functions for symbolic computation in Julia ###################
function coeff2(f::Basic, e::Array{Basic,1}, N)::Basic #coefficient of multivariate Polynomial for univariate case just use diff
    # f is function of plinomials of variable e and N is index vector of power series. for example,
    # if f=e(1)^j*e(2)^k -> I=[j,k]
    v = f
    for i in eachindex(N)
        if N[i] == 0
            v *= 1
        else
            v = diff(v, e[i], N[i]) / factorial(N[i])
        end
    end
    for j in eachindex(N)
        v = subs(v, e[j] => 0)
    end
    return v
end

function recursive_apply(f, x::Basic)::Basic
    cls = SymEngine.get_symengine_class(x)
    args = [recursive_apply(f, arg) for arg in get_args(x)]
    if cls == :Add
        return sum(args)
    elseif cls == :Mul
        return prod(args)
    elseif cls == :Pow
        return args[1]^args[2]
    elseif cls in [:RealDouble, :RealMPFR, :ComplexDouble, :ComplexMPFR]
        return f(x)
    else
        return x
    end
end

function chop_expr(x::Basic, zero_tol::Real)::Basic #chopping very small coefficients to simplify the equaion expression
    return recursive_apply(x -> (abs(N(x)) <= zero_tol) ? zero(typeof(x)) : x, x)
end

function recursive_real(x::Basic)::Basic # Real part of complex symbols. All symbol variables are assumed to be real.
    # To define a complex symbol variable c, express it as c=a+b*I.
    return recursive_apply(real, x)
end

function recursive_imag(x::Basic)::Basic # Imaginary part of complex symbols. All symbol variables are assumed to be real.
    # To define a complex symbol variable c, express it as c=a+b*I.
    return recursive_apply(imag, x)
end

function verify(
    f::Array{Basic},
    T::Array{Basic},
    v::Array{Basic},
    norder::Integer,
    cm::Integer,
    m3::Integer,
    n::Integer,
)::Array{Basic}
    # verify the result of Centermanifold computation upto norder
    # f: function input, T: computed centermanifold, v: state-space variables, norder: oreder of computation,
    # cm: dimension of centermanifold, m3: dimension of stable-unstable manifold
    G2 = Array{Basic}(undef, m3)
    dhf = Array{Basic}(undef, m3)
    dh = Array{Basic}(undef, m3)
    F = Array{Basic}(undef, cm)
    veri = Array{Basic}(undef, m3)
    veri_o = Array{Basic}(undef, m3, norder - 1)
    @vars ϵ
    # Compute F
    for i = 1:cm
        F[i] = f[i] - Eign[i] * v[i]
        for m = 1:n
            if m > cm
                F[i] = subs(F[i], v[m] => T[m-cm])
                F[i] = expand(F[i])
                F[i] = chop_expr(F[i], 1e-10)
            end
        end
    end
    # Compute dh*f
    for i = 1:m3
        dhf[i] = 0
        for l = 1:cm
            dhf[i] += diff(T[i], v[l]) * F[l]
        end
    end
    # Compute g-dh*f
    for k = 1:m3
        G[k] = f[cm+k] - Eign[cm+k] * v[cm+k]
        for m = 1:n
            if m > cm
                G[k] = subs(G[k], v[m] => T[m-cm])
            end
        end
        G2[k] = G[k] - dhf[k]
    end
    # compute dh(x)J_cx-J_s*h(x)
    for i = 1:m3
        dh[i] = 0
        for l = 1:cm
            dh[i] += diff(T[i], v[l]) * Eign[l] * v[l]
        end
        dh[i] = dh[i] - Eign[i+cm] * T[i]
    end
    # veri[i]= g-dh*f - {dh(x)J_cx-J_s*h(x)}
    # We have to check order j=2,3,…,norder of veri[i] is zero
    for i = 1:m3
        veri[i] = dh[i] - G2[i]
        veri[i] = expand(veri[i])
        veri[i] = chop_expr(veri[i], 1e-10)
    end
    # returning value veri_o[i,j] is j-th order of veri[i]
    for j = 2:norder
        for i = 1:m3
            for m = 1:cm
                veri[i] = subs(veri[i], v[m] => ϵ * v[m])
            end
            veri_o[i, j-1] = subs(diff(veri[i], ϵ, j), ϵ => 0) / factorial(j)
        end
    end
    return veri_o
end

#==
Note that complex conjugate pair of eigenvalue corresponding to center subspace is transfromed to +-i.
This is equivalent to rescailing time.

################### Define the index for recursive computation ###################
To compute centermanifold of higher dimension n>3, you should modify this part.
index of coefficient arrays are [k,l,m,n] for centermanifold of dimension 3
k is for k-th part of function g=[g_1,g_2,…,g_k,…,g_m3], l,m,n is for coefficient of polynomial of v_1^l* v_2^m* v_3^n
of power series of g_k
==#

function calculate_normal_form(
    f::Vector{Basic},
    v::Vector{Basic},
    norder::Integer,
    cm::Integer,
    m1::Integer,
    m2::Integer,
    m3::Integer,
    n::Integer,
    ω_h::AbstractFloat,
    zero_tol::AbstractFloat
)
    G = Array{Basic}(undef, m3);
    T = Array{Basic}(undef, m3);
    d = Array{Basic}(undef, m3);
    F = Array{Basic}(undef, cm);

    a = Array{Basic}(undef, cm)
    RD = Array{Basic}(undef, cm, norder); # Reduced dynamics on the center manifold by order
    RD_sum = Array{Basic}(undef, cm); # Reduced dynamics on the center manifold summation of all orders

    @vars ϵ, ν
    # ϵ is symbolic variable which is used in the algorithm which does not have any physical meaning

    if cm == 2 # set multidimensional array for saving coefficints of the polynomial
    # for cm>=3 you should modify this part
        B = OffsetArray{Basic}(undef, 1:m3, 0:norder, 0:norder)
        H = OffsetArray{Basic}(undef, 1:m3, 0:norder, 0:norder)
    elseif cm == 3
        B = OffsetArray{Basic}(undef, 1:m3, 0:norder, 0:norder, 0:norder)
        H = OffsetArray{Basic}(undef, 1:m3, 0:norder, 0:norder, 0:norder)
        A = OffsetArray{Basic}(undef, 1:2, 0:norder, 0:norder, 0:norder)
    end

    Eign = Array{Basic}(undef, n); # Array to save eigenvalues

    ########## Collect Eigenvalues ############
    # Eigenvalues are collected from the input function
    for i = 1:n
        Eign[i] = diff(f[i], v[i])
        for k = 1:n
            Eign[i] = subs(Eign[i], v[k] => 0)
        end
    end

    ########## Initializing the centermanifold and function d (i.e. Dh(x)f(x,h(x))) ######
    for k = 1:m3
        T[k] = 0
        d[k] = 0
    end

    ############ Compute g(x,h(x)) ##############
    for j = 2:norder
        for k = 1:m3
            G[k] = f[cm+k] - Eign[cm+k] * v[cm+k]
            for m = 1:n
                if m > cm
                    G[k] = subs(G[k], v[m] => T[m-cm])
                end
            end
            G[k] = G[k] - d[k]
            for m = 1:cm
                G[k] = subs(G[k], v[m] => ϵ * v[m])
            end
            G[k] = subs(diff(G[k], ϵ, j), ϵ => 0) / factorial(j) # Collect the j-th order of the G[k]
            G[k] = expand(G[k])
            G[k] = chop_expr(G[k], zero_tol)

            if cm == 2
                ind_l = ((k, l, j - l) for l = 0:j)
            elseif cm == 3
                ind_l = ((k, l, m, j - l - m) for l = 0:j for m = 0:j-l)
            end
            for arg in ind_l
                xs = one(Basic)
                δ = zero(Basic)
                B[CartesianIndex(arg)] = coeff2(G[arg[1]], v, arg[2:end])
                for i = 1:cm
                    B[CartesianIndex(arg)] =
                        subs(B[CartesianIndex(arg)], v[i] => 0)
                end
                for i = 1:cm
                    δ += Eign[i] * arg[1+i]
                    xs *= v[i]^arg[1+i]
                end
                δ = δ - Eign[cm+arg[1]]
                H[CartesianIndex(arg)] = B[CartesianIndex(arg)] / δ
                H[CartesianIndex(arg)] = expand(H[CartesianIndex(arg)])
                H[CartesianIndex(arg)] =
                    chop_expr(H[CartesianIndex(arg)], zero_tol)
                T[arg[1]] += H[CartesianIndex(arg)] * xs
                T[arg[1]] = expand(T[arg[1]])
            end
        end
    ############ Compute Dh(x)f(x,h(x)) ##############
        for i = 1:cm
            F[i] = f[i] - Eign[i] * v[i]
            for m = 1:n
                if m > cm
                    F[i] = subs(F[i], v[m] => T[m-cm])
                    F[i] = expand(F[i])
                    F[i] = chop_expr(F[i], zero_tol)
                end
            end
            d[i] = 0
        end
        for i = 1:m3, l = 1:cm
            d[i] += diff(T[i], v[l]) * F[l]
            d[i] = expand(d[i])
            d[i] = chop_expr(d[i], zero_tol)
        end
    end
    ########## Computing reduced dynamics on the centermanifold ##########
    for k = 1:cm, j = 1:norder
        if j == 1
            RD[k, j] = Eign[k] * v[k]
        else
            RD[k, j] = f[k] - Eign[k] * v[k]
            for m = 1:n
                if m > cm
                    RD[k, j] = subs(RD[k, j], v[m] => T[m-cm])
                end
            end
            for m = 1:cm
                RD[k, j] = subs(RD[k, j], v[m] => ϵ * v[m])
            end
            RD[k, j] = subs(diff(RD[k, j], ϵ, j), ϵ => 0) / factorial(j) # Collect the j-th order of Reduced dynmics
            RD[k, j] = expand(RD[k, j])
            RD[k, j] = chop_expr(RD[k, j], zero_tol)
        end
    end
    for n = 1:cm
        RD_sum[n] = sum(RD[n, :])
    end
    ########## Simplest Normal form computation of Hopf bifurcation ##########

    A[1, 1, 0, 1] = recursive_real(coeff2(RD_sum[2], v, [1, 1, 0]));
    A[1, 2, 1, 0] = recursive_real(coeff2(RD_sum[2], v, [0, 2, 1]));
    A[1, 2, 0, 0] = recursive_real(coeff2(RD_sum[2], v, [0, 2, 0]));
    A[2, 1, 1, 0] = recursive_imag(coeff2(RD_sum[2], v, [0, 1, 1]));
    A[2, 2, 0, 0] = recursive_imag(coeff2(RD_sum[2], v, [0, 2, 0]));
    A[1, 1, 1, 0] = recursive_real(coeff2(RD_sum[2], v, [0, 1, 1]));

    S = A[1, 2, 1, 0] - A[1, 2, 0, 0] * A[2, 1, 1, 0] - A[2, 2, 0, 0] * A[1, 1, 1, 0];
    S = expand(S);

    return (CM=T,RD=RD_sum,SLC=S,LCO_R2=-A[1, 1, 0, 1] / S* ν)
end

# Include example systems
include("Systems.jl")
include("Flutter_opt.jl")

end # module


using JFVM
import PyPlot, JSON

include("../functions/rel_perms_real.jl")

function polymer_viscosity(cs, cp)
    a1 = 4.0
    a2 = 0.0
    a3 = 6.0
    sp = -0.24
    μ_w = 0.001 # Pa.s
    return μ_w*(1.0+cp.*cs.^sp.*(a1+a2*cp+a3*cp.*cp))
end

function dpolymer_viscosity_dcs(cs, cp)
    a1 = 4.0
    a2 = 0.0
    a3 = 6.0
    sp = -0.24
    μ_w = 0.001 # Pa.s
    return sp*μ_w*cp.*cs.^(sp-1).*(a1+a2*cp+a3*cp.*cp)
end

function dpolymer_viscosity_dcp(cs, cp)
    a1 = 4.0
    a2 = 0.0
    a3 = 6.0
    sp = -0.24
    μ_w = 0.001 # Pa.s
    return μ_w*cs.^sp.*(a1+2*a2*cp+3*a3*cp.*cp)
end

# domain
Lx = 1.0
Ly = 1.0
Nx = 300
Ny = 10
m = createMesh1D(Nx, Lx)

# (petro)physical properties
mu_oil = 10e-3 # Pa.s
perm = 0.1e-12
poros = 0.3

# rel-perms
krw0 = 0.2
kro0 = 0.8
n_o  = 2.0
n_w  = 2.0
swc  = 0.08
sor = 0.2

KRW  = sw -> krw.(sw, krw0, sor, swc, n_w)
dKRW = sw -> dkrwdsw.(sw, krw0, sor, swc, n_w)
KRO  = sw -> kro.(sw, kro0, sor, swc, n_o)
dKRO = sw -> dkrodsw.(sw, kro0, sor, swc, n_o)

# initial and boundary conditions
p_res = 100e5 # Pa reservoir (and right boundary pressure)
sw_0 = swc # inital water saturation
sw_inj = 1.0 # injection water saturation
cp_0 = 0.0 # initial polymer concentration kg/m^3
cp_inj = 1.0 # injected polymer concentration
cs_0 = 10.0  # initial salt concentration kg/m^3
cs_inj = 5.0 # injected salt concentration

φ = createCellVariable(m, poros) # porosity
k = createCellVariable(m, perm) # m^2 permeability
k_face   = harmonicMean(k) # harmonic averaging for perm
u_inj = 1e-5 # m/s injection Darcy velocity
pv_inj = 0.3 # injected pore volumes
t_final = pv_inj*Lx/(u_inj/poros) # s final time
dt0 = t_final/Nx # s time step
dt  = t_final/Nx # variable time step

BCp = createBC(m)   # pressure boundary
BCp.left.a[:] = perm/polymer_viscosity(cs_inj, cp_inj)
BCp.left.b[:] = 0.0
BCp.left.c[:] = -u_inj
BCp.right.a[:] = 0.0
BCp.right.b[:] = 1.0
BCp.right.c[:] = p_res
M_bc_p, RHS_bc_p = boundaryConditionTerm(BCp)

# saturation boundary
BCs = createBC(m)
BCs.left.a[:] = 0.0
BCs.left.b[:] = 1.0
BCs.left.c[:] = sw_inj # injected water saturation
M_bc_s, RHS_bc_s = boundaryConditionTerm(BCs)

# salt concentration boundary
BCcs = createBC(m)
BCcs.left.a[:] = 0.0
BCcs.left.b[:] = 1.0
BCcs.left.c[:] = cs_inj # injected salt concentration
M_bc_cs, RHS_bc_cs = boundaryConditionTerm(BCcs)

# polymer concentration boundary
BCcp = createBC(m)
BCcp.left.a[:] = 0.0
BCcp.left.b[:] = 1.0
BCcp.left.c[:] = cp_inj # injected polymer concentration
M_bc_cp, RHS_bc_cp = boundaryConditionTerm(BCcp)

# initial conditions
p_init  = createCellVariable(m, p_res)
sw_init = createCellVariable(m, sw_0)
cs_init = createCellVariable(m, cs_0)
cp_init = createCellVariable(m, cp_0)
p_val  = createCellVariable(m, p_res)
sw_val = createCellVariable(m, sw_0)
cs_val = createCellVariable(m, cs_0)
cp_val = createCellVariable(m, cp_0)

uw       = -gradientTerm(p_val)
# M_bc_p, RHS_bc_p = boundaryConditionTerm(BCp)
# M_bc_s, RHS_bc_s = boundaryConditionTerm(BCs)
# M_bc_c, RHS_bc_c = boundaryConditionTerm(BCcs)
# M_bc_t, RHS_bc_t = boundaryConditionTerm(BCt)
tol_s = 1e-7
tol_c = 1e-7
max_int_loop = 40
t = 0.0
FL = fluxLimiter("SUPERBEE") # flux limiter

while t<t_final
    error_s = 2*tol_s
    error_cp = 2*tol_c
    loop_countor = 0
    while (error_s>tol_s) || (error_cp>tol_c)
        loop_countor += 1
        if loop_countor > max_int_loop
            sw_val = copyCell(sw_init)
            cp_val  = copyCell(cp_init)
            cs_val  = copyCell(cs_init)
            p_val  = copyCell(p_init)
            dt = dt/3.0
            break
        end
        ∇p0      = gradientTerm(p_val)
        ∇s0      = gradientTerm(sw_val)

        cs_face   = upwindMean(cs_val, uw)
        cp_face   = upwindMean(cp_val, uw)
        sw_face       = upwindMean(sw_val, uw)

        mu_water_face = faceEval(polymer_viscosity, cs_face, cp_face)
        d_mu_cs_face = faceEval(dpolymer_viscosity_dcs, cs_face, cp_face)
        d_mu_cp_face = faceEval(dpolymer_viscosity_dcp, cs_face, cp_face)
        d_cs_mu_dcs = (mu_water_face-d_mu_cs_face.*cs_face)./(mu_water_face.*mu_water_face)
        d_cs_mu_dcp = (-d_mu_cp_face.*cs_face)./(mu_water_face.*mu_water_face)
        d_cp_mu_dcp = (mu_water_face-d_mu_cp_face.*cp_face)./(mu_water_face.*mu_water_face)
        d_cp_mu_dcs = (-d_mu_cs_face.*cp_face)./(mu_water_face.*mu_water_face)
        d_mu_dcs = -d_mu_cs_face./(mu_water_face.*mu_water_face)
        d_mu_dcp = -d_mu_cp_face./(mu_water_face.*mu_water_face)

        # println("visc")
        krw_face      = faceEval(KRW,sw_face)
        kro_face      = faceEval(KRO,sw_face)

        λ_w_face      = k_face.*krw_face./mu_water_face
        λ_o_face      = k_face.*kro_face/mu_oil
        uw            = -λ_w_face.*∇p0
        uo            = -λ_o_face.*∇p0
        ut            = uw + uo

        dλ_w_face     = k_face.*faceEval(dKRW,sw_face)./mu_water_face
        dλ_o_face     = k_face.*faceEval(dKRO,sw_face)./mu_oil

        # water mass balance (wmb)
        M_t_s_wmb, RHS_t_s_wmb = transientTerm(sw_init, dt, φ)
        M_d_p_wmb = diffusionTerm(-λ_w_face)
        # M_a_s_wmb = convectionUpwindTerm(-dλ_w_face.*∇p0, ut)
        M_a_s_wmb, RHS_a_s_wmb = convectionTvdTerm(-dλ_w_face.*∇p0, sw_val, FL, ut)
        # M_a_cs_wmb = convectionUpwindTerm(-k_face.*krw_face.*d_mu_dcs.*∇p0, ut)
        M_a_cs_wmb, RHS_a_cs_wmb = convectionTvdTerm(-k_face.*krw_face.*d_mu_dcs.*∇p0, cs_val, FL, ut)
        # M_a_cp_wmb = convectionUpwindTerm(-k_face.*krw_face.*d_mu_dcp.*∇p0, ut)
        M_a_cp_wmb, RHS_a_cp_wmb = convectionTvdTerm(-k_face.*krw_face.*d_mu_dcp.*∇p0, cp_val, FL, ut)

        # oil mass balance (omb)
        M_t_s_omb, RHS_t_s_omb = transientTerm(sw_init, dt, -φ)
        M_d_p_omb = diffusionTerm(-λ_o_face)
        # M_a_s_omb = convectionUpwindTerm(-(dλ_o_face.*∇p0), ut)
        M_a_s_omb, RHS_a_s_omb = convectionTvdTerm(-(dλ_o_face.*∇p0), sw_val, FL, ut)

        # salt mass balance
        M_t_cs_smb, RHS_t_cs_smb = transientTerm(cs_init, dt, sw_val.*φ)
        M_t_s_smb, RHS_t_s_smb = transientTerm(sw_init, dt, cs_val.*φ)
        # M_a_cs_smb, RHS_a_cs_smb = convectionTvdTerm(uw, cs_val, FL, ut)
        # M_a_s_smb = convectionUpwindTerm(-dλ_w_face.*cs_face.*∇p0, ut)
        M_a_s_smb, RHS_a_s_smb = convectionTvdTerm(-dλ_w_face.*cs_face.*∇p0, sw_val, FL, ut)
        M_d_p_smb = diffusionTerm(-λ_w_face.*cs_face)
        # M_a_cs_smb = convectionUpwindTerm(-k_face.*krw_face.*d_cs_mu_dcs.*∇p0, ut)
        # M_a_cp_smb = convectionUpwindTerm(-k_face.*krw_face.*d_cs_mu_dcp.*∇p0, ut)
        M_a_cs_smb, RHS_a_cs_smb = convectionTvdTerm(-k_face.*krw_face.*d_cs_mu_dcs.*∇p0, cs_val, FL, ut)
        M_a_cp_smb, RHS_a_cp_smb = convectionTvdTerm(-k_face.*krw_face.*d_cs_mu_dcp.*∇p0, cp_val, FL, ut)

        # polymer mass balance
        M_t_cp_pmb, RHS_t_cp_pmb = transientTerm(cp_init, dt, sw_val.*φ)
        M_t_s_pmb, RHS_t_s_pmb = transientTerm(sw_init, dt, cp_val.*φ)
        # M_a_s_pmb = convectionUpwindTerm(-dλ_w_face.*cp_face.*∇p0, ut)
        M_a_s_pmb, RHS_a_s_pmb = convectionTvdTerm(-dλ_w_face.*cp_face.*∇p0, sw_val, FL, ut)
        # M_a_cp, RHS_a_cp = convectionTvdTerm(uw, cp_val, FL, ut)
        M_d_p_pmb = diffusionTerm(-λ_w_face.*cp_face)
        # M_a_cs_pmb = convectionUpwindTerm(-k_face.*krw_face.*d_cp_mu_dcs.*∇p0, ut)
        # M_a_cp_pmb = convectionUpwindTerm(-k_face.*krw_face.*d_cp_mu_dcp.*∇p0, ut)
        M_a_cs_pmb, RHS_a_cs_pmb = convectionTvdTerm(-k_face.*krw_face.*d_cp_mu_dcs.*∇p0, cs_val, FL, ut)
        M_a_cp_pmb, RHS_a_cp_pmb = convectionTvdTerm(-k_face.*krw_face.*d_cp_mu_dcp.*∇p0, cp_val, FL, ut)

        # create the PDE system M [p;s;c]=RHS
        # x_val = [p_val.value[:]; sw_val.value[:]; c_val.value[:]]
        M = [M_bc_p+M_d_p_wmb   M_t_s_wmb+M_a_s_wmb  M_a_cs_wmb  M_a_cp_wmb;
            M_d_p_omb   M_bc_s+M_t_s_omb+M_a_s_omb  zeros(Nx+2,Nx+2)  zeros(Nx+2,Nx+2);
            M_d_p_smb  M_a_s_smb+M_t_s_smb  M_a_cs_smb+M_t_cs_smb+M_bc_cs  M_a_cp_smb;
            M_d_p_pmb  M_a_s_pmb+M_t_s_pmb  M_a_cs_pmb  M_t_cp_pmb+M_bc_cp+M_a_cp_pmb]

        RHS = [RHS_bc_p+RHS_t_s_wmb+RHS_a_s_wmb+RHS_a_cs_wmb+RHS_a_cp_wmb+M_a_s_wmb*sw_val.value[:]+M_a_cs_wmb*cs_val.value[:]+M_a_cp_wmb*cp_val.value[:];
               RHS_bc_s+RHS_t_s_omb+RHS_a_s_omb+(M_a_s_omb)*sw_val.value[:];
               RHS_t_cs_smb+RHS_t_s_smb+RHS_bc_cs+RHS_a_s_smb+RHS_a_cs_smb+RHS_a_cp_smb+M_a_s_smb*sw_val.value[:]+M_a_cs_smb*cs_val.value[:]+M_a_cp_smb*cp_val.value[:];
               RHS_t_cp_pmb+RHS_t_s_pmb+RHS_bc_cp+RHS_a_s_pmb+RHS_a_cs_pmb+RHS_a_cp_pmb+M_a_s_pmb*sw_val.value[:]+M_a_cs_pmb*cs_val.value[:]+M_a_cp_pmb*cp_val.value[:]]

        # x_sol = solveLinearPDE(m, M, RHS)
        x_sol = 0
        try
            x_sol = M\RHS
        catch
            loop_countor = max_int_loop
            continue
        end

        p_new = x_sol[1:Nx+2]
        s_new = x_sol[(1:Nx+2)+(Nx+2)]
        cs_new = x_sol[(1:Nx+2)+2*(Nx+2)]
        cp_new = x_sol[(1:Nx+2)+3*(Nx+2)]

        error_s = sum(abs, s_new[2:end-1]-internalCells(sw_val))
        error_cp = sum(abs, cp_new[2:end-1]-internalCells(cp_val))

        #   error_c = sum(abs, internalCells(cp_val)-internalCells(cp_val))
        # println(error_s)
        # println(error_c)
        # w_p = 0.9
        # w_sw = 0.8
        # w_c = 0.8
        # p_val.value[:] = w_p*p_new[:]+(1-w_p)*p_val.value[:]
        # sw_val.value[:] = w_sw*s_new[:]+(1-w_sw)*sw_val.value[:]
        # c_val.value[:] = w_c*c_new[:]+(1-w_c)*c_val.value[:]
        p_val.value[:] = p_new[:]
        dsw_apple = 0.2
        eps_apple = sqrt(eps()) # 1e-5
        for i in eachindex(s_new)
            if s_new[i]>=(1-sor)
                if sw_val.value[i]<(1-sor-eps_apple)
                    sw_val.value[i] = 1-sor-eps_apple
                else
                    sw_val.value[i] = 1-sor
                end
            elseif s_new[i]<=swc
                if sw_val.value[i]> swc+eps_apple
                    sw_val.value[i] = swc+eps_apple
                else
                    sw_val.value[i] = swc
                end
            elseif abs(s_new[i]-sw_val.value[i])>dsw_apple
                sw_val.value[i] += dsw_apple*sign(s_new[i]-sw_val.value[i])
            else
                sw_val.value[i] = s_new[i]
            end
        end
        sw_val = createCellVariable(m, internalCells(sw_val), BCs)
        cs_new[cs_new.<0.0] = 0.0
        cp_new[cp_new.<0.0] = 0.0

        cs_val.value[:] = cs_new[:]
        cp_val.value[:] = cp_new[:]

        cs_val = createCellVariable(m, cs_new[2:end-1], BCcs)
        cp_val = createCellVariable(m, cp_new[2:end-1], BCcp)

        # GR.plot(sw_val.value)
        # sw_val.value[:] = w_sw*s_new[:]+(1-w_sw)*sw_val.value[:]
    end
    if loop_countor<max_int_loop
        p_init = copyCell(p_val)
        sw_init = copyCell(sw_val)
        cp_init = copyCell(cp_val)
        cs_init = copyCell(cs_val)

        t +=dt
        dt = dt0
        println("progress: $(t/(t_final)*100) [%]")
    end
end
# JLD.save("results2D/all_data_heterogen_wf.jld", "p", p_val, "c", c_val,
# "c_oil", c_oil, "sw", sw_val, "c_tracer", c_t_val)
# figure(1)
figure()
visualizeCells(sw_init)
# title("Sw")
visualizeCells(cs_init)
# title("c_DME")
visualizeCells(cp_init)
# title("c_tracer")
# savefig("profile.png")
# figure(2)
# plot(t_s, rec_fact)
# savefig("recovery.png")
# figure(3)
# plot(t_s, dp_hist)
# savefig("dp.png")
# PyPlot.pcolormesh(sw_val.value[2:end-1, 2:end-1]', cmap="YlGnBu", shading = "gouraud", vmin=0, vmax=1)


module TestTrophicRank
    using Test
    using BioEnergeticFoodWebs

    food_chain = [0 0 0;0 0 1; 1 0 0]
    @assert trophic_rank(food_chain) == vec([1.0 3.0 2.0])

    food_chain = [0 1 0; 0 0 1; 0 0 0]
    @assert trophic_rank(food_chain) == vec([3.0 2.0 1.0])

    omnivory = [0 1 1; 0 0 1; 0 0 0]
    @assert trophic_rank(omnivory) == vec([2.5 2.0 1.0])

    diamond = [0 1 1 0; 0 0 0 1; 0 0 0 1; 0 0 0 0]
    @assert trophic_rank(diamond) == vec([3.0 2.0 2.0 1.0])

    chain_four = [0 1 0 0; 0 0 1 0; 0 0 0 1; 0 0 0 0]
    @assert trophic_rank(chain_four) == vec([4.0 3.0 2.0 1.0])

    @assert trophic_rank(zeros(Int64, (3, 3))) == ones(3)

end


using Plasma

TD = 30000 # eV
Te = 10000 # eV

D = species.D
e = species.e

D_D = Distribution(Maxwellian(TD, D.m), D)
D_e = Distribution(Maxwellian(Te, e.m), e)
G = Geometry()

plasma = ElectrostaticPlasma([D_D, D_e], G)

sol = Plasma.solve(plasma, dim=1, GPU=false)

Plasma.plot(sol)

## 2D with custom P and species
using Plasma

Tα = 70000 # eV

α = Species(1.602176634e-19, 6.6446562e-27)

function HotCarrier(T)
    Kb = 8.617333262145e-5
    P(x,v) = exp(-v/(Kb*T))
end

Dα = Distribution(HotCarrier(Tα), α)
G = Geometry() # TODO define a custom geometry

plasma = ElectrostaticPlasma([Dα], G)

Plasma.solve(plasma, dim=2) # with GPU

Plasma.plot(sol)

## 3D CollisionlessPlasma
using Plasma

TD = 15000 # eV
TT = 15000 # eV
Te = 13000 # eV

e = species.e
T = species.T
D = species.D

De = Distribution(Maxwellian(Te, e.m), e)
DT = Distribution(Maxwellian(TT, T.m), T)
DD = Distribution(Maxwellian(TD, D.m), D)
G = Geometry()

plasma = CollisionlessPlasma([De,DT,DD], G)

Plasma.solve(plasma)

Plasma.plot(sol)


using Printf
using JLD2
using DataDeps
using OceanTurb
using CairoMakie
using WindMixing
using OceanParameterizations
using Images, FileIO
using ImageTransformations

include("modified_pacalowski_philander_model.jl")

train_files = ["wind_-5.5e-4_new"]
𝒟 = WindMixing.data(train_files, scale_type=ZeroMeanUnitVarianceScaling, enforce_surface_fluxes=false)

NDE_profile_path = joinpath(pwd(), "final_results", "18sim_old", "test_$(train_files[1])", "profiles_fluxes_oceananigans.jld2")
file = jldopen(NDE_profile_path)
data = file["NDE_profile"]
close(file)

Nz = 32
Lz = 256
u₀ = 𝒟.u.coarse[:,1]
v₀ = 𝒟.v.coarse[:,1]
T₀ = 𝒟.T.coarse[:,1]
Fu = 𝒟.uw.coarse[end,1]
Fθ = 𝒟.wT.coarse[end,1]
f₀ = 1f-4
g = 9.80655f0
α = 2f-4
∂T₀∂z_bottom = (T₀[end-1] - T₀[end]) / (Lz / Nz)

times = 𝒟.t

# constants_kpp = (f=f₀, α=α, g=g, Nz=Nz, H=Lz)
# BCs_kpp = (uw=(top=𝒟.uw.coarse[end, 1], bottom=𝒟.uw.coarse[1, 1]),
#             vw=(top=𝒟.vw.coarse[end, 1], bottom=𝒟.uw.coarse[1, 1]),
#             wT=(top=𝒟.wT.coarse[end, 1], bottom=𝒟.wT.coarse[1, 1]))
# ICs_kpp = (u=u₀, v=v₀, T=T₀)

# column_model_1D_kpp(constants_kpp, BCs_kpp, ICs_kpp, times, OceanTurb.KPP.Parameters())

constants_pp = OceanTurb.Constants(Float64, f=f₀)
parameters_pp = PacanowskiPhilander.Parameters()
model_pp = PacanowskiPhilander.Model(N=Nz, L=Lz, stepper=:BackwardEuler, constants=constants_pp, parameters=parameters_pp)

model_pp.bcs[1].top = OceanTurb.FluxBoundaryCondition(Fu)
model_pp.bcs[3].top = OceanTurb.FluxBoundaryCondition(Fθ)

model_pp.solution[1].data[1:Nz] .= u₀
model_pp.solution[2].data[1:Nz] .= v₀
model_pp.solution[3].data[1:Nz] .= T₀

Δt = 1.0
Nt = length(times)

u_solution_pp = zeros(Nz, Nt)
v_solution_pp = zeros(Nz, Nt)
T_solution_pp = zeros(Nz, Nt)

uw_solution_pp = zeros(Nz+1, Nt)
vw_solution_pp = zeros(Nz+1, Nt)
wT_solution_pp = zeros(Nz+1, Nt)

Ri_solution_pp = zeros(Nz+1, Nt)

function get_diffusive_flux(field_index, model)
    flux = FaceField(model.grid)
    field = model.solution[field_index]
    K = model.timestepper.eqn.K[field_index]
    for i in interiorindices(flux)
        @inbounds flux[i] = - K(model, i) * ∂z(field, i)
    end
    return flux
end

function get_richardson_number_profile(model)
    Ri = FaceField(model.grid)
    for i in interiorindices(Ri)
        @inbounds Ri[i] = local_richardson(model, i)
    end
    return Ri
end

for n in 1:Nt
    OceanTurb.run_until!(model_pp, Δt, times[n])
    @info "Time = $(times[n])"

    u_solution_pp[:, n] .= model_pp.solution[1][1:Nz]
    v_solution_pp[:, n] .= model_pp.solution[2][1:Nz]
    T_solution_pp[:, n] .= model_pp.solution[3][1:Nz]

    uw_solution_pp[:, n] .= get_diffusive_flux(1, model_pp)[1:Nz+1]
    vw_solution_pp[:, n] .= get_diffusive_flux(2, model_pp)[1:Nz+1]
    wT_solution_pp[:, n] .= get_diffusive_flux(3, model_pp)[1:Nz+1]

    uw_solution_pp[Nz+1, n] = Fu
    wT_solution_pp[Nz+1, n] = Fθ

    # Ri_solution_pp[:, n] = get_richardson_number_profile(model_pp)[1:Nz+1]
end

D_face = Dᶠ(32, 256 / 32)

function local_richardson_profile(u, v, T, g, α)
    ∂u∂z = D_face * u
    ∂v∂z = D_face * v
    ∂T∂z = D_face * T
    Bz = g .* α .* ∂T∂z
    S² = ∂u∂z .^ 2 .+ ∂v∂z .^ 2
    return Bz ./ S²
end

for i in 1:size(Ri_solution_pp, 2)
    Ri_solution_pp[:,i] = local_richardson_profile(u_solution_pp[:,i], v_solution_pp[:,i], T_solution_pp[:,i], g, α)
end

# Get rid of ∞ and super large values.
Ri_solution_pp = clamp.(Ri_solution_pp, -1, 2)

u_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "u.png"))))
v_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "v.png"))))
T_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "T.png"))))
uw_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "uw.png"))))
vw_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "vw.png"))))
wT_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "wT.png"))))
Ri_img = rotr90(load(assetpath(joinpath(pwd(), "images_plot", "Ri.png"))))
z_img = load(assetpath(joinpath(pwd(), "images_plot", "z.png")))

axis_images = (
    u = u_img,
    v = v_img,
    T = T_img,
    z = z_img,
    uw = uw_img,
    vw = vw_img,
    wT = wT_img,
    Ri = Ri_img,
)

FILE_PATH = joinpath(pwd(), "final_results", "base_kpp_pp")
fps=60

times_days = data["t"] ./ 86400

frame = Node(1)

time_point = @lift [times_days[$frame]]

u_data = [
    data["truth_u"],
    # data["test_u_modified_pacanowski_philander"],
    data["test_u_kpp"],
    # data["test_u"],
    u_solution_pp,
]

v_data = [
    data["truth_v"],
    # data["test_v_modified_pacanowski_philander"],
    data["test_v_kpp"],
    # data["test_v"],
    v_solution_pp,
]

T_data = [
    data["truth_T"],
    # data["test_T_modified_pacanowski_philander"],
    data["test_T_kpp"],
    # data["test_T"],
    T_solution_pp,
]

uw_data = [
    data["truth_uw"],
    # data["test_uw_modified_pacanowski_philander"],
    data["test_uw_kpp"],
    # data["test_uw"],
    uw_solution_pp,
]

vw_data = [
    data["truth_vw"],
    # data["test_vw_modified_pacanowski_philander"],
    data["test_vw_kpp"],
    # data["test_vw"],
    vw_solution_pp,
]

wT_data = [
    data["truth_wT"],
    # data["test_wT_modified_pacanowski_philander"],
    data["test_wT_kpp"],
    # data["test_wT"],
    wT_solution_pp,
]

uw_data .*= 1f4
vw_data .*= 1f4
wT_data .*= 1f5

Ri_data = [
    clamp.(data["truth_Ri"], -1, 2),
    # clamp.(data["test_Ri_modified_pacanowski_philander"], -1, 2),
    clamp.(data["test_Ri_kpp"], -1, 2),
    # clamp.(data["test_Ri"], -1, 2),
    Ri_solution_pp
]

u_frames = [@lift data[:,$frame] for data in u_data]
v_frames = [@lift data[:,$frame] for data in v_data]
T_frames = [@lift data[:,$frame] for data in T_data]

uw_frames = [@lift data[:,$frame] for data in uw_data]
vw_frames = [@lift data[:,$frame] for data in vw_data]
wT_frames = [@lift data[:,$frame] for data in wT_data]

Ri_frames = [@lift data[:,$frame] for data in Ri_data]

# losses_point_frames = [@lift [data[$frame]] for data in losses_data]

@inline function find_lims(datasets)
    return maximum(maximum.(datasets)), minimum(minimum.(datasets))
end

u_max, u_min = find_lims(u_data)
v_max, v_min = find_lims(v_data)
T_max, T_min = find_lims(T_data)

uw_max, uw_min = find_lims(uw_data)
vw_max, vw_min = find_lims(vw_data)
wT_max, wT_min = find_lims(wT_data)

# losses_max, losses_min = find_lims(losses_data)

train_parameters = data["train_parameters"]
ν₀ = train_parameters.ν₀
ν₋ = train_parameters.ν₋
ΔRi = train_parameters.ΔRi
Riᶜ = train_parameters.Riᶜ
Pr = train_parameters.Pr
loss_scalings = train_parameters.loss_scalings

BC_str = @sprintf "Momentum Flux = %.1e m² s⁻², Temperature Flux = %.1e m s⁻¹ °C" data["truth_uw"][end, 1] maximum(data["truth_wT"][end, :])

plot_title = @lift "Traditional Parameterisations: $BC_str, Time = $(round(times_days[$frame], digits=2)) days"

# fig = Figure(resolution=(1920, 1080))
fig = Figure(resolution=(1920, 960))

# colors = distinguishable_colors(length(uw_data)+1, [RGB(1,1,1), RGB(0,0,0)], dropseed=true)
colors = distinguishable_colors(length(uw_data), [RGB(1,1,1), RGB(0,0,0)], dropseed=true)

u_img = axis_images.u
v_img = axis_images.v
T_img = axis_images.T
uw_img = axis_images.uw
vw_img = axis_images.uw
wT_img = axis_images.uw
Ri_img = axis_images.uw
z_img = axis_images.z

zc = data["depth_profile"]
zf = data["depth_flux"]
zf_interior = zf[2:end-1]

rel_size = 30
# aspect = 1920 / 1080
aspect = 2

ax_u = fig[1, 2] = CairoMakie.Axis(fig)
ax_v = fig[1, 4] = CairoMakie.Axis(fig)

T_layout = fig[1:4, 5] = GridLayout()
colsize!(fig.layout, 5, CairoMakie.Relative(0.4))

ax_T = T_layout[1, 2] = CairoMakie.Axis(fig)
y_ax_T = T_layout[1,1] = CairoMakie.Axis(fig, aspect=DataAspect())
x_ax_T = T_layout[2,2] = CairoMakie.Axis(fig, aspect=DataAspect())

ax_Ri = fig[1, 7] = CairoMakie.Axis(fig)
ax_uw = fig[3, 2] = CairoMakie.Axis(fig)
ax_vw = fig[3, 4] = CairoMakie.Axis(fig)
ax_wT = fig[3, 7] = CairoMakie.Axis(fig)

y_ax_u = CairoMakie.Axis(fig[1,1], aspect=DataAspect())
y_ax_v = CairoMakie.Axis(fig[1,3], aspect=DataAspect())
y_ax_Ri = CairoMakie.Axis(fig[1,6], aspect=DataAspect())
y_ax_uw = CairoMakie.Axis(fig[3,1], aspect=DataAspect())
y_ax_vw = CairoMakie.Axis(fig[3,3], aspect=DataAspect())
y_ax_wT = CairoMakie.Axis(fig[3,6], aspect=DataAspect())

x_ax_u = CairoMakie.Axis(fig[2,2], aspect=DataAspect())
x_ax_v = CairoMakie.Axis(fig[2,4], aspect=DataAspect())
x_ax_Ri = CairoMakie.Axis(fig[2,7], aspect=DataAspect())
x_ax_uw = CairoMakie.Axis(fig[4,2], aspect=DataAspect())
x_ax_vw = CairoMakie.Axis(fig[4,4], aspect=DataAspect())
x_ax_wT = CairoMakie.Axis(fig[4,7], aspect=DataAspect())

hidedecorations!(y_ax_u)
hidedecorations!(y_ax_v)
hidedecorations!(y_ax_Ri)
hidedecorations!(y_ax_uw)
hidedecorations!(y_ax_vw)
hidedecorations!(y_ax_wT)

hidedecorations!(x_ax_u)
hidedecorations!(x_ax_v)
hidedecorations!(x_ax_Ri)
hidedecorations!(x_ax_uw)
hidedecorations!(x_ax_vw)
hidedecorations!(x_ax_wT)

hidespines!(y_ax_u)
hidespines!(y_ax_v)
hidespines!(y_ax_T)
hidespines!(y_ax_Ri)
hidespines!(y_ax_uw)
hidespines!(y_ax_vw)
hidespines!(y_ax_wT)

hidespines!(x_ax_u)
hidespines!(x_ax_v)
hidespines!(x_ax_T)
hidespines!(x_ax_Ri)
hidespines!(x_ax_uw)
hidespines!(x_ax_vw)
hidespines!(x_ax_wT)

image!(x_ax_u, axis_images.u)
image!(x_ax_v, axis_images.v)
image!(x_ax_Ri, axis_images.Ri)
image!(x_ax_uw, axis_images.uw)
image!(x_ax_vw, axis_images.vw)
image!(x_ax_wT, axis_images.wT)

image!(y_ax_u, axis_images.z)
image!(y_ax_v, axis_images.z)
image!(y_ax_Ri, axis_images.z)
image!(y_ax_uw, axis_images.z)
image!(y_ax_vw, axis_images.z)
image!(y_ax_wT, axis_images.z)

hidedecorations!(y_ax_T)
hidedecorations!(x_ax_T)
image!(x_ax_T, axis_images.T)
image!(y_ax_T, axis_images.z)

colsize!(T_layout, 1, CairoMakie.Relative(1 / rel_size * 1.5))
rowsize!(T_layout, 2, CairoMakie.Relative(1 / rel_size / aspect * 1.75))
colgap!(T_layout, Relative(1 / rel_size / aspect / 2))
rowgap!(T_layout, Relative(1 / rel_size / aspect))

rowsize!(fig.layout, 2, CairoMakie.Relative(1 / rel_size))
rowsize!(fig.layout, 4, CairoMakie.Relative(1 / rel_size))
colsize!(fig.layout, 1, CairoMakie.Relative(1 / rel_size / aspect))
colsize!(fig.layout, 3, CairoMakie.Relative(1 / rel_size / aspect))
colsize!(fig.layout, 6, CairoMakie.Relative(1 / rel_size / aspect))

colgap!(fig.layout, Relative(1 / rel_size / aspect / 2))
rowgap!(fig.layout, Relative(1 / rel_size / aspect))

alpha=0.5
truth_linewidth = 7
linewidth = 3

CairoMakie.xlims!(ax_u, u_min, u_max)
CairoMakie.xlims!(ax_v, v_min, v_max)
CairoMakie.xlims!(ax_T, T_min, T_max)
CairoMakie.xlims!(ax_uw, uw_min, uw_max)
CairoMakie.xlims!(ax_vw, vw_min, vw_max)
CairoMakie.xlims!(ax_wT, wT_min, wT_max)
CairoMakie.xlims!(ax_Ri, -1, 2)
# CairoMakie.xlims!(ax_losses, times[1], times[end])

CairoMakie.ylims!(ax_u, minimum(zc), 0)
CairoMakie.ylims!(ax_v, minimum(zc), 0)
CairoMakie.ylims!(ax_T, minimum(zc), 0)
CairoMakie.ylims!(ax_uw, minimum(zf), 0)
CairoMakie.ylims!(ax_vw, minimum(zf), 0)
CairoMakie.ylims!(ax_wT, minimum(zf), 0)
CairoMakie.ylims!(ax_Ri, minimum(zf), 0)
# CairoMakie.ylims!(ax_losses, losses_min, losses_max)

u_lines = [
     lines!(ax_u, u_frames[1], zc, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_u, u_frames[i], zc, linewidth=linewidth, color=colors[i]) for i in 2:length(u_data)]
]

v_lines = [
     lines!(ax_v, v_frames[1], zc, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_v, v_frames[i], zc, linewidth=linewidth, color=colors[i]) for i in 2:length(v_data)]
]

T_lines = [
     lines!(ax_T, T_frames[1], zc, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_T, T_frames[i], zc, linewidth=linewidth, color=colors[i]) for i in 2:length(T_data)]
]

uw_lines = [
     lines!(ax_uw, uw_frames[1], zf, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_uw, uw_frames[i], zf, linewidth=linewidth, color=colors[i]) for i in 2:length(uw_data)]
]

vw_lines = [
     lines!(ax_vw, vw_frames[1], zf, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_vw, vw_frames[i], zf, linewidth=linewidth, color=colors[i]) for i in 2:length(vw_data)]
]

wT_lines = [
    lines!(ax_wT, wT_frames[1], zf, linewidth=truth_linewidth, color=(colors[1], alpha));
   [lines!(ax_wT, wT_frames[i], zf, linewidth=linewidth, color=colors[i]) for i in 2:length(wT_data)]
]

Ri_lines = [
     lines!(ax_Ri, Ri_frames[1], zf, linewidth=truth_linewidth, color=(colors[1], alpha));
    [lines!(ax_Ri, Ri_frames[i], zf, linewidth=linewidth, color=colors[i]) for i in 2:length(Ri_data)]
]

axislegend(ax_T, T_lines, ["Oceananigans.jl Large Eddy Simulation", "K-Profile Parameterisation", "Pacanowski-Philander Model"], "Data Type", position = :rb)

supertitle = fig[0, :] = Label(fig, plot_title, textsize=25)

trim!(fig.layout)

print_frame = maximum([1, Int(floor(length(times)/20))])

function print_progress(n, n_total, print_frame, type)
    if n % print_frame == 0
        @info "Animating $(type) frame $n/$n_total"
    end
end

@info "Starting Animation"

CairoMakie.record(fig, "$FILE_PATH.mp4", 1:length(times_days), framerate=fps, compression=1) do n
    print_progress(n, length(times_days), print_frame, "mp4")
    frame[] = n
end


# A collection of iterated maps

LogisticMap(r) = IteratedMap{:LogisticMap,1}(x->r*x*(1-x))

LogisticMap(r,x) = LogisticMap(r)(x)

HenonMap(p) = IteratedMap{:HenonMap,2}(x->[1-p[1]*x[1]^2+x[2],p[2]*x[1]])
HenonMap(p,x) = HenonMap(p)(x)

DuffingMap(p) = IteratedMap{:DuffingMap,2}(x->[x[2],-p[2]*x[1]+p[1]*x[2]-x[2]^3])
DuffingMap(p,x) = DuffingMap(p)(x)

TentMap(μ) = IteratedMap{:TentMap,1}(x->x<0.5?μ*x:μ*(1-x))
TentMap(μ,x) = TentMap(μ)(x)