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

using Markdown
using InteractiveUtils

# ╔═╡ f1d0ed8f-bfd3-4210-84fd-f189bba6ebf9
begin
	using PlutoUI
	PlutoUI.TableOfContents(title = "Contents")
end

# ╔═╡ 08eadc12-f8e9-11eb-1094-b7adc6936aa1
md"""
# Permutations
"""

# ╔═╡ c4d2ff29-2579-4079-a6cb-29015ef5ecd2
md"""
## Objectives
- _**Create a custom type**_ to represent a permutation
- _**Utilise multiple dispatch to create new methods for inbuilt functions**_ with `Permutation` objects as arguments
- Generate the elements of groups from generators in `Permutation` form
"""

# ╔═╡ db919e43-6dc8-4e53-8081-d9312446117a
md"""
## What is multiple dispatch?
Multiple dispatch is a feature central to the design of Julia with incredibly powerful consequences. In a most basic sense, it allows for the same function name to be used for multiple different operations, such as `-` being used for unary minus (``x \to -x``) as well as for subtraction. More generally, if the types of the inputs are specified differently for different declarations of the function, the function can act differently depending on the inputs that it gets. An example of this inbuilt into Julia is `a * b`, which multiplies `a` and `b` if they are numbers, and concatenates `a` and `b` if they are strings:
"""

# ╔═╡ 524dee4b-6687-44de-af79-b99d3724f9ac
3 * 4

# ╔═╡ a16202d3-406a-4b43-807e-b60c673d4132
"3" * "4"

# ╔═╡ 41fc4769-5b67-4503-b1d1-2854a0d9c3a2
md"""
These, or similar examples, are duplicated notation common to most programming languages, but what Julia allows that others don't is for your own custom functions to utilise multiple dispatch as well. For example, say that I want to check if an input is four or not, but I don't mind if the input is a number or a string. To do this, I follow the argument with double colons and the type to determine when this method will run (here `Number` is an abstract type, which cannot itself be the type of an object, but encompasses all numeric types into one keyword). Note that such notation of arguments can still be used for functions with only a single methods, where it can be helpful for debugging purposes.
"""

# ╔═╡ acf690f4-b65b-4576-9145-1e894d5342f2
isfour(x::String) = (lowercase(x) ∈ ["4", "four"])

# ╔═╡ 9f193c51-b6ea-477f-ab0e-1bc5a4be4b86
md"""
`isfour` is now defined for a `Number` and a `String` input, but not currently for anything else. If I decide that nothing else can be considered as four, then I can create a third method using the keyword `Any` (the abstract type representing all possible types) which will deal with everything else:
"""

# ╔═╡ 29b729bb-cd1a-4ee1-ae2e-26b55ce744a4
isfour(x::Any) = false

# ╔═╡ 740f8275-ee0c-4aa2-b0c1-76d46fad3ecd
md"""
Of course, an input like `4` is also of type `Any`, but Julia's type system recognises that `Number` is more restrictive and so is of higher precedence, hence any `Number` input to `isfour` triggers the specialised method for that type.

More of the power of multiple dispatch is illustrated later in this case study.
"""

# ╔═╡ 73d93917-b72e-468b-970c-f34ec6d51de4
md"""
## Creating the `Permutation` type

### Declaring `Permutation` as a new type
I begin by declaring a new type with the keyword `struct`. The most basic syntax to do this is:
```julia
struct Permutation
    image::Vector{Int64}
end
```

although I will want more functionality than this. The new type has a single field `image` which is a vector of integers, representing the images of ``1, 2, \dots n`` under the permutation (where ``n`` is the maximum integer permuted). For example, the permutation ``(123)(45)`` has image ``(2, 3, 1, 5, 4)^T``, and can be created as an object by the syntax:
```julia
σ = Permutation([2,3,1,5,4])
```

Its field will also be accessible through dot notation by:
```julia
σ.image
```
"""

# ╔═╡ c18ae801-b186-44cf-9e42-83188c4781aa
md"""
There are currently some issues with the new type, primarily that at the moment any vector of integers could be the image of a `Permutation`, so I need to be able to check that an image is valid, i.e.
- The image contains all of the integers ``1, 2, ..., n`` for some ``n`` exactly once
- The image contains nothing else

I will also impose an additional condition, which is that the largest integer `n` must not be fixed, since there is no point in storing it. However, this can be amended within the code, rather than rejecting such an image vector.

The way to impose these conditions is to use an inner constructor, which is a method within the `struct` block.
```julia
struct Permutation
    image::Vector{Int64}

	function Permutation(image::Vector{Int64})
		# Inner constructor code
	end
end
```

To check that the vector `image` contains the right elements, I sort it and compare it to the vector `1:n`, where `n` is the length of `image`.
```julia
sort(image) == 1:length(image)
```

Then, I find the integers which aren't fixed by using Boolean indexing, which is indexing an array with an array of `true`s and `false`s of the same size, and taking the elements which correspond to `true`. The maximum of them is then found, with `1` added to the list since `maximum` cannot find the maximum element of an empty vector.
```julia
m = maximum([image[image .!= 1:length(image)]..., 1])
```

Now, I can make use of the inner constructor specific keyword `new` to construct the permutation, by ordering the arguments as they would be if I were calling the `Permutation` function. With an error message if the conditions are not met by `image`, this completes my inner constructor
```julia
function Permutation(image::Vector{Int64})
    sortedimage = sort(image)
 	m = maximum([image[image .!= 1:length(image)]..., 1])
    sortedimage == 1:last(sortedimage) && return new(image[1:m])
    error("not a valid permutation")
end
```
"""

# ╔═╡ 99fe42de-2afa-409f-9416-5c7671e4d0f4
md"""
One last thing I want to adjust before I am done with the `Permutation` type is to be able to construct a permutation by just giving it a list of integers and not having to wrap them up in vector form myself. This I do by utilising multiple dispatch, with an outer constructor, which is simply a new method with different inputs to construct a `Permutation` object.
```julia
Permutation(imagevals::Int64...) = Permutation([x for x ∈ imagevals])
```

Putting this all together, I get (including some code from later sections, due to the reactivity of the Pluto notebook in which this is written):
"""

# ╔═╡ b9c535e5-3317-4291-9925-696bc6cb0252
begin
	struct Permutation
    	image::Vector{Int64}

	    function Permutation(image::Vector{Int64})
    	    sortedimage = sort(image)
        	m = maximum([image[image .!= 1:length(image)]..., 1])
	        sortedimage == 1:last(sortedimage) && return new(image[1:m])
    	    error("not a valid permutation")
	    end
	end

	Permutation(imagevals::Int64...) = Permutation([x for x ∈ imagevals])

	# See section Evaluating a Permutation as a function below
	maxarg(σ::Permutation) = length(σ.image)
	(σ::Permutation)(x::Int64) = x ∈ 1:maxarg(σ) ? σ.image[x] : x

	# See section Generating groups with Permutations
	import Base.:(==)
	==(σ::Permutation, τ::Permutation) = (σ.image == τ.image)

	nothing
end

# ╔═╡ 1d56249e-d86a-4d1e-b7d2-6d63a621bfa8
isfour(x::Number) = (x == 4)

# ╔═╡ 9ab255ff-b40a-4f16-bed8-fe14e4495502
isfour(4.0)

# ╔═╡ 98750bd0-84bd-4cb1-900b-9ee26142e4ec
isfour("Twelve")

# ╔═╡ 4e4508db-47c0-462f-9cec-93e52bbffd7c
isfour([4])

# ╔═╡ de16f61d-853a-43ed-adf0-fb77ea3c8f4a
σ = Permutation(4,5,1,3,2); σ.image

# ╔═╡ 322f2dbc-f8f4-4f62-a9bb-ea6b00a333ae
md"""
### Evaluating a `Permutation` as a function
I now have a `Permutation` object, but it doesn't do anything yet. The next step is to be able to use a Permutation object as the function that it represents, i.e. evaluate it at an integer. I can do this by defining a function with its name as a `Permutation` parameter:
```julia
(σ::Permutation)(x::Int64) = σ.image[x]
```

However this will give an error for any integer not between `1` and the maximum unfixed integer (`m` in the inner constructor), since no such index exists. Instead, I want it to fix all other integers, which will make composing permutations later much easier, so I amend the code, defining a new function `maxarg` to go with it:
```julia
maxarg(σ::Permutation) = length(σ.image)
(σ::Permutation)(x::Int64) = x ∈ 1:maxarg(σ) ? σ.image[x] : x
```

This is included in the block above with the constructors. I can test the result of this on the sample permutation `σ` defined above
"""

# ╔═╡ ff8ab6b4-34de-44eb-b98a-2f5e0f1c48fc
σ(4)

# ╔═╡ 5ce89616-326b-4404-9f2f-707c91835fa3
maxarg(σ)

# ╔═╡ 9617f0f2-1b45-418f-8818-26df1abde6cb
md"""
### Creating a custom display format for a `Permutation` object
At the moment, whenever a `Permutation` object is returned, it displays in the default form for a new `struct`, which is the name of the type followed by a list of its fields (in this case `σ` would look like `Permutation([4, 5, 1, 3, 2])`). This isn’t very useful, so be more understandable (and more in keeping with standard mathematical notation), I want to customise this displayed form, which I can do with multiple dispatch by adding a method to the function `show` specifically for the `Permutation` type.

The format that I want to display the permutation in is as its disjoint cycle decomposition. This will first require a function to calculate the orbit generated by acting repeatedly on a given element:
"""

# ╔═╡ ee9d45f3-0a25-47f6-8445-8ba733c3a8f6
function orbit(σ::Permutation, x::Int64)
    orb = [x]
    y = σ(x)
    while y != x
        push!(orb,y)
        y = σ(y)
    end
    return orb
end

# ╔═╡ b119e9f0-b8ba-42a3-a76c-82ed2dc170a4
md"""
To calculate the disjoint cycle decomposition, I build up the output as follows:
- I start out with an empty `Vector{Vector{Int64}}` (that is, a vector whose elements are vectors of integers)
```julia
decomp = Vector{Int64}[]
```

- A vector `unaccounted` tracks which values in the range `1:maxarg(σ)` are yet to be added to the decomposition. I will iterate until this vector is entirely false
```julia
unaccounted = trues(maxarg(σ))

while any(unaccounted)
    # code to iterate
end
```

- Inside the loop, I will look for the first value which is unaccounted for in the decomposition so far, calculate its orbit, and then update `unaccounted` accordingly
```julia
x = findfirst(unaccounted)
xorbit = orbit(σ,x)
unaccounted[xorbit] .= false
```

- Then, if the orbit is non-trivial, I add it to the decomposition
```julia
length(xorbit) > 1 && push!(decomp,xorbit)
```

This results in `decomp` being a list of cycles exactly determining the disjoint cycle decomposition.
"""

# ╔═╡ e4ce9583-5c18-4df7-a5fb-a46670fe40af
function dcd(σ::Permutation)
    decomp = Vector{Int64}[]
    unaccounted = trues(maxarg(σ))
    while any(unaccounted)
        x = findfirst(unaccounted)
        xorbit = orbit(σ,x)
        unaccounted[xorbit] .= false
        length(xorbit) > 1 && push!(decomp,xorbit)
    end
    return decomp
end

# ╔═╡ b2ac585a-d84a-4d19-afe5-ad6d288e847f
md"""
I then need to consider how to build up the string from this decomposition
- Each cycle from the decomposition will be expressed as an open parenthesis, followed by the list of integers in order, separated by spaces, and then closed with another parenthesis. To do this, I list out strings that I can then concatenate with `*`.
```julia
*( "( ", ["$y " for y ∈ x]... , ")" )
```

- This needs to be repeated for all of the cycles in the decomposition, with an additional empty string added in case the disjoint cycle decomposition is empty.
```julia
toprint = *([ *( "( ", ["$y " for y ∈ x]... , ")" ) for x ∈ dcd(σ)]..., "")
```

- Finally, if (and only if) the permutation is the identity permutation, then the string will be empty at this point. Instead, I want to use the symbol `ι`.
```julia
toprint == "" && (toprint = "ι")
```

In order to write my new method for `show`, I first need to `import` it, ensuring that I do not overwrite the inbuilt methods that it has. Then, I construct my function, mirroring the syntax used for `show` in the source code as it is good practice.
"""

# ╔═╡ 7a96d2c8-aa9c-405f-b34f-1f3c7c53b6c6
begin
	import Base.show

	function show(io::IO,σ::Permutation)
	    toprint = *([ *( "( ", ["$y " for y ∈ x]... , ")" ) for x ∈ dcd(σ)]..., "")
    	toprint == "" && (toprint = "ι")
	    print(io, toprint)
	end
end

# ╔═╡ 12346644-107b-4c43-aaca-3e404fecc8b3
md"""
Permutations now have a much more aesthetic and useful display style:
"""

# ╔═╡ 0afa030a-58ad-44f6-a10f-0ae3baabe6f8
σ

# ╔═╡ 1e8659dc-27d5-429d-8431-d5cfa9006411
md"""
### Adding `Permutation` arithmetic
I now want to add some arithmetic for combining `Permutation` objects, starting with a shortcut for the identity. This will be represented by the constant `ι` (mirroring the way it is displayed by `show`), and is given by:
"""

# ╔═╡ 068a9d15-ce46-4aa5-9c13-3f79e70b0eda
const ι = Permutation(1)

# ╔═╡ b3151f52-6dd7-4f63-a158-becb5d1aec2e
md"""
Then, I add two methods to the inbuilt function one allowing it to be obtained either by passing a permutation, or the `Permutation` type. For these functions, the input variables are not given names, only their types specified, since the type is the only relevant property about the input, as needed for multiple dispatch.
"""

# ╔═╡ d614739d-b22f-4991-9722-e8371732bf4b
begin
	import Base.one
	one(::Permutation) = ι
	one(::Type{Permutation}) = ι
end

# ╔═╡ 363481da-d185-4846-82ab-c227e9f668f8
md"""
Next, I want to be able to compose permutations, for which I will write new `Permutation`-specific methods for the inbuilt operator `∘`. In order for an arbitrary number of permutations to be composable, I use an inductive definition which mirrors the definition of composition of functions [from the Julia source code](https://github.com/JuliaLang/julia/blob/master/base/operators.jl#L1096), that is:
```julia
∘(f) = f
∘(f, g) = ComposedFunction(f, g)
∘(f, g, h...) = ∘(f ∘ g, h...)
```

Hence, my methods for `∘` are:
"""

# ╔═╡ 8f925ec6-7cd8-49fe-a467-807711b3425b
begin
	import Base.∘
	∘(σ::Permutation) = σ
	∘(σ::Permutation, τ::Permutation) =
        Permutation([σ(τ(x)) for x ∈ 1:max(maxarg(σ),maxarg(τ))])
	∘(σ::Permutation, τ::Permutation, υ::Permutation...) = ∘(σ ∘ τ, υ...)
end

# ╔═╡ 4552e49c-af66-4fd3-a827-28284215f4f9
md"""
Thirdly, I want an inverse function to be able to find the inverse of a permutation. To do this, I need to find the index of each of `1:maxarg(σ)` in the image vector and let that be the image of the new permutation, which is done by:
"""

# ╔═╡ 0e710b07-b5db-4ea7-986b-8eb7fc0946b1
begin
	import Base.inv
	inv(σ::Permutation) = Permutation([findfirst(σ.image .== x) for x ∈ 1:maxarg(σ)])
end

# ╔═╡ 7c082457-980c-4961-bb10-d91b7248d0f6
md"""
Finally, I want to be able to exponentiate by any integer (including zero and negative integers), which can be done using some logic combined with the three operations above:
"""

# ╔═╡ 59075aa4-dd35-47af-b4d1-40e89b2e599f
begin
	import Base.^

	function ^(σ::Permutation,n::Int64)
	    n == 0 && return one(Permutation)
	    n < 0 && ((n,σ) = (-n,inv(σ)))
	    return ∘(fill(σ,n)...)
	end
end

# ╔═╡ 764fb693-f4ac-4327-93a3-2171621ad039
md"""
I can now test these out:
"""

# ╔═╡ 75e5491a-6536-44b0-a5b9-10cf7f9e5c66
ι

# ╔═╡ 9d43f21c-c60a-4bf8-b1a0-bf3b56b3a97b
σ

# ╔═╡ b421f174-d554-47a8-a72d-4d5f3ec6ba2e
τ = Permutation(4,5,2,3,1)

# ╔═╡ ff82e088-aa60-49fe-bb12-fe39b0430d51
σ ∘ τ

# ╔═╡ 7c0c4e99-ebff-409b-87f3-919a6870cccc
σ^-1

# ╔═╡ 029a0c9f-5d32-409d-b6ad-9e0d4e19e854
τ^5

# ╔═╡ c75f7100-778e-44c6-9d43-675d11175673
md"""
## Generating groups with `Permutation` objects

### Constructing symmetric and alternating groups
The first groups that I would like to be able to list the elements of are symmetric or alternating groups. Before this however, I will define two more functions which will be useful to me in their construction.

First, I would like a way to quickly construct a transposition:
"""

# ╔═╡ 0598e94b-a043-4106-bb7b-d032eb67be0b
function transposition(m::Int64, n::Int64)
    image = collect(1:max(m,n))
    image[m], image[n] = n, m
    return Permutation(image)
end

# ╔═╡ 00a61c28-230d-41b0-81ac-e2f0869d6640
transposition(3,4)

# ╔═╡ 32b52a58-6086-4c2f-a833-e1a7efb69d29
transposition(3,3)

# ╔═╡ 84bd5d92-00eb-4e66-a1c5-c86626b8f7d2
md"""
Note that `transposition(m,m)` returns the identity, which turns out to be exactly what I want.

To find the alternating group, I will need to consider the parity of elements. The easiest way to find the parity of a permutation is to consider it as a product of cycles and consider their parities. I already have such a decomposition, given by `dcd`, so using the `iseven` function, and `sum` to count the number of `true` elements of an array, I can write such a function:
"""

# ╔═╡ 1b10a694-a310-4a4e-9831-654e3e9691d1
parity(σ::Permutation) = (-1)^sum([iseven(length(x)) for x ∈ dcd(σ)])

# ╔═╡ 9812e6fb-9308-4c27-9d57-9fbf71474343
parity(ι)

# ╔═╡ d69ce761-f489-4728-a19d-783b84b436a5
parity(σ)

# ╔═╡ a6be1358-652d-4084-9c8e-f3089fb99169
parity(τ)

# ╔═╡ 41cf3bc2-c70c-4053-bed0-74a11f8a5ad5
md"""
Now, I have all the tools to create the `n`th symmetric group as a vector of `Permutation` objects, which I will do recursively using the inductive formula:
```math
S_1 = \{ \iota \}, \qquad S_n = \{ \sigma \circ (m n) : \sigma \in S_{n-1}, m \in \{ 1, 2, \dots, n \} \}
```
where as in the `transposition` function above, ``(n n) = \iota``.

- First, I will check that `n` is positive to avoid non-terminating loops (and also because the group doesn’t make sense otherwise)
```julia
n ≥ 1 || error("symmetric group must have a positive parameter")
```

- Then, I implement the base case of the trivial group when `n == 1`:
```julia
n == 1 && return [ι]
```

- For the recursive step, I copy `symmetricgroup(n-1)` into a vector `n` times, and then set up the corresponding vector of transpositions that I will multiply each by, using array filling and concatenation. Then, I compose the two elementwise:
```julia
permutations = vcat(fill(symmetricgroup(n-1),n)...)
transpositions = vcat([fill(transposition(i,n),factorial(n-1)) for i ∈ 1:n]...)
return permutations .∘ transpositions
```

The entire function is:
"""

# ╔═╡ c70a297d-adb7-4f70-b6bb-949a3d49d3c0
function symmetricgroup(n::Int64)
    n ≥ 1 || error("symmetric group must have a positive parameter")
    n == 1 && return [ι]

    permutations = vcat(fill(symmetricgroup(n-1),n)...)
    transpositions = vcat([fill(transposition(i,n),factorial(n-1)) for i ∈ 1:n]...)
    return permutations .∘ transpositions
end

# ╔═╡ eb5640fa-0814-497c-b157-da718df1c129
md"""
Now, finding the `n`th alternating group is as simple as finding the `n`th symmetric group and taking the even permutations:
"""

# ╔═╡ d546b70a-eeaa-45db-8b65-d36cb13b67f6
function alternatinggroup(n::Int64)
    Sₙ = symmetricgroup(n)
    return Sₙ[parity.(Sₙ) .== 1]
end

# ╔═╡ 5dea6f25-65fe-4304-a7e7-2d4b830f4724
symmetricgroup(5)

# ╔═╡ 13cacc56-7fcc-4a71-9cc5-be298295939b
alternatinggroup(5)

# ╔═╡ 46eeeed1-a8e4-49cc-8d97-6ddfb1fa2b1e
md"""
### Generating arbitrary finite groups
Instead of just symmetric and alternating groups, I would like to be able to generate a group given a list of `Permutation` objects as generators. By Cayley's theorem, this suffices to construct any finite group, assuming that you can find an appropriate list of generators. This will require one final use of multiple dispatch, which will be for me to define equality of permutations:
```julia
import Base.:(==)
==(σ::Permutation, τ::Permutation) = (σ.image == τ.image)
```

This snippet of code is included in the block with the definition of the `Permutation` type at the end of the section [Declaring `Permutation` as a new type](#73d93917-b72e-468b-970c-f34ec6d51de4), since otherwise the import of `==` here conflicts with its use in the definition of `Permutation` in the Pluto notebook format.

I can now write a function to generate a group from a list of generators. This I will do by repeatedly checking closure under multiplication, which is not a particularly clever or efficient algorithm, but it will work which is all I am looking for in this case.
"""

# ╔═╡ 3e8b7d82-1d63-4aec-8ed6-0b1f11c5cb33
function generate(σs::Permutation...)
    G = [σs...]
    while true
        newelements = Permutation[]

        for σ ∈ G, τ ∈ G
            υ = σ ∘ τ
            υ ∉ G && υ ∉ newelements && push!(newelements, υ)
        end

        if isempty(newelements)
            return G
        else
            G = vcat(G,newelements)
        end
    end
end

# ╔═╡ 79584487-dd19-4014-97e7-7fcc0582efa8
md"""
An interesting point of note here is that nowhere in the function do I use the syntax `==`, and yet the function would not work without it. This demonstrates further the power of multiple dispatch, since the operator `∉` uses `==` to determine whether or not `υ` is an element of `G` (or `newelements`), and I have defined what `==` means for a `Permutation`, so `∉` will automatically work for permutations without me explicitly telling it how to.

Some demonstrations of the `generate` function are shown below, with `C₅`, `D₁₀`, and `G` the group generated by `σ` and `τ`, which happens to be `S₅`:
"""

# ╔═╡ 4725c3bb-a1c1-49cc-9bf7-a7e443db085c
C₅ = generate(Permutation(2,3,4,5,1))

# ╔═╡ 9e84f004-fd88-4a25-bfba-0797d2df0942
D₁₀ = generate(Permutation(2,3,4,5,1),Permutation(5,4,3,2,1))

# ╔═╡ bd677d6a-3b42-497d-bcfb-58543dee7ebd
G = generate(σ,τ)

# ╔═╡ 076885f1-61b4-4913-8530-fd65c531f419
S₅ = symmetricgroup(5)

# ╔═╡ 738df784-eeb7-4d3b-b19f-6317f2459c6d
all([g ∈ G for g ∈ S₅]) && all([g ∈ S₅ for g ∈ G])

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
PlutoUI = "~0.7.9"
"""

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

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

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

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

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

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

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

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

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "477bf42b4d1496b454c10cce46645bb5b8a0cf2c"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.0.2"

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

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

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

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

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

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

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

# ╔═╡ Cell order:
# ╟─08eadc12-f8e9-11eb-1094-b7adc6936aa1
# ╟─c4d2ff29-2579-4079-a6cb-29015ef5ecd2
# ╟─db919e43-6dc8-4e53-8081-d9312446117a
# ╠═524dee4b-6687-44de-af79-b99d3724f9ac
# ╠═a16202d3-406a-4b43-807e-b60c673d4132
# ╟─41fc4769-5b67-4503-b1d1-2854a0d9c3a2
# ╠═1d56249e-d86a-4d1e-b7d2-6d63a621bfa8
# ╠═acf690f4-b65b-4576-9145-1e894d5342f2
# ╠═9ab255ff-b40a-4f16-bed8-fe14e4495502
# ╠═98750bd0-84bd-4cb1-900b-9ee26142e4ec
# ╟─9f193c51-b6ea-477f-ab0e-1bc5a4be4b86
# ╠═29b729bb-cd1a-4ee1-ae2e-26b55ce744a4
# ╠═4e4508db-47c0-462f-9cec-93e52bbffd7c
# ╟─740f8275-ee0c-4aa2-b0c1-76d46fad3ecd
# ╟─73d93917-b72e-468b-970c-f34ec6d51de4
# ╟─c18ae801-b186-44cf-9e42-83188c4781aa
# ╟─99fe42de-2afa-409f-9416-5c7671e4d0f4
# ╠═b9c535e5-3317-4291-9925-696bc6cb0252
# ╠═de16f61d-853a-43ed-adf0-fb77ea3c8f4a
# ╟─322f2dbc-f8f4-4f62-a9bb-ea6b00a333ae
# ╠═ff8ab6b4-34de-44eb-b98a-2f5e0f1c48fc
# ╠═5ce89616-326b-4404-9f2f-707c91835fa3
# ╟─9617f0f2-1b45-418f-8818-26df1abde6cb
# ╠═ee9d45f3-0a25-47f6-8445-8ba733c3a8f6
# ╟─b119e9f0-b8ba-42a3-a76c-82ed2dc170a4
# ╠═e4ce9583-5c18-4df7-a5fb-a46670fe40af
# ╟─b2ac585a-d84a-4d19-afe5-ad6d288e847f
# ╠═7a96d2c8-aa9c-405f-b34f-1f3c7c53b6c6
# ╟─12346644-107b-4c43-aaca-3e404fecc8b3
# ╠═0afa030a-58ad-44f6-a10f-0ae3baabe6f8
# ╟─1e8659dc-27d5-429d-8431-d5cfa9006411
# ╠═068a9d15-ce46-4aa5-9c13-3f79e70b0eda
# ╟─b3151f52-6dd7-4f63-a158-becb5d1aec2e
# ╠═d614739d-b22f-4991-9722-e8371732bf4b
# ╟─363481da-d185-4846-82ab-c227e9f668f8
# ╠═8f925ec6-7cd8-49fe-a467-807711b3425b
# ╟─4552e49c-af66-4fd3-a827-28284215f4f9
# ╠═0e710b07-b5db-4ea7-986b-8eb7fc0946b1
# ╟─7c082457-980c-4961-bb10-d91b7248d0f6
# ╠═59075aa4-dd35-47af-b4d1-40e89b2e599f
# ╟─764fb693-f4ac-4327-93a3-2171621ad039
# ╠═75e5491a-6536-44b0-a5b9-10cf7f9e5c66
# ╠═9d43f21c-c60a-4bf8-b1a0-bf3b56b3a97b
# ╠═b421f174-d554-47a8-a72d-4d5f3ec6ba2e
# ╠═ff82e088-aa60-49fe-bb12-fe39b0430d51
# ╠═7c0c4e99-ebff-409b-87f3-919a6870cccc
# ╠═029a0c9f-5d32-409d-b6ad-9e0d4e19e854
# ╟─c75f7100-778e-44c6-9d43-675d11175673
# ╠═0598e94b-a043-4106-bb7b-d032eb67be0b
# ╠═00a61c28-230d-41b0-81ac-e2f0869d6640
# ╠═32b52a58-6086-4c2f-a833-e1a7efb69d29
# ╟─84bd5d92-00eb-4e66-a1c5-c86626b8f7d2
# ╠═1b10a694-a310-4a4e-9831-654e3e9691d1
# ╠═9812e6fb-9308-4c27-9d57-9fbf71474343
# ╠═d69ce761-f489-4728-a19d-783b84b436a5
# ╠═a6be1358-652d-4084-9c8e-f3089fb99169
# ╟─41cf3bc2-c70c-4053-bed0-74a11f8a5ad5
# ╠═c70a297d-adb7-4f70-b6bb-949a3d49d3c0
# ╟─eb5640fa-0814-497c-b157-da718df1c129
# ╠═d546b70a-eeaa-45db-8b65-d36cb13b67f6
# ╠═5dea6f25-65fe-4304-a7e7-2d4b830f4724
# ╠═13cacc56-7fcc-4a71-9cc5-be298295939b
# ╟─46eeeed1-a8e4-49cc-8d97-6ddfb1fa2b1e
# ╠═3e8b7d82-1d63-4aec-8ed6-0b1f11c5cb33
# ╟─79584487-dd19-4014-97e7-7fcc0582efa8
# ╠═4725c3bb-a1c1-49cc-9bf7-a7e443db085c
# ╠═9e84f004-fd88-4a25-bfba-0797d2df0942
# ╠═bd677d6a-3b42-497d-bcfb-58543dee7ebd
# ╠═076885f1-61b4-4913-8530-fd65c531f419
# ╠═738df784-eeb7-4d3b-b19f-6317f2459c6d
# ╟─f1d0ed8f-bfd3-4210-84fd-f189bba6ebf9
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


# Matern 3/2 ARD covariance function

"""
    Mat32Ard <: MaternARD

ARD Matern 3/2 kernel (covariance)
```math
k(x,x') = σ²(1 + √3|x-x'|/L)\\exp(- √3|x-x'|/L)
```
with length scale ``ℓ = (ℓ₁, ℓ₂, …)`` and signal standard deviation ``σ`` where
``L = diag(ℓ₁, ℓ₂, …)``.
"""
mutable struct Mat32Ard{T} <: MaternARD where {T<:Real}
    "Inverse squared length scale"
    iℓ2::Vector{T}
    "Signal variance"
    σ2::T
    "Priors for kernel parameters"
    priors::Array
end

"""
Matern 3/2 ARD covariance function

    Mat32Ard(ll::Vector{T}, lσ::T)

# Arguments
  - `ll::Vector{Real}`: vector of length scales (given on log scale)
  - `lσ::Real`: signal standard deviation (given on log scale)
"""
Mat32Ard(ll::Vector{T}, lσ::T) where T = Mat32Ard{T}(exp.(-2 .* ll), exp(2 * lσ), [])

function set_params!(mat::Mat32Ard, hyp::AbstractVector)
    length(hyp) == num_params(mat) || throw(ArgumentError("Mat32 kernel has $(num_params(mat)) parameters, received $(length(hyp))."))
    @views @. mat.iℓ2 = exp(-2 * hyp[1:(end-1)])
    mat.σ2 = exp(2 * hyp[end])
end

get_params(mat::Mat32Ard{T}) where T = T[-log.(mat.iℓ2) / 2; log(mat.σ2) / 2]
get_param_names(mat::Mat32Ard) = [get_param_names(mat.iℓ2, :ll); :lσ]
num_params(mat::Mat32Ard) = length(mat.iℓ2) + 1

cov(mat::Mat32Ard, r::Number) =
    (s = √3 * r; mat.σ2 * (1 + s) * exp(-s))

dk_dll(mat::Mat32Ard, r::Real, wdiffp::Real) = 3 * mat.σ2 * wdiffp * exp(-√3 * r)


mult(x, y) = x * y
typeof(mult) #> Function
m = mult
m(6, 6) #> 36

plustwo = function (x)
              x + 2
          end
plustwo(3) #> 5

3+4 #> 7
+(3,4) #> 7

function numerical_derivative(f, x, dx=0.01)
  derivative = (f(x+dx) - f(x-dx))/(2*dx)
  return derivative
end

f = x -> 2x^2 + 30x + 9
println(numerical_derivative(f, 1, 0.001)) #> 33.99999999999537

function derivative(f)
    return function(x)
        # pick a small value for h
        h = x == 0 ? sqrt(eps(Float64)) : sqrt(eps(Float64)) * x
        xph = x + h
        dx = xph - x
        f1 = f(xph) # evaluate f at x + h
        f0 = f(x) # evaluate f at x
        return (f1 - f0) / dx  # divide by h
    end
end

# closure:
function counter()
    n = 0
    () -> n += 1, () -> n = 0
end

(addOne, reset) = counter()
addOne()
addOne()
addOne()
reset()

# currying:
function add(x)
    return function f(y)
        return x + y
    end
end
add(1)(2) #> 3
add(x) = f(y) = x + y
add(x) = y -> x + y

# nested:
function afun(x)
    z = x * 2
    function b(z)
        z += 1
    end
    b(z)
end

d = 5
afun(d) #> 11

# recursive:
sum(n) =  n > 1 ? sum(n-1) + n : n
sum(100) #> 5050
fib(n) = n < 2 ? n : fib(n-1) + fib(n-2)
println(fib(25)) #> 75025


### A Pluto.jl notebook ###
# v0.17.7

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 736856ce-f490-11eb-3349-057c86edfe7e
begin
    using Pkg
	Pkg.activate(joinpath(Pkg.devdir(), "MLCourse"))
    using MLJ, MLJLinearModels, Plots, DataFrames, Random
end

# ╔═╡ c693088f-7f80-4cdd-b9b5-65a50da732ac
begin
    using MLCourse
    import MLCourse: poly, Polynomial
    MLCourse.list_notebooks(@__FILE__)
end

# ╔═╡ e1077092-f72a-42af-b6e0-a616f938cba8
using PlutoUI; PlutoUI.TableOfContents()

# ╔═╡ 7eb6a060-f948-4d85-881a-4909e74c15bd
md"# Validation Set Approach

In the following cell we define a `data_generator` that creates data sets with a polynomial relationship of degree 3 between the input x and the average output f(x).
"

# ╔═╡ 7dd7e9a7-9245-4c64-af0c-8f7d2f62b2bf
begin
    f(x) = 0.3 + 2x - 0.8x^2 - 0.4x^3
    function data_generator(; n = 500, seed = 12)
        rng = MersenneTwister(seed)
        x = randn(rng, n)
        DataFrame(x = x, y = f.(x) .+ randn(rng, n))
    end
end;

# ╔═╡ 5049407c-bf93-41a5-97bd-863a69d7016a
md"## Test Error Estimation

We will estimate here the test error with the validation set approach for fixed hyper-parameters."

# ╔═╡ 46827e21-0eea-4ec9-83ed-05e41dda1502
data1 = data_generator(seed = 1);

# ╔═╡ 6cfc32b3-40c0-4149-ba67-2ec67b3938f3
md"split seed = $(@bind split_seed Slider(1:20, show_value = true))"

# ╔═╡ 1320b9b0-d6fc-4e61-8d62-de3988b8b21d
idxs = shuffle(MersenneTwister(split_seed), 1:500) # shuffle the indices

# ╔═╡ 5dc21667-d94f-4c2a-b217-80b9a2262d8c
begin
	data1_train = data1[idxs[1:250], :]
	data1_test = data1[idxs[251:end], :]
	mach = machine(Polynomial(degree = 8) |> LinearRegressor(),
	               select(data1_train, :x), data1_train.y)
	fit!(mach, verbosity = 0)
	(training_error = rmse(predict(mach), data1_train.y),
	 test_error = rmse(predict(mach, select(data1_test, :x)), data1_test.y))
end

# ╔═╡ 1fe3f6d1-6336-417e-885e-de3fc5821bdf
let
	scatter(data1_train.x, data1_train.y, label = "training set")
	scatter!(data1_test.x, data1_test.y, label = "test set")
	xgrid = -4:.1:4
	pred = predict(mach, DataFrame(x = xgrid))
	plot!(xgrid, pred, label = "fit", ylim = (-6, 5), xlim = (-4, 4))
end

# ╔═╡ b79fae90-87f3-4f89-933b-d82e76c94d81
md"We see that the resulting curve and the test error depends strongly on the training set, i.e. the indices that are selected for the training set. The training error is always close to the irreducible error in this case."

# ╔═╡ 7d30f043-2915-4bba-aad5-cb7dbe76a3e5
md"## Hyper-parameter tuning

Now we will use a nested validation set approach to find the optimal degree in polynomial regression and estimate the test error. For this we will split the data into training, validation and test set. For each hyper-parameter value we will fit the parameters on the training set and estimate the test error on the validation set. When we have found the best hyper-parameters based on the test error computed on the validation set, we will estimate the test error of the best model using the test set."

# ╔═╡ ffffa3d0-bffb-4ff8-9966-c3312f952ac5
begin
	function data_split(data;
                        shuffle = false,
                        idx_train = 1:50,
                        idx_valid = 51:100,
                        idx_test = 101:500)
        idxs = if shuffle
                randperm(size(data, 1))
            else
                1:size(data, 1)
            end
        (train = data[idxs[idx_train], :],
         valid = data[idxs[idx_valid], :],
         test = data[idxs[idx_test], :])
    end
    function fit_and_evaluate(model, data)
        mach = fit!(machine(model, select(data.train, :x), data.train.y),
                    verbosity = 0)
        (train = rmse(predict(mach, select(data.train, :x)), data.train.y),
         valid = rmse(predict(mach, select(data.valid, :x)), data.valid.y),
         test = rmse(predict(mach, select(data.test, :x)), data.test.y))
    end
end

# ╔═╡ 2e02ac0d-c4d0-47ba-be57-445adeb6ab8b
losses1 = [fit_and_evaluate(Polynomial(; degree) |> LinearRegressor(),
	                        data_split(data1))
           for degree in 1:10]

# ╔═╡ 91eecd2b-af18-4a63-9684-28950e604d1a
let validlosses = getproperty.(losses1, :valid), i = argmin(validlosses)
    plot(1:10, validlosses, label = nothing)
    scatter!([i], [validlosses[i]], label = nothing,
             xlabel = "degree", ylabel = "validation loss")
end

# ╔═╡ b45b32b1-8a65-4823-b3bb-f0b7cc57604b
md"For this seed of the random number generator the optimal degree (x coordinate of the red point in the figure above) found by the validation set approach is close to the actual degree of the data generating process. The estimated test loss of the model with degree $(argmin((x -> x.valid).(losses1))) is $(losses1[argmin((x -> x.valid).(losses1))].test)."

# ╔═╡ 23e91f72-0124-48d4-8b53-bd5da67b7ac7
losses1[4]

# ╔═╡ 4461d5b0-bc9b-4c89-aa76-70524a5caa7c
md"Let us repeat the process above for other splits of the data."

# ╔═╡ ac3c7e84-6c47-4dc1-b862-de4cfb05dad9
losses = [fit_and_evaluate(Polynomial(; degree) |> LinearRegressor(),
                           data_split(data1, shuffle = true))
          for degree in 1:10, _ in 1:20]

# ╔═╡ 26362233-b006-423d-8fb5-7cd9150405b4
let validlosses = getproperty.(losses, :valid), i = argmin(validlosses, dims = 1)
    plot(validlosses, label = nothing, ylims = (0, 10))
    scatter!((x -> x[1]).(i), validlosses[i], label = nothing,
             xlabel = "degree", ylabel = "validation loss")
end

# ╔═╡ 5cff19c7-3426-4d91-b555-f059f5f41886
md"Here we see a large variability of the optimal degree (x coordinates of all the points in the figure above).

In the following cell we compare the validation set estimate of the test error to the test set estimation of the test error. We can confirm that the validation set estimate is much lower than the test set estimate."

# ╔═╡ df16313f-76be-4d8b-88e7-7cc618ff49d4
let validlosses = getproperty.(losses, :valid),
	testlosses = getproperty.(losses, :test),
	winners = argmin(validlosses, dims = 1)
	(mean_valid_winner = mean(validlosses[winners]),
	 mean_test_winner = mean(testlosses[winners]))
end

# ╔═╡ 46bffe01-1a7c-405b-b3ec-bfe7570b8a3c
md"# Cross-Validation

Here we would like to find the optimal hyper-parameters with cross-validation.
We will not estimate the final test error of the model with the best hyper-parameters.

In the following cell we implement cross-validation."

# ╔═╡ b3aed705-5072-4d9b-bb8f-865ac1561bf6
begin
    function cross_validation_sets(idx, K)
        n = length(idx)
        r = n ÷ K
        [let idx_valid = idx[(i-1)*r+1:(i == K ? n : i*r)]
             (idx_valid = idx_valid, idx_train = setdiff(idx, idx_valid))
         end
         for i in 1:K]
    end
    function cross_validation(model, data; K = 5, shuffle = false)
		idxs = 1:size(data, 1)
		if shuffle
            idxs = Random.shuffle(idxs)
		end
        losses = [fit_and_evaluate(model,
                                   data_split(data;
									          idx_test = [], # no test set
									          idxs...)) # training and validation
                  for idxs in cross_validation_sets(idxs, K)]
        (train = mean(getproperty.(losses, :train)),
         valid = mean(getproperty.(losses, :valid)))
    end
end

# ╔═╡ 6916273a-d116-4173-acaa-2bcac1d1753b
cross_validation_sets(1:100, 4)

# ╔═╡ c14c9e49-a993-456e-9115-97da86f8e498
losses1_cv10 = [cross_validation(Polynomial(; degree) |> LinearRegressor(),
		                         data1,
		                         K = 10) for degree in 1:10]

# ╔═╡ 4b435518-2f12-4921-bb1f-fdd049ddfaed
let validlosses = getproperty.(losses1_cv10, :valid), i = argmin(validlosses)
    plot(1:10, validlosses, label = nothing)
    scatter!([i], [validlosses[i]], label = nothing,
             xlabel = "degree", ylabel = "validation loss")
end

# ╔═╡ 1e584a38-2fef-4877-87f6-92237d71c4b3
losses_cv10 = [cross_validation(Polynomial(; degree) |> LinearRegressor(),
                                data1,
                                K = 10, shuffle = true)
               for degree in 1:10, _ in 1:20]

# ╔═╡ a7c88b3f-92cb-4253-a889-c78683722c1d
let validlosses = getproperty.(losses_cv10, :valid), i = argmin(validlosses, dims = 1)
    plot(validlosses, label = nothing, ylims = (0.5, 2))
    scatter!((x -> x[1]).(i), validlosses[i], label = nothing,
             xlabel = "degree", ylabel = "validation loss")
end

# ╔═╡ 0b81c3d5-277a-4fe6-889b-550e2f83c39d
md"With cross-validation we find often degree = 3 as the optimal one."

# ╔═╡ 5747d116-1e61-45f0-a87b-89372c6f270f
md"### Tuning Hyper-parameters with cross-validation and estimating the test error with the validation set approach

Instead of using all data for hyper-parameter search using cross-validation, we could set some data aside as a test set and use this to estimate the test error of the model with the optimal hyper-parameters."

# ╔═╡ 8204d9da-855d-4c66-b4dd-2c18a17b539b
losses2_cv10 = [cross_validation(Polynomial(; degree) |> LinearRegressor(),
		                         data1[1:100, :],
		                         K = 10) for degree in 1:10]

# ╔═╡ 96b3c070-0c72-44c7-ac98-ff6cb7e82380
md"The optimal degree found with cross-validation on the first 100 data points is:"

# ╔═╡ bae635b1-6eba-4ba4-80e9-547770612843
argmin((x -> x.valid).(losses2_cv10))

# ╔═╡ d9b04582-3a7a-4e19-a0da-d363810883f7
md"Let us refit the model with the optimal degree on all data that was used for cross-validation and estimate it's test error on data points 101 to 500."

# ╔═╡ 9042e811-9a34-42b8-8069-877f3b3f1e75
let
	model = Polynomial(degree = 4) |> LinearRegressor()
	mach = machine(model, select(data1[1:100, :], :x), data1[1:100, :y])
	fit!(mach, verbosity = 0)
	testset = data1[101:end, :]
	(test_error = rmse(predict(mach, select(testset, :x)), testset.y),)
end

# ╔═╡ e2658128-4053-4484-8e1e-229eceb755ab
md"# Resampling Strategies with MLJ"

# ╔═╡ eaacf529-1727-4743-941b-360c53088b1d
md"Instead of our own cross-validation function we can also use the builtin functions of MLJ. The `train_test_pairs` function does basically the same as our `cross_validation_sets` function."

# ╔═╡ 967a6e08-f0e6-45b7-988b-d0df237f3ddf
MLJ.MLJBase.train_test_pairs(CV(nfolds = 10), 1:100)

# ╔═╡ 535a11d9-833e-4237-949e-9a3e289c600b
md"`MLJ` has the very useful function `evaluate!`. Have a look at the Live docs to learn more about this function. In the following you see an example where we perform 10-fold cross-validation on the `rmse` measure to estimate the expected test error for a polynomial regressor of degree 4."

# ╔═╡ 8ae790b3-3987-4f41-8e21-adbb71081eb9
evaluate!(machine(Polynomial(degree = 4) |> LinearRegressor(),
	              select(data1, :x), data1.y),
          resampling = CV(nfolds = 10), measure = rmse)

# ╔═╡ 134df37f-b737-4c80-a5f9-1149aeec970c
md"We can use this now to find the best degree with 10-fold cross-validation for 100 different seeds. You will see that - even though not perfect - the hyper-parameters found with 10-fold cross-validation have lower variance and are usually closer to the true value than with the validation set approach."

# ╔═╡ abb71af3-8aae-4806-9d5a-d144c15d22ef
losses_mlj_cv10 = [evaluate!(machine(Polynomial(; degree) |> LinearRegressor(),
                                     select(data1, :x), data1.y),
                             resampling = CV(nfolds = 10, shuffle = true),
                             measure = rmse,
                             verbosity = 0).measurement[]
                   for degree in 1:10, seed in 1:20]

# ╔═╡ 6599d2b4-68f4-4c22-8e40-bf3722597692
let validlosses  = losses_mlj_cv10, i = argmin(validlosses, dims = 1)
    plot(validlosses, label = nothing, ylims = (0.5, 2))
    scatter!((x -> x[1]).(i), validlosses[i], label = nothing,
             xlabel = "degree", ylabel = "validation loss")
end

# ╔═╡ fc07009a-eeac-4c6b-9559-1e99ae68a6c3
md"MLJ implements the following resampling strategies. With `Holdout` we use the validation set approach."

# ╔═╡ 42511b31-6d70-495e-8362-01a29143b96e
subtypes(ResamplingStrategy)

# ╔═╡ 29683c99-6a6a-4f65-bea2-d592895d887e
md"# Model Tuning

Finding good hyper-parameters (tuning) is such an important step in the process of finding good machine learning models that there exist some nice utility functions to tune the hyper-parameters ([tuning section in the MLJ manual](https://alan-turing-institute.github.io/MLJ.jl/dev/tuning_models/)). In the cell below you see an example where the degree of polynomial regression is automatically tuned on a given dataset by performing 10-fold cross-validation on all degrees on a \"grid\", i.e. of all degrees from 1 to 17 are tested."

# ╔═╡ f93f20db-4fed-481f-b085-ca744b68fa8f
begin
    model = Polynomial() |> LinearRegressor()
    data2 = data_generator(seed = 2, n = 100)
    self_tuning_model = TunedModel(model = model,
                                   resampling = CV(nfolds = 10),
                                   tuning = Grid(),
                                   range = range(model, :(polynomial.degree),
                                                 values = 1:17),
                                   measure = rmse)
    self_tuning_mach = machine(self_tuning_model, select(data2, :x), data2.y)
	fit!(self_tuning_mach, verbosity = 0)
end

# ╔═╡ d72609a1-f93f-4ca6-8759-727662233e97
md"The self-tuned model actually found the true degree 3 for this data set, as we can see from the report of the tuned machine. The estimated test error can also be found in the report under `best_history_entry.measurement`."

# ╔═╡ 59245b3e-ddfc-46c4-ba44-86ce191672ae
report(self_tuning_mach)

# ╔═╡ f8b48f50-09cb-498a-89d1-9ac9b5722d0c
md"The result of the tuned machine looks quite good here: the reducible error (average distance between red and green curve) looks small."

# ╔═╡ 47a2c14b-a02d-43d7-ac8a-4d95a0d91fa8
let x = -3:.1:3
    scatter(data2.x, data2.y, label = "data", legend = :bottomleft)
    plot!(f, label = "generator", w = 2, xlims = (-3, 3))
    plot!(x, predict(self_tuning_mach, DataFrame(x = x)), label = "self tuning fit", w = 2)
end

# ╔═╡ cf3841ba-3963-4716-8fc9-8cce1dc4a5fa
md"# Nested Cross-Validation

If we want to tune the hyper-parameters with cross-validation and care about the estimated test error of the model found with cross-validation, we can use nested cross-validation, where hyper-parameters are optimized with cross-validation for each fold of the outer cross-validation. This happens when we `evaluate!` with cross-validation a self-tuning machine that itself uses cross-validation for hyper-parameter tuning. If we used `resampling = Holdout(fraction_train = 0.5)` we would use the validation set approach to estimate the test error with 50% of the data in the training set. If we used `resampling = Holdout(fraction_train = 0.5)` in the self-tuning machine, we would use the validation set approach for tuning the hyper-parameters.
"

# ╔═╡ 3e9c6f1c-4cb6-48d8-8119-07a4b03c2e4b
nested_cv = evaluate!(machine(self_tuning_model, select(data1, :x), data1.y),
                      resampling = CV(nfolds = 5), measure = rmse)

# ╔═╡ eba90af0-bf67-48f8-8d3c-f831b4fc289a
md"Have a look at the best models found for each fold:"

# ╔═╡ c8d26ab6-86ec-4b37-9540-0f785fd8cdc2
nested_cv.report_per_fold

# ╔═╡ ee89c448-1e69-4fd1-a4b8-7297a09f2685
md"# Exercises

## Conceptual
1. We review k-fold cross-validation.
    - Explain how k-fold cross-validation is implemented.
    - What are the advantages and disadvantages of k-fold cross-validation relative to:
        - The validation set approach?
        - LOOCV?
3. You are given a model with unknown hyper-parameters and the goal is to find the optimal hyper-parameters and estimate the test error of the optimal model. Suppose of colleague of yours argues that this should be done in the following way: \"Take all the data and run cross-validation to find the best hyper-parameters. Taking all data is better, than taking only a subset, because one has more data to estimate the parameters and estimate the test error on the validation sets. Once the best hyper-parameters are found: estimate the test error by running again cross-validation on all the data with the hyper-parameter fixed to the best value.\" Do you agree with your colleague? If not, explain where you disagree and where you agree.
2. Suppose you receive the following email of a colleague of yours. Write an answer to this email.
```
    Hi

    In my internship I am doing this research project with a company that
    manufactures constituents for medical devices. We need to assure high
    quality and detect the pieces that are not fully functional. To do so we
    perform 23 measurements on each piece. My goal is to find a machine
    learning method that detects defective pieces automatically. Both, my
    training set and my test set consist of 4000 pieces that were fine and
    30 defective ones.  I ran logistic regression on that data and found a
    training error of 0.1% and a test error of 0.3%, which is already pretty
    good, I think, no? But with kNN classification it is even better. For
    k = 7 I found a test error of 0.05% which is by far lower than all test
    errors I obtained with other values of k. I was really impressed. Now
    we have a method that predicts with 99.95% accuracy whether a piece is
    defective or not!

    Because you are taking this machine learning class now, I wanted to ask you
    for advice. Does it all make sense to you what I described? If not, do you
    have any suggestion to get even better results?

    Ã bientôt
    Jamie
```

## Applied
1. Take the `classification_data` in our notebook on
   \"flexibility and bias-variance-decomposition notebook\" and find with 10-fold
   cross-validation the optimal number ``k`` of neighbors of kNN
   classification, using the AUC measure. Hint: `MLJ` has the builtin function `auc`.
   Plot the validation AUC for ``k = 1, \ldots, 50``.
3. With the same data as in the previous exercise, use the `MLJ` function `estimate!` and a self tuning machine to estimate with the validation set approach the test error of kNN classifier whose hyper-parameter is tuned with 5 fold cross-validation. Use one quarter of the data for the test set.
2. In this exercise you apply our \"recipe for supervised learning\" (see slides). The goal is to predict the miles a car can drive per gallon fuel (mpg) as a function of its horsepower. You can download a dataset with `using OpenML; cars = DataFrame(OpenML.load(455))`. In the cleaning step we will remove all rows that contain missing values (you can use the function `dropmissing`). We select the machine learning methods polynomial regression and k nearest neighbors regression and we take as measure the `rmse`. Make sure to go trough the steps 2, 5, 9 of the recipe. Plot the predictions of the best method you found.
"

# 1. Perform k-nearest neighbors regression on data generated with our `data_generator`
#    defined in the first cell of this notebook and find the optimal number k of neighbors
#    with k-fold cross validation.

# ╔═╡ 0651292e-3f4e-4263-8235-4caa563403ec
MLCourse.footer()

# ╔═╡ Cell order:
# ╠═736856ce-f490-11eb-3349-057c86edfe7e
# ╟─7eb6a060-f948-4d85-881a-4909e74c15bd
# ╠═7dd7e9a7-9245-4c64-af0c-8f7d2f62b2bf
# ╟─5049407c-bf93-41a5-97bd-863a69d7016a
# ╠═46827e21-0eea-4ec9-83ed-05e41dda1502
# ╟─6cfc32b3-40c0-4149-ba67-2ec67b3938f3
# ╠═1320b9b0-d6fc-4e61-8d62-de3988b8b21d
# ╠═5dc21667-d94f-4c2a-b217-80b9a2262d8c
# ╟─1fe3f6d1-6336-417e-885e-de3fc5821bdf
# ╟─b79fae90-87f3-4f89-933b-d82e76c94d81
# ╟─7d30f043-2915-4bba-aad5-cb7dbe76a3e5
# ╠═ffffa3d0-bffb-4ff8-9966-c3312f952ac5
# ╠═2e02ac0d-c4d0-47ba-be57-445adeb6ab8b
# ╟─91eecd2b-af18-4a63-9684-28950e604d1a
# ╟─b45b32b1-8a65-4823-b3bb-f0b7cc57604b
# ╠═23e91f72-0124-48d4-8b53-bd5da67b7ac7
# ╟─4461d5b0-bc9b-4c89-aa76-70524a5caa7c
# ╠═ac3c7e84-6c47-4dc1-b862-de4cfb05dad9
# ╟─26362233-b006-423d-8fb5-7cd9150405b4
# ╟─5cff19c7-3426-4d91-b555-f059f5f41886
# ╠═df16313f-76be-4d8b-88e7-7cc618ff49d4
# ╟─46bffe01-1a7c-405b-b3ec-bfe7570b8a3c
# ╠═b3aed705-5072-4d9b-bb8f-865ac1561bf6
# ╠═6916273a-d116-4173-acaa-2bcac1d1753b
# ╠═c14c9e49-a993-456e-9115-97da86f8e498
# ╟─4b435518-2f12-4921-bb1f-fdd049ddfaed
# ╠═1e584a38-2fef-4877-87f6-92237d71c4b3
# ╟─a7c88b3f-92cb-4253-a889-c78683722c1d
# ╟─0b81c3d5-277a-4fe6-889b-550e2f83c39d
# ╟─5747d116-1e61-45f0-a87b-89372c6f270f
# ╠═8204d9da-855d-4c66-b4dd-2c18a17b539b
# ╟─96b3c070-0c72-44c7-ac98-ff6cb7e82380
# ╠═bae635b1-6eba-4ba4-80e9-547770612843
# ╟─d9b04582-3a7a-4e19-a0da-d363810883f7
# ╠═9042e811-9a34-42b8-8069-877f3b3f1e75
# ╟─e2658128-4053-4484-8e1e-229eceb755ab
# ╟─eaacf529-1727-4743-941b-360c53088b1d
# ╠═967a6e08-f0e6-45b7-988b-d0df237f3ddf
# ╟─535a11d9-833e-4237-949e-9a3e289c600b
# ╠═8ae790b3-3987-4f41-8e21-adbb71081eb9
# ╟─134df37f-b737-4c80-a5f9-1149aeec970c
# ╠═abb71af3-8aae-4806-9d5a-d144c15d22ef
# ╟─6599d2b4-68f4-4c22-8e40-bf3722597692
# ╟─fc07009a-eeac-4c6b-9559-1e99ae68a6c3
# ╠═42511b31-6d70-495e-8362-01a29143b96e
# ╟─29683c99-6a6a-4f65-bea2-d592895d887e
# ╠═f93f20db-4fed-481f-b085-ca744b68fa8f
# ╟─d72609a1-f93f-4ca6-8759-727662233e97
# ╠═59245b3e-ddfc-46c4-ba44-86ce191672ae
# ╟─f8b48f50-09cb-498a-89d1-9ac9b5722d0c
# ╟─47a2c14b-a02d-43d7-ac8a-4d95a0d91fa8
# ╟─cf3841ba-3963-4716-8fc9-8cce1dc4a5fa
# ╠═3e9c6f1c-4cb6-48d8-8119-07a4b03c2e4b
# ╟─eba90af0-bf67-48f8-8d3c-f831b4fc289a
# ╠═c8d26ab6-86ec-4b37-9540-0f785fd8cdc2
# ╟─ee89c448-1e69-4fd1-a4b8-7297a09f2685
# ╟─c693088f-7f80-4cdd-b9b5-65a50da732ac
# ╟─e1077092-f72a-42af-b6e0-a616f938cba8
# ╟─0651292e-3f4e-4263-8235-4caa563403ec


macro myevalpoly(z,a...)
    isempty(a) && error("You forgot to pass coefficients!")
    ex = :($(a[length(a)]))
    for i in 1:length(a)-1
       ex = :($ex * $(z) + $(a[length(a)-i]) )
    end
    println(ex)
    ex
end

@myevalpoly 7 2 3 4 5
@evalpoly 7 2 3 4 5

##
function root2coeff(z::AbstractVector{T}) where T
    N = length(z)
    co = zeros(T, N+1)
    # The last coefficient is always one
    co[end] = 1
    # The outer loop adds one root at a time
    for j in 1:N, i in j:-1:1
        co[end-i] -= z[j]*co[end-i+1]
    end
    co
end
@show typemax(Int), typemax(Int128)
root2coeff(1:20)
root2coeff(Int128(1):20)

##
using LinearAlgebra
function poly_roots(z)
    len = length(z)
    # construct the ones part
    mat = diagm(-1 => ones(len-2))
    # insert coefficients
    mat[:, end] = -z[1:end-1]
    eigvals(mat)
end

## Calculate all the roots and plot it
using Random
Random.seed!(2020)
function wilkinson_poly_roots(n=100)
    # original coefficients
    coeff = root2coeff(Int128(1):20)
    rts = Vector{Complex{Float64}}[]
    # add perturbation
    for i in 1:n
        pert_coeff = coeff.*(1 .+ rand(21)*1e-6)
        push!(rts, poly_roots(pert_coeff))
    end
    rts
end
using Plots; gr()
function plt_wilkinson_roots(rts)
    # plot roots without perturbation
    plt = scatter(1:20, zeros(20), color = :green, markersize = 5, legend = false, grid = false, aspect_ratio = 1)
    for i in eachindex(rts)
        # plot roots with perturbation
        scatter!(plt, real.(rts[i]), imag.(rts[i]), color = :red, mscolor = :red, markersize = 2, shape = :star)
    end
    plt
end
wilkinson_poly_roots(100) |> plt_wilkinson_roots


# ## Getting started

using MLJ, RDatasets, Random
MLJ.color_off() # hide

data = dataset("datasets", "USArrests")
names(data)

# Let's have a look at the mean and standard deviation of each feature:

describe(data, :mean, :std)

# Let's extract the numerical component and coerce

X = select(data, Not(:State))
X = coerce(X, :UrbanPop=>Continuous);

# ## PCA pipeline
#
# PCA is usually best done after standardization but we won't do it here:

@load PCA pkg=MultivariateStats

pca_mdl = PCA(pratio=1)
pca = machine(pca_mdl, X)
fit!(pca)

W = transform(pca, X);

# W is the PCA'd data; here we've used default settings for PCA and it has recovered 2 components:

schema(W).names

# Let's inspect the fit:

r = report(pca)
cumsum(r.principalvars ./ r.tvar)

# In the second line we look at the explained variance with 1 then 2 PCA features and it seems that with 2 we almost completely recover all of the variance.

# ## More interesting data...

# Instead of just playing with toy data, let's load the orange juice data and extract only the columns corresponding to price data:

data = dataset("ISLR", "OJ")

X = select(data, [:PriceCH, :PriceMM, :DiscCH, :DiscMM, :SalePriceMM,
                  :SalePriceCH, :PriceDiff, :PctDiscMM, :PctDiscCH]);

# ### PCA pipeline

Random.seed!(1515)

@pipeline SPCA(std = Standardizer(),
               pca = PCA(pratio=1-1e-4))
spca_mdl = SPCA()
spca = machine(spca_mdl, X)
fit!(spca)
W = transform(spca, X)
names(W)

# What kind of variance can we explain?

r  = report(spca).reports[1]
cs = cumsum(r.principalvars ./ r.tvar)

# Let's visualise this

using PyPlot

figure(figsize=(8,6))

bar(1:length(cs), cs)
plot(1:length(cs), cs, color="red", marker="o")

xlabel("Number of PCA features", fontsize=14)
ylabel("Ratio of explained variance", fontsize=14)

savefig(joinpath(@OUTPUT, "ISL-lab-10-g1.svg")) # hide

# \figalt{PCA explained variance}{ISL-lab-10-g1.svg}

# So 4 PCA features are enough to recover most of the variance.

# ### Clustering

Random.seed!(1515)

@load KMeans pkg=Clustering
@pipeline SPCA2(std = Standardizer(),
                pca = PCA(),
                km = KMeans(k=3))

spca2_mdl = SPCA2()
spca2 = machine(spca2_mdl, X)
fit!(spca2)

assignments = report(spca2).reports[1].assignments
mask1 = assignments .== 1
mask2 = assignments .== 2
mask3 = assignments .== 3;

# Now we can  try visualising this

using PyPlot

figure(figsize=(8, 6))
for (m, c) in zip((mask1, mask2, mask3), ("red", "green", "blue"))
    plot(W[m, 1], W[m, 2], ls="none", marker=".", markersize=10, color=c)
end

xlabel("PCA-1", fontsize=13)
ylabel("PCA-2", fontsize=13)
legend(["Group 1", "Group 2", "Group 3"], fontsize=13)

savefig(joinpath(@OUTPUT, "ISL-lab-10-cluster.svg")) # hide

# \fig{ISL-lab-10-cluster.svg}
PyPlot.close_figs() # hide


using Test
import DynamicPolynomials: @polyvar
# using Revise
using PolynomialBasis

PB = PolynomialBasis

@polyvar x

function allequal(v1,v2,tol)
    return allequal(v1,v2;tol=tol)
end

function allequal(v1,v2;tol=1e3eps())
    np = length(v1)
    f = length(v2) == np
    return f && all([isapprox(v1[i],v2[i],atol=tol) for i = 1:np])
end

basis = PB.CubicHermiteBasis()
@test allequal(basis.points,[-1.0,1.0])
@test allequal(basis(-1.0),[1.0,0.,0.,0.])
@test allequal(basis(1.0),[0.0,0.,1.,0.])
@test allequal(PB.derivative(basis,-1.0),[0.,1.,0.,0.])
@test allequal(PB.derivative(basis,1.0),[0.,0.,0.,1.])
@test allequal(PB.gradient(basis,-1.0),[0.,1.,0.,0.])
@test allequal(PB.gradient(basis,1.0),[0.,0.,0.,1.])

hess = PB.hessian(basis,-1.0)
@test allequal(hess,[-1.5,-2.0,1.5,-1.0])
hess = PB.hessian(basis,1.0)
@test allequal(hess,[1.5,1,-1.5,2])


#
# Estimate-pi.jl
#
sumsq(x,y) = x*x + y*y;

N=1000000;
x = 0;
for i = 1:N
  if sumsq(rand(), rand()) < 1.0
    x += 1
  end
end
@printf "Estimate of PI for %d trials is %8.5f\n" N 4.0*(x / N);


export TFTS, TAAFT

"""
    TFTS(fϵ::Real)

A truncated Fourier transform surrogate[^Nakamura2006] (TFTS).

TFTS surrogates are generated by leaving some frequencies untouched when performing the
phase shuffling step (as opposed to randomizing all frequencies, like for
[`RandomFourier`](@ref) surrogates).

These surrogates were designed to deal with data with irregular fluctuations superimposed
over long term trends (by preserving low frequencies)[^Nakamura2006]. Hence, TFTS surrogates
can be used to test the null hypothesis that the signal is a stationary linear system
generated the irregular fluctuations part of the signal[^Nakamura2006].

## Controlling the truncation of the spectrum

The truncation parameter `fϵ ∈ [-1, 0) ∪ (0, 1]` controls which parts of the spectrum are preserved.

- If `fϵ > 0`, then `fϵ` indicates the ratio of high frequency domain to the entire frequency domain.
    For example, `fϵ = 0.5` preserves 50% of the frequency domain (randomizing the higher
    frequencies, leaving low frequencies intact).
- If `fϵ < 0`, then `fϵ` indicates ratio of low frequency domain to the entire frequency domain.
    For example, `fϵ = -0.2` preserves 20% of the frequency domain (leaving higher frequencies intact,
    randomizing the lower frequencies).
- If `fϵ ± 1`, then all frequencies are randomized. The method is then equivalent to
    [`RandomFourier`](@ref).

The appropriate value of `fϵ` strongly depends on the data and time series length, and must be
manually determined[^Nakamura2006], for example by comparing periodograms for the time series and
the surrogates.

[^Nakamura2006]: Nakamura, Tomomichi, Michael Small, and Yoshito Hirata. "Testing for nonlinearity in irregular fluctuations with long-term trends." Physical Review E 74.2 (2006): 026205.
"""
struct TFTS <: Surrogate
    fϵ::Real

    function TFTS(fϵ::Real)
        if !(0 < fϵ ≤ 1) && !(-1 ≤ fϵ < 0)
            throw(ArgumentError("`fϵ` must be on the interval [-1, 0) ∪ (0, 1] (positive if preserving high frequencies, negative if preserving low frequencies)"))
        end
        new(fϵ)
    end
end

function surrogenerator(x, method::TFTS, rng = Random.default_rng())
    # Pre-plan Fourier transforms
    forward = plan_rfft(x)
    inverse = plan_irfft(forward*x, length(x))

    # Pre-compute 𝓕
    𝓕 = forward * x

    # Polar coordinate representation of the Fourier transform
    rx = abs.(𝓕)
    ϕx = angle.(𝓕)
    n = length(𝓕)

    # These are updated during iteration procedure
    𝓕new = Vector{Complex{Float64}}(undef, length(𝓕))
    𝓕s = Vector{Complex{Float64}}(undef, length(𝓕))
    ϕs = Vector{Complex{Float64}}(undef, length(𝓕))

    init = (forward = forward, inverse = inverse,
        rx = rx, ϕx = ϕx, n = n,
        𝓕new = 𝓕new, 𝓕s = 𝓕s, ϕs = ϕs)

    return SurrogateGenerator(method, x, similar(x), init, rng)
end

function (sg::SurrogateGenerator{<:TFTS})()
    x, s = sg.x, sg.s
    fϵ = sg.method.fϵ
    L = length(x)

    init_fields = (:forward, :inverse,
        :rx, :ϕx, :n,
        :𝓕new, :𝓕s, :ϕs)

    forward, inverse,
        rx, ϕx, n,
        𝓕new, 𝓕s, ϕs = getfield.(Ref(sg.init), init_fields)

    # Surrogate starts out as a random permutation of x
    s .= x[StatsBase.sample(sg.rng, 1:L, L; replace = false)]
    𝓕s .= forward * s
    ϕs .= angle.(𝓕s)

    # Updated spectrum is the old amplitudes with the mixed phases.
    if fϵ > 0
        # Frequencies are ordered from lowest when taking the Fourier
        # transform, so by keeping the 1:n_ni first phases intact,
        # we are only randomizing the high-frequency components of the
        # signal.
        n_preserve = ceil(Int, abs(fϵ * n))
        ϕs[1:n_preserve] .= ϕx[1:n_preserve]
    elseif fϵ < 0
        # Do the exact opposite to preserve high-frequencies
        n_preserve = ceil(Int, abs(fϵ * n))
        ϕs[end-n_preserve+1:end] .= ϕx[end-n_preserve+1:end]
    end

    𝓕new .= rx .* exp.(ϕs .* 1im)
    s .= inverse * 𝓕new

    return s
end

"""
    TAAFT(fϵ)

An truncated version of the amplitude-adjusted-fourier-transform surrogate[^Theiler1991][^Nakamura2006].

The truncation parameter and phase randomization procedure is identical to [`TFTS`](@ref), but here an
additional step of rescaling back to the original data is performed. This preserves the
amplitude distribution of the original data.

[^Theiler1991]: J. Theiler, S. Eubank, A. Longtin, B. Galdrikian, J. Farmer, Testing for nonlinearity in time series: The method of surrogate data, Physica D 58 (1–4) (1992) 77–94.
[^Nakamura2006]: Nakamura, Tomomichi, Michael Small, and Yoshito Hirata. "Testing for nonlinearity in irregular fluctuations with long-term trends." Physical Review E 74.2 (2006): 026205.
"""
struct TAAFT <: Surrogate
    fϵ::Real

    function TAAFT(fϵ::Real)
        fϵ != 0 || throw(ArgumentError("`fϵ` must be on the interval [-1, 0) ∪ (0, 1] (positive if preserving high frequencies, negative if preserving low frequencies)"))
        new(fϵ)
    end
end

function surrogenerator(x, method::TAAFT, rng = Random.default_rng())
    init = (
        gen = surrogenerator(x, TFTS(method.fϵ), rng),
        x_sorted = sort(x),
        idxs = collect(1:length(x)),
        perm = zeros(Int, length(x)),
    )

    s = similar(x)
    return SurrogateGenerator(method, x, s, init, rng)
end

function (taaft::SurrogateGenerator{<:TAAFT})()
    sg = taaft.init.gen
    x_sorted, idxs, perm = taaft.init.x_sorted, taaft.init.idxs, taaft.init.perm

    x, s = sg.x, sg.s
    fϵ = sg.method.fϵ
    L = length(x)

    init_fields = (:forward, :inverse,
        :rx, :ϕx, :n,
        :𝓕new, :𝓕s, :ϕs)

    forward, inverse,
        rx, ϕx, n,
        𝓕new, 𝓕s, ϕs = getfield.(Ref(sg.init), init_fields)

    # Surrogate starts out as a random permutation of x
    StatsBase.sample!(sg.rng, idxs, perm, replace = false)
    permuted_x_into_s!(s, x, perm)
    𝓕s .= forward * s
    ϕs .= angle.(𝓕s)

    # Updated spectrum is the old amplitudes with the mixed phases.
    if fϵ > 0
        # Frequencies are ordered from lowest when taking the Fourier
        # transform, so by keeping the 1:n_ni first phases intact,
        # we are only randomizing the high-frequency components of the
        # signal.
        n_preserve = ceil(Int, abs(fϵ * n))
        ϕs[1:n_preserve] .= @view ϕx[1:n_preserve]
    elseif fϵ < 0
        # Do the exact opposite to preserve high-frequencies
        n_preserve = ceil(Int, abs(fϵ * n))
        ϕs[end-n_preserve+1:end] .= @view ϕx[end-n_preserve+1:end]
    end

    𝓕new .= rx .* exp.(ϕs .* 1im)
    s .= inverse * 𝓕new

    s[sortperm(s)] .= x_sorted
    return s
end

function permuted_x_into_s!(s, x, perm)
    k = 1
    for i in perm
        s[k] = x[i]
        k += 1
    end
end


#=

    Seguir tutorial de instalação do Julia demonstrado no README.md
    Como rodar esse script: $ julia bhaskara.jl

    Objetivo desse script: realizar uma operação básica de bhaskara.
    Referência: https://brasilescola.uol.com.br/matematica/formula-bhaskara.htm

=#

#=
    Declaração de variáveis a, b e c respectivamente
    Como visto, variáveis no Julia não são tipadas como em algumas linguagens,
    ou seja, você pode declará-las e o inicializador delas assumirá o valor
    que você colocar como atribuição.
    Ex.: Em a = 1 (a será um inteiro) / Em a = 'Olá' (a será uma string)
=#
a = 1
b = 0
c = -16

#   Definindo uma variável chamada 'delta' que assumirá o papel de
#   uma função onde algo será executado e o resultado será armazenado
delta = function (a, b, c)
    return b * b - 4 * a * c
end

#   Definindo uma variável chamada 'raiz1' que assumirá o papel de
#   uma função onde algo será executado e o resultado será armazenado
raiz1 = function (a, b, c)
    return (-b + Base.Math.sqrt(delta(a, b, c))) / 2 * a;
end

#   Definindo uma variável chamada 'raiz2' que assumirá o papel de
#   uma função onde algo será executado e o resultado será armazenado
raiz2 = function (a, b, c)
    return (-b - Base.Math.sqrt(delta(a, b, c))) / 2 * a;
end

#   Printando resultado das funções atribuídos a variável respectiva.
#   Comentário para demonstrar o valor que deve ser apresentado.
#   Seria interessante colocar um assert.equals para ficar mais nítida a compreensão.

println(delta(a, b, c)) # Deve imprimir: 64
println(raiz1(a, b, c)) # Deve imprimir: 4
println(raiz2(a, b, c)) # Deve imprimir: -4


# Utility functions for space data analysis.

using ImageFiltering, OffsetArrays, FFTW, Statistics, LinearAlgebra

export mag2db, sma, spectrum

"Convert `x` from magnitude to decibels."
mag2db(x) = 10. * log10.(x)

"""
    sma(x, n=(5,5))

Simple moving average of `x` with `n[1]` leading and `n[2]` trailing points.
Default boundary is "replicate", meaning that the border pixels extend beyond
the boundaries.
"""
function sma(x, n::Tuple{Int,Int}=(5,5))
   len = sum(n) + 1
   kernel = OffsetArray(fill(1/len, len), -n[1]:n[2])
   imfilter(x, kernel)
end

"""
    sma(x, n=5)

Simple moving box average of the vector data `x` with box length 'n'.
One-sided average on the left and right edge with replicate border.
"""
function sma(x::Vector, n::Int=5)
   iseven(n) && @warn "Even box length detected!"
   sma(x, ((n-1)÷2, (n-1)÷2))
end

"""
    spectrum(x, Fs)

Return the frequency and amplitude for the single-sided spectrum of vector `x`
given sample frequency of `Fs`.
"""
function spectrum(x::Vector, Fs)
   L = length(x)
   x̃ = fft(x)
   # Compute the two-sided spectrum amplitude P2.
   P₂ = @. abs(x̃/L)
   # Compute the single-sided spectrum amplitude P1.
   P₁ = P₂[1:floor(Int, L/2+1)]
   P₁[2:end-1] .= 2 .* P₁[2:end-1]

   f = Fs * (0:(L/2)) / L

   f, P₁
end


using Optim, Random
obj = build_loss_objective(prob1,Tsit5(),L2Loss(t,data),
                           maxiters=10000,verbose=false)

### Optim Method

println("Use Optim Brent to fit the parameter")
result = Optim.optimize(obj, 1.0, 10.0)
@test_broken result.minimizer[1] ≈ 1.5 atol=3e-1

println("Use Optim BFGS to fit the parameter")
result = Optim.optimize(obj, [1.0], Optim.BFGS())
@test result.minimizer[1] ≈ 1.5 atol=3e-1
#sol_optimized2 = solve(prob)
#plot!(sol_optimized2,leg=false)

cost_function2 = build_loss_objective(prob2,Tsit5(),L2Loss(t,data),
                                      maxiters=10000,verbose=false)
result_bfgs = Optim.optimize(cost_function2, [1.0,2.5], Optim.BFGS())
@test_broken result_bfgs.minimizer ≈ [1.5;3.0] atol=3e-1

Random.seed!(200)
cost_function3 = build_loss_objective(prob3,Tsit5(),L2Loss(t,data),
                                      maxiters=10000,verbose=false)
result_bfgs = Optim.optimize(cost_function3, [1.3,0.8,2.8,1.2], Optim.BFGS())
@test result_bfgs.minimizer ≈ [1.5;1.0;3.0;1.0] atol=5e-1

#=
println("Use Optim NelderMead to fit the parameter")
result_neldermead = Optim.optimize(cost_function2, [1.0,2.5], Optim.NelderMead())
@test result_neldermead.minimizer ≈ [1.5;3.0] atol=3e-1

println("Use Optim SimulatedAnnealing to fit the parameter")
options = Optim.Options(iterations=100_000)
result_simulatedannealing = Optim.optimize(obj, [1.0], Optim.SimulatedAnnealing(),options)
@test result_simulatedannealing.minimizer ≈ [1.5] atol=3e-1

result_simulatedannealing = Optim.optimize(cost_function2, [1.0,2.5], Optim.SimulatedAnnealing(),options)
@test result_simulatedannealing.minimizer ≈ [1.5;3.0] atol=3e-1

println("Use Optim Particle Swarm Optimization to fit the parameter")
initial_x = [1.0]
lower = [0.0]
upper = [5.0]
n_particles = 4
options = Optim.Options(iterations=100)
result_pso = Optim.optimize(obj, initial_x, Optim.ParticleSwarm(lower, upper, n_particles), options)
@test result_pso.minimizer[1] ≈ 1.5 atol=3e-1

initial_x = [1.0,2.5]
lower = [0.0,2.0]
upper = [5.0,10.0]
n_particles = 4
options = Optim.Options(iterations=100)
result_pso = Optim.optimize(cost_function2, initial_x, Optim.ParticleSwarm(lower, upper, n_particles), options)
@test result_pso.minimizer ≈ [1.5;3.0] atol=3e-1
=#


import SpecialFunctions
import SpecialFunctions: besselj

@testset "Bessel Roots" begin
    tol = 1e-11

    # Check if besselj(nu, besselroots(nu, n) ) is small:

    for nu = 0.:0.1:5.
        let n = 10
            @test norm( besselj.(nu, besselroots(nu, n) ), Inf ) < tol
        end
    end
end


function haversine(p1, p2; R=6371.0)
    lon1, lat1 = p1
    lon2, lat2 = p2
    φ1 = lat1 * π/180.0
    φ2 = lat2 * π/180.0
    Δφ = (lat2-lat1) * π/180.0
    Δλ = (lon2-lon1) * π/180.0
    a = sin(Δφ/2.0)^2.0 + cos(φ1)*cos(φ2) * sin(Δλ)^2.0
    c = 2.0 * atan(sqrt(a), sqrt(1.0-a))
    return R*c
end

"""
    slidingwindow(L::LT, f::FT, d::IT) where {LT <: SimpleSDMLayer, FT <: Function, IT <: Number}

This function will replace the value at any cell by applying the function `f` to
the array of cells values that are within a distance `d` (in kilometers) from
the focal cell. This is, for example, useful to use an average to smooth out the
layers. The distance is estimated using the haversine distance, assuming that
the radius of the Earth is 6371.0 km. This means that the size of the window
will vary a little bit across latitudes, but this is far better than using a
number of cells, which would have dramatic consequences near the poles.

It *always* returns a `SimpleSDMResponse`, and the cells containing `nothing`
will also not contain a value in the output. This is *different* from the
behavior of `coarsen`, which tends to expand the area of the layer in which we
have data.

This function is currently relatively slow. Performance improvements will arrive
at some point.
"""
function slidingwindow(layer::LT, f::FT, d::IT) where {LT <: SimpleSDMLayer, FT <: Function, IT <: Number}
    return_type = typeof(f(collect(layer)[1:min(3, length(layer))]))
    newgrid = convert(Matrix{Union{Nothing,return_type}}, fill(nothing, size(layer)))
    N = SimpleSDMResponse(newgrid, layer)
    pixels = []
    for lat in latitudes(layer)
        for lon in longitudes(layer)
            if !isnothing(layer[lon,lat])
                push!(pixels, (lon, lat) => layer[lon,lat])
            end
        end
    end

    for p1 in pixels
        ok = filter(p2 -> haversine(p2.first, p1.first) < 100.0, pixels)
        val = [p2.second for p2 in ok]
        N[p1.first...] = f(val)
    end

    #internal_types = unique(typeof.(N.grid))
    #N.grid = convert(Matrix{Union{internal_types...}}, N.grid)
    N = typeof(layer) <: SimpleSDMPredictor ? convert(SimpleSDMPredictor, N) : N

    return N
end


using LinearAlgebra
using Test
using Random

using ProximalOperators
using ProximalAlgorithms

@testset "Lasso small (strongly convex, $T)" for T in [Float32, Float64]

    Random.seed!(777)
    dim = 5
    μf = T(1)
    Lf = T(10)

    x_star = convert(Vector{T}, 1.5 * rand(T, dim) .- 0.5)

    lam = (μf + Lf) / 2
    @test typeof(lam) == T

    D = Diagonal(sqrt(μf) .+ (sqrt(Lf) - sqrt(μf)) * rand(T, dim))
    D[1] = sqrt(μf)
    D[end] = sqrt(Lf)
    Q = qr(rand(T, (dim, dim))).Q
    A = Q * D * Q'
    b = A * x_star + lam * inv(A') * sign.(x_star)

    f = LeastSquares(A, b)
    h = NormL1(lam)

    TOL = T(1e-4)

    @testset "SFISTA" begin

        # SFISTA

        x0 =  A \ b
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.SFISTA(tol = TOL)
        y, it = solver(x0, f = f, h = h, Lf = Lf, μf = μf)
        @test eltype(y) == T
        @test norm(y - x_star) <= TOL
        @test it < 200
        @test x0 == x0_backup

    end

end


"""
    F = JopErf(spc)

where `F` is the error function operator with domain and range given by `spc::JetSpace`.
we use 'erf'(z) = 2/√π ∫_{0}^{z} \exp{-t^2}dt,
the derivative of which is 2/√π  \exp{-z^2}.
We expect the domain and range to be real.
"""
JopErf(spc::JetSpace) = JopNl(dom = spc, rng = spc, f! = JopErf_f!, df! = JopErf_df!)
export JopErf

# TODO: assert denominator in linearization is > 0

JopErf_f!(d::AbstractArray{T}, m::AbstractArray{T}) where {T<:Real} = d .= erf.(m)
JopErf_df!(δd::AbstractArray{T}, δm::AbstractArray{T}; mₒ) where {T<:Real} = δd .= (2/sqrt( π )) .* exp.(-mₒ.^2) .* δm


"""
    clebsch_gordan_condon_shortley(j₁, m₁, j₂, m₂, j₃, m₃=m₁+m₂)

Calculate the vector coupling coefficient `⟨j₁ m₁, j₂ m₂|j₃ m₃⟩`
according to the Condon–Shortley phase convention and using the
definition of Eq. (8.1.12) in Varshalovich (1988).
"""
clebsch_gordan_condon_shortley(j₁, m₁, j₂, m₂, j₃, m₃=m₁+m₂) =
    powneg1(Int(j₁ - j₂ + m₃))*∏(j₃)*wigner3j(j₁, j₂, j₃,
                                              m₁, m₂, -m₃)

export clebsch_gordan_condon_shortley


using NSGAIII
using Base.Test

const C1 = [2 5 4 7 ; 3 3 5 7 ; 3 8 4 2 ; 6 5 2 5]
const C2 = [3 3 6 2 ; 5 3 7 3 ; 5 2 7 4 ; 4 6 3 5]
const C3 = [4 2 5 3 ; 5 3 4 3 ; 4 3 5 2 ; 6 4 7 3]

z(x, C) = sum(inds->C[inds...], enumerate(x))
z(x::Vector{Int}) = z(x, C1), z(x, C2), z(x, C3)

res = unique(nsga(500, 100, ()->randperm(4), z, 15))
display(res)
@test length(res) == 7

const d = RealCoding(6, [-10], [10])
z1(x) = x^2
z2(x) = (x-2)^2
z(bits) = begin
    x = decode(bits, d)[1]
    z1(x), z2(x)
end
seed = encode.([-10 + rand()*20 for _ =1:100] ,d)
res = nsga(500, 200, ()->rand(Bool, d.nbbitstotal), z, 5, seed = seed)

@test maximum(x -> x.y[1], res) >= 3.99
@test minimum(x -> x.y[1], res) <= 1e-4
@test maximum(x -> x.y[2], res) >= 3.99
@test minimum(x -> x.y[2], res) <= 1e-4


using Distributions, Turing, Plots, StatsPlots, MCMCChains, DataFrames

function BiExponential(theta1::Real, theta2::Real, p::Real)
    return MixtureModel(Exponential[
                             Exponential(theta1),
                             Exponential(theta2),
                            ], [p, 1 - p])
end

function generate_samples(theta1::Real, theta2::Real, p::Real, n::Int64)
    m = BiExponential(theta1, theta2, p)
    rand(m, n)
end

@model bi_exponential(samples::Array{Real, 1}) = begin
    theta1 ~ Gamma(1, 10)
    theta2 ~ truncated(Normal(), theta1, Inf)
    p ~ Beta(1, 1)
    m = BiExponential(theta1, theta2, p)
    for i in 1:length(samples)
        samples[i] ~ m
    end
end

function svf(x::Array{Real, 1}, digits::Int64)
    x = round.(x, digits=digits)
    unqx = unique(x) |> sort
    counts = map(x_ -> sum(x .== x_), unqx)
    DataFrame(x = unqx, count = counts, svr = 1. .- cumsum(counts ./ sum(counts)))
end

function predict(chains::Chains, x::Array{Real, 1})
    theta1 = mean(chains["theta1"])
    theta2 = mean(chains["theta2"])
    p = mean(chains["p"])
    becdf = cdf(BiExponential(theta1, theta2, p), x)
    return 1. .- becdf
end

samples = Array{Real, 1}(generate_samples(0.1, 5, 0.9, 2000))
chains = sample(bi_exponential(samples), NUTS(), MCMCThreads(), 1000, 4)

svf_samples = svf(samples, 2)
pred_samples = predict(chains, Array{Real, 1}(svf_samples.x))
scatter(svf_samples.x, log10.(svf_samples.svr), markersize=2)
plot!(svf_samples.x, log10.(pred_samples))
ylims!(-3, 0)


Store(α, rmax, Mmax) = Store(α, rmax, Mmax, Mmax)

function Store(α::T, rmax::Integer, Mmax::Integer, ppImax::Integer
              ) where T <: AbstractFloat
    if Mmax < rmax
        throw(ArgumentError("Mmax must be at least as big as rmax"))
    end
    if Mmax > ppImax
        throw(ArgumentError("ppImax must be at least as big as Mmax"))
    end
    unitlegendre = Vector{Matrix{T}}(undef, Mmax)
    legendre = Vector{Matrix{T}}(undef, Mmax)
    jacobi1 = Vector{Matrix{T}}(undef, Mmax)
    jacobi2 = Vector{Matrix{T}}(undef, Mmax)
    jacobi3 = Vector{Matrix{T}}(undef, Mmax)
    jacobi4 = Vector{Matrix{T}}(undef, Mmax)
    jacobi5 = Vector{Matrix{T}}(undef, Mmax)
    for M = 1:Mmax
        unitlegendre[M] = Array{T}(undef, M, 2)
        legendre[M] = Array{T}(undef, M, 2)
        x, w = GaussQuadrature.legendre(T, M)
        legendre[M][:,1], legendre[M][:,2] = x, w
        unitlegendre[M][:,1] = ( x .+ 1 ) / 2
        unitlegendre[M][:,2] = w / 2
        jacobi1[M] = Array{T}(undef, M, 2)
        jacobi1[M][:,1], jacobi1[M][:,2] = GaussQuadrature.jacobi(M, α-1,
                                                                  zero(T))
        jacobi2[M] = Array{T}(undef, M, 2)
        jacobi2[M][:,1], jacobi2[M][:,2] = GaussQuadrature.jacobi(M, zero(T),
                                                                  α-1)
        jacobi3[M] = Array{T}(undef, M, 2)
        jacobi3[M][:,1], jacobi3[M][:,2] = GaussQuadrature.jacobi(M, α, zero(T))
        jacobi4[M] = Array{T}(undef, M, 2)
        jacobi4[M][:,1], jacobi4[M][:,2] = GaussQuadrature.jacobi(M, zero(T), α)
        jacobi5[M] = Array{T}(undef, M, 2)
        jacobi5[M][:,1], jacobi5[M][:,2] = GaussQuadrature.jacobi(M, one(T),
                                                                  α-1)
    end
    Ψ  = Array{T}(undef, rmax, ppImax)
    dΨ = Array{T}(undef, rmax, Mmax)
    A = Vector{T}(undef, rmax)
    B = Vector{T}(undef, rmax)
    C = Array{T}(undef, rmax, rmax)
    return Store(α, rmax, Mmax, unitlegendre, legendre,
                 jacobi1, jacobi2, jacobi3, jacobi4, jacobi5,
                 ppImax, Ψ, dΨ, A, B, C)
end

function rule(storerule::Matrix{T}) where T <: AbstractFloat
    return storerule[:,1], storerule[:,2]
end

function rules(storerules::Vector{Matrix{T}}) where T <: AbstractFloat
    M = length(storerules)
    pt = Vector{Vector{T}}(undef, M)
    wt = Vector{Vector{T}}(undef, M)
    for m = 1:M
        pt[m] = view(storerules[m], :, 1)
        wt[m] = view(storerules[m], :, 2)
    end
    return pt, wt
end


## Packages
using NLsolve
using Plots
using Parameters

## Files
include("Tauchen.jl")
include("NewtonRoot.jl")

## Parameters
Model = @with_kw (θ = 0.35, # capital share
    δ = 0.0464, # depreciation rate
    γ_z = 0.016, # growth rate of productivity
    γ_n = 0.015, # growth rate of population
    β = 0.9722, # discount rate
    β_hat = β*(1+γ_n), # detrended discount rate
    ψ = 2.24, # labor coefficient
    ρ = 0.2,# autocorr of productivity process
    σ = 0.5, # stand. dev. for productivity process
    nz = 5, # number of states for productivity
    μ_z = 0.0, # mean of productivity process
    zGrid = exp.(Tauchen(μ_z,ρ,σ,nz)[1]), # grid for productivity process
    P_z = Tauchen(μ_z,ρ,σ,nz)[2]) # transition matrix

## Get the Steady-State level of capital from a Nonlinear solver
function getSS(md)
    @unpack θ,ψ,δ,β_hat,γ_n,γ_z = md

    function ee!(eq,x)
        k=(x[1])
        h=(x[2])
        eq[1]=β_hat*(θ*k^(θ-1)*h^(1-θ)+1-δ)-(1+γ_n)*(1+γ_z)
        eq[2]=(-(ψ/(1-h)))+((1-θ)*(k^θ)*(h^(-θ)))/((k^θ)*(h^(1-θ))-(1+γ_n)*(1+γ_z)*k+(1-δ)*k)
    end

    S = nlsolve(ee!, [0.5,0.5],ftol = :1.0e-9, method = :trust_region , autoscale = true);
    kss = S.zero[1];
    hss = S.zero[2];
    lss = 1-hss;
    css = (kss^θ)*((exp(0)*hss)^(1-θ))-(1+γ_n)*(1+γ_z)*kss+(1-δ)*kss;

    return kss,hss,lss,css
end

# Solve h from intratemporal equation over the grids
function geth(md,nk,kGrid)

    @unpack θ,ψ,γ_n,γ_z,δ,nz,zGrid = md

    h_star = zeros(nk,nk,nz);

    for i = 1:nk
        for j = 1:nk
            for k = 1:nz
                intra(h) = (ψ/(1-θ))*((kGrid[i]^θ)*((zGrid[k]*h)^(1-θ))-(1+γ_n)*(1+γ_z)*kGrid[j]+(1-δ)*kGrid[i])+((h-1)*(kGrid[i]^θ)*(zGrid[k]^(1-θ))*(h^(-θ)))
                h = NewtonRoot(h -> intra(h),0.5)
                if h > 1
                    h_star[i,j,k] = 0.9999;
                else
                    h_star[i,j,k] = h;
                end
            end
        end
    end

    return h_star

end

# Precalculate return value
function getrt(md,nk,kGrid)

    @unpack θ,ψ,γ_n,γ_z,δ,nz,zGrid = md

    h_star = geth(md,nk,kGrid)
    rt = zeros(nk,nk,nz);

    for i = 1:nk
        for j = 1:nk
            for k = 1:nz
                c = (kGrid[i]^θ)*((zGrid[k]*h_star[i,j,k])^(1-θ))-(1+γ_n)*(1+γ_z)*kGrid[j]+(1-δ)*kGrid[i];
                if c <= 0
                    rt[i,j,k] = -100.0;
                else
                    rt[i,j,k] = log(c)+ψ*log(1-h_star[i,j,k]);
                end
            end
        end
    end

    return rt,h_star

end

# Main loop
function VFI(md,nk,kGrid)
    @unpack θ,ψ,γ_n,γ_z,δ,nz,zGrid,β_hat,P_z = md

    #Preallocate memories
    V_old = zeros(nk,nz);
    V_new = zeros(nk,nz);
    argmax = Array{Int64,2}(undef,nk,nz);

    kpol = zeros(nk,nz);
    hpol = zeros(nk,nz);
    cpol = zeros(nk,nz);

    maxDiff = 10.0;
    tol = 1.0e-10;
    iter = 0;

    rt,h_star = getrt(md,nk,kGrid);

    while maxDiff > tol
        iter = iter+1;
        println("Current iteration: $iter")

        EV = V_old*P_z';
        for i = 1:nk, k = 1:nz
            V_new[i,k],argmax[i,k] = findmax(rt[i,:,k].+β_hat*EV[:,k]);
        end

        maxDiff = maximum(abs.(V_new-V_old));
        println("maxDiff = $maxDiff")
        V_old = copy(V_new);
    end

    for i = 1:nk,j = 1:nz
        kpol[i,j] = kGrid[argmax[i,j]];
        hpol[i,j] = h_star[i,argmax[i,j],j];
        cpol[i,j] = (kGrid[i]^θ)*(zGrid[j]*hpol[i,j])^(1-θ)-(1+γ_n)*(1+γ_z)*kpol[i,j]+(1-δ)*kGrid[i];
    end

    return V_new,kpol,hpol,cpol
end


using Plots

Δ = .1
N = 100
τ = N*Δ/5
ran = 0.0:2pi/N:2pi
function fillxs(Ts = 1,N = 100, ini=0.0, fin=100)
    xs = []
    Δ = fin/N
    τ = N*Δ/5

    for j in 1:Ts
        for i in ini:Δ:τ-Δ
            push!(xs,1)
        end

        for i in τ:Δ:fin
            push!(xs,0)
        end

    end

    return xs
end

function Xforward(x,n)
    N = length(x)
    j = im
    b = 2pi/N
    sum = 0.0
    for k in 1:N
        sum += x[k] * exp(-j * k * b * n)
    end

    return abs(1/N * sum)
end

xs = Int.(fillxs())

ys =  [Xforward(xs,i) for i in 0:1:N]

scatter(xs, title = "Train signal")
scatter(ys, title = "DFT")

function xback(x,n)
    N = length(x)
    j = im
    b = 2pi/N
    sum = 0.0
    for k in 1:N
        sum += x[k] * exp(j * k * b * n)
    end

    return abs(sum)
end

ys2 = [xback(xs,i) for i in 0:1:N]

scatter(ys2, title = "DFT back")


# # Estimate Binomial draw probabilility

using DynamicHMCModels

Random.seed!(1356779)

# Define a structure to hold the data.

Base.@kwdef struct BernoulliProblem
    "Total number of draws in the data."
    n::Int
    "Number of draws ' == 1' "
    obs::Vector{Int}
end;

# Write a function to return properly dimensioned transformation.

make_transformation(model::BernoulliProblem) =
    as((p = as𝕀, ))

# Add data

model = BernoulliProblem(; n = 9, obs = rand(Binomial(9, 2/3), 3))

# Make the type callable with the parameters *as a single argument*.

function (model::BernoulliProblem)(θ)
    @unpack n, obs = model        # extract the data
    @unpack p = θ
    loglikelihood(Binomial(n, p), obs)
end

# Use a flat priors (the default, omitted) for α

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P);

# Sample chain

results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000;
  reporter = NoProgressReport()
)

posterior = P.transformation.(results.chain)

# Create Particles NamedTuple object

println()
p = as_particles(posterior)
p |> display

println()
DynamicHMC.Diagnostics.EBFMI(results.tree_statistics) |> display

println()
DynamicHMC.Diagnostics.summarize_tree_statistics(results.tree_statistics) |> display


module Lecture2

using Test

include("../../lib/KTBC.jl")
import .KTBC.CreateKTBC

include("../../lib/SpecialVectors.jl")
using .SpecialVectors: constants, linear, squares, rampAtK, sines, cosines, exponentials

include("../../lib/DifferencesMatrix.jl")
using .DifferencesMatrix: backwardDiff, forwardDiff

dims = 4

K, T, B, C = CreateKTBC(dims)

Kₐ = convert(Array{Int}, K)

cV = constants(dims)

@test (K * cV)[2:3] == [
  0.0
  0.0
]

lV = linear(dims)

@test (K * lV)[2:3] == [
  0.0
  0.0
]

sV = squares(dims)

@test (-K * sV)[2:3] == [
 2.0
 2.0
]

rV = rampAtK(2, dims)

@test (-K * rV)[1:3] == [
  0.0
  1.0
  0.0
]

t = π / 4
multiplier = round(2 * cos(t) - 2, digits=2)

sinV = round.(sines(t, dims), digits=2)
@test round.((-K * sinV)[1:3], digits=1) == round.((multiplier * sinV)[1:3], digits=1)

cosV = round.(cosines(t, dims), digits=2)
@test round.((-K * cosV)[2:3], digits=1) == round.((multiplier * cosV)[2:3], digits=1)

expV = round.(exponentials(t, dims), digits=2)
@test round.((-K * expV)[2:3], digits=1) == round.((multiplier * expV)[2:3], digits=1)

# Matlab experiment
function mathlabExperiment(n::Int)
    h = 1 / (n + 1)
    u = cos.(π * [1:n...] * h / 2)
    c = (π / 2)^2
    f = c * u

    K, T, B, C = CreateKTBC(n)

    U = (h^2) * T \ f
    e = 1 - U[1]
    g = [c / 2; f]

    T = backwardDiff(n + 1) * forwardDiff(n + 1)

    V = (h^2) * T \ g
    E = 1 - V[1]

    return f, e, T, V, E
end

f, e, T, V, E = mathlabExperiment(3)
f, e, T, V, E = mathlabExperiment(7)
f, e, T, V, E = mathlabExperiment(15)

end


# A simple example to demonstrate the use of gridtune

using MLBase
using MultivariateStats

## prepare data

n_tr = 20  # number of training samples
n_te = 10  # number of testing samples
d = 5      # dimension of observations

theta = randn(d)
X_tr = randn(n_tr, d)
y_tr = X_tr * theta + 0.1 * randn(n_tr)
X_te = randn(n_te, d)
y_te = X_te * theta + 0.1 * randn(n_te)

## tune the model

function estfun(regcoef, bias)
    s = ridge(X_tr, y_tr, regcoef; bias=bias)
    return bias ? (s[1:end-1], s[end]) : (s, 0.0)
end

evalfun(m) = msd(X_te * m[1] + m[2], y_te)

r = gridtune(estfun, evalfun,
            ("regcoef", [1.0e-3, 1.0e-2, 1.0e-1, 1.0]),
            ("bias", (true, false));
            ord=Reverse,
            verbose=true)

best_model, best_cfg, best_score = r

## print results

θ, b = best_model
println("Best model:")
println("  θ = $(θ')"),
println("  b = $b")
println("Best config: regcoef = $(best_cfg[1]), bias = $(best_cfg[2])")
println("Best score: $(best_score)")


using DiffEqBayes, CmdStan, DynamicHMC

using Distributions, BenchmarkTools
using OrdinaryDiffEq, RecursiveArrayTools, ParameterizedFunctions
using Plots

gr(fmt=:png)

f = @ode_def LotkaVolterraTest begin
    dx = a*x - b*x*y
    dy = -c*y + d*x*y
end a b c d

u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1,0]

prob = ODEProblem(f,u0,tspan,p)
sol = solve(prob,Tsit5())

t = collect(range(1,stop=10,length=10))
sig = 0.49
data = convert(Array, VectorOfArray([(sol(t[i]) + sig*randn(2)) for i in 1:length(t)]))

scatter(t, data[1,:], lab="#prey (data)")
scatter!(t, data[2,:], lab="#predator (data)")
plot!(sol)

priors = [Truncated(Normal(1.5,0.5),0.5,2.5),Truncated(Normal(1.2,0.5),0,2),Truncated(Normal(3.0,0.5),1,4),Truncated(Normal(1.0,0.5),0,2)]

@btime bayesian_result_stan = stan_inference(prob,t,data,priors,num_samples=10_000,printsummary=false)

@btime bayesian_result_turing = turing_inference(prob,Tsit5(),t,data,priors,num_samples=10_000)

@btime bayesian_result_dynamichmc = dynamichmc_inference(prob,Tsit5(),t,data,priors,num_samples=10_000)

using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


using Plots
using DSP
using FFTW
include("RSProject.jl")

function write_binary_file(data,filename::String,dtype::DataType=Matrix{ComplexF32})
  f = open(filename,"w");
  data = convert(dtype,data)
  write(f,data)
  close(f)
end

"""
    exportAndDisplay(u,a,iter,K,m,name)
Exports and displays PCFM algorith results.
  ...
# Arguments
- `u::Vector`: Window Function PSD.
- `a::Integer`: Log base. If zero use non log version.
- `iter::Integer`: Number of iterations.
- `K::Integer`: Oversampling factor.
- `m::Integer`: Size.
- `name::String`: Name of window function.
...
"""
function exportAndDisplay(u,a,iter,K,m,name)
    u[findall(<(-50), 10*log10.(u))] .= 10^-5
    u = abs.(u).^2
    result = RSProject.funPcfm(u,a,iter,K)
    (B,Bb,x) = RSProject.funPcfmHelper(m,K)
    s = exp.(im.*B*result)
    # write_binary_file(s,"data/"*name*".bin")
    sb = vcat(s, zeros(m-1,1))
    sbf =  fftshift(fft(sb))
    sbf = sbf ./maximum(abs.(sbf))
    display(plot(10*log10.(abs.(sbf).^2),ylim=(-50, 0)))
    display(plot!(10*log10.(u),ylim=(-50, 0)))
end
m = 512
K=3
a = 10
iter = 1000
#Guassian Window
u = gaussian((2*m-1,1),0.1; padding = 0, zerophase = false)
exportAndDisplay(u,a,iter,K,m,"Gaussian")
u[findall(<(-50), 10*log10.(u))] .= 10^-5
u = abs.(u).^2
time = @time begin
  result = RSProject.profm(sqrt.(u),iter)
end
#time2 = @time begin
  #result2 = RSProject.funPcfm(u,a,iter,K)
#end
#(B,Bb,x) = RSProject.funPcfmHelper(m,K)
#s = exp.(im.*B*result2)
#sb = vcat(s, zeros(m-1,1))
#sbf =  fftshift(fft(sb))
#sbf = sbf ./maximum(abs.(sbf))
resPlot = abs.(fftshift(fft(result)))
resPlot = resPlot./maximum(resPlot)
display(plot(abs.(resPlot).^2))
#display(plot!(abs.(sbf).^2))
#exportAndDisplay(u,a,iter,K,m,"Guassian")
#Hanning Window
#u = hanning((2*m-1,1); padding = 0, zerophase = false)
#exportAndDisplay(u,a,iter,K,m,"Hanning")
#Hamming Window
#u = hamming((2*m-1,1); padding = 0, zerophase = false)
#exportAndDisplay(u,a,iter,K,m,"Hamming")
#Tukey Window
# u = tukey((2*m-1,1), 0.5; padding = 0, zerophase = false)
# u[findall(<(-50), 10*log10.(u))] .= 10^-5
# u = abs.(u).^2
# time = @time begin
#   result = profm(sqrt.(u),iter)
# end
# time2 = @time begin
#   result2 = funPcfm(u,a,iter,K)
# end
# (B,Bb,x) = funPcfmHelper(m,K)
# s = exp.(im.*B*result2)
# sb = vcat(s, zeros(m-1,1))
# sbf =  fftshift(fft(sb))
# sbf = sbf ./maximum(abs.(sbf))
# resPlot = abs.(fftshift(fft(result)))
# resPlot = resPlot./maximum(resPlot)
# display(plot(abs.(resPlot).^2))

# display(plot!(abs.(sbf).^2))
# #exportAndDisplay(u,a,iter,K,m,"Tukey")
# #Rectangular Window

# #u = rect((2*m-1,1); padding = 0, zerophase = false)

# #exportAndDisplay(u,a,iter,K,m,"Rectangular")
# #Triangular Window
# #u = triang((2*m-1,1); padding = 0, zerophase = false)
#u = tukey((2*m-1,1), 0.5; padding = 0, zerophase = false)
#exportAndDisplay(u,a,iter,K,m,"Tukey")
#Rectangular Window
#u = rect((2*m-1,1); padding = 0, zerophase = false)
#exportAndDisplay(u,a,iter,K,m,"Rectangular")
#Triangular Window
#u = triang((2*m-1,1); padding = 0, zerophase = false)
# #exportAndDisplay(u,a,iter,K,m,"Triangular")


# coding: utf-8

# In[1]:

#fractal is one of the interesting topics in geometry
#it is usually described by a recursive function
#voila,here we are!
using Plots

# In[2]:

#define starting and ending points
#and the order of hilbert curve
n=4
starting_point=(0,0)
ending_point=(2^n-1,0)

# In[3]:

#the code for hilbert curve may look intimidating at the first glance
#it is extremely simple and straight forward
#julia version is slightly different
#as we cannot use yield like python
#coroutine in julia is too complicated
#maybe another day...
function hilbert_curve(coordinates,point1,point2,n)

    #unpack
    x_start,y_start=point1
    x_end,y_end=point2

    #the function consists four different parts
    #which are 4 different rotations plus flips of a second order hilbert curve
    #0 degree second order hilbert curve
    if x_start<x_end && y_start==y_end

        #compute the grid length
        L=x_end-x_start

        #keypoints are starting and ending points
        #of 4 different first order hilbert curves
        #on a second order hilbert curve
        keypoints=[(x_start,y_start),(x_start,y_start-(L-1)/2),
        (x_start,y_start-(L-1)/2-1),(x_start+(L-1)/2,y_start-(L-1)/2-1),
        (x_start+(L-1)/2+1,y_start-(L-1)/2-1),(x_start+L,y_start-(L-1)/2-1),
        (x_start+L,y_start-(L-1)/2),(x_start+L,y_start)]

        #base case
        if n==1
            append!(coordinates,keypoints)
        else

            #recursion
            for i in 1:2:8
                hilbert_curve(coordinates,
                    keypoints[i],keypoints[i+1],n-1)
            end
        end
    end

    #180 degree second order hilbert curve
    if x_start>x_end && y_start==y_end
        L=x_start-x_end
        keypoints=[(x_start,y_start),(x_start,y_start+(L-1)/2),
        (x_start,y_start+(L-1)/2+1),(x_start-(L-1)/2,y_start+(L-1)/2+1),
        (x_start-(L-1)/2-1,y_start+(L-1)/2+1),(x_start-L,y_start+(L-1)/2+1),
        (x_start-L,y_start+(L-1)/2),(x_start-L,y_start)]
        if n==1
            append!(coordinates,keypoints)
        else
            for i in 1:2:8
                hilbert_curve(coordinates,
                    keypoints[i],keypoints[i+1],n-1)
            end
        end
    end

    #clockwise 90 degree horizontal flipped second order hilbert curve
    if x_start==x_end && y_start>y_end
        L=y_start-y_end
        keypoints=[(x_start,y_start),(x_start+(L-1)/2,y_start),
        (x_start+(L-1)/2+1,y_start),(x_start+(L-1)/2+1,y_start-(L-1)/2),
        (x_start+(L-1)/2+1,y_start-(L-1)/2-1),(x_start+(L-1)/2+1,y_start-L),
        (x_start+(L-1)/2,y_start-L),(x_start,y_start-L)]
        if n==1
            append!(coordinates,keypoints)
        else
            for i in 1:2:8
                hilbert_curve(coordinates,
                    keypoints[i],keypoints[i+1],n-1)
            end
        end
    end

    #clockwise 270 degree horizontal flipped second order hilbert curve
    if x_start==x_end && y_start<y_end
        L=y_end-y_start
        keypoints=[(x_start,y_start),(x_start-(L-1)/2,y_start),
        (x_start-(L-1)/2-1,y_start),(x_start-(L-1)/2-1,y_start+(L-1)/2),
        (x_start-(L-1)/2-1,y_start+(L-1)/2+1),(x_start-(L-1)/2-1,y_start+L),
        (x_start-(L-1)/2,y_start+L),(x_start,y_start+L)]
        if n==1
            append!(coordinates,keypoints)
        else
            for i in 1:2:8
                hilbert_curve(coordinates,
                    keypoints[i],keypoints[i+1],n-1)
            end
        end
    end
    coordinates
end

# In[4]:

#get coordinates
coordinates=hilbert_curve([],
    starting_point,ending_point,n);

# In[5]:

#viz
gr(size=(250,250))
fig=plot(legend=false,grid=false,axis=false,ticks=false)
plot!([i[1] for i in coordinates],
        [i[2] for i in coordinates])
fig


# Binary search for zero of continuous function f
# with one zero-crossing in [lo, hi]
function binary_search(f, lo, hi; ϵ = 1e-10, maxiter = 100)
    if f(lo) > 0
        @warn "f($(lo))=$(f(lo)) should be negative at low end";
        return lo;
    end
    if f(hi) < 0
        @warn "f($(hi))=$(f(hi)) should be positive at high end";
        return hi;
    end

    # @assert f(lo) < +ϵ "f($(lo))=$(f(lo)) should be negative at low end";
    # @assert f(hi) > -ϵ "f($(hi))=$(f(hi)) should be positive at high end";

    for i in 1:maxiter
        mid = (lo+hi)/2;
        if mid == lo || mid == hi
            # not going to get any better ...
            #@warn "binary_search collapsed after $i iterations.\nf($(lo)) = $(f(lo))\nf($(hi)) = $(f(hi))";
            return mid;
        end
        fmid = f(mid);
        if fmid < -ϵ
            lo = mid
        elseif fmid > ϵ
            hi = mid;
        else
            return mid;
        end
    end
    @warn "binary_search did not reach tolerance $ϵ in $maxiter iterations.\nf($(lo)) = $(f(lo))\nf($(hi)) = $(f(hi)),\nmid would be $((lo+hi)/2)";
    return (lo+hi)/2;
end


"""
    legendre_1term_raise_lm_accuracy.jl

Simple tests to explore the numerical accuracy of the 1-term recurrence relation which
boosts (m,m) -> (m+1,m+1) in the associated Legendre polynomials.

## Tests

- `recur0`: Baseline version which is a direct translation of the math.
- `recur1`: Uses FMA instructions to compute (1-x^2). Eliminating intermediate rounding
  is expected to increase the accuracy of the entire recursion.
- `recur2`: Same as `recur1` but also iterates upward in degree/order in pairs. Stepping
  from even to odd ``(m,m) -> (m+1,m+1)`` multiplies by ``y = \\sqrt{1-x^2}`` and the
  coefficient ``μ_m``. The second half then steps from even to even
  ``(m,m) -> (m+2,m+2)`` by multiplying by ``y^2 = 1-x^2`` (without taking the square
  root) and the pair of coefficients ``μ_m`` and ``μ_{m+1}``.
- `recur3`: Same as `recur2` but replaces the FMA calculation of ``1-x^2`` with the
  product ``(1-x)(1+x)``.

## Conclusions
- Without FMA, the product ``(1-x)(1+x)`` is better for ``|x| ≈ 1`` than ``1 - x^2``.
- FMA is the best option for computing ``y = \\sqrt(1-x^2)`` and should be preferred.
- The pair-wise iteration improves numerical accuracy as ``m \\to \\finty``.
"""

using PyPlot, PyCall

function recur0(x, nstep)
    y = sqrt(one(x) - x^2)
    vals = zeros(typeof(x), 2nstep)
    z = inv(sqrt(4oftype(x, π)))
    @inbounds for ii = 1:2nstep
        μ = sqrt(one(x) + one(x)/2(ii+1))
        z = -μ * y * z
        vals[ii] = z
    end
    return vals
end

function recur1(x, nstep)
    y = sqrt(-fma(x, x, -one(x)))
    vals = zeros(typeof(x), 2nstep)
    z = inv(sqrt(4oftype(x, π)))
    @inbounds for ii = 1:2nstep
        μ = sqrt(one(x) + one(x)/2(ii+1))
        z = -μ * y * z
        vals[ii] = z
    end
    return vals
end

function recur2(x, nstep)
    y² = -fma(x, x, -one(x))
    y¹ = sqrt(y²)
    vals = zeros(typeof(x), 2nstep)
    z = inv(sqrt(4oftype(x, π)))
    @inbounds for ii = 1:2:2nstep
        μ₁ = sqrt(one(x) + one(x)/2(ii+1))
        μ₂ = sqrt(one(x) + one(x)/2(ii+2))
        vals[ii] = -μ₁ * y¹ * z
        z = μ₁ * μ₂ * y² * z
        vals[ii+1] = z
    end
    return vals
end

function recur3(x, nstep)
    y² = (one(x) - x) * (one(x) + x)
    y¹ = sqrt(y²)
    vals = zeros(typeof(x), 2nstep)
    z = inv(sqrt(4oftype(x, π)))
    @inbounds for ii = 1:2:2nstep
        μ₁ = sqrt(one(x) + one(x)/2(ii+1))
        μ₂ = sqrt(one(x) + one(x)/2(ii+2))
        vals[ii] = -μ₁ * y¹ * z
        z = μ₁ * μ₂ * y² * z
        vals[ii+1] = z
    end
    return vals
end

# Half of the lmax/mmax to iterate to.
nstep = 512
# Range of arguments, taken more densely from x ≈ 1 where catastrophic cancellation is
# expected without FMA.
x₀ = 1 .- exp2.(range(-16, -eps(), length=2560))

# Reference values calculated with extended precision.
ref = Float64.(reduce(hcat, setprecision(512) do
    recur1.(big.(x₀), nstep)
end))

v0 = reduce(hcat, recur0.(x₀, nstep))
v1 = reduce(hcat, recur1.(x₀, nstep))
v2 = reduce(hcat, recur2.(x₀, nstep))
v3 = reduce(hcat, recur3.(x₀, nstep))

colors = pyimport("matplotlib.colors")
fig, axs = subplots(2, 2, num = "legendre_1term_raise_lm_accuracy", clear = true)
function plot_compare(ax, v; title=nothing, kws...)
    lnorm = colors.LogNorm(vmin = 1, vmax=1e3)
    ulps = abs.(v .- ref) ./ eps.(ref)
    img = ax.pcolormesh(1 .- x₀, 1:2nstep, ulps,
                    norm=lnorm;
                    kws...)
    ax.set_xscale("log")
    ax.set_xlabel("argument \$(1-x)\$")
    ax.set_ylabel("degree/order")
    cbar = ax.figure.colorbar(img, ax=ax)
    cbar.set_label("ulps");
    title !== nothing && ax.set_title(title)
    ax.text(0.05, 0.9, "ulp max = $(round(maximum(ulps), digits=2))", transform=ax.transAxes)
    nothing
end
plot_compare(axs[1,1], v0, title="v0 — Baseline")
plot_compare(axs[1,2], v1, title="v1 — FMA \$(1 - x^2)\$")
plot_compare(axs[2,1], v2, title="v2 — FMA + pairwise iteration")
plot_compare(axs[2,2], v3, title="v3 — \$(1-x)(1+x)\$ + pairwise iteration")