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

using Markdown
using InteractiveUtils

# ╔═╡ d941d6c2-55bf-11eb-0002-35c7474e4050
using NiLang, Test

# ╔═╡ 2061b434-0ad1-46eb-a0c7-1a5f432bfa62
begin
	twocol(left, right; llabel="forward", rlabel="backward") = HTML("

	$llabel
	$rlabel

$(html(left))

$(html(right))

")
	example(str) = HTML("""
$str
""")
	title1(str) = HTML("""

'_'))>$str

""")
	title2(str) = HTML("""
'_'))>$str
""")
	titleref(str) = HTML("""'_'))">$str""")
	using PlutoUI: TableOfContents
	using Pkg
	pkgversion(m::Module) = Pkg.TOML.parsefile(NiLang.project_relative_path("Project.toml"))["version"]
	hightlight(str) = HTML("$str")
end;

# ╔═╡ 8c2c4fa6-172f-4dde-a279-5d0aecfdbe46
module M
using NiLang

# define two functions
function new_forward(x)
	if x > 0
		return x * 2
	elseif x < 0
		return x / 2
	end
end

function new_backward(x)
	if x > 0
		return x / 2
	elseif x < 0
		return x * 2
	end
end

# declare them as reversible to each other
@dual new_forward new_backward

# The following is need only when your function is differentiable
using NiLang.AD: GVar
function new_backward(x::GVar)
	if x.x > 0
		GVar(new_backward(x.x), x.g * 2)
	elseif x.x < 0
		GVar(new_backward(x.x), x.g / 2)
	end
end

function new_forward(x::GVar)
	if x.x > 0
		GVar(new_forward(x.x), x.g / 2)
	elseif x.x < 0
		GVar(new_forward(x.x), x.g * 2)
	end
end
end

# ╔═╡ 3199a048-7b39-40f8-8183-6a54cccd91b6
using BenchmarkTools

# ╔═╡ 0e1ba158-a6bc-401c-9ba7-ed78020ad068
using Base.Threads

# ╔═╡ a4e76427-f051-4b29-915a-fdfce3a299bb
html"""

  <svg class="octicon octicon-mark-github v-align-middle" height="32" viewBox="0 0 16 16" version="1.1" width="32" aria-hidden="true"><path fill-rule="evenodd" d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.27.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0016 8c0-4.42-3.58-8-8-8z"></path></svg>

Download this notebook

"""

# ╔═╡ c2c7b4d4-f8c9-4ebf-8da2-0103f03136e7
md"# NiLang's (v$(pkgversion(NiLang))) Documentation

NiLang is a embeded domain specific language (eDSL) in Julia, so one need to install [Julia](https://julialang.org/) first. Before reading this documentation, you need to know basic Julia grammar, and how to install and use packages.
Also, it might be good to read the [README](https://github.com/GiggleLiu/NiLang.jl) first. In this tutorial, we focus on

* NiLang grammar and design patterns,
* Automatic differentiation based on reversible programming.
"

# ╔═╡ 12f07cc7-979c-43c3-9dc9-36ea1463c1f6
md"The symbols used in this notebook"

# ╔═╡ 611b577f-4722-42bf-8f8e-aeb2fb30be71
md"""
|  symbol  |  meaning  | how to type |
| -------  | --------- | ----------- |
| ← | allocate |  \leftarrow + TAB |
| → | deallocate | \rightarrow + TAB |
| ↔ | exchange | \leftrightarrow + TAB |
| ∅ | empty variable | \emptyset + TAB |
| ~ | inverse | ~ |
"""

# ╔═╡ 605872cf-f3fd-462e-a2b1-7d1c5ae45efd
title1("Getting started")

# ╔═╡ fb3dee44-5fa9-4773-8b7f-a83c44358545
md"
After installing NiLang in a Julia REPL by typing `]add NiLang`, one can use NiLang and use the macro `@i` to define a reversible function ``f_1: (x, y) → (x+5y, y)``."

# ╔═╡ 70088425-6779-4a2d-ba6d-b0a34c8e93a6
@i @inline function f1(x, y; constant)
	x += y * constant
end

# ╔═╡ af738f89-3214-429c-9c7d-18a6ea0d9401
f1(1.0, 2.0; constant=5.0)  # call

# ╔═╡ 48d7ebc1-5def-4a57-9ec1-3fc370a4543f
 (~f1)(11.0, 2.0; constant = 5.0)  # uncall

# ╔═╡ f0e94247-f615-472b-8218-3fa287b38aa1
md"A NiLang function defines a ``\mathbb{R}^n\rightarrow\mathbb{R}^n`` (notice inputs and outputs have the same shapes) mapping, it can take both keyword and positional arguments, where positional arguments can only be used as constants. There is no `return` statement because this function returns input non-keyword variables automatically, it is forbiden to write `return` statement inside a NiLang function.
One can aslo put macros like `@inline` after NiLang's `@i` macro. NiLang's macro will render the body of the function first and pass it to other macros.
"

# ╔═╡ 2581aa33-1dc5-40b1-aa9f-6a11cc750c93
md"`x += y * constant` is an instruction that defines a bijective mapping ``\mathbb{R}^3\rightarrow\mathbb{R}^3``. All NiLang instructions change the variable inplace. Here, we accumulate the result to `x` rather than using `y *= constant` to modify the variable directly. This is because in a regular number system, one can easily use the zero element as the eraser to erase all information, which will cause irreversibility."

# ╔═╡ 60575978-081a-4bca-a3ed-2b51cd6abc92
md"One can differentiate a NiLang function with the `NiLang.gradient(f, args; iloss, kwargs...)` where `args` and `kwargs` are positional and keyword arguments for `f`, and `iloss` is the index of the loss variable."

# ╔═╡ f98305cb-4ba2-404a-a5c3-65510e059504
NiLang.AD.gradient(f1, (1.0, 2.0); iloss=1, constant=5.0)

# ╔═╡ e8cd6667-597f-458b-8465-1822e09a7891
md"Here, we specify the first variable as the one that stores the loss. We get
```math
\begin{cases}
\frac{\partial x+5y}{\partial x}=1\\
\frac{\partial x+5y}{\partial y}=5\\
\end{cases}
```
"

# ╔═╡ 20145d75-004a-4c2f-b7ff-c400ca846d42
let
	content = md"""The above macro generates two functions, one is `f` and another is `~f` (or `Inv(f)`). The `x += y * constant` is translated to function `(x, y, constant) = PlusEq(*)(x, y, constant)`, where the function `PlusEq(*)` is bijective.

```julia
julia> using MacroTools

julia> MacroTools.prettify(@macroexpand @i function f(x, y; constant)
           x += y * constant
       end)
quote
    $(Expr(:meta, :doc))
    function $(Expr(:where, :(f(x, y; constant))))
        hare = wrap_tuple(((PlusEq)(*))(x, y, constant))
        x = hare[1]
        y = hare[2]
        constant = hare[3]
        (x, y)
    end
    if (NiLangCore)._typeof(f) != _typeof(~f)
        function $(Expr(:where, :((newt::_typeof(~f))(x, y; constant))))
            boar = wrap_tuple(((MinusEq)(*))(x, y, constant))
            x = boar[1]
            y = boar[2]
            constant = boar[3]
            (x, y)
        end
    end
end
```
"""
	HTML("How does the compiler work?$(html(content))")
end

# ╔═╡ c682a17f-600f-4034-bfe3-a851ab645c10
title1("Instructions and operands")

# ╔═╡ 5239dfe2-ea6d-4e07-a1b1-90954fe8ddc9
md"""
The basic form of an instruction is `y ⊙= f(args...)`, where `⊙` can be `+`, `-`, `*`, `/` or `⊻`, where `*=` and `/=` are only reversible in the logarithmic number systems. See section $(titleref("Integers, floating-point numbers, fixed-point numbers and logarithmic numbers")) for details.

A function/instruction can be used in NiLang only if its reverse is defined. A NiLang function is differentiable if the function body is composed of differentiable instructions, differentiable NiLang functions and NiLang control flows. A list of differentiable NiLang instructions are

| instruction | output   |
| ----------- | ---------- |
| ``{\rm FLIP}(y)`` | ``\sim y`` |
| ``{\rm NEG}(y)`` | ``-y`` |
| ``{\rm INC}(y)`` | ``y+1`` |
| ``{\rm DEC}(y)`` | ``y-1`` |
| ``{\rm INV}(y)`` | ``y ^ {-1}`` |
| ``{\rm HADAMARD}(x, y)`` | ``\frac{1}{\sqrt{2}}(x+y), \frac{1}{\sqrt{2}}(x-y)``
| ``{\rm SWAP}(a, b)`` | ``b, a`` |
| ``{\rm ROT}(a, b, \theta)`` | ``a \cos\theta - b\sin\theta, b \cos\theta + a\sin\theta, \theta`` |
| ``{\rm IROT}(a, b, \theta)`` | ``a \cos\theta + b\sin\theta, b \cos\theta - a\sin\theta, \theta`` |
| ``y \mathrel{\{+,-\}}= f_{+-}(args...)`` | ``y\{+, -\}f_{+-}(args...), args...`` |
| ``y \mathrel{\{*, /\}}= f_{*/}(args...)`` | ``y\{*, /\}f_{*/}(args...), args...`` |

Functions ``f_{+-} ∈ \rm \{identity, +, -, *, /, ^\wedge, abs, abs2, sqrt, exp, log, sin, sinh, asin, cos, cosh,`` ``\rm acos, tan, tanh, atan, sincos, convert\}`` and ``f_{*/}∈\rm \{identity, +, -, *, /, ^\wedge, convert\}.``
Functions `FLIP`, `NEG`, `INV`, `HADAMARD`, `SWAP` and `y ⊻= f_{⊻}(args...)` are self-reversible (or reflexive). {`ROT`, `IROT`} and {`INC`, `DEC`}, {`y += f_{+-}(args...)`, `y -= f_{+-}(args...)`} and {`y *= f_{*/}(args...)`, `y /= f_{*/}(args...)`} are pair-wise reversible.

For Jacobians and Hessians defined on these instructions, please check this [blog post](https://giggleliu.github.io/2020/01/18/jacobians.html).
This set of instructions is extensible, see section $(titleref("Extending the instruction set")) for an example.
"""

# ╔═╡ e4d86a5a-e820-4a70-8a87-08bac416291b
md"The operands of an instruction can be a composition of the following expressions"

# ╔═╡ c307b6a4-906d-4be7-9fd7-57c942aded51
md"""

| expression | meaning |
| ---------- | ------- |
| x | change a variable |
| x.field | change a field |
| x.:1    | change a tuple element |
| x'      | change the adjoint |
| -x      | change the inverse |
| x[i]    | change an array/dict element |
| (x, y, z)  | change multiple variables |
| @fields x | change the fields of an object |
| A{Float64}(x, y) | wrap, update fields and unwrap |
| x \|> subarray(:,i) | change the view of an array |
| x \|> f | change a field map or a bijective mapping |
| @const 3 | target is a constant that can not be changed |
| @skip! f(x) | an expression that can not be assigned back |
"""

# ╔═╡ 47b502d4-e8af-4d58-9067-9700784ea435
md"In the following example, one modifies the negated real part of a complex number in a vector inside a tuple directly."

# ╔═╡ faecc0a7-55d7-42a1-8e9f-7e30143eef9c
@i function dataview_func1(x::Vector, y::Tuple, θ)
	(-x[2].re, y.:1) += sincos(θ)
end

# ╔═╡ 4cbe69d6-b68f-4bda-a0dd-209f9ee54f18
dataview_func1([1+2.0im, 3+4im], (1.0,2.0), π/6)

# ╔═╡ 7dc82f28-77bf-40da-b520-800ed1bc80c9
md"""One can check the real part of the second element of `x` is decreased by `sin(π/6)`, while the first element of the tuple is increased by `cos(π/6)`. A NiLang instruction is different to a regular Julia instruction in that $(hightlight("A NiLang instruction changes variables inplace")), eventhough `ComplexF64`, `-x` and `tuple` are all considered as immutable in Julia."""

# ╔═╡ 9f5f9de3-9558-4c18-9d98-b77d19b570ec
example("Complex valued log")

# ╔═╡ 6dfcfa19-f78f-4dac-89f7-d3c5dbe17987
@i function complex_log(y!::Complex{T}, x::Complex{T}) where T
    n ← zero(T)
    n += abs(x)

    y!.re += log(n)
    y!.im += angle(x)

    n -= abs(x)
    n → zero(T)
end; @test complex_log(0.0im, 3.0+2.0im)[1] ≈ log(3.0+2.0im)

# ╔═╡ dad1c6c0-d61b-4f9f-a71e-e683fe143aaa
title2("Broadcasting")

# ╔═╡ 3e4a3916-8fe4-4262-bcd3-3014822717a3
md"The interfaces is similar to julia's native broadcast, but expanded to native NiLang loops."

# ╔═╡ 34906208-e6f1-4a67-860a-a7b056a86dde
@i function complex_log_broadcast!(ys!, xs)
	complex_log.(ys!, xs)
end; @test (x=[1.0+2im, 3.0+4im, 4.0+5im]; complex_log_broadcast!(zeros(ComplexF64, 3), x)[1] ≈ log.(x))

# ╔═╡ 4601df35-679f-465d-9191-c18748b2fd83
title2("Avoid shared read-write")

# ╔═╡ 4fc72b9d-19a2-40f1-a4a8-5e97d3d5e529
md"Shared read-write is not allowed"

# ╔═╡ a0fde16f-8454-4f5c-a29c-a9e415c0c311
# `y -= y` effectively clears the content in `y`, this is why shared read-write is so dangerous.
@test_throws LoadError macroexpand(NiLang, :(@i function shared_readwrite_error(y)
	y -= y
end))

# ╔═╡ 633ff8f3-8d93-4f73-bec2-c42070e6ece9
md"NiLang is more restrictive, it forbids shared read too. This is in purpose, shared read will become shared write to the gradient field in the back-propagation of gradients - the main goal of NiLang."

# ╔═╡ 57f2d890-b5a0-47b7-9e3d-af4d03b10605
# `shared read is also forbidden`
@test_throws LoadError macroexpand(NiLang, :(@i function shared_read_error(x, y)
	x -= y * y  # should be written as `x -= y^2`
end))

# ╔═╡ 34063cd0-171e-46ce-80dd-52a341fa50a1
md"The correct way of avoiding shared read is renaming one of the variable."

# ╔═╡ 5d5d01db-8ff9-434c-8771-1fec6393e1fb
@i function avoid_shared_read(x, y)
	tmp ← zero(y)
	tmp += y
	x -= y * tmp
	tmp -= y
	tmp → zero(y)
end

# ╔═╡ 10d85a50-f2f9-403e-8f6c-baef61cf702a
avoid_shared_read(0.0, 3.0)

# ╔═╡ b52648bf-a28a-48af-8912-31729d943ce0
md"Shared read-write issue is more tricky when one uses NiLang to write kernel function in a parallel program (multi-threading, MPI and CUDA). See $(titleref(\"Multi-threading and CUDA\")) for details."

# ╔═╡ f45db10f-a836-40f3-9d8d-054ea6540e87
title2("Protect constant variables")

# ╔═╡ 1903563e-ccc2-44d9-9dbe-e5dede275b3c
md"""
To achieve the goal of "everything is mutable", NiLang assigns the output value back to the input variable after a call. Sometime it can cause issues for special variables like function, type parameters and constants (including the results generated from a function call). One can use `@const` (assert a variable is a constant) or `@skip!` (skip assigning back) to avoid such complications.
"""

# ╔═╡ 583f2585-15a3-47c6-a70e-e2f002754028
md"When using a function (e.g. `exp` as shown bellow) as an variable, one should be careful about the scope issue."

# ╔═╡ 60e6ff80-3593-4ae4-a273-914847f692db
@i function func_arg(y, f, x)
	y += f(x)
end

# ╔═╡ 9e5cfd68-b58d-4d83-aae2-447e5f805c97
@i function use_func_arg(y, x)
	func_arg(y, exp, x)
end

# ╔═╡ dc85a942-cf52-4405-ad03-32a768e1b6e7
@test_throws UndefVarError use_func_arg(0.0, 3.0)

# ╔═╡ 5cdd346b-10a5-485c-ba78-4c0b3cb0e02f
md"We see an error, but why calling `f2g` causes an error? If one check the generated code with `macroexpand`, one will see the `exp` is assigned in the local scope. The compiler takes it as a local variable and compaints that `exp` is not defined."

# ╔═╡ af9287b7-6131-46f6-beb8-6885e55e1975
macroexpand(NiLang, :(@i function use_func_arg(y, x)
	f2g(y, exp, x)
end)) |> NiLangCore.rmlines

# ╔═╡ e20eeabf-1c80-431e-8cfc-4d1b79c52b5a
@i function avoid_assignback(y, x)
	func_arg(y, (@const exp), x)
end; @test avoid_assignback(0.0, 3.0)[1] == exp(3)

# ╔═╡ 90d30eea-53de-48a0-9700-ff35681fdf38
md"Type parameter can not be assigned back too."

# ╔═╡ 390f58a5-6f5f-4d3a-bb16-ba04e43a07e7
@test_throws ErrorException Core.eval(NiLang, :(@i function type_arg(t::Type{T}, x) where T
	x += one(T)
end))

# ╔═╡ 2b57443e-a516-434b-be86-80616a98e2f5
@i function avoid_type_assignback(t::Type{T}, x) where T
	x += one(@skip! T)
end; @test avoid_type_assignback(Float64, 0.0)[2] == 1.0

# ╔═╡ fc2e27f9-b7ba-44cc-a953-6745548ad733
md"A function call that returning a constant should also be decorated with the `@const` or `@skip!` macro."

# ╔═╡ fc744931-360b-4478-9f77-c50f048de243
@test_throws LoadError macroexpand(NiLang, :(@i function funccall_arg(y, x::Matrix) where T
	y += size(x, 1) * size(x, 2)
end))

# ╔═╡ 9a152b36-f377-44da-9700-ca9e05e365ff
@i function avoid_funccall_assignback(y, x::Matrix) where T
	y += (@const size(x, 1)) * (@const size(x, 2))
end; @test avoid_funccall_assignback(0, randn(3,4))[1] == 12

# ╔═╡ b6dcd18c-606f-4340-b2ec-163e8bad03f5
title1("Variable manipulation")

# ╔═╡ a1a29f34-f8a9-4e9f-9afe-7d0096771440
title2("Allocate and deallocate a variable")

# ╔═╡ 90bd6ad4-3dd8-4e7c-b445-aed1d248a2ec
md"""
One can allocate a new variable `x` like `x ← constant` and deallocate a variable with known value with `x → known value`. They are reversible to each other with the following relation.

$(
twocol(md"
```julia
x ← constant
```", md"
```julia
x → constant
```
")
)

For example
"""

# ╔═╡ c0259e48-1973-486c-a828-1fcd3e4331c6
@i function alloc_func1()
	tmp ← 1
	# some code that uses `tmp` for computing and restores it to `1`
	tmp → 1
end

# ╔═╡ 8bbffa31-04a6-49ca-b36f-4d4140d75992
md"allocate multiple variable of the same type at one time"

# ╔═╡ a6f18c34-80ee-4b52-9ff8-f3c1b1d80f90
@i function power12(y, x::T) where T
	@zeros T a b c  # three variable of type `T`
	a += x^2
	b += a^2
	c += b*a
	y += c^2
	c -= b*a
	b -= a^2
	a -= x^2
	@safe @show a b c x y
	~@zeros T a b c
end

# ╔═╡ a694132b-4f52-467f-8bc4-dc32fe2812db
@test power12(0, 2)[1] == 4096

# ╔═╡ 8c2c82f2-1240-4f2f-830e-ee8021c1a41a
md"One can copy and push a value into a stack and use it later. It inverse operation will pop out a variable and assert its value."

# ╔═╡ 6203cf10-f8cc-4fb9-b814-7552b68c01dc
twocol(md"
```julia
stack[end+1] ← variable
```", md"
```julia
stack[end] → variable
```
")

# ╔═╡ f97a6bab-b9f9-4b95-98a9-381c51397526
@i function stack_push_and_pop!(stack, x, y, z)
	z += y
	stack[end+1] ← x  # copy a variable into a stack
	stack[end+1] ← y
	stack[end] → z  # pop a variable from a stack, `z` must have the same value as the variable.
end

# ╔═╡ 2a2970f4-ab01-486b-89a2-6ff96f734018
md"A less recommended approach is using the global stacks in NiLang, since NiLang is an eDSL, it can not guarantee the access order. Available global stacks are `FLOAT64_STACK`, `COMPLEXF64_STACK`, `INT64_STACK` and their 32 bit counter parts, as well as a `BOOL_STACK`."

# ╔═╡ 0b80d9be-53d7-4bf3-a558-659607af4709
@i function stack_push_and_pop!(x, y)
	GLOBAL_STACK[end+1] ← x  # copy a variable into a stack
	FLOAT64_STACK[end+1] ← y
end

# ╔═╡ 4ca48a2e-43da-457a-8e9f-6476097e4d7b
let
	stack = FastStack{Float64}(1000) # a preallocated stack of size 1000
	stack_push_and_pop!(stack, 5.0, 1.0, 0.0)  # you will get a stack of size 1
	@test length(stack) == 1
end

# ╔═╡ 92362fda-bae2-4e35-bfe4-dcaea853d50b
let
	NiLang.empty_global_stacks!()  # empty stacks
	stack_push_and_pop!(5.0, 1.0)
	@test length(GLOBAL_STACK) == 1
	@test length(FLOAT64_STACK) == 1
end

# ╔═╡ db9e7940-39f1-4ccf-ac70-146a521daa6e
md"one can also allocate and deallocate on dicts"

# ╔═╡ 93936612-1447-4114-b864-aba43adef4bd
md"""
The forward pass of dictionary allocation adds an entry to the dictionary (cause key error if the key already exists), the backward pass. The backward pass checks the variable in the dictionary is consistent with the asserted variable, and delete the key.

$(
twocol(md"
```julia
dict[key] ← variable
```", md"
```julia
dict[key] → variable
```
")
)
"""

# ╔═╡ b2add70c-c5d6-4e0f-a153-43e21a197181
@i function dict_assign(dict::Dict)
	var1 ← 3.14

	# copy a new variable to a dict
	dict["data"] ← var1

	# deallocate the original variables
	var1 → dict["data"]
end

# ╔═╡ 5c03d5a5-99f0-4efd-9a32-ce6d7c2b266c
dict_assign(Dict())

# ╔═╡ 2349e3ea-3053-42a4-b9d9-f97a76e4abd7
title2("Exchange two variables")

# ╔═╡ 269f18ee-3cd8-466a-a522-7c624503e31b
let
	expr = twocol(md"
```julia
var1 ↔ var2
```", md"
```julia
var1 ↔ var2
```
")
md"""One can exchange two variables is using `↔`.
$expr
"""
end

# ╔═╡ 89139719-c478-4066-9452-f9893f36d561
@i function exchange_func1(x, y)
	x ↔ y
end

# ╔═╡ d620c5ee-7d9c-4d3f-9e87-0c828dfab9ca
exchange_func1(3, 5)

# ╔═╡ 255d01b9-a873-4e63-9298-9d8f073348b0
md"One can also make a \"link\" by exchanging a variable with an empty variable such as `var::∅` and `stack[end+1]`. The forward pass push a variable to the stack and deallocate variable. The backward pass pops a variable and asserts its value."

# ╔═╡ f6cf1729-766c-4ed7-b004-c8c8ec6c7e07
let
	expr = twocol(md"
```julia
stack[end+1] ↔ var2
var1::∅ ↔ var2
```", md"
```julia
stack[end] ↔ var2::∅
var1 ↔ var2::∅
```
")
end

# ╔═╡ 3645d672-423f-4ac8-805f-0452793fee5a
@i function exchange_func2(x, y)
	anc ← 0.0
	anc += x * y
	anc ↔ z::∅  # declare `z` as an empty variable
	# after exchange, `anc` is empty and deallocated automatically.
	z -= x * y
	z → 0.0
end

# ╔═╡ c2a0024e-11dd-4ef7-8346-4374d98cafc0
exchange_func2(3, 4)

# ╔═╡ b20004e9-3c73-4dfb-8fd5-f377786fd53b
md"When exchanging with a stack top + 1, it means push and deallocate."

# ╔═╡ 5c1952b1-5016-4c87-b23c-8e6a235bf8cd
@i function stack_exchange(stack, y, x)
	stack[end+1] ↔ y  # push a variable into a stack and deallocate `y`
	y ← 1.2  # since `y` is deallocated, you can assign any value to it
	stack[end] ↔ anc::∅  # pop a variable to `anc`
	anc ↔ x    # exchange `anc` and x
	stack[end+1] ↔ anc  # push `anc` back to stack
end

# ╔═╡ 8e4470ee-01da-4547-b091-c4f65cd729b0
let
	stack = FastStack{Float64}(1000) # a preallocated stack
	stack_exchange(stack, 2.0, 3.0)  # you will get a stack of size 2
	@test length(stack) == 1 && stack.data[1] == 3.0
end

# ╔═╡ 0863bd06-cc70-4dde-b3b2-0a466805a356
md"""$(title2("Integers, floating-point numbers, fixed-point numbers and logarithmic numbers"))

A fixed-point zero with 43 fraction bits can be declared as `x ← Fixed43(0.0)` or `x ← zero(Fixed43)`, a logarithmic one can be declared as `x ← ULogarithmic(1.0)`, `x ← one(ULogarithmic{Float64})` or `x ← ULogarithmic(Fixed43(1.0))`. A complex number zero can be defined as `x ← Complex(0.0, 0.0)`.

| Number Type | +=  | *= | ⊻= | Source |
| ---- | ---- | ---- | ---- | ---- |
| boolean | - |  - | ✓ | JuliaLang |
| integer | ✓ |  × | ✓ | JuliaLang |
| floating-point number | ✓ (rounding error) |  × | - | JuliaLang |
| fixed-point number | ✓ |  × | - | [FixedPointNumbers.jl](https://github.com/JuliaMath/FixedPointNumbers.jl)
| logarithmic number | × |  ✓ | - |  [LogarithmicNumbers.jl](https://github.com/cjdoris/LogarithmicNumbers.jl)

* `✓`: an operation has its reverse when operating on a number type.
* `×`: an operation does not have a reverse when operating on a number type.
* `-`: an operation does not apply on a number type.

The `+=` operation is not regoriously reversible on floating point numbers, but we ignore the rounding errors in NiLang and use the reversibility check to detect the potential too-large rounding errors. Whether logarithmic number has rounding errors depends on its content type. If it uses floating point numbers as storage, then yes, otherwise if it uses fixed point number as the content type, then no.

One can use the `y ⊙= convert(x)` statement to convert `x` to the target type `typeof(y)` and accumulate it to `y`. Here `⊙=` can be one of `+=`, `*=` and `⊻=` that has its reverse on type `typeof(y)`.
"""

# ╔═╡ a0bae195-04e1-4642-9e14-fe4691e0906b
md"Fixed point numbers and Floating point numbers are reversible under the `+=` operation"

# ╔═╡ 20d6e8a0-2cf5-48ad-9549-60506b42b970
@i function fixed_pluseq(y1::T, x::T) where T<:Union{Fixed43, AbstractFloat}
	y1 += x
end; @test fixed_pluseq(Fixed43(0.5), Fixed43(0.6))[1] === Fixed43(1.1) && fixed_pluseq(0.5, 0.6)[1] === 1.1

# ╔═╡ 7dee5748-ed73-4e13-aa80-7a50efbc8449
@i function fixed_muleq(y1::T, x::T) where T<:Union{Fixed43, AbstractFloat}
	y1 *= x
end; @test_throws MethodError fixed_muleq(Fixed43(1.0), Fixed43(2.0))

# ╔═╡ 4c719e9b-641e-404e-9ab7-59e89135f3ba
md"Logarithmic numbers are reversible under the `*=` operation"

# ╔═╡ 77947e00-42c3-4c9e-b62a-b4b29489db43
@i function ulog_pluseq(y1::ULogarithmic{T}, x::ULogarithmic{T}) where T
	y1 += x
end; @test_throws MethodError ulog_pluseq(ULogarithmic(1.0), ULogarithmic(3.0))

# ╔═╡ 2614127d-34fb-4c3d-b678-42693f3c9341
@i function ulog_muleq(y1::ULogarithmic{T}, x::ULogarithmic{T}) where T
	y1 *= x
end; @test ulog_muleq(ULogarithmic(1.0), ULogarithmic(3.0))[1] === ULogarithmic(3.0)

# ╔═╡ 8f169235-3bd1-4cc4-a083-79736d306ad5
example("computing x ^ 3 with logarithmic numbers")

# ╔═╡ dfc9d305-5bce-4555-bfa3-d8d61fe4ca09
@i function power3(y::ULogarithmic{T}, x::T) where T
	for i=1:3
		y *= convert(x)
	end
end; @test power3(ULogarithmic(1.0), 3.0)[1] ≈ 27

# ╔═╡ f6bfa015-c101-45e8-995c-2bb6a3b7dc7d
title2("Types and arrays")

# ╔═╡ 8651d7ec-6bcd-4dbe-a062-c4bde32e5e91
md"
A Julia type can be accessed in NiLang if its default constructor is not overloaded, because NiLang requires the default constructor to \"modify\" a field of a immutable type.
"

# ╔═╡ edaa9fdb-3af8-4554-a701-0e3bff2107a5
md"It is also possbile to extract the fields directly."

# ╔═╡ 7551a880-340e-4e3f-815b-188e73f7eb9a
@i function complex_add(y::Complex{T}, x::Complex{T}) where T
    ((a, b), (c, d))::∅ ↔ ((@fields x), (@fields y))
	a += c
	b += d
    ((a, b), (c, d)) ↔ ((@fields x), (@fields y))::∅
end

# ╔═╡ 0489e51b-781f-4441-bb7f-ff3bd2e848ad
@test complex_add(1+2im, 3+4im) == (1+2im, 4+6im)

# ╔═╡ 7b0d30d6-39ff-4f6e-b13c-0ddbfcb576e5
md"Type cast is also possible"

# ╔═╡ 042297d8-6ab3-4ae6-b6e7-3b1ab2d5553b
@i function add4(a, b, c, d)
	complex_add(Complex{}(a, b), Complex{}(c, d))  # do not omit `{}`
end

# ╔═╡ 57d65a36-bfa8-4dc2-8e11-d87fa1324122
@test add4(1, 2, 3, 4) == (1, 2, 4, 6)

# ╔═╡ c21d81c3-981f-4472-ad61-d1661bfe5c4e
example("Implementing \"axpy\" function")

# ╔═╡ 99d6fe7b-d704-48f3-b115-2b3159a78068
md"`axpy` function is defined as ``\vec y += a \vec x``. One can modify the `Array` directly"

# ╔═╡ 1950ff70-54eb-4ece-a26d-a23fd0e90f5a
@i function arrayaxpy!(y!::Vector{T}, a::T, x::Vector{T}) where T
    for i=1:length(x)
		y![i] += a * x[i]
	end
end; @test arrayaxpy!(zeros(10), 2.0, collect(1.0:10.0))[1] ≈ collect(2.0:2.0:20.0)

# ╔═╡ 21458f81-9007-46f8-92e0-7a17c60beb36
md"To modify an element of a `Tuple`, we need to use a different style to avoid confusion with array"

# ╔═╡ 7813f4ce-6e98-45f3-94a8-7f5981129f2b
@i function tupleaxpy!(y!::NTuple{N,T}, a::T, x::NTuple{N,T}) where {N, T}
    for i=1:length(x)
		(y! |> tget(i)) += a * (x |> tget(i))
	end
end; @test tupleaxpy!((0,0,0), 2, (1,2,3))[1] == (2, 4, 6)

# ╔═╡ 59ec7cb7-6011-456d-9f57-a55bb8ea51a0
md"Here `data |> tget(i)` represents the `i`th field of the tuple (note it is not allowed to write `data.:i`)."

# ╔═╡ aacf63a2-9708-40db-8928-049621a7bbc4
md"## Control flows"

# ╔═╡ ad0097e7-c8ad-457a-82a9-18b998a9e9fb
md"""
#### If statement

The condition expression in `if` statement contains two parts, one is precondition and another is postcondition.

$(
twocol(md"
```julia
if (precondition[, postcondition])
	...
end
```
", md"
```julia
if (postcondition[, precondition])
	~(...)
end
```
")
)

where `...` are statements and `~(...)` are the backward execution of them.
"""

# ╔═╡ 94cd1345-3132-4882-86fe-d2429f610d1d
md"""If no postcondition is provided, it means the precondition is the same as the postcondition. It is translated to `if (cond, cond) ... end`. `elseif` is also supported."""

# ╔═╡ 4a558bd3-6e42-4c61-bd23-888b7f33ae25
@i function f7a(x)
	if x > 1
		x -= 1
	end
end; @test_throws InvertibilityError f7a(1.2)

# ╔═╡ 004f727a-e0c8-49cb-8858-dfdf4d3ac57a
@i function f7b(x, branch_keeper)
	branch_keeper ⊻= x > 1
	if (x > 1, branch_keeper)
		x += 1
	end
end; @test f7b(1.2, false)[1] == 2.2

# ╔═╡ 4c03cde9-b643-40ff-b275-f1795f88949e
title2("For statement")

# ╔═╡ fae7c74e-d25e-4c1e-ac97-199e6dae3365
md"""
The reversible `for` statement is similar to its irreversible counterpart.

$(
twocol(md"
```julia
for iter = start:step:stop
	...
end
```
", md"
```julia
for iter = stop:-step:start
	~(...)
end
```
")
)
"""

# ╔═╡ 2f2b24ea-66d0-4b3b-a460-53b6b3f28ef0
md"The iterator length should not be changed during the iteration."

# ╔═╡ e9dbb64a-27b9-443a-b917-69d55c290235
@i function f7c(x, y)
	for i=1:length(x)
		POP!(x, y[i])
	end
end; @test_throws InvertibilityError f7c([1,2,3], [0,0,0])

# ╔═╡ 0d56ce96-81a5-4102-acbc-7d88f80adcb3
md"There is an `InvertibilityError` because the length of the `x` has been changed. The inverse execution will give incorrect result."

# ╔═╡ 95c41bd1-e50a-42e8-93c3-3b754a458c13
title2("While statement")

# ╔═╡ b8629aeb-6c9a-44ed-87a1-9ab22d9485ed
md"""
The reversible `while` statement starts with the `@from` macro.

$(
twocol(md"
```julia
@from condition1 while condition2
	...
end
```
", md"
```julia
@from !(condition2) while !(condition1)
    ~(...)
end
```
")
)
"""

# ╔═╡ 3b211406-041f-4b41-acae-3958e4a37224
md"where the `condition1` in the forward pass is a condition that holds before entering the loop body, but broken at the first iteration, while `condition2` is just a normal while condition. In the backward pass, `!(condition1)` becomes the criteria to break the loop."

# ╔═╡ 62522772-cb59-4d13-acdd-d5067b223910
@i function f7d(x, i)
	@from i==0 while i<10
		i += 1
		x += 1
	end
end

# ╔═╡ 2ba68a0f-6e36-4ea2-a91d-6af43741bad1
f7d(1, 0)

# ╔═╡ 75cebaf1-38de-475f-892e-346fd2b46f6f
@test (~f7d)(11, 10) == (1, 0)

# ╔═╡ 72dcf2fe-eb48-4dee-8121-efafc87637e3
title2("Compute-copy-uncompute statement")

# ╔═╡ 84321198-93d4-4d22-8c0f-a5a10b884e1f
md"The *compute-copy-uncompute* statement is a widely used design pattern in reversible programming.
We compute the forward pass for the result, then we copy the result to the output variable, and run the backward pass to erase intermediate results.

For example, to compute `y = x * exp(k)`, we might write the following code"

# ╔═╡ 32244789-afbf-4215-97cd-15483f438eee
@i function f7e(y, x, k)
	expk ← zero(k)
	expk += exp(k)
	y += x * expk
	# uncompute the ancilla and deallocate it
	expk -= exp(k)
	expk → zero(k)
end; @test f7e(0.0, 2.0, 3.0)[1] ≈ 2.0 * exp(3.0)

# ╔═╡ 4b7f0baf-0316-4da7-9ded-50c064ddbaa3
md"It is equivalent to the following statement that generates the backward pass automatically for you."

# ╔═╡ 0e02952c-7589-4606-b006-16a9f3e52ae1
@i function f7f(y, x, k)
	# record the forward pass
	@routine begin
		expk ← zero(k)
		expk += exp(k)
	end
	y += x * expk
	# reverse execute the recorded the program
	~@routine
end; @test f7f(0.0, 2.0, 3.0)[1] ≈ 2.0 * exp(3.0)

# ╔═╡ f0904d3f-1bf1-459c-9959-b53c0f774e3f
example("Computing Fibonacci numbers")

# ╔═╡ 19bb2af5-2a67-453d-82b0-7d3059b1fa47
md"The sequence of Fibonacci numbers are: 1, 1, 2, 3, 5, 8"

# ╔═╡ 5b5858bf-63ac-4e31-a516-055a9cd18ffe
@i function rfib(out!, n::T) where T
    @routine begin
		n1 ← zero(T)
		n2 ← zero(T)
        n1 += n - 1
        n2 += n - 2
    end
    if (value(n) <= 2, ~)
        out! += 1
    else
        rfib(out!, n1)
        rfib(out!, n2)
    end
    ~@routine
end

# ╔═╡ 95060588-f24b-4eeb-9b0b-ed7159962a3c
@test rfib(0, 6)[1] == 8

# ╔═╡ c4cd9f88-9cd6-4364-b016-78f90aba6a66
title1("Extending the instruction set")
# They are translated to `y += f(args...)` is translated to `PlusEq(f)(y, args...)`, `y -= f(args...)` is translated to `MinusEq(f)(y, args...)`, `y *= f(args...)` is translated to `MulEq(f)(y, args...)`, `y /= f(args...)` is translated to `DivEq(f)(y, args...)` and `y ⊻= f(args...)` is translated to `XorEq(f)(y, args...)`.

# ╔═╡ f6049c78-7468-47ce-a4a5-84fab34d115a
title2("How to create a new elementary reversible function")

# ╔═╡ b0f73825-bbb1-448c-b491-bf634fdd398a
md"To define a pair of elementary functions that **reverse to each other**,
1. declare two functions `f` and `g` that each of them defines a mapping ``\mathbb{R}^n \rightarrow \mathbb{R}^n``
2. use `@dual f g` to tell NiLang they are reversible to each other.
3. if you want to make `f` and `g` differentiable, you can specify backward rules on these two function by defining two mappings on ``\mathbb{G}^n\rightarrow \mathbb{G}^n``, where ``\mathbb{G}`` is a 2-tuple of ``\mathbb{R}`` (or `NiLang.AD.GVar`) in NiLang. It is similar to `ForwardDiff.Dual` (check [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl)) but defined for the backward pass.

To define a **self-reversible** elementary function
1. declare a functions `f` that defines a mapping ``\mathbb{R}^n \rightarrow \mathbb{R}^n``
2. use `@selfdual f`
3. define the backward rule on `f` to make it differentiable.
"

# ╔═╡ 648cdcd6-f4f5-461f-a525-4b350cae9eb0
example("defining a new elementary function")

# ╔═╡ d6b1abd6-749d-4591-99e8-64aaa9199ab5
md"""
One can use the invertibility checker to check if the function is really reversible (under a certain tolerance `NiLangCore.GLOBAL_ATOL[]` = $(NiLangCore.GLOBAL_ATOL[])).
"""

# ╔═╡ f502b8c1-9b80-4e67-80e8-a64ddb88fb0f
@test NiLang.check_inv(M.new_forward, (3.0,))

# ╔═╡ 0bce342e-9a8e-4005-8b88-82da2d2c7163
md"""
To check of the gradients are properly defined, one can use `NiLang.AD.check_grad`
"""

# ╔═╡ fd9cf757-2698-4886-9f0a-c6c23ff0d331
@test NiLang.AD.check_grad(M.new_forward, (3.0,); iloss=1)

# ╔═╡ dda6652a-d063-4511-8041-e869bb88ca26
@test NiLang.AD.check_grad(M.new_backward, (3.0,); iloss=1)

# ╔═╡ a7d47e83-7f44-49d0-a43d-e01316fc6eba
title1("Performance Tips")

# ╔═╡ eca3efef-f35b-4623-8af8-0b830a55566d
md"The following trick still work in NiLang
* Removing boundary check with `@inbounds` (it works on FastStack),
* Add `@inline` before short functions,
* Add `@simd` before a for loop,

Other tricks like type stability are introduced in the [Julia documentation](https://docs.julialang.org/en/v1/manual/performance-tips/).
"

# ╔═╡ 45985244-adbf-4d6d-9732-a963cca62212
title2("Remove reversibility check")

# ╔═╡ 83d7e75f-7273-4c6a-bec1-a2180ebc3fb9
md"This can be done by putting an `@invcheckoff` before a code block."

# ╔═╡ 7fb05c65-f47c-430a-b588-c2f9bade40a9
example("computing the exp function by Taylor expansion")

# ╔═╡ 14c0caa1-51ea-448c-a7dc-d06e34dd0895
md"Note: this is not a clever implementation. There is an approach of defining it without allocation."

# ╔═╡ 457d07bb-e999-413e-8f29-58714670296f
@i function exp_with_reversibility_check(y::T, x::T) where T
	@routine begin
		N ← 1_000
		anc ← zeros(T, N)
		anc[1] += 1
		anc_y ← T(2.0)
		for i=2:N
			@routine begin
				temp ← zero(T)
				temp += x * anc[i-1]
			end
			anc[i] += temp / i
			anc_y += anc[i]
			~@routine
		end
	end
	y += anc_y
	~@routine
end

# ╔═╡ ac53eac0-1a59-4407-8bf6-3d8b966a9bff
@benchmark exp_with_reversibility_check(0.0, 1.0) seconds=0.3

# ╔═╡ 85c8ac7b-54f5-47dc-bd50-e78ffd6cf1cf
@i function exp_without_reversibility_check(y::T, x::T) where T
	@routine @invcheckoff begin
		N ← 1_000
		anc ← zeros(T, N)
		anc[1] += 1
		anc_y ← T(2.0)
		for i=2:N
			@routine begin
				temp ← zero(T)
				temp += x * anc[i-1]
			end
			anc[i] += temp / i
			anc_y += anc[i]
			~@routine
		end
	end
	y += anc_y
	~@routine
end

# ╔═╡ 95c55847-0591-4f7f-b9a1-aa974ccfef69
@benchmark exp_without_reversibility_check(0.0, 1.0) seconds=0.3

# ╔═╡ 91f8cfc6-e261-4945-8506-eed8caa607c2
title1("Multi-threading and CUDA")

# ╔═╡ c82b3b5c-c4e2-4bf6-b4ec-0d05ba9a669b
@i function multi_thread_exp(y::Vector, x::Vector)
	# check the size of `x` and `y`. `@assert` is not a valid statement in NiLang, so one should decorate it with `@safe` to tell the compiler, doing this is safe, do not check this statement.
	@safe @assert length(x) == length(y)
	@threads for i=1:length(y)
		y[i] += exp(x[i])
	end
end; @test multi_thread_exp(zeros(3), [1.0, 2.0, 3.0])[1] ≈ [exp(1.0), exp(2.0), exp(3.0)]

# ╔═╡ 32d75270-60d7-4326-a4ff-8674d0fbd491
md"With [CUDA](https://github.com/JuliaGPU/CUDA.jl), one can also define parallel reversile and differentiable GPU device functions. Use the broadcasting version `y += x` as an example."

# ╔═╡ 4b8834c1-8bb3-49f2-ae9e-1dbb8832d7f0
@i function addkernel(target, source)
    @invcheckoff begin
		@routine b ← (blockIdx().x-1) * blockDim().x + threadIdx().x
		if (b <= length(target), ~)
			@inbounds target[b] += source[b]
		end
		~@routine
	end
end

# ╔═╡ d8f0ae56-e643-48a1-86ee-1cd907ecb662
md"One can launch the kernel function in NiLang with `@cuda`"

# ╔═╡ e72b9dc2-dfac-4631-b114-01ec14297427
md"""
```julia
using CUDA

@i function :(+=)(identity)(target::CuArray, source::CuArray)
    @safe @assert length(target) == length(source)
    @cuda threads=256 blocks=ceil(Int,length(target)/256) addkernel(target, source)
end
```
"""

# ╔═╡ 8c93a773-edc0-4ec2-88ef-1b58b7deddc5
title2("Shared read-write in parallel computing and autodiff")

# ╔═╡ 16d08950-0575-4a4b-afc8-11ddca3198c7
md"The parallel code may suffer from the shared read issue when computing gradients.
Let's take a look at a parallel code that computes the loss ``\mathcal{L} = \sum x \vec z``."

# ╔═╡ 7c594d19-59fc-433a-bffa-c63bad46869e
@i function shared_read(loss::Real, y::Vector, x::Real, z::Vector)
	@safe @assert length(z) == length(y)
	@threads for i=1:length(y)
		y[i] += x * z[i]
	end
	for i=1:length(y)
		loss += y[i]
	end
end; @test shared_read(0.0, zeros(3), 2.0, [1.0, 2.0, 3.0])[1] ≈ 12

# ╔═╡ 345b344d-afda-4ce1-a0e7-ce6063a69206
md"However, when computing the gradients, the gradient on `x` will not be computed correctly."

# ╔═╡ c55eb045-daca-42f9-a357-095edef24644
let
	z = randn(100)
	_, gy, gx, gz = NiLang.AD.gradient(shared_read, (0.0, zeros(100), 2.0, z); iloss=1)
	@test gx ≈ sum(z)
end

# ╔═╡ c6903b65-a8b8-4aef-8c43-c822077b9d0e
md"The error is expected. Because the variable `x` is shared by multiple threads, when updating the gradient field of `x` in the backward pass, all threads will try to update the same gradient field, this is famous [race condition](https://en.wikipedia.org/wiki/Race_condition) in parallel computing."

# ╔═╡ f80353d6-0dfe-4b0a-a1af-655d344473bf
title1("Resources")

# ╔═╡ 4ca276fb-859d-4c5a-81c3-8e4b28922fa4
title2("Help and Discussion")

# ╔═╡ 3d020209-b8dd-4605-9329-78a985f0a6a3
md"""
`reversible-computing` channel of [Julia slack](https://julialang.org/slack/) and [Julia Zulip](https://julialang.zulipchat.com/register/).
"""

# ╔═╡ c7ec3496-79ea-4956-976c-b88dd22207c7
title2("Learning")

# ╔═╡ 8530a9d1-5a27-4a1d-883f-4b033a6f8fe4
md"""
1. Reversible computing book: Perumalla, Kalyan S. Introduction to reversible computing. CRC Press, 2013.
2. Our paper: Liu, Jin-Guo, and Taine Zhao. "Differentiate Everything with a Reversible Programming Language." arXiv:2003.04617 (2020).
"""

# ╔═╡ 7ce31932-0447-4445-99aa-7ebced7d0bad
TableOfContents()

# ╔═╡ Cell order:
# ╟─2061b434-0ad1-46eb-a0c7-1a5f432bfa62
# ╟─a4e76427-f051-4b29-915a-fdfce3a299bb
# ╟─c2c7b4d4-f8c9-4ebf-8da2-0103f03136e7
# ╟─12f07cc7-979c-43c3-9dc9-36ea1463c1f6
# ╟─611b577f-4722-42bf-8f8e-aeb2fb30be71
# ╟─605872cf-f3fd-462e-a2b1-7d1c5ae45efd
# ╟─fb3dee44-5fa9-4773-8b7f-a83c44358545
# ╠═d941d6c2-55bf-11eb-0002-35c7474e4050
# ╠═70088425-6779-4a2d-ba6d-b0a34c8e93a6
# ╠═af738f89-3214-429c-9c7d-18a6ea0d9401
# ╠═48d7ebc1-5def-4a57-9ec1-3fc370a4543f
# ╟─f0e94247-f615-472b-8218-3fa287b38aa1
# ╟─2581aa33-1dc5-40b1-aa9f-6a11cc750c93
# ╟─60575978-081a-4bca-a3ed-2b51cd6abc92
# ╠═f98305cb-4ba2-404a-a5c3-65510e059504
# ╟─e8cd6667-597f-458b-8465-1822e09a7891
# ╟─20145d75-004a-4c2f-b7ff-c400ca846d42
# ╟─c682a17f-600f-4034-bfe3-a851ab645c10
# ╟─5239dfe2-ea6d-4e07-a1b1-90954fe8ddc9
# ╟─e4d86a5a-e820-4a70-8a87-08bac416291b
# ╟─c307b6a4-906d-4be7-9fd7-57c942aded51
# ╟─47b502d4-e8af-4d58-9067-9700784ea435
# ╠═faecc0a7-55d7-42a1-8e9f-7e30143eef9c
# ╠═4cbe69d6-b68f-4bda-a0dd-209f9ee54f18
# ╟─7dc82f28-77bf-40da-b520-800ed1bc80c9
# ╟─9f5f9de3-9558-4c18-9d98-b77d19b570ec
# ╠═6dfcfa19-f78f-4dac-89f7-d3c5dbe17987
# ╟─dad1c6c0-d61b-4f9f-a71e-e683fe143aaa
# ╟─3e4a3916-8fe4-4262-bcd3-3014822717a3
# ╠═34906208-e6f1-4a67-860a-a7b056a86dde
# ╟─4601df35-679f-465d-9191-c18748b2fd83
# ╟─4fc72b9d-19a2-40f1-a4a8-5e97d3d5e529
# ╠═a0fde16f-8454-4f5c-a29c-a9e415c0c311
# ╟─633ff8f3-8d93-4f73-bec2-c42070e6ece9
# ╠═57f2d890-b5a0-47b7-9e3d-af4d03b10605
# ╟─34063cd0-171e-46ce-80dd-52a341fa50a1
# ╠═5d5d01db-8ff9-434c-8771-1fec6393e1fb
# ╠═10d85a50-f2f9-403e-8f6c-baef61cf702a
# ╟─b52648bf-a28a-48af-8912-31729d943ce0
# ╟─f45db10f-a836-40f3-9d8d-054ea6540e87
# ╟─1903563e-ccc2-44d9-9dbe-e5dede275b3c
# ╟─583f2585-15a3-47c6-a70e-e2f002754028
# ╠═60e6ff80-3593-4ae4-a273-914847f692db
# ╠═9e5cfd68-b58d-4d83-aae2-447e5f805c97
# ╠═dc85a942-cf52-4405-ad03-32a768e1b6e7
# ╟─5cdd346b-10a5-485c-ba78-4c0b3cb0e02f
# ╠═af9287b7-6131-46f6-beb8-6885e55e1975
# ╠═e20eeabf-1c80-431e-8cfc-4d1b79c52b5a
# ╟─90d30eea-53de-48a0-9700-ff35681fdf38
# ╠═390f58a5-6f5f-4d3a-bb16-ba04e43a07e7
# ╠═2b57443e-a516-434b-be86-80616a98e2f5
# ╟─fc2e27f9-b7ba-44cc-a953-6745548ad733
# ╠═fc744931-360b-4478-9f77-c50f048de243
# ╠═9a152b36-f377-44da-9700-ca9e05e365ff
# ╟─b6dcd18c-606f-4340-b2ec-163e8bad03f5
# ╟─a1a29f34-f8a9-4e9f-9afe-7d0096771440
# ╟─90bd6ad4-3dd8-4e7c-b445-aed1d248a2ec
# ╠═c0259e48-1973-486c-a828-1fcd3e4331c6
# ╟─8bbffa31-04a6-49ca-b36f-4d4140d75992
# ╠═a6f18c34-80ee-4b52-9ff8-f3c1b1d80f90
# ╠═a694132b-4f52-467f-8bc4-dc32fe2812db
# ╟─8c2c82f2-1240-4f2f-830e-ee8021c1a41a
# ╟─6203cf10-f8cc-4fb9-b814-7552b68c01dc
# ╠═f97a6bab-b9f9-4b95-98a9-381c51397526
# ╠═4ca48a2e-43da-457a-8e9f-6476097e4d7b
# ╟─2a2970f4-ab01-486b-89a2-6ff96f734018
# ╠═0b80d9be-53d7-4bf3-a558-659607af4709
# ╠═92362fda-bae2-4e35-bfe4-dcaea853d50b
# ╟─db9e7940-39f1-4ccf-ac70-146a521daa6e
# ╟─93936612-1447-4114-b864-aba43adef4bd
# ╠═b2add70c-c5d6-4e0f-a153-43e21a197181
# ╠═5c03d5a5-99f0-4efd-9a32-ce6d7c2b266c
# ╟─2349e3ea-3053-42a4-b9d9-f97a76e4abd7
# ╟─269f18ee-3cd8-466a-a522-7c624503e31b
# ╠═89139719-c478-4066-9452-f9893f36d561
# ╠═d620c5ee-7d9c-4d3f-9e87-0c828dfab9ca
# ╟─255d01b9-a873-4e63-9298-9d8f073348b0
# ╟─f6cf1729-766c-4ed7-b004-c8c8ec6c7e07
# ╠═3645d672-423f-4ac8-805f-0452793fee5a
# ╠═c2a0024e-11dd-4ef7-8346-4374d98cafc0
# ╟─b20004e9-3c73-4dfb-8fd5-f377786fd53b
# ╠═5c1952b1-5016-4c87-b23c-8e6a235bf8cd
# ╠═8e4470ee-01da-4547-b091-c4f65cd729b0
# ╟─0863bd06-cc70-4dde-b3b2-0a466805a356
# ╟─a0bae195-04e1-4642-9e14-fe4691e0906b
# ╠═20d6e8a0-2cf5-48ad-9549-60506b42b970
# ╠═7dee5748-ed73-4e13-aa80-7a50efbc8449
# ╟─4c719e9b-641e-404e-9ab7-59e89135f3ba
# ╠═77947e00-42c3-4c9e-b62a-b4b29489db43
# ╠═2614127d-34fb-4c3d-b678-42693f3c9341
# ╟─8f169235-3bd1-4cc4-a083-79736d306ad5
# ╠═dfc9d305-5bce-4555-bfa3-d8d61fe4ca09
# ╟─f6bfa015-c101-45e8-995c-2bb6a3b7dc7d
# ╟─8651d7ec-6bcd-4dbe-a062-c4bde32e5e91
# ╟─edaa9fdb-3af8-4554-a701-0e3bff2107a5
# ╠═7551a880-340e-4e3f-815b-188e73f7eb9a
# ╠═0489e51b-781f-4441-bb7f-ff3bd2e848ad
# ╟─7b0d30d6-39ff-4f6e-b13c-0ddbfcb576e5
# ╠═042297d8-6ab3-4ae6-b6e7-3b1ab2d5553b
# ╠═57d65a36-bfa8-4dc2-8e11-d87fa1324122
# ╟─c21d81c3-981f-4472-ad61-d1661bfe5c4e
# ╟─99d6fe7b-d704-48f3-b115-2b3159a78068
# ╠═1950ff70-54eb-4ece-a26d-a23fd0e90f5a
# ╟─21458f81-9007-46f8-92e0-7a17c60beb36
# ╠═7813f4ce-6e98-45f3-94a8-7f5981129f2b
# ╟─59ec7cb7-6011-456d-9f57-a55bb8ea51a0
# ╟─aacf63a2-9708-40db-8928-049621a7bbc4
# ╟─ad0097e7-c8ad-457a-82a9-18b998a9e9fb
# ╟─94cd1345-3132-4882-86fe-d2429f610d1d
# ╠═4a558bd3-6e42-4c61-bd23-888b7f33ae25
# ╠═004f727a-e0c8-49cb-8858-dfdf4d3ac57a
# ╟─4c03cde9-b643-40ff-b275-f1795f88949e
# ╟─fae7c74e-d25e-4c1e-ac97-199e6dae3365
# ╟─2f2b24ea-66d0-4b3b-a460-53b6b3f28ef0
# ╠═e9dbb64a-27b9-443a-b917-69d55c290235
# ╟─0d56ce96-81a5-4102-acbc-7d88f80adcb3
# ╟─95c41bd1-e50a-42e8-93c3-3b754a458c13
# ╟─b8629aeb-6c9a-44ed-87a1-9ab22d9485ed
# ╟─3b211406-041f-4b41-acae-3958e4a37224
# ╠═62522772-cb59-4d13-acdd-d5067b223910
# ╠═2ba68a0f-6e36-4ea2-a91d-6af43741bad1
# ╠═75cebaf1-38de-475f-892e-346fd2b46f6f
# ╟─72dcf2fe-eb48-4dee-8121-efafc87637e3
# ╟─84321198-93d4-4d22-8c0f-a5a10b884e1f
# ╠═32244789-afbf-4215-97cd-15483f438eee
# ╟─4b7f0baf-0316-4da7-9ded-50c064ddbaa3
# ╠═0e02952c-7589-4606-b006-16a9f3e52ae1
# ╟─f0904d3f-1bf1-459c-9959-b53c0f774e3f
# ╟─19bb2af5-2a67-453d-82b0-7d3059b1fa47
# ╠═5b5858bf-63ac-4e31-a516-055a9cd18ffe
# ╠═95060588-f24b-4eeb-9b0b-ed7159962a3c
# ╟─c4cd9f88-9cd6-4364-b016-78f90aba6a66
# ╟─f6049c78-7468-47ce-a4a5-84fab34d115a
# ╟─b0f73825-bbb1-448c-b491-bf634fdd398a
# ╟─648cdcd6-f4f5-461f-a525-4b350cae9eb0
# ╠═8c2c4fa6-172f-4dde-a279-5d0aecfdbe46
# ╟─d6b1abd6-749d-4591-99e8-64aaa9199ab5
# ╠═f502b8c1-9b80-4e67-80e8-a64ddb88fb0f
# ╟─0bce342e-9a8e-4005-8b88-82da2d2c7163
# ╠═fd9cf757-2698-4886-9f0a-c6c23ff0d331
# ╠═dda6652a-d063-4511-8041-e869bb88ca26
# ╟─a7d47e83-7f44-49d0-a43d-e01316fc6eba
# ╟─eca3efef-f35b-4623-8af8-0b830a55566d
# ╟─45985244-adbf-4d6d-9732-a963cca62212
# ╟─83d7e75f-7273-4c6a-bec1-a2180ebc3fb9
# ╟─7fb05c65-f47c-430a-b588-c2f9bade40a9
# ╟─14c0caa1-51ea-448c-a7dc-d06e34dd0895
# ╠═457d07bb-e999-413e-8f29-58714670296f
# ╠═3199a048-7b39-40f8-8183-6a54cccd91b6
# ╠═ac53eac0-1a59-4407-8bf6-3d8b966a9bff
# ╠═85c8ac7b-54f5-47dc-bd50-e78ffd6cf1cf
# ╠═95c55847-0591-4f7f-b9a1-aa974ccfef69
# ╟─91f8cfc6-e261-4945-8506-eed8caa607c2
# ╠═0e1ba158-a6bc-401c-9ba7-ed78020ad068
# ╠═c82b3b5c-c4e2-4bf6-b4ec-0d05ba9a669b
# ╟─32d75270-60d7-4326-a4ff-8674d0fbd491
# ╠═4b8834c1-8bb3-49f2-ae9e-1dbb8832d7f0
# ╟─d8f0ae56-e643-48a1-86ee-1cd907ecb662
# ╟─e72b9dc2-dfac-4631-b114-01ec14297427
# ╟─8c93a773-edc0-4ec2-88ef-1b58b7deddc5
# ╟─16d08950-0575-4a4b-afc8-11ddca3198c7
# ╠═7c594d19-59fc-433a-bffa-c63bad46869e
# ╟─345b344d-afda-4ce1-a0e7-ce6063a69206
# ╠═c55eb045-daca-42f9-a357-095edef24644
# ╟─c6903b65-a8b8-4aef-8c43-c822077b9d0e
# ╟─f80353d6-0dfe-4b0a-a1af-655d344473bf
# ╟─4ca276fb-859d-4c5a-81c3-8e4b28922fa4
# ╟─3d020209-b8dd-4605-9329-78a985f0a6a3
# ╟─c7ec3496-79ea-4956-976c-b88dd22207c7
# ╟─8530a9d1-5a27-4a1d-883f-4b033a6f8fe4
# ╟─7ce31932-0447-4445-99aa-7ebced7d0bad


text:
function mc_policy_eval!(mdp, V, action_probability::Function;
                         iterations = 100,
                         Nsamps = 320,
                         verbose = true,
                         max_steps_per_ep = 1000,
                         opt = ADAM(1e-3),
                         batchsize = 32,
                         shuffle = true,
                         rng::AbstractRNG = Random.GLOBAL_RNG,
                         a_and_p = action_and_probability,
                         logger = TBLogger("log/", tb_increment),
                         eval_fn = (args...)->0,
                         log_freq = 10,
                         exploration_policy = nothing,
                         ϵstart = 0.,
                         ϵend = 0.,
                         ϵsteps = iterations/2)
    log = !isnothing(logger)
    loss = (x, y) -> Flux.mse(V(x), y)
    best_avgR = 0
    steps_per_iter = Int(Nsamps / batchsize)

    if log
        t = 0
        isnothing(exploration_policy) && (exploration_policy = ISPolicy(V, mdp, action_probability, rng))
        S, G, avgR = sample_episodes(mdp, exploration_policy, Nsamps,
                                 max_steps = max_steps_per_ep, a_and_p = a_and_p, ϵ = ϵstart)
        with_logger(logger) do
            log_value(logger, "eval_reward", eval_fn(exploration_policy), step = t)
            log_value(logger, "avg_reward", avgR, step = t)
            log_value(logger, "loss", loss(S,G), step = t)
            log_value(logger, "eps", ϵstart, step = t)
            end
    end
    for iter in 1:iterations
        t = steps_per_iter * iter
        verbose && println("iteration: ", iter)
        # Setup the exploration policy. If none is provided, use the current policy
        isnothing(exploration_policy) && (exploration_policy = ISPolicy(V, mdp, action_probability, rng))

        # Sample episodes. S is the states, G is the per-state return (with IS weights), R is the undiscounted return for each episode
        ϵ = max(ϵend, ϵstart + (ϵend - ϵstart)*(iter / ϵsteps))
        S, G, avgR = sample_episodes(mdp, exploration_policy, Nsamps,
                                 max_steps = max_steps_per_ep, a_and_p = a_and_p, ϵ = ϵ)

        # If the returns are higher (i.e. more failures found) then save the model
        log && avgR > best_avgR && @save string(logger.logdir, "/best_model.bson") V

        # Load the data into a DataLoader and train
        data = Flux.Data.DataLoader((S, G), batchsize=batchsize, shuffle = shuffle)
        Flux.train!(loss, Flux.params(V), data, opt)

        # If we are logging then record the loss and the avg return
        (!log || iter % log_freq != 0) && continue
        with_logger(logger) do
          log_value(logger, "eval_reward", eval_fn(exploration_policy), step = t)
          log_value(logger, "avg_reward", avgR, step = t)
          log_value(logger, "loss", loss(S,G), step = t)
          log_value(logger, "eps", ϵ, step = t)
          @save string(logger.logdir, "/last_model.bson") V
        end
    end
end

ext: jl


text:
include("utils.jl")
using Manifolds:
    default_differential_backend,
    _derivative,
    _derivative!,
    differential,
    differential!,
    gradient,
    gradient!,
    _gradient,
    _gradient!,
    _jacobian,
    _jacobian!,
    set_default_differential_backend!

# differentiation
using Manifolds:
    AbstractDiffBackend,
    AbstractRiemannianDiffBackend,
    ExplicitEmbeddedBackend,
    TangentDiffBackend,
    RiemannianProjectionBackend

import Manifolds: gradient

struct TestRiemannianBackend <: AbstractRiemannianDiffBackend end
function Manifolds.gradient(::AbstractManifold, f, p, ::TestRiemannianBackend)
    return collect(1.0:length(p))
end

using FiniteDifferences
using LinearAlgebra: Diagonal, dot

@testset "Differentiation backend" begin
    fd51 = Manifolds.FiniteDifferencesBackend()
    @testset "default_differential_backend" begin
        @test default_differential_backend() isa Manifolds.FiniteDifferencesBackend

        @test length(fd51.method.grid) == 5
        # check method order
        @test typeof(fd51.method).parameters[2] == 1
        fd71 = Manifolds.FiniteDifferencesBackend(central_fdm(7, 1))
        @test set_default_differential_backend!(fd71) == fd71
        @test default_differential_backend() == fd71
    end

    @testset "gradient" begin
        set_default_differential_backend!(fd51)
        r2 = Euclidean(2)

        c1(t) = [sin(t), cos(t)]
        f1(x) = x[1] + x[2]^2
        function f1!(y, x)
            y .= x[1] + x[2]^2
            return y
        end
        f2(x) = 3 * x[1] * x[2] + x[2]^3
        @test _jacobian(c1, 0.0) ≈ [1.0; 0.0]

        @testset for backend in [fd51]
            set_default_differential_backend!(backend)
            @test _derivative(c1, 0.0) ≈ [1.0, 0.0]
            X = [-1.0, -1.0]
            @test _derivative!(c1, X, 0.0) === X
            @test isapprox(X, [1.0, 0.0])
        end
        @testset for backend in [fd51]
            set_default_differential_backend!(backend)
            X = [-1.0, -1.0]
            @test _gradient(f1, [1.0, -1.0]) ≈ [1.0, -2.0]
            @test _gradient!(f1, X, [1.0, -1.0]) === X
            @test X ≈ [1.0, -2.0]
        end
        set_default_differential_backend!(Manifolds.NoneDiffBackend())
        @testset for backend in [fd51]
            @test _derivative(c1, 0.0, backend) ≈ [1.0, 0.0]
            @test _gradient(f1, [1.0, -1.0], backend) ≈ [1.0, -2.0]
            @test _jacobian(c1, 0.0, backend) ≈ [1.0; 0.0]
            jac = [NaN; NaN]
            _jacobian!(c1, jac, 0.0, backend)
            @test jac ≈ [1.0; 0.0]
        end

        set_default_differential_backend!(fd51)
    end
end

rb_onb_default = TangentDiffBackend(
    default_differential_backend(),
    Manifolds.ExponentialRetraction(),
    Manifolds.LogarithmicInverseRetraction(),
    DefaultOrthonormalBasis(),
    DefaultOrthonormalBasis(),
)

rb_onb_fd51 = TangentDiffBackend(Manifolds.FiniteDifferencesBackend())

rb_onb_default2 = TangentDiffBackend(
    default_differential_backend();
    basis_arg=CachedBasis(
        DefaultOrthonormalBasis(),
        [[0.0, -1.0, 0.0], [sqrt(2) / 2, 0.0, -sqrt(2) / 2]],
    ),
    basis_val=CachedBasis(
        DefaultOrthonormalBasis(),
        [[0.0, -1.0, 0.0], [sqrt(2) / 2, 0.0, -sqrt(2) / 2]],
    ),
)

rb_proj = Manifolds.RiemannianProjectionBackend(default_differential_backend())

@testset "Riemannian differentials" begin
    s2 = Sphere(2)
    p = [0.0, 0.0, 1.0]
    q = [1.0, 0.0, 0.0]
    c1(t) = geodesic(s2, q, p, t)

    Xval = [-sqrt(2) / 2, 0.0, sqrt(2) / 2]
    @test isapprox(s2, c1(π / 4), differential(s2, c1, π / 4, rb_onb_default), Xval)
    X = similar(p)
    differential!(s2, c1, X, π / 4, rb_onb_default)
    @test isapprox(s2, c1(π / 4), X, Xval)

    @testset for backend in [rb_onb_fd51]
        @test isapprox(s2, c1(π / 4), differential(s2, c1, π / 4, backend), Xval)
        X = similar(p)
        differential!(s2, c1, X, π / 4, backend)
        @test isapprox(s2, c1(π / 4), X, Xval)
    end
end

@testset "Riemannian gradients" begin
    s2 = Sphere(2)
    f1(p) = p[1]

    q = [sqrt(2) / 2, 0, sqrt(2) / 2]
    X = similar(q)
    for backend in [rb_onb_default, rb_proj]
        @test isapprox(s2, q, gradient(s2, f1, q, backend), [0.5, 0.0, -0.5])
        @test gradient!(s2, f1, X, q, backend) === X
        @test isapprox(s2, q, X, [0.5, 0.0, -0.5])
    end
    X = similar(q)
    for backend in [rb_onb_default, rb_proj]
        gradient!(s2, f1, X, q, backend)
        @test isapprox(s2, q, X, [0.5, 0.0, -0.5])
    end

    # Test the gradient fallback
    @test gradient(s2, f1, q, TestRiemannianBackend()) == [1.0, 2.0, 3.0]
    X = similar(q)
    @test gradient!(s2, f1, X, q, TestRiemannianBackend()) === X
    @test X == [1.0, 2.0, 3.0]
end

@testset "Riemannian Jacobians" begin
    s2 = Sphere(2)
    f1(p) = p

    q = [sqrt(2) / 2, 0, sqrt(2) / 2]
    X = similar(q)
    @test isapprox(
        s2,
        q,
        Manifolds.jacobian(s2, s2, f1, q, rb_onb_default),
        [1.0 0.0; 0.0 1.0],
    )

    q2 = [1.0, 0.0, 0.0]
    f2(X) = [0.0 0.0 0.0; 0.0 2.0 -1.0; 0.0 -3.0 1.0] * X
    Tq2s2 = TangentSpaceAtPoint(s2, q2)
    @test isapprox(
        Manifolds.jacobian(Tq2s2, Tq2s2, f2, zero_vector(s2, q2), rb_onb_default),
        [2.0 -1.0; -3.0 1.0],
    )

    q3 = [0.0, 1.0, 0.0]
    f3(X) = [0.0 2.0 1.0; 0.0 0.0 0.0; 0.0 5.0 1.0] * X
    Tq3s2 = TangentSpaceAtPoint(s2, q3)
    @test isapprox(
        Manifolds.jacobian(Tq2s2, Tq3s2, f3, zero_vector(s2, q2), rb_onb_default),
        [-2.0 -1.0; 5.0 1.0],
    )
end

@testset "EmbeddedBackend" begin
    A = [1 0 0; 0 2 0; 0 0 3.0]
    p = 1 / sqrt(2.0) .* [1.0, 1.0, 0.0]

    cost = p -> p' * A * p # Euclidean cost of Rayleigh
    grad = (M, p) -> 2A * p
    grad! = (M, X, p) -> X .= 2A * p
    M = Euclidean(3)
    E = ExplicitEmbeddedBackend(M, gradient=grad, (gradient!)=grad!)
    S = Sphere(2)
    r_grad = (S, p) -> project(S, p, grad(M, p))
    Xt = r_grad(S, p)
    @test is_vector(S, p, Xt, true; atol=1e-14)

    R = RiemannianProjectionBackend(E)
    X = gradient(S, cost, p, R)
    @test isapprox(S, p, X, Xt)
    gradient!(S, cost, X, p, R)
    @test isapprox(S, p, X, Xt)
    # Errors with empty gradient
    Ee = ExplicitEmbeddedBackend(M)
    Re = RiemannianProjectionBackend(Ee)
    @test_throws MissingException gradient(S, cost, p, Re)
    @test_throws MissingException gradient!(S, cost, X, p, Re)
end

@testset "Default Errors for the ODEExponentialRetraction" begin
    @test_throws DomainError ODEExponentialRetraction(
        ProjectionRetraction(),
        CachedBasis(DefaultOrthogonalBasis(), []), # not yet supported
    )
    @test_throws DomainError ODEExponentialRetraction(
        ExponentialRetraction(), # not possible
        DefaultOrthogonalBasis(),
    )
    @test_throws DomainError ODEExponentialRetraction(
        ExponentialRetraction(), # not possible
        CachedBasis(DefaultOrthogonalBasis(), []), # combination not possible
    )
end


text:
# Fallbacks
@inline diffusivity_mixing_length(m, i) = mixing_length(m, i)
@inline dissipation_length(m, i) = mixing_length(m, i)

#
# A "simple" mixing length
#

Base.@kwdef struct SimpleMixingLength{T} <: AbstractParameters
    Cᴸᵇ :: T = 3.9302
end

@inline function mixing_length(m::Model{<:SimpleMixingLength}, i)
    # Two mixing lengths:

    # 1. Based on distance from surface (and limited by grid spacing)
    @inbounds ℓᶻ = - m.grid.zc[i]

    # 2. Based on stratification
    ℓᵇ = nan2inf(m.mixing_length.Cᴸᵇ * sqrt_e(m, i) / oncell(sqrt_∂B∂z, m, i))

    # Take hard minimum between the two:
    ℓ = min(ℓᶻ, ℓᵇ)

    # Limit mixing length to half cell width
    ℓ = max(ℓ, Δf(m.grid, i) / 2)

    return ℓ
end

#
# Mixing length model due to Ignacio Lopez-Gomez + Clima
#

Base.@kwdef struct EquilibriumMixingLength{T} <: AbstractParameters
    Cᴸᵟ :: T = 1.0
    Cᴸʷ :: T = 1.688 # Limits to Von-Karman constant for stress-driven turbulence
    Cᴸᵇ :: T = 1.664
end

@inline function mixing_length(m::Model{<:EquilibriumMixingLength}, i)
    @inbounds e = m.solution.e[i] # TKE at cell i

    # "Smallest" diffusivity limited by grid spacing.
    ℓᵟ = m.mixing_length.Cᴸᵟ * Δf(m.grid, i)

    if e <= 0 # shortcut when TKE is zero (diffusivity is zero in this case regardless).
        return ℓᵟ
    else
        # For notational convenience
        Cᴸᵇ = m.mixing_length.Cᴸᵇ
        Cᴸʷ = m.mixing_length.Cᴸʷ
        Cᴰ  = m.tke_equation.Cᴰ

        # Extract buoyancy frequency and shear squared.
        N² = oncell_∂B∂z(m, i)
        S² = oncell(shear_squared, m, i)

        # Length scale associated with a production-buoyancy flux-dissipation balance
        # in the TKE budget. Valid only for a linear equation of state with a shared diffusivity
        # for salinity and temperature.
        #
        # Note that ω² = Cᴷu * S² - Cᴷc * N² can be negative, in which case this model predicts ℓᵀᴷᴱ=0.
        ω² = Cᴷu(m, i) * S² - Cᴷc(m, i) * N²
        ℓᵀᴷᴱ = maxsqrt(Cᴰ * e / ω²)

        # Length-scale limitation by strong stratification.
        ℓᵇ = N² <= 0 ? Inf : Cᴸᵇ * √(e / N²)

        # Near-wall length scale:
        ℓʷ = @inbounds - Cᴸʷ * m.grid.zc[i] * u★(m) / √e

        # Hard minimum for now
        ℓ = min(ℓᵀᴷᴱ, ℓᵇ, ℓʷ)

        return max(ℓ, ℓᵟ) # limits mixing length to be larger than grid spacing
    end
end


text:
#
# This file is part of module
#   ModiaMedia (ModiaMedia/src/ModiaMedia.jl)
#

"""
    data = SimpleIdealGasMediumData(;cp_const=nothing, R_gas=nothing, MM_const=nothing,
                                     eta_const=nothing, lambda_const=nothing, T_min=nothing,
                                     T_max=nothing, T0=nothing)

Generate a `SimpleIdealGasMediumData` object containing the data
for a SimpleIdealGas medium.
"""
mutable struct SimpleIdealGasMediumData
    cp_const::Float64       # Constant specific heat capacity at constant pressure
    cv_const::Float64       # Constant specific heat capacity at constant volume
    R_gas::Float64          # Medium specific gas constant
    MM_const::Float64       # Molar mass
    eta_const::Float64      # Constant dynamic viscosity
    lambda_const::Float64   # Constant thermal conductivity
    T_min::Float64          # Minimum temperature valid for medium model
    T_max::Float64          # Maximum temperature valid for medium model
    T0::Float64             # Zero enthalpy temperature

    function SimpleIdealGasMediumData(;cp_const=nothing, R_gas=nothing, MM_const=nothing,
                                       eta_const=nothing, lambda_const=nothing, T_min=nothing, T_max=nothing, T0=nothing)
         cv_const = cp_const - R_gas
         new(cp_const, cv_const, R_gas, MM_const, eta_const, lambda_const,
             T_min, T_max, T0)
    end
end

"""
    medium = SimpleIdealGasMedium(; mediumName     = nothing,
                                    reference_p    = 101325,
                                    reference_T    = 298.15,
                                    p_default      = 101325,
                                    T_default      = 293.15,
                                    fluidConstants = nothing,
                                    data           = nothing)

Generate a `SimpleIdealGasMedium <: PureSubstance` medium object.
"""
struct SimpleIdealGasMedium <: PureSubstance
    infos::FluidInfos
    fluidConstants::SVector{1,BasicFluidConstants}
    fluidLimits::FluidLimits
    data::SimpleIdealGasMediumData

    function SimpleIdealGasMedium(; mediumName=nothing,
                                    reference_p=101325,
                                    reference_T=298.15,
                                    p_default=101325,
                                    T_default=293.15,
                                    fluidConstants=nothing,
                                    data=nothing)

        infos = FluidInfos(mediumName           = mediumName,
                           substanceNames       = [mediumName],
                           extraPropertiesNames = fill("",0),
                           ThermoStates         = IndependentVariables_pT,
                           baseProperties       = :BaseProperties_SimpleIdealGasMedium,
                           singleState          = false,
                           reducedX             = true,
                           fixedX               = false,
                           reference_p          = reference_p,
                           reference_T          = reference_T,
                           reference_X          = fill(1.0,1),
                           p_default            = p_default,
                           T_default            = T_default,
                           h_default            = specificEnthalpy(data, ThermodynamicState_pT(p_default,T_default)))

        fluidLimits = FluidLimits(TMIN = data.T_min,
                                  TMAX = data.T_max,
                                  HMIN = specificEnthalpy(data, ThermodynamicState_pT(p_default,data.T_min)),
                                  HMAX = specificEnthalpy(data, ThermodynamicState_pT(p_default,data.T_max)))

        new(infos, fill(fluidConstants,1), fluidLimits, data)
    end
end

setState_pTX(m::SimpleIdealGasMedium,p,T,X) = ThermodynamicState_pT(p,T)
setState_phX(m::SimpleIdealGasMedium,p,h,X) = ThermodynamicState_pT(p,m.data.T0+h/m.data.cp_const)
setState_psX(m::SimpleIdealGasMedium,p,s,X) = ThermodynamicState_pT(p,exp(s/m.data.cp_const + log(m.infos.reference_T)+m.infos.R_gas*log(p/m.infos.reference_p)))
setState_dTX(m::SimpleIdealGasMedium,d,T,X) = ThermodynamicState_pT(d*m.infos.R_gas*T,T)

specificEnthalpy(data::SimpleIdealGasMediumData, state::ThermodynamicState_pT)::Float64 = data.cp_const*(state.T - data.T0)
pressure(               m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = state.p
temperature(            m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = state.T
density(                m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = state.p/(m.data.R_gas*state.T)
specificEnthalpy(       m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = specificEnthalpy(m.data,state)
specificInternalEnergy( m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = m.data.cp_const*(state.T-m.data.T0)-m.data.R_gas*state.T
specificEntropy(        m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = m.data.cp_const*log(state.T/m.data.T0)-m.data.R_gas*log(state.p/m.infos.reference_p)
specificGibbsEnergy(    m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = m.data.cp_const*(state.T-m.data.T0)-state.T*specificEntropy(m,state)
specificHelmholtzEnergy(m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = specificInternalEnergy(state)-state.T*specificEntropy(state)
dynamicViscosity(       m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = m.data.eta_const
thermalConductivity(    m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = m.data.lambda_const
specificHeatCapacityCp( m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = m.data.cp_const
specificHeatCapacityCv( m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = m.data.cv_const
isentropicExponent(     m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = m.data.cp_const/m.data.cv_const
velocityOfSound(        m::SimpleIdealGasMedium, state::ThermodynamicState_pT)::Float64 = sqrt(m.data.cp_const/m.data.cv_const*m.data.R_gas*state.T)

density_pT(      m::SimpleIdealGasMedium,p,T)::Float64 = p/(m.data.R_gas*T)
density_pT_der_1(m::SimpleIdealGasMedium,p,T)::Float64 = 0.0
density_pT_der_2(m::SimpleIdealGasMedium,p,T)::Float64 = 1.0/(m.data.R_gas*T)
density_pT_der_3(m::SimpleIdealGasMedium,p,T)::Float64 = -p/(m.data.R_gas*T*T)

specificInternalEnergy_T(      m::SimpleIdealGasMedium,T)::Float64 = m.data.cp_const*(T-m.data.T0)-m.data.R_gas*T
specificInternalEnergy_T_der_1(m::SimpleIdealGasMedium,T)::Float64 = 0.0
specificInternalEnergy_T_der_2(m::SimpleIdealGasMedium,T)::Float64 = m.data.cp_const - m.data.R_gas

function standardCharacteristics(m::SimpleIdealGasMedium)::Dict{AbstractString,Any}
    p_ref = m.infos.reference_p
    T     = collect( range(m.fluidLimits.TMIN, stop=min(1600.0, m.fluidLimits.TMAX), length=501) )
    p     = [0.5e5, 1.0e5, 2.0e5]
    nT    = length(T)
    np    = length(p)
    h     = zeros(nT)
    u     = zeros(nT)
    cp    = zeros(nT)
    cv    = zeros(nT)
    d     = zeros(nT,np)

    for i in 1:nT
        state = setState_pT(m,p_ref,T[i])
        h[i]  = specificEnthalpy(      m, state)
        u[i]  = specificInternalEnergy(m, state)
        cp[i] = specificHeatCapacityCp(m, state)
        cv[i] = specificHeatCapacityCv(m, state)
    end

    for j in 1:np
        for i in 1:nT
            d[i,j] = to_DensityDisplayUnit( density(m, ThermodynamicState_pT(p[j],T[i]) ) )
        end
    end

    mediumDict = Dict{AbstractString,Any}()
    mediumDict["T"]  = uconvert.(u"°C", T*1U"K")
    mediumDict["h"]  = h*1U"J/kg"
    mediumDict["u"]  = u*1U"J/kg"
    mediumDict["cp"] = cp*1U"J/(kg*K)"
    mediumDict["cv"] = cv*1U"J/(kg*K)"
    mediumDict["d(p=0.5 bar)"] = d[:,1]*1U"g/cm^3"
    mediumDict["d(p=1.0 bar)"] = d[:,2]*1U"g/cm^3"
    mediumDict["d(p=2.0 bar)"] = d[:,3]*1U"g/cm^3"
    return mediumDict
end

function standardPlot(m::SimpleIdealGasMedium; figure=1)
    mediumDict = standardCharacteristics(m)
    ModiaMath.plot(mediumDict, [("h", "u"), ("cp","cv"), ("d(p=0.5 bar)" ,
                                                          "d(p=1.0 bar)" ,
                                                          "d(p=2.0 bar)")], xAxis="T", heading=m.infos.mediumName, figure=figure)
end


text:
export RCWAplot

"""
    plot(::DeltaRCWASolution{1})

Visualize the RCWA solution in terms of the incident and scattered fields

Keyword arguments:
- combine{Bool} :: whether or not to plot the total fields or their components
- method{Symbol} :: choose between :series and :fft (default, fallback)
- part{Function} :: A function that gives a real from a complex (default: real)
"""
RCWAplot

@recipe function f(sol::DeltaRCWASolution{1}; combine=false, part=real, method=:fft)
    # Nz = sol.modes.dims[1][1]
    # dz = sol.modes.dims[1][2] / Nz
    # z⃗ = range(0, step=dz, length=Nz)
    pw, stack, I₁, I₂, O₁, O₂ = sol.pw, sol.stack, sol.I₁, sol.I₂, sol.O₁, sol.O₂
    x⃗ = pw.x⃗[1]
    z⃗₂ = x⃗
    z⃗₁ = -reverse(z⃗₂)
    kz₁, kz₂ = _get_kz.(Ref(pw), stack.media[[1, end]])
    if method == :series
        AO₂ = zeros(length.((x⃗, z⃗₂)))
        AI₁ = zeros(length.((x⃗, z⃗₁)))
        AO₁ = zeros(length.((x⃗, z⃗₁)))
        AI₂ = zeros(length.((x⃗, z⃗₂)))
        k⃗ = pw.k⃗[1]
        for j in eachindex(x⃗)
            for i in eachindex(z⃗₂)
                AO₂[j, i] = part.(sum(exp.((z⃗₂[i] * im) .*  kz₂ .+ (x⃗[j] * im) .* k⃗) .* O₂))
                AI₂[j, i] = part.(sum(exp.((z⃗₂[i] * im) .* -kz₂ .+ (x⃗[j] * im) .* k⃗) .* I₂))
                AI₁[j, i] = part.(sum(exp.((z⃗₁[i] * im) .*  kz₁ .+ (x⃗[j] * im) .* k⃗) .* I₁))
                AO₁[j, i] = part.(sum(exp.((z⃗₁[i] * im) .* -kz₁ .+ (x⃗[j] * im) .* k⃗) .* O₁))
            end
        end
    else
        # using bfft here because it matches the summation
        # but in reality ifft (resp. bfft) are off by a factor of
        # sqrt(length(x⃗)) (resp. 1/sqrt(length(x⃗))) when it comes to preserving
        # the norm/energy of the frequency domain representation.
        # Note that fft doesn't preserve the norm (increases by sqrt(length(x⃗)))
        # Ultimately this is just a choice of convention that can be modified.
            AO₂ = part.(bfft(exp.( kz₂ * transpose(im * z⃗₂)) .* O₂, 1))
            AI₂ = part.(bfft(exp.(-kz₂ * transpose(im * z⃗₂)) .* I₂, 1))
            AI₁ = part.(bfft(exp.( kz₁ * transpose(im * z⃗₁)) .* I₁, 1))
            AO₁ = part.(bfft(exp.(-kz₁ * transpose(im * z⃗₁)) .* O₁, 1))
    end
    if combine
       @series begin
            seriestype := :heatmap
            title := "port 1 | port 2"
            xguide := "z"
            yguide := "x"
            vcat(z⃗₁, z⃗₂), x⃗, cat(AO₁ .+ AI₁, AO₂ .+ AI₂; dims=2)
        end
    else
        layout := 4
        # plot_title := String(part) * " part of $(typeof(sol.pol)) field"
        # subplots are in row-major order
        @series begin
            seriestype := :heatmap
            subplot := 1
            title := "Scattered, port 1"
            xaxis := nothing
            yguide := "x"
            z⃗₁, x⃗, AO₁
        end
        @series begin
            seriestype := :heatmap
            subplot := 2
            title := "Incident, port 2"
            xaxis := nothing
            yaxis := nothing
            z⃗₂, x⃗, AI₂
        end
        @series begin
            seriestype := :heatmap
            subplot := 3
            title := "Incident, port 1"
            xguide := "z"
            yguide := "x"
            z⃗₁, x⃗, AI₁
        end
        @series begin
            seriestype := :heatmap
            subplot := 4
            title := "Scattered, port 2"
            yaxis := nothing
            xguide := "z"
            z⃗₂, x⃗, AO₂
        end
    end
end


text:
function dance(data, n=1)
    p = collect('a':'p')
    seen = []
    for i in 1:n
        cur = String(p)
        if cur in seen
            return seen[n % (i-1) + 1]
        end
        push!(seen, cur)
        for d in data
            if d[1] == 's'
                p = circshift(p, parse(Int, d[2:end]))
            elseif d[1] == 'x'
                pos = parse.(Int, matchall(r"\d+", d)) + 1
                p[pos[1]], p[pos[2]] = p[pos[2]], p[pos[1]]
            else
                pos = indexin([d[2], d[4]], p)
                p[pos[1]], p[pos[2]] = d[4], d[2]
            end
        end
    end
    return p
end
data = readdlm("input.txt", ',')
println(String(dance(data)))
println(String(dance(data, 1000000000)))


text:
# The basic symbolic operators that are included automatically

function sym_dot_op(a,b)
    if (size(a) == size(b)) && (ndims(a) == 1)
        return [transpose(a)*b];
    else
        printerr("Usupported dimensions for: dot(a,b), sizes: a="*string(size(a))*", b="*string(size(b))*")");
    end
end

function sym_inner_op(a,b)
    if size(a) == size(b)
        c = 0;
        for i=1:length(a)
            c += a[i]*b[i];
        end
        return [c];
    else
        printerr("Unequal dimensions for: inner(a,b), sizes: a="*string(size(a))*", b="*string(size(b))*")");
    end
end

function sym_cross_op(a,b)
    if size(a) == size(b)
        if size(a) == (1,) # scalar
            # return [a[1]*b[1]]; # Should we allow this?
        elseif ndims(a) == 1 # vector
            if length(a) == 2 # 2D
                # return [a[1]*b[2] - a[2]*b[1]]; # Should we allow this?
            elseif length(a) == 3 # 3D
                return [a[2]*b[3]-a[3]*b[2], a[3]*b[1]-a[1]*b[3], a[1]*b[2]-a[2]*b[1]];
            end
        else
            printerr("Unsupported dimensions for: cross(a,b), sizes: a="*string(size(a))*", b="*string(size(b))*")");
        end
    else
        printerr("Unequal dimensions for: cross(a,b), sizes: a="*string(size(a))*", b="*string(size(b))*")");
    end
end

function sym_transpose_op(a)
    if size(a) == (1,) # scalar
        return a;
    elseif ndims(a) == 1 # vector
        return reshape(a,1,length(a));
    else
        return permutedims(a);
    end
end

#########################################################################
# Special ops for DG

# Surface integral
function sym_surface_op(ex)
    newex = copy(ex);
    SURFACEINTEGRAL = symbols("SURFACEINTEGRAL");
    if typeof(ex) <: Array
        for i=1:length(ex)
            newex[i] = sym_surface_op(newex[i]);
        end
    elseif typeof(ex) == Basic
        return SURFACEINTEGRAL*ex;
    elseif typeof(ex) <: Number
        return SURFACEINTEGRAL*ex;
    end

    return newex;
end

function sym_ave_op(var)
    #prefix = "DGAVERAGE_";
    side1 = "DGSIDE1_";
    side2 = "DGSIDE2_";
    if typeof(var) <: Array
        result = copy(var);
        for i=1:length(result)
            result[i] = sym_ave_op(var[i]);
        end
    elseif typeof(var) == Basic
        result = (symbols(side1*string(var)) + symbols(side2*string(var))) * Basic(0.5);
    elseif typeof(var) <: Number
        result = Basic(var);
    end
    return result;
end

function sym_jump_op(var)
    #prefix = "DGJUMP_";
    side1 = "DGSIDE1_";
    side2 = "DGSIDE2_";
    norm1 = sym_normal_op(1);
    norm2 = sym_normal_op(2);
    if typeof(var) <: Array
        # make into a vector according to normal
        if length(var) == 1
            result = sym_jump_op(var[1]);
        else
            # matrix rows for vec components, cols for normal components
            result = sym_jump_op(var[1]);
            for i=2:length(var)
                result = hcat(result, sym_jump_op(var[i]));
            end
            permutedims(result)
        end

    elseif typeof(var) == Basic
        # Note: norm is a vector
        result = symbols(side2*string(var)) .* norm2 .+ symbols(side1*string(var)) .*norm1;
    elseif typeof(var) <: Number
        result = Basic(0) .* norm1;
    end

    return result;
end

# function sym_ave_normdotgrad_op(var)
#     prefix = "DGAVENORMDOTGRAD_";
#     if typeof(var) <: Array
#         result = copy(var);
#         for i=1:length(result)
#             result[i] = sym_ave_normdotgrad_op(var[i]);
#         end
#     elseif typeof(var) == Basic
#         result = symbols(prefix*string(var));
#     elseif typeof(var) <: Number
#         result = Basic(0);
#     end
#     return result;
# end

# function sym_jump_normdotgrad_op(var)
#     prefix = "DGJUMPNORMDOTGRAD_";
#     if typeof(var) <: Array
#         result = copy(var);
#         for i=1:length(result)
#             result[i] = sym_jump_normdotgrad_op(var[i]);
#         end
#     elseif typeof(var) == Basic
#         result = symbols(prefix*string(var));
#     elseif typeof(var) <: Number
#         result = Basic(0);
#     end
#     return result;
# end

function sym_normal_op()
    _FACENORMAL_1 = symbols("_FACENORMAL1_1");
    _FACENORMAL_2 = symbols("_FACENORMAL1_2");
    _FACENORMAL_3 = symbols("_FACENORMAL1_3");
    d = config.dimension;
    if d==1
        return [_FACENORMAL_1];
    elseif d==2
        return [_FACENORMAL_1, _FACENORMAL_2];
    elseif d==3
        return [_FACENORMAL_1, _FACENORMAL_2, _FACENORMAL_3];
    end
end

function sym_normal_op(side)
    _FACENORMAL_1 = symbols("_FACENORMAL"*string(side)*"_1");
    _FACENORMAL_2 = symbols("_FACENORMAL"*string(side)*"_2");
    _FACENORMAL_3 = symbols("_FACENORMAL"*string(side)*"_3");
    d = config.dimension;
    if d==1
        return [_FACENORMAL_1];
    elseif d==2
        return [_FACENORMAL_1, _FACENORMAL_2];
    elseif d==3
        return [_FACENORMAL_1, _FACENORMAL_2, _FACENORMAL_3];
    end
end

#########################################################################
# derivative ops
#########################################################################

# Multiplies the expression by TIMEDERIV
# Dt(a*b+c)  ->  TIMEDERIV*(a*b+c)
function sym_Dt_op(ex)
    newex = copy(ex);
    TIMEDERIV = symbols("TIMEDERIV");
    if typeof(ex) <: Array
        for i=1:length(ex)
            newex[i] = sym_Dt_op(newex[i]);
        end
    elseif typeof(ex) == Basic
        return TIMEDERIV*ex;
    elseif typeof(ex) <: Number
        newex = 0;
    end

    return newex;
end

# Applies a derivative prefix. wrt is the axis index
# sym_deriv_string(u_12, 1) -> D1_u_12
function sym_deriv_string(var, wrt)
    # since numbers may have been converted to Basic[]
    if typeof(wrt) <: Number
        swrt = string(wrt);
    elseif typeof(wrt) <: Array
        swrt = string(wrt[1]);
    end

    prefix = "D"*swrt*"_";

    return symbols(prefix*string(var));
end

function sym_deriv_op(u, wrt)
    if typeof(u) <: Array
        result = copy(u);
        for i=1:length(result)
            result[i] = sym_deriv_op(u[i], wrt);
        end
    elseif typeof(u) == Basic
        result = sym_deriv_string(u, wrt);
    elseif typeof(u) <: Number
        result = 0;
    end

    return result;
end

function sym_grad_op(u)
    result = Array{Basic,1}(undef,0);
    if typeof(u) <: Array
        d = config.dimension;
        rank = 0;
        if ndims(u) == 1 && length(u) > 1
            rank = 1;
        elseif ndims(u) == 2
            rank = 2;
        end

        if rank == 0
            # result is a vector
            for i=1:d
                push!(result, sym_deriv_string(u[1], i));
            end
        elseif rank == 1
            # result is a tensor
            for i=1:d
                for j=1:d
                    push!(result, sym_deriv_string(u[i], j));
                end
            end
            result = reshape(result, d,d);
        elseif rank == 2
            # not yet ready
            # printerr("unsupported operator, grad(tensor)");
            # return nothing;
            for i=1:d
                for j=1:d
                    for k=1:d
                        push!(result, sym_deriv_string(u[i+(j-1)*d], k));
                    end
                end
            end
            result = reshape(result, d,d,d);
        end
    elseif typeof(u) == Basic
        # result is a vector
        d = config.dimension;
        for i=1:d
            push!(result, sym_deriv_string(u, i));
        end
    elseif typeof(u) <: Number
        return zeros(config.dimension);
    end

    return result;
end

function sym_div_op(u)
    result = Array{Basic,1}(undef,0);
    if typeof(u) <: Array
        d = config.dimension;
        rank = 0;
        if ndims(u) == 1 && length(u) > 1
            rank = 1;
        elseif ndims(u) == 2
            rank = 2;
        end

        if rank == 0
            # Not allowed
            printerr("unsupported operator, div(scalar)");
            return nothing;
        elseif rank == 1
            # result is a scalar
            if d==1
                result = [sym_deriv_string(u[1], 1)];
            else
                ex = :(a+b);
                ex.args = [:+];
                for i=1:d
                    push!(ex.args, sym_deriv_string(u[i], i))
                end
                result = [Basic(ex)];
            end
        elseif rank == 2
            # not yet ready
            printerr("unsupported operator, div(tensor)");
            return nothing;
        end
    elseif typeof(u) <: Number
        # Not allowed
        printerr("unsupported operator, div(number)");
        return nothing;
    end

    return result;
end

function sym_curl_op(u)
    result = Array{Basic,1}(undef,0);
    if typeof(u) <: Array
        d = config.dimension;
        rank = 0;
        if ndims(u) == 1 && sz[1] > 1
            rank = 1;
        elseif ndims(u) == 2
            rank = 2;
        end

        if rank == 0
            # Not allowed
            printerr("unsupported operator, curl(scalar)");
            return nothing;
        elseif rank == 1
            # result is a vector
            if d==1
                result = [sym_deriv_string(u[1], 1)];
            else
                #TODO
                printerr("curl not ready");
                return nothing;
            end
        elseif rank == 2
            # not yet ready
            printerr("unsupported operator, curl(tensor)");
            return nothing;
        end
    elseif typeof(u) <: Number
        # Not allowed
        printerr("unsupported operator, curl(number)");
        return nothing;
    end

    return result;
end

function sym_laplacian_op(u)
    # simply use the above ops
    return sym_div_op(sym_grad_op(u));
end

#################################################################################
# The finite volume operators.
function sym_left_op(f)
    # Input will be an array of Basic. Output should be in a similar array.
    if typeof(f) <: Array
        result = copy(f);
        for i=1:length(result)
            result[i] = sym_left_op(f[i]);
        end

    elseif typeof(f) == Basic
        # Apply a CELL1 tag to signal to the code generator which side of the face the value is from.
        result = apply_flag_to_all_symbols("CELL1", f);

    elseif typeof(f) <: Number
        # If the input was just a constant, the result will just be that constant.
        result = Basic(f);
    end
    return result;
end

function sym_right_op(f)
    # Input will be an array of Basic. Output should be in a similar array.
    if typeof(f) <: Array
        result = copy(f);
        for i=1:length(result)
            result[i] = sym_right_op(f[i]);
        end

    elseif typeof(f) == Basic
        # Apply a CELL2 tag to signal to the code generator which side of the face the value is from.
        result = apply_flag_to_all_symbols("CELL2", f);

    elseif typeof(f) <: Number
        # If the input was just a constant, the result will just be that constant.
        result = Basic(f);
    end
    return result;
end

function sym_central_op(f)
    # Input will be an array of Basic. Output should be in a similar array.
    if typeof(f) <: Array
        result = copy(f);
        for i=1:length(result)
            result[i] = sym_right_op(f[i]);
        end

    elseif typeof(f) == Basic
        # Apply a CELL2 tag to signal to the code generator which side of the face the value is from.
        result = apply_flag_to_all_symbols("CENTRAL", f);

    elseif typeof(f) <: Number
        # If the input was just a constant, the result will just be that constant.
        result = Basic(f);
    end
    return result;
end

function sym_neighborhood_op(f)
    # Input will be an array of Basic. Output should be in a similar array.
    if typeof(f) <: Array
        result = copy(f);
        for i=1:length(result)
            result[i] = sym_neighborhood_op(f[i]);
        end

    elseif typeof(f) == Basic
        # Apply a CELL2 tag to signal to the code generator which side of the face the value is from.
        result = apply_flag_to_all_symbols("NEIGHBORHOOD", f);

    elseif typeof(f) <: Number
        # If the input was just a constant, the result will just be that constant.
        result = Basic(f);
    end
    return result;
end

function sym_upwind_op(v, f)
    # Input will be an array of Basic. Output should be in a similar array.
    if typeof(f) <: Array
        result = copy(f);
        for i=1:length(result)
            result[i] = sym_upwind_op(v, f[i]);
        end

    elseif typeof(f) == Basic
        # Apply a CELLn tag to signal to the code generator which side of the face the value is from.
        side1 = apply_flag_to_all_symbols("CELL1", f);
        side2 = apply_flag_to_all_symbols("CELL2", f);
        # The formula for an adjustable first order upwind scheme.
        # F(u) = 0.5*(side1+side2)*(v.normal) + 0.5*(side1-side2)*abs(v.normal)*(1-alpha)
        result = conditional(isgreaterthan(sym_dot_op(v, sym_normal_op())[1], Basic(0)),
                    (sym_dot_op(v, sym_normal_op()) .* side1)[1],
                    # else
                    (sym_dot_op(v, sym_normal_op()) .* side2)[1]) # Note that the arguments to conditional and isgreaterthan are not in arrays.

    elseif typeof(f) <: Number
        # If the input was just a constant, the result will just be that constant.
        result = Basic(f);
    end
    return result;
end

function sym_upwindA_op(v, f, alpha=Basic(0))
    # Input will be an array of Basic. Output should be in a similar array.
    if typeof(f) <: Array
        result = copy(f);
        for i=1:length(result)
            result[i] = sym_upwindA_op(v, f[i], alpha);
        end

    elseif typeof(f) == Basic
        # Apply a CELLn tag to signal to the code generator which side of the face the value is from.
        side1 = apply_flag_to_all_symbols("CELL1", f);
        side2 = apply_flag_to_all_symbols("CELL2", f);
        # The formula for an adjustable first order upwind scheme.
        # F(u) = 0.5*(side1+side2)*(v.normal) + 0.5*(side1-side2)*abs(v.normal)*(1-alpha)
        result = Basic(0.5) .* ( (sym_dot_op(v, sym_normal_op())) .* (side1 + side2) .+ abs.(sym_dot_op(v, sym_normal_op())[1]) .* (Basic(1.0)-alpha) .* (side1 - side2) );
        result = result[1]; # Since dot and normal returned arrays, result was placed in an array, but this is done above redundantly.

    elseif typeof(f) <: Number
        # If the input was just a constant, the result will just be that constant.
        result = Basic(f);
    end
    return result;
end

function sym_burgerGodunov_op(u, f)
    if typeof(f) <: Array
        result = copy(f);
        for i=1:length(result)
            result[i] = sym_burgerGodunov_op(u[i], f[i]);
        end
    elseif typeof(f) == Basic
        uside1 = apply_flag_to_all_symbols("CELL1", u);
        uside2 = apply_flag_to_all_symbols("CELL2", u);
        fside1 = apply_flag_to_all_symbols("CELL1", f);
        fside2 = apply_flag_to_all_symbols("CELL2", f);
        result = conditional(isgreaterthan(sym_normal_op()[1], Basic(0)), # if normal[1] > 0
                    conditional(isgreaterthan(uside1, uside2), # if u1 < u2
                        symbolmax(fside1, fside2), # max(f1,f2)
                        # else
                        conditional(isgreaterthan(uside1*uside2, Basic(0)), # if u1*u2 > 0 (same sign)
                            symbolmin(fside1, fside2), # min(f1,f2)
                            # else
                            Basic(0))), # 0
                    # elseif normal < 0 (just reverse 1 and 2)
                    conditional(isgreaterthan(uside2, uside1), # if u2 < u1
                        symbolmax(fside2, fside1), # max(f1,f2)
                        # else
                        conditional(isgreaterthan(uside2*uside1, Basic(0)), # if u1*u2 > 0 (same sign)
                            symbolmin(fside2, fside1), # min(f1,f2)
                            # else
                            Basic(0)))) # 0
    elseif typeof(f) <: Number
        result = Basic(f);
    end
    return result;
end

# Load them into the global arrays
op_names = [:dot, :inner, :cross, :transpose, :surface, :ave, :jump, :normal,
            :Dt, :deriv, :grad, :div, :curl, :laplacian,
            :left, :right, :central, :neighborhood, :upwind, :upwindA, :burgerGodunov];
_handles = [sym_dot_op, sym_inner_op, sym_cross_op, sym_transpose_op, sym_surface_op, sym_ave_op, sym_jump_op,
            sym_normal_op, sym_Dt_op, sym_deriv_op, sym_grad_op,
            sym_div_op, sym_curl_op, sym_laplacian_op,
            sym_left_op, sym_right_op, sym_central_op, sym_neighborhood_op, sym_upwind_op, sym_upwindA_op, sym_burgerGodunov_op];
for i=1:length(op_names)
    push!(ops, SymOperator(op_names[i], _handles[i]));
end

max_stars_repo_name: paralab/Finch

max_issues_repo_name: paralab/Finch

max_forks_repo_name: paralab/Finch


text:
using TransformUtils
using SensorFeatureTracking
using CoordinateTransformations
using Images
using FixedPointNumbers
using Colors
using StaticArrays

cd(dirname(@__FILE__))

function warpKeypoint(keypoints::Keypoints, At::AffineMap)
	# convert CartesianIndex to SVector, applie affine map At to it and round and convert result back to CartesianIndex (keypoints)
	map(kp -> (kp[1] > 0 < kp[2])? CartesianIndex(round.(Int32,[At(SVector(kp[1],kp[2]))...])...) : CartesianIndex(0,0), keypoints )

# CartesianIndex(round.(Int32,b[1])...)
end

function invwarpKeypoint(keypoints::Keypoints, At::AffineMap)
	At = inv(At)
	# convert CartesianIndex to SVector, applie affine map At to it and round and convert result back to CartesianIndex (keypoints)
	map(kp -> (kp[1] > 0 < kp[2])? CartesianIndex(round.(Int32,[At(SVector(kp[1],kp[2]))...])...) : CartesianIndex(0,0), keypoints )

# CartesianIndex(round.(Int32,b[1])...)
end

function padandcutoffsetImg(img, ro, co)
	blankImg = zeros(Gray{N0f8},ro,co)

	ro_start =  max(indices(img)[1].start, 1)
	ro_stop  = 	min(indices(img)[1].stop, ro)

	co_start =  max(indices(img)[2].start, 1)
	co_stop  = 	min(indices(img)[2].stop, co)

	return blankImg[ro_start:ro_stop,co_start:co_stop] = img[ro_start:ro_stop,co_start:co_stop]

end

## Extract example data
# Extract example imu data
imudata = Vector{IMU_DATA}()

imu_data,imu_fields = readcsv("../Data/testSequence_IMU.csv",header=true)

for i = 1:size(imu_data,1)
	push!(imudata,IMU_DATA(Int64(imu_data[i,1]),imu_data[i,9:11],imu_data[i,3:5]))
end

# read image log file and get timestaps
timestamps = readcsv("../Data/testSequence_CAM.csv", Int64, skipstart=4)[:,2]
N = size(timestamps,1)

# constants setup
index = PInt64(1)

fx = 524.040
	fy = 524.040
	cy = 319.254
	cx = 251.227
	cam = CameraModel(640,480,[fx, fy],[cx, cy], 0., [0])
# this was calculated using HornAbsoluteOrientation, see hornEstimateTransformCam_IMU example
# iQc = Quaternion()
# 	iQc.s = 0.620782
# 	iQc.v = [-0.26786, 0.454148, 0.580198]
# 	iRc = convert(SO3,iQc).R
# 	cRi = iRc'

iQc = Quaternion()
	iQc.s = 0.6475888324459562
	iQc.v = [-0.261166, 0.533419, 0.477373]
	iRc = convert(SO3,iQc).R
	cRi = iRc'

# calculate delta affinities

predAffineMaps = Vector{AffineMap}(N)

for i = 1:N#length(timestamps)
	ctime = timestamps[i]
	(valid, H) = predictHomographyIMU!(index,ctime,imudata,cam.K,cam.Ki, cRi = cRi)
	At = predictAffinity(H)

	#add all predicted affinities to matrix
	predAffineMaps[i] = At
end

# spot check if images line up
i = 65
At = predAffineMaps[i]
im1 = load("../Data/testSequence/image_$(i-1).jpg")
im2 = load("../Data/testSequence/image_$(i).jpg")
range = 10:470,10:630#200:250,330:380
@show At
colorview(RGB, paddedviews(0,im2,warpedview(im1,At),zeroarray)...)

## Feature tracking compensated with the predicted affinity maps
imu_only = false
startFrame = 50
numFrames = 50
# joinpath(datadir,
im1 = load("../Data/testSequence/image_$(startFrame).jpg")
im2 = load("../Data/testSequence/image_$(startFrame+1).jpg")

(ro,co) = size(im1)
feats = getApproxBestShiTomasi(im1,nfeatures=200, stepguess=0.95)

flow = BlockTracker(feats, search_size = 5, matchFunction = compute_ssd)
# grid_features!(flow,co,ro)

image1 = deepcopy(im1)
oldfeats = deepcopy(flow.features)

# warp keypoints and image
flow.features[:] = invwarpKeypoint(flow.features, predAffineMaps[startFrame+1])
im1_w = warp(im1, predAffineMaps[startFrame+1])

!imu_only && @time block_tracker!(flow, padandcutoffsetImg(im1_w, ro, co), im2)

blankImg = zeros(Gray{N0f8},ro,co)
map((ft1,ft2)->drawfeatureLine!(blankImg, Feature(ft1),Feature(ft2)),oldfeats,flow.features)

colorview(RGB, paddedviews(0,im2,(im1_w),blankImg)...)

#
tic()
for frame_idx = startFrame+2:startFrame+numFrames -1
    println("Processing Frame: ", frame_idx)
    im1 = deepcopy(im2)
    im2 = load("../Data/testSequence/image_$(frame_idx).jpg")

	oldfeats = deepcopy(flow.features)
	# warp keypoints and image
	flow.features[:] = invwarpKeypoint(flow.features, predAffineMaps[frame_idx])
	im1_w = warp(im1, predAffineMaps[frame_idx])

	!imu_only && block_tracker!(flow, padandcutoffsetImg(im1_w, ro, co), im2)

    map((ft1,ft2)->drawfeatureLine!(blankImg, Feature(ft1),Feature(ft2)),flow.features, oldfeats)

end
toc()

# overlay tracking traces with image
# o_img = (reinterpret(UInt8,blankImg)) .| (reinterpret(UInt8,im2))
# b_img = reinterpret(Gray{N0f8},o_img)
o_img = (reinterpret(UInt8,blankImg)) .| (reinterpret(UInt8,im2))
r_img = reinterpret(Gray{N0f8},o_img)
n_img = .~(reinterpret(UInt8,blankImg)) .& (reinterpret(UInt8,im2))
n_img = reinterpret(Gray{N0f8},n_img)
#
# colorview(RGB, n_img, n_img, b_img)
# colorview(RGB, n_img, g_img, n_img)
colorview(RGB, r_img, n_img, n_img)

##

max_stars_repo_name: dehann/SensorFeatureTracking

max_issues_repo_name: dehann/SensorFeatureTracking

max_forks_repo_name: dehann/SensorFeatureTracking


text:
Pkg.add("Gadfly")
using Gadfly

include("Jlo.jl")

x = linspace(-1,1,100)
g = Jlo.gaussian(x,0,1)

f0 = 1
y = sin(2*pi*f0*x)
plot(x=x,y=y)

plot(x = fs(x), y =abs(ft(y)))

max_stars_repo_name: awojdyla/jlo

max_issues_repo_name: awojdyla/jlo

max_forks_repo_name: awojdyla/jlo


text:
#Imports plot recipes (so that they can be expanded).
import RecipesBase.plot, RecipesBase.plot!

#Plots for single dimensional bifurcation diagrams.
function plot(bif_diagram::bif_diagram,var::Int64=1)
    plot(); return plot!(bif_diagram,var);
end
function plot!(bif_diagram::bif_diagram,var::Int64=1)
    for i = 1:length(bif_diagram.param_values), j = 1:length(bif_diagram.fixed_points[i])
        scatter!((bif_diagram.param_values[i],bif_diagram.fixed_points[i][j][var]),color=bif_diagram.colors[i][j],label="")
    end
    return plot!()
end

#Plots for two dimensional bifurcation diagrams.
function plot(bif_diagram_2d::bif_diagram_2d,var::Int64=1)
    plot(); return plot!(bif_diagram_2d,var);
end
function plot!(bif_diagram_2d::bif_diagram_2d,var::Int64=1)
    for i = 1:length(bif_diagram_2d.param_values1), j = 1:length(bif_diagram_2d.param_values2), k = 1:length(bif_diagram_2d.fixed_points[i,j])
        scatter!((bif_diagram_2d.param_values1[i],bif_diagram_2d.param_values2[j],bif_diagram_2d.fixed_points[i,j][k][var]),color=bif_diagram_2d.colors[i,j][k],label="")
    end
    return plot!()
end

#These bifurcation diagrams takes a variable input as a symbol, and plots that variable in the bifurcation diagram.
function plot(bif_diagram::bif_diagram,var::Symbol)
    !in(var,bif_diagram.reactant_syms)&&error("The reactant ",var," do not seem to be a part of the original reaction network.")
    plot(bif_diagram,findfirst(bif_diagram.reactant_syms.==var))
end
function plot!(bif_diagram::bif_diagram,var::Symbol)
    !in(var,bif_diagram.reactant_syms)&&error("The reactant ",var," do not seem to be a part of the original reaction network.")
    plot!(bif_diagram,findfirst(bif_diagram.reactant_syms.==var))
end
function plot(bif_diagram_2d::bif_diagram_2d,var::Symbol)
    !in(var,bif_diagram_2d.reactant_syms)&&error("The reactant ",var," do not seem to be a part of the original reaction network.")
    plot(bif_diagram_2d,findfirst(bif_diagram_2d.reactant_syms.==var))
end
function plot!(bif_diagram_2d::bif_diagram_2d,var::Symbol)
    !in(var,bif_diagram_2d.reactant_syms)&&error("The reactant ",var," do not seem to be a part of the original reaction network.")
    plot!(bif_diagram_2d,findfirst(bif_diagram_2d.reactant_syms.==var))
end

#Creates plot_bifurcation_diagram function, which both calculates and plots a bifurcation diagram.
plot_bifurcation_diagram(var::Union{Int64,Symbol},args...) = plot(bifurcation_diagram(args...),var)
plot_bifurcation_diagram!(var::Union{Int64,Symbol},args...) = plot!(bifurcation_diagram(args...),var)
plot_bifurcation_diagram_2d(var::Union{Int64,Symbol},args...) = plot(bifurcation_diagram_2d(args...),var)
plot_bifurcation_diagram_2d!(var::Union{Int64,Symbol},args...) = plot!(bifurcation_diagram_2d(args...),var)

max_stars_repo_name: TorkelE/BRNEquilibrate

max_issues_repo_name: TorkelE/BRNEquilibrate

max_forks_repo_name: TorkelE/BRNEquilibrate


text:
using MemoryHunter
using Test

@testset "MemoryHunter.jl" begin
    # Write your tests here.
end


text:
using Random
using HDF5

import Base: iterate, length, IteratorEltype

export dataset

export D4RLDataSet

"""
Represents an `Iterable` dataset with the following fields:

# Fields
- `dataset::Dict{Symbol, Any}`: representation of the dataset as a Dictionary with style as `style`.
- `repo::String`: the repository from which the dataset is taken.
- `dataset_size::Int`, the number of samples in the dataset.
- `batch_size::Int`: the size of the batches returned by `iterate`.
- `style::Tuple{Symbol}`: the style of the `Iterator` that is returned, check out: [`SARTS`](@ref), [`SART`](@ref) and [`SA`](@ref) for types supported out of the box.
- `rng<:AbstractRNG`.
- `meta::Dict`: the metadata provided along with the dataset.
- `is_shuffle::Bool`: determines if the batches returned by `iterate` are shuffled.
"""
struct D4RLDataSet{T<:AbstractRNG} <: RLDataSet
    dataset::Dict{Symbol, Any}
    repo::String
    dataset_size::Integer
    batch_size::Integer
    style::Tuple
    rng::T
    meta::Dict
    is_shuffle::Bool
end

# TO-DO: include other functionality like is_sequential, will be implemented soon
# TO-DO: enable the users providing their own paths to datasets if they already have it
# TO-DO: add additional env arg to do complete verify function
"""
    dataset(dataset; <keyword arguments>)

Create a dataset enclosed in a [`D4RLDataSet`](@ref) `Iterable` type. Contain other related metadata
for the `dataset` that is passed. The returned type is an infinite or a finite `Iterator`
respectively depending upon whether `is_shuffle` is `true` or `false`. For more information regarding
the dataset, refer to [D4RL](https://github.com/rail-berkeley/d4rl). Check out d4rl_pybullet_dataset_params() or d4rl_dataset_params().

# Arguments

- `dataset::String`: name of the datset.
- `repo::String="d4rl"`: name of the repository of the dataset. can be "d4rl" or "d4rl-pybullet".
- `style::Tuple{Symbol}=SARTS`: the style of the `Iterator` that is returned. can be [`SARTS`](@ref), [`SART`](@ref) or [`SA`](@ref).
- `rng<:AbstractRNG=StableRNG(123)`.
- `is_shuffle::Bool=true`: determines if the dataset is shuffled or not.
- `batch_size::Int=256`: batch_size that is yielded by the iterator.

!!! note

[`FLOW`](https://flow-project.github.io/) and [`CARLA`](https://github.com/rail-berkeley/d4rl/wiki/CARLA-Setup) supported by [D4RL](https://github.com/rail-berkeley/d4rl) have not
been tested in this package yet.
"""
function dataset(
    dataset::String;
    repo::String="d4rl",
    style::NTuple=SARTS,
    rng::AbstractRNG=MersenneTwister(123),
    is_shuffle::Bool=true,
    batch_size::Int=256
)

    try
        @datadep_str repo*"-"*dataset
    catch e
        if isa(e, KeyError)
            throw("Invalid params, check out d4rl_pybullet_dataset_params() or d4rl_dataset_params()")
        end
    end

    path = @datadep_str repo*"-"*dataset

    @assert length(readdir(path)) == 1
    file_name = readdir(path)[1]

    data = h5open(path*"/"*file_name, "r") do file
        read(file)
    end

    # sanity checks on data
    d4rl_verify(data)

    dataset = Dict{Symbol, Any}()
    meta = Dict{String, Any}()

    N_samples = size(data["observations"])[2]

    for (key, d_key) in zip(["observations", "actions", "rewards", "terminals"], Symbol.(["state", "action", "reward", "terminal"]))
            dataset[d_key] = data[key]
    end

    for key in keys(data)
        if !(key in ["observations", "actions", "rewards", "terminals"])
            meta[key] = data[key]
        end
    end

    return D4RLDataSet(dataset, repo, N_samples, batch_size, style, rng, meta, is_shuffle)

end

function iterate(ds::D4RLDataSet, state = 0)
    rng = ds.rng
    batch_size = ds.batch_size
    size = ds.dataset_size
    is_shuffle = ds.is_shuffle
    style = ds.style

    if is_shuffle
        inds = rand(rng, 1:size, batch_size)
        map((x)-> if x <= size x else 1 end, inds)
    else
        if (state+1) * batch_size <= size
            inds = state*batch_size+1:(state+1)*batch_size
        else
            return nothing
        end
        state += 1
    end

    batch = (state = copy(ds.dataset[:state][:, inds]),
    action = copy(ds.dataset[:action][:, inds]),
    reward = copy(ds.dataset[:reward][inds]),
    terminal = copy(ds.dataset[:terminal][inds]))

    if style == SARTS
        batch = merge(batch, (next_state = copy(ds.dataset[:state][:, (1).+(inds)]),))
    end

    return batch, state
end

take(ds::D4RLDataSet, n::Integer) = take(ds.dataset, n)
length(ds::D4RLDataSet) = ds.dataset_size
IteratorEltype(::Type{D4RLDataSet}) = EltypeUnknown() # see if eltype can be known (not sure about carla and adroit)

function d4rl_verify(data::Dict{String, Any})
    for key in ["observations", "actions", "rewards", "terminals"]
        @assert (key in keys(data)) "Expected keys not present in data"
    end
    N_samples = size(data["observations"])[2]
    @assert size(data["rewards"]) == (N_samples,) || size(data["rewards"]) == (1, N_samples)
    @assert size(data["terminals"]) == (N_samples,) || size(data["terminals"]) == (1, N_samples)
end